diff --git a/CHANGELOG.md b/CHANGELOG.md index 1e8ac8a..7352d89 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,10 +8,12 @@ Write each change in both `### English` and `### 中文` under `## Unreleased`. ### English - Recommend Docker Compose as the supported install and managed-update path in the README and deployment guide +- Keep an account on a model route after a later catalog refresh omits an ID it already served, and report pool-wide quota cooldown as `insufficient_quota` instead of a generic rate limit ### 中文 - 在 README 与部署说明中明确推荐 Docker Compose 作为官方安装与托管更新路径 +- 账号已成功服务过的模型，在后续目录刷新漏掉该 ID 时仍可继续路由；全池额度冷却改为返回 `insufficient_quota`，而不是笼统的限流错误 ## 0.2.45 - 2026-09-05 diff --git a/internal/accounts/pool.go b/internal/accounts/pool.go index 03b1f8f..4d1fe4a 100644 --- a/internal/accounts/pool.go +++ b/internal/accounts/pool.go @@ -71,6 +71,11 @@ type Item struct { // spelling so Trae config_name case is preserved. Models []string ModelsAt time.Time + // ProvenModels are public IDs this account has actually served. A + // later catalog refresh must not drop a model that just succeeded: + // WorkBuddy CLI snapshots can omit a live ID and would otherwise + // leave only quota-cooled empty-catalog accounts on that route. + ProvenModels []string // ModelDownUntil is per-model cooldown. One model hitting a limit must // not take the whole account offline for other models. ModelDownUntil map[string]time.Time @@ -138,10 +143,9 @@ func routeModel(model string) string { return id } -func itemHasModel(item Item, publicModel string) bool { - want := routeModel(publicModel) +func itemHasCatalogModel(item Item, want string) bool { if want == "" || item.Models == nil { - return true + return false } for _, model := range item.Models { if CanonicalModelID(model) == want { @@ -151,6 +155,82 @@ func itemHasModel(item Item, publicModel string) bool { return false } +func itemHasProvenModel(item Item, want string) bool { + if want == "" { + return false + } + for _, model := range item.ProvenModels { + if CanonicalModelID(model) == want { + return true + } + } + return false +} + +func rememberProvenModel(item *Item, model string) { + if item == nil { + return + } + want := routeModel(model) + if want == "" || itemHasProvenModel(*item, want) { + return + } + native := NativeModelID(*item, model) + if strings.TrimSpace(native) == "" { + native = model + } + item.ProvenModels = append(item.ProvenModels, native) +} + +func dropProvenModel(item *Item, want string) { + if item == nil || want == "" || len(item.ProvenModels) == 0 { + return + } + next := item.ProvenModels[:0] + for _, model := range item.ProvenModels { + if CanonicalModelID(model) != want { + next = append(next, model) + } + } + if len(next) == 0 { + item.ProvenModels = nil + return + } + item.ProvenModels = next +} + +// itemCouldServeModel reports whether this account belongs on a model route +// at all, including unknown-catalog accounts that are currently cooling. +// Cooling empty-catalog accounts must stay visible as retry hints so a +// quota pool is not reported as model_not_available. +func itemCouldServeModel(item Item, publicModel string) bool { + want := routeModel(publicModel) + if want == "" { + return true + } + if itemHasCatalogModel(item, want) || itemHasProvenModel(item, want) { + return true + } + return item.Models == nil +} + +func itemHasModel(item Item, publicModel string) bool { + want := routeModel(publicModel) + if want == "" { + return true + } + if itemHasCatalogModel(item, want) || itemHasProvenModel(item, want) { + return true + } + if item.Models == nil { + // Unknown catalog fail-open only for accounts that can send now. + // A quota-cooled account must not occupy this model just because + // its catalog fetch failed. + return !itemDown(item, time.Now()) + } + return false +} + // NativeModelID returns the provider-native catalog spelling for a public // model. Trae config_name is case-sensitive; routing matches on the // canonical form, but the upstream request must keep the original ID. @@ -164,6 +244,11 @@ func NativeModelID(item Item, publicModel string) string { return model } } + for _, model := range item.ProvenModels { + if CanonicalModelID(model) == want { + return model + } + } return strings.TrimSpace(publicModel) } @@ -205,7 +290,7 @@ func itemReady(item Item) bool { return item.Ready == nil || *item.Ready } -func routeMatches(item Item, q RouteQuery) bool { +func routeBaseMatches(item Item, q RouteQuery) bool { if !itemReady(item) { return false } @@ -221,12 +306,17 @@ func routeMatches(item Item, q RouteQuery) bool { if q.RegionFilter != "" && itemRegion(item) != strings.ToLower(strings.TrimSpace(q.RegionFilter)) { return false } - if !itemHasModel(item, q.PublicModel) { - return false - } return true } +func routeMatches(item Item, q RouteQuery) bool { + return routeBaseMatches(item, q) && itemHasModel(item, q.PublicModel) +} + +func routeHintMatches(item Item, q RouteQuery) bool { + return routeBaseMatches(item, q) && itemCouldServeModel(item, q.PublicModel) +} + type Pool struct { mu sync.Mutex items []Item @@ -401,7 +491,7 @@ func (p *Pool) PickRoute(q RouteQuery) (Item, bool) { } if pinned != nil { if _, skip := q.Excluded[pinned.ID]; !skip { - if routeModel(q.PublicModel) != "" && pinned.Models != nil && !itemHasModel(*pinned, q.PublicModel) { + if routeModel(q.PublicModel) != "" && !itemCouldServeModel(*pinned, q.PublicModel) { return Item{}, false } if routeMatches(*pinned, q) && !itemDown(*pinned, now) && @@ -433,8 +523,11 @@ func (p *Pool) PickRoute(q RouteQuery) (Item, bool) { // the caller can report a classified retry-after error. var best Item found := false - for _, i := range eligible { + for i := range p.items { item := p.items[i] + if !routeHintMatches(item, q) { + continue + } if !itemDown(item, now) && !itemModelDown(item, q, now) { // Saturated only — skip; not a candidate for the retry hint. continue @@ -844,6 +937,7 @@ func (p *Pool) MarkOK(id, model string) { p.items[i].LastKind = "" p.items[i].BackoffLevel = 0 } + rememberProvenModel(&p.items[i], model) } else { // Account-wide success: the account proved it can serve traffic. p.items[i].DownUntil = time.Time{} @@ -995,16 +1089,19 @@ func (p *Pool) RemoveModel(id, model string) { p.mu.Lock() defer p.mu.Unlock() for i := range p.items { - if p.items[i].ID != id || p.items[i].Models == nil { + if p.items[i].ID != id { continue } - next := make([]string, 0, len(p.items[i].Models)) - for _, existing := range p.items[i].Models { - if CanonicalModelID(existing) != want { - next = append(next, existing) + if p.items[i].Models != nil { + next := make([]string, 0, len(p.items[i].Models)) + for _, existing := range p.items[i].Models { + if CanonicalModelID(existing) != want { + next = append(next, existing) + } } + p.items[i].Models = next } - p.items[i].Models = next + dropProvenModel(&p.items[i], want) return } } @@ -1100,6 +1197,9 @@ func (item Item) clone() Item { if item.Models != nil { out.Models = append([]string(nil), item.Models...) } + if item.ProvenModels != nil { + out.ProvenModels = append([]string(nil), item.ProvenModels...) + } if item.Quota != nil { q := *item.Quota out.Quota = &q @@ -1192,6 +1292,9 @@ func (p *Pool) Upsert(item Item) { item.Models = p.items[i].Models item.ModelsAt = p.items[i].ModelsAt } + if item.ProvenModels == nil { + item.ProvenModels = p.items[i].ProvenModels + } p.items[i] = item return } diff --git a/internal/accounts/pool_test.go b/internal/accounts/pool_test.go index 8f6d395..b0db3c8 100644 --- a/internal/accounts/pool_test.go +++ b/internal/accounts/pool_test.go @@ -520,3 +520,37 @@ func TestAccountBackoffIgnoresModelForAccountFailures(t *testing.T) { t.Fatalf("account auth must not create a model kind: %+v", item.ModelLastKind) } } + +func TestQuotaCooledEmptyCatalogDoesNotBlockProvenAccount(t *testing.T) { + p := NewPool(nil, nil) + p.Upsert(Item{ID: "quota", Provider: "workbuddy", Region: "cn"}) + p.Upsert(Item{ID: "ready", Provider: "workbuddy", Region: "cn"}) + p.MarkClassified("quota", Classified{Kind: KindQuota, Cooldown: time.Hour, Message: "额度已用尽"}) + p.MarkOK("ready", "deepseek-v4-flash") + p.MergeModels("ready", []string{"glm-5.2"}) + + item, ok := p.PickRoute(RouteQuery{PublicModel: "deepseek-v4-flash", ProviderFilter: "workbuddy"}) + if !ok || item.ID != "ready" { + t.Fatalf("proven ready account must win over quota-cooled empty catalog, got %+v ok=%v", item, ok) + } + if retry := p.RetryAfter(item, "deepseek-v4-flash"); retry > 0 { + t.Fatalf("ready account must be dispatchable, retry-after=%v", retry) + } +} + +func TestQuotaCooledEmptyCatalogSurfacesQuotaHint(t *testing.T) { + p := NewPool(nil, nil) + p.Upsert(Item{ID: "quota", Provider: "workbuddy", Region: "cn"}) + p.MarkClassified("quota", Classified{Kind: KindQuota, Cooldown: time.Hour, Message: "额度已用尽"}) + + item, ok := p.PickRoute(RouteQuery{PublicModel: "deepseek-v4-flash", ProviderFilter: "workbuddy"}) + if !ok || item.ID != "quota" { + t.Fatalf("quota-cooled empty catalog must surface as a retry hint, got %+v ok=%v", item, ok) + } + if retry := p.RetryAfter(item, "deepseek-v4-flash"); retry <= 0 { + t.Fatal("quota hint must carry retry-after") + } + if p.LenRoute(RouteQuery{PublicModel: "deepseek-v4-flash", ProviderFilter: "workbuddy"}) != 0 { + t.Fatal("quota-cooled empty catalog must not count as a live candidate") + } +} diff --git a/internal/executor/chat.go b/internal/executor/chat.go index 757ce4a..bd17467 100644 --- a/internal/executor/chat.go +++ b/internal/executor/chat.go @@ -342,16 +342,7 @@ func (e ChatExecutor) pick(prefer, providerFilter, regionFilter, publicModel str if e.Pool != nil { if item, ok := e.Pool.PickRoute(query); ok { if retryAfter := e.Pool.RetryAfter(item, publicModel); retryAfter > 0 { - failover := true - message := "all accounts are cooling down" - if e.Pool.CooldownScope(item, publicModel) == "model" && strings.TrimSpace(publicModel) != "" { - message = fmt.Sprintf("model %s is cooling down on all available accounts", publicModel) - } - return accounts.Item{}, &providers.Error{ - Kind: accounts.KindRateLimit, Status: 429, Code: "rate_limit", - Type: "api_error", Message: message, Cooldown: retryAfter, - RetryAfter: retryAfter, Failover: &failover, - } + return accounts.Item{}, coolingPickError(item, publicModel, retryAfter) } return item, nil } @@ -393,6 +384,33 @@ func (e ChatExecutor) pick(prefer, providerFilter, regionFilter, publicModel str return accounts.Item{}, fmt.Errorf("no worker accounts configured") } +func coolingPickError(item accounts.Item, publicModel string, retryAfter time.Duration) error { + failover := true + kind := accounts.KindRateLimit + code := "rate_limit" + typ := "api_error" + message := "all accounts are cooling down" + if !item.DownUntil.IsZero() && time.Now().Before(item.DownUntil) && item.LastKind == accounts.KindQuota { + kind = accounts.KindQuota + code = "insufficient_quota" + typ = "insufficient_quota" + failover = false + if strings.TrimSpace(publicModel) != "" { + message = fmt.Sprintf("all accounts that can serve %s are on quota cooldown", publicModel) + } else { + message = "all accounts are on quota cooldown" + } + } else if strings.TrimSpace(publicModel) != "" { + if until, ok := item.ModelDownUntil[accounts.CanonicalModelID(publicModel)]; ok && !until.IsZero() { + message = fmt.Sprintf("model %s is cooling down on all available accounts", publicModel) + } + } + return &providers.Error{ + Kind: kind, Status: 429, Code: code, Type: typ, Message: message, + Cooldown: retryAfter, RetryAfter: retryAfter, Failover: &failover, + } +} + func (e ChatExecutor) attemptsFor(providerFilter, regionFilter, publicModel string, allowed []string) int { maxAttempts := e.MaxAttempts if maxAttempts <= 0 { @@ -457,6 +475,13 @@ func (e ChatExecutor) ObserveStreamFailure(accountID string, err error, model st // The request body was rejected; the account itself is healthy. return } + if classified.Kind == accounts.KindModelNotAvailable { + // Streaming 200 headers already recorded this model as proven. + // MarkClassified ignores catalog misses, so drop the stale ID + // here or the next pick will keep sending it. + e.Pool.RemoveModel(accountID, model) + return + } if classified.Kind == accounts.KindQuota { // Quota is classified with Failover=false because the account is not // at fault on a normal request path. Here the response already went diff --git a/internal/executor/chat_test.go b/internal/executor/chat_test.go index 68beb49..03adb98 100644 --- a/internal/executor/chat_test.go +++ b/internal/executor/chat_test.go @@ -518,6 +518,112 @@ func TestChatNonStreamDoesNotDispatchWhileAccountCooling(t *testing.T) { } } +func TestChatNonStreamUsesProvenAccountOverQuotaCatalog(t *testing.T) { + var hits atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + _ = json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{{"finish_reason": "stop", "message": map[string]any{"content": "ok"}}}, + "usage": map[string]any{"source": "upstream"}, + }) + })) + defer srv.Close() + + pool := accounts.NewPool(nil, nil) + pool.Upsert(accounts.Item{ID: "quota", URL: "http://127.0.0.1:1", Provider: "workbuddy", Region: "cn", Runtime: "child_process"}) + pool.Upsert(accounts.Item{ID: "ready", URL: srv.URL, Provider: "workbuddy", Region: "cn", Runtime: "child_process"}) + pool.MarkClassified("quota", accounts.Classified{Kind: accounts.KindQuota, Cooldown: time.Hour, Message: "额度已用尽"}) + pool.MarkOK("ready", "deepseek-v4-flash") + pool.MergeModels("ready", []string{"glm-5.2"}) + + ex := NewChatExecutor(pool, "") + ex.HTTPClient = srv.Client() + got, err := ex.ChatNonStream(context.Background(), translate.ChatRequest{ + Model: "deepseek-v4-flash", + Messages: []translate.ChatMessage{{Role: "user", Content: "hi"}}, + }, "", "workbuddy") + if err != nil { + t.Fatal(err) + } + if got.AccountID != "ready" || hits.Load() != 1 { + t.Fatalf("got %+v hits=%d", got, hits.Load()) + } +} + +func TestChatNonStreamExpiredQuotaDoesNotMaskModelRateLimit(t *testing.T) { + pool := accounts.NewPool(nil, nil) + pool.Upsert(accounts.Item{ + ID: "a", URL: "http://127.0.0.1:1", Provider: "workbuddy", Region: "cn", Runtime: "child_process", + DownUntil: time.Now().Add(-time.Minute), LastKind: accounts.KindQuota, + }) + pool.MarkClassified("a", accounts.Classified{ + Kind: accounts.KindRateLimit, Cooldown: time.Hour, Failover: true, + Model: "deepseek-v4-flash", Message: "too many requests", + }) + ex := NewChatExecutor(pool, "") + _, err := ex.ChatNonStream(context.Background(), translate.ChatRequest{ + Model: "deepseek-v4-flash", + Messages: []translate.ChatMessage{{Role: "user", Content: "hi"}}, + }, "", "workbuddy") + if err == nil { + t.Fatal("expected model cooling error") + } + var classified *providers.Error + if !errors.As(err, &classified) || classified.Kind != accounts.KindRateLimit || classified.Code != "rate_limit" { + t.Fatalf("expired quota must not mask a live model rate limit, got %#v", err) + } + if classified.Failover == nil || !*classified.Failover { + t.Fatal("model rate-limit cooling must remain failoverable") + } +} + +func TestChatNonStreamQuotaCoolingIsQuotaNotRateLimit(t *testing.T) { + pool := accounts.NewPool(nil, nil) + pool.Upsert(accounts.Item{ID: "quota", URL: "http://127.0.0.1:1", Provider: "workbuddy", Region: "cn", Runtime: "child_process"}) + pool.MarkClassified("quota", accounts.Classified{Kind: accounts.KindQuota, Cooldown: time.Hour, Message: "额度已用尽"}) + ex := NewChatExecutor(pool, "") + _, err := ex.ChatNonStream(context.Background(), translate.ChatRequest{ + Model: "deepseek-v4-flash", + Messages: []translate.ChatMessage{{Role: "user", Content: "hi"}}, + }, "", "workbuddy") + if err == nil { + t.Fatal("expected quota cooling error") + } + var classified *providers.Error + if !errors.As(err, &classified) || classified.Kind != accounts.KindQuota || classified.Code != "insufficient_quota" { + t.Fatalf("expected quota cooling error, got %#v", err) + } + if !strings.Contains(classified.Message, "quota cooldown") { + t.Fatalf("quota cooling error should say quota, got %q", classified.Message) + } +} + +func TestObserveStreamFailureDropsProvenModel(t *testing.T) { + pool := accounts.NewPool(nil, nil) + pool.Upsert(accounts.Item{ID: "ready", Provider: "workbuddy", Region: "cn", Runtime: "child_process"}) + pool.MarkOK("ready", "deepseek-v4-flash") + pool.MergeModels("ready", []string{"glm-5.2"}) + item, ok := pool.ByID("ready") + if !ok || len(item.ProvenModels) != 1 { + t.Fatalf("expected proven model before stream failure, got %+v", item) + } + + ex := NewChatExecutor(pool, "") + ex.ObserveStreamFailure("ready", &providers.Error{ + Kind: accounts.KindModelNotAvailable, + Status: 400, + Message: "model not available", + }, "deepseek-v4-flash") + + item, _ = pool.ByID("ready") + if len(item.ProvenModels) != 0 { + t.Fatalf("stream catalog miss must drop proven model, got %v", item.ProvenModels) + } + if _, ok := pool.PickRoute(accounts.RouteQuery{PublicModel: "deepseek-v4-flash", ProviderFilter: "workbuddy"}); ok { + t.Fatal("account must leave the omitted-model route after a stream catalog miss") + } +} + func TestObserveStreamFailureWithoutModelTakesAccountDown(t *testing.T) { pool := accounts.NewPool([]string{"http://127.0.0.1:1"}, []string{"acc-quota"}) pool.Upsert(accounts.Item{ID: "acc-quota"})