fix: classify transient account exhaustion as 503 instead of 404

This commit is contained in:
wucm667
2026-07-17 22:10:18 +08:00
parent 57914967cb
commit 2594950993
15 changed files with 331 additions and 36 deletions
+6 -5
View File
@@ -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
@@ -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}}
+54 -11
View File
@@ -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 {
@@ -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)
}
}
@@ -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")
}
@@ -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
@@ -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")
}
@@ -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
@@ -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
@@ -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
}
@@ -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
}
@@ -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.
@@ -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")
}
@@ -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})
}
@@ -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)
}