feat(ai-cloudflare): send defaultHeaders as extraHeaders on binding runs

.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 }`.

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:

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 {

packages/ai-cloudflare/src/utils/config.ts

@@ -20,6 +20,11 @@ export type FetchLike = NonNullable<ClientOptions['fetch']>
 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<string, string>
   accountId?: never
   apiKey?: never
 }

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<Response> {
  * 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)
   }

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,
     )
   }
 

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 })

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."
+      }
+    }
+  ]
+}

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,

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<string, unknown>,
+          options: { extraHeaders?: Record<string, string> },
+        ) =>
+          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 })
+      },
+    },
+  },
+})

testing/e2e/tests/cloudflare-binding-headers-wire.spec.ts

@@ -0,0 +1,35 @@
+import { test, expect } from './fixtures'
+
+type JournalEntry = {
+  headers?: Record<string, string>
+}
+
+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<JournalEntry>
+    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([])
+  })
+})