diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 1d5010679..8f190e44b 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/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; }; diff --git a/src/agents/pi-embedded-runner/run/params.ts b/src/agents/pi-embedded-runner/run/params.ts index ac5753bf6..2065868c5 100644 --- a/src/agents/pi-embedded-runner/run/params.ts +++ b/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; }; diff --git a/src/agents/pi-embedded-runner/run/types.ts b/src/agents/pi-embedded-runner/run/types.ts index ffef0fb8a..f36f88b61 100644 --- a/src/agents/pi-embedded-runner/run/types.ts +++ b/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; type ModelRegistry = ReturnType; @@ -58,6 +59,7 @@ export type EmbeddedRunAttemptParams = { thinkLevel: ThinkLevel; verboseLevel?: VerboseLevel; reasoningLevel?: ReasoningLevel; + secrets?: SecretScrubber; toolResultFormat?: ToolResultFormat; execOverrides?: Pick; bashElevated?: ExecElevatedDefaults; diff --git a/src/agents/pi-embedded-subscribe.handlers.messages.ts b/src/agents/pi-embedded-subscribe.handlers.messages.ts index 1f515e113..664c7c2f4 100644 --- a/src/agents/pi-embedded-subscribe.handlers.messages.ts +++ b/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, diff --git a/src/agents/pi-embedded-subscribe.handlers.types.ts b/src/agents/pi-embedded-subscribe.handlers.types.ts index 4a464c5e2..141a84e52 100644 --- a/src/agents/pi-embedded-subscribe.handlers.types.ts +++ b/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 = diff --git a/src/agents/pi-embedded-subscribe.scrubbing.test.ts b/src/agents/pi-embedded-subscribe.scrubbing.test.ts new file mode 100644 index 000000000..113077416 --- /dev/null +++ b/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: "My secret is reasoning-secret-123", + 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"); + }); +}); diff --git a/src/agents/pi-embedded-subscribe.ts b/src/agents/pi-embedded-subscribe.ts index a4a4b906a..857bab2e0 100644 --- a/src/agents/pi-embedded-subscribe.ts +++ b/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)); diff --git a/src/agents/pi-embedded-subscribe.types.ts b/src/agents/pi-embedded-subscribe.types.ts index 766ff7f18..e313ea408 100644 --- a/src/agents/pi-embedded-subscribe.types.ts +++ b/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; onAgentEvent?: (evt: { stream: string; data: Record }) => void | Promise; enforceFinalTag?: boolean; + secrets?: SecretScrubber; }; export type { BlockReplyChunking } from "./pi-embedded-block-chunker.js"; diff --git a/src/security/scrubber.test.ts b/src/security/scrubber.test.ts new file mode 100644 index 000000000..6b3929fb9 --- /dev/null +++ b/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); + }); + }); +}); diff --git a/src/security/scrubber.ts b/src/security/scrubber.ts new file mode 100644 index 000000000..e15b6b26e --- /dev/null +++ b/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(); + 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(); + 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); +}