trace-agent-server 0.4.0__tar.gz → 0.6.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (51) hide show
  1. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/PKG-INFO +1 -1
  2. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/pyproject.toml +1 -1
  3. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/api/runs.py +24 -10
  4. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/config.py +2 -0
  5. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/db.py +2 -1
  6. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/main.py +36 -1
  7. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/schemas.py +4 -0
  8. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/events.py +6 -2
  9. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/graph.py +42 -13
  10. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/otlp.py +18 -4
  11. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/run_lifecycle.py +33 -3
  12. trace_agent_server-0.6.0/src/trace_agent_server/services/stored_redaction.py +57 -0
  13. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server.egg-info/PKG-INFO +1 -1
  14. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server.egg-info/SOURCES.txt +5 -0
  15. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_audit_guardrails_api.py +16 -0
  16. trace_agent_server-0.6.0/tests/test_concurrent_writers.py +85 -0
  17. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_framework_observations.py +39 -0
  18. trace_agent_server-0.6.0/tests/test_graph.py +60 -0
  19. trace_agent_server-0.6.0/tests/test_inactivity_sweep.py +50 -0
  20. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_otlp_ingest.py +19 -0
  21. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_run_lifecycle.py +72 -0
  22. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_runs_api.py +4 -1
  23. trace_agent_server-0.6.0/tests/test_stored_redaction.py +148 -0
  24. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/README.md +0 -0
  25. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/setup.cfg +0 -0
  26. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/__init__.py +0 -0
  27. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/api/__init__.py +0 -0
  28. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/api/otlp.py +0 -0
  29. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/cli.py +0 -0
  30. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/models.py +0 -0
  31. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/__init__.py +0 -0
  32. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/analytics.py +0 -0
  33. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/audit_guardrails.py +0 -0
  34. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/compare.py +0 -0
  35. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/dashboard.py +0 -0
  36. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/explanation.py +0 -0
  37. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/observability.py +0 -0
  38. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/otlp_decode.py +0 -0
  39. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/provider.py +0 -0
  40. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server/services/run_summary.py +0 -0
  41. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server.egg-info/dependency_links.txt +0 -0
  42. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server.egg-info/entry_points.txt +0 -0
  43. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server.egg-info/requires.txt +0 -0
  44. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/src/trace_agent_server.egg-info/top_level.txt +0 -0
  45. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_dashboard.py +0 -0
  46. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_events.py +0 -0
  47. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_otlp_decode.py +0 -0
  48. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_providers.py +0 -0
  49. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_sdk_wrappers.py +0 -0
  50. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_sse.py +0 -0
  51. {trace_agent_server-0.4.0 → trace_agent_server-0.6.0}/tests/test_turns.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: trace-agent-server
3
- Version: 0.4.0
3
+ Version: 0.6.0
4
4
  Summary: TraceAgent FastAPI backend server.
5
5
  Author-email: Enrique Javier Villar Cea <villarcea.enriquejavier@gmail.com>, Javier Rodriguez Castellano <javirodriguezcastellano@gmail.com>
6
6
  License: MIT
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "trace-agent-server"
7
- version = "0.4.0"
7
+ version = "0.6.0"
8
8
  description = "TraceAgent FastAPI backend server."
9
9
  readme = "README.md"
