diff --git a/devlog/_plan/260918_lane_a_bug_train/050_cohort_pool_affinity.md b/devlog/_plan/260918_lane_a_bug_train/050_cohort_pool_affinity.md new file mode 100644 index 00000000000..86e5f6ee4d0 --- /dev/null +++ b/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. diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 40e0c2a669f..45b7df580de 100644 --- a/src/codex/auth-context.ts +++ b/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. diff --git a/src/codex/lineage.ts b/src/codex/lineage.ts index ff5233c04bc..4788efbe274 100644 --- a/src/codex/lineage.ts +++ b/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; diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 5f768f554b1..e1d1ea9e931 100644 --- a/src/codex/routing.ts +++ b/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, ); diff --git a/structure/providers/openai-tiers.md b/structure/providers/openai-tiers.md index f10d28fe1ed..1d5e01c639a 100644 --- a/structure/providers/openai-tiers.md +++ b/structure/providers/openai-tiers.md @@ -73,35 +73,60 @@ support them. > Decision record: [ADR-0084](../decisions/ADR-0084-public-provider-contract.md) -Pool affinity keys every Codex V2 thread as itself. A request carrying `thread-id` maps to an -opaque `app:HMAC(session-id ?? x-codex-parent-thread-id, thread-id)` under a random process-local -key, so a root is keyed exactly as before while each child holds an independent binding instead of -collapsing onto the raw parent id shared by all its siblings. A request naming only -`x-codex-parent-thread-id` rides that parent's own lane as `app:HMAC(parent, parent)`: one parent -id still maps to exactly one lane, and no caller-supplied identifier reaches Pool state. Which -requests bind at all is unchanged -- a bare `thread-id` with neither a session nor a parent has no -family anchor and stays unbound. Components are trimmed and bounded at 512 bytes, missing or +Pool affinity keys a Codex V2 conversation TREE as one cohort. The binding unit is the cohort the +client already declares, because upstream keys its prompt cache the same way: `prompt_cache_key` +is `responses_metadata.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 +spawned from that root. Two requests carrying the same `prompt_cache_key` are therefore served by +the same account. While the key was per thread, a tree could split across accounts while every +member kept sending one cache key, so the split member asserted a warm prefix that was +deterministically cold on its account: no failure, only a full prompt replay and the token spend +(#4780). + +A request carrying `session-id` maps to an opaque `app:HMAC(session, session)` under a random +process-local key, and every root, child and grandchild of that tree resolves to it. Without a +session the cohort is read from the parent's lineage record, falling back to +`app:HMAC(parent, parent)` when this scope has not seen that parent -- which is the same key that +parent derives for itself, so a chain of parent-only turns converges rather than splitting at +every depth. No caller-supplied identifier reaches Pool state. Which requests bind at all is +unchanged -- a bare `thread-id` with neither a session nor a parent has no family anchor and +stays unbound. Components are trimmed and bounded at 512 bytes, missing or oversized values stay unbound, raw identifiers and durable hashes are never stored, and account-qualified selectors skip both lookup and mutation. Selection, subagent fallback preview, and terminal outcome accounting carry the same key so route planning cannot preview one account and authenticate another. A transient-failure streak does not delete the live binding that actually selected the account; the request is served by another account while the binding is kept. +**Cohort keying is not a revert of #4546 wp8, and reading it as one will flip it straight back.** +wp8 fixed a different defect: a child bound under the RAW parent id, an identity unrelated to the +root's own `app:HMAC(session, thread)` binding, so siblings 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 incoherence, because the root's own binding IS the cohort key -- there is one +identity for the tree rather than two competing ones, and the grandchild-orphan property is pinned +as a test rather than left as an argument. What wp8 additionally gave each thread, a binding of +its own, is what #4780 deliberately gives up. + +The cost is explicit: a tree gains cache locality and loses per-thread placement independence. All +members share 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. + `src/codex/lineage.ts` owns that derivation and records the family relation behind it: each thread's own conversation key, its immediate parent's thread id, and the transitive root, so a grandchild resolves to the same root as its parent. Records are held per authenticated scope (an HMAC of the caller's Authorization under the same process-local key), and bounded in both dimensions -- idle TTL and an LRU cap on records per scope, and an LRU cap on scopes. -First placement is the only routing decision that consults lineage. A thread that has never bound -starts on the account CURRENTLY serving its parent, which includes a live transient detour rather -than the parent's stale home, then on a compatible sibling's current serving account, then on -ordinary cold placement. An ineligible or dead family account contributes nothing, because a stale -home is worse than no hint. The child then holds an ordinary binding of its own: a later move of -the parent does not drag it, and the hint does not move the shared active-account cursor. The same -module exposes the root lookup other layers use for cost attribution and a lineage-backed -worker/interactive answer; admission's header-only classification is unchanged and does not yet -read it. +First placement is largely subsumed by cohort keying: a member of a tree any other member has +already bound resolves to that same binding, so there is nothing to place. `pickLineageServingAccount` +remains for the one case cohort keying cannot unify -- a session-less chain whose parent this scope +has not recorded -- and is 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. An ineligible or dead family +account still contributes nothing, because a stale home is worse than no hint, and the hint does +not move the shared active-account cursor. The same module exposes the root lookup other layers +use for cost attribution and a lineage-backed worker/interactive answer; admission's header-only +classification is unchanged and does not yet read it. > Decision record: [ADR-0085](../decisions/ADR-0085-public-provider-contract.md) diff --git a/tests/codex-integration/codex-auth-context.test.ts b/tests/codex-integration/codex-auth-context.test.ts index b3d8d83ae4c..bf4da6cfd41 100644 --- a/tests/codex-integration/codex-auth-context.test.ts +++ b/tests/codex-integration/codex-auth-context.test.ts @@ -915,19 +915,19 @@ describe("Codex auth context", () => { const resolved = await resolveCodexAuthContext(headers, cfg, "pool"); expect(resolved).toMatchObject({ kind: "pool", accountId: "pool-a" }); if (resolved.kind !== "pool") throw new Error("expected pool context"); - // The parent used to BE the key, so every child of one parent shared a single binding - // entry and none of them could hold one of their own. A child now keys as its own - // conversation; the parent qualifies placement, not identity. + // The parent used to BE the key (#4546 wp8); since #4780 the tree is the binding unit and + // the session is the cohort. What matters here is that the key stays opaque, derives from + // the session rather than any caller-supplied identifier, and is stable across turns. expect(resolved.affinityKey?.startsWith("app:")).toBe(true); expect(resolved.affinityKey).not.toContain("canonical-parent-thread"); expect(resolved.affinityKey).not.toContain("desktop-session-private"); expect(resolved.affinityKey).not.toContain("desktop-thread-private"); - // Stable across turns that drop the parent header: the key is the session/thread pair. + // Stable across turns that drop the parent header: the cohort is the session. expect(resolved.affinityKey).toBe(codexPoolAffinityKey(new Headers({ "session-id": "desktop-session-private", "thread-id": "desktop-thread-private", }))); - // And distinct from the parent's own lane, which is what a parent-only request rides. + // And distinct from a session-less parent-only lane, which anchors on the parent instead. expect(resolved.affinityKey).not.toBe(codexPoolAffinityKey(new Headers({ "x-codex-parent-thread-id": "canonical-parent-thread", }))); diff --git a/tests/codex-integration/codex-lineage-placement.test.ts b/tests/codex-integration/codex-lineage-placement.test.ts index f41e519c8eb..496c49e2929 100644 --- a/tests/codex-integration/codex-lineage-placement.test.ts +++ b/tests/codex-integration/codex-lineage-placement.test.ts @@ -142,32 +142,15 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { } }); - test("every thread keys as itself, and the unbound set is unchanged", () => { - const rootKey = codexPoolAffinityKey(rootHeaders())!; - const childKey = codexPoolAffinityKey(childHeaders("child-1"))!; - const grandchildKey = codexPoolAffinityKey(childHeaders("grand-1", "child-1"))!; - for (const key of [rootKey, childKey, grandchildKey]) { - expect(key.startsWith("app:")).toBe(true); - } - // The three used to be two: both children collapsed onto the raw parent id. - expect(new Set([rootKey, childKey, grandchildKey]).size).toBe(3); - // A child keys as its own conversation whether or not this turn names the parent, which is - // what lets it hold a binding of its own across a fan-out. - expect(childKey).toBe(codexPoolAffinityKey(new Headers({ "session-id": "sess", "thread-id": "child-1" }))); - // A request naming only a parent rides that parent's lane. With the session in hand that lane - // is derivable, and it IS the parent's own key -- no record required. - expect(codexPoolAffinityKey(new Headers({ - "session-id": "sess", "x-codex-parent-thread-id": "root", - }))).toBe(rootKey); - // Without the session and without a recorded parent there is nothing to reproduce it from, - // so the bare parent lane is its own key. The recorded case is the test below. - expect(codexPoolAffinityKey(new Headers({ "x-codex-parent-thread-id": "root" }))).not.toBe(rootKey); - // Unchanged from before #4546: which requests bind at all did not move. A bare thread-id - // with neither a session nor a parent still has no family anchor and stays unbound. - expect(codexPoolAffinityKey(new Headers({ "thread-id": "lone" }))).toBeUndefined(); - expect(codexPoolAffinityKey(new Headers())).toBeUndefined(); - expect(codexPoolAffinityKey(new Headers({ "x-codex-parent-thread-id": "p".repeat(513) }))).toBeUndefined(); - }); + /** + * The per-thread keying this unit introduced was retired by #4780, which makes the tree the + * binding unit. What that test asserted now lives in the cohort block at the end of this file: + * the key shape and the unchanged unbound set are pinned there, and the grandchild-orphan + * property wp8 was written to prevent is pinned there too, as a property rather than as a + * consequence of per-thread keys. The placement tests below are rewritten for the same reason. + * + * A request naming only a parent still rides that parent's lane; that case keeps its own test. + */ test("lineage resolves the root transitively and stays inside its auth scope", () => { const root = recordCodexThreadLineage(rootHeaders(), NOW)!; @@ -197,8 +180,11 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { }); test("the table is bounded in both dimensions, not just per scope", () => { + // One cohort per index. Members of ONE tree now share a conversation key, so a fixture that + // varied only the thread id would write every record under a single key and the eviction + // probe below would read the newest record back through the oldest key. const keyFor = (index: number) => recordCodexThreadLineage( - new Headers({ "session-id": "bulk", "thread-id": `bulk-${index}` }), NOW, + new Headers({ "session-id": `bulk-${index}`, "thread-id": `bulk-${index}` }), NOW, )!.conversationKey; const oldest = keyFor(0); for (let index = 1; index <= CODEX_LINEAGE_MAX_ENTRIES; index += 1) keyFor(index); @@ -221,32 +207,7 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { expect(codexThreadLineageLookup(heldKey, codexLineageScopeKey(held), NOW)).toBeUndefined(); }); - test("a child with no binding starts on the parent's account under its OWN key", () => { - const config = makeConfig(); - updateAccountQuota("a", 10); - updateAccountQuota("b", 20); - updateAccountQuota("c", 30); - const root = recordCodexThreadLineage(rootHeaders(), NOW)!; - expect(resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW)) - .toMatchObject({ status: "selected", accountId: "a" }); - - const child = recordCodexThreadLineage(childHeaders("child-1"), NOW)!; - // The reason carries the proof here: a cold pick would also have chosen the coolest - // account. The tests below make the ACCOUNT itself the discriminator. - expect(resolveCodexAccountForThreadDetailed( - child.conversationKey, config, NOW, undefined, undefined, undefined, child, - )).toMatchObject({ - status: "selected", - accountId: "a", - affinity: { move: "new_bind", reason: "lineage_parent" }, - }); - // An independent binding, not a root-wide pin: the child's next turn reuses its own entry - // without consulting the family again. - expect(resolveCodexAccountForThreadDetailed(child.conversationKey, config, NOW + 1)) - .toMatchObject({ status: "selected", accountId: "a", affinity: { move: "reused", reason: "healthy" } }); - }); - - test("a new child follows the account ACTUALLY serving the parent, detour included", () => { + test("a child is served where its cohort is being served, detour included", () => { const config = makeConfig(); updateAccountQuota("a", 10); updateAccountQuota("b", 20); @@ -263,18 +224,21 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { affinity: { move: "detour", reason: "transient" }, }); - // The child starts where the parent is being served NOW (b), not at its stale home (a). + // The child shares the cohort's binding, so it takes the same detour rather than starting + // anywhere of its own. Under per-thread keys this needed a placement hint to reach b; the + // cohort key makes it the same binding, so there is nothing to place. const child = recordCodexThreadLineage(childHeaders("child-1"), NOW + 2)!; + expect(child.conversationKey).toBe(root.conversationKey); expect(resolveCodexAccountForThreadDetailed( child.conversationKey, config, NOW + 2, undefined, undefined, undefined, child, )).toMatchObject({ status: "selected", accountId: "b", - affinity: { move: "new_bind", reason: "lineage_parent" }, + affinity: { move: "detour", reason: "transient" }, }); }); - test("a later move of the parent does not drag an already-bound child", () => { + test("a move of the cohort carries every member, including an already-active child", () => { const config = makeConfig(); updateAccountQuota("a", 10); updateAccountQuota("b", 20); @@ -284,56 +248,46 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { streakTransientFailures(config, "a", NOW); resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW + 1); - // Which account serves the parent at any moment is the quota strategy's business, not this - // layer's. What this layer promises is relative, so it is asserted relative to what actually - // happened rather than against account names predicted from a fixture nobody ran. const parentAtPlacement = resolveCodexAccountForThreadDetailed( root.conversationKey, config, NOW + 2, ).accountId; - // The child binds to the account actually SERVING its parent, detour included. const child = recordCodexThreadLineage(childHeaders("child-1"), NOW + 2)!; const childPlacement = resolveCodexAccountForThreadDetailed( child.conversationKey, config, NOW + 2, undefined, undefined, undefined, child, ); expect(childPlacement).toMatchObject({ status: "selected" }); expect(childPlacement.accountId).toBe(parentAtPlacement); - const childBoundTo = childPlacement.accountId; - // Now the parent moves for its OWN reason: a quota refusal retires its binding. This is the - // parent's move, not the family's. + // A quota refusal retires the binding. Under per-thread keys this was the PARENT's move and + // the asymmetry test pinned that an already-bound child was not dragged by it. #4780 retires + // that asymmetry deliberately: there is one binding, so the move is the cohort's and every + // member is on the other side of it. That is the cost of cohort cache locality, and it is + // the behaviour the invariant "same prompt_cache_key, same account" requires. updateAccountQuota("c", 5); recordCodexUpstreamOutcome(config, "a", 429, { now: NOW + 3 }); const parentAfterMove = resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW + 3); expect(parentAfterMove).toMatchObject({ status: "selected" }); - // Where the parent lands is the quota strategy's decision and may legitimately be the same - // account the child already holds, so nothing is asserted about the destination here. - // THE ASYMMETRY, which is the whole point of this test: the already-bound child is untouched - // by the parent's move. It reuses its own binding rather than being dragged. const childAfterParentMoved = resolveCodexAccountForThreadDetailed( child.conversationKey, config, NOW + 4, ); - expect(childAfterParentMoved).toMatchObject({ - status: "selected", - affinity: { move: "reused", reason: "healthy" }, - }); - expect(childAfterParentMoved.accountId).toBe(childBoundTo); + expect(childAfterParentMoved).toMatchObject({ status: "selected" }); + // Where the cohort lands is the quota strategy's decision, so this is asserted relative to + // what actually happened rather than against an account name predicted from the fixture. + expect(childAfterParentMoved.accountId).toBe(parentAfterMove.accountId); - // A NEW child, however, reads the parent's CURRENT account rather than the one its sibling - // holds, which is the other half of the same rule. + // A later sibling joins the same binding rather than being placed against it. const lateChild = recordCodexThreadLineage(childHeaders("child-2"), NOW + 5)!; + expect(lateChild.conversationKey).toBe(root.conversationKey); const latePlacement = resolveCodexAccountForThreadDetailed( lateChild.conversationKey, config, NOW + 5, undefined, undefined, undefined, lateChild, ); - expect(latePlacement).toMatchObject({ - status: "selected", - affinity: { move: "new_bind", reason: "lineage_parent" }, - }); + expect(latePlacement).toMatchObject({ status: "selected" }); expect(latePlacement.accountId).toBe(parentAfterMove.accountId); }); - test("a compatible sibling places the child when the parent is not eligible", () => { + test("a cohort whose account became ineligible rebinds off it, once, for everyone", () => { const config = makeConfig(); updateAccountQuota("a", 10); updateAccountQuota("b", 20); @@ -342,37 +296,34 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { expect(resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW)) .toMatchObject({ status: "selected", accountId: "a" }); - // The parent keeps its binding on a, but a is no longer eligible to serve anyone. A stale - // home is worse than no hint, so the parent contributes nothing here. + // a is no longer eligible to serve anyone. A stale home is worse than no hint, so the + // cohort must leave it rather than keep resolving there. config.pausedCodexAccountIds = ["a"]; const sibling = recordCodexThreadLineage(childHeaders("child-1"), NOW + 1)!; const siblingPlacement = resolveCodexAccountForThreadDetailed( sibling.conversationKey, config, NOW + 1, undefined, undefined, undefined, sibling, ); expect(siblingPlacement).toMatchObject({ status: "selected" }); - // The point is the NEGATIVE: a paused parent is a stale home and must contribute nothing. - // Which account the ordinary rule then picks belongs to the quota strategy. + // The point is the NEGATIVE: a paused account must contribute nothing. Which account the + // ordinary rule then picks belongs to the quota strategy. expect(siblingPlacement.affinity?.reason).not.toBe("lineage_parent"); expect(siblingPlacement.accountId).not.toBe("a"); - // Make an unrelated cold thread prefer a DIFFERENT account, so the orphan landing on its - // sibling's account cannot be explained by the ordinary cold rule agreeing by accident. + // Make an unrelated cold thread prefer a DIFFERENT account, so a later member landing with + // its cohort cannot be explained by the ordinary cold rule agreeing by accident. updateAccountQuota("c", 1); const coldPick = resolveCodexAccountForThreadDetailed("unrelated-cold-thread", config, NOW + 2); expect(coldPick).toMatchObject({ status: "selected" }); const orphan = recordCodexThreadLineage(childHeaders("child-2"), NOW + 2)!; - expect(orphan.siblingConversationKeys).toContain(sibling.conversationKey); + expect(orphan.conversationKey).toBe(sibling.conversationKey); const orphanPlacement = resolveCodexAccountForThreadDetailed( orphan.conversationKey, config, NOW + 2, undefined, undefined, undefined, orphan, ); - // The orphan follows its SIBLING, which is the reachable half of the family when the parent - // is not eligible. Asserted against the sibling's actual placement rather than an account - // name predicted from the quota fixture. + // It lands with its cohort rather than taking the cold pick, which is what the shared + // binding buys. Asserted against the sibling's actual placement rather than an account name + // predicted from the quota fixture. expect(orphanPlacement.accountId).toBe(siblingPlacement.accountId); - expect(orphanPlacement).toMatchObject({ - status: "selected", - affinity: { move: "new_bind", reason: "lineage_sibling" }, - }); + expect(orphanPlacement).toMatchObject({ status: "selected" }); }); test("no known family account falls back to ordinary cold placement", () => { @@ -455,7 +406,7 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { }); }); - test("a child follows the parent's MODEL detour, not a home account that cannot serve it", () => { + test("a cohort's members are detoured together when the binding cannot serve the model", () => { const config = makeConfig(); updateAccountQuota("a", 10); updateAccountQuota("b", 20); @@ -472,16 +423,24 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { expect(resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW, undefined, roster, modelId)) .toMatchObject({ status: "selected", accountId: "b" }); - // Make c the cold pick inside the roster, so b is reachable only through the detour. + // Move the roster's preferred account, so what is asserted cannot be satisfied by a stale + // constant. Under per-thread keys the child was an unbound thread that had to be PLACED on + // the parent's current detour target. Under cohort keying it shares the binding, so the + // property is simply that both members are served in the same place at the same moment. updateAccountQuota("b", 40); const child = recordCodexThreadLineage(childHeaders("child-1"), NOW + 1)!; - expect(resolveCodexAccountForThreadDetailed( + expect(child.conversationKey).toBe(root.conversationKey); + const rootServed = resolveCodexAccountForThreadDetailed( + root.conversationKey, config, NOW + 1, undefined, roster, modelId, + ); + const childServed = resolveCodexAccountForThreadDetailed( child.conversationKey, config, NOW + 1, undefined, roster, modelId, child, - )).toMatchObject({ - status: "selected", - accountId: "b", - affinity: { move: "new_bind", reason: "lineage_parent" }, - }); + ); + expect(rootServed).toMatchObject({ status: "selected" }); + expect(childServed).toMatchObject({ status: "selected" }); + // The binding sits on a, which this roster cannot serve, so both are detoured off it. + expect(childServed.accountId).not.toBe("a"); + expect(childServed.accountId).toBe(rootServed.accountId); }); test("a preview reads the family only for a request that may own Pool state", () => { @@ -499,10 +458,17 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { expect(previewCodexPoolLineage(callerOwned, config, { requestScopedMainCredential: true })) .toBeUndefined(); - // Read-only: the record belongs to the resolution that binds. A preview must not leave one - // behind for a request that turns out to own no Pool state at all. + // Read-only: the record belongs to the resolution that binds, so even an ELIGIBLE preview + // leaves nothing behind. Probed on a cohort nothing has recorded, because a member of an + // already-recorded tree would answer from its root's entry now that they share one key. + const unseen = new Headers({ + "session-id": "unseen-sess", + "thread-id": "unseen-child", + "x-codex-parent-thread-id": "unseen-root", + }); + expect(previewCodexPoolLineage(unseen, config)).toBeDefined(); expect(codexThreadLineageLookup( - codexPoolAffinityKey(child)!, codexLineageScopeKey(child), NOW, + codexPoolAffinityKey(unseen)!, codexLineageScopeKey(unseen), NOW, )).toBeUndefined(); }); @@ -518,3 +484,148 @@ describe("codex thread lineage and first placement (#4546 wp8)", () => { expect(codexLineageWorkflowLane(new Headers({ "thread-id": "child-9" }), NOW)).toBe("worker"); }); }); + +/** + * #4780. The binding unit is the COHORT the client already declares, not the thread. + * + * Upstream keys its prompt cache on something the whole tree shares. `prompt_cache_key()` + * returns `responses_metadata.session_id`, or `{source}:{parent_thread_id}` for an internal + * session; `AgentControl.session_id` "is equal to the root thread's ID" and that one control + * handle is shared with every sub-agent spawned from the root; and the upstream suite asserts + * root and child carrying DIFFERENT thread ids while sending the SAME `promptCacheKey`. + * openai/codex#44862 went further on 2026-09-11, making an ephemeral fork inherit its parent's + * session id for exactly this reason. + * + * While the proxy keyed per thread, two requests could carry an identical `prompt_cache_key` and + * be served by different accounts. The split-off member's key asserts a warm prefix that is + * deterministically cold on its account, so the prompt is replayed in full and the cache can + * never hit. Nothing fails; only tokens burn. + * + * THIS IS NOT A REVERT OF wp8. wp8 fixed a different defect: a child keyed under the RAW parent + * id, producing one shared entry unrelated to the root's own `app:HMAC(session, thread)` + * binding, so a grandchild keying on its own parent landed on a key nobody had ever bound. A + * cohort key has no such incoherence, because the root's own binding IS the cohort key. The + * orphan test below exists to prove that rather than assert it. + */ +describe("cohort pool affinity (#4780)", () => { + beforeEach(() => { + installScratchHome(); + clearThreadAccountMap(); + clearCodexUpstreamHealth(); + clearCodexThreadLineageForTests(); + clearPoolRotationState(); + clearAccountQuota(); + for (const id of ACCOUNT_IDS) saveTestCredential(id); + }); + + afterEach(async () => { + try { + clearAccountQuota(); + clearCodexUpstreamHealth(); + clearThreadAccountMap(); + clearCodexThreadLineageForTests(); + clearPoolRotationState(); + } finally { + await removeScratchHome(); + } + }); + + test("one conversation tree resolves to one affinity key", () => { + // The three ids upstream would send an identical prompt_cache_key for. + const rootKey = codexPoolAffinityKey(rootHeaders())!; + const childKey = codexPoolAffinityKey(childHeaders("child-1"))!; + const siblingKey = codexPoolAffinityKey(childHeaders("child-2"))!; + const grandchildKey = codexPoolAffinityKey(childHeaders("grand-1", "child-1"))!; + + expect(rootKey.startsWith("app:")).toBe(true); + expect(new Set([rootKey, childKey, siblingKey, grandchildKey]).size).toBe(1); + + // A different tree is a different cohort, so this is a cohort key and not a constant. + expect(codexPoolAffinityKey(new Headers({ + "session-id": "other-sess", "thread-id": "root", + }))).not.toBe(rootKey); + // The raw session id is never the key; it is still HMAC'd under the process-local secret. + expect(rootKey).not.toContain("sess"); + }); + + test("a grandchild never lands on a key nobody bound", () => { + // wp8's defect, restated as the property that must hold under cohort keying. The root binds + // first; every later member of the tree must resolve to the key the root is already on. + const root = recordCodexThreadLineage(rootHeaders(), NOW)!; + const child = recordCodexThreadLineage(childHeaders("child-1"), NOW)!; + const grandchild = recordCodexThreadLineage(childHeaders("grand-1", "child-1"), NOW)!; + expect(child.conversationKey).toBe(root.conversationKey); + expect(grandchild.conversationKey).toBe(root.conversationKey); + expect(grandchild.rootSessionKey).toBe(root.conversationKey); + + // The session-less chain is the case that could still split, because each depth would + // otherwise anchor on its own parent. A recorded parent carries the cohort down. + const bare = (threadId: string, parentId: string) => new Headers({ + "thread-id": threadId, "x-codex-parent-thread-id": parentId, + }); + const bareChild = recordCodexThreadLineage(bare("b-child", "b-root"), NOW)!; + const bareGrandchild = recordCodexThreadLineage(bare("b-grand", "b-child"), NOW)!; + expect(bareGrandchild.conversationKey).toBe(bareChild.conversationKey); + // ...and it is a cohort of its own, not folded into the session-keyed tree above. + expect(bareChild.conversationKey).not.toBe(root.conversationKey); + }); + + test("which requests bind at all is unchanged", () => { + // Deliberately untouched by #4780: only the VALUE of the key moves, never the set of + // requests that produce one. A bare thread-id still has no family anchor. + expect(codexPoolAffinityKey(new Headers({ "thread-id": "lone" }))).toBeUndefined(); + expect(codexPoolAffinityKey(new Headers())).toBeUndefined(); + expect(codexPoolAffinityKey(new Headers({ "x-codex-parent-thread-id": "p".repeat(513) }))) + .toBeUndefined(); + expect(codexPoolAffinityKey(new Headers({ "session-id": "s".repeat(513), "thread-id": "t" }))) + .toBeUndefined(); + }); + + test("a cohort key stays inside its authenticated scope", () => { + // Two callers presenting the same session must not share a binding; the scope HMAC is what + // keeps that true, and cohort keying must not have widened it. + const mine = rootHeaders(); + const theirs = rootHeaders(); + theirs.set("authorization", "Bearer someone-else"); + expect(codexLineageScopeKey(theirs)).not.toBe(codexLineageScopeKey(mine)); + + const root = recordCodexThreadLineage(mine, NOW)!; + expect(codexThreadLineageLookup(root.conversationKey, codexLineageScopeKey(theirs), NOW)) + .toBeUndefined(); + }); + + test("a tree shares one binding, so a move of any member moves the tree", () => { + // The behaviour change this issue asks for, stated as the tradeoff it is: the tree gains + // cache locality and gives up per-thread placement independence. A member cannot be served + // by an account other than the one its cohort is bound to. + const config = makeConfig(); + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + + const root = recordCodexThreadLineage(rootHeaders(), NOW)!; + expect(resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW)) + .toMatchObject({ status: "selected", accountId: "a" }); + + // A child arriving later is not an unbound thread any more: its cohort is already on a. + const child = recordCodexThreadLineage(childHeaders("child-1"), NOW + 1)!; + expect(resolveCodexAccountForThreadDetailed( + child.conversationKey, config, NOW + 1, undefined, undefined, undefined, child, + )).toMatchObject({ + status: "selected", + accountId: "a", + affinity: { move: "reused" }, + }); + + // And when the cohort moves, every member moves with it, which is the whole point: the + // prompt_cache_key they all send keeps naming one account. + streakTransientFailures(config, "a", NOW + 2); + const moved = resolveCodexAccountForThreadDetailed( + child.conversationKey, config, NOW + 3, undefined, undefined, undefined, child, + ); + expect(moved).toMatchObject({ status: "selected" }); + expect(moved.accountId).not.toBe("a"); + expect(resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW + 3)) + .toMatchObject({ status: "selected", accountId: moved.accountId }); + }); +});