From 643efdabb64e40a697a69a6f55797066bb4f14d4 Mon Sep 17 00:00:00 2001 From: rene Date: Tue, 26 May 2026 17:31:58 +0800 Subject: [PATCH] feat: add multi-key support with automatic rate-limit rotation Support multiple API keys via `api_keys` array in config or `OCGO_API_KEYS` env var. On HTTP 429, automatically rotate to the next key and retry. Key index persisted to `~/.config/ocgo/key-index` for restart safety. Backward compatible with existing single-key `api_key` configs. Fixes #115 Co-Authored-By: Claude Opus 4.7 --- cmd/ocgo/main.go | 166 +++++++++++++++++++++++++++++++++++------------ 1 file changed, 125 insertions(+), 41 deletions(-) diff --git a/cmd/ocgo/main.go b/cmd/ocgo/main.go index a0f5c4f..cdfbd19 100644 --- a/cmd/ocgo/main.go +++ b/cmd/ocgo/main.go @@ -32,10 +32,11 @@ const ( var version = "dev" type Config struct { - APIKey string `json:"api_key"` - Host string `json:"host"` - Port int `json:"port"` - Model string `json:"model"` + APIKey string `json:"api_key,omitempty"` + APIKeys []string `json:"api_keys,omitempty"` + Host string `json:"host"` + Port int `json:"port"` + Model string `json:"model"` } type AnthropicRequest struct { @@ -148,30 +149,38 @@ func main() { } func setupCmd() *cobra.Command { - var key string + var apiKeys string cmd := &cobra.Command{ Use: "setup", - Short: "Save your OpenCode Go API key", + Short: "Save your OpenCode Go API key(s)", RunE: func(cmd *cobra.Command, args []string) error { - if strings.TrimSpace(key) == "" { - key = os.Getenv("OCGO_API_KEY") + if apiKeys == "" { + apiKeys, _ = cmd.Flags().GetString("api-key") } - if strings.TrimSpace(key) == "" { - fmt.Print("OpenCode Go API key: ") + if strings.TrimSpace(apiKeys) == "" { + apiKeys = os.Getenv("OCGO_API_KEYS") + } + if strings.TrimSpace(apiKeys) == "" { + apiKeys = os.Getenv("OCGO_API_KEY") + } + if strings.TrimSpace(apiKeys) == "" { + fmt.Print("OpenCode Go API key(s) (comma-separated): ") line, err := bufio.NewReader(os.Stdin).ReadString('\n') if err != nil && line == "" { return err } - key = line + apiKeys = line } - cfg := Config{APIKey: strings.TrimSpace(key), Host: defaultHost, Port: defaultPort} - if cfg.APIKey == "" { - return errors.New("API key cannot be empty") + keys := parseKeys(strings.TrimSpace(apiKeys)) + if len(keys) == 0 { + return errors.New("at least one API key is required") } + cfg := Config{APIKeys: keys, Host: defaultHost, Port: defaultPort} return saveConfig(cfg) }, } - cmd.Flags().StringVar(&key, "api-key", "", "OpenCode Go API key") + cmd.Flags().StringVar(&apiKeys, "api-keys", "", "OpenCode Go API key(s), comma-separated") + cmd.Flags().String("api-key", "", "OpenCode Go API key (single)") return cmd } @@ -413,14 +422,7 @@ func proxyMessages(w http.ResponseWriter, r *http.Request, cfg Config) { return } body, _ := json.Marshal(or) - req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, openAIURL, bytes.NewReader(body)) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - req.Header.Set("Authorization", "Bearer "+cfg.APIKey) - req.Header.Set("Content-Type", "application/json") - resp, err := (&http.Client{Timeout: 10 * time.Minute}).Do(req) + resp, err := cfg.postWithRetry(r.Context(), openAIURL, "application/json", body) if err != nil { http.Error(w, err.Error(), http.StatusBadGateway) return @@ -464,14 +466,7 @@ func proxyChatCompletions(w http.ResponseWriter, r *http.Request, cfg Config) { } } } - req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, openAIURL, bytes.NewReader(body)) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - req.Header.Set("Authorization", "Bearer "+cfg.APIKey) - req.Header.Set("Content-Type", "application/json") - resp, err := (&http.Client{Timeout: 10 * time.Minute}).Do(req) + resp, err := cfg.postWithRetry(r.Context(), openAIURL, "application/json", body) if err != nil { http.Error(w, err.Error(), http.StatusBadGateway) return @@ -498,14 +493,7 @@ func proxyResponses(w http.ResponseWriter, r *http.Request, cfg Config) { return } body, _ := json.Marshal(or) - req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, openAIURL, bytes.NewReader(body)) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - req.Header.Set("Authorization", "Bearer "+cfg.APIKey) - req.Header.Set("Content-Type", "application/json") - resp, err := (&http.Client{Timeout: 10 * time.Minute}).Do(req) + resp, err := cfg.postWithRetry(r.Context(), openAIURL, "application/json", body) if err != nil { http.Error(w, err.Error(), http.StatusBadGateway) return @@ -1522,6 +1510,76 @@ func startServerProcess(detached bool, model string) (*exec.Cmd, error) { func configDir() string { home, _ := os.UserHomeDir(); return filepath.Join(home, ".config", "ocgo") } func configFile() string { return filepath.Join(configDir(), "config.json") } func pidFile() string { return filepath.Join(configDir(), "ocgo.pid") } +func keyIndexFile() string { return filepath.Join(configDir(), "key-index") } + +func readKeyIndex() int { + b, err := os.ReadFile(keyIndexFile()) + if err != nil { + return 0 + } + var idx int + fmt.Sscanf(string(b), "%d", &idx) + return idx +} + +func writeKeyIndex(idx int) error { + return os.WriteFile(keyIndexFile(), []byte(fmt.Sprint(idx)), 0600) +} + +func (cfg Config) activeKeys() []string { + if len(cfg.APIKeys) > 0 { + return cfg.APIKeys + } + if cfg.APIKey != "" { + return []string{cfg.APIKey} + } + return nil +} + +func (cfg Config) currentKey() string { + keys := cfg.activeKeys() + if len(keys) == 0 { + return "" + } + idx := readKeyIndex() % len(keys) + return keys[idx] +} + +func (cfg Config) rotateKey() string { + keys := cfg.activeKeys() + if len(keys) <= 1 { + return cfg.currentKey() + } + idx := readKeyIndex() + idx = (idx + 1) % len(keys) + _ = writeKeyIndex(idx) + return keys[idx] +} + +func (cfg Config) postWithRetry(ctx context.Context, url, contentType string, body []byte) (*http.Response, error) { + client := &http.Client{Timeout: 10 * time.Minute} + keys := cfg.activeKeys() + for i := 0; i < len(keys); i++ { + key := cfg.currentKey() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+key) + req.Header.Set("Content-Type", contentType) + resp, err := client.Do(req) + if err != nil { + return nil, err + } + if resp.StatusCode == http.StatusTooManyRequests && i < len(keys)-1 { + resp.Body.Close() + cfg.rotateKey() + continue + } + return resp, nil + } + panic("unreachable") +} func codexConfigFile() string { home, _ := os.UserHomeDir() @@ -1674,6 +1732,12 @@ func versionParts(v string) [3]int { } func saveConfig(cfg Config) error { + // Write api_key for backward compatibility with older ocgo versions + if len(cfg.APIKeys) > 0 { + cfg.APIKey = cfg.APIKeys[0] + } else { + cfg.APIKey = "" + } if err := os.MkdirAll(configDir(), 0755); err != nil { return err } @@ -1686,12 +1750,21 @@ func saveConfig(cfg Config) error { } func loadConfig() (Config, error) { - cfg := Config{Host: defaultHost, Port: defaultPort, APIKey: os.Getenv("OCGO_API_KEY")} + cfg := Config{Host: defaultHost, Port: defaultPort} b, err := os.ReadFile(configFile()) if err == nil { _ = json.Unmarshal(b, &cfg) } - if cfg.APIKey == "" { + if envKey := os.Getenv("OCGO_API_KEY"); envKey != "" { + cfg.APIKey = envKey + } + if envKeys := os.Getenv("OCGO_API_KEYS"); envKeys != "" { + cfg.APIKeys = parseKeys(envKeys) + } + if len(cfg.APIKeys) == 0 && cfg.APIKey != "" { + cfg.APIKeys = []string{cfg.APIKey} + } + if len(cfg.APIKeys) == 0 { return cfg, errors.New("missing API key; run: ocgo setup") } if cfg.Host == "" { @@ -1703,6 +1776,17 @@ func loadConfig() (Config, error) { return cfg, nil } +func parseKeys(s string) []string { + parts := strings.Split(s, ",") + var keys []string + for _, p := range parts { + if k := strings.TrimSpace(p); k != "" { + keys = append(keys, k) + } + } + return keys +} + func readPID() (int, error) { b, err := os.ReadFile(pidFile()) if err != nil {