Skip to content
Sign in

flagon-io/g1t

public
0

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

193 lines9,312 bytesCodeBlame
1import assert from "node:assert/strict";
2import { createServer } from "node:http";
3import { test } from "node:test";
4
5import {
6 type AgentRouting,
7 type ChangeSize,
8 DEFAULT_ROUTING,
9 canReachModel,
10 changeSize,
11 chooseTier,
12 gatewaySession,
13 lastAttemptFailed,
14 modelEnv,
15 parseRouting,
16} from "./model-env.ts";
17
18const routes: AgentRouting = {
19 ...DEFAULT_ROUTING,
20 tiers: {
21 small: { modelName: "Claude Haiku 4.5", model: "claude-haiku-4-5-20251001" },
22 large: { modelName: "Claude Sonnet 5.5", model: "claude-sonnet-5-5" },
23 },
24};
25const tags = { repo: "acme/site", pull: 12 };
26const direct = { ANTHROPIC_API_KEY: "sk-test", AI_GATEWAY_ID: "", CLOUDFLARE_ACCOUNT_ID: "acct" };
27
28/** `ANTHROPIC_CUSTOM_HEADERS` as the harness reads it: one header per line. */
29function customHeaders(vars: Record<string, string>): Record<string, string> {
30 return Object.fromEntries(
31 vars.ANTHROPIC_CUSTOM_HEADERS.split("\n").map((line) => {
32 const at = line.indexOf(": ");
33 return [line.slice(0, at), line.slice(at + 2)];
34 }),
35 );
36}
37
38const small: ChangeSize = { files: 3, lines: 80, sensitive: [] };
39
40test("planning and catching up run on the small tier, making a change on the large", () => {
41 assert.equal(chooseTier("plan", {}, routes), "small");
42 assert.equal(chooseTier("update", {}, routes), "small");
43 assert.equal(chooseTier("implement", {}, routes), "large");
44 assert.equal(chooseTier("implement", { change: small }, routes), "large");
45});
46
47test("a review is small only for a small change that touches nothing sensitive", () => {
48 assert.equal(chooseTier("review", { change: small }, routes), "small");
49 assert.equal(chooseTier("review", { change: { ...small, lines: 200, files: 10 } }, routes), "small");
50 assert.equal(chooseTier("review", { change: { ...small, lines: 201 } }, routes), "large");
51 assert.equal(chooseTier("review", { change: { ...small, files: 11 } }, routes), "large");
52 assert.equal(chooseTier("review", { change: { ...small, sensitive: ["CI workflows"] } }, routes), "large");
53 assert.equal(chooseTier("review", { change: small, labels: ["Security"] }, routes), "large");
54 assert.equal(chooseTier("review", { change: small, labels: ["docs"] }, routes), "small");
55});
56
57test("a review of a change g1t cannot size runs on the large tier", () => {
58 assert.equal(chooseTier("review", {}, routes), "large");
59 assert.equal(chooseTier("review", { change: null }, routes), "large");
60 assert.equal(chooseTier("review", { change: { files: 0, lines: 0, sensitive: [] } }, routes), "large");
61});
62
63test("a retry after a failed attempt goes up to the large tier", () => {
64 assert.equal(chooseTier("plan", { retry: true }, routes), "large");
65 assert.equal(chooseTier("update", { retry: true }, routes), "large");
66 assert.equal(chooseTier("review", { change: small, retry: true }, routes), "large");
67});
68
69test("a retry is the same work again after its latest attempt failed", () => {
70 assert.equal(lastAttemptFailed([]), false);
71 assert.equal(lastAttemptFailed([{ status: "failed" }, { status: "succeeded" }]), true);
72 assert.equal(lastAttemptFailed([{ status: "succeeded" }, { status: "failed" }]), false);
73 assert.equal(lastAttemptFailed([{ status: "stopped", halted: "budget" }]), true);
74 assert.equal(lastAttemptFailed([{ status: "stopped", halted: null }]), false);
75 assert.equal(lastAttemptFailed([{ status: "running" }]), false);
76 assert.equal(lastAttemptFailed([{ status: "failed", title: "Add search" }], "Add search "), true);
77 assert.equal(lastAttemptFailed([{ status: "failed", title: "Add search" }], "Add billing"), false);
78});
79
80test("the configuration decides the tiers, the rules and the limits", () => {
81 const parsed = parseRouting(
82 JSON.stringify({
83 tiers: { small: { modelName: "Small", model: "small-1" } },
84 tasks: { update: "large" },
85 smallChange: { lines: 50 },
86 }),
87 );
88 assert.deepEqual(parsed.tiers.small, { modelName: "Small", model: "small-1" });
89 assert.deepEqual(parsed.tiers.large, DEFAULT_ROUTING.tiers.large);
90 assert.equal(chooseTier("update", {}, parsed), "large");
91 assert.equal(chooseTier("plan", {}, parsed), "small");
92 assert.equal(chooseTier("review", { change: { ...small, lines: 51 } }, parsed), "large");
93 assert.equal(chooseTier("review", { change: { ...small, files: 10, lines: 50 } }, parsed), "small");
94 assert.deepEqual(parseRouting(undefined), DEFAULT_ROUTING);
95 assert.deepEqual(parseRouting("not json"), DEFAULT_ROUTING);
96});
97
98test("a change's size is its files and the lines added and removed", () => {
99 assert.deepEqual(
100 changeSize([{ additions: 10, deletions: 2 }, { additions: 0, deletions: 5 }], ["secrets"]),
101 { files: 2, lines: 17, sensitive: ["secrets"] },
102 );
103});
104
105test("the tier decides the model, and the harness's small tasks use the small tier", () => {
106 const large = modelEnv(direct, routes, "implement", "large", tags);
107 assert.equal(large.ANTHROPIC_MODEL, "claude-sonnet-5-5");
108 assert.equal(large.AGENT_MODEL_NAME, "Claude Sonnet 5.5");
109 assert.equal(large.ANTHROPIC_SMALL_FAST_MODEL, "claude-haiku-4-5-20251001");
110 assert.equal(large.ANTHROPIC_DEFAULT_HAIKU_MODEL, "claude-haiku-4-5-20251001");
111 const review = modelEnv(direct, routes, "review", "small", tags);
112 assert.equal(review.ANTHROPIC_MODEL, "claude-haiku-4-5-20251001");
113 assert.equal(review.AGENT_MODEL_NAME, "Claude Haiku 4.5");
114});
115
116test("without a gateway, requests go to the provider directly", () => {
117 const vars = modelEnv(direct, routes, "implement", "large", tags);
118 assert.equal(vars.ANTHROPIC_BASE_URL, undefined);
119 assert.equal(vars.ANTHROPIC_CUSTOM_HEADERS, undefined);
120 assert.equal(vars.ANTHROPIC_API_KEY, "sk-test");
121});
122
123test("with a gateway, requests go through it and say what they are for", () => {
124 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t" }, routes, "review", "large", tags);
125 assert.equal(vars.ANTHROPIC_BASE_URL, "https://gateway.ai.cloudflare.com/v1/acct/g1t/anthropic");
126 assert.deepEqual(customHeaders(vars), {
127 "cf-aig-metadata": '{"task":"review","tier":"large","repo":"acme/site","pull":12}',
128 });
129});
130
131test("a run straight to the gateway carries its session, so billing can settle it", () => {
132 const session = gatewaySession();
133 assert.match(session, /^rs_[0-9a-f]{24}$/);
134 assert.notEqual(gatewaySession(), session);
135 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t" }, routes, "implement", "small", { ...tags, session });
136 const metadata = JSON.parse(customHeaders(vars)["cf-aig-metadata"]);
137 assert.equal(metadata.session, session);
138 // The gateway keeps at most five metadata entries.
139 assert.ok(Object.keys(metadata).length <= 5);
140});
141
142test("an authenticated gateway is sent its token", () => {
143 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t", AI_GATEWAY_TOKEN: "tok" }, routes, "implement", "large", tags);
144 assert.equal(customHeaders(vars)["cf-aig-authorization"], "Bearer tok");
145 assert.equal(vars.ANTHROPIC_API_KEY, "sk-test");
146});
147
148test("when the gateway holds the provider's key, the sandbox never gets it", () => {
149 const gatewayOnly = { AI_GATEWAY_ID: "g1t", AI_GATEWAY_TOKEN: "tok", CLOUDFLARE_ACCOUNT_ID: "acct" };
150 assert.equal(canReachModel(gatewayOnly), true);
151 assert.equal(modelEnv(gatewayOnly, routes, "implement", "large", tags).ANTHROPIC_API_KEY, "tok");
152 assert.equal(canReachModel({ AI_GATEWAY_ID: "g1t", CLOUDFLARE_ACCOUNT_ID: "acct" }), false);
153 assert.equal(canReachModel(direct), true);
154});
155
156// A stand-in for the gateway: what a sandbox's agent sends, given these
157// variables, arrives at the gateway's path with both credentials.
158test("a request built from these variables reaches a gateway as expected", async () => {
159 const seen: { url?: string; key?: string; gateway?: string; metadata?: string; model?: string } = {};
160 const server = createServer((request, response) => {
161 let body = "";
162 request.on("data", (chunk) => (body += chunk));
163 request.on("end", () => {
164 seen.url = request.url;
165 seen.key = request.headers["x-api-key"] as string;
166 seen.gateway = request.headers["cf-aig-authorization"] as string;
167 seen.metadata = request.headers["cf-aig-metadata"] as string;
168 seen.model = JSON.parse(body).model;
169 response.setHeader("content-type", "application/json");
170 response.end(JSON.stringify({ type: "message", content: [] }));
171 });
172 });
173 await new Promise<void>((resolve) => server.listen(0, resolve));
174 const { port } = server.address() as { port: number };
175
176 const vars = modelEnv({ ...direct, AI_GATEWAY_ID: "g1t", AI_GATEWAY_TOKEN: "tok" }, routes, "implement", "large", tags);
177 // Same path as the real gateway, on the stand-in's address.
178 const base = vars.ANTHROPIC_BASE_URL.replace("https://gateway.ai.cloudflare.com", `http://localhost:${port}`);
179 await fetch(`${base}/v1/messages`, {
180 method: "POST",
181 headers: { "x-api-key": vars.ANTHROPIC_API_KEY, ...customHeaders(vars), "content-type": "application/json" },
182 body: JSON.stringify({ model: vars.ANTHROPIC_MODEL, max_tokens: 1, messages: [] }),
183 });
184 server.close();
185
186 assert.deepEqual(seen, {
187 url: "/v1/acct/g1t/anthropic/v1/messages",
188 key: "sk-test",
189 gateway: "Bearer tok",
190 metadata: '{"task":"implement","tier":"large","repo":"acme/site","pull":12}',
191 model: "claude-sonnet-5-5",
192 });
193});