From dab1360311d201ab64e46b98d844b44b43d5b56b Mon Sep 17 00:00:00 2001 From: Zexi Li Date: Mon, 27 Jul 2026 12:16:12 +0800 Subject: [PATCH] fix(aiproxy): avoid permanent ai_key blacklist and clarify skip reasons (#25225) Enter cooldown when health score hits zero so keys can recover, and include concrete skip reasons in resolve errors. --- pkg/aiproxy/models/ai_key_health.go | 36 +++- pkg/aiproxy/models/ai_key_health_test.go | 215 +++++++++++++++++++++++ pkg/aiproxy/models/ai_key_resolve.go | 72 +++++++- 3 files changed, 314 insertions(+), 9 deletions(-) create mode 100644 pkg/aiproxy/models/ai_key_health_test.go diff --git a/pkg/aiproxy/models/ai_key_health.go b/pkg/aiproxy/models/ai_key_health.go index ffd8728623..592227c431 100644 --- a/pkg/aiproxy/models/ai_key_health.go +++ b/pkg/aiproxy/models/ai_key_health.go @@ -33,6 +33,13 @@ type aiKeyHealthState struct { cooldownUntil time.Time } +// aiKeyHealthSnapshot is a read-only view for error messages (not exported as API). +type aiKeyHealthSnapshot struct { + score int + inCooldown bool + remainingSec int +} + var ( aiKeyHealthMu sync.RWMutex aiKeyHealth = map[string]*aiKeyHealthState{} @@ -57,6 +64,27 @@ func getAiKeyHealth(keyId string) *aiKeyHealthState { return st } +// aiKeyHealthInfo returns a read-only snapshot for diagnostics / error text. +func aiKeyHealthInfo(keyId string) aiKeyHealthSnapshot { + if keyId == "" { + return aiKeyHealthSnapshot{score: aiKeyHealthMaxScore} + } + st := getAiKeyHealth(keyId) + now := time.Now() + aiKeyHealthMu.RLock() + defer aiKeyHealthMu.RUnlock() + info := aiKeyHealthSnapshot{score: st.score} + if !st.cooldownUntil.IsZero() && now.Before(st.cooldownUntil) { + info.inCooldown = true + sec := int(st.cooldownUntil.Sub(now).Seconds()) + if sec < 1 { + sec = 1 + } + info.remainingSec = sec + } + return info +} + // dynamicAiKeyWeightMultiplier returns 0-100 applied to configured ai_key.weight (100 = full weight). func dynamicAiKeyWeightMultiplier(keyId string) int { if keyId == "" { @@ -75,7 +103,12 @@ func dynamicAiKeyWeightMultiplier(keyId string) int { st.score = aiKeyHealthMaxScore / 2 } } + // score<=0 without an active cooldown would permanently exclude the key; + // start a cooldown so it can recover via the path above after the period. if st.score <= 0 { + if st.cooldownUntil.IsZero() { + st.cooldownUntil = now.Add(aiKeyHealthCooldownPeriod) + } return 0 } if st.score > aiKeyHealthMaxScore { @@ -113,7 +146,8 @@ func RecordAiKeyFailure(keyId string, statusCode int) { if st.score < 0 { st.score = 0 } - if st.consecutiveFails >= aiKeyHealthCooldownAfter { + // Enter cooldown on streak or when score is exhausted (avoids permanent blacklist). + if st.consecutiveFails >= aiKeyHealthCooldownAfter || st.score <= 0 { st.cooldownUntil = time.Now().Add(aiKeyHealthCooldownPeriod) st.score = 0 } diff --git a/pkg/aiproxy/models/ai_key_health_test.go b/pkg/aiproxy/models/ai_key_health_test.go new file mode 100644 index 0000000000..72e6f3e5fe --- /dev/null +++ b/pkg/aiproxy/models/ai_key_health_test.go @@ -0,0 +1,215 @@ +// Copyright 2019 Yunion +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package models + +import ( + "strings" + "testing" + "time" + + api "yunion.io/x/onecloud/pkg/apis/aiproxy" +) + +func resetAiKeyHealthForTest() { + aiKeyHealthMu.Lock() + aiKeyHealth = map[string]*aiKeyHealthState{} + aiKeyHealthMu.Unlock() +} + +func TestRecordAiKeyFailure_ConsecutiveTriggersCooldown(t *testing.T) { + resetAiKeyHealthForTest() + const id = "key-consec" + for i := 0; i < aiKeyHealthCooldownAfter; i++ { + RecordAiKeyFailure(id, 429) + } + if mul := dynamicAiKeyWeightMultiplier(id); mul != 0 { + t.Fatalf("expected multiplier 0 during cooldown, got %d", mul) + } + info := aiKeyHealthInfo(id) + if !info.inCooldown { + t.Fatal("expected inCooldown after consecutive failures") + } + if info.score != 0 { + t.Fatalf("expected score 0, got %d", info.score) + } +} + +func TestRecordAiKeyFailure_IntermittentScoreExhaustionEntersCooldown(t *testing.T) { + resetAiKeyHealthForTest() + const id = "key-intermittent" + // fail -25 / success +10 with consecutiveFails reset: score can hit 0 without 3 consecutive fails. + for { + info := aiKeyHealthInfo(id) + if info.score <= 0 || info.inCooldown { + break + } + RecordAiKeyFailure(id, 500) + info = aiKeyHealthInfo(id) + if info.score <= 0 || info.inCooldown { + break + } + RecordAiKeySuccess(id) + if dynamicAiKeyWeightMultiplier(id) <= 0 { + t.Fatal("unexpected zero multiplier after success") + } + } + info := aiKeyHealthInfo(id) + if !info.inCooldown { + t.Fatalf("expected cooldown after score exhaustion, got score=%d inCooldown=%v", info.score, info.inCooldown) + } + if mul := dynamicAiKeyWeightMultiplier(id); mul != 0 { + t.Fatalf("expected multiplier 0 during cooldown, got %d", mul) + } +} + +func TestDynamicAiKeyWeightMultiplier_RecoversAfterCooldown(t *testing.T) { + resetAiKeyHealthForTest() + const id = "key-recover" + RecordAiKeyFailure(id, 429) + RecordAiKeyFailure(id, 429) + RecordAiKeyFailure(id, 429) + st := getAiKeyHealth(id) + aiKeyHealthMu.Lock() + st.cooldownUntil = time.Now().Add(-time.Second) + aiKeyHealthMu.Unlock() + + mul := dynamicAiKeyWeightMultiplier(id) + if mul < aiKeyHealthMaxScore/2 { + t.Fatalf("expected recovered multiplier >= %d, got %d", aiKeyHealthMaxScore/2, mul) + } + info := aiKeyHealthInfo(id) + if info.inCooldown { + t.Fatal("expected cooldown cleared after expiry") + } + if info.score < aiKeyHealthMaxScore/2 { + t.Fatalf("expected score >= %d, got %d", aiKeyHealthMaxScore/2, info.score) + } +} + +func TestDynamicAiKeyWeightMultiplier_StuckScoreStartsCooldown(t *testing.T) { + resetAiKeyHealthForTest() + const id = "key-stuck" + st := getAiKeyHealth(id) + aiKeyHealthMu.Lock() + st.score = 0 + st.cooldownUntil = time.Time{} + aiKeyHealthMu.Unlock() + + if mul := dynamicAiKeyWeightMultiplier(id); mul != 0 { + t.Fatalf("expected 0, got %d", mul) + } + info := aiKeyHealthInfo(id) + if !info.inCooldown { + t.Fatal("expected fallback cooldown for stuck score=0") + } + + aiKeyHealthMu.Lock() + st.cooldownUntil = time.Now().Add(-time.Second) + aiKeyHealthMu.Unlock() + mul := dynamicAiKeyWeightMultiplier(id) + if mul < aiKeyHealthMaxScore/2 { + t.Fatalf("expected recovery after fallback cooldown, got %d", mul) + } +} + +func TestRecordAiKeySuccess_ClearsCooldown(t *testing.T) { + resetAiKeyHealthForTest() + const id = "key-success" + RecordAiKeyFailure(id, 401) + RecordAiKeyFailure(id, 401) + RecordAiKeyFailure(id, 401) + RecordAiKeySuccess(id) + info := aiKeyHealthInfo(id) + if info.inCooldown { + t.Fatal("success should clear cooldown") + } + if mul := dynamicAiKeyWeightMultiplier(id); mul <= 0 { + t.Fatalf("expected positive multiplier after success, got %d", mul) + } +} + +func TestAiKeySkipReason(t *testing.T) { + resetAiKeyHealthForTest() + + t.Run("already tried", func(t *testing.T) { + k := &SAiKey{Secret: "sk-test"} + k.Id = "id-tried" + k.Name = "tried-key" + reason := aiKeySkipReason(k, "deepseek-v4-pro", map[string]bool{"id-tried": true}) + if !strings.Contains(reason, "already tried") { + t.Fatalf("got %q", reason) + } + }) + + t.Run("cooldown", func(t *testing.T) { + k := &SAiKey{Secret: "sk-test", Weight: 1} + k.Id = "id-cd" + k.Name = "cd-key" + RecordAiKeyFailure(k.Id, 429) + RecordAiKeyFailure(k.Id, 429) + RecordAiKeyFailure(k.Id, 429) + reason := aiKeySkipReason(k, "deepseek-v4-pro", nil) + if !strings.Contains(reason, "cooldown") || !strings.Contains(reason, "remaining") { + t.Fatalf("got %q", reason) + } + }) + + t.Run("routing", func(t *testing.T) { + resetAiKeyHealthForTest() + k := &SAiKey{ + Secret: "sk-test", + Weight: 1, + Routing: &api.SAiKeyRouting{ + AllowedModelKeys: []string{"other-model"}, + }, + } + k.Id = "id-route" + k.Name = "route-key" + reason := aiKeySkipReason(k, "deepseek-v4-pro", nil) + if !strings.Contains(reason, "model not allowed by routing") { + t.Fatalf("got %q", reason) + } + }) + + t.Run("empty secret", func(t *testing.T) { + k := &SAiKey{Secret: " "} + k.Name = "empty-key" + reason := aiKeySkipReason(k, "m", nil) + if !strings.Contains(reason, "empty secret") { + t.Fatalf("got %q", reason) + } + }) + + t.Run("usable", func(t *testing.T) { + resetAiKeyHealthForTest() + k := &SAiKey{Secret: "sk-ok", Weight: 1} + k.Id = "id-ok" + k.Name = "ok-key" + if reason := aiKeySkipReason(k, "deepseek-v4-pro", nil); reason != "" { + t.Fatalf("expected empty reason, got %q", reason) + } + }) +} + +func TestFormatAiKeySkipReasons_Truncates(t *testing.T) { + reasons := make([]string, maxAiKeySkipReasonsInError+3) + for i := range reasons { + reasons[i] = "r" + } + out := formatAiKeySkipReasons(reasons) + if !strings.Contains(out, "and 3 more") { + t.Fatalf("got %q", out) + } +} diff --git a/pkg/aiproxy/models/ai_key_resolve.go b/pkg/aiproxy/models/ai_key_resolve.go index 1ac97b9ae3..0a87074e2e 100644 --- a/pkg/aiproxy/models/ai_key_resolve.go +++ b/pkg/aiproxy/models/ai_key_resolve.go @@ -16,6 +16,7 @@ package models import ( "crypto/rand" + "fmt" "math/big" "strings" @@ -25,6 +26,8 @@ import ( "yunion.io/x/onecloud/pkg/httperrors" ) +const maxAiKeySkipReasonsInError = 8 + func effectiveAiKeyRoutingWeight(r *api.SAiKeyRouting) int { if r == nil || r.Weight <= 0 { return 0 @@ -84,6 +87,58 @@ func aiKeyRoutingAcceptsModel(r *api.SAiKeyRouting, reqModel string) bool { return true } +func aiKeyLabel(k *SAiKey) string { + if k == nil { + return "?" + } + if n := strings.TrimSpace(k.Name); n != "" { + return n + } + if id := strings.TrimSpace(k.Id); id != "" { + return id + } + return "?" +} + +// aiKeySkipReason returns why an ai_key cannot be used for modelKey, or "" if usable. +func aiKeySkipReason(k *SAiKey, modelKey string, exclude map[string]bool) string { + if k == nil { + return "?: nil ai_key" + } + label := aiKeyLabel(k) + if strings.TrimSpace(k.Secret) == "" { + return label + ": empty secret" + } + if exclude != nil && exclude[k.Id] { + return label + ": already tried" + } + if baseAiKeyWeight(k) <= 0 { + return label + ": weight=0" + } + if effectiveAiKeyWeight(k) <= 0 { + info := aiKeyHealthInfo(k.Id) + if info.inCooldown { + return fmt.Sprintf("%s: cooldown %ds remaining (health_score=%d)", label, info.remainingSec, info.score) + } + return fmt.Sprintf("%s: health_score=%d", label, info.score) + } + if !aiKeyRoutingAcceptsModel(k.Routing, modelKey) { + return label + ": model not allowed by routing (allowed_model_keys/blocked_model_keys)" + } + return "" +} + +func formatAiKeySkipReasons(reasons []string) string { + if len(reasons) == 0 { + return "" + } + if len(reasons) <= maxAiKeySkipReasonsInError { + return strings.Join(reasons, "; ") + } + shown := strings.Join(reasons[:maxAiKeySkipReasonsInError], "; ") + return fmt.Sprintf("%s; and %d more", shown, len(reasons)-maxAiKeySkipReasonsInError) +} + func pickWeightedAiKey(candidates []*SAiKey) *SAiKey { if len(candidates) == 0 { return nil @@ -145,21 +200,18 @@ func resolveUpstreamAPIKeyExcluding(prov *SAiProvider, modelKey string, exclude candidates := make([]*SAiKey, 0, len(keys)) hasSecretKey := false + skipReasons := make([]string, 0, len(keys)) for i := range keys { k := &keys[i] if strings.TrimSpace(k.Secret) == "" { continue } hasSecretKey = true - if exclude != nil && exclude[k.Id] { + if reason := aiKeySkipReason(k, modelKey, exclude); reason != "" { + skipReasons = append(skipReasons, reason) continue } - if effectiveAiKeyWeight(k) <= 0 { - continue - } - if aiKeyRoutingAcceptsModel(k.Routing, modelKey) { - candidates = append(candidates, k) - } + candidates = append(candidates, k) } if len(candidates) > 0 { chosen := pickWeightedAiKey(candidates) @@ -173,7 +225,11 @@ func resolveUpstreamAPIKeyExcluding(prov *SAiProvider, modelKey string, exclude }, nil } if hasSecretKey { - return nil, errors.Wrapf(httperrors.ErrInvalidStatus, "no available ai_key for catalog model %q (check weight, cooldown, allowed_model_keys)", modelKey) + detail := formatAiKeySkipReasons(skipReasons) + if detail != "" { + return nil, errors.Wrapf(httperrors.ErrInvalidStatus, "no available ai_key for catalog model %q: %s", modelKey, detail) + } + return nil, errors.Wrapf(httperrors.ErrInvalidStatus, "no available ai_key for catalog model %q", modelKey) } return nil, errors.Wrap(httperrors.ErrInvalidStatus, "add an enabled ai_key with secret for this provider") }