Fix Codex session affinity and WebSocket continuation reuse

src/http/routes/proxy.ts

@@ -8,7 +8,11 @@ import {
 } from "../../db/repositories/request-usage";
 import { getRoutableProviderAccount } from "../../domain/providers/provider-service";
 import { prepareClaudeProxyRequest } from "../../providers/proxies/claude-proxy";
-import { prepareCodexProxyRequest } from "../../providers/proxies/codex-proxy";
+import {
+  deriveCodexSessionId,
+  prepareCodexProxyRequest,
+  readCodexSessionId,
+} from "../../providers/proxies/codex-proxy";
 import { tryProxyCodexWebSocket } from "../../providers/proxies/codex-websocket";
 import { prepareCopilotProxyRequest } from "../../providers/proxies/copilot-proxy";
 import type { UsageRequestSource } from "../../usage/request-outcome";
@@ -221,6 +225,13 @@ const proxyRequest = async (
 
   switch (route.provider) {
     case "codex": {
+      const codexSessionId = readCodexSessionId(requestBodyJson, headers);
+      const codexUpstreamSessionId = codexSessionId
+        ? await deriveCodexSessionId(
+            `${apiKeyId}:${account.id}`,
+            codexSessionId
+          )
+        : null;
       const codexProxy = prepareCodexProxyRequest({
         headers,
         accessToken: account.accessToken,
@@ -229,6 +240,7 @@ const proxyRequest = async (
           account.metadata?.provider === "codex" ? account.metadata : null,
         bodyText: requestBody,
         bodyJson: requestBodyJson,
+        sessionId: codexUpstreamSessionId,
         onTokenUsage: usageRecorder.onTokenUsage,
       });
       upstreamUrl = codexProxy.upstreamUrl;
@@ -239,6 +251,8 @@ const proxyRequest = async (
         headers,
         bodyJson: codexProxy.bodyJson,
         accountKey: `${apiKeyId}:${account.id}`,
+        sessionId: codexSessionId,
+        upstreamSessionId: codexUpstreamSessionId,
         onTokenUsage: usageRecorder.onTokenUsage,
         signal: context.req.raw.signal,
       });

src/providers/proxies/codex-proxy.ts

@@ -21,8 +21,53 @@ import {
 const trimString = (value: unknown): string =>
   typeof value === "string" ? value.trim() : "";
 
+const toHex = (bytes: Uint8Array): string =>
+  Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join("");
+
+export const readCodexSessionId = (
+  body: unknown,
+  headers: Headers
+): string | null => {
+  const bodySessionId = isObjectRecord(body)
+    ? trimString(body.prompt_cache_key)
+    : "";
+  return (
+    bodySessionId ||
+    trimString(headers.get("session_id")) ||
+    trimString(headers.get("session-id")) ||
+    trimString(headers.get("x-session-affinity")) ||
+    null
+  );
+};
+
+export const deriveCodexSessionId = async (
+  accountKey: string,
+  sessionId: string
+): Promise<string> => {
+  const data = new TextEncoder().encode(`${accountKey}:${sessionId}`);
+  const digest = await crypto.subtle.digest("SHA-256", data);
+  return `kleis_${toHex(new Uint8Array(digest)).slice(0, 48)}`;
+};
+
+export const applyCodexSessionHeaders = (
+  headers: Headers,
+  sessionId: string
+): void => {
+  clearCodexSessionHeaders(headers);
+  headers.set("session-id", sessionId);
+  headers.set("x-client-request-id", sessionId);
+};
+
+export const clearCodexSessionHeaders = (headers: Headers): void => {
+  headers.delete("session_id");
+  headers.delete("session-id");
+  headers.delete("x-session-affinity");
+  headers.delete("x-client-request-id");
+};
+
 export const transformCodexBodyJson = (
-  bodyJson: unknown
+  bodyJson: unknown,
+  sessionId?: string | null
 ): JsonObject | null => {
   if (!isObjectRecord(bodyJson)) {
     return null;
@@ -51,6 +96,7 @@ export const transformCodexBodyJson = (
   return {
     ...nextBody,
     instructions,
+    ...(sessionId ? { prompt_cache_key: sessionId } : {}),
   };
 };
 
@@ -74,6 +120,7 @@ type CodexProxyPreparationInput = {
   metadata: CodexAccountMetadata | null;
   bodyText: string;
   bodyJson: unknown;
+  sessionId?: string | null;
   onTokenUsage?: ((usage: TokenUsage) => void) | null;
 };
 
@@ -89,7 +136,7 @@ export const prepareCodexProxyRequest = (
 ): CodexProxyPreparationResult => {
   const isStreamingRequest =
     readBooleanField(input.bodyJson, "stream") === true;
-  const bodyJson = transformCodexBodyJson(input.bodyJson);
+  const bodyJson = transformCodexBodyJson(input.bodyJson, input.sessionId);
 
   input.headers.set("authorization", `Bearer ${input.accessToken}`);
   if (!input.headers.get("originator")) {
@@ -100,6 +147,10 @@ export const prepareCodexProxyRequest = (
   if (accountId) {
     input.headers.set(CODEX_ACCOUNT_ID_HEADER, accountId);
   }
+  clearCodexSessionHeaders(input.headers);
+  if (input.sessionId) {
+    applyCodexSessionHeaders(input.headers, input.sessionId);
+  }
 
   return {
     upstreamUrl: CODEX_RESPONSE_ENDPOINT,