feat(agent): per-turn attribution columns and the capture policy for persisted tool I/O

[GabrielDrapor]

Two changes that both write to Postgres and are therefore independent of any trace flag.

Attribution. ai_messages gains turn_kind (rca_execution | rca_followup | chat | detector | digest, set per turn at write time), execution_id (FK to detector_rca_executions) and initiator_user_id; ai_sessions gains execution_id. Today all of this is inferred from ai_sessions.user_id IS NULL. The migration fills turn_kind deterministically from the existing kind column — a column rename in effect, not a backfill — so metering never needs a dual read path.

Capture policy for persisted tool I/O, one module applied identically to span attributes and tool_step metadata:

• args captured for every tool, after redaction (ghp_/gho_, sk-, AKIA, bearer tokens, .env-style assignments)
• output captured only for an allowlist — download_traces, download_session, submit_result, i.e. data the customer already owns in TraceRoot. bash, read, write, git_clone, check_github_access record exit code, byte count and duration only
• redact, then truncate — truncating first can split a token and defeat the pattern
• budgets: 8 KB per step, 256 KB per run; over budget, later steps degrade to size-only

Lands before any emit, because the step rows write to Postgres immediately and are not behind the flag.

Closes #2062. Part of #2058.

[trident-sentinel]

PR overview

This pull request adds per-turn attribution, durable streaming message persistence, shared tool-I/O redaction, and per-session chat stream isolation.

One security concern remains open: an authenticated user can cause multi-megabyte tool arguments to bypass the capture budgets and consume Postgres storage and write capacity. Exploitation requires an authenticated project user and is limited to resource exhaustion.

Open issues (1)

• Low: Unbounded tool arguments — frontend/packages/core/src/lib/capture-policy.ts:71

_{Scanned with Semgrep · TruffleHog · Trident review. View in Trident}

Fixed/addressed: 1 · PR risk: 4/10

[cubic-dev-ai]

1 issue found across 12 files

Confidence score: 3/5

• In frontend/ee/agent/src/stream-persister.ts, non-allowlisted tool rows omit exit code and duration while retaining only byte count and isError, which can make execution outcomes and performance unavailable for bash, read, and writ...; pass status and timing through capture and persist those fields.

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="frontend/ee/agent/src/stream-persister.ts">

<violation number="1" location="frontend/ee/agent/src/stream-persister.ts:66">
P2: For non-allowlisted tools, these rows record only byte count and `isError`; they omit exit code and duration. Pass execution status and timing through the capture path and persist those fields so `bash`, `read`, `write`, `git_clone`, and `check_github_access` retain the required size-only diagnostics.</violation>
</file>

Architecture diagram

sequenceDiagram
    participant Client as Client
    participant API as Agent API (index.ts)
    participant SM as SessionManager
    participant CP as CapturePolicy
    participant DB as Postgres (Prisma)
    participant Worker as Worker Processors
    participant Trace as TraceRoot SDK

    Note over Client,DB: Per-Turn Attribution Flow (NEW)
    
    Client->>API: POST /sessions (with optional executionId)
    API->>DB: createSession(executionId)
    DB-->>API: Session created
    
    Client->>API: POST /messages (prompt)
    API->>DB: Find session (with executionId)
    DB-->>API: Session
    
    alt Session.userId === null
        alt Request has x-user-id
            Note over API: Attribution: rca_followup
        else
            Note over API: Attribution: rca_execution
        end
    else
        Note over API: Attribution: chat
    end
    
    API->>SM: appendMessage("user", prompt, attribution)
    SM->>DB: Create AIMessage with turnKind, executionId, initiatorUserId
    SM->>SM: deriveAttribution() fallback
    DB-->>SM: Message persisted
    SM-->>API: Message record
    
    Note over API,DB: Agent Run + Tool Capture
    
    API->>API: Create StreamPersister with attribution
    
    loop Agent events
        API->>API: Process tool_execution_start
        API->>API: Store args in pendingToolArgs
        API->>API: Process tool_execution_end
        API->>API: applyCapturePolicy(toolName, args, result)
        API->>CP: Redact args (deep)
        API->>CP: Check output allowlist
        
        alt Tool in OUTPUT_ALLOWLIST
            CP->>CP: Redact result, truncate to 8KB
            CP-->>API: Keep result + outputBytes
        else Tool not allowlisted
            CP-->>API: Withhold result, keep outputBytes only
        end
        
        API->>SM: appendMessage("tool_step", metadata, attribution)
        SM->>DB: Insert tool_step row with captured I/O
    end
    
    API->>SM: appendMessage("assistant", response, attribution)
    SM->>DB: Insert assistant message
    
    Note over Worker,DB: Detector/Digest Attribution (CHANGED)
    
    Worker->>DB: Insert detector messages with turnKind="detector"
    Worker->>DB: Insert digest messages with turnKind="digest"
    
    Note over DB: Migration Backfill
    
    DB->>DB: Add turn_kind columns
    DB->>DB: Backfill turn_kind from legacy kind column
    DB->>DB: Set initiator_user_id from session
    DB->>DB: Add indexes for turnKind queries

Loading

<sub>Tip: instead of fixing issues one by one fix them all with cubic

Re-trigger cubic</sub>

[cubic-dev-ai]

P2: For non-allowlisted tools, these rows record only byte count and isError; they omit exit code and duration. Pass execution status and timing through the capture path and persist those fields so bash, read, write, git_clone, and check_github_access retain the required size-only diagnostics.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At frontend/ee/agent/src/stream-persister.ts, line 66:

<comment>For non-allowlisted tools, these rows record only byte count and `isError`; they omit exit code and duration. Pass execution status and timing through the capture path and persist those fields so `bash`, `read`, `write`, `git_clone`, and `check_github_access` retain the required size-only diagnostics.</comment>

<file context>
@@ -50,11 +54,18 @@ export class StreamPersister {
-        result: event.result,
+        args: captured.args,
+        ...(captured.result !== undefined ? { result: captured.result } : {}),
+        outputBytes: captured.outputBytes,
+        ...(captured.truncated ? { truncated: true } : {}),
+        ...(captured.withheld ? { withheld: captured.withheld } : {}),
</file context>

[GabrielDrapor]

Half right, and the half that holds is fixed: duration is now recorded for every step, which matters most exactly where output is withheld — it and the byte count are all that survives. Exit code is not available: pi's tool_execution_end carries isError and the result, no status field, so there is nothing to pass through. Surfacing it would need the same SDK-side channel tracked for the tool span id (traceroot-ai/traceroot-ts#150).

[GabrielDrapor]

⚠️ This PR currently carries #1960 / #1961's commits, and should not be reviewed or merged until they land.

The stack was restructured so agent self-trace does not share a stack with the chat-session work. #2067–#2070 are now independent and based on main. This PR cannot be: the attribution work modifies the exact code in frontend/ee/agent/src/index.ts that #1961's stream persister introduces — attribution is applied to the rows that persister writes, so without it there is nothing to attribute. That is a real functional dependency, not a patch-context one (verified by cherry-picking onto main without #1961: conflict in index.ts).

So the 10 commits from #1960 and #1961 ride along here for now, and the only commits belonging to this PR are the last two:

• feat(agent): per-turn attribution columns on ai_messages
• feat(agent): capture policy for persisted tool I/O

Once #1961 merges to main, this branch gets rebased and those 10 commits disappear from the diff. Marked draft until then. The same applies to everything above it in the stack (#2072–#2074).

[trident-sentinel]

Low: Unbounded tool arguments

applyCapturePolicy redacts input.args but applies perStepBytes and perRunBytes only to the serialized tool result, while StreamPersister writes captured.args directly into tool_step metadata at line 64. An authenticated project user can send a message to /api/v1/projects/p/sessions/s/messages that causes a tool call with a multi-megabyte argument such as { "command": "A".repeat(10000000) }, and the argument is persisted even after the run's 256 KB budget is exhausted. This lets the attacker consume excessive Postgres row storage and write capacity through oversized tool-I/O records. Bound the redacted serialized arguments and include their bytes in the per-step and per-run budgets before persisting them.

[GabrielDrapor]

Same finding as the P2 on line 71 and fixed with it: args are bounded and charged to the step/run budget rather than written through unbounded. Thanks for pinning the exact write path in StreamPersister — that is what made the omission obvious.

[cubic-dev-ai]

1 issue found across 10 files (changes from recent commits).

Confidence score: 3/5

• frontend/packages/core/src/lib/capture-policy.ts: capArgs resets the 8 KB limit for each argument leaf, so a step with multiple leaves can exceed the intended per-step capture ceiling and produce oversized output; track one remaining byte budget across all arguments before applying the output limit.

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="frontend/packages/core/src/lib/capture-policy.ts">

<violation number="1" location="frontend/packages/core/src/lib/capture-policy.ts:125">
P2: A step with multiple argument leaves can exceed the 8 KB per-step capture ceiling because `capArgs` resets the limit for every leaf, then the output gets a separate limit. Track one per-step remaining byte budget across the complete persisted args and result, including structural JSON bytes, before updating the run total.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Fix all with cubic | Re-trigger cubic</sub>

[cubic-dev-ai]

P2: A step with multiple argument leaves can exceed the 8 KB per-step capture ceiling because capArgs resets the limit for every leaf, then the output gets a separate limit. Track one per-step remaining byte budget across the complete persisted args and result, including structural JSON bytes, before updating the run total.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At frontend/packages/core/src/lib/capture-policy.ts, line 125:

<comment>A step with multiple argument leaves can exceed the 8 KB per-step capture ceiling because `capArgs` resets the limit for every leaf, then the output gets a separate limit. Track one per-step remaining byte budget across the complete persisted args and result, including structural JSON bytes, before updating the run total.</comment>

<file context>
@@ -68,16 +89,48 @@ export function applyCapturePolicy(
+    if (typeof value === "string") {
+      const remaining = budget.perRunBytes - state.spentBytes;
+      if (remaining <= 0) return "[withheld: budget]";
+      const { text } = truncateTo(value, Math.min(budget.perStepBytes, remaining));
+      state.spentBytes += Buffer.byteLength(text, "utf8");
+      return text;
</file context>

[GabrielDrapor]

Valid — the per-leaf reset let an args object with N string fields keep N × perStepBytes, and the result then got its own allowance on top. One remaining-byte budget is now shared across every leaf and the result. Writing the test found a related off-by-three: with less room left than the truncation marker itself, truncateTo still appended it and went over; it keeps nothing in that case now.