diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index d69ce45ac9..4b18cc3a87 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -470,6 +470,7 @@ "codex-integration-record.test.ts": "codex-integration", "codex-journal.test.ts": "codex-integration", "codex-legacy-config-keys.test.ts": "codex-integration", + "codex-lineage-placement.test.ts": "codex-integration", "codex-log-guard-coderabbit.test.ts": "codex-integration", "codex-log-guard-doctor-coderabbit.test.ts": "codex-integration", "codex-log-guard-doctor-protection.test.ts": "codex-integration", diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 7b0eb72ecc..54431e0b7c 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -1,5 +1,5 @@ import type { PoolQuotaWriter } from "./quota-types"; -import { createHash, createHmac, randomBytes, timingSafeEqual } from "node:crypto"; +import { createHash, timingSafeEqual } from "node:crypto"; import { CodexCredentialGenerationConflictError, CodexCredentialRefreshLockTimeoutError, @@ -40,6 +40,12 @@ import { resolveCodexAccountForThreadDetailed, type CodexAffinityDecision, } from "./routing"; +import { + codexConversationIdentity, + recordCodexThreadLineage, + resolveCodexThreadLineage, + type CodexThreadLineage, +} from "./lineage"; import { entitledCodexAccountIdsForModel, isDirectCallerEntitledToCodexModel, @@ -53,7 +59,6 @@ import { CODEX_UNKNOWN_USAGE_SCORE, getAccountQuota, parseUsageQuota, parseMainP import type { CodexAccountMode, OcxConfig, OcxProviderConfig } from "../types"; import { FORWARD_HEADERS } from "../adapters/openai-responses"; import { captureConfigGeneration } from "../lib/state-store-sweeper"; -import { retainedUtf8Bytes } from "../lib/admission"; import { extractAccountId, extractEmail } from "../oauth/chatgpt"; import { getMainAccountHardLockStatus, isMainAccountHardLocked } from "./main-account-hard-lock"; import { @@ -70,9 +75,6 @@ import type { DataPlaneAdmission } from "../server/auth-cors"; import { getMainReserveAuthorization, isMainReserveAuthorizationLive, nativeUserIdClaims, type MainReserveAuthorization } from "./reserve-availability"; import { UpstreamRetryEvidenceError } from "../lib/upstream-retry"; -const CODEX_AFFINITY_COMPONENT_MAX_BYTES = 512; -const CODEX_APP_AFFINITY_KEY = randomBytes(32); - /** * A request-owned bearer cannot inspect the physical main credential for its plan, but cached * WHAM usage is still valid routing evidence for the same logical main account. Score it with @@ -88,32 +90,90 @@ function requestOwnedMainPinHasQuotaHeadroom(config: OcxConfig): boolean { return usage >= CODEX_UNKNOWN_USAGE_SCORE || usage < threshold; } -function boundedCodexAffinityComponent(value: string | null): string | undefined { - const normalized = value?.trim(); - if (!normalized) return undefined; - if (retainedUtf8Bytes(normalized) > CODEX_AFFINITY_COMPONENT_MAX_BYTES) return undefined; - return normalized; +/** + * Every thread keys as ITSELF, never as its parent (#4546, wp8). + * + * 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. + * + * 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. + * + * 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. + */ +export function codexPoolAffinityKey(headers: Headers, now = Date.now()): string | undefined { + // `now` is threaded rather than read inside because a parent-only turn resolves its key + // through the recorded lineage, and that record is TTL-bounded: a caller working against a + // fixed clock would otherwise see a live record as expired and fall back to a key the parent + // never bound under. + return codexConversationIdentity(headers, now)?.conversationKey; +} + +/** What a caller needs to know to answer the Pool-state question below before auth has run. */ +export interface CodexPoolStateEligibility { + /** An exact account selector from the route, i.e. `options.accountId` here. */ + readonly accountId?: string; + readonly modelId?: string; + readonly admission?: Pick; + /** The caller presented its own forwardable ChatGPT credential for this route. */ + readonly requestScopedMainCredential?: boolean; +} + +/** The one expression both the resolution below and any preview must agree on. */ +function poolStateEligible( + fixedAccountId: string | undefined, + requestScopedMainCredential: boolean, +): boolean { + return fixedAccountId === undefined && !requestScopedMainCredential; } /** - * Preserve Codex's parent-thread affinity when present. Desktop App requests can omit that - * header while retaining a stable session/thread pair, so derive an opaque process-local key - * only from the complete bounded pair. Raw identifiers and durable hashes never enter Pool state. + * May this request own Pool affinity state at all? + * + * Two credentials authenticate outside the Pool: an exact account selector (including the + * Reserve pin) and a request-owned main bearer, which exists for one request and must never + * fold into durable account state. Neither may read or write a binding, so neither may read + * or write LINEAGE either. + * + * Exported so that a preview asks the question with the code that answers it, instead of a + * restatement that can drift. It drifted once already: preview read a family relation from raw + * request headers before this function had decided anything, so it could follow a Pool family + * binding while the resolution below deliberately created no affinity -- and model fallback then + * evaluated eligibility against an account the request would never be authenticated as. + */ +export function codexPoolStateEligible( + headers: Headers, + policy: CodexAuthPolicyConfig | undefined, + options: CodexPoolStateEligibility = {}, +): boolean { + const reserve = requiresReserveAuthorization(policy, options.modelId, options.admission); + return poolStateEligible( + reserve ? MAIN_CODEX_ACCOUNT_ID : options.accountId, + options.requestScopedMainCredential === true && hasCallerCodexBearer(headers), + ); +} + +/** + * The lineage a PREVIEW is allowed to see: read-only, and only for a request that may hold Pool + * state. Recording is left to the resolution that actually binds, so a preview can never leave a + * record behind for a request that turned out to own no Pool state at all. */ -export function codexPoolAffinityKey(headers: Headers): string | undefined { - const parentThreadId = boundedCodexAffinityComponent(headers.get("x-codex-parent-thread-id")); - if (parentThreadId) return parentThreadId; - - const sessionId = boundedCodexAffinityComponent(headers.get("session-id")); - const threadId = boundedCodexAffinityComponent(headers.get("thread-id")); - if (!sessionId || !threadId) return undefined; - - return `app:${createHmac("sha256", CODEX_APP_AFFINITY_KEY) - .update("opencodex-app-pool-affinity-v1\0") - .update(sessionId) - .update("\0") - .update(threadId) - .digest("base64url")}`; +export function previewCodexPoolLineage( + headers: Headers, + policy: CodexAuthPolicyConfig | undefined, + options: CodexPoolStateEligibility = {}, +): CodexThreadLineage | undefined { + return codexPoolStateEligible(headers, policy, options) + ? resolveCodexThreadLineage(headers) + : undefined; } export type CodexAuthContext = @@ -798,9 +858,15 @@ export async function resolveCodexAuthContext( // A caller bearer can still accompany a request that selects a configured Pool account. Do not // let that request read, delete, or create a file-main affinity binding while deciding whether a // stored account is available; only the stored credential selected below may own Pool state. - const affinityKey = fixedAccountId === undefined && !requestScopedMainCredential + const affinityKey = poolStateEligible(fixedAccountId, requestScopedMainCredential) ? codexPoolAffinityKey(headers) : undefined; + // The thread's family relation, recorded under the same condition as the key itself. A + // first-placing child consults it; a request-owned or fixed credential never enters Pool + // state, so it never enters lineage either. + const lineage = affinityKey !== undefined + ? recordCodexThreadLineage(headers) + : undefined; // Why this request is on this account, carried to the request log so a move reads as an event // instead of something inferred from account labels across lines (#4546). let affinityDecision: CodexAffinityDecision | undefined; @@ -873,6 +939,7 @@ export async function resolveCodexAuthContext( quotaScope, selectionOptions, options.modelId, + lineage, ); if (resolution.status === "expired") throw new CodexThreadAffinityExpiredError(resolution.accountId); const selected = resolution.status === "selected" ? resolution.accountId : null; diff --git a/src/codex/lineage.ts b/src/codex/lineage.ts new file mode 100644 index 0000000000..ff5233c04b --- /dev/null +++ b/src/codex/lineage.ts @@ -0,0 +1,458 @@ +/** + * Codex V2 conversation lineage: root, parent, child, grandchild (#4546, wp8). + * + * The pool affinity key used to prefer `x-codex-parent-thread-id`, which collapsed two different + * identities into one map. A root bound under `app:HMAC(session, thread)` while every child bound + * under the RAW parent id, so siblings shared one binding entry unrelated to the root's, and a + * grandchild keyed on its own parent landed on a key nobody had ever bound. The proxy therefore + * treated one workflow as unrelated strangers even while the provider saw a single prompt-cache + * family. + * + * This module records the real relation -- each thread's own conversation key, its immediate + * parent's thread id, and the transitive root -- scoped per authenticated caller, bounded, and + * process-local like the binding map it feeds. + * + * 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. + * - 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. + * - WORKER CLASSIFICATION, exposed but not rewired. Admission classifies header-only today: a + * request naming a parent plus a distinct `thread-id` is worker traffic, and a request without + * `thread-id` is interactive even when it belongs to a recorded fan-out. + * {@link codexLineageWorkflowLane} is the lineage-backed answer a later lane consumes. + * + * Scope is an HMAC of the caller's Authorization header under a process-local key, the same + * posture as the affinity key itself. Two callers presenting identical thread ids can never read + * each other's lineage, and no raw identifier or durable hash is stored. + * + * LIFETIME, stated plainly because the word "affinity" invites the opposite assumption: none of + * this survives the process. The binding map is in memory, and the HMAC key above is fresh random + * bytes taken at module load, so a restart does not merely forget the table -- it makes yesterday's + * keys unreproducible. This is a warm-start hint for the life of one proxy process, never durable + * account ownership, and nothing here should be read as a promise to a conversation that outlives + * a restart. + * + * The one upgrade that is neither a fresh start nor an untouched process is a code swap under a + * live conversation, where the binding map is still populated with entries made under the + * pre-#4546 RAW parent key. Silently rebinding those cold is the exact defect this module exists + * to prevent, so a request that names a parent carries {@link CodexThreadLineage.legacyConversationKey} + * -- the key the old rule would have returned -- and routing adopts that binding once under the new + * key and retires the legacy entry. It is a one-way migration, not a second lookup path. + */ +import { createHmac, randomBytes } from "node:crypto"; +import { retainedUtf8Bytes } from "../lib/admission"; + +const CODEX_LINEAGE_COMPONENT_MAX_BYTES = 512; +const CODEX_LINEAGE_KEY = randomBytes(32); + +/** + * Mirrors `CODEX_THREAD_AFFINITY_IDLE_TTL_MS` in ./routing. Deliberately duplicated rather than + * imported: lineage is a leaf module, and a value import from the routing module that consumes it + * would turn an erased type-only edge into a real cycle. + */ +export const CODEX_LINEAGE_IDLE_TTL_MS = 24 * 60 * 60_000; +/** Records per authenticated scope, on the order of the binding map's own 2048-entry cap. */ +export const CODEX_LINEAGE_MAX_ENTRIES = 2048; +/** Distinct authenticated callers retained. Without this the scope map is the unbounded one. */ +export const CODEX_LINEAGE_MAX_SCOPES = 64; +/** Sibling hints kept per parent, most recently used first. */ +export const CODEX_LINEAGE_MAX_SIBLINGS = 8; + +const LOCAL_LINEAGE_SCOPE = "local"; + +export type CodexWorkflowLane = "worker" | "interactive"; + +interface CodexLineageRecord { + threadId: string; + /** This thread's own pool binding key, byte-identical to what `codexPoolAffinityKey` returns. */ + conversationKey: string; + /** Immediate parent's raw thread id, retained once seen even if a later turn omits the header. */ + parentThreadId?: string; + /** Topmost ancestor's conversation key: a grandchild resolves to the root's, not its parent's. */ + rootSessionKey: string; + lastUsedAt: number; +} + +interface CodexLineageScope { + /** threadId -> record, iterated oldest-first so TTL pruning and eviction stay amortised O(1). */ + records: Map; + /** conversationKey -> threadId, so cost attribution is a lookup instead of a scan. */ + threadIdByConversationKey: Map; + /** parentThreadId -> child thread ids, most recent first, capped. */ + childThreadIdsByParent: Map; + lastUsedAt: number; +} + +/** + * What placement and cost attribution are allowed to see. `parentConversationKey` is the parent's + * OWN binding key -- either recorded, or derived from the shared session when the parent has not + * been seen yet -- never the raw header value the old affinity key returned. + */ +export interface CodexThreadLineage { + readonly conversationKey: string; + readonly rootSessionKey: string; + readonly parentThreadId?: string; + readonly parentConversationKey?: string; + /** + * The key the pre-#4546 rule would have returned for this request -- the RAW parent id -- when + * that differs from the key it binds under now. Present so routing can adopt a binding left by + * the old rule exactly once; see the lifetime note at the top of this file. + */ + readonly legacyConversationKey?: string; + /** Siblings under the same declared parent, most recently used first. */ + readonly siblingConversationKeys: readonly string[]; +} + +/** The identity the pool affinity key and the lineage record are both derived from. */ +export interface CodexConversationIdentity { + readonly conversationKey: string; + /** Thread id this request records under: its own, or the parent's on a parent-only request. */ + readonly recordThreadId: string; + readonly sessionId?: string; + readonly parentThreadId?: string; + /** Raw parent id, i.e. the key the pre-#4546 rule returned for this request. */ + readonly legacyConversationKey?: string; + /** True only when the request names a parent distinct from its own thread. */ + readonly declaresParent: boolean; +} + +const lineageByScope = new Map(); + +/** A record is only evidence while it is live; an idle-expired one answers like no record. */ +function liveLineageRecord( + scope: CodexLineageScope | undefined, + threadId: string, + now: number, +): CodexLineageRecord | undefined { + const record = scope?.records.get(threadId); + return record !== undefined && now - record.lastUsedAt <= CODEX_LINEAGE_IDLE_TTL_MS + ? record + : undefined; +} + +function boundedLineageComponent(value: string | null): string | undefined { + const normalized = value?.trim(); + if (!normalized) return undefined; + if (retainedUtf8Bytes(normalized) > CODEX_LINEAGE_COMPONENT_MAX_BYTES) return undefined; + return normalized; +} + +/** + * The single derivation behind both the pool affinity key and lineage records. Keeping it here, + * rather than duplicated at the call site, is what guarantees a record's `conversationKey` is + * byte-identical to the key the thread actually binds under. + */ +export function codexConversationKeyFor(familyId: string, threadId: string): string { + return `app:${createHmac("sha256", CODEX_LINEAGE_KEY) + .update("opencodex-app-pool-affinity-v1\0") + .update(familyId) + .update("\0") + .update(threadId) + .digest("base64url")}`; +} + +/** + * 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: + * + * - 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). + * + * 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. + */ +export function codexConversationIdentity( + headers: Headers, + now = Date.now(), +): CodexConversationIdentity | undefined { + const threadId = boundedLineageComponent(headers.get("thread-id")); + const sessionId = boundedLineageComponent(headers.get("session-id")); + const parentThreadId = boundedLineageComponent(headers.get("x-codex-parent-thread-id")); + + if (threadId === undefined) { + if (parentThreadId === undefined) return undefined; + const recorded = liveLineageRecord( + lineageByScope.get(codexLineageScopeKey(headers)), + parentThreadId, + now, + ); + return { + conversationKey: recorded?.conversationKey + ?? codexConversationKeyFor(sessionId ?? parentThreadId, parentThreadId), + recordThreadId: parentThreadId, + ...(sessionId !== undefined ? { sessionId } : {}), + legacyConversationKey: parentThreadId, + declaresParent: false, + }; + } + const familyId = sessionId ?? parentThreadId; + if (familyId === undefined) return undefined; + return { + conversationKey: codexConversationKeyFor(familyId, threadId), + recordThreadId: threadId, + ...(sessionId !== undefined ? { sessionId } : {}), + ...(parentThreadId !== undefined ? { parentThreadId } : {}), + ...(parentThreadId !== undefined ? { legacyConversationKey: parentThreadId } : {}), + declaresParent: parentThreadId !== undefined && parentThreadId !== threadId, + }; +} + +/** + * Which authenticated caller this request's lineage belongs to. The bearer is never stored; an + * unauthenticated (loopback-trusted) request lands in the single local scope. + */ +export function codexLineageScopeKey(headers: Headers): string { + const authorization = headers.get("authorization")?.trim(); + if (!authorization) return LOCAL_LINEAGE_SCOPE; + return `auth:${createHmac("sha256", CODEX_LINEAGE_KEY) + .update("opencodex-lineage-scope-v1\0") + .update(authorization) + .digest("base64url")}`; +} + +function dropLineageRecord(scope: CodexLineageScope, threadId: string): void { + const record = scope.records.get(threadId); + if (record === undefined) return; + scope.records.delete(threadId); + if (scope.threadIdByConversationKey.get(record.conversationKey) === threadId) { + scope.threadIdByConversationKey.delete(record.conversationKey); + } + if (record.parentThreadId === undefined) return; + const siblings = scope.childThreadIdsByParent.get(record.parentThreadId); + if (siblings === undefined) return; + const remaining = siblings.filter(id => id !== threadId); + if (remaining.length === 0) scope.childThreadIdsByParent.delete(record.parentThreadId); + else scope.childThreadIdsByParent.set(record.parentThreadId, remaining); +} + +/** Records are held in least-recently-used order, so the expired ones are a prefix. */ +function pruneLineageScope(scope: CodexLineageScope, now: number): void { + for (const [threadId, record] of scope.records) { + if (now - record.lastUsedAt <= CODEX_LINEAGE_IDLE_TTL_MS) break; + dropLineageRecord(scope, threadId); + } + while (scope.records.size > CODEX_LINEAGE_MAX_ENTRIES) { + const oldest = scope.records.keys().next(); + if (oldest.done === true) break; + dropLineageRecord(scope, oldest.value); + } +} + +function pruneLineageScopes(now: number): void { + for (const [scopeKey, scope] of lineageByScope) { + if (now - scope.lastUsedAt <= CODEX_LINEAGE_IDLE_TTL_MS) break; + lineageByScope.delete(scopeKey); + } + while (lineageByScope.size > CODEX_LINEAGE_MAX_SCOPES) { + const oldest = lineageByScope.keys().next(); + if (oldest.done === true) break; + lineageByScope.delete(oldest.value); + } +} + +function touchLineageScope(scopeKey: string, now: number): CodexLineageScope { + const existing = lineageByScope.get(scopeKey); + const scope: CodexLineageScope = existing ?? { + records: new Map(), + threadIdByConversationKey: new Map(), + childThreadIdsByParent: new Map(), + lastUsedAt: now, + }; + if (existing !== undefined) lineageByScope.delete(scopeKey); + scope.lastUsedAt = now; + lineageByScope.set(scopeKey, scope); + pruneLineageScopes(now); + pruneLineageScope(scope, now); + return scope; +} + +/** + * The lineage view for one identity inside one scope, computed without writing anything. + * + * A parent seen for the first time through one of its children is derived rather than invented: + * the child knows the shared session, so HMAC(session, parent) reproduces the key the parent + * binds under. Once the parent has actually been recorded, its own key wins. + * + * Depth is transitive by construction -- a grandchild inherits its parent's resolved root instead + * of re-deriving one hop -- so a workflow never scatters across several roots. + */ +function lineageFor( + scope: CodexLineageScope | undefined, + identity: CodexConversationIdentity, + now: number, +): CodexThreadLineage { + const previous = liveLineageRecord(scope, identity.recordThreadId, now); + // A turn that omits the parent header does not orphan a thread whose parent is already known. + // That retention is the whole of the lineage-backed worker answer below. + const parentThreadId = identity.declaresParent + ? identity.parentThreadId + : previous?.parentThreadId; + + const parentRecord = parentThreadId !== undefined + ? liveLineageRecord(scope, parentThreadId, now) + : undefined; + const parentConversationKey = parentThreadId === undefined + ? undefined + : parentRecord?.conversationKey + ?? codexConversationKeyFor(identity.sessionId ?? parentThreadId, parentThreadId); + const rootSessionKey = parentConversationKey === undefined + ? identity.conversationKey + : parentRecord?.rootSessionKey ?? parentConversationKey; + + const siblingConversationKeys: string[] = []; + if (parentThreadId !== undefined && scope !== undefined) { + for (const siblingThreadId of scope.childThreadIdsByParent.get(parentThreadId) ?? []) { + if (siblingThreadId === identity.recordThreadId) continue; + const sibling = liveLineageRecord(scope, siblingThreadId, now); + if (sibling !== undefined) siblingConversationKeys.push(sibling.conversationKey); + } + } + + // Only a key the old rule would have produced AND that this request no longer uses is a + // migration candidate. A root's key is unchanged, so it never carries one. + const legacyConversationKey = identity.legacyConversationKey !== undefined + && identity.legacyConversationKey !== identity.conversationKey + ? identity.legacyConversationKey + : undefined; + + return { + conversationKey: identity.conversationKey, + rootSessionKey, + ...(parentThreadId !== undefined ? { parentThreadId } : {}), + ...(parentConversationKey !== undefined ? { parentConversationKey } : {}), + ...(legacyConversationKey !== undefined ? { legacyConversationKey } : {}), + siblingConversationKeys, + }; +} + +/** + * Read this request's lineage without recording it. + * + * A preview must see what the final resolution will see, but it must not be the thing that + * creates the record: preview runs before auth has decided whether this request may hold Pool + * state at all, and a record written there would outlive a decision to hold none. + */ +export function resolveCodexThreadLineage( + headers: Headers, + now = Date.now(), +): CodexThreadLineage | undefined { + const identity = codexConversationIdentity(headers, now); + if (identity === undefined) return undefined; + return lineageFor(lineageByScope.get(codexLineageScopeKey(headers)), identity, now); +} + +/** Record this request's thread relation and return the resolved lineage. */ +export function recordCodexThreadLineage( + headers: Headers, + now = Date.now(), +): CodexThreadLineage | undefined { + const identity = codexConversationIdentity(headers, now); + if (identity === undefined) return undefined; + const scope = touchLineageScope(codexLineageScopeKey(headers), now); + const lineage = lineageFor(scope, identity, now); + const parentThreadId = lineage.parentThreadId; + + // Re-insert rather than mutate: the records map doubles as the LRU order. + dropLineageRecord(scope, identity.recordThreadId); + scope.records.set(identity.recordThreadId, { + threadId: identity.recordThreadId, + conversationKey: identity.conversationKey, + ...(parentThreadId !== undefined ? { parentThreadId } : {}), + rootSessionKey: lineage.rootSessionKey, + lastUsedAt: now, + }); + scope.threadIdByConversationKey.set(identity.conversationKey, identity.recordThreadId); + if (parentThreadId !== undefined) { + const siblings = (scope.childThreadIdsByParent.get(parentThreadId) ?? []) + .filter(id => id !== identity.recordThreadId); + siblings.unshift(identity.recordThreadId); + scope.childThreadIdsByParent.set(parentThreadId, siblings.slice(0, CODEX_LINEAGE_MAX_SIBLINGS)); + } + pruneLineageScope(scope, now); + + return lineage; +} + +/** + * Cost-attribution lookup for other layers: which root workflow owns this conversation key. + * Scoped like the records, so a caller can only ever resolve inside its own scope, and read-only + * -- reading a lineage for accounting must not extend its lifetime. + */ +export function codexThreadLineageLookup( + conversationKey: string, + scopeKey: string, + now = Date.now(), +): { conversationKey: string; rootSessionKey: string; parentThreadId?: string } | undefined { + const scope = lineageByScope.get(scopeKey); + if (scope === undefined) return undefined; + const threadId = scope.threadIdByConversationKey.get(conversationKey); + if (threadId === undefined) return undefined; + const record = scope.records.get(threadId); + if (record === undefined || now - record.lastUsedAt > CODEX_LINEAGE_IDLE_TTL_MS) return undefined; + return { + conversationKey: record.conversationKey, + rootSessionKey: record.rootSessionKey, + ...(record.parentThreadId !== undefined ? { parentThreadId: record.parentThreadId } : {}), + }; +} + +/** + * The root a request's spend belongs to, for a caller holding headers rather than a key. An + * unrecorded conversation is its own root, so this never answers null for a bindable request and + * an accounting layer has no reason to invent one. + */ +export function codexLineageRootForRequest(headers: Headers, now = Date.now()): string | undefined { + const identity = codexConversationIdentity(headers, now); + if (identity === undefined) return undefined; + return codexThreadLineageLookup(identity.conversationKey, codexLineageScopeKey(headers), now) + ?.rootSessionKey + ?? identity.conversationKey; +} + +/** + * The lineage-backed worker/interactive answer. + * + * Admission classifies header-only today: a request is worker traffic only when it names a parent + * AND a distinct `thread-id`, so a fan-out turn that stopped sending the parent header reads as + * interactive. This keeps that rule and adds what the headers could not say -- a thread already + * recorded with a parent is worker traffic. Nothing here changes admission; a later lane consumes + * it. + */ +export function codexLineageWorkflowLane(headers: Headers, now = Date.now()): CodexWorkflowLane { + const threadId = boundedLineageComponent(headers.get("thread-id")); + const parentThreadId = boundedLineageComponent(headers.get("x-codex-parent-thread-id")); + if (parentThreadId !== undefined && threadId !== undefined && threadId !== parentThreadId) { + return "worker"; + } + if (threadId === undefined) return "interactive"; + const record = lineageByScope.get(codexLineageScopeKey(headers))?.records.get(threadId); + return record !== undefined + && now - record.lastUsedAt <= CODEX_LINEAGE_IDLE_TTL_MS + && record.parentThreadId !== undefined + ? "worker" + : "interactive"; +} + +/** Test-only reset; production state is process-local and dies with the process. */ +export function clearCodexThreadLineageForTests(): void { + lineageByScope.clear(); +} diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 8fd3582fbe..d260122afa 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -37,6 +37,7 @@ import { captureConfigGeneration, type GenerationContext } from "../lib/state-st import { isCanonicalOpenAiForwardProvider } from "../providers/openai-tiers"; import { retainedUtf8Bytes } from "../lib/admission"; import { recordUpstreamHostFailure } from "./upstream-host-health"; +import type { CodexThreadLineage } from "./lineage"; import { clearAllCodexPoolRefreshFailures, isCodexPoolRefreshCooling } from "./pool-refresh-backoff"; @@ -98,7 +99,11 @@ export type CodexAffinityReason = | "quota_avoided" | "generation" | "expired" - | "model_lane"; + | "model_lane" + /** First placement followed the parent's CURRENT serving account (#4546, wp8). */ + | "lineage_parent" + /** First placement followed a compatible sibling's current serving account. */ + | "lineage_sibling"; export interface CodexAffinityDecision { move: CodexAffinityMove; @@ -1795,6 +1800,122 @@ function transientDetourAccount( : pickAlternateCodexAccount(config, entry.accountId, now, quotaScope, selectionOptions); } +/** + * Which account is ACTUALLY answering for one conversation key right now (#4546, wp8). + * + * First placement reads this, not the binding alone: a parent parked on a transient detour + * is being served by the detour, so a new child placed "where the parent lives" would miss + * the warm account by one hop. A dead binding, an expired hold, and an ineligible serving + * account all answer null -- the caller then tries a sibling, then falls back to cold + * placement, which is the correct order because a stale home is worse than no hint. + * + * "Right now" includes the MODEL lane. A parent whose home account is not entitled to this + * model is being served through a model-scoped detour, which is the same "serving, not stale + * home" case one level further in: reading only the ordinary binding would hand the child an + * account this request cannot use, and it would then start cold on the very model whose + * warm account the family already found. The detour scope embeds the model and the quota + * scope, so the entry consulted here is compatible by construction. + */ +function lineageServingAccountId( + conversationKey: string, + config: OcxConfig, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + modelId?: string, +): string | null { + const entry = (modelId !== undefined + ? getModelDetourAffinity(conversationKey, modelId, quotaScope) + : undefined) + ?? getThreadAffinity(conversationKey, quotaScope); + if (!entry || isThreadAffinityExpired(entry, now) || !isThreadAffinityGenerationLive(entry)) { + return null; + } + const holdLive = entry.transientHoldSince !== undefined && !isTransientHoldExpired(entry, now); + const serving = holdLive && entry.transientDetourAccountId !== undefined + ? entry.transientDetourAccountId + : entry.accountId; + return isCodexAccountSelectable(config, serving, now, quotaScope, selectionOptions) + && !hasUnrecoveredCodexQuotaRefusal(serving, quotaScope) + && !shouldFailover(config, serving, now) + && !isCodexAccountSoftAvoided(serving, now) + ? serving + : null; +} + +/** + * First placement only: where a child with NO binding of its own should start. Parent's + * current serving account first, then a compatible sibling's -- "compatible" meaning the + * same quota-scope slot, since a Reserve sibling says nothing about the shared lane. The + * child still binds under its own key; this is a hint for turn one, not a root-wide pin. + */ +function pickLineageServingAccount( + config: OcxConfig, + lineage: CodexThreadLineage, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + modelId?: string, +): { accountId: string; reason: CodexAffinityReason } | null { + if (lineage.parentConversationKey !== undefined) { + const parent = lineageServingAccountId( + lineage.parentConversationKey, config, now, quotaScope, selectionOptions, modelId, + ); + if (parent) return { accountId: parent, reason: "lineage_parent" }; + for (const siblingKey of lineage.siblingConversationKeys) { + const sibling = lineageServingAccountId( + siblingKey, config, now, quotaScope, selectionOptions, modelId, + ); + if (sibling) return { accountId: sibling, reason: "lineage_sibling" }; + } + } + return null; +} + +/** + * Move one scope's binding from the pre-#4546 RAW parent key onto the key this thread uses now. + * + * Bindings and the key that derives them are process-local, so an ordinary restart already + * discards every binding and there is nothing to migrate. The case this exists for is the + * narrow one: a code swap under a live conversation, where the map still holds entries made by + * the old rule. Rebinding those cold is precisely the defect the lineage work exists to prevent, + * so the conversation keeps its account and the legacy entry is retired in the same step. + * + * One way, once. The legacy entry is deleted even when it was dead on arrival, because nothing + * can reach it again under the new rule and an orphan only spends an LRU slot a live + * conversation needs. Only the account moves: a transient hold describes a failure happening + * right now, and the ordinary path re-derives it on this very request. + */ +function adoptLegacyAffinityForScope( + threadId: string, + legacyKey: string, + now: number, + scope: ThreadAffinityScope, +): void { + if (getThreadAffinityForScope(threadId, scope) !== undefined) return; + const legacy = getThreadAffinityForScope(legacyKey, scope); + if (legacy === undefined) return; + if (!isThreadAffinityExpired(legacy, now) && isThreadAffinityGenerationLive(legacy)) { + bindThreadAffinityForScope(threadId, legacy.accountId, now, scope); + } + deleteThreadAffinityForScope(legacyKey, scope); +} + +/** Both lanes of the legacy migration: the ordinary binding and this request's model detour. */ +function adoptLegacyLineageAffinity( + threadId: string, + lineage: CodexThreadLineage | undefined, + now: number, + quotaScope?: CodexQuotaScope, + modelId?: string, +): void { + const legacyKey = lineage?.legacyConversationKey; + if (legacyKey === undefined || legacyKey === threadId) return; + adoptLegacyAffinityForScope(threadId, legacyKey, now, threadAffinityScope(quotaScope)); + const detourScope = modelDetourAffinityScope(modelId, quotaScope); + if (detourScope) adoptLegacyAffinityForScope(threadId, legacyKey, now, detourScope); +} + /** Earliest future shared short/weekly reset; missing evidence and ties use usage order. */ function pickResetFirstCodexAccount( config: OcxConfig, @@ -2433,8 +2554,9 @@ export function resolveCodexAccountForThread( config: OcxConfig, now = Date.now(), quotaScope?: CodexQuotaScope, + lineage?: CodexThreadLineage, ): string | null { - const resolution = resolveCodexAccountForThreadDetailed(threadId, config, now, quotaScope); + const resolution = resolveCodexAccountForThreadDetailed(threadId, config, now, quotaScope, undefined, undefined, lineage); return resolution.status === "selected" ? resolution.accountId : null; } @@ -2710,6 +2832,7 @@ export function previewCodexAccountForRequest( quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, modelId?: string, + lineage?: CodexThreadLineage, ): string | null { // A request-scoped model detour keeps its own serving-account affinity. Preview // reads it before the ordinary lane, but never repairs or deletes it. Roster @@ -2735,6 +2858,30 @@ export function previewCodexAccountForRequest( ); if (ordinaryPreview) return ordinaryPreview; + // A conversation carried across an in-process swap is still bound under the pre-#4546 raw + // parent key, and resolve adopts that binding rather than rebinding cold. Preview has to name + // the same account. Read-only, as everything here is: it neither adopts nor retires the entry. + if (threadId && !entry && lineage?.legacyConversationKey !== undefined) { + const legacyPreview = previewReusableAffinityAccount( + getThreadAffinity(lineage.legacyConversationKey, quotaScope), + config, + now, + quotaScope, + selectionOptions, + ); + if (legacyPreview) return legacyPreview; + } + + // First placement mirrors resolve: a child with no binding previews the account actually + // serving its parent (or a compatible sibling), so the subagent fallback does not decide + // against a cold pick the real request would never make. Read-only: nothing binds here. + if (threadId && !entry && lineage) { + const lineagePreview = pickLineageServingAccount( + config, lineage, now, quotaScope, selectionOptions, modelId, + ); + if (lineagePreview) return lineagePreview.accountId; + } + const strategyPick = pickUnboundStrategyAccount( config, threadId, @@ -2793,6 +2940,7 @@ export function resolveCodexAccountForThreadDetailed( quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, modelId?: string, + lineage?: CodexThreadLineage, ): CodexThreadResolution { // An entitlement roster constrains only this model request. It must not rewrite // the operator's shared active/pin choice or the task's ordinary-model affinity. @@ -2820,6 +2968,13 @@ export function resolveCodexAccountForThreadDetailed( ) ); + // A conversation that was live across an in-process code swap is still bound under the + // pre-#4546 raw parent key. Adopt that binding onto this thread's key BEFORE anything below + // reads an entry, so the conversation arrives here as an ordinary bound thread instead of a + // cold one: every branch that follows -- detour reuse, transient hold, quota re-eval -- + // should treat it as the continuing conversation it is. No-op on a fresh process. + if (threadId) adoptLegacyLineageAffinity(threadId, lineage, now, quotaScope, modelId); + if (threadId && modelScopedSelection) { const detourEntry = getModelDetourAffinity(threadId, modelId, quotaScope); if (detourEntry) { @@ -3000,6 +3155,38 @@ export function resolveCodexAccountForThreadDetailed( // arrives) is the reason this request is starting cold, so it outranks having found nothing. releaseReason ??= peekPendingReleaseReason(threadId); + // FIRST PLACEMENT for a child thread (#4546, wp8). A child with no binding of its own used + // to bind under the raw parent id -- an entry unrelated to the root's real binding -- or + // land cold while its parent was being served warm somewhere. Consult the family's CURRENT + // serving account first (detour included), then a compatible sibling's, and only then fall + // through to cold placement. The child binds under its OWN key below: this is a warm start, + // not a root-wide pin, so a later move of the parent never drags the child with it. + // + // Guarded on `entry === undefined`, which is strictly narrower than "has no usable binding": + // a thread whose binding was just released above still holds its own history and re-decides + // through the ordinary path. Only a thread that has never bound takes a family hint. That + // also makes `preserveExistingModelScopedAffinity` unreachable here -- it is only ever set + // while reusing an existing model-detour entry -- so this binds through the ordinary lane. + if (threadId && entry === undefined && lineage) { + const lineagePick = pickLineageServingAccount( + config, lineage, now, quotaScope, selectionOptions, modelId, + ); + if (lineagePick) { + bindThreadAffinity(threadId, lineagePick.accountId, now, quotaScope); + // Deliberately no promoteActiveCodexAccount: a family hint places THIS request, it does + // not move the operator-visible shared cursor for unrelated new threads. + return { + status: "selected", + accountId: lineagePick.accountId, + // A pending release still outranks the hint as the reported reason, and consuming it + // here is what stops the next request reporting the same release a second time. + affinity: releaseReason === undefined + ? { move: "new_bind", reason: lineagePick.reason } + : affinityAfterRelease(threadId, releaseReason), + }; + } + } + // A request-scoped roster may still contain unhealthy candidates. Non-quota strategies return // before the quota/failover helpers below, so prefer only shared-healthy roster members here; // otherwise RR/fill-first can immediately re-pick a known failing account even when another diff --git a/src/server/request-log-conversation.ts b/src/server/request-log-conversation.ts index fd77f05798..c69f5773ba 100644 --- a/src/server/request-log-conversation.ts +++ b/src/server/request-log-conversation.ts @@ -64,13 +64,16 @@ export function sessionIdHeaderFromRequest(headers: Headers): string | null { /** * Fixed-size logical turn lane (#820). * - * A lane must be as SPECIFIC as the identity available, which is the opposite of what - * `codexPoolAffinityKey` wants. Affinity deliberately prefers the parent thread so a whole - * subagent fan-out pins to one account; a lane keyed that way would put every parallel - * subagent of one parent into a single lane and reject all but the first with 503 — the - * fan-out is the normal case, not an abuse. + * A lane must be as SPECIFIC as the identity available. `codexPoolAffinityKey` used to be the + * opposite: it preferred the parent thread, so a whole subagent fan-out shared one entry, and a + * lane keyed that way would have put every parallel subagent of one parent into a single lane + * and rejected all but the first with 503 — the fan-out is the normal case, not an abuse. + * Since #4546 affinity keys every thread as ITSELF and reads the parent only as a first-placement + * hint, so the two now agree on the unit. They still derive it differently: a lane is a digest an + * operator can match against what the client sent, while an affinity key is an opaque HMAC + * precisely so no caller-supplied identifier ends up in Pool state. * - * So the parent is a QUALIFIER, never the lane on its own when a child thread exists: the + * The parent stays a QUALIFIER here, never the lane on its own when a child thread exists: the * pair separates siblings while still keeping one conversation's overlapping turns together. */ export function sessionLaneIdFromRequest(headers: Headers): string | undefined { @@ -256,4 +259,3 @@ export function getOrAllocateRequestSessionLane(req: Request): string { export function linkRequestSessionLane(sourceReq: Request, targetReq: Request): void { requestAllocatedSessionLanes.set(targetReq, getOrAllocateRequestSessionLane(sourceReq)); } - diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index cab98df8db..63119f6f72 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -173,6 +173,7 @@ import { createCodexReserveDispatchGuard, unwrapUpstreamRetryEvidenceError, codexPoolAffinityKey, + previewCodexPoolLineage, CodexAccountCooldownError, CodexAuthContextError, CodexMainProfileDrainingError, @@ -2306,6 +2307,75 @@ type ResponsesAuthResolution = | { ok: true; authCtx: CodexAuthContext; headers: Headers; callerAuthHeaders: Headers; substituteMainCredential: boolean } | { ok: false; response: Response }; +/** + * The caller credential the final Codex auth resolution will be given, as far as the ROUTE + * decides it: a route change that may cross a credential domain drops the raw caller credential, + * and a trusted Claude-main handoff replaces it. + * + * Shared with the lineage preview in `handleResponsesInner`, which has to read a conversation's + * family under the same authenticated scope the resolution will record it under -- that scope is + * an HMAC of exactly this Authorization header. Two copies of this rule would put preview and + * final auth in different scopes the first time one of them changed. + */ +function codexRouteCredentialDomainHeaders( + req: Request, + route: RouteResult, + options: HandleResponsesOptions, + credentialDomainWasRewritten: boolean, +): Headers { + const trustedClaudeMainForFinalRoute = options.stripClaudeMainAuthForNoncanonicalForward === true + && isCanonicalOpenAiForwardProvider(route.provider) + ? options.trustedClaudeMainAuth : undefined; + if (trustedClaudeMainForFinalRoute) { + const claudeMainHeaders = new Headers(req.headers); + claudeMainHeaders.set("authorization", trustedClaudeMainForFinalRoute.authorization); + if (trustedClaudeMainForFinalRoute.chatgptAccountId) { + claudeMainHeaders.set("chatgpt-account-id", trustedClaudeMainForFinalRoute.chatgptAccountId); + } else { + claudeMainHeaders.delete("chatgpt-account-id"); + } + return claudeMainHeaders; + } + // Route-changing recursion retains typed admission, never an unscoped raw + // caller credential. Bearer admission is substituted or stripped below. + const routeMayChangeCredentialDomain = options.comboAttempt === true + || route.routeKind === "policy" + || credentialDomainWasRewritten; + if (routeMayChangeCredentialDomain && options.admission?.source !== "bearer") { + const scoped = new Headers(req.headers); + scoped.delete("authorization"); + scoped.delete("chatgpt-account-id"); + return scoped; + } + return req.headers; +} + +/** + * Does this route substitute OUR stored main credential, and does the caller own the credential + * this request will authenticate with? + * + * Both answers are needed twice: by the resolution below, and by the lineage preview, which must + * not follow a Pool family binding for a request whose credential never enters Pool state. One + * implementation, because two copies of this predicate disagreeing is the divergence the preview + * gate exists to prevent. The reasoning behind the substitution test itself is at its use site + * below (#1686, #2132). + */ +function codexRouteCredentialOwnership( + authInputHeaders: Headers, + config: OcxConfig, + route: RouteResult, + options: HandleResponsesOptions, +): { substituteMainCredential: boolean; requestScopedMainCredential: boolean } { + const substituteMainCredential = options.admission?.source === "bearer" + && (route.codexAccountMode !== undefined || isCanonicalOpenAiForwardProvider(route.provider)); + return { + substituteMainCredential, + requestScopedMainCredential: route.codexAccountMode !== undefined + && !substituteMainCredential + && hasForwardableCodexBearer(authInputHeaders, config), + }; +} + /** * Resolve Codex auth for a route. On unusable contexts, releases any probe lease * before returning the 401 (nothing reaches upstream). @@ -2318,30 +2388,12 @@ async function resolveResponsesCodexAuth( credentialDomainWasRewritten = false, ): Promise { try { - const routeMayChangeCredentialDomain = options.comboAttempt === true - || route.routeKind === "policy" - || credentialDomainWasRewritten; - const trustedClaudeMainForFinalRoute = options.stripClaudeMainAuthForNoncanonicalForward === true - && isCanonicalOpenAiForwardProvider(route.provider) - ? options.trustedClaudeMainAuth : undefined; - let authInputHeaders = req.headers; - // Route-changing recursion retains typed admission, never an unscoped raw - // caller credential. Bearer admission is substituted or stripped below. - if (routeMayChangeCredentialDomain && options.admission?.source !== "bearer" - && !trustedClaudeMainForFinalRoute) { - authInputHeaders = new Headers(req.headers); - authInputHeaders.delete("authorization"); - authInputHeaders.delete("chatgpt-account-id"); - } - if (trustedClaudeMainForFinalRoute) { - authInputHeaders = new Headers(authInputHeaders); - authInputHeaders.set("authorization", trustedClaudeMainForFinalRoute.authorization); - if (trustedClaudeMainForFinalRoute.chatgptAccountId) { - authInputHeaders.set("chatgpt-account-id", trustedClaudeMainForFinalRoute.chatgptAccountId); - } else { - authInputHeaders.delete("chatgpt-account-id"); - } - } + let authInputHeaders = codexRouteCredentialDomainHeaders( + req, + route, + options, + credentialDomainWasRewritten, + ); // A caller-auth transport that is not canonical OpenAI (keyless Cursor) consumes the // caller's Authorization as its own upstream token. Keep that contract only for a clean // single bearer with NO ChatGPT-domain marker. A bearer marked for the ChatGPT domain — @@ -2405,12 +2457,13 @@ async function resolveResponsesCodexAuth( // bug; the transport is the authority, because the transport is what actually carries the // header. A key-authenticated routed provider is still not canonical-forward, so #2132's // no-ChatGPT-login install keeps working. - const substituteMainCredential = options.admission?.source === "bearer" - && (route.codexAccountMode !== undefined || isCanonicalOpenAiForwardProvider(route.provider)); + const { substituteMainCredential, requestScopedMainCredential } = codexRouteCredentialOwnership( + authInputHeaders, + config, + route, + options, + ); const stripAuthorization = options.admission?.source === "bearer" && !substituteMainCredential; - const requestScopedMainCredential = route.codexAccountMode !== undefined - && !substituteMainCredential - && hasForwardableCodexBearer(authInputHeaders, config); if (route.codexAccountMode === "direct" && !substituteMainCredential) { validateForwardAdmissionCredential(authInputHeaders, config); } @@ -4044,6 +4097,33 @@ async function handleResponsesInner( let subagentQuotaFailureModel = parsed.modelId; const parentThreadId = req.headers.get("x-codex-parent-thread-id")?.trim() ?? null; const poolAffinityKey = codexPoolAffinityKey(req.headers) ?? null; + // Preview has to see the same lineage resolve does. Without it, a child's first turn is + // previewed as a cold pick and resolved onto the family account, and the subagent fallback + // then decides model eligibility against an account the request will never use. + // + // "The same" means both halves of the question the final resolution asks. The Authorization + // it will be given, because the lineage scope is an HMAC of exactly that header; and its own + // Pool-state predicate, because a fixed account selector and a request-owned credential + // deliberately create no affinity at all -- previewing a family binding for one of those would + // hand model fallback an account this request can never authenticate as. Read-only: the record + // is written by the resolution that binds, never by a preview that may own no Pool state. + const previewAuthHeaders = codexRouteCredentialDomainHeaders( + req, + route, + options, + credentialDomainWasRewritten, + ); + const poolLineage = previewCodexPoolLineage(previewAuthHeaders, options.codexAuthPolicy ?? config, { + accountId: route.codexAccountId, + modelId: route.modelId, + admission: options.admission, + requestScopedMainCredential: codexRouteCredentialOwnership( + previewAuthHeaders, + config, + route, + options, + ).requestScopedMainCredential, + }); try { if ( @@ -4082,6 +4162,7 @@ async function handleResponsesInner( codexQuotaScopeForModel(modelId), { ...previewSelectionOptions, modelEligibleAccountIds }, modelId, + poolLineage, ); const previewAccountId = route.codexAccountId ?? subagentFallbackAccountPreview( route.modelId, @@ -4227,6 +4308,7 @@ async function handleResponsesInner( codexQuotaScopeForModel(modelId), { ...recoverySelectionOptions, modelEligibleAccountIds }, modelId, + poolLineage, ); const recoveryPreviewAccountId = subagentFallbackAccountPreview( parsed.modelId, diff --git a/structure/providers/openai-tiers.md b/structure/providers/openai-tiers.md index a1824746dc..3407c03888 100644 --- a/structure/providers/openai-tiers.md +++ b/structure/providers/openai-tiers.md @@ -71,16 +71,36 @@ support them. > Decision record: [ADR-0084](../decisions/ADR-0084-public-provider-contract.md) -Pool affinity preserves the existing `x-codex-parent-thread-id` supplied by ordinary Codex clients. -The parent id is trimmed and bounded under the same 512-byte component limit as the Desktop -fallback. When Codex Desktop omits it or sends an unusable value, the complete bounded `session-id` -plus `thread-id` pair is mapped to an opaque HMAC under a random process-local key. Missing or -oversized components remain unbound, raw identifiers and durable hashes are never stored, and +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 +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. +`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. + > Decision record: [ADR-0085](../decisions/ADR-0085-public-provider-contract.md) An explicit `Retry-After` or an unclassified quota 429 is account-wide. A reset-derived native-model diff --git a/tests/codex-integration/codex-auth-context.test.ts b/tests/codex-integration/codex-auth-context.test.ts index b488bc7855..b3d8d83ae4 100644 --- a/tests/codex-integration/codex-auth-context.test.ts +++ b/tests/codex-integration/codex-auth-context.test.ts @@ -21,6 +21,7 @@ import { materializeCodexUpstreamAuth, CodexMainSubstitutionUnavailableError, isCodexAuthContextUsable, + codexPoolAffinityKey, resolveCodexAuthContext, shouldMarkAccountNeedsReauthForCodexAuthFailure, stripCodexRuntimeProviderFields, @@ -896,7 +897,7 @@ describe("Codex auth context", () => { }); }); - test("the canonical parent-thread affinity stays authoritative over Desktop fallback headers", async () => { + test("a parent-bearing Desktop request keys as its own thread, not as its parent (#4546 wp8)", async () => { const cfg = config(); cfg.autoSwitchThreshold = 0; saveCodexAccountCredential("pool-a", { @@ -912,11 +913,24 @@ describe("Codex auth context", () => { }); const resolved = await resolveCodexAuthContext(headers, cfg, "pool"); - expect(resolved).toMatchObject({ - kind: "pool", - accountId: "pool-a", - affinityKey: "canonical-parent-thread", - }); + 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. + 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. + 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. + expect(resolved.affinityKey).not.toBe(codexPoolAffinityKey(new Headers({ + "x-codex-parent-thread-id": "canonical-parent-thread", + }))); }); test("an oversized parent-thread id falls back to the bounded Desktop pair", async () => { diff --git a/tests/codex-integration/codex-lineage-placement.test.ts b/tests/codex-integration/codex-lineage-placement.test.ts new file mode 100644 index 0000000000..f41e519c8e --- /dev/null +++ b/tests/codex-integration/codex-lineage-placement.test.ts @@ -0,0 +1,520 @@ +/** + * Codex V2 lineage and FIRST PLACEMENT (#4546, wp8). + * + * Two defects are pinned here. Keying: every child of one parent used to bind under the RAW + * parent id, one shared entry unrelated to the root's own binding, so no child could hold a + * binding of its own and a grandchild keyed on a key nobody had bound. Placement: a child with + * no binding started cold even while its parent was being served warm somewhere. + * + * The asymmetry is the point and has its own test below. A family hint decides where a child + * STARTS; it is not a root-wide pin, so a later move of the parent must leave an already-bound + * child exactly where it is. + * + * The fixture mirrors tests/codex-integration/codex-pool-rotation.test.ts: quota strategy, three + * accounts, and an explicit usage order, so every expected account is the one a cold pick would + * NOT have produced wherever that distinction carries the proof. + */ +import { describe, expect, test, beforeEach, afterEach } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + clearCodexUpstreamHealth, + clearThreadAccountMap, + recordCodexUpstreamOutcome, + resolveCodexAccountForThreadDetailed, +} from "../../src/codex/routing"; +import { codexPoolAffinityKey, previewCodexPoolLineage } from "../../src/codex/auth-context"; +import { + CODEX_LINEAGE_IDLE_TTL_MS, + CODEX_LINEAGE_MAX_ENTRIES, + CODEX_LINEAGE_MAX_SCOPES, + clearCodexThreadLineageForTests, + codexLineageRootForRequest, + codexLineageScopeKey, + codexLineageWorkflowLane, + codexThreadLineageLookup, + recordCodexThreadLineage, +} from "../../src/codex/lineage"; +import { clearPoolRotationState } from "../../src/codex/pool-rotation"; +import { saveCodexAccountCredential } from "../../src/codex/account-store"; +import { clearAccountQuota, updateAccountQuota } from "../../src/codex/auth-api"; +import { flushConfigDirHardeningForTests } from "../../src/config/paths"; +import { setAsyncIcaclsRunnerForTests, setIcaclsRunnerForTests } from "../../src/lib/windows-secret-acl"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; +import type { OcxConfig } from "../../src/types"; + +let TEST_DIR = ""; +let previousOpencodexHome: string | undefined; +let previousCodexHome: string | undefined; + +const ICACLS_OK = { success: true, exitCode: 0, timedOut: false, stdout: "" }; +const ACCOUNT_IDS = ["a", "b", "c"] as const; +const NOW = 1_700_000_000_000; + +function installScratchHome(): void { + previousOpencodexHome = process.env.OPENCODEX_HOME; + previousCodexHome = process.env.CODEX_HOME; + TEST_DIR = mkdtempSync(join(tmpdir(), "ocx-lineage-")); + setIcaclsRunnerForTests(() => ICACLS_OK); + setAsyncIcaclsRunnerForTests(async () => ICACLS_OK); + process.env.OPENCODEX_HOME = TEST_DIR; + process.env.CODEX_HOME = TEST_DIR; +} + +async function removeScratchHome(): Promise { + const ownedDirectory = TEST_DIR; + TEST_DIR = ""; + try { + await flushConfigDirHardeningForTests(); + } finally { + setIcaclsRunnerForTests(null); + setAsyncIcaclsRunnerForTests(null); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + if (previousCodexHome === undefined) delete process.env.CODEX_HOME; + else process.env.CODEX_HOME = previousCodexHome; + if (ownedDirectory) removeTreeWithRetry(ownedDirectory); + } +} + +function saveTestCredential(id: string): void { + saveCodexAccountCredential(id, { + accessToken: `access-${id}`, + refreshToken: `refresh-${id}`, + expiresAt: Date.now() + 5 * 60_000, + chatgptAccountId: `acct-${id}`, + }); +} + +/** Quota strategy with an explicit usage order, so every cold pick below is predictable. */ +function makeConfig(overrides: Partial = {}): OcxConfig { + return { + providers: {}, + codexAccounts: ACCOUNT_IDS.map(id => ({ id, email: `${id}@example.test`, isMain: false })), + accountPoolStrategy: "quota", + activeCodexAccountId: "a", + autoSwitchThreshold: 80, + upstreamFailoverThreshold: 3, + ...overrides, + } as OcxConfig; +} + +/** + * The session id is deliberately NOT the thread id. Codex's own root sends the same string for + * both, and a fixture that copies it makes HMAC(parent, parent) accidentally equal the root's + * key -- which is exactly the coincidence that hid the parent-only defect pinned below. + */ +const rootHeaders = () => new Headers({ "session-id": "sess", "thread-id": "root" }); +const childHeaders = (threadId: string, parentId = "root") => new Headers({ + "session-id": "sess", + "thread-id": threadId, + "x-codex-parent-thread-id": parentId, +}); + +/** One transient streak: the binding stays put while this request is sent elsewhere. */ +function streakTransientFailures(config: OcxConfig, accountId: string, now: number): void { + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, accountId, 503, { now }); + } +} + +describe("codex thread lineage and first placement (#4546 wp8)", () => { + 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("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(); + }); + + test("lineage resolves the root transitively and stays inside its auth scope", () => { + const root = recordCodexThreadLineage(rootHeaders(), NOW)!; + const child = recordCodexThreadLineage(childHeaders("child-1"), NOW)!; + const grandchild = recordCodexThreadLineage(childHeaders("grand-1", "child-1"), NOW)!; + expect(root.rootSessionKey).toBe(root.conversationKey); + expect(child.parentConversationKey).toBe(root.conversationKey); + expect(child.rootSessionKey).toBe(root.rootSessionKey); + // Transitive: the grandchild's spend belongs to the ROOT workflow, not to child-1. + expect(grandchild.parentConversationKey).toBe(child.conversationKey); + expect(grandchild.rootSessionKey).toBe(root.rootSessionKey); + + const scope = codexLineageScopeKey(rootHeaders()); + expect(codexThreadLineageLookup(grandchild.conversationKey, scope, NOW)).toMatchObject({ + rootSessionKey: root.rootSessionKey, + parentThreadId: "child-1", + }); + expect(codexLineageRootForRequest(childHeaders("grand-1", "child-1"), NOW)).toBe(root.rootSessionKey); + // Another authenticated caller presenting identical thread ids sees nothing of this scope. + const otherScope = codexLineageScopeKey(new Headers({ authorization: "Bearer other" })); + expect(otherScope).not.toBe(scope); + expect(codexThreadLineageLookup(grandchild.conversationKey, otherScope, NOW)).toBeUndefined(); + // Idle expiry bounds the table exactly like the binding map it feeds. + expect(codexThreadLineageLookup( + grandchild.conversationKey, scope, NOW + CODEX_LINEAGE_IDLE_TTL_MS + 1, + )).toBeUndefined(); + }); + + test("the table is bounded in both dimensions, not just per scope", () => { + const keyFor = (index: number) => recordCodexThreadLineage( + new Headers({ "session-id": "bulk", "thread-id": `bulk-${index}` }), NOW, + )!.conversationKey; + const oldest = keyFor(0); + for (let index = 1; index <= CODEX_LINEAGE_MAX_ENTRIES; index += 1) keyFor(index); + const newest = keyFor(CODEX_LINEAGE_MAX_ENTRIES + 1); + const localScope = codexLineageScopeKey(new Headers()); + expect(codexThreadLineageLookup(oldest, localScope, NOW)).toBeUndefined(); + expect(codexThreadLineageLookup(newest, localScope, NOW)).toBeDefined(); + + // The scope map is the one an untrusted caller could grow without the cap below. + const held = new Headers({ authorization: "Bearer held", "session-id": "s", "thread-id": "t" }); + const heldKey = recordCodexThreadLineage(held, NOW)!.conversationKey; + expect(codexThreadLineageLookup(heldKey, codexLineageScopeKey(held), NOW)).toBeDefined(); + for (let index = 0; index <= CODEX_LINEAGE_MAX_SCOPES; index += 1) { + recordCodexThreadLineage(new Headers({ + authorization: `Bearer caller-${index}`, + "session-id": "s", + "thread-id": "t", + }), NOW); + } + 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", () => { + 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" }); + + // The binding is HELD on a while the request itself detours to b. + streakTransientFailures(config, "a", NOW); + expect(resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW + 1)).toMatchObject({ + status: "selected", + accountId: "b", + affinity: { move: "detour", reason: "transient" }, + }); + + // The child starts where the parent is being served NOW (b), not at its stale home (a). + const child = recordCodexThreadLineage(childHeaders("child-1"), NOW + 2)!; + expect(resolveCodexAccountForThreadDetailed( + child.conversationKey, config, NOW + 2, undefined, undefined, undefined, child, + )).toMatchObject({ + status: "selected", + accountId: "b", + affinity: { move: "new_bind", reason: "lineage_parent" }, + }); + }); + + test("a later move of the parent does not drag an already-bound child", () => { + const config = makeConfig(); + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + const root = recordCodexThreadLineage(rootHeaders(), NOW)!; + resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW); + 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. + 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); + + // 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. + const lateChild = recordCodexThreadLineage(childHeaders("child-2"), NOW + 5)!; + 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.accountId).toBe(parentAfterMove.accountId); + }); + + test("a compatible sibling places the child when the parent is not eligible", () => { + 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" }); + + // 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. + 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. + 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. + 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); + 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. + expect(orphanPlacement.accountId).toBe(siblingPlacement.accountId); + expect(orphanPlacement).toMatchObject({ + status: "selected", + affinity: { move: "new_bind", reason: "lineage_sibling" }, + }); + }); + + test("no known family account falls back to ordinary cold placement", () => { + const config = makeConfig(); + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + // The parent was never seen and holds no binding, so lineage cannot help. The request takes + // exactly the pick an unrelated new thread would. + const child = recordCodexThreadLineage(childHeaders("child-1"), NOW)!; + const resolution = resolveCodexAccountForThreadDetailed( + child.conversationKey, config, NOW, undefined, undefined, undefined, child, + ); + expect(resolution).toMatchObject({ status: "selected", accountId: "a" }); + expect(resolution.affinity?.reason).not.toBe("lineage_parent"); + expect(resolution.affinity?.reason).not.toBe("lineage_sibling"); + }); + + test("a parent-only turn continues the parent's conversation, session id or not", () => { + 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" }); + + // This turn carries nothing but the parent id, so only the recorded relation can reproduce + // the key the parent bound under. HMAC(parent, parent) would be a different key, and this + // conversation would start cold on every such turn while replacing the parent's record. + const parentOnly = new Headers({ "x-codex-parent-thread-id": "root" }); + expect(codexPoolAffinityKey(parentOnly, NOW + 1)).toBe(root.conversationKey); + + const followUp = recordCodexThreadLineage(parentOnly, NOW + 1)!; + expect(followUp.conversationKey).toBe(root.conversationKey); + expect(resolveCodexAccountForThreadDetailed( + followUp.conversationKey, config, NOW + 1, undefined, undefined, undefined, followUp, + )).toMatchObject({ + status: "selected", + accountId: "a", + affinity: { move: "reused", reason: "healthy" }, + }); + + // And recording it left the parent's record intact rather than overwriting it. + expect(codexThreadLineageLookup(root.conversationKey, codexLineageScopeKey(parentOnly), NOW + 1)) + .toMatchObject({ conversationKey: root.conversationKey, rootSessionKey: root.rootSessionKey }); + }); + + test("a binding left under the old raw-parent key is adopted, not rebound cold", () => { + const config = makeConfig(); + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + // What a code swap under a live conversation leaves behind: a binding made by the pre-#4546 + // rule, under the RAW parent id. c is where it sits, and c is not where a cold pick goes. + config.pausedCodexAccountIds = ["a", "b"]; + expect(resolveCodexAccountForThreadDetailed("root", config, NOW)) + .toMatchObject({ status: "selected", accountId: "c" }); + config.pausedCodexAccountIds = []; + config.activeCodexAccountId = "a"; + + const child = recordCodexThreadLineage(childHeaders("child-1"), NOW + 1)!; + expect(child.legacyConversationKey).toBe("root"); + // The conversation keeps its account AND its status as a bound thread. A cold rebind here is + // the exact defect this unit exists to prevent, so "reused" is the assertion, not "c". + expect(resolveCodexAccountForThreadDetailed( + child.conversationKey, config, NOW + 1, undefined, undefined, undefined, child, + )).toMatchObject({ + status: "selected", + accountId: "c", + affinity: { move: "reused", reason: "healthy" }, + }); + + // One way, once: nothing answers on the legacy key any more, so a request arriving there + // binds fresh instead of finding the account it just handed over. + expect(resolveCodexAccountForThreadDetailed("root", config, NOW + 2)).toMatchObject({ + status: "selected", + accountId: "a", + affinity: { move: "new_bind" }, + }); + }); + + test("a child follows the parent's MODEL detour, not a home account that cannot serve it", () => { + const config = makeConfig(); + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + const modelId = "native-gated-model"; + const roster = { modelEligibleAccountIds: new Set(["b", "c"]) }; + + const root = recordCodexThreadLineage(rootHeaders(), NOW)!; + // The parent's home account is a, chosen with no model roster in play. + expect(resolveCodexAccountForThreadDetailed(root.conversationKey, config, NOW)) + .toMatchObject({ status: "selected", accountId: "a" }); + // a is not entitled to this model, so the parent is now SERVED through a model detour on b + // while its ordinary binding stays on a. + 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. + updateAccountQuota("b", 40); + const child = recordCodexThreadLineage(childHeaders("child-1"), NOW + 1)!; + expect(resolveCodexAccountForThreadDetailed( + child.conversationKey, config, NOW + 1, undefined, roster, modelId, child, + )).toMatchObject({ + status: "selected", + accountId: "b", + affinity: { move: "new_bind", reason: "lineage_parent" }, + }); + }); + + test("a preview reads the family only for a request that may own Pool state", () => { + const config = makeConfig(); + const root = recordCodexThreadLineage(rootHeaders(), NOW)!; + const child = childHeaders("child-1"); + + expect(previewCodexPoolLineage(child, config)?.parentConversationKey).toBe(root.conversationKey); + // An exact account selector authenticates outside the Pool and creates no affinity, so a + // preview that followed the family here would decide model fallback against an account the + // request will never be. + expect(previewCodexPoolLineage(child, config, { accountId: "b" })).toBeUndefined(); + const callerOwned = childHeaders("child-2"); + callerOwned.set("authorization", "Bearer caller-owned-credential"); + 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. + expect(codexThreadLineageLookup( + codexPoolAffinityKey(child)!, codexLineageScopeKey(child), NOW, + )).toBeUndefined(); + }); + + test("worker classification stays header-first and gains the lineage-backed answer", () => { + // Header-only rule preserved: a parent plus a distinct thread-id is worker traffic. + expect(codexLineageWorkflowLane(childHeaders("child-1"), NOW)).toBe("worker"); + // A bare thread-id with no recorded family is interactive, matching today's admission. + expect(codexLineageWorkflowLane(new Headers({ "thread-id": "lone" }), NOW)).toBe("interactive"); + expect(codexLineageWorkflowLane(new Headers(), NOW)).toBe("interactive"); + // The lineage-backed half: a thread recorded with a parent is worker traffic even when THIS + // request's headers no longer declare one. + recordCodexThreadLineage(childHeaders("child-9"), NOW); + expect(codexLineageWorkflowLane(new Headers({ "thread-id": "child-9" }), NOW)).toBe("worker"); + }); +}); diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index cef91aa256..acbe6b76b9 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -304,6 +304,7 @@ "codex-integration-record.test.ts": "codex-integration", "codex-journal.test.ts": "codex-integration", "codex-legacy-config-keys.test.ts": "codex-integration", + "codex-lineage-placement.test.ts": "codex-integration", "codex-log-guard-coderabbit.test.ts": "codex-integration", "codex-log-guard-doctor-coderabbit.test.ts": "codex-integration", "codex-log-guard-doctor-protection.test.ts": "codex-integration", diff --git a/tests/responses/responses-pool-401-refresh.test.ts b/tests/responses/responses-pool-401-refresh.test.ts index 88dbfc0afc..b1043cc2dd 100644 --- a/tests/responses/responses-pool-401-refresh.test.ts +++ b/tests/responses/responses-pool-401-refresh.test.ts @@ -4,6 +4,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { createHash } from "node:crypto"; import { clearAccountNeedsReauth, isAccountNeedsReauth } from "../../src/codex/auth-api"; +import { codexPoolAffinityKey } from "../../src/codex/auth-context"; import { clearCodexUpstreamHealth, clearThreadAccountMap, @@ -605,7 +606,13 @@ describe("ordinary pool 401 refresh and replay (#2887)", () => { // reported as expired, which is the behavior the missing handoff produces. // The binding lives under the model's quota scope, so resolution must be asked in that // same scope; a scopeless read looks in the legacy bucket and finds nothing. - expect(resolveCodexAccountForThreadDetailed(THREAD_ID, cfg, Date.now(), "shared")).toMatchObject({ + // Since #4546 a thread keys as ITSELF through an opaque HMAC, and the parent header is a + // first-placement hint rather than the key. A parent-only turn therefore binds under the + // derived key, not under the raw parent id this suite used to read back. + const affinedKey = codexPoolAffinityKey( + new Headers({ "x-codex-parent-thread-id": THREAD_ID }), + )!; + expect(resolveCodexAccountForThreadDetailed(affinedKey, cfg, Date.now(), "shared")).toMatchObject({ status: "selected", accountId: ACCOUNT_ID, }); diff --git a/tests/routing/subagent-fallback-handle-responses.test.ts b/tests/routing/subagent-fallback-handle-responses.test.ts index 41c47092c7..a9fdcf6348 100644 --- a/tests/routing/subagent-fallback-handle-responses.test.ts +++ b/tests/routing/subagent-fallback-handle-responses.test.ts @@ -1595,8 +1595,12 @@ describe("native fallback account preview", () => { } // And both must actually forward it into the preview call, not merely accept it. + // The guarantee is that BOTH sites forward the eligible set, which is what recovery lost. + // `modelId` is no longer the final argument -- #4546 appends the resolved pool lineage so + // preview and final resolution agree on a child's first turn -- so anything after it is + // allowed here rather than pinning the argument count. const forwarded = source.match( - /\{ \.\.\.(previewSelectionOptions|recoverySelectionOptions), modelEligibleAccountIds \},\s*modelId,\s*\)/g, + /\{ \.\.\.(previewSelectionOptions|recoverySelectionOptions), modelEligibleAccountIds \},\s*modelId,[^)]*\)/g, ) ?? []; expect(forwarded).toHaveLength(2); });