diff --git a/.changeset/cloudflare-binding-headers.md b/.changeset/cloudflare-binding-headers.md new file mode 100644 index 0000000000..39efe98bda --- /dev/null +++ b/.changeset/cloudflare-binding-headers.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-cloudflare': minor +--- + +Add `defaultHeaders` to the binding config. Every binding run sends them as the Workers AI `extraHeaders` option, next to `gateway`, so you can set `x-session-affinity` for prompt caching when you use `{ binding: env.AI }`. diff --git a/docs/adapters/cloudflare.md b/docs/adapters/cloudflare.md index 7f0fc20247..abd6117f6a 100644 --- a/docs/adapters/cloudflare.md +++ b/docs/adapters/cloudflare.md @@ -166,6 +166,41 @@ For Vue, Solid, Svelte, Angular, and the other frameworks, see [Quick Start](../ You now have a streaming chat on Workers AI. [Tools](../tools/tools) and [Structured Outputs](../structured-outputs/overview) work the same as with every adapter. +## Reuse the prompt cache + +Every chat turn sends the whole conversation again. Workers AI can serve the repeated start of the prompt from its cache, but only when the turns of one conversation reach the same model instance. Send an `x-session-affinity` header with one value per conversation to keep them together. See [Workers AI prompt caching](https://developers.cloudflare.com/workers-ai/features/prompt-caching/). + +Set the header with `defaultHeaders`: + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { createCloudflareText } from "@tanstack/ai-cloudflare"; +import type { Ai } from "@cloudflare/workers-types"; + +interface Env { + AI: Ai; +} + +export default { + async fetch(request: Request, env: Env) { + const { messages, conversationId } = await request.json(); + + const stream = chat({ + adapter: createCloudflareText("@cf/zai-org/glm-5.3-flash", { + binding: env.AI, + defaultHeaders: { "x-session-affinity": `ses_${conversationId}` }, + }), + messages, + threadId: conversationId, + }); + + return toServerSentEventsResponse(stream); + }, +}; +``` + +The binding sends `defaultHeaders` as the `extraHeaders` run option, next to `gateway`. The REST config takes the same `defaultHeaders` field. + ## Route through AI Gateway AI Gateway gives you caching, logs, rate limits, and one bill for all providers. To send requests through it, add `gateway` to any config. Use `"default"` for the gateway that Cloudflare makes for your account: diff --git a/packages/ai-cloudflare/src/adapters/text.ts b/packages/ai-cloudflare/src/adapters/text.ts index ad60711b8c..3f74ee854e 100644 --- a/packages/ai-cloudflare/src/adapters/text.ts +++ b/packages/ai-cloudflare/src/adapters/text.ts @@ -61,7 +61,7 @@ function clientOptions(config: CloudflareTextConfig): ClientOptions { return { // The binding authenticates by itself; the SDK only requires a value. apiKey: 'cloudflare-binding', - fetch: createBindingFetch(config.binding, config.gateway), + fetch: createBindingFetch(config), } } const { diff --git a/packages/ai-cloudflare/src/utils/config.ts b/packages/ai-cloudflare/src/utils/config.ts index 48fffb36fe..b7efc96be7 100644 --- a/packages/ai-cloudflare/src/utils/config.ts +++ b/packages/ai-cloudflare/src/utils/config.ts @@ -20,6 +20,11 @@ export type FetchLike = NonNullable export interface CloudflareBindingConfig { binding: Ai gateway?: CloudflareGatewayOptions + /** + * Workers AI request headers, such as `x-session-affinity` for prompt + * caching. Sent on every run as the binding's `extraHeaders` option. + */ + defaultHeaders?: Record accountId?: never apiKey?: never } diff --git a/packages/ai-cloudflare/src/utils/fetch.ts b/packages/ai-cloudflare/src/utils/fetch.ts index 48eec5f9c1..4c06d04a67 100644 --- a/packages/ai-cloudflare/src/utils/fetch.ts +++ b/packages/ai-cloudflare/src/utils/fetch.ts @@ -1,5 +1,9 @@ import type { Ai } from '@cloudflare/workers-types' -import type { CloudflareGatewayOptions, FetchLike } from './config' +import type { + CloudflareBindingConfig, + CloudflareGatewayOptions, + FetchLike, +} from './config' /** * Workers AI streams end with a usage-only trailer (`{"response":"","usage": @@ -107,11 +111,15 @@ export async function normalizeResponse(response: Response): Promise { * Makes `env.AI` look like an OpenAI-compatible HTTP endpoint to the OpenAI * SDK: the JSON request body becomes `binding.run(model, inputs)` and the * raw inference `Response` (OpenAI-format JSON or SSE) is handed back. + * + * The SDK's own request headers (auth, user agent, `x-stainless-*`) are + * dropped; only the configured `defaultHeaders` reach Workers AI. */ -export function createBindingFetch( - binding: Ai, - gateway?: CloudflareGatewayOptions, -): FetchLike { +export function createBindingFetch({ + binding, + gateway, + defaultHeaders, +}: CloudflareBindingConfig): FetchLike { // `Ai` is typed against the bundled model catalog; the adapter accepts any // model id, so widen the binding to the open catalog shape for this call. const run = binding.run.bind(binding) as ( @@ -126,6 +134,7 @@ export function createBindingFetch( const response = (await run(model, inputs, { returnRawResponse: true, ...(gateway && { gateway }), + ...(defaultHeaders && { extraHeaders: defaultHeaders }), })) as Response return await normalizeResponse(response) } diff --git a/packages/ai-cloudflare/src/utils/run.ts b/packages/ai-cloudflare/src/utils/run.ts index c0b40c0d1f..efa9102ae6 100644 --- a/packages/ai-cloudflare/src/utils/run.ts +++ b/packages/ai-cloudflare/src/utils/run.ts @@ -41,10 +41,16 @@ export async function runModel( }, } : inputs + const { gateway, defaultHeaders } = config return await run( model, bindingInputs, - config.gateway ? { gateway: config.gateway } : undefined, + gateway || defaultHeaders + ? { + ...(gateway && { gateway }), + ...(defaultHeaders && { extraHeaders: defaultHeaders }), + } + : undefined, ) } diff --git a/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts b/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts index 127473f6ba..633a997eae 100644 --- a/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts +++ b/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts @@ -169,6 +169,32 @@ describe('text adapter (binding)', () => { expect(finished.usage).toMatchObject({ promptTokens: 41, totalTokens: 45 }) }) + it('sends defaultHeaders as extraHeaders next to the gateway', async () => { + const { binding, run } = fakeBinding(() => sse(chatStreamEvents)) + const adapter = createCloudflareText('@cf/test', { + binding, + gateway: { id: 'g1' }, + defaultHeaders: { 'x-session-affinity': 'ses_abc' }, + }) + + await collect( + adapter.chatStream({ + threadId: 'thread-1', + model: '@cf/test', + messages: [{ role: 'user', content: 'Hi' }], + logger, + }), + ) + + // Only the configured headers: none of the OpenAI SDK's own headers + // (authorization, user-agent, x-stainless-*) reach the binding. + expect(run.mock.calls[0]![2]).toEqual({ + returnRawResponse: true, + gateway: { id: 'g1' }, + extraHeaders: { 'x-session-affinity': 'ses_abc' }, + }) + }) + const plainSse = sse(chatStreamEvents) it.each([ [ @@ -412,6 +438,22 @@ describe('embedding adapter', () => { ]) }) + it('sends defaultHeaders as extraHeaders', async () => { + const { binding, run } = fakeBinding(() => ({ data: [[0.1]] })) + const adapter = createCloudflareEmbedding('@cf/baai/bge-m3', { + binding, + defaultHeaders: { 'x-session-affinity': 'ses_abc' }, + }) + await adapter.createEmbeddings({ + model: '@cf/baai/bge-m3', + input: ['a'], + logger, + }) + expect(run.mock.calls[0]![2]).toEqual({ + extraHeaders: { 'x-session-affinity': 'ses_abc' }, + }) + }) + it('rejects dimensions and a mismatched output shape', async () => { const { binding } = fakeBinding(() => ({ response: 'wrong shape' })) const adapter = createCloudflareEmbedding('@cf/baai/bge-m3', { binding }) diff --git a/testing/e2e/fixtures/chat/cloudflare-binding-headers-wire.json b/testing/e2e/fixtures/chat/cloudflare-binding-headers-wire.json new file mode 100644 index 0000000000..89fd1fe34e --- /dev/null +++ b/testing/e2e/fixtures/chat/cloudflare-binding-headers-wire.json @@ -0,0 +1,12 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[cloudflare-binding-headers-wire] affinity check" + }, + "response": { + "content": "Affinity checked." + } + } + ] +} diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index e5ec31ba17..7703bae630 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -78,6 +78,7 @@ import { Route as ApiCallOwnershipWireRouteImport } from './routes/api.call-owne import { Route as ApiChatRouteImport } from './routes/api.chat' import { Route as ApiChatCompletionsIncompleteStreamRouteImport } from './routes/api.chat-completions-incomplete-stream' import { Route as ApiClientToolResumeStoreRouteImport } from './routes/api.client-tool-resume-store' +import { Route as ApiCloudflareBindingHeadersWireRouteImport } from './routes/api.cloudflare-binding-headers-wire' import { Route as ApiCloudflareBindingWireRouteImport } from './routes/api.cloudflare-binding-wire' import { Route as ApiCodingToolsRouteImport } from './routes/api.coding-tools' import { Route as ApiCodingWorkspaceRouteImport } from './routes/api.coding-workspace' @@ -568,6 +569,12 @@ const ApiClientToolResumeStoreRoute = path: '/api/client-tool-resume-store', getParentRoute: () => rootRouteImport, } as any) +const ApiCloudflareBindingHeadersWireRoute = + ApiCloudflareBindingHeadersWireRouteImport.update({ + id: '/api/cloudflare-binding-headers-wire', + path: '/api/cloudflare-binding-headers-wire', + getParentRoute: () => rootRouteImport, + } as any) const ApiCloudflareBindingWireRoute = ApiCloudflareBindingWireRouteImport.update({ id: '/api/cloudflare-binding-wire', @@ -1297,6 +1304,7 @@ export interface FileRoutesByFullPath { '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute + '/api/cloudflare-binding-headers-wire': typeof ApiCloudflareBindingHeadersWireRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/coding-tools': typeof ApiCodingToolsRoute '/api/coding-workspace': typeof ApiCodingWorkspaceRoute @@ -1494,6 +1502,7 @@ export interface FileRoutesByTo { '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute + '/api/cloudflare-binding-headers-wire': typeof ApiCloudflareBindingHeadersWireRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/coding-tools': typeof ApiCodingToolsRoute '/api/coding-workspace': typeof ApiCodingWorkspaceRoute @@ -1692,6 +1701,7 @@ export interface FileRoutesById { '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute + '/api/cloudflare-binding-headers-wire': typeof ApiCloudflareBindingHeadersWireRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/coding-tools': typeof ApiCodingToolsRoute '/api/coding-workspace': typeof ApiCodingWorkspaceRoute @@ -1891,6 +1901,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/chat-completions-incomplete-stream' | '/api/client-tool-resume-store' + | '/api/cloudflare-binding-headers-wire' | '/api/cloudflare-binding-wire' | '/api/coding-tools' | '/api/coding-workspace' @@ -2088,6 +2099,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/chat-completions-incomplete-stream' | '/api/client-tool-resume-store' + | '/api/cloudflare-binding-headers-wire' | '/api/cloudflare-binding-wire' | '/api/coding-tools' | '/api/coding-workspace' @@ -2285,6 +2297,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/chat-completions-incomplete-stream' | '/api/client-tool-resume-store' + | '/api/cloudflare-binding-headers-wire' | '/api/cloudflare-binding-wire' | '/api/coding-tools' | '/api/coding-workspace' @@ -2483,6 +2496,7 @@ export interface RootRouteChildren { ApiChatRoute: typeof ApiChatRoute ApiChatCompletionsIncompleteStreamRoute: typeof ApiChatCompletionsIncompleteStreamRoute ApiClientToolResumeStoreRoute: typeof ApiClientToolResumeStoreRoute + ApiCloudflareBindingHeadersWireRoute: typeof ApiCloudflareBindingHeadersWireRoute ApiCloudflareBindingWireRoute: typeof ApiCloudflareBindingWireRoute ApiCodingToolsRoute: typeof ApiCodingToolsRoute ApiCodingWorkspaceRoute: typeof ApiCodingWorkspaceRoute @@ -3092,6 +3106,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiClientToolResumeStoreRouteImport parentRoute: typeof rootRouteImport } + '/api/cloudflare-binding-headers-wire': { + id: '/api/cloudflare-binding-headers-wire' + path: '/api/cloudflare-binding-headers-wire' + fullPath: '/api/cloudflare-binding-headers-wire' + preLoaderRoute: typeof ApiCloudflareBindingHeadersWireRouteImport + parentRoute: typeof rootRouteImport + } '/api/cloudflare-binding-wire': { id: '/api/cloudflare-binding-wire' path: '/api/cloudflare-binding-wire' @@ -4109,6 +4130,7 @@ const rootRouteChildren: RootRouteChildren = { ApiChatCompletionsIncompleteStreamRoute: ApiChatCompletionsIncompleteStreamRoute, ApiClientToolResumeStoreRoute: ApiClientToolResumeStoreRoute, + ApiCloudflareBindingHeadersWireRoute: ApiCloudflareBindingHeadersWireRoute, ApiCloudflareBindingWireRoute: ApiCloudflareBindingWireRoute, ApiCodingToolsRoute: ApiCodingToolsRoute, ApiCodingWorkspaceRoute: ApiCodingWorkspaceRoute, diff --git a/testing/e2e/src/routes/api.cloudflare-binding-headers-wire.ts b/testing/e2e/src/routes/api.cloudflare-binding-headers-wire.ts new file mode 100644 index 0000000000..cad4ef6ff9 --- /dev/null +++ b/testing/e2e/src/routes/api.cloudflare-binding-headers-wire.ts @@ -0,0 +1,67 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, createChatOptions } from '@tanstack/ai' +import { createCloudflareText } from '@tanstack/ai-cloudflare' +import type { CloudflareBindingConfig } from '@tanstack/ai-cloudflare' + +const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' + +/** + * Drives the Cloudflare text adapter in binding mode. The stand-in `env.AI` + * does what the real binding does with `returnRawResponse` and + * `extraHeaders`: it posts the inputs with those headers and returns the raw + * response. The companion spec reads aimock's request journal to check that + * `defaultHeaders` reach Workers AI and the OpenAI SDK's own headers do not. + */ +export const Route = createFileRoute('/api/cloudflare-binding-headers-wire')({ + server: { + handlers: { + POST: async ({ request }) => { + const testId = new URL(request.url).searchParams.get('testId') ?? '' + + const run = async ( + model: string, + inputs: Record, + options: { extraHeaders?: Record }, + ) => + await fetch(`${LLMOCK_DEFAULT_BASE}/v1/chat/completions`, { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-test-id': testId, + ...options.extraHeaders, + }, + body: JSON.stringify({ model, ...inputs }), + }) + // ponytail: only `run` is exercised; the real binding type has more. + const binding = { run } as unknown as CloudflareBindingConfig['binding'] + + const adapter = createCloudflareText('@cf/zai-org/glm-5.3-flash', { + binding, + defaultHeaders: { 'x-session-affinity': 'ses_e2e' }, + }) + + try { + for await (const _ of chat({ + ...createChatOptions({ adapter }), + threadId: 'cloudflare-binding-headers-wire', + messages: [ + { + role: 'user', + content: '[cloudflare-binding-headers-wire] affinity check', + }, + ], + })) { + // Drain the stream. + } + } catch (error) { + return Response.json({ + ok: false, + error: error instanceof Error ? error.message : String(error), + }) + } + + return Response.json({ ok: true }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/cloudflare-binding-headers-wire.spec.ts b/testing/e2e/tests/cloudflare-binding-headers-wire.spec.ts new file mode 100644 index 0000000000..7d26773cde --- /dev/null +++ b/testing/e2e/tests/cloudflare-binding-headers-wire.spec.ts @@ -0,0 +1,35 @@ +import { test, expect } from './fixtures' + +type JournalEntry = { + headers?: Record +} + +test.describe('cloudflare — binding defaultHeaders on the wire', () => { + test('sends x-session-affinity and drops OpenAI SDK headers', async ({ + request, + aimockPort, + testId, + }) => { + const response = await request.post( + `/api/cloudflare-binding-headers-wire?testId=${encodeURIComponent(testId)}`, + ) + expect(response.ok()).toBe(true) + expect(await response.json()).toEqual({ ok: true }) + + const journalResponse = await request.get( + `http://127.0.0.1:${aimockPort}/v1/_requests`, + ) + const entries = (await journalResponse.json()) as Array + const captured = entries.find( + (entry) => entry.headers?.['x-test-id'] === testId, + ) + expect(captured).toBeDefined() + expect(captured?.headers?.['x-session-affinity']).toBe('ses_e2e') + expect(captured?.headers).not.toHaveProperty('authorization') + expect( + Object.keys(captured?.headers ?? {}).filter((name) => + name.startsWith('x-stainless-'), + ), + ).toEqual([]) + }) +})