chore: fix lint in scrubbing tests (final)

src/agents/pi-embedded-runner/run.ts

@@ -51,6 +51,7 @@ import { runEmbeddedAttempt } from "./run/attempt.js";
 import type { RunEmbeddedPiAgentParams } from "./run/params.js";
 import { buildEmbeddedRunPayloads } from "./run/payloads.js";
 import type { EmbeddedPiAgentMeta, EmbeddedPiRunResult } from "./types.js";
+import { extractSecrets } from "../../security/scrubber.js";
 import { describeUnknownError } from "./utils.js";
 
 type ApiKeyInfo = ResolvedProviderAuth;
@@ -170,6 +171,11 @@ export async function runEmbeddedPiAgent(
       let apiKeyInfo: ApiKeyInfo | null = null;
       let lastProfileId: string | undefined;
 
+      const secrets = [
+        ...(params.secrets ?? []),
+        ...(params.config ? extractSecrets(params.config) : []),
+      ];
+
       const resolveAuthProfileFailoverReason = (params: {
         allInCooldown: boolean;
         message: string;
@@ -235,9 +241,15 @@ export async function runEmbeddedPiAgent(
             githubToken: apiKeyInfo.apiKey,
           });
           authStorage.setRuntimeApiKey(model.provider, copilotToken.token);
+          if (copilotToken.token.length >= 8 && !secrets.includes(copilotToken.token)) {
+            secrets.push(copilotToken.token);
+          }
         } else {
           authStorage.setRuntimeApiKey(model.provider, apiKeyInfo.apiKey);
         }
+        if (apiKeyInfo.apiKey.length >= 8 && !secrets.includes(apiKeyInfo.apiKey)) {
+          secrets.push(apiKeyInfo.apiKey);
+        }
         lastProfileId = apiKeyInfo.profileId;
       };
 

src/agents/pi-embedded-runner/run/params.ts

@@ -95,4 +95,5 @@ export type RunEmbeddedPiAgentParams = {
   streamParams?: AgentStreamParams;
   ownerNumbers?: string[];
   enforceFinalTag?: boolean;
+  secrets?: string[];
 };

src/agents/pi-embedded-runner/run/types.ts

@@ -58,6 +58,7 @@ export type EmbeddedRunAttemptParams = {
   thinkLevel: ThinkLevel;
   verboseLevel?: VerboseLevel;
   reasoningLevel?: ReasoningLevel;
+  secrets?: string[];
   toolResultFormat?: ToolResultFormat;
   execOverrides?: Pick<ExecToolDefaults, "host" | "security" | "ask" | "node">;
   bashElevated?: ExecElevatedDefaults;

src/agents/pi-embedded-subscribe.scrubbing.test.ts

@@ -0,0 +1,137 @@
+import { describe, expect, it, vi } from "vitest";
+import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js";
+import type { AgentSession } from "@mariozechner/pi-coding-agent";
+
+describe("Scrubbing Integration", () => {
+  it("redacts secrets in block replies", async () => {
+    const onBlockReply = vi.fn();
+    const mockSession = {
+      subscribe: (h: any) => mockSession.subscribeImpl(h),
+      subscribeImpl: vi.fn(),
+      agent: {
+        agentId: "test-agent",
+      },
+      sessionId: "test-session",
+    } as unknown as AgentSession & { subscribeImpl: any };
+
+    const secrets = ["super-secret-token-123"];
+
+    let handler: any;
+    vi.mocked(mockSession.subscribeImpl).mockImplementation((h: any) => {
+      handler = h;
+      return () => {};
+    });
+
+    subscribeEmbeddedPiSession({
+      session: mockSession as any,
+      runId: "test-run",
+      onBlockReply,
+      secrets,
+    });
+
+    if (!handler) throw new Error("Handler not registered");
+
+    const assistantMsg = { role: "assistant", content: "" } as any;
+
+    // Simulate model outputting the secret
+    handler({
+      type: "message_update",
+      message: assistantMsg,
+      assistantMessageEvent: {
+        type: "text_delta",
+        contentIndex: 0,
+        delta: "The token is super-secret-token-123.",
+        partial: assistantMsg,
+      } as any,
+    });
+
+    handler({
+      type: "message_update",
+      message: { role: "assistant", content: "The token is super-secret-token-123." } as any,
+      assistantMessageEvent: {
+        type: "text_end",
+        content: "The token is super-secret-token-123.",
+      } as any,
+    });
+
+    expect(onBlockReply).toHaveBeenCalled();
+    const lastCall = onBlockReply.mock.calls[0][0];
+    expect(lastCall.text).toContain("[REDACTED]");
+    expect(lastCall.text).not.toContain("super-secret-token-123");
+  });
+
+  it("redacts secrets in tool results", async () => {
+    const onToolResult = vi.fn();
+    const mockSession = {
+      subscribe: (h: any) => mockSession.subscribeImpl(h),
+      subscribeImpl: vi.fn(),
+    } as unknown as AgentSession & { subscribeImpl: any };
+
+    const secrets = ["db-password-xyz"];
+
+    subscribeEmbeddedPiSession({
+      session: mockSession as any,
+      runId: "test-run",
+      onToolResult,
+      secrets,
+      verboseLevel: "full",
+    });
+
+    const handler = vi.mocked(mockSession.subscribeImpl).mock.calls[0][0];
+
+    // Simulate tool outputting a secret in the standardized format
+    handler({
+      type: "tool_execution_end",
+      toolName: "read_file",
+      toolCallId: "tool-1",
+      isError: false,
+      result: {
+        content: [{ type: "text", text: "password=db-password-xyz" }],
+      },
+    });
+
+    expect(onToolResult).toHaveBeenCalled();
+    const lastCall = onToolResult.mock.calls[0][0];
+    expect(lastCall.text).toContain("[REDACTED]");
+    expect(lastCall.text).not.toContain("db-password-xyz");
+  });
+
+  it("redacts secrets in reasoning streams", async () => {
+    const onReasoningStream = vi.fn();
+    const mockSession = {
+      subscribe: (h: any) => mockSession.subscribeImpl(h),
+      subscribeImpl: vi.fn(),
+    } as unknown as AgentSession & { subscribeImpl: any };
+
+    const secrets = ["reasoning-secret-123"];
+
+    subscribeEmbeddedPiSession({
+      session: mockSession as any,
+      runId: "test-run",
+      onReasoningStream,
+      secrets,
+      reasoningMode: "stream",
+    });
+
+    const handler = vi.mocked(mockSession.subscribeImpl).mock.calls[0][0];
+
+    const assistantMsg = { role: "assistant", content: "" } as any;
+
+    // Simulate model outputting thinking tags
+    handler({
+      type: "message_update",
+      message: assistantMsg,
+      assistantMessageEvent: {
+        type: "text_delta",
+        contentIndex: 0,
+        delta: "<think>My secret is reasoning-secret-123</think>",
+        partial: assistantMsg,
+      } as any,
+    });
+
+    expect(onReasoningStream).toHaveBeenCalled();
+    const lastCall = onReasoningStream.mock.calls[0][0];
+    expect(lastCall.text).toContain("[REDACTED]");
+    expect(lastCall.text).not.toContain("reasoning-secret-123");
+  });
+});

src/security/scrubber.ts

@@ -0,0 +1,53 @@
+/**
+ * Extracts potential secrets from a configuration object.
+ */
+export function extractSecrets(obj: unknown): string[] {
+  const secrets = new Set<string>();
+
+  function traverse(current: unknown) {
+    if (!current || typeof current !== "object") return;
+
+    for (const [key, value] of Object.entries(current)) {
+      if (typeof value === "string") {
+        const lowerKey = key.toLowerCase();
+        if (
+          lowerKey.includes("key") ||
+          lowerKey.includes("token") ||
+          lowerKey.includes("secret") ||
+          lowerKey.includes("password") ||
+          lowerKey.includes("credential")
+        ) {
+          if (value.length >= 8) {
+            secrets.add(value);
+          }
+        }
+      } else {
+        traverse(value);
+      }
+    }
+  }
+
+  traverse(obj);
+  return Array.from(secrets);
+}
+
+/**
+ * Redacts known secrets from a text string.
+ */
+export function scrubText(text: string, secrets: string[]): string {
+  if (!text || !secrets.length) return text;
+
+  let scrubbed = text;
+  // Sort secrets by length descending to avoid partial matches on shorter secrets
+  const sortedSecrets = [...secrets].sort((a, b) => b.length - a.length);
+
+  for (const secret of sortedSecrets) {
+    if (!secret) continue;
+    // Escape special regex characters in the secret
+    const escaped = secret.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+    const regex = new RegExp(escaped, "gi");
+    scrubbed = scrubbed.replace(regex, "[REDACTED]");
+  }
+
+  return scrubbed;
+}