fix(codex): key pool affinity on the conversation cohort, not the thread

devlog/_plan/260918_lane_a_bug_train/050_cohort_pool_affinity.md

@@ -0,0 +1,96 @@
+# #4780 — the binding unit becomes the conversation cohort
+
+## What was verified before anything changed
+
+The issue was filed rather than patched because it is a design decision. Before implementing it,
+every claim in it was re-derived at the current head and against the pinned upstream clone.
+
+At `origin/dev`, `codexPoolAffinityKey` resolves to `codexConversationIdentity(...).conversationKey`,
+which was `HMAC(session ?? parent, thread)` — per thread, deliberately, and the function's own
+comment said so.
+
+Upstream keys the prompt cache on the tree, not the thread. In
+`/Users/jun/Developer/codex/121_openai-codex` at `095da4b7e`:
+
+- `codex-rs/core/src/client.rs` — `prompt_cache_key()` returns `responses_metadata.session_id`,
+  or `{source}:{parent_thread_id}` for an internal session. Never the thread's own id.
+- `codex-rs/core/src/agent/control.rs` — the comment is explicit: `session_id` "is equal to the
+  root thread's ID", and that one `AgentControl` is "shared with every sub-agent spawned from
+  that root".
+- `codex-rs/core/tests/suite/prompt_cache_key.rs` — asserts `differentThreadIds: true` while
+  root and child both send `promptCacheKey == expected_session_id`.
+
+The clone's HEAD is 2026-09-08, so it does **not** contain openai/codex#44862 ("Preserve parent
+cache affinity for ephemeral forks"), which merged on 2026-09-11 and makes an ephemeral fork
+inherit its parent's session id for exactly this reason. Checking that was the point of the
+exercise: the issue body is not stale, and upstream has moved further toward cohort keying since
+it was written.
+
+So two requests could carry an identical `prompt_cache_key` and be served by different accounts.
+The split member asserts a warm prefix that is deterministically cold on its account. Nothing
+fails — the prompt is replayed in full and the tokens burn, which is why no log line reports it.
+
+## This is not a revert of #4546 wp8
+
+Stated here, in the code, in `structure/`, and in the PR, because a reader who concludes
+otherwise will flip it straight back.
+
+wp8 fixed a genuine incoherence: a child bound under the RAW parent id, which is a different
+identity from the root's own `app:HMAC(session, thread)` binding. Siblings therefore shared an
+entry the root was not on, and a grandchild keying on its own parent landed on a key nobody had
+ever bound. A cohort key cannot produce that, because the root's own binding IS the cohort key:
+one identity for the tree instead of two competing ones.
+
+What wp8 additionally gave each thread — a binding of its own — is what this gives up, knowingly.
+
+## The cost, stated rather than buried
+
+A tree gains cache locality and loses per-thread placement independence. Every member shares one
+binding, so a fan-out cannot spread across accounts, and when that account is exhausted or retired
+the whole tree moves together. That is correct for cache affinity and it is a behaviour change,
+not a refinement. It also interacts with the send-budget and placement work #4546 introduced,
+since a tree is now one binding for accounting as well as for routing — the issue flags that, and
+it remains true.
+
+## Order of work
+
+1. The invariant first, as a regression: root, child and grandchild sharing a session resolve to
+   one key; the unbound set is unchanged; the key stays inside its authenticated scope; and a
+   move of the cohort carries every member.
+2. The orphan guard, because the wp8 argument above has to be proved rather than asserted: a
+   grandchild must never land on a key nobody bound. Covered for the session case and for the
+   session-less chain, which is the one that could still split if each depth anchored on its own
+   parent.
+3. The derivation, then the first-placement hook.
+4. `structure/providers/openai-tiers.md`.
+
+## The session-less chain
+
+The only case cohort keying cannot read off the headers. There is nothing that names the tree, so
+the cohort is read from the parent's lineage record and falls back to `HMAC(parent, parent)` when
+this scope has not seen that parent — the same key that parent derives for itself. Without the
+record lookup a chain of parent-only turns would anchor on a different ancestor at every depth and
+split the cohort again, which is the wp8 shape arriving through the other door.
+
+## First placement
+
+Largely dead now: a member of a tree any other member has bound resolves to that binding, so there
+is nothing to place. `pickLineageServingAccount` is kept for the session-less unrecorded-parent
+case and gated on the parent's key actually differing from the request's own, so it never re-asks
+a question the binding lookup already answered. Removing the mechanism outright would have taken
+the `lineage_parent`/`lineage_sibling` affinity reasons and their diagnostics with it, which is a
+wider blast radius than this change needs.
+
+## Retired coverage, named
+
+Four placement tests encoded the per-thread contract and could not survive it changing. Each is
+either rewritten for the cohort equivalent or explicitly retired with its replacement named:
+
+- "every thread keys as itself" — retired; the key shape, the unbound set and the grandchild
+  orphan property are pinned in the new cohort block instead.
+- "a child with no binding starts on the parent's account under its OWN key" — retired; the
+  cohort block asserts the child joins the existing binding.
+- "a later move of the parent does not drag an already-bound child" — inverted deliberately, and
+  the comment says so: the asymmetry is what #4780 gives up.
+- "a compatible sibling places the child" — rewritten as the cohort rebinding off an ineligible
+  account, keeping the negative assertion that a stale home contributes nothing.

src/codex/auth-context.ts

@@ -135,20 +135,22 @@ export function requestOwnedMainPinState(
 }
 
 /**
- * Every thread keys as ITSELF, never as its parent (#4546, wp8).
+ * One conversation tree keys as ONE cohort (#4780).
  *
- * The old rule preferred `x-codex-parent-thread-id`, so every child of one parent bound under
- * the RAW parent id -- one shared entry, unrelated to the root's own `app:HMAC(session, thread)`
- * binding -- and a grandchild keyed on its own parent landed on a key nobody had ever bound.
- * A child therefore started cold while its parent was being served warm somewhere, and no
- * child could hold a binding of its own.
+ * Upstream keys its prompt cache on something the whole tree shares -- `prompt_cache_key` is the
+ * session id, or `{source}:{parent_thread_id}` for an internal session -- and one `AgentControl`
+ * whose `session_id` is the root thread's id is shared with every sub-agent. While the proxy
+ * keyed per thread, two requests could carry an identical `prompt_cache_key` and be served by
+ * different accounts, so the split member asserted a warm prefix that was deterministically cold
+ * on its account. Nothing failed; the prompt was replayed in full and the tokens burned.
  *
- * Now a request with a `thread-id` keys as HMAC(session ?? parent, thread). A root is
- * unchanged, a child gets an independent key, and a request naming only a parent rides the
- * parent's lane under HMAC(parent, parent) -- the same one-to-one lane it always had, minus
- * the caller-supplied identifier that used to sit in Pool state. Which requests produce no
- * key at all is unchanged. First placement for a child is what consults the family, through
- * `recordCodexThreadLineage` below and the placement hook in ./routing.
+ * THIS IS NOT A REVERT OF #4546 wp8. wp8 fixed a different defect -- a child bound under the RAW
+ * parent id, an identity unrelated to the root's own binding, so siblings shared an entry the
+ * root was not on and a grandchild landed on a key nobody had bound. A cohort key has no such
+ * incoherence, because the root's own binding IS the cohort key. What wp8 additionally gave each
+ * thread, a binding of its own, is what this deliberately gives up.
+ *
+ * Which requests produce no key at all is unchanged, through both units.
  *
  * The derivation itself lives in ./lineage so a lineage record's conversation key and the key
  * the thread actually binds under can never drift apart.

src/codex/lineage.ts

@@ -14,9 +14,11 @@
  *
  * What it is for, and what it is not for:
  *
- * - FIRST PLACEMENT. A child with no binding of its own may start where its family is already
- *   warm; see `pickLineageServingAccount` in ./routing. Once bound, the child is an ordinary
- *   binding, so a later move of the parent does not drag it.
+ * - FIRST PLACEMENT. Largely subsumed by cohort keying (#4780): a member of a tree that any
+ *   other member has already bound resolves to that same binding, so there is nothing to place.
+ *   `pickLineageServingAccount` in ./routing remains for the one case cohort keying cannot
+ *   unify -- a session-less chain whose parent is not recorded in this scope -- and is gated on
+ *   the parent's key actually differing from this request's.
  * - COST ATTRIBUTION. {@link codexThreadLineageLookup} and {@link codexLineageRootForRequest}
  *   answer which root workflow a conversation belongs to, so a grandchild's spend aggregates
  *   onto the root. No budget is implemented here.
@@ -155,30 +157,83 @@ export function codexConversationKeyFor(familyId: string, threadId: string): str
     .digest("base64url")}`;
 }
 
+/**
+ * The cohort anchor's key: the same string on both sides of the derivation.
+ *
+ * A cohort is identified by one value the whole tree shares, so the key is that value keyed
+ * against itself rather than against a member. Using {@link codexConversationKeyFor} keeps one
+ * derivation in the module, which is what guarantees a lineage record's `conversationKey` stays
+ * byte-identical to the key the thread actually binds under.
+ */
+function cohortKeyFromAnchor(sessionId: string | undefined, fallbackAnchor: string): string {
+  const anchor = sessionId ?? fallbackAnchor;
+  return codexConversationKeyFor(anchor, anchor);
+}
+
+/**
+ * The cohort this request belongs to, or undefined when it names no cohort at all.
+ *
+ * The session IS the cohort, which is exactly how upstream keys the prompt cache:
+ * `prompt_cache_key()` returns `responses_metadata.session_id` (or `{source}:{parent_thread_id}`
+ * for an internal session), and `AgentControl.session_id` is the root thread's id, shared with
+ * every sub-agent spawned from that root. Two requests carrying the same `prompt_cache_key` must
+ * not be served by different accounts, and keying on the session is what makes that structural
+ * rather than a hint (#4780).
+ *
+ * Without a session there is nothing in the headers that names the tree, so the cohort is
+ * whatever the parent is already bound to. That lookup is what keeps a chain of parent-only
+ * turns converging on one key: anchoring each depth on its own parent would split the cohort
+ * again at every hop. When the parent has not been seen in this scope the parent id anchors it,
+ * which is the same key that parent derives for itself.
+ */
+function cohortConversationKey(
+  headers: Headers,
+  sessionId: string | undefined,
+  parentThreadId: string | undefined,
+  now: number,
+): string | undefined {
+  if (sessionId !== undefined) return cohortKeyFromAnchor(sessionId, sessionId);
+  if (parentThreadId === undefined) return undefined;
+  const recorded = liveLineageRecord(
+    lineageByScope.get(codexLineageScopeKey(headers)),
+    parentThreadId,
+    now,
+  );
+  return recorded?.conversationKey ?? cohortKeyFromAnchor(undefined, parentThreadId);
+}
+
 /**
  * Resolve a request's conversation identity, or undefined when it carries no bindable thread
  * identity at all.
  *
- * The set of requests that produce NO key is deliberately unchanged from the pre-#4546 rule: a
- * bare `thread-id` with neither a session nor a parent stays unbound, exactly as the Desktop
- * fallback required both halves of its pair. Only the VALUE moves, and only for requests that
- * name a parent:
+ * The set of requests that produce NO key is deliberately unchanged, through both #4546 and
+ * #4780: a bare `thread-id` with neither a session nor a parent stays unbound, exactly as the
+ * Desktop fallback required both halves of its pair. Only the VALUE moves.
+ *
+ * Every member of one tree resolves to the SAME key, because the cohort is the binding unit
+ * (#4780):
+ *
+ * - root (`session-id` + `thread-id`) -> the session's cohort key;
+ * - child and grandchild (parent + own `thread-id`) -> the same cohort key, since they carry the
+ *   same session;
+ * - parent-only (no `thread-id`) -> the same cohort key, derived from the session or, without
+ *   one, read from the parent's record.
  *
- * - root (`session-id` + `thread-id`) -> HMAC(session, thread), unchanged;
- * - child (parent + own `thread-id`) -> HMAC(session ?? parent, thread), previously the raw parent
- *   id, which is what made siblings share one entry and made a child's first turn land on a key
- *   the root had never bound;
- * - parent-only (no `thread-id`) -> the parent's OWN recorded key when this scope has one, and
- *   otherwise HMAC(session ?? parent, parent).
+ * THIS IS NOT A REVERT OF #4546 wp8, and reading it as one would flip it straight back. wp8
+ * fixed a real incoherence: a child keyed under the RAW parent id, which is a different identity
+ * from the root's own `app:HMAC(session, thread)` binding, so siblings shared an entry unrelated
+ * to the root's and a grandchild keying on its own parent landed on a key nobody had ever bound.
+ * A cohort key cannot produce that, because the root's own binding IS the cohort key: there is
+ * one identity for the tree rather than two competing ones. What wp8 additionally gave each
+ * thread -- a binding of its own -- is what #4780 deliberately gives up, and the reason is that
+ * upstream never agreed to it: `prompt_cache_key` is keyed on the session the whole tree shares,
+ * so a proxy that splits the tree makes every split member assert a warm prefix that is cold on
+ * its account.
  *
- * That last case is the one with a trap in it. A parent-only turn belongs to the parent's
- * conversation, so it has to land on the binding the parent is already using -- but the parent's
- * key is HMAC(session, thread), and HMAC(parent, parent) reproduces it only when the session id
- * and the thread id are the same string. Codex's own root happens to satisfy that, which is
- * exactly why deriving the key looks correct until a caller whose session differs from its thread
- * starts a COLD conversation on every parent-only turn and overwrites the parent's record on the
- * way through. So the recorded key wins, the session-derived key is the fallback that reproduces
- * it when the parent has not been seen in this scope, and the raw parent id is never the answer.
+ * The parent-only case keeps the trap it always had. Such a turn belongs to the parent's
+ * conversation, so it must land on the binding the parent is already using. With a session in
+ * hand that is immediate, since both derive the same cohort key. Without one, the recorded key
+ * wins and the parent id anchors the fallback; the raw parent id is never the answer.
  */
 export function codexConversationIdentity(
   headers: Headers,
@@ -190,24 +245,20 @@ export function codexConversationIdentity(
 
   if (threadId === undefined) {
     if (parentThreadId === undefined) return undefined;
-    const recorded = liveLineageRecord(
-      lineageByScope.get(codexLineageScopeKey(headers)),
-      parentThreadId,
-      now,
-    );
+    const parentOnlyKey = cohortConversationKey(headers, sessionId, parentThreadId, now);
+    if (parentOnlyKey === undefined) return undefined;
     return {
-      conversationKey: recorded?.conversationKey
-        ?? codexConversationKeyFor(sessionId ?? parentThreadId, parentThreadId),
+      conversationKey: parentOnlyKey,
       recordThreadId: parentThreadId,
       ...(sessionId !== undefined ? { sessionId } : {}),
       legacyConversationKey: parentThreadId,
       declaresParent: false,
     };
   }
-  const familyId = sessionId ?? parentThreadId;
-  if (familyId === undefined) return undefined;
+  const conversationKey = cohortConversationKey(headers, sessionId, parentThreadId, now);
+  if (conversationKey === undefined) return undefined;
   return {
-    conversationKey: codexConversationKeyFor(familyId, threadId),
+    conversationKey,
     recordThreadId: threadId,
     ...(sessionId !== undefined ? { sessionId } : {}),
     ...(parentThreadId !== undefined ? { parentThreadId } : {}),
@@ -313,7 +364,7 @@ function lineageFor(
   const parentConversationKey = parentThreadId === undefined
     ? undefined
     : parentRecord?.conversationKey
-      ?? codexConversationKeyFor(identity.sessionId ?? parentThreadId, parentThreadId);
+      ?? cohortKeyFromAnchor(identity.sessionId, parentThreadId);
   const rootSessionKey = parentConversationKey === undefined
     ? identity.conversationKey
     : parentRecord?.rootSessionKey ?? parentConversationKey;

src/codex/routing.ts

@@ -392,7 +392,14 @@ function pickLineageServingAccount(
   selectionOptions?: CodexAccountUsabilityOptions,
   modelId?: string,
 ): { accountId: string; reason: CodexAffinityReason } | null {
-  if (lineage.parentConversationKey !== undefined) {
+  // Cohort keying (#4780) makes a tree share one key, so for a same-session family the parent's
+  // key IS this request's and the lookup below would re-ask a question the caller already
+  // answered by finding no binding entry. What remains is the case cohort keying cannot unify:
+  // a session-less chain whose parent this scope has not recorded, where the keys differ.
+  if (
+    lineage.parentConversationKey !== undefined
+    && lineage.parentConversationKey !== lineage.conversationKey
+  ) {
     const parent = lineageServingAccountId(
       lineage.parentConversationKey, config, now, quotaScope, selectionOptions, modelId,
     );