Skip to content

Commit a7a99d3

Browse files
committed
fix(xai): prompt-cache affinity — x-grok-conv-id from prompt_cache_key + grok reasoning_content replay
- resolveProviderTransport gains promptCacheKey: sha256-32hex x-grok-conv-id sticky-routing header in BOTH auth modes (docs.x.ai best-practices; CLIProxyAPI precedent). Case-insensitive user-header override, blank keys never emit. - 429 key-rotation paths re-resolve the transport so conv-id/subscription headers survive retries. - xai preset preserveReasoningContentModels for grok reasoning models: xAI documents dropped reasoning_content as the top cause of multi-turn cache misses (64M-token melt report, issue #112 adjacent). - devlog record kept locally (devlog/ is gitignored).
1 parent ff09370 commit a7a99d3

4 files changed

Lines changed: 156 additions & 7 deletions

File tree

src/providers/registry.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,11 @@ export const PROVIDER_REGISTRY: readonly ProviderRegistryEntry[] = [
236236
models: ["grok-4.5", "grok-4.3", "grok-4.20-multi-agent-0309", "grok-4.20-0309-reasoning", "grok-4.20-0309-non-reasoning", "grok-build-0.1", "grok-composer-2.5-fast"],
237237
defaultModel: "grok-4.5",
238238
noReasoningModels: ["grok-4.20-0309-non-reasoning", "grok-build-0.1", "grok-composer-2.5-fast"],
239+
// Replay assistant reasoning_content for grok reasoning models: xAI documents dropped
240+
// reasoning_content as the top cause of prompt-cache misses on multi-turn conversations
241+
// (docs.x.ai prompt-caching/multi-turn, verified 2026-07-13 — devlog/_plan/260713_grok_caching).
242+
// Models that never emit reasoning simply have no thinking parts to replay (no-op).
243+
preserveReasoningContentModels: ["grok-4.5", "grok-4.3", "grok-4.20-multi-agent-0309", "grok-4.20-0309-reasoning"],
239244
// grok-4.5 reasoning is always-on with low/medium/high control (no off tier upstream).
240245
modelReasoningEfforts: { "grok-4.5": ["low", "medium", "high"] },
241246
modelContextWindows: {

src/providers/xai-transport.ts

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { createHash } from "node:crypto";
12
import type { OcxProviderConfig } from "../types";
23

34
/**
@@ -18,21 +19,60 @@ const XAI_GROK_CLI_HEADERS: Readonly<Record<string, string>> = {
1819
"x-xai-token-auth": "xai-grok-cli",
1920
};
2021

22+
/**
23+
* Sticky-routing hint for xAI's automatic prefix cache. xAI routes requests
24+
* carrying the same `x-grok-conv-id` to the same server, which is where the
25+
* prompt cache lives (docs.x.ai prompt-caching best-practices; verified
26+
* 2026-07-13, devlog/_plan/260713_grok_caching). Codex clients send a stable
27+
* per-conversation `prompt_cache_key`; hash it so the raw session id never
28+
* leaves the proxy.
29+
*/
30+
export const XAI_CONV_ID_HEADER = "x-grok-conv-id";
31+
32+
function hasHeaderCaseInsensitive(headers: Record<string, string> | undefined, name: string): boolean {
33+
if (!headers) return false;
34+
const target = name.toLowerCase();
35+
return Object.keys(headers).some(key => key.toLowerCase() === target);
36+
}
37+
38+
export function deriveXaiConvId(promptCacheKey: string): string {
39+
return createHash("sha256").update(promptCacheKey).digest("hex").slice(0, 32);
40+
}
41+
2142
/**
2243
* Resolve the effective xAI transport without mutating persisted config.
2344
* User-provided headers are preserved and may advance the compatibility
2445
* version without waiting for an opencodex release.
46+
*
47+
* `promptCacheKey` (the client's stable conversation key) additionally pins
48+
* cache-affinity routing via `x-grok-conv-id` in BOTH auth modes. Blank or
49+
* whitespace-only keys are ignored so unrelated requests can never collapse
50+
* onto one shared conv id, and any user-configured header (any case) wins.
2551
*/
2652
export function resolveProviderTransport(
2753
providerName: string,
2854
provider: OcxProviderConfig,
55+
promptCacheKey?: string,
2956
): OcxProviderConfig {
30-
if (providerName !== "xai" || provider.authMode !== "oauth") return provider;
57+
if (providerName !== "xai") return provider;
58+
const cacheKey = promptCacheKey?.trim();
59+
const convIdHeaders: Record<string, string> =
60+
cacheKey && !hasHeaderCaseInsensitive(provider.headers, XAI_CONV_ID_HEADER)
61+
? { [XAI_CONV_ID_HEADER]: deriveXaiConvId(cacheKey) }
62+
: {};
63+
if (provider.authMode !== "oauth") {
64+
if (Object.keys(convIdHeaders).length === 0) return provider;
65+
return {
66+
...provider,
67+
headers: { ...convIdHeaders, ...(provider.headers ?? {}) },
68+
};
69+
}
3170
return {
3271
...provider,
3372
baseUrl: XAI_GROK_CLI_BASE_URL,
3473
headers: {
3574
...XAI_GROK_CLI_HEADERS,
75+
...convIdHeaders,
3676
...(provider.headers ?? {}),
3777
},
3878
};

src/server/responses.ts

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -613,7 +613,7 @@ export async function handleResponses(
613613
return formatErrorResponse(401, "authentication_error", err instanceof Error ? err.message : String(err));
614614
}
615615
}
616-
route.provider = resolveProviderTransport(route.providerName, route.provider);
616+
route.provider = resolveProviderTransport(route.providerName, route.provider, parsed.options.promptCacheKey);
617617

618618
// Vision sidecar: the routed model can't see images (provider.noVisionModels). Describe each
619619
// attached image through the selected sidecar backend and replace it with text BEFORE the main
@@ -888,9 +888,11 @@ export async function handleResponses(
888888
on429: retryAfter => {
889889
const rotated = rotateKeyOn429(config, route.providerName, retryAfter, Date.now(), route.provider.apiKey);
890890
if (!rotated) return null;
891-
route.provider = rotated;
891+
// Re-resolve the auth-mode transport so the conv-id / subscription headers derived at
892+
// line ~616 survive the key rotation (rotated providers come from raw config).
893+
route.provider = resolveProviderTransport(route.providerName, rotated, parsed.options.promptCacheKey);
892894
return resolveAdapter(
893-
resolveWireProtocolOverride(route.providerName, route.modelId, rotated),
895+
resolveWireProtocolOverride(route.providerName, route.modelId, route.provider),
894896
config.cacheRetention,
895897
);
896898
},
@@ -944,9 +946,11 @@ export async function handleResponses(
944946
// Release the failed response's socket before retrying; unread bodies otherwise linger
945947
// until runtime cleanup (one per rotated key under a rate-limit storm).
946948
try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ }
947-
route.provider = rotated;
949+
// Same transport re-resolution as the streaming on429 path: keep conv-id + subscription
950+
// headers on the retried request instead of silently reverting to the raw config provider.
951+
route.provider = resolveProviderTransport(route.providerName, rotated, parsed.options.promptCacheKey);
948952
const retryAdapter = resolveAdapter(
949-
resolveWireProtocolOverride(route.providerName, route.modelId, rotated),
953+
resolveWireProtocolOverride(route.providerName, route.modelId, route.provider),
950954
config.cacheRetention,
951955
);
952956
const retryRequest = await retryAdapter.buildRequest(parsed, { headers: selectedForwardHeaders });

tests/xai-transport.test.ts

Lines changed: 101 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,13 @@ import { createOpenAIChatAdapter } from "../src/adapters/openai-chat";
33
import { buildModelsRequest } from "../src/oauth";
44
import {
55
resolveProviderTransport,
6+
deriveXaiConvId,
7+
XAI_CONV_ID_HEADER,
68
XAI_GROK_CLI_BASE_URL,
79
XAI_GROK_CLIENT_VERSION,
810
} from "../src/providers/xai-transport";
9-
import type { OcxParsedRequest, OcxProviderConfig } from "../src/types";
11+
import { getProviderRegistryEntry } from "../src/providers/registry";
12+
import type { OcxAssistantMessage, OcxParsedRequest, OcxProviderConfig } from "../src/types";
1013

1114
function provider(authMode: "oauth" | "key"): OcxProviderConfig {
1215
return {
@@ -83,3 +86,100 @@ describe("xAI auth-mode transport selection", () => {
8386
});
8487
});
8588
});
89+
90+
describe("xAI prompt-cache conv-id affinity", () => {
91+
test("promptCacheKey derives a stable hashed x-grok-conv-id in oauth mode", () => {
92+
const effective = resolveProviderTransport("xai", provider("oauth"), "codex-session-abc");
93+
const again = resolveProviderTransport("xai", provider("oauth"), "codex-session-abc");
94+
95+
expect(effective.headers?.[XAI_CONV_ID_HEADER]).toBe(deriveXaiConvId("codex-session-abc"));
96+
expect(effective.headers?.[XAI_CONV_ID_HEADER]).toMatch(/^[0-9a-f]{32}$/);
97+
// Stable across requests (cache affinity) and never the raw session id.
98+
expect(again.headers?.[XAI_CONV_ID_HEADER]).toBe(effective.headers?.[XAI_CONV_ID_HEADER]);
99+
expect(effective.headers?.[XAI_CONV_ID_HEADER]).not.toContain("codex-session-abc");
100+
});
101+
102+
test("key mode gains conv-id affinity without touching baseUrl or CLI headers", () => {
103+
const configured = provider("key");
104+
const effective = resolveProviderTransport("xai", configured, "codex-session-abc");
105+
const request = createOpenAIChatAdapter(effective).buildRequest(parsed());
106+
107+
expect(effective.baseUrl).toBe("https://api.x.ai/v1");
108+
expect(request.url).toBe("https://api.x.ai/v1/chat/completions");
109+
expect(effective.headers?.[XAI_CONV_ID_HEADER]).toBe(deriveXaiConvId("codex-session-abc"));
110+
expect(effective.headers?.["x-grok-client-identifier"]).toBeUndefined();
111+
expect(effective.headers?.["x-xai-token-auth"]).toBeUndefined();
112+
});
113+
114+
test("missing, empty, and whitespace-only cache keys never emit a conv-id", () => {
115+
const noKeyOauth = resolveProviderTransport("xai", provider("oauth"));
116+
const emptyKeyOauth = resolveProviderTransport("xai", provider("oauth"), "");
117+
const blankKeyOauth = resolveProviderTransport("xai", provider("oauth"), " ");
118+
const configuredKey = provider("key");
119+
const emptyKeyApi = resolveProviderTransport("xai", configuredKey, "");
120+
121+
expect(noKeyOauth.headers?.[XAI_CONV_ID_HEADER]).toBeUndefined();
122+
expect(emptyKeyOauth.headers?.[XAI_CONV_ID_HEADER]).toBeUndefined();
123+
expect(blankKeyOauth.headers?.[XAI_CONV_ID_HEADER]).toBeUndefined();
124+
// Key mode without a conv-id stays the exact configured object (no clone churn).
125+
expect(emptyKeyApi).toBe(configuredKey);
126+
});
127+
128+
test("user-configured conv-id header wins in any casing (no duplicate header pair)", () => {
129+
const lower = provider("oauth");
130+
lower.headers = { [XAI_CONV_ID_HEADER]: "user-pinned" };
131+
const mixed = provider("key");
132+
mixed.headers = { "X-Grok-Conv-Id": "user-pinned-mixed" };
133+
134+
const lowerResolved = resolveProviderTransport("xai", lower, "codex-session-abc");
135+
const mixedResolved = resolveProviderTransport("xai", mixed, "codex-session-abc");
136+
137+
expect(lowerResolved.headers?.[XAI_CONV_ID_HEADER]).toBe("user-pinned");
138+
// Mixed casing: the generated lowercase header must be suppressed entirely.
139+
expect(mixedResolved.headers?.["X-Grok-Conv-Id"]).toBe("user-pinned-mixed");
140+
expect(mixedResolved.headers?.[XAI_CONV_ID_HEADER]).toBeUndefined();
141+
const convIdKeys = Object.keys(mixedResolved.headers ?? {}).filter(k => k.toLowerCase() === XAI_CONV_ID_HEADER);
142+
expect(convIdKeys).toHaveLength(1);
143+
});
144+
});
145+
146+
describe("xAI reasoning_content cache preservation", () => {
147+
test("registry preset replays reasoning_content for grok reasoning models only", () => {
148+
const entry = getProviderRegistryEntry("xai");
149+
expect(entry?.preserveReasoningContentModels).toEqual([
150+
"grok-4.5",
151+
"grok-4.3",
152+
"grok-4.20-multi-agent-0309",
153+
"grok-4.20-0309-reasoning",
154+
]);
155+
for (const noReasoning of entry?.noReasoningModels ?? []) {
156+
expect(entry?.preserveReasoningContentModels).not.toContain(noReasoning);
157+
}
158+
});
159+
160+
test("assistant thinking parts round-trip as reasoning_content on grok-4.5 history", () => {
161+
// xAI docs: dropped reasoning_content is the top cause of multi-turn cache misses
162+
// (docs.x.ai prompt-caching/multi-turn, 2026-07-13).
163+
const prov: OcxProviderConfig = {
164+
...provider("oauth"),
165+
preserveReasoningContentModels: getProviderRegistryEntry("xai")?.preserveReasoningContentModels ?? [],
166+
};
167+
const assistant: OcxAssistantMessage = {
168+
role: "assistant",
169+
content: [
170+
{ type: "thinking", thinking: "cached chain" },
171+
{ type: "text", text: "answer" },
172+
],
173+
timestamp: 0,
174+
};
175+
const req: OcxParsedRequest = {
176+
modelId: "grok-4.5",
177+
context: { messages: [{ role: "user", content: "q1", timestamp: 0 }, assistant, { role: "user", content: "q2", timestamp: 0 }] },
178+
stream: false,
179+
options: {},
180+
};
181+
const body = JSON.parse(createOpenAIChatAdapter(prov).buildRequest(req).body as string) as { messages: Array<Record<string, unknown>> };
182+
const replayed = body.messages.find(m => m.role === "assistant");
183+
expect(replayed?.reasoning_content).toBe("cached chain");
184+
});
185+
});

0 commit comments

Comments
 (0)