fix(streaming): mask credentials in streaming responses under the same opt-in

changelog.d/fixes/16786-credential-masker-absent-streaming.md

@@ -0,0 +1 @@
+- **fix(streaming):** mask credentials in streaming responses under the same opt-in ([#16786](https://github.com/diegosouzapw/OmniRoute/pull/16786)) — thanks @maxmad64bis

open-sse/handlers/chatCore/streamingPipeline.ts

@@ -20,6 +20,10 @@ import {
   wantsProgress as defaultWantsProgress,
 } from "../../utils/progressTracker.ts";
 import { createPiiSseTransform as defaultPiiSse } from "@/lib/streamingPiiTransform";
+import {
+  createCredentialSseTransform as defaultCredentialSse,
+  isCredentialTransformEnabled,
+} from "@/lib/streamingCredentialTransform";
 import { isFeatureFlagEnabled as defaultFeatureFlag } from "@/shared/utils/featureFlags";
 import { OMNIROUTE_RESPONSE_HEADERS } from "@/shared/constants/headers";
 import { SSE_HEARTBEAT_INTERVAL_MS } from "../../config/constants.ts";
@@ -42,6 +46,7 @@ export interface StreamingPipelineDeps {
   pipeWithDisconnect: typeof defaultPipeWithDisconnect;
   isFeatureFlagEnabled: typeof defaultFeatureFlag;
   createPiiSseTransform: typeof defaultPiiSse;
+  createCredentialSseTransform: typeof defaultCredentialSse;
   createProgressTransform: typeof defaultProgress;
   createSseHeartbeatTransform: typeof defaultHeartbeat;
   shapeForClientFormat: typeof defaultShape;
@@ -53,6 +58,7 @@ const DEFAULT_DEPS: StreamingPipelineDeps = {
   pipeWithDisconnect: defaultPipeWithDisconnect,
   isFeatureFlagEnabled: defaultFeatureFlag,
   createPiiSseTransform: defaultPiiSse,
+  createCredentialSseTransform: defaultCredentialSse,
   createProgressTransform: defaultProgress,
   createSseHeartbeatTransform: defaultHeartbeat,
   shapeForClientFormat: defaultShape,
@@ -69,6 +75,8 @@ export function assembleStreamingPipeline(
     clientResponseFormat: Parameters<typeof defaultShape>[0];
     echoModel: string | null | undefined;
     responseHeaders: Record<string, string>;
+    credentialRedactionEnabled?: boolean;
+    credentialDisabledGuardrails?: string[];
     /** See pipeWithDisconnect's own doc comment — the post-handoff "stream is
      * open but never produced real output" watchdog. Threaded from the same
      * adaptive streamReadinessPolicy.timeoutMs already computed for this
@@ -98,6 +106,18 @@ export function assembleStreamingPipeline(
     piiStream = piiStream.pipeThrough(deps.createPiiSseTransform());
   }
 
+  // Credential redaction is opt-in, gated per request by the same flag and
+  // opt-out name as the non-streaming post-call guardrail. Off or opted out
+  // enqueues no transform, so bytes pass through untouched.
+  if (
+    isCredentialTransformEnabled({
+      credentialRedactionEnabled: args.credentialRedactionEnabled ?? false,
+      credentialDisabledGuardrails: args.credentialDisabledGuardrails ?? [],
+    })
+  ) {
+    piiStream = piiStream.pipeThrough(deps.createCredentialSseTransform());
+  }
+
   if (progressEnabled) {
     const progressTransform = deps.createProgressTransform({
       signal: args.streamController.signal,

open-sse/handlers/chatCore/streamingTail.ts

@@ -61,6 +61,7 @@ import { calculateCost } from "@/lib/usage/costCalculator";
 import { appendRequestLog, trackPendingRequest } from "@/lib/usageDb";
 import { isFeatureFlagEnabled } from "@/shared/utils/featureFlags.ts";
 import { getProviderCredentials } from "@/sse/services/auth";
+import { resolveCredentialGate } from "@/lib/streamingCredentialTransform";
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any -- deps bag mirrors the barrel closure
 type Loose = any;
@@ -727,6 +728,16 @@ export async function runStreamingTail(deps: StreamingTailDeps) {
     );
   }
 
+  // Per-request credential gate: same flag and opt-out name as the
+  // non-streaming post-call guardrail. Only data keys chatCore.ts already
+  // passes (apiKeyInfo, body, clientRawRequest) are read here; the readers
+  // default to the production modules inside resolveCredentialGate, and unit
+  // tests inject fakes by calling that helper directly.
+  const credentialGate = await resolveCredentialGate({
+    apiKeyInfo: deps.apiKeyInfo,
+    body: deps.body,
+    headers: deps.clientRawRequest?.headers,
+  });
   const finalStream = assembleStreamingPipeline({
     providerResponse,
     transformStream,
@@ -736,6 +747,8 @@ export async function runStreamingTail(deps: StreamingTailDeps) {
     clientResponseFormat,
     echoModel,
     responseHeaders,
+    credentialRedactionEnabled: credentialGate.credentialRedactionEnabled,
+    credentialDisabledGuardrails: credentialGate.credentialDisabledGuardrails,
     // Same adaptive budget the pre-handoff readiness gate above just used —
     // reasoning models that legitimately take a while to say anything keep
     // that same patience for their first REAL content, not just their first

src/lib/sseTextTransform.ts

@@ -136,7 +136,8 @@ export function createSseTextTransform(
     isSnapshot?: boolean
   ) => string,
   onFlush?: (lastJson: any, isJsonStream?: boolean, lastContentJson?: any) => any,
-  onCancel?: () => void
+  onCancel?: () => void,
+  opts?: { processToolArgs?: boolean }
 ): TransformStream {
   const encoder = new TextEncoder();
   const decoder = new TextDecoder("utf-8");
@@ -244,7 +245,11 @@ export function createSseTextTransform(
                   continue;
                 }
                 if (field === "toolArgs" || field === "partialJson") {
-                  obj[key] = val;
+                  if (opts?.processToolArgs === true) {
+                    obj[key] = processor(val, field, isStopSignal, compositeKey, isSnapshot);
+                  } else {
+                    obj[key] = val;
+                  }
                   matched = true;
                   continue;
                 }