134 lines
4.6 KiB
Python
134 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from agent_integration.client import OpenAgentEvent
|
|
from arr_processing.agent_trace import (
|
|
JsonlAgentTraceStore,
|
|
capture_failure_record,
|
|
project_trace_event,
|
|
sanitize_trace_summary,
|
|
trace_run_id,
|
|
)
|
|
|
|
|
|
class AgentTraceProjectionTests(unittest.TestCase):
|
|
def test_trace_projection_keeps_summaries_but_drops_content_and_sensitive_values(self) -> None:
|
|
raw_summary = (
|
|
"调用完成 submission_grant=grant-secret "
|
|
"Authorization:Bearer token-secret "
|
|
"https://example.test/private?q=1 /Users/private/work/result.json "
|
|
"guest@example.test 8613800138000 <FULL_NAME>Guest</FULL_NAME>"
|
|
)
|
|
event = OpenAgentEvent(
|
|
event="trace",
|
|
data={
|
|
"event": "step.updated",
|
|
"run_id": "run-trace-001",
|
|
"ts": "2026-07-29T13:00:00Z",
|
|
"data": {"summary": raw_summary},
|
|
"org_id": "must-not-escape",
|
|
},
|
|
)
|
|
|
|
projected = project_trace_event(event, attempt_no=1, sequence=2)
|
|
|
|
self.assertIsNotNone(projected)
|
|
assert projected is not None
|
|
self.assertEqual(projected["code"], "AGENT_STEP_UPDATED")
|
|
rendered = json.dumps(projected, ensure_ascii=False)
|
|
for private in (
|
|
"grant-secret",
|
|
"token-secret",
|
|
"example.test",
|
|
"/Users/private",
|
|
"guest@example.test",
|
|
"8613800138000",
|
|
"FULL_NAME",
|
|
"must-not-escape",
|
|
):
|
|
self.assertNotIn(private, rendered)
|
|
for marker in ("[REDACTED]", "[URL]", "[PATH]", "[EMAIL]", "[NUMBER]", "[XML]"):
|
|
self.assertIn(marker, rendered)
|
|
|
|
hidden = project_trace_event(
|
|
OpenAgentEvent(
|
|
event="trace",
|
|
data={
|
|
"event": "message.final",
|
|
"run_id": "run-trace-001",
|
|
"text": "signed payload must not escape",
|
|
},
|
|
),
|
|
attempt_no=1,
|
|
sequence=3,
|
|
)
|
|
self.assertIsNone(hidden)
|
|
|
|
def test_run_id_is_accepted_only_from_run_started(self) -> None:
|
|
started = OpenAgentEvent(
|
|
event="trace",
|
|
data={"event": "run.started", "run_id": "run-trace-001"},
|
|
)
|
|
completed = OpenAgentEvent(
|
|
event="trace",
|
|
data={"event": "run.completed", "run_id": "run-trace-001"},
|
|
)
|
|
|
|
self.assertEqual(trace_run_id(started), "run-trace-001")
|
|
self.assertIsNone(trace_run_id(completed))
|
|
|
|
def test_capture_failure_exposes_only_exception_class(self) -> None:
|
|
secret = "network secret payload"
|
|
record = capture_failure_record(
|
|
attempt_no=2,
|
|
sequence=9,
|
|
error=RuntimeError(secret),
|
|
now=datetime(2026, 7, 29, 13, 0, tzinfo=timezone.utc),
|
|
)
|
|
|
|
rendered = json.dumps(record, ensure_ascii=False)
|
|
self.assertIn("RuntimeError", rendered)
|
|
self.assertNotIn(secret, rendered)
|
|
|
|
def test_empty_summary_is_not_logged(self) -> None:
|
|
self.assertIsNone(sanitize_trace_summary(" \n\t "))
|
|
|
|
|
|
class JsonlAgentTraceStoreTests(unittest.TestCase):
|
|
def test_round_trip_is_append_only_and_invalid_lines_are_reported(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
store = JsonlAgentTraceStore(Path(temporary) / "trace")
|
|
record = project_trace_event(
|
|
OpenAgentEvent(
|
|
event="trace",
|
|
data={
|
|
"event": "run.started",
|
|
"run_id": "run-trace-001",
|
|
"ts": "2026-07-29T13:00:00Z",
|
|
},
|
|
),
|
|
attempt_no=1,
|
|
sequence=1,
|
|
)
|
|
assert record is not None
|
|
store.append("arrjob-trace-001", record)
|
|
path = store._path("arrjob-trace-001") # pylint: disable=protected-access
|
|
with path.open("a", encoding="utf-8") as target:
|
|
target.write('{"raw":"must not be accepted"}\n')
|
|
|
|
records = store.read("arrjob-trace-001")
|
|
|
|
self.assertEqual(records[0]["code"], "AGENT_RUN_STARTED")
|
|
self.assertEqual(records[-1]["code"], "AGENT_TRACE_RECORD_INVALID")
|
|
self.assertEqual(records[-1]["details"]["invalid_records"], 1)
|
|
self.assertNotIn("must not be accepted", json.dumps(records))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|