diff --git a/backend/internal/repository/account_repo_codex_lifecycle.go b/backend/internal/repository/account_repo_codex_lifecycle.go index 4bd08b1d1668..d6a02e9bebbc 100644 --- a/backend/internal/repository/account_repo_codex_lifecycle.go +++ b/backend/internal/repository/account_repo_codex_lifecycle.go @@ -117,6 +117,11 @@ func (r *accountRepository) SeedCodexTicketRenewals(ctx context.Context, models if err != nil { return err } + // The scheduler metadata projection intentionally omits ticket material, so + // the authoritative full-account snapshot must be refreshed after a restart. + // Otherwise a valid durable ticket can exist while the cached account still + // appears ticketless to the full-account gate. + syncAccountIDs := make([]int64, 0, len(accounts)) for i := range accounts { if !service.CodexTicketAccountEligible(&accounts[i]) { continue @@ -128,8 +133,12 @@ func (r *accountRepository) SeedCodexTicketRenewals(ctx context.Context, models if err := r.seedCodexTicket(ctx, accounts[i].ID, model, length, now); err != nil { return err } + syncAccountIDs = append(syncAccountIDs, accounts[i].ID) } } + if len(syncAccountIDs) > 0 { + r.syncSchedulerAccountSnapshots(ctx, syncAccountIDs) + } return nil } diff --git a/backend/internal/repository/account_repo_codex_lifecycle_integration_test.go b/backend/internal/repository/account_repo_codex_lifecycle_integration_test.go index 5890d8fd71b4..8f3986ff22e5 100644 --- a/backend/internal/repository/account_repo_codex_lifecycle_integration_test.go +++ b/backend/internal/repository/account_repo_codex_lifecycle_integration_test.go @@ -154,3 +154,58 @@ func TestCodexTicketRepositoryInterruptedRenewalIsConsumed(t *testing.T) { require.NoError(t, integrationDB.QueryRowContext(ctx, `SELECT phase FROM openai_codex_ticket_runtime WHERE account_id=$1 AND model=$2`, a.ID, model).Scan(&phase)) require.Equal(t, "stopped", phase) } + +type codexTicketSnapshotRecorder struct { + service.SchedulerCache + mu sync.Mutex + accountIDs []int64 +} + +func (r *codexTicketSnapshotRecorder) SetAccount(_ context.Context, account *service.Account) error { + if r == nil || account == nil { + return nil + } + r.mu.Lock() + defer r.mu.Unlock() + r.accountIDs = append(r.accountIDs, account.ID) + return nil +} + +func (r *codexTicketSnapshotRecorder) contains(accountID int64) bool { + r.mu.Lock() + defer r.mu.Unlock() + for _, id := range r.accountIDs { + if id == accountID { + return true + } + } + return false +} + +func TestCodexTicketSeedRefreshesAuthoritativeSchedulerSnapshot(t *testing.T) { + ctx := context.Background() + client := testEntClient(t) + now := time.Now().UTC().Truncate(time.Second) + model := "gpt-6-astra" + account := mustCreateAccount(t, client, &service.Account{ + Name: t.Name(), + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Credentials: map[string]any{"access_token": "seed-refresh-token", "chatgpt_account_id": "seed-refresh-principal"}, + Extra: map[string]any{}, + }) + t.Cleanup(func() { + _, _ = integrationDB.ExecContext(ctx, `DELETE FROM accounts WHERE id=$1`, account.ID) + }) + account.Extra["codex_turn_ticket:"+model] = successfulTicket(account, model, now) + raw, err := json.Marshal(account.Extra) + require.NoError(t, err) + _, err = integrationDB.ExecContext(ctx, `UPDATE accounts SET extra=$2::jsonb WHERE id=$1`, account.ID, string(raw)) + require.NoError(t, err) + + recorder := &codexTicketSnapshotRecorder{} + repo := newAccountRepositoryWithSQL(client, integrationDB, recorder) + require.NoError(t, repo.SeedCodexTicketRenewals(ctx, []string{model}, 292, now)) + require.True(t, recorder.contains(account.ID), "startup seeding must refresh the full authoritative scheduler snapshot") +} diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index 55fd88bf2d00..80bc7ac57ca5 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -913,6 +913,11 @@ func (c *schedulerCache) mgetChunked(ctx context.Context, keys []string) ([]any, return out, nil } +// buildSchedulerMetadataAccount builds the partial candidate projection used by +// the scheduler snapshot. It intentionally omits credential-bound material +// (including Codex turn-state tickets) and the full credential identity. Gates +// that depend on those fields must run on the authoritative account returned by +// GetAccount/recheckSelectedOpenAIAccountFromDB, never on this projection. func buildSchedulerMetadataAccount(account service.Account) service.Account { return service.Account{ ID: account.ID, @@ -1007,8 +1012,9 @@ func filterSchedulerCredentials(credentials map[string]any) map[string]any { if len(credentials) == 0 { return nil } - // Candidate-list admission evaluates the account override before hydrating - // the full account. Preserve routing and capability fields used by provider gates. + // Candidate-list admission evaluates projection-safe routing and capability + // fields before hydrating the full account. Credential-bound identity and + // ticket fields are deliberately excluded; see buildSchedulerMetadataAccount. keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold", "base_url", "openai_capabilities"} filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/service/openai_account_runtime_block_fastpath.go b/backend/internal/service/openai_account_runtime_block_fastpath.go index 1f94202d1abe..af02c10ed39a 100644 --- a/backend/internal/service/openai_account_runtime_block_fastpath.go +++ b/backend/internal/service/openai_account_runtime_block_fastpath.go @@ -1572,11 +1572,34 @@ func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlocked(account *Acc } func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlockedContext(ctx context.Context, account *Account, requestedModel string, requireCompact ...bool) bool { + return s.isOpenAIAccountRequestRuntimeBlockedContextInternal(ctx, account, requestedModel, true, requireCompact...) +} + +// isOpenAIAccountCandidateRuntimeBlockedContext evaluates only gates that can +// be proven from the scheduler's partial candidate projection. Credential-bound +// gates, including Codex turn-state tickets, are intentionally deferred to the +// authoritative account returned by resolveFreshSchedulableOpenAIAccount or +// recheckSelectedOpenAIAccountFromDB. Running those gates on a projection whose +// credentials and Extra were filtered would turn a valid ticket into a +// fail-closed runtime block. +func (s *OpenAIGatewayService) isOpenAIAccountCandidateRuntimeBlockedContext(ctx context.Context, account *Account, requestedModel string, requireCompact ...bool) bool { + return s.isOpenAIAccountRequestRuntimeBlockedContextInternal(ctx, account, requestedModel, false, requireCompact...) +} + +func (s *OpenAIGatewayService) isOpenAIAccountRequestRuntimeBlockedContextInternal( + ctx context.Context, + account *Account, + requestedModel string, + evaluateCodexTicket bool, + requireCompact ...bool, +) bool { if s == nil || account == nil { return false } // Only OAuth/setup-token accounts can own tickets; API-key paths stay unchanged. - if isOpenAICodexTicketAccount(account) { + // Scheduler candidate projections intentionally omit ticket material and the + // identity fields needed to validate its owner, so that check is authoritative-only. + if evaluateCodexTicket && isOpenAICodexTicketAccount(account) { compact := len(requireCompact) > 0 && requireCompact[0] outboundModel := s.openAICodexTicketOutboundModel(account, requestedModel, compact) if s.openAICodexTicketBlocksAccount(account, outboundModel) { diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index ed1984400a8b..0a0092b3e4bf 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -1520,7 +1520,10 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( filterStats.exclude("platform_mismatch") continue } - if s.service.isOpenAIAccountRequestRuntimeBlockedContext(ctx, account, openAIRequestedModelForAccount(ctx, account, req.RequestedModel), req.RequireCompact) { + // This is the partial scheduler projection. Ticket and other + // credential-bound gates are evaluated on the authoritative account + // during the fresh/recheck steps below. + if s.service.isOpenAIAccountCandidateRuntimeBlockedContext(ctx, account, openAIRequestedModelForAccount(ctx, account, req.RequestedModel), req.RequireCompact) { filterStats.exclude("runtime_blocked") filterStats.observeRuntimeCooldown(s.service, account.ID) continue @@ -1531,7 +1534,7 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( filterStats.exclude("privacy_not_set") continue } - if compatible, reason := s.isAccountRequestCompatibleReason(ctx, account, req); !compatible { + if compatible, reason := s.isCandidateAccountRequestCompatibleReason(ctx, account, req); !compatible { filterStats.exclude(reason) continue } @@ -1836,6 +1839,18 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(ctx context.C // openAISelectionFilterStats so that "no available accounts" errors state why // each candidate was dropped instead of failing silently (#4599). func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx context.Context, account *Account, req OpenAIAccountScheduleRequest) (bool, string) { + return s.isAccountRequestCompatibleReasonWithRuntimeGate(ctx, account, req, true) +} + +// isCandidateAccountRequestCompatibleReason is the projection-safe variant used +// while filtering the partial scheduler snapshot. The authoritative variant runs +// again after the account is hydrated/rechecked, so ticket and other +// credential-bound gates are never lost. +func (s *defaultOpenAIAccountScheduler) isCandidateAccountRequestCompatibleReason(ctx context.Context, account *Account, req OpenAIAccountScheduleRequest) (bool, string) { + return s.isAccountRequestCompatibleReasonWithRuntimeGate(ctx, account, req, false) +} + +func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReasonWithRuntimeGate(ctx context.Context, account *Account, req OpenAIAccountScheduleRequest, authoritative bool) (bool, string) { if account == nil { return false, "account_nil" } @@ -1846,8 +1861,16 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx con return false, "provider_identity_mismatch" } requestedModel := openAIRequestedModelForAccount(ctx, account, req.RequestedModel) - if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlockedContext(ctx, account, requestedModel, req.RequireCompact) { - return false, "runtime_blocked" + if s != nil && s.service != nil { + blocked := false + if authoritative { + blocked = s.service.isOpenAIAccountRequestRuntimeBlockedContext(ctx, account, requestedModel, req.RequireCompact) + } else { + blocked = s.service.isOpenAIAccountCandidateRuntimeBlockedContext(ctx, account, requestedModel, req.RequireCompact) + } + if blocked { + return false, "runtime_blocked" + } } // The scheduler's candidate filter already checks the account's ordinary // active/schedulable state. This gate is intentionally only the per-model diff --git a/backend/internal/service/openai_codex_ticket_scheduler_test.go b/backend/internal/service/openai_codex_ticket_scheduler_test.go new file mode 100644 index 000000000000..19a27800c8de --- /dev/null +++ b/backend/internal/service/openai_codex_ticket_scheduler_test.go @@ -0,0 +1,181 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func schedulerTicketAccount(id int64, groupID int64, withTicket bool) *Account { + account := &Account{ + ID: id, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + GroupIDs: []int64{groupID}, + Credentials: map[string]any{ + "access_token": "scheduler-ticket-token", + "chatgpt_account_id": "scheduler-principal", + "model_mapping": map[string]any{"gpt-6-astra": "gpt-6-astra"}, + }, + } + if withTicket { + account.Extra = map[string]any{ + openAICodexTicketExtraKey("gpt-6-astra"): &openAICodexTicket{ + AccountID: id, + Model: "gpt-6-astra", + Identity: CodexTicketAccountIdentity(account), + State: fakeCodexTicketState(292), + Length: 292, + CapturedAt: time.Now(), + ExpiresAt: time.Now().Add(time.Hour), + }, + } + } + return account +} + +func schedulerTicketMetadataProjection(account *Account) *Account { + projected := *account + projected.Credentials = map[string]any{ + "model_mapping": map[string]any{"gpt-6-astra": "gpt-6-astra"}, + } + projected.Extra = nil + return &projected +} + +func schedulerTicketConfig() *config.Config { + cfg := &config.Config{ + Gateway: config.GatewayConfig{ + OpenAICodexTicket: config.OpenAICodexTicketConfig{ + Enabled: true, + TargetLength: 292, + TTLSeconds: 3600, + FailClosed: true, + Models: []string{"gpt-6-astra"}, + }, + }, + } + cfg.Gateway.Scheduling.LoadBatchEnabled = true + return cfg +} + +func TestCodexTicketCandidateGateDefersCredentialChecks(t *testing.T) { + svc := ticketTestService(t, schedulerTicketConfig().Gateway.OpenAICodexTicket, nil) + ctx := context.Background() + full := schedulerTicketAccount(41, 101, true) + + require.False(t, svc.isOpenAIAccountRequestRuntimeBlockedContext(ctx, full, "gpt-6-astra")) + + // The projection has neither the ticket nor the identity credentials. The + // candidate gate must not turn that partial view into a fail-closed block. + projected := schedulerTicketMetadataProjection(full) + require.False(t, svc.isOpenAIAccountCandidateRuntimeBlockedContext(ctx, projected, "gpt-6-astra")) + + // A projection that carries a ticket without the owner identity is still + // rejected by the authoritative gate; this preserves principal binding. + projectedWithTicket := *projected + projectedWithTicket.Extra = full.Extra + require.False(t, svc.isOpenAIAccountCandidateRuntimeBlockedContext(ctx, &projectedWithTicket, "gpt-6-astra")) + require.True(t, svc.isOpenAIAccountRequestRuntimeBlockedContext(ctx, &projectedWithTicket, "gpt-6-astra")) + + // Missing, expired, wrong-principal and wrong-length tickets remain blocked + // once the authoritative account is available. + ticketless := *full + ticketless.Extra = nil + require.False(t, svc.isOpenAIAccountCandidateRuntimeBlockedContext(ctx, &ticketless, "gpt-6-astra")) + require.True(t, svc.isOpenAIAccountRequestRuntimeBlockedContext(ctx, &ticketless, "gpt-6-astra")) + + changedPrincipal := *full + changedPrincipal.Credentials = map[string]any{ + "access_token": "scheduler-ticket-token", + "chatgpt_account_id": "another-principal", + } + require.True(t, svc.isOpenAIAccountRequestRuntimeBlockedContext(ctx, &changedPrincipal, "gpt-6-astra")) + + expiredTicket := *full.Extra[openAICodexTicketExtraKey("gpt-6-astra")].(*openAICodexTicket) + expiredTicket.ExpiresAt = time.Now().Add(-time.Minute) + expired := *full + expired.Extra = map[string]any{openAICodexTicketExtraKey("gpt-6-astra"): &expiredTicket} + require.True(t, svc.isOpenAIAccountRequestRuntimeBlockedContext(ctx, &expired, "gpt-6-astra")) + + // Models outside the configured ticket scope keep their existing behavior. + require.False(t, svc.isOpenAIAccountRequestRuntimeBlockedContext(ctx, full, "gpt-5.5")) +} + +func TestSelectAccountWithSchedulerDefersTicketGateToAuthoritativeAccount(t *testing.T) { + for _, tc := range []struct { + name string + advanced bool + }{ + {name: "legacy_load_aware"}, + {name: "advanced_scheduler", advanced: true}, + } { + t.Run(tc.name, func(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + ctx := context.Background() + groupID := int64(92001) + full := schedulerTicketAccount(42, groupID, true) + projected := schedulerTicketMetadataProjection(full) + snapshotCache := &openAISnapshotCacheStub{ + snapshotAccounts: []*Account{projected}, + accountsByID: map[int64]*Account{full.ID: full}, + } + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{*full}}, + cfg: schedulerTicketConfig(), + schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache}, + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + if tc.advanced { + svc.rateLimitService = newOpenAIAdvancedSchedulerRateLimitService("true") + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, &groupID, "", "", "gpt-6-astra", nil, OpenAIUpstreamTransportAny, false, + ) + + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, full.ID, selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + require.NotNil(t, selection.Account.Extra[openAICodexTicketExtraKey("gpt-6-astra")]) + }) + } +} + +func TestSelectAccountWithSchedulerKeepsFailClosedForTicketlessAccounts(t *testing.T) { + resetOpenAIAdvancedSchedulerSettingCacheForTest() + ctx := context.Background() + groupID := int64(92002) + ticketless := schedulerTicketAccount(43, groupID, false) + projected := schedulerTicketMetadataProjection(ticketless) + snapshotCache := &openAISnapshotCacheStub{ + snapshotAccounts: []*Account{projected}, + accountsByID: map[int64]*Account{ticketless.ID: ticketless}, + } + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{*ticketless}}, + cfg: schedulerTicketConfig(), + schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache}, + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + } + + selection, _, err := svc.SelectAccountWithScheduler( + ctx, &groupID, "", "", "gpt-6-astra", nil, OpenAIUpstreamTransportAny, false, + ) + + require.Nil(t, selection) + require.Error(t, err) + require.True(t, errors.Is(err, ErrNoAvailableAccounts), "fail-closed semantics must remain on the authoritative recheck") +} diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index aea1c6abe471..4c18b5c6658e 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -1298,7 +1298,10 @@ func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Contex continue } accountRequestedModel := openAIRequestedModelForAccount(ctx, acc, requestedModel) - if s.isOpenAIAccountRequestRuntimeBlockedContext(ctx, acc, accountRequestedModel, requireCompact) { + // This candidate list is the partial scheduler projection. Credential-bound + // gates run after resolveFreshSchedulableOpenAIAccount/recheck reads the + // authoritative account; doing so here would misclassify valid tickets. + if s.isOpenAIAccountCandidateRuntimeBlockedContext(ctx, acc, accountRequestedModel, requireCompact) { filterStats.exclude("runtime_blocked") filterStats.observeRuntimeCooldown(s, acc.ID) continue