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<ResponsesRequest, ResponsesResponse> = {
extractTexts(request: ResponsesRequest): TextSpan[] {
return collectText(request).map((item, index) => ({
): ResponsesResponse {
let result = response;
for (const item of collectText(response)) {
- result = setAtPath(result, item.path, restorePlaceholders(item.value, context, formatValue));
+ const restored = isResponsesFunctionCallArguments(response, item.path)
+ ? restoreSerializedFunctionCallArguments(item.value, context, formatValue)
+ : restorePlaceholders(item.value, context, formatValue);
+ result = setAtPath(result, item.path, restored);
}
return result;
},
output: [{ content: [{ type: "output_text", text: "Your key is [protected]sk-secret" }] }],
});
});
+
+ test("Responses function-call arguments remain valid serialized JSON", () => {
+ 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" }] });
+ });
});
piiContext?: PlaceholderContext;
secretsContext?: PlaceholderContext;
config: MaskingConfig;
+ formatValue?: (original: string) => string;
}
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 {
return processedText;
}
+ hasPending(): boolean {
+ return Boolean(this.piiBuffer || this.secretsBuffer);
+ }
+
flush(): string {
let flushed = "";
return result;
}
+function dataEvent(event: Record<string, unknown>): 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, unknown> = {},
+): string {
+ return dataEvent({
+ type: "response.function_call_arguments.delta",
+ item_id: itemId,
+ output_index: outputIndex,
+ delta,
+ ...metadata,
+ });
+}
+
+type DataEvent = Record<string, unknown>;
+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<Record<string, unknown>>, 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", () => {
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"]);
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<Uint8Array>({
+ 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<Uint8Array>({
+ 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<string, unknown>).delta).toBe('{"name":"[[PER');
+ expect(JSON.parse((timeline[1] as Record<string, unknown>).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<string, unknown>).delta).toBe('{"name":"[[PERSON');
+ expect(timeline.at(-1)).toBe("[DONE]");
});
});
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<string, unknown>;
+ ready: boolean;
+}
+
+type TimelineFrame = EventTimelineFrame | { line: string; ready: boolean };
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+ return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
export function createResponsesUnmaskingStream(
stream: ReadableStream<Uint8Array>,
piiContext: PlaceholderContext | undefined,
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<string, ArgumentStreamState>();
+ 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, unknown>): 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<string, unknown>): 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<string, unknown>, 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<string, unknown>): 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<string, unknown>): Record<string, unknown> {
+ releaseArgumentsBefore(payload);
const result = payload as ResponsesResponse;
const spans = responsesExtractor.extractTexts(result);
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<string, unknown>);
+ } 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({
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));
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]";
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}`;
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}`;
import { logRequest } from "../logging/logger";
import type { PlaceholderContext } from "../masking/context";
import {
+ isResponsesFunctionCallArguments,
type ResponsesRequest,
type ResponsesResponse,
responsesExtractor,
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<string, string>();
+ 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 }];
});
git.99rst.org / sgasser-llm-shield.git / commitdiff