Sign in

pulls/pr_01m47d15m3e54sn21z27rpy5n9

public

pr_01m47d15m3e54sn21z27rpy5n9/services/runner/src/model-env.test.ts

99 lines4,662 bytesCodeBlame
1import assert from "node:assert/strict";
2import { createServer } from "node:http";
3import { test } from "node:test";
4
5import { type AgentRoutes, canReachModel, modelEnv } from "./model-env.ts";
6
7const routes: AgentRoutes = {
8 implement: { modelName: "Claude Sonnet 5.5", model: "claude-sonnet-5-5" },
9 review: { modelName: "Claude Opus 5.5", model: "claude-opus-5-5" },
10 update: { modelName: "Claude Sonnet 5.5", model: "claude-sonnet-5-5" },
11 plan: { modelName: "Claude Sonnet 5.5", model: "claude-sonnet-5-5" },
12};
13const tags = { repo: "acme/site", pull: 12 };
14const direct = { ANTHROPIC_API_KEY: "sk-test", AI_GATEWAY_ID: "", CLOUDFLARE_ACCOUNT_ID: "acct" };
15
16/** `ANTHROPIC_CUSTOM_HEADERS` as the harness reads it: one header per line. */
17function customHeaders(vars: Record<string, string>): Record<string, string> {
18 return Object.fromEntries(
19 vars.ANTHROPIC_CUSTOM_HEADERS.split("\n").map((line) => {
20 const at = line.indexOf(": ");
21 return [line.slice(0, at), line.slice(at + 2)];
22 }),
23 );
24}
25
26test("the kind of work decides the model", () => {
27 assert.equal(modelEnv(direct, routes, "implement", tags).ANTHROPIC_MODEL, "claude-sonnet-5-5");
28 const review = modelEnv(direct, routes, "review", tags);
29 assert.equal(review.ANTHROPIC_MODEL, "claude-opus-5-5");
30 assert.equal(review.AGENT_MODEL_NAME, "Claude Opus 5.5");
31});
32
33test("without a gateway, requests go to the provider directly", () => {
34 const vars = modelEnv(direct, routes, "implement", tags);
35 assert.equal(vars.ANTHROPIC_BASE_URL, undefined);
36 assert.equal(vars.ANTHROPIC_CUSTOM_HEADERS, undefined);
37 assert.equal(vars.ANTHROPIC_API_KEY, "sk-test");
38});
39
40test("with a gateway, requests go through it and say what they are for", () => {
41 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t" }, routes, "review", tags);
42 assert.equal(vars.ANTHROPIC_BASE_URL, "https://gateway.ai.cloudflare.com/v1/acct/g1t/anthropic");
43 assert.deepEqual(customHeaders(vars), {
44 "cf-aig-metadata": '{"task":"review","repo":"acme/site","pull":12}',
45 });
46});
47
48test("an authenticated gateway is sent its token", () => {
49 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t", AI_GATEWAY_TOKEN: "tok" }, routes, "implement", tags);
50 assert.equal(customHeaders(vars)["cf-aig-authorization"], "Bearer tok");
51 assert.equal(vars.ANTHROPIC_API_KEY, "sk-test");
52});
53
54test("when the gateway holds the provider's key, the sandbox never gets it", () => {
55 const gatewayOnly = { AI_GATEWAY_ID: "g1t", AI_GATEWAY_TOKEN: "tok", CLOUDFLARE_ACCOUNT_ID: "acct" };
56 assert.equal(canReachModel(gatewayOnly), true);
57 assert.equal(modelEnv(gatewayOnly, routes, "implement", tags).ANTHROPIC_API_KEY, "tok");
58 assert.equal(canReachModel({ AI_GATEWAY_ID: "g1t", CLOUDFLARE_ACCOUNT_ID: "acct" }), false);
59 assert.equal(canReachModel(direct), true);
60});
61
62// A stand-in for the gateway: what a sandbox's agent sends, given these
63// variables, arrives at the gateway's path with both credentials.
64test("a request built from these variables reaches a gateway as expected", async () => {
65 const seen: { url?: string; key?: string; gateway?: string; metadata?: string; model?: string } = {};
66 const server = createServer((request, response) => {
67 let body = "";
68 request.on("data", (chunk) => (body += chunk));
69 request.on("end", () => {
70 seen.url = request.url;
71 seen.key = request.headers["x-api-key"] as string;
72 seen.gateway = request.headers["cf-aig-authorization"] as string;
73 seen.metadata = request.headers["cf-aig-metadata"] as string;
74 seen.model = JSON.parse(body).model;
75 response.setHeader("content-type", "application/json");
76 response.end(JSON.stringify({ type: "message", content: [] }));
77 });
78 });
79 await new Promise<void>((resolve) => server.listen(0, resolve));
80 const { port } = server.address() as { port: number };
81
82 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t", AI_GATEWAY_TOKEN: "tok" }, routes, "implement", tags);
83 // Same path as the real gateway, on the stand-in's address.
84 const base = vars.ANTHROPIC_BASE_URL.replace("https://gateway.ai.cloudflare.com", `http://localhost:${port}`);
85 await fetch(`${base}/v1/messages`, {
86 method: "POST",
87 headers: { "x-api-key": vars.ANTHROPIC_API_KEY, ...customHeaders(vars), "content-type": "application/json" },
88 body: JSON.stringify({ model: vars.ANTHROPIC_MODEL, max_tokens: 1, messages: [] }),
89 });
90 server.close();
91
92 assert.deepEqual(seen, {
93 url: "/v1/acct/g1t/anthropic/v1/messages",
94 key: "sk-test",
95 gateway: "Bearer tok",
96 metadata: '{"task":"implement","repo":"acme/site","pull":12}',
97 model: "claude-sonnet-5-5",
98 });
99});