10
10
  authors = [
@@ -12,7 +12,7 @@ from sqlalchemy.orm import Session
12
12
 
13
13
  from trace_agent_server.db import get_session
14
14
  from trace_agent_server.models import Artifact, Decision, Event, Run, Turn
15
- from trace_agent_sdk.redaction import stable_hash, summarize_value
15
+ from trace_agent_sdk.redaction import redact_keys, stable_hash, summarize_value
16
16
  from trace_agent_server.schemas import (
17
17
  ArtifactOut,
18
18
  DashboardResponse,
@@ -36,6 +36,7 @@ from trace_agent_server.schemas import (
36
36
  TurnCreateResponse,
37
37
  TurnOut,
38
38
  )
39
+ from trace_agent_server.services.stored_redaction import redact_messages, redact_tool_calls
39
40
  from trace_agent_server.services.analytics import build_similar_runs, build_tool_graph, build_tool_metrics
40
41
  from trace_agent_server.services.compare import compare_runs
41
42
  from trace_agent_server.services.events import create_artifact, record_event
@@ -295,6 +296,8 @@ def create_turn(
295
296
  raise HTTPException(status_code=400, detail="Input blocked by guardrails.")
296
297
 
297
298
  tool_redaction = {tool.name: tool.redact_fields for tool in payload.tools}
299
+ # What is stored is a masked copy; the provider and the agent keep the real data.
300
+ stored_messages = redact_messages(sanitized_messages, tool_redaction)
298
301
  if payload.tools:
299
302
  run.run_metadata = {
300
303
  **run.run_metadata,
@@ -315,7 +318,7 @@ def create_turn(
315
318
  provider=payload.provider,
316
319
  model=payload.model,
317
320
  payload_hash=stable_hash(sanitized_messages),
318
- payload_full={"messages": sanitized_messages} if capture_full else None,
321
+ payload_full={"messages": stored_messages} if capture_full else None,
319
322
  metadata={
320
323
  "registered_tool_count": len(payload.tools),
321
324
  "tool_names": [tool.name for tool in payload.tools],
@@ -355,6 +358,10 @@ def create_turn(
355
358
  run.provider = payload.provider
356
359
  run.model = payload.model
357
360
 
361
+ stored_assistant_message = (
362
+ redact_messages([response.assistant_message], tool_redaction)[0] if response.assistant_message else None
363
+ )
364
+
358
365
  # Create Turn record with full messages and token usage
359
366
  usage = response.metadata.get("usage", {}) if response.metadata else {}
360
367
  turn = Turn(
@@ -362,9 +369,12 @@ def create_turn(
362
369
  seq=run.last_event_seq + 1,
363
370
  provider=payload.provider,
364
371
  model=payload.model,
365
- messages=sanitized_messages,
366
- assistant_message=response.assistant_message,
367
- tool_calls=[{"id": c.id, "name": c.name, "arguments": c.arguments, "step_id": c.step_id} for c in response.tool_calls],
372
+ messages=stored_messages,
373
+ assistant_message=stored_assistant_message,
374
+ tool_calls=redact_tool_calls(
375
+ [{"id": c.id, "name": c.name, "arguments": c.arguments, "step_id": c.step_id} for c in response.tool_calls],
376
+ tool_redaction,
377
+ ),
368
378
  prompt_tokens=usage.get("prompt_tokens"),
369
379
  completion_tokens=usage.get("completion_tokens"),
370
380
  total_tokens=usage.get("total_tokens"),
@@ -409,7 +419,7 @@ def create_turn(
409
419
  provider=payload.provider,
410
420
  model=payload.model,
411
421
  payload_hash=stable_hash(response.assistant_message or [tool_call.__dict__ for tool_call in response.tool_calls]),
412
- payload_full={"assistant_message": response.assistant_message, "tool_calls": turn.tool_calls} if capture_full else None,
422
+ payload_full={"assistant_message": stored_assistant_message, "tool_calls": turn.tool_calls} if capture_full else None,
413
423
  turn_id=turn.id,
414
424
  metadata=response_metadata,
415
425
  )
@@ -430,7 +440,7 @@ def create_turn(
430
440
  model=payload.model,
431
441
  tool_name=call.name,
432
442
  payload_hash=stable_hash(call.arguments),
433
- payload_full={"arguments": call.arguments} if capture_full else None,
443
+ payload_full={"arguments": redact_keys(call.arguments, tool_redaction.get(call.name, []))} if capture_full else None,
434
444
  turn_id=turn.id,
435
445
  metadata={"call_id": call.id},
436
446
  )
@@ -710,7 +720,7 @@ def store_observations(
710
720
  model=observation.model,
711
721
  error_code=observation.error_code,
712
722
  payload_hash=stable_hash(observation.metadata),
713
- payload_full=observation.metadata if capture_full else None,
723
+ payload_full={**observation.metadata, **(observation.content or {})} if capture_full else None,
714
724
  metadata={
715
725
  "provider": observation.provider,
716
726
  "model": observation.model,
@@ -799,10 +809,12 @@ def store_observations(
799
809
  def finish_run(
800
810
  run_id: str,
801
811
  payload: FinishRunRequest,
812
+ request: Request,
802
813
  session: Session = Depends(get_session),
803
814
  ) -> dict[str, Any]:
804
815
  run = _get_run_or_404(session, run_id)
805
- derivation = complete_run(session, run, payload.final_output, payload.artifacts)
816
+ capture_full = getattr(request.app.state.settings, "capture_full_payloads", True)
817
+ derivation = complete_run(session, run, payload.final_output, payload.artifacts, capture_full=capture_full)
806
818
  session.commit()
807
819
  return {"status": "completed", "run_id": run.id, "decisions": len(derivation["decisions"])}
808
820
 
@@ -962,7 +974,7 @@ def get_artifacts(run_id: str, session: Session = Depends(get_session)) -> list[
962
974
 
963
975
 
964
976
  @router.get("/runs/{run_id}/audit")
965
- def get_run_audit(run_id: str, session: Session = Depends(get_session)) -> dict[str, Any]:
977
+ def get_run_audit(run_id: str, request: Request, session: Session = Depends(get_session)) -> dict[str, Any]:
966
978
  """Return a structured audit report for a run based on guardrail events."""
967
979
  run = _get_run_or_404(session, run_id)
968
980
  events = list(session.scalars(select(Event).where(Event.run_id == run_id).order_by(Event.seq)))
@@ -1020,6 +1032,8 @@ def get_run_audit(run_id: str, session: Session = Depends(get_session)) -> dict[
1020
1032
 
1021
1033
  return {
1022
1034
  "run_id": run_id,
1035
+ # Without guardrails nothing is checked, so the score would only be a meaningless 100.
1036
+ "guardrails_enabled": bool(getattr(request.app.state.settings, "audit_enable_guardrails", False)),
1023
1037
  "security_score": security_score,
1024
1038
  "pii_hit_count": pii_hit_count,
1025
1039
  "injection_hit_count": injection_hit_count,
@@ -45,6 +45,8 @@ class Settings(BaseSettings):
45
45
  capture_full_payloads: bool = True
46
46
  otlp_enabled: bool = True
47
47
  otlp_max_body_bytes: int = 10_000_000
48
+ # A run with no event for this long is closed as failed (0 turns it off).
49
+ run_inactivity_timeout_seconds: int = 3600
48
50
 
49
51
  model_config = SettingsConfigDict(
50
52
  env_prefix="TRACE_AGENT_",
@@ -12,7 +12,8 @@ class Base(DeclarativeBase):
12
12
 
13
13
 
14
14
  def create_session_factory(database_url: str) -> tuple[object, sessionmaker[Session]]:
15
- connect_args = {"check_same_thread": False} if database_url.startswith("sqlite") else {}
15
+ # SQLite lets one writer at a time: a request that finds the database busy waits for it instead of failing.
16
+ connect_args = {"check_same_thread": False, "timeout": 30} if database_url.startswith("sqlite") else {}
16
17
  engine = create_engine(database_url, connect_args=connect_args, future=True)
17
18
  session_factory = sessionmaker(bind=engine, autoflush=False, autocommit=False, expire_on_commit=False)
18
19
  return engine, session_factory
@@ -1,5 +1,10 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import asyncio
4
+ from collections.abc import AsyncIterator
5
+ from contextlib import asynccontextmanager
6
+ from datetime import timedelta
7
+
3
8
  from fastapi import FastAPI
4
9
  from fastapi.middleware.cors import CORSMiddleware
5
10
 
@@ -9,6 +14,7 @@ from trace_agent_server.config import Settings, get_settings
9
14
  from trace_agent_server.db import Base, create_session_factory
10
15
  from trace_agent_server.services.audit_guardrails import AuditGuardrails
11
16
  from trace_agent_server.services.observability import AppObservability, configure_opentelemetry
17
+ from trace_agent_server.services.run_lifecycle import expire_inactive_runs
12
18
  from trace_agent_server.services.provider import (
13
19
  AnthropicProviderAdapter,
14
20
  BaseProviderAdapter,
@@ -41,9 +47,38 @@ def build_provider_registry(settings: Settings) -> dict[str, BaseProviderAdapter
41
47
  return registry
42
48
 
43
49
 
50
+ def _close_inactive_runs(app: FastAPI, timeout: timedelta) -> None:
51
+ with app.state.session_factory() as session:
52
+ expire_inactive_runs(session, timeout=timeout)
53
+
54
+
55
+ @asynccontextmanager
56
+ async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
57
+ """While the server is up, close the runs whose agent disappeared without finishing them."""
58
+ seconds = app.state.settings.run_inactivity_timeout_seconds
59
+ task: asyncio.Task[None] | None = None
60
+ if seconds > 0:
61
+ timeout = timedelta(seconds=seconds)
62
+
63
+ async def sweep() -> None:
64
+ while True:
65
+ await asyncio.sleep(min(60.0, max(1.0, seconds / 4)))
66
+ try:
67
+ await asyncio.to_thread(_close_inactive_runs, app, timeout)
68
+ except Exception: # a failed sweep must never take the server down; try again next time
69
+ pass
70
+
71
+ task = asyncio.create_task(sweep())
72
+ try:
73
+ yield
74
+ finally:
75
+ if task is not None:
76
+ task.cancel()
77
+
78
+
44
79
  def create_app(settings: Settings | None = None) -> FastAPI:
45
80
  current_settings = settings or get_settings()
46
- app = FastAPI(title=current_settings.app_name)
81
+ app = FastAPI(title=current_settings.app_name, lifespan=_lifespan)
47
82
  app.add_middleware(
48
83
  CORSMiddleware,
49
84
  allow_origins=current_settings.cors_origins,
@@ -61,6 +61,8 @@ class ToolResultInput(BaseModel):
61
61
  status: Literal["succeeded", "failed"]
62
62
  args: dict[str, Any] = Field(default_factory=dict)
63
63
  output_summary: str | None = None
64
+ # What the tool returned, whole (masked by the sender). Stored only with full payloads on.
65
+ output: Any | None = None
64
66
  output_hash: str | None = None
65
67
  error_code: str | None = None
66
68
  error_summary: str | None = None
@@ -113,6 +115,8 @@ class ObservationInput(BaseModel):
113
115
  provider: str | None = None
114
116
  input_summary: str | None = None
115
117
  artifact: ArtifactInput | None = None
118
+ # What a model call was asked and answered (input_messages / output_messages). Stored only with full payloads on.
119
+ content: dict[str, Any] | None = None
116
120
  metadata: dict[str, Any] = Field(default_factory=dict)
117
121
 
118
122
 
@@ -4,6 +4,7 @@ import asyncio
4
4
  from datetime import datetime, timezone
5
5
  from typing import Any
6
6
 
7
+ from sqlalchemy import update
7
8
  from sqlalchemy.orm import Session
8
9
 
9
10
  from trace_agent_server.models import Artifact, Event, Run
@@ -66,8 +67,11 @@ def record_event(
66
67
  timestamp: datetime | None = None,
67
68
  ) -> Event:
68
69
  """Record an event; ``timestamp`` overrides "now" for sources that report when it happened."""
69
- run.last_event_seq += 1
70
- run.updated_at = utcnow()
70
+ # The database increments the counter, it is not computed here: two requests writing to the same run at once
71
+ # (the agent recording by hand while the capability posts from its thread) would read the same value and
72
+ # collide on (run_id, seq). The second request waits for the first one's commit and gets the next number.
73
+ session.execute(update(Run).where(Run.id == run.id).values(last_event_seq=Run.last_event_seq + 1, updated_at=utcnow()))
74
+ session.refresh(run, ["last_event_seq"])
71
75
  event = Event(
72
76
  run_id=run.id,
73
77
  turn_id=turn_id,
@@ -5,36 +5,65 @@ from typing import Any
5
5
  from trace_agent_server.models import Artifact, Decision, Event
6
6
 
7
7
 
8
+ def _steps(events: list[Event]) -> list[list[Event]]:
9
+ """Events grouped by the step they belong to (a tool call's start and end, a model call's request and
10
+ response), in the order the steps began. An event without a step is a step of its own."""
11
+ steps: list[list[Event]] = []
12
+ position: dict[str, int] = {}
13
+ for event in events:
14
+ if event.step_id and event.step_id in position:
15
+ steps[position[event.step_id]].append(event)
16
+ continue
17
+ if event.step_id:
18
+ position[event.step_id] = len(steps)
19
+ steps.append([event])
20
+ return steps
21
+
22
+
8
23
  def build_execution_graph(events: list[Event], artifacts: list[Artifact]) -> dict[str, Any]:
24
+ """One node per step, not per event: a tool is one node with its result, not a 'started' and a 'finished'."""
9
25
  nodes: list[dict[str, Any]] = []
10
26
  edges: list[dict[str, Any]] = []
27
+ node_of_event: dict[int, str] = {}
11
28
  x = 0.0
12
29
  y = 0.0
13
30
 
14
- for index, event in enumerate(events):
31
+ steps = _steps(events)
32
+ for index, step in enumerate(steps):
33
+ last = step[-1]
34
+ node_id = f"event-{last.id}"
35
+ for event in step:
36
+ node_of_event[event.id] = node_id
37
+ failed = any(event.status == "failed" for event in step)
38
+ command = next((e.event_metadata.get("command") for e in step if e.type.startswith("command.") and e.event_metadata), None)
15
39
  nodes.append(
16
40
  {
17
- "id": f"event-{event.id}",
41
+ "id": node_id,
18
42
  "type": "event",
19
- "label": event.type,
43
+ "label": last.type,
20
44
  "position": {"x": x, "y": y},
21
45
  "data": {
22
- "summary": event.summary,
23
- "eventId": event.id,
24
- "toolName": event.tool_name,
25
- "status": event.status,
26
- "evidenceIds": [event.id],
46
+ "summary": last.summary,
47
+ # The last event of the step (its result) is the one the inspector opens.
48
+ "eventId": last.id,
49
+ "toolName": last.tool_name,
50
+ "status": "failed" if failed else last.status,
51
+ "durationMs": last.duration_ms,
52
+ "stepKind": step[0].type.split(".")[0],
53
+ "command": command,
54
+ "evidenceIds": [event.id for event in step],
27
55
  },
28
56
  }
29
57
  )
30
58
  if index > 0:
59
+ previous = nodes[index - 1]
31
60
  edges.append(
32
61
  {
33
- "id": f"edge-event-{events[index - 1].id}-{event.id}",
34
- "source": f"event-{events[index - 1].id}",
35
- "target": f"event-{event.id}",
62
+ "id": f"edge-{previous['id']}-{node_id}",
63
+ "source": previous["id"],
64
+ "target": node_id,
36
65
  "label": None,
37
- "animated": event.type in {"tool.failed", "run.failed"},
66
+ "animated": failed or last.type == "run.failed",
38
67
  }
39
68
  )
40
69
  x += 220
@@ -62,7 +91,7 @@ def build_execution_graph(events: list[Event], artifacts: list[Artifact]) -> dic
62
91
  edges.append(
63
92
  {
64
93
  "id": f"edge-artifact-{artifact.id}",
65
- "source": f"event-{artifact.source_event_id}",
94
+ "source": node_of_event.get(artifact.source_event_id, f"event-{artifact.source_event_id}"),
66
95
  "target": f"artifact-{artifact.id}",
67
96
  "label": "artifact",
68
97
  "animated": False,
@@ -120,6 +120,17 @@ def _first_user_text(input_messages: Any) -> str | None:
120
120
  return None
121
121
 
122
122
 
123
+ def _with_instructions(input_messages: Any, instructions: Any) -> Any:
124
+ """The instructions of the agent come beside the messages in the GenAI conventions: show them first, as a system message."""
125
+ parts = _parse_json(instructions)
126
+ if isinstance(parts, str):
127
+ parts = [{"type": "text", "content": parts}]
128
+ if not isinstance(parts, list) or not parts:
129
+ return input_messages
130
+ messages = input_messages if isinstance(input_messages, list) else []
131
+ return [{"role": "system", "parts": parts}, *messages]
132
+
133
+
123
134
  def _requested_tools(output_messages: Any) -> list[str]:
124
135
  names: list[str] = []
125
136
  for message in output_messages if isinstance(output_messages, list) else []:
@@ -196,7 +207,7 @@ def _ingest_run(session: Session, run_id: str, spans: list[SpanRecord], capture_
196
207
  _track_agent_calls(run, fresh)
197
208
  closing = _closing_span(run, fresh) if was_open else None
198
209
  if closing is not None:
199
- _close_run(session, run, closing)
210
+ _close_run(session, run, closing, capture_full)
200
211
  elif not was_open:
201
212
  # Late spans of a finished run: refresh its decisions and flags.
202
213
  derive_explanation(session, run)
@@ -328,7 +339,9 @@ def _model_steps(session: Session, run: Run, span: SpanRecord, capture_full: boo
328
339
  payload = None
329
340
  if capture_full:
330
341
  payload = {
331
- "input_messages": _parse_json(attributes.get("gen_ai.input.messages")),
342
+ "input_messages": _with_instructions(
343
+ _parse_json(attributes.get("gen_ai.input.messages")), attributes.get("gen_ai.system_instructions")
344
+ ),
332
345
  "output_messages": output_messages,
333
346
  }
334
347
  record_event(
@@ -360,6 +373,7 @@ def _tool_steps(session: Session, run: Run, span: SpanRecord, capture_full: bool
360
373
  status="failed" if failed else "succeeded",
361
374
  args=redact_payload(arguments if isinstance(arguments, dict) else {"input": arguments} if arguments is not None else {}),
362
375
  output_summary=summarize_value(_readable(result_text)) if result_text is not None else None,
376
+ output=_parse_json(result_text) if result_text is not None and not failed else None,
363
377
  output_hash=stable_hash(result_text if not failed else {"error": message}),
364
378
  error_code=code,
365
379
  error_summary=message,
@@ -448,9 +462,9 @@ def _closing_span(run: Run, spans: list[SpanRecord]) -> SpanRecord | None:
448
462
  return max(agents, key=lambda s: s.end)
449
463
 
450
464
 
451
- def _close_run(session: Session, run: Run, span: SpanRecord) -> None:
465
+ def _close_run(session: Session, run: Run, span: SpanRecord, capture_full: bool) -> None:
452
466
  if span.status_code == "ERROR":
453
467
  code, message = _error_of(span)
454
468
  fail_run_with(session, run, code, message, ended_at=span.end)
455
469
  else:
456
- complete_run(session, run, span.attributes.get("final_result"), ended_at=span.end)
470
+ complete_run(session, run, span.attributes.get("final_result"), ended_at=span.end, capture_full=capture_full)
@@ -3,9 +3,10 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  from collections.abc import Sequence
6
- from datetime import datetime
6
+ from datetime import datetime, timedelta
7
7
  from typing import Any
8
8
 
9
+ from sqlalchemy import select
9
10
  from sqlalchemy.orm import Session
10
11
 
11
12
  from trace_agent_sdk.redaction import stable_hash, summarize_value
@@ -85,7 +86,11 @@ def record_tool_result(
85
86
  tool_name=result.name,
86
87
  duration_ms=result.duration_ms,
87
88
  payload_hash=result.output_hash,
88
- payload_full={"output_summary": result.output_summary} if capture_full else None,
89
+ payload_full=(
90
+ {"output_summary": result.output_summary, **({"output": result.output} if result.output is not None else {})}
91
+ if capture_full
92
+ else None
93
+ ),
89
94
  metadata={"call_id": result.call_id, "artifact_count": len(result.artifacts)},
90
95
  timestamp=ended_at,
91
96
  )
@@ -146,8 +151,12 @@ def complete_run(
146
151
  artifacts: Sequence[ArtifactInput] = (),
147
152
  *,
148
153
  ended_at: datetime | None = None,
154
+ capture_full: bool = False,
149
155
  ) -> dict[str, Any]:
150
- """Mark a run completed, record its final output and derive its explanation."""
156
+ """Mark a run completed, record its final output and derive its explanation.
157
+
158
+ The summary of the output is always kept; the whole output only with ``capture_full``.
159
+ """
151
160
  final_summary = summarize_value(final_output) if final_output is not None else "Run finished without final output."
152
161
  run.final_output_summary = final_summary
153
162
  run.status = "completed"
@@ -160,6 +169,7 @@ def complete_run(
160
169
  summary=final_summary,
161
170
  status="completed",
162
171
  payload_hash=stable_hash(final_output),
172
+ payload_full={"final_output": final_output} if capture_full and final_output is not None else None,
163
173
  timestamp=ended_at,
164
174
  )
165
175
  for artifact_input in artifacts:
@@ -213,3 +223,23 @@ def fail_run_with(
213
223
  timestamp=ended_at,
214
224
  )
215
225
  derive_explanation(session, run)
226
+
227
+
228
+ def expire_inactive_runs(session: Session, *, timeout: timedelta, now: datetime | None = None) -> int:
229
+ """Fail the runs that are still running but have had no event for ``timeout``.
230
+
231
+ An agent that dies without calling finish or fail would otherwise stay "running" forever and keep
232
+ counting as active. Returns how many runs were closed.
233
+ """
234
+ now = now or utcnow()
235
+ stale = list(session.scalars(select(Run).where(Run.status == "running", Run.updated_at < now - timeout)))
236
+ minutes = max(1, int(timeout.total_seconds() // 60))
237
+ for run in stale:
238
+ fail_run_with(
239
+ session,
240
+ run,
241
+ "RunAbandoned",
242
+ f"No activity for {minutes} minutes: the agent stopped without finishing the run.",
243
+ )
244
+ session.commit()
245
+ return len(stale)
@@ -0,0 +1,57 @@
1
+ """Copies of a model turn with the sensitive keys of each tool masked, for what the server stores.
2
+
3
+ The provider and the agent still get the real data; only what is saved (events, turns) is masked.
4
+ A tool names its sensitive keys with ``redact_fields``; they apply to that tool's arguments and result.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import json
10
+ from typing import Any
11
+
12
+ from trace_agent_sdk.redaction import redact_keys
13
+
14
+
15
+ def _redact_json_text(text: Any, fields: list[str]) -> Any:
16
+ if not fields or not isinstance(text, str):
17
+ return text
18
+ try:
19
+ parsed = json.loads(text)
20
+ except ValueError:
21
+ return text
22
+ return json.dumps(redact_keys(parsed, fields), ensure_ascii=False)
23
+
24
+
25
+ def redact_messages(messages: list[dict[str, Any]], tool_redaction: dict[str, list[str]]) -> list[dict[str, Any]]:
26
+ """The conversation with the arguments and results of each tool masked."""
27
+ if not any(tool_redaction.values()):
28
+ return messages
29
+
30
+ tool_names: dict[str, str] = {}
31
+ for message in messages:
32
+ for call in message.get("tool_calls") or []:
33
+ tool_names[str(call.get("id"))] = str((call.get("function") or {}).get("name"))
34
+
35
+ redacted: list[dict[str, Any]] = []
36
+ for message in messages:
37
+ message = dict(message)
38
+ if message.get("role") == "tool":
39
+ fields = tool_redaction.get(tool_names.get(str(message.get("tool_call_id")), ""), [])
40
+ message["content"] = _redact_json_text(message.get("content"), fields)
41
+ if message.get("tool_calls"):
42
+ calls = []
43
+ for call in message["tool_calls"]:
44
+ function = dict(call.get("function") or {})
45
+ function["arguments"] = _redact_json_text(function.get("arguments"), tool_redaction.get(str(function.get("name")), []))
46
+ calls.append({**call, "function": function})
47
+ message["tool_calls"] = calls
48
+ redacted.append(message)
49
+ return redacted
50
+
51
+
52
+ def redact_tool_calls(tool_calls: list[dict[str, Any]], tool_redaction: dict[str, list[str]]) -> list[dict[str, Any]]:
53
+ """Tool calls ({"name", "arguments", ...}) with the arguments of each tool masked."""
54
+ return [
55
+ {**call, "arguments": redact_keys(call.get("arguments"), tool_redaction.get(str(call.get("name")), []))}
56
+ for call in tool_calls
57
+ ]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: trace-agent-server
3
- Version: 0.4.0
3
+ Version: 0.6.0
4
4
  Summary: TraceAgent FastAPI backend server.
5
5
  Author-email: Enrique Javier Villar Cea <villarcea.enriquejavier@gmail.com>, Javier Rodriguez Castellano <javirodriguezcastellano@gmail.com>
6
6
  License: MIT
@@ -30,10 +30,14 @@ src/trace_agent_server/services/otlp_decode.py
30
30
  src/trace_agent_server/services/provider.py
31
31
  src/trace_agent_server/services/run_lifecycle.py
32
32
  src/trace_agent_server/services/run_summary.py
33
+ src/trace_agent_server/services/stored_redaction.py
33
34
  tests/test_audit_guardrails_api.py
35
+ tests/test_concurrent_writers.py
34
36
  tests/test_dashboard.py
35
37
  tests/test_events.py
36
38
  tests/test_framework_observations.py
39
+ tests/test_graph.py
40
+ tests/test_inactivity_sweep.py
37
41
  tests/test_otlp_decode.py
38
42
  tests/test_otlp_ingest.py
39
43
  tests/test_providers.py
@@ -41,4 +45,5 @@ tests/test_run_lifecycle.py
41
45
  tests/test_runs_api.py
42
46
  tests/test_sdk_wrappers.py
43
47
  tests/test_sse.py
48
+ tests/test_stored_redaction.py
44
49
  tests/test_turns.py
@@ -137,3 +137,19 @@ def test_metrics_endpoint_is_available_when_enabled(tmp_path) -> None:
137
137
  body = metrics.text
138
138
  assert "TRACE_AGENT_http_requests_total" in body
139
139
  assert "TRACE_AGENT_http_request_duration_seconds" in body
140
+
141
+
142
+ def test_audit_says_guardrails_are_off_by_default(tmp_path) -> None:
143
+ client = build_client(tmp_path, CountingProvider())
144
+
145
+ report = client.get(f"/api/runs/{_create_run(client)}/audit").json()
146
+
147
+ assert report["guardrails_enabled"] is False
148
+
149
+
150
+ def test_audit_says_guardrails_are_on_when_enabled(tmp_path) -> None:
151
+ client = build_client(tmp_path, CountingProvider(), audit_enable_guardrails=True)
152
+
153
+ report = client.get(f"/api/runs/{_create_run(client)}/audit").json()
154
+
155
+ assert report["guardrails_enabled"] is True
@@ -0,0 +1,85 @@
1
+ from __future__ import annotations
2
+
3
+ import threading
4
+
5
+ import pytest
6
+ from fastapi.testclient import TestClient
7
+
8
+ from trace_agent_server.config import Settings
9
+ from trace_agent_server.main import create_app
10
+ from trace_agent_server.models import Event, Run
11
+ from trace_agent_server.services.events import record_event
12
+
13
+ WRITERS = 6
14
+ EVENTS_EACH = 25
15
+
16
+
17
+ @pytest.fixture
18
+ def app(tmp_path):
19
+ return create_app(Settings(database_url=f"sqlite:///{tmp_path / 'writers.db'}", cors_origins=["http://localhost:5173"]))
20
+
21
+
22
+ class TestSeveralWritersOnOneRun:
23
+ """The capability posts from a background thread while the agent records by hand: both write to the same run."""
24
+
25
+ def test_events_recorded_at_the_same_time_never_collide(self, app) -> None:
26
+ with app.state.session_factory() as session:
27
+ run = Run(agent_name="agent", goal="goal")
28
+ session.add(run)
29
+ session.commit()
30
+ run_id = run.id
31
+
32
+ errors: list[BaseException] = []
33
+ start = threading.Barrier(WRITERS)
34
+
35
+ def write(writer: int) -> None:
36
+ try:
37
+ start.wait()
38
+ for index in range(EVENTS_EACH):
39
+ # Each request loads the run, records an event and commits, as the endpoints do.
40
+ with app.state.session_factory() as session:
41
+ run = session.get(Run, run_id)
42
+ record_event(session, run, actor="sdk", event_type="tool.started", summary=f"writer {writer} event {index}")
43
+ session.commit()
44
+ except BaseException as exc: # noqa: BLE001 - reported by the test below
45
+ errors.append(exc)
46
+
47
+ threads = [threading.Thread(target=write, args=(writer,)) for writer in range(WRITERS)]
48
+ for thread in threads:
49
+ thread.start()
50
+ for thread in threads:
51
+ thread.join(timeout=120)
52
+
53
+ assert errors == []
54
+ with app.state.session_factory() as session:
55
+ sequence = [seq for (seq,) in session.query(Event.seq).filter_by(run_id=run_id).order_by(Event.seq)]
56
+ assert session.get(Run, run_id).last_event_seq == WRITERS * EVENTS_EACH
57
+ assert sequence == list(range(1, WRITERS * EVENTS_EACH + 1))
58
+
59
+ def test_requests_to_the_same_run_at_the_same_time_all_succeed(self, app) -> None:
60
+ client = TestClient(app)
61
+ run_id = client.post("/api/runs", json={"agent_name": "agent", "goal": "goal"}).json()["run_id"]
62
+ failures: list[int] = []
63
+ start = threading.Barrier(4)
64
+
65
+ def post(writer: int) -> None:
66
+ start.wait()
67
+ for index in range(10):
68
+ response = client.post(
69
+ f"/api/runs/{run_id}/observations",
70
+ json={"observations": [{"kind": "command", "command": f"echo {writer}-{index}", "step_id": f"c-{writer}-{index}", "status": "succeeded", "exit_code": 0}]},
71
+ )
72
+ if response.status_code != 200:
73
+ failures.append(response.status_code)
74
+
75
+ threads = [threading.Thread(target=post, args=(writer,)) for writer in range(4)]
76
+ for thread in threads:
77
+ thread.start()
78
+ for thread in threads:
79
+ thread.join(timeout=120)
80
+
81
+ assert failures == []
82
+ with app.state.session_factory() as session:
83
+ seqs = [seq for (seq,) in session.query(Event.seq).filter_by(run_id=run_id).order_by(Event.seq)]
84
+ assert seqs == list(range(1, len(seqs) + 1))
85
+ assert len(seqs) == 1 + 4 * 10 * 2 # run.started, then command.started and command.succeeded for each command
@@ -120,6 +120,45 @@ class TestLlmCallUsage:
120
120
  assert execution["assistant_turn_count"] == 2
121
121
 
122
122
 
123
+ class TestLlmCallContent:
124
+ content = {
125
+ "input_messages": [{"role": "user", "parts": [{"type": "text", "content": "¿Por qué bajó el OEE?"}]}],
126
+ "output_messages": [{"role": "assistant", "parts": [{"type": "text", "content": "Por las paradas."}]}],
127
+ }
128
+
129
+ def completed_event(self, client, run_id: str) -> Event:
130
+ with client.app.state.session_factory() as session:
131
+ event = session.query(Event).filter_by(run_id=run_id, type="llm.completed").one()
132
+ session.expunge(event)
133
+ return event
134
+
135
+ def test_what_the_model_was_asked_and_answered_is_stored_with_the_event(self, client) -> None:
136
+ run_id = start_run(client)
137
+ observe(client, run_id, kind="llm_call", step_id="l1", status="succeeded", provider="openai", model="m",
138
+ content=self.content, metadata={"prompt_tokens": 3})
139
+
140
+ event = self.completed_event(client, run_id)
141
+ assert event.payload_full["input_messages"] == self.content["input_messages"]
142
+ assert event.payload_full["output_messages"] == self.content["output_messages"]
143
+ assert event.payload_full["prompt_tokens"] == 3
144
+ # The content is not copied into the metadata that is always kept and listed.
145
+ assert "input_messages" not in event.event_metadata
146
+
147
+ def test_it_is_not_stored_when_full_payloads_are_off(self, tmp_path) -> None:
148
+ app = create_app(Settings(database_url=f"sqlite:///{tmp_path / 'off.db'}", cors_origins=["http://localhost:5173"], capture_full_payloads=False))
149
+ client = TestClient(app)
150
+ run_id = start_run(client)
151
+ observe(client, run_id, kind="llm_call", step_id="l1", status="succeeded", provider="openai", model="m", content=self.content)
152
+
153
+ assert self.completed_event(client, run_id).payload_full is None
154
+
155
+ def test_a_call_without_content_is_stored_as_before(self, client) -> None:
156
+ run_id = start_run(client)
157
+ observe(client, run_id, kind="llm_call", step_id="l1", status="succeeded", provider="openai", model="m", metadata={"prompt_tokens": 3})
158
+
159
+ assert self.completed_event(client, run_id).payload_full == {"prompt_tokens": 3}
160
+
161
+
123
162
  class TestTokenUsageOfFrameworkRuns:
124
163
  def model_call(self, client, run_id, step, **tokens) -> None:
125
164
  observe(client, run_id, kind="llm_call", step_id=step, status="succeeded", provider="openai", model="m", metadata=tokens)
@@ -0,0 +1,60 @@
1
+ from __future__ import annotations
2
+
3
+ from trace_agent_server.models import Artifact, Event
4
+ from trace_agent_server.services.graph import build_execution_graph
5
+
6
+
7
+ def event(event_id: int, event_type: str, step_id: str | None, *, status: str | None = None, tool_name: str | None = None,
8
+ summary: str = "", metadata: dict | None = None, duration_ms: int | None = None) -> Event:
9
+ return Event(id=event_id, run_id="r", seq=event_id, actor="sdk", type=event_type, summary=summary, status=status,
10
+ step_id=step_id, tool_name=tool_name, event_metadata=metadata or {}, duration_ms=duration_ms)
11
+
12
+
13
+ class TestExecutionGraphSteps:
14
+ def test_a_tool_is_one_node_with_its_result_not_a_started_and_a_finished_node(self) -> None:
15
+ events = [
16
+ event(1, "run.started", None),
17
+ event(2, "tool.started", "t1", tool_name="search", status="completed"),
18
+ event(3, "tool.succeeded", "t1", tool_name="search", status="succeeded", summary="3 results", duration_ms=40),
19
+ event(4, "run.completed", None),
20
+ ]
21
+
22
+ graph = build_execution_graph(events, [])
23
+
24
+ assert [n["data"]["eventId"] for n in graph["nodes"]] == [1, 3, 4]
25
+ tool = graph["nodes"][1]
26
+ assert (tool["label"], tool["data"]["summary"], tool["data"]["durationMs"], tool["data"]["toolName"]) == ("tool.succeeded", "3 results", 40, "search")
27
+ assert tool["data"]["evidenceIds"] == [2, 3]
28
+ assert tool["data"]["stepKind"] == "tool"
29
+
30
+ def test_the_nodes_are_chained_in_the_order_the_steps_began(self) -> None:
31
+ events = [event(1, "llm.started", "l1"), event(2, "tool.started", "t1"), event(3, "llm.completed", "l1"), event(4, "tool.succeeded", "t1")]
32
+
33
+ graph = build_execution_graph(events, [])
34
+
35
+ assert [n["data"]["eventId"] for n in graph["nodes"]] == [3, 4] # each node is named after the last event of its step
36
+ assert [(e["source"], e["target"]) for e in graph["edges"]] == [("event-3", "event-4")]
37
+
38
+ def test_a_step_that_failed_at_any_point_is_failed_and_its_edge_animated(self) -> None:
39
+ events = [event(1, "tool.started", "t1", status="completed"), event(2, "tool.failed", "t1", status="failed"), event(3, "tool.started", "t2")]
40
+
41
+ graph = build_execution_graph(events, [])
42
+
43
+ assert graph["nodes"][0]["data"]["status"] == "failed"
44
+ assert graph["edges"][0]["animated"] is False # the edge into the next step follows that step, which did not fail
45
+
46
+ def test_the_command_of_a_command_step_is_kept_for_its_title(self) -> None:
47
+ events = [event(1, "command.started", "c1", metadata={"command": "python -c # Leer\nprint(1)"}), event(2, "command.succeeded", "c1", metadata={"command": "python -c # Leer\nprint(1)"})]
48
+
49
+ graph = build_execution_graph(events, [])
50
+
51
+ assert graph["nodes"][0]["data"]["command"] == "python -c # Leer\nprint(1)"
52
+ assert graph["nodes"][0]["data"]["stepKind"] == "command"
53
+
54
+ def test_an_artifact_hangs_from_the_node_of_the_step_that_made_it(self) -> None:
55
+ events = [event(1, "tool.started", "t1"), event(2, "tool.succeeded", "t1")]
56
+ artifact = Artifact(id=9, run_id="r", kind="file", label="report", summary="the report", source_event_id=1)
57
+
58
+ graph = build_execution_graph(events, [artifact])
59
+
60
+ assert ("event-2", "artifact-9") in [(e["source"], e["target"]) for e in graph["edges"]]
@@ -0,0 +1,50 @@
1
+ from __future__ import annotations
2
+
3
+ import time
4
+
5
+ from fastapi.testclient import TestClient
6
+
7
+ from trace_agent_server.config import Settings
8
+ from trace_agent_server.main import create_app
9
+
10
+
11
+ def build_app(tmp_path, timeout_seconds: int):
12
+ return create_app(
13
+ Settings(
14
+ database_url=f"sqlite:///{tmp_path / 'sweep.db'}",
15
+ cors_origins=["http://localhost:5173"],
16
+ run_inactivity_timeout_seconds=timeout_seconds,
17
+ )
18
+ )
19
+
20
+
21
+ def wait_for_status(client: TestClient, run_id: str, status: str, seconds: float) -> str:
22
+ deadline = time.time() + seconds
23
+ current = ""
24
+ while time.time() < deadline:
25
+ current = client.get(f"/api/runs/{run_id}").json()["status"]
26
+ if current == status:
27
+ break
28
+ time.sleep(0.25)
29
+ return current
30
+
31
+
32
+ class TestInactivitySweep:
33
+ def test_a_running_server_closes_a_run_whose_agent_disappeared(self, tmp_path) -> None:
34
+ with TestClient(build_app(tmp_path, timeout_seconds=1)) as client:
35
+ run_id = client.post("/api/runs", json={"agent_name": "a", "goal": "g"}).json()["run_id"]
36
+
37
+ assert wait_for_status(client, run_id, "failed", seconds=8) == "failed"
38
+
39
+ events = client.get(f"/api/runs/{run_id}/events").json()
40
+ events = events if isinstance(events, list) else events["events"]
41
+ assert events[-1]["type"] == "run.failed"
42
+ assert events[-1]["error_code"] == "RunAbandoned"
43
+
44
+ def test_a_timeout_of_zero_turns_the_sweep_off(self, tmp_path) -> None:
45
+ with TestClient(build_app(tmp_path, timeout_seconds=0)) as client:
46
+ run_id = client.post("/api/runs", json={"agent_name": "a", "goal": "g"}).json()["run_id"]
47
+
48
+ time.sleep(2.5)
49
+
50
+ assert client.get(f"/api/runs/{run_id}").json()["status"] == "running"
@@ -190,6 +190,25 @@ class TestRunCreationAndMapping:
190
190
  assert completed.event_metadata["prompt_tokens"] == 60
191
191
  assert completed.event_metadata["requested_tools"] == ["get_weather"]
192
192
 
193
+ def test_what_the_model_was_asked_and_answered_is_kept_with_the_instructions_first(self, client) -> None:
194
+ spans = agent_spans()
195
+ spans[1].attributes["gen_ai__system_instructions"] = json.dumps([{"type": "text", "content": "Be brief."}])
196
+
197
+ post(client, export(*spans))
198
+
199
+ _, events = load(client)
200
+ payload = next(e for e in events if e.type == "llm.completed").payload_full
201
+ assert payload["input_messages"][0] == {"role": "system", "parts": [{"type": "text", "content": "Be brief."}]}
202
+ assert payload["input_messages"][1]["role"] == "user"
203
+ assert payload["output_messages"][0]["parts"][0]["name"] == "get_weather"
204
+
205
+ def test_without_instructions_the_messages_are_kept_as_they_came(self, client) -> None:
206
+ post(client, export(*agent_spans()))
207
+
208
+ _, events = load(client)
209
+ payload = next(e for e in events if e.type == "llm.completed").payload_full
210
+ assert [m["role"] for m in payload["input_messages"]] == ["user"]
211
+
193
212
  def test_tool_arguments_and_result_are_recorded(self, client) -> None:
194
213
  post(client, export(*agent_spans()))
195
214
 
@@ -13,6 +13,7 @@ from trace_agent_server.services.events import record_event
13
13
  from trace_agent_server.services.run_lifecycle import (
14
14
  apply_llm_usage,
15
15
  complete_run,
16
+ expire_inactive_runs,
16
17
  fail_run_with,
17
18
  record_tool_result,
18
19
  token_count,
@@ -99,6 +100,36 @@ class TestRecordToolResult:
99
100
  assert all(e.payload_full is None for e in events_of(session, run))
100
101
 
101
102
 
103
+ class TestFinalOutputAndToolOutput:
104
+ def test_the_whole_final_output_is_kept_with_full_payloads_on_and_only_the_summary_without(self, session, run) -> None:
105
+ long_output = {"recomendacion": "x" * 600, "causa": "avería"}
106
+
107
+ complete_run(session, run, long_output, capture_full=True)
108
+ completed = [e for e in events_of(session, run) if e.type == "run.completed"][0]
109
+
110
+ assert completed.payload_full == {"final_output": long_output}
111
+ assert len(completed.summary) <= 220
112
+
113
+ def test_without_full_payloads_only_the_summary_is_stored(self, session, run) -> None:
114
+ complete_run(session, run, {"a": 1})
115
+
116
+ assert [e for e in events_of(session, run) if e.type == "run.completed"][0].payload_full is None
117
+
118
+ def test_the_whole_tool_output_is_kept_next_to_its_summary(self, session, run) -> None:
119
+ result = ToolResultInput(call_id="c1", name="check", step_id="s1", status="succeeded", args={}, output_summary="{...}", output={"veredicto": "descartado", "ventana": [60, 120]})
120
+
121
+ record_tool_result(session, run, result, capture_full=True)
122
+
123
+ done = [e for e in events_of(session, run) if e.type == "tool.succeeded"][0]
124
+ assert done.payload_full == {"output_summary": "{...}", "output": {"veredicto": "descartado", "ventana": [60, 120]}}
125
+
126
+ def test_a_tool_without_output_is_stored_as_before(self, session, run) -> None:
127
+ record_tool_result(session, run, ToolResultInput(call_id="c1", name="check", step_id="s1", status="succeeded", args={}, output_summary="ok"), capture_full=True)
128
+
129
+ done = [e for e in events_of(session, run) if e.type == "tool.succeeded"][0]
130
+ assert done.payload_full == {"output_summary": "ok"}
131
+
132
+
102
133
  class TestRunEnd:
103
134
  def test_complete_run_sets_status_summary_and_end_time(self, session, run) -> None:
104
135
  derivation = complete_run(session, run, {"answer": 42}, ended_at=START)
@@ -115,3 +146,44 @@ class TestRunEnd:
115
146
  assert (run.status, run.error_count, run.final_output_summary) == ("failed", 1, "boom")
116
147
  last = events_of(session, run)[-1]
117
148
  assert (last.type, last.error_code) == ("run.failed", "ValueError")
149
+
150
+
151
+ class TestExpireInactiveRuns:
152
+ def stale_run(self, session, *, minutes_ago: int, status: str = "running") -> Run:
153
+ run = Run(agent_name="agent", goal="goal", status=status)
154
+ session.add(run)
155
+ session.flush()
156
+ run.updated_at = START - timedelta(minutes=minutes_ago)
157
+ session.flush()
158
+ return run
159
+
160
+ def test_a_run_with_no_activity_is_failed_with_a_clear_reason(self, session) -> None:
161
+ run = self.stale_run(session, minutes_ago=90)
162
+
163
+ closed = expire_inactive_runs(session, timeout=timedelta(minutes=60), now=START)
164
+
165
+ assert closed == 1
166
+ assert run.status == "failed"
167
+ assert "No activity for 60 minutes" in run.final_output_summary
168
+ last = events_of(session, run)[-1]
169
+ assert (last.type, last.error_code) == ("run.failed", "RunAbandoned")
170
+
171
+ def test_a_recent_run_is_left_running(self, session) -> None:
172
+ run = self.stale_run(session, minutes_ago=5)
173
+
174
+ assert expire_inactive_runs(session, timeout=timedelta(minutes=60), now=START) == 0
175
+ assert run.status == "running"
176
+
177
+ def test_finished_runs_are_not_touched(self, session) -> None:
178
+ run = self.stale_run(session, minutes_ago=500, status="completed")
179
+
180
+ assert expire_inactive_runs(session, timeout=timedelta(minutes=60), now=START) == 0
181
+ assert run.status == "completed"
182
+
183
+ def test_an_active_run_keeps_going_as_events_arrive(self, session) -> None:
184
+ run = self.stale_run(session, minutes_ago=90)
185
+ record_event(session, run, actor="sdk", event_type="tool.started", summary="still working", timestamp=START)
186
+ run.updated_at = START
187
+
188
+ assert expire_inactive_runs(session, timeout=timedelta(minutes=60), now=START) == 0
189
+ assert run.status == "running"
@@ -282,7 +282,10 @@ def test_successful_run_generates_timeline_graph_and_explanation(tmp_path) -> No
282
282
  assert all(decision["alternatives"] == [] for decision in explanation["decisions"])
283
283
 
284
284
  graph = test_client.get(f"/api/runs/{run.run_id}/graph?view=execution").json()
285
- assert len(graph["nodes"]) >= len(event_types)
285
+ # One node per step: a tool's start and end are one node, and every event is evidence of some node.
286
+ assert 0 < len(graph["nodes"]) < len(event_types)
287
+ evidence = {event_id for node in graph["nodes"] for event_id in node["data"]["evidenceIds"]}
288
+ assert len(evidence) == len(event_types)
286
289
 
287
290
  analytics = test_client.get(f"/api/runs/{run.run_id}/analytics").json()
288
291
  assert analytics["tool_graph"]["view"] == "tool_chain"
@@ -0,0 +1,148 @@
1
+ from __future__ import annotations
2
+
3
+ import json
4
+
5
+ import pytest
6
+ from fastapi.testclient import TestClient
7
+
8
+ from trace_agent_sdk import TraceAgentClient
9
+ from trace_agent_server.config import Settings
10
+ from trace_agent_server.main import create_app
11
+ from trace_agent_server.models import Event, Turn
12
+ from trace_agent_server.services.provider import ProviderToolCall, ProviderTurnRequest, ProviderTurnResponse
13
+ from trace_agent_server.services.stored_redaction import redact_messages, redact_tool_calls
14
+
15
+ SECRET = "SECRET-KEY-123"
16
+ CARD = "4111111111111111"
17
+
18
+
19
+ class ChargingProvider:
20
+ """Asks for the `charge` tool once, then answers. Keeps what it was sent."""
21
+
22
+ provider_name = "openai"
23
+
24
+ def __init__(self) -> None:
25
+ self.seen: list[list[dict]] = []
26
+
27
+ def generate_turn(self, request: ProviderTurnRequest) -> ProviderTurnResponse:
28
+ self.seen.append(request.messages)
29
+ if not any(m.get("role") == "tool" for m in request.messages):
30
+ return ProviderTurnResponse(
31
+ assistant_message={
32
+ "role": "assistant",
33
+ "content": None,
34
+ "tool_calls": [
35
+ {"id": "call_1", "type": "function", "function": {"name": "charge", "arguments": json.dumps({"customer": "ana", "api_key": SECRET})}}
36
+ ],
37
+ },
38
+ tool_calls=[ProviderToolCall(id="call_1", name="charge", arguments={"customer": "ana", "api_key": SECRET}, step_id="step-1")],
39
+ summary="calls charge",
40
+ )
41
+ return ProviderTurnResponse(assistant_message={"role": "assistant", "content": "Charged."}, summary="Charged.")
42
+
43
+
44
+ @pytest.fixture
45
+ def provider() -> ChargingProvider:
46
+ return ChargingProvider()
47
+
48
+
49
+ @pytest.fixture
50
+ def client(tmp_path, provider) -> TestClient:
51
+ app = create_app(Settings(database_url=f"sqlite:///{tmp_path / 'test.db'}", cors_origins=["http://localhost:5173"]))
52
+ app.state.provider_registry = {"openai": provider}
53
+ return TestClient(app)
54
+
55
+
56
+ def everything_stored(client: TestClient) -> str:
57
+ with client.app.state.session_factory() as session:
58
+ events = [(e.summary, e.payload_full, e.event_metadata) for e in session.query(Event).all()]
59
+ turns = [(t.messages, t.assistant_message, t.tool_calls) for t in session.query(Turn).all()]
60
+ return json.dumps([events, turns], default=str)
61
+
62
+
63
+ class TestToolRedactionOnStoredData:
64
+ def run_agent(self, client: TestClient, redact_fields: list[str]):
65
+ run = TraceAgentClient("http://testserver", http_client=client).start_run("billing", "Charge ana")
66
+ received: dict = {}
67
+
68
+ @run.tool(name="charge", description="Charge a customer.", redact_fields=redact_fields)
69
+ def charge(customer: str, api_key: str) -> dict:
70
+ received["args"] = {"customer": customer, "api_key": api_key}
71
+ return {"ok": True, "customer": customer, "card": CARD}
72
+
73
+ result = run.create_model_turn([{"role": "user", "content": "Charge ana"}], model="m")
74
+ run.finish(result.assistant_message["content"])
75
+ return received, result
76
+
77
+ def test_nothing_sensitive_is_stored_anywhere(self, client, provider) -> None:
78
+ self.run_agent(client, ["api_key", "card"])
79
+
80
+ stored = everything_stored(client)
81
+ assert SECRET not in stored
82
+ assert CARD not in stored
83
+ assert "[redacted]" in stored
84
+
85
+ def test_the_tool_and_the_model_still_get_the_real_data(self, client, provider) -> None:
86
+ received, _ = self.run_agent(client, ["api_key", "card"])
87
+
88
+ assert received["args"] == {"customer": "ana", "api_key": SECRET}
89
+ tool_message = next(m for m in provider.seen[-1] if m.get("role") == "tool")
90
+ assert CARD in tool_message["content"]
91
+
92
+ def test_the_agent_is_handed_the_real_tool_call_arguments(self, client, provider) -> None:
93
+ _, result = self.run_agent(client, ["api_key", "card"])
94
+
95
+ assert result.executed_tools[0]["status"] == "succeeded"
96
+
97
+ def test_without_redact_fields_the_data_is_stored_as_before(self, client, provider) -> None:
98
+ self.run_agent(client, [])
99
+
100
+ assert SECRET in everything_stored(client)
101
+
102
+ def test_other_values_are_kept_in_full(self, client, provider) -> None:
103
+ self.run_agent(client, ["api_key"])
104
+
105
+ stored = everything_stored(client)
106
+ assert "ana" in stored
107
+ assert CARD in stored # `card` was not listed, so it is not masked
108
+
109
+
110
+ class TestRedactMessages:
111
+ def test_masks_the_arguments_of_a_call_and_the_result_of_its_tool(self) -> None:
112
+ messages = [
113
+ {"role": "user", "content": "hi"},
114
+ {"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "charge", "arguments": json.dumps({"api_key": "k", "n": 1})}}]},
115
+ {"role": "tool", "tool_call_id": "c1", "content": json.dumps({"card": "4111", "ok": True})},
116
+ ]
117
+
118
+ masked = redact_messages(messages, {"charge": ["api_key", "card"]})
119
+
120
+ assert json.loads(masked[1]["tool_calls"][0]["function"]["arguments"]) == {"api_key": "[redacted]", "n": 1}
121
+ assert json.loads(masked[2]["content"]) == {"card": "[redacted]", "ok": True}
122
+ assert messages[2]["content"] == json.dumps({"card": "4111", "ok": True}) # the original is untouched
123
+
124
+ def test_a_tool_without_redact_fields_and_plain_text_are_left_alone(self) -> None:
125
+ messages = [
126
+ {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "other", "arguments": '{"a": 1}'}}]},
127
+ {"role": "tool", "tool_call_id": "c1", "content": "plain text"},
128
+ ]
129
+
130
+ assert redact_messages(messages, {"charge": ["api_key"]}) == messages
131
+ assert redact_messages(messages, {}) is messages
132
+
133
+ def test_accents_survive_the_masking(self) -> None:
134
+ messages = [
135
+ {"role": "assistant", "tool_calls": [{"id": "c1", "function": {"name": "t", "arguments": json.dumps({"pw": "x", "msg": "bajó"}, ensure_ascii=False)}}]},
136
+ ]
137
+
138
+ masked = redact_messages(messages, {"t": ["pw"]})
139
+
140
+ assert "bajó" in masked[0]["tool_calls"][0]["function"]["arguments"]
141
+
142
+
143
+ class TestRedactToolCalls:
144
+ def test_masks_the_arguments_of_the_listed_tools(self) -> None:
145
+ calls = [{"id": "c1", "name": "charge", "arguments": {"api_key": "k", "n": 1}, "step_id": "s"}]
146
+
147
+ assert redact_tool_calls(calls, {"charge": ["api_key"]})[0]["arguments"] == {"api_key": "[redacted]", "n": 1}
148
+ assert calls[0]["arguments"]["api_key"] == "k"