diff --git a/packages/ai/src/providers/transform-messages.ts b/packages/ai/src/providers/transform-messages.ts index 0751c4ece..4b85aa45d 100644 --- a/packages/ai/src/providers/transform-messages.ts +++ b/packages/ai/src/providers/transform-messages.ts @@ -1,5 +1,14 @@ import { renderDemotedThinking } from "../dialect/demotion"; -import type { Api, AssistantMessage, Message, Model, ToolCall, ToolResultMessage, UserMessage } from "../types"; +import type { + Api, + AssistantMessage, + DeveloperMessage, + Message, + Model, + ToolCall, + ToolResultMessage, + UserMessage, +} from "../types"; import { isDemotedThinking, kDemotedThinking } from "../utils/block-symbols"; const enum ToolCallStatus { @@ -286,6 +295,122 @@ function normalizeAnthropicTargetToolCallId( * - Preserves tool call structure (unlike converting to text summaries) * - Injects synthetic "aborted" tool results */ +const SENSITIVE_TOKEN_RE = + /\b(gh[opusr]_[a-zA-Z0-9_*]{36,}|github_pat_[a-zA-Z0-9_*]{36,}|glpat-[a-zA-Z0-9_*-]{20,}|sk-proj-[a-zA-Z0-9_*]{36,}|sk-ant-[a-zA-Z0-9_*]{36,}|sk-[a-zA-Z0-9_*]{48,})(?![a-zA-Z0-9_*])/g; + +export function redactSensitiveCredentials(text: string): string { + return text.replace(SENSITIVE_TOKEN_RE, (_match, token) => { + if (token.startsWith("gh")) { + return "[github_token_redacted]"; + } + if (token.startsWith("gl")) { + return "[gitlab_token_redacted]"; + } + if (token.startsWith("sk-ant-")) { + return "[anthropic_token_redacted]"; + } + if (token.startsWith("sk")) { + return "[openai_token_redacted]"; + } + return "[token_redacted]"; + }); +} + +function redactSensitiveInObject(val: unknown): unknown { + if (typeof val === "string") { + return redactSensitiveCredentials(val); + } + if (Array.isArray(val)) { + return val.map(redactSensitiveInObject); + } + if (val !== null && typeof val === "object") { + const res: Record = {}; + for (const [k, v] of Object.entries(val)) { + res[k] = redactSensitiveInObject(v); + } + return res; + } + return val; +} + +function redactSensitiveCredentialsInMessages(messages: Message[]): Message[] { + return messages.map((msg): Message => { + if (msg.role === "user" || msg.role === "developer") { + const userMsg = msg as UserMessage | DeveloperMessage; + if (typeof userMsg.content === "string") { + const redacted = redactSensitiveCredentials(userMsg.content); + if (redacted === userMsg.content) return msg; + return { ...userMsg, content: redacted } as Message; + } + const contentArray = userMsg.content; + let changed = false; + const content = contentArray.map((block): UserMessage["content"][number] => { + if (block.type === "text") { + const redacted = redactSensitiveCredentials(block.text); + if (redacted !== block.text) { + changed = true; + return { ...block, text: redacted }; + } + } + return block; + }); + return (changed ? { ...userMsg, content } : userMsg) as Message; + } + + if (msg.role === "toolResult") { + const toolResultMsg = msg as ToolResultMessage; + let changed = false; + const content = toolResultMsg.content.map((block): ToolResultMessage["content"][number] => { + if (block.type === "text") { + const redacted = redactSensitiveCredentials(block.text); + if (redacted !== block.text) { + changed = true; + return { ...block, text: redacted }; + } + } + return block; + }); + return (changed ? { ...toolResultMsg, content } : toolResultMsg) as Message; + } + + if (msg.role === "assistant") { + const assistantMsg = msg as AssistantMessage; + let changed = false; + const content = assistantMsg.content.map((block): AssistantMessage["content"][number] => { + if (block.type === "text") { + const redacted = redactSensitiveCredentials(block.text); + if (redacted !== block.text) { + changed = true; + return { ...block, text: redacted }; + } + } else if (block.type === "thinking") { + const redacted = redactSensitiveCredentials(block.thinking); + if (redacted !== block.thinking) { + changed = true; + return { ...block, thinking: redacted }; + } + } else if (block.type === "toolCall") { + if (block.arguments) { + const redactedArgs = redactSensitiveInObject(block.arguments); + if (JSON.stringify(redactedArgs) !== JSON.stringify(block.arguments)) { + changed = true; + const castArgs = + redactedArgs && typeof redactedArgs === "object" && !Array.isArray(redactedArgs) + ? (redactedArgs as Record) + : undefined; + return { ...block, arguments: castArgs } as AssistantMessage["content"][number]; + } + } + } + return block; + }); + return (changed ? { ...assistantMsg, content } : assistantMsg) as Message; + } + + return msg; + }); +} + export function transformMessages( messages: Message[], model: Model, @@ -294,6 +419,10 @@ export function transformMessages( duplicateToolCallIdSuffixPrefix = "_dup", targetCompat: Model["compat"] = model.compat, ): Message[] { + // Redact sensitive credential-like patterns from all outbound messages + // to prevent security block errors from LLM providers (e.g. invalid_prompt). + messages = redactSensitiveCredentialsInMessages(messages); + // Drop assistant `toolCall` blocks with empty/whitespace `id` or `name` // (and their matched `toolResult` messages) before anything else looks at // the history. Replays of these would 400 every provider — see diff --git a/packages/ai/test/transform-messages-redact-sensitive.test.ts b/packages/ai/test/transform-messages-redact-sensitive.test.ts new file mode 100644 index 000000000..358cd793e --- /dev/null +++ b/packages/ai/test/transform-messages-redact-sensitive.test.ts @@ -0,0 +1,105 @@ +import { describe, expect, it } from "bun:test"; +import { transformMessages } from "@oh-my-pi/pi-ai/providers/transform-messages"; +import type { AssistantMessage, Message, Model, ToolCall, ToolResultMessage } from "@oh-my-pi/pi-ai/types"; +import { buildModel } from "@oh-my-pi/pi-catalog/build"; + +function makeModel(): Model<"openai-responses"> { + return buildModel({ + api: "openai-responses", + name: "GPT Test", + id: "gpt-test", + provider: "openai", + baseUrl: "https://api.openai.com/v1", + contextWindow: 8192, + maxTokens: 2048, + input: ["text"], + reasoning: false, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + }); +} + +describe("transformMessages redact sensitive credentials", () => { + it("redacts already-masked and real tokens from outbound messages", () => { + const messages: Message[] = [ + { + role: "user", + content: "Token: gho_************************************", + timestamp: Date.now(), + }, + { + role: "assistant", + content: [ + { + type: "text", + text: "I found this key: sk-proj-************************************", + }, + { + type: "toolCall", + id: "call_x", + name: "bash", + arguments: { + command: "echo gho_************************************", + }, + }, + ], + api: "openai-responses", + provider: "openai", + model: "gpt-test", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "toolUse", + timestamp: Date.now(), + }, + { + role: "toolResult", + toolCallId: "call_x", + toolName: "bash", + content: [{ type: "text", text: "Token is ghp_************************************ inside output" }], + isError: false, + timestamp: Date.now(), + }, + ]; + + const transformed = transformMessages(messages, makeModel()); + + // 1. Verify user message is redacted + const userMsg = transformed[0]; + expect(userMsg.role).toBe("user"); + expect(userMsg.content).toBe("Token: [github_token_redacted]"); + + // 2. Verify assistant message text and toolCall arguments are redacted + const assistantMsg = transformed[1]; + expect(assistantMsg.role).toBe("assistant"); + const castAssistantMsg = assistantMsg as AssistantMessage; + const assistantContent = castAssistantMsg.content; + const textBlock = assistantContent[0]; + expect(textBlock.type).toBe("text"); + if (textBlock.type === "text") { + expect(textBlock.text).toBe("I found this key: [openai_token_redacted]"); + } + + const toolCallBlock = assistantContent[1]; + expect(toolCallBlock.type).toBe("toolCall"); + + // 3. Verify toolResult message is redacted + const resultMsg = transformed[2]; + expect(resultMsg.role).toBe("toolResult"); + const toolResultMsg = resultMsg as ToolResultMessage; + const toolResultBlock = toolResultMsg.content[0]; + expect(toolResultBlock.type).toBe("text"); + if (toolResultBlock.type === "text") { + expect(toolResultBlock.text).toBe("Token is [github_token_redacted] inside output"); + } + if (toolCallBlock.type === "toolCall") { + const toolCall = toolCallBlock as ToolCall; + const commandArg = toolCall.arguments?.command; + expect(commandArg).toBe("echo [github_token_redacted]"); + } + }); +});