From d28c6a17fbc9bb4920f0ff4d96749d4d188fbb29 Mon Sep 17 00:00:00 2001 From: Thiago Butignon Date: Wed, 28 Jan 2026 18:38:34 -0300 Subject: [PATCH] chore: fix lint in scrubbing tests (final) --- src/agents/pi-embedded-runner/run.ts | 12 ++ src/agents/pi-embedded-runner/run/params.ts | 1 + src/agents/pi-embedded-runner/run/types.ts | 1 + .../pi-embedded-subscribe.scrubbing.test.ts | 137 ++++++++++++++++++ src/security/scrubber.ts | 53 +++++++ 5 files changed, 204 insertions(+) create mode 100644 src/agents/pi-embedded-subscribe.scrubbing.test.ts create mode 100644 src/security/scrubber.ts diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 870453f38..007703677 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 { 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; }; diff --git a/src/agents/pi-embedded-runner/run/params.ts b/src/agents/pi-embedded-runner/run/params.ts index b21a5e3fc..12f684a77 100644 --- a/src/agents/pi-embedded-runner/run/params.ts +++ b/src/agents/pi-embedded-runner/run/params.ts @@ -95,4 +95,5 @@ export type RunEmbeddedPiAgentParams = { streamParams?: AgentStreamParams; ownerNumbers?: string[]; enforceFinalTag?: boolean; + secrets?: string[]; }; diff --git a/src/agents/pi-embedded-runner/run/types.ts b/src/agents/pi-embedded-runner/run/types.ts index 92bb3ff46..ec174f386 100644 --- a/src/agents/pi-embedded-runner/run/types.ts +++ b/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; bashElevated?: ExecElevatedDefaults; 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..31779aea0 --- /dev/null +++ b/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: "My secret is reasoning-secret-123", + 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"); + }); +}); diff --git a/src/security/scrubber.ts b/src/security/scrubber.ts new file mode 100644 index 000000000..cf1b6bbd9 --- /dev/null +++ b/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(); + + 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; +}