feat(agent): capture policy for persisted tool I/O (redact, allowlist, truncate, budget)

The redactSecrets gh*_ pattern keeps the plan's documented fix ($1 capture
group preserving the prefix) instead of the initial pattern + normalizer,
since the normalizer is a no-op once the prefix is preserved on match.
Updated the pre-existing "records tool args..." persister test to use
download_traces (allowlisted) instead of get_traces so its result assertion
still holds under the new policy — get_traces is not allowlisted and would
now have its result withheld.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018DopjNH21gLaDtWeeMDEHX

Author: GabrielDrapor

frontend/ee/agent/src/__tests__/stream-persister.test.ts

@@ -17,17 +17,26 @@ const thinkingDelta = (delta: string): AgentEvent =>
     assistantMessageEvent: { type: "thinking_delta", delta } as never,
   }) as AgentEvent;
 
-const toolStart = (id: string, args: Record<string, unknown> = {}): AgentEvent => ({
+const toolStart = (
+  id: string,
+  args: Record<string, unknown> = {},
+  toolName = "get_traces",
+): AgentEvent => ({
   type: "tool_execution_start",
   toolCallId: id,
-  toolName: "get_traces",
+  toolName,
   args,
 });
 
-const toolEnd = (id: string, result: unknown = "ok", isError = false): AgentEvent => ({
+const toolEnd = (
+  id: string,
+  result: unknown = "ok",
+  isError = false,
+  toolName = "get_traces",
+): AgentEvent => ({
   type: "tool_execution_end",
   toolCallId: id,
-  toolName: "get_traces",
+  toolName,
   result,
   isError,
 });
@@ -96,17 +105,20 @@ describe("StreamPersister", () => {
 
   it("records tool args from start and result from end in metadata", async () => {
     const { persister, calls } = makePersister();
-    persister.onEvent(toolStart("t1", { query: "errors" }));
-    persister.onEvent(toolEnd("t1", { rows: [] }, true));
+    // download_traces is capture-policy-allowlisted, so its result is kept
+    // (as a redacted/bounded string) rather than withheld — see Task 9.
+    persister.onEvent(toolStart("t1", { query: "errors" }, "download_traces"));
+    persister.onEvent(toolEnd("t1", { rows: [] }, true, "download_traces"));
     await persister.finish();
 
     expect(calls).toHaveLength(1);
     expect(calls[0].role).toBe("tool_step");
     expect(calls[0].metadata).toEqual({
       toolCallId: "t1",
-      toolName: "get_traces",
+      toolName: "download_traces",
       args: { query: "errors" },
-      result: { rows: [] },
+      result: '{"rows":[]}',
+      outputBytes: 11,
       isError: true,
     });
   });
@@ -204,4 +216,44 @@ describe("StreamPersister", () => {
     expect(errorSpy).toHaveBeenCalled();
     errorSpy.mockRestore();
   });
+
+  it("withholds bash output but keeps its size; keeps download_traces output", async () => {
+    const calls: AppendCall[] = [];
+    const p = new StreamPersister(async (role, content, metadata) => {
+      calls.push({ role, content, metadata });
+    });
+    p.onEvent({
+      type: "tool_execution_start",
+      toolCallId: "1",
+      toolName: "bash",
+      args: { command: "cat secrets" },
+    });
+    p.onEvent({
+      type: "tool_execution_end",
+      toolCallId: "1",
+      toolName: "bash",
+      result: "ghp_" + "x".repeat(40),
+      isError: false,
+    });
+    p.onEvent({
+      type: "tool_execution_start",
+      toolCallId: "2",
+      toolName: "download_traces",
+      args: {},
+    });
+    p.onEvent({
+      type: "tool_execution_end",
+      toolCallId: "2",
+      toolName: "download_traces",
+      result: '{"spans":[]}',
+      isError: false,
+    });
+    await p.finish();
+    const bash = calls.find((c) => c.metadata?.toolName === "bash")!.metadata!;
+    expect(bash.result).toBeUndefined();
+    expect(bash.outputBytes).toBe(44);
+    expect(bash.withheld).toBe("not-allowlisted");
+    const dl = calls.find((c) => c.metadata?.toolName === "download_traces")!.metadata!;
+    expect(dl.result).toBe('{"spans":[]}');
+  });
 });

frontend/ee/agent/src/stream-persister.ts

@@ -1,4 +1,5 @@
 import type { AgentEvent } from "@earendil-works/pi-agent-core";
+import { applyCapturePolicy } from "@traceroot/core";
 import type { TokenUsageData, TurnAttribution } from "./session.js";
 
 /** Signature of SessionManager.appendMessage — injected so the persister is testable. */
@@ -31,6 +32,8 @@ export class StreamPersister {
   private thinking = "";
   /** args by toolCallId, captured at tool_execution_start (end events lack args) */
   private pendingToolArgs = new Map<string, Record<string, unknown>>();
+  /** Per-run capture-policy budget accumulator (see applyCapturePolicy). */
+  private readonly captureState = { spentBytes: 0 };
 
   constructor(private readonly append: AppendMessageFn) {}
 
@@ -51,11 +54,18 @@ export class StreamPersister {
     if (event.type === "tool_execution_end") {
       const args = this.pendingToolArgs.get(event.toolCallId) ?? {};
       this.pendingToolArgs.delete(event.toolCallId);
+      const captured = applyCapturePolicy(
+        { toolName: event.toolName, args, result: event.result },
+        this.captureState,
+      );
       this.enqueue("tool_step", "", {
         toolCallId: event.toolCallId,
         toolName: event.toolName,
-        args,
-        result: event.result,
+        args: captured.args,
+        ...(captured.result !== undefined ? { result: captured.result } : {}),
+        outputBytes: captured.outputBytes,
+        ...(captured.truncated ? { truncated: true } : {}),
+        ...(captured.withheld ? { withheld: captured.withheld } : {}),
         isError: event.isError,
       });
     }

frontend/packages/core/src/__tests__/capture-policy.test.ts

@@ -0,0 +1,63 @@
+import { describe, expect, it } from "vitest";
+import { applyCapturePolicy, redactSecrets } from "../lib/capture-policy.ts";
+
+describe("redactSecrets", () => {
+  it.each([
+    ["ghp_abcdefghijklmnopqrstuvwxyz0123456789", "ghp_[REDACTED]"],
+    ["gho_abcdefghijklmnopqrstuvwxyz0123456789", "gho_[REDACTED]"],
+    ["sk-proj-abcdefghijklmnopqrstuvwxyz0123456789", "sk-[REDACTED]"],
+    ["AKIAIOSFODNN7EXAMPLE", "AKIA[REDACTED]"],
+    ["Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.abc.def", "Authorization: Bearer [REDACTED]"],
+    ["OPENAI_API_KEY=abc123def456", "OPENAI_API_KEY=[REDACTED]"],
+  ])("redacts %s", (input, expected) => {
+    expect(redactSecrets(input)).toBe(expected);
+  });
+  it("leaves ordinary text alone", () => {
+    expect(redactSecrets("git log --oneline -10")).toBe("git log --oneline -10");
+  });
+});
+
+describe("applyCapturePolicy", () => {
+  it("keeps result for allowlisted tools, redacted and truncated", () => {
+    const state = { spentBytes: 0 };
+    const r = applyCapturePolicy(
+      {
+        toolName: "download_traces",
+        args: { traceId: "t" },
+        result: "token ghp_" + "x".repeat(40) + " " + "y".repeat(9000),
+      },
+      state,
+    );
+    expect(r.result).toContain("ghp_[REDACTED]");
+    expect(r.result!.length).toBeLessThanOrEqual(8_192 + 1);
+    expect(r.truncated).toBe(true);
+    expect(r.withheld).toBeNull();
+    expect(state.spentBytes).toBeGreaterThan(0);
+  });
+  it("withholds result for non-allowlisted tools but reports its size", () => {
+    const r = applyCapturePolicy(
+      { toolName: "bash", args: { command: "ls" }, result: "a".repeat(500) },
+      { spentBytes: 0 },
+    );
+    expect(r.result).toBeUndefined();
+    expect(r.outputBytes).toBe(500);
+    expect(r.withheld).toBe("not-allowlisted");
+  });
+  it("degrades to size-only once the per-run budget is spent", () => {
+    const state = { spentBytes: 262_144 };
+    const r = applyCapturePolicy({ toolName: "download_traces", args: {}, result: "small" }, state);
+    expect(r.result).toBeUndefined();
+    expect(r.withheld).toBe("budget");
+  });
+  it("redacts inside args too", () => {
+    const r = applyCapturePolicy(
+      {
+        toolName: "bash",
+        args: { command: "curl -H 'Authorization: Bearer abc.def.ghi'" },
+        result: "",
+      },
+      { spentBytes: 0 },
+    );
+    expect(JSON.stringify(r.args)).toContain("Bearer [REDACTED]");
+  });
+});

frontend/packages/core/src/index.ts

@@ -48,6 +48,9 @@ export * from "./model-pricing/index.ts";
 // Shared types
 export * from "./types/index.ts";
 
+// Capture policy for persisted tool I/O (agent StreamPersister + SDK captureToolIo)
+export * from "./lib/capture-policy.ts";
+
 // NOTE: pi-ai Model resolver lives at `@traceroot/core/model-resolver` (subpath).
 // We do NOT re-export it here — pulling pi-ai into the main barrel would bundle
 // Node-only code (`node:fs`, etc.) into the Next.js client. Server-side

frontend/packages/core/src/lib/capture-policy.ts

@@ -0,0 +1,83 @@
+/**
+ * One policy for what tool I/O may be persisted — used by the agent's StreamPersister
+ * (ai_messages.tool_step rows) and by the SDK's captureToolIo hook (spans), so both
+ * stores hold exactly the same content. Order matters: redact, then allowlist, then
+ * truncate, then budget — truncating first could split a token and defeat a pattern.
+ */
+export interface CaptureBudget {
+  perStepBytes: number;
+  perRunBytes: number;
+}
+export const DEFAULT_CAPTURE_BUDGET: CaptureBudget = { perStepBytes: 8_192, perRunBytes: 262_144 };
+
+/** Tools whose output is data the customer already owns inside TraceRoot. */
+export const OUTPUT_ALLOWLIST: ReadonlySet<string> = new Set([
+  "download_traces",
+  "download_session",
+  "submit_result",
+]);
+
+const PATTERNS: Array<[RegExp, string]> = [
+  [/\b(gh[pousr]_)[A-Za-z0-9]{20,}/g, "$1[REDACTED]"],
+  [/\bsk-[A-Za-z0-9_-]{16,}/g, "sk-[REDACTED]"],
+  [/\bAKIA[0-9A-Z]{12,}/g, "AKIA[REDACTED]"],
+  [/(Bearer\s+)[A-Za-z0-9._~+/=-]{8,}/g, "$1[REDACTED]"],
+  [/\b([A-Z][A-Z0-9_]{2,}(?:KEY|TOKEN|SECRET|PASSWORD))=([^\s'"]+)/g, "$1=[REDACTED]"],
+];
+
+export function redactSecrets(text: string): string {
+  let out = text;
+  for (const [re, rep] of PATTERNS) out = out.replace(re, rep);
+  return out;
+}
+
+function redactDeep(value: unknown): unknown {
+  if (typeof value === "string") return redactSecrets(value);
+  if (Array.isArray(value)) return value.map(redactDeep);
+  if (value && typeof value === "object") {
+    return Object.fromEntries(
+      Object.entries(value as Record<string, unknown>).map(([k, v]) => [k, redactDeep(v)]),
+    );
+  }
+  return value;
+}
+
+function toText(value: unknown): string {
+  if (typeof value === "string") return value;
+  try {
+    return JSON.stringify(value) ?? "";
+  } catch {
+    return String(value);
+  }
+}
+
+function truncateTo(text: string, bytes: number): { text: string; truncated: boolean } {
+  const buf = Buffer.from(text, "utf8");
+  if (buf.length <= bytes) return { text, truncated: false };
+  return { text: buf.subarray(0, bytes).toString("utf8") + "…", truncated: true };
+}
+
+export function applyCapturePolicy(
+  input: { toolName: string; args: unknown; result: unknown },
+  state: { spentBytes: number },
+  budget: CaptureBudget = DEFAULT_CAPTURE_BUDGET,
+): {
+  args: unknown;
+  result?: string;
+  outputBytes: number;
+  truncated: boolean;
+  withheld: "not-allowlisted" | "budget" | null;
+} {
+  const args = redactDeep(input.args);
+  const raw = toText(input.result);
+  const outputBytes = Buffer.byteLength(raw, "utf8");
+  if (!OUTPUT_ALLOWLIST.has(input.toolName)) {
+    return { args, outputBytes, truncated: false, withheld: "not-allowlisted" };
+  }
+  if (state.spentBytes >= budget.perRunBytes) {
+    return { args, outputBytes, truncated: false, withheld: "budget" };
+  }
+  const { text, truncated } = truncateTo(redactSecrets(raw), budget.perStepBytes);
+  state.spentBytes += Buffer.byteLength(text, "utf8");
+  return { args, result: text, outputBytes, truncated, withheld: null };
+}