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 {
});
});
+ 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 = {
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",
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<string>();
+ 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
*
): 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) }
+ : {}),
+ },
+ };
+ }),
};
},
};
return result;
}
+interface ToolCallDelta {
+ index?: number;
+ id?: string;
+ type?: string;
+ function?: {
+ name?: string;
+ arguments?: string;
+ [key: string]: unknown;
+ };
+ [key: string]: unknown;
+}
+
+interface ChoiceDelta {
+ index?: number;
+ delta: {
+ content?: unknown;
+ tool_calls?: ToolCallDelta[];
+ [key: string]: unknown;
+ };
+ finish_reason?: string | null;
+ [key: string]: unknown;
+}
+
+interface ParsedStreamEvent {
+ choices: ChoiceDelta[];
+ [key: string]: unknown;
+}
+
+function createSSEEvent(event: ParsedStreamEvent): string {
+ return `data: ${JSON.stringify(event)}\n\n`;
+}
+
+function createToolArgumentEvent(
+ argumentsFragment: string,
+ {
+ choiceIndex = 0,
+ toolCallIndex = 0,
+ }: {
+ choiceIndex?: number;
+ toolCallIndex?: number;
+ } = {},
+): string {
+ return createSSEEvent({
+ choices: [
+ {
+ index: choiceIndex,
+ delta: {
+ tool_calls: [
+ {
+ index: toolCallIndex,
+ function: { arguments: argumentsFragment },
+ },
+ ],
+ },
+ finish_reason: null,
+ },
+ ],
+ });
+}
+
+function parseDataPayloads(output: string): string[] {
+ return output
+ .split("\n")
+ .filter((line) => 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();
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");
+ });
});
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<string, unknown>;
+
+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<string, unknown> {
+ return typeof value === "object" && value !== null;
+}
+
+function getSyntheticEventMetadata(event: Record<string, unknown>): 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<Uint8Array>,
piiContext: PlaceholderContext | undefined,
const encoder = new TextEncoder();
let lineBuffer = "";
const restorer = new StreamRestorer({ piiContext, secretsContext, config });
+ const argumentChannelsByPosition = new Map<string, ToolArgumentChannel>();
+ 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<string, unknown>): 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<string, unknown>;
+ 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<string, unknown>) {
+ 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") {
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`));
}
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;
for (const line of lines) {
processLine(line);
+ if (terminated) break;
+ }
+
+ if (terminated) {
+ await reader.cancel().catch(() => undefined);
+ controller.close();
+ break;
}
}
} catch (error) {
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
-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<PIIDetectionResult>>(() => 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", {
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<string, unknown> | 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<Record<string, unknown>>;
+ 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<string, unknown> | 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<Record<string, unknown>>;
+ 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<string, unknown> | 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<Record<string, unknown>>;
+ 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<string, unknown> | 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<Record<string, unknown>>;
+ 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", () => {
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 {
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"),
git.99rst.org / sgasser-llm-shield.git / commitdiff