From 44dd4b25086380b7870ffa87a62943f15e52c0a8 Mon Sep 17 00:00:00 2001 From: myk1yt Date: Sat, 31 Oct 2026 01:16:23 +0900 Subject: [PATCH 1/2] fix(xai): send body prompt_cache_key for Responses cache affinity What: add a task-scoped `prompt_cache_key` to the xAI Responses API request body in both `createMessage` and `completePrompt`, gated on the model's `supportsPromptCache` flag and clamped to 64 characters. Why: xAI's prompt cache is stored per server, so every turn of one conversation must be routed to the same server to reuse it. For the Responses API the documented affinity lever is the body field `prompt_cache_key`; the `x-grok-conv-id` header is the Chat Completions / gRPC equivalent and is not read on this path, so it was deliberately NOT added. Before this change `xai.ts` sent no affinity identifier at all, leaving server-side cache routing to load-balancer luck. Impact: xAI turns of one task now pin to one cache server, raising cache-read tokens and lowering cost on long conversations. The key is omitted (never sent as an empty string) when no identifier is available or the model does not support prompt caching. Usage accounting is unchanged: the shared normalizer already reads `input_tokens_details.cached_tokens`. --- src/api/providers/__tests__/xai.spec.ts | 146 ++++++++++++++++++++++++ src/api/providers/xai.ts | 48 ++++++++ 2 files changed, 194 insertions(+) diff --git a/src/api/providers/__tests__/xai.spec.ts b/src/api/providers/__tests__/xai.spec.ts index 7d66f6985d..15e400ab87 100644 --- a/src/api/providers/__tests__/xai.spec.ts +++ b/src/api/providers/__tests__/xai.spec.ts @@ -33,6 +33,24 @@ import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { makeCreateMessageMetadata } from "../../../test-utils/api" import { clearAllMocks } from "../../../test-utils/reset" +/** + * Flip only `supportsPromptCache` on the handler's resolved model, so the rest of + * the catalog entry (maxTokens, temperature, reasoning) stays as the real one. + * + * The xAI catalog is declared `as const`, so every entry's `supportsPromptCache` + * is the literal `true` and no real uncached model exists to point at. The double + * assertion is confined to this test double and is the only way to reach the + * `supportsPromptCache === false` arm; the production code is untouched by it. + */ +function stubSupportsPromptCache(handler: XAIHandler, supportsPromptCache: boolean): void { + const realModel = handler.getModel() + vitest + .spyOn(handler, "getModel") + .mockReturnValue({ ...realModel, info: { ...realModel.info, supportsPromptCache } } as unknown as ReturnType< + XAIHandler["getModel"] + >) +} + describe("XAIHandler", () => { let handler: XAIHandler @@ -372,6 +390,8 @@ describe("XAIHandler", () => { model: xaiDefaultModelId, input: [{ role: "user", content: [{ type: "input_text", text: "test prompt" }] }], store: false, + // No task metadata exists on this path, so the handler session scopes the cache. + prompt_cache_key: expect.any(String), }, undefined, ) @@ -722,4 +742,130 @@ describe("XAIHandler", () => { expect(removeEventListenerSpy).toHaveBeenCalledTimes(1) expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", listener) }) + + describe("prompt cache affinity", () => { + const lastRequestBody = () => mockResponsesCreate.mock.calls[mockResponsesCreate.mock.calls.length - 1][0] + + it("should send the task id as the body prompt_cache_key", async () => { + mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([])) + + const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId: "task-abc" })) + await stream.next() + + expect(lastRequestBody().prompt_cache_key).toBe("task-abc") + }) + + it("should keep the key stable across two turns of the same task", async () => { + mockResponsesCreate.mockResolvedValue(asyncStreamFrom([])) + + const first = handler.createMessage("first turn", [], makeCreateMessageMetadata({ taskId: "task-stable" })) + await first.next() + const second = handler.createMessage( + "second turn", + [], + makeCreateMessageMetadata({ taskId: "task-stable" }), + ) + await second.next() + + // Two separate requests, one value: routing is pinned to one cache server. + expect(mockResponsesCreate).toHaveBeenCalledTimes(2) + expect(mockResponsesCreate.mock.calls[0][0].prompt_cache_key).toBe("task-stable") + expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe( + mockResponsesCreate.mock.calls[0][0].prompt_cache_key, + ) + }) + + it("should give distinct keys to distinct tasks", async () => { + mockResponsesCreate.mockResolvedValue(asyncStreamFrom([])) + + for (const taskId of ["task-one", "task-two"]) { + const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId })) + await stream.next() + } + + expect(mockResponsesCreate.mock.calls[0][0].prompt_cache_key).toBe("task-one") + expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe("task-two") + }) + + it("should fall back to the handler session id when no task id is supplied", async () => { + mockResponsesCreate.mockResolvedValue(asyncStreamFrom([])) + + const stream = handler.createMessage("test prompt", []) + await stream.next() + + const key = lastRequestBody().prompt_cache_key + expect(typeof key).toBe("string") + expect(key.length).toBeGreaterThan(0) + + // The same handler keeps the same fallback key on the next turn. + const second = handler.createMessage("test prompt", []) + await second.next() + expect(lastRequestBody().prompt_cache_key).toBe(key) + }) + + it("should clamp an over-length task id to the 64-character limit", async () => { + mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([])) + + const overLongTaskId = "t".repeat(80) + const stream = handler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ taskId: overLongTaskId }), + ) + await stream.next() + + const key = lastRequestBody().prompt_cache_key + expect(key).toHaveLength(64) + expect(key).toBe(overLongTaskId.slice(0, 64)) + }) + + it("should never send an empty-string prompt_cache_key", async () => { + mockResponsesCreate.mockResolvedValue(asyncStreamFrom([])) + + const stream = handler.createMessage("test prompt", [], makeCreateMessageMetadata({ taskId: "" })) + await stream.next() + + const key = lastRequestBody().prompt_cache_key + // An empty task id falls back to the session id instead of sending "". + expect(key).not.toBe("") + expect(typeof key).toBe("string") + }) + + it("should omit the key when the model does not support prompt caching", async () => { + const uncachedHandler = new XAIHandler({ apiModelId: xaiDefaultModelId }) + stubSupportsPromptCache(uncachedHandler, false) + + mockResponsesCreate.mockResolvedValueOnce(asyncStreamFrom([])) + const stream = uncachedHandler.createMessage( + "test prompt", + [], + makeCreateMessageMetadata({ taskId: "task-x" }), + ) + await stream.next() + + expect(lastRequestBody()).not.toHaveProperty("prompt_cache_key") + }) + + it("completePrompt should scope the prompt_cache_key to the handler session", async () => { + mockResponsesCreate.mockResolvedValue({ output_text: "response" }) + + await handler.completePrompt("first prompt") + await handler.completePrompt("second prompt") + + const firstKey = mockResponsesCreate.mock.calls[0][0].prompt_cache_key + expect(typeof firstKey).toBe("string") + expect(firstKey.length).toBeGreaterThan(0) + expect(mockResponsesCreate.mock.calls[1][0].prompt_cache_key).toBe(firstKey) + }) + + it("completePrompt should omit the key when the model does not support prompt caching", async () => { + const uncachedHandler = new XAIHandler({ apiModelId: xaiDefaultModelId }) + stubSupportsPromptCache(uncachedHandler, false) + + mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" }) + await uncachedHandler.completePrompt("test prompt") + + expect(mockResponsesCreate.mock.calls[0][0]).not.toHaveProperty("prompt_cache_key") + }) + }) }) diff --git a/src/api/providers/xai.ts b/src/api/providers/xai.ts index f09f3df252..fda3241437 100644 --- a/src/api/providers/xai.ts +++ b/src/api/providers/xai.ts @@ -1,5 +1,6 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" +import { v7 as uuidv7 } from "uuid" import { type XAIModelId, xaiDefaultModelId, xaiModels, ApiProviderError } from "@roo-code/types" import { TelemetryService } from "@roo-code/telemetry" @@ -26,10 +27,15 @@ import { isMcpTool } from "../../utils/mcp-name" const XAI_DEFAULT_TEMPERATURE = 0 +// xAI caps `prompt_cache_key` at 64 characters, the same bound the OpenAI Responses +// API enforces on the field. Task ids are uuid v7 (36 chars), so this is a guard. +const XAI_MAX_PROMPT_CACHE_KEY_LENGTH = 64 + export class XAIHandler extends BaseProvider implements SingleCompletionHandler { protected options: ApiHandlerOptions private client: OpenAI private readonly providerName = "xAI" + private readonly sessionId: string = uuidv7() constructor(options: ApiHandlerOptions) { super() @@ -45,6 +51,38 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler }) } + /** + * Cache routing affinity key for the Responses API. + * + * xAI's prompt cache lives on individual servers, so every turn of one + * conversation has to reach the same server to reuse it. For the Responses API + * the documented lever is the body field `prompt_cache_key`; the + * `x-grok-conv-id` header is the Chat Completions / gRPC equivalent and is not + * read on this path. + * + * The task id is the longest-lived identifier available here: `Task.taskId` is + * assigned once and restored from history on resume, so it survives restarts + * and provider switches, which a per-handler session id does not. The handler + * session id is the fallback for callers that send no task id, and is the sole + * key for `completePrompt`, which carries no task metadata at all. + * + * Returns undefined (never an empty string) when there is nothing to key on or + * when the model does not support prompt caching, so the field is omitted + * rather than sent blank. + */ + private getPromptCacheKey(model: ReturnType, taskId?: string): string | undefined { + if (!model.info.supportsPromptCache) { + return undefined + } + + const key = taskId || this.sessionId + if (!key) { + return undefined + } + + return key.length > XAI_MAX_PROMPT_CACHE_KEY_LENGTH ? key.slice(0, XAI_MAX_PROMPT_CACHE_KEY_LENGTH) : key + } + override getModel() { const id = this.options.apiModelId && this.options.apiModelId in xaiModels @@ -121,6 +159,12 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler include: ["reasoning.encrypted_content"], } + // Route every turn of this task to the server holding its cache. + const promptCacheKey = this.getPromptCacheKey(model, metadata?.taskId) + if (promptCacheKey) { + requestBody.prompt_cache_key = promptCacheKey + } + if (model.maxTokens) { requestBody.max_output_tokens = model.maxTokens } @@ -202,11 +246,15 @@ export class XAIHandler extends BaseProvider implements SingleCompletionHandler requestOptions.timeout = toSdkTimeoutMs(timeoutMs) } + const promptCacheKey = this.getPromptCacheKey(model) const response = await this.client.responses.create( { model: model.id, input: [{ role: "user", content: [{ type: "input_text", text: prompt }] }], store: false, + // A completion is its own cache scope: no task metadata exists on + // this path, so the handler session id is the whole key. + ...(promptCacheKey ? { prompt_cache_key: promptCacheKey } : {}), }, Object.keys(requestOptions).length > 0 ? requestOptions : undefined, ) From 9b641cf1d344a13306e24dbf9ddbb01e1a5f51f9 Mon Sep 17 00:00:00 2001 From: myk1yt Date: Sat, 31 Oct 2026 01:45:53 +0900 Subject: [PATCH 2/2] docs(api): attribute the 64-char prompt_cache_key cap honestly What: reword the XAI_MAX_PROMPT_CACHE_KEY_LENGTH comment. Why: review seat (PR #2256 mlarge P3) caught that xAI's docs do not state a 64-char bound for prompt_cache_key - the bound is OpenAI's documented one; we apply it as a conservative guard. Impact: comment-only; zero code change. --- src/api/providers/xai.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/api/providers/xai.ts b/src/api/providers/xai.ts index fda3241437..93f3e0b51f 100644 --- a/src/api/providers/xai.ts +++ b/src/api/providers/xai.ts @@ -27,8 +27,9 @@ import { isMcpTool } from "../../utils/mcp-name" const XAI_DEFAULT_TEMPERATURE = 0 -// xAI caps `prompt_cache_key` at 64 characters, the same bound the OpenAI Responses -// API enforces on the field. Task ids are uuid v7 (36 chars), so this is a guard. +// Conservative length guard for `prompt_cache_key`: OpenAI's documented cap is 64 +// chars; xAI does not document a bound for this field, so we apply the same cap. +// Task ids are uuid v7 (36 chars), so this never truncates in practice. const XAI_MAX_PROMPT_CACHE_KEY_LENGTH = 64 export class XAIHandler extends BaseProvider implements SingleCompletionHandler {