From: Stefan Gasser Date: Tue, 25 Aug 2026 09:27:27 +0000 (+0200) Subject: fix(responses): restore tool arguments safely (#170) X-Git-Tag: v0.9.3~1 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=c0b660b533780de251cee601c16e21a7e46b9850;p=sgasser-llm-shield.git fix(responses): restore tool arguments safely (#170) --- diff --git a/src/masking/extractors/responses.ts b/src/masking/extractors/responses.ts index f9df568..3c5478d 100644 --- a/src/masking/extractors/responses.ts +++ b/src/masking/extractors/responses.ts @@ -110,6 +110,51 @@ function pathFromString(path: string): Array { return result; } +function getAtPath(value: unknown, path: readonly (string | number)[]): unknown { + let current = value; + for (const part of path) { + if (Array.isArray(current) && typeof part === "number") { + current = current[part]; + } else if (isRecord(current) && typeof part === "string") { + current = current[part]; + } else { + return undefined; + } + } + return current; +} + +export function isResponsesFunctionCallArguments( + response: ResponsesResponse, + path: string | readonly (string | number)[], +): boolean { + const parts = typeof path === "string" ? pathFromString(path) : path; + if (parts.at(-1) !== "arguments") return false; + + const parent = getAtPath(response, parts.slice(0, -1)); + return ( + isRecord(parent) && + (parent.type === "function_call" || parent.type === "response.function_call_arguments.done") + ); +} + +export function restoreSerializedFunctionCallArguments( + text: string, + context: PlaceholderContext, + formatValue?: (original: string) => string, +): string { + try { + JSON.parse(text); + } catch { + return text; + } + + return restorePlaceholders(text, context, (original) => { + const restored = formatValue ? formatValue(original) : original; + return JSON.stringify(restored).slice(1, -1); + }); +} + export const responsesExtractor: RequestExtractor = { extractTexts(request: ResponsesRequest): TextSpan[] { return collectText(request).map((item, index) => ({ @@ -135,7 +180,10 @@ export const responsesExtractor: RequestExtractor { + const person = 'Jane "JJ" \\vault\nline\t\u0001 café 😀 [[marker-like]]'; + const secret = 'sk-"quoted"\\path\r\nnext'; + const response: ResponsesResponse = { + output: [ + { + type: "function_call", + arguments: JSON.stringify({ + nested: { person: "[[PERSON_1]]" }, + values: ["before", "[[API_KEY_SK_1]]"], + }), + }, + ], + }; + + const result = restoreResponse(response, responsesExtractor, markerConfig, { + piiContext: context({ "[[PERSON_1]]": person }), + secretsContext: context({ "[[API_KEY_SK_1]]": secret }), + }); + const output = result.output as Array<{ arguments: string }>; + + expect(JSON.parse(output[0].arguments)).toEqual({ + nested: { person: `[protected]${person}` }, + values: ["before", `[protected]${secret}`], + }); + }); + + test("preserves malformed serialized function-call arguments", () => { + const argumentsText = '{"person":"[[PERSON_1]]"'; + const response: ResponsesResponse = { + output: [{ type: "function_call", arguments: argumentsText }], + }; + + const result = restoreResponse(response, responsesExtractor, defaultConfig, { + piiContext: context({ "[[PERSON_1]]": 'Jane "JJ"' }), + }); + const output = result.output as Array<{ arguments: string }>; + + expect(output[0].arguments).toBe(argumentsText); + }); + + test("does not treat arbitrary arguments fields as serialized function JSON", () => { + const response: ResponsesResponse = { + output: [{ type: "message", arguments: "Contact [[PERSON_1]]" }], + }; + + expect( + restoreResponse(response, responsesExtractor, defaultConfig, { + piiContext: context({ "[[PERSON_1]]": "Jane" }), + }), + ).toEqual({ output: [{ type: "message", arguments: "Contact Jane" }] }); + }); }); diff --git a/src/masking/stream-restorer.ts b/src/masking/stream-restorer.ts index e9297ef..345b4bd 100644 --- a/src/masking/stream-restorer.ts +++ b/src/masking/stream-restorer.ts @@ -7,6 +7,7 @@ export interface StreamRestorerOptions { piiContext?: PlaceholderContext; secretsContext?: PlaceholderContext; config: MaskingConfig; + formatValue?: (original: string) => string; } export class StreamRestorer { @@ -15,7 +16,7 @@ export class StreamRestorer { private readonly formatValue: ((original: string) => string) | undefined; constructor(private readonly options: StreamRestorerOptions) { - this.formatValue = createRestoreFormatter(options.config); + this.formatValue = options.formatValue ?? createRestoreFormatter(options.config); } restoreChunk(text: string): string { @@ -46,6 +47,10 @@ export class StreamRestorer { return processedText; } + hasPending(): boolean { + return Boolean(this.piiBuffer || this.secretsBuffer); + } + flush(): string { let flushed = ""; diff --git a/src/protocols/responses/stream-transformer.test.ts b/src/protocols/responses/stream-transformer.test.ts index 150a931..52b5f50 100644 --- a/src/protocols/responses/stream-transformer.test.ts +++ b/src/protocols/responses/stream-transformer.test.ts @@ -46,8 +46,58 @@ async function consumeStream(stream: ReadableStream): Promise): string { + return `data: ${JSON.stringify(event)}\n\n`; +} + function codexDelta(text: string): string { - return `data: ${JSON.stringify({ type: "response.output_text.delta", delta: text })}\n\n`; + return dataEvent({ type: "response.output_text.delta", delta: text }); +} + +function functionCallDelta( + itemId: string, + outputIndex: number, + delta: string, + metadata: Record = {}, +): string { + return dataEvent({ + type: "response.function_call_arguments.delta", + item_id: itemId, + output_index: outputIndex, + delta, + ...metadata, + }); +} + +type DataEvent = Record; +type DataTimelineEntry = DataEvent | "[DONE]"; + +function dataTimeline(stream: string): DataTimelineEntry[] { + return stream + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => (line === "data: [DONE]" ? "[DONE]" : JSON.parse(line.slice(6)))); +} + +function dataEvents(stream: string): DataEvent[] { + return dataTimeline(stream).filter((event): event is DataEvent => event !== "[DONE]"); +} + +function metadataOccurrences(timeline: DataTimelineEntry[], trace: string): number { + return timeline.filter( + (event) => + event !== "[DONE]" && (event.metadata as { trace?: string } | undefined)?.trace === trace, + ).length; +} + +function streamedArguments(events: Array>, itemId: string): string { + return events + .filter( + (event) => + event.type === "response.function_call_arguments.delta" && event.item_id === itemId, + ) + .map((event) => event.delta) + .join(""); } describe("createResponsesUnmaskingStream", () => { @@ -91,6 +141,153 @@ describe("createResponsesUnmaskingStream", () => { expect(result).toContain("[protected]Jane used [protected]sk-secret"); }); + test("restores fragmented function arguments with JSON-safe escaping", async () => { + const person = 'Jane "JJ" \\vault\nline\t\u0001 café 😀 [[marker-like]]'; + const secret = 'sk-"quoted"\\path\r\nnext'; + const piiContext = context({ "[[PERSON_1]]": person }); + const secretsContext = context({ "[[API_KEY_SK_1]]": secret }); + const source = createSSEStream([ + functionCallDelta("fc_1", 2, '{"nested":{"person":"[[PER', { + sequence_number: 10, + trace: "keep-me", + }), + functionCallDelta("fc_1", 2, 'SON_1]]"},"values":["[[API_KEY_', { sequence_number: 11 }), + functionCallDelta("fc_1", 2, 'SK_1]]"]}', { sequence_number: 12 }), + ]); + + const result = await consumeStream( + createResponsesUnmaskingStream( + source, + piiContext, + { ...defaultConfig, show_markers: true }, + secretsContext, + ), + ); + const events = dataEvents(result); + + expect(events.map(({ delta: _delta, ...event }) => event)).toEqual([ + { + type: "response.function_call_arguments.delta", + item_id: "fc_1", + output_index: 2, + sequence_number: 10, + trace: "keep-me", + }, + { + type: "response.function_call_arguments.delta", + item_id: "fc_1", + output_index: 2, + sequence_number: 11, + }, + { + type: "response.function_call_arguments.delta", + item_id: "fc_1", + output_index: 2, + sequence_number: 12, + }, + ]); + expect(JSON.parse(streamedArguments(events, "fc_1"))).toEqual({ + nested: { person: `[protected]${person}` }, + values: [`[protected]${secret}`], + }); + }); + + test("keeps interleaved function-call restoration state independent", async () => { + const piiContext = context({ + "[[EMAIL_ADDRESS_1]]": "one@example.com", + "[[EMAIL_ADDRESS_2]]": "two@example.com", + }); + const source = createSSEStream([ + functionCallDelta("fc_a", 0, '{"email":"[[EMAIL_'), + functionCallDelta("fc_b", 1, '{"email":"[[EMAIL_ADDRESS_2]]"}'), + functionCallDelta("fc_a", 0, 'ADDRESS_1]]"}'), + ]); + + const result = await consumeStream( + createResponsesUnmaskingStream(source, piiContext, defaultConfig), + ); + const events = dataEvents(result); + + expect(JSON.parse(streamedArguments(events, "fc_a"))).toEqual({ + email: "one@example.com", + }); + expect(JSON.parse(streamedArguments(events, "fc_b"))).toEqual({ + email: "two@example.com", + }); + expect(events.map((event) => [event.item_id, event.output_index])).toEqual([ + ["fc_a", 0], + ["fc_b", 1], + ["fc_a", 0], + ]); + }); + + test("restores finalized function arguments in protocol completion events", async () => { + const piiContext = context({ "[[PERSON_1]]": 'Jane "JJ"' }); + const serialized = JSON.stringify({ person: "[[PERSON_1]]" }); + const source = createSSEStream([ + dataEvent({ + type: "response.function_call_arguments.done", + item_id: "fc_done", + output_index: 0, + name: "lookup", + arguments: serialized, + sequence_number: 1, + metadata: { trace: "keep-done-metadata" }, + }), + dataEvent({ + type: "response.output_item.done", + output_index: 1, + item: { + id: "fc_item", + type: "function_call", + call_id: "call_item", + name: "lookup", + arguments: serialized, + status: "completed", + }, + sequence_number: 2, + }), + dataEvent({ + type: "response.completed", + response: { + id: "resp_1", + output: [ + { + id: "fc_completed", + type: "function_call", + call_id: "call_completed", + name: "lookup", + arguments: serialized, + }, + ], + }, + sequence_number: 3, + }), + ]); + + const events = dataEvents( + await consumeStream(createResponsesUnmaskingStream(source, piiContext, defaultConfig)), + ); + + const { arguments: doneArguments, ...doneEvent } = events[0]; + expect(doneEvent).toEqual({ + type: "response.function_call_arguments.done", + item_id: "fc_done", + output_index: 0, + name: "lookup", + sequence_number: 1, + metadata: { trace: "keep-done-metadata" }, + }); + expect(JSON.parse(doneArguments as string)).toEqual({ person: 'Jane "JJ"' }); + expect(JSON.parse((events[1].item as { arguments: string }).arguments)).toEqual({ + person: 'Jane "JJ"', + }); + const completedCall = (events[2].response as { output: Array<{ arguments: string }> }) + .output[0]; + expect(JSON.parse(completedCall.arguments)).toEqual({ person: 'Jane "JJ"' }); + expect(events.map((event) => event.sequence_number)).toEqual([1, 2, 3]); + }); + test("passes malformed JSON and done events through", async () => { const source = createSSEStream(["data: not-json\n\n", "data: [DONE]\n\n"]); @@ -102,15 +299,287 @@ describe("createResponsesUnmaskingStream", () => { expect(result).toContain("data: [DONE]"); }); - test("emits Codex-compatible final flush events", async () => { + test("flushes output text without replaying upstream event metadata", async () => { const piiContext = context({ "[[EMAIL_ADDRESS_1]]": "jane@example.com" }); - const source = createSSEStream([codexDelta("Email [[EMAIL")]); + const source = createSSEStream([ + dataEvent({ + type: "response.output_text.delta", + item_id: "msg_1", + output_index: 3, + content_index: 0, + delta: "Email [[EMAIL", + sequence_number: 7, + metadata: { trace: "text-once" }, + }), + "data: [DONE]\n\n", + ]); const result = await consumeStream( createResponsesUnmaskingStream(source, piiContext, defaultConfig), ); + const timeline = dataTimeline(result); + + expect(timeline).toEqual([ + { + type: "response.output_text.delta", + item_id: "msg_1", + output_index: 3, + content_index: 0, + delta: "Email ", + sequence_number: 7, + metadata: { trace: "text-once" }, + }, + { type: "response.output_text.delta", delta: "[[EMAIL" }, + "[DONE]", + ]); + expect(metadataOccurrences(timeline, "text-once")).toBe(1); + expect( + timeline.filter((event) => event !== "[DONE]" && event.sequence_number === 7), + ).toHaveLength(1); + }); + + test("flushes pending data and closes promptly when the source stays open after DONE", async () => { + const piiContext = context({ "[[PERSON_1]]": "Jane" }); + const encoder = new TextEncoder(); + let sourceCancelled = false; + let closeSource = () => {}; + const sourcePayload = [ + dataEvent({ + type: "response.output_text.delta", + delta: "[", + item_id: "msg_1", + output_index: 0, + sequence_number: 60, + metadata: { trace: "text-once" }, + }), + functionCallDelta("fc_1", 1, '{"name":"[[PERSON_1]]","tail":"[', { + sequence_number: 70, + metadata: { trace: "argument-once" }, + }), + "data: [DONE]\n\n", + "data: [DONE]\n\n", + dataEvent({ + type: "response.completed", + response: { id: "resp_late", output: [] }, + sequence_number: 80, + }), + ].join(""); + const source = new ReadableStream({ + start(controller) { + closeSource = () => controller.close(); + controller.enqueue(encoder.encode(sourcePayload)); + }, + cancel() { + sourceCancelled = true; + }, + }); + + let result: string | undefined; + const completion = consumeStream( + createResponsesUnmaskingStream(source, piiContext, defaultConfig), + ).then((value) => { + result = value; + }); + await Promise.race([completion, Bun.sleep(250)]); + const terminalReadCompleted = result !== undefined; + if (!terminalReadCompleted) { + closeSource(); + await completion; + } + + expect(terminalReadCompleted).toBe(true); + const timeline = dataTimeline(result!); + + expect(timeline).toEqual([ + { + type: "response.output_text.delta", + delta: "", + item_id: "msg_1", + output_index: 0, + sequence_number: 60, + metadata: { trace: "text-once" }, + }, + { + type: "response.function_call_arguments.delta", + item_id: "fc_1", + output_index: 1, + delta: '{"name":"Jane","tail":"[', + sequence_number: 70, + metadata: { trace: "argument-once" }, + }, + { type: "response.output_text.delta", delta: "[" }, + "[DONE]", + ]); + expect(sourceCancelled).toBe(true); + expect(timeline.map((event) => (event === "[DONE]" ? event : event.sequence_number))).toEqual([ + 60, + 70, + undefined, + "[DONE]", + ]); + expect(metadataOccurrences(timeline, "text-once")).toBe(1); + expect(metadataOccurrences(timeline, "argument-once")).toBe(1); + expect( + timeline + .filter( + (event): event is DataEvent => + event !== "[DONE]" && event.type === "response.output_text.delta", + ) + .map((event) => event.delta) + .join(""), + ).toBe("["); + expect(timeline.at(-1)).toBe("[DONE]"); + }); + + test("keeps terminal output when source cancellation rejects", async () => { + const piiContext = context({ "[[PERSON_1]]": "Jane" }); + const encoder = new TextEncoder(); + let cancelAttempts = 0; + const sourcePayload = [ + dataEvent({ + type: "response.output_text.delta", + delta: "[", + item_id: "msg_cancel_error", + output_index: 0, + sequence_number: 90, + }), + functionCallDelta("fc_cancel_error", 1, '{"name":"[[PERSON_1]]","tail":"[', { + sequence_number: 100, + }), + "data: [DONE]\n\n", + "data: [DONE]\n\n", + dataEvent({ + type: "response.completed", + response: { id: "resp_late", output: [] }, + sequence_number: 110, + }), + ].join(""); + const source = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(sourcePayload)); + }, + cancel() { + cancelAttempts++; + throw new Error("cancel failed"); + }, + }); + const transformed = createResponsesUnmaskingStream(source, piiContext, defaultConfig); + + await Bun.sleep(20); + const timeline = dataTimeline(await consumeStream(transformed)); + + expect(timeline).toEqual([ + { + type: "response.output_text.delta", + delta: "", + item_id: "msg_cancel_error", + output_index: 0, + sequence_number: 90, + }, + { + type: "response.function_call_arguments.delta", + item_id: "fc_cancel_error", + output_index: 1, + delta: '{"name":"Jane","tail":"[', + sequence_number: 100, + }, + { type: "response.output_text.delta", delta: "[" }, + "[DONE]", + ]); + expect(cancelAttempts).toBe(1); + expect(timeline.at(-1)).toBe("[DONE]"); + }); + + test("flushes interleaved unfinished calls once at stream end", async () => { + const piiContext = context({ "[[PERSON_1]]": "Jane" }); + const source = createSSEStream([ + functionCallDelta("fc_a", 0, '{"name":"[[PER', { + sequence_number: 4, + metadata: { trace: "stream-end-a-once" }, + }), + functionCallDelta("fc_b", 1, '{"name":"[[PERSON', { + sequence_number: 5, + metadata: { trace: "stream-end-b-once" }, + }), + ]); + + const timeline = dataTimeline( + await consumeStream(createResponsesUnmaskingStream(source, piiContext, defaultConfig)), + ); + + expect( + timeline.map((event) => + event === "[DONE]" + ? event + : [event.sequence_number, event.item_id, event.output_index, event.delta], + ), + ).toEqual([ + [4, "fc_a", 0, '{"name":"[[PER'], + [5, "fc_b", 1, '{"name":"[[PERSON'], + ]); + for (const trace of ["stream-end-a-once", "stream-end-b-once"]) { + expect(metadataOccurrences(timeline, trace)).toBe(1); + } + }); + + test("rewrites a buffered argument delta once before its done event", async () => { + const piiContext = context({ "[[PERSON_1]]": "Jane" }); + const source = createSSEStream([ + functionCallDelta("fc_done", 3, '{"name":"[[PER', { + sequence_number: 40, + metadata: { trace: "delta-once" }, + }), + dataEvent({ + type: "response.function_call_arguments.done", + item_id: "fc_done", + output_index: 3, + name: "save_name", + arguments: JSON.stringify({ name: "[[PERSON_1]]" }), + sequence_number: 41, + metadata: { trace: "done-once" }, + }), + "data: [DONE]\n\n", + ]); + + const result = await consumeStream( + createResponsesUnmaskingStream(source, piiContext, defaultConfig), + ); + const timeline = dataTimeline(result); + + expect(timeline.map((event) => (event === "[DONE]" ? event : event.sequence_number))).toEqual([ + 40, + 41, + "[DONE]", + ]); + expect(metadataOccurrences(timeline, "delta-once")).toBe(1); + expect(metadataOccurrences(timeline, "done-once")).toBe(1); + expect((timeline[0] as Record).delta).toBe('{"name":"[[PER'); + expect(JSON.parse((timeline[1] as Record).arguments as string)).toEqual({ + name: "Jane", + }); + expect(timeline.at(-1)).toBe("[DONE]"); + }); + + test("rewrites a buffered argument delta once before stream end", async () => { + const piiContext = context({ "[[PERSON_1]]": "Jane" }); + const source = createSSEStream([ + functionCallDelta("fc_end", 4, '{"name":"[[PERSON', { + sequence_number: 50, + metadata: { trace: "stream-end-once" }, + }), + "data: [DONE]\n\n", + ]); + + const timeline = dataTimeline( + await consumeStream(createResponsesUnmaskingStream(source, piiContext, defaultConfig)), + ); - expect(result).toContain('"type":"response.output_text.delta"'); - expect(result).toContain('"delta":"[[EMAIL"'); + expect(timeline.map((event) => (event === "[DONE]" ? event : event.sequence_number))).toEqual([ + 50, + "[DONE]", + ]); + expect(metadataOccurrences(timeline, "stream-end-once")).toBe(1); + expect((timeline[0] as Record).delta).toBe('{"name":"[[PERSON'); + expect(timeline.at(-1)).toBe("[DONE]"); }); }); diff --git a/src/protocols/responses/stream-transformer.ts b/src/protocols/responses/stream-transformer.ts index 6c3d301..cff0c6a 100644 --- a/src/protocols/responses/stream-transformer.ts +++ b/src/protocols/responses/stream-transformer.ts @@ -1,8 +1,30 @@ import type { MaskingConfig } from "../../config"; import type { PlaceholderContext } from "../../masking/context"; -import { type ResponsesResponse, responsesExtractor } from "../../masking/extractors/responses"; +import { + isResponsesFunctionCallArguments, + type ResponsesResponse, + responsesExtractor, + restoreSerializedFunctionCallArguments, +} from "../../masking/extractors/responses"; +import { createRestoreFormatter } from "../../masking/restore-policy"; import { StreamRestorer } from "../../masking/stream-restorer"; +interface ArgumentStreamState { + restorer: StreamRestorer; + pendingFrames: EventTimelineFrame[]; +} + +interface EventTimelineFrame { + event: Record; + ready: boolean; +} + +type TimelineFrame = EventTimelineFrame | { line: string; ready: boolean }; + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + export function createResponsesUnmaskingStream( stream: ReadableStream, piiContext: PlaceholderContext | undefined, @@ -12,13 +34,138 @@ export function createResponsesUnmaskingStream( const decoder = new TextDecoder(); const encoder = new TextEncoder(); let lineBuffer = ""; - const restorer = new StreamRestorer({ + let streamTerminated = false; + const textRestorer = new StreamRestorer({ piiContext, secretsContext, config: maskingConfig, }); + const pendingArguments = new Map(); + const timeline: TimelineFrame[] = []; + const restoreFormatter = createRestoreFormatter(maskingConfig); + const jsonStringFormatter = (original: string) => { + const restored = restoreFormatter ? restoreFormatter(original) : original; + return JSON.stringify(restored).slice(1, -1); + }; + + function argumentKey(payload: Record): string | undefined { + return typeof payload.item_id === "string" && typeof payload.output_index === "number" + ? JSON.stringify([payload.item_id, payload.output_index]) + : undefined; + } + + function createArgumentRestorer(): StreamRestorer { + return new StreamRestorer({ + piiContext, + secretsContext, + config: maskingConfig, + formatValue: jsonStringFormatter, + }); + } + + function restoreCompleteArguments(text: string): string { + let result = text; + if (piiContext) { + result = restoreSerializedFunctionCallArguments(result, piiContext, restoreFormatter); + } + if (secretsContext) { + result = restoreSerializedFunctionCallArguments(result, secretsContext, restoreFormatter); + } + return result; + } + + function releaseArgument(key: string | undefined): void { + if (!key) return; + + const state = pendingArguments.get(key); + if (!state) return; + + pendingArguments.delete(key); + const delta = state.restorer.flush(); + const lastFrame = state.pendingFrames.at(-1); + if (delta && lastFrame) { + lastFrame.event.delta = `${String(lastFrame.event.delta)}${delta}`; + } + for (const frame of state.pendingFrames) frame.ready = true; + } + + function releaseAllArgumentFrames(): void { + for (const key of pendingArguments.keys()) releaseArgument(key); + } - function unmaskPayload(payload: unknown): unknown { + function releaseArgumentsBefore(payload: Record): void { + if ( + payload.type === "response.completed" || + payload.type === "response.incomplete" || + payload.type === "response.failed" + ) { + releaseAllArgumentFrames(); + } else if (payload.type === "response.function_call_arguments.done") { + releaseArgument(argumentKey(payload)); + } else if ( + payload.type === "response.output_item.done" && + isRecord(payload.item) && + payload.item.type === "function_call" + ) { + releaseArgument( + argumentKey({ item_id: payload.item.id, output_index: payload.output_index }), + ); + } + } + + function enqueueEvent(event: Record, ready = true): EventTimelineFrame { + const frame = { event, ready }; + timeline.push(frame); + return frame; + } + + function enqueueLine(line: string): void { + timeline.push({ line, ready: true }); + } + + function enqueueTextFlush(): void { + const delta = textRestorer.flush(); + if (!delta) return; + + enqueueEvent({ type: "response.output_text.delta", delta }); + enqueueLine(""); + } + + function drainTimeline(): string { + let output = ""; + while (timeline[0]?.ready) { + const frame = timeline.shift()!; + output += `${"event" in frame ? `data: ${JSON.stringify(frame.event)}` : frame.line}\n`; + } + return output; + } + + function enqueueArgumentDelta(payload: Record): void { + const key = argumentKey(payload); + if (!key || typeof payload.delta !== "string") { + enqueueEvent(payload); + return; + } + + const state = pendingArguments.get(key) ?? { + restorer: createArgumentRestorer(), + pendingFrames: [], + }; + const delta = state.restorer.restoreChunk(payload.delta); + const hasPending = state.restorer.hasPending(); + const frame = enqueueEvent({ ...payload, delta }, !hasPending); + + if (hasPending) { + state.pendingFrames.push(frame); + pendingArguments.set(key, state); + } else { + for (const pendingFrame of state.pendingFrames) pendingFrame.ready = true; + pendingArguments.delete(key); + } + } + + function unmaskPayload(payload: Record): Record { + releaseArgumentsBefore(payload); const result = payload as ResponsesResponse; const spans = responsesExtractor.extractTexts(result); @@ -30,26 +177,45 @@ export function createResponsesUnmaskingStream( result, spans.map((span) => ({ ...span, - maskedText: restorer.restoreChunk(span.text), + maskedText: isResponsesFunctionCallArguments(result, span.path) + ? restoreCompleteArguments(span.text) + : textRestorer.restoreChunk(span.text), })), ); } function processLine(line: string): string { + if (streamTerminated) return ""; + if (!line.startsWith("data: ")) { - return `${line}\n`; + enqueueLine(line); + return drainTimeline(); } const data = line.slice(6); if (data === "[DONE]") { - return "data: [DONE]\n"; + streamTerminated = true; + releaseAllArgumentFrames(); + enqueueTextFlush(); + enqueueLine("data: [DONE]"); + enqueueLine(""); + return drainTimeline(); } try { - return `data: ${JSON.stringify(unmaskPayload(JSON.parse(data)))}\n`; + const payload = JSON.parse(data) as unknown; + if (!isRecord(payload)) { + enqueueEvent(payload as Record); + } else if (payload.type === "response.function_call_arguments.delta") { + enqueueArgumentDelta(payload); + } else { + enqueueEvent(unmaskPayload(payload)); + } } catch { - return `${line}\n`; + enqueueLine(line); } + + return drainTimeline(); } return new ReadableStream({ @@ -68,24 +234,30 @@ export function createResponsesUnmaskingStream( let output = ""; for (const line of lines) { output += processLine(line); + if (streamTerminated) break; } if (output) { controller.enqueue(encoder.encode(output)); } + + if (streamTerminated) { + try { + await reader.cancel(); + } catch { + // Cancellation cannot replace an already completed terminal response. + } + break; + } } lineBuffer += decoder.decode(); let finalOutput = lineBuffer ? processLine(lineBuffer) : ""; lineBuffer = ""; - const flushed = restorer.flush(); - if (flushed) { - finalOutput += `data: ${JSON.stringify({ - type: "response.output_text.delta", - delta: flushed, - })}\n\n`; - } + releaseAllArgumentFrames(); + enqueueTextFlush(); + finalOutput += drainTimeline(); if (finalOutput) { controller.enqueue(encoder.encode(finalOutput)); diff --git a/src/routes/codex.test.ts b/src/routes/codex.test.ts index 1f5d428..e760808 100644 --- a/src/routes/codex.test.ts +++ b/src/routes/codex.test.ts @@ -316,6 +316,65 @@ describe("Codex proxy", () => { expect(parsed.delta).toBe(`Key ${secret}`); }); + test("unmasks streamed function-call arguments without breaking inner JSON", async () => { + const person = 'Jane "JJ" \\vault\nline\t\u0001 café 😀'; + const input = `Person ${person}`; + const entity = { + entity_type: "PERSON", + start: input.indexOf(person), + end: input.indexOf(person) + person.length, + score: 0.99, + }; + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[entity]], + allEntities: [entity], + scanTimeMs: 3, + }); + const deltas = [ + { delta: '{"person":"[[PER', sequence_number: 1 }, + { delta: 'SON_1]]"}', sequence_number: 2 }, + ]; + const sse = `${deltas + .map( + (delta) => + `data: ${JSON.stringify({ + type: "response.function_call_arguments.delta", + item_id: "fc_1", + output_index: 0, + ...delta, + })}\n\n`, + ) + .join("")}data: [DONE]\n\n`; + globalThis.fetch = (async (_input: string | URL | Request, _init?: RequestInit) => + Promise.resolve( + new Response(sse, { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }), + )) as typeof fetch; + + const res = await app.request("/codex/responses", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.5", input, stream: true }), + headers: { + Authorization: "Bearer chatgpt-token", + "Content-Type": "application/json", + }, + }); + const text = await res.text(); + const restoredDeltas = text + .split("\n") + .filter((line) => line.startsWith("data: ") && line !== "data: [DONE]") + .map((line) => JSON.parse(line.slice(6)) as { delta: string }) + .map((event) => event.delta) + .join(""); + + expect(res.status).toBe(200); + expect(JSON.parse(restoredDeltas)).toEqual({ person }); + expect(text.trimEnd().endsWith("data: [DONE]")).toBe(true); + }); + test("adds markers to streamed Codex secrets when show_markers is true", async () => { config.masking.show_markers = true; config.masking.marker_text = "[protected]"; diff --git a/src/routes/openai-responses.test.ts b/src/routes/openai-responses.test.ts index ee2530d..66b7353 100644 --- a/src/routes/openai-responses.test.ts +++ b/src/routes/openai-responses.test.ts @@ -173,6 +173,37 @@ describe("POST /openai/v1/responses", () => { expect(body).not.toContain("[[EMAIL_ADDRESS_1]]"); }); + test("restores non-streaming function-call arguments as valid JSON", async () => { + const person = 'Jane "JJ" \\vault\nline\t\u0001 café 😀'; + const input = `Person ${person}`; + mockAnalyzeRequest.mockResolvedValueOnce(emailDetection(input, person)); + globalThis.fetch = (async (_target: string | URL | Request, _init?: RequestInit) => + Response.json({ + id: "resp_function", + output: [ + { + id: "fc_1", + type: "function_call", + call_id: "call_1", + name: "save_person", + arguments: JSON.stringify({ nested: { person: "[[EMAIL_ADDRESS_1]]" } }), + }, + ], + })) as typeof fetch; + + const response = await app.request("/openai/v1/responses", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ model: "openai/gpt-test", input }), + }); + const body = (await response.json()) as { + output: Array<{ arguments: string }>; + }; + + expect(response.status).toBe(200); + expect(JSON.parse(body.output[0].arguments)).toEqual({ nested: { person } }); + }); + test("remasks known values in restored assistant history", async () => { const email = "history@example.com"; const userText = `My email is ${email}`; @@ -217,6 +248,293 @@ describe("POST /openai/v1/responses", () => { expect(serialized.match(/\[\[EMAIL_ADDRESS_1\]\]/g)).toHaveLength(2); }); + test("remasks JSON-escaped known values in echoed function-call history", async () => { + const person = 'Ava "Snow" \\path\nline 雪'; + const userText = `Remember ${person}`; + const start = userText.indexOf(person); + const entity = { + entity_type: "PERSON", + start, + end: start + person.length, + score: 0.99, + }; + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[entity], [], [], []], + allEntities: [entity], + scanTimeMs: 2, + }); + + let upstreamBody: { + input: Array<{ type: string; arguments?: string; content?: Array<{ text: string }> }>; + } | null = null; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ id: "resp_function_history", output: [] }); + }) as typeof fetch; + + const response = await app.request("/openai/v1/responses", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "openai/gpt-test", + input: [ + { + type: "message", + role: "user", + content: [{ type: "input_text", text: userText }], + }, + { + type: "function_call", + id: "fc_history", + call_id: "call_history", + name: "save_person", + arguments: JSON.stringify({ name: person, city: "Bolzano" }), + }, + { + type: "custom_event", + arguments: JSON.stringify({ name: person }), + }, + { + type: "function_call", + id: "fc_malformed", + call_id: "call_malformed", + name: "save_malformed", + arguments: `{"name":"${JSON.stringify(person).slice(1, -1)}`, + }, + ], + }), + }); + + expect(response.status).toBe(200); + expect(upstreamBody).not.toBeNull(); + const input = upstreamBody!.input; + expect(input[0].content?.[0].text).toBe("Remember [[PERSON_1]]"); + const argumentsText = input[1].arguments!; + expect(JSON.parse(argumentsText)).toEqual({ name: "[[PERSON_1]]", city: "Bolzano" }); + expect(argumentsText).toContain("[[PERSON_1]]"); + expect(argumentsText).not.toContain(JSON.stringify(person).slice(1, -1)); + expect(JSON.parse(input[2].arguments!)).toEqual({ name: person }); + expect(input[3].arguments).toBe('{"name":"[[PERSON_1]]'); + }); + + test("keeps colliding raw and escaped function-argument values distinct", async () => { + const first = "A\nB"; + const second = "A\\nB"; + const userText = `First ${first}; second ${second}`; + const firstStart = userText.indexOf(first); + const secondStart = userText.indexOf(second, firstStart + first.length); + const firstEntity = { + entity_type: "PERSON", + start: firstStart, + end: firstStart + first.length, + score: 0.99, + }; + const secondEntity = { + entity_type: "LOCATION", + start: secondStart, + end: secondStart + second.length, + score: 0.99, + }; + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[firstEntity, secondEntity], []], + allEntities: [firstEntity, secondEntity], + scanTimeMs: 2, + }); + + let upstreamBody: { + input: Array<{ type: string; arguments?: string; content?: Array<{ text: string }> }>; + } | null = null; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ id: "resp_collision_history", output: [] }); + }) as typeof fetch; + + const response = await app.request("/openai/v1/responses", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "openai/gpt-test", + input: [ + { + type: "message", + role: "user", + content: [{ type: "input_text", text: userText }], + }, + { + type: "function_call", + id: "fc_collision", + call_id: "call_collision", + name: "save_values", + arguments: JSON.stringify({ first, second }), + }, + ], + }), + }); + + expect(response.status).toBe(200); + expect(upstreamBody).not.toBeNull(); + const input = upstreamBody!.input; + expect(input[0].content?.[0].text).toBe("First [[PERSON_1]]; second [[LOCATION_1]]"); + expect(JSON.parse(input[1].arguments!)).toEqual({ + first: "[[PERSON_1]]", + second: "[[LOCATION_1]]", + }); + }); + + test("remasks a known value represented as a JSON number", async () => { + const phone = "3471234567"; + const city = "Berlin"; + const userText = `Phone ${phone}; city ${city}`; + const phoneStart = userText.indexOf(phone); + const cityStart = userText.indexOf(city); + const phoneEntity = { + entity_type: "PHONE_NUMBER", + start: phoneStart, + end: phoneStart + phone.length, + score: 0.99, + }; + const cityEntity = { + entity_type: "LOCATION", + start: cityStart, + end: cityStart + city.length, + score: 0.99, + }; + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[phoneEntity, cityEntity], []], + allEntities: [phoneEntity, cityEntity], + scanTimeMs: 2, + }); + + let upstreamBody: { + input: Array<{ type: string; arguments?: string; content?: Array<{ text: string }> }>; + } | null = null; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ id: "resp_numeric_history", output: [] }); + }) as typeof fetch; + + const response = await app.request("/openai/v1/responses", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "openai/gpt-test", + input: [ + { + type: "message", + role: "user", + content: [{ type: "input_text", text: userText }], + }, + { + type: "function_call", + id: "fc_numeric", + call_id: "call_numeric", + name: "save_contact", + arguments: JSON.stringify({ + phone: Number(phone), + city, + attempts: 3, + verified: true, + note: null, + }), + }, + ], + }), + }); + + expect(response.status).toBe(200); + expect(upstreamBody).not.toBeNull(); + const input = upstreamBody!.input; + expect(input[0].content?.[0].text).toBe("Phone [[PHONE_NUMBER_1]]; city [[LOCATION_1]]"); + expect(JSON.parse(input[1].arguments!)).toEqual({ + phone: "[[PHONE_NUMBER_1]]", + city: "[[LOCATION_1]]", + attempts: 3, + verified: true, + note: null, + }); + }); + + test("preserves unrelated serialized JSON bytes while remasking", async () => { + const phone = "3471234567"; + const city = "Berlin"; + const userText = `Phone ${phone}; city ${city}`; + const phoneStart = userText.indexOf(phone); + const cityStart = userText.indexOf(city); + const phoneEntity = { + entity_type: "PHONE_NUMBER", + start: phoneStart, + end: phoneStart + phone.length, + score: 0.99, + }; + const cityEntity = { + entity_type: "LOCATION", + start: cityStart, + end: cityStart + city.length, + score: 0.99, + }; + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[phoneEntity, cityEntity], []], + allEntities: [phoneEntity, cityEntity], + scanTimeMs: 2, + }); + + const argumentsText = `{ + "phone" : "3471234567", + "Berlin": "office", + "big": 9007199254740993, + "decimal": 1.0, + "exponent": 1e+3, + "duplicate": "first", + "duplicate": "second" +}`; + const expectedArguments = `{ + "phone" : "[[PHONE_NUMBER_1]]", + "[[LOCATION_1]]": "office", + "big": 9007199254740993, + "decimal": 1.0, + "exponent": 1e+3, + "duplicate": "first", + "duplicate": "second" +}`; + let upstreamBody: { + input: Array<{ type: string; arguments?: string }>; + } | null = null; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ id: "resp_number_format_history", output: [] }); + }) as typeof fetch; + + const response = await app.request("/openai/v1/responses", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "openai/gpt-test", + input: [ + { + type: "message", + role: "user", + content: [{ type: "input_text", text: userText }], + }, + { + type: "function_call", + id: "fc_number_format", + call_id: "call_number_format", + name: "save_contact", + arguments: argumentsText, + }, + ], + }), + }); + + expect(response.status).toBe(200); + expect(upstreamBody).not.toBeNull(); + expect(upstreamBody!.input[1].arguments).toBe(expectedArguments); + }); + test("blocks stateful options when values were masked", async () => { const email = "state@example.com"; const input = `Email ${email}`; diff --git a/src/routes/openai-responses.ts b/src/routes/openai-responses.ts index a6a1852..3dec403 100644 --- a/src/routes/openai-responses.ts +++ b/src/routes/openai-responses.ts @@ -7,6 +7,7 @@ import { formatMaskedRequestForLog } from "../logging/log-content"; import { logRequest } from "../logging/logger"; import type { PlaceholderContext } from "../masking/context"; import { + isResponsesFunctionCallArguments, type ResponsesRequest, type ResponsesResponse, responsesExtractor, @@ -301,10 +302,72 @@ function remaskKnownValues( if (replacements.size === 0) return request; const ordered = [...replacements].sort(([a], [b]) => b.length - a.length); + const escapePattern = (value: string) => value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + const rawPattern = new RegExp(ordered.map(([value]) => escapePattern(value)).join("|"), "g"); + const serializedReplacements = new Map(); + for (const [original, placeholder] of ordered) { + for (const variant of [JSON.stringify(original).slice(1, -1), original]) { + if (!serializedReplacements.has(variant)) serializedReplacements.set(variant, placeholder); + } + } + const serializedOrdered = [...serializedReplacements].sort(([a], [b]) => b.length - a.length); + const serializedPattern = new RegExp( + serializedOrdered.map(([value]) => escapePattern(value)).join("|"), + "g", + ); + const jsonTokenPattern = + /"(?:\\[\s\S]|[^"\\])*"|-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?|true|false|null/g; + const remaskJsonString = (token: string): string => { + const decoded = JSON.parse(token) as string; + rawPattern.lastIndex = 0; + const matches = [...decoded.matchAll(rawPattern)]; + if (matches.length === 0) return token; + + const sourceBoundaries = [1]; + for (let index = 1; index < token.length - 1; ) { + index += token[index] === "\\" ? (token[index + 1] === "u" ? 6 : 2) : 1; + sourceBoundaries.push(index); + } + + let remasked = '"'; + let sourceIndex = 1; + for (const match of matches) { + const start = sourceBoundaries[match.index!]; + const end = sourceBoundaries[match.index! + match[0].length]; + remasked += token.slice(sourceIndex, start); + remasked += JSON.stringify(replacements.get(match[0])!).slice(1, -1); + sourceIndex = end; + } + return remasked + token.slice(sourceIndex); + }; const changed = responsesExtractor.extractTexts(request).flatMap((span) => { let maskedText = span.text; - for (const [original, placeholder] of ordered) { - maskedText = maskedText.split(original).join(placeholder); + if (isResponsesFunctionCallArguments(request, span.path)) { + try { + JSON.parse(span.text); + let jsonChanged = false; + maskedText = span.text.replace(jsonTokenPattern, (token) => { + if (token.startsWith('"')) { + const remasked = remaskJsonString(token); + if (remasked !== token) jsonChanged = true; + return remasked; + } + const placeholder = replacements.get(token); + if (!placeholder) return token; + jsonChanged = true; + return JSON.stringify(placeholder); + }); + if (!jsonChanged) maskedText = span.text; + } catch { + maskedText = maskedText.replace( + serializedPattern, + (value) => serializedReplacements.get(value)!, + ); + } + } else { + for (const [original, placeholder] of ordered) { + maskedText = maskedText.split(original).join(placeholder); + } } return maskedText === span.text ? [] : [{ ...span, maskedText }]; });