fix(opencode-go): carry session affinity through Pi native chat

docs-site/src/content/docs/guides/providers.md

@@ -379,6 +379,20 @@ free-experimentation model.
 | Cloudflare AI Gateway | `https://gateway.ai.cloudflare.com/v1/{account-id}/{gateway}/anthropic` |
 | …and more | opencode zen, Vercel AI Gateway, Venice, NanoGPT, Synthetic, Qianfan, Alibaba, Parallel, ZenMux, LiteLLM |
 
+**OpenCode Go** requires a stable session identifier for routing. OpenCodex derives
+its Go session header from Codex thread/session headers, or from a client's
+`x-opencode-session` header when Codex headers are absent. This applies to direct
+Chat Completions requests and requests bridged to Responses. Even an `ocx_`-prefixed
+inbound value is treated as client input and
+hashed into Go affinity; the internal bridge carries the original value, so native
+Chat, bridged Chat, and Responses derive the same result. Explicit provider-config
+session headers are operator overrides and are sent unchanged. Clients must keep the
+identifier stable within a conversation and distinct across conversations; requests
+without a session identifier cannot receive automatic session affinity.
+Generated Pi provider configurations enable `compat.sendSessionAffinityHeaders`
+so Pi sends its per-session identity to the proxy. Existing manually managed Pi
+configurations can set this option on their `opencodex` provider as well.
+
 **OpenCode Zen** (`opencode-zen`) and the keyless **OpenCode Free** preset share
 `https://opencode.ai/zen/v1`. Free models on that gateway often hit a short-window burst
 limit around 15–20 requests/minute (community-measured; OpenCode does not publish RPM).

src/clients/config-export.ts

@@ -695,6 +695,7 @@ export interface PiProviderBlock {
   baseUrl: string;
   api: string;
   apiKey: string;
+  compat?: { sendSessionAffinityHeaders: boolean };
   models: PiModelEntry[];
 }
 
