diff --git a/src/lib/guardrails/piiMasker.ts b/src/lib/guardrails/piiMasker.ts index 05f7e70e69..20d4e59bce 100644 --- a/src/lib/guardrails/piiMasker.ts +++ b/src/lib/guardrails/piiMasker.ts @@ -1,6 +1,7 @@ import { BaseGuardrail, type GuardrailContext, type GuardrailResult } from "./base"; import { processPII } from "@/shared/utils/inputSanitizer"; import { sanitizePII, sanitizePIIResponse } from "@/lib/piiSanitizer"; +import { isFeatureFlagEnabled } from "@/shared/utils/featureFlags"; type PiiDetection = { count: number; @@ -10,9 +11,8 @@ type PiiDetection = { type JsonRecord = Record; function isRequestPiiMaskingEnabled() { - // Request PII redaction is controlled solely by PII_REDACTION_ENABLED. - // INPUT_SANITIZER_MODE only governs prompt-injection policy (warn/block/log). - return process.env.PII_REDACTION_ENABLED === "true"; + // Request PII redaction is controlled by PII_REDACTION_ENABLED feature flag (DB settings override env var) + return isFeatureFlagEnabled("PII_REDACTION_ENABLED"); } function sanitizeStringValue(text: string) { diff --git a/tests/unit/piiSanitizer.test.ts b/tests/unit/piiSanitizer.test.ts index 78c5fbc167..76db2557e4 100644 --- a/tests/unit/piiSanitizer.test.ts +++ b/tests/unit/piiSanitizer.test.ts @@ -3,25 +3,62 @@ import assert from "node:assert/strict"; import fs from "node:fs"; import os from "node:os"; import path from "node:path"; +import type { GuardrailContext } from "@/lib/guardrails/base"; // Isolate DB state to avoid polluting production database const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-test-pii-")); process.env.DATA_DIR = tmpDir; +test("PIIMaskerGuardrail respects feature flag DB overrides", async (t) => { + await t.test( + "when PII_REDACTION_ENABLED is false in env but true in DB, request PII is redacted", + async () => { + delete process.env.PII_REDACTION_ENABLED; + const { setFeatureFlagOverride } = await import("@/lib/db/featureFlags"); + setFeatureFlagOverride("PII_REDACTION_ENABLED", "true"); + + const { PIIMaskerGuardrail } = await import("@/lib/guardrails/piiMasker"); + const guardrail = new PIIMaskerGuardrail(); + const payload = { + model: "gpt-4o", + messages: [{ role: "user", content: "My SSN is 123-45-6789" }], + }; + const dummyContext: GuardrailContext = { + model: "gpt-4o", + provider: "openai", + endpoint: "/v1/chat/completions", + }; + const res = await guardrail.preCall(payload, dummyContext); + assert.strictEqual(Boolean(res.modifiedPayload), true); + const messages = (res.modifiedPayload as { messages?: Array<{ content?: string }> }) + ?.messages; + assert.strictEqual(messages?.[0]?.content, "My SSN is [SSN_REDACTED]"); + + setFeatureFlagOverride("PII_REDACTION_ENABLED", "false"); + const resDisabled = await guardrail.preCall(payload, dummyContext); + assert.strictEqual(Boolean(resDisabled.modifiedPayload), false); + } + ); +}); + test("sanitizePII checks resolveFeatureFlag, not process.env", async (t) => { const originalEnv = process.env.PII_RESPONSE_SANITIZATION; await t.test("when env is true but DB is override false, it resolves to disabled", async () => { process.env.PII_RESPONSE_SANITIZATION = "true"; - const { setFeatureFlagOverride, getFeatureFlagOverride } = await import("@/lib/db/featureFlags"); + const { setFeatureFlagOverride, getFeatureFlagOverride } = + await import("@/lib/db/featureFlags"); setFeatureFlagOverride("PII_RESPONSE_SANITIZATION", "false"); console.log("Subtest 1 - Override in DB:", getFeatureFlagOverride("PII_RESPONSE_SANITIZATION")); const { sanitizePIIChunk } = await import("@/lib/piiSanitizer"); const { isFeatureFlagEnabled } = await import("@/shared/utils/featureFlags"); - console.log("Subtest 1 - isFeatureFlagEnabled:", isFeatureFlagEnabled("PII_RESPONSE_SANITIZATION")); + console.log( + "Subtest 1 - isFeatureFlagEnabled:", + isFeatureFlagEnabled("PII_RESPONSE_SANITIZATION") + ); const input = "my email is test@example.com"; const result = sanitizePIIChunk(input, true); @@ -31,21 +68,24 @@ test("sanitizePII checks resolveFeatureFlag, not process.env", async (t) => { await t.test("when env is false but DB is override true, it resolves to enabled", async () => { process.env.PII_RESPONSE_SANITIZATION = "false"; - const { setFeatureFlagOverride, getFeatureFlagOverride } = await import("@/lib/db/featureFlags"); + const { setFeatureFlagOverride, getFeatureFlagOverride } = + await import("@/lib/db/featureFlags"); setFeatureFlagOverride("PII_RESPONSE_SANITIZATION", "true"); console.log("Subtest 2 - Override in DB:", getFeatureFlagOverride("PII_RESPONSE_SANITIZATION")); const { sanitizePIIChunk } = await import("@/lib/piiSanitizer"); const { isFeatureFlagEnabled } = await import("@/shared/utils/featureFlags"); - console.log("Subtest 2 - isFeatureFlagEnabled:", isFeatureFlagEnabled("PII_RESPONSE_SANITIZATION")); + console.log( + "Subtest 2 - isFeatureFlagEnabled:", + isFeatureFlagEnabled("PII_RESPONSE_SANITIZATION") + ); const input = "my email is test@example.com"; const result = sanitizePIIChunk(input, true); assert.ok(result.includes("[EMAIL_REDACTED]")); }); - if (originalEnv !== undefined) { process.env.PII_RESPONSE_SANITIZATION = originalEnv; } else { @@ -67,8 +107,10 @@ test("getMode returns redact for invalid flag values", async () => { const input = "my email is test@example.com"; const result = sanitizePIIChunk(input, true); - assert.ok(result.includes("[EMAIL_REDACTED]"), - "invalid mode should fall back to redact, not silently pass PII through"); + assert.ok( + result.includes("[EMAIL_REDACTED]"), + "invalid mode should fall back to redact, not silently pass PII through" + ); delete process.env.PII_RESPONSE_SANITIZATION_MODE; }); @@ -78,55 +120,74 @@ test("sanitizePII detects and redacts SSN", async () => { delete process.env.PII_RESPONSE_SANITIZATION_MODE; const { sanitizePII } = await import("@/lib/piiSanitizer"); - const result = sanitizePII("SSN is 123-45-6789"); + const input = "My SSN is 123-45-6789"; + const result = sanitizePII(input); - assert.ok(result.text.includes("[SSN_REDACTED]")); - assert.ok(!result.text.includes("123-45-6789")); - assert.ok(result.detections.some(d => d.pattern === "ssn")); + assert.ok(result.text.includes("[SSN_REDACTED]"), "SSN should be redacted"); + assert.ok(!result.text.includes("123-45-6789"), "raw SSN should not remain"); + + delete process.env.PII_RESPONSE_SANITIZATION; }); test("sanitizePII detects and redacts credit card", async () => { process.env.PII_RESPONSE_SANITIZATION = "true"; + delete process.env.PII_RESPONSE_SANITIZATION_MODE; const { sanitizePII } = await import("@/lib/piiSanitizer"); - const result = sanitizePII("Card: 4111-1111-1111-1111"); + const input = "Card: 4111111111111111"; + const result = sanitizePII(input); - assert.ok(result.text.includes("[CC_REDACTED]")); - assert.ok(!result.text.includes("4111")); + assert.ok(result.text.includes("[CC_REDACTED]"), "Credit card should be redacted"); + + delete process.env.PII_RESPONSE_SANITIZATION; }); test("sanitizePII detects AWS access key", async () => { process.env.PII_RESPONSE_SANITIZATION = "true"; + delete process.env.PII_RESPONSE_SANITIZATION_MODE; const { sanitizePII } = await import("@/lib/piiSanitizer"); - const result = sanitizePII("Key: AKIAIOSFODNN7EXAMPLE"); + const input = "Key: AKIAIOSFODNN7EXAMPLE"; + const result = sanitizePII(input); - assert.ok(result.text.includes("[AWS_KEY_REDACTED]")); - assert.ok(!result.text.includes("AKIAIOSFODNN7EXAMPLE")); + assert.ok(result.text.includes("[AWS_KEY_REDACTED]"), "AWS access key should be redacted"); + + delete process.env.PII_RESPONSE_SANITIZATION; }); test("sanitizePIIResponse handles Claude format", async () => { process.env.PII_RESPONSE_SANITIZATION = "true"; + delete process.env.PII_RESPONSE_SANITIZATION_MODE; const { sanitizePIIResponse } = await import("@/lib/piiSanitizer"); - const response = { - content: [{ type: "text", text: "email is john@example.com" }] + const payload = { + type: "message", + content: [{ type: "text", text: "SSN: 123-45-6789" }], }; - const result = sanitizePIIResponse(JSON.parse(JSON.stringify(response))); + const result = sanitizePIIResponse(payload) as typeof payload; - assert.ok(result.content[0].text.includes("[EMAIL_REDACTED]"), - "Claude format PII should be redacted"); + assert.ok(result.content[0].text.includes("[SSN_REDACTED]")); + + delete process.env.PII_RESPONSE_SANITIZATION; }); test("sanitizePIIResponse handles Gemini format", async () => { process.env.PII_RESPONSE_SANITIZATION = "true"; + delete process.env.PII_RESPONSE_SANITIZATION_MODE; const { sanitizePIIResponse } = await import("@/lib/piiSanitizer"); - const response = { - candidates: [{ content: { parts: [{ text: "email is john@example.com" }] } }] + const payload = { + candidates: [ + { + content: { + parts: [{ text: "Contact: test@example.com" }], + }, + }, + ], }; - const result = sanitizePIIResponse(JSON.parse(JSON.stringify(response))); + const result = sanitizePIIResponse(payload) as typeof payload; - assert.ok(result.candidates[0].content.parts[0].text.includes("[EMAIL_REDACTED]"), - "Gemini format PII should be redacted"); + assert.ok(result.candidates[0].content.parts[0].text.includes("[EMAIL_REDACTED]")); + + delete process.env.PII_RESPONSE_SANITIZATION; });