pr_01m47d24b0e6n91zwymwxg0vpx/services/models/src/route.test.ts

57 lines2,812 bytesCodeBlame
1import assert from "node:assert/strict";
2import { test } from "node:test";
3
4import type { ModelUpstream } from "@g1t/contracts";
5
6import { presentedToken, upstreamRequest } from "./route.ts";
7
8const run = { workspace: "acme", repo: "acme/web", number: 7, task: "implement", session: "ms_abc", baseUrl: null, apiKey: null, authHeader: null, api: "anthropic" as const, model: null, official: false, provider: "g1t" };
9const hosted = { AI_GATEWAY_ID: "g1t", CLOUDFLARE_ACCOUNT_ID: "acct", AI_GATEWAY_TOKEN: "gw-token" };
10
11function incoming(): Headers {
12 return new Headers({
13 "x-api-key": "g1tm_run_token",
14 "anthropic-version": "2023-06-01",
15 "anthropic-beta": "tools-2024",
16 "content-type": "application/json",
17 });
18}
19
20test("the run's token is read from either header", () => {
21 assert.equal(presentedToken(new Headers({ "x-api-key": "g1tm_a" })), "g1tm_a");
22 assert.equal(presentedToken(new Headers({ authorization: "Bearer g1tm_b" })), "g1tm_b");
23 assert.equal(presentedToken(new Headers()), null);
24});
25
26test("g1t's runs go through its gateway, tagged, without the sandbox's token", () => {
27 const upstream: ModelUpstream = { ...run, route: "g1t" };
28 const { url, headers } = upstreamRequest(upstream, hosted, "/v1/messages?beta=true", incoming());
29 assert.equal(url, "https://gateway.ai.cloudflare.com/v1/acct/g1t/anthropic/v1/messages?beta=true");
30 assert.equal(headers.get("cf-aig-authorization"), "Bearer gw-token");
31 assert.equal(headers.get("x-api-key"), null);
32 assert.equal(headers.get("anthropic-beta"), "tools-2024");
33 assert.deepEqual(JSON.parse(headers.get("cf-aig-metadata") ?? "{}"), {
34 task: "implement",
35 workspace: "acme",
36 repo: "acme/web",
37 pull: 7,
38 session: "ms_abc",
39 });
40});
41
42test("a workspace's own Anthropic key goes to Anthropic, and only there", () => {
43 const upstream: ModelUpstream = { ...run, route: "anthropic", baseUrl: "https://api.anthropic.com", apiKey: "sk-ant-theirs", authHeader: "x-api-key" };
44 const { url, headers } = upstreamRequest(upstream, hosted, "/v1/messages", incoming());
45 assert.equal(url, "https://api.anthropic.com/v1/messages");
46 assert.equal(headers.get("x-api-key"), "sk-ant-theirs");
47 assert.equal(headers.get("cf-aig-authorization"), null);
48 assert.equal(headers.get("cf-aig-metadata"), null);
49});
50
51test("a workspace's own endpoint gets the key the way it asks for it", () => {
52 const upstream: ModelUpstream = { ...run, route: "endpoint", baseUrl: "https://llm.acme.dev/anthropic/", apiKey: "theirs", authHeader: "authorization" };
53 const { url, headers } = upstreamRequest(upstream, hosted, "/v1/messages", incoming());
54 assert.equal(url, "https://llm.acme.dev/anthropic/v1/messages");
55 assert.equal(headers.get("authorization"), "Bearer theirs");
56 assert.equal(headers.get("x-api-key"), null);
57});