From: Stefan Gasser Date: Tue, 25 Aug 2026 09:21:36 +0000 (+0200) Subject: fix(anthropic): restore tool input arguments (#169) X-Git-Tag: v0.9.3~2 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=a84dcf92fb08649aa64ae6223f93702f3e885a65;p=sgasser-llm-shield.git fix(anthropic): restore tool input arguments (#169) * fix(anthropic): restore tool input arguments * fix(anthropic): harden tool input restoration --- diff --git a/src/masking/extractors/anthropic.test.ts b/src/masking/extractors/anthropic.test.ts index 11576ea..f1736ff 100644 --- a/src/masking/extractors/anthropic.test.ts +++ b/src/masking/extractors/anthropic.test.ts @@ -5,7 +5,7 @@ import type { AnthropicRequest, AnthropicResponse, } from "../../providers/anthropic/types"; -import { anthropicExtractor } from "./anthropic"; +import { anthropicExtractor, remaskAnthropicToolUseHistory } from "./anthropic"; /** Helper to create a minimal request from messages */ function createRequest( @@ -16,6 +16,61 @@ function createRequest( } describe("Anthropic Text Extractor", () => { + describe("remaskAnthropicToolUseHistory", () => { + test("remasks exact known numeric primitives without changing unrelated values", () => { + const phone = 3901234567; + const request = createRequest([ + { + role: "assistant", + content: [ + { + type: "tool_use", + id: "tool_phone", + name: "lookup_contact", + input: { + phone, + asString: String(phone), + nested: [phone, { phone, unrelated: 3901234568 }, 42.5, true, false, null], + }, + vendor_metadata: { stable: true }, + }, + ], + }, + ]); + const context: PlaceholderContext = { + mapping: { "[[PHONE_NUMBER_1]]": String(phone) }, + reverseMapping: {}, + counters: { PHONE_NUMBER: 1 }, + }; + + const result = remaskAnthropicToolUseHistory(request, context); + + expect(result.messages[0]).toEqual({ + role: "assistant", + content: [ + { + type: "tool_use", + id: "tool_phone", + name: "lookup_contact", + input: { + phone: "[[PHONE_NUMBER_1]]", + asString: "[[PHONE_NUMBER_1]]", + nested: [ + "[[PHONE_NUMBER_1]]", + { phone: "[[PHONE_NUMBER_1]]", unrelated: 3901234568 }, + 42.5, + true, + false, + null, + ], + }, + vendor_metadata: { stable: true }, + }, + ], + }); + }); + }); + describe("extractTexts", () => { test("extracts text from string content", () => { const request = createRequest([ @@ -651,6 +706,75 @@ describe("Anthropic Text Extractor", () => { expect((result.content[0] as { text: string }).text).toBe("Hello [protected]John"); }); + test("recursively unmasks string leaves in tool_use input", () => { + const response: AnthropicResponse = { + id: "msg_tool", + type: "message", + role: "assistant", + content: [ + { type: "text", text: "Calling for [[PERSON_1]]" }, + { + type: "tool_use", + id: "tool_1", + name: "create_record", + input: { + contact: "[[PERSON_1]] <[[EMAIL_ADDRESS_1]]>", + nested: { + values: [ + "token [[SECRET_API_KEY_1]]", + 42, + true, + null, + { note: "Email: [[EMAIL_ADDRESS_1]]" }, + ], + }, + "[[PERSON_1]]": "keys stay masked", + }, + caller: { source: "fixture" }, + }, + ], + model: "claude-3-sonnet-20240229", + stop_reason: "tool_use", + stop_sequence: null, + usage: { input_tokens: 10, output_tokens: 5 }, + }; + const context: PlaceholderContext = { + mapping: { + "[[PERSON_1]]": "Alice", + "[[EMAIL_ADDRESS_1]]": "alice@example.com", + "[[SECRET_API_KEY_1]]": "sk-test-value", + }, + reverseMapping: {}, + counters: {}, + }; + + const result = anthropicExtractor.unmaskResponse( + response, + context, + (value) => `[protected]${value}`, + ); + const toolUse = result.content[1]; + + expect(result.content[0]).toEqual({ type: "text", text: "Calling for [protected]Alice" }); + expect(toolUse).toEqual({ + ...response.content[1], + input: { + contact: "[protected]Alice <[protected]alice@example.com>", + nested: { + values: [ + "token [protected]sk-test-value", + 42, + true, + null, + { note: "Email: [protected]alice@example.com" }, + ], + }, + "[[PERSON_1]]": "keys stay masked", + }, + }); + expect(response.content[1]).not.toBe(toolUse); + }); + test("handles multiple text blocks", () => { const response: AnthropicResponse = { id: "msg_123", diff --git a/src/masking/extractors/anthropic.ts b/src/masking/extractors/anthropic.ts index 3f9d6b2..2d86eec 100644 --- a/src/masking/extractors/anthropic.ts +++ b/src/masking/extractors/anthropic.ts @@ -10,7 +10,7 @@ * System spans use messageIndex -1 to distinguish from message spans. */ -import type { PlaceholderContext } from "../../masking/context"; +import { type PlaceholderContext, restorePlaceholders } from "../../masking/context"; import type { AnthropicRequest, AnthropicResponse, @@ -18,12 +18,107 @@ import type { TextBlock, ThinkingBlock, ToolResultBlock, + ToolUseBlock, } from "../../providers/anthropic/types"; import type { MaskedSpan, RequestExtractor, TextSpan } from "../types"; /** System content uses messageIndex -1 */ const SYSTEM_MESSAGE_INDEX = -1; +export function unmaskAnthropicToolInput( + value: unknown, + context: PlaceholderContext, + formatValue?: (original: string) => string, +): unknown { + if (typeof value === "string") { + return restorePlaceholders(value, context, formatValue); + } + + if (Array.isArray(value)) { + return value.map((item) => unmaskAnthropicToolInput(item, context, formatValue)); + } + + if (value !== null && typeof value === "object") { + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [ + key, + unmaskAnthropicToolInput(item, context, formatValue), + ]), + ); + } + + return value; +} + +function remaskAnthropicToolInput( + value: unknown, + contexts: Array, +): unknown { + if (typeof value === "string") { + let result = value; + const replacements = contexts + .flatMap((context) => (context ? Object.entries(context.mapping) : [])) + .sort(([, left], [, right]) => right.length - left.length); + + for (const [placeholder, original] of replacements) { + if (original) result = result.split(original).join(placeholder); + } + + return result; + } + + if (typeof value === "number" && Number.isFinite(value)) { + const original = String(value); + for (const context of contexts) { + for (const [placeholder, knownValue] of Object.entries(context?.mapping ?? {})) { + if (knownValue === original) return placeholder; + } + } + } + + if (Array.isArray(value)) { + return value.map((item) => remaskAnthropicToolInput(item, contexts)); + } + + if (value !== null && typeof value === "object") { + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [key, remaskAnthropicToolInput(item, contexts)]), + ); + } + + return value; +} + +export function remaskAnthropicToolUseHistory( + request: AnthropicRequest, + piiContext: PlaceholderContext | undefined, + secretsContext?: PlaceholderContext, +): AnthropicRequest { + const contexts = [piiContext, secretsContext]; + if (!contexts.some((context) => context && Object.keys(context.mapping).length > 0)) { + return request; + } + + return { + ...request, + messages: request.messages.map((message) => { + if (message.role !== "assistant" || !Array.isArray(message.content)) return message; + + return { + ...message, + content: message.content.map((block) => + block.type === "tool_use" + ? { + ...block, + input: remaskAnthropicToolInput(block.input, contexts) as Record, + } + : block, + ), + }; + }), + }; +} + /** * Extract text from a single content block */ @@ -257,20 +352,24 @@ export const anthropicExtractor: RequestExtractor string, ): AnthropicResponse { - const unmaskText = (text: string): string => { - let result = text; - for (const [placeholder, original] of Object.entries(context.mapping)) { - const value = formatValue ? formatValue(original) : original; - result = result.replaceAll(placeholder, value); - } - return result; - }; - return { ...response, content: response.content.map((block) => { if (block.type === "text") { - return { ...block, text: unmaskText((block as TextBlock).text) }; + return { + ...block, + text: restorePlaceholders((block as TextBlock).text, context, formatValue), + }; + } + if (block.type === "tool_use") { + const toolUse = block as ToolUseBlock; + return { + ...toolUse, + input: unmaskAnthropicToolInput(toolUse.input, context, formatValue) as Record< + string, + unknown + >, + }; } return block; }), diff --git a/src/providers/anthropic/stream-transformer.test.ts b/src/providers/anthropic/stream-transformer.test.ts index c4f95cd..435a962 100644 --- a/src/providers/anthropic/stream-transformer.test.ts +++ b/src/providers/anthropic/stream-transformer.test.ts @@ -65,6 +65,52 @@ function createTextDelta(text: string, index = 0): string { }); } +function createBlockStop(index: number, metadata: Record = {}): string { + return createAnthropicEvent("content_block_stop", { + type: "content_block_stop", + index, + ...metadata, + }); +} + +function createInputJsonDelta( + partialJson: string, + index = 0, + eventMetadata: Record = {}, + deltaMetadata: Record = {}, +): string { + return createAnthropicEvent("content_block_delta", { + type: "content_block_delta", + index, + ...eventMetadata, + delta: { type: "input_json_delta", partial_json: partialJson, ...deltaMetadata }, + }); +} + +function parseDataEvents(result: string): Array<{ + type: string; + index?: number; + delta?: { type: string; partial_json?: string; text?: string; [key: string]: unknown }; + [key: string]: unknown; +}> { + return result + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => JSON.parse(line.slice(6))); +} + +function concatenateToolJson(result: string, index: number): string { + return parseDataEvents(result) + .filter( + (event) => + event.type === "content_block_delta" && + event.index === index && + event.delta?.type === "input_json_delta", + ) + .map((event) => event.delta?.partial_json ?? "") + .join(""); +} + describe("createAnthropicUnmaskingStream", () => { test("unmasks complete placeholder in single chunk", async () => { const context = createMaskingContext(); @@ -284,38 +330,267 @@ describe("createAnthropicUnmaskingStream", () => { expect(result).toContain("Bob"); }); - test("handles tool_use deltas (input_json_delta)", async () => { + test("unmasks a complete placeholder in input_json_delta", async () => { const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Alice"; - const toolUseDelta = createAnthropicEvent("content_block_delta", { - type: "content_block_delta", - index: 0, - delta: { type: "input_json_delta", partial_json: '{"arg": "value"}' }, - }); - const source = createSSEStream([toolUseDelta]); + const source = createSSEStream([ + createInputJsonDelta('{"name":"[[PERSON_1]]"}'), + createBlockStop(0), + ]); const unmaskedStream = createAnthropicUnmaskingStream(source, context, defaultConfig); const result = await consumeStream(unmaskedStream); - // input_json_delta should pass through unchanged - expect(result).toContain("input_json_delta"); - expect(result).toContain("arg"); - expect(result).toContain("value"); + expect(JSON.parse(concatenateToolJson(result, 0))).toEqual({ name: "Alice" }); }); - test("handles content_block_stop events", async () => { + test("preserves tool JSON bytes when no mapped placeholder is present", async () => { const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Alice"; + const delta = createInputJsonDelta('{ "count": 1e+2, "enabled": true }', 2); + const stop = createBlockStop(2); - const blockStop = createAnthropicEvent("content_block_stop", { - type: "content_block_stop", - index: 0, + const result = await consumeStream( + createAnthropicUnmaskingStream(createSSEStream([delta, stop]), context, defaultConfig), + ); + + expect(result).toBe(delta + stop); + }); + + test("restores a placeholder without normalizing untouched JSON tokens", async () => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Alice"; + const originalJson = + '{ "big": 9007199254740993, "fixed": 1.0, "exp": 1e+3, "dup": "first", "dup": "last", "name": "[[PERSON_1]]" }'; + const expectedJson = originalJson.replace("[[PERSON_1]]", "Alice"); + const chunks = [ + createInputJsonDelta(originalJson.slice(0, 45), 8), + createInputJsonDelta(originalJson.slice(45, 91), 8), + createInputJsonDelta(originalJson.slice(91), 8), + createBlockStop(8), + ]; + + const result = await consumeStream( + createAnthropicUnmaskingStream(createSSEStream(chunks), context, defaultConfig), + ); + + expect(concatenateToolJson(result, 8)).toBe(expectedJson); + }); + + test("restores Unicode-escaped placeholders without rewriting surrounding escapes", async () => { + const context = createMaskingContext(); + const restoredValue = 'Quote " slash \\ line\n snowman ☃'; + context.mapping["[[PERSON_1]]"] = restoredValue; + const encodedPlaceholder = "\\u005b\\u005bPERSON_1\\u005d\\u005d"; + const originalJson = + `{ "encoded": "${encodedPlaceholder}", ` + + '"untouched": "quote \\" slash \\\\ newline\\n snowman \\u2603" }'; + const serializedValue = JSON.stringify(restoredValue).slice(1, -1); + const expectedJson = originalJson.replace(encodedPlaceholder, serializedValue); + + const result = await consumeStream( + createAnthropicUnmaskingStream( + createSSEStream([createInputJsonDelta(originalJson, 9), createBlockStop(9)]), + context, + defaultConfig, + ), + ); + + expect(concatenateToolJson(result, 9)).toBe(expectedJson); + expect(JSON.parse(concatenateToolJson(result, 9))).toEqual({ + encoded: restoredValue, + untouched: 'quote " slash \\ newline\n snowman ☃', }); - const source = createSSEStream([blockStop]); + }); - const unmaskedStream = createAnthropicUnmaskingStream(source, context, defaultConfig); - const result = await consumeStream(unmaskedStream); + test("unmasks placeholders and JSON strings split across input_json_delta events", async () => { + const context = createMaskingContext(); + context.mapping["[[EMAIL_ADDRESS_1]]"] = "alice@example.com"; + + const source = createSSEStream([ + createInputJsonDelta('{"nested":{"email":"before [[EMAIL_', 3), + createInputJsonDelta('ADDRESS_1]] after","items":["x",', 3), + createInputJsonDelta('"[[EMAIL_ADDRESS_1]]"]}}', 3), + createBlockStop(3), + ]); + + const result = await consumeStream( + createAnthropicUnmaskingStream(source, context, defaultConfig), + ); + + expect(JSON.parse(concatenateToolJson(result, 3))).toEqual({ + nested: { + email: "before alice@example.com after", + items: ["x", "alice@example.com"], + }, + }); + expect( + parseDataEvents(result).filter((event) => event.delta?.type === "input_json_delta"), + ).toHaveLength(3); + }); + + test("keeps independent state for interleaved tool block indexes and preserves event metadata", async () => { + const piiContext = createMaskingContext(); + piiContext.mapping["[[PERSON_1]]"] = "Alice"; + const secretsContext = createMaskingContext(); + secretsContext.mapping["[[SECRET_API_KEY_1]]"] = "sk-test"; + + const chunks = [ + createInputJsonDelta('{"owner":"[[PER', 1, { trace: "first" }, { ordinal: 1 }), + createInputJsonDelta('{"key":"[[SECRET_', 2, { trace: "second" }, { ordinal: 2 }), + createTextDelta("Visible [[PERSON_1]]", 0), + createInputJsonDelta('SON_1]]"}', 1, { trace: "third" }, { ordinal: 3 }), + createBlockStop(1, { stop_metadata: "one" }), + createInputJsonDelta('API_KEY_1]]"}', 2, { trace: "fourth" }, { ordinal: 4 }), + createBlockStop(2, { stop_metadata: "two" }), + ]; + + const result = await consumeStream( + createAnthropicUnmaskingStream( + createSSEStream(chunks), + piiContext, + defaultConfig, + secretsContext, + ), + ); + const events = parseDataEvents(result); + + expect(events.map((event) => [event.type, event.index])).toEqual([ + ["content_block_delta", 1], + ["content_block_delta", 2], + ["content_block_delta", 0], + ["content_block_delta", 1], + ["content_block_stop", 1], + ["content_block_delta", 2], + ["content_block_stop", 2], + ]); + expect(JSON.parse(concatenateToolJson(result, 1))).toEqual({ owner: "Alice" }); + expect(JSON.parse(concatenateToolJson(result, 2))).toEqual({ key: "sk-test" }); + expect(events[0].trace).toBe("first"); + expect(events[0].delta?.ordinal).toBe(1); + expect(events[3].trace).toBe("third"); + expect(events[4].stop_metadata).toBe("one"); + expect(events[6].stop_metadata).toBe("two"); + expect(events[2].delta?.text).toBe("Visible Alice"); + expect( + result + .split("\n") + .filter((line) => line.startsWith("event: ")) + .map((line) => line.slice(7)), + ).toEqual(events.map((event) => event.type)); + }); + + test("preserves interleaved frame identity while restoring independent tool blocks", async () => { + const piiContext = createMaskingContext(); + piiContext.mapping["[[PERSON_1]]"] = "Alice"; + const secretsContext = createMaskingContext(); + secretsContext.mapping["[[SECRET_API_KEY_1]]"] = "sk-test"; + const first = + 'event: content_block_delta\r\ndata: { "type" : "content_block_delta", "index" : 11, "trace" : "first", "delta" : { "type" : "input_json_delta", "partial_json" : "{\\"owner\\":\\"[[PERSON_1]]\\"}", "ordinal" : 1 } }\r\n\r\n'; + const second = + 'event: content_block_delta\r\ndata: { "type" : "content_block_delta", "index" : 12, "trace" : "second", "delta" : { "type" : "input_json_delta", "partial_json" : "{\\"key\\":\\"[[SECRET_API_KEY_1]]\\"}", "ordinal" : 2 } }\r\n\r\n'; + const firstStop = createBlockStop(11, { trace: "first-stop" }); + const secondStop = createBlockStop(12, { trace: "second-stop" }); + + const result = await consumeStream( + createAnthropicUnmaskingStream( + createSSEStream([first, second, firstStop, secondStop]), + piiContext, + defaultConfig, + secretsContext, + ), + ); + + expect(result).toBe( + first.replace("[[PERSON_1]]", "Alice") + + second.replace("[[SECRET_API_KEY_1]]", "sk-test") + + firstStop + + secondStop, + ); + }); + + test("preserves event and delta metadata bytes when tool JSON changes", async () => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Alice"; + const delta = + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":13,"big":9007199254740993,"fixed":1.0,"unicode":"\\u2603","delta":{"type":"input_json_delta","partial_json":"{\\"name\\":\\"[[PERSON_1]]\\"}","exp":1e+3},"tail":true}\n\n'; + const stop = createBlockStop(13, { stable: "metadata" }); + + const result = await consumeStream( + createAnthropicUnmaskingStream(createSSEStream([delta, stop]), context, defaultConfig), + ); + + expect(result).toBe(delta.replace("[[PERSON_1]]", "Alice") + stop); + }); + + test("serializes restored PII and secrets with markers as valid JSON", async () => { + const piiContext = createMaskingContext(); + piiContext.mapping["[[PERSON_1]]"] = + 'Quote " slash \\ line\n tab\t nul\u0000 snowman ☃ [[prefix'; + const secretsContext = createMaskingContext(); + secretsContext.mapping["[[SECRET_API_KEY_1]]"] = "secret\r\b\fvalue"; + + const source = createSSEStream([ + createInputJsonDelta('{"person":"[[PERSON_', 4), + createInputJsonDelta('1]]","secret":"[[SECRET_API_', 4), + createInputJsonDelta('KEY_1]]"}', 4), + createBlockStop(4), + ]); + const result = await consumeStream( + createAnthropicUnmaskingStream(source, piiContext, markerConfig, secretsContext), + ); + + expect(JSON.parse(concatenateToolJson(result, 4))).toEqual({ + person: '[protected]Quote " slash \\ line\n tab\t nul\u0000 snowman ☃ [[prefix', + secret: "[protected]secret\r\b\fvalue", + }); + }); + + test("applies PII then secrets inside tool JSON like non-stream restoration", async () => { + const piiContext = createMaskingContext(); + piiContext.mapping["[[PERSON_1]]"] = "Alice [[SECRET_API_KEY_1]]"; + const secretsContext = createMaskingContext(); + secretsContext.mapping["[[SECRET_API_KEY_1]]"] = "sk-test"; + const source = createSSEStream([ + createInputJsonDelta('{"value":"[[PERSON_1]]"}', 14), + createBlockStop(14), + ]); + + const result = await consumeStream( + createAnthropicUnmaskingStream(source, piiContext, defaultConfig, secretsContext), + ); + + expect(JSON.parse(concatenateToolJson(result, 14))).toEqual({ + value: "Alice sk-test", + }); + }); + + test("restores complete tool JSON on clean stream end without a stop event", async () => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Alice"; + const finalEvent = createInputJsonDelta('{"name":"[[PERSON_1]]"}', 7).trimEnd(); + + const result = await consumeStream( + createAnthropicUnmaskingStream(createSSEStream([finalEvent]), context, defaultConfig), + ); + + expect(JSON.parse(concatenateToolJson(result, 7))).toEqual({ name: "Alice" }); + }); + + test("passes through malformed accumulated tool JSON unchanged on stop and end", async () => { + const context = createMaskingContext(); + context.mapping["[[PERSON_1]]"] = "Alice"; + const stopped = createInputJsonDelta('{"name":"[[PERSON_1]]"', 5); + const stop = createBlockStop(5); + const ended = createInputJsonDelta('{"other":"[[PERSON_1]]"', 6).trimEnd(); + const sourceText = stopped + stop + ended; + + const result = await consumeStream( + createAnthropicUnmaskingStream(createSSEStream([sourceText]), context, defaultConfig), + ); - expect(result).toContain("content_block_stop"); + expect(result).toBe(sourceText); }); test("handles message_delta events", async () => { diff --git a/src/providers/anthropic/stream-transformer.ts b/src/providers/anthropic/stream-transformer.ts index c39852f..d7abf96 100644 --- a/src/providers/anthropic/stream-transformer.ts +++ b/src/providers/anthropic/stream-transformer.ts @@ -1,10 +1,357 @@ // Anthropic SSE differs from OpenAI: event lines identify message/content events, -// and text arrives as content_block_delta data with delta.type === "text_delta". +// and content arrives as text or partial serialized tool-input JSON deltas. 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 { ContentBlockDeltaEvent, TextDelta } from "./types"; + +interface ParsedEvent { + type?: unknown; + index?: unknown; + delta?: { + type?: unknown; + text?: unknown; + partial_json?: unknown; + [key: string]: unknown; + }; + [key: string]: unknown; +} + +interface DataLine { + data: string; + start: number; + end: number; + newline: string; +} + +interface PendingFrame { + output: string | undefined; +} + +interface PendingToolFrame extends PendingFrame { + frame: string; + dataLine: DataLine; + event: ParsedEvent; + fragmentLength: number; +} + +interface ToolBlockState { + json: string; + frames: PendingToolFrame[]; +} + +function findDataLine(frame: string): DataLine | undefined { + const match = /^data: (.*)(\r?\n|$)/m.exec(frame); + if (!match) return undefined; + + return { + data: match[1], + start: match.index, + end: match.index + match[0].length, + newline: match[2], + }; +} + +function replaceDataLine(frame: string, dataLine: DataLine, data: string): string { + return `${frame.slice(0, dataLine.start)}data: ${data}${dataLine.newline}${frame.slice(dataLine.end)}`; +} + +function replacePartialJsonValue( + data: string, + originalFragment: string, + restoredFragment: string, +): string | undefined { + const originalValue = JSON.stringify(originalFragment); + const restoredValue = JSON.stringify(restoredFragment); + const property = /"partial_json"\s*:\s*/g; + let match = property.exec(data); + + while (match) { + const valueStart = match.index + match[0].length; + if (data.startsWith(originalValue, valueStart)) { + return ( + data.slice(0, valueStart) + restoredValue + data.slice(valueStart + originalValue.length) + ); + } + match = property.exec(data); + } + + return undefined; +} + +function distributeJson(json: string, frames: PendingToolFrame[]): string[] { + let offset = 0; + + return frames.map((frame, index) => { + if (index === frames.length - 1) { + return json.slice(offset); + } + + const fragment = json.slice(offset, offset + frame.fragmentLength); + offset += frame.fragmentLength; + return fragment; + }); +} + +interface StringReplacement { + start: number; + end: number; + value: string; +} + +interface SerializedCharacter { + start: number; + end: number; + value: string; +} + +function decodeSerializedString(value: string): SerializedCharacter[] { + const characters: SerializedCharacter[] = []; + let position = 0; + + while (position < value.length) { + if (value[position] !== "\\") { + characters.push({ start: position, end: position + 1, value: value[position] }); + position++; + continue; + } + + const escapeCode = value[position + 1]; + if (escapeCode === "u") { + characters.push({ + start: position, + end: position + 6, + value: String.fromCharCode(Number.parseInt(value.slice(position + 2, position + 6), 16)), + }); + position += 6; + continue; + } + + const escapedValues: Record = { + '"': '"', + "\\": "\\", + "/": "/", + b: "\b", + f: "\f", + n: "\n", + r: "\r", + t: "\t", + }; + characters.push({ start: position, end: position + 2, value: escapedValues[escapeCode] }); + position += 2; + } + + return characters; +} + +function replaceSerializedPlaceholder( + value: string, + placeholder: string, + replacement: string, +): string { + const characters = decodeSerializedString(value); + const decoded = characters.map((character) => character.value).join(""); + const matches: Array<{ start: number; end: number }> = []; + let matchStart = decoded.indexOf(placeholder); + + while (matchStart !== -1) { + matches.push({ + start: characters[matchStart].start, + end: characters[matchStart + placeholder.length - 1].end, + }); + matchStart = decoded.indexOf(placeholder, matchStart + placeholder.length); + } + + if (matches.length === 0) return value; + + let result = ""; + let unchangedStart = 0; + for (const match of matches) { + result += value.slice(unchangedStart, match.start) + replacement; + unchangedStart = match.end; + } + return result + value.slice(unchangedStart); +} + +function restoreSerializedString( + value: string, + contexts: Array, + formatValue: ((original: string) => string) | undefined, +): string { + let result = value; + for (const context of contexts) { + const replacements = Object.entries(context?.mapping ?? {}).sort( + ([left], [right]) => right.length - left.length, + ); + + for (const [placeholder, original] of replacements) { + const serialized = JSON.stringify(formatValue ? formatValue(original) : original).slice( + 1, + -1, + ); + result = replaceSerializedPlaceholder(result, placeholder, serialized); + } + } + + return result; +} + +function restoreJsonStringValues( + json: string, + contexts: Array, + formatValue: ((original: string) => string) | undefined, +): string | undefined { + let position = 0; + const replacements: StringReplacement[] = []; + + function skipWhitespace() { + while ( + json[position] === " " || + json[position] === "\t" || + json[position] === "\n" || + json[position] === "\r" + ) { + position++; + } + } + + function parseString(restoreValue: boolean) { + if (json[position] !== '"') throw new Error("Expected JSON string"); + position++; + const contentStart = position; + + while (position < json.length) { + const character = json[position]; + + if (character === '"') { + if (restoreValue) { + const original = json.slice(contentStart, position); + const restored = restoreSerializedString(original, contexts, formatValue); + if (restored !== original) { + replacements.push({ start: contentStart, end: position, value: restored }); + } + } + position++; + return; + } + + if (character === "\\") { + const escapeCode = json[position + 1]; + if (escapeCode === "u") { + if (!/^[0-9a-fA-F]{4}$/.test(json.slice(position + 2, position + 6))) { + throw new Error("Invalid JSON Unicode escape"); + } + position += 6; + continue; + } + if (!escapeCode || !'"\\/bfnrt'.includes(escapeCode)) { + throw new Error("Invalid JSON escape"); + } + position += 2; + continue; + } + + if (character.charCodeAt(0) < 0x20) throw new Error("Invalid JSON control character"); + position++; + } + + throw new Error("Incomplete JSON string"); + } + + function parseArray() { + position++; + skipWhitespace(); + if (json[position] === "]") { + position++; + return; + } + + while (true) { + parseValue(); + skipWhitespace(); + if (json[position] === "]") { + position++; + return; + } + if (json[position] !== ",") throw new Error("Invalid JSON array"); + position++; + skipWhitespace(); + } + } + + function parseObject() { + position++; + skipWhitespace(); + if (json[position] === "}") { + position++; + return; + } + + while (true) { + parseString(false); + skipWhitespace(); + if (json[position] !== ":") throw new Error("Invalid JSON object"); + position++; + parseValue(); + skipWhitespace(); + if (json[position] === "}") { + position++; + return; + } + if (json[position] !== ",") throw new Error("Invalid JSON object"); + position++; + skipWhitespace(); + } + } + + function parseValue() { + skipWhitespace(); + const character = json[position]; + + if (character === '"') { + parseString(true); + return; + } + if (character === "{") { + parseObject(); + return; + } + if (character === "[") { + parseArray(); + return; + } + + for (const literal of ["true", "false", "null"]) { + if (json.startsWith(literal, position)) { + position += literal.length; + return; + } + } + + const number = /^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?/.exec(json.slice(position))?.[0]; + if (!number) throw new Error("Invalid JSON value"); + position += number.length; + } + + try { + parseValue(); + skipWhitespace(); + if (position !== json.length) return undefined; + } catch { + return undefined; + } + + if (replacements.length === 0) return undefined; + + let restored = ""; + let unchangedStart = 0; + for (const replacement of replacements) { + restored += json.slice(unchangedStart, replacement.start) + replacement.value; + unchangedStart = replacement.end; + } + return restored + json.slice(unchangedStart); +} export function createAnthropicUnmaskingStream( source: ReadableStream, @@ -14,89 +361,187 @@ export function createAnthropicUnmaskingStream( ): ReadableStream { const decoder = new TextDecoder(); const encoder = new TextEncoder(); - let lineBuffer = ""; - const restorer = new StreamRestorer({ piiContext, secretsContext, config }); + let sseBuffer = ""; + const textRestorer = new StreamRestorer({ piiContext, secretsContext, config }); + const formatValue = createRestoreFormatter(config); + const toolBlocks = new Map(); + const pendingFrames: PendingFrame[] = []; return new ReadableStream({ async start(controller) { const reader = source.getReader(); - try { - while (true) { - const { done, value } = await reader.read(); + function flushReadyFrames() { + while (pendingFrames[0]?.output !== undefined) { + const pending = pendingFrames.shift(); + if (pending?.output) { + controller.enqueue(encoder.encode(pending.output)); + } + } + } + + function restoreChangedToolJson(json: string): string | undefined { + return restoreJsonStringValues(json, [piiContext, secretsContext], formatValue); + } - if (done) { - const flushed = restorer.flush(); - - // Send flushed content as final text delta - if (flushed) { - const finalEvent: ContentBlockDeltaEvent = { - type: "content_block_delta", - index: 0, - delta: { type: "text_delta", text: flushed }, - }; - controller.enqueue( - encoder.encode( - `event: content_block_delta\ndata: ${JSON.stringify(finalEvent)}\n\n`, - ), - ); + function finalizeToolBlock(index: number) { + const state = toolBlocks.get(index); + if (!state) return; + + const restoredJson = restoreChangedToolJson(state.json); + if (restoredJson === undefined) { + for (const pending of state.frames) { + pending.output = pending.frame; + } + } else { + const fragments = distributeJson(restoredJson, state.frames); + const restoredData = state.frames.map((pending, frameIndex) => + replacePartialJsonValue( + pending.dataLine.data, + pending.event.delta?.partial_json as string, + fragments[frameIndex], + ), + ); + + if (restoredData.some((data) => data === undefined)) { + for (const pending of state.frames) { + pending.output = pending.frame; } + toolBlocks.delete(index); + return; + } - controller.close(); - break; + for (let frameIndex = 0; frameIndex < state.frames.length; frameIndex++) { + const pending = state.frames[frameIndex]; + pending.output = replaceDataLine( + pending.frame, + pending.dataLine, + restoredData[frameIndex] as string, + ); } + } - lineBuffer += decoder.decode(value, { stream: true }); - const lines = lineBuffer.split("\n"); - lineBuffer = lines.pop() || ""; + toolBlocks.delete(index); + } - for (const line of lines) { - // Pass through event type lines - if (line.startsWith("event: ")) { - controller.enqueue(encoder.encode(`${line}\n`)); - continue; - } + function processFrame(frame: string) { + const dataLine = findDataLine(frame); + if (!dataLine) { + pendingFrames.push({ output: frame }); + flushReadyFrames(); + return; + } - // Process data lines - if (line.startsWith("data: ")) { - const data = line.slice(6); - - try { - const parsed = JSON.parse(data) as { type: string; delta?: { type: string } }; - - // Only process text deltas - if (parsed.type === "content_block_delta" && parsed.delta?.type === "text_delta") { - const event = parsed as ContentBlockDeltaEvent; - const textDelta = event.delta as TextDelta; - const processedText = restorer.restoreChunk(textDelta.text); - - // Only emit if we have content - if (processedText) { - const modifiedEvent = { - ...parsed, - delta: { ...textDelta, text: processedText }, - }; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(modifiedEvent)}\n`)); - } - } else { - // Pass through other events unchanged - controller.enqueue(encoder.encode(`data: ${data}\n`)); - } - } catch { - // Pass through unparseable data - controller.enqueue(encoder.encode(`${line}\n`)); - } - continue; - } + let parsed: ParsedEvent; + try { + parsed = JSON.parse(dataLine.data) as ParsedEvent; + } catch { + pendingFrames.push({ output: frame }); + flushReadyFrames(); + return; + } - // Pass through empty lines and other content - if (line.trim() === "") { - controller.enqueue(encoder.encode("\n")); - } else { - controller.enqueue(encoder.encode(`${line}\n`)); - } + if ( + parsed.type === "content_block_delta" && + typeof parsed.index === "number" && + parsed.delta?.type === "input_json_delta" && + typeof parsed.delta.partial_json === "string" + ) { + const pending: PendingToolFrame = { + output: undefined, + frame, + dataLine, + event: parsed, + fragmentLength: parsed.delta.partial_json.length, + }; + const state = toolBlocks.get(parsed.index) ?? { json: "", frames: [] }; + state.json += parsed.delta.partial_json; + state.frames.push(pending); + toolBlocks.set(parsed.index, state); + pendingFrames.push(pending); + flushReadyFrames(); + return; + } + + if (parsed.type === "content_block_stop" && typeof parsed.index === "number") { + finalizeToolBlock(parsed.index); + } + + if ( + parsed.type === "content_block_delta" && + parsed.delta?.type === "text_delta" && + typeof parsed.delta.text === "string" + ) { + const processedText = textRestorer.restoreChunk(parsed.delta.text); + + if (processedText) { + pendingFrames.push({ + output: replaceDataLine( + frame, + dataLine, + JSON.stringify({ + ...parsed, + delta: { ...parsed.delta, text: processedText }, + }), + ), + }); + } else { + pendingFrames.push({ + output: frame.slice(0, dataLine.start) + frame.slice(dataLine.end), + }); } + } else { + pendingFrames.push({ output: frame }); + } + + flushReadyFrames(); + } + + function processCompleteFrames() { + while (true) { + const separator = /\r?\n\r?\n/.exec(sseBuffer); + if (!separator) return; + + const frameEnd = separator.index + separator[0].length; + processFrame(sseBuffer.slice(0, frameEnd)); + sseBuffer = sseBuffer.slice(frameEnd); } + } + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + sseBuffer += decoder.decode(value, { stream: true }); + processCompleteFrames(); + } + + sseBuffer += decoder.decode(); + processCompleteFrames(); + if (sseBuffer) { + processFrame(sseBuffer); + sseBuffer = ""; + } + + for (const index of [...toolBlocks.keys()]) { + finalizeToolBlock(index); + } + flushReadyFrames(); + + const flushed = textRestorer.flush(); + if (flushed) { + const finalEvent = { + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text: flushed }, + }; + controller.enqueue( + encoder.encode(`event: content_block_delta\ndata: ${JSON.stringify(finalEvent)}\n\n`), + ); + } + + controller.close(); } catch (error) { controller.error(error); } finally { diff --git a/src/routes/anthropic.test.ts b/src/routes/anthropic.test.ts index be327ed..dcb4f39 100644 --- a/src/routes/anthropic.test.ts +++ b/src/routes/anthropic.test.ts @@ -1,12 +1,306 @@ -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 { filterAllowlistedEntities, type PIIDetectionResult, PIIDetector } from "../pii/detect"; import { AnthropicRequestSchema } from "../providers/anthropic/types"; -import { anthropicRoutes } from "./anthropic"; + +const noPII: PIIDetectionResult = { + hasPII: false, + spanEntities: [], + allEntities: [], + scanTimeMs: 0, +}; +const mockAnalyzeRequest = mock<() => Promise>(() => Promise.resolve(noPII)); +const mockLogRequest = mock(() => {}); + +mock.module("../pii/detect", () => ({ + PIIDetector, + filterAllowlistedEntities, + getPIIDetector: () => ({ + analyzeRequest: mockAnalyzeRequest, + detectPII: mock(() => Promise.resolve([])), + healthCheck: mock(() => Promise.resolve(true)), + }), +})); + +mock.module("../logging/logger", () => ({ + logRequest: mockLogRequest, +})); + +const { anthropicRoutes } = await import("./anthropic"); const app = new Hono(); app.route("/anthropic", anthropicRoutes); +const originalFetch = globalThis.fetch; +const config = getConfig(); +const originalMode = config.mode; +const originalPiiScanRoles = [...config.pii_detection.scan_roles]; +const originalSecretsEnabled = config.secrets_detection.enabled; +const originalSecretsAction = config.secrets_detection.action; +const originalSecretsScanRoles = [...config.secrets_detection.scan_roles]; + +afterEach(() => { + globalThis.fetch = originalFetch; + config.mode = originalMode; + config.pii_detection.scan_roles = [...originalPiiScanRoles]; + config.secrets_detection.enabled = originalSecretsEnabled; + config.secrets_detection.action = originalSecretsAction; + config.secrets_detection.scan_roles = [...originalSecretsScanRoles]; + mockAnalyzeRequest.mockClear(); + mockAnalyzeRequest.mockResolvedValue(noPII); + mockLogRequest.mockClear(); +}); + describe("POST /anthropic/v1/messages", () => { + test("remasks known PII in assistant tool-use history before forwarding", async () => { + const email = "jane@example.com"; + const userText = `Email ${email}`; + const entity = { + entity_type: "EMAIL_ADDRESS", + start: userText.indexOf(email), + end: userText.indexOf(email) + email.length, + score: 0.99, + }; + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[entity], [], []], + allEntities: [entity], + scanTimeMs: 2, + }); + + let upstreamBody: Record | undefined; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ + id: "msg_test", + type: "message", + role: "assistant", + content: [], + model: "claude-test", + stop_reason: "end_turn", + stop_sequence: null, + usage: { input_tokens: 1, output_tokens: 1 }, + }); + }) as typeof fetch; + + config.mode = "mask"; + config.pii_detection.scan_roles = ["user", "tool", "function", "mcp"]; + const response = await app.request("/anthropic/v1/messages", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "claude-test", + max_tokens: 128, + messages: [ + { role: "user", content: userText }, + { + role: "assistant", + content: [ + { type: "text", text: `I will use ${email}` }, + { type: "thinking", thinking: `Consider ${email}`, signature: "sig_test" }, + { + type: "tool_use", + id: "tool_123", + name: "send_email", + input: { email }, + cache_control: { type: "ephemeral" }, + }, + ], + }, + { + role: "user", + content: [{ type: "tool_result", tool_use_id: "tool_123", content: "sent" }], + }, + ], + }), + }); + + expect(response.status).toBe(200); + expect(config.pii_detection.scan_roles).not.toContain("assistant"); + expect(upstreamBody).toEqual({ + model: "claude-test", + max_tokens: 128, + messages: [ + { role: "user", content: "Email [[EMAIL_ADDRESS_1]]" }, + { + role: "assistant", + content: [ + { type: "text", text: `I will use ${email}` }, + { type: "thinking", thinking: `Consider ${email}`, signature: "sig_test" }, + { + type: "tool_use", + id: "tool_123", + name: "send_email", + input: { email: "[[EMAIL_ADDRESS_1]]" }, + cache_control: { type: "ephemeral" }, + }, + ], + }, + { + role: "user", + content: [{ type: "tool_result", tool_use_id: "tool_123", content: "sent" }], + }, + ], + }); + }); + + test("recursively remasks known secrets in assistant tool-use history", async () => { + const secret = "sk-proj-abc123def456ghi789jkl012mno345pqr678stu901vwx"; + let upstreamBody: + | { + messages: Array<{ role: string; content: unknown }>; + } + | undefined; + globalThis.fetch = (async (_target: string | URL | Request, init?: RequestInit) => { + upstreamBody = JSON.parse(String(init?.body)); + return Response.json({ + id: "msg_secret", + type: "message", + role: "assistant", + content: [], + model: "claude-test", + stop_reason: "end_turn", + stop_sequence: null, + usage: { input_tokens: 1, output_tokens: 1 }, + }); + }) as typeof fetch; + + config.mode = "mask"; + config.secrets_detection.enabled = true; + config.secrets_detection.action = "mask"; + config.secrets_detection.scan_roles = ["user", "tool", "function", "mcp"]; + const response = await app.request("/anthropic/v1/messages", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "claude-test", + max_tokens: 128, + messages: [ + { role: "user", content: `Use key ${secret}` }, + { + role: "assistant", + content: [ + { + type: "tool_use", + id: "tool_secret", + name: "configure_service", + input: { + auth: { + values: [secret, { header: `Bearer ${secret}` }], + enabled: true, + retries: 3, + fallback: null, + }, + }, + vendor_metadata: { stable: true }, + }, + ], + }, + { + role: "user", + content: [{ type: "tool_result", tool_use_id: "tool_secret", content: "configured" }], + }, + ], + }), + }); + + expect(response.status).toBe(200); + expect(config.secrets_detection.scan_roles).not.toContain("assistant"); + expect(upstreamBody?.messages).toEqual([ + { role: "user", content: "Use key [[API_KEY_SK_1]]" }, + { + role: "assistant", + content: [ + { + type: "tool_use", + id: "tool_secret", + name: "configure_service", + input: { + auth: { + values: ["[[API_KEY_SK_1]]", { header: "Bearer [[API_KEY_SK_1]]" }], + enabled: true, + retries: 3, + fallback: null, + }, + }, + vendor_metadata: { stable: true }, + }, + ], + }, + { + role: "user", + content: [{ type: "tool_result", tool_use_id: "tool_secret", content: "configured" }], + }, + ]); + }); + + test("keeps the mocked streaming provider boundary masked outbound and lossless inbound", async () => { + const email = "stream@example.com"; + const userText = `Email ${email}`; + const entity = { + entity_type: "EMAIL_ADDRESS", + start: userText.indexOf(email), + end: userText.indexOf(email) + email.length, + score: 0.99, + }; + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[entity]], + allEntities: [entity], + scanTimeMs: 2, + }); + const toolJson = '{ "big": 9007199254740993, "fixed": 1.0, "email": "[[EMAIL_ADDRESS_1]]" }'; + const delta = `event: content_block_delta\ndata: ${JSON.stringify({ + type: "content_block_delta", + index: 0, + delta: { type: "input_json_delta", partial_json: toolJson }, + })}\n\n`; + const stop = `event: content_block_stop\ndata: ${JSON.stringify({ + type: "content_block_stop", + index: 0, + })}\n\n`; + let upstreamUrl: string | undefined; + let upstreamBody: Record | undefined; + globalThis.fetch = (async (target: string | URL | Request, init?: RequestInit) => { + const request = target instanceof Request ? target : new Request(target, init); + upstreamUrl = request.url; + upstreamBody = JSON.parse(await request.clone().text()); + return new Response(delta + stop, { headers: { "Content-Type": "text/event-stream" } }); + }) as typeof fetch; + + config.mode = "mask"; + const response = await app.request("/anthropic/v1/messages", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "claude-test", + max_tokens: 128, + stream: true, + messages: [{ role: "user", content: userText }], + }), + }); + const responseText = await response.text(); + const streamedToolJson = responseText + .split("\n") + .filter((line) => line.startsWith("data: ")) + .map((line) => JSON.parse(line.slice(6))) + .filter((event) => event.delta?.type === "input_json_delta") + .map((event) => event.delta.partial_json) + .join(""); + + expect(response.status).toBe(200); + expect(response.headers.get("Content-Type")).toContain("text/event-stream"); + expect(upstreamUrl).toBe(`${config.providers.anthropic!.base_url}/v1/messages`); + expect(upstreamBody).toEqual({ + model: "claude-test", + max_tokens: 128, + stream: true, + messages: [{ role: "user", content: "Email [[EMAIL_ADDRESS_1]]" }], + }); + expect(streamedToolJson).toBe(toolJson.replace("[[EMAIL_ADDRESS_1]]", email)); + }); + test("returns 400 for missing messages", async () => { const res = await app.request("/anthropic/v1/messages", { method: "POST", diff --git a/src/routes/anthropic.ts b/src/routes/anthropic.ts index 19853f3..c52198c 100644 --- a/src/routes/anthropic.ts +++ b/src/routes/anthropic.ts @@ -5,7 +5,7 @@ import { getConfig } from "../config"; import { formatMaskedRequestForLog } from "../logging/log-content"; import { logRequest } from "../logging/logger"; import type { PlaceholderContext } from "../masking/context"; -import { anthropicExtractor } from "../masking/extractors/anthropic"; +import { anthropicExtractor, remaskAnthropicToolUseHistory } from "../masking/extractors/anthropic"; import { restoreResponse } from "../masking/restorer"; import type { PIIDetectResult } from "../pii/request"; import { @@ -121,10 +121,15 @@ anthropicRoutes.post( }); } + const providerRequest = remaskAnthropicToolUseHistory( + privacy.request, + privacy.piiMaskingContext, + secretsResult.maskingContext, + ); const maskedContent = - piiResult.hasPII || secretsResult.masked ? formatRequestForLog(privacy.request) : undefined; + piiResult.hasPII || secretsResult.masked ? formatRequestForLog(providerRequest) : undefined; - return sendToAnthropic(c, privacy.request, { + return sendToAnthropic(c, providerRequest, { startTime, piiResult, piiMaskingContext: privacy.piiMaskingContext,