diff --git a/packages/ai/src/api/cloudflare-gateway-binding.ts b/packages/ai/src/api/cloudflare-gateway-binding.ts new file mode 100644 index 000000000..c2eaa5d48 --- /dev/null +++ b/packages/ai/src/api/cloudflare-gateway-binding.ts @@ -0,0 +1,192 @@ +/** + * AI Gateway transport over the Workers AI binding. + * + * pi's Cloudflare AI Gateway support speaks HTTPS + * (`gateway.ai.cloudflare.com/v1/{account}/{gateway}/{provider}/...`, see `api/cloudflare.ts`), + * which needs a Cloudflare API token even when the caller is a Worker in the gateway's own + * account. + * + * In order to solve for this problem, `createGatewayBindingFetch` returns a {@link FetchFunction} + * that translates requests under a gateway HTTPS prefix into calls to the Workers AI binding's + * universal endpoint, `env.AI.gateway(id).run({provider, endpoint, headers, query})`. + * Binding calls are pre-authenticated in-account and return the provider's native wire format as a + * regular (streaming) `Response`, so API implementations behave identically over either + * transport. + * + * The result is the transport for one gateway-bound client, not a general-purpose fetch: + * requests it cannot serve — URLs outside the prefix, or in-prefix requests the universal + * endpoint cannot express (non-POST, non-JSON body) — reject with a descriptive error. + * Transport selection is the caller's job, per client: route such traffic over HTTPS with + * real gateway auth instead of through this shim. + */ + +import type { FetchFunction } from "../types.ts"; + +/** + * Structural type for the Workers AI binding's gateway surface (`env.AI`), so this + * module does not depend on `@cloudflare/workers-types`. Any real `Ai` binding satisfies it. + */ +export interface AiGatewayBinding { + gateway(id: string): AiGatewayBindingGateway; +} + +export interface AiGatewayBindingGateway { + run(data: AiGatewayUniversalRequestLike, options?: { signal?: AbortSignal }): Promise; +} + +/** One universal-endpoint request entry, as accepted by `AiGateway.run()`. */ +export interface AiGatewayUniversalRequestLike { + provider: string; + endpoint: string; + headers: Record; + query: unknown; +} + +/** + * Placeholder value for auth headers on binding-routed requests. API implementations + * require an API key or a recognized auth header (`authorization`, `x-api-key`, + * `cf-aig-authorization`) before dispatch; binding calls are pre-authenticated, so pass + * `cf-aig-authorization: Bearer ${CLOUDFLARE_GATEWAY_BINDING_AUTH_SENTINEL}` to satisfy + * the check. The shim strips `cf-aig-authorization` before calling the binding. Pair it with + * `Authorization: null` / `x-api-key: null` so the SDKs' placeholder auth headers never reach + * the gateway, which would treat a request-supplied auth header as a BYOK provider key that + * overrides its stored keys — the same as it would over HTTPS. + */ +export const CLOUDFLARE_GATEWAY_BINDING_AUTH_SENTINEL = "cloudflare-gateway-binding"; + +export interface GatewayBindingFetchOptions { + /** The Workers AI binding (e.g. `env.AI`). */ + binding: AiGatewayBinding; + /** + * Gateway HTTPS prefix every request must fall under, without a trailing slash: + * `https://gateway.ai.cloudflare.com/v1/{accountId}/{gatewayName}`. + */ + baseUrl: string; + /** Gateway name passed to `binding.gateway()`. Must match the `baseUrl` gateway. */ + gateway: string; +} + +// Never forwarded to the binding: hop-by-hop/derived headers, and gateway auth +// (binding calls are pre-authenticated; the sentinel must not reach the wire). +const STRIP_HEADERS = new Set(["content-length", "host", "cf-aig-authorization"]); + +type FetchInput = Parameters[0]; + +/** + * Create a `fetch` that routes AI Gateway requests through the Workers AI binding. + * See the module docs for behavior and composition notes. + */ +export function createGatewayBindingFetch(options: GatewayBindingFetchOptions): FetchFunction { + const { binding, gateway } = options; + // Prefix matching runs on URL-normalized components (origin + pathname), not raw strings: + // dot segments resolve away and fragments drop, matching what real fetch would put on the + // wire, so a lexical variant can't split provider/endpoint differently than HTTPS would. + const base = new URL(options.baseUrl); + const basePath = base.pathname.endsWith("/") ? base.pathname : `${base.pathname}/`; + + return async (input: FetchInput, init?: RequestInit): Promise => { + const request = input instanceof Request ? input : undefined; + const url = request ? request.url : input.toString(); + const method = (init?.method ?? request?.method ?? "GET").toUpperCase(); + let parsed: URL | undefined; + try { + parsed = new URL(url); + } catch { + parsed = undefined; + } + // Out-of-prefix URLs are a configuration bug, not passthrough traffic: silently + // forwarding would ship the auth sentinel to whatever host the URL names. + if (parsed === undefined || parsed.origin !== base.origin || !parsed.pathname.startsWith(basePath)) { + throw new Error( + `createGatewayBindingFetch: ${method} ${url} is outside the configured gateway ` + + `prefix (${base.origin}${basePath}); this fetch only serves its gateway-bound client`, + ); + } + + // In-prefix requests the universal endpoint cannot express always reject: forwarding + // them over HTTPS would send the sentinel to the gateway and fail with a misleading + // auth error instead of naming the real problem. Callers that need such endpoints + // route them over HTTPS with real gateway auth themselves. + const unexpressible = (reason: string): never => { + throw new Error( + `createGatewayBindingFetch: cannot express ${method} ${url} as a universal ` + + `gateway request (${reason}); route it over HTTPS with gateway auth instead`, + ); + }; + if (method !== "POST") return unexpressible("only POST is supported"); + + const rest = parsed.pathname.slice(basePath.length); + const slash = rest.indexOf("/"); + if (slash <= 0) { + return unexpressible("missing provider/endpoint path"); + } + const provider = rest.slice(0, slash); + // Keep the query string on the endpoint — it's part of what HTTPS would have sent. + const endpoint = rest.slice(slash + 1) + parsed.search; + + const bodyText = await readBodyText(request, init); + let query: unknown; + try { + query = bodyText === undefined ? undefined : JSON.parse(bodyText); + } catch { + return unexpressible("non-JSON body"); + } + if (query === undefined) { + return unexpressible("missing body"); + } + + const headers = collectHeaders(request, init); + // Per the fetch spec an explicit `signal: null` in init clears a Request input's signal. + const signal = init?.signal ?? (init && "signal" in init && init.signal === null ? undefined : request?.signal); + return binding.gateway(gateway).run({ provider, endpoint, headers, query }, signal ? { signal } : {}); + }; +} + +async function readBodyText(request: Request | undefined, init?: RequestInit): Promise { + const body = init?.body; + if (body === undefined || body === null) { + // Per the fetch spec an explicit `body: null` in init clears a Request input's body. + if (init && "body" in init && body === null) return undefined; + if (request && request.body !== null) return request.clone().text(); + return undefined; + } + if (typeof body === "string") return body; + if (body instanceof Uint8Array) return new TextDecoder().decode(body); + if (body instanceof ArrayBuffer) return new TextDecoder().decode(new Uint8Array(body)); + // URLSearchParams, FormData, Blob, ReadableStream in init: read via a Request wrapper. + // Consuming a one-shot stream here is fine — unexpressible requests reject rather than + // replay, so nothing downstream needs the body again. + return new Request("http://body.local", { + method: "POST", + body, + // The fetch spec requires `duplex: "half"` to construct a Request with a stream body + // (Node's undici enforces it; it is ignored for the replayable body types). TypeScript's + // RequestInit does not declare the field yet, hence the cast. + duplex: "half", + } as RequestInit).text(); +} + +// Entry header names are lowercased so case-variant duplicates collapse and stripping is +// uniform. Per the fetch spec, `init.headers` replaces a Request input's headers entirely. +function collectHeaders(request: Request | undefined, init?: RequestInit): Record { + const result: Record = {}; + const add = (key: string, value: string) => { + const name = key.toLowerCase(); + if (!STRIP_HEADERS.has(name)) result[name] = value; + }; + const headers = init?.headers; + if (headers === undefined) { + if (request) { + for (const [key, value] of request.headers) add(key, value); + } + } else if (headers instanceof Headers) { + for (const [key, value] of headers) add(key, value); + } else if (Array.isArray(headers)) { + for (const [key, value] of headers) add(key, value); + } else { + for (const [key, value] of Object.entries(headers)) { + if (value !== undefined) add(key, String(value)); + } + } + return result; +} diff --git a/packages/ai/test/cloudflare-gateway-binding.test.ts b/packages/ai/test/cloudflare-gateway-binding.test.ts new file mode 100644 index 000000000..d59611a88 --- /dev/null +++ b/packages/ai/test/cloudflare-gateway-binding.test.ts @@ -0,0 +1,332 @@ +import { describe, expect, it } from "vitest"; +import { + type AiGatewayUniversalRequestLike, + CLOUDFLARE_GATEWAY_BINDING_AUTH_SENTINEL, + createGatewayBindingFetch, +} from "../src/api/cloudflare-gateway-binding.ts"; +import { streamSimple as streamOpenAICompletions } from "../src/api/openai-completions.ts"; +import type { Model } from "../src/types.ts"; + +const BASE_URL = "https://gateway.ai.cloudflare.com/v1/account-id/my-gateway"; + +interface CapturedRun { + gatewayId: string; + data: AiGatewayUniversalRequestLike; + options: { signal?: AbortSignal } | undefined; +} + +function fakeBinding(response?: Response) { + const runs: CapturedRun[] = []; + const binding = { + gateway: (gatewayId: string) => ({ + run: (data: AiGatewayUniversalRequestLike, options?: { signal?: AbortSignal }) => { + runs.push({ gatewayId, data, options }); + return Promise.resolve(response ?? new Response("{}")); + }, + }), + }; + return { binding, runs }; +} + +describe("createGatewayBindingFetch", () => { + it("derives provider and endpoint from gateway passthrough URLs", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await fetchFn(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + body: JSON.stringify({ model: "claude" }), + }); + await fetchFn(`${BASE_URL}/openai/responses`, { + method: "POST", + body: JSON.stringify({ model: "gpt" }), + }); + await fetchFn(`${BASE_URL}/workers-ai/v1/chat/completions`, { + method: "POST", + body: JSON.stringify({ model: "@cf/meta/llama" }), + }); + + expect(runs.map((run) => [run.data.provider, run.data.endpoint])).toEqual([ + ["anthropic", "v1/messages"], + ["openai", "responses"], + ["workers-ai", "v1/chat/completions"], + ]); + expect(runs.map((run) => run.gatewayId)).toEqual(["my-gateway", "my-gateway", "my-gateway"]); + expect(runs[0].data.query).toEqual({ model: "claude" }); + }); + + it("keeps the query string in the endpoint", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await fetchFn(`${BASE_URL}/openai/responses?beta=true`, { + method: "POST", + body: "{}", + }); + + expect(runs[0].data.endpoint).toBe("responses?beta=true"); + }); + + it("lowercases header names so case-variant duplicates collapse", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await fetchFn(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + headers: { "Anthropic-Version": "2023-06-01" }, + body: "{}", + }); + + expect(runs[0].data.headers).toEqual({ "anthropic-version": "2023-06-01" }); + }); + + it("lets init headers replace a Request input's headers, per the fetch spec", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await fetchFn( + new Request(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + headers: { "x-from-request": "yes" }, + body: "{}", + }), + { headers: { "x-from-init": "yes" } }, + ); + + expect(runs[0].data.headers["x-from-init"]).toBe("yes"); + expect(runs[0].data.headers["x-from-request"]).toBeUndefined(); + }); + + it("strips gateway auth and derived headers, forwards the rest", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await fetchFn(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + headers: { + "Content-Type": "application/json", + "Content-Length": "17", + "CF-AIG-Authorization": `Bearer ${CLOUDFLARE_GATEWAY_BINDING_AUTH_SENTINEL}`, + "cf-aig-metadata": '{"user":"42"}', + "anthropic-version": "2023-06-01", + "x-api-key": "provider-key", + }, + body: "{}", + }); + + const headers = Object.fromEntries( + Object.entries(runs[0].data.headers).map(([key, value]) => [key.toLowerCase(), value]), + ); + expect(headers["cf-aig-authorization"]).toBeUndefined(); + expect(headers["content-length"]).toBeUndefined(); + expect(headers["cf-aig-metadata"]).toBe('{"user":"42"}'); + expect(headers["anthropic-version"]).toBe("2023-06-01"); + // Provider auth headers pass through: that is how request-supplied (BYOK) keys ride. + expect(headers["x-api-key"]).toBe("provider-key"); + }); + + it("accepts Request inputs and forwards their headers and body", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await fetchFn( + new Request(`${BASE_URL}/openai/chat/completions`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ stream: true }), + }), + ); + + expect(runs).toHaveLength(1); + expect(runs[0].data.provider).toBe("openai"); + expect(runs[0].data.endpoint).toBe("chat/completions"); + expect(runs[0].data.query).toEqual({ stream: true }); + expect(runs[0].data.headers["content-type"]).toBe("application/json"); + }); + + it("forwards the abort signal", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + const controller = new AbortController(); + + await fetchFn(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + body: "{}", + signal: controller.signal, + }); + + expect(runs[0].options?.signal).toBe(controller.signal); + }); + + it("lets an explicit `signal: null` in init clear a Request input's signal, per the fetch spec", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + const controller = new AbortController(); + + await fetchFn( + new Request(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + body: "{}", + signal: controller.signal, + }), + { signal: null }, + ); + + expect(runs).toHaveLength(1); + expect(runs[0].options?.signal).toBeUndefined(); + }); + + it("returns the binding response untouched, including streaming bodies", async () => { + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data: {}\n\n")); + controller.close(); + }, + }); + const bindingResponse = new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream", "cf-aig-log-id": "log-1" }, + }); + const { binding } = fakeBinding(bindingResponse); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + const response = await fetchFn(`${BASE_URL}/workers-ai/v1/chat/completions`, { + method: "POST", + body: "{}", + }); + + expect(response).toBe(bindingResponse); + expect(response.headers.get("cf-aig-log-id")).toBe("log-1"); + expect(await response.text()).toBe("data: {}\n\n"); + }); + + it("rejects in-prefix requests the universal endpoint cannot express", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await expect(fetchFn(`${BASE_URL}/anthropic/v1/messages`, { method: "GET" })).rejects.toThrow( + "cannot express GET", + ); + await expect(fetchFn(`${BASE_URL}/anthropic/v1/messages`, { method: "POST", body: "not json" })).rejects.toThrow( + "non-JSON body", + ); + await expect(fetchFn(`${BASE_URL}/anthropic`, { method: "POST", body: "{}" })).rejects.toThrow( + "missing provider/endpoint path", + ); + expect(runs).toHaveLength(0); + }); + + it("rejects URLs outside the gateway prefix: transport selection is the caller's", async () => { + // Silent passthrough would ship the auth sentinel to whatever host the URL names; a + // misconfigured baseUrl must fail loudly instead. + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + await expect( + fetchFn("https://api.openai.com/v1/chat/completions", { method: "POST", body: "{}" }), + ).rejects.toThrow("outside the configured gateway prefix"); + // Same origin, different path (another account's gateway) is just as out-of-prefix. + await expect( + fetchFn("https://gateway.ai.cloudflare.com/v1/other-account/my-gateway/anthropic/v1/messages", { + method: "POST", + body: "{}", + }), + ).rejects.toThrow("outside the configured gateway prefix"); + expect(runs).toHaveLength(0); + }); + + it("matches and splits on the URL-normalized path, as real fetch would send it", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + + // Dot segments normalize away before the provider/endpoint split, so a lexical variant + // routes exactly like its normal form (raw string prefixing would split it differently). + await fetchFn(`${BASE_URL}/anthropic/../anthropic/v1/./messages`, { + method: "POST", + body: JSON.stringify({ model: "claude" }), + }); + expect(runs.map((run) => [run.data.provider, run.data.endpoint])).toEqual([["anthropic", "v1/messages"]]); + + // A dot-segment URL that resolves outside the prefix is rejected even though it starts + // with the prefix as a raw string. + await expect( + fetchFn(`${BASE_URL}/../other-gateway/anthropic/v1/messages`, { method: "POST", body: "{}" }), + ).rejects.toThrow("outside the configured gateway prefix"); + expect(runs).toHaveLength(1); + }); + + it("consumes a one-shot stream body for the JSON probe", async () => { + const { binding, runs } = fakeBinding(); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + const streamOf = (text: string) => + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(text)); + controller.close(); + }, + }); + + // JSON stream body: consumed once, reaches the binding as the parsed query. + await fetchFn(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + body: streamOf('{"model":"claude"}'), + duplex: "half", + } as RequestInit); + expect(runs).toHaveLength(1); + expect(runs[0].data.query).toEqual({ model: "claude" }); + + // Non-JSON stream body: rejects like any other non-JSON body (never replayed). + await expect( + fetchFn(`${BASE_URL}/anthropic/v1/messages`, { + method: "POST", + body: streamOf("not json"), + duplex: "half", + } as RequestInit), + ).rejects.toThrow("non-JSON body"); + expect(runs).toHaveLength(1); + }); + + it("keeps SDK placeholder auth out of entries when paired with null auth headers", async () => { + // The full header contract from the module docs: the sentinel satisfies pi's request-auth + // check, and the explicit nulls make the OpenAI SDK delete its own `Authorization: Bearer + // unused` placeholder before the request reaches the shim. + const { binding, runs } = fakeBinding( + Response.json({ error: { type: "bad_request", message: "stubbed" } }, { status: 400 }), + ); + const fetchFn = createGatewayBindingFetch({ binding, baseUrl: BASE_URL, gateway: "my-gateway" }); + const model: Model<"openai-completions"> = { + id: "test-model", + name: "Test Model", + api: "openai-completions", + provider: "openai", + baseUrl: `${BASE_URL}/openai`, + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 10_000, + maxTokens: 1_000, + }; + + const result = await streamOpenAICompletions( + model, + { messages: [{ role: "user", content: "hello", timestamp: 1 }] }, + { + headers: { + "cf-aig-authorization": `Bearer ${CLOUDFLARE_GATEWAY_BINDING_AUTH_SENTINEL}`, + Authorization: null, + "x-api-key": null, + }, + fetch: fetchFn, + maxRetries: 0, + }, + ).result(); + + expect(result.stopReason).toBe("error"); + expect(runs).toHaveLength(1); + expect(runs[0].data.provider).toBe("openai"); + const headerNames = Object.keys(runs[0].data.headers); + expect(headerNames).not.toContain("authorization"); + expect(headerNames).not.toContain("x-api-key"); + expect(headerNames).not.toContain("cf-aig-authorization"); + }); +});