| 1 | import type { AppContext } from "../auth/principal.js"; |
| 2 | import { Directory } from "../lib/directory.js"; |
| 3 | import { taskPersonAllowed } from "./access.js"; |
| 4 | import { isProvider, ModelKeys, PROVIDERS, type Provider } from "./keys.js"; |
| 5 | import { bookSpend, proxyPayer, settleSpend, spendMonth } from "./org-policy.js"; |
| 6 | import { freeSlot, MAX_IN_FLIGHT, takeSlot } from "./inflight.js"; |
| 7 | import { overLimit } from "../lib/rate-limit.js"; |
| 8 | import { meterUsage, NO_USAGE, prepareRequest, readBody, unreadCharge, weighted } from "./usage.js"; |
| 9 | |
| 10 | /** The largest model request the proxy forwards (it reads each one to charge it). */ |
| 11 | const MAX_REQUEST = 32 * 1024 * 1024; |
| 12 | |
| 13 | /** |
| 14 | * The model proxy: `/-/model/<provider>/v1/...`. A task's runtime calls it |
| 15 | * with the run's token (`dmt_...`, as its API key) and the proxy forwards the |
| 16 | * request to the provider with the key of whoever pays (the person's own, or |
| 17 | * under its policy an organization's: org-policy.ts), decrypted here for this |
| 18 | * request only. So no key ever enters a sandbox, container or terminal, and the |
| 19 | * token stops working when the run ends. |
| 20 | * |
| 21 | * It runs before the app's identity middleware: its token isn't an account |
| 22 | * credential and works nowhere else. |
| 23 | */ |
| 24 | |
| 25 | const UPSTREAM: Record<Provider, string> = { anthropic: "https://api.anthropic.com", openai: "https://api.openai.com", openrouter: "https://openrouter.ai/api" }; |
| 26 | |
| 27 | /** |
| 28 | * Only the model endpoints the runtimes use. The person's key could reach the |
| 29 | * rest of their provider account (stored files, batches, assistants), which an |
| 30 | * agent misled by repository content must not. |
| 31 | */ |
| 32 | const ALLOWED: Record<Provider, Array<[method: string, path: RegExp]>> = { |
| 33 | anthropic: [ |
| 34 | ["POST", /^\/v1\/messages$/], |
| 35 | ["POST", /^\/v1\/messages\/count_tokens$/], |
| 36 | ["GET", /^\/v1\/models(\/[A-Za-z0-9._:-]+)?$/], |
| 37 | ], |
| 38 | openai: [ |
| 39 | ["POST", /^\/v1\/responses$/], |
| 40 | // Codex compacts long conversations, and reads a background response back by id. |
| 41 | ["POST", /^\/v1\/responses\/compact$/], |
| 42 | ["GET", /^\/v1\/responses\/[A-Za-z0-9_-]+$/], |
| 43 | ["POST", /^\/v1\/chat\/completions$/], |
| 44 | ["GET", /^\/v1\/models(\/[A-Za-z0-9._:-]+)?$/], |
| 45 | ], |
| 46 | // OpenRouter: its Anthropic-compatible Messages endpoint (Claude Code), the OpenAI-compatible ones (Codex, |
| 47 | // custom commands) and the model list. Never keys, credits, generation stats or anything else on the account. |
| 48 | openrouter: [ |
| 49 | ["POST", /^\/v1\/messages$/], |
| 50 | ["POST", /^\/v1\/messages\/count_tokens$/], |
| 51 | ["POST", /^\/v1\/responses$/], |
| 52 | ["POST", /^\/v1\/chat\/completions$/], |
| 53 | ["POST", /^\/v1\/completions$/], |
| 54 | ["GET", /^\/v1\/models$/], |
| 55 | ], |
| 56 | }; |
| 57 | |
| 58 | export const proxyAllows = (provider: Provider, method: string, path: string): boolean => ALLOWED[provider].some(([m, re]) => m === method && re.test(path)); |
| 59 | |
| 60 | /** Request headers that never go upstream: credentials, cookies and the hop's own. */ |
| 61 | const DROP = new Set(["authorization", "x-api-key", "cookie", "host", "content-length", "connection", "cf-connecting-ip", "cf-ipcountry", "cf-ray", "cf-visitor", "x-forwarded-for", "x-forwarded-proto", "x-real-ip"]); |
| 62 | |
| 63 | /** OpenRouter's app attribution and request options: only Diffinery sets these, never the sandbox. */ |
| 64 | const ATTRIBUTION = /^(x-openrouter-.*|x-title|http-referer|referer)$/i; |
| 65 | |
| 66 | /** Headers that pick an organization or project within a provider account (OpenAI's). */ |
| 67 | const ACCOUNT_SCOPE = /^openai-(organization|project)$/i; |
| 68 | |
| 69 | const err = (status: number, message: string) => |
| 70 | new Response(JSON.stringify({ type: "error", error: { type: status === 401 ? "authentication_error" : status === 404 ? "not_found_error" : status === 429 ? "rate_limit_error" : "permission_error", message } }), { status, headers: { "content-type": "application/json" } }); |
| 71 | |
| 72 | export function proxyToken(req: Request): string | null { |
| 73 | const key = req.headers.get("x-api-key"); |
| 74 | if (key?.startsWith("dmt_")) return key.trim(); |
| 75 | const m = /^Bearer\s+(dmt_\S+)$/i.exec(req.headers.get("authorization") ?? ""); |
| 76 | return m ? m[1]! : null; |
| 77 | } |
| 78 | |
| 79 | export async function modelProxy(c: AppContext): Promise<Response> { |
| 80 | const provider = c.req.param("provider") ?? ""; |
| 81 | if (!isProvider(provider)) return err(404, "unknown model provider"); |
| 82 | const url = new URL(c.req.url); |
| 83 | const rest = url.pathname.slice(`/-/model/${provider}`.length); |
| 84 | if (!proxyAllows(provider, c.req.method, rest)) return err(403, `${c.req.method} ${rest} isn't available to agent tasks: only the ${PROVIDERS[provider].name} model endpoints are`); |
| 85 | const token = proxyToken(c.req.raw); |
| 86 | if (!token) return err(401, "a task's model token is required"); |
| 87 | const repoId = token.split("_")[1]; |
| 88 | if (!repoId || !/^[A-Za-z0-9-]{8,64}$/.test(repoId)) return err(401, "malformed model token"); |
| 89 | const env = c.env; |
| 90 | // Check the token before reading a body of up to 32 MB, and only wake a project that exists. Wrong |
| 91 | // tokens count against the sign-in limit per address. |
| 92 | const known = await env.DB.prepare("SELECT 1 FROM repos WHERE id = ?").bind(repoId).first(); |
| 93 | const project = known ? env.PROJECT.get(env.PROJECT.idFromName(repoId)) : null; |
| 94 | if (!project || !(await project.modelTokenLive(token))) { |
| 95 | if (await overLimit(env, [["RL_AUTH", `ip:${c.req.header("cf-connecting-ip") ?? "unknown"}`]])) return err(429, "too many requests with invalid model tokens"); |
| 96 | return err(401, "this model token isn't valid: its task run has ended"); |
| 97 | } |
| 98 | // Read and prepare the request first: its estimate is booked against the run before it goes upstream. |
| 99 | let body: string | undefined; |
| 100 | let prepared: ReturnType<typeof prepareRequest> = { body: "", estimate: null }; |
| 101 | if (!["GET", "HEAD"].includes(c.req.method)) { |
| 102 | if (Number(c.req.header("content-length") ?? 0) > MAX_REQUEST) return err(413, "the request is too large"); |
| 103 | const bytes = await readBody(c.req.raw.body, MAX_REQUEST); |
| 104 | if (!bytes) return err(413, "the request is too large"); |
| 105 | prepared = prepareRequest(c.req.method, rest, new TextDecoder().decode(bytes), bytes.byteLength); |
| 106 | body = prepared.body; |
| 107 | } |
| 108 | const { estimate } = prepared; |
| 109 | // The leases a limit on requests at once gave this request (inflight.ts), freed once it is settled. |
| 110 | const slots: string[] = []; |
| 111 | const free = () => { |
| 112 | for (const id of slots.splice(0)) c.executionCtx.waitUntil(freeSlot(env, id).catch((e) => console.error("freeing a model request's slot failed", e))); |
| 113 | }; |
| 114 | // Everything from here, the provider's answer and its settlement included, runs under one waitUntil registered |
| 115 | // now: if the sandbox drops its connection while the Worker waits, the request is still settled and its slots freed. |
| 116 | let settling: Promise<unknown> = Promise.resolve(); |
| 117 | const work = (async (): Promise<Response> => { |
| 118 | const run = await project.modelAuth(token, { provider, estimate, reviewRefuse: prepared.reviewRefuse, model: prepared.model, routing: prepared.routing }); |
| 119 | if (!run) return err(401, "this model token isn't valid: its task run has ended"); |
| 120 | if ("limit" in run) return err(429, run.limit); |
| 121 | if ("refuse" in run) return err(403, run.refuse); |
| 122 | if ("busy" in run) { |
| 123 | const res = err(429, run.busy); |
| 124 | res.headers.set("retry-after", "60"); |
| 125 | return res; |
| 126 | } |
| 127 | // Runs held to their bookings (reviews, coordinators and the tasks they start, and runs on an organization's key, whose |
| 128 | // monthly allowance depends on them) send the capped, tool-stripped body their booking assumes. |
| 129 | const limited = run.bounded || run.orgKey; |
| 130 | if (limited && prepared.reviewBody) body = prepared.reviewBody; |
| 131 | else if (prepared.jailBody) body = prepared.jailBody; |
| 132 | // A booked request that never reaches the provider, or that it answers with an error, costs nothing. |
| 133 | // `spend` is the booking against the person's monthly allowance on an organization's key, once made. |
| 134 | let spend: { org: string; tokens: number; month: string } | null = null; |
| 135 | const refund = (unsent = true) => { |
| 136 | if (estimate || unsent) c.executionCtx.waitUntil(project.chargeUsage(run.task, run.run, NO_USAGE, estimate, unsent).catch((e) => console.error("refunding a model request failed", e))); |
| 137 | if (spend) c.executionCtx.waitUntil(settleSpend(env, spend.org, run.personId, -spend.tokens, spend.month).catch((e) => console.error("refunding an organization's allowance failed", e))); |
| 138 | free(); |
| 139 | }; |
| 140 | const busy = (why: string) => { |
| 141 | refund(); |
| 142 | const res = err(429, why); |
| 143 | res.headers.set("retry-after", "5"); |
| 144 | return res; |
| 145 | }; |
| 146 | // A stored response can't be read back on these runs: on an organization's key it could be anyone's in that account. |
| 147 | if (limited && c.req.method === "GET" && rest.startsWith("/v1/responses")) { |
| 148 | refund(); |
| 149 | return err(403, `reading stored responses isn't available to agent reviews, runs with a token budget and runs on an organization's key`); |
| 150 | } |
| 151 | // A run held to a budget has at most MAX_IN_FLIGHT generating requests at once. |
| 152 | if (run.bounded && estimate) { |
| 153 | const slot = await takeSlot(env, `run:${run.repoId}:${run.task}:${run.run}`); |
| 154 | if (!slot) return busy(`this run already has ${MAX_IN_FLIGHT} model requests on their way; try again when one is answered`); |
| 155 | slots.push(slot); |
| 156 | } |
| 157 | const stop = (why: string) => c.executionCtx.waitUntil(project.stop(run.task, "Diffinery", `Stopped: ${why}.`).then(() => undefined, () => undefined)); |
| 158 | // The person must still be allowed to run the task, on every request; otherwise the task stops. |
| 159 | const dir = new Directory(env); |
| 160 | const [repo, agent] = await Promise.all([dir.repoById(run.repoId), dir.account(run.agentId)]); |
| 161 | if (!repo || !agent || !(await taskPersonAllowed(dir, run.personId, agent, repo))) { |
| 162 | refund(); |
| 163 | stop("the person who started the task can no longer run it"); |
| 164 | return err(403, "the person who started this task can no longer run it"); |
| 165 | } |
| 166 | // Whoever pays must still be allowed to, on every request and uncached (phase 5b): a member removed from the |
| 167 | // organization, a repository transferred out of it, or a policy or key changed applies at the next request. |
| 168 | const payer = await proxyPayer(env, dir, { repo, personId: run.personId, keyOwner: run.keyOwner, provider }); |
| 169 | if ("refuse" in payer) { |
| 170 | refund(); |
| 171 | stop(payer.refuse); |
| 172 | return err(403, payer.refuse); |
| 173 | } |
| 174 | if (payer.org) { |
| 175 | // Under an allowance, a person has at most MAX_IN_FLIGHT generating requests at once on the organization's keys, |
| 176 | // across all their runs: what an image or PDF bills isn't in the booking, so this bounds how far they can go past it. |
| 177 | if (payer.org.cap !== null && estimate) { |
| 178 | const slot = await takeSlot(env, `org:${payer.org.id}:${run.personId}`); |
| 179 | if (!slot) return busy(`${MAX_IN_FLIGHT} of your model requests on ${payer.org.login}'s keys are already on their way; try again when one is answered`); |
| 180 | slots.push(slot); |
| 181 | } |
| 182 | const tokens = estimate ? weighted(estimate) : 0; |
| 183 | const month = spendMonth(); |
| 184 | if (!(await bookSpend(env, payer.org.id, run.personId, tokens, payer.org.cap))) { |
| 185 | const why = `this month's allowance of ${payer.org.cap!.toLocaleString("en-US")} tokens per person on ${payer.org.login}'s keys is used up`; |
| 186 | refund(); |
| 187 | stop(why); |
| 188 | return err(403, why); |
| 189 | } |
| 190 | spend = { org: payer.org.id, tokens, month }; |
| 191 | } |
| 192 | // Decrypted here, for this request only. |
| 193 | const key = await new ModelKeys(env).reveal(run.keyOwner, provider); |
| 194 | if (!key) { |
| 195 | refund(); |
| 196 | return err(403, payer.org ? `${payer.org.login} has no ${PROVIDERS[provider].name} key any more` : `no ${PROVIDERS[provider].name} key in Settings → Model access`); |
| 197 | } |
| 198 | const headers = new Headers(); |
| 199 | // On an organization's key, the run doesn't choose the organization or project within the provider's account. |
| 200 | for (const [k, v] of c.req.raw.headers) if (!DROP.has(k.toLowerCase()) && !ATTRIBUTION.test(k) && !(payer.org && ACCOUNT_SCOPE.test(k))) headers.set(k, v); |
| 201 | if (provider === "anthropic") headers.set("x-api-key", key); |
| 202 | else headers.set("authorization", `Bearer ${key}`); |
| 203 | if (provider === "openrouter") { |
| 204 | // OpenRouter's app attribution headers: Diffinery, not whatever the sandbox claims. |
| 205 | headers.set("http-referer", env.PUBLIC_URL); |
| 206 | headers.set("x-title", "Diffinery"); |
| 207 | } |
| 208 | const base = (provider === "anthropic" ? env.MODEL_UPSTREAM_ANTHROPIC : provider === "openai" ? env.MODEL_UPSTREAM_OPENAI : env.MODEL_UPSTREAM_OPENROUTER) || UPSTREAM[provider]; |
| 209 | // If the provider can't be reached, the booking stays: it may have started the answer. |
| 210 | const res = await fetch(`${base}${rest}${url.search}`, { method: c.req.method, headers, body, redirect: "manual" }); |
| 211 | const out = new Headers(res.headers); |
| 212 | out.delete("set-cookie"); |
| 213 | if (!res.body || !res.ok) { |
| 214 | refund(false); |
| 215 | return new Response(res.body, { status: res.status, headers: out }); |
| 216 | } |
| 217 | if (!estimate) { |
| 218 | free(); |
| 219 | return new Response(res.body, { status: res.status, headers: out }); |
| 220 | } |
| 221 | // Settle the booking to the tokens the provider reports, read from a copy of the response as it streams back. |
| 222 | // An answer that reports none, or with a part too large to read for usage, is charged its estimate, or on a run |
| 223 | // held to its bookings unreadCharge; one whose reading fails is never settled, and so stays charged its estimate. |
| 224 | const [agentCopy, meter] = res.body.tee(); |
| 225 | const booked = spend; |
| 226 | const sent = body === undefined ? 0 : new TextEncoder().encode(body).byteLength; |
| 227 | settling = |
| 228 | meterUsage(meter, res.headers.get("content-type") ?? "") |
| 229 | .then(async ({ usage, unreadable }) => { |
| 230 | const u = limited && (unreadable || !usage) ? unreadCharge(estimate, sent) : (usage ?? estimate); |
| 231 | await project.chargeUsage(run.task, run.run, u, estimate); |
| 232 | if (booked) await settleSpend(env, booked.org, run.personId, weighted(u) - booked.tokens, booked.month); |
| 233 | }) |
| 234 | .catch((e) => console.error("reading model usage failed", e)) |
| 235 | .finally(free); |
| 236 | return new Response(agentCopy, { status: res.status, headers: out }); |
| 237 | })().catch((e) => { |
| 238 | free(); |
| 239 | throw e; |
| 240 | }); |
| 241 | c.executionCtx.waitUntil(work.then(() => settling, () => undefined)); |
| 242 | return work; |
| 243 | } |