From: Stefan Gasser Date: Tue, 25 Aug 2026 09:16:15 +0000 (+0200) Subject: fix(openai): restore tool call arguments (#168) X-Git-Tag: v0.9.3~3 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=0e0a8d9d7dc0b98327a0fd03a5be9c1bd624ae68;p=sgasser-llm-shield.git fix(openai): restore tool call arguments (#168) * fix(openai): restore tool call arguments * fix(openai): remask tool call history --- diff --git a/src/masking/extractors/openai.test.ts b/src/masking/extractors/openai.test.ts index 2674784..23262e2 100644 --- a/src/masking/extractors/openai.test.ts +++ b/src/masking/extractors/openai.test.ts @@ -1,7 +1,12 @@ import { describe, expect, test } from "bun:test"; import type { PlaceholderContext } from "../../masking/context"; -import type { OpenAIMessage, OpenAIRequest, OpenAIResponse } from "../../providers/openai/types"; -import { openaiExtractor } from "./openai"; +import { + type OpenAIMessage, + type OpenAIRequest, + type OpenAIResponse, + OpenAIResponseSchema, +} from "../../providers/openai/types"; +import { openaiExtractor, remaskOpenAIToolCallArguments } from "./openai"; /** Helper to create a minimal request from messages */ function createRequest(messages: OpenAIMessage[]): OpenAIRequest { @@ -175,6 +180,159 @@ describe("OpenAI Text Extractor", () => { }); }); + describe("remaskOpenAIToolCallArguments", () => { + test("preserves JSON bytes except known values and chooses the longest overlap", () => { + const argumentsText = String.raw`{ "big":9007199254740993, "decimal":1.0, "exponent":1e+3, "dup":"jane@example.com", "dup":"example.com", "quote":"say \"hi\"", "backslash":"C:\\temp", "unicode":"caf\u00e9 \uD83D\uDE00", "unrelated":"quote \" slash \\ newline \n snow 雪", "literal":"\\nbreak", "actual":"\nbreak" }`; + const request = createRequest([ + { role: "user", content: "Continue" }, + { + role: "assistant", + content: "Unscanned assistant metadata stays unchanged", + tool_calls: [ + { + id: "call_fidelity", + type: "function", + provider_metadata: { trace: "keep" }, + function: { + name: "store_values", + arguments: argumentsText, + }, + }, + ], + // biome-ignore lint/suspicious/noExplicitAny: testing passthrough field preservation + } as any, + { role: "tool", tool_call_id: "call_fidelity", content: "Done" }, + ]); + // biome-ignore lint/suspicious/noExplicitAny: testing passthrough field preservation + (request as any).metadata = { request_id: "keep" }; + const context: PlaceholderContext = { + mapping: { + "[[DOMAIN_1]]": "example.com", + "[[EMAIL_ADDRESS_1]]": "jane@example.com", + "[[QUOTE_1]]": 'say "hi"', + "[[BACKSLASH_1]]": "C:\\temp", + "[[UNICODE_1]]": "café 😀", + "[[NEWLINE_1]]": "\nbreak", + }, + reverseMapping: { + "example.com": "[[DOMAIN_1]]", + "jane@example.com": "[[EMAIL_ADDRESS_1]]", + 'say "hi"': "[[QUOTE_1]]", + "C:\\temp": "[[BACKSLASH_1]]", + "café 😀": "[[UNICODE_1]]", + "\nbreak": "[[NEWLINE_1]]", + }, + counters: { + DOMAIN: 1, + EMAIL_ADDRESS: 1, + QUOTE: 1, + BACKSLASH: 1, + UNICODE: 1, + NEWLINE: 1, + }, + }; + + const result = remaskOpenAIToolCallArguments(request, [context]); + // biome-ignore lint/suspicious/noExplicitAny: inspecting passthrough tool-call fields + const toolCall = (result.messages[1] as any).tool_calls[0]; + + expect(toolCall.function.arguments).toBe( + String.raw`{ "big":9007199254740993, "decimal":1.0, "exponent":1e+3, "dup":"[[EMAIL_ADDRESS_1]]", "dup":"[[DOMAIN_1]]", "quote":"[[QUOTE_1]]", "backslash":"[[BACKSLASH_1]]", "unicode":"[[UNICODE_1]]", "unrelated":"quote \" slash \\ newline \n snow 雪", "literal":"\\nbreak", "actual":"[[NEWLINE_1]]" }`, + ); + expect(JSON.parse(toolCall.function.arguments)).toEqual({ + big: 9007199254740992, + decimal: 1, + exponent: 1000, + dup: "[[DOMAIN_1]]", + quote: "[[QUOTE_1]]", + backslash: "[[BACKSLASH_1]]", + unicode: "[[UNICODE_1]]", + unrelated: 'quote " slash \\ newline \n snow 雪', + literal: "\\nbreak", + actual: "[[NEWLINE_1]]", + }); + expect(toolCall.provider_metadata).toEqual({ trace: "keep" }); + expect(result.messages[1].content).toBe("Unscanned assistant metadata stays unchanged"); + // biome-ignore lint/suspicious/noExplicitAny: inspecting passthrough request metadata + expect((result as any).metadata).toEqual({ request_id: "keep" }); + }); + + test("remasks an exact known numeric primitive without changing other number lexemes", () => { + const phone = "3901234567"; + const argumentsText = `{"phone":${phone},"asString":"${phone}","larger":13901234567,"decimal":${phone}.0,"exponent":${phone}e0,"big":9007199254740993,"one":1.0,"thousand":1e+3}`; + const request = createRequest([ + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_numeric", + type: "function", + function: { name: "lookup_phone", arguments: argumentsText }, + }, + ], + // biome-ignore lint/suspicious/noExplicitAny: testing passthrough tool-call arguments + } as any, + ]); + const context: PlaceholderContext = { + mapping: { "[[PHONE_NUMBER_1]]": phone }, + reverseMapping: { [phone]: "[[PHONE_NUMBER_1]]" }, + counters: { PHONE_NUMBER: 1 }, + }; + + const result = remaskOpenAIToolCallArguments(request, [context]); + // biome-ignore lint/suspicious/noExplicitAny: inspecting passthrough tool-call arguments + const remasked = (result.messages[0] as any).tool_calls[0].function.arguments; + + expect(remasked).toBe( + '{"phone":"[[PHONE_NUMBER_1]]","asString":"[[PHONE_NUMBER_1]]","larger":13901234567,"decimal":3901234567.0,"exponent":3901234567e0,"big":9007199254740993,"one":1.0,"thousand":1e+3}', + ); + }); + + test("remasks safe malformed JSON while leaving unsafe string bytes unchanged", () => { + const context: PlaceholderContext = { + mapping: { "[[EMAIL_ADDRESS_1]]": "jane@example.com" }, + reverseMapping: { "jane@example.com": "[[EMAIL_ADDRESS_1]]" }, + counters: { EMAIL_ADDRESS: 1 }, + }; + const safeMalformed = '{"email":"jane@example.com",}'; + const safeUnterminated = '{"email":"jane@example.com'; + const unsafeMalformed = '{"email":"jane@example.com\\'; + const request = createRequest([ + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_safe", + type: "function", + function: { name: "safe", arguments: safeMalformed }, + }, + { + id: "call_unterminated", + type: "function", + function: { name: "unterminated", arguments: safeUnterminated }, + }, + { + id: "call_unsafe", + type: "function", + function: { name: "unsafe", arguments: unsafeMalformed }, + }, + ], + // biome-ignore lint/suspicious/noExplicitAny: testing passthrough tool-call arguments + } as any, + ]); + + const result = remaskOpenAIToolCallArguments(request, [context]); + // biome-ignore lint/suspicious/noExplicitAny: inspecting passthrough tool-call arguments + const toolCalls = (result.messages[0] as any).tool_calls; + + expect(toolCalls[0].function.arguments).toBe('{"email":"[[EMAIL_ADDRESS_1]]",}'); + expect(toolCalls[1].function.arguments).toBe('{"email":"[[EMAIL_ADDRESS_1]]'); + expect(toolCalls[2].function.arguments).toBe(unsafeMalformed); + }); + }); + describe("unmaskResponse", () => { test("unmasks placeholders in response content", () => { const response: OpenAIResponse = { @@ -211,6 +369,80 @@ describe("OpenAI Text Extractor", () => { expect(result.choices[0].message.content).toBe("Hello John, your email is john@example.com"); }); + test("unmasks placeholders in tool call function arguments as valid JSON", () => { + const originalValue = 'Quote " slash \\ controls \b\f\n\r\t \u0001 Unicode 雪 😀'; + const response: OpenAIResponse = { + id: "test-id", + object: "chat.completion", + created: 123456, + model: "gpt-4", + choices: [ + { + index: 0, + message: { + role: "assistant", + content: "Hello [[PERSON_1]]", + tool_calls: [ + { + id: "call_123", + type: "function", + function: { + name: "lookup_person", + arguments: JSON.stringify({ person: "[[PERSON_1]]" }), + }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }; + const context: PlaceholderContext = { + mapping: { "[[PERSON_1]]": originalValue }, + reverseMapping: { [originalValue]: "[[PERSON_1]]" }, + counters: { PERSON: 1 }, + }; + + const result = openaiExtractor.unmaskResponse(response, context); + const toolCalls = result.choices[0].message.tool_calls as Array<{ + id: string; + type: string; + function: { name: string; arguments: string }; + }>; + + expect(result.choices[0].message.content).toBe(`Hello ${originalValue}`); + expect(toolCalls[0]).toEqual({ + id: "call_123", + type: "function", + function: { + name: "lookup_person", + arguments: expect.any(String), + }, + }); + expect(JSON.parse(toolCalls[0].function.arguments)).toEqual({ person: originalValue }); + }); + + test.each([ + "tool_calls", + "function_call", + ] as const)("accepts the OpenAI %s finish reason", (finishReason) => { + expect( + OpenAIResponseSchema.safeParse({ + id: "test-id", + object: "chat.completion", + created: 123456, + model: "gpt-4", + choices: [ + { + index: 0, + message: { role: "assistant", content: null }, + finish_reason: finishReason, + }, + ], + }).success, + ).toBeTrue(); + }); + test("applies formatValue function when provided", () => { const response: OpenAIResponse = { id: "test-id", diff --git a/src/masking/extractors/openai.ts b/src/masking/extractors/openai.ts index 5005211..a59db9d 100644 --- a/src/masking/extractors/openai.ts +++ b/src/masking/extractors/openai.ts @@ -39,6 +39,288 @@ function unmaskContent( return content; } +function unmaskToolCalls( + toolCalls: unknown[], + context: PlaceholderContext, + formatValue?: (original: string) => string, +): unknown[] { + return toolCalls.map((toolCall) => { + if (typeof toolCall !== "object" || toolCall === null || !("function" in toolCall)) { + return toolCall; + } + + const functionCall = toolCall.function; + if ( + typeof functionCall !== "object" || + functionCall === null || + !("arguments" in functionCall) || + typeof functionCall.arguments !== "string" + ) { + return toolCall; + } + + return { + ...toolCall, + function: { + ...functionCall, + arguments: restorePlaceholders(functionCall.arguments, context, (original) => + JSON.stringify(formatValue ? formatValue(original) : original).slice(1, -1), + ), + }, + }; + }); +} + +interface KnownValueReplacement { + original: string; + serializedPlaceholder: string; +} + +interface ReplacementMatch extends KnownValueReplacement { + start: number; + end: number; +} + +function knownValueReplacements( + contexts: readonly (PlaceholderContext | undefined)[], +): KnownValueReplacement[] { + const originals = new Set(); + const replacements: KnownValueReplacement[] = []; + + for (const context of contexts) { + if (!context) continue; + for (const [placeholder, original] of Object.entries(context.mapping)) { + if (original.length === 0 || originals.has(original)) continue; + originals.add(original); + replacements.push({ + original, + serializedPlaceholder: JSON.stringify(placeholder).slice(1, -1), + }); + } + } + + return replacements.sort((a, b) => b.original.length - a.original.length); +} + +function findJsonStringEnd(serialized: string, start: number): number { + let index = start + 1; + while (index < serialized.length) { + if (serialized[index] === '"') return index; + if (serialized[index] === "\\") index++; + index++; + } + return -1; +} + +function decodeJsonStringToken( + token: string, +): { decoded: string; rawBoundaries: number[] } | undefined { + let decoded: unknown; + try { + decoded = JSON.parse(token); + } catch { + return undefined; + } + if (typeof decoded !== "string") return undefined; + + const rawBoundaries = [1]; + let rawIndex = 1; + while (rawIndex < token.length - 1) { + if (token[rawIndex] === "\\") { + rawIndex += token[rawIndex + 1] === "u" ? 6 : 2; + } else { + rawIndex++; + } + rawBoundaries.push(rawIndex); + } + + return rawBoundaries.length === decoded.length + 1 ? { decoded, rawBoundaries } : undefined; +} + +function findReplacementMatches( + decoded: string, + replacements: KnownValueReplacement[], +): ReplacementMatch[] { + const candidates: ReplacementMatch[] = []; + + for (const replacement of replacements) { + let start = decoded.indexOf(replacement.original); + while (start !== -1) { + candidates.push({ + ...replacement, + start, + end: start + replacement.original.length, + }); + start = decoded.indexOf(replacement.original, start + 1); + } + } + + candidates.sort((a, b) => b.original.length - a.original.length || a.start - b.start); + + const selected: ReplacementMatch[] = []; + for (const candidate of candidates) { + if (selected.some((match) => candidate.start < match.end && match.start < candidate.end)) { + continue; + } + selected.push(candidate); + } + + return selected.sort((a, b) => a.start - b.start); +} + +function remaskJsonStringToken(token: string, replacements: KnownValueReplacement[]): string { + const decodedToken = decodeJsonStringToken(token); + if (!decodedToken) return token; + + const matches = findReplacementMatches(decodedToken.decoded, replacements); + if (matches.length === 0) return token; + + const parts: string[] = []; + let rawIndex = 0; + for (const match of matches) { + const rawStart = decodedToken.rawBoundaries[match.start]; + const rawEnd = decodedToken.rawBoundaries[match.end]; + parts.push(token.slice(rawIndex, rawStart), match.serializedPlaceholder); + rawIndex = rawEnd; + } + parts.push(token.slice(rawIndex)); + return parts.join(""); +} + +function isJsonWhitespace(character: string | undefined): boolean { + return character === " " || character === "\t" || character === "\n" || character === "\r"; +} + +function findJsonNumberEnd(serialized: string, start: number): number { + const previous = serialized[start - 1]; + if ( + start > 0 && + previous !== "[" && + previous !== "," && + previous !== ":" && + !isJsonWhitespace(previous) + ) { + return -1; + } + + const match = /^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?/.exec(serialized.slice(start)); + if (!match) return -1; + + const end = start + match[0].length; + const next = serialized[end]; + return end === serialized.length || + next === "," || + next === "]" || + next === "}" || + isJsonWhitespace(next) + ? end + : -1; +} + +function remaskSerializedJsonValues( + serialized: string, + replacements: KnownValueReplacement[], +): string { + const parts: string[] = []; + let copyFrom = 0; + let scanIndex = 0; + + while (scanIndex < serialized.length) { + const character = serialized[scanIndex]; + if (character !== '"') { + const numberEnd = + character === "-" || (character >= "0" && character <= "9") + ? findJsonNumberEnd(serialized, scanIndex) + : -1; + if (numberEnd !== -1) { + const token = serialized.slice(scanIndex, numberEnd); + const replacement = replacements.find(({ original }) => original === token); + if (replacement) { + parts.push( + serialized.slice(copyFrom, scanIndex), + `"${replacement.serializedPlaceholder}"`, + ); + copyFrom = numberEnd; + } + scanIndex = numberEnd; + continue; + } + scanIndex++; + continue; + } + + const end = findJsonStringEnd(serialized, scanIndex); + if (end === -1) { + const closedToken = `${serialized.slice(scanIndex)}"`; + const remasked = remaskJsonStringToken(closedToken, replacements); + if (remasked !== closedToken) { + parts.push(serialized.slice(copyFrom, scanIndex), remasked.slice(0, -1)); + copyFrom = serialized.length; + } + break; + } + + const token = serialized.slice(scanIndex, end + 1); + const remasked = remaskJsonStringToken(token, replacements); + if (remasked !== token) { + parts.push(serialized.slice(copyFrom, scanIndex), remasked); + copyFrom = end + 1; + } + scanIndex = end + 1; + } + + if (parts.length === 0) return serialized; + parts.push(serialized.slice(copyFrom)); + return parts.join(""); +} + +export function remaskOpenAIToolCallArguments( + request: OpenAIRequest, + contexts: readonly (PlaceholderContext | undefined)[], +): OpenAIRequest { + const replacements = knownValueReplacements(contexts); + + if (replacements.length === 0) return request; + + let requestChanged = false; + const messages = request.messages.map((message) => { + if (message.role !== "assistant" || !Array.isArray(message.tool_calls)) return message; + + let messageChanged = false; + const toolCalls = message.tool_calls.map((toolCall) => { + if (typeof toolCall !== "object" || toolCall === null || !("function" in toolCall)) { + return toolCall; + } + + const functionCall = toolCall.function; + if ( + typeof functionCall !== "object" || + functionCall === null || + !("arguments" in functionCall) || + typeof functionCall.arguments !== "string" + ) { + return toolCall; + } + + const remasked = remaskSerializedJsonValues(functionCall.arguments, replacements); + + if (remasked === functionCall.arguments) return toolCall; + + messageChanged = true; + return { + ...toolCall, + function: { ...functionCall, arguments: remasked }, + }; + }); + + if (!messageChanged) return message; + requestChanged = true; + return { ...message, tool_calls: toolCalls }; + }); + + return requestChanged ? { ...request, messages } : request; +} + /** * OpenAI request extractor * @@ -123,13 +405,20 @@ export const openaiExtractor: RequestExtractor = ): OpenAIResponse { return { ...response, - choices: response.choices.map((choice) => ({ - ...choice, - message: { - ...choice.message, - content: unmaskContent(choice.message.content, context, formatValue), - }, - })), + choices: response.choices.map((choice) => { + const toolCalls = choice.message.tool_calls; + + return { + ...choice, + message: { + ...choice.message, + content: unmaskContent(choice.message.content, context, formatValue), + ...(Array.isArray(toolCalls) + ? { tool_calls: unmaskToolCalls(toolCalls, context, formatValue) } + : {}), + }, + }; + }), }; }, }; diff --git a/src/providers/openai/stream-transformer.test.ts b/src/providers/openai/stream-transformer.test.ts index 8471856..276695d 100644 --- a/src/providers/openai/stream-transformer.test.ts +++ b/src/providers/openai/stream-transformer.test.ts @@ -44,6 +44,94 @@ async function consumeStream(stream: ReadableStream): Promise line.startsWith("data: ")) + .map((line) => line.slice(6)); +} + +function parseJsonEvents(output: string): ParsedStreamEvent[] { + return parseDataPayloads(output) + .filter((payload) => payload !== "[DONE]" && payload !== "not-json") + .map((payload) => JSON.parse(payload) as ParsedStreamEvent); +} + +function collectToolArguments( + events: ParsedStreamEvent[], + choiceIndex: number, + toolCallIndex: number, +): string { + return events + .flatMap((event) => event.choices) + .filter((choice) => choice.index === choiceIndex) + .flatMap((choice) => choice.delta.tool_calls ?? []) + .filter((toolCall) => toolCall.index === toolCallIndex) + .map((toolCall) => toolCall.function?.arguments) + .filter((value): value is string => typeof value === "string") + .join(""); +} + describe("createUnmaskingStream", () => { test("unmasks complete placeholder in single chunk", async () => { const context = createMaskingContext(); @@ -235,4 +323,403 @@ describe("createUnmaskingStream", () => { expect(result).toContain("[protected]sk-secret"); expect(result).not.toContain("[[API_KEY_SK_1]]"); }); + + test.each([ + ["a complete placeholder", ['{"person":"[[PERSON_1]]"}']], + ["a placeholder split across deltas", ['{"person":"[[PERS', 'ON_1]]"}']], + ["an argument JSON string split across deltas", ['{"per', 'son":"[[PERSON_1]]"}']], + ])("restores tool arguments with %s", async (_description, fragments) => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Taylor Example"; + const chunks = fragments.map((argumentsFragment) => + createToolArgumentEvent(argumentsFragment, { + toolCallIndex: 2, + }), + ); + chunks.push("data: [DONE]\n\n"); + + const output = await consumeStream( + createUnmaskingStream(createSSEStream(chunks), context, defaultConfig), + ); + const events = parseJsonEvents(output); + + expect(JSON.parse(collectToolArguments(events, 0, 2))).toEqual({ + person: "Taylor Example", + }); + expect(parseDataPayloads(output).at(-1)).toBe("[DONE]"); + expect(events).toHaveLength(fragments.length); + }); + + test("buffers an argument event split mid-JSON across source chunks", async () => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Taylor Example"; + const event = createToolArgumentEvent('{"person":"[[PERSON_1]]"}'); + const splitAt = event.indexOf("PERSON_1"); + const source = createSSEStream([event.slice(0, splitAt), event.slice(splitAt)]); + + const output = await consumeStream(createUnmaskingStream(source, context, defaultConfig)); + const events = parseJsonEvents(output); + + expect(JSON.parse(collectToolArguments(events, 0, 0))).toEqual({ + person: "Taylor Example", + }); + expect(events).toHaveLength(1); + }); + + test("keeps PII and secret arguments valid JSON with markers and escaped values", async () => { + const piiContext = createMaskingContext(); + const secretsContext = createMaskingContext(); + const person = 'Quote " slash \\ controls \b\f\n\r\t \u0001 Unicode 雪 😀'; + const secret = 'secret "\\\n\u0002雪'; + piiContext.mapping["[[PERSON_1]]"] = person; + secretsContext.mapping["[[API_KEY_SK_1]]"] = secret; + const fragments = ['{"person":"[[PERSON_', '1]]","secret":"[[API_KEY_', 'SK_1]]"}']; + const chunks = fragments.map((argumentsFragment) => createToolArgumentEvent(argumentsFragment)); + + const output = await consumeStream( + createUnmaskingStream(createSSEStream(chunks), piiContext, markerConfig, secretsContext), + ); + const restoredArguments = collectToolArguments(parseJsonEvents(output), 0, 0); + + expect(JSON.parse(restoredArguments)).toEqual({ + person: `[protected]${person}`, + secret: `[protected]${secret}`, + }); + }); + + test("isolates interleaved choices and tool calls while preserving event metadata", async () => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Taylor Example"; + context.mapping["[[LOCATION_1]]"] = "Paris"; + const chunks = [ + createSSEEvent({ + id: "chunk-interleaved-1", + object: "chat.completion.chunk", + created: 123, + model: "mock-model", + vendor_event: "first", + choices: [ + { + index: 2, + delta: { + role: "assistant", + tool_calls: [ + { + index: 1, + id: "call_person_2", + type: "function", + vendor_tool: "person", + function: { name: "person", arguments: '{"value":"[[PER' }, + }, + { + index: 0, + id: "call_location_2", + type: "function", + function: { name: "location", arguments: '{"value":"[[LOC' }, + }, + ], + }, + finish_reason: null, + logprobs: null, + vendor_choice: "choice-2", + }, + { + index: 0, + delta: { + tool_calls: [ + { + index: 3, + id: "call_person_0", + type: "function", + function: { name: "person", arguments: '{"value":"[[PERS' }, + }, + ], + }, + finish_reason: null, + }, + ], + }), + createSSEEvent({ + id: "chunk-interleaved-2", + object: "chat.completion.chunk", + created: 124, + model: "mock-model", + vendor_event: "second", + choices: [ + { + index: 0, + delta: { + tool_calls: [{ index: 3, function: { arguments: 'ON_1]]"}' } }], + }, + finish_reason: null, + }, + { + index: 2, + delta: { + tool_calls: [ + { index: 0, function: { arguments: 'ATION_1]]"}' } }, + { index: 1, function: { arguments: 'SON_1]]"}' } }, + ], + }, + finish_reason: null, + }, + ], + }), + createSSEEvent({ + id: "chunk-usage", + choices: [], + usage: { prompt_tokens: 4, completion_tokens: 3, total_tokens: 7 }, + vendor_event: "usage", + }), + "data: [DONE]\n\n", + ]; + + const output = await consumeStream( + createUnmaskingStream(createSSEStream(chunks), context, defaultConfig), + ); + const events = parseJsonEvents(output); + + expect(JSON.parse(collectToolArguments(events, 2, 1))).toEqual({ + value: "Taylor Example", + }); + expect(JSON.parse(collectToolArguments(events, 2, 0))).toEqual({ value: "Paris" }); + expect(JSON.parse(collectToolArguments(events, 0, 3))).toEqual({ + value: "Taylor Example", + }); + expect(events).toHaveLength(3); + expect(events.map((event) => event.id)).toEqual([ + "chunk-interleaved-1", + "chunk-interleaved-2", + "chunk-usage", + ]); + expect(events[0].vendor_event).toBe("first"); + expect(events[0].choices.map((choice) => choice.index)).toEqual([2, 0]); + expect(events[0].choices[0]).toMatchObject({ + logprobs: null, + vendor_choice: "choice-2", + }); + expect(events[0].choices[0].delta.tool_calls?.map((toolCall) => toolCall.index)).toEqual([ + 1, 0, + ]); + expect(events[0].choices[0].delta.tool_calls?.[0]).toMatchObject({ + id: "call_person_2", + type: "function", + vendor_tool: "person", + function: { name: "person" }, + }); + expect(events[2]).toEqual({ + id: "chunk-usage", + choices: [], + usage: { prompt_tokens: 4, completion_tokens: 3, total_tokens: 7 }, + vendor_event: "usage", + }); + expect(parseDataPayloads(output).at(-1)).toBe("[DONE]"); + }); + + test.each([ + [ + "choice index", + { + delta: { + tool_calls: [ + { + index: 0, + function: { arguments: '{"person":"[[PERSON_1]]"}' }, + }, + ], + }, + finish_reason: null, + }, + ], + [ + "tool-call index", + { + index: 0, + delta: { + tool_calls: [ + { + function: { arguments: '{"person":"[[PERSON_1]]"}' }, + }, + ], + }, + finish_reason: null, + }, + ], + ] satisfies Array< + [string, ChoiceDelta] + >)("passes through tool arguments without an explicit %s", async (_description, choice) => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Taylor Example"; + const input = createSSEEvent({ id: "chunk-missing-index", choices: [choice] }); + + const output = await consumeStream( + createUnmaskingStream(createSSEStream([input, "data: [DONE]\n\n"]), context, defaultConfig), + ); + + expect(parseJsonEvents(output)[0].choices[0].delta.tool_calls?.[0].function?.arguments).toBe( + '{"person":"[[PERSON_1]]"}', + ); + }); + + test("flushes incomplete tool arguments before finish and DONE without changing malformed data", async () => { + const context = createMaskingContext(); + const chunks = [ + createToolArgumentEvent('{"raw":"value[[UNKNOWN', { + choiceIndex: 1, + toolCallIndex: 4, + }), + "data: not-json\n\n", + createSSEEvent({ + id: "chunk-finish", + object: "chat.completion.chunk", + created: 457, + model: "mock-model", + choices: [ + { + index: 1, + delta: {}, + finish_reason: "tool_calls", + vendor_finish: true, + }, + ], + }), + "data: [DONE]\n\n", + ]; + + const output = await consumeStream( + createUnmaskingStream(createSSEStream(chunks), context, defaultConfig), + ); + const payloads = parseDataPayloads(output); + const events = parseJsonEvents(output); + + expect(payloads).toHaveLength(4); + expect(payloads[1]).toBe("not-json"); + expect(collectToolArguments(events, 1, 4)).toBe('{"raw":"value[[UNKNOWN'); + expect(events[1]).toMatchObject({ + id: "chunk-finish", + object: "chat.completion.chunk", + created: 457, + model: "mock-model", + choices: [ + { + index: 1, + delta: { tool_calls: [{ index: 4, function: { arguments: "[[UNKNOWN" } }] }, + finish_reason: "tool_calls", + vendor_finish: true, + }, + ], + }); + expect(payloads.at(-1)).toBe("[DONE]"); + }); + + test("flushes incomplete tool arguments on clean EOF without inventing DONE", async () => { + const context = createMaskingContext(); + const chunks = [ + createSSEEvent({ + id: "chunk-eof", + object: "chat.completion.chunk", + created: 789, + model: "mock-model", + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 2, + id: "call_eof", + type: "function", + function: { name: "extract", arguments: '{"raw":"value[' }, + }, + ], + }, + finish_reason: null, + }, + ], + }), + ]; + + const output = await consumeStream( + createUnmaskingStream(createSSEStream(chunks), context, defaultConfig), + ); + const payloads = parseDataPayloads(output); + const events = parseJsonEvents(output); + + expect(payloads).toHaveLength(2); + expect(payloads).not.toContain("[DONE]"); + expect(collectToolArguments(events, 0, 2)).toBe('{"raw":"value['); + expect(events[1]).toMatchObject({ + id: "chunk-eof", + object: "chat.completion.chunk", + created: 789, + model: "mock-model", + choices: [ + { + index: 0, + delta: { tool_calls: [{ index: 2, function: { arguments: "[" } }] }, + finish_reason: null, + }, + ], + }); + }); + + test("flushes pending tool arguments before a CRLF terminal marker", async () => { + const context = createMaskingContext(); + const event = createToolArgumentEvent('{"raw":"value[').replaceAll("\n", "\r\n"); + + const output = await consumeStream( + createUnmaskingStream( + createSSEStream([event, "data: [DONE]\r\n\r\n"]), + context, + defaultConfig, + ), + ); + const payloads = parseDataPayloads(output); + + expect(payloads.at(-1)).toBe("[DONE]"); + const events = payloads.slice(0, -1).map((payload) => JSON.parse(payload) as ParsedStreamEvent); + expect(collectToolArguments(events, 0, 0)).toBe('{"raw":"value['); + }); + + test("stops after the first DONE and flushes tool arguments before it", async () => { + const context = createMaskingContext(); + const chunks = [ + createSSEEvent({ + id: "chunk-arguments", + object: "chat.completion.chunk", + created: 123, + model: "mock-model", + choices: [ + { + index: 0, + delta: { + tool_calls: [{ index: 1, function: { arguments: '{"raw":"value[' } }], + }, + finish_reason: null, + }, + ], + }), + "data:[DONE]\n\n", + "data:[DONE]\n\n", + createSSEEvent({ + id: "chunk-late-usage", + choices: [], + usage: { prompt_tokens: 4, completion_tokens: 3, total_tokens: 7 }, + }), + ]; + + const output = await consumeStream( + createUnmaskingStream(createSSEStream(chunks), context, defaultConfig), + ); + const timeline = parseDataPayloads(output); + const events = timeline + .filter((payload) => payload !== "[DONE]") + .map((payload) => JSON.parse(payload) as ParsedStreamEvent); + + expect(timeline.at(-1)).toBe("[DONE]"); + expect(timeline.filter((payload) => payload === "[DONE]")).toHaveLength(1); + expect(collectToolArguments(events, 0, 1)).toBe('{"raw":"value['); + expect(events.map((event) => event.id)).toEqual(["chunk-arguments", "chunk-arguments"]); + expect(output).not.toContain("chunk-late-usage"); + }); }); diff --git a/src/providers/openai/stream-transformer.ts b/src/providers/openai/stream-transformer.ts index f783588..2e77c62 100644 --- a/src/providers/openai/stream-transformer.ts +++ b/src/providers/openai/stream-transformer.ts @@ -1,8 +1,44 @@ import type { MaskingConfig } from "../../config"; import type { PlaceholderContext } from "../../masking/context"; +import { createRestoreFormatter } from "../../masking/restore-policy"; import { StreamRestorer } from "../../masking/stream-restorer"; import type { OpenAIContentPart } from "../../utils/content"; +type EventMetadata = Record; + +const SYNTHETIC_EVENT_METADATA_KEYS = [ + "id", + "object", + "created", + "model", + "system_fingerprint", + "service_tier", +] as const; + +interface ToolArgumentChannel { + choiceIndex: number; + toolCallIndex: number; + metadata: EventMetadata; + restorer: StreamRestorer; +} + +interface FlushedArgumentChannel { + channel: ToolArgumentChannel; + text: string; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function getSyntheticEventMetadata(event: Record): EventMetadata { + const metadata: EventMetadata = {}; + for (const key of SYNTHETIC_EVENT_METADATA_KEYS) { + if (key in event) metadata[key] = event[key]; + } + return metadata; +} + export function createUnmaskingStream( source: ReadableStream, piiContext: PlaceholderContext | undefined, @@ -13,31 +49,234 @@ export function createUnmaskingStream( const encoder = new TextEncoder(); let lineBuffer = ""; const restorer = new StreamRestorer({ piiContext, secretsContext, config }); + const argumentChannelsByPosition = new Map(); + const formatValue = createRestoreFormatter(config); + const formatArgumentValue = (original: string) => + JSON.stringify(formatValue ? formatValue(original) : original).slice(1, -1); + const formatArgumentContext = (context: PlaceholderContext | undefined) => + context + ? { + ...context, + mapping: Object.fromEntries( + Object.entries(context.mapping).map(([placeholder, original]) => [ + placeholder, + formatArgumentValue(original), + ]), + ), + } + : undefined; + const argumentPiiContext = formatArgumentContext(piiContext); + const argumentSecretsContext = formatArgumentContext(secretsContext); + const argumentConfig = { ...config, show_markers: false }; + + function getArgumentChannel( + choiceIndex: number, + toolCallIndex: number, + metadata: EventMetadata, + ): ToolArgumentChannel { + const positionKey = `${choiceIndex}:${toolCallIndex}`; + let channel = argumentChannelsByPosition.get(positionKey); + + if (!channel) { + channel = { + choiceIndex, + toolCallIndex, + metadata, + restorer: new StreamRestorer({ + piiContext: argumentPiiContext, + secretsContext: argumentSecretsContext, + config: argumentConfig, + }), + }; + argumentChannelsByPosition.set(positionKey, channel); + } else { + channel.metadata = metadata; + } + + return channel; + } + + function restoreToolArguments(event: Record): boolean { + let restoredArguments = false; + if (!Array.isArray(event.choices)) return restoredArguments; + + const metadata = getSyntheticEventMetadata(event); + event.choices = event.choices.map((choiceValue) => { + if ( + !isRecord(choiceValue) || + !isRecord(choiceValue.delta) || + typeof choiceValue.index !== "number" + ) { + return choiceValue; + } + + const choiceIndex = choiceValue.index; + const finishing = + choiceValue.finish_reason !== null && choiceValue.finish_reason !== undefined; + let toolCalls = Array.isArray(choiceValue.delta.tool_calls) + ? choiceValue.delta.tool_calls.map((toolCallValue) => { + if ( + !isRecord(toolCallValue) || + typeof toolCallValue.index !== "number" || + !isRecord(toolCallValue.function) || + typeof toolCallValue.function.arguments !== "string" + ) { + return toolCallValue; + } + + restoredArguments = true; + const channel = getArgumentChannel(choiceIndex, toolCallValue.index, metadata); + let restored = channel.restorer.restoreChunk(toolCallValue.function.arguments); + if (finishing) restored += channel.restorer.flush(); + + return { + ...toolCallValue, + function: { ...toolCallValue.function, arguments: restored }, + }; + }) + : undefined; + + if (finishing) { + const flushedChannels = takeFlushedArgumentChannels(choiceIndex); + if (flushedChannels.length > 0) { + restoredArguments = true; + toolCalls = mergeFlushedArguments(toolCalls ?? [], flushedChannels); + } + } + + if (!toolCalls) return choiceValue; + + return { + ...choiceValue, + delta: { ...choiceValue.delta, tool_calls: toolCalls }, + }; + }); + + return restoredArguments; + } + + function takeFlushedArgumentChannels(choiceIndex?: number): FlushedArgumentChannel[] { + const flushedChannels: FlushedArgumentChannel[] = []; + + for (const channel of argumentChannelsByPosition.values()) { + if (choiceIndex !== undefined && channel.choiceIndex !== choiceIndex) continue; + const text = channel.restorer.flush(); + if (text) flushedChannels.push({ channel, text }); + } + + return flushedChannels; + } + + function mergeFlushedArguments( + toolCalls: unknown[], + flushedChannels: FlushedArgumentChannel[], + ): unknown[] { + const merged = [...toolCalls]; + + for (const { channel, text } of flushedChannels) { + const existingIndex = merged.findIndex( + (toolCall) => isRecord(toolCall) && toolCall.index === channel.toolCallIndex, + ); + + if (existingIndex === -1) { + merged.push({ + index: channel.toolCallIndex, + function: { arguments: text }, + }); + continue; + } + + const toolCall = merged[existingIndex] as Record; + const functionCall = isRecord(toolCall.function) ? toolCall.function : {}; + const argumentsPrefix = + typeof functionCall.arguments === "string" ? functionCall.arguments : ""; + merged[existingIndex] = { + ...toolCall, + function: { ...functionCall, arguments: argumentsPrefix + text }, + }; + } + + return merged; + } return new ReadableStream({ async start(controller) { const reader = source.getReader(); + let terminated = false; + + function enqueueEvent(event: Record) { + controller.enqueue(encoder.encode(`data: ${JSON.stringify(event)}\n\n`)); + } + + function flushArgumentChannels() { + for (const { channel, text } of takeFlushedArgumentChannels()) { + enqueueEvent({ + ...channel.metadata, + choices: [ + { + index: channel.choiceIndex, + delta: { + tool_calls: [ + { + index: channel.toolCallIndex, + function: { arguments: text }, + }, + ], + }, + finish_reason: null, + }, + ], + }); + } + } + + function flushContent() { + const flushed = restorer.flush(); + + if (flushed) { + enqueueEvent({ + id: `flush-${Date.now()}`, + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + choices: [ + { + index: 0, + delta: { content: flushed }, + finish_reason: null, + }, + ], + }); + } + } function processLine(line: string) { - if (line.startsWith("data: ")) { - const data = line.slice(6); + if (terminated) return; + const dataLine = line.endsWith("\r") ? line.slice(0, -1) : line; + const data = dataLine.startsWith("data: ") + ? dataLine.slice(6) + : dataLine.startsWith("data:") + ? dataLine.slice(5) + : undefined; + if (data !== undefined) { if (data === "[DONE]") { + flushArgumentChannels(); + flushContent(); + terminated = true; controller.enqueue(encoder.encode("data: [DONE]\n\n")); return; } try { const parsed = JSON.parse(data); + const hasRestoredToolArguments = isRecord(parsed) && restoreToolArguments(parsed); const content = parsed.choices?.[0]?.delta?.content; if (typeof content === "string" && content !== "") { const text = restorer.restoreChunk(content); + parsed.choices[0].delta.content = text; - if (text) { - parsed.choices[0].delta.content = text; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); - } + if (text || hasRestoredToolArguments) enqueueEvent(parsed); } else if (Array.isArray(content)) { const processedContent = content.flatMap((part: OpenAIContentPart) => { if (part.type !== "text" || typeof part.text !== "string") { @@ -53,10 +292,10 @@ export function createUnmaskingStream( return [{ ...part, text }]; }); - if (processedContent.length > 0) { - parsed.choices[0].delta.content = processedContent; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); - } + parsed.choices[0].delta.content = processedContent; + if (processedContent.length > 0 || hasRestoredToolArguments) enqueueEvent(parsed); + } else if (hasRestoredToolArguments) { + enqueueEvent(parsed); } else { controller.enqueue(encoder.encode(`data: ${data}\n\n`)); } @@ -80,22 +319,9 @@ export function createUnmaskingStream( lineBuffer = ""; } - const flushed = restorer.flush(); - - if (flushed) { - const finalEvent = { - id: `flush-${Date.now()}`, - object: "chat.completion.chunk", - created: Math.floor(Date.now() / 1000), - choices: [ - { - index: 0, - delta: { content: flushed }, - finish_reason: null, - }, - ], - }; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(finalEvent)}\n\n`)); + if (!terminated) { + flushArgumentChannels(); + flushContent(); } controller.close(); break; @@ -107,6 +333,13 @@ export function createUnmaskingStream( for (const line of lines) { processLine(line); + if (terminated) break; + } + + if (terminated) { + await reader.cancel().catch(() => undefined); + controller.close(); + break; } } } catch (error) { diff --git a/src/providers/openai/types.ts b/src/providers/openai/types.ts index 7c2e28c..46eceb7 100644 --- a/src/providers/openai/types.ts +++ b/src/providers/openai/types.ts @@ -56,7 +56,9 @@ export const OpenAIResponseSchema = z.object({ z.object({ index: z.number(), message: OpenAIMessageSchema.passthrough(), - finish_reason: z.enum(["stop", "length", "content_filter"]).nullable(), + finish_reason: z + .enum(["stop", "length", "content_filter", "tool_calls", "function_call"]) + .nullable(), }), ), usage: z diff --git a/src/routes/openai.test.ts b/src/routes/openai.test.ts index 3cf49a8..433a5f0 100644 --- a/src/routes/openai.test.ts +++ b/src/routes/openai.test.ts @@ -1,11 +1,66 @@ -import { describe, expect, test } from "bun:test"; +import { afterEach, describe, expect, mock, test } from "bun:test"; import { Hono } from "hono"; +import { getConfig } from "../config"; +import { getLogger, Logger, normalizeRequestSource } from "../logging/logger"; +import { filterAllowlistedEntities, type PIIDetectionResult, PIIDetector } from "../pii/detect"; import { OpenAIRequestSchema } from "../providers/openai/types"; -import { openaiRoutes } from "./openai"; + +const noPII: PIIDetectionResult = { + hasPII: false, + spanEntities: [], + allEntities: [], + scanTimeMs: 0, +}; +const mockAnalyzeRequest = mock<() => Promise>(() => Promise.resolve(noPII)); +const mockLogRequest = mock(() => {}); + +mock.module("../pii/detect", () => ({ + PIIDetector, + filterAllowlistedEntities, + getPIIDetector: () => ({ + analyzeRequest: mockAnalyzeRequest, + detectPII: mock(() => Promise.resolve([])), + healthCheck: mock(() => Promise.resolve(true)), + }), +})); + +mock.module("../logging/logger", () => ({ + getLogger, + Logger, + logRequest: mockLogRequest, + normalizeRequestSource, +})); + +const { openaiRoutes } = await import("./openai"); const app = new Hono(); app.route("/openai", openaiRoutes); +const originalFetch = globalThis.fetch; +const config = getConfig(); +const originalMode = config.mode; +const originalLocal = config.local ? { ...config.local } : undefined; +const originalSecretsDetection = { + enabled: config.secrets_detection.enabled, + action: config.secrets_detection.action, + entities: [...config.secrets_detection.entities], + scan_roles: [...config.secrets_detection.scan_roles], +}; + +afterEach(() => { + globalThis.fetch = originalFetch; + config.mode = originalMode; + if (originalLocal) config.local = { ...originalLocal }; + else delete config.local; + config.secrets_detection.enabled = originalSecretsDetection.enabled; + config.secrets_detection.action = originalSecretsDetection.action; + config.secrets_detection.entities = [...originalSecretsDetection.entities]; + config.secrets_detection.scan_roles = [...originalSecretsDetection.scan_roles]; + mockAnalyzeRequest.mockClear(); + mockAnalyzeRequest.mockResolvedValue(noPII); + mockLogRequest.mockClear(); +}); + describe("POST /openai/v1/chat/completions", () => { test("returns 400 for missing messages", async () => { const res = await app.request("/openai/v1/chat/completions", { @@ -43,6 +98,345 @@ describe("POST /openai/v1/chat/completions", () => { expect(res.status).toBe(400); }); + + test("remasks restored PII in assistant tool-call history before forwarding", async () => { + config.mode = "mask"; + const email = "jane@example.com"; + const userContent = `Email ${email}`; + const emailStart = userContent.indexOf(email); + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [ + [ + { + entity_type: "EMAIL_ADDRESS", + start: emailStart, + end: emailStart + email.length, + score: 0.99, + }, + ], + [], + ], + allEntities: [ + { + entity_type: "EMAIL_ADDRESS", + start: emailStart, + end: emailStart + email.length, + score: 0.99, + }, + ], + scanTimeMs: 2, + }); + + let upstreamBody: Record | undefined; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ + id: "chatcmpl_test", + object: "chat.completion", + created: 123, + model: "gpt-test", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Done" }, + finish_reason: "stop", + }, + ], + }); + }) as typeof fetch; + + const response = await app.request("/openai/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "gpt-test", + messages: [ + { role: "user", content: userContent }, + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_123", + type: "function", + function: { + name: "lookup_email", + arguments: JSON.stringify({ email }), + }, + }, + ], + }, + { role: "tool", tool_call_id: "call_123", content: "Lookup complete" }, + ], + }), + }); + + expect(response.status).toBe(200); + const messages = upstreamBody?.messages as Array>; + expect(messages[0].content).toBe("Email [[EMAIL_ADDRESS_1]]"); + expect( + (messages[1].tool_calls as Array<{ function: { arguments: string } }>)[0].function.arguments, + ).toBe('{"email":"[[EMAIL_ADDRESS_1]]"}'); + expect(JSON.stringify(upstreamBody)).not.toContain(email); + }); + + test("remasks tool-call history and restores a streaming Chat response", async () => { + config.mode = "mask"; + const email = "stream@example.com"; + const userContent = `Email ${email}`; + const emailStart = userContent.indexOf(email); + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [ + [ + { + entity_type: "EMAIL_ADDRESS", + start: emailStart, + end: emailStart + email.length, + score: 0.99, + }, + ], + [], + ], + allEntities: [ + { + entity_type: "EMAIL_ADDRESS", + start: emailStart, + end: emailStart + email.length, + score: 0.99, + }, + ], + scanTimeMs: 2, + }); + + let upstreamBody: Record | undefined; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return new Response( + `data: ${JSON.stringify({ + choices: [ + { + index: 0, + delta: { content: "Email [[EMAIL_ADDRESS_1]]" }, + finish_reason: null, + }, + ], + })}\n\ndata: [DONE]\n\n`, + { headers: { "Content-Type": "text/event-stream" } }, + ); + }) as typeof fetch; + + const response = await app.request("/openai/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "gpt-test", + stream: true, + messages: [ + { role: "user", content: userContent }, + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_stream", + type: "function", + function: { + name: "lookup_email", + arguments: JSON.stringify({ email }), + }, + }, + ], + }, + { role: "tool", tool_call_id: "call_stream", content: "Lookup complete" }, + ], + }), + }); + + const messages = upstreamBody?.messages as Array>; + expect(response.status).toBe(200); + expect(response.headers.get("Content-Type")).toContain("text/event-stream"); + expect(upstreamBody?.stream).toBe(true); + expect( + (messages[1].tool_calls as Array<{ function: { arguments: string } }>)[0].function.arguments, + ).toBe('{"email":"[[EMAIL_ADDRESS_1]]"}'); + const responseBody = await response.text(); + expect(responseBody).toContain(email); + expect(responseBody).not.toContain("[[EMAIL_ADDRESS_1]]"); + }); + + test("remasks restored secrets in assistant tool-call history before forwarding", async () => { + config.mode = "mask"; + config.secrets_detection.enabled = true; + config.secrets_detection.action = "mask"; + config.secrets_detection.entities = ["API_KEY_SK"]; + config.secrets_detection.scan_roles = ["user", "tool", "function", "mcp"]; + const secret = "sk-proj-abc123def456ghi789jkl012mno345pqr678stu901vwx"; + + let upstreamBody: Record | undefined; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ + id: "chatcmpl_secret", + object: "chat.completion", + created: 123, + model: "gpt-test", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Done" }, + finish_reason: "stop", + }, + ], + }); + }) as typeof fetch; + + const response = await app.request("/openai/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "gpt-test", + messages: [ + { role: "user", content: `Key ${secret}` }, + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_secret", + type: "function", + function: { + name: "store_key", + arguments: JSON.stringify({ key: secret }), + }, + }, + ], + }, + { role: "tool", tool_call_id: "call_secret", content: "Stored" }, + ], + }), + }); + + expect(response.status).toBe(200); + const messages = upstreamBody?.messages as Array>; + expect(messages[0].content).toBe("Key [[API_KEY_SK_1]]"); + expect( + (messages[1].tool_calls as Array<{ function: { arguments: string } }>)[0].function.arguments, + ).toBe('{"key":"[[API_KEY_SK_1]]"}'); + expect(JSON.stringify(upstreamBody)).not.toContain(secret); + }); + + test("keeps PII clear for local tool history while preserving secrets masking", async () => { + config.mode = "route"; + config.local = { + type: "openai", + base_url: "http://local.test/v1", + model: "local-model", + }; + config.secrets_detection.enabled = true; + config.secrets_detection.action = "mask"; + config.secrets_detection.entities = ["API_KEY_SK"]; + config.secrets_detection.scan_roles = ["user", "tool", "function", "mcp"]; + const email = "local@example.com"; + const secret = "sk-proj-localtest123456789012345678901234567890123456789"; + const userContent = `Email ${email} Key ${secret}`; + const emailStart = userContent.indexOf(email); + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [ + [ + { + entity_type: "EMAIL_ADDRESS", + start: emailStart, + end: emailStart + email.length, + score: 0.99, + }, + ], + [], + ], + allEntities: [ + { + entity_type: "EMAIL_ADDRESS", + start: emailStart, + end: emailStart + email.length, + score: 0.99, + }, + ], + scanTimeMs: 2, + }); + + let upstreamUrl: string | undefined; + let upstreamBody: Record | undefined; + globalThis.fetch = (async (target: string | URL | Request, init?: RequestInit) => { + upstreamUrl = target instanceof Request ? target.url : String(target); + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ + id: "chatcmpl_local", + object: "chat.completion", + created: 123, + model: "local-model", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Done" }, + finish_reason: "stop", + }, + ], + }); + }) as typeof fetch; + + const modeDescriptor = Object.getOwnPropertyDescriptor(config, "mode"); + if (!modeDescriptor) throw new Error("Expected mode property descriptor"); + let modeReads = 0; + // Let the real pipeline supply a PII context, then exercise the local destination policy. + Object.defineProperty(config, "mode", { + configurable: true, + enumerable: modeDescriptor.enumerable, + get: () => (modeReads++ === 0 ? "mask" : "route"), + }); + + let response: Response; + try { + response = await app.request("/openai/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "gpt-test", + messages: [ + { role: "user", content: userContent }, + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_local", + type: "function", + function: { + name: "store_contact", + arguments: JSON.stringify({ email, key: secret }), + }, + }, + ], + }, + { role: "tool", tool_call_id: "call_local", content: "Stored" }, + ], + }), + }); + } finally { + Object.defineProperty(config, "mode", modeDescriptor); + } + + expect(response.status).toBe(200); + expect(response.headers.get("X-PasteGuard-Provider")).toBe("local"); + expect(upstreamUrl).toBe("http://local.test/v1/chat/completions"); + const messages = upstreamBody?.messages as Array>; + expect(messages[0].content).toBe(`Email ${email} Key [[API_KEY_SK_1]]`); + expect( + (messages[1].tool_calls as Array<{ function: { arguments: string } }>)[0].function.arguments, + ).toBe(`{"email":"${email}","key":"[[API_KEY_SK_1]]"}`); + expect(JSON.stringify(upstreamBody)).not.toContain(secret); + }); }); describe("Zod schema preserves unknown fields", () => { diff --git a/src/routes/openai.ts b/src/routes/openai.ts index 43a3efa..9017e35 100644 --- a/src/routes/openai.ts +++ b/src/routes/openai.ts @@ -6,7 +6,7 @@ import { getConfig, type MaskingConfig } from "../config"; import { formatMaskedRequestForLog } from "../logging/log-content"; import { logRequest } from "../logging/logger"; import type { PlaceholderContext } from "../masking/context"; -import { openaiExtractor } from "../masking/extractors/openai"; +import { openaiExtractor, remaskOpenAIToolCallArguments } from "../masking/extractors/openai"; import { restoreResponse } from "../masking/restorer"; import type { PIIDetectResult } from "../pii/request"; import { @@ -78,34 +78,32 @@ openaiRoutes.post( throw new Error("PII detection result missing from privacy pipeline"); } - if (config.mode === "mask") { - return sendToOpenAI(c, request, { - request: privacy.request, - piiResult, - piiMaskingContext: privacy.piiMaskingContext, - secretsResult, - startTime, - authHeader: c.req.header("Authorization"), - }); - } - // Route mode: send to local if PII/secrets detected, otherwise OpenAI const shouldRouteLocal = - piiResult.hasPII || - (secretsResult.detection?.detected && config.secrets_detection.action === "route_local"); + config.mode === "route" && + (piiResult.hasPII || + (secretsResult.detection?.detected && config.secrets_detection.action === "route_local")); if (shouldRouteLocal) { + const localRequest = remaskOpenAIToolCallArguments(privacy.requestAfterSecrets, [ + secretsResult.maskingContext, + ]); return sendToLocal(c, request, { - request: privacy.requestAfterSecrets, + request: localRequest, piiResult, secretsResult, startTime, }); } + const upstreamRequest = remaskOpenAIToolCallArguments( + config.mode === "mask" ? privacy.request : privacy.requestAfterSecrets, + [secretsResult.maskingContext, privacy.piiMaskingContext], + ); return sendToOpenAI(c, request, { - request: privacy.requestAfterSecrets, + request: upstreamRequest, piiResult, + piiMaskingContext: privacy.piiMaskingContext, secretsResult, startTime, authHeader: c.req.header("Authorization"),