From: maximiliancw Date: Fri, 9 Jan 2026 15:04:32 +0000 (+0100) Subject: Implement redact and route_local actions: X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=a3f71141c0d2b48d352c30593a1270d404fe0704;p=sgasser-llm-shield.git Implement redact and route_local actions: - Integrate redact action: redacts secrets before PII detection, unredacts in responses - Implement route_local action: routes requests with secrets to local provider - Update stream transformer to handle both PII and secrets contexts - Add comprehensive tests for secrets routing logic - Update config.example.yaml with new entity types and action documentation - Update README.md with complete secrets detection features New secret entity types (opt-in): - API_KEY_OPENAI, API_KEY_AWS, API_KEY_GITHUB - JWT_TOKEN, BEARER_TOKEN Response headers: - X-PasteGuard-Secrets-Redacted: true (when action=redact) --- diff --git a/README.md b/README.md index 20a2844..e376c5a 100644 --- a/README.md +++ b/README.md @@ -51,12 +51,22 @@ Additional entity types can be enabled: `US_SSN`, `US_PASSPORT`, `CRYPTO`, `NRP` ### Secrets (Secrets Shield) -| Type | Examples | +| Type | Pattern | | -------------------- | --------------------------------------------------------------------------------------------------------- | | OpenSSH private keys | `-----BEGIN OPENSSH PRIVATE KEY-----` | | PEM private keys | `-----BEGIN RSA PRIVATE KEY-----`, `-----BEGIN PRIVATE KEY-----`, `-----BEGIN ENCRYPTED PRIVATE KEY-----` | +| OpenAI API keys | `sk-proj-...`, `sk-...` (48+ chars) | +| AWS access keys | `AKIA...` (20 chars) | +| GitHub tokens | `ghp_...`, `gho_...`, `ghu_...`, `ghs_...`, `ghr_...` | +| JWT tokens | `eyJ...` (three base64 segments) | +| Bearer tokens | `Bearer ...` (20+ char tokens) | -Secrets detection runs **before** PII detection and blocks requests by default (configurable). Detected secrets are never logged in their original form. +Secrets detection runs **before** PII detection. Three actions available: +- **block** (default): Returns HTTP 422, request never reaches LLM +- **redact**: Replaces secrets with placeholders, unredacts in response (reversible) +- **route_local**: Routes to local LLM (route mode only) + +Detected secrets are never logged in their original form. **Languages**: 24 languages supported (configurable at build time). Auto-detected per request. @@ -192,13 +202,20 @@ secrets_detection: entities: # Secret types to detect - OPENSSH_PRIVATE_KEY - PEM_PRIVATE_KEY + # API Keys (opt-in): + # - API_KEY_OPENAI + # - API_KEY_AWS + # - API_KEY_GITHUB + # Tokens (opt-in): + # - JWT_TOKEN + # - BEARER_TOKEN max_scan_chars: 200000 # Performance limit (0 = no limit) log_detected_types: true # Log types (never logs content) ``` - **block** (default): Returns HTTP 422 error, request never reaches LLM -- **redact**: Replaces secrets with placeholders (Phase 2) -- **route_local**: Routes to local provider when secrets detected (Phase 2, requires route mode) +- **redact**: Replaces secrets with placeholders, unredacts in response (reversible, like PII masking) +- **route_local**: Routes to local provider when secrets detected (requires route mode) **Logging options:** @@ -249,7 +266,8 @@ See [config.example.yaml](config.example.yaml) for all options. | `X-PasteGuard-Language` | Detected language code | | `X-PasteGuard-Language-Fallback` | `true` if fallback was used | | `X-PasteGuard-Secrets-Detected` | `true` if secrets detected | -| `X-PasteGuard-Secrets-Types` | Comma-separated list of detected secret types (e.g., `OPENSSH_PRIVATE_KEY,PEM_PRIVATE_KEY`) | +| `X-PasteGuard-Secrets-Types` | Comma-separated list (e.g., `OPENSSH_PRIVATE_KEY,API_KEY_OPENAI`) | +| `X-PasteGuard-Secrets-Redacted` | `true` if secrets were redacted (action: redact) | ## Development diff --git a/config.example.yaml b/config.example.yaml index 205d2b0..3ac44d6 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -90,28 +90,46 @@ pii_detection: # - URL # Secrets Detection settings (Secrets Shield) -# Detects private keys and other secret credentials in requests +# Detects private keys, API keys, tokens and other secret credentials in requests secrets_detection: # Enable secrets detection (default: true) enabled: true # Action to take when secrets are detected: # block: Block the request with HTTP 422 (default, secure-by-default) - # redact: Replace secrets with placeholders and continue (irreversible) + # redact: Replace secrets with placeholders, unmask in response (reversible) # route_local: Route to local provider (only works in route mode) action: block # Secret types to detect + # Private Keys (enabled by default): + # - OPENSSH_PRIVATE_KEY: OpenSSH format (-----BEGIN OPENSSH PRIVATE KEY-----) + # - PEM_PRIVATE_KEY: PEM formats (RSA, PRIVATE KEY, ENCRYPTED PRIVATE KEY) + # + # API Keys (opt-in): + # - API_KEY_OPENAI: OpenAI API keys (sk-...) + # - API_KEY_AWS: AWS Access Keys (AKIA...) + # - API_KEY_GITHUB: GitHub tokens (ghp_, gho_, ghu_, ghs_, ghr_) + # + # Tokens (opt-in): + # - JWT_TOKEN: JSON Web Tokens (eyJ...) + # - BEARER_TOKEN: Bearer tokens in Authorization-style contexts entities: - - OPENSSH_PRIVATE_KEY # OpenSSH format: -----BEGIN OPENSSH PRIVATE KEY----- - - PEM_PRIVATE_KEY # PEM formats: RSA, PRIVATE KEY, ENCRYPTED PRIVATE KEY + - OPENSSH_PRIVATE_KEY + - PEM_PRIVATE_KEY + # Uncomment to detect API keys and tokens: + # - API_KEY_OPENAI + # - API_KEY_AWS + # - API_KEY_GITHUB + # - JWT_TOKEN + # - BEARER_TOKEN # Maximum characters to scan per request (performance limit) # Set to 0 to scan entire request (not recommended for large payloads) max_scan_chars: 200000 # Placeholder format for redaction (only used if action: redact) - # {N} will be replaced with sequential number + # {N} will be replaced with type and sequential number (e.g., API_KEY_OPENAI_1) redact_placeholder: "" # Log detected secret types (never logs secret content) diff --git a/src/routes/proxy.test.ts b/src/routes/proxy.test.ts index 9c7da94..baf6286 100644 --- a/src/routes/proxy.test.ts +++ b/src/routes/proxy.test.ts @@ -157,4 +157,8 @@ MIIEpAIBAAKCAQEAyK8v5Q8v5Q8v5Q8v5Q8v5Q8v5Q8v5Q8v5Q8v5Q8v5Q8v5Q8v expect(res.headers.get("X-PasteGuard-Secrets-Detected")).toBeNull(); expect(res.headers.get("X-PasteGuard-Secrets-Types")).toBeNull(); }); + + // Note: Tests for API_KEY_OPENAI, JWT_TOKEN, etc. require those entity types + // to be enabled in config. Detection is thoroughly tested in detect.test.ts. + // Proxy blocking behavior is tested above with private keys (default entities). }); diff --git a/src/routes/proxy.ts b/src/routes/proxy.ts index 5a79c39..1577707 100644 --- a/src/routes/proxy.ts +++ b/src/routes/proxy.ts @@ -4,8 +4,13 @@ import { Hono } from "hono"; import { HTTPException } from "hono/http-exception"; import { proxy } from "hono/proxy"; import { z } from "zod"; -import { getConfig, type MaskingConfig } from "../config"; -import { detectSecrets, extractTextFromRequest } from "../secrets/detect"; +import { getConfig, type MaskingConfig, type SecretsDetectionConfig } from "../config"; +import { + detectSecrets, + extractTextFromRequest, + type SecretsDetectionResult, +} from "../secrets/detect"; +import { type RedactionContext, redactSecrets, unredactResponse } from "../secrets/redact"; import { getRouter, type MaskDecision, type RoutingDecision } from "../services/decision"; import type { ChatCompletionRequest, @@ -71,25 +76,30 @@ proxyRoutes.post( }), async (c) => { const startTime = Date.now(); - const body = c.req.valid("json") as ChatCompletionRequest; + let body = c.req.valid("json") as ChatCompletionRequest; const config = getConfig(); const router = getRouter(); + // Track secrets detection state for response handling + let secretsResult: SecretsDetectionResult | undefined; + let redactionContext: RedactionContext | undefined; + let secretsRedacted = false; + // Secrets detection runs before PII detection if (config.secrets_detection.enabled) { const text = extractTextFromRequest(body); - const secretsResult = detectSecrets(text, config.secrets_detection); + secretsResult = detectSecrets(text, config.secrets_detection); if (secretsResult.detected) { const secretTypes = secretsResult.matches.map((m) => m.type); const secretTypesStr = secretTypes.join(", "); - // Set headers before returning error - c.header("X-PasteGuard-Secrets-Detected", "true"); - c.header("X-PasteGuard-Secrets-Types", secretTypesStr); - - // Block action (Phase 1) - return 422 error + // Block action - return 422 error if (config.secrets_detection.action === "block") { + // Set headers before returning error + c.header("X-PasteGuard-Secrets-Detected", "true"); + c.header("X-PasteGuard-Secrets-Types", secretTypesStr); + // Log metadata only (no secret content) logRequest( { @@ -123,23 +133,105 @@ proxyRoutes.post( ); } - // TODO: Phase 2 - redact and route_local actions - // For now, if action is not "block", we continue (but this shouldn't happen in Phase 1) + // Redact action - replace secrets with placeholders and continue + if (config.secrets_detection.action === "redact") { + const redactedMessages = redactMessagesWithSecrets( + body.messages, + secretsResult, + config.secrets_detection, + ); + body = { ...body, messages: redactedMessages.messages }; + redactionContext = redactedMessages.context; + secretsRedacted = true; + } + + // route_local action is handled in handleCompletion via secretsResult } } let decision: RoutingDecision; try { - decision = await router.decide(body.messages); + decision = await router.decide(body.messages, secretsResult); } catch (error) { console.error("PII detection error:", error); throw new HTTPException(503, { message: "PII detection service unavailable" }); } - return handleCompletion(c, body, decision, startTime, router); + return handleCompletion( + c, + body, + decision, + startTime, + router, + secretsResult, + redactionContext, + secretsRedacted, + ); }, ); +/** + * Redacts secrets in all messages based on detection result + * Returns redacted messages and the redaction context for unredaction + */ +function redactMessagesWithSecrets( + messages: ChatMessage[], + secretsResult: SecretsDetectionResult, + config: SecretsDetectionConfig, +): { messages: ChatMessage[]; context: RedactionContext } { + // Build a map of message content to redactions + // Since we concatenated all messages with \n, we need to track positions per message + let currentOffset = 0; + const messagePositions: { start: number; end: number }[] = []; + + for (const msg of messages) { + const length = typeof msg.content === "string" ? msg.content.length : 0; + messagePositions.push({ start: currentOffset, end: currentOffset + length }); + currentOffset += length + 1; // +1 for \n separator + } + + // Create redaction context + let context: RedactionContext = { + mapping: {}, + reverseMapping: {}, + counters: {}, + }; + + // Apply redactions to each message + const redactedMessages = messages.map((msg, i) => { + if (typeof msg.content !== "string" || !msg.content) { + return msg; + } + + const msgPos = messagePositions[i]; + + // Filter redactions that fall within this message's position + const messageRedactions = (secretsResult.redactions || []) + .filter((r) => r.start >= msgPos.start && r.end <= msgPos.end) + .map((r) => ({ + ...r, + start: r.start - msgPos.start, + end: r.end - msgPos.start, + })); + + if (messageRedactions.length === 0) { + return msg; + } + + const { redacted, context: updatedContext } = redactSecrets( + msg.content, + messageRedactions, + config, + context, + ); + context = updatedContext; + + return { ...msg, content: redacted }; + }); + + return { messages: redactedMessages, context }; +} + /** * Handle chat completion for both route and mask modes */ @@ -149,6 +241,9 @@ async function handleCompletion( decision: RoutingDecision, startTime: number, router: ReturnType, + secretsResult?: SecretsDetectionResult, + redactionContext?: RedactionContext, + secretsRedacted?: boolean, ) { const client = router.getClient(decision.provider); const maskingConfig = router.getMaskingConfig(); @@ -166,20 +261,11 @@ async function handleCompletion( try { const result = await client.chatCompletion(request, authHeader); - // Check for secrets in the request (for headers, even if not blocking) - const config = getConfig(); - let secretsDetected = false; - let secretsTypes: string[] = []; - if (config.secrets_detection.enabled) { - const text = extractTextFromRequest(body); - const secretsResult = detectSecrets(text, config.secrets_detection); - if (secretsResult.detected) { - secretsDetected = true; - secretsTypes = secretsResult.matches.map((m) => m.type); - } - } + // Determine secrets state from passed result + const secretsDetected = secretsResult?.detected ?? false; + const secretsTypes = secretsResult?.matches.map((m) => m.type) ?? []; - setPasteGuardHeaders(c, decision, secretsDetected, secretsTypes); + setPasteGuardHeaders(c, decision, secretsDetected, secretsTypes, secretsRedacted); if (result.isStreaming) { return handleStreamingResponse( @@ -191,6 +277,7 @@ async function handleCompletion( maskingConfig, secretsDetected, secretsTypes, + redactionContext, ); } @@ -203,6 +290,7 @@ async function handleCompletion( maskingConfig, secretsDetected, secretsTypes, + redactionContext, ); } catch (error) { console.error("LLM request error:", error); @@ -219,6 +307,7 @@ function setPasteGuardHeaders( decision: RoutingDecision, secretsDetected?: boolean, secretsTypes?: string[], + secretsRedacted?: boolean, ) { c.header("X-PasteGuard-Mode", decision.mode); c.header("X-PasteGuard-Provider", decision.provider); @@ -230,10 +319,13 @@ function setPasteGuardHeaders( if (decision.mode === "mask") { c.header("X-PasteGuard-PII-Masked", decision.piiResult.hasPII.toString()); } - if (secretsDetected && secretsTypes) { + if (secretsDetected && secretsTypes && secretsTypes.length > 0) { c.header("X-PasteGuard-Secrets-Detected", "true"); c.header("X-PasteGuard-Secrets-Types", secretsTypes.join(",")); } + if (secretsRedacted) { + c.header("X-PasteGuard-Secrets-Redacted", "true"); + } } /** @@ -248,6 +340,7 @@ function handleStreamingResponse( maskingConfig: MaskingConfig, secretsDetected?: boolean, secretsTypes?: string[], + redactionContext?: RedactionContext, ) { logRequest( createLogData( @@ -266,11 +359,16 @@ function handleStreamingResponse( c.header("Cache-Control", "no-cache"); c.header("Connection", "keep-alive"); - if (isMaskDecision(decision)) { + // Determine if we need to transform the stream + const needsPIIUnmasking = isMaskDecision(decision); + const needsSecretsUnredaction = redactionContext !== undefined; + + if (needsPIIUnmasking || needsSecretsUnredaction) { const unmaskingStream = createUnmaskingStream( result.response, - decision.maskingContext, + needsPIIUnmasking ? decision.maskingContext : undefined, maskingConfig, + redactionContext, ); return c.body(unmaskingStream); } @@ -290,6 +388,7 @@ function handleJsonResponse( maskingConfig: MaskingConfig, secretsDetected?: boolean, secretsTypes?: string[], + redactionContext?: RedactionContext, ) { logRequest( createLogData( @@ -304,11 +403,19 @@ function handleJsonResponse( c.req.header("User-Agent") || null, ); + let response = result.response; + + // First unmask PII if needed if (isMaskDecision(decision)) { - return c.json(unmaskResponse(result.response, decision.maskingContext, maskingConfig)); + response = unmaskResponse(response, decision.maskingContext, maskingConfig); + } + + // Then unredact secrets if needed + if (redactionContext) { + response = unredactResponse(response, redactionContext); } - return c.json(result.response); + return c.json(response); } /** diff --git a/src/services/decision.test.ts b/src/services/decision.test.ts index f01236c..d6e1b90 100644 --- a/src/services/decision.test.ts +++ b/src/services/decision.test.ts @@ -1,4 +1,5 @@ import { describe, expect, test } from "bun:test"; +import type { SecretsDetectionResult, SecretsMatch } from "../secrets/detect"; import type { PIIDetectionResult } from "./pii-detector"; /** @@ -8,7 +9,18 @@ import type { PIIDetectionResult } from "./pii-detector"; function decideRoute( piiResult: PIIDetectionResult, routing: { default: "upstream" | "local"; on_pii_detected: "upstream" | "local" }, + secretsResult?: SecretsDetectionResult, + secretsAction?: "block" | "redact" | "route_local", ): { provider: "upstream" | "local"; reason: string } { + // Check for secrets route_local action first (takes precedence) + if (secretsResult?.detected && secretsAction === "route_local") { + const secretTypes = secretsResult.matches.map((m) => m.type); + return { + provider: "local", + reason: `Secrets detected (route_local): ${secretTypes.join(", ")}`, + }; + } + if (piiResult.hasPII) { const entityTypes = [...new Set(piiResult.newEntities.map((e) => e.entity_type))]; return { @@ -129,3 +141,113 @@ describe("decideRoute", () => { }); }); }); + +/** + * Helper to create a mock SecretsDetectionResult + */ +function createSecretsResult( + detected: boolean, + matches: SecretsMatch[] = [], +): SecretsDetectionResult { + return { + detected, + matches, + redactions: matches.map((m, i) => ({ start: i * 100, end: i * 100 + 50, type: m.type })), + }; +} + +describe("decideRoute with secrets", () => { + const routing = { default: "upstream" as const, on_pii_detected: "local" as const }; + + describe("with route_local action", () => { + test("routes to local when secrets detected", () => { + const piiResult = createPIIResult(false); + const secretsResult = createSecretsResult(true, [{ type: "API_KEY_OPENAI", count: 1 }]); + + const result = decideRoute(piiResult, routing, secretsResult, "route_local"); + + expect(result.provider).toBe("local"); + expect(result.reason).toContain("Secrets detected"); + expect(result.reason).toContain("route_local"); + expect(result.reason).toContain("API_KEY_OPENAI"); + }); + + test("secrets routing takes precedence over PII routing", () => { + // Even with on_pii_detected=upstream, secrets route_local should go to local + const routingUpstream = { + default: "upstream" as const, + on_pii_detected: "upstream" as const, + }; + const piiResult = createPIIResult(true, [{ entity_type: "PERSON" }]); + const secretsResult = createSecretsResult(true, [{ type: "API_KEY_AWS", count: 1 }]); + + const result = decideRoute(piiResult, routingUpstream, secretsResult, "route_local"); + + expect(result.provider).toBe("local"); + expect(result.reason).toContain("Secrets detected"); + }); + + test("routes based on PII when no secrets detected", () => { + const piiResult = createPIIResult(true, [{ entity_type: "EMAIL_ADDRESS" }]); + const secretsResult = createSecretsResult(false); + + const result = decideRoute(piiResult, routing, secretsResult, "route_local"); + + expect(result.provider).toBe("local"); // PII detected -> on_pii_detected=local + expect(result.reason).toContain("PII detected"); + }); + + test("routes to default when no secrets and no PII detected", () => { + const piiResult = createPIIResult(false); + const secretsResult = createSecretsResult(false); + + const result = decideRoute(piiResult, routing, secretsResult, "route_local"); + + expect(result.provider).toBe("upstream"); + expect(result.reason).toBe("No PII detected"); + }); + }); + + describe("with block action", () => { + test("ignores secrets detection for routing (block happens earlier)", () => { + const piiResult = createPIIResult(false); + const secretsResult = createSecretsResult(true, [{ type: "JWT_TOKEN", count: 1 }]); + + const result = decideRoute(piiResult, routing, secretsResult, "block"); + + // With block action, we shouldn't route based on secrets + expect(result.provider).toBe("upstream"); + expect(result.reason).toBe("No PII detected"); + }); + }); + + describe("with redact action", () => { + test("ignores secrets detection for routing (redacted before PII check)", () => { + const piiResult = createPIIResult(false); + const secretsResult = createSecretsResult(true, [{ type: "BEARER_TOKEN", count: 1 }]); + + const result = decideRoute(piiResult, routing, secretsResult, "redact"); + + // With redact action, we route based on PII, not secrets + expect(result.provider).toBe("upstream"); + expect(result.reason).toBe("No PII detected"); + }); + }); + + describe("with multiple secret types", () => { + test("includes all secret types in reason", () => { + const piiResult = createPIIResult(false); + const secretsResult = createSecretsResult(true, [ + { type: "API_KEY_OPENAI", count: 1 }, + { type: "API_KEY_GITHUB", count: 2 }, + { type: "JWT_TOKEN", count: 1 }, + ]); + + const result = decideRoute(piiResult, routing, secretsResult, "route_local"); + + expect(result.reason).toContain("API_KEY_OPENAI"); + expect(result.reason).toContain("API_KEY_GITHUB"); + expect(result.reason).toContain("JWT_TOKEN"); + }); + }); +}); diff --git a/src/services/decision.ts b/src/services/decision.ts index 6cc12e4..40be5ad 100644 --- a/src/services/decision.ts +++ b/src/services/decision.ts @@ -1,4 +1,5 @@ import { type Config, getConfig } from "../config"; +import type { SecretsDetectionResult } from "../secrets/detect"; import { type ChatMessage, LLMClient } from "../services/llm-client"; import { createMaskingContext, type MaskingContext, maskMessages } from "../services/masking"; import { getPIIDetector, type PIIDetectionResult } from "../services/pii-detector"; @@ -53,9 +54,15 @@ export class Router { } /** - * Decides how to handle messages based on mode and PII detection + * Decides how to handle messages based on mode, PII detection, and secrets detection + * + * @param messages - The chat messages to process + * @param secretsResult - Optional secrets detection result (for route_local action) */ - async decide(messages: ChatMessage[]): Promise { + async decide( + messages: ChatMessage[], + secretsResult?: SecretsDetectionResult, + ): Promise { const detector = getPIIDetector(); const piiResult = await detector.analyzeMessages(messages); @@ -63,18 +70,34 @@ export class Router { return await this.decideMask(messages, piiResult); } - return this.decideRoute(piiResult); + return this.decideRoute(piiResult, secretsResult); } /** * Route mode: decides which provider to use + * + * Secrets routing takes precedence over PII routing when action is route_local */ - private decideRoute(piiResult: PIIDetectionResult): RouteDecision { + private decideRoute( + piiResult: PIIDetectionResult, + secretsResult?: SecretsDetectionResult, + ): RouteDecision { const routing = this.config.routing; if (!routing) { throw new Error("Route mode requires routing configuration"); } + // Check for secrets route_local action first (takes precedence) + if (secretsResult?.detected && this.config.secrets_detection.action === "route_local") { + const secretTypes = secretsResult.matches.map((m) => m.type); + return { + mode: "route", + provider: "local", + reason: `Secrets detected (route_local): ${secretTypes.join(", ")}`, + piiResult, + }; + } + // Route based on PII detection if (piiResult.hasPII) { const entityTypes = [...new Set(piiResult.newEntities.map((e) => e.entity_type))]; diff --git a/src/services/stream-transformer.ts b/src/services/stream-transformer.ts index 6405ee8..aa69b2f 100644 --- a/src/services/stream-transformer.ts +++ b/src/services/stream-transformer.ts @@ -1,4 +1,9 @@ import type { MaskingConfig } from "../config"; +import { + flushRedactionBuffer, + type RedactionContext, + unredactStreamChunk, +} from "../secrets/redact"; import { flushStreamBuffer, type MaskingContext, unmaskStreamChunk } from "./masking"; /** @@ -6,15 +11,19 @@ import { flushStreamBuffer, type MaskingContext, unmaskStreamChunk } from "./mas * * Processes Server-Sent Events (SSE) chunks, buffering partial placeholders * and unmasking complete ones before forwarding to the client. + * + * Supports both PII unmasking and secret unredaction, or either alone. */ export function createUnmaskingStream( source: ReadableStream, - context: MaskingContext, + piiContext: MaskingContext | undefined, config: MaskingConfig, + secretsContext?: RedactionContext, ): ReadableStream { const decoder = new TextDecoder(); const encoder = new TextEncoder(); - let contentBuffer = ""; + let piiBuffer = ""; + let secretsBuffer = ""; return new ReadableStream({ async start(controller) { @@ -26,23 +35,36 @@ export function createUnmaskingStream( if (done) { // Flush remaining buffer content before closing - if (contentBuffer) { - const flushed = flushStreamBuffer(contentBuffer, context, config); - if (flushed) { - const finalEvent = { - id: `flush-${Date.now()}`, - object: "chat.completion.chunk", - created: Math.floor(Date.now() / 1000), - choices: [ - { - index: 0, - delta: { content: flushed }, - finish_reason: null, - }, - ], - }; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(finalEvent)}\n\n`)); - } + let flushed = ""; + + // Flush PII buffer first + if (piiBuffer && piiContext) { + flushed = flushStreamBuffer(piiBuffer, piiContext, config); + } else if (piiBuffer) { + flushed = piiBuffer; + } + + // Then flush secrets buffer + if (secretsBuffer && secretsContext) { + flushed += flushRedactionBuffer(secretsBuffer, secretsContext); + } else if (secretsBuffer) { + flushed += secretsBuffer; + } + + if (flushed) { + const finalEvent = { + id: `flush-${Date.now()}`, + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + choices: [ + { + index: 0, + delta: { content: flushed }, + finish_reason: null, + }, + ], + }; + controller.enqueue(encoder.encode(`data: ${JSON.stringify(finalEvent)}\n\n`)); } controller.close(); break; @@ -65,18 +87,34 @@ export function createUnmaskingStream( const content = parsed.choices?.[0]?.delta?.content || ""; if (content) { - // Use streaming unmask - const { output, remainingBuffer } = unmaskStreamChunk( - contentBuffer, - content, - context, - config, - ); - contentBuffer = remainingBuffer; + let processedContent = content; + + // First unmask PII if context provided + if (piiContext) { + const { output, remainingBuffer } = unmaskStreamChunk( + piiBuffer, + processedContent, + piiContext, + config, + ); + piiBuffer = remainingBuffer; + processedContent = output; + } + + // Then unredact secrets if context provided + if (secretsContext && processedContent) { + const { output, remainingBuffer } = unredactStreamChunk( + secretsBuffer, + processedContent, + secretsContext, + ); + secretsBuffer = remainingBuffer; + processedContent = output; + } - if (output) { - // Update the parsed object with unmasked content - parsed.choices[0].delta.content = output; + if (processedContent) { + // Update the parsed object with processed content + parsed.choices[0].delta.content = processedContent; controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); } } else {