| 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 | |
| 6 | /** |
| 7 | * The model proxy: `/-/model/<provider>/v1/...`. A task's runtime calls it |
| 8 | * with the run's token (`dmt_...`, as its API key) and the proxy forwards the |
| 9 | * request to the provider with the person's own key, added here. So the key |
| 10 | * never enters the sandbox, and the token stops working when the run ends. |
| 11 | * |
| 12 | * It runs before the app's identity middleware: its token isn't an account |
| 13 | * credential and works nowhere else. |
| 14 | */ |
| 15 | |
| 16 | const UPSTREAM: Record<Provider, string> = { anthropic: "https://api.anthropic.com", openai: "https://api.openai.com" }; |
| 17 | |
| 18 | /** |
| 19 | * Only the model endpoints the runtimes use. The person's key could reach the |
| 20 | * rest of their provider account (stored files, batches, assistants), which an |
| 21 | * agent misled by repository content must not. |
| 22 | */ |
| 23 | const ALLOWED: Record<Provider, Array<[method: string, path: RegExp]>> = { |
| 24 | anthropic: [ |
| 25 | ["POST", /^\/v1\/messages$/], |
| 26 | ["POST", /^\/v1\/messages\/count_tokens$/], |
| 27 | ["GET", /^\/v1\/models(\/[A-Za-z0-9._:-]+)?$/], |
| 28 | ], |
| 29 | openai: [ |
| 30 | ["POST", /^\/v1\/responses$/], |
| 31 | // Codex compacts long conversations, and reads a background response back by id. |
| 32 | ["POST", /^\/v1\/responses\/compact$/], |
| 33 | ["GET", /^\/v1\/responses\/[A-Za-z0-9_-]+$/], |
| 34 | ["POST", /^\/v1\/chat\/completions$/], |
| 35 | ["GET", /^\/v1\/models(\/[A-Za-z0-9._:-]+)?$/], |
| 36 | ], |
| 37 | }; |
| 38 | |
| 39 | export const proxyAllows = (provider: Provider, method: string, path: string): boolean => ALLOWED[provider].some(([m, re]) => m === method && re.test(path)); |
| 40 | |
| 41 | /** Request headers that never go upstream: credentials, cookies and the hop's own. */ |
| 42 | 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"]); |
| 43 | |
| 44 | const err = (status: number, message: string) => |
| 45 | new Response(JSON.stringify({ type: "error", error: { type: status === 401 ? "authentication_error" : status === 404 ? "not_found_error" : "permission_error", message } }), { status, headers: { "content-type": "application/json" } }); |
| 46 | |
| 47 | export function proxyToken(req: Request): string | null { |
| 48 | const key = req.headers.get("x-api-key"); |
| 49 | if (key?.startsWith("dmt_")) return key.trim(); |
| 50 | const m = /^Bearer\s+(dmt_\S+)$/i.exec(req.headers.get("authorization") ?? ""); |
| 51 | return m ? m[1]! : null; |
| 52 | } |
| 53 | |
| 54 | export async function modelProxy(c: AppContext): Promise<Response> { |
| 55 | const provider = c.req.param("provider") ?? ""; |
| 56 | if (!isProvider(provider)) return err(404, "unknown model provider"); |
| 57 | const url = new URL(c.req.url); |
| 58 | const rest = url.pathname.slice(`/-/model/${provider}`.length); |
| 59 | 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`); |
| 60 | const token = proxyToken(c.req.raw); |
| 61 | if (!token) return err(401, "a task's model token is required"); |
| 62 | const repoId = token.split("_")[1]; |
| 63 | if (!repoId || !/^[A-Za-z0-9-]{8,64}$/.test(repoId)) return err(401, "malformed model token"); |
| 64 | const env = c.env; |
| 65 | const project = env.PROJECT.get(env.PROJECT.idFromName(repoId)); |
| 66 | const run = await project.modelAuth(token); |
| 67 | if (!run) return err(401, "this model token isn't valid: its task run has ended"); |
| 68 | if (!run.providers.includes(provider)) return err(403, `this task doesn't use ${PROVIDERS[provider].name}`); |
| 69 | // The person must still be allowed to run the task, on every request; otherwise the task stops. |
| 70 | const dir = new Directory(env); |
| 71 | const [repo, agent] = await Promise.all([dir.repoById(run.repoId), dir.account(run.agentId)]); |
| 72 | if (!repo || !agent || !(await taskPersonAllowed(dir, run.personId, agent, repo))) { |
| 73 | c.executionCtx.waitUntil(project.stop(run.task, "Diffinery", "Stopped: the person who started the task can no longer run it.").then(() => undefined, () => undefined)); |
| 74 | return err(403, "the person who started this task can no longer run it"); |
| 75 | } |
| 76 | const key = await new ModelKeys(env).reveal(run.personId, provider); |
| 77 | if (!key) return err(403, `no ${PROVIDERS[provider].name} key in Settings → Model access`); |
| 78 | const headers = new Headers(); |
| 79 | for (const [k, v] of c.req.raw.headers) if (!DROP.has(k.toLowerCase())) headers.set(k, v); |
| 80 | if (provider === "anthropic") headers.set("x-api-key", key); |
| 81 | else headers.set("authorization", `Bearer ${key}`); |
| 82 | const base = (provider === "anthropic" ? env.MODEL_UPSTREAM_ANTHROPIC : env.MODEL_UPSTREAM_OPENAI) || UPSTREAM[provider]; |
| 83 | const res = await fetch(`${base}${rest}${url.search}`, { method: c.req.method, headers, body: ["GET", "HEAD"].includes(c.req.method) ? undefined : c.req.raw.body, redirect: "manual" }); |
| 84 | const out = new Headers(res.headers); |
| 85 | out.delete("set-cookie"); |
| 86 | return new Response(res.body, { status: res.status, headers: out }); |
| 87 | } |