diff --git a/.changeset/tanstack-ai-byok-alias.md b/.changeset/tanstack-ai-byok-alias.md new file mode 100644 index 0000000000..4763f60652 --- /dev/null +++ b/.changeset/tanstack-ai-byok-alias.md @@ -0,0 +1,5 @@ +--- +"@cloudflare/tanstack-ai": minor +--- + +Add `byokAlias` on the shared `AiGatewayConfig` so every gateway adapter (OpenAI, Anthropic, Gemini, Grok, OpenRouter, Workers AI) can select a stored BYOK key via `cf-aig-byok-alias`. `createGatewayFetch` forwards it on REST; Gemini maps the same field through `httpOptions.headers`. The AI binding ignores the header for third-party models. diff --git a/docs/concepts/binding-vs-rest.md b/docs/concepts/binding-vs-rest.md index c99e86cd62..75656c2e2b 100644 --- a/docs/concepts/binding-vs-rest.md +++ b/docs/concepts/binding-vs-rest.md @@ -35,7 +35,8 @@ const provider = createWorkersAI({ | ------------------------------------------------- | ------- | ---- | | Chat / generate / stream | ✅ | ✅ | | Image / embeddings / transcription / TTS / rerank | ✅ | ✅ | -| Gateway routing (caching, metadata, BYOK) | ✅ | ✅ | +| Gateway routing (caching, metadata) | ✅ | ✅ | +| BYOK stored-key alias (`cf-aig-byok-alias`) | ❌ | ✅ | | Server-side fallback | ✅ | ✅ | | **Resumable streaming** _(coming soon)_ | ✅ | ❌ | diff --git a/docs/tanstack-ai/README.md b/docs/tanstack-ai/README.md index 8502e884d7..6ffa5da19f 100644 --- a/docs/tanstack-ai/README.md +++ b/docs/tanstack-ai/README.md @@ -166,6 +166,19 @@ const adapter = createOpenAiChat("gpt-5", { }); ``` +`byokAlias` (`cf-aig-byok-alias`) is credentials / REST only — including +Gemini (`createGeminiChat` via `httpOptions.headers`). The AI binding does +not select a non-`default` stored key for third-party models: + +```ts +const adapter = createOpenAiChat("gpt-5", { + accountId: env.CF_ACCOUNT_ID, + gatewayId: "my-gateway", + cfApiKey: env.CF_AIG_TOKEN, + byokAlias: "development", +}); +``` + See the [package README](../../packages/tanstack-ai/README.md) for the full list of provider factories (chat / summarize / image / transcription / TTS / video) and the four Workers AI configuration modes. diff --git a/packages/tanstack-ai/README.md b/packages/tanstack-ai/README.md index 63a13935a5..bb68ab7b42 100644 --- a/packages/tanstack-ai/README.md +++ b/packages/tanstack-ai/README.md @@ -234,6 +234,23 @@ const adapter = createOpenAiChat("gpt-4o", { }); ``` +**BYOK alias (credentials / REST only):** + +`cf-aig-byok-alias` is a provider-passthrough header. The AI binding does not +honor it for third-party models — only the `default` stored key is consulted +there. Same option on `createAnthropicChat` / `createGrokChat` / +`createOpenRouterChat` / `createWorkersAiChat` (gateway REST) / +`createGeminiChat` (credentials / `httpOptions.headers`). + +```typescript +const adapter = createOpenRouterChat("openai/gpt-4o", { + accountId: "your-account-id", + gatewayId: "your-gateway-id", + cfApiKey: "your-cf-api-key", + byokAlias: "development", +}); +``` + ### Workers AI through Gateway ```typescript diff --git a/packages/tanstack-ai/src/adapters/gemini.ts b/packages/tanstack-ai/src/adapters/gemini.ts index 7eec834aa6..1aa53cad23 100644 --- a/packages/tanstack-ai/src/adapters/gemini.ts +++ b/packages/tanstack-ai/src/adapters/gemini.ts @@ -66,6 +66,9 @@ function buildGeminiGatewayConfig(config: GeminiGatewayConfig) { if (typeof config.metadata === "object") { headers["cf-aig-metadata"] = JSON.stringify(config.metadata); } + if (typeof config.byokAlias === "string" && config.byokAlias.length > 0) { + headers["cf-aig-byok-alias"] = config.byokAlias; + } const apiKey = config.apiKey ?? config.cfApiKey; diff --git a/packages/tanstack-ai/src/utils/create-fetcher.ts b/packages/tanstack-ai/src/utils/create-fetcher.ts index 8f5b4b9bad..31d8531f05 100644 --- a/packages/tanstack-ai/src/utils/create-fetcher.ts +++ b/packages/tanstack-ai/src/utils/create-fetcher.ts @@ -61,6 +61,13 @@ export interface AiGatewayConfig { cacheTtl?: number; customCacheKey?: string; metadata?: Record; + /** + * BYOK stored-key alias (`cf-aig-byok-alias`). Honored on credentials / + * REST / provider-passthrough only. The AI binding ignores this header + * for third-party models. + * See https://developers.cloudflare.com/ai-gateway/configuration/bring-your-own-keys/ + */ + byokAlias?: string; } export type AiGatewayAdapterConfig = (AiGatewayBindingConfig | AiGatewayCredentialsConfig) & @@ -258,6 +265,11 @@ export function createGatewayFetch( ...(config.metadata && typeof config.metadata === "object" ? { metadata: config.metadata as GatewayMetadata } : {}), + ...("gatewayId" in config && + typeof config.byokAlias === "string" && + config.byokAlias.length > 0 + ? { byokAlias: config.byokAlias } + : {}), }); const request: { diff --git a/packages/tanstack-ai/test/gateway-adapters.test.ts b/packages/tanstack-ai/test/gateway-adapters.test.ts index f0757700c2..6c0e8ab8cc 100644 --- a/packages/tanstack-ai/test/gateway-adapters.test.ts +++ b/packages/tanstack-ai/test/gateway-adapters.test.ts @@ -304,6 +304,18 @@ describe("Gemini gateway adapters", () => { expect(config.httpOptions.headers["cf-aig-metadata"]).toBe(JSON.stringify({ env: "test" })); }); + it("createGeminiChat passes byokAlias via httpOptions.headers", async () => { + const { createGeminiChat } = await import("../src/adapters/gemini"); + const configWithAlias: GeminiGatewayConfig = { + ...geminiConfig, + byokAlias: "development", + }; + createGeminiChat("gemini-2.5-flash" as any, configWithAlias); + + const [config] = mockGeminiTextCtor.mock.calls[0]!; + expect(config.httpOptions.headers["cf-aig-byok-alias"]).toBe("development"); + }); + it("createGeminiImage with credentials config", async () => { const { createGeminiImage } = await import("../src/adapters/gemini"); createGeminiImage("imagen-4.0-generate-001" as any, geminiConfig); @@ -542,6 +554,31 @@ describe("OpenRouter gateway adapters", () => { const [config] = mockOpenRouterTextCtor.mock.calls[0]!; expect(config.apiKey).toBe("unused"); }); + + it("createOpenRouterChat forwards byokAlias through the HTTPClient fetcher", async () => { + const originalFetch = globalThis.fetch; + const mockFetch = vi.fn(async (..._args: unknown[]) => new Response("ok")); + globalThis.fetch = mockFetch as any; + + try { + const { createOpenRouterChat } = await import("../src/adapters/openrouter"); + createOpenRouterChat("openai/gpt-4o", { + ...credentialsConfig, + byokAlias: "development", + }); + + const [config] = mockOpenRouterTextCtor.mock.calls[0]!; + await config.httpClient.fetcher("https://openrouter.ai/api/v1/chat/completions", { + method: "POST", + body: JSON.stringify({ model: "openai/gpt-4o", messages: [] }), + }); + + const [, init] = mockFetch.mock.calls[0]!; + expect((init as any).headers["cf-aig-byok-alias"]).toBe("development"); + } finally { + globalThis.fetch = originalFetch; + } + }); }); describe("gateway fetch integration", () => { @@ -569,6 +606,37 @@ describe("gateway fetch integration", () => { } }); + it("createOpenAiChat and createAnthropicChat forward byokAlias on REST", async () => { + const originalFetch = globalThis.fetch; + const mockFetch = vi.fn(async (..._args: unknown[]) => new Response("ok")); + globalThis.fetch = mockFetch as any; + + try { + const { createOpenAiChat } = await import("../src/adapters/openai"); + const { createAnthropicChat } = await import("../src/adapters/anthropic"); + const withAlias = { ...credentialsConfig, byokAlias: "development" }; + + createOpenAiChat("gpt-4o" as any, withAlias); + await mockOpenAITextCtor.mock.calls[0]![0].fetch( + "https://api.openai.com/v1/chat/completions", + { body: JSON.stringify({ model: "gpt-4o", messages: [] }) }, + ); + + createAnthropicChat("claude-sonnet-4-5" as any, withAlias); + await mockAnthropicTextCtor.mock.calls[0]![0].fetch( + "https://api.anthropic.com/v1/messages", + { body: JSON.stringify({ model: "claude-sonnet-4-5", messages: [] }) }, + ); + + expect(mockFetch).toHaveBeenCalledTimes(2); + for (const [, init] of mockFetch.mock.calls) { + expect((init as any).headers["cf-aig-byok-alias"]).toBe("development"); + } + } finally { + globalThis.fetch = originalFetch; + } + }); + it("binding config produces fetch that calls binding.run", async () => { const { createAnthropicChat } = await import("../src/adapters/anthropic"); createAnthropicChat("claude-sonnet-4-5" as any, bindingConfig); @@ -598,6 +666,7 @@ describe("gateway fetch integration", () => { cacheTtl: 300, customCacheKey: "my-key", metadata: { env: "test" }, + byokAlias: "development", }); const [config] = mockGrokTextCtor.mock.calls[0]!; @@ -612,6 +681,8 @@ describe("gateway fetch integration", () => { expect(body.headers["cf-aig-cache-ttl"]).toBe("300"); expect(body.headers["cf-aig-cache-key"]).toBe("my-key"); expect(body.headers["cf-aig-metadata"]).toBe(JSON.stringify({ env: "test" })); + expect(body.headers["cf-aig-byok-alias"]).toBe("development"); + expect((init as any).headers["cf-aig-byok-alias"]).toBe("development"); } finally { globalThis.fetch = originalFetch; } diff --git a/packages/tanstack-ai/test/gateway-fetch.test.ts b/packages/tanstack-ai/test/gateway-fetch.test.ts index 4e2fced930..fe941fd0b4 100644 --- a/packages/tanstack-ai/test/gateway-fetch.test.ts +++ b/packages/tanstack-ai/test/gateway-fetch.test.ts @@ -212,6 +212,24 @@ describe("createGatewayFetch", () => { const [, init] = mockFetch.mock.calls[0]!; expect(init.headers["cf-aig-authorization"]).toBeUndefined(); }); + + it("should set cf-aig-byok-alias on the outer REST request", async () => { + const fetcher = createGatewayFetch("openrouter", { + ...credentialsConfig, + byokAlias: "development", + }); + + await fetcher("https://openrouter.ai/api/v1/chat/completions", { + method: "POST", + body: JSON.stringify({}), + }); + + const [, init] = mockFetch.mock.calls[0]!; + expect(init.headers["cf-aig-byok-alias"]).toBe("development"); + const body = JSON.parse(init.body); + expect(body.provider).toBe("openrouter"); + expect(body.headers["cf-aig-byok-alias"]).toBe("development"); + }); }); describe("cache headers", () => { @@ -304,6 +322,23 @@ describe("createGatewayFetch", () => { expect(request.headers["cf-aig-cache-ttl"]).toBeUndefined(); expect(request.headers["cf-aig-cache-key"]).toBeUndefined(); expect(request.headers["cf-aig-metadata"]).toBeUndefined(); + expect(request.headers["cf-aig-byok-alias"]).toBeUndefined(); + }); + + it("does not set cf-aig-byok-alias on the binding path", async () => { + const config: AiGatewayAdapterConfig = { + binding: mockBinding, + byokAlias: "development", + }; + const fetcher = createGatewayFetch("openai", config); + + await fetcher("https://api.openai.com/v1/chat/completions", { + method: "POST", + body: JSON.stringify({}), + }); + + const request = mockBinding.run.mock.calls[0]![0]; + expect(request.headers["cf-aig-byok-alias"]).toBeUndefined(); }); });