62 lines
2.8 KiB
TypeScript
62 lines
2.8 KiB
TypeScript
/**
|
|
* TDD for risk-gate mask/restore round-trip.
|
|
* Run: node --import tsx/esm --test tests/unit/compression/riskGateStep.test.ts
|
|
*/
|
|
import { describe, it } from "node:test";
|
|
import assert from "node:assert/strict";
|
|
|
|
import { applyRiskMask, restoreRiskBlocks } from "../../../open-sse/services/compression/riskGate/riskGateStep.ts";
|
|
import {
|
|
DEFAULT_COMPRESSION_CONFIG,
|
|
type CompressionConfig,
|
|
} from "../../../open-sse/services/compression/types.ts";
|
|
|
|
const PEM = "-----BEGIN PRIVATE KEY-----\nMIIBVQbody\n-----END PRIVATE KEY-----";
|
|
|
|
describe("applyRiskMask / restoreRiskBlocks", () => {
|
|
it("masks risky spans in message content and round-trips byte-identically", () => {
|
|
const body = { messages: [{ role: "user", content: `here is a key:\n${PEM}\nthanks` }] };
|
|
const { maskedBody, blocks, stats } = applyRiskMask(body, { enabled: true });
|
|
const maskedText = (maskedBody.messages as Array<{ content: string }>)[0].content;
|
|
assert.ok(!maskedText.includes("BEGIN PRIVATE KEY"), "secret removed from masked body");
|
|
assert.ok(maskedText.includes("OMNI_CAVEMAN"), "uses SENTINEL placeholder family");
|
|
assert.equal(stats.spansProtected, 1);
|
|
assert.equal(stats.categories.private_key, 1);
|
|
|
|
const restored = restoreRiskBlocks(maskedBody, blocks);
|
|
assert.equal((restored.messages as Array<{ content: string }>)[0].content, body.messages[0].content);
|
|
});
|
|
|
|
it("is a no-op when nothing risky is present", () => {
|
|
const body = { messages: [{ role: "user", content: "just normal prose, nothing to see" }] };
|
|
const { maskedBody, blocks, stats } = applyRiskMask(body, { enabled: true });
|
|
assert.equal(blocks.length, 0);
|
|
assert.equal(stats.spansProtected, 0);
|
|
assert.deepEqual(maskedBody, body);
|
|
});
|
|
|
|
it("masks text parts inside array (multimodal) content", () => {
|
|
const body = {
|
|
messages: [{ role: "user", content: [{ type: "text", text: `key:\n${PEM}` }, { type: "image_url", image_url: { url: "x" } }] }],
|
|
};
|
|
const { maskedBody, blocks } = applyRiskMask(body, { enabled: true });
|
|
const part = (maskedBody.messages as Array<{ content: Array<{ type: string; text?: string }> }>)[0].content[0];
|
|
assert.ok(!part.text!.includes("BEGIN PRIVATE KEY"));
|
|
const restored = restoreRiskBlocks(maskedBody, blocks);
|
|
assert.equal(
|
|
(restored.messages as Array<{ content: Array<{ text?: string }> }>)[0].content[0].text,
|
|
`key:\n${PEM}`
|
|
);
|
|
});
|
|
});
|
|
|
|
describe("risk-gate config defaults", () => {
|
|
it("ships disabled by default", () => {
|
|
assert.equal(DEFAULT_COMPRESSION_CONFIG.riskGate?.enabled ?? false, false);
|
|
});
|
|
it("accepts a riskGate field on CompressionConfig", () => {
|
|
const cfg: CompressionConfig = { ...DEFAULT_COMPRESSION_CONFIG, riskGate: { enabled: true } };
|
|
assert.equal(cfg.riskGate?.enabled, true);
|
|
});
|
|
});
|