From 25949509938ed2ea1915eb83f836f725ce48cf2d Mon Sep 17 00:00:00 2001 From: wucm667 Date: Fri, 17 Jul 2026 22:10:18 +0800 Subject: [PATCH] fix: classify transient account exhaustion as 503 instead of 404 --- backend/internal/handler/no_account_error.go | 11 +-- .../internal/handler/no_account_error_test.go | 14 ++++ backend/internal/repository/account_repo.go | 65 +++++++++++++--- .../account_repo_model_availability_test.go | 58 ++++++++++++++ backend/internal/server/api_contract_test.go | 4 + backend/internal/service/account_service.go | 7 ++ .../service/account_service_delete_test.go | 4 + .../service/gateway_model_availability.go | 56 +++++++++---- .../gateway_model_availability_test.go | 78 ++++++++++++++++++- .../service/gateway_multiplatform_test.go | 28 +++++++ .../service/gemini_multiplatform_test.go | 3 + .../openai_gateway_model_availability.go | 28 +++++-- .../service/ratelimit_session_window_test.go | 3 + .../scheduler_snapshot_batch_query_test.go | 4 + ...er_snapshot_full_rebuild_lifecycle_test.go | 4 + 15 files changed, 331 insertions(+), 36 deletions(-) create mode 100644 backend/internal/repository/account_repo_model_availability_test.go diff --git a/backend/internal/handler/no_account_error.go b/backend/internal/handler/no_account_error.go index 001cef611d..842935aa2d 100644 --- a/backend/internal/handler/no_account_error.go +++ b/backend/internal/handler/no_account_error.go @@ -38,11 +38,12 @@ type noAccountErrorClassification struct { // The classifier intentionally does not consume the original error: the // selection layer never tells us *why* the pool came up empty (rate-limited // vs. unsupported model are both wrapped as ErrNoAvailableAccounts). Instead -// we re-check pool composition through DiagnoseModelAvailabilityForPlatform, -// which only inspects model_mapping configuration and ignores transient -// state. That guarantees a 404 is only returned when no operator action -// short of editing the account's model_mapping could make this request -// succeed. +// we re-check pool composition through DiagnoseModelAvailabilityForPlatform. +// Its dedicated database query considers only persistent eligibility +// (active status + schedulable setting) and model_mapping, bypassing scheduler +// snapshots and transient filters. That guarantees a 404 is only returned +// when persistent account/group/model configuration must change before the +// request can succeed. // // routingModel is the model name that account selection actually compared // against (i.e. after group-level dispatch mapping). displayModel is the diff --git a/backend/internal/handler/no_account_error_test.go b/backend/internal/handler/no_account_error_test.go index 174da82cc7..8c6784e84d 100644 --- a/backend/internal/handler/no_account_error_test.go +++ b/backend/internal/handler/no_account_error_test.go @@ -154,6 +154,20 @@ func TestClassifyNoAccountError_HasModelSupport_KeepsRoutingMessageGenerationToC require.False(t, cls.ModelNotFound) } +func TestClassifyNoAccountError_ModelSupportedOnlyByRateLimitedAccount_Returns503(t *testing.T) { + c := newTestGinContextWithRequest() + // The diagnoser's configured-state lookup still sees the model-supporting + // account even though normal scheduling has excluded it during cooldown. + fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}} + apiKey := &service.APIKey{GroupID: ptrInt64(7)} + + cls := classifyNoAccountErrorFromGin(c, fd, apiKey, "claude-opus-4-8", "claude-opus-4-8", service.PlatformAnthropic) + + require.Equal(t, http.StatusServiceUnavailable, cls.Status) + require.Equal(t, "api_error", cls.ErrType) + require.False(t, cls.ModelNotFound, "temporary account cooldown must remain retryable") +} + func TestClassifyNoAccountError_NoAccountsInPool_Stays503(t *testing.T) { c := newTestGinContextWithRequest() fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: false, HasModelSupport: false}} diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index f014155bb2..940083d6e4 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -1877,6 +1877,46 @@ func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Con }) } +// ListModelAvailabilityCandidates returns the persistently configured account +// pool used to decide whether a model is supported. Unlike scheduling queries, +// it intentionally ignores transient runtime state (rate limits, overload, +// temporary unschedulability, and expiry windows). +func (r *accountRepository) ListModelAvailabilityCandidates( + ctx context.Context, + groupID *int64, + platforms []string, + includeGrouped bool, +) ([]service.Account, error) { + if len(platforms) == 0 { + return []service.Account{}, nil + } + if groupID != nil { + return r.queryAccountsByGroup(ctx, *groupID, accountGroupQueryOptions{ + status: service.StatusActive, + schedulable: true, + ignoreTransientState: true, + platforms: platforms, + }) + } + + preds := []dbpredicate.Account{ + dbaccount.StatusEQ(service.StatusActive), + dbaccount.SchedulableEQ(true), + dbaccount.PlatformIn(platforms...), + } + if !includeGrouped { + preds = append(preds, dbaccount.Not(dbaccount.HasAccountGroups())) + } + accounts, err := r.client.Account.Query(). + Where(preds...). + Order(dbent.Asc(dbaccount.FieldPriority)). + All(ctx) + if err != nil { + return nil, err + } + return r.accountsToService(ctx, accounts) +} + func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { now := time.Now() _, err := r.client.Account.Update(). @@ -2662,9 +2702,10 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates } type accountGroupQueryOptions struct { - status string - schedulable bool - platforms []string // 允许的多个平台，空切片表示不进行平台过滤 + status string + schedulable bool + ignoreTransientState bool + platforms []string // 允许的多个平台，空切片表示不进行平台过滤 } func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID int64, opts accountGroupQueryOptions) ([]service.Account, error) { @@ -2681,14 +2722,16 @@ func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID in preds = append(preds, dbaccount.PlatformIn(opts.platforms...)) } if opts.schedulable { - now := time.Now() - preds = append(preds, - dbaccount.SchedulableEQ(true), - tempUnschedulablePredicate(), - notExpiredPredicate(now), - dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)), - dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)), - ) + preds = append(preds, dbaccount.SchedulableEQ(true)) + if !opts.ignoreTransientState { + now := time.Now() + preds = append(preds, + tempUnschedulablePredicate(), + notExpiredPredicate(now), + dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)), + dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)), + ) + } } if len(preds) > 0 { diff --git a/backend/internal/repository/account_repo_model_availability_test.go b/backend/internal/repository/account_repo_model_availability_test.go new file mode 100644 index 0000000000..917e73b0c4 --- /dev/null +++ b/backend/internal/repository/account_repo_model_availability_test.go @@ -0,0 +1,58 @@ +package repository + +import ( + "context" + "strings" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + dbent "github.com/Wei-Shaw/sub2api/ent" + _ "github.com/Wei-Shaw/sub2api/ent/runtime" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" + + "entgo.io/ent/dialect" + entsql "entgo.io/ent/dialect/sql" +) + +func TestListModelAvailabilityCandidates_GroupQueryIgnoresTransientState(t *testing.T) { + var capturedSQL string + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(captureEntQueryMatcher{actual: &capturedSQL})) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + driver := entsql.OpenDB(dialect.Postgres, db) + client := dbent.NewClient(dbent.Driver(driver)) + t.Cleanup(func() { _ = client.Close() }) + repo := newAccountRepositoryWithSQL(client, db, nil) + + mock.ExpectQuery("model availability candidates"). + WillReturnRows(sqlmock.NewRows([]string{"id"})) + groupID := int64(42) + accounts, err := repo.ListModelAvailabilityCandidates( + context.Background(), + &groupID, + []string{service.PlatformAnthropic}, + false, + ) + require.NoError(t, err) + require.Empty(t, accounts) + require.NoError(t, mock.ExpectationsWereMet()) + + normalized := normalizeSQLWhitespace(capturedSQL) + _, whereClause, found := strings.Cut(normalized, " WHERE ") + require.True(t, found, "expected WHERE clause in query: %s", normalized) + whereClause, _, _ = strings.Cut(whereClause, " ORDER BY ") + for _, configuredPredicate := range []string{"group_id", "status", "schedulable", "platform"} { + require.Contains(t, whereClause, configuredPredicate) + } + for _, transientPredicate := range []string{ + "rate_limit_reset_at", + "overload_until", + "temp_unschedulable_until", + "expires_at", + "auto_pause_on_expired", + } { + require.NotContains(t, whereClause, transientPredicate, "configured-state diagnosis must not filter transient predicate %q", transientPredicate) + } +} diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index 955e9df1be..aa170d4b62 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -1874,6 +1874,10 @@ func (s *stubAccountRepo) ListSchedulableUngroupedByPlatforms(ctx context.Contex return nil, errors.New("not implemented") } +func (s *stubAccountRepo) ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]service.Account, error) { + return nil, errors.New("not implemented") +} + func (s *stubAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { return errors.New("not implemented") } diff --git a/backend/internal/service/account_service.go b/backend/internal/service/account_service.go index b7fe0a68cc..283425c230 100644 --- a/backend/internal/service/account_service.go +++ b/backend/internal/service/account_service.go @@ -91,6 +91,13 @@ type AccountRepository interface { ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) + // ListModelAvailabilityCandidates returns accounts that are enabled by + // persistent configuration (active + schedulable) for model-support + // diagnosis. It deliberately does not filter transient runtime state such + // as rate-limit, overload, temporary-unschedulable, or expiry windows. + // When groupID is nil, includeGrouped controls whether the query scans all + // matching accounts or only accounts without a group binding. + ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error diff --git a/backend/internal/service/account_service_delete_test.go b/backend/internal/service/account_service_delete_test.go index ee6163239e..8d93e6a932 100644 --- a/backend/internal/service/account_service_delete_test.go +++ b/backend/internal/service/account_service_delete_test.go @@ -159,6 +159,10 @@ func (s *accountRepoStub) ListSchedulableUngroupedByPlatforms(ctx context.Contex panic("unexpected ListSchedulableUngroupedByPlatforms call") } +func (s *accountRepoStub) ListModelAvailabilityCandidates(ctx context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error) { + panic("unexpected ListModelAvailabilityCandidates call") +} + func (s *accountRepoStub) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { panic("unexpected SetRateLimited call") } diff --git a/backend/internal/service/gateway_model_availability.go b/backend/internal/service/gateway_model_availability.go index f34a591658..bc78a7859c 100644 --- a/backend/internal/service/gateway_model_availability.go +++ b/backend/internal/service/gateway_model_availability.go @@ -3,17 +3,20 @@ package service import ( "context" "strings" + + "github.com/Wei-Shaw/sub2api/internal/config" ) // ModelAvailabilityDiagnosis describes whether the requested model can be -// served by any configured account in the group, ignoring transient state -// (rate limits, quota auto-pause, runtime blocks). Handlers use this on the -// "no available accounts" error path to distinguish 404 model_not_found from -// 503 service_unavailable. +// served by any persistently eligible account in the group (active with its +// schedulable setting enabled), ignoring transient state such as rate limits, +// overload, temporary unschedulability, and runtime blocks. Handlers use this +// on the "no available accounts" error path to distinguish 404 +// model_not_found from 503 service_unavailable. type ModelAvailabilityDiagnosis struct { - // HasAccountsInPool is true if the group has at least one schedulable - // account on the queried platform (or, for Anthropic/Gemini, on the - // platform plus mixed-scheduled Antigravity accounts). + // HasAccountsInPool is true if the group has at least one persistently + // eligible account on the queried platform (or, for Anthropic/Gemini, on + // the platform plus mixed-scheduled Antigravity accounts). HasAccountsInPool bool // HasModelSupport is true if at least one account's model mapping admits // the requested model. @@ -33,10 +36,11 @@ type ModelAvailabilityDiagnoser interface { ) ModelAvailabilityDiagnosis } -// DiagnoseModelAvailabilityForPlatform inspects schedulable accounts of the -// given platform and returns whether the requested model is configured to be -// served by any of them. It deliberately ignores schedulability, rate limits, -// quotas, and runtime blocks — those are transient. +// DiagnoseModelAvailabilityForPlatform inspects accounts enabled for scheduling +// by persistent configuration and returns whether the requested model is +// configured to be served by any of them. The dedicated repository query +// bypasses scheduler snapshots and deliberately ignores transient rate-limit, +// overload, temporary-unschedulable, expiry, quota, and runtime-block state. // // Safe to call on the error path: returns {true,true} on any internal failure // or when the inputs preclude meaningful diagnosis (empty model, etc.), so @@ -61,9 +65,30 @@ func (s *GatewayService) DiagnoseModelAvailabilityForPlatform( return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} } - // hasForcePlatform=false so Anthropic/Gemini also surface mixed-scheduled - // Antigravity accounts, matching what selection would consider. - accounts, _, err := s.listSchedulableAccounts(ctx, groupID, platform, false) + if s.accountRepo == nil { + return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} + } + + useMixed := platform == PlatformAnthropic || platform == PlatformGemini + platforms := []string{platform} + if useMixed { + platforms = append(platforms, PlatformAntigravity) + } + + queryGroupID := groupID + includeGrouped := false + if useMixed { + // Preserve the generic scheduler's scope rules: an explicit group wins + // for mixed scheduling, while group-less simple mode scans all accounts. + if groupID == nil && s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + includeGrouped = true + } + } else if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + queryGroupID = nil + includeGrouped = true + } + + accounts, err := s.accountRepo.ListModelAvailabilityCandidates(ctx, queryGroupID, platforms, includeGrouped) if err != nil { // Conservative fallback: pretend everything is fine so the caller // returns 503 (we don't want to flip to 404 just because a lookup @@ -73,6 +98,9 @@ func (s *GatewayService) DiagnoseModelAvailabilityForPlatform( diag := ModelAvailabilityDiagnosis{} for i := range accounts { + if useMixed && accounts[i].Platform == PlatformAntigravity && !accounts[i].IsMixedSchedulingEnabled() { + continue + } diag.HasAccountsInPool = true if s.isModelSupportedByAccountWithContext(ctx, &accounts[i], requestedModel) { diag.HasModelSupport = true diff --git a/backend/internal/service/gateway_model_availability_test.go b/backend/internal/service/gateway_model_availability_test.go index bcca0e5e0b..ed32f620e8 100644 --- a/backend/internal/service/gateway_model_availability_test.go +++ b/backend/internal/service/gateway_model_availability_test.go @@ -5,6 +5,7 @@ package service import ( "context" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -117,6 +118,7 @@ func TestDiagnoseModelAvailabilityForPlatform_WildcardMappingMatches(t *testing. } func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSignal(t *testing.T) { + groupID := int64(42) repo := &mockAccountRepoForPlatform{ accounts: []Account{ { @@ -124,6 +126,9 @@ func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSig Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, + AccountGroups: []AccountGroup{ + {GroupID: groupID}, + }, Credentials: map[string]any{"model_mapping": map[string]any{"gpt-5": "gpt-5"}}, }, { @@ -131,6 +136,9 @@ func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSig Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true, + AccountGroups: []AccountGroup{ + {GroupID: groupID}, + }, Credentials: map[string]any{"model_mapping": map[string]any{"gpt-5-mini": "gpt-5-mini"}}, }, }, @@ -141,12 +149,80 @@ func TestDiagnoseModelAvailabilityForPlatform_NoMatchingModel_ReturnsNotFoundSig } svc := &GatewayService{accountRepo: repo, cfg: testConfig()} - diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), nil, "gpt-5.1-codex-mini", PlatformOpenAI) + diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "gpt-5.1-codex-mini", PlatformOpenAI) require.True(t, diag.HasAccountsInPool, "group has OpenAI accounts") require.False(t, diag.HasModelSupport, "no account mapping admits the requested model — handler should return 404") } +func TestDiagnoseModelAvailabilityForPlatform_RateLimitedSupportingAccountRemainsConfigured(t *testing.T) { + groupID := int64(42) + cooldownUntil := time.Now().Add(time.Hour) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 1, + Platform: PlatformAnthropic, + Status: StatusActive, + Schedulable: true, + RateLimitResetAt: &cooldownUntil, + OverloadUntil: &cooldownUntil, + TempUnschedulableUntil: &cooldownUntil, + AccountGroups: []AccountGroup{{GroupID: groupID}}, + Credentials: map[string]any{ + "model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"}, + }, + }, + }, + accountsByID: map[int64]*Account{}, + } + require.False(t, repo.accounts[0].IsSchedulable(), "test account must be excluded from normal scheduling while cooling down") + svc := &GatewayService{ + accountRepo: repo, + cfg: testConfig(), + schedulerSnapshot: &SchedulerSnapshotService{}, // diagnosis must bypass the transient-only snapshot + } + + diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "claude-opus-4-8", PlatformAnthropic) + + require.True(t, diag.HasAccountsInPool) + require.True(t, diag.HasModelSupport, "a configured model remains supported while every matching account is temporarily cooling down") +} + +func TestOpenAIDiagnoseModelAvailabilityForPlatform_RateLimitedSupportingAccountRemainsConfigured(t *testing.T) { + groupID := int64(43) + cooldownUntil := time.Now().Add(time.Hour) + repo := &mockAccountRepoForPlatform{ + accounts: []Account{ + { + ID: 2, + Platform: PlatformOpenAI, + Status: StatusActive, + Schedulable: true, + RateLimitResetAt: &cooldownUntil, + OverloadUntil: &cooldownUntil, + TempUnschedulableUntil: &cooldownUntil, + AccountGroups: []AccountGroup{{GroupID: groupID}}, + Credentials: map[string]any{ + "model_mapping": map[string]any{"claude-opus-4-8": "claude-opus-4-8"}, + }, + }, + }, + accountsByID: map[int64]*Account{}, + } + require.False(t, repo.accounts[0].IsSchedulable(), "test account must be excluded from normal scheduling while cooling down") + svc := &OpenAIGatewayService{ + accountRepo: repo, + cfg: testConfig(), + schedulerSnapshot: &SchedulerSnapshotService{}, // diagnosis must bypass the transient-only snapshot + } + + diag := svc.DiagnoseModelAvailabilityForPlatform(context.Background(), &groupID, "claude-opus-4-8", PlatformOpenAI) + + require.True(t, diag.HasAccountsInPool) + require.True(t, diag.HasModelSupport, "OpenAI-compatible diagnosis must keep transiently limited supporting accounts in the configured pool") +} + func TestDiagnoseModelAvailabilityForPlatform_WrongPlatformFiltersOut(t *testing.T) { // Group has only Anthropic accounts; user routes to OpenAI gateway. // Diagnosis must NOT see Anthropic accounts (listSchedulableAccounts filters diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index 35a60c8124..7b5d967cf5 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -156,6 +156,34 @@ func (m *mockAccountRepoForPlatform) ListSchedulableUngroupedByPlatform(ctx cont func (m *mockAccountRepoForPlatform) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) { return m.ListSchedulableByPlatforms(ctx, platforms) } +func (m *mockAccountRepoForPlatform) ListModelAvailabilityCandidates(_ context.Context, groupID *int64, platforms []string, includeGrouped bool) ([]Account, error) { + platformSet := make(map[string]struct{}, len(platforms)) + for _, platform := range platforms { + platformSet[platform] = struct{}{} + } + result := make([]Account, 0, len(m.accounts)) + for _, acc := range m.accounts { + if _, ok := platformSet[acc.Platform]; !ok || acc.Status != StatusActive || !acc.Schedulable { + continue + } + if groupID != nil { + inGroup := false + for _, accountGroup := range acc.AccountGroups { + if accountGroup.GroupID == *groupID { + inGroup = true + break + } + } + if !inGroup { + continue + } + } else if !includeGrouped && (len(acc.AccountGroups) > 0 || len(acc.GroupIDs) > 0) { + continue + } + result = append(result, acc) + } + return result, nil +} func (m *mockAccountRepoForPlatform) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { return nil } diff --git a/backend/internal/service/gemini_multiplatform_test.go b/backend/internal/service/gemini_multiplatform_test.go index 7d5ed0ec9e..e4c29852e2 100644 --- a/backend/internal/service/gemini_multiplatform_test.go +++ b/backend/internal/service/gemini_multiplatform_test.go @@ -147,6 +147,9 @@ func (m *mockAccountRepoForGemini) ListSchedulableUngroupedByPlatform(ctx contex func (m *mockAccountRepoForGemini) ListSchedulableUngroupedByPlatforms(ctx context.Context, platforms []string) ([]Account, error) { return m.ListSchedulableByPlatforms(ctx, platforms) } +func (m *mockAccountRepoForGemini) ListModelAvailabilityCandidates(ctx context.Context, _ *int64, platforms []string, _ bool) ([]Account, error) { + return m.ListSchedulableByPlatforms(ctx, platforms) +} func (m *mockAccountRepoForGemini) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { return nil } diff --git a/backend/internal/service/openai_gateway_model_availability.go b/backend/internal/service/openai_gateway_model_availability.go index 1bbc3ee7ca..a665052e60 100644 --- a/backend/internal/service/openai_gateway_model_availability.go +++ b/backend/internal/service/openai_gateway_model_availability.go @@ -3,13 +3,16 @@ package service import ( "context" "strings" + + "github.com/Wei-Shaw/sub2api/internal/config" ) // DiagnoseModelAvailabilityForPlatform reports whether the requested model -// is configured to be served by any OpenAI-compatible account in the group -// for the given platform (e.g. PlatformOpenAI, PlatformGrok). The platform -// scopes the candidate pool so distinct OpenAI-compatible platforms do not -// cross-contaminate diagnosis results. +// is configured to be served by any persistently eligible OpenAI-compatible +// account in the group for the given platform (e.g. PlatformOpenAI, +// PlatformGrok). The platform scopes the candidate pool so distinct +// OpenAI-compatible platforms do not cross-contaminate diagnosis results. +// The query bypasses scheduler snapshots and ignores transient runtime state. // // Safe to call on the error path: returns {true,true} on any internal // failure or when the inputs preclude meaningful diagnosis (empty model, @@ -27,8 +30,23 @@ func (s *OpenAIGatewayService) DiagnoseModelAvailabilityForPlatform( if requestedModel == "" { return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} } + if s.accountRepo == nil { + return ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true} + } - accounts, err := s.listSchedulableAccounts(ctx, groupID, platform) + platform = normalizeOpenAICompatiblePlatform(platform) + queryGroupID := groupID + includeGrouped := false + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + queryGroupID = nil + includeGrouped = true + } + accounts, err := s.accountRepo.ListModelAvailabilityCandidates( + ctx, + queryGroupID, + []string{platform}, + includeGrouped, + ) if err != nil { // Conservative fallback so the caller keeps returning 503; we do not // want a transient lookup failure to flip into 404 model_not_found. diff --git a/backend/internal/service/ratelimit_session_window_test.go b/backend/internal/service/ratelimit_session_window_test.go index 279a31ccdc..b484121143 100644 --- a/backend/internal/service/ratelimit_session_window_test.go +++ b/backend/internal/service/ratelimit_session_window_test.go @@ -143,6 +143,9 @@ func (m *sessionWindowMockRepo) ListSchedulableUngroupedByPlatform(context.Conte func (m *sessionWindowMockRepo) ListSchedulableUngroupedByPlatforms(context.Context, []string) ([]Account, error) { panic("unexpected") } +func (m *sessionWindowMockRepo) ListModelAvailabilityCandidates(context.Context, *int64, []string, bool) ([]Account, error) { + panic("unexpected") +} func (m *sessionWindowMockRepo) SetRateLimited(context.Context, int64, time.Time) error { panic("unexpected") } diff --git a/backend/internal/service/scheduler_snapshot_batch_query_test.go b/backend/internal/service/scheduler_snapshot_batch_query_test.go index 221610bdbf..b957891d36 100644 --- a/backend/internal/service/scheduler_snapshot_batch_query_test.go +++ b/backend/internal/service/scheduler_snapshot_batch_query_test.go @@ -56,6 +56,10 @@ func (r *batchAccountQueryRepo) ListSchedulableUngroupedByPlatforms(_ context.Co return r.run(batchAccountQueryKey{platform: platforms[0], mixed: true}) } +func (r *batchAccountQueryRepo) ListModelAvailabilityCandidates(context.Context, *int64, []string, bool) ([]Account, error) { + panic("unexpected ListModelAvailabilityCandidates call") +} + func (r *batchAccountQueryRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) { return r.run(batchAccountQueryKey{platform: platform}) } diff --git a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go index 35ce99a8d1..fa918a0862 100644 --- a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go +++ b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go @@ -214,6 +214,10 @@ func (r *fullRebuildAccountRepo) ListSchedulableUngroupedByPlatforms(_ context.C return r.record(0, firstPlatform(platforms)) } +func (r *fullRebuildAccountRepo) ListModelAvailabilityCandidates(context.Context, *int64, []string, bool) ([]Account, error) { + panic("unexpected ListModelAvailabilityCandidates call") +} + func (r *fullRebuildAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) { return r.record(0, platform) }