|
| 1 | +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; |
| 2 | +import plugin from "../src/index.ts"; |
| 3 | +import { applyZdrHeader, isZdrEnabled, ZDR_HEADER } from "../src/hooks/zdr.ts"; |
| 4 | + |
| 5 | +const URL_CC = "https://api.commandcode.ai/provider/v1/chat/completions"; |
| 6 | + |
| 7 | +function makeCtx() { |
| 8 | + const hooks: Record<string, ((event: any) => Promise<void>)[]> = { |
| 9 | + "http.request": [], |
| 10 | + "http.response": [], |
| 11 | + retry: [], |
| 12 | + }; |
| 13 | + const ctx: any = { |
| 14 | + session: { |
| 15 | + hook: async (name: any, cb: any) => { |
| 16 | + hooks[name]?.push(cb); |
| 17 | + return { dispose: async () => {} }; |
| 18 | + }, |
| 19 | + synthetic: async () => ({}), |
| 20 | + }, |
| 21 | + model: { transform: async () => {} }, |
| 22 | + tool: { transform: async () => {} }, |
| 23 | + command: { transform: async () => {} }, |
| 24 | + }; |
| 25 | + return { ctx, hooks }; |
| 26 | +} |
| 27 | + |
| 28 | +async function runHook(hooks: Record<string, ((event: any) => Promise<void>)[]>, request: Request) { |
| 29 | + const event: any = { request, sessionID: "ses_test" }; |
| 30 | + for (const cb of hooks["http.request"] ?? []) { |
| 31 | + await cb(event); |
| 32 | + } |
| 33 | + return event.request as Request; |
| 34 | +} |
| 35 | + |
| 36 | +describe("isZdrEnabled", () => { |
| 37 | + test("only the documented opt-in value enables ZDR", () => { |
| 38 | + expect(isZdrEnabled({ CMD_ZDR: "1" })).toBe(true); |
| 39 | + expect(isZdrEnabled({ CMD_ZDR: "0" })).toBe(false); |
| 40 | + expect(isZdrEnabled({ CMD_ZDR: "true" })).toBe(false); |
| 41 | + expect(isZdrEnabled({})).toBe(false); |
| 42 | + }); |
| 43 | +}); |
| 44 | + |
| 45 | +describe("applyZdrHeader", () => { |
| 46 | + test("sets the header when opted in", () => { |
| 47 | + const headers = new Headers(); |
| 48 | + applyZdrHeader(headers, { CMD_ZDR: "1" }); |
| 49 | + expect(headers.get(ZDR_HEADER)).toBe("1"); |
| 50 | + }); |
| 51 | + |
| 52 | + test("strips a caller-supplied header when not opted in", () => { |
| 53 | + const headers = new Headers({ "x-cmd-zdr": "1" }); |
| 54 | + applyZdrHeader(headers, {}); |
| 55 | + expect(headers.get(ZDR_HEADER)).toBeNull(); |
| 56 | + }); |
| 57 | + |
| 58 | + test("strips any casing of the header when not opted in", () => { |
| 59 | + const headers = new Headers({ "X-CMD-ZDR": "1" }); |
| 60 | + applyZdrHeader(headers, { CMD_ZDR: "0" }); |
| 61 | + expect(headers.get(ZDR_HEADER)).toBeNull(); |
| 62 | + }); |
| 63 | + |
| 64 | + test("keeps exactly one header value when opted in", () => { |
| 65 | + const headers = new Headers(); |
| 66 | + headers.set("x-cmd-zdr", "0"); |
| 67 | + applyZdrHeader(headers, { CMD_ZDR: "1" }); |
| 68 | + expect(headers.get(ZDR_HEADER)).toBe("1"); |
| 69 | + }); |
| 70 | +}); |
| 71 | + |
| 72 | +describe("plugin ZDR hook", () => { |
| 73 | + const previous = process.env.CMD_ZDR; |
| 74 | + |
| 75 | + afterEach(() => { |
| 76 | + if (previous === undefined) delete process.env.CMD_ZDR; |
| 77 | + else process.env.CMD_ZDR = previous; |
| 78 | + }); |
| 79 | + |
| 80 | + test("emits x-cmd-zdr on Command Code requests when CMD_ZDR=1", async () => { |
| 81 | + process.env.CMD_ZDR = "1"; |
| 82 | + const { ctx, hooks } = makeCtx(); |
| 83 | + await plugin.setup(ctx); |
| 84 | + |
| 85 | + const request = new Request(URL_CC, { method: "POST", body: "{}" }); |
| 86 | + const outgoing = await runHook(hooks, request); |
| 87 | + expect(outgoing.headers.get(ZDR_HEADER)).toBe("1"); |
| 88 | + }); |
| 89 | + |
| 90 | + test("strips a spoofed header when CMD_ZDR is not set", async () => { |
| 91 | + delete process.env.CMD_ZDR; |
| 92 | + const { ctx, hooks } = makeCtx(); |
| 93 | + await plugin.setup(ctx); |
| 94 | + |
| 95 | + const request = new Request(URL_CC, { |
| 96 | + method: "POST", |
| 97 | + body: "{}", |
| 98 | + headers: { "x-cmd-zdr": "1" }, |
| 99 | + }); |
| 100 | + const outgoing = await runHook(hooks, request); |
| 101 | + expect(outgoing.headers.get(ZDR_HEADER)).toBeNull(); |
| 102 | + }); |
| 103 | + |
| 104 | + test("never touches a lookalike host that only shares the path", async () => { |
| 105 | + process.env.CMD_ZDR = "1"; |
| 106 | + const { ctx, hooks } = makeCtx(); |
| 107 | + await plugin.setup(ctx); |
| 108 | + |
| 109 | + const request = new Request("https://api.commandcode.ai.evil.com/provider/v1/chat/completions", { |
| 110 | + method: "POST", |
| 111 | + body: "{}", |
| 112 | + }); |
| 113 | + const outgoing = await runHook(hooks, request); |
| 114 | + expect(outgoing.headers.get(ZDR_HEADER)).toBeNull(); |
| 115 | + }); |
| 116 | +}); |
0 commit comments