diff --git a/src/api/providers/__tests__/xai.spec.ts b/src/api/providers/__tests__/xai.spec.ts index e8f2f547a5..ce345de291 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 @@ -565,6 +583,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), }) }) @@ -912,4 +932,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 e3c50adb33..04fa2dd4fe 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,16 @@ import { isMcpTool } from "../../utils/mcp-name" const XAI_DEFAULT_TEMPERATURE = 0 +// 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 { protected options: ApiHandlerOptions private client: OpenAI private readonly providerName = "xAI" + private readonly sessionId: string = uuidv7() constructor(options: ApiHandlerOptions) { super() @@ -45,6 +52,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 @@ -169,6 +208,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 + } + // Model params are always resolved by getModel(); send them unconditionally. requestBody.max_output_tokens = model.maxTokens requestBody.temperature = model.temperature @@ -255,11 +300,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, )