--- /dev/null
+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");
+ });
+});
--- /dev/null
+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;
+}
--- /dev/null
+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<unknown, TestResponse> = {
+ 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<string, string>): 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" }] }],
+ });
+ });
+});
--- /dev/null
+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<TRequest, TResponse>(
+ response: TResponse,
+ extractor: RequestExtractor<TRequest, TResponse>,
+ 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;
+}
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
*
--- /dev/null
+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<string, string>): 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",
+ );
+ });
+});
--- /dev/null
+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;
+ }
+}
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 {
});
});
-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([
});
});
-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";
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 = `<p>Contact [[PERSON_1]] at [[EMAIL_ADDRESS_1]]</p>`;
- const result = unmask(htmlResponse, context, defaultConfig);
-
- expect(result).toBe("<p>Contact Dr. Sarah Chen at sarah.chen@hospital.org</p>");
- });
-
- 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 = `
- <div class="profile">
- <h1>[[PERSON_1]]</h1>
- <a href="mailto:[[EMAIL_ADDRESS_1]]">[[EMAIL_ADDRESS_1]]</a>
- </div>
- `;
-
- 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";
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", () => {
-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";
);
}
-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[],
};
}
-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<TRequest> {
request: TRequest;
context: PlaceholderContext;
existingContext,
);
}
-
-export function unmaskResponse<TRequest, TResponse>(
- response: TResponse,
- context: PlaceholderContext,
- config: MaskingConfig,
- extractor: RequestExtractor<TRequest, TResponse>,
-): TResponse {
- return extractor.unmaskResponse(response, context, getFormatValue(config));
-}
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(
): ReadableStream<Uint8Array> {
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) {
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) {
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) {
--- /dev/null
+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<string, string>): PlaceholderContext {
+ const ctx = createPlaceholderContext();
+ ctx.mapping = mapping;
+ return ctx;
+}
+
+function createSSEStream(chunks: string[]): ReadableStream<Uint8Array> {
+ 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<Uint8Array>): Promise<string> {
+ 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"');
+ });
+});
--- /dev/null
+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<Uint8Array>,
+ piiContext: PlaceholderContext | undefined,
+ maskingConfig: MaskingConfig,
+ secretsContext?: PlaceholderContext,
+): ReadableStream<Uint8Array> {
+ 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();
+ }
+ },
+ });
+}
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<Uint8Array>,
piiContext: PlaceholderContext | undefined,
): ReadableStream<Uint8Array> {
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) {
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)) {
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) {
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 = {
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 {
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,
}),
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
);
}
- // Step 1: Process secrets
- const secretsResult = processSecretsRequest(
- request,
- config.secrets_detection,
- anthropicExtractor,
- );
+ let privacy: PrivacyPipelineResult<AnthropicRequest>;
+ 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<AnthropicRequest>,
+ 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 ||
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,
});
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 ---
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);
}
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,
}),
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<CodexResponsesRequest>;
+ 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 =
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),
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(
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<Uint8Array>,
- piiContext: PlaceholderContext | undefined,
- maskingConfig: MaskingConfig,
- secretsContext?: PlaceholderContext,
-): ReadableStream<Uint8Array> {
- 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();
- }
- },
- });
-}
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";
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,
}),
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<OpenAIRequest>;
+ 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"),
if (shouldRouteLocal) {
return sendToLocal(c, request, {
- request,
+ request: privacy.requestAfterSecrets,
piiResult,
secretsResult,
startTime,
}
return sendToOpenAI(c, request, {
- request,
+ request: privacy.requestAfterSecrets,
piiResult,
secretsResult,
startTime,
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 ---
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);
}
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 {
});
});
-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 = `
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", () => {
* 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";
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
*/
};
}
-/**
- * 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<TRequest, TResponse>(
- response: TResponse,
- context: PlaceholderContext,
- config: MaskingConfig,
- extractor: RequestExtractor<TRequest, TResponse>,
-): TResponse {
- return extractor.unmaskResponse(response, context, getFormatValue(config));
-}
-
/**
* Result of masking a request
*/
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]]";
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]]");
+ });
+});
-import type { TextSpan } from "../masking/types";
+import type { Config } from "../config";
+import type { RequestExtractor, TextSpan } from "../masking/types";
export interface LogContentDecision {
maskedContent?: string;
return lines.length > 0 ? lines.join("\n").slice(0, 20000) : undefined;
}
+export function formatMaskedRequestForLog<TRequest, TResponse>(
+ request: TRequest,
+ extractor: RequestExtractor<TRequest, TResponse>,
+ config: Pick<Config, "pii_detection" | "masking" | "secrets_detection">,
+): 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;
--- /dev/null
+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<PIIDetectionResult>
+>(() =>
+ 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);
+ }
+ });
+});
--- /dev/null
+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<Config, "mode" | "secrets_detection">;
+
+export interface PrivacyPipelineResult<TRequest> {
+ requestAfterSecrets: TRequest;
+ request: TRequest;
+ secretsResult: SecretsProcessResult<TRequest>;
+ piiResult?: PIIDetectResult;
+ piiMaskingContext?: PlaceholderContext;
+}
+
+export class PrivacyPipelineDetectionError extends Error {
+ constructor(
+ message: string,
+ public readonly request: unknown,
+ public readonly secretsResult: SecretsProcessResult<unknown>,
+ options?: { cause?: unknown },
+ ) {
+ super(message);
+ this.name = "PrivacyPipelineDetectionError";
+ this.cause = options?.cause;
+ }
+}
+
+export async function processPrivacyPipeline<TRequest, TResponse>(
+ request: TRequest,
+ config: PrivacyPipelineConfig,
+ extractor: RequestExtractor<TRequest, TResponse>,
+): Promise<PrivacyPipelineResult<TRequest>> {
+ 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,
+ };
+}
git.99rst.org / sgasser-llm-shield.git / commitdiff