diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 881e34760..53d0ad464 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -12,6 +12,7 @@ import { type Message, type MessageAttribution, type Model, + type Tool, type Usage, } from "@oh-my-pi/pi-ai"; import { clampThinkingLevelForModel } from "@oh-my-pi/pi-catalog/model-thinking"; @@ -19,7 +20,7 @@ import { countTokens } from "@oh-my-pi/pi-natives"; import { logger, prompt } from "@oh-my-pi/pi-utils"; import { type AgentTelemetry, instrumentedCompleteSimple } from "../telemetry"; import { ThinkingLevel } from "../thinking"; -import type { AgentMessage, AgentTool } from "../types"; +import type { AgentMessage } from "../types"; import type { CompactionEntry, SessionEntry } from "./entries"; import { type ConvertToLlm, convertToLlm, createBranchSummaryMessage, createCustomMessage } from "./messages"; import { @@ -690,7 +691,7 @@ export interface HandoffOptions { /** Live agent system prompt — passed verbatim so providers hit the cached prefix. */ systemPrompt: string[]; /** Live agent tool list — same purpose. Forced to `toolChoice: "none"`. */ - tools?: AgentTool[]; + tools?: Tool[]; customInstructions?: string; convertToLlm?: ConvertToLlm; initiatorOverride?: MessageAttribution; diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index b46c01a61..3f9855c9b 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -9,7 +9,6 @@ import { type ThinkingLevel, } from "@oh-my-pi/pi-agent-core"; import { - type Context, type CredentialDisabledEvent, type Message, type Model, @@ -100,6 +99,7 @@ import { deobfuscateSessionContext, loadSecrets, obfuscateMessages, + obfuscateProviderContext, SecretObfuscator, } from "./secrets"; import { AgentSession } from "./session/agent-session"; @@ -2139,10 +2139,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} if (!obfuscator?.hasSecrets()) return converted; return obfuscateMessages(obfuscator, converted); }; - const obfuscateProviderContext = (context: Context): Context => { - if (!obfuscator?.hasSecrets()) return context; - return obfuscator.obfuscateObject(context); - }; const transformContext = async (messages: AgentMessage[], _signal?: AbortSignal) => { const withContext = await extensionRunner.emitContext(messages); @@ -2221,7 +2217,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const openrouterRoutingPreset = settings.get("providers.openrouterVariant"); const openrouterVariant = openrouterRoutingPreset && openrouterRoutingPreset !== "default" ? openrouterRoutingPreset : undefined; - return streamSimple(streamModel, obfuscateProviderContext(context), { + return streamSimple(streamModel, obfuscator ? obfuscateProviderContext(obfuscator, context) : context, { ...streamOptions, openrouterVariant: streamOptions?.openrouterVariant ?? openrouterVariant, }); diff --git a/packages/coding-agent/src/secrets/index.ts b/packages/coding-agent/src/secrets/index.ts index 420150c25..550557658 100644 --- a/packages/coding-agent/src/secrets/index.ts +++ b/packages/coding-agent/src/secrets/index.ts @@ -4,7 +4,14 @@ import { YAML } from "bun"; import type { SecretEntry } from "./obfuscator"; import { compileSecretRegex } from "./regex"; -export { deobfuscateSessionContext, obfuscateMessages, type SecretEntry, SecretObfuscator } from "./obfuscator"; +export { + deobfuscateSessionContext, + obfuscateMessages, + obfuscateProviderContext, + obfuscateProviderTools, + type SecretEntry, + SecretObfuscator, +} from "./obfuscator"; /** * Load secrets from project-local and global secrets.yml files. diff --git a/packages/coding-agent/src/secrets/obfuscator.ts b/packages/coding-agent/src/secrets/obfuscator.ts index 72c7e8425..ef0845f79 100644 --- a/packages/coding-agent/src/secrets/obfuscator.ts +++ b/packages/coding-agent/src/secrets/obfuscator.ts @@ -1,4 +1,5 @@ -import type { Message } from "@oh-my-pi/pi-ai"; +import type { Context, Message, Tool } from "@oh-my-pi/pi-ai"; +import { toolWireSchema } from "@oh-my-pi/pi-ai/utils/schema"; import type { SessionContext } from "../session/session-manager"; import { compileSecretRegex } from "./regex"; @@ -222,6 +223,31 @@ export function obfuscateMessages(obfuscator: SecretObfuscator, messages: Messag return obfuscator.obfuscateObject(messages); } +/** Obfuscate provider request context without walking live tool schema instances. */ +export function obfuscateProviderContext(obfuscator: SecretObfuscator | undefined, context: Context): Context { + if (!obfuscator?.hasSecrets()) return context; + return { + ...context, + systemPrompt: obfuscator.obfuscateObject(context.systemPrompt), + messages: obfuscator.obfuscateObject(context.messages), + tools: obfuscateProviderTools(obfuscator, context.tools), + }; +} + +/** Convert tool schemas to wire JSON Schema before obfuscating provider-visible strings. */ +export function obfuscateProviderTools( + obfuscator: SecretObfuscator | undefined, + tools: Tool[] | undefined, +): Tool[] | undefined { + if (!tools || !obfuscator?.hasSecrets()) return tools; + return tools.map(tool => ({ + ...tool, + description: obfuscator.obfuscate(tool.description), + parameters: obfuscator.obfuscateObject(toolWireSchema(tool)), + customFormat: tool.customFormat ? obfuscator.obfuscateObject(tool.customFormat) : undefined, + })); +} + // ═══════════════════════════════════════════════════════════════════════════ // Helpers // ═══════════════════════════════════════════════════════════════════════════ @@ -252,7 +278,7 @@ function deepWalkStrings(obj: T, transform: (s: string) => string): T { }); return (changed ? result : obj) as unknown as T; } - if (obj !== null && typeof obj === "object") { + if (obj !== null && typeof obj === "object" && isPlainRecord(obj)) { let changed = false; const result: Record = {}; for (const key of Object.keys(obj)) { @@ -265,3 +291,8 @@ function deepWalkStrings(obj: T, transform: (s: string) => string): T { } return obj; } + +function isPlainRecord(obj: object): obj is Record { + const prototype = Object.getPrototypeOf(obj); + return prototype === Object.prototype || prototype === null; +} diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 7f9be3cb0..73c7c9069 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -184,7 +184,7 @@ import planModeToolDecisionReminderPrompt from "../prompts/system/plan-mode-tool import ttsrInterruptTemplate from "../prompts/system/ttsr-interrupt.md" with { type: "text" }; import ttsrToolReminderTemplate from "../prompts/system/ttsr-tool-reminder.md" with { type: "text" }; import { type AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry"; -import { deobfuscateSessionContext, type SecretObfuscator } from "../secrets/obfuscator"; +import { deobfuscateSessionContext, obfuscateProviderTools, type SecretObfuscator } from "../secrets/obfuscator"; import { invalidateHostMetadata } from "../ssh/connection-manager"; import { AUTO_THINKING, @@ -6383,7 +6383,7 @@ export class AgentSession { apiKey, { systemPrompt: this.#obfuscateForProvider(this.#baseSystemPrompt), - tools: this.#obfuscateForProvider(this.agent.state.tools), + tools: obfuscateProviderTools(this.#obfuscator, this.agent.state.tools), customInstructions, convertToLlm: messages => this.#convertToLlmForSideRequest(messages), initiatorOverride: "agent", diff --git a/packages/coding-agent/test/secrets-obfuscator.test.ts b/packages/coding-agent/test/secrets-obfuscator.test.ts index 96d07c760..48cf89e4e 100644 --- a/packages/coding-agent/test/secrets-obfuscator.test.ts +++ b/packages/coding-agent/test/secrets-obfuscator.test.ts @@ -3,9 +3,14 @@ */ import { describe, expect, it } from "bun:test"; -import type { Message } from "@oh-my-pi/pi-ai"; -import { obfuscateMessages, SecretObfuscator } from "@oh-my-pi/pi-coding-agent/secrets/obfuscator"; +import type { Context, Message } from "@oh-my-pi/pi-ai"; +import { + obfuscateMessages, + obfuscateProviderContext, + SecretObfuscator, +} from "@oh-my-pi/pi-coding-agent/secrets/obfuscator"; import { compileSecretRegex } from "@oh-my-pi/pi-coding-agent/secrets/regex"; +import { z } from "zod"; describe("compileSecretRegex", () => { it("adds global flag when not provided", () => { @@ -65,6 +70,7 @@ describe("SecretObfuscator regex behavior", () => { const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]); const payload = { systemPrompt: [`workspace contains ${secret}`], + messages: [], tools: [ { name: "handoff", @@ -77,11 +83,36 @@ describe("SecretObfuscator regex behavior", () => { ], }; - const obfuscated = obfuscator.obfuscateObject(payload); + const obfuscated = obfuscateProviderContext(obfuscator, payload); const serialized = JSON.stringify(obfuscated); expect(serialized).not.toContain(secret); - expect(obfuscator.deobfuscateObject(obfuscated)).toEqual(payload); + expect(obfuscator.deobfuscateObject(obfuscated).tools?.[0]?.description).toEqual(payload.tools[0]?.description); + }); + + it("redacts Zod tool schemas without cloning the live schema instance", () => { + const secret = "SUPER_SECRET_TOKEN_12345"; + const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]); + const parameters = z.object({ + note: z.string().describe(`write ${secret}`), + }); + const context: Context = { + messages: [], + tools: [ + { + name: "extension_tool", + description: `preserve ${secret}`, + parameters, + }, + ], + }; + + const obfuscated = obfuscateProviderContext(obfuscator, context); + + expect(obfuscator.obfuscateObject(parameters)).toBe(parameters); + expect(context.tools?.[0]?.parameters).toBe(parameters); + expect(obfuscated.tools?.[0]?.parameters).not.toBe(parameters); + expect(JSON.stringify(obfuscated)).not.toContain(secret); }); it("obfuscates system reminders and assistant tool calls in messages", () => {