diff --git a/internal/proxy/codex_http_request_plan.go b/internal/proxy/codex_http_request_plan.go index 05f3350..5f12a41 100644 --- a/internal/proxy/codex_http_request_plan.go +++ b/internal/proxy/codex_http_request_plan.go @@ -182,8 +182,9 @@ type CodexPreparedHTTPRequest struct { Frozen *CodexFrozenRequest Lifecycle CodexHTTPRequestLifecycle - leaseHandle *CodexLeaseRequestHandle - receipt *codexTurnReceiptHandle + leaseHandle *CodexLeaseRequestHandle + receipt *codexTurnReceiptHandle + portableQuotaRetry bool } // CodexHTTPRequestPlanErrorCode classifies preparation failures without @@ -576,6 +577,13 @@ func (factory *CodexHTTPRequestPlanFactory) buildOnce(ctx context.Context, input AcceptedRevision: input.AcceptedRevision, Now: now, } + // A continuation cannot move its response or turn state to another account. + // Keep its exhausted binding available for one quota probe instead of + // removing it from the candidate set and reporting a routing failure. + boundQuotaProbe := boundAccountKey != "" && containsCodexHTTPRequestAccountKey(snapshot.QuotaExhaustedAccountKeys, boundAccountKey) + if boundQuotaProbe { + dispatchInput.UnavailableAccountKeys = excludeCodexHTTPRequestAccountKeys(dispatchUnavailable, []codex.AccountKey{boundAccountKey}) + } dispatch, err := factory.buildDispatch(ctx, dispatchInput) if err != nil { emitCodexTrace(ctx, CodexTraceEvent{Phase: "route_selection", Outcome: "error", Reason: string(codexRequestFailureReason(err))}) @@ -651,27 +659,37 @@ func (factory *CodexHTTPRequestPlanFactory) buildOnce(ctx context.Context, input quotaExhaustionProbe := containsCodexHTTPRequestAccountKey(snapshot.QuotaExhaustedAccountKeys, choice.AccountKey) if policyDecision.Status == PolicyDecisionSelected { available := excludeCodexHTTPRequestAccountKeys(policyDecision.Allowed, dispatchUnavailable) - if len(available) != 0 { - policyDecision.Allowed = available - } else if quotaExhaustionProbe { + if quotaExhaustionProbe { policyDecision.Allowed = []codex.AccountKey{choice.AccountKey} + } else if len(available) != 0 { + policyDecision.Allowed = available } } + if boundAccountKey == "" && codexHTTPRequestAccountUnavailablePortable(protocol) { + dispatch.accountUnavailablePortable = true + } if len(resetInventory.Accounts) > 1 { resetPlan := dispatch if boundAccountKey != "" { resetInput := dispatchInput + resetInput.UnavailableAccountKeys = dispatchUnavailable + resetInput.ProbeUnavailableWhenAll = false resetInput.Inventory = resetInventory resetInput.AffinityAccountKey = "" resetInput.AffinityEffectiveModel = "" resetInput.BoundAccountKey = "" resetPlan, err = factory.buildDispatch(ctx, resetInput) if err != nil { - return result, newCodexHTTPRequestPlanError(CodexHTTPRequestPlanDispatch, err) + var exhausted *CachedUsageLimitError + if !errors.As(err, &exhausted) { + return result, newCodexHTTPRequestPlanError(CodexHTTPRequestPlanDispatch, err) + } + // No reset route remains, but the bound attempt is still valid. + resetPlan = CodexFrozenDispatchPlan{} } } dispatch = dispatch.withAccountUnavailableResetCandidates(resetPlan, choice) - if accountUnavailablePortable { + if accountUnavailablePortable && !quotaExhaustionProbe { dispatch = dispatch.withAccountUnavailableFallbacks(resetPlan, choice) } } @@ -729,6 +747,7 @@ func (factory *CodexHTTPRequestPlanFactory) buildOnce(ctx context.Context, input } emitCodexTrace(ctx, CodexTraceEvent{Phase: "lease_begin", Outcome: "success", AccountHint: codexTraceAccountHint(choice.AccountKey)}) + result.portableQuotaRetry = codexHTTPRequestAccountUnavailablePortable(protocol) result.Dispatch = dispatch result.Frozen = frozen result.leaseHandle = handle diff --git a/internal/proxy/codex_http_request_plan_test.go b/internal/proxy/codex_http_request_plan_test.go index 77311b9..5e7ead3 100644 --- a/internal/proxy/codex_http_request_plan_test.go +++ b/internal/proxy/codex_http_request_plan_test.go @@ -2693,3 +2693,110 @@ func assertCodexHTTPRequestPlanError(t *testing.T, err error, code CodexHTTPRequ } } } + +func TestCodexHTTPRequestPlanFactoryProbesQuotaExhaustedBoundAccount(t *testing.T) { + now := time.Unix(1_700_000_000, 0).UTC() + for _, test := range []struct { + name string + previous bool + turnState bool + expected bool + restartable bool + allExhausted bool + authenticated bool + zeroAlternate bool + }{ + {name: "previous response", previous: true}, + {name: "zero capacity alternate", previous: true, zeroAlternate: true}, + {name: "authenticated continuation", previous: true, authenticated: true}, + {name: "turn state", turnState: true}, + {name: "expected binding", expected: true}, + {name: "restartable portable", restartable: true}, + {name: "all exhausted continuation", previous: true, allExhausted: true}, + } { + t.Run(test.name, func(t *testing.T) { + inventory := codex.Inventory{Accounts: []codex.LogicalAccount{ + frozenDispatchTestLogicalAccount("account-a", frozenDispatchCandidate("account-a", "candidate-a", "revision-a", codex.SourceSystem, false, now.Add(time.Hour))), + frozenDispatchTestLogicalAccount("account-b", frozenDispatchCandidate("account-b", "candidate-b", "revision-b", codex.SourceExternal, false, now.Add(time.Hour))), + }} + identity := CodexJournalRecordIdentity{LaneDigest: "lane", TurnDigest: "turn", ModeEpoch: 1, Authoritative: true} + snapshot := CodexLeaseRouteSnapshot{ + Classification: CodexRestoredLaneCurrent, JournalGeneration: 2, + BoundAccountKey: "account-a", BoundIdentity: identity, BoundRecordGeneration: 1, + BoundChoice: RouteChoice{AccountKey: "account-a", EffectiveModel: "gpt-5", RequiredBuckets: []CapacityBucket{CapacityBucketBase}}, + RestartableFailedHead: test.restartable, + UnavailableAccountKeys: []codex.AccountKey{"account-a"}, + QuotaExhaustedAccountKeys: []codex.AccountKey{"account-a"}, + } + if test.allExhausted { + snapshot.UnavailableAccountKeys = append(snapshot.UnavailableAccountKeys, "account-b") + snapshot.QuotaExhaustedAccountKeys = append(snapshot.QuotaExhaustedAccountKeys, "account-b") + } + runtime := &codexHTTPRequestPlanTestRuntime{handle: &CodexLeaseRequestHandle{account: "account-a"}} + factory := &CodexHTTPRequestPlanFactory{ + Inventory: &codexHTTPRequestPlanTestInventory{inventory: inventory}, + Routes: &codexHTTPRequestPlanTestSnapshotter{snapshot: snapshot}, Runtime: runtime, + DefaultAccountKey: "account-a", Authority: CodexLeaseAuthorityPolicy{ModeEpoch: 1, Authoritative: true}, + Now: func() time.Time { return now }, + } + if test.zeroAlternate { + capacity := NewCodexCapacityLedger(func() time.Time { return now }, time.Hour) + frozenDispatchObserveCapacity(t, capacity, "account-b", CapacityBucketBase, 0, now) + factory.Capacity = capacity + } + key := []byte("01234567890123456789012345678901") + factory.SessionPolicy = NewSessionPolicyResolver(key, routingPolicyV2ForTest(RoutingPolicyV1{ + SchemaVersion: 1, AuthorityGeneration: 1, RoutingGeneration: 1, EffectiveGeneration: 1, + Pools: []AccountPoolV1{{Name: "team", Members: []codex.AccountKey{"account-a", "account-b"}}}, + SessionBindings: []SessionBindingV1{{SessionDigest: keyedSessionDigest(key, []byte("session")), Pool: "team"}}, + })) + permits := &sessionPolicyPermitRecorder{} + factory.DispatchPermits = permits + ctx := withRuntimeCallerAuthority(context.Background(), RuntimeCallerAuthorityV1{Domain: NormalCallerLocal, SubjectID: "local-caller", ConsumptionDigest: strings.Repeat("a", 64)}) + if test.authenticated { + ctx = withRuntimeCallerAuthority(ctx, RuntimeCallerAuthorityV1{Domain: NormalCallerCodex, SubjectID: "account-a", IndexEpoch: 1, ConsumptionDigest: strings.Repeat("a", 64)}) + ctx = withRuntimeCallerIdentity(ctx, "account-a\x00candidate-a\x00revision-a") + } + input := CodexHTTPRequestPlanInput{Encoded: frozenRequestBody("gpt-5", CodexRequestTurn, "private-body")} + if test.previous { + input.Encoded = []byte(strings.TrimSuffix(string(input.Encoded), "}") + `,"previous_response_id":"response-a"}`) + } + if test.turnState { + input.Headers = http.Header{"X-Codex-Turn-State": {"private-turn-state"}} + } + if test.expected { + input.ExpectedBound = &CodexLeaseBoundExpectation{Identity: identity, AccountKey: "account-a", RecordGeneration: 1} + } + prepared, err := factory.Build(ctx, input) + if err != nil { + t.Fatal(err) + } + defer prepared.Frozen.Release() + if prepared.portableQuotaRetry != (!test.previous && !test.turnState) { + t.Fatalf("portable quota retry = %t", prepared.portableQuotaRetry) + } + if len(permits.requests) != 1 || !slices.Equal(permits.requests[0].AllowedAccounts, []codex.AccountKey{"account-a"}) { + t.Fatalf("bound probe permit scope = %#v, want only account-a", permits.requests) + } + accounts := prepared.Dispatch.Accounts() + if len(accounts) != 1 || accounts[0].Choice().AccountKey != "account-a" { + t.Fatalf("bound probe accounts = %#v, want only account-a", accounts) + } + if !runtime.plan.QuotaExhaustionProbe { + t.Fatal("bound account was not a quota probe") + } + for _, slot := range runtime.plan.Slots { + if slot.AccountKey != "account-a" { + t.Fatalf("quota probe contains alternate slot: %#v", slot) + } + } + wantReset := []codex.AccountKey{"account-b"} + if test.allExhausted || test.zeroAlternate { + wantReset = nil + } + if got := prepared.Dispatch.AccountUnavailableResetCandidates(); !slices.Equal(got, wantReset) { + t.Fatalf("reset candidates = %v, want %v", got, wantReset) + } + }) + } +} diff --git a/internal/proxy/codex_http_request_session.go b/internal/proxy/codex_http_request_session.go index d8ec328..240a13b 100644 --- a/internal/proxy/codex_http_request_session.go +++ b/internal/proxy/codex_http_request_session.go @@ -391,6 +391,8 @@ type CodexHTTPRequestSessionResult struct { Choice RouteChoice Attempt CandidateAttempt Lifecycle CodexHTTPRequestLifecycle + + quotaExhausted bool } // CodexHTTPAttemptSlotPlan is one raw-free bridge entry for the durable lease @@ -741,7 +743,7 @@ accountsLoop: result.Attempt = retainedAttempt return result, nil } - if (authRejected || hardRejected) && plan.TerminalError() != nil && codexHTTPRequestCanRecordAccountUnavailable(plan, result.Lifecycle) { + if authRejected && plan.TerminalError() != nil && codexHTTPRequestCanRecordAccountUnavailable(plan, result.Lifecycle) { discardCodexHTTPRequestResponse(ctx, response) result.Response = nil next, finishErr := codexHTTPRequestRecordAccountUnavailable(ctx, result.Lifecycle, 0, hardRejected) @@ -763,6 +765,7 @@ accountsLoop: return result, finishErr } result.Lifecycle = next + result.quotaExhausted = hardRejected && codexHTTPRequestCanRecordAccountUnavailable(plan, result.Lifecycle) return result, nil } } diff --git a/internal/proxy/codex_lease_v2_cas.go b/internal/proxy/codex_lease_v2_cas.go index 00d1bde..649af96 100644 --- a/internal/proxy/codex_lease_v2_cas.go +++ b/internal/proxy/codex_lease_v2_cas.go @@ -759,7 +759,7 @@ func (store *CodexLeaseStore) buildCodexLeaseRecordAfterImage(old CodexJournalRe if old.HasEncryptedState && !input.HasEncryptedState { return CodexJournalRecordV2{}, 0, false, fmt.Errorf("%w: encrypted-state authority was cleared", ErrCodexLeaseInvalidMutation) } - if old.HasTurnState && !input.HasTurnState { + if old.HasTurnState && !input.HasTurnState && !bindingReassignment { return CodexJournalRecordV2{}, 0, false, fmt.Errorf("%w: turn-state authority was cleared", ErrCodexLeaseInvalidMutation) } if !old.HasTurnState && input.HasTurnState && !input.TurnStateLatchCurrent { @@ -767,10 +767,10 @@ func (store *CodexLeaseStore) buildCodexLeaseRecordAfterImage(old CodexJournalRe } validLatchMigration := migrateTurnStateLatch && beginRequest && old.HasTurnState && !old.TurnStateLatchCurrent && input.TurnStateLatchCurrent && old.EverAdmitted && constantTimeCodexLeaseDigestEqual(old.AccountHash, input.AccountHash) && codexLeaseRuntimeCanBeginRequest(old) - if old.TurnStateLatchCurrent != input.TurnStateLatchCurrent && !validLatchMigration && !(old.HasTurnState == false && input.HasTurnState && input.TurnStateLatchCurrent) { + if old.TurnStateLatchCurrent != input.TurnStateLatchCurrent && !validLatchMigration && !(bindingReassignment && !input.HasTurnState && !input.TurnStateLatchCurrent) && !(old.HasTurnState == false && input.HasTurnState && input.TurnStateLatchCurrent) { return CodexJournalRecordV2{}, 0, false, fmt.Errorf("%w: turn-state latch marker changed outside admission or migration", ErrCodexLeaseInvalidMutation) } - if old.HasResponseAnchor && (!input.HasResponseAnchor || input.CorrelationHash == "") { + if old.HasResponseAnchor && (!input.HasResponseAnchor || input.CorrelationHash == "") && !bindingReassignment { return CodexJournalRecordV2{}, 0, false, fmt.Errorf("%w: response anchor was cleared", ErrCodexLeaseInvalidMutation) } result.RecordGeneration = old.RecordGeneration + 1 diff --git a/internal/proxy/codex_lease_v2_runtime.go b/internal/proxy/codex_lease_v2_runtime.go index 3f53179..f1bf185 100644 --- a/internal/proxy/codex_lease_v2_runtime.go +++ b/internal/proxy/codex_lease_v2_runtime.go @@ -777,7 +777,8 @@ func (runtime *CodexLeaseRuntime) BeginRequestContext(ctx context.Context, plan return nil, err } handle.newTurn = newTurn - handle.relatchTurnStateOnAdmission = runtime.canRelatchAuthenticatedTurnState(current.Record, selected.AccountKey, plan.Evidence, plan.authenticatedCallerContinuity) + handle.relatchTurnStateOnAdmission = runtime.canRelatchAuthenticatedTurnState(current.Record, selected.AccountKey, plan.Evidence, plan.authenticatedCallerContinuity) || + (current.Record.HasTurnState && codexLeasePortableUnavailableContinuation(current.Record, plan.Evidence)) return handle, nil } @@ -1488,6 +1489,9 @@ func (handle *CodexLeaseRequestHandle) applyAccountUnavailableRebind(desired *Co return ErrCodexLeaseTransition } desired.AccountHash = accountHash + // The previous account's response ID cannot authorise this account. + desired.CorrelationHash = "" + desired.HasResponseAnchor = false return nil } @@ -2030,6 +2034,7 @@ func (runtime *CodexLeaseRuntime) validateRequestContinuity(restored CodexRestor if newTurn && authenticatedCallerContinuity && (!found || !authority.Record.Authoritative) { return true, nil } + portableUnavailable := found && !newTurn && codexLeasePortableUnavailableContinuation(authority.Record, evidence) if !newTurn && found { ingress := firstIngressContinuity(ingressContinuity) canMigrateLatch := runtime.canMigrateAuthenticatedTurnStateLatch(authority.Record, selected, evidence, authenticatedCallerContinuity) @@ -2037,7 +2042,7 @@ func (runtime *CodexLeaseRuntime) validateRequestContinuity(restored CodexRestor return false, &codexContinuityError{reason: codexContinuityTurnStateMismatch} } missingAuthenticatedState := authenticatedCallerContinuity && authority.Record.HasTurnState && !evidence.HasTurnState - if authority.Record.HasTurnState != evidence.HasTurnState && !missingAuthenticatedState { + if authority.Record.HasTurnState != evidence.HasTurnState && !missingAuthenticatedState && !portableUnavailable { return false, &codexContinuityError{reason: codexContinuityTurnStatePresenceMismatch} } if evidence.HasTurnState && !constantTimeCodexLeaseDigestEqual(authority.Record.TurnStateHash, runtime.store.hash("turn-state", evidence.TurnState)) && @@ -2051,7 +2056,7 @@ func (runtime *CodexLeaseRuntime) validateRequestContinuity(restored CodexRestor } } requiresAccount := authenticatedCallerContinuity || evidence.PreviousResponseID != "" || evidence.HasTurnState || (found && !newTurn && codexLeaseRecordRequiresAccount(authority.Record)) - if found && codexLeaseCurrentAttemptState(authority.Record) == CodexAttemptAccountUnavailable && evidence.PreviousResponseID == "" && !evidence.HasTurnState { + if portableUnavailable { requiresAccount = false } if requiresAccount && (!found || authority.Record.AccountHash == "" || !constantTimeCodexLeaseDigestEqual(authority.Record.AccountHash, runtime.store.hash("account", string(selected)))) { @@ -2060,6 +2065,12 @@ func (runtime *CodexLeaseRuntime) validateRequestContinuity(restored CodexRestor return requiresAccount, nil } +// A full create after a drained account rejection can establish new provider +// state. Keep the old admission evidence until the replacement is admitted. +func codexLeasePortableUnavailableContinuation(record CodexJournalRecordV2, evidence CodexLeaseRequestEvidence) bool { + return evidence.PreviousResponseID == "" && !evidence.HasTurnState && codexLeaseAccountUnavailableCanBeginRequest(record) +} + func (runtime *CodexLeaseRuntime) canMigrateAuthenticatedTurnStateLatch(record CodexJournalRecordV2, selected codex.AccountKey, evidence CodexLeaseRequestEvidence, authenticatedCallerContinuity bool) bool { return runtime != nil && runtime.store != nil && authenticatedCallerContinuity && evidence.HasTurnState && record.HasTurnState && !record.TurnStateLatchCurrent && record.AccountHash != "" && constantTimeCodexLeaseDigestEqual(record.AccountHash, runtime.store.hash("account", string(selected))) && @@ -2139,6 +2150,11 @@ func (handle *CodexLeaseRequestHandle) applyAdmissionEvidence(record *CodexJourn return fmt.Errorf("%w: invalid HTTP admission evidence", ErrCodexLeaseInvalidMutation) } if !evidence.HasTurnState { + if codexLeaseAccountUnavailableAdmission(handle.record, *record) { + record.TurnStateHash = "" + record.HasTurnState = false + record.TurnStateLatchCurrent = false + } return nil } if record.HasTurnState && !handle.relatchTurnStateOnAdmission { diff --git a/internal/proxy/codex_lease_v2_runtime_test.go b/internal/proxy/codex_lease_v2_runtime_test.go index d8eebfa..8babfce 100644 --- a/internal/proxy/codex_lease_v2_runtime_test.go +++ b/internal/proxy/codex_lease_v2_runtime_test.go @@ -4393,3 +4393,108 @@ func waitCodexLeaseRuntimeGateRefs(t *testing.T, gates *codexAccountGateSet, acc } t.Fatalf("account gate %q did not reach %d references", account, want) } + +func TestCodexLeaseRuntimePortableQuotaResetReplacesTurnStateOnAdmission(t *testing.T) { + t.Parallel() + for _, replacementState := range []string{"state-b", ""} { + t.Run("replacement_state="+replacementState, func(t *testing.T) { + coordinator, _, _ := openCodexLeaseRuntimeTestCoordinator(t) + runtimeLease := newCodexLeaseRuntimeTest(t, coordinator) + plan := codexLeaseRuntimeTestPlan("turn", []CodexLeaseAttemptSlotPlan{{AccountKey: "account-a", CandidateID: "initial-a", Kind: CodexAttemptSlotDirect}}) + plan.Accounts = []codex.AccountKey{"account-a", "account-b"} + handle, err := runtimeLease.BeginRequest(plan) + if err != nil { + t.Fatal(err) + } + handle, err = handle.MarkDispatched() + if err != nil { + t.Fatal(err) + } + handle, err = handle.AdmitHTTP2xxContext(context.Background(), CodexHTTPAdmissionEvidence{TurnState: "state-a", HasTurnState: true}) + if err != nil { + t.Fatal(err) + } + handle, err = handle.ProviderCompleted(CodexHTTPCompletionEvidence{CodexHTTPResponseEvidence: CodexHTTPResponseEvidence{ResponseAnchor: "response-a", HasResponseAnchor: true}, EndTurn: false}) + if err != nil { + t.Fatal(err) + } + handle, err = handle.Drain() + if err != nil { + t.Fatal(err) + } + + // An authenticated full create can omit the provider's latched turn state. + bound := plan + bound.RequiresAccountContinuity = true + bound.authenticatedCallerContinuity = true + bound.ExpectedBound = &CodexLeaseBoundExpectation{Identity: handle.identity, AccountKey: "account-a", RecordGeneration: handle.record.RecordGeneration} + handle, err = runtimeLease.BeginRequest(bound) + if err != nil { + t.Fatal(err) + } + handle, err = handle.MarkDispatched() + if err != nil { + t.Fatal(err) + } + handle, err = handle.RecordQuotaExhausted(0) + if err != nil { + t.Fatal(err) + } + + fresh := plan + fresh.Slots = []CodexLeaseAttemptSlotPlan{{AccountKey: "account-b", CandidateID: "retry-b", Kind: CodexAttemptSlotDirect}} + for _, evidence := range []CodexLeaseRequestEvidence{ + {PreviousResponseID: "response-a"}, + {TurnState: "state-a", HasTurnState: true}, + } { + invalid := fresh + invalid.Evidence = evidence + if _, err := runtimeLease.BeginRequest(invalid); !errors.Is(err, ErrCodexContinuity) { + t.Fatalf("provider continuation migrated: %v", err) + } + } + handle, err = runtimeLease.BeginRequest(fresh) + if err != nil { + t.Fatalf("portable quota reset: %v", err) + } + if handle.record.TurnStateHash != coordinator.store.hash("turn-state", "state-a") { + t.Fatal("prepared retry changed admitted evidence") + } + handle, err = handle.MarkDispatched() + if err != nil { + t.Fatal(err) + } + handle, err = handle.AdmitHTTP2xxContext(context.Background(), CodexHTTPAdmissionEvidence{TurnState: replacementState, HasTurnState: replacementState != ""}) + if err != nil { + t.Fatal(err) + } + if handle.AccountKey() != "account-b" { + t.Fatal("new admission retained exhausted account evidence") + } + handle, err = handle.ProviderFailed(CodexHTTPResponseEvidence{}) + if err != nil { + t.Fatal(err) + } + handle, err = handle.Drain() + if err != nil { + t.Fatal(err) + } + oldAnchor := fresh + oldAnchor.Evidence = CodexLeaseRequestEvidence{PreviousResponseID: "response-a", TurnState: replacementState, HasTurnState: replacementState != ""} + if _, err := runtimeLease.BeginRequest(oldAnchor); !errors.Is(err, ErrCodexContinuity) { + t.Fatalf("old account response anchor accepted after replacement: %v", err) + } + if replacementState == "" { + fresh.Evidence = CodexLeaseRequestEvidence{TurnState: "state-a", HasTurnState: true} + if _, err := runtimeLease.BeginRequest(fresh); !errors.Is(err, ErrCodexContinuity) { + t.Fatalf("old account state accepted after stateless replacement: %v", err) + } + } else { + fresh.Evidence = CodexLeaseRequestEvidence{TurnState: replacementState, HasTurnState: true} + if _, err := runtimeLease.BeginRequest(fresh); err != nil { + t.Fatalf("replacement continuation: %v", err) + } + } + }) + } +} diff --git a/internal/proxy/codex_native_http.go b/internal/proxy/codex_native_http.go index 1e5186e..e4cc351 100644 --- a/internal/proxy/codex_native_http.go +++ b/internal/proxy/codex_native_http.go @@ -209,113 +209,131 @@ func (handler *CodexNativeHTTPHandler) TryServe(writer http.ResponseWriter, requ } func (handler *CodexNativeHTTPHandler) serveEncoded(writer http.ResponseWriter, request *http.Request, compact bool, encoded []byte, claim *CodexRetainedHTTPRequestClaim) (bool, string) { - trace := codexInstalledHTTPTraceFromContext(request.Context()) - emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "planning", Outcome: "started"}) - prepared, err := func() (CodexPreparedHTTPRequest, error) { - ctx, release := handler.requestPlanningContext(request.Context()) - defer release() - input := CodexHTTPRequestPlanInput{Encoded: encoded, Headers: request.Header} - if claim != nil { - expected := claim.ExpectedBound - input.ExpectedBound = &expected - input.retainedPlanning = claim.planning - } - return handler.planner.Build(ctx, input) - }() - clearBytes(encoded) - if err != nil { - if writeCodexCapacityError(writer, err) { + defer clearBytes(encoded) + // Freeze the retry budget from the first plan so a changing inventory cannot + // keep one downstream request retrying indefinitely. + remainingQuotaRetries := -1 + for { + trace := codexInstalledHTTPTraceFromContext(request.Context()) + emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "planning", Outcome: "started"}) + prepared, err := func() (CodexPreparedHTTPRequest, error) { + ctx, release := handler.requestPlanningContext(request.Context()) + defer release() + input := CodexHTTPRequestPlanInput{Encoded: encoded, Headers: request.Header} + if claim != nil { + expected := claim.ExpectedBound + input.ExpectedBound = &expected + input.retainedPlanning = claim.planning + } + return handler.planner.Build(ctx, input) + }() + if err != nil { + if writeCodexCapacityError(writer, err) { + return true, "" + } + status := http.StatusServiceUnavailable + errorType := "api_error" + message := "Codex native HTTP routing unavailable" + failure := CodexHTTPRequestPlanFailure{Stage: codexHTTPRequestPlanUnknown, Reason: CodexRequestFailureUnknown} + var planErr *CodexHTTPRequestPlanError + if errors.As(err, &planErr) { + failure.Stage = safeCodexHTTPRequestPlanErrorCode(planErr.Code) + failure.Reason = safeCodexRequestFailureReason(planErr.Reason) + if planErr.Code == CodexHTTPRequestPlanInspect { + status = http.StatusBadRequest + errorType = "invalid_request_error" + message = "invalid Codex Responses request" + } + } + noteCodexObservation(request.Context(), codexObservationFields{Decision: "plan_failed", Reason: string(failure.Reason)}) + emitCodexTrace(request.Context(), CodexTraceEvent{ + Phase: "planning", Stage: string(failure.Stage), Outcome: "error", Reason: string(failure.Reason), StatusCode: status, + }) + if handler.reportPlanFailure != nil { + handler.reportPlanFailure(failure) + } + writeError(writer, status, errorType, message) return true, "" } - status := http.StatusServiceUnavailable - errorType := "api_error" - message := "Codex native HTTP routing unavailable" - failure := CodexHTTPRequestPlanFailure{Stage: codexHTTPRequestPlanUnknown, Reason: CodexRequestFailureUnknown} - var planErr *CodexHTTPRequestPlanError - if errors.As(err, &planErr) { - failure.Stage = safeCodexHTTPRequestPlanErrorCode(planErr.Code) - failure.Reason = safeCodexRequestFailureReason(planErr.Reason) - if planErr.Code == CodexHTTPRequestPlanInspect { - status = http.StatusBadRequest - errorType = "invalid_request_error" - message = "invalid Codex Responses request" - } + + if remainingQuotaRetries < 0 { + remainingQuotaRetries = len(prepared.Dispatch.AccountUnavailableResetCandidates()) + } + model := "" + if accounts := prepared.Dispatch.Accounts(); len(accounts) > 0 { + model = accounts[0].Choice().EffectiveModel } - noteCodexObservation(request.Context(), codexObservationFields{Decision: "plan_failed", Reason: string(failure.Reason)}) emitCodexTrace(request.Context(), CodexTraceEvent{ - Phase: "planning", Stage: string(failure.Stage), Outcome: "error", Reason: string(failure.Reason), StatusCode: status, + Phase: "planning", Outcome: "success", Candidates: codexTraceDispatchCandidates(prepared.Dispatch), }) - if handler.reportPlanFailure != nil { - handler.reportPlanFailure(failure) - } - writeError(writer, status, errorType, message) - return true, "" - } - - model := "" - if accounts := prepared.Dispatch.Accounts(); len(accounts) > 0 { - model = accounts[0].Choice().EffectiveModel - } - emitCodexTrace(request.Context(), CodexTraceEvent{ - Phase: "planning", Outcome: "success", Candidates: codexTraceDispatchCandidates(prepared.Dispatch), - }) - template := handler.requestTemplate(request, compact) - emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "session", Outcome: "started"}) - result, err := handler.session.Do( - request.Context(), - template, - prepared.Dispatch, - prepared.Frozen, - prepared.Lifecycle, - ) - if err != nil { - if writeCodexCapacityError(writer, err) { + canReplanQuota := prepared.portableQuotaRetry && prepared.Lifecycle != nil && prepared.Lifecycle.EverAdmitted() + template := handler.requestTemplate(request, compact) + emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "session", Outcome: "started"}) + result, err := handler.session.Do( + request.Context(), + template, + prepared.Dispatch, + prepared.Frozen, + prepared.Lifecycle, + ) + if err != nil { + if writeCodexCapacityError(writer, err) { + return true, model + } + failure := classifyCodexNativeHTTPSessionFailure(err) + noteCodexObservation(request.Context(), codexObservationFields{Decision: "session_failed", Reason: failure.reason}) + event := CodexTraceEvent{Phase: "session", Stage: failure.stage, Outcome: "error", Reason: failure.reason, StatusCode: http.StatusBadGateway} + if failure.roundTrip != nil { + event.ErrorClass = string(failure.roundTrip.reason) + } + emitCodexTrace(request.Context(), event) + if handler.reportSessionFailure != nil { + handler.reportSessionFailure(failure) + } + writeError(writer, http.StatusBadGateway, "api_error", "Codex upstream request failed") return true, model } - failure := classifyCodexNativeHTTPSessionFailure(err) - noteCodexObservation(request.Context(), codexObservationFields{Decision: "session_failed", Reason: failure.reason}) - event := CodexTraceEvent{Phase: "session", Stage: failure.stage, Outcome: "error", Reason: failure.reason, StatusCode: http.StatusBadGateway} - if failure.roundTrip != nil { - event.ErrorClass = string(failure.roundTrip.reason) + if result.Response == nil || result.Response.Body == nil { + failure := codexNativeHTTPSessionFailure{stage: "response_validate", reason: "response_unavailable"} + noteCodexObservation(request.Context(), codexObservationFields{Decision: "session_failed", Reason: failure.reason}) + emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "session", Stage: failure.stage, Outcome: "error", Reason: failure.reason, StatusCode: http.StatusBadGateway}) + if handler.reportSessionFailure != nil { + handler.reportSessionFailure(failure) + } + writeError(writer, http.StatusBadGateway, "api_error", "Codex upstream response unavailable") + return true, model } - emitCodexTrace(request.Context(), event) - if handler.reportSessionFailure != nil { - handler.reportSessionFailure(failure) + emitCodexTrace(request.Context(), CodexTraceEvent{ + Phase: "upstream_response", Outcome: codexTraceHTTPOutcome(result.Response.StatusCode), + UpstreamStatus: result.Response.StatusCode, AccountHint: codexTraceAccountHint(result.Choice.AccountKey), Attempt: result.Attempt.Ordinal, + }) + + // The session has durably rejected this account before any response bytes + // reached the client. A full create can now acquire a replacement binding. + if result.quotaExhausted && canReplanQuota && claim == nil && remainingQuotaRetries > 0 && request.Context().Err() == nil { + closeCodexHTTPResponseBody(result.Response.Body) + remainingQuotaRetries-- + emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "failover", Outcome: "replan", Reason: "capacity_exhausted", Retry: true, Failover: true}) + continue } - writeError(writer, http.StatusBadGateway, "api_error", "Codex upstream request failed") - return true, model - } - if result.Response == nil || result.Response.Body == nil { - failure := codexNativeHTTPSessionFailure{stage: "response_validate", reason: "response_unavailable"} - noteCodexObservation(request.Context(), codexObservationFields{Decision: "session_failed", Reason: failure.reason}) - emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "session", Stage: failure.stage, Outcome: "error", Reason: failure.reason, StatusCode: http.StatusBadGateway}) - if handler.reportSessionFailure != nil { - handler.reportSessionFailure(failure) + + if result.Response.StatusCode < http.StatusOK || result.Response.StatusCode >= http.StatusMultipleChoices { + defer closeCodexHTTPResponseBody(result.Response.Body) + relayErr := relayCodexHTTPResponse(writer, result.Response, false) + emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "relay", Outcome: codexTraceRelayOutcome(relayErr), UpstreamStatus: result.Response.StatusCode, Reason: codexTraceErrorReason(relayErr)}) + trace.relayedResponse(false, true, relayErr) + return true, model } - writeError(writer, http.StatusBadGateway, "api_error", "Codex upstream response unavailable") - return true, model - } - emitCodexTrace(request.Context(), CodexTraceEvent{ - Phase: "upstream_response", Outcome: codexTraceHTTPOutcome(result.Response.StatusCode), - UpstreamStatus: result.Response.StatusCode, AccountHint: codexTraceAccountHint(result.Choice.AccountKey), Attempt: result.Attempt.Ordinal, - }) - if result.Response.StatusCode < http.StatusOK || result.Response.StatusCode >= http.StatusMultipleChoices { - defer closeCodexHTTPResponseBody(result.Response.Body) - relayErr := relayCodexHTTPResponse(writer, result.Response, false) + mode := codexHTTPResponseModeSSE + if compact { + mode = codexHTTPResponseModeCompact + } + relayErr := relayCodexAcceptedHTTPResponse(request.Context(), writer, result.Response, mode, result.Lifecycle) emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "relay", Outcome: codexTraceRelayOutcome(relayErr), UpstreamStatus: result.Response.StatusCode, Reason: codexTraceErrorReason(relayErr)}) - trace.relayedResponse(false, true, relayErr) + trace.relayedResponse(true, false, relayErr) return true, model } - - mode := codexHTTPResponseModeSSE - if compact { - mode = codexHTTPResponseModeCompact - } - relayErr := relayCodexAcceptedHTTPResponse(request.Context(), writer, result.Response, mode, result.Lifecycle) - emitCodexTrace(request.Context(), CodexTraceEvent{Phase: "relay", Outcome: codexTraceRelayOutcome(relayErr), UpstreamStatus: result.Response.StatusCode, Reason: codexTraceErrorReason(relayErr)}) - trace.relayedResponse(true, false, relayErr) - return true, model } func (handler *CodexNativeHTTPHandler) requestPlanningContext(requestContext context.Context) (context.Context, func()) { diff --git a/internal/proxy/codex_normal_transport_integration_test.go b/internal/proxy/codex_normal_transport_integration_test.go index 14c9ca1..5ae8921 100644 --- a/internal/proxy/codex_normal_transport_integration_test.go +++ b/internal/proxy/codex_normal_transport_integration_test.go @@ -105,6 +105,7 @@ type normalTransportGateBackend struct { httpAuthRecovered bool httpAuthActive bool httpAuthStatus int + httpTurnState bool wsConnections int failures chan error firstWSClosed chan struct{} @@ -253,6 +254,9 @@ func (backend *normalTransportGateBackend) serveHTTP(writer http.ResponseWriter, } default: writer.Header().Set("Content-Type", "text/event-stream") + if backend.httpTurnState { + writer.Header().Set("X-Codex-Turn-State", "state-"+accountID) + } encryptedState := "" if normalTransportGateNonPortableWebSocketScenario(backend.scenario) { encryptedState = `,"encrypted_content":"opaque-normal-transport-state"` @@ -1497,6 +1501,83 @@ func TestNormalProxyTransportHTTPHardLimitMigratesBeforeLeak(t *testing.T) { harness.backend.assertNoFailure(t) } +func TestNormalProxyTransportHTTPAdmittedQuotaFailureRetriesWithinRequest(t *testing.T) { + harness := newNormalTransportGateCodexCallerHarness(t, normalTransportGateHTTPSuccess) + harness.backend.httpTurnState = true + metadata := CodexTurnMetadata{SessionID: "automatic-quota-session", ThreadID: "automatic-quota-thread", TurnID: "automatic-quota-turn", RequestKind: CodexRequestTurn} + encoded := normalTransportGateHTTPBody(t, metadata) + status, body := normalTransportGateHTTPCall(t, harness, encoded) + if status != http.StatusOK { + t.Fatalf("seed = %d %q", status, body) + } + harness.backend.scenario = normalTransportGateHTTPHardLimit + status, body = normalTransportGateHTTPCall(t, harness, encoded) + if status != http.StatusOK || bytes.Contains(body, []byte("usage_limit_reached")) || !bytes.Contains(body, []byte(`"type":"response.completed"`)) { + t.Fatalf("same-turn quota recovery = %d %q, want automatic B/200", status, body) + } + status, body = normalTransportGateHTTPCall(t, harness, encoded) + if status != http.StatusOK { + t.Fatalf("continuation after recovery = %d %q", status, body) + } + receipts := normalTransportGateReceipts(harness.backend.snapshot(), "http") + if len(receipts) != 4 || receipts[0].accountID != "validation-upstream-a" || receipts[1].status != http.StatusTooManyRequests || receipts[2].accountID != "validation-upstream-b" || receipts[3].accountID != "validation-upstream-b" || receipts[1].payload != receipts[2].payload { + t.Fatalf("receipts = %#v, want A/200, A/429, identical B/200, B/200", receipts) + } + harness.backend.assertNoFailure(t) +} + +func TestNormalProxyTransportHTTPAdmittedQuotaRetryBounds(t *testing.T) { + for _, test := range []struct { + name string + scenario normalTransportGateScenario + continuation bool + wantAttempts int + }{ + {"all exhausted", normalTransportGateHTTPAllHardLimit, false, 4}, + {"soft limit", normalTransportGateHTTPSoftLimit, false, 2}, + {"response ID continuation", normalTransportGateHTTPHardLimit, true, 2}, + } { + t.Run(test.name, func(t *testing.T) { + harness := newNormalTransportGateCodexCallerHarness(t, normalTransportGateHTTPSuccess) + harness.backend.httpTurnState = true + metadata := CodexTurnMetadata{SessionID: "bounded-quota-session", ThreadID: "bounded-quota-thread", TurnID: "bounded-quota-turn", RequestKind: CodexRequestTurn} + encoded := normalTransportGateHTTPBody(t, metadata) + status, body := normalTransportGateHTTPCall(t, harness, encoded) + if status != http.StatusOK { + t.Fatalf("seed = %d %q", status, body) + } + harness.backend.scenario = test.scenario + if test.continuation { + encoded = normalTransportGateHTTPContinuationBody(t, metadata, "normal-transport-http") + } + status, body = normalTransportGateHTTPCall(t, harness, encoded) + if status != http.StatusTooManyRequests { + t.Fatalf("terminal status = %d %q, want429", status, body) + } + receipts := normalTransportGateReceipts(harness.backend.snapshot(), "http") + if len(receipts) != test.wantAttempts { + t.Fatalf("attempts = %d, want %d: %#v", len(receipts), test.wantAttempts, receipts) + } + harness.backend.assertNoFailure(t) + }) + } +} + +func TestNormalProxyTransportHTTPFreshQuotaExhaustionDoesNotReprobe(t *testing.T) { + harness := newNormalTransportGateCodexCallerHarness(t, normalTransportGateHTTPAllHardLimit) + harness.httpPlanner.DefaultAccountKey = "" + metadata := CodexTurnMetadata{SessionID: "fresh-quota-session", ThreadID: "fresh-quota-thread", TurnID: "fresh-quota-turn", RequestKind: CodexRequestTurn} + status, body := normalTransportGateHTTPCall(t, harness, normalTransportGateHTTPBody(t, metadata)) + if status != http.StatusTooManyRequests { + t.Fatalf("status=%d %q, want429", status, body) + } + receipts := normalTransportGateReceipts(harness.backend.snapshot(), "http") + if len(receipts) != 2 { + t.Fatalf("attempts=%d, want2: %#v", len(receipts), receipts) + } + harness.backend.assertNoFailure(t) +} + func TestNormalProxyTransportRequiredAffinityHardLimitReplaysPortableRequest(t *testing.T) { for _, test := range []struct { name string