Merge pull request #1866 from uta-a/fix/codex-session-affinity

fix(core): send session_id header for Codex Responses prompt cache affinity
2026-10-31 03:52:16 +08:00 · 2026-10-09 11:55:08 +08:00
parent 861c09f91f b6f4d6a9fc
commit 4b31d1d6cd
2 changed files with 231 additions and 6 deletions
@@ -25,6 +25,9 @@ export type ResponsesSessionAffinityInput = {

 const sessionIdHeaderNames = ["x-claude-code-session-id", "x-claude-session-id"];
 const codexUpstreamUrlMarkers = ["chatgpt.com/backend-api/codex", "/backend-api/codex"];
+const codexSessionHeaderName = "session_id";
+const codexSessionIdMaxLength = 256;
+const metadataUserIdSessionPattern = /_session_([^_\s]+)$/;

 /**
 * Copies the Claude Code session identity onto outbound OpenAI Responses
@@ -33,8 +36,8 @@ const codexUpstreamUrlMarkers = ["chatgpt.com/backend-api/codex", "/backend-api/
 * on body fields hash each turn onto a different channel and the next hop
 * rejects channel-bound `encrypted_content` continuations. A caller-supplied
 * non-empty `prompt_cache_key` always wins; other protocols and non-JSON
- * bodies pass through untouched. Codex upstreams reject unknown body fields
- * with `400 Unsupported parameter`, so affinity is skipped for them.
+ * bodies pass through untouched. Codex upstreams are handled separately by
+ * `applyCodexSessionAffinity`.
 */
 export function applyResponsesSessionAffinity(input: ResponsesSessionAffinityInput): UpstreamRequest {
  const upstreamRequest = input.upstreamRequest;
@@ -43,7 +46,7 @@ export function applyResponsesSessionAffinity(input: ResponsesSessionAffinityInp
    return upstreamRequest;
  }
  if (isCodexResponsesUpstream(upstreamRequest.url, input.targetProviderConfig)) {
-    return upstreamRequest;
+    return applyCodexSessionAffinity(input);
  }
  const body = upstreamRequest.body;
  if ((upstreamRequest.bodyEncoding ?? "json") !== "json" || !isRecord(body)) {
@@ -70,6 +73,84 @@ export function applyResponsesSessionAffinity(input: ResponsesSessionAffinityInp
  };
 }

+/**
+ * The Codex backend routes prompt caching by the `session_id` request header
+ * (as the Codex CLI sends it); a body `prompt_cache_key` alone does not hit
+ * the cache. It may reject `metadata`, so only the header and
+ * `prompt_cache_key` are added. An existing `session_id`/`session-id` header
+ * and a caller-supplied `prompt_cache_key` always win. The header is added
+ * even for non-JSON bodies.
+ */
+function applyCodexSessionAffinity(input: ResponsesSessionAffinityInput): UpstreamRequest {
+  const upstreamRequest = input.upstreamRequest;
+  const sessionId = resolveCodexSessionId(input.request?.headers, inboundMetadataUserId(input.request?.body));
+  if (!sessionId) {
+    return upstreamRequest;
+  }
+
+  let headers = upstreamRequest.headers;
+  if (!hasCodexSessionHeader(headers)) {
+    headers = { ...headers, [codexSessionHeaderName]: sessionId };
+  }
+  let body = upstreamRequest.body;
+  if ((upstreamRequest.bodyEncoding ?? "json") === "json" && isRecord(body) && !stringValue(body.prompt_cache_key)) {
+    body = { ...body, prompt_cache_key: sessionId };
+  }
+  if (headers === upstreamRequest.headers && body === upstreamRequest.body) {
+    return upstreamRequest;
+  }
+  return { ...upstreamRequest, body, headers };
+}
+
+/**
+ * Codex session id: the Claude Code session header, else the session id
+ * embedded in the inbound `metadata.user_id` (a JSON string with a
+ * `session_id` field, or the legacy `..._session_<id>` form). Values that are
+ * unsafe as an HTTP header value are dropped.
+ */
+function resolveCodexSessionId(
+  headers: Record<string, HeaderValue> | undefined,
+  inboundUserId: string | undefined
+): string | undefined {
+  const candidate = resolveResponsesSessionKey(headers, undefined) ?? metadataUserIdSessionId(inboundUserId);
+  return sanitizeCodexSessionId(candidate);
+}
+
+function metadataUserIdSessionId(userId: string | undefined): string | undefined {
+  if (!userId) {
+    return undefined;
+  }
+  if (userId.startsWith("{")) {
+    try {
+      const parsed: unknown = JSON.parse(userId);
+      return isRecord(parsed) ? stringValue(parsed.session_id) : undefined;
+    } catch {
+      return undefined;
+    }
+  }
+  return metadataUserIdSessionPattern.exec(userId)?.[1];
+}
+
+function sanitizeCodexSessionId(value: string | undefined): string | undefined {
+  const trimmed = value?.trim();
+  if (!trimmed || trimmed.length > codexSessionIdMaxLength) {
+    return undefined;
+  }
+  for (const character of trimmed) {
+    const code = character.charCodeAt(0);
+    if (code < 0x20 || code > 0x7e) {
+      return undefined;
+    }
+  }
+  return trimmed;
+}
+
+function hasCodexSessionHeader(headers: Record<string, string> | undefined): boolean {
+  return Object.keys(headers ?? {}).some((name) =>
+    name.trim().toLowerCase().replace(/-/g, "_") === codexSessionHeaderName
+  );
+}
+
 /**
 * First non-empty of the Claude Code session headers (case-insensitive),
 * falling back to the inbound Anthropic `metadata.user_id`.
@@ -91,22 +91,166 @@ test("x-claude-session-id and inbound metadata.user_id are fallback key sources"
  assert.equal(resolveResponsesSessionKey(undefined, undefined), undefined);
 });

-test("codex upstreams receive neither prompt_cache_key nor metadata", () => {
+function codexInput(overrides = {}) {
  const input = responsesInput({
    targetProviderConfig: {
      name: "codex-api::openai_responses",
      type: "openai_responses"
-    }
+    },
+    ...overrides
  });
  input.upstreamRequest.url = "https://chatgpt.com/backend-api/codex/responses";
+  return input;
+}
+
+test("codex upstreams receive the session_id header and prompt_cache_key but no metadata", () => {
+  const input = codexInput();
+  input.request.headers = { "x-claude-code-session-id": "S" };
+
+  const result = applyResponsesSessionAffinity(input);
+
+  assert.equal(result.headers.session_id, "S");
+  assert.equal(result.headers["content-type"], "application/json");
+  assert.equal(result.body.prompt_cache_key, "S");
+  assert.equal(result.body.metadata, undefined);
+  assert.equal(input.upstreamRequest.headers.session_id, undefined);
+  assert.equal(input.upstreamRequest.body.prompt_cache_key, undefined);
+});
+
+test("codex upstreams keep an existing session_id header regardless of spelling", () => {
+  for (const name of ["session_id", "Session_Id", "Session-Id", "session-id"]) {
+    const input = codexInput();
+    input.upstreamRequest.headers = { ...input.upstreamRequest.headers, [name]: "existing" };
+
+    const result = applyResponsesSessionAffinity(input);
+
+    assert.equal(result.headers[name], "existing");
+    assert.deepEqual(
+      Object.keys(result.headers).filter((header) => /^session[-_]id$/i.test(header)),
+      [name]
+    );
+    assert.equal(result.body.prompt_cache_key, "session-1111-2222");
+  }
+});
+
+test("codex upstreams keep a caller-supplied prompt_cache_key", () => {
+  const input = codexInput();
+  input.upstreamRequest.body.prompt_cache_key = "caller-key";
+
+  const result = applyResponsesSessionAffinity(input);
+
+  assert.equal(result.body.prompt_cache_key, "caller-key");
+  assert.equal(result.headers.session_id, "session-1111-2222");
+});
+
+test("codex upstreams pass through untouched without a session source", () => {
+  const input = codexInput({
+    request: {
+      body: { messages: [] },
+      headers: { "content-type": "application/json" }
+    }
+  });

  const result = applyResponsesSessionAffinity(input);

  assert.equal(result, input.upstreamRequest);
-  assert.equal(result.body.prompt_cache_key, undefined);
+});
+
+test("codex upstreams fall back to the session id embedded in metadata.user_id", () => {
+  const jsonInput = codexInput();
+  jsonInput.request.headers = {};
+  jsonInput.request.body.metadata = {
+    user_id: JSON.stringify({ account_uuid: "", device_id: "device-1", session_id: "json-session" })
+  };
+  const jsonResult = applyResponsesSessionAffinity(jsonInput);
+  assert.equal(jsonResult.headers.session_id, "json-session");
+  assert.equal(jsonResult.body.prompt_cache_key, "json-session");
+  assert.equal(jsonResult.body.metadata, undefined);
+
+  const legacyInput = codexInput();
+  legacyInput.request.headers = {};
+  legacyInput.request.body.metadata = {
+    user_id: "user_abc123_account__session_0f8e9a1b-1111-2222-3333-444455556666"
+  };
+  const legacyResult = applyResponsesSessionAffinity(legacyInput);
+  assert.equal(legacyResult.headers.session_id, "0f8e9a1b-1111-2222-3333-444455556666");
+  assert.equal(legacyResult.body.prompt_cache_key, "0f8e9a1b-1111-2222-3333-444455556666");
+});
+
+test("codex upstreams skip the session_id header when metadata.user_id carries no session id", () => {
+  const input = codexInput();
+  input.request.headers = {};
+  input.request.body.metadata = { user_id: "metadata-user" };
+
+  const result = applyResponsesSessionAffinity(input);
+
+  assert.equal(result, input.upstreamRequest);
+});
+
+test("codex upstreams reject session ids that are unsafe as header values", () => {
+  const input = codexInput();
+  input.request.headers = { "x-claude-code-session-id": "bad\r\nx-injected: 1" };
+  input.request.body.metadata = undefined;
+
+  const result = applyResponsesSessionAffinity(input);
+
+  assert.equal(result, input.upstreamRequest);
+});
+
+test("codex upstreams receive the session_id header even for non-JSON bodies", () => {
+  const input = codexInput();
+  input.upstreamRequest = {
+    ...input.upstreamRequest,
+    body: Buffer.from("{}"),
+    bodyEncoding: "bytes"
+  };
+
+  const result = applyResponsesSessionAffinity(input);
+
+  assert.equal(result.headers.session_id, "session-1111-2222");
+  assert.equal(result.body, input.upstreamRequest.body);
+});
+
+test("codex detection via provider baseurl also applies the session_id header", () => {
+  const input = codexInput({
+    targetProviderConfig: {
+      baseurl: "https://chatgpt.com/backend-api/codex",
+      type: "openai_responses"
+    }
+  });
+  input.upstreamRequest.url = "https://proxy.example/v1/responses";
+
+  const result = applyResponsesSessionAffinity(input);
+
+  assert.equal(result.headers.session_id, "session-1111-2222");
  assert.equal(result.body.metadata, undefined);
 });

+test("non-codex openai_responses upstreams do not gain a session_id header", () => {
+  const result = applyResponsesSessionAffinity(responsesInput());
+
+  assert.equal(result.headers.session_id, undefined);
+});
+
+test("gateway boundary plugin forwards the session_id header to codex upstreams", async () => {
+  const hooks = createGatewayPlugin().providerHooks;
+  const hookNames = hooks.map(({ key }) => key);
+  assert.ok(hookNames.indexOf("ccr-upstream-header-sanitizer") < hookNames.indexOf("ccr-responses-session-affinity"));
+
+  let input = codexInput();
+  input.request.headers = { "x-claude-code-session-id": "S" };
+  for (const providerHook of hooks) {
+    if (!providerHook.transformRequest) continue;
+    const result = await providerHook.transformRequest(input);
+    assert.equal(result.ok, true);
+    input = { ...input, upstreamRequest: result.value };
+  }
+
+  assert.equal(input.upstreamRequest.headers.session_id, "S");
+  assert.equal(input.upstreamRequest.body.prompt_cache_key, "S");
+  assert.equal(input.upstreamRequest.body.metadata, undefined);
+});
+
 test("non-codex openai_responses upstreams keep the affinity injection", () => {
  const input = responsesInput();