wehub-resource-sync
350 lines
14 KiB
TypeScript
350 lines
14 KiB
TypeScript
/**
|
|
* Tests for secrets regex parsing, compilation, and obfuscation.
|
|
*/
|
|
|
|
import { describe, expect, it } from "bun:test";
|
|
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
|
import type { AssistantMessage, Context, Message } from "@oh-my-pi/pi-ai";
|
|
import {
|
|
deobfuscateAgentMessages,
|
|
deobfuscateToolArguments,
|
|
obfuscateMessages,
|
|
obfuscateProviderContext,
|
|
SecretObfuscator,
|
|
} from "@oh-my-pi/pi-coding-agent/secrets/obfuscator";
|
|
import { compileSecretRegex } from "@oh-my-pi/pi-coding-agent/secrets/regex";
|
|
import { type } from "arktype";
|
|
|
|
describe("compileSecretRegex", () => {
|
|
it("adds global flag when not provided", () => {
|
|
const regex = compileSecretRegex("api[_-]?key\\s*=\\s*\\w+", "i");
|
|
expect(regex.source).toBe("api[_-]?key\\s*=\\s*\\w+");
|
|
expect(regex.flags).toBe("gi");
|
|
});
|
|
|
|
it("defaults to global flag when no flags provided", () => {
|
|
const regex = compileSecretRegex("api[_-]?key\\s*=\\s*\\w+");
|
|
expect(regex.source).toBe("api[_-]?key\\s*=\\s*\\w+");
|
|
expect(regex.flags).toBe("g");
|
|
});
|
|
|
|
it("rejects invalid regex pattern", () => {
|
|
expect(() => compileSecretRegex("(")).toThrow();
|
|
});
|
|
it("rejects invalid regex flags", () => {
|
|
expect(() => compileSecretRegex("x", "zz")).toThrow();
|
|
});
|
|
});
|
|
|
|
describe("SecretObfuscator regex behavior", () => {
|
|
it("obfuscates and deobfuscates regex matches with flags", () => {
|
|
const obfuscator = new SecretObfuscator([{ type: "regex", content: "api[_-]?key\\s*=\\s*\\w+", flags: "i" }]);
|
|
const original = "API_KEY=abc and api-key=def";
|
|
const obfuscated = obfuscator.obfuscate(original);
|
|
expect(obfuscated).not.toEqual(original);
|
|
expect(obfuscator.deobfuscate(obfuscated)).toEqual(original);
|
|
});
|
|
|
|
it("supports bare regex patterns without explicit flags", () => {
|
|
const obfuscator = new SecretObfuscator([{ type: "regex", content: "api[_-]?key\\s*=\\s*\\w+" }]);
|
|
const text = "api_key=abc and API_KEY=def";
|
|
const obfuscated = obfuscator.obfuscate(text);
|
|
expect(obfuscated).not.toEqual(text);
|
|
expect(obfuscator.deobfuscate(obfuscated)).toEqual(text);
|
|
});
|
|
it("deobfuscates placeholders through tool-call arguments", () => {
|
|
const obfuscator = new SecretObfuscator([{ type: "regex", content: "api[_-]?key\\s*=\\s*\\w+", flags: "i" }]);
|
|
const original = { cmd: "API_KEY=abc and api-key=def", status: "ok", nested: { note: "API_KEY=zzz" } };
|
|
const obfuscated = {
|
|
cmd: obfuscator.obfuscate(original.cmd),
|
|
status: original.status,
|
|
nested: { note: obfuscator.obfuscate(original.nested.note) },
|
|
};
|
|
expect(JSON.stringify(obfuscated)).not.toContain("API_KEY=abc");
|
|
expect(deobfuscateToolArguments(obfuscator, obfuscated)).toEqual(original);
|
|
});
|
|
|
|
it("obfuscates conversation messages but leaves the system prompt untouched", () => {
|
|
const secret = "SUPER_SECRET_TOKEN_12345";
|
|
const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]);
|
|
const context: Context = {
|
|
systemPrompt: [`workspace contains ${secret}`],
|
|
messages: [{ role: "user", content: `use ${secret}`, timestamp: 1 }],
|
|
};
|
|
|
|
const obfuscated = obfuscateProviderContext(obfuscator, context);
|
|
|
|
// Conversation messages are redacted (and round-trip back to the secret)...
|
|
expect(JSON.stringify(obfuscated.messages)).not.toContain(secret);
|
|
expect(obfuscator.deobfuscate(JSON.stringify(obfuscated.messages))).toContain(secret);
|
|
// ...but the author-controlled system prompt passes through by reference.
|
|
expect(obfuscated.systemPrompt).toBe(context.systemPrompt);
|
|
});
|
|
|
|
it("leaves tool schemas untouched in provider context (no clone, no redaction)", () => {
|
|
const secret = "SUPER_SECRET_TOKEN_12345";
|
|
const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]);
|
|
const parameters = type({
|
|
note: "string",
|
|
}).describe(`write ${secret}`);
|
|
const context: Context = {
|
|
messages: [],
|
|
tools: [
|
|
{
|
|
name: "extension_tool",
|
|
description: `preserve ${secret}`,
|
|
parameters,
|
|
},
|
|
],
|
|
};
|
|
|
|
const obfuscated = obfuscateProviderContext(obfuscator, context);
|
|
|
|
expect(obfuscated.tools).toBe(context.tools);
|
|
expect(obfuscated.tools?.[0]?.parameters).toBe(parameters);
|
|
});
|
|
|
|
it("redacts only user, tool-result, and user-attributed developer messages", () => {
|
|
const secret = "SUPER_SECRET_TOKEN_12345";
|
|
const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]);
|
|
const userMsg: Message = { role: "user", content: `user says ${secret}`, timestamp: 1 };
|
|
const systemDeveloperMsg: Message = { role: "developer", content: `system reminder ${secret}`, timestamp: 1 };
|
|
const fileMentionMsg: Message = {
|
|
role: "developer",
|
|
content: `<file>${secret}</file>`,
|
|
attribution: "user",
|
|
timestamp: 1,
|
|
};
|
|
const assistantMsg: Message = {
|
|
role: "assistant",
|
|
content: [
|
|
{
|
|
type: "toolCall",
|
|
id: "call_1",
|
|
name: "handoff",
|
|
arguments: { note: secret },
|
|
intent: `handoff ${secret}`,
|
|
},
|
|
],
|
|
api: "test",
|
|
provider: "test",
|
|
model: "test",
|
|
usage: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "toolUse",
|
|
timestamp: 1,
|
|
};
|
|
const toolResultMsg: Message = {
|
|
role: "toolResult",
|
|
toolCallId: "call_1",
|
|
toolName: "read",
|
|
content: [{ type: "text", text: `tool output ${secret}` }],
|
|
isError: false,
|
|
timestamp: 1,
|
|
};
|
|
|
|
const obfuscated = obfuscateMessages(obfuscator, [
|
|
userMsg,
|
|
systemDeveloperMsg,
|
|
fileMentionMsg,
|
|
assistantMsg,
|
|
toolResultMsg,
|
|
]);
|
|
|
|
// User, user-attributed developer, and tool results are redacted.
|
|
expect(JSON.stringify(obfuscated[0])).not.toContain(secret);
|
|
expect(JSON.stringify(obfuscated[2])).not.toContain(secret);
|
|
expect(JSON.stringify(obfuscated[4])).not.toContain(secret);
|
|
// System developer reminders and assistant output pass through untouched (same reference).
|
|
expect(obfuscated[1]).toBe(systemDeveloperMsg);
|
|
expect(obfuscated[3]).toBe(assistantMsg);
|
|
});
|
|
|
|
it("never rewrites inline image bytes", () => {
|
|
const secret = "SUPER_SECRET_TOKEN_12345";
|
|
const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]);
|
|
// A base64 payload that literally contains the secret substring must survive byte-identical;
|
|
// rewriting it would corrupt the data URL (the Codex "invalid base64" failure).
|
|
const imageData = `iVBORw0KGgo${secret}AAAASUVORK5CYII=`;
|
|
const message: Message = {
|
|
role: "toolResult",
|
|
toolCallId: "call_1",
|
|
toolName: "read",
|
|
content: [
|
|
{ type: "text", text: `read ${secret}` },
|
|
{ type: "image", data: imageData, mimeType: "image/png" },
|
|
],
|
|
isError: false,
|
|
timestamp: 1,
|
|
};
|
|
|
|
const [obfuscated] = obfuscateMessages(obfuscator, [message]) as [typeof message];
|
|
const blocks = obfuscated.content;
|
|
const image = blocks[1];
|
|
const text = blocks[0];
|
|
// Image bytes untouched...
|
|
expect(image.type === "image" && image.data).toBe(imageData);
|
|
// ...while the adjacent text is redacted.
|
|
expect(text.type === "text" && text.text.includes(secret)).toBe(false);
|
|
});
|
|
|
|
it("ignores configured plain secrets shorter than 8 characters", () => {
|
|
const obfuscator = new SecretObfuscator([{ type: "plain", content: "esp" }]);
|
|
expect(obfuscator.hasSecrets()).toBe(false);
|
|
expect(obfuscator.obfuscate("the response despite whitespace")).toBe("the response despite whitespace");
|
|
});
|
|
|
|
it("ignores regex matches shorter than 8 characters", () => {
|
|
const obfuscator = new SecretObfuscator([{ type: "regex", content: "esp" }]);
|
|
expect(obfuscator.obfuscate("the response despite whitespace")).toBe("the response despite whitespace");
|
|
});
|
|
});
|
|
|
|
describe("SecretObfuscator cross-turn cache stability", () => {
|
|
// The provider prompt cache is content-addressed: convertToLlm / transformProviderContext
|
|
// re-run obfuscation over the WHOLE message array every turn, so a non-deterministic
|
|
// placeholder for the same secret would rewrite already-sent prefix bytes and bust the
|
|
// cache (cacheWrite @ $6.25/M vs cacheRead @ $0.50/M on opus). These tests pin the
|
|
// determinism that makes obfuscation cache-safe so a future change cannot silently
|
|
// reintroduce per-turn cache invalidation.
|
|
it("produces byte-identical output when re-obfuscating the same content across turns", () => {
|
|
const secret = "SUPER_SECRET_TOKEN_12345";
|
|
const obfuscator = new SecretObfuscator([
|
|
{ type: "plain", content: secret },
|
|
{ type: "regex", content: "tok_[a-z0-9]+" },
|
|
]);
|
|
const messages: Message[] = [{ role: "user", content: `use ${secret} and tok_abc123`, timestamp: 1 }];
|
|
|
|
const turn1 = JSON.stringify(obfuscateMessages(obfuscator, messages));
|
|
const turn2 = JSON.stringify(obfuscateMessages(obfuscator, messages));
|
|
|
|
expect(turn1).not.toContain(secret);
|
|
expect(turn1).not.toContain("tok_abc123");
|
|
// Identical bytes on the second pass → the cached prefix stays valid.
|
|
expect(turn2).toEqual(turn1);
|
|
});
|
|
|
|
it("keeps earlier message placeholders stable when a later message reveals a new regex secret", () => {
|
|
const obfuscator = new SecretObfuscator([{ type: "regex", content: "tok_[a-z0-9]+" }]);
|
|
const early: Message[] = [{ role: "user", content: "first uses tok_aaaa", timestamp: 1 }];
|
|
|
|
// Turn N: only the early message exists; tok_aaa mints a fresh placeholder.
|
|
const earlyTurnN = JSON.stringify(obfuscateMessages(obfuscator, early));
|
|
expect(earlyTurnN).not.toContain("tok_aaaa");
|
|
|
|
// A later turn reveals a brand-new secret. Lazy regex discovery assigns it a fresh
|
|
// index — this MUST NOT shift the placeholder already minted for tok_aaa.
|
|
const later: Message[] = [{ role: "user", content: "later uses tok_bbbb", timestamp: 2 }];
|
|
const laterOut = JSON.stringify(obfuscateMessages(obfuscator, later));
|
|
expect(laterOut).not.toContain("tok_bbbb");
|
|
|
|
// Re-obfuscate the early message after the new discovery: identical bytes → the
|
|
// already-cached prefix for the early message stays valid.
|
|
const earlyTurnNPlus1 = JSON.stringify(obfuscateMessages(obfuscator, early));
|
|
expect(earlyTurnNPlus1).toEqual(earlyTurnN);
|
|
});
|
|
});
|
|
|
|
describe("deobfuscateAgentMessages (display restore)", () => {
|
|
it("restores assistant text and tool calls while leaving raw user text and thinking untouched", () => {
|
|
const secret = "DISPLAY_SECRET_TOKEN_123";
|
|
const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]);
|
|
const placeholder = obfuscator.obfuscate(secret);
|
|
expect(placeholder).not.toBe(secret);
|
|
|
|
const userMsg: AgentMessage = { role: "user", content: `literal ${placeholder} token`, timestamp: 1 };
|
|
const assistantMsg: AgentMessage = {
|
|
role: "assistant",
|
|
content: [
|
|
{ type: "text", text: `answer ${placeholder}` },
|
|
{ type: "thinking", thinking: `reason ${placeholder}` },
|
|
{
|
|
type: "toolCall",
|
|
id: "call_1",
|
|
name: "read",
|
|
arguments: { path: `path ${placeholder}` },
|
|
intent: `intent ${placeholder}`,
|
|
},
|
|
],
|
|
api: "test",
|
|
provider: "test",
|
|
model: "test",
|
|
usage: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "toolUse",
|
|
timestamp: 2,
|
|
};
|
|
const branchSummary: AgentMessage = {
|
|
role: "branchSummary",
|
|
summary: `branch ${placeholder}`,
|
|
fromId: "x",
|
|
timestamp: 3,
|
|
};
|
|
const compactionSummary: AgentMessage = {
|
|
role: "compactionSummary",
|
|
summary: `compact ${placeholder}`,
|
|
shortSummary: `short ${placeholder}`,
|
|
tokensBefore: 0,
|
|
timestamp: 4,
|
|
};
|
|
|
|
const restored = deobfuscateAgentMessages(obfuscator, [userMsg, assistantMsg, branchSummary, compactionSummary]);
|
|
|
|
// Assistant text and tool-call args/intent are restored to the real secret.
|
|
const restoredAssistant = restored[1] as AssistantMessage;
|
|
const assistantJson = JSON.stringify(restoredAssistant.content);
|
|
expect(assistantJson).toContain(secret);
|
|
expect(assistantJson).not.toContain(`answer ${placeholder}`);
|
|
expect(assistantJson).not.toContain(`path ${placeholder}`);
|
|
expect(assistantJson).not.toContain(`intent ${placeholder}`);
|
|
// Opaque thinking is never walked: placeholder-shaped bytes survive unchanged.
|
|
expect(assistantJson).toContain(`reason ${placeholder}`);
|
|
expect(assistantJson).not.toContain(`reason ${secret}`);
|
|
// Model-generated summaries are restored.
|
|
expect((restored[2] as { summary: string }).summary).toBe(`branch ${secret}`);
|
|
expect((restored[3] as { summary: string; shortSummary?: string }).summary).toBe(`compact ${secret}`);
|
|
expect((restored[3] as { summary: string; shortSummary?: string }).shortSummary).toBe(`short ${secret}`);
|
|
// The user message is persisted raw and never walked: a literal placeholder-shaped token
|
|
// survives byte-identical (same reference) rather than being turned into the secret.
|
|
expect(restored[0]).toBe(userMsg);
|
|
});
|
|
|
|
it("restores compactionSummary block text while leaving snapcompact image bytes intact", () => {
|
|
const secret = "BLOCKS_SECRET_TOKEN_456";
|
|
const obfuscator = new SecretObfuscator([{ type: "plain", content: secret }]);
|
|
const placeholder = obfuscator.obfuscate(secret);
|
|
const imageData = `frame${secret}bytes==`;
|
|
const message: AgentMessage = {
|
|
role: "compactionSummary",
|
|
summary: `summary ${placeholder}`,
|
|
tokensBefore: 0,
|
|
blocks: [
|
|
{ type: "text", text: `archived ${placeholder}` },
|
|
{ type: "image", data: imageData, mimeType: "image/png" },
|
|
],
|
|
timestamp: 1,
|
|
};
|
|
|
|
const [restored] = deobfuscateAgentMessages(obfuscator, [message]) as [typeof message];
|
|
const blocks = restored.blocks ?? [];
|
|
const text = blocks[0];
|
|
const image = blocks[1];
|
|
// Archived text is restored to the real secret...
|
|
expect(text.type === "text" && text.text).toBe(`archived ${secret}`);
|
|
// ...while the snapcompact image bytes pass through untouched.
|
|
expect(image.type === "image" && image.data).toBe(imageData);
|
|
});
|
|
});
|