| 1 | /** |
| 2 | * The model proxy: every model request a g1t sandbox makes comes through |
| 3 | * here, at `https://models.g1t.sh/anthropic`. |
| 4 | * |
| 5 | * A sandbox holds a token for its one run, never a key. The proxy looks the |
| 6 | * token up and forwards the request with the credentials for that run: |
| 7 | * g1t's AI Gateway when g1t pays, or one of the workspace's own providers |
| 8 | * when it does. A provider that speaks OpenAI's API gets the request |
| 9 | * translated, and its answer translated back. So a sandbox that is tricked |
| 10 | * into printing its environment gives away a token that stops working when |
| 11 | * the run ends (the runner closes its session then, and lookups are kept |
| 12 | * only seconds), and nothing of the workspace's. |
| 13 | * |
| 14 | * Responses stream through. What each answer used is read from a copy as |
| 15 | * it passes and reported to billing afterwards, counted per run for usage |
| 16 | * views. |
| 17 | * |
| 18 | * The same address is the AI Gateway for a workspace's own code: a request |
| 19 | * with one of the workspace's access tokens (`g1t_…`) instead of a run's |
| 20 | * goes to `gateway.ts`, and is logged and charged to the workspace. |
| 21 | */ |
| 22 | import { |
| 23 | type GatewayModel, |
| 24 | type ModelUpstream, |
| 25 | type ServiceBinding, |
| 26 | type User, |
| 27 | billingClient, |
| 28 | identityClient, |
| 29 | integrationsClient, |
| 30 | } from "@g1t/contracts"; |
| 31 | |
| 32 | import { type AnthropicRequest, StreamTranslator, errorFromChat, estimateTokens, fromChat, toChat } from "./openai"; |
| 33 | import { isAnswer, tokenReport } from "./report"; |
| 34 | import { |
| 35 | type Caller, |
| 36 | anthropicError, |
| 37 | callerOf, |
| 38 | errorMessage, |
| 39 | gatewayRecord, |
| 40 | gatewayRoute, |
| 41 | hostedRequest, |
| 42 | requestId, |
| 43 | sessionOf, |
| 44 | unoffered, |
| 45 | unpriced, |
| 46 | } from "./gateway"; |
| 47 | import { type HostedRouting, presentedToken, upstreamRequest } from "./route"; |
| 48 | import { NO_TOKENS, type Tokens, measure } from "./usage"; |
| 49 | |
| 50 | interface Env extends HostedRouting { |
| 51 | INTEGRATIONS: ServiceBinding; |
| 52 | BILLING: ServiceBinding; |
| 53 | IDENTITY: ServiceBinding; |
| 54 | } |
| 55 | |
| 56 | /** |
| 57 | * How long a looked-up token is trusted before it is looked up again. Short, |
| 58 | * because a run's token is closed the moment the run ends (and a connection |
| 59 | * may be removed mid-run): the proxy refuses it again within this long. |
| 60 | */ |
| 61 | const REMEMBER_MS = 10_000; |
| 62 | const remembered = new Map<string, { upstream: ModelUpstream | null; until: number }>(); |
| 63 | |
| 64 | async function lookUp(env: Env, token: string): Promise<ModelUpstream | null> { |
| 65 | const now = Date.now(); |
| 66 | const hit = remembered.get(token); |
| 67 | if (hit && hit.until > now) return hit.upstream; |
| 68 | const upstream = await integrationsClient(env.INTEGRATIONS).modelUpstream(token); |
| 69 | if (remembered.size > 5_000) remembered.clear(); |
| 70 | remembered.set(token, { upstream, until: now + REMEMBER_MS }); |
| 71 | return upstream; |
| 72 | } |
| 73 | |
| 74 | /** An error in the shape Anthropic's API uses, which the harness understands. */ |
| 75 | function refuse(status: number, message: string): Response { |
| 76 | return Response.json( |
| 77 | { type: "error", error: { type: status === 401 ? "authentication_error" : "not_found_error", message } }, |
| 78 | { status }, |
| 79 | ); |
| 80 | } |
| 81 | |
| 82 | /** |
| 83 | * Passes an answer through and, once it has all gone by, tells billing what |
| 84 | * it used. Reporting happens after the answer, and a report that fails is |
| 85 | * dropped: the answer never waits on it or breaks for it. |
| 86 | */ |
| 87 | function counted(answer: Response, upstream: ModelUpstream, env: Env, ctx: ExecutionContext): Response { |
| 88 | const { response, tokens, model } = measure(answer); |
| 89 | ctx.waitUntil( |
| 90 | (async () => { |
| 91 | const report = tokenReport(upstream, await model, await tokens); |
| 92 | if (report) await billingClient(env.BILLING).recordTokens(report); |
| 93 | })().catch(() => undefined), |
| 94 | ); |
| 95 | return response; |
| 96 | } |
| 97 | |
| 98 | /** Sends an Anthropic request to a provider that speaks OpenAI's API. */ |
| 99 | async function viaChat(upstream: ModelUpstream, path: string, request: Request): Promise<Response> { |
| 100 | const body = (await request.json()) as AnthropicRequest; |
| 101 | const model = upstream.model ?? body.model ?? ""; |
| 102 | if (path.startsWith("/v1/messages/count_tokens")) { |
| 103 | return Response.json({ input_tokens: estimateTokens(body) }); |
| 104 | } |
| 105 | if (!path.startsWith("/v1/messages")) return refuse(404, `${path} has no counterpart at this provider.`); |
| 106 | |
| 107 | const headers = new Headers({ "content-type": "application/json" }); |
| 108 | if (upstream.gatewayToken) headers.set("cf-aig-authorization", `Bearer ${upstream.gatewayToken}`); |
| 109 | if (upstream.apiKey) { |
| 110 | // `authorization` means a bearer token; any other header takes the key as it is. |
| 111 | const header = upstream.authHeader ?? "authorization"; |
| 112 | headers.set(header, header === "authorization" ? `Bearer ${upstream.apiKey}` : upstream.apiKey); |
| 113 | } |
| 114 | const answer = await fetch(`${(upstream.baseUrl ?? "").replace(/\/+$/, "")}/chat/completions`, { |
| 115 | method: "POST", |
| 116 | headers, |
| 117 | body: JSON.stringify(toChat(body, model, { official: upstream.official, provider: upstream.provider })), |
| 118 | }); |
| 119 | if (!answer.ok) { |
| 120 | return Response.json(errorFromChat(answer.status, await answer.text()), { status: answer.status }); |
| 121 | } |
| 122 | if (!body.stream) return Response.json(fromChat((await answer.json()) as Record<string, unknown>, model)); |
| 123 | |
| 124 | const translator = new StreamTranslator(model); |
| 125 | const decoder = new TextDecoder(); |
| 126 | const encoder = new TextEncoder(); |
| 127 | const translated = answer.body!.pipeThrough( |
| 128 | new TransformStream<Uint8Array, Uint8Array>({ |
| 129 | transform(chunk, controller) { |
| 130 | const out = translator.push(decoder.decode(chunk, { stream: true })); |
| 131 | if (out) controller.enqueue(encoder.encode(out)); |
| 132 | }, |
| 133 | flush(controller) { |
| 134 | const out = translator.push(decoder.decode()) + translator.finish(); |
| 135 | if (out) controller.enqueue(encoder.encode(out)); |
| 136 | }, |
| 137 | }), |
| 138 | ); |
| 139 | return new Response(translated, { |
| 140 | headers: { "content-type": "text/event-stream", "cache-control": "no-cache" }, |
| 141 | }); |
| 142 | } |
| 143 | |
| 144 | // --- The AI Gateway ------------------------------------------------------------ |
| 145 | |
| 146 | /** |
| 147 | * What gateway requests look up, remembered as briefly as a run's token |
| 148 | * is: a deleted token, a key added under Integrations or credit just bought |
| 149 | * takes effect within `REMEMBER_MS`. The catalogue changes rarely. |
| 150 | */ |
| 151 | const CATALOGUE_MS = 5 * 60_000; |
| 152 | const callers = new Map<string, { value: User | null; until: number }>(); |
| 153 | const ownKeys = new Map<string, { value: ModelUpstream | null; until: number }>(); |
| 154 | const admitted = new Map<string, { value: string | null; until: number }>(); |
| 155 | let catalogue: { models: GatewayModel[]; until: number } | null = null; |
| 156 | |
| 157 | async function cached<T>(cache: Map<string, { value: T; until: number }>, key: string, read: () => Promise<T>): Promise<T> { |
| 158 | const now = Date.now(); |
| 159 | const hit = cache.get(key); |
| 160 | if (hit && hit.until > now) return hit.value; |
| 161 | const value = await read(); |
| 162 | if (cache.size > 5_000) cache.clear(); |
| 163 | cache.set(key, { value, until: now + REMEMBER_MS }); |
| 164 | return value; |
| 165 | } |
| 166 | |
| 167 | async function offered(env: Env): Promise<GatewayModel[]> { |
| 168 | if (catalogue && catalogue.until > Date.now()) return catalogue.models; |
| 169 | const models = await billingClient(env.BILLING).gatewayModels(); |
| 170 | catalogue = { models, until: Date.now() + CATALOGUE_MS }; |
| 171 | return models; |
| 172 | } |
| 173 | |
| 174 | /** A workspace's own request, sent with one of its access tokens. */ |
| 175 | async function gateway(request: Request, env: Env, ctx: ExecutionContext, token: string, path: string): Promise<Response> { |
| 176 | const started = Date.now(); |
| 177 | const who = callerOf(await cached(callers, token, () => identityClient(env.IDENTITY).userForAccessToken(token))); |
| 178 | if (!("caller" in who)) return anthropicError(who.status, who.type, who.message); |
| 179 | const caller: Caller = who.caller; |
| 180 | const route = gatewayRoute(path); |
| 181 | if (!route || request.method !== "POST") { |
| 182 | return anthropicError(404, "not_found_error", "The AI Gateway answers POST /v1/messages and POST /v1/messages/count_tokens."); |
| 183 | } |
| 184 | let body: Record<string, unknown>; |
| 185 | try { |
| 186 | body = (await request.json()) as Record<string, unknown>; |
| 187 | } catch { |
| 188 | return anthropicError(400, "invalid_request_error", "The request body is not JSON."); |
| 189 | } |
| 190 | const model = typeof body.model === "string" ? body.model : ""; |
| 191 | const streamed = body.stream === true; |
| 192 | const id = requestId(); |
| 193 | // Every request is logged once it is known whose it is. Counting tokens |
| 194 | // is a question about a request, not one, and is not. |
| 195 | const log = (status: number, ownKey: boolean, tokens: Tokens = NO_TOKENS, answeredBy: string | null = null, error: string | null = null) => { |
| 196 | if (route !== "messages") return Promise.resolve(); |
| 197 | const record = gatewayRecord({ id, caller, model: answeredBy ?? model, tokens, status, ownKey, streamed, durationMs: Date.now() - started, error }); |
| 198 | return billingClient(env.BILLING) |
| 199 | .recordGateway(record) |
| 200 | .then(() => undefined) |
| 201 | .catch(() => undefined); |
| 202 | }; |
| 203 | |
| 204 | const own = await cached(ownKeys, caller.workspace, () => integrationsClient(env.INTEGRATIONS).gatewayUpstream(caller.workspace)); |
| 205 | let target: { url: string; headers: Headers }; |
| 206 | if (own) { |
| 207 | // The workspace's own key: nothing to admit or charge. |
| 208 | target = upstreamRequest(own, env, path, request.headers); |
| 209 | } else { |
| 210 | const why = unoffered(body.model, await offered(env)) ?? unpriced(body); |
| 211 | if (why) { |
| 212 | ctx.waitUntil(log(400, false, NO_TOKENS, null, why)); |
| 213 | return anthropicError(400, "invalid_request_error", why); |
| 214 | } |
| 215 | const refusal = await cached(admitted, caller.workspace, async () => { |
| 216 | const answer = await billingClient(env.BILLING).gatewayAdmit(caller.workspace); |
| 217 | return answer.ok ? null : answer.error.message; |
| 218 | }); |
| 219 | if (refusal) { |
| 220 | ctx.waitUntil(log(402, false, NO_TOKENS, null, refusal)); |
| 221 | return anthropicError(402, "billing_error", refusal); |
| 222 | } |
| 223 | target = hostedRequest(env, path, request.headers, caller, sessionOf(caller.tokenId, new Date())); |
| 224 | } |
| 225 | target.headers.delete("content-length"); |
| 226 | const answer = await fetch(target.url, { method: "POST", headers: target.headers, body: JSON.stringify(body) }); |
| 227 | if (route === "count_tokens") return answer; |
| 228 | if (!answer.ok) { |
| 229 | const text = await answer.clone().text(); |
| 230 | ctx.waitUntil(log(answer.status, own != null, NO_TOKENS, null, errorMessage(answer.status, text))); |
| 231 | return answer; |
| 232 | } |
| 233 | const { response, tokens, model: answered } = measure(answer); |
| 234 | ctx.waitUntil( |
| 235 | (async () => { |
| 236 | const used = await tokens; |
| 237 | const by = await answered; |
| 238 | // On g1t's key the model asked for is the one priced, whatever dated |
| 239 | // name the provider answers with; on the workspace's own, the one |
| 240 | // that answered. |
| 241 | await log(answer.status, own != null, used, own ? by : null); |
| 242 | })().catch(() => undefined), |
| 243 | ); |
| 244 | return response; |
| 245 | } |
| 246 | |
| 247 | export default { |
| 248 | async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> { |
| 249 | const url = new URL(request.url); |
| 250 | if (url.pathname === "/" || url.pathname === "") { |
| 251 | return new Response("g1t's model proxy and AI Gateway. See https://docs.g1t.sh/guides/ai-gateway/\n"); |
| 252 | } |
| 253 | if (!url.pathname.startsWith("/anthropic/")) return refuse(404, "Requests go to /anthropic/v1/…."); |
| 254 | const token = presentedToken(request.headers); |
| 255 | // A workspace's own access token: the AI Gateway. |
| 256 | if (token?.startsWith("g1t_")) return gateway(request, env, ctx, token, url.pathname.slice("/anthropic".length) + url.search); |
| 257 | if (!token?.startsWith("g1tm_")) return refuse(401, "This needs a g1t run's model token, or a workspace's access token for the AI Gateway."); |
| 258 | const upstream = await lookUp(env, token); |
| 259 | if (!upstream) return refuse(401, "This run's model token has expired, or its model connection was removed."); |
| 260 | |
| 261 | const path = url.pathname.slice("/anthropic".length) + url.search; |
| 262 | // Both routes answer in Anthropic's shape, so one reading counts either. |
| 263 | const answer = (response: Response) => (isAnswer(url.pathname.slice("/anthropic".length)) ? counted(response, upstream, env, ctx) : response); |
| 264 | if (upstream.api === "openai") return answer(await viaChat(upstream, path, request)); |
| 265 | |
| 266 | const { url: target, headers } = upstreamRequest(upstream, env, path, request.headers); |
| 267 | // A route that names a model gets it for every request of the run, |
| 268 | // including the harness's small background ones. |
| 269 | let body: BodyInit | null = request.method === "GET" || request.method === "HEAD" ? null : request.body; |
| 270 | if (upstream.model && body && path.startsWith("/v1/messages")) { |
| 271 | const parsed = (await request.json()) as Record<string, unknown>; |
| 272 | body = JSON.stringify({ ...parsed, model: upstream.model }); |
| 273 | headers.delete("content-length"); |
| 274 | } |
| 275 | return answer(await fetch(target, { method: request.method, headers, body })); |
| 276 | }, |
| 277 | } satisfies ExportedHandler<Env>; |