Files
2026-07-13 13:39:12 +08:00

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);
});
});