From: Stefan Gasser Date: Sun, 28 Jun 2026 18:41:49 +0000 (+0200) Subject: Consolidate masking restoration into shared helpers (#124) X-Git-Tag: v0.7.3~6 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=c191a9e90f498ffbfe34f031cd20be54133d1046;p=sgasser-llm-shield.git Consolidate masking restoration into shared helpers (#124) * Add shared restoration helpers * Share stream restoration across providers * Extract request privacy pipeline * Remove deferred TODO note * Remove dead unmask wrappers after restoration refactor pii/mask and secrets/mask no longer re-export unmask/stream/flush/response wrappers; callers go through the shared StreamRestorer and restoreResponse. Drop the unused unmask re-export from masking/service and the restoreText helper. Simplify processPrivacyPipeline to derive PII masking from config.mode and drop the unused originalRequest/piiMasked result fields and the maskPII option. --- diff --git a/src/masking/restore-policy.test.ts b/src/masking/restore-policy.test.ts new file mode 100644 index 0000000..ab9ea75 --- /dev/null +++ b/src/masking/restore-policy.test.ts @@ -0,0 +1,32 @@ +import { describe, expect, test } from "bun:test"; +import type { MaskingConfig } from "../config"; +import { createRestoreFormatter } from "./restore-policy"; + +const defaultConfig: MaskingConfig = { + show_markers: false, + marker_text: "[protected]", + allowlist: [], + denylist: [], +}; + +describe("createRestoreFormatter", () => { + test("returns undefined when markers are disabled", () => { + expect(createRestoreFormatter(defaultConfig)).toBeUndefined(); + }); + + test("prefixes restored values when markers are enabled", () => { + const formatter = createRestoreFormatter({ ...defaultConfig, show_markers: true }); + + expect(formatter?.("secret")).toBe("[protected]secret"); + }); + + test("uses configured marker text", () => { + const formatter = createRestoreFormatter({ + ...defaultConfig, + show_markers: true, + marker_text: "[masked] ", + }); + + expect(formatter?.("jane@example.com")).toBe("[masked] jane@example.com"); + }); +}); diff --git a/src/masking/restore-policy.ts b/src/masking/restore-policy.ts new file mode 100644 index 0000000..579baf8 --- /dev/null +++ b/src/masking/restore-policy.ts @@ -0,0 +1,7 @@ +import type { MaskingConfig } from "../config"; + +export type RestoreFormatter = (original: string) => string; + +export function createRestoreFormatter(config: MaskingConfig): RestoreFormatter | undefined { + return config.show_markers ? (original: string) => `${config.marker_text}${original}` : undefined; +} diff --git a/src/masking/restorer.test.ts b/src/masking/restorer.test.ts new file mode 100644 index 0000000..3ed2b77 --- /dev/null +++ b/src/masking/restorer.test.ts @@ -0,0 +1,142 @@ +import { describe, expect, test } from "bun:test"; +import type { MaskingConfig } from "../config"; +import type { AnthropicResponse } from "../providers/anthropic/types"; +import type { OpenAIResponse } from "../providers/openai/types"; +import { createPlaceholderContext, type PlaceholderContext } from "./context"; +import { anthropicExtractor } from "./extractors/anthropic"; +import { type CodexResponsesResponse, codexExtractor } from "./extractors/codex"; +import { openaiExtractor } from "./extractors/openai"; +import { restoreResponse } from "./restorer"; +import type { RequestExtractor } from "./types"; + +interface TestResponse { + text: string; +} + +const defaultConfig: MaskingConfig = { + show_markers: false, + marker_text: "[protected]", + allowlist: [], + denylist: [], +}; + +const markerConfig: MaskingConfig = { ...defaultConfig, show_markers: true }; + +const extractor: RequestExtractor = { + extractTexts: () => [], + applyMasked: (request) => request, + unmaskResponse: (response, context, formatValue) => { + let text = response.text; + for (const [placeholder, original] of Object.entries(context.mapping)) { + text = text.split(placeholder).join(formatValue ? formatValue(original) : original); + } + return { ...response, text }; + }, +}; + +function context(mapping: Record): PlaceholderContext { + const ctx = createPlaceholderContext(); + ctx.mapping = mapping; + return ctx; +} + +describe("restoreResponse", () => { + test("returns response unchanged with no contexts", () => { + const response = { text: "Hello [[PERSON_1]]" }; + + expect(restoreResponse(response, extractor, defaultConfig, {})).toEqual(response); + }); + + test("restores PII only", () => { + const response = { text: "Hello [[PERSON_1]]" }; + + expect( + restoreResponse(response, extractor, defaultConfig, { + piiContext: context({ "[[PERSON_1]]": "Jane" }), + }), + ).toEqual({ text: "Hello Jane" }); + }); + + test("restores secrets only", () => { + const response = { text: "Key [[API_KEY_SK_1]]" }; + + expect( + restoreResponse(response, extractor, defaultConfig, { + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }), + ).toEqual({ text: "Key sk-secret" }); + }); + + test("restores PII before secrets with the same marker policy", () => { + const response = { text: "[[PERSON_1]] used [[API_KEY_SK_1]]" }; + + expect( + restoreResponse( + response, + extractor, + { ...defaultConfig, show_markers: true }, + { + piiContext: context({ "[[PERSON_1]]": "Jane" }), + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }, + ), + ).toEqual({ text: "[protected]Jane used [protected]sk-secret" }); + }); +}); + +describe("restoreResponse applies markers through each provider extractor", () => { + test("OpenAI response", () => { + const response: OpenAIResponse = { + id: "test", + object: "chat.completion", + created: 0, + model: "gpt-4", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Your key is [[API_KEY_SK_1]]" }, + finish_reason: "stop", + }, + ], + }; + + const result = restoreResponse(response, openaiExtractor, markerConfig, { + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }); + + expect(result.choices[0].message.content).toBe("Your key is [protected]sk-secret"); + }); + + test("Anthropic response", () => { + const response: AnthropicResponse = { + id: "msg_test", + type: "message", + role: "assistant", + content: [{ type: "text", text: "Your key is [[API_KEY_SK_1]]" }], + model: "claude-3-5-sonnet", + stop_reason: "end_turn", + stop_sequence: null, + usage: { input_tokens: 10, output_tokens: 5 }, + }; + + const result = restoreResponse(response, anthropicExtractor, markerConfig, { + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }); + + expect(result.content[0]).toEqual({ type: "text", text: "Your key is [protected]sk-secret" }); + }); + + test("Codex response", () => { + const response: CodexResponsesResponse = { + output: [{ content: [{ type: "output_text", text: "Your key is [[API_KEY_SK_1]]" }] }], + }; + + const result = restoreResponse(response, codexExtractor, markerConfig, { + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }); + + expect(result).toEqual({ + output: [{ content: [{ type: "output_text", text: "Your key is [protected]sk-secret" }] }], + }); + }); +}); diff --git a/src/masking/restorer.ts b/src/masking/restorer.ts new file mode 100644 index 0000000..a531213 --- /dev/null +++ b/src/masking/restorer.ts @@ -0,0 +1,29 @@ +import type { MaskingConfig } from "../config"; +import type { PlaceholderContext } from "./context"; +import { createRestoreFormatter } from "./restore-policy"; +import type { RequestExtractor } from "./types"; + +export interface RestoreContexts { + piiContext?: PlaceholderContext; + secretsContext?: PlaceholderContext; +} + +export function restoreResponse( + response: TResponse, + extractor: RequestExtractor, + config: MaskingConfig, + contexts: RestoreContexts, +): TResponse { + const formatValue = createRestoreFormatter(config); + let result = response; + + if (contexts.piiContext) { + result = extractor.unmaskResponse(result, contexts.piiContext, formatValue); + } + + if (contexts.secretsContext) { + result = extractor.unmaskResponse(result, contexts.secretsContext, formatValue); + } + + return result; +} diff --git a/src/masking/service.ts b/src/masking/service.ts index 76df74b..6088686 100644 --- a/src/masking/service.ts +++ b/src/masking/service.ts @@ -100,21 +100,6 @@ export function createMaskingContext(): PlaceholderContext { return createPlaceholderContext(); } -/** - * Unmasks text by replacing placeholders with original values - * - * @param text - Text containing placeholders - * @param context - Masking context with mappings - * @param formatValue - Optional function to format restored values - */ -export function unmask( - text: string, - context: PlaceholderContext, - formatValue?: (original: string) => string, -): string { - return restorePlaceholders(text, context, formatValue); -} - /** * Processes a stream chunk, buffering partial placeholders * diff --git a/src/masking/stream-restorer.test.ts b/src/masking/stream-restorer.test.ts new file mode 100644 index 0000000..5a2650c --- /dev/null +++ b/src/masking/stream-restorer.test.ts @@ -0,0 +1,85 @@ +import { describe, expect, test } from "bun:test"; +import type { MaskingConfig } from "../config"; +import { createPlaceholderContext, type PlaceholderContext } from "./context"; +import { StreamRestorer } from "./stream-restorer"; + +const defaultConfig: MaskingConfig = { + show_markers: false, + marker_text: "[protected]", + allowlist: [], + denylist: [], +}; + +function context(mapping: Record): PlaceholderContext { + const ctx = createPlaceholderContext(); + ctx.mapping = mapping; + return ctx; +} + +describe("StreamRestorer", () => { + test("passes chunks through unchanged with no contexts", () => { + const restorer = new StreamRestorer({ config: defaultConfig }); + + expect(restorer.restoreChunk("Hello [[PERSON_1]]")).toBe("Hello [[PERSON_1]]"); + expect(restorer.flush()).toBe(""); + }); + + test("restores PII placeholders split across chunks", () => { + const restorer = new StreamRestorer({ + config: defaultConfig, + piiContext: context({ "[[EMAIL_ADDRESS_1]]": "jane@example.com" }), + }); + + expect(restorer.restoreChunk("Email [[EMAIL")).toBe("Email "); + expect(restorer.restoreChunk("_ADDRESS_1]] sent")).toBe("jane@example.com sent"); + expect(restorer.flush()).toBe(""); + }); + + test("restores secret placeholders split across chunks", () => { + const restorer = new StreamRestorer({ + config: defaultConfig, + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }); + + expect(restorer.restoreChunk("Key [[API_KEY")).toBe("Key "); + expect(restorer.restoreChunk("_SK_1]] ready")).toBe("sk-secret ready"); + expect(restorer.flush()).toBe(""); + }); + + test("keeps PII and secrets buffers independent", () => { + const restorer = new StreamRestorer({ + config: defaultConfig, + piiContext: context({ "[[PERSON_1]]": "Jane" }), + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }); + + expect(restorer.restoreChunk("[[PERSON")).toBe(""); + expect(restorer.restoreChunk("_1]] uses [[API")).toBe("Jane uses "); + expect(restorer.restoreChunk("_KEY_SK_1]]")).toBe("sk-secret"); + expect(restorer.flush()).toBe(""); + }); + + test("flushes buffered PII then secrets", () => { + const restorer = new StreamRestorer({ + config: defaultConfig, + piiContext: context({ "[[PERSON_1]]": "Jane" }), + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }); + + expect(restorer.restoreChunk("[[PERSON")).toBe(""); + expect(restorer.restoreChunk("_1]][[API_KEY")).toBe("Jane"); + expect(restorer.flush()).toBe("[[API_KEY"); + }); + + test("applies markers to restored PII and secrets", () => { + const restorer = new StreamRestorer({ + config: { ...defaultConfig, show_markers: true }, + piiContext: context({ "[[PERSON_1]]": "Jane" }), + secretsContext: context({ "[[API_KEY_SK_1]]": "sk-secret" }), + }); + + expect(restorer.restoreChunk("[[PERSON_1]] and [[API_KEY_SK_1]]")).toBe( + "[protected]Jane and [protected]sk-secret", + ); + }); +}); diff --git a/src/masking/stream-restorer.ts b/src/masking/stream-restorer.ts new file mode 100644 index 0000000..e9297ef --- /dev/null +++ b/src/masking/stream-restorer.ts @@ -0,0 +1,68 @@ +import type { MaskingConfig } from "../config"; +import type { PlaceholderContext } from "./context"; +import { createRestoreFormatter } from "./restore-policy"; +import { flushMaskingBuffer, unmaskStreamChunk } from "./service"; + +export interface StreamRestorerOptions { + piiContext?: PlaceholderContext; + secretsContext?: PlaceholderContext; + config: MaskingConfig; +} + +export class StreamRestorer { + private piiBuffer = ""; + private secretsBuffer = ""; + private readonly formatValue: ((original: string) => string) | undefined; + + constructor(private readonly options: StreamRestorerOptions) { + this.formatValue = createRestoreFormatter(options.config); + } + + restoreChunk(text: string): string { + let processedText = text; + + if (this.options.piiContext) { + const { output, remainingBuffer } = unmaskStreamChunk( + this.piiBuffer, + processedText, + this.options.piiContext, + this.formatValue, + ); + this.piiBuffer = remainingBuffer; + processedText = output; + } + + if (this.options.secretsContext && processedText) { + const { output, remainingBuffer } = unmaskStreamChunk( + this.secretsBuffer, + processedText, + this.options.secretsContext, + this.formatValue, + ); + this.secretsBuffer = remainingBuffer; + processedText = output; + } + + return processedText; + } + + flush(): string { + let flushed = ""; + + if (this.options.piiContext && this.piiBuffer) { + flushed = flushMaskingBuffer(this.piiBuffer, this.options.piiContext, this.formatValue); + this.piiBuffer = ""; + } + + if (this.options.secretsContext && this.secretsBuffer) { + flushed += flushMaskingBuffer( + this.secretsBuffer, + this.options.secretsContext, + this.formatValue, + ); + this.secretsBuffer = ""; + } + + return flushed; + } +} diff --git a/src/pii/mask.test.ts b/src/pii/mask.test.ts index ae71712..8d293b6 100644 --- a/src/pii/mask.test.ts +++ b/src/pii/mask.test.ts @@ -1,32 +1,10 @@ import { describe, expect, test } from "bun:test"; -import type { MaskingConfig } from "../config"; +import { restorePlaceholders } from "../masking/context"; import { openaiExtractor } from "../masking/extractors/openai"; -import type { OpenAIMessage, OpenAIRequest, OpenAIResponse } from "../providers/openai/types"; +import type { OpenAIMessage, OpenAIRequest } from "../providers/openai/types"; import { createPIIResultFromSpans } from "../test-utils/detection-results"; import type { PIIEntity } from "./detect"; -import { - createMaskingContext, - flushMaskingBuffer, - mask, - maskRequest, - unmask, - unmaskResponse, - unmaskStreamChunk, -} from "./mask"; - -const defaultConfig: MaskingConfig = { - show_markers: false, - marker_text: "[protected]", - allowlist: [], - denylist: [], -}; - -const configWithMarkers: MaskingConfig = { - show_markers: true, - marker_text: "[protected]", - allowlist: [], - denylist: [], -}; +import { mask, maskRequest } from "./mask"; /** Helper to create a minimal request from messages */ function createRequest(messages: OpenAIMessage[]): OpenAIRequest { @@ -64,54 +42,6 @@ describe("PII placeholder format", () => { }); }); -describe("marker feature", () => { - test("adds markers when show_markers is true", () => { - const context = createMaskingContext(); - context.mapping["[[EMAIL_ADDRESS_1]]"] = "john@example.com"; - - const result = unmask("Email: [[EMAIL_ADDRESS_1]]", context, configWithMarkers); - expect(result).toBe("Email: [protected]john@example.com"); - }); - - test("no markers when show_markers is false", () => { - const context = createMaskingContext(); - context.mapping["[[EMAIL_ADDRESS_1]]"] = "john@example.com"; - - const result = unmask("Email: [[EMAIL_ADDRESS_1]]", context, defaultConfig); - expect(result).toBe("Email: john@example.com"); - }); - - test("markers work with streaming", () => { - const context = createMaskingContext(); - context.mapping["[[PERSON_1]]"] = "John Doe"; - - const { output } = unmaskStreamChunk("", "Hello [[PERSON_1]]!", context, configWithMarkers); - expect(output).toBe("Hello [protected]John Doe!"); - }); - - test("markers work with response unmasking", () => { - const context = createMaskingContext(); - context.mapping["[[PERSON_1]]"] = "John Doe"; - - const response: OpenAIResponse = { - id: "test", - object: "chat.completion", - created: 1234567890, - model: "gpt-4", - choices: [ - { - index: 0, - message: { role: "assistant", content: "Hello [[PERSON_1]]" }, - finish_reason: "stop", - }, - ], - }; - - const result = unmaskResponse(response, context, configWithMarkers, openaiExtractor); - expect(result.choices[0].message.content).toBe("Hello [protected]John Doe"); - }); -}); - describe("maskRequest with PIIDetectionResult", () => { test("masks multiple messages using detection result", () => { const request = createRequest([ @@ -160,46 +90,6 @@ describe("maskRequest with PIIDetectionResult", () => { }); }); -describe("streaming with PII placeholders", () => { - test("buffers partial [[TYPE placeholder", () => { - const context = createMaskingContext(); - context.mapping["[[EMAIL_ADDRESS_1]]"] = "test@test.com"; - - const { output, remainingBuffer } = unmaskStreamChunk( - "", - "Hello [[EMAIL_ADD", - context, - defaultConfig, - ); - - expect(output).toBe("Hello "); - expect(remainingBuffer).toBe("[[EMAIL_ADD"); - }); - - test("completes buffered placeholder across chunks", () => { - const context = createMaskingContext(); - context.mapping["[[EMAIL_ADDRESS_1]]"] = "test@test.com"; - - const { output, remainingBuffer } = unmaskStreamChunk( - "[[EMAIL_ADD", - "RESS_1]] there", - context, - defaultConfig, - ); - - expect(output).toBe("test@test.com there"); - expect(remainingBuffer).toBe(""); - }); - - test("flushes remaining buffer at end of stream", () => { - const context = createMaskingContext(); - context.mapping["[[EMAIL_ADDRESS_1]]"] = "test@test.com"; - - const flushed = flushMaskingBuffer("[[EMAIL_ADD", context, defaultConfig); - expect(flushed).toBe("[[EMAIL_ADD"); - }); -}); - describe("PII conflict resolution", () => { test("handles overlapping entities with same start - keeps longer", () => { const text = "Given Eric's feedback"; @@ -255,77 +145,13 @@ describe("mask -> unmask roundtrip", () => { expect(masked).not.toContain("+49123456789"); const llmResponse = `I see ${masked.match(/\[\[PERSON_1\]\]/)?.[0]}, email ${masked.match(/\[\[EMAIL_ADDRESS_1\]\]/)?.[0]}`; - const unmasked = unmask(llmResponse, context, defaultConfig); + const unmasked = restorePlaceholders(llmResponse, context); expect(unmasked).toContain("Hans Müller"); expect(unmasked).toContain("hans@firma.de"); }); }); -describe("HTML context handling", () => { - test("unmasks placeholders in HTML without encoding issues", () => { - const context = createMaskingContext(); - context.mapping["[[PERSON_1]]"] = "Dr. Sarah Chen"; - context.mapping["[[EMAIL_ADDRESS_1]]"] = "sarah.chen@hospital.org"; - - const htmlResponse = `

Contact [[PERSON_1]] at [[EMAIL_ADDRESS_1]]

`; - const result = unmask(htmlResponse, context, defaultConfig); - - expect(result).toBe("

Contact Dr. Sarah Chen at sarah.chen@hospital.org

"); - }); - - test("works with complex HTML structures", () => { - const context = createMaskingContext(); - context.mapping["[[PERSON_1]]"] = "Dr. Sarah Chen"; - context.mapping["[[EMAIL_ADDRESS_1]]"] = "sarah@hospital.org"; - - const complexHtml = ` -
-

[[PERSON_1]]

- [[EMAIL_ADDRESS_1]] -
- `; - - const result = unmask(complexHtml, context, defaultConfig); - - expect(result).toContain("Dr. Sarah Chen"); - expect(result).toContain("sarah@hospital.org"); - expect(result).not.toContain("[["); - }); -}); - -describe("unmaskResponse", () => { - test("unmasks all choices in response", () => { - const context = createMaskingContext(); - context.mapping["[[EMAIL_ADDRESS_1]]"] = "test@test.com"; - context.mapping["[[PERSON_1]]"] = "John Doe"; - - const response: OpenAIResponse = { - id: "chatcmpl-123", - object: "chat.completion", - created: 1234567890, - model: "gpt-4", - choices: [ - { - index: 0, - message: { - role: "assistant", - content: "Contact [[PERSON_1]] at [[EMAIL_ADDRESS_1]]", - }, - finish_reason: "stop", - }, - ], - usage: { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 }, - }; - - const result = unmaskResponse(response, context, defaultConfig, openaiExtractor); - - expect(result.choices[0].message.content).toBe("Contact John Doe at test@test.com"); - expect(result.id).toBe("chatcmpl-123"); - expect(result.model).toBe("gpt-4"); - }); -}); - describe("edge cases", () => { test("handles unicode in masked text", () => { const text = "Kontakt: François Müller"; @@ -334,14 +160,14 @@ describe("edge cases", () => { const { masked, context } = mask(text, entities); expect(masked).toBe("Kontakt: [[PERSON_1]]"); - const unmasked = unmask(masked, context, defaultConfig); + const unmasked = restorePlaceholders(masked, context); expect(unmasked).toBe("Kontakt: François Müller"); }); test("handles empty text", () => { const { masked, context } = mask("", []); expect(masked).toBe(""); - expect(unmask("", context, defaultConfig)).toBe(""); + expect(restorePlaceholders("", context)).toBe(""); }); test("reuses placeholder for duplicate values", () => { diff --git a/src/pii/mask.ts b/src/pii/mask.ts index befdb11..dec33f9 100644 --- a/src/pii/mask.ts +++ b/src/pii/mask.ts @@ -1,18 +1,10 @@ -import type { MaskingConfig } from "../config"; import { resolveConflicts } from "../masking/conflict-resolver"; import { incrementAndGenerate } from "../masking/context"; import { generatePlaceholder as generatePlaceholderFromFormat, PII_PLACEHOLDER_FORMAT, } from "../masking/placeholders"; -import { - flushMaskingBuffer as flushBuffer, - type MaskSpansResult, - maskSpans, - type PlaceholderContext, - unmaskStreamChunk as unmaskChunk, - unmask as unmaskText, -} from "../masking/service"; +import { type MaskSpansResult, maskSpans, type PlaceholderContext } from "../masking/service"; import type { RequestExtractor, TextSpan } from "../masking/types"; import type { PIIDetectionResult, PIIEntity } from "./detect"; @@ -29,10 +21,6 @@ function generatePlaceholder(entityType: string, context: PlaceholderContext): s ); } -function getFormatValue(config: MaskingConfig): ((original: string) => string) | undefined { - return config.show_markers ? (original: string) => `${config.marker_text}${original}` : undefined; -} - export function mask( text: string, entities: PIIEntity[], @@ -56,27 +44,6 @@ export function mask( }; } -export function unmask(text: string, context: PlaceholderContext, config: MaskingConfig): string { - return unmaskText(text, context, getFormatValue(config)); -} - -export function unmaskStreamChunk( - buffer: string, - newChunk: string, - context: PlaceholderContext, - config: MaskingConfig, -): { output: string; remainingBuffer: string } { - return unmaskChunk(buffer, newChunk, context, getFormatValue(config)); -} - -export function flushMaskingBuffer( - buffer: string, - context: PlaceholderContext, - config: MaskingConfig, -): string { - return flushBuffer(buffer, context, getFormatValue(config)); -} - export interface MaskRequestResult { request: TRequest; context: PlaceholderContext; @@ -119,12 +86,3 @@ function maskSpansWithEntities( existingContext, ); } - -export function unmaskResponse( - response: TResponse, - context: PlaceholderContext, - config: MaskingConfig, - extractor: RequestExtractor, -): TResponse { - return extractor.unmaskResponse(response, context, getFormatValue(config)); -} diff --git a/src/providers/anthropic/stream-transformer.ts b/src/providers/anthropic/stream-transformer.ts index 491edf9..c39852f 100644 --- a/src/providers/anthropic/stream-transformer.ts +++ b/src/providers/anthropic/stream-transformer.ts @@ -3,8 +3,7 @@ import type { MaskingConfig } from "../../config"; import type { PlaceholderContext } from "../../masking/context"; -import { flushMaskingBuffer, unmaskStreamChunk } from "../../pii/mask"; -import { flushSecretsMaskingBuffer, unmaskSecretsStreamChunk } from "../../secrets/mask"; +import { StreamRestorer } from "../../masking/stream-restorer"; import type { ContentBlockDeltaEvent, TextDelta } from "./types"; export function createAnthropicUnmaskingStream( @@ -15,9 +14,8 @@ export function createAnthropicUnmaskingStream( ): ReadableStream { const decoder = new TextDecoder(); const encoder = new TextEncoder(); - let piiBuffer = ""; - let secretsBuffer = ""; let lineBuffer = ""; + const restorer = new StreamRestorer({ piiContext, secretsContext, config }); return new ReadableStream({ async start(controller) { @@ -28,20 +26,7 @@ export function createAnthropicUnmaskingStream( const { done, value } = await reader.read(); if (done) { - // Flush remaining buffers - let flushed = ""; - - if (piiBuffer && piiContext) { - flushed = flushMaskingBuffer(piiBuffer, piiContext, config); - } else if (piiBuffer) { - flushed = piiBuffer; - } - - if (secretsBuffer && secretsContext) { - flushed += flushSecretsMaskingBuffer(secretsBuffer, secretsContext, config); - } else if (secretsBuffer) { - flushed += secretsBuffer; - } + const flushed = restorer.flush(); // Send flushed content as final text delta if (flushed) { @@ -83,31 +68,7 @@ export function createAnthropicUnmaskingStream( if (parsed.type === "content_block_delta" && parsed.delta?.type === "text_delta") { const event = parsed as ContentBlockDeltaEvent; const textDelta = event.delta as TextDelta; - let processedText = textDelta.text; - - // Unmask PII - if (piiContext && processedText) { - const { output, remainingBuffer } = unmaskStreamChunk( - piiBuffer, - processedText, - piiContext, - config, - ); - piiBuffer = remainingBuffer; - processedText = output; - } - - // Unmask secrets - if (secretsContext && processedText) { - const { output, remainingBuffer } = unmaskSecretsStreamChunk( - secretsBuffer, - processedText, - secretsContext, - config, - ); - secretsBuffer = remainingBuffer; - processedText = output; - } + const processedText = restorer.restoreChunk(textDelta.text); // Only emit if we have content if (processedText) { diff --git a/src/providers/codex/stream-transformer.test.ts b/src/providers/codex/stream-transformer.test.ts new file mode 100644 index 0000000..4c02370 --- /dev/null +++ b/src/providers/codex/stream-transformer.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, test } from "bun:test"; +import type { MaskingConfig } from "../../config"; +import { createPlaceholderContext, type PlaceholderContext } from "../../masking/context"; +import { createCodexUnmaskingStream } from "./stream-transformer"; + +const defaultConfig: MaskingConfig = { + show_markers: false, + marker_text: "[protected]", + allowlist: [], + denylist: [], +}; + +function context(mapping: Record): PlaceholderContext { + const ctx = createPlaceholderContext(); + ctx.mapping = mapping; + return ctx; +} + +function createSSEStream(chunks: string[]): ReadableStream { + const encoder = new TextEncoder(); + let index = 0; + + return new ReadableStream({ + pull(controller) { + if (index < chunks.length) { + controller.enqueue(encoder.encode(chunks[index])); + index++; + } else { + controller.close(); + } + }, + }); +} + +async function consumeStream(stream: ReadableStream): Promise { + const reader = stream.getReader(); + const decoder = new TextDecoder(); + let result = ""; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + result += decoder.decode(value, { stream: true }); + } + + return result; +} + +function codexDelta(text: string): string { + return `data: ${JSON.stringify({ type: "response.output_text.delta", delta: text })}\n\n`; +} + +describe("createCodexUnmaskingStream", () => { + test("restores complete placeholders", async () => { + const piiContext = context({ "[[EMAIL_ADDRESS_1]]": "jane@example.com" }); + const source = createSSEStream([codexDelta("Email [[EMAIL_ADDRESS_1]]")]); + + const result = await consumeStream( + createCodexUnmaskingStream(source, piiContext, defaultConfig), + ); + + expect(result).toContain("Email jane@example.com"); + }); + + test("buffers placeholders split across SSE events", async () => { + const piiContext = context({ "[[EMAIL_ADDRESS_1]]": "jane@example.com" }); + const source = createSSEStream([codexDelta("Email [[EMAIL_"), codexDelta("ADDRESS_1]] done")]); + + const result = await consumeStream( + createCodexUnmaskingStream(source, piiContext, defaultConfig), + ); + + expect(result).toContain("jane@example.com done"); + expect(result).not.toContain("[[EMAIL_ADDRESS_1]]"); + }); + + test("restores PII and secrets with markers", async () => { + const piiContext = context({ "[[PERSON_1]]": "Jane" }); + const secretsContext = context({ "[[API_KEY_SK_1]]": "sk-secret" }); + const source = createSSEStream([codexDelta("[[PERSON_1]] used [[API_KEY_SK_1]]")]); + + const result = await consumeStream( + createCodexUnmaskingStream( + source, + piiContext, + { ...defaultConfig, show_markers: true }, + secretsContext, + ), + ); + + expect(result).toContain("[protected]Jane used [protected]sk-secret"); + }); + + test("passes malformed JSON and done events through", async () => { + const source = createSSEStream(["data: not-json\n\n", "data: [DONE]\n\n"]); + + const result = await consumeStream( + createCodexUnmaskingStream(source, undefined, defaultConfig), + ); + + expect(result).toContain("data: not-json"); + expect(result).toContain("data: [DONE]"); + }); + + test("emits Codex-compatible final flush events", async () => { + const piiContext = context({ "[[EMAIL_ADDRESS_1]]": "jane@example.com" }); + const source = createSSEStream([codexDelta("Email [[EMAIL")]); + + const result = await consumeStream( + createCodexUnmaskingStream(source, piiContext, defaultConfig), + ); + + expect(result).toContain('"type":"response.output_text.delta"'); + expect(result).toContain('"delta":"[[EMAIL"'); + }); +}); diff --git a/src/providers/codex/stream-transformer.ts b/src/providers/codex/stream-transformer.ts new file mode 100644 index 0000000..5aa93be --- /dev/null +++ b/src/providers/codex/stream-transformer.ts @@ -0,0 +1,102 @@ +import type { MaskingConfig } from "../../config"; +import type { PlaceholderContext } from "../../masking/context"; +import { type CodexResponsesResponse, codexExtractor } from "../../masking/extractors/codex"; +import { StreamRestorer } from "../../masking/stream-restorer"; + +export function createCodexUnmaskingStream( + stream: ReadableStream, + piiContext: PlaceholderContext | undefined, + maskingConfig: MaskingConfig, + secretsContext?: PlaceholderContext, +): ReadableStream { + const decoder = new TextDecoder(); + const encoder = new TextEncoder(); + let lineBuffer = ""; + const restorer = new StreamRestorer({ + piiContext, + secretsContext, + config: maskingConfig, + }); + + function unmaskPayload(payload: unknown): unknown { + const result = payload as CodexResponsesResponse; + const spans = codexExtractor.extractTexts(result); + + if (spans.length === 0) { + return result; + } + + return codexExtractor.applyMasked( + result, + spans.map((span) => ({ + ...span, + maskedText: restorer.restoreChunk(span.text), + })), + ); + } + + function processLine(line: string): string { + if (!line.startsWith("data: ")) { + return `${line}\n`; + } + + const data = line.slice(6); + if (data === "[DONE]") { + return "data: [DONE]\n"; + } + + try { + return `data: ${JSON.stringify(unmaskPayload(JSON.parse(data)))}\n`; + } catch { + return `${line}\n`; + } + } + + return new ReadableStream({ + async start(controller) { + const reader = stream.getReader(); + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + lineBuffer += decoder.decode(value, { stream: true }); + const lines = lineBuffer.split("\n"); + lineBuffer = lines.pop() ?? ""; + + let output = ""; + for (const line of lines) { + output += processLine(line); + } + + if (output) { + controller.enqueue(encoder.encode(output)); + } + } + + lineBuffer += decoder.decode(); + let finalOutput = lineBuffer ? processLine(lineBuffer) : ""; + lineBuffer = ""; + + const flushed = restorer.flush(); + if (flushed) { + finalOutput += `data: ${JSON.stringify({ + type: "response.output_text.delta", + delta: flushed, + })}\n\n`; + } + + if (finalOutput) { + controller.enqueue(encoder.encode(finalOutput)); + } + + controller.close(); + } catch (error) { + controller.error(error); + } finally { + reader.releaseLock(); + } + }, + }); +} diff --git a/src/providers/openai/stream-transformer.ts b/src/providers/openai/stream-transformer.ts index bec613d..f783588 100644 --- a/src/providers/openai/stream-transformer.ts +++ b/src/providers/openai/stream-transformer.ts @@ -1,46 +1,8 @@ import type { MaskingConfig } from "../../config"; import type { PlaceholderContext } from "../../masking/context"; -import { flushMaskingBuffer, unmaskStreamChunk } from "../../pii/mask"; -import { flushSecretsMaskingBuffer, unmaskSecretsStreamChunk } from "../../secrets/mask"; +import { StreamRestorer } from "../../masking/stream-restorer"; import type { OpenAIContentPart } from "../../utils/content"; -function unmaskTextContent( - text: string, - piiBuffer: string, - piiContext: PlaceholderContext | undefined, - config: MaskingConfig, - secretsBuffer: string, - secretsContext?: PlaceholderContext, -): { text: string; piiBuffer: string; secretsBuffer: string } { - let processedText = text; - let nextPiiBuffer = piiBuffer; - let nextSecretsBuffer = secretsBuffer; - - if (piiContext) { - const { output, remainingBuffer } = unmaskStreamChunk( - nextPiiBuffer, - processedText, - piiContext, - config, - ); - nextPiiBuffer = remainingBuffer; - processedText = output; - } - - if (secretsContext && processedText) { - const { output, remainingBuffer } = unmaskSecretsStreamChunk( - nextSecretsBuffer, - processedText, - secretsContext, - config, - ); - nextSecretsBuffer = remainingBuffer; - processedText = output; - } - - return { text: processedText, piiBuffer: nextPiiBuffer, secretsBuffer: nextSecretsBuffer }; -} - export function createUnmaskingStream( source: ReadableStream, piiContext: PlaceholderContext | undefined, @@ -49,9 +11,8 @@ export function createUnmaskingStream( ): ReadableStream { const decoder = new TextDecoder(); const encoder = new TextEncoder(); - let piiBuffer = ""; - let secretsBuffer = ""; let lineBuffer = ""; + const restorer = new StreamRestorer({ piiContext, secretsContext, config }); return new ReadableStream({ async start(controller) { @@ -71,19 +32,10 @@ export function createUnmaskingStream( const content = parsed.choices?.[0]?.delta?.content; if (typeof content === "string" && content !== "") { - const unmasked = unmaskTextContent( - content, - piiBuffer, - piiContext, - config, - secretsBuffer, - secretsContext, - ); - piiBuffer = unmasked.piiBuffer; - secretsBuffer = unmasked.secretsBuffer; - - if (unmasked.text) { - parsed.choices[0].delta.content = unmasked.text; + const text = restorer.restoreChunk(content); + + if (text) { + parsed.choices[0].delta.content = text; controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); } } else if (Array.isArray(content)) { @@ -92,22 +44,13 @@ export function createUnmaskingStream( return [part]; } - const unmasked = unmaskTextContent( - part.text, - piiBuffer, - piiContext, - config, - secretsBuffer, - secretsContext, - ); - piiBuffer = unmasked.piiBuffer; - secretsBuffer = unmasked.secretsBuffer; - - if (!unmasked.text) { + const text = restorer.restoreChunk(part.text); + + if (!text) { return []; } - return [{ ...part, text: unmasked.text }]; + return [{ ...part, text }]; }); if (processedContent.length > 0) { @@ -137,19 +80,7 @@ export function createUnmaskingStream( lineBuffer = ""; } - let flushed = ""; - - if (piiBuffer && piiContext) { - flushed = flushMaskingBuffer(piiBuffer, piiContext, config); - } else if (piiBuffer) { - flushed = piiBuffer; - } - - if (secretsBuffer && secretsContext) { - flushed += flushSecretsMaskingBuffer(secretsBuffer, secretsContext, config); - } else if (secretsBuffer) { - flushed += secretsBuffer; - } + const flushed = restorer.flush(); if (flushed) { const finalEvent = { diff --git a/src/routes/anthropic.ts b/src/routes/anthropic.ts index db3e8e9..580a575 100644 --- a/src/routes/anthropic.ts +++ b/src/routes/anthropic.ts @@ -4,7 +4,7 @@ import { Hono } from "hono"; import { getConfig } from "../config"; import type { PlaceholderContext } from "../masking/context"; import { anthropicExtractor } from "../masking/extractors/anthropic"; -import { unmaskResponse as unmaskPIIResponse } from "../pii/mask"; +import { restoreResponse } from "../masking/restorer"; import { callAnthropic } from "../providers/anthropic/client"; import { createAnthropicUnmaskingStream } from "../providers/anthropic/stream-transformer"; import { @@ -13,15 +13,15 @@ import { type AnthropicResponse, } from "../providers/anthropic/types"; import { callLocalAnthropic } from "../providers/local"; -import { unmaskSecretsResponse } from "../secrets/mask"; -import { formatMaskedSpansForLog, logScanRoles } from "../services/log-content"; +import { formatMaskedRequestForLog } from "../services/log-content"; import { logRequest } from "../services/logger"; -import { detectPII, maskPII, type PIIDetectResult } from "../services/pii"; +import type { PIIDetectResult } from "../services/pii"; import { - processSecretsRequest, - type SecretsProcessResult, - secretPlaceholders, -} from "../services/secrets"; + PrivacyPipelineDetectionError, + type PrivacyPipelineResult, + processPrivacyPipeline, +} from "../services/privacy-pipeline"; +import type { SecretsProcessResult } from "../services/secrets"; import { createLogData, errorFormats, @@ -52,7 +52,7 @@ anthropicRoutes.post( }), async (c) => { const startTime = Date.now(); - let request = c.req.valid("json") as AnthropicRequest; + const request = c.req.valid("json") as AnthropicRequest; const config = getConfig(); // Route mode requires local provider @@ -82,32 +82,31 @@ anthropicRoutes.post( ); } - // Step 1: Process secrets - const secretsResult = processSecretsRequest( - request, - config.secrets_detection, - anthropicExtractor, - ); + let privacy: PrivacyPipelineResult; + try { + privacy = await processPrivacyPipeline(request, config, anthropicExtractor); + } catch (error) { + if (error instanceof PrivacyPipelineDetectionError) { + console.error("PII detection error:", error.cause ?? error); + return respondDetectionError( + c, + error.request as AnthropicRequest, + error.secretsResult as SecretsProcessResult, + startTime, + ); + } + throw error; + } + const { secretsResult, piiResult } = privacy; if (secretsResult.blocked) { return respondBlocked(c, request, secretsResult, startTime); } - // Apply secrets masking to request - if (secretsResult.masked) { - request = secretsResult.request; + if (!piiResult) { + throw new Error("PII detection result missing from privacy pipeline"); } - // Step 2: Detect PII and configured denylist terms - let piiResult: PIIDetectResult; - try { - piiResult = await detectPII(request, anthropicExtractor, secretPlaceholders(secretsResult)); - } catch (error) { - console.error("PII detection error:", error); - return respondDetectionError(c, request, secretsResult, startTime); - } - - // Step 3: Route mode - send to local if PII or secrets detected const shouldRouteToLocal = config.mode === "route" && (piiResult.hasPII || @@ -115,31 +114,20 @@ anthropicRoutes.post( if (shouldRouteToLocal) { return sendToLocal(c, request, { - request, + request: privacy.requestAfterSecrets, startTime, piiResult, secretsResult, }); } - // Step 4: Mask mode - mask PII if found, send to Anthropic - let piiMaskingContext: PlaceholderContext | undefined; - let maskedContent: string | undefined; - - if (piiResult.hasPII) { - const masked = maskPII(request, piiResult.detection, anthropicExtractor); - request = masked.request; - piiMaskingContext = masked.maskingContext; - maskedContent = formatRequestForLog(request); - } else if (secretsResult.masked) { - maskedContent = formatRequestForLog(request); - } + const maskedContent = + piiResult.hasPII || secretsResult.masked ? formatRequestForLog(privacy.request) : undefined; - // Step 5: Send to Anthropic - return sendToAnthropic(c, request, { + return sendToAnthropic(c, privacy.request, { startTime, piiResult, - piiMaskingContext, + piiMaskingContext: privacy.piiMaskingContext, secretsResult, maskedContent, }); @@ -193,15 +181,7 @@ interface LocalOptions { function formatRequestForLog(request: AnthropicRequest): string | undefined { const config = getConfig(); - return formatMaskedSpansForLog( - anthropicExtractor.extractTexts(request), - logScanRoles({ - piiRoles: config.pii_detection.scan_roles, - piiActive: config.pii_detection.enabled || config.masking.denylist.length > 0, - secretRoles: config.secrets_detection.scan_roles, - secretsActive: config.secrets_detection.enabled, - }), - ); + return formatMaskedRequestForLog(request, anthropicExtractor, config); } // --- Response handlers --- @@ -412,15 +392,10 @@ function respondJson( secretsContext: PlaceholderContext | undefined, ) { const config = getConfig(); - let result = response; - - if (piiMaskingContext) { - result = unmaskPIIResponse(result, piiMaskingContext, config.masking, anthropicExtractor); - } - - if (secretsContext) { - result = unmaskSecretsResponse(result, secretsContext, config.masking, anthropicExtractor); - } + const result = restoreResponse(response, anthropicExtractor, config.masking, { + piiContext: piiMaskingContext, + secretsContext, + }); return c.json(result); } diff --git a/src/routes/codex.ts b/src/routes/codex.ts index a84a5a8..336928e 100644 --- a/src/routes/codex.ts +++ b/src/routes/codex.ts @@ -3,32 +3,25 @@ import type { Context } from "hono"; import { Hono } from "hono"; import { proxy } from "hono/proxy"; import { z } from "zod"; -import { getConfig, type MaskingConfig } from "../config"; +import { getConfig } from "../config"; import type { PlaceholderContext } from "../masking/context"; import { type CodexResponsesRequest, type CodexResponsesResponse, codexExtractor, } from "../masking/extractors/codex"; -import { - flushMaskingBuffer, - unmaskResponse as unmaskPIIResponse, - unmaskStreamChunk, -} from "../pii/mask"; +import { restoreResponse } from "../masking/restorer"; +import { createCodexUnmaskingStream } from "../providers/codex/stream-transformer"; import { ProviderError } from "../providers/errors"; -import { - flushSecretsMaskingBuffer, - unmaskSecretsResponse, - unmaskSecretsStreamChunk, -} from "../secrets/mask"; -import { formatMaskedSpansForLog, logScanRoles } from "../services/log-content"; +import { formatMaskedRequestForLog } from "../services/log-content"; import { logRequest } from "../services/logger"; -import { detectPII, maskPII, type PIIDetectResult } from "../services/pii"; +import type { PIIDetectResult } from "../services/pii"; import { - processSecretsRequest, - type SecretsProcessResult, - secretPlaceholders, -} from "../services/secrets"; + PrivacyPipelineDetectionError, + type PrivacyPipelineResult, + processPrivacyPipeline, +} from "../services/privacy-pipeline"; +import type { SecretsProcessResult } from "../services/secrets"; import { createLogData, errorFormats, @@ -68,23 +61,27 @@ codexRoutes.post( }), async (c) => { const startTime = Date.now(); - let request = c.req.valid("json") as CodexResponsesRequest; + const request = c.req.valid("json") as CodexResponsesRequest; const config = getConfig(); - const secretsResult = processSecretsRequest(request, config.secrets_detection, codexExtractor); + let privacy: PrivacyPipelineResult; + try { + privacy = await processPrivacyPipeline(request, config, codexExtractor); + } catch (error) { + if (error instanceof PrivacyPipelineDetectionError) { + console.error("PII detection error:", error.cause ?? error); + return respondDetectionError(c, error.request as CodexResponsesRequest, startTime); + } + throw error; + } + + const { secretsResult, piiResult } = privacy; if (secretsResult.blocked) { return respondBlocked(c, request, secretsResult, startTime); } - if (secretsResult.masked) { - request = secretsResult.request; - } - let piiResult: PIIDetectResult; - try { - piiResult = await detectPII(request, codexExtractor, secretPlaceholders(secretsResult)); - } catch (error) { - console.error("PII detection error:", error); - return respondDetectionError(c, request, startTime); + if (!piiResult) { + throw new Error("PII detection result missing from privacy pipeline"); } const shouldBlockRouteMode = @@ -96,13 +93,10 @@ codexRoutes.post( return respondRouteModeBlocked(c, request, piiResult, secretsResult, startTime); } - const piiMasked = - config.mode === "mask" ? maskPII(request, piiResult.detection, codexExtractor) : undefined; - return sendToCodex(c, request, { - request: piiMasked?.request ?? request, + request: privacy.request, piiResult, - piiMaskingContext: piiMasked?.maskingContext, + piiMaskingContext: privacy.piiMaskingContext, secretsResult, startTime, headers: getForwardHeaders(c), @@ -173,13 +167,7 @@ function getForwardHeaders(c: Context): Record { function formatCodexForLog(request: CodexResponsesRequest): string | undefined { const config = getConfig(); - const scanRoles = logScanRoles({ - piiRoles: config.pii_detection.scan_roles, - piiActive: config.pii_detection.enabled || config.masking.denylist.length > 0, - secretRoles: config.secrets_detection.scan_roles, - secretsActive: config.secrets_detection.enabled, - }); - return formatMaskedSpansForLog(codexExtractor.extractTexts(request), scanRoles); + return formatMaskedRequestForLog(request, codexExtractor, config); } function respondBlocked( @@ -378,136 +366,10 @@ function respondJson( secretsContext?: PlaceholderContext, maskingConfig = getConfig().masking, ) { - let result = response; - - if (piiContext) { - result = unmaskPIIResponse(result, piiContext, maskingConfig, codexExtractor); - } - if (secretsContext) { - result = unmaskSecretsResponse(result, secretsContext, maskingConfig, codexExtractor); - } + const result = restoreResponse(response, codexExtractor, maskingConfig, { + piiContext, + secretsContext, + }); return c.json(result); } - -function createCodexUnmaskingStream( - stream: ReadableStream, - piiContext: PlaceholderContext | undefined, - maskingConfig: MaskingConfig, - secretsContext?: PlaceholderContext, -): ReadableStream { - const decoder = new TextDecoder(); - const encoder = new TextEncoder(); - let piiBuffer = ""; - let secretsBuffer = ""; - let lineBuffer = ""; - - function unmaskPayload(payload: unknown): unknown { - let result = payload as CodexResponsesResponse; - - if (piiContext) { - const spans = codexExtractor.extractTexts(result); - result = codexExtractor.applyMasked( - result, - spans.map((span) => { - const { output, remainingBuffer } = unmaskStreamChunk( - piiBuffer, - span.text, - piiContext, - maskingConfig, - ); - piiBuffer = remainingBuffer; - return { ...span, maskedText: output }; - }), - ); - } - - if (secretsContext) { - const spans = codexExtractor.extractTexts(result); - result = codexExtractor.applyMasked( - result, - spans.map((span) => { - const { output, remainingBuffer } = unmaskSecretsStreamChunk( - secretsBuffer, - span.text, - secretsContext, - maskingConfig, - ); - secretsBuffer = remainingBuffer; - return { ...span, maskedText: output }; - }), - ); - } - - return result; - } - - function processLine(line: string): string { - if (!line.startsWith("data: ")) { - return `${line}\n`; - } - - const data = line.slice(6); - if (data === "[DONE]") { - return "data: [DONE]\n"; - } - - try { - return `data: ${JSON.stringify(unmaskPayload(JSON.parse(data)))}\n`; - } catch { - return `${line}\n`; - } - } - - return new ReadableStream({ - async start(controller) { - const reader = stream.getReader(); - - try { - while (true) { - const { done, value } = await reader.read(); - if (done) break; - - lineBuffer += decoder.decode(value, { stream: true }); - const lines = lineBuffer.split("\n"); - lineBuffer = lines.pop() ?? ""; - - let output = ""; - for (const line of lines) { - output += processLine(line); - } - - if (output) { - controller.enqueue(encoder.encode(output)); - } - } - - lineBuffer += decoder.decode(); - let finalOutput = lineBuffer ? processLine(lineBuffer) : ""; - lineBuffer = ""; - - if (piiContext && piiBuffer) { - finalOutput += `data: ${JSON.stringify({ - type: "response.output_text.delta", - delta: flushMaskingBuffer(piiBuffer, piiContext, maskingConfig), - })}\n\n`; - } - if (secretsContext && secretsBuffer) { - finalOutput += `data: ${JSON.stringify({ - type: "response.output_text.delta", - delta: flushSecretsMaskingBuffer(secretsBuffer, secretsContext, maskingConfig), - })}\n\n`; - } - if (finalOutput) { - controller.enqueue(encoder.encode(finalOutput)); - } - - controller.close(); - } catch (error) { - controller.error(error); - } finally { - reader.releaseLock(); - } - }, - }); -} diff --git a/src/routes/openai.ts b/src/routes/openai.ts index 05ed5ca..48269ad 100644 --- a/src/routes/openai.ts +++ b/src/routes/openai.ts @@ -5,7 +5,7 @@ import { proxy } from "hono/proxy"; import { getConfig, type MaskingConfig } from "../config"; import type { PlaceholderContext } from "../masking/context"; import { openaiExtractor } from "../masking/extractors/openai"; -import { unmaskResponse as unmaskPIIResponse } from "../pii/mask"; +import { restoreResponse } from "../masking/restorer"; import { callLocal } from "../providers/local"; import { callOpenAI, getOpenAIInfo, type ProviderResult } from "../providers/openai/client"; import { createUnmaskingStream } from "../providers/openai/stream-transformer"; @@ -14,15 +14,15 @@ import { OpenAIRequestSchema, type OpenAIResponse, } from "../providers/openai/types"; -import { unmaskSecretsResponse } from "../secrets/mask"; -import { formatMaskedSpansForLog, logScanRoles } from "../services/log-content"; +import { formatMaskedRequestForLog } from "../services/log-content"; import { logRequest } from "../services/logger"; -import { detectPII, maskPII, type PIIDetectResult } from "../services/pii"; +import type { PIIDetectResult } from "../services/pii"; import { - processSecretsRequest, - type SecretsProcessResult, - secretPlaceholders, -} from "../services/secrets"; + PrivacyPipelineDetectionError, + type PrivacyPipelineResult, + processPrivacyPipeline, +} from "../services/privacy-pipeline"; +import type { SecretsProcessResult } from "../services/secrets"; import { createLogData, errorFormats, @@ -53,37 +53,35 @@ openaiRoutes.post( }), async (c) => { const startTime = Date.now(); - let request = c.req.valid("json") as OpenAIRequest; + const request = c.req.valid("json") as OpenAIRequest; const config = getConfig(); - // Step 1: Process secrets - const secretsResult = processSecretsRequest(request, config.secrets_detection, openaiExtractor); + let privacy: PrivacyPipelineResult; + try { + privacy = await processPrivacyPipeline(request, config, openaiExtractor); + } catch (error) { + if (error instanceof PrivacyPipelineDetectionError) { + console.error("PII detection error:", error.cause ?? error); + return respondDetectionError(c, error.request as OpenAIRequest, startTime); + } + throw error; + } + + const { secretsResult, piiResult } = privacy; if (secretsResult.blocked) { return respondBlocked(c, request, secretsResult, startTime); } - // Apply secrets masking to request - if (secretsResult.masked) { - request = secretsResult.request; - } - - // Step 2: Detect PII and configured denylist terms - let piiResult: PIIDetectResult; - try { - piiResult = await detectPII(request, openaiExtractor, secretPlaceholders(secretsResult)); - } catch (error) { - console.error("PII detection error:", error); - return respondDetectionError(c, request, startTime); + if (!piiResult) { + throw new Error("PII detection result missing from privacy pipeline"); } - // Step 3: Process based on mode if (config.mode === "mask") { - const piiMasked = maskPII(request, piiResult.detection, openaiExtractor); return sendToOpenAI(c, request, { - request: piiMasked.request, + request: privacy.request, piiResult, - piiMaskingContext: piiMasked.maskingContext, + piiMaskingContext: privacy.piiMaskingContext, secretsResult, startTime, authHeader: c.req.header("Authorization"), @@ -97,7 +95,7 @@ openaiRoutes.post( if (shouldRouteLocal) { return sendToLocal(c, request, { - request, + request: privacy.requestAfterSecrets, piiResult, secretsResult, startTime, @@ -105,7 +103,7 @@ openaiRoutes.post( } return sendToOpenAI(c, request, { - request, + request: privacy.requestAfterSecrets, piiResult, secretsResult, startTime, @@ -152,15 +150,7 @@ interface LocalOptions { function formatRequestForLog(request: OpenAIRequest): string | undefined { const config = getConfig(); - return formatMaskedSpansForLog( - openaiExtractor.extractTexts(request), - logScanRoles({ - piiRoles: config.pii_detection.scan_roles, - piiActive: config.pii_detection.enabled || config.masking.denylist.length > 0, - secretRoles: config.secrets_detection.scan_roles, - secretsActive: config.secrets_detection.enabled, - }), - ); + return formatMaskedRequestForLog(request, openaiExtractor, config); } // --- Response handlers --- @@ -375,14 +365,10 @@ function respondJson( piiContext?: PlaceholderContext, secretsContext?: PlaceholderContext, ) { - let result = response; - - if (piiContext) { - result = unmaskPIIResponse(result, piiContext, maskingConfig, openaiExtractor); - } - if (secretsContext) { - result = unmaskSecretsResponse(result, secretsContext, maskingConfig, openaiExtractor); - } + const result = restoreResponse(response, openaiExtractor, maskingConfig, { + piiContext, + secretsContext, + }); return c.json(result); } diff --git a/src/secrets/mask.test.ts b/src/secrets/mask.test.ts index 7ffd291..256f71f 100644 --- a/src/secrets/mask.test.ts +++ b/src/secrets/mask.test.ts @@ -1,34 +1,13 @@ import { describe, expect, test } from "bun:test"; -import type { MaskingConfig } from "../config"; -import { anthropicExtractor } from "../masking/extractors/anthropic"; -import { type CodexResponsesResponse, codexExtractor } from "../masking/extractors/codex"; +import { restorePlaceholders } from "../masking/context"; import { openaiExtractor } from "../masking/extractors/openai"; -import type { AnthropicResponse } from "../providers/anthropic/types"; -import type { OpenAIMessage, OpenAIRequest, OpenAIResponse } from "../providers/openai/types"; +import type { OpenAIMessage, OpenAIRequest } from "../providers/openai/types"; import { createSecretsResultFromSpans } from "../test-utils/detection-results"; import type { SecretLocation } from "./detect"; -import { - createSecretsMaskingContext, - flushSecretsMaskingBuffer, - maskRequest, - maskSecrets, - unmaskSecrets, - unmaskSecretsResponse, - unmaskSecretsStreamChunk, -} from "./mask"; +import { createSecretsMaskingContext, maskRequest, maskSecrets } from "./mask"; const sampleSecret = "sk-proj-abc123def456ghi789jkl012mno345pqr678stu901vwx"; const stripeSecret = "sk_live_abc123def456ghi789jkl012"; -const defaultConfig: MaskingConfig = { - show_markers: false, - marker_text: "[protected]", - allowlist: [], - denylist: [], -}; -const markerConfig: MaskingConfig = { - ...defaultConfig, - show_markers: true, -}; /** Helper to create a minimal request from messages */ function createRequest(messages: OpenAIMessage[]): OpenAIRequest { @@ -152,59 +131,6 @@ describe("maskRequest with MessageSecretsResult", () => { }); }); -describe("streaming with secrets placeholders", () => { - test("buffers partial [[ placeholder", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const { output, remainingBuffer } = unmaskSecretsStreamChunk( - "", - "Key: [[API_KEY", - context, - defaultConfig, - ); - - expect(output).toBe("Key: "); - expect(remainingBuffer).toBe("[[API_KEY"); - }); - - test("completes buffered placeholder across chunks", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const { output, remainingBuffer } = unmaskSecretsStreamChunk( - "[[API_KEY", - "_SK_1]] done", - context, - defaultConfig, - ); - - expect(output).toBe(`${sampleSecret} done`); - expect(remainingBuffer).toBe(""); - }); - - test("flushes incomplete buffer as-is", () => { - const context = createSecretsMaskingContext(); - const result = flushSecretsMaskingBuffer("[[API_KEY", context, defaultConfig); - expect(result).toBe("[[API_KEY"); - }); - - test("adds markers to streamed secret restoration when show_markers is true", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const { output, remainingBuffer } = unmaskSecretsStreamChunk( - "", - "Key: [[API_KEY_SK_1]]", - context, - markerConfig, - ); - - expect(output).toBe(`Key: [protected]${sampleSecret}`); - expect(remainingBuffer).toBe(""); - }); -}); - describe("mask -> unmask roundtrip", () => { test("preserves original data through roundtrip", () => { const originalText = ` @@ -225,131 +151,9 @@ Please store them securely. expect(masked).not.toContain(sampleSecret); expect(masked).toContain("[[API_KEY_SK_1]]"); - const restored = unmaskSecrets(masked, context, defaultConfig); + const restored = restorePlaceholders(masked, context); expect(restored).toBe(originalText); }); - - test("does not add markers when show_markers is false", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const restored = unmaskSecrets("Key: [[API_KEY_SK_1]]", context, defaultConfig); - - expect(restored).toBe(`Key: ${sampleSecret}`); - }); -}); - -describe("unmaskSecretsResponse", () => { - test("unmasks all OpenAI choices in response", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const response: OpenAIResponse = { - id: "test", - object: "chat.completion", - created: Date.now(), - model: "gpt-4", - choices: [ - { - index: 0, - message: { - role: "assistant", - content: "Your key is [[API_KEY_SK_1]]", - }, - finish_reason: "stop", - }, - ], - }; - - const result = unmaskSecretsResponse(response, context, defaultConfig, openaiExtractor); - expect(result.choices[0].message.content).toBe(`Your key is ${sampleSecret}`); - }); - - test("adds markers to OpenAI response secrets when show_markers is true", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const response: OpenAIResponse = { - id: "test", - object: "chat.completion", - created: Date.now(), - model: "gpt-4", - choices: [ - { - index: 0, - message: { - role: "assistant", - content: "Your key is [[API_KEY_SK_1]]", - }, - finish_reason: "stop", - }, - ], - }; - - const result = unmaskSecretsResponse(response, context, markerConfig, openaiExtractor); - expect(result.choices[0].message.content).toBe(`Your key is [protected]${sampleSecret}`); - }); - - test("adds markers to Anthropic response secrets when show_markers is true", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const response: AnthropicResponse = { - id: "msg_test", - type: "message", - role: "assistant", - content: [{ type: "text", text: "Your key is [[API_KEY_SK_1]]" }], - model: "claude-3-5-sonnet", - stop_reason: "end_turn", - stop_sequence: null, - usage: { input_tokens: 10, output_tokens: 5 }, - }; - - const result = unmaskSecretsResponse(response, context, markerConfig, anthropicExtractor); - expect(result.content[0]).toEqual({ - type: "text", - text: `Your key is [protected]${sampleSecret}`, - }); - }); - - test("adds markers to Codex response secrets when show_markers is true", () => { - const context = createSecretsMaskingContext(); - context.mapping["[[API_KEY_SK_1]]"] = sampleSecret; - - const response: CodexResponsesResponse = { - output: [{ content: [{ type: "output_text", text: "Your key is [[API_KEY_SK_1]]" }] }], - }; - - const result = unmaskSecretsResponse(response, context, markerConfig, codexExtractor); - expect(result).toEqual({ - output: [ - { content: [{ type: "output_text", text: `Your key is [protected]${sampleSecret}` }] }, - ], - }); - }); - - test("preserves response structure", () => { - const context = createSecretsMaskingContext(); - const response: OpenAIResponse = { - id: "test-id", - object: "chat.completion", - created: 12345, - model: "gpt-4-turbo", - choices: [ - { - index: 0, - message: { role: "assistant", content: "Hello" }, - finish_reason: "stop", - }, - ], - usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, - }; - - const result = unmaskSecretsResponse(response, context, defaultConfig, openaiExtractor); - expect(result.id).toBe("test-id"); - expect(result.model).toBe("gpt-4-turbo"); - expect(result.usage).toEqual({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }); - }); }); describe("edge cases", () => { diff --git a/src/secrets/mask.ts b/src/secrets/mask.ts index 702c8a2..5f0ad77 100644 --- a/src/secrets/mask.ts +++ b/src/secrets/mask.ts @@ -2,18 +2,10 @@ * Secrets masking */ -import type { MaskingConfig } from "../config"; import { resolveOverlaps } from "../masking/conflict-resolver"; import { incrementAndGenerate } from "../masking/context"; import { generateSecretPlaceholder } from "../masking/placeholders"; -import { - createMaskingContext, - flushMaskingBuffer as flushBuffer, - maskSpans, - type PlaceholderContext, - unmaskStreamChunk as unmaskChunk, - unmask as unmaskText, -} from "../masking/service"; +import { createMaskingContext, maskSpans, type PlaceholderContext } from "../masking/service"; import type { RequestExtractor, TextSpan } from "../masking/types"; import type { MessageSecretsResult, SecretLocation } from "./detect"; @@ -37,10 +29,6 @@ function generatePlaceholder(secretType: string, context: PlaceholderContext): s return incrementAndGenerate(secretType, context, generateSecretPlaceholder); } -function getFormatValue(config: MaskingConfig): ((original: string) => string) | undefined { - return config.show_markers ? (original: string) => `${config.marker_text}${original}` : undefined; -} - /** * Masks secrets in text, replacing them with placeholders */ @@ -67,52 +55,6 @@ export function maskSecrets( }; } -/** - * Unmasks text by replacing placeholders with original secrets - */ -export function unmaskSecrets( - text: string, - context: PlaceholderContext, - config: MaskingConfig, -): string { - return unmaskText(text, context, getFormatValue(config)); -} - -/** - * Streaming unmask helper - processes chunks and unmasks when complete placeholders are found - */ -export function unmaskSecretsStreamChunk( - buffer: string, - newChunk: string, - context: PlaceholderContext, - config: MaskingConfig, -): { output: string; remainingBuffer: string } { - return unmaskChunk(buffer, newChunk, context, getFormatValue(config)); -} - -/** - * Flushes remaining buffer at end of stream - */ -export function flushSecretsMaskingBuffer( - buffer: string, - context: PlaceholderContext, - config: MaskingConfig, -): string { - return flushBuffer(buffer, context, getFormatValue(config)); -} - -/** - * Unmasks secrets in a response using an extractor - */ -export function unmaskSecretsResponse( - response: TResponse, - context: PlaceholderContext, - config: MaskingConfig, - extractor: RequestExtractor, -): TResponse { - return extractor.unmaskResponse(response, context, getFormatValue(config)); -} - /** * Result of masking a request */ diff --git a/src/services/log-content.test.ts b/src/services/log-content.test.ts index ef037f4..9876556 100644 --- a/src/services/log-content.test.ts +++ b/src/services/log-content.test.ts @@ -1,5 +1,12 @@ import { describe, expect, test } from "bun:test"; -import { formatMaskedSpansForLog, logScanRoles, shouldLogMaskedContent } from "./log-content"; +import { openaiExtractor } from "../masking/extractors/openai"; +import type { OpenAIRequest } from "../providers/openai/types"; +import { + formatMaskedRequestForLog, + formatMaskedSpansForLog, + logScanRoles, + shouldLogMaskedContent, +} from "./log-content"; describe("shouldLogMaskedContent", () => { const maskedWithSecret = "My key is [[API_KEY_SK_1]] and email [[EMAIL_ADDRESS_1]]"; @@ -196,3 +203,45 @@ describe("logScanRoles", () => { expect(result).toBeUndefined(); }); }); + +describe("formatMaskedRequestForLog", () => { + test("uses shared scan-role intersection for provider requests", () => { + const request: OpenAIRequest = { + model: "gpt-4", + messages: [ + { role: "system", content: "System [[EMAIL_ADDRESS_1]]" }, + { role: "user", content: "User [[EMAIL_ADDRESS_2]]" }, + { role: "assistant", content: "Assistant [[EMAIL_ADDRESS_3]]" }, + ], + }; + + const result = formatMaskedRequestForLog(request, openaiExtractor, { + pii_detection: { + enabled: true, + detector_url: "http://localhost:8080", + phone_regions: [], + score_threshold: 0.7, + entities: ["EMAIL_ADDRESS"], + scan_roles: ["user", "assistant"], + }, + masking: { + show_markers: false, + marker_text: "[protected]", + allowlist: [], + denylist: [], + }, + secrets_detection: { + enabled: true, + action: "mask", + entities: ["API_KEY_SK"], + max_scan_chars: 200000, + log_detected_types: true, + scan_roles: ["user", "tool"], + }, + }); + + expect(result).toContain("[user prompt] User [[EMAIL_ADDRESS_2]]"); + expect(result).not.toContain("System [[EMAIL_ADDRESS_1]]"); + expect(result).not.toContain("Assistant [[EMAIL_ADDRESS_3]]"); + }); +}); diff --git a/src/services/log-content.ts b/src/services/log-content.ts index cbb5bc3..1192849 100644 --- a/src/services/log-content.ts +++ b/src/services/log-content.ts @@ -1,4 +1,5 @@ -import type { TextSpan } from "../masking/types"; +import type { Config } from "../config"; +import type { RequestExtractor, TextSpan } from "../masking/types"; export interface LogContentDecision { maskedContent?: string; @@ -26,6 +27,22 @@ export function formatMaskedSpansForLog( return lines.length > 0 ? lines.join("\n").slice(0, 20000) : undefined; } +export function formatMaskedRequestForLog( + request: TRequest, + extractor: RequestExtractor, + config: Pick, +): string | undefined { + return formatMaskedSpansForLog( + extractor.extractTexts(request), + logScanRoles({ + piiRoles: config.pii_detection.scan_roles, + piiActive: config.pii_detection.enabled || config.masking.denylist.length > 0, + secretRoles: config.secrets_detection.scan_roles, + secretsActive: config.secrets_detection.enabled, + }), + ); +} + export function logScanRoles(opts: { piiRoles: readonly string[]; piiActive: boolean; diff --git a/src/services/privacy-pipeline.test.ts b/src/services/privacy-pipeline.test.ts new file mode 100644 index 0000000..6b9ee1b --- /dev/null +++ b/src/services/privacy-pipeline.test.ts @@ -0,0 +1,137 @@ +import { afterEach, describe, expect, mock, test } from "bun:test"; +import { openaiExtractor } from "../masking/extractors/openai"; +import type { PIIDetectionResult } from "../pii/detect"; +import type { OpenAIRequest } from "../providers/openai/types"; +import type { PrivacyPipelineConfig } from "./privacy-pipeline"; + +const sampleSecret = "sk-proj-abc123def456ghi789jkl012mno345pqr678stu901vwx"; + +const mockAnalyzeRequest = mock< + ( + request: unknown, + extractor: unknown, + knownPlaceholders: readonly string[], + ) => Promise +>(() => + Promise.resolve({ + hasPII: false, + spanEntities: [], + allEntities: [], + scanTimeMs: 0, + }), +); + +mock.module("../pii/detect", () => ({ + getPIIDetector: () => ({ + analyzeRequest: mockAnalyzeRequest, + detectPII: mock(() => Promise.resolve([])), + healthCheck: mock(() => Promise.resolve(true)), + }), +})); + +const { PrivacyPipelineDetectionError, processPrivacyPipeline } = await import( + "./privacy-pipeline" +); + +const baseConfig: PrivacyPipelineConfig = { + mode: "mask", + secrets_detection: { + enabled: true, + action: "mask", + entities: ["API_KEY_SK"], + max_scan_chars: 200000, + log_detected_types: true, + scan_roles: ["user", "tool", "function", "mcp"], + }, +}; + +function request(content: string): OpenAIRequest { + return { + model: "gpt-4", + messages: [{ role: "user", content }], + }; +} + +afterEach(() => { + mockAnalyzeRequest.mockReset(); + mockAnalyzeRequest.mockResolvedValue({ + hasPII: false, + spanEntities: [], + allEntities: [], + scanTimeMs: 0, + }); +}); + +describe("processPrivacyPipeline", () => { + test("runs secrets before PII and passes secret placeholders to detection", async () => { + const input = request(`Key ${sampleSecret} email jane@example.com`); + + const result = await processPrivacyPipeline(input, baseConfig, openaiExtractor); + + expect(result.requestAfterSecrets.messages[0].content).toContain("[[API_KEY_SK_1]]"); + expect(result.requestAfterSecrets.messages[0].content).not.toContain(sampleSecret); + expect(mockAnalyzeRequest).toHaveBeenCalledTimes(1); + + const [detectedRequest, , knownPlaceholders] = mockAnalyzeRequest.mock.calls[0]; + expect((detectedRequest as OpenAIRequest).messages[0].content).toContain("[[API_KEY_SK_1]]"); + expect(knownPlaceholders).toEqual(["[[API_KEY_SK_1]]"]); + }); + + test("returns privacy facts without route decisions", async () => { + const result = await processPrivacyPipeline(request("Hello"), baseConfig, openaiExtractor); + + expect(result.secretsResult.masked).toBe(false); + expect(result.piiResult?.hasPII).toBe(false); + expect("shouldRouteLocal" in result).toBe(false); + expect("shouldBlock" in result).toBe(false); + }); + + test("masks PII in mask mode and returns its restoration context", async () => { + mockAnalyzeRequest.mockResolvedValueOnce({ + hasPII: true, + spanEntities: [[{ entity_type: "EMAIL_ADDRESS", start: 6, end: 22, score: 0.99 }]], + allEntities: [{ entity_type: "EMAIL_ADDRESS", start: 6, end: 22, score: 0.99 }], + scanTimeMs: 3, + }); + + const result = await processPrivacyPipeline( + request("Email jane@example.com"), + baseConfig, + openaiExtractor, + ); + + expect(result.request.messages[0].content).toBe("Email [[EMAIL_ADDRESS_1]]"); + expect(result.piiMaskingContext?.mapping["[[EMAIL_ADDRESS_1]]"]).toBe("jane@example.com"); + }); + + test("does not call PII detection when secrets block the request", async () => { + const result = await processPrivacyPipeline( + request(`Key ${sampleSecret}`), + { + ...baseConfig, + secrets_detection: { ...baseConfig.secrets_detection, action: "block" }, + }, + openaiExtractor, + ); + + expect(result.secretsResult.blocked).toBe(true); + expect(result.piiResult).toBeUndefined(); + expect(mockAnalyzeRequest).not.toHaveBeenCalled(); + }); + + test("preserves the post-secrets request on PII detection errors", async () => { + mockAnalyzeRequest.mockRejectedValueOnce(new Error("detector down")); + + try { + await processPrivacyPipeline(request(`Key ${sampleSecret}`), baseConfig, openaiExtractor); + throw new Error("Expected processPrivacyPipeline to throw"); + } catch (error) { + if (!(error instanceof PrivacyPipelineDetectionError)) { + throw error; + } + + expect((error.request as OpenAIRequest).messages[0].content).toContain("[[API_KEY_SK_1]]"); + expect(error.secretsResult.masked).toBe(true); + } + }); +}); diff --git a/src/services/privacy-pipeline.ts b/src/services/privacy-pipeline.ts new file mode 100644 index 0000000..2b469ec --- /dev/null +++ b/src/services/privacy-pipeline.ts @@ -0,0 +1,70 @@ +import type { Config } from "../config"; +import type { PlaceholderContext } from "../masking/context"; +import type { RequestExtractor } from "../masking/types"; +import { detectPII, maskPII, type PIIDetectResult } from "./pii"; +import { processSecretsRequest, type SecretsProcessResult, secretPlaceholders } from "./secrets"; + +export type PrivacyPipelineConfig = Pick; + +export interface PrivacyPipelineResult { + requestAfterSecrets: TRequest; + request: TRequest; + secretsResult: SecretsProcessResult; + piiResult?: PIIDetectResult; + piiMaskingContext?: PlaceholderContext; +} + +export class PrivacyPipelineDetectionError extends Error { + constructor( + message: string, + public readonly request: unknown, + public readonly secretsResult: SecretsProcessResult, + options?: { cause?: unknown }, + ) { + super(message); + this.name = "PrivacyPipelineDetectionError"; + this.cause = options?.cause; + } +} + +export async function processPrivacyPipeline( + request: TRequest, + config: PrivacyPipelineConfig, + extractor: RequestExtractor, +): Promise> { + const secretsResult = processSecretsRequest(request, config.secrets_detection, extractor); + let workingRequest = secretsResult.masked ? secretsResult.request : request; + const requestAfterSecrets = workingRequest; + + if (secretsResult.blocked) { + return { requestAfterSecrets, request: workingRequest, secretsResult }; + } + + let piiResult: PIIDetectResult; + try { + piiResult = await detectPII(workingRequest, extractor, secretPlaceholders(secretsResult)); + } catch (error) { + throw new PrivacyPipelineDetectionError( + "PII detection service unavailable", + workingRequest, + secretsResult, + { cause: error }, + ); + } + + let piiMaskingContext: PlaceholderContext | undefined; + + if (config.mode === "mask") { + const masked = maskPII(workingRequest, piiResult.detection, extractor); + workingRequest = masked.request; + piiMaskingContext = masked.maskingContext; + } + + return { + requestAfterSecrets, + request: workingRequest, + secretsResult, + piiResult, + piiMaskingContext, + }; +}