@@ -816,7 +817,7 @@ export interface GajaeGeneratedConfig {
  * model. The rest of this contract (omitting `cost`) is still ours rather than
  * a claim about Pi's acceptance.
  */
-function buildPiClientConfig(ctx: ExportContext): PiGeneratedConfig {
+function buildPiClientConfig(ctx: ExportContext, sendSessionAffinityHeaders = false): PiGeneratedConfig {
   const models: PiModelEntry[] = [];
   for (const model of normalizeExportModels(ctx.models)) {
     // Text is the one modality every routed model supports; anything richer must come
@@ -859,6 +860,7 @@ function buildPiClientConfig(ctx: ExportContext): PiGeneratedConfig {
         baseUrl: ctx.baseUrl,
         api: PI_API_DIALECT,
         apiKey: LOOPBACK_API_KEY_PLACEHOLDER,
+        ...(sendSessionAffinityHeaders ? { compat: { sendSessionAffinityHeaders: true } } : {}),
         models,
       },
     },
@@ -1031,7 +1033,7 @@ function buildOpencodeContribution(ctx: ExportContext): ManagedContribution {
 }
 
 function buildPiContribution(ctx: ExportContext): ManagedContribution {
-  const doc = buildPiClientConfig(ctx);
+  const doc = buildPiClientConfig(ctx, true);
   return singleFragment("pi", ["providers", OPENCODE_PROVIDER_ID], doc.providers[OPENCODE_PROVIDER_ID]);
 }
 
@@ -1127,7 +1129,7 @@ export const EXPORT_CLIENTS: Record<ExportClientId, ExportClientSpec> = {
     destination: env => piConfigPath(env),
     apiKeyEnv: "",
     exportHint: "Pi reads a non-secret placeholder from models.json; loopback needs no key.",
-    build: buildPiClientConfig,
+    build: ctx => buildPiClientConfig(ctx, true),
     format: "json",
     summarize: summarizePi,
     buildContribution: buildPiContribution,

src/server/chat-completions.ts

@@ -25,6 +25,8 @@ import { estimateTokens } from "../lib/token-estimate";
 import { NoEligiblePolicyCandidateError, UnknownRoutingPolicyError, routeModel } from "../router";
 import { evidenceFromBody } from "../routing/request-evidence";
 import { resolveWireProtocolOverride } from "./adapter-resolve";
+import { resolveOpenCodeGoTransport } from "../providers/opencode-go-transport";
+import { normalizeLogConversationId, sessionLaneIdFromRequest } from "./request-log-conversation";
 import type { OcxConfig } from "../types";
 import { readJsonRequestBody } from "./request-decompress";
 import {
@@ -136,6 +138,8 @@ async function handleChatCompletionsWithBudget(
   let chatNativeRoute: ReturnType<typeof routeModel> | null = null;
   try {
     const route = routeModel(config, chatBody.model as string, evidenceFromBody(chatBody));
+    route.provider = resolveOpenCodeGoTransport(route.provider,
+      sessionLaneIdFromRequest(req.headers) ?? normalizeLogConversationId(req.headers.get("x-opencode-session")));
     // Settle the wire once so every branch below reads the adapter this model will
     // actually use, not the provider-wide default (#404).
     route.provider = resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, "chat");
@@ -237,6 +241,9 @@ async function handleChatCompletionsWithBudget(
     return chatCompletionsErrorResponse(400, CODEX_RESERVE_HELPER_UNSUPPORTED_MESSAGE, "invalid_request_error");
   }
   const headers = new Headers({ "content-type": "application/json" });
+  // Internal bridge metadata; the Go resolver scopes and hashes it before upstream use.
+  const openCodeSession = req.headers.get("x-opencode-session");
+  if (openCodeSession) headers.set("x-opencode-session", openCodeSession);
   for (const name of FORWARD_HEADERS) {
     if (name === "authorization" && !directRoute) continue;
     const value = req.headers.get(name);

src/server/responses/core.ts

@@ -2304,7 +2304,8 @@ async function applyFinalRouteRequestNormalization(args: {
 
   // Settle the wire once so logging, fast-mode, auth, and sidecars read the adapter
   // this request will actually use (#404).
-  route.provider = resolveOpenCodeGoTransport(route.provider, sessionLaneIdFromRequest(req.headers));
+  route.provider = resolveOpenCodeGoTransport(route.provider,
+    sessionLaneIdFromRequest(req.headers) ?? normalizeLogConversationId(req.headers.get("x-opencode-session")));
   route.provider = resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire);
   if (preserveAnthropicResponseModel) parsed._responseModelId = responseModelId;
   logCtx.model = route.modelId;

tests/clients/prime-client.test.ts

@@ -37,18 +37,14 @@ function context(): ExportContext {
 }
 
 describe("Prime Agent client config", () => {
-  /**
-   * The load-bearing claim of this client: Prime Agent is the pi coding agent
-   * under a different brand, so it reads the SAME models.json contract rather
-   * than a lookalike. Locking the two documents together is what keeps that
-   * claim true — if a future Pi-only change diverges, this fails here instead
-   * of silently shipping Prime users a config their agent rejects.
-   */
-  test("generates byte-for-byte the document Pi generates", () => {
-    const prime = buildClientConfigText("prime", context());
-    const pi = buildClientConfigText("pi", context());
-    expect(prime.format).toBe("json");
-    expect(prime.text).toBe(pi.text);
+  test("shares Pi's model contract without opting Prime into session headers", () => {
+    const prime = buildClientConfig("prime", context()) as PiGeneratedConfig;
+    const pi = buildClientConfig("pi", context()) as PiGeneratedConfig;
+    expect(pi.providers[OPENCODE_PROVIDER_ID]!.compat).toEqual({ sendSessionAffinityHeaders: true });
+    delete pi.providers[OPENCODE_PROVIDER_ID]!.compat;
+    expect(prime).toEqual(pi);
+    expect(buildClientContribution("prime", context()).fragments[0]!.value)
+      .toEqual(prime.providers[OPENCODE_PROVIDER_ID]);
   });
 
   test("adds only providers.opencodex, wired to the loopback proxy", () => {

tests/config/client-config-export.test.ts

@@ -11,6 +11,7 @@ import {
   LOOPBACK_API_KEY_PLACEHOLDER,
   SCHEMA_REQUIRED_OUTPUT_BUDGET,
   buildClientConfig,
+  buildClientContribution,
   buildClientConfigText,
   isExportClientId,
   normalizeExportModels,
@@ -315,6 +316,8 @@ describe("Pi serializer (accept criterion 2)", () => {
     expect(provider.baseUrl).toBe(BASE_URL);
     expect(provider.api).toBe("openai-completions");
     expect(provider.apiKey).toBe(LOOPBACK_API_KEY_PLACEHOLDER);
+    expect(provider.compat?.sendSessionAffinityHeaders).toBe(true);
+    expect(buildClientContribution("pi", ctx()).fragments[0]!.value).toEqual(provider);
   });
 
   test("cost is omitted on every entry — zeros would assert routed models are free", () => {
@@ -899,7 +902,7 @@ describe("EXPORT_CLIENTS registry", () => {
 `);
   });
 
-  test("pi bytes are unchanged, to the last newline", () => {
+  test("pi bytes include session affinity, to the last newline", () => {
     const built = buildClientConfigText("pi", ctx({ config: cfg() }));
     expect(built.format).toBe("json");
     expect(built.text).toBe(`{
@@ -908,6 +911,9 @@ describe("EXPORT_CLIENTS registry", () => {
       "baseUrl": "http://127.0.0.1:10100/v1",
       "api": "openai-completions",
       "apiKey": "opencodex-loopback",
+      "compat": {
+        "sendSessionAffinityHeaders": true
+      },
       "models": [
         {
           "id": "anthropic/claude-opus-5",

tests/providers/opencode-go-session-header.test.ts

@@ -1,8 +1,10 @@
 import { afterEach, describe, expect, test } from "bun:test";
 import { providerConfigSeed } from "../../src/providers/derive";
-import { resolveOpenCodeGoTransport } from "../../src/providers/opencode-go-transport";
+import { deriveOpenCodeGoSessionId, resolveOpenCodeGoTransport } from "../../src/providers/opencode-go-transport";
 import { getProviderRegistryEntry } from "../../src/providers/registry";
 import { handleResponses } from "../../src/server/responses/core";
+import { handleChatCompletions } from "../../src/server/chat-completions";
+import { normalizeLogConversationId } from "../../src/server/request-log-conversation";
 import type { OcxConfig, OcxProviderConfig } from "../../src/types";
 
 const MUSE_MODEL = "muse-spark-1.3-contributor";
@@ -52,6 +54,8 @@ async function captureRequest(input: {
   model?: string;
   child?: string;
   provider?: OcxProviderConfig;
+  nativeChat?: boolean;
+  headers?: Record<string, string>;
 } = {}): Promise<{ url: string; headers: Headers }> {
   const providerName = input.providerName ?? "opencode-go";
   const model = input.model ?? MUSE_MODEL;
@@ -65,10 +69,18 @@ async function captureRequest(input: {
   const config = {
     providers: { [providerName]: input.provider ?? opencodeGo() },
   } as unknown as OcxConfig;
-  const response = await handleResponses(
+  const response = input.nativeChat ? await handleChatCompletions(
+    new Request("http://localhost/v1/chat/completions", {
+      method: "POST",
+      headers: input.headers ?? codexHeaders(input.child),
+      body: JSON.stringify({ model: `${providerName}/${model}`, messages: [{ role: "user", content: "ping" }], stream: false }),
+    }),
+    config,
+    { model: "", provider: "" },
+  ) : await handleResponses(
     new Request("http://localhost/v1/responses", {
       method: "POST",
-      headers: codexHeaders(input.child),
+      headers: input.headers ?? codexHeaders(input.child),
       body: JSON.stringify({ model: `${providerName}/${model}`, input: "ping", stream: false }),
     }),
     config,
@@ -77,6 +89,7 @@ async function captureRequest(input: {
   );
 
   expect(response.status).toBe(200);
+  await response.text();
   expect(requests).toHaveLength(1);
   return requests[0]!;
 }
@@ -85,6 +98,65 @@ describe("OpenCode Go session affinity (#3344)", () => {
   const originalFetch = globalThis.fetch;
   afterEach(() => { globalThis.fetch = originalFetch; });
 
+  test("native Chat ingress preserves stable Go affinity and separates conversations", async () => {
+    const provider = opencodeGo();
+    const input = { nativeChat: true, model: "omen-alpha", provider };
+    const first = await captureRequest(input);
+    const continued = await captureRequest(input);
+    const sibling = await captureRequest({ ...input, child: "child-thread-b" });
+    expect(first.url).toBe("https://opencode.ai/zen/go/v1/chat/completions");
+    expect(first.headers.get(SESSION_HEADER)).toMatch(/^ocx_[0-9a-f]{32}$/);
+    expect(continued.headers.get(SESSION_HEADER)).toBe(first.headers.get(SESSION_HEADER));
+    expect(sibling.headers.get(SESSION_HEADER)).not.toBe(first.headers.get(SESSION_HEADER));
+    expect(provider.headers?.[SESSION_HEADER]).toBeUndefined();
+  });
+
+  test("native Chat honors configured session headers on renamed Go providers", async () => {
+    const captured = await captureRequest({
+      nativeChat: true, model: "omen-alpha", providerName: "renamed-go",
+      provider: opencodeGo({ headers: { "X-OpenCode-Session": "operator-session" } }),
+    });
+    expect(captured.headers.get(SESSION_HEADER)).toBe("operator-session");
+  });
+
+  test("uses a Pi session header without Codex headers on native and bridged Chat", async () => {
+    const headers = { "content-type": "application/json", "x-opencode-session": "pi-conversation-a" };
+    const chat = await captureRequest({ nativeChat: true, model: "omen-alpha", headers });
+    const bridged = await captureRequest({ nativeChat: true, model: MUSE_MODEL, headers });
+    expect(chat.headers.get(SESSION_HEADER)).toMatch(/^ocx_[0-9a-f]{32}$/);
+    expect(chat.headers.get(SESSION_HEADER)).not.toContain("pi-conversation-a");
+    expect(bridged.headers.get(SESSION_HEADER)).toBe(chat.headers.get(SESSION_HEADER));
+  });
+
+  for (const session of ["client-session-a", "ocx_0123456789abcdef0123456789abcdef"]) {
+    test(`treats inbound ${session.startsWith("ocx_") ? "ocx-prefixed" : "raw"} identity as client input on every ingress`, async () => {
+      const headers = { "content-type": "application/json", [SESSION_HEADER]: session };
+      const expected = deriveOpenCodeGoSessionId(normalizeLogConversationId(session)!);
+      const native = await captureRequest({ nativeChat: true, model: "omen-alpha", headers });
+      const bridged = await captureRequest({ nativeChat: true, model: MUSE_MODEL, headers });
+      const responses = await captureRequest({ model: MUSE_MODEL, headers });
+      expect(native.url).toEndWith("/chat/completions");
+      expect(bridged.url).toEndWith("/responses");
+      for (const request of [native, bridged, responses]) {
+        expect(request.headers.get(SESSION_HEADER)).toBe(expected);
+        expect(request.headers.get(SESSION_HEADER)).not.toBe(session);
+      }
+      const override = await captureRequest({
+        nativeChat: true, model: "omen-alpha", headers,
+        provider: opencodeGo({ headers: { "X-OpenCode-Session": session } }),
+      });
+      expect(override.headers.get(SESSION_HEADER)).toBe(session);
+    });
+  }
+
+  test("native Chat does not send Go affinity to an unrelated destination", async () => {
+    const captured = await captureRequest({
+      nativeChat: true, model: "omen-alpha", providerName: "custom-go",
+      provider: opencodeGo({ baseUrl: "https://opencode.ai.evil.test/zen/go/v1" }),
+    });
+    expect(captured.headers.has(SESSION_HEADER)).toBe(false);
+  });
+
   test("sends one stable opaque session header on Responses and Chat wires", async () => {
     const responses = await captureRequest({ model: MUSE_MODEL });
     const chat = await captureRequest({ model: CHAT_MODEL });