feat(codex): give V2 threads real lineage and place a child on its parent's serving account (#4546)

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",

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<DataPlaneAdmission, "source">;
+  /** 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;