smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(coding-agent): preserved tool schemas while redacting

Browse Source 
Converted provider-facing tool parameters to wire JSON Schema before redaction so live Zod instances are not deep-cloned into plain objects.

Fixes #2146
This commit is contained in:
roboomp
2026-06-10 08:26:02 +02:00
committed by can1357
parent 2a90108893
commit aa4cd0ab2d
6 changed files with 83 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/agent/src/compaction/compaction.ts
+3 -2
View File
@@ -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<any>[];
tools?: Tool[];
customInstructions?: string;
convertToLlm?: ConvertToLlm;
initiatorOverride?: MessageAttribution;
packages/coding-agent/src/sdk.ts
+2 -6
View File
@@ -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,
});
packages/coding-agent/src/secrets/index.ts
+8 -1
View File
@@ -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.
packages/coding-agent/src/secrets/obfuscator.ts
+33 -2
View File
@@ -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<T>(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<string, unknown> = {};
for (const key of Object.keys(obj)) {
@@ -265,3 +291,8 @@ function deepWalkStrings<T>(obj: T, transform: (s: string) => string): T {
}
return obj;
}
function isPlainRecord(obj: object): obj is Record<string, unknown> {
const prototype = Object.getPrototypeOf(obj);
return prototype === Object.prototype || prototype === null;
}
packages/coding-agent/src/session/agent-session.ts
+2 -2
View File
@@ -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",
packages/coding-agent/test/secrets-obfuscator.test.ts
+35 -4
View File
@@ -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", () => {