| 1 | import assert from "node:assert/strict"; |
| 2 | import { test } from "node:test"; |
| 3 | |
| 4 | import type { ModelUpstream } from "@g1t/contracts"; |
| 5 | |
| 6 | import { presentedToken, upstreamRequest } from "./route.ts"; |
| 7 | |
| 8 | const 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" }; |
| 9 | const hosted = { AI_GATEWAY_ID: "g1t", CLOUDFLARE_ACCOUNT_ID: "acct", AI_GATEWAY_TOKEN: "gw-token" }; |
| 10 | |
| 11 | function 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 | |
| 20 | test("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 | |
| 26 | test("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 | |
| 42 | test("a run's tier is tagged at the gateway, so spend can be read per tier", () => { |
| 43 | const upstream: ModelUpstream = { ...run, route: "g1t", tier: "small" }; |
| 44 | const { headers } = upstreamRequest(upstream, hosted, "/v1/messages", incoming()); |
| 45 | const metadata = JSON.parse(headers.get("cf-aig-metadata") ?? "{}"); |
| 46 | assert.deepEqual(metadata, { |
| 47 | task: "implement", |
| 48 | tier: "small", |
| 49 | repo: "acme/web", |
| 50 | pull: 7, |
| 51 | session: "ms_abc", |
| 52 | }); |
| 53 | // The gateway keeps five entries; the session, which billing settles |
| 54 | // by, must be one of them. |
| 55 | assert.ok(Object.keys(metadata).length <= 5); |
| 56 | }); |
| 57 | |
| 58 | test("a workspace's own Anthropic key goes to Anthropic, and only there", () => { |
| 59 | const upstream: ModelUpstream = { ...run, route: "anthropic", baseUrl: "https://api.anthropic.com", apiKey: "sk-ant-theirs", authHeader: "x-api-key" }; |
| 60 | const { url, headers } = upstreamRequest(upstream, hosted, "/v1/messages", incoming()); |
| 61 | assert.equal(url, "https://api.anthropic.com/v1/messages"); |
| 62 | assert.equal(headers.get("x-api-key"), "sk-ant-theirs"); |
| 63 | assert.equal(headers.get("cf-aig-authorization"), null); |
| 64 | assert.equal(headers.get("cf-aig-metadata"), null); |
| 65 | }); |
| 66 | |
| 67 | test("a workspace's own endpoint gets the key the way it asks for it", () => { |
| 68 | const upstream: ModelUpstream = { ...run, route: "endpoint", baseUrl: "https://llm.acme.dev/anthropic/", apiKey: "theirs", authHeader: "authorization" }; |
| 69 | const { url, headers } = upstreamRequest(upstream, hosted, "/v1/messages", incoming()); |
| 70 | assert.equal(url, "https://llm.acme.dev/anthropic/v1/messages"); |
| 71 | assert.equal(headers.get("authorization"), "Bearer theirs"); |
| 72 | assert.equal(headers.get("x-api-key"), null); |
| 73 | }); |