From 4481897ed472c6411a9d7487fe9024b04645046d Mon Sep 17 00:00:00 2001 From: 6Kmfi6HP <179412085+6Kmfi6HP@users.noreply.github.com> Date: Sun, 27 Sep 2026 17:10:52 +0200 Subject: [PATCH 1/8] feat(keypool): multi-key rotation, failover and admin Keys tab - Key pool with round_robin/weighted/sticky strategies, zen/go group filter, string-shorthand keys, atomic RR selection - Failover in callOpenCodeEndpoint: 429/5xx cooldown, 401/billing 15min cooldown, pool failover budget 1+max_retries, per-attempt key switch with key_id-only logging - Admin: Keys tab with inline edit, batch add via POST /api/key_parse canonical grammar, per-key status; GET /api/key_status - Per-key usage stats in stats.json; saveConfig tightened to 0600 Tests: go test ./... PASS, go vet clean, gofmt clean --- config.example.json | 15 +- docs/keypool-design.md | 82 +++++++ internal/app/admin.go | 18 ++ internal/app/admin_keys_status.go | 104 +++++++++ internal/app/admin_pool_test.go | 175 +++++++++++++++ internal/app/config.go | 21 +- internal/app/keypool.go | 342 ++++++++++++++++++++++++++++++ internal/app/keypool_test.go | 276 ++++++++++++++++++++++++ internal/app/main.go | 2 + internal/app/opencode.go | 78 +++++-- internal/app/server.go | 1 + internal/app/web/admin.html | 287 ++++++++++++++++++++++++- internal/domain/types.go | 130 ++++++++++++ internal/stats/stats.go | 77 +++++++ internal/stats/stats_key_test.go | 95 +++++++++ 15 files changed, 1686 insertions(+), 17 deletions(-) create mode 100644 docs/keypool-design.md create mode 100644 internal/app/admin_keys_status.go create mode 100644 internal/app/admin_pool_test.go create mode 100644 internal/app/keypool.go create mode 100644 internal/app/keypool_test.go create mode 100644 internal/stats/stats_key_test.go diff --git a/config.example.json b/config.example.json index c60e861..fef4a13 100644 --- a/config.example.json +++ b/config.example.json @@ -14,5 +14,18 @@ "nemotron-3-ultra-free": 1000000, "nemotron-3.5-lightning-free": 1000000 }, - "protocol_rules": [] + "protocol_rules": [], + "_comment_key_pool": "Optional upstream key pool: plain strings in keys[] are shorthand for {\"key\": ...}. Strategy: round_robin | weighted | sticky. Remove this comment key; copy the key_pool block below and fill real keys.", + "key_pool": { + "enabled": false, + "strategy": "round_robin", + "max_retries": 2, + "retry_on": [429, 500, 502, 503, 504], + "cooldown_secs": 60, + "blacklist_after": 3, + "keys": [ + {"id": "k1", "key": "sk-placeholder-1", "group": "", "weight": 1, "enabled": true, "note": "account A"}, + "sk-placeholder-2" + ] + } } diff --git a/docs/keypool-design.md b/docs/keypool-design.md new file mode 100644 index 0000000..5c3e9be --- /dev/null +++ b/docs/keypool-design.md @@ -0,0 +1,82 @@ +# KeyPool 多 Key 管理 / 负载均衡 / Failover 设计（P0 冻结） + +日期：2026-09-27。参考 gpt-load（Channel→Group→AccessKey、原子轮询、冷却拉黑），按本项目做减法： +单上游（opencode zen/go）、单二进制单依赖、无 DB，全部基于 `config.json` + 内存。 + +## 1. 配置契约（`config.json`，面板 `key_pool` 全字段对等可设） + +```jsonc +"key_pool": { + "enabled": false, + "strategy": "round_robin", // round_robin | weighted | sticky，三选一 + "max_retries": 2, // 换 key 重试上限；总 attempt = 1+max_retries + "retry_on": [429, 500, 502, 503, 504], + "cooldown_secs": 60, // 429/5xx/传输错误后的冷却窗口 + "blacklist_after": 3, // 连续失败 N 次 → 长冷却 15min（内存黑名单，重启清空） + "keys": [ + {"id": "k1", "key": "sk-...", "group": "", "weight": 1, "enabled": true, "note": "账号A"} + ] +} +``` + +* `keys[].group`：`""`（默认）= zen/go 双表面可用；`"zen"`/`"go"` 限制表面。 + 按 `auth.shouldUseGoEndpoint(modelID)` 决定本次 surface，只选 group 匹配的 key。 +* `keys[]` 允许字符串简写：`"sk-..."` ≡ `{"key":"sk-..."}`（仿 `ModelAliasList` 的宽容解析）。 +* 缺 `id` → `apply` 时归一化 `k1,k2…`；重复 id → 面板 POST 400，文件加载宽容去重。 +* 零值 = 现状：`enabled:false` 或 keys 为空 → 客户端 `Bearer` 直通，public 保持 public。 + +## 2. 生效规则（可预测，文档化） + +* `key_pool.enabled && 有可用候选 && auth.Mode != Public` → 上游 key 强制走池选择， + 客户端 token 仅保留于日志 `source`，不再作为上游 `Authorization`。 + `go:`/`zen:` 前缀仍控制 surface（即 group 过滤），token 部分被池 key 代替。 +* `auth.Mode == Public` → 永远 public，不动用池内付费 key（防匿名蹭付费额度）。 +* 池启用但本次 surface 无候选 → 回退客户端直通（不断服），记 `pool_fallback` 日志。 + +## 3. 选择器（`internal/app/keypool.go`，内存状态，仿 `socks5RRIndex` 原子模式） + +* `keyRRIndex atomic.Uint64`：`round_robin` 取 `Add % len`；`weighted` 取 `Add % totalWeight` 走权重区间； + `sticky` 取 `fnv32a(stickyKeyForRequest 同源串) % len`（与 egress sticky 同键，保证 prompt cache 亲和）。 +* 候选过滤：`enabled && group匹配 && now > cooldownUntil`；全冷却 → 放行最早过期的那把（不断服）。 +* 状态表（`keypoolMu` 守卫）：`{cooldownUntil, consecutiveFails}`；成功清零。 + +## 4. Failover（改 `opencode.go:582` 重试循环内两点，不建新子系统） + +* 每 attempt：`attemptAuth, keyID := selectPoolKey(auth, modelID, attempt)`， + 以值拷贝传入 `selectUpstreamTarget` + `buildOCRequestWithSubpath` + `invalidateUpstreamTarget` + （三处都收 `auth` 值类型，race-free；sticky egress 按池 key 绑定）。 +* 记账（`ReportKeyResult`）： + | 结果 | 动作 | + |---|---| + | 2xx | `consecutiveFails=0` | + | 429/5xx/`retry_on`/传输错误 | `fails++`，冷却 `cooldown_secs`；`fails>=blacklist_after` → 15min | + | 401 非账单 | 长冷却 15min（token 疑似失效） | + | 账单错误（`isNonRetryableUpstreamError`=true：402/403/401+credits 类型） | 长冷却 15min **+ 换下一把 key 继续**（原来直接 break，现在池内 failover；池耗尽才返回） | +* attempt 上限：`1+max_retries`（默认 3，与现有 `maxUpstreamRetries=3` 对齐）。 +* 流式：首字节写出后不再换 key，只记账（零拷贝透传不动）。 + +## 5. Per-key 统计（`stats.json` 复用异步 delta 落盘模式） + +* `stats.TokenStatsData.Keys map[keyID]*KeyStats{requests, errors, last_status, last_error, last_used_unix}`； + 冷却/失败计数为纯内存视图，随 `GET /api/config` 的 `key_pool_status` 下发，不持久化。 +* 面板 Keys Tab 展示：id/note/group/weight/启用/冷却剩余/requests/errors/last_status。 + +## 6. 管理面板（与配置文件 1:1） + +* 新 Tab「API Keys」：策略表单（enabled/strategy/max_retries/retry_on/cooldown/blacklist_after）+ + 单条添加行 + **批量添加 textarea**（每行 `key | 名称|key | 名称|key|group|weight`，`#` 注释，容错跳过并报告计数）+ + keys 表（启用开关/权重/分组/备注/删除）+ 状态表。 +* `saveConfig()` 提交 `key_pool`；`loadConfig()` 回填；`Ctrl+S` 沿用。 +* `GET /api/config` 返回池 key **明文**（与 `socks5_proxies` 密码同先例，管理员已鉴权）； + 日志/状态一律只出现 `id` 或掩码，永不打完整 key。 +* POST 校验失败 → 400 且不落盘不生效：strategy 非法、空 key、重复 id、`weight<1`。 + +## 7. 明确不做（P5 条件触发：key>100 或多机共享） + +DB/SQLite、密钥加密落盘（现阶段靠文件权限，见 §8）、下游 AccessKey 签发（本项目无下游凭证体系， +加签发 = 改信任模型）、主从同步、per-key 模型路由（已有 `protocol_rules`）。 + +## 8. 安全备注 + +* `saveConfig` 写权限从 `0644` 收紧为 `0600`（单运维者工具；key 明文落盘的最低补偿）。 +* 运维备份 `config.json` 即备份全部 key；`config.example.json` 只放占位示例。 diff --git a/internal/app/admin.go b/internal/app/admin.go index aed5811..7e3730c 100644 --- a/internal/app/admin.go +++ b/internal/app/admin.go @@ -69,6 +69,12 @@ func adminConfigHandler(w http.ResponseWriter, r *http.Request) { cfg.UpstreamBaseURLs = upstreamBaseURLs socks5StickyRT := socks5Sticky socks5Mu.RUnlock() + // key_pool (plaintext keys, same precedent as socks5_proxies + // passwords): admin is authenticated; the torn-read caveat above + // applies equally here. + keypoolMu.RLock() + keyPoolRT := keypoolCfg + keypoolMu.RUnlock() w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]any{ "model_alias": cfg.ModelAlias, @@ -85,6 +91,8 @@ func adminConfigHandler(w http.ResponseWriter, r *http.Request) { "socks5_sticky": socks5StickyRT, "text_only_models": textOnlyModelsRT, "protocol_rules": getProtocolRules(), + "key_pool": keyPoolRT, + "key_pool_status": keyPoolStatus(), "log_level": logging.LevelString(), "log_bodies": logging.BodiesEnabled(), }) @@ -107,6 +115,16 @@ func adminConfigHandler(w http.ResponseWriter, r *http.Request) { return } } + // key_pool 严格校验：strategy 非法、空 key、重复 id、weight<1 + // 即 400，且不落盘不生效。Keys==nil 表示字段缺席，保持原值。 + if payload.KeyPool.Keys != nil { + if err := validateKeyPool(payload.KeyPool); err != nil { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + json.NewEncoder(w).Encode(map[string]string{"error": "invalid key_pool: " + err.Error()}) + return + } + } if err := saveConfig(configPath, payload.AppConfig); err != nil { http.Error(w, `{"error":"Failed to save config"}`, http.StatusInternalServerError) return diff --git a/internal/app/admin_keys_status.go b/internal/app/admin_keys_status.go new file mode 100644 index 0000000..8cf9940 --- /dev/null +++ b/internal/app/admin_keys_status.go @@ -0,0 +1,104 @@ +package app + +import ( + "encoding/json" + "net/http" + "strconv" + "strings" +) + +// keyPoolStatusHandler serves GET /api/key_status behind requireAuth, +// returning the live pool snapshot plus the pool toggle/strategy. +func keyPoolStatusHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + keypoolMu.RLock() + enabled := keypoolCfg.Enabled + strategy := keypoolCfg.Strategy + keypoolMu.RUnlock() + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "keys": keyPoolStatus(), + "pool": map[string]any{"enabled": enabled, "strategy": strategy}, + }) +} + +// keyPoolParseHandler serves POST /api/key_parse behind requireAuth: the +// single canonical implementation of the batch textarea grammar. The panel +// delegates parsing here so client and server can never drift. +func keyPoolParseHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + var payload struct { + Text string `json:"text"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + http.Error(w, `{"error":"Invalid JSON"}`, http.StatusBadRequest) + return + } + if len(payload.Text) > 1024*1024 { + http.Error(w, `{"error":"batch text too large (max 1MiB)"}`, http.StatusBadRequest) + return + } + keys, added, skipped := parseBatchKeys(payload.Text) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "keys": keys, + "added": added, + "skipped": skipped, + }) +} + +// parseBatchKeys parses pasted batch key lines into pool entries. +// +// Grammar per line: key [| note] [| group] [| weight]. +// Lines are trimmed; blank lines and `#` comments are silently ignored +// (not counted as skipped). Group is lowercased and kept only for zen/go, +// otherwise cleared. Weight defaults to 1 when absent; a present but +// unparseable or <1 weight invalidates the line. Lines with an empty key +// or more than 4 columns are invalid. skipped counts only invalid +// (non-blank, non-comment) lines. +func parseBatchKeys(text string) (keys []UpstreamKey, added, skipped int) { + for _, line := range strings.Split(text, "\n") { + line = strings.TrimSpace(line) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + cols := strings.Split(line, "|") + if len(cols) > 4 { + skipped++ + continue + } + for i := range cols { + cols[i] = strings.TrimSpace(cols[i]) + } + if cols[0] == "" { + skipped++ + continue + } + entry := UpstreamKey{Key: cols[0], Weight: 1} + if len(cols) > 1 { + entry.Note = cols[1] + } + if len(cols) > 2 { + if g := strings.ToLower(cols[2]); g == "zen" || g == "go" { + entry.Group = g + } + } + if len(cols) > 3 && cols[3] != "" { + w, err := strconv.Atoi(cols[3]) + if err != nil || w < 1 { + skipped++ + continue + } + entry.Weight = w + } + keys = append(keys, entry) + added++ + } + return keys, added, skipped +} diff --git a/internal/app/admin_pool_test.go b/internal/app/admin_pool_test.go new file mode 100644 index 0000000..479f169 --- /dev/null +++ b/internal/app/admin_pool_test.go @@ -0,0 +1,175 @@ +package app + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + + "github.com/6Kmfi6HP/opencode2api/internal/config" +) + +func postConfig(t *testing.T, payload map[string]any) *httptest.ResponseRecorder { + t.Helper() + body, _ := json.Marshal(payload) + req := httptest.NewRequest(http.MethodPost, "/api/config", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + adminConfigHandler(rec, req) + return rec +} + +func TestAdmin_KeyPoolValidation(t *testing.T) { + oldSnap := config.Get() + oldPool := func() KeyPool { keypoolMu.RLock(); defer keypoolMu.RUnlock(); return keypoolCfg }() + configMu.Lock() + oldCP := configPath + configMu.Unlock() + t.Cleanup(func() { + config.Update(func(s *config.Snapshot) { *s = oldSnap }) + setKeyPool(oldPool) + configMu.Lock() + configPath = oldCP + configMu.Unlock() + }) + tmp := t.TempDir() + cfgPath := filepath.Join(tmp, "config.json") + configMu.Lock() + configPath = cfgPath + configMu.Unlock() + + cases := []struct { + name string + keyPool map[string]any + }{ + {"invalid strategy", map[string]any{"strategy": "bogus", "keys": []any{"sk-a"}}}, + {"duplicate id", map[string]any{"keys": []any{ + map[string]any{"id": "k1", "key": "sk-a"}, + map[string]any{"id": "k1", "key": "sk-b"}, + }}}, + {"empty key", map[string]any{"keys": []any{map[string]any{"id": "k1", "key": ""}}}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + rec := postConfig(t, map[string]any{"key_pool": tc.keyPool}) + if rec.Code != http.StatusBadRequest { + t.Fatalf("POST status = %d, want 400; body=%s", rec.Code, rec.Body.String()) + } + if _, err := os.Stat(cfgPath); !os.IsNotExist(err) { + t.Fatalf("config file must not be written on validation failure: %v", err) + } + }) + } +} + +func TestParseBatchKeys(t *testing.T) { + text := "sk-a | first note | zen | 3\n" + + "# a comment\n" + + "\n" + + "sk-b\n" + + "sk-c|note-c|go|2\n" + + "| no key\n" + + "sk-d|x|y|z|extra\n" + + "sk-e|||0\n" + keys, added, skipped := parseBatchKeys(text) + if added != 3 || skipped != 3 { + t.Fatalf("added=%d skipped=%d, want 3/3", added, skipped) + } + if len(keys) != 3 { + t.Fatalf("len(keys)=%d, want 3", len(keys)) + } + if keys[0].Key != "sk-a" || keys[0].Note != "first note" || keys[0].Group != "zen" || keys[0].Weight != 3 { + t.Fatalf("keys[0]=%+v, want sk-a/first note/zen/3", keys[0]) + } + if keys[1].Key != "sk-b" || keys[1].Weight != 1 { + t.Fatalf("keys[1]=%+v, want sk-b weight 1 default", keys[1]) + } + if keys[2].Key != "sk-c" || keys[2].Group != "go" || keys[2].Weight != 2 { + t.Fatalf("keys[2]=%+v, want sk-c/go/2", keys[2]) + } +} + +func TestKeyStatusEndpoint(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, Strategy: "round_robin", Keys: []UpstreamKey{ + {ID: "k1", Key: "sk-a"}, + {ID: "k2", Key: "sk-b", Group: "zen"}, + }}) + req := httptest.NewRequest(http.MethodGet, "/api/key_status", nil) + rec := httptest.NewRecorder() + keyPoolStatusHandler(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("GET status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + var got struct { + Keys []KeyPoolEntryStatus `json:"keys"` + Pool struct { + Enabled bool `json:"enabled"` + Strategy string `json:"strategy"` + } `json:"pool"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if len(got.Keys) != 2 { + t.Fatalf("len(keys)=%d, want 2; body=%s", len(got.Keys), rec.Body.String()) + } + if !got.Pool.Enabled || got.Pool.Strategy != "round_robin" { + t.Fatalf("pool=%+v, want enabled/round_robin", got.Pool) + } +} + +func TestKeyParseEndpoint(t *testing.T) { + text := "sk-a | note a | zen | 2\n# comment\n\nsk-b\n| bad\n" + body, _ := json.Marshal(map[string]any{"text": text}) + req := httptest.NewRequest(http.MethodPost, "/api/key_parse", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + keyPoolParseHandler(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("POST status = %d, want 200; body=%s", rec.Code, rec.Body.String()) + } + var got struct { + Keys []UpstreamKey `json:"keys"` + Added int `json:"added"` + Skipped int `json:"skipped"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if got.Added != 2 || got.Skipped != 1 || len(got.Keys) != 2 { + t.Fatalf("got added=%d skipped=%d keys=%d, want 2/1/2; body=%s", got.Added, got.Skipped, len(got.Keys), rec.Body.String()) + } + if got.Keys[0].Key != "sk-a" || got.Keys[0].Group != "zen" || got.Keys[0].Weight != 2 { + t.Fatalf("keys[0]=%+v, want sk-a/zen/2", got.Keys[0]) + } + // 方法错误 → 405。 + reqGet := httptest.NewRequest(http.MethodGet, "/api/key_parse", nil) + recGet := httptest.NewRecorder() + keyPoolParseHandler(recGet, reqGet) + if recGet.Code != http.StatusMethodNotAllowed { + t.Fatalf("GET status = %d, want 405", recGet.Code) + } +} + +func TestSaveConfig_TightensPerms(t *testing.T) { + tmp := t.TempDir() + cfgPath := filepath.Join(tmp, "config.json") + // 预置旧 0644 文件：保存后必须收紧为 0600。 + if err := os.WriteFile(cfgPath, []byte("{}"), 0o644); err != nil { + t.Fatal(err) + } + if err := saveConfig(cfgPath, AppConfig{KeyPool: KeyPool{Enabled: true}}); err != nil { + t.Fatal(err) + } + fi, err := os.Stat(cfgPath) + if err != nil { + t.Fatal(err) + } + if fi.Mode().Perm() != 0o600 { + t.Fatalf("config perm = %o, want 600", fi.Mode().Perm()) + } +} diff --git a/internal/app/config.go b/internal/app/config.go index baf579d..258dfd0 100644 --- a/internal/app/config.go +++ b/internal/app/config.go @@ -57,7 +57,17 @@ func saveConfig(path string, cfg AppConfig) error { return err } } - return os.WriteFile(path, data, 0o644) + // Key 池明文落盘：新文件 0600；已存在的旧文件（0644 等）同样收紧， + // 避免仅靠 umask 残留可读权限。 + if err := os.WriteFile(path, data, 0o600); err != nil { + return err + } + if fi, err := os.Stat(path); err == nil && fi.Mode().Perm() != 0o600 { + if err := os.Chmod(path, 0o600); err != nil { + slog.Warn("config chmod failed", "path", path, "error", err) + } + } + return nil } func compileKeywordRules(rules []domain.ModelKeywordRule) ([]domain.ModelKeywordRule, []compiledKeywordRule) { @@ -210,6 +220,15 @@ func applyConfig(cfg AppConfig) { if cfg.ProtocolRules != nil { setProtocolRules(compileProtocolRulesLenient(cfg.ProtocolRules)) } + + // Key pool: normalize (default strategy round_robin, weight>=1, + // k1.. ids, dedup ids) and swap into runtime. A nil Keys slice means + // the section was absent, so keep prior runtime state (same precedent + // as the socks5 nil-guard above); normalizeKeyPool itself also + // preserves nil as nil for torn-read parity. + if cfg.KeyPool.Keys != nil { + setKeyPool(cfg.KeyPool) + } } // stripContextSuffix splits a model ID into its base and context suffix. diff --git a/internal/app/keypool.go b/internal/app/keypool.go new file mode 100644 index 0000000..5ec9280 --- /dev/null +++ b/internal/app/keypool.go @@ -0,0 +1,342 @@ +package app + +import ( + "fmt" + "hash/fnv" + "net/http" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" +) + +const ( + keyPoolDefaultCooldownSecs = 60 + keyPoolDefaultBlacklistAfter = 3 + keyPoolLongCooldown = 15 * time.Minute +) + +// keypoolMaxAttempts returns total pool attempts (1+max_retries, default 3). +func keypoolMaxAttempts() int { + keypoolMu.RLock() + maxRetries := keypoolCfg.MaxRetries + keypoolMu.RUnlock() + if maxRetries <= 0 { + maxRetries = 2 + } + return 1 + maxRetries +} + +// keypoolAttemptsExhausted reports whether the pool failover budget is spent. +func keypoolAttemptsExhausted(attempt int) bool { + return attempt+1 >= keypoolMaxAttempts() +} + +type keypoolEntryState struct { + cooldownUntil time.Time + consecutiveFails int +} + +var ( + keypoolMu sync.RWMutex + keypoolCfg KeyPool + keypoolEntries []UpstreamKey + keypoolState = map[string]*keypoolEntryState{} + keypoolRRIndex atomic.Uint64 +) + +// normalizeKeyPool assigns k1,k2… ids to entries missing one, drops empty-key +// and duplicate-id entries (first wins), and applies weight/group defaults. +func normalizeKeyPool(p KeyPool) KeyPool { + out := p + out.Keys = nil + if strings.TrimSpace(out.Strategy) == "" { + out.Strategy = "round_robin" + } + seen := map[string]bool{} + nextAuto := 1 + for _, k := range p.Keys { + if strings.TrimSpace(k.Key) == "" { + continue + } + k = k.Normalized() + if strings.TrimSpace(k.ID) == "" { + for { + candidate := "k" + strconv.Itoa(nextAuto) + nextAuto++ + if !seen[candidate] { + k.ID = candidate + break + } + } + } + if seen[k.ID] { + continue + } + seen[k.ID] = true + out.Keys = append(out.Keys, k) + } + return out +} + +// validateKeyPool strictly validates a key_pool section for admin POST: +// unknown strategy, empty key, duplicate id, or weight<1 fails with an +// indexed error. Missing ids are assigned k1,k2… at apply time, so only +// explicit ids participate in the duplicate check. A nil Keys slice means +// the section was absent and is accepted as-is (callers skip validation +// in that case, matching the protocol_rules nil-guard precedent). +func validateKeyPool(p KeyPool) error { + switch s := strings.ToLower(strings.TrimSpace(p.Strategy)); s { + case "", "round_robin", "weighted", "sticky": + default: + return fmt.Errorf("strategy must be round_robin, weighted, or sticky; got %q", p.Strategy) + } + seen := map[string]int{} + for i, k := range p.Keys { + if strings.TrimSpace(k.Key) == "" { + return fmt.Errorf("keys[%d]: key must not be empty", i) + } + if k.Weight < 1 { + return fmt.Errorf("keys[%d]: weight must be >= 1", i) + } + id := strings.TrimSpace(k.ID) + if id == "" { + continue + } + if first, dup := seen[id]; dup { + return fmt.Errorf("keys[%d]: duplicate id %q (first at keys[%d])", i, id, first) + } + seen[id] = i + } + return nil +} + +// setKeyPool replaces entries and clears runtime state. +func setKeyPool(p KeyPool) { + p = normalizeKeyPool(p) + keypoolMu.Lock() + keypoolCfg = p + keypoolEntries = p.Keys + keypoolState = map[string]*keypoolEntryState{} + keypoolMu.Unlock() + keypoolRRIndex.Store(0) +} + +func poolEnabled() bool { + keypoolMu.RLock() + defer keypoolMu.RUnlock() + return keypoolCfg.Enabled && len(keypoolEntries) > 0 +} + +// stickySessionBase mirrors stickyKeyForRequest's token part without +// importing body context: account token wins, else the public fallback. +func stickySessionBase(auth UpstreamAuth) string { + if auth.Token != "" { + return "tok:" + auth.Token + } + return stickyPublicFallback +} + +func keyPoolGroupOK(group string, goSurface bool) bool { + switch group { + case "": + return true + case "go": + return goSurface + case "zen": + return !goSurface + default: + return false + } +} + +// selectPoolKey picks a pooled key for this request. +// ok=false means caller falls back to client-token passthrough +// (pool disabled, public auth, or no surface candidate). +func selectPoolKey(auth UpstreamAuth, modelID string, attempt ...int) (UpstreamAuth, string, bool) { + _ = attempt + if auth.Mode == AuthRoutePublic { + return auth, "", false + } + keypoolMu.RLock() + enabled := keypoolCfg.Enabled + strategy := keypoolCfg.Strategy + entries := keypoolEntries + keypoolMu.RUnlock() + if !enabled || len(entries) == 0 { + return auth, "", false + } + goSurface := auth.shouldUseGoEndpoint(modelID) + now := time.Now() + + keypoolMu.RLock() + candidates := make([]UpstreamKey, 0, len(entries)) + for _, e := range entries { + if !e.IsEnabled() || !keyPoolGroupOK(e.Group, goSurface) { + continue + } + candidates = append(candidates, e) + } + states := make(map[string]keypoolEntryState, len(candidates)) + for _, c := range candidates { + if s, ok := keypoolState[c.ID]; ok { + states[c.ID] = *s + } + } + keypoolMu.RUnlock() + if len(candidates) == 0 { + return auth, "", false + } + + avail := candidates[:0:0] + for _, c := range candidates { + if s, ok := states[c.ID]; !ok || !now.Before(s.cooldownUntil) { + avail = append(avail, c) + } + } + pool := avail + if len(pool) == 0 { + // All cooling: serve the earliest-expiring one, never hard-fail. + earliest := candidates[0] + for _, c := range candidates[1:] { + if states[c.ID].cooldownUntil.Before(states[earliest.ID].cooldownUntil) { + earliest = c + } + } + pool = []UpstreamKey{earliest} + } + + var picked UpstreamKey + switch strings.ToLower(strings.TrimSpace(strategy)) { + case "weighted": + total := 0 + for _, c := range pool { + total += c.Weight + } + if total < 1 { + total = len(pool) + } + slot := int(keypoolRRIndex.Add(1)-1) % total + for _, c := range pool { + w := c.Weight + if w < 1 { + w = 1 + } + if slot < w { + picked = c + break + } + slot -= w + } + if picked.ID == "" { + picked = pool[len(pool)-1] + } + case "sticky": + h := fnv.New32a() + _, _ = h.Write([]byte(stickySessionBase(auth))) + picked = pool[int(h.Sum32()%uint32(len(pool)))] + default: // round_robin + picked = pool[int(keypoolRRIndex.Add(1)-1)%len(pool)] + } + + out := auth + out.Token = picked.Key + return out, picked.ID, true +} + +// reportKeyResult records an upstream attempt for failover accounting. +// Optional body enables billing-error detection (isNonRetryableUpstreamError +// parses the payload); without it, 402/403 fall through to retry_on rules. +func reportKeyResult(id string, status int, transportErr error, body ...[]byte) { + keypoolMu.RLock() + cooldownSecs := keypoolCfg.CooldownSecs + blacklistAfter := keypoolCfg.BlacklistAfter + retryOn := keypoolCfg.RetryOn + keypoolMu.RUnlock() + if cooldownSecs <= 0 { + cooldownSecs = keyPoolDefaultCooldownSecs + } + if blacklistAfter <= 0 { + blacklistAfter = keyPoolDefaultBlacklistAfter + } + now := time.Now() + + keypoolMu.Lock() + defer keypoolMu.Unlock() + s := keypoolState[id] + if s == nil { + s = &keypoolEntryState{} + keypoolState[id] = s + } + if status >= 200 && status < 300 { + s.consecutiveFails = 0 + return + } + var payload []byte + if len(body) > 0 { + payload = body[0] + } + if status == http.StatusUnauthorized || isNonRetryableUpstreamError(status, payload) { + // Invalid token or billing failure: long cooldown. + s.consecutiveFails++ + s.cooldownUntil = now.Add(keyPoolLongCooldown) + return + } + retryable := transportErr != nil || status == 429 || (status >= 500 && status < 600) + if !retryable && len(retryOn) > 0 { + for _, code := range retryOn { + if code == status { + retryable = true + break + } + } + } + if !retryable { + return + } + s.consecutiveFails++ + if s.consecutiveFails >= blacklistAfter { + s.cooldownUntil = now.Add(keyPoolLongCooldown) + return + } + s.cooldownUntil = now.Add(time.Duration(cooldownSecs) * time.Second) +} + +// KeyPoolEntryStatus is the admin snapshot row for one pooled key. +type KeyPoolEntryStatus struct { + ID string `json:"id"` + Group string `json:"group,omitempty"` + Weight int `json:"weight"` + Enabled bool `json:"enabled"` + Note string `json:"note,omitempty"` + ConsecutiveFails int `json:"consecutive_fails"` + InCooldown bool `json:"in_cooldown"` + CooldownRemainingSecs int64 `json:"cooldown_remaining_secs"` +} + +// keyPoolStatus snapshots pool entries with memory-only cooldown state. +func keyPoolStatus() []KeyPoolEntryStatus { + now := time.Now() + keypoolMu.RLock() + defer keypoolMu.RUnlock() + out := make([]KeyPoolEntryStatus, 0, len(keypoolEntries)) + for _, e := range keypoolEntries { + row := KeyPoolEntryStatus{ + ID: e.ID, + Group: e.Group, + Weight: e.Weight, + Enabled: e.IsEnabled(), + Note: e.Note, + } + if s, ok := keypoolState[e.ID]; ok { + row.ConsecutiveFails = s.consecutiveFails + if now.Before(s.cooldownUntil) { + row.InCooldown = true + row.CooldownRemainingSecs = int64(s.cooldownUntil.Sub(now).Seconds()) + } + } + out = append(out, row) + } + return out +} diff --git a/internal/app/keypool_test.go b/internal/app/keypool_test.go new file mode 100644 index 0000000..cacf875 --- /dev/null +++ b/internal/app/keypool_test.go @@ -0,0 +1,276 @@ +package app + +import ( + "encoding/json" + "errors" + "testing" +) + +func resetPool(t *testing.T) { + t.Helper() + setKeyPool(KeyPool{}) + t.Cleanup(func() { setKeyPool(KeyPool{}) }) +} + +func TestKeyPool_RoundRobin(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, Keys: []UpstreamKey{{Key: "k-a"}, {Key: "k-b"}, {Key: "k-c"}}}) + auth := UpstreamAuth{Mode: AuthRouteAuto, Token: "client"} + var got []string + for range 6 { + _, id, ok := selectPoolKey(auth, "m") + if !ok { + t.Fatal("want ok") + } + got = append(got, id) + } + want := []string{"k1", "k2", "k3", "k1", "k2", "k3"} + for i := range want { + if got[i] != want[i] { + t.Fatalf("order[%d] = %s, want %s (%v)", i, got[i], want[i], got) + } + } + // Pooled auth replaces client token. + a, _, _ := selectPoolKey(auth, "m") + if a.Token == "client" { + t.Fatal("pooled auth must replace client token") + } +} + +func TestKeyPool_Weighted(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, Strategy: "weighted", Keys: []UpstreamKey{ + {ID: "a", Key: "ka", Weight: 3}, + {ID: "b", Key: "kb", Weight: 1}, + }}) + auth := UpstreamAuth{Mode: AuthRouteAuto, Token: "c"} + counts := map[string]int{} + for range 8 { + _, id, ok := selectPoolKey(auth, "m") + if !ok { + t.Fatal("want ok") + } + counts[id]++ + } + if counts["a"] != 6 || counts["b"] != 2 { + t.Fatalf("weighted distribution = %v, want a=6 b=2", counts) + } +} + +func TestKeyPool_Sticky(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, Strategy: "sticky", Keys: []UpstreamKey{{Key: "ka"}, {Key: "kb"}}}) + a1 := UpstreamAuth{Mode: AuthRouteAuto, Token: "user1"} + a2 := UpstreamAuth{Mode: AuthRouteAuto, Token: "user2"} + _, id1a, _ := selectPoolKey(a1, "m") + _, id1b, _ := selectPoolKey(a1, "m") + if id1a != id1b { + t.Fatal("sticky must return same key for same session") + } + if got := stickySessionBase(a1); got != "tok:user1" { + t.Fatalf("stickySessionBase = %q", got) + } + if got := stickySessionBase(UpstreamAuth{}); got != stickyPublicFallback { + t.Fatalf("public fallback = %q", got) + } + _ = a2 +} + +func TestKeyPool_Disabled(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: false, Keys: []UpstreamKey{{Key: "ka"}}}) + if poolEnabled() { + t.Fatal("pool must be disabled") + } + if _, _, ok := selectPoolKey(UpstreamAuth{Mode: AuthRouteAuto, Token: "c"}, "m"); ok { + t.Fatal("disabled pool must return ok=false") + } + setKeyPool(KeyPool{Enabled: true}) + if poolEnabled() { + t.Fatal("empty pool must report disabled") + } +} + +func TestKeyPool_CooldownSkip(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, CooldownSecs: 60, Keys: []UpstreamKey{{ID: "a", Key: "ka"}, {ID: "b", Key: "kb"}}}) + reportKeyResult("a", 429, nil) + auth := UpstreamAuth{Mode: AuthRouteAuto, Token: "c"} + for range 4 { + _, id, ok := selectPoolKey(auth, "m") + if !ok || id != "b" { + t.Fatalf("cooling key must be skipped, got %q ok=%v", id, ok) + } + } + // All cooling → earliest expiry still serves. + reportKeyResult("b", 500, nil) + if _, _, ok := selectPoolKey(auth, "m"); !ok { + t.Fatal("all-cooldown must fall back to earliest expiry, not fail") + } + // Success clears fails. + reportKeyResult("a", 200, nil) + st := keyPoolStatus() + for _, r := range st { + if r.ID == "a" && r.ConsecutiveFails != 0 { + t.Fatalf("success must clear fails: %+v", r) + } + } + // 401 → long cooldown. + reportKeyResult("a", 401, nil) + st = keyPoolStatus() + for _, r := range st { + if r.ID == "a" && !r.InCooldown { + t.Fatal("401 must trigger long cooldown") + } + } + // Billing body → long cooldown. + reportKeyResult("b", 402, nil, []byte(`{"error":{"type":"CreditsError","message":"x"}}`)) + // transport error with blacklist_after=1 → long cooldown path exercised. + setKeyPool(KeyPool{Enabled: true, CooldownSecs: 60, BlacklistAfter: 1, Keys: []UpstreamKey{{ID: "a", Key: "ka"}}}) + reportKeyResult("a", 0, errors.New("dial")) + st = keyPoolStatus() + if !st[0].InCooldown { + t.Fatal("transport error must cool down") + } +} + +func TestKeyPool_GroupFilter(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, Keys: []UpstreamKey{ + {ID: "z", Key: "kz", Group: "zen"}, + {ID: "g", Key: "kg", Group: "go"}, + {ID: "both", Key: "kb"}, + }}) + zenAuth := UpstreamAuth{Mode: AuthRouteZen, Token: "c"} + goAuth := UpstreamAuth{Mode: AuthRouteGo, Token: "c"} + // zen surface: go-only key must never be picked. + for range 10 { + _, id, ok := selectPoolKey(zenAuth, "some-model") + if !ok { + t.Fatal("want ok") + } + if id == "g" { + t.Fatal("zen surface picked go-only key") + } + } + // Seed catalogs: "go-only-model" exists only in the go catalog. + oldModels, oldGo := modelsCache, goModelsCache + modelMu.Lock() + modelsCache = []ModelInfo{{ID: "shared-model"}} + goModelsCache = []ModelInfo{{ID: "shared-model"}, {ID: "go-only-model"}} + modelMu.Unlock() + t.Cleanup(func() { + modelMu.Lock() + modelsCache, goModelsCache = oldModels, oldGo + modelMu.Unlock() + }) + goModel := "go-only-model" + for range 10 { + _, id, ok := selectPoolKey(goAuth, goModel) + if !ok { + t.Fatal("want ok") + } + if id == "z" { + t.Fatal("go surface picked zen-only key") + } + } +} + +func TestKeyPool_PublicNeverPooled(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, Keys: []UpstreamKey{{Key: "ka"}}}) + if _, _, ok := selectPoolKey(UpstreamAuth{Mode: AuthRoutePublic}, "m"); ok { + t.Fatal("public must never use pool") + } +} + +func TestKeyPool_StringShorthand(t *testing.T) { + var p KeyPool + if err := json.Unmarshal([]byte(`{"enabled":true,"keys":["sk-a",{"id":"x","key":"sk-b"}]}`), &p); err != nil { + t.Fatal(err) + } + if len(p.Keys) != 2 || p.Keys[0].Key != "sk-a" || p.Keys[1].ID != "x" { + t.Fatalf("shorthand unmarshal: %+v", p.Keys) + } +} + +func TestKeyPool_NormalizeDedup(t *testing.T) { + p := normalizeKeyPool(KeyPool{Keys: []UpstreamKey{ + {Key: "ka"}, + {Key: "kb"}, + {ID: "k1", Key: "dup"}, + {Key: ""}, + {ID: "z", Key: "kz", Weight: 0, Group: " ZEN "}, + }}) + if len(p.Keys) != 3 { + t.Fatalf("want 3 entries, got %+v", p.Keys) + } + ids := map[string]bool{} + for _, k := range p.Keys { + if ids[k.ID] { + t.Fatalf("dup id %q", k.ID) + } + ids[k.ID] = true + if k.Weight < 1 { + t.Fatalf("weight default missing: %+v", k) + } + } + if p.Keys[2].Group != "zen" { + t.Fatalf("group normalize: %+v", p.Keys[2]) + } +} + +func TestKeyPool_FailoverBillingLongCooldown(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, CooldownSecs: 60, Keys: []UpstreamKey{{ID: "a", Key: "ka"}, {ID: "b", Key: "kb"}}}) + reportKeyResult("a", 402, nil, []byte(`{"error":{"type":"CreditsError","message":"insufficient credits"}}`)) + st := keyPoolStatus() + for _, r := range st { + if r.ID == "a" { + if !r.InCooldown { + t.Fatal("billing error must trigger cooldown") + } + if r.CooldownRemainingSecs < 500 { + t.Fatalf("billing error must use 15min cooldown, got %d secs", r.CooldownRemainingSecs) + } + } + } +} + +func TestKeyPool_Failover429CooldownSecs(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, CooldownSecs: 60, Keys: []UpstreamKey{{ID: "a", Key: "ka"}}}) + reportKeyResult("a", 429, nil) + st := keyPoolStatus() + if !st[0].InCooldown { + t.Fatal("429 must trigger cooldown") + } + if st[0].CooldownRemainingSecs > 65 || st[0].CooldownRemainingSecs <= 0 { + t.Fatalf("429 must use cooldown_secs window, got %d", st[0].CooldownRemainingSecs) + } +} + +func TestKeyPool_FailoverSuccessClearsFails(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, CooldownSecs: 60, Keys: []UpstreamKey{{ID: "a", Key: "ka"}}}) + reportKeyResult("a", 500, nil) + reportKeyResult("a", 200, nil) + st := keyPoolStatus() + if st[0].ConsecutiveFails != 0 { + t.Fatalf("success must clear fails: %+v", st[0]) + } +} + +func TestKeyPool_FailoverAttemptsExhausted(t *testing.T) { + resetPool(t) + setKeyPool(KeyPool{Enabled: true, MaxRetries: 2, Keys: []UpstreamKey{{ID: "a", Key: "ka"}}}) + if keypoolMaxAttempts() != 3 { + t.Fatalf("default pool attempts must be 3, got %d", keypoolMaxAttempts()) + } + if keypoolAttemptsExhausted(0) || keypoolAttemptsExhausted(1) { + t.Fatal("attempts 0,1 must not be exhausted with max 3") + } + if !keypoolAttemptsExhausted(2) { + t.Fatal("attempt 2 must be exhausted with max 3") + } +} diff --git a/internal/app/main.go b/internal/app/main.go index 8e497d0..05d726e 100644 --- a/internal/app/main.go +++ b/internal/app/main.go @@ -24,6 +24,8 @@ const MatchRegex = domain.MatchRegex type ModelAliasList = domain.ModelAliasList type Socks5Proxy = domain.Socks5Proxy +type UpstreamKey = domain.UpstreamKey +type KeyPool = domain.KeyPool type ClaudeRequest = domain.ClaudeRequest type ClaudeMessage = domain.ClaudeMessage type ClaudeContent = domain.ClaudeContent diff --git a/internal/app/opencode.go b/internal/app/opencode.go index 73233b7..b968dea 100644 --- a/internal/app/opencode.go +++ b/internal/app/opencode.go @@ -19,6 +19,7 @@ import ( "github.com/6Kmfi6HP/opencode2api/internal/logging" "github.com/6Kmfi6HP/opencode2api/internal/modelsdev" "github.com/6Kmfi6HP/opencode2api/internal/random" + statsx "github.com/6Kmfi6HP/opencode2api/internal/stats" ) // ======================== 随机 ID ======================== @@ -552,6 +553,20 @@ func maxAttemptsForUpstreamStatus(status int) int { return maxUpstreamRetries } +func truncateKeyErr(s string) string { + if len(s) > 200 { + return s[:200] + } + return s +} + +func truncateKeyErrBytes(b []byte) []byte { + if len(b) > 200 { + return b[:200] + } + return b +} + // callOpenCodeEndpoint 统一封装所有对上游 /zen/v1/* 和 /zen/go/v1/* 端点的 HTTP 调用， // 包含重试机制、SOCKS5 会话粘性与轮换、多域名轮换、错误归一与结构化日志输出。 func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamBody []byte, modelID string, auth UpstreamAuth) (io.ReadCloser, int, http.Header, error) { @@ -578,6 +593,9 @@ func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamB if max401Retries > maxAttempts { maxAttempts = max401Retries } + if poolEnabled() && keypoolMaxAttempts() > maxAttempts { + maxAttempts = keypoolMaxAttempts() + } for attempt := 0; attempt < maxAttempts; attempt++ { // 仅"裸进程内回退值"才现造一次会话,避免同一请求的重试在 @@ -586,26 +604,34 @@ func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamB if strings.TrimSpace(ocSession) == "" { ocSession = newOCSessionID() } + attemptAuth, keyID, pooled := selectPoolKey(auth, modelID, attempt) + targetAuth := auth + if pooled { + targetAuth = attemptAuth + } upstreamHeaders := upstreamHeadersFromContext(ctx) - baseURL, client := selectUpstreamTarget(auth, bodyMap, upstreamHeaders, normalizedTransportScope(ocSession)) + baseURL, client := selectUpstreamTarget(targetAuth, bodyMap, upstreamHeaders, normalizedTransportScope(ocSession)) lastBaseURL = baseURL - up, err := buildOCRequestWithSubpath(modelID, bodyMap, auth, useGoEndpoint, baseURL, endpointSubpath, ocSession) + up, err := buildOCRequestWithSubpath(modelID, bodyMap, targetAuth, useGoEndpoint, baseURL, endpointSubpath, ocSession) if err != nil { return nil, 500, nil, err } - up = up.WithContext(ctx) attemptStart := time.Now() resp, err := client.Do(up) durationMs := time.Since(attemptStart).Milliseconds() if err != nil { lastErr = err lastStatus = 0 + if pooled { + reportKeyResult(keyID, 0, err) + statsx.RecordKeyUsage(keyID, 0, truncateKeyErr(err.Error())) + } retryReason := "transport_error" canRetry := attempt+1 < maxUpstreamRetries if !canRetry { retryReason = "" } - log.Info("upstream_attempt", + args := []any{ "try_model", modelID, "base_url", baseURL, "surface", surface, @@ -614,24 +640,36 @@ func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamB "attempt_index", attempt, "retry_reason", retryReason, "error", err.Error(), - ) + } + if pooled { + args = append(args, "key_id", keyID) + } + log.Info("upstream_attempt", args...) if canRetry { client.CloseIdleConnections() - invalidateUpstreamTarget(auth, bodyMap, upstreamHeaders, sessionFromRequestContext(ctx, ocSessionID)) + invalidateUpstreamTarget(targetAuth, bodyMap, upstreamHeaders, sessionFromRequestContext(ctx, ocSessionID)) retryCount++ continue } break } if resp.StatusCode >= 200 && resp.StatusCode < 300 { - log.Info("upstream_attempt", + if pooled { + reportKeyResult(keyID, resp.StatusCode, nil) + statsx.RecordKeyUsage(keyID, resp.StatusCode, "") + } + args := []any{ "try_model", modelID, "base_url", baseURL, "surface", surface, "status", resp.StatusCode, "duration_ms", durationMs, "attempt_index", attempt, - ) + } + if pooled { + args = append(args, "key_id", keyID) + } + log.Info("upstream_attempt", args...) log.Info("upstream_result", "models_tried", []string{modelID}, "base_url", baseURL, @@ -644,6 +682,10 @@ func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamB errBody, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024)) resp.Body.Close() logging.UpstreamError(ctx, modelID, resp.StatusCode, errBody, baseURL) + if pooled { + reportKeyResult(keyID, resp.StatusCode, nil, errBody) + statsx.RecordKeyUsage(keyID, resp.StatusCode, string(truncateKeyErrBytes(errBody))) + } nonRetryable := isNonRetryableUpstreamError(resp.StatusCode, errBody) canRetry := !nonRetryable && shouldRetryUpstreamStatus(resp.StatusCode) && attempt+1 < maxAttemptsForUpstreamStatus(resp.StatusCode) retryReason := "" @@ -651,9 +693,17 @@ func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamB retryReason = fmt.Sprintf("status_%d", resp.StatusCode) } if nonRetryable { - retryReason = "non_retryable_upstream" + if pooled && !keypoolAttemptsExhausted(attempt) { + retryReason = "pool_failover" + canRetry = true + } else { + retryReason = "non_retryable_upstream" + } + } else if pooled && resp.StatusCode == http.StatusTooManyRequests && !canRetry && !keypoolAttemptsExhausted(attempt) { + retryReason = "pool_failover" + canRetry = true } - log.Info("upstream_attempt", + args := []any{ "try_model", modelID, "base_url", baseURL, "surface", surface, @@ -661,7 +711,11 @@ func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamB "duration_ms", durationMs, "attempt_index", attempt, "retry_reason", retryReason, - ) + } + if pooled { + args = append(args, "key_id", keyID) + } + log.Info("upstream_attempt", args...) lastBody = errBody lastStatus = resp.StatusCode lastHeader = resp.Header @@ -671,7 +725,7 @@ func callOpenCodeEndpoint(ctx context.Context, endpointSubpath string, upstreamB } // 免费层 429 按出口 IP 限流,5xx 也可能是出口问题: // 重试前切断 sticky,让同一会话换到下一个出口。 - invalidateUpstreamTarget(auth, bodyMap, upstreamHeaders, sessionFromRequestContext(ctx, ocSessionID)) + invalidateUpstreamTarget(targetAuth, bodyMap, upstreamHeaders, sessionFromRequestContext(ctx, ocSessionID)) client.CloseIdleConnections() retryCount++ } diff --git a/internal/app/server.go b/internal/app/server.go index 6de8eff..bf2033b 100644 --- a/internal/app/server.go +++ b/internal/app/server.go @@ -132,6 +132,7 @@ func buildMux() *http.ServeMux { mux.HandleFunc("/logout", logging.Middleware(logoutHandler)) mux.HandleFunc("/api/config", logging.Middleware(requireAuth(adminConfigHandler))) mux.HandleFunc("/api/stats", logging.Middleware(requireAuth(adminStatsHandler))) + mux.HandleFunc("/api/key_parse", logging.Middleware(requireAuth(keyPoolParseHandler))) mux.HandleFunc("/api/reload", logging.Middleware(requireAuth(reloadHandler))) mux.HandleFunc("/health", logging.Middleware(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) diff --git a/internal/app/web/admin.html b/internal/app/web/admin.html index 3f694e6..e053676 100644 --- a/internal/app/web/admin.html +++ b/internal/app/web/admin.html @@ -905,6 +905,10 @@ 网络与代理 + + +
+ + +
+
+ +
+ +
+
+
+ + Key 池列表 +
+
+
+ + + + + + + + + + + + + + +
启用IDKEY 掩码分组权重备注状态操作
+
+

冷却 / 失败计数为内存视图，重启后清空。

+
+ - +
@@ -1425,6 +1534,8 @@