feat: complete refactor of secret scrubber addressing review findings

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

@@ -51,7 +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 { SecretScrubber } from "../../security/scrubber.js";
 import { describeUnknownError } from "./utils.js";
 
 type ApiKeyInfo = ResolvedProviderAuth;
@@ -171,10 +171,10 @@ export async function runEmbeddedPiAgent(
       let apiKeyInfo: ApiKeyInfo | null = null;
       let lastProfileId: string | undefined;
 
-      const secrets = [
-        ...(params.secrets ?? []),
-        ...(params.config ? extractSecrets(params.config) : []),
-      ];
+      const secrets = params.secrets ?? new SecretScrubber();
+      if (params.config) {
+        secrets.extractFromConfig(params.config);
+      }
 
       const resolveAuthProfileFailoverReason = (params: {
         allInCooldown: boolean;
@@ -241,15 +241,11 @@ 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);
-          }
+          secrets.add(copilotToken.token);
         } else {
           authStorage.setRuntimeApiKey(model.provider, apiKeyInfo.apiKey);
         }
-        if (apiKeyInfo.apiKey.length >= 8 && !secrets.includes(apiKeyInfo.apiKey)) {
-          secrets.push(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,5 +96,5 @@ export type RunEmbeddedPiAgentParams = {
   streamParams?: AgentStreamParams;
   ownerNumbers?: string[];
   enforceFinalTag?: boolean;
-  secrets?: string[];
+  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,7 +59,7 @@ export type EmbeddedRunAttemptParams = {
   thinkLevel: ThinkLevel;
   verboseLevel?: VerboseLevel;
   reasoningLevel?: ReasoningLevel;
-  secrets?: string[];
+  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

@@ -1,6 +1,7 @@
 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 () => {
@@ -14,7 +15,7 @@ describe("Scrubbing Integration", () => {
       sessionId: "test-session",
     } as unknown as AgentSession & { subscribeImpl: any };
 
-    const secrets = ["super-secret-token-123"];
+    const secrets = new SecretScrubber(["super-secret-token-123"]);
 
     let handler: any;
     vi.mocked(mockSession.subscribeImpl).mockImplementation((h: any) => {
@@ -67,7 +68,7 @@ describe("Scrubbing Integration", () => {
       subscribeImpl: vi.fn(),
     } as unknown as AgentSession & { subscribeImpl: any };
 
-    const secrets = ["db-password-xyz"];
+    const secrets = new SecretScrubber(["db-password-xyz"]);
 
     subscribeEmbeddedPiSession({
       session: mockSession as any,
@@ -103,7 +104,7 @@ describe("Scrubbing Integration", () => {
       subscribeImpl: vi.fn(),
     } as unknown as AgentSession & { subscribeImpl: any };
 
-    const secrets = ["reasoning-secret-123"];
+    const secrets = new SecretScrubber(["reasoning-secret-123"]);
 
     subscribeEmbeddedPiSession({
       session: mockSession as any,

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: [REDACTED]");
+      expect(scrubbed).toContain("Pass: [REDACTED]");
+      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: [REDACTED], token: [REDACTED]");
+    });
+  });
+
+  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("[REDACTED]");
+    });
+
+    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 [REDACTED] 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 "[REDACTED]123"
+      expect(scrubber.scrub(text)).toBe("My pass is [REDACTED]");
+    });
+  });
+
+  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

@@ -1,53 +1,126 @@
 /**
- * Extracts potential secrets from a configuration object.
+ * Heuristics for identifying sensitive keys in configuration.
+ * Using suffix matching to avoid false positives like 'keyboard' or 'monkey'.
  */
-export function extractSecrets(obj: unknown): string[] {
-  const secrets = new Set<string>();
+const SENSITIVE_KEY_SUFFIXES = [
+  "key",
+  "token",
+  "secret",
+  "password",
+  "credential",
+  "passphrase",
+  "webhook",
+];
 
-  function traverse(current: unknown) {
-    if (!current || typeof current !== "object") return;
+const MIN_SECRET_LENGTH = 8;
+const MAX_SECRET_LENGTH = 4096;
 
-    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);
-          }
+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;
         }
-      } else {
-        traverse(value);
       }
     }
+
+    if (changed) {
+      this.rebuildRegex();
+    }
   }
 
-  traverse(obj);
-  return Array.from(secrets);
+  /**
+   * 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, "[REDACTED]");
+  }
+
+  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");
+  }
 }
 
 /**
- * Redacts known secrets from a text string.
+ * 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 {
-  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;
+  const scrubber = new SecretScrubber(secrets);
+  return scrubber.scrub(text);
 }