Merge bcf21fbeaa into da71eaebd2

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 { SecretScrubber } 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 ?? new SecretScrubber();
+      if (params.config) {
+        secrets.extractFromConfig(params.config);
+      }
+
       const resolveAuthProfileFailoverReason = (params: {
         allInCooldown: boolean;
         message: string;
@@ -235,9 +241,11 @@ export async function runEmbeddedPiAgent(
             githubToken: apiKeyInfo.apiKey,
           });
           authStorage.setRuntimeApiKey(model.provider, copilotToken.token);
+          secrets.add(copilotToken.token);
         } else {
           authStorage.setRuntimeApiKey(model.provider, apiKeyInfo.apiKey);
         }
+        secrets.add(apiKeyInfo.apiKey);
         lastProfileId = apiKeyInfo.profileId;
       };
 

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

@@ -5,6 +5,7 @@ import type { AgentStreamParams } from "../../../commands/agent/types.js";
 import type { enqueueCommand } from "../../../process/command-queue.js";
 import type { ExecElevatedDefaults, ExecToolDefaults } from "../../bash-tools.js";
 import type { BlockReplyChunking, ToolResultFormat } from "../../pi-embedded-subscribe.js";
+import type { SecretScrubber } from "../../../security/scrubber.js";
 import type { SkillSnapshot } from "../../skills.js";
 
 // Simplified tool definition for client-provided tools (OpenResponses hosted tools)
@@ -95,4 +96,5 @@ export type RunEmbeddedPiAgentParams = {
   streamParams?: AgentStreamParams;
   ownerNumbers?: string[];
   enforceFinalTag?: boolean;
+  secrets?: SecretScrubber;
 };

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

@@ -11,6 +11,7 @@ import type { BlockReplyChunking, ToolResultFormat } from "../../pi-embedded-sub
 import type { SkillSnapshot } from "../../skills.js";
 import type { SessionSystemPromptReport } from "../../../config/sessions/types.js";
 import type { ClientToolDefinition } from "./params.js";
+import type { SecretScrubber } from "../../../security/scrubber.js";
 
 type AuthStorage = ReturnType<typeof discoverAuthStorage>;
 type ModelRegistry = ReturnType<typeof discoverModels>;
@@ -58,6 +59,7 @@ export type EmbeddedRunAttemptParams = {
   thinkLevel: ThinkLevel;
   verboseLevel?: VerboseLevel;
   reasoningLevel?: ReasoningLevel;
+  secrets?: SecretScrubber;
   toolResultFormat?: ToolResultFormat;
   execOverrides?: Pick<ExecToolDefaults, "host" | "security" | "ask" | "node">;
   bashElevated?: ExecElevatedDefaults;

src/agents/pi-embedded-subscribe.handlers.messages.ts

@@ -118,22 +118,22 @@ export function handleMessageUpdate(
         runId: ctx.params.runId,
         stream: "assistant",
         data: {
-          text: cleanedText,
-          delta: deltaText,
+          text: ctx.scrub(cleanedText),
+          delta: ctx.scrub(deltaText),
           mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
         },
       });
       void ctx.params.onAgentEvent?.({
         stream: "assistant",
         data: {
-          text: cleanedText,
-          delta: deltaText,
+          text: ctx.scrub(cleanedText),
+          delta: ctx.scrub(deltaText),
           mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
         },
       });
       if (ctx.params.onPartialReply && ctx.state.shouldEmitPartialReplies) {
         void ctx.params.onPartialReply({
-          text: cleanedText,
+          text: ctx.scrub(cleanedText),
           mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
         });
       }
@@ -198,7 +198,7 @@ export function handleMessageEnd(
   const maybeEmitReasoning = () => {
     if (!shouldEmitReasoning || !formattedReasoning) return;
     ctx.state.lastReasoningSent = formattedReasoning;
-    void onBlockReply?.({ text: formattedReasoning });
+    void onBlockReply?.({ text: ctx.scrub(formattedReasoning) });
   };
 
   if (shouldEmitReasoningBeforeAnswer) maybeEmitReasoning();
@@ -239,7 +239,7 @@ export function handleMessageEnd(
           // Emit if there's content OR audioAsVoice flag (to propagate the flag).
           if (cleanedText || (mediaUrls && mediaUrls.length > 0) || audioAsVoice) {
             void onBlockReply({
-              text: cleanedText,
+              text: ctx.scrub(cleanedText),
               mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
               audioAsVoice,
               replyToId,
@@ -270,7 +270,7 @@ export function handleMessageEnd(
       } = tailResult;
       if (cleanedText || (mediaUrls && mediaUrls.length > 0) || audioAsVoice) {
         void onBlockReply({
-          text: cleanedText,
+          text: ctx.scrub(cleanedText),
           mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
           audioAsVoice,
           replyToId,

src/agents/pi-embedded-subscribe.handlers.types.ts

@@ -94,6 +94,7 @@ export type EmbeddedPiSubscribeContext = {
   noteCompactionRetry: () => void;
   resolveCompactionRetry: () => void;
   maybeResolveCompactionWait: () => void;
+  scrub: (text: string) => string;
 };
 
 export type EmbeddedPiSubscribeEvent =

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

@@ -0,0 +1,138 @@
+import { describe, expect, it, vi } from "vitest";
+import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js";
+import type { AgentSession } from "@mariozechner/pi-coding-agent";
+import { SecretScrubber } from "../security/scrubber.js";
+
+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 = new SecretScrubber(["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("your-key-here");
+    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 = new SecretScrubber(["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("your-key-here");
+    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 = new SecretScrubber(["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("your-key-here");
+    expect(lastCall.text).not.toContain("reasoning-secret-123");
+  });
+});

src/agents/pi-embedded-subscribe.ts

@@ -229,7 +229,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     if (!cleanedText && (!mediaUrls || mediaUrls.length === 0)) return;
     try {
       void params.onToolResult({
-        text: cleanedText,
+        text: params.secrets ? params.secrets.scrub(cleanedText) : cleanedText,
         mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
       });
     } catch {
@@ -246,7 +246,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     if (!cleanedText && (!mediaUrls || mediaUrls.length === 0)) return;
     try {
       void params.onToolResult({
-        text: cleanedText,
+        text: params.secrets ? params.secrets.scrub(cleanedText) : cleanedText,
         mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
       });
     } catch {
@@ -390,7 +390,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     // Skip empty payloads, but always emit if audioAsVoice is set (to propagate the flag)
     if (!cleanedText && (!mediaUrls || mediaUrls.length === 0) && !audioAsVoice) return;
     void params.onBlockReply({
-      text: cleanedText,
+      text: ctx.scrub(cleanedText),
       mediaUrls: mediaUrls?.length ? mediaUrls : undefined,
       audioAsVoice,
       replyToId,
@@ -422,7 +422,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     if (formatted === state.lastStreamedReasoning) return;
     state.lastStreamedReasoning = formatted;
     void params.onReasoningStream({
-      text: formatted,
+      text: ctx.scrub(formatted),
     });
   };
 
@@ -463,6 +463,7 @@ export function subscribeEmbeddedPiSession(params: SubscribeEmbeddedPiSessionPar
     noteCompactionRetry,
     resolveCompactionRetry,
     maybeResolveCompactionWait,
+    scrub: (text: string) => (params.secrets ? params.secrets.scrub(text) : text),
   };
 
   const unsubscribe = params.session.subscribe(createEmbeddedPiSessionEventHandler(ctx));

src/agents/pi-embedded-subscribe.types.ts

@@ -1,5 +1,6 @@
 import type { AgentSession } from "@mariozechner/pi-coding-agent";
 
+import type { SecretScrubber } from "../security/scrubber.js";
 import type { ReasoningLevel, VerboseLevel } from "../auto-reply/thinking.js";
 import type { BlockReplyChunking } from "./pi-embedded-block-chunker.js";
 
@@ -31,6 +32,7 @@ export type SubscribeEmbeddedPiSessionParams = {
   onAssistantMessageStart?: () => void | Promise<void>;
   onAgentEvent?: (evt: { stream: string; data: Record<string, unknown> }) => void | Promise<void>;
   enforceFinalTag?: boolean;
+  secrets?: SecretScrubber;
 };
 
 export type { BlockReplyChunking } from "./pi-embedded-block-chunker.js";

src/security/scrubber.test.ts

@@ -0,0 +1,74 @@
+import { describe, expect, it } from "vitest";
+import { SecretScrubber } from "./scrubber.js";
+
+describe("SecretScrubber", () => {
+  describe("Heuristics", () => {
+    it("identifies secrets with suffix matching", () => {
+      const scrubber = new SecretScrubber();
+      scrubber.extractFromConfig({
+        my_api_key: "sensitive-key-123",
+        db_password: "sensitive-password-123",
+        keyboard_layout: "not-sensitive-123", // false positive if using includes
+        monkey_patch: "not-sensitive-456", // false positive if using includes
+      });
+
+      const text =
+        "Key: sensitive-key-123, Pass: sensitive-password-123, Layout: not-sensitive-123";
+      const scrubbed = scrubber.scrub(text);
+
+      expect(scrubbed).toContain("Key: your-key-here");
+      expect(scrubbed).toContain("Pass: your-key-here");
+      expect(scrubbed).toContain("Layout: not-sensitive-123");
+    });
+
+    it("identifies common camelCase sensitive keys", () => {
+      const scrubber = new SecretScrubber();
+      scrubber.extractFromConfig({
+        apiKey: "sensitive-api-key",
+        authToken: "sensitive-auth-token",
+      });
+
+      const text = "api: sensitive-api-key, token: sensitive-auth-token";
+      const scrubbed = scrubber.scrub(text);
+
+      expect(scrubbed).toBe("api: your-key-here, token: your-key-here");
+    });
+  });
+
+  describe("Stability", () => {
+    it("handles circular references gracefully", () => {
+      const scrubber = new SecretScrubber();
+      const config: any = { api_key: "secret-circular" };
+      config.self = config;
+
+      expect(() => scrubber.extractFromConfig(config)).not.toThrow();
+      expect(scrubber.scrub("secret-circular")).toBe("your-key-here");
+    });
+
+    it("escapes regex special characters in secrets", () => {
+      const scrubber = new SecretScrubber(["$ecret.with*chars+"]);
+      const text = "My secret is $ecret.with*chars+ now.";
+      expect(scrubber.scrub(text)).toBe("My secret is your-key-here now.");
+    });
+
+    it("handles overlapping secrets by redacting longer ones first", () => {
+      const scrubber = new SecretScrubber(["password", "password123"]);
+      const text = "My pass is password123";
+      // If "password" matched first, result would be "your-key-here123"
+      expect(scrubber.scrub(text)).toBe("My pass is your-key-here");
+    });
+  });
+
+  describe("Performance & Limits", () => {
+    it("ignores secrets that are too short", () => {
+      const scrubber = new SecretScrubber(["short"]);
+      expect(scrubber.scrub("short")).toBe("short");
+    });
+
+    it("ignores secrets that are too long", () => {
+      const longSecret = "a".repeat(5000);
+      const scrubber = new SecretScrubber([longSecret]);
+      expect(scrubber.scrub(longSecret)).toBe(longSecret);
+    });
+  });
+});

src/security/scrubber.ts

@@ -0,0 +1,126 @@
+/**
+ * Heuristics for identifying sensitive keys in configuration.
+ * Using suffix matching to avoid false positives like 'keyboard' or 'monkey'.
+ */
+const SENSITIVE_KEY_SUFFIXES = [
+  "key",
+  "token",
+  "secret",
+  "password",
+  "credential",
+  "passphrase",
+  "webhook",
+];
+
+const MIN_SECRET_LENGTH = 8;
+const MAX_SECRET_LENGTH = 4096;
+
+export class SecretScrubber {
+  private secrets = new Set<string>();
+  private regex: RegExp | null = null;
+
+  constructor(initialSecrets: string[] = []) {
+    this.add(initialSecrets);
+  }
+
+  /**
+   * Adds new secrets to the scrubber.
+   */
+  add(secret: string | string[]): void {
+    const toAdd = Array.isArray(secret) ? secret : [secret];
+    let changed = false;
+
+    for (const s of toAdd) {
+      if (s && s.length >= MIN_SECRET_LENGTH && s.length <= MAX_SECRET_LENGTH) {
+        if (!this.secrets.has(s)) {
+          this.secrets.add(s);
+          changed = true;
+        }
+      }
+    }
+
+    if (changed) {
+      this.rebuildRegex();
+    }
+  }
+
+  /**
+   * Extracts secrets from a configuration object using heuristics.
+   */
+  extractFromConfig(config: unknown): void {
+    const visited = new Set<unknown>();
+    const found: string[] = [];
+
+    const traverse = (current: unknown) => {
+      if (!current || typeof current !== "object" || visited.has(current)) {
+        return;
+      }
+      visited.add(current);
+
+      for (const [key, value] of Object.entries(current)) {
+        if (typeof value === "string") {
+          const lowerKey = key.toLowerCase();
+          const isSensitive = SENSITIVE_KEY_SUFFIXES.some(
+            (suffix) =>
+              lowerKey === suffix || lowerKey.endsWith(`_${suffix}`) || lowerKey.endsWith(suffix),
+          );
+
+          // Additional check for common camelCase/PascalCase sensitive keys without underscore
+          const isCommonSensitive = /api[Kk]ey|authToken|botToken/.test(key);
+
+          if (isSensitive || isCommonSensitive) {
+            found.push(value);
+          }
+        } else if (value && typeof value === "object") {
+          traverse(value);
+        }
+      }
+    };
+
+    traverse(config);
+    this.add(found);
+  }
+
+  /**
+   * Redacts all known secrets from the given text.
+   */
+  scrub(text: string): string {
+    if (!text || !this.regex || this.secrets.size === 0) {
+      return text;
+    }
+    return text.replace(this.regex, "your-key-here");
+  }
+
+  private rebuildRegex(): void {
+    if (this.secrets.size === 0) {
+      this.regex = null;
+      return;
+    }
+
+    // Sort secrets by length descending to ensure longer secrets match before their substrings
+    const sorted = Array.from(this.secrets).sort((a, b) => b.length - a.length);
+    const escaped = sorted.map((s) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"));
+    this.regex = new RegExp(escaped.join("|"), "gi");
+  }
+}
+
+/**
+ * Convenience function for one-off extraction and scrubbing.
+ * @deprecated Use SecretScrubber class for production code to avoid repeated regex builds.
+ */
+export function extractSecrets(obj: unknown): string[] {
+  const scrubber = new SecretScrubber();
+  scrubber.extractFromConfig(obj);
+  // We return the internals for compatibility with existing tests/code if needed,
+  // but the class is preferred.
+  return (scrubber as any).secrets ? Array.from((scrubber as any).secrets) : [];
+}
+
+/**
+ * Convenience function for one-off scrubbing.
+ * @deprecated Use SecretScrubber class for production code.
+ */
+export function scrubText(text: string, secrets: string[]): string {
+  const scrubber = new SecretScrubber(secrets);
+  return scrubber.scrub(text);
+}