From fe4ae4989c4db599ed8b59a176fc6809b9271acb Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sun, 26 Jul 2026 14:07:27 +0800 Subject: [PATCH] chore(pluginhost): refactor and remove unused interceptors and executor methods - Removed deprecated interceptor and executor-related methods, including `callRequestInterceptor`, `callResponseInterceptor`, and `callStreamChunkInterceptor`. - Consolidated unused logic and pruned redundant imports to streamline `adapters.go`. - No functional changes. --- .../api/handlers/management/auth_files.go | 2263 ----- .../handlers/management/auth_files_crud.go | 550 ++ .../handlers/management/auth_files_fields.go | 684 ++ .../management/auth_files_oauth_callback.go | 220 + .../management/auth_files_provider_oauth.go | 875 ++ internal/api/server.go | 1801 ---- internal/api/server_keepalive.go | 89 + internal/api/server_management.go | 312 + internal/api/server_middleware.go | 172 + internal/api/server_options.go | 135 + internal/api/server_reload.go | 276 + internal/api/server_routes.go | 896 ++ internal/config/config.go | 1906 ----- internal/config/config_defaults.go | 7 + internal/config/config_load.go | 176 + internal/config/config_normalization.go | 297 + internal/config/config_types.go | 580 ++ internal/config/config_validation.go | 79 + internal/config/config_yaml.go | 815 ++ internal/logging/request_logger.go | 1964 ----- .../logging/request_logger_body_source.go | 256 + internal/logging/request_logger_format.go | 720 ++ internal/logging/request_logger_home.go | 246 + internal/logging/request_logger_streaming.go | 380 + internal/logging/request_logger_writer.go | 413 + internal/pluginhost/adapters.go | 1884 ----- internal/pluginhost/adapters_auth.go | 149 + internal/pluginhost/adapters_executors.go | 922 ++ internal/pluginhost/adapters_interceptors.go | 492 ++ .../pluginhost/adapters_usage_translation.go | 369 + .../runtime/executor/antigravity_executor.go | 2701 ------ .../executor/antigravity_executor_auth.go | 320 + .../executor/antigravity_executor_credits.go | 795 ++ .../executor/antigravity_executor_execute.go | 738 ++ .../executor/antigravity_executor_request.go | 449 + .../executor/antigravity_executor_stream.go | 319 + .../executor/antigravity_executor_tokens.go | 188 + internal/runtime/executor/claude_executor.go | 2497 ------ .../runtime/executor/claude_executor_auth.go | 49 + .../executor/claude_executor_cloaking.go | 960 +++ .../executor/claude_executor_execute.go | 213 + .../executor/claude_executor_request.go | 908 ++ .../executor/claude_executor_stream.go | 323 + .../executor/claude_executor_tokens.go | 135 + internal/runtime/executor/codex_executor.go | 2350 +----- .../runtime/executor/codex_executor_auth.go | 110 + .../executor/codex_executor_execute.go | 295 + .../executor/codex_executor_reasoning.go | 788 ++ .../executor/codex_executor_request.go | 482 ++ .../runtime/executor/codex_executor_stream.go | 217 + .../executor/codex_executor_terminal.go | 373 + .../runtime/executor/codex_executor_tokens.go | 175 + .../executor/codex_websockets_connection.go | 240 + .../executor/codex_websockets_errors.go | 199 + .../executor/codex_websockets_execute.go | 323 + .../executor/codex_websockets_executor.go | 2211 ----- .../executor/codex_websockets_request.go | 323 + .../executor/codex_websockets_session.go | 788 ++ .../executor/codex_websockets_stream.go | 429 + internal/runtime/executor/xai_executor.go | 2848 ------- .../runtime/executor/xai_executor_auth.go | 76 + .../runtime/executor/xai_executor_execute.go | 382 + .../runtime/executor/xai_executor_media.go | 141 + .../runtime/executor/xai_executor_request.go | 1145 +++ .../runtime/executor/xai_executor_response.go | 881 ++ .../runtime/executor/xai_executor_stream.go | 174 + .../runtime/executor/xai_executor_tokens.go | 149 + sdk/api/handlers/handlers.go | 1783 ---- sdk/api/handlers/handlers_context.go | 159 + sdk/api/handlers/handlers_errors.go | 145 + sdk/api/handlers/handlers_execution.go | 296 + sdk/api/handlers/handlers_interceptors.go | 377 + sdk/api/handlers/handlers_routing.go | 336 + sdk/api/handlers/handlers_stream.go | 542 ++ .../openai/openai_responses_websocket.go | 1717 ---- .../openai_responses_websocket_forward.go | 552 ++ .../openai_responses_websocket_prewarm.go | 173 + .../openai_responses_websocket_requests.go | 506 ++ .../openai_responses_websocket_session.go | 237 + .../openai_responses_websocket_timeline.go | 317 + sdk/cliproxy/auth/conductor.go | 7379 ----------------- sdk/cliproxy/auth/conductor_cooldown.go | 1684 ++++ sdk/cliproxy/auth/conductor_execution.go | 1221 +++ sdk/cliproxy/auth/conductor_home.go | 1116 +++ sdk/cliproxy/auth/conductor_home_execution.go | 195 + sdk/cliproxy/auth/conductor_lifecycle.go | 262 + sdk/cliproxy/auth/conductor_models.go | 830 ++ sdk/cliproxy/auth/conductor_refresh.go | 510 ++ sdk/cliproxy/auth/conductor_selection.go | 1361 +++ sdk/cliproxy/auth/conductor_stream.go | 309 + sdk/cliproxy/service.go | 3536 -------- sdk/cliproxy/service_auth.go | 432 + sdk/cliproxy/service_config.go | 287 + sdk/cliproxy/service_executors.go | 458 + sdk/cliproxy/service_home.go | 743 ++ sdk/cliproxy/service_lifecycle.go | 369 + sdk/cliproxy/service_models.go | 987 +++ sdk/cliproxy/service_plugins.go | 357 + 98 files changed, 37959 insertions(+), 36839 deletions(-) create mode 100644 internal/api/handlers/management/auth_files_crud.go create mode 100644 internal/api/handlers/management/auth_files_fields.go create mode 100644 internal/api/handlers/management/auth_files_oauth_callback.go create mode 100644 internal/api/handlers/management/auth_files_provider_oauth.go create mode 100644 internal/api/server_keepalive.go create mode 100644 internal/api/server_management.go create mode 100644 internal/api/server_middleware.go create mode 100644 internal/api/server_options.go create mode 100644 internal/api/server_reload.go create mode 100644 internal/api/server_routes.go create mode 100644 internal/config/config_defaults.go create mode 100644 internal/config/config_load.go create mode 100644 internal/config/config_normalization.go create mode 100644 internal/config/config_types.go create mode 100644 internal/config/config_validation.go create mode 100644 internal/config/config_yaml.go create mode 100644 internal/logging/request_logger_body_source.go create mode 100644 internal/logging/request_logger_format.go create mode 100644 internal/logging/request_logger_home.go create mode 100644 internal/logging/request_logger_streaming.go create mode 100644 internal/logging/request_logger_writer.go create mode 100644 internal/pluginhost/adapters_auth.go create mode 100644 internal/pluginhost/adapters_executors.go create mode 100644 internal/pluginhost/adapters_interceptors.go create mode 100644 internal/pluginhost/adapters_usage_translation.go create mode 100644 internal/runtime/executor/antigravity_executor_auth.go create mode 100644 internal/runtime/executor/antigravity_executor_credits.go create mode 100644 internal/runtime/executor/antigravity_executor_execute.go create mode 100644 internal/runtime/executor/antigravity_executor_request.go create mode 100644 internal/runtime/executor/antigravity_executor_stream.go create mode 100644 internal/runtime/executor/antigravity_executor_tokens.go create mode 100644 internal/runtime/executor/claude_executor_auth.go create mode 100644 internal/runtime/executor/claude_executor_cloaking.go create mode 100644 internal/runtime/executor/claude_executor_execute.go create mode 100644 internal/runtime/executor/claude_executor_request.go create mode 100644 internal/runtime/executor/claude_executor_stream.go create mode 100644 internal/runtime/executor/claude_executor_tokens.go create mode 100644 internal/runtime/executor/codex_executor_auth.go create mode 100644 internal/runtime/executor/codex_executor_execute.go create mode 100644 internal/runtime/executor/codex_executor_reasoning.go create mode 100644 internal/runtime/executor/codex_executor_request.go create mode 100644 internal/runtime/executor/codex_executor_stream.go create mode 100644 internal/runtime/executor/codex_executor_terminal.go create mode 100644 internal/runtime/executor/codex_executor_tokens.go create mode 100644 internal/runtime/executor/codex_websockets_connection.go create mode 100644 internal/runtime/executor/codex_websockets_errors.go create mode 100644 internal/runtime/executor/codex_websockets_execute.go create mode 100644 internal/runtime/executor/codex_websockets_request.go create mode 100644 internal/runtime/executor/codex_websockets_session.go create mode 100644 internal/runtime/executor/codex_websockets_stream.go create mode 100644 internal/runtime/executor/xai_executor_auth.go create mode 100644 internal/runtime/executor/xai_executor_execute.go create mode 100644 internal/runtime/executor/xai_executor_media.go create mode 100644 internal/runtime/executor/xai_executor_request.go create mode 100644 internal/runtime/executor/xai_executor_response.go create mode 100644 internal/runtime/executor/xai_executor_stream.go create mode 100644 internal/runtime/executor/xai_executor_tokens.go create mode 100644 sdk/api/handlers/handlers_context.go create mode 100644 sdk/api/handlers/handlers_errors.go create mode 100644 sdk/api/handlers/handlers_execution.go create mode 100644 sdk/api/handlers/handlers_interceptors.go create mode 100644 sdk/api/handlers/handlers_routing.go create mode 100644 sdk/api/handlers/handlers_stream.go create mode 100644 sdk/api/handlers/openai/openai_responses_websocket_forward.go create mode 100644 sdk/api/handlers/openai/openai_responses_websocket_prewarm.go create mode 100644 sdk/api/handlers/openai/openai_responses_websocket_requests.go create mode 100644 sdk/api/handlers/openai/openai_responses_websocket_session.go create mode 100644 sdk/api/handlers/openai/openai_responses_websocket_timeline.go create mode 100644 sdk/cliproxy/auth/conductor_cooldown.go create mode 100644 sdk/cliproxy/auth/conductor_execution.go create mode 100644 sdk/cliproxy/auth/conductor_home.go create mode 100644 sdk/cliproxy/auth/conductor_home_execution.go create mode 100644 sdk/cliproxy/auth/conductor_lifecycle.go create mode 100644 sdk/cliproxy/auth/conductor_models.go create mode 100644 sdk/cliproxy/auth/conductor_refresh.go create mode 100644 sdk/cliproxy/auth/conductor_selection.go create mode 100644 sdk/cliproxy/auth/conductor_stream.go create mode 100644 sdk/cliproxy/service_auth.go create mode 100644 sdk/cliproxy/service_config.go create mode 100644 sdk/cliproxy/service_executors.go create mode 100644 sdk/cliproxy/service_home.go create mode 100644 sdk/cliproxy/service_lifecycle.go create mode 100644 sdk/cliproxy/service_models.go create mode 100644 sdk/cliproxy/service_plugins.go diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index 6e4c3bde0..80801d7e7 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -1,20 +1,11 @@ package management import ( - "bytes" - "context" - "crypto/sha256" - "encoding/hex" "encoding/json" "errors" "fmt" - "io" - "mime/multipart" - "net" - "net/http" "os" "path/filepath" - "runtime" "sort" "strconv" "strings" @@ -22,43 +13,16 @@ import ( "time" "github.com/gin-gonic/gin" - "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/antigravity" - "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude" "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex" - "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/kimi" - xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" - "github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" - "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer" - sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" ) var lastRefreshKeys = []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"} -const ( - anthropicCallbackPort = 54545 - codexCallbackPort = 1455 -) - -type callbackForwarder struct { - provider string - server *http.Server - done chan struct{} -} - -type codexOAuthService interface { - GenerateAuthURL(state string, pkceCodes *codex.PKCECodes) (string, error) - ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *codex.PKCECodes) (*codex.CodexAuthBundle, error) - CreateTokenStorage(bundle *codex.CodexAuthBundle) *codex.CodexTokenStorage -} - var ( callbackForwardersMu sync.Mutex callbackForwarders = make(map[int]*callbackForwarder) @@ -122,201 +86,6 @@ func parseLastRefreshValue(v any) (time.Time, bool) { return time.Time{}, false } -func isWebUIRequest(c *gin.Context) bool { - raw := strings.TrimSpace(c.Query("is_webui")) - if raw == "" { - return false - } - switch strings.ToLower(raw) { - case "1", "true", "yes", "on": - return true - default: - return false - } -} - -func startCallbackForwarder(port int, provider, targetBase string) (*callbackForwarder, error) { - callbackForwardersMu.Lock() - prev := callbackForwarders[port] - if prev != nil { - delete(callbackForwarders, port) - } - callbackForwardersMu.Unlock() - - if prev != nil { - stopForwarderInstance(port, prev) - } - - addr := fmt.Sprintf("0.0.0.0:%d", port) - ln, err := net.Listen("tcp", addr) - if err != nil { - return nil, fmt.Errorf("failed to listen on %s: %w", addr, err) - } - - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - target := targetBase - if raw := r.URL.RawQuery; raw != "" { - if strings.Contains(target, "?") { - target = target + "&" + raw - } else { - target = target + "?" + raw - } - } - w.Header().Set("Cache-Control", "no-store") - http.Redirect(w, r, target, http.StatusFound) - }) - - srv := &http.Server{ - Handler: handler, - ReadHeaderTimeout: 5 * time.Second, - WriteTimeout: 5 * time.Second, - } - done := make(chan struct{}) - - go func() { - if errServe := srv.Serve(ln); errServe != nil && !errors.Is(errServe, http.ErrServerClosed) { - log.WithError(errServe).Warnf("callback forwarder for %s stopped unexpectedly", provider) - } - close(done) - }() - - forwarder := &callbackForwarder{ - provider: provider, - server: srv, - done: done, - } - - callbackForwardersMu.Lock() - callbackForwarders[port] = forwarder - callbackForwardersMu.Unlock() - - log.Infof("callback forwarder for %s listening on %s", provider, addr) - - return forwarder, nil -} - -func stopCallbackForwarderInstance(port int, forwarder *callbackForwarder) { - if forwarder == nil { - return - } - callbackForwardersMu.Lock() - if current := callbackForwarders[port]; current == forwarder { - delete(callbackForwarders, port) - } - callbackForwardersMu.Unlock() - - stopForwarderInstance(port, forwarder) -} - -func stopForwarderInstance(port int, forwarder *callbackForwarder) { - if forwarder == nil || forwarder.server == nil { - return - } - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - - if err := forwarder.server.Shutdown(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) { - log.WithError(err).Warnf("failed to shut down callback forwarder on port %d", port) - } - - select { - case <-forwarder.done: - case <-time.After(2 * time.Second): - } - - log.Infof("callback forwarder on port %d stopped", port) -} - -func (h *Handler) managementCallbackURL(path string) (string, error) { - if h == nil || h.cfg == nil || h.cfg.Port <= 0 { - return "", fmt.Errorf("server port is not configured") - } - if !strings.HasPrefix(path, "/") { - path = "/" + path - } - scheme := "http" - if h.cfg.TLS.Enable { - scheme = "https" - } - return fmt.Sprintf("%s://127.0.0.1:%d%s", scheme, h.cfg.Port, path), nil -} - -func pluginAuthProviderFromPath(path string) (string, bool) { - path = strings.TrimSpace(path) - const prefix = "/v0/management/" - const suffix = "-auth-url" - if !strings.HasPrefix(path, prefix) || !strings.HasSuffix(path, suffix) { - return "", false - } - provider := strings.TrimSuffix(strings.TrimPrefix(path, prefix), suffix) - provider = strings.ToLower(strings.TrimSpace(provider)) - if provider == "" { - return "", false - } - for _, r := range provider { - switch { - case r >= 'a' && r <= 'z': - case r >= '0' && r <= '9': - case r == '-': - default: - return "", false - } - } - return provider, true -} - -func (h *Handler) ServePluginAuthURL(c *gin.Context) bool { - if h == nil || c == nil || c.Request == nil || c.Request.URL == nil { - return false - } - h.mu.Lock() - host := h.pluginHost - h.mu.Unlock() - if host == nil { - return false - } - provider, ok := pluginAuthProviderFromPath(c.Request.URL.Path) - if !ok || !host.HasAuthProvider(provider) { - return false - } - - ctx := PopulateAuthContext(context.Background(), c) - baseURL, errBaseURL := h.managementCallbackURL("/v0/management/oauth-callback") - if errBaseURL != nil { - log.WithError(errBaseURL).Error("failed to compute plugin auth callback URL") - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) - return true - } - resp, handled, errStart := host.StartLogin(ctx, provider, baseURL) - if !handled { - return false - } - if errStart != nil { - log.WithError(errStart).Error("failed to start plugin auth login") - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) - return true - } - state := strings.TrimSpace(resp.State) - if state == "" { - log.WithField("provider", provider).Error("plugin auth provider returned empty state") - c.JSON(http.StatusBadGateway, gin.H{"error": "invalid oauth state"}) - return true - } - if errState := ValidateOAuthState(state); errState != nil { - log.WithError(errState).WithField("provider", provider).Error("plugin auth provider returned invalid state") - c.JSON(http.StatusBadGateway, gin.H{"error": "invalid oauth state"}) - return true - } - if errRegister := RegisterPluginOAuthSession(state, provider, resp.Metadata); errRegister != nil { - log.WithError(errRegister).WithField("provider", provider).Error("failed to register plugin oauth session") - c.JSON(http.StatusBadGateway, gin.H{"error": "failed to generate authorization url"}) - return true - } - c.JSON(http.StatusOK, gin.H{"status": "ok", "url": resp.URL, "state": state}) - return true -} - func (h *Handler) ListAuthFiles(c *gin.Context) { if h == nil { c.JSON(500, gin.H{"error": "handler not initialized"}) @@ -776,2035 +545,3 @@ func isUnsafeAuthFileName(name string) bool { } return false } - -// Download single auth file by name -func (h *Handler) DownloadAuthFile(c *gin.Context) { - name := strings.TrimSpace(c.Query("name")) - if isUnsafeAuthFileName(name) { - c.JSON(400, gin.H{"error": "invalid name"}) - return - } - if !strings.HasSuffix(strings.ToLower(name), ".json") { - c.JSON(400, gin.H{"error": "name must end with .json"}) - return - } - full := filepath.Join(h.cfg.AuthDir, name) - data, err := os.ReadFile(full) - if err != nil { - if os.IsNotExist(err) { - c.JSON(404, gin.H{"error": "file not found"}) - } else { - c.JSON(500, gin.H{"error": fmt.Sprintf("failed to read file: %v", err)}) - } - return - } - c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", name)) - c.Data(200, "application/json", data) -} - -// Upload auth file: multipart or raw JSON with ?name= -func (h *Handler) UploadAuthFile(c *gin.Context) { - if h.authManager == nil { - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) - return - } - ctx := c.Request.Context() - - fileHeaders, errMultipart := h.multipartAuthFileHeaders(c) - if errMultipart != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("invalid multipart form: %v", errMultipart)}) - return - } - if len(fileHeaders) == 1 { - if _, errUpload := h.storeUploadedAuthFile(ctx, fileHeaders[0]); errUpload != nil { - if errors.Is(errUpload, errAuthFileMustBeJSON) { - c.JSON(http.StatusBadRequest, gin.H{"error": "file must be .json"}) - return - } - c.JSON(http.StatusInternalServerError, gin.H{"error": errUpload.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"status": "ok"}) - return - } - if len(fileHeaders) > 1 { - uploaded := make([]string, 0, len(fileHeaders)) - failed := make([]gin.H, 0) - for _, file := range fileHeaders { - name, errUpload := h.storeUploadedAuthFile(ctx, file) - if errUpload != nil { - failureName := "" - if file != nil { - failureName = filepath.Base(file.Filename) - } - msg := errUpload.Error() - if errors.Is(errUpload, errAuthFileMustBeJSON) { - msg = "file must be .json" - } - failed = append(failed, gin.H{"name": failureName, "error": msg}) - continue - } - uploaded = append(uploaded, name) - } - if len(failed) > 0 { - c.JSON(http.StatusMultiStatus, gin.H{ - "status": "partial", - "uploaded": len(uploaded), - "files": uploaded, - "failed": failed, - }) - return - } - c.JSON(http.StatusOK, gin.H{"status": "ok", "uploaded": len(uploaded), "files": uploaded}) - return - } - if c.ContentType() == "multipart/form-data" { - c.JSON(http.StatusBadRequest, gin.H{"error": "no files uploaded"}) - return - } - name := strings.TrimSpace(c.Query("name")) - if isUnsafeAuthFileName(name) { - c.JSON(400, gin.H{"error": "invalid name"}) - return - } - if !strings.HasSuffix(strings.ToLower(name), ".json") { - c.JSON(400, gin.H{"error": "name must end with .json"}) - return - } - data, err := io.ReadAll(c.Request.Body) - if err != nil { - c.JSON(400, gin.H{"error": "failed to read body"}) - return - } - if err = h.writeAuthFile(ctx, filepath.Base(name), data); err != nil { - c.JSON(500, gin.H{"error": err.Error()}) - return - } - c.JSON(200, gin.H{"status": "ok"}) -} - -// Delete auth files: single by name or all -func (h *Handler) DeleteAuthFile(c *gin.Context) { - if h.authManager == nil { - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) - return - } - ctx := c.Request.Context() - if all := c.Query("all"); all == "true" || all == "1" || all == "*" { - entries, err := os.ReadDir(h.cfg.AuthDir) - if err != nil { - c.JSON(500, gin.H{"error": fmt.Sprintf("failed to read auth dir: %v", err)}) - return - } - deleted := 0 - for _, e := range entries { - if e.IsDir() { - continue - } - name := e.Name() - if !strings.HasSuffix(strings.ToLower(name), ".json") { - continue - } - full := filepath.Join(h.cfg.AuthDir, name) - if !filepath.IsAbs(full) { - if abs, errAbs := filepath.Abs(full); errAbs == nil { - full = abs - } - } - if err = os.Remove(full); err == nil { - if errDel := h.deleteTokenRecord(ctx, full); errDel != nil { - c.JSON(500, gin.H{"error": errDel.Error()}) - return - } - deleted++ - h.removeAuth(ctx, full) - } - } - c.JSON(200, gin.H{"status": "ok", "deleted": deleted}) - return - } - - names, errNames := requestedAuthFileNamesForDelete(c) - if errNames != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": errNames.Error()}) - return - } - if len(names) == 0 { - c.JSON(400, gin.H{"error": "invalid name"}) - return - } - if len(names) == 1 { - if _, status, errDelete := h.deleteAuthFileByName(ctx, names[0]); errDelete != nil { - c.JSON(status, gin.H{"error": errDelete.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"status": "ok"}) - return - } - - deletedFiles := make([]string, 0, len(names)) - failed := make([]gin.H, 0) - for _, name := range names { - deletedName, _, errDelete := h.deleteAuthFileByName(ctx, name) - if errDelete != nil { - failed = append(failed, gin.H{"name": name, "error": errDelete.Error()}) - continue - } - deletedFiles = append(deletedFiles, deletedName) - } - if len(failed) > 0 { - c.JSON(http.StatusMultiStatus, gin.H{ - "status": "partial", - "deleted": len(deletedFiles), - "files": deletedFiles, - "failed": failed, - }) - return - } - c.JSON(http.StatusOK, gin.H{"status": "ok", "deleted": len(deletedFiles), "files": deletedFiles}) -} - -func (h *Handler) multipartAuthFileHeaders(c *gin.Context) ([]*multipart.FileHeader, error) { - if h == nil || c == nil || c.ContentType() != "multipart/form-data" { - return nil, nil - } - form, err := c.MultipartForm() - if err != nil { - return nil, err - } - if form == nil || len(form.File) == 0 { - return nil, nil - } - - keys := make([]string, 0, len(form.File)) - for key := range form.File { - keys = append(keys, key) - } - sort.Strings(keys) - - headers := make([]*multipart.FileHeader, 0) - for _, key := range keys { - headers = append(headers, form.File[key]...) - } - return headers, nil -} - -func (h *Handler) storeUploadedAuthFile(ctx context.Context, file *multipart.FileHeader) (string, error) { - if file == nil { - return "", fmt.Errorf("no file uploaded") - } - name := filepath.Base(strings.TrimSpace(file.Filename)) - if !strings.HasSuffix(strings.ToLower(name), ".json") { - return "", errAuthFileMustBeJSON - } - src, err := file.Open() - if err != nil { - return "", fmt.Errorf("failed to open uploaded file: %w", err) - } - defer src.Close() - - data, err := io.ReadAll(src) - if err != nil { - return "", fmt.Errorf("failed to read uploaded file: %w", err) - } - if err := h.writeAuthFile(ctx, name, data); err != nil { - return "", err - } - return name, nil -} - -func (h *Handler) writeAuthFile(ctx context.Context, name string, data []byte) error { - dst := filepath.Join(h.cfg.AuthDir, filepath.Base(name)) - if !filepath.IsAbs(dst) { - if abs, errAbs := filepath.Abs(dst); errAbs == nil { - dst = abs - } - } - auth, err := h.buildAuthFromFileData(dst, data) - if err != nil { - return err - } - if errWrite := os.WriteFile(dst, data, 0o600); errWrite != nil { - return fmt.Errorf("failed to write file: %w", errWrite) - } - if err := h.upsertAuthRecord(ctx, auth); err != nil { - return err - } - return nil -} - -func requestedAuthFileNamesForDelete(c *gin.Context) ([]string, error) { - if c == nil { - return nil, nil - } - names := uniqueAuthFileNames(c.QueryArray("name")) - if len(names) > 0 { - return names, nil - } - - body, err := io.ReadAll(c.Request.Body) - if err != nil { - return nil, fmt.Errorf("failed to read body") - } - body = bytes.TrimSpace(body) - if len(body) == 0 { - return nil, nil - } - - var objectBody struct { - Name string `json:"name"` - Names []string `json:"names"` - } - if body[0] == '[' { - var arrayBody []string - if err := json.Unmarshal(body, &arrayBody); err != nil { - return nil, fmt.Errorf("invalid request body") - } - return uniqueAuthFileNames(arrayBody), nil - } - if err := json.Unmarshal(body, &objectBody); err != nil { - return nil, fmt.Errorf("invalid request body") - } - - out := make([]string, 0, len(objectBody.Names)+1) - if strings.TrimSpace(objectBody.Name) != "" { - out = append(out, objectBody.Name) - } - out = append(out, objectBody.Names...) - return uniqueAuthFileNames(out), nil -} - -func uniqueAuthFileNames(names []string) []string { - if len(names) == 0 { - return nil - } - seen := make(map[string]struct{}, len(names)) - out := make([]string, 0, len(names)) - for _, name := range names { - name = strings.TrimSpace(name) - if name == "" { - continue - } - if _, ok := seen[name]; ok { - continue - } - seen[name] = struct{}{} - out = append(out, name) - } - return out -} - -func (h *Handler) deleteAuthFileByName(ctx context.Context, name string) (string, int, error) { - name = strings.TrimSpace(name) - if isUnsafeAuthFileName(name) { - return "", http.StatusBadRequest, fmt.Errorf("invalid name") - } - - targetPath := filepath.Join(h.cfg.AuthDir, filepath.Base(name)) - targetID := "" - if targetAuth := h.findAuthForDelete(name); targetAuth != nil { - if !isPluginVirtualSourceDelete(name, targetAuth) { - return filepath.Base(name), http.StatusConflict, errPluginVirtualAuth - } - targetID = strings.TrimSpace(targetAuth.ID) - if path := strings.TrimSpace(authAttribute(targetAuth, "path")); path != "" { - targetPath = path - } - } - if !filepath.IsAbs(targetPath) { - if abs, errAbs := filepath.Abs(targetPath); errAbs == nil { - targetPath = abs - } - } - if errRemove := os.Remove(targetPath); errRemove != nil { - if os.IsNotExist(errRemove) { - return filepath.Base(name), http.StatusNotFound, errAuthFileNotFound - } - return filepath.Base(name), http.StatusInternalServerError, fmt.Errorf("failed to remove file: %w", errRemove) - } - if errDeleteRecord := h.deleteTokenRecord(ctx, targetPath); errDeleteRecord != nil { - return filepath.Base(name), http.StatusInternalServerError, errDeleteRecord - } - h.removeAuthsForPath(ctx, targetPath, targetID) - return filepath.Base(name), http.StatusOK, nil -} - -func isPluginVirtualSourceDelete(name string, auth *coreauth.Auth) bool { - if !coreauth.IsPluginVirtualAuth(auth) { - return true - } - sourcePath := strings.TrimSpace(authAttribute(auth, coreauth.AttributeVirtualSource)) - if sourcePath == "" { - sourcePath = strings.TrimSpace(authAttribute(auth, "path")) - } - if sourcePath == "" { - return false - } - return strings.EqualFold(filepath.Base(strings.TrimSpace(name)), filepath.Base(sourcePath)) -} - -func (h *Handler) findAuthForDelete(name string) *coreauth.Auth { - if h == nil || h.authManager == nil { - return nil - } - name = strings.TrimSpace(name) - if name == "" { - return nil - } - if auth, ok := h.authManager.GetByID(name); ok { - return auth - } - auths := h.authManager.List() - for _, auth := range auths { - if auth == nil { - continue - } - if strings.TrimSpace(auth.FileName) == name { - return auth - } - if filepath.Base(strings.TrimSpace(authAttribute(auth, "path"))) == name { - return auth - } - } - return nil -} - -func (h *Handler) authIDForPath(path string) string { - path = strings.TrimSpace(path) - if path == "" { - return "" - } - path = filepath.Clean(path) - if !filepath.IsAbs(path) { - if abs, errAbs := filepath.Abs(path); errAbs == nil { - path = abs - } - } - id := path - if h != nil && h.cfg != nil { - authDir := strings.TrimSpace(h.cfg.AuthDir) - if resolvedAuthDir, errResolve := util.ResolveAuthDir(authDir); errResolve == nil && resolvedAuthDir != "" { - authDir = resolvedAuthDir - } - if authDir != "" { - authDir = filepath.Clean(authDir) - if !filepath.IsAbs(authDir) { - if abs, errAbs := filepath.Abs(authDir); errAbs == nil { - authDir = abs - } - } - if rel, errRel := filepath.Rel(authDir, path); errRel == nil && rel != "" { - id = rel - } - } - } - // On Windows, normalize ID casing to avoid duplicate auth entries caused by case-insensitive paths. - if runtime.GOOS == "windows" { - id = strings.ToLower(id) - } - return id -} - -func (h *Handler) registerAuthFromFile(ctx context.Context, path string, data []byte) error { - if h.authManager == nil { - return nil - } - auth, err := h.buildAuthFromFileData(path, data) - if err != nil { - return err - } - return h.upsertAuthRecord(ctx, auth) -} - -func (h *Handler) buildAuthFromFileData(path string, data []byte) (*coreauth.Auth, error) { - if path == "" { - return nil, fmt.Errorf("auth path is empty") - } - if data == nil { - var err error - data, err = os.ReadFile(path) - if err != nil { - return nil, fmt.Errorf("failed to read auth file: %w", err) - } - } - metadata := make(map[string]any) - if err := json.Unmarshal(data, &metadata); err != nil { - return nil, fmt.Errorf("invalid auth file: %w", err) - } - provider, _ := metadata["type"].(string) - if provider == "" { - provider = "unknown" - } - label := provider - if email, ok := metadata["email"].(string); ok && email != "" { - label = email - } - lastRefresh, hasLastRefresh := extractLastRefreshTimestamp(metadata) - - authID := h.authIDForPath(path) - if authID == "" { - authID = path - } - auth := (*coreauth.Auth)(nil) - if h != nil && h.cfg != nil { - sctx := &synthesizer.SynthesisContext{ - Config: h.cfg, - AuthDir: h.cfg.AuthDir, - Now: time.Now(), - IDGenerator: synthesizer.NewStableIDGenerator(), - } - if generated := synthesizer.SynthesizeAuthFile(sctx, path, data); len(generated) > 0 && generated[0] != nil { - auth = generated[0].Clone() - } - } - if auth == nil { - auth = &coreauth.Auth{ - ID: authID, - Provider: provider, - Label: label, - Status: coreauth.StatusActive, - Attributes: map[string]string{ - "path": path, - "source": path, - }, - Metadata: metadata, - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - } - } - auth.ID = authID - auth.FileName = filepath.Base(path) - if hasLastRefresh { - auth.LastRefreshedAt = lastRefresh - } - if h != nil && h.authManager != nil { - if existing, ok := h.authManager.GetByID(authID); ok { - auth.CreatedAt = existing.CreatedAt - if !hasLastRefresh { - auth.LastRefreshedAt = existing.LastRefreshedAt - } - auth.NextRefreshAfter = existing.NextRefreshAfter - auth.Runtime = existing.Runtime - } - } - coreauth.ApplyCustomHeadersFromMetadata(auth) - return auth, nil -} - -func (h *Handler) upsertAuthRecord(ctx context.Context, auth *coreauth.Auth) error { - if h == nil || h.authManager == nil || auth == nil { - return nil - } - if existing, ok := h.authManager.GetByID(auth.ID); ok { - auth.CreatedAt = existing.CreatedAt - _, err := h.authManager.Update(ctx, auth) - return err - } - _, err := h.authManager.Register(ctx, auth) - return err -} - -// PatchAuthFileStatus toggles the disabled state of an auth file -func (h *Handler) PatchAuthFileStatus(c *gin.Context) { - if h.authManager == nil { - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) - return - } - - var req struct { - Name string `json:"name"` - AuthIndex string `json:"auth_index"` - Disabled *bool `json:"disabled"` - } - if err := c.ShouldBindJSON(&req); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) - return - } - - name := strings.TrimSpace(req.Name) - authIndex := strings.TrimSpace(req.AuthIndex) - if name == "" { - c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) - return - } - if req.Disabled == nil { - c.JSON(http.StatusBadRequest, gin.H{"error": "disabled is required"}) - return - } - - ctx := c.Request.Context() - - targetAuth, _ := h.lookupAuthFile(name, authIndex) - if targetAuth == nil { - c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"}) - return - } - if coreauth.IsPluginVirtualAuth(targetAuth) { - // Allow status changes only when targeting the source auth file name, matching delete semantics. - // Expanded virtual project auths still cannot be modified independently. - if !isPluginVirtualSourceDelete(name, targetAuth) { - c.JSON(http.StatusConflict, gin.H{"error": errPluginVirtualAuth.Error()}) - return - } - if errPatch := h.patchPluginVirtualSourceStatus(ctx, targetAuth, *req.Disabled); errPatch != nil { - status := http.StatusInternalServerError - if errors.Is(errPatch, errAuthFileNotFound) || os.IsNotExist(errPatch) { - status = http.StatusNotFound - } - c.JSON(status, gin.H{"error": errPatch.Error()}) - return - } - c.JSON(http.StatusOK, gin.H{"status": "ok", "disabled": *req.Disabled}) - return - } - - if coreauth.IsConfigAPIKeyAuth(targetAuth) { - h.mu.Lock() - handled, errToggle := toggleConfigAPIKeyExcludedAll(h.cfg, targetAuth, *req.Disabled) - if errToggle != nil { - h.mu.Unlock() - c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update config api key: %v", errToggle)}) - return - } - if !handled { - h.mu.Unlock() - c.JSON(http.StatusNotFound, gin.H{"error": "config api key entry not found"}) - return - } - cfgSnapshot, okSnapshot := h.saveConfigAndSnapshotLocked(c) - h.mu.Unlock() - if !okSnapshot { - return - } - h.reloadConfigAfterManagementSave(ctx, cfgSnapshot) - if h.tokenStore != nil { - _ = h.tokenStore.Delete(ctx, targetAuth.ID) - } - c.JSON(http.StatusOK, gin.H{ - "status": "ok", - "disabled": *req.Disabled, - "via": "config:excluded-models", - "excluded_pattern": configAPIKeyDisablePattern, - }) - return - } - - applyAuthDisabledState(targetAuth, *req.Disabled) - if _, err := h.authManager.Update(ctx, targetAuth); err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update auth: %v", err)}) - return - } - - c.JSON(http.StatusOK, gin.H{"status": "ok", "disabled": *req.Disabled}) -} - -// patchPluginVirtualSourceStatus toggles disabled on a plugin multi-auth source file and all -// runtime auths expanded from it. Virtual project children cannot be toggled independently. -func (h *Handler) patchPluginVirtualSourceStatus(ctx context.Context, targetAuth *coreauth.Auth, disabled bool) error { - if h == nil || h.authManager == nil || targetAuth == nil { - return fmt.Errorf("core auth manager unavailable") - } - sourcePath := strings.TrimSpace(authAttribute(targetAuth, coreauth.AttributeVirtualSource)) - if sourcePath == "" { - sourcePath = strings.TrimSpace(authAttribute(targetAuth, "path")) - } - if sourcePath == "" { - return errPluginVirtualAuth - } - if errWrite := setSourceAuthFileDisabled(sourcePath, disabled); errWrite != nil { - if os.IsNotExist(errWrite) { - return errAuthFileNotFound - } - return fmt.Errorf("failed to update source auth file: %w", errWrite) - } - now := time.Now() - for _, auth := range h.authManager.List() { - if auth == nil { - continue - } - if !sameAuthFilePath(authAttribute(auth, "path"), sourcePath) && - !sameAuthFilePath(authAttribute(auth, coreauth.AttributeVirtualSource), sourcePath) { - continue - } - applyAuthDisabledState(auth, disabled) - auth.UpdatedAt = now - if _, errUpdate := h.authManager.Update(ctx, auth); errUpdate != nil { - return fmt.Errorf("failed to update auth %s: %w", auth.ID, errUpdate) - } - } - return nil -} - -func setSourceAuthFileDisabled(path string, disabled bool) error { - path = strings.TrimSpace(path) - if path == "" { - return fmt.Errorf("source auth path is empty") - } - data, errRead := os.ReadFile(path) - if errRead != nil { - return errRead - } - metadata := make(map[string]any) - if len(bytes.TrimSpace(data)) > 0 { - if errUnmarshal := json.Unmarshal(data, &metadata); errUnmarshal != nil { - return fmt.Errorf("invalid auth file: %w", errUnmarshal) - } - } - if metadata == nil { - metadata = make(map[string]any) - } - metadata["disabled"] = disabled - raw, errMarshal := json.Marshal(metadata) - if errMarshal != nil { - return fmt.Errorf("marshal auth file: %w", errMarshal) - } - if errWrite := os.WriteFile(path, raw, 0o600); errWrite != nil { - return errWrite - } - return nil -} - -func applyAuthDisabledState(auth *coreauth.Auth, disabled bool) { - if auth == nil { - return - } - auth.Disabled = disabled - if disabled { - auth.Status = coreauth.StatusDisabled - auth.StatusMessage = "disabled via management API" - } else { - auth.Status = coreauth.StatusActive - auth.StatusMessage = "" - } - auth.UpdatedAt = time.Now() - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["disabled"] = disabled -} - -// PatchAuthFileFields updates arbitrary metadata fields of an auth file. -func (h *Handler) PatchAuthFileFields(c *gin.Context) { - if h.authManager == nil { - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) - return - } - - var req map[string]json.RawMessage - decoder := json.NewDecoder(c.Request.Body) - decoder.UseNumber() - if err := decoder.Decode(&req); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) - return - } - - nameRaw, ok := req["name"] - if !ok { - c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) - return - } - var nameValue string - if err := json.Unmarshal(nameRaw, &nameValue); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) - return - } - name := strings.TrimSpace(nameValue) - if name == "" { - c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) - return - } - delete(req, "name") - - ctx := c.Request.Context() - - // Find auth by name or ID - var targetAuth *coreauth.Auth - if auth, ok := h.authManager.GetByID(name); ok { - targetAuth = auth - } else { - auths := h.authManager.List() - for _, auth := range auths { - if auth.FileName == name { - targetAuth = auth - break - } - } - } - - if targetAuth == nil { - c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"}) - return - } - if coreauth.IsPluginVirtualAuth(targetAuth) { - c.JSON(http.StatusConflict, gin.H{"error": errPluginVirtualAuth.Error()}) - return - } - - changed := false - touchedRoots := make(map[string]struct{}, len(req)) - for key, rawValue := range req { - fieldPath := strings.TrimSpace(key) - if fieldPath == "" { - c.JSON(http.StatusBadRequest, gin.H{"error": "field name is required"}) - return - } - value, errDecode := decodeAuthFileFieldValue(rawValue) - if errDecode != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("invalid field %s", fieldPath)}) - return - } - if targetAuth.Metadata == nil { - targetAuth.Metadata = make(map[string]any) - } - - if fieldPath == "headers" { - applyAuthFileHeadersPatch(targetAuth, value) - } else if errSet := setAuthFileMetadataValue(targetAuth.Metadata, fieldPath, value); errSet != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": errSet.Error()}) - return - } - if root := rootAuthFileField(fieldPath); root != "" { - touchedRoots[root] = struct{}{} - } - changed = true - } - if changed { - syncAuthFileMetadataFields(targetAuth, touchedRoots) - } - - if !changed { - c.JSON(http.StatusBadRequest, gin.H{"error": "no fields to update"}) - return - } - - targetAuth.UpdatedAt = time.Now() - - if _, err := h.authManager.Update(ctx, targetAuth); err != nil { - c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update auth: %v", err)}) - return - } - - c.JSON(http.StatusOK, gin.H{"status": "ok"}) -} - -func decodeAuthFileFieldValue(raw json.RawMessage) (any, error) { - decoder := json.NewDecoder(bytes.NewReader(raw)) - decoder.UseNumber() - var value any - if err := decoder.Decode(&value); err != nil { - return nil, err - } - return value, nil -} - -func rootAuthFileField(path string) string { - path = strings.TrimSpace(path) - if path == "" { - return "" - } - if idx := strings.Index(path, "."); idx >= 0 { - return strings.TrimSpace(path[:idx]) - } - return path -} - -func setAuthFileMetadataValue(metadata map[string]any, path string, value any) error { - if metadata == nil { - return fmt.Errorf("metadata is nil") - } - parts := strings.Split(path, ".") - current := metadata - for i, rawPart := range parts { - part := strings.TrimSpace(rawPart) - if part == "" { - return fmt.Errorf("invalid field path: %s", path) - } - if i == len(parts)-1 { - current[part] = value - return nil - } - next, ok := current[part].(map[string]any) - if !ok { - next = make(map[string]any) - current[part] = next - } - current = next - } - return nil -} - -func applyAuthFileHeadersPatch(auth *coreauth.Auth, value any) { - if auth == nil { - return - } - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - headersPatch, ok := authFileHeadersStringMap(value) - if !ok { - auth.Metadata["headers"] = value - return - } - - existingHeaders := coreauth.ExtractCustomHeadersFromMetadata(auth.Metadata) - nextHeaders := make(map[string]string, len(existingHeaders)) - for key, val := range existingHeaders { - nextHeaders[key] = val - } - for key, value := range headersPatch { - name := strings.TrimSpace(key) - if name == "" { - continue - } - val := strings.TrimSpace(value) - if val == "" { - delete(nextHeaders, name) - continue - } - nextHeaders[name] = val - } - - if len(nextHeaders) == 0 { - delete(auth.Metadata, "headers") - return - } - metaHeaders := make(map[string]any, len(nextHeaders)) - for key, value := range nextHeaders { - metaHeaders[key] = value - } - auth.Metadata["headers"] = metaHeaders -} - -func authFileHeadersStringMap(value any) (map[string]string, bool) { - switch typed := value.(type) { - case map[string]string: - return typed, true - case map[string]any: - out := make(map[string]string, len(typed)) - for key, rawValue := range typed { - value, ok := rawValue.(string) - if !ok { - return nil, false - } - out[key] = value - } - return out, true - default: - return nil, false - } -} - -func syncAuthFileMetadataFields(auth *coreauth.Auth, touchedRoots map[string]struct{}) { - if auth == nil || len(touchedRoots) == 0 { - return - } - if _, ok := touchedRoots["prefix"]; ok { - if prefix, okString := auth.Metadata["prefix"].(string); okString { - auth.Prefix = strings.TrimSpace(prefix) - } - } - if _, ok := touchedRoots["proxy_url"]; ok { - if proxyURL, okString := auth.Metadata["proxy_url"].(string); okString { - auth.ProxyURL = strings.TrimSpace(proxyURL) - } - } - if _, ok := touchedRoots["headers"]; ok { - syncAuthFileHeaderAttributes(auth) - } - if _, ok := touchedRoots["priority"]; ok { - syncAuthFilePriorityAttribute(auth) - } - if _, ok := touchedRoots["note"]; ok { - syncAuthFileNoteAttribute(auth) - } - if _, ok := touchedRoots["websockets"]; ok { - syncAuthFileWebsocketsAttribute(auth) - } - if _, ok := touchedRoots["disabled"]; ok { - syncAuthFileDisabledState(auth) - } -} - -func syncAuthFileHeaderAttributes(auth *coreauth.Auth) { - if auth == nil { - return - } - if auth.Attributes == nil { - auth.Attributes = make(map[string]string) - } - for key := range auth.Attributes { - if strings.HasPrefix(key, "header:") { - delete(auth.Attributes, key) - } - } - for name, value := range coreauth.ExtractCustomHeadersFromMetadata(auth.Metadata) { - auth.Attributes["header:"+name] = value - } -} - -func syncAuthFilePriorityAttribute(auth *coreauth.Auth) { - if auth == nil { - return - } - if auth.Attributes == nil { - auth.Attributes = make(map[string]string) - } - priority, ok := authFileIntValue(auth.Metadata["priority"]) - if !ok { - delete(auth.Attributes, "priority") - return - } - if priority == 0 { - delete(auth.Attributes, "priority") - return - } - auth.Attributes["priority"] = strconv.Itoa(priority) -} - -func authFileIntValue(value any) (int, bool) { - switch typed := value.(type) { - case int: - return typed, true - case int64: - return int(typed), true - case float64: - return int(typed), true - case json.Number: - if i, err := typed.Int64(); err == nil { - return int(i), true - } - case string: - if i, err := strconv.Atoi(strings.TrimSpace(typed)); err == nil { - return i, true - } - } - return 0, false -} - -func syncAuthFileNoteAttribute(auth *coreauth.Auth) { - if auth == nil { - return - } - if auth.Attributes == nil { - auth.Attributes = make(map[string]string) - } - note, ok := auth.Metadata["note"].(string) - if !ok { - delete(auth.Attributes, "note") - return - } - note = strings.TrimSpace(note) - if note == "" { - delete(auth.Attributes, "note") - return - } - auth.Attributes["note"] = note -} - -func syncAuthFileWebsocketsAttribute(auth *coreauth.Auth) { - if auth == nil { - return - } - if auth.Attributes == nil { - auth.Attributes = make(map[string]string) - } - websockets, ok := authFileBoolValue(auth.Metadata["websockets"]) - if !ok { - delete(auth.Attributes, "websockets") - return - } - auth.Attributes["websockets"] = strconv.FormatBool(websockets) -} - -func authFileBoolValue(value any) (bool, bool) { - switch typed := value.(type) { - case bool: - return typed, true - case string: - parsed, errParse := strconv.ParseBool(strings.TrimSpace(typed)) - if errParse == nil { - return parsed, true - } - } - return false, false -} - -func syncAuthFileDisabledState(auth *coreauth.Auth) { - if auth == nil { - return - } - disabled, ok := authFileBoolValue(auth.Metadata["disabled"]) - if !ok { - return - } - auth.Disabled = disabled - if disabled { - auth.Status = coreauth.StatusDisabled - if strings.TrimSpace(auth.StatusMessage) == "" { - auth.StatusMessage = "disabled via management API" - } - return - } - auth.Status = coreauth.StatusActive - auth.StatusMessage = "" -} - -func (h *Handler) removeAuth(ctx context.Context, id string) { - if h == nil || h.authManager == nil { - return - } - id = strings.TrimSpace(id) - if id == "" { - return - } - if _, ok := h.authManager.GetByID(id); ok { - h.authManager.Remove(ctx, id) - return - } - authID := h.authIDForPath(id) - if authID == "" { - return - } - h.authManager.Remove(ctx, authID) -} - -func (h *Handler) removeAuthsForPath(ctx context.Context, path string, fallbackID string) { - if h == nil || h.authManager == nil { - return - } - removed := false - for _, auth := range h.authManager.List() { - if auth == nil { - continue - } - if sameAuthFilePath(authAttribute(auth, "path"), path) || sameAuthFilePath(authAttribute(auth, coreauth.AttributeVirtualSource), path) { - h.removeAuth(ctx, auth.ID) - removed = true - } - } - if removed { - return - } - if strings.TrimSpace(fallbackID) != "" { - h.removeAuth(ctx, fallbackID) - return - } - h.removeAuth(ctx, path) -} - -func sameAuthFilePath(left, right string) bool { - left = cleanAuthFilePath(left) - right = cleanAuthFilePath(right) - if left == "" || right == "" { - return false - } - if runtime.GOOS == "windows" { - return strings.EqualFold(left, right) - } - return left == right -} - -func cleanAuthFilePath(path string) string { - path = strings.TrimSpace(path) - if path == "" { - return "" - } - if abs, errAbs := filepath.Abs(path); errAbs == nil && strings.TrimSpace(abs) != "" { - path = abs - } - return filepath.Clean(path) -} - -func (h *Handler) deleteTokenRecord(ctx context.Context, path string) error { - if strings.TrimSpace(path) == "" { - return fmt.Errorf("auth path is empty") - } - store := h.tokenStoreWithBaseDir() - if store == nil { - return fmt.Errorf("token store unavailable") - } - return store.Delete(ctx, path) -} - -func (h *Handler) tokenStoreWithBaseDir() coreauth.Store { - if h == nil { - return nil - } - store := h.tokenStore - if store == nil { - store = sdkAuth.GetTokenStore() - h.tokenStore = store - } - if h.cfg != nil { - if dirSetter, ok := store.(interface{ SetBaseDir(string) }); ok { - dirSetter.SetBaseDir(h.cfg.AuthDir) - } - } - return store -} - -func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (string, error) { - if record == nil { - return "", fmt.Errorf("token record is nil") - } - store := h.tokenStoreWithBaseDir() - if store == nil { - return "", fmt.Errorf("token store unavailable") - } - if h.postAuthHook != nil { - if err := h.postAuthHook(ctx, record); err != nil { - return "", fmt.Errorf("post-auth hook failed: %w", err) - } - } - savedPath, errSave := store.Save(ctx, record) - if errSave != nil { - return savedPath, errSave - } - if h.postAuthPersistHook != nil { - if errHook := h.postAuthPersistHook(ctx, record); errHook != nil { - return savedPath, fmt.Errorf("post-auth persist hook failed: %w", errHook) - } - } - return savedPath, nil -} - -func (h *Handler) RequestAnthropicToken(c *gin.Context) { - ctx := context.Background() - ctx = PopulateAuthContext(ctx, c) - - fmt.Println("Initializing Claude authentication...") - - // Generate PKCE codes - pkceCodes, err := claude.GeneratePKCECodes() - if err != nil { - log.Errorf("Failed to generate PKCE codes: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"}) - return - } - - // Generate random state parameter - state, err := misc.GenerateRandomState() - if err != nil { - log.Errorf("Failed to generate state parameter: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"}) - return - } - - // Initialize Claude auth service - anthropicAuth := claude.NewClaudeAuth(h.cfg) - - // Generate authorization URL (then override redirect_uri to reuse server port) - authURL, state, err := anthropicAuth.GenerateAuthURL(state, pkceCodes) - if err != nil { - log.Errorf("Failed to generate authorization URL: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) - return - } - - RegisterOAuthSession(state, "anthropic") - - isWebUI := isWebUIRequest(c) - var forwarder *callbackForwarder - if isWebUI { - targetURL, errTarget := h.managementCallbackURL("/anthropic/callback") - if errTarget != nil { - log.WithError(errTarget).Error("failed to compute anthropic callback target") - c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"}) - return - } - var errStart error - if forwarder, errStart = startCallbackForwarder(anthropicCallbackPort, "anthropic", targetURL); errStart != nil { - log.WithError(errStart).Error("failed to start anthropic callback forwarder") - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"}) - return - } - } - - go func() { - if isWebUI { - defer stopCallbackForwarderInstance(anthropicCallbackPort, forwarder) - } - - // Helper: wait for callback file - waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-anthropic-%s.oauth", state)) - waitForFile := func(path string, timeout time.Duration) (map[string]string, error) { - deadline := time.Now().Add(timeout) - for { - if !IsOAuthSessionPending(state, "anthropic") { - return nil, errOAuthSessionNotPending - } - if time.Now().After(deadline) { - SetOAuthSessionError(state, "Timeout waiting for OAuth callback") - return nil, fmt.Errorf("timeout waiting for OAuth callback") - } - data, errRead := os.ReadFile(path) - if errRead == nil { - var m map[string]string - _ = json.Unmarshal(data, &m) - _ = os.Remove(path) - return m, nil - } - time.Sleep(500 * time.Millisecond) - } - } - - fmt.Println("Waiting for authentication callback...") - // Wait up to 5 minutes - resultMap, errWait := waitForFile(waitFile, 5*time.Minute) - if errWait != nil { - if errors.Is(errWait, errOAuthSessionNotPending) { - return - } - authErr := claude.NewAuthenticationError(claude.ErrCallbackTimeout, errWait) - log.Error(claude.GetUserFriendlyMessage(authErr)) - return - } - if errStr := resultMap["error"]; errStr != "" { - oauthErr := claude.NewOAuthError(errStr, "", http.StatusBadRequest) - log.Error(claude.GetUserFriendlyMessage(oauthErr)) - SetOAuthSessionError(state, "Bad request") - return - } - if resultMap["state"] != state { - authErr := claude.NewAuthenticationError(claude.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, resultMap["state"])) - log.Error(claude.GetUserFriendlyMessage(authErr)) - SetOAuthSessionError(state, "State code error") - return - } - - // Parse code (Claude may append state after '#') - rawCode := resultMap["code"] - code := strings.Split(rawCode, "#")[0] - - // Exchange code for tokens using internal auth service - bundle, errExchange := anthropicAuth.ExchangeCodeForTokens(ctx, code, state, pkceCodes) - if errExchange != nil { - authErr := claude.NewAuthenticationError(claude.ErrCodeExchangeFailed, errExchange) - log.Errorf("Failed to exchange authorization code for tokens: %v", authErr) - SetOAuthSessionError(state, "Failed to exchange authorization code for tokens") - return - } - - // Create token storage - tokenStorage := anthropicAuth.CreateTokenStorage(bundle) - record := &coreauth.Auth{ - ID: fmt.Sprintf("claude-%s.json", tokenStorage.Email), - Provider: "claude", - FileName: fmt.Sprintf("claude-%s.json", tokenStorage.Email), - Storage: tokenStorage, - Metadata: map[string]any{"email": tokenStorage.Email}, - } - if errGuard := guardOAuthSessionPendingForSave(state, "anthropic"); errGuard != nil { - return - } - savedPath, errSave := h.saveTokenRecord(ctx, record) - if errSave != nil { - log.Errorf("Failed to save authentication tokens: %v", errSave) - SetOAuthSessionError(state, "Failed to save authentication tokens") - return - } - - fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) - if bundle.APIKey != "" { - fmt.Println("API key obtained and saved") - } - fmt.Println("You can now use Claude services through this CLI") - CompleteOAuthSession(state) - }() - - c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) -} - -func (h *Handler) RequestCodexToken(c *gin.Context) { - ctx := context.Background() - ctx = PopulateAuthContext(ctx, c) - - fmt.Println("Initializing Codex authentication...") - - // Generate PKCE codes - pkceCodes, err := codex.GeneratePKCECodes() - if err != nil { - log.Errorf("Failed to generate PKCE codes: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"}) - return - } - - // Generate random state parameter - state, err := misc.GenerateRandomState() - if err != nil { - log.Errorf("Failed to generate state parameter: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"}) - return - } - - // Initialize Codex auth service - openaiAuth := newCodexOAuthService(h.cfg) - - // Generate authorization URL - authURL, err := openaiAuth.GenerateAuthURL(state, pkceCodes) - if err != nil { - log.Errorf("Failed to generate authorization URL: %v", err) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) - return - } - - RegisterOAuthSession(state, "codex") - - isWebUI := isWebUIRequest(c) - var forwarder *callbackForwarder - if isWebUI { - targetURL, errTarget := h.managementCallbackURL("/codex/callback") - if errTarget != nil { - log.WithError(errTarget).Error("failed to compute codex callback target") - c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"}) - return - } - var errStart error - if forwarder, errStart = startCallbackForwarder(codexCallbackPort, "codex", targetURL); errStart != nil { - log.WithError(errStart).Error("failed to start codex callback forwarder") - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"}) - return - } - } - - go func() { - if isWebUI { - defer stopCallbackForwarderInstance(codexCallbackPort, forwarder) - } - - // Wait for callback file - waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-codex-%s.oauth", state)) - deadline := time.Now().Add(5 * time.Minute) - var code string - for { - if !IsOAuthSessionPending(state, "codex") { - return - } - if time.Now().After(deadline) { - authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, fmt.Errorf("timeout waiting for OAuth callback")) - log.Error(codex.GetUserFriendlyMessage(authErr)) - SetOAuthSessionError(state, "Timeout waiting for OAuth callback") - return - } - if data, errR := os.ReadFile(waitFile); errR == nil { - var m map[string]string - _ = json.Unmarshal(data, &m) - _ = os.Remove(waitFile) - if errStr := m["error"]; errStr != "" { - oauthErr := codex.NewOAuthError(errStr, "", http.StatusBadRequest) - log.Error(codex.GetUserFriendlyMessage(oauthErr)) - SetOAuthSessionError(state, "Bad Request") - return - } - if m["state"] != state { - authErr := codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, m["state"])) - SetOAuthSessionError(state, "State code error") - log.Error(codex.GetUserFriendlyMessage(authErr)) - return - } - code = m["code"] - break - } - time.Sleep(500 * time.Millisecond) - } - - log.Debug("Authorization code received, exchanging for tokens...") - // Exchange code for tokens using internal auth service - bundle, errExchange := openaiAuth.ExchangeCodeForTokens(ctx, code, pkceCodes) - if errExchange != nil { - authErr := codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, errExchange) - SetOAuthSessionError(state, oauthSessionErrorWithCause("Failed to exchange authorization code for tokens", errExchange)) - log.Errorf("Failed to exchange authorization code for tokens: %v", authErr) - return - } - - // Extract additional info for filename generation - claims, _ := codex.ParseJWTToken(bundle.TokenData.IDToken) - planType := "" - hashAccountID := "" - if claims != nil { - planType = strings.TrimSpace(claims.CodexAuthInfo.ChatgptPlanType) - if accountID := claims.GetAccountID(); accountID != "" { - digest := sha256.Sum256([]byte(accountID)) - hashAccountID = hex.EncodeToString(digest[:])[:8] - } - } - - // Create token storage and persist - tokenStorage := openaiAuth.CreateTokenStorage(bundle) - fileName := codex.CredentialFileName(tokenStorage.Email, planType, hashAccountID, true) - record := &coreauth.Auth{ - ID: fileName, - Provider: "codex", - FileName: fileName, - Storage: tokenStorage, - Metadata: map[string]any{ - "email": tokenStorage.Email, - "account_id": tokenStorage.AccountID, - }, - } - if errGuard := guardOAuthSessionPendingForSave(state, "codex"); errGuard != nil { - return - } - savedPath, errSave := h.saveTokenRecord(ctx, record) - if errSave != nil { - SetOAuthSessionError(state, "Failed to save authentication tokens") - log.Errorf("Failed to save authentication tokens: %v", errSave) - return - } - fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) - if bundle.APIKey != "" { - fmt.Println("API key obtained and saved") - } - fmt.Println("You can now use Codex services through this CLI") - CompleteOAuthSession(state) - }() - - c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) -} - -func (h *Handler) RequestAntigravityToken(c *gin.Context) { - ctx := context.Background() - ctx = PopulateAuthContext(ctx, c) - - fmt.Println("Initializing Antigravity authentication...") - - authSvc := antigravity.NewAntigravityAuth(h.cfg, nil) - - state, errState := misc.GenerateRandomState() - if errState != nil { - log.Errorf("Failed to generate state parameter: %v", errState) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"}) - return - } - - redirectURI := fmt.Sprintf("http://localhost:%d/oauth-callback", antigravity.CallbackPort) - authURL := authSvc.BuildAuthURL(state, redirectURI) - - RegisterOAuthSession(state, "antigravity") - - isWebUI := isWebUIRequest(c) - var forwarder *callbackForwarder - if isWebUI { - targetURL, errTarget := h.managementCallbackURL("/antigravity/callback") - if errTarget != nil { - log.WithError(errTarget).Error("failed to compute antigravity callback target") - c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"}) - return - } - var errStart error - if forwarder, errStart = startCallbackForwarder(antigravity.CallbackPort, "antigravity", targetURL); errStart != nil { - log.WithError(errStart).Error("failed to start antigravity callback forwarder") - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"}) - return - } - } - - go func() { - if isWebUI { - defer stopCallbackForwarderInstance(antigravity.CallbackPort, forwarder) - } - - waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-antigravity-%s.oauth", state)) - deadline := time.Now().Add(5 * time.Minute) - var authCode string - for { - if !IsOAuthSessionPending(state, "antigravity") { - return - } - if time.Now().After(deadline) { - log.Error("oauth flow timed out") - SetOAuthSessionError(state, "OAuth flow timed out") - return - } - if data, errReadFile := os.ReadFile(waitFile); errReadFile == nil { - var payload map[string]string - _ = json.Unmarshal(data, &payload) - _ = os.Remove(waitFile) - if errStr := strings.TrimSpace(payload["error"]); errStr != "" { - log.Errorf("Authentication failed: %s", errStr) - SetOAuthSessionError(state, "Authentication failed") - return - } - if payloadState := strings.TrimSpace(payload["state"]); payloadState != "" && payloadState != state { - log.Errorf("Authentication failed: state mismatch") - SetOAuthSessionError(state, "Authentication failed: state mismatch") - return - } - authCode = strings.TrimSpace(payload["code"]) - if authCode == "" { - log.Error("Authentication failed: code not found") - SetOAuthSessionError(state, "Authentication failed: code not found") - return - } - break - } - time.Sleep(500 * time.Millisecond) - } - - tokenResp, errToken := authSvc.ExchangeCodeForTokens(ctx, authCode, redirectURI) - if errToken != nil { - log.Errorf("Failed to exchange token: %v", errToken) - SetOAuthSessionError(state, "Failed to exchange token") - return - } - - accessToken := strings.TrimSpace(tokenResp.AccessToken) - if accessToken == "" { - log.Error("antigravity: token exchange returned empty access token") - SetOAuthSessionError(state, "Failed to exchange token") - return - } - - email, errInfo := authSvc.FetchUserInfo(ctx, accessToken) - if errInfo != nil { - log.Errorf("Failed to fetch user info: %v", errInfo) - SetOAuthSessionError(state, "Failed to fetch user info") - return - } - email = strings.TrimSpace(email) - if email == "" { - log.Error("antigravity: user info returned empty email") - SetOAuthSessionError(state, "Failed to fetch user info") - return - } - - projectID := "" - if accessToken != "" { - fetchedProjectID, errProject := authSvc.FetchProjectID(ctx, accessToken) - if errProject != nil { - log.Warnf("antigravity: failed to fetch project ID: %v", errProject) - } else { - projectID = fetchedProjectID - log.Infof("antigravity: obtained project ID %s", util.HideAPIKey(projectID)) - } - } - - now := time.Now() - metadata := map[string]any{ - "type": "antigravity", - "access_token": tokenResp.AccessToken, - "refresh_token": tokenResp.RefreshToken, - "expires_in": tokenResp.ExpiresIn, - "timestamp": now.UnixMilli(), - "expired": now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339), - } - if email != "" { - metadata["email"] = email - } - if projectID != "" { - metadata["project_id"] = projectID - } - - fileName := antigravity.CredentialFileName(email) - label := strings.TrimSpace(email) - if label == "" { - label = "antigravity" - } - - record := &coreauth.Auth{ - ID: fileName, - Provider: "antigravity", - FileName: fileName, - Label: label, - Metadata: metadata, - } - if errGuard := guardOAuthSessionPendingForSave(state, "antigravity"); errGuard != nil { - return - } - savedPath, errSave := h.saveTokenRecord(ctx, record) - if errSave != nil { - log.Errorf("Failed to save token to file: %v", errSave) - SetOAuthSessionError(state, "Failed to save token to file") - return - } - - CompleteOAuthSession(state) - fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) - if projectID != "" { - fmt.Printf("Using GCP project: %s\n", util.HideAPIKey(projectID)) - } - fmt.Println("You can now use Antigravity services through this CLI") - }() - - c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) -} - -func (h *Handler) RequestXAIToken(c *gin.Context) { - ctx := context.Background() - ctx = PopulateAuthContext(ctx, c) - - fmt.Println("Initializing xAI authentication...") - - state := fmt.Sprintf("xai-%d", time.Now().UnixNano()) - authSvc := xaiauth.NewXAIAuth(h.cfg) - - deviceFlow, errStartDeviceFlow := authSvc.StartDeviceFlow(ctx) - if errStartDeviceFlow != nil { - log.Errorf("Failed to start xAI device flow: %v", errStartDeviceFlow) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start device authorization flow"}) - return - } - authURL := strings.TrimSpace(deviceFlow.VerificationURIComplete) - if authURL == "" { - authURL = strings.TrimSpace(deviceFlow.VerificationURI) - } - - RegisterOAuthSession(state, "xai") - - go func() { - pollCtx, cancelPoll := context.WithCancel(ctx) - defer cancelPoll() - go watchOAuthSessionCancel(pollCtx, cancelPoll, state, "xai") - - fmt.Println("Waiting for xAI authentication...") - bundle, errWaitForAuthorization := authSvc.WaitForAuthorization(pollCtx, deviceFlow) - if errWaitForAuthorization != nil { - if !IsOAuthSessionPending(state, "xai") { - return - } - log.Errorf("xAI authentication failed: %v", errWaitForAuthorization) - SetOAuthSessionError(state, oauthSessionErrorWithCause("Authentication failed", errWaitForAuthorization)) - return - } - if !IsOAuthSessionPending(state, "xai") { - return - } - - tokenStorage := authSvc.CreateTokenStorage(bundle) - if tokenStorage == nil || strings.TrimSpace(tokenStorage.AccessToken) == "" { - log.Error("xAI token exchange returned empty access token") - SetOAuthSessionError(state, "Failed to exchange token") - return - } - - fileName := xaiauth.CredentialFileName(tokenStorage.Email, tokenStorage.Subject) - label := strings.TrimSpace(tokenStorage.Email) - if label == "" { - label = "xAI" - } - - metadata := map[string]any{ - "type": "xai", - "access_token": tokenStorage.AccessToken, - "refresh_token": tokenStorage.RefreshToken, - "id_token": tokenStorage.IDToken, - "token_type": tokenStorage.TokenType, - "expires_in": tokenStorage.ExpiresIn, - "expired": tokenStorage.Expire, - "last_refresh": tokenStorage.LastRefresh, - "base_url": tokenStorage.BaseURL, - "token_endpoint": tokenStorage.TokenEndpoint, - "auth_kind": "oauth", - } - if tokenStorage.Email != "" { - metadata["email"] = tokenStorage.Email - } - if tokenStorage.Subject != "" { - metadata["sub"] = tokenStorage.Subject - } - - record := &coreauth.Auth{ - ID: fileName, - Provider: "xai", - FileName: fileName, - Label: label, - Storage: tokenStorage, - Metadata: metadata, - Attributes: map[string]string{ - "auth_kind": "oauth", - "base_url": tokenStorage.BaseURL, - }, - } - if errGuard := guardOAuthSessionPendingForSave(state, "xai"); errGuard != nil { - return - } - savedPath, errSave := h.saveTokenRecord(ctx, record) - if errSave != nil { - log.Errorf("Failed to save xAI token to file: %v", errSave) - SetOAuthSessionError(state, "Failed to save token to file") - return - } - - CompleteOAuthSession(state) - fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) - fmt.Println("You can now use xAI services through this CLI") - }() - - response := gin.H{"status": "ok", "url": authURL, "state": state, "flow": "device"} - if userCode := strings.TrimSpace(deviceFlow.UserCode); userCode != "" { - response["user_code"] = userCode - } - if deviceFlow.ExpiresIn > 0 { - response["expires_in"] = deviceFlow.ExpiresIn - } else { - response["expires_in"] = int(xaiauth.MaxPollDuration / time.Second) - } - c.JSON(200, response) -} - -func (h *Handler) RequestKimiToken(c *gin.Context) { - ctx := context.Background() - ctx = PopulateAuthContext(ctx, c) - - fmt.Println("Initializing Kimi authentication...") - - state := fmt.Sprintf("kmi-%d", time.Now().UnixNano()) - // Initialize Kimi auth service - kimiAuth := kimi.NewKimiAuth(h.cfg) - - // Generate authorization URL - deviceFlow, errStartDeviceFlow := kimiAuth.StartDeviceFlow(ctx) - if errStartDeviceFlow != nil { - log.Errorf("Failed to generate authorization URL: %v", errStartDeviceFlow) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) - return - } - authURL := deviceFlow.VerificationURIComplete - if authURL == "" { - authURL = deviceFlow.VerificationURI - } - - RegisterOAuthSession(state, "kimi") - - go func() { - pollCtx, cancelPoll := context.WithCancel(ctx) - defer cancelPoll() - go watchOAuthSessionCancel(pollCtx, cancelPoll, state, "kimi") - - fmt.Println("Waiting for authentication...") - authBundle, errWaitForAuthorization := kimiAuth.WaitForAuthorization(pollCtx, deviceFlow) - if errWaitForAuthorization != nil { - if !IsOAuthSessionPending(state, "kimi") { - return - } - SetOAuthSessionError(state, oauthSessionErrorWithCause("Authentication failed", errWaitForAuthorization)) - fmt.Printf("Authentication failed: %v\n", errWaitForAuthorization) - return - } - if !IsOAuthSessionPending(state, "kimi") { - return - } - - // Create token storage - tokenStorage := kimiAuth.CreateTokenStorage(authBundle) - - metadata := map[string]any{ - "type": "kimi", - "access_token": authBundle.TokenData.AccessToken, - "refresh_token": authBundle.TokenData.RefreshToken, - "token_type": authBundle.TokenData.TokenType, - "scope": authBundle.TokenData.Scope, - "timestamp": time.Now().UnixMilli(), - } - if authBundle.TokenData.ExpiresAt > 0 { - expired := time.Unix(authBundle.TokenData.ExpiresAt, 0).UTC().Format(time.RFC3339) - metadata["expired"] = expired - } - if strings.TrimSpace(authBundle.DeviceID) != "" { - metadata["device_id"] = strings.TrimSpace(authBundle.DeviceID) - } - - fileName := fmt.Sprintf("kimi-%d.json", time.Now().UnixMilli()) - record := &coreauth.Auth{ - ID: fileName, - Provider: "kimi", - FileName: fileName, - Label: "Kimi User", - Storage: tokenStorage, - Metadata: metadata, - } - if errGuard := guardOAuthSessionPendingForSave(state, "kimi"); errGuard != nil { - return - } - savedPath, errSave := h.saveTokenRecord(ctx, record) - if errSave != nil { - log.Errorf("Failed to save authentication tokens: %v", errSave) - SetOAuthSessionError(state, "Failed to save authentication tokens") - return - } - - fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) - fmt.Println("You can now use Kimi services through this CLI") - CompleteOAuthSession(state) - }() - - response := gin.H{"status": "ok", "url": authURL, "state": state, "flow": "device"} - if userCode := strings.TrimSpace(deviceFlow.UserCode); userCode != "" { - response["user_code"] = userCode - } - if deviceFlow.ExpiresIn > 0 { - response["expires_in"] = deviceFlow.ExpiresIn - } - c.JSON(200, response) -} - -// watchOAuthSessionCancel cancels pollCtx once the OAuth session is no longer pending. -func watchOAuthSessionCancel(pollCtx context.Context, cancel context.CancelFunc, state, provider string) { - if cancel == nil { - return - } - ticker := time.NewTicker(2 * time.Second) - defer ticker.Stop() - for { - select { - case <-pollCtx.Done(): - return - case <-ticker.C: - if !IsOAuthSessionPending(state, provider) { - cancel() - return - } - } - } -} - -// CancelAuthSession cancels a pending OAuth session identified by state. -// Protected by management auth. Safe for both callback and device-code flows: -// waiters check IsOAuthSessionPending and exit without saving credentials. -func (h *Handler) CancelAuthSession(c *gin.Context) { - state := strings.TrimSpace(c.Query("state")) - if state == "" { - c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "missing state"}) - return - } - if err := ValidateOAuthState(state); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid state"}) - return - } - cancelled := CancelOAuthSession(state) - c.JSON(http.StatusOK, gin.H{"status": "ok", "cancelled": cancelled}) -} - -func (h *Handler) GetAuthStatus(c *gin.Context) { - state := strings.TrimSpace(c.Query("state")) - if state == "" { - c.JSON(http.StatusOK, gin.H{"status": "ok"}) - return - } - if err := ValidateOAuthState(state); err != nil { - c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid state"}) - return - } - - provider, status, isPlugin, metadata, completed, ok := GetOAuthSessionDetails(state) - if !ok { - c.JSON(http.StatusOK, gin.H{"status": "error", "error": "unknown or expired state"}) - return - } - if completed { - c.JSON(http.StatusOK, gin.H{"status": "ok"}) - return - } - if status != "" { - c.JSON(http.StatusOK, gin.H{"status": "error", "error": status}) - return - } - h.mu.Lock() - host := h.pluginHost - h.mu.Unlock() - if isPlugin && host != nil && host.HasAuthProvider(provider) { - ctx := PopulateAuthContext(context.Background(), c) - resp, handled, errPoll := host.PollLogin(ctx, provider, state, metadata) - if handled { - if errPoll != nil { - message := strings.TrimSpace(errPoll.Error()) - if message == "" { - message = "Authentication failed" - } - SetOAuthSessionError(state, message) - c.JSON(http.StatusOK, gin.H{"status": "error", "error": message}) - return - } - switch resp.Status { - case "", pluginapi.AuthLoginStatusPending: - c.JSON(http.StatusOK, gin.H{"status": "wait"}) - return - case pluginapi.AuthLoginStatusError: - message := strings.TrimSpace(resp.Message) - if message == "" { - message = "Authentication failed" - } - SetOAuthSessionError(state, message) - c.JSON(http.StatusOK, gin.H{"status": "error", "error": message}) - return - case pluginapi.AuthLoginStatusSuccess: - records := pluginLoginPollAuths(host, resp) - if len(records) == 0 { - SetOAuthSessionError(state, "Authentication failed") - c.JSON(http.StatusOK, gin.H{"status": "error", "error": "Authentication failed"}) - return - } - if errSave := h.savePluginLoginRecords(ctx, records); errSave != nil { - log.WithError(errSave).WithField("provider", provider).Error("failed to save plugin auth tokens") - SetOAuthSessionError(state, "Failed to save authentication tokens") - c.JSON(http.StatusOK, gin.H{"status": "error", "error": "Failed to save authentication tokens"}) - return - } - CompleteOAuthSession(state) - c.JSON(http.StatusOK, gin.H{"status": "ok"}) - return - default: - c.JSON(http.StatusOK, gin.H{"status": "wait"}) - return - } - } - } - c.JSON(http.StatusOK, gin.H{"status": "wait"}) -} - -func pluginLoginPollAuths(host *pluginhost.Host, resp pluginapi.AuthLoginPollResponse) []*coreauth.Auth { - if host == nil { - return nil - } - authDatas := resp.Auths - if len(authDatas) == 0 { - authDatas = []pluginapi.AuthData{resp.Auth} - } - records := make([]*coreauth.Auth, 0, len(authDatas)) - for _, authData := range authDatas { - record := host.AuthDataToCoreAuth(authData, "", "") - if record == nil { - return nil - } - records = append(records, record) - } - return records -} - -func (h *Handler) savePluginLoginRecords(ctx context.Context, records []*coreauth.Auth) error { - savedPaths := make([]string, 0, len(records)) - for _, record := range records { - savedPath, errSave := h.saveTokenRecord(ctx, record) - if strings.TrimSpace(savedPath) != "" { - savedPaths = append(savedPaths, savedPath) - } - if errSave != nil { - h.rollbackSavedTokenRecords(ctx, savedPaths) - return errSave - } - } - return nil -} - -func (h *Handler) rollbackSavedTokenRecords(ctx context.Context, savedPaths []string) { - for i := len(savedPaths) - 1; i >= 0; i-- { - path := strings.TrimSpace(savedPaths[i]) - if path == "" { - continue - } - if errDelete := h.deleteTokenRecord(ctx, path); errDelete != nil { - log.WithError(errDelete).WithField("path", path).Warn("failed to roll back plugin auth token") - } - h.removeAuthsForPath(ctx, path, path) - } -} - -// PopulateAuthContext extracts request info and adds it to the context -func PopulateAuthContext(ctx context.Context, c *gin.Context) context.Context { - info := &coreauth.RequestInfo{ - Query: c.Request.URL.Query(), - Headers: c.Request.Header, - } - return coreauth.WithRequestInfo(ctx, info) -} diff --git a/internal/api/handlers/management/auth_files_crud.go b/internal/api/handlers/management/auth_files_crud.go new file mode 100644 index 000000000..c7416334e --- /dev/null +++ b/internal/api/handlers/management/auth_files_crud.go @@ -0,0 +1,550 @@ +package management + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "mime/multipart" + "net/http" + "os" + "path/filepath" + "runtime" + "sort" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +// Download single auth file by name +func (h *Handler) DownloadAuthFile(c *gin.Context) { + name := strings.TrimSpace(c.Query("name")) + if isUnsafeAuthFileName(name) { + c.JSON(400, gin.H{"error": "invalid name"}) + return + } + if !strings.HasSuffix(strings.ToLower(name), ".json") { + c.JSON(400, gin.H{"error": "name must end with .json"}) + return + } + full := filepath.Join(h.cfg.AuthDir, name) + data, err := os.ReadFile(full) + if err != nil { + if os.IsNotExist(err) { + c.JSON(404, gin.H{"error": "file not found"}) + } else { + c.JSON(500, gin.H{"error": fmt.Sprintf("failed to read file: %v", err)}) + } + return + } + c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", name)) + c.Data(200, "application/json", data) +} + +// Upload auth file: multipart or raw JSON with ?name= +func (h *Handler) UploadAuthFile(c *gin.Context) { + if h.authManager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) + return + } + ctx := c.Request.Context() + + fileHeaders, errMultipart := h.multipartAuthFileHeaders(c) + if errMultipart != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("invalid multipart form: %v", errMultipart)}) + return + } + if len(fileHeaders) == 1 { + if _, errUpload := h.storeUploadedAuthFile(ctx, fileHeaders[0]); errUpload != nil { + if errors.Is(errUpload, errAuthFileMustBeJSON) { + c.JSON(http.StatusBadRequest, gin.H{"error": "file must be .json"}) + return + } + c.JSON(http.StatusInternalServerError, gin.H{"error": errUpload.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + return + } + if len(fileHeaders) > 1 { + uploaded := make([]string, 0, len(fileHeaders)) + failed := make([]gin.H, 0) + for _, file := range fileHeaders { + name, errUpload := h.storeUploadedAuthFile(ctx, file) + if errUpload != nil { + failureName := "" + if file != nil { + failureName = filepath.Base(file.Filename) + } + msg := errUpload.Error() + if errors.Is(errUpload, errAuthFileMustBeJSON) { + msg = "file must be .json" + } + failed = append(failed, gin.H{"name": failureName, "error": msg}) + continue + } + uploaded = append(uploaded, name) + } + if len(failed) > 0 { + c.JSON(http.StatusMultiStatus, gin.H{ + "status": "partial", + "uploaded": len(uploaded), + "files": uploaded, + "failed": failed, + }) + return + } + c.JSON(http.StatusOK, gin.H{"status": "ok", "uploaded": len(uploaded), "files": uploaded}) + return + } + if c.ContentType() == "multipart/form-data" { + c.JSON(http.StatusBadRequest, gin.H{"error": "no files uploaded"}) + return + } + name := strings.TrimSpace(c.Query("name")) + if isUnsafeAuthFileName(name) { + c.JSON(400, gin.H{"error": "invalid name"}) + return + } + if !strings.HasSuffix(strings.ToLower(name), ".json") { + c.JSON(400, gin.H{"error": "name must end with .json"}) + return + } + data, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(400, gin.H{"error": "failed to read body"}) + return + } + if err = h.writeAuthFile(ctx, filepath.Base(name), data); err != nil { + c.JSON(500, gin.H{"error": err.Error()}) + return + } + c.JSON(200, gin.H{"status": "ok"}) +} + +// Delete auth files: single by name or all +func (h *Handler) DeleteAuthFile(c *gin.Context) { + if h.authManager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) + return + } + ctx := c.Request.Context() + if all := c.Query("all"); all == "true" || all == "1" || all == "*" { + entries, err := os.ReadDir(h.cfg.AuthDir) + if err != nil { + c.JSON(500, gin.H{"error": fmt.Sprintf("failed to read auth dir: %v", err)}) + return + } + deleted := 0 + for _, e := range entries { + if e.IsDir() { + continue + } + name := e.Name() + if !strings.HasSuffix(strings.ToLower(name), ".json") { + continue + } + full := filepath.Join(h.cfg.AuthDir, name) + if !filepath.IsAbs(full) { + if abs, errAbs := filepath.Abs(full); errAbs == nil { + full = abs + } + } + if err = os.Remove(full); err == nil { + if errDel := h.deleteTokenRecord(ctx, full); errDel != nil { + c.JSON(500, gin.H{"error": errDel.Error()}) + return + } + deleted++ + h.removeAuth(ctx, full) + } + } + c.JSON(200, gin.H{"status": "ok", "deleted": deleted}) + return + } + + names, errNames := requestedAuthFileNamesForDelete(c) + if errNames != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": errNames.Error()}) + return + } + if len(names) == 0 { + c.JSON(400, gin.H{"error": "invalid name"}) + return + } + if len(names) == 1 { + if _, status, errDelete := h.deleteAuthFileByName(ctx, names[0]); errDelete != nil { + c.JSON(status, gin.H{"error": errDelete.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + return + } + + deletedFiles := make([]string, 0, len(names)) + failed := make([]gin.H, 0) + for _, name := range names { + deletedName, _, errDelete := h.deleteAuthFileByName(ctx, name) + if errDelete != nil { + failed = append(failed, gin.H{"name": name, "error": errDelete.Error()}) + continue + } + deletedFiles = append(deletedFiles, deletedName) + } + if len(failed) > 0 { + c.JSON(http.StatusMultiStatus, gin.H{ + "status": "partial", + "deleted": len(deletedFiles), + "files": deletedFiles, + "failed": failed, + }) + return + } + c.JSON(http.StatusOK, gin.H{"status": "ok", "deleted": len(deletedFiles), "files": deletedFiles}) +} + +func (h *Handler) multipartAuthFileHeaders(c *gin.Context) ([]*multipart.FileHeader, error) { + if h == nil || c == nil || c.ContentType() != "multipart/form-data" { + return nil, nil + } + form, err := c.MultipartForm() + if err != nil { + return nil, err + } + if form == nil || len(form.File) == 0 { + return nil, nil + } + + keys := make([]string, 0, len(form.File)) + for key := range form.File { + keys = append(keys, key) + } + sort.Strings(keys) + + headers := make([]*multipart.FileHeader, 0) + for _, key := range keys { + headers = append(headers, form.File[key]...) + } + return headers, nil +} + +func (h *Handler) storeUploadedAuthFile(ctx context.Context, file *multipart.FileHeader) (string, error) { + if file == nil { + return "", fmt.Errorf("no file uploaded") + } + name := filepath.Base(strings.TrimSpace(file.Filename)) + if !strings.HasSuffix(strings.ToLower(name), ".json") { + return "", errAuthFileMustBeJSON + } + src, err := file.Open() + if err != nil { + return "", fmt.Errorf("failed to open uploaded file: %w", err) + } + defer src.Close() + + data, err := io.ReadAll(src) + if err != nil { + return "", fmt.Errorf("failed to read uploaded file: %w", err) + } + if err := h.writeAuthFile(ctx, name, data); err != nil { + return "", err + } + return name, nil +} + +func (h *Handler) writeAuthFile(ctx context.Context, name string, data []byte) error { + dst := filepath.Join(h.cfg.AuthDir, filepath.Base(name)) + if !filepath.IsAbs(dst) { + if abs, errAbs := filepath.Abs(dst); errAbs == nil { + dst = abs + } + } + auth, err := h.buildAuthFromFileData(dst, data) + if err != nil { + return err + } + if errWrite := os.WriteFile(dst, data, 0o600); errWrite != nil { + return fmt.Errorf("failed to write file: %w", errWrite) + } + if err := h.upsertAuthRecord(ctx, auth); err != nil { + return err + } + return nil +} + +func requestedAuthFileNamesForDelete(c *gin.Context) ([]string, error) { + if c == nil { + return nil, nil + } + names := uniqueAuthFileNames(c.QueryArray("name")) + if len(names) > 0 { + return names, nil + } + + body, err := io.ReadAll(c.Request.Body) + if err != nil { + return nil, fmt.Errorf("failed to read body") + } + body = bytes.TrimSpace(body) + if len(body) == 0 { + return nil, nil + } + + var objectBody struct { + Name string `json:"name"` + Names []string `json:"names"` + } + if body[0] == '[' { + var arrayBody []string + if err := json.Unmarshal(body, &arrayBody); err != nil { + return nil, fmt.Errorf("invalid request body") + } + return uniqueAuthFileNames(arrayBody), nil + } + if err := json.Unmarshal(body, &objectBody); err != nil { + return nil, fmt.Errorf("invalid request body") + } + + out := make([]string, 0, len(objectBody.Names)+1) + if strings.TrimSpace(objectBody.Name) != "" { + out = append(out, objectBody.Name) + } + out = append(out, objectBody.Names...) + return uniqueAuthFileNames(out), nil +} + +func uniqueAuthFileNames(names []string) []string { + if len(names) == 0 { + return nil + } + seen := make(map[string]struct{}, len(names)) + out := make([]string, 0, len(names)) + for _, name := range names { + name = strings.TrimSpace(name) + if name == "" { + continue + } + if _, ok := seen[name]; ok { + continue + } + seen[name] = struct{}{} + out = append(out, name) + } + return out +} + +func (h *Handler) deleteAuthFileByName(ctx context.Context, name string) (string, int, error) { + name = strings.TrimSpace(name) + if isUnsafeAuthFileName(name) { + return "", http.StatusBadRequest, fmt.Errorf("invalid name") + } + + targetPath := filepath.Join(h.cfg.AuthDir, filepath.Base(name)) + targetID := "" + if targetAuth := h.findAuthForDelete(name); targetAuth != nil { + if !isPluginVirtualSourceDelete(name, targetAuth) { + return filepath.Base(name), http.StatusConflict, errPluginVirtualAuth + } + targetID = strings.TrimSpace(targetAuth.ID) + if path := strings.TrimSpace(authAttribute(targetAuth, "path")); path != "" { + targetPath = path + } + } + if !filepath.IsAbs(targetPath) { + if abs, errAbs := filepath.Abs(targetPath); errAbs == nil { + targetPath = abs + } + } + if errRemove := os.Remove(targetPath); errRemove != nil { + if os.IsNotExist(errRemove) { + return filepath.Base(name), http.StatusNotFound, errAuthFileNotFound + } + return filepath.Base(name), http.StatusInternalServerError, fmt.Errorf("failed to remove file: %w", errRemove) + } + if errDeleteRecord := h.deleteTokenRecord(ctx, targetPath); errDeleteRecord != nil { + return filepath.Base(name), http.StatusInternalServerError, errDeleteRecord + } + h.removeAuthsForPath(ctx, targetPath, targetID) + return filepath.Base(name), http.StatusOK, nil +} + +func isPluginVirtualSourceDelete(name string, auth *coreauth.Auth) bool { + if !coreauth.IsPluginVirtualAuth(auth) { + return true + } + sourcePath := strings.TrimSpace(authAttribute(auth, coreauth.AttributeVirtualSource)) + if sourcePath == "" { + sourcePath = strings.TrimSpace(authAttribute(auth, "path")) + } + if sourcePath == "" { + return false + } + return strings.EqualFold(filepath.Base(strings.TrimSpace(name)), filepath.Base(sourcePath)) +} + +func (h *Handler) findAuthForDelete(name string) *coreauth.Auth { + if h == nil || h.authManager == nil { + return nil + } + name = strings.TrimSpace(name) + if name == "" { + return nil + } + if auth, ok := h.authManager.GetByID(name); ok { + return auth + } + auths := h.authManager.List() + for _, auth := range auths { + if auth == nil { + continue + } + if strings.TrimSpace(auth.FileName) == name { + return auth + } + if filepath.Base(strings.TrimSpace(authAttribute(auth, "path"))) == name { + return auth + } + } + return nil +} + +func (h *Handler) authIDForPath(path string) string { + path = strings.TrimSpace(path) + if path == "" { + return "" + } + path = filepath.Clean(path) + if !filepath.IsAbs(path) { + if abs, errAbs := filepath.Abs(path); errAbs == nil { + path = abs + } + } + id := path + if h != nil && h.cfg != nil { + authDir := strings.TrimSpace(h.cfg.AuthDir) + if resolvedAuthDir, errResolve := util.ResolveAuthDir(authDir); errResolve == nil && resolvedAuthDir != "" { + authDir = resolvedAuthDir + } + if authDir != "" { + authDir = filepath.Clean(authDir) + if !filepath.IsAbs(authDir) { + if abs, errAbs := filepath.Abs(authDir); errAbs == nil { + authDir = abs + } + } + if rel, errRel := filepath.Rel(authDir, path); errRel == nil && rel != "" { + id = rel + } + } + } + // On Windows, normalize ID casing to avoid duplicate auth entries caused by case-insensitive paths. + if runtime.GOOS == "windows" { + id = strings.ToLower(id) + } + return id +} + +func (h *Handler) registerAuthFromFile(ctx context.Context, path string, data []byte) error { + if h.authManager == nil { + return nil + } + auth, err := h.buildAuthFromFileData(path, data) + if err != nil { + return err + } + return h.upsertAuthRecord(ctx, auth) +} + +func (h *Handler) buildAuthFromFileData(path string, data []byte) (*coreauth.Auth, error) { + if path == "" { + return nil, fmt.Errorf("auth path is empty") + } + if data == nil { + var err error + data, err = os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("failed to read auth file: %w", err) + } + } + metadata := make(map[string]any) + if err := json.Unmarshal(data, &metadata); err != nil { + return nil, fmt.Errorf("invalid auth file: %w", err) + } + provider, _ := metadata["type"].(string) + if provider == "" { + provider = "unknown" + } + label := provider + if email, ok := metadata["email"].(string); ok && email != "" { + label = email + } + lastRefresh, hasLastRefresh := extractLastRefreshTimestamp(metadata) + + authID := h.authIDForPath(path) + if authID == "" { + authID = path + } + auth := (*coreauth.Auth)(nil) + if h != nil && h.cfg != nil { + sctx := &synthesizer.SynthesisContext{ + Config: h.cfg, + AuthDir: h.cfg.AuthDir, + Now: time.Now(), + IDGenerator: synthesizer.NewStableIDGenerator(), + } + if generated := synthesizer.SynthesizeAuthFile(sctx, path, data); len(generated) > 0 && generated[0] != nil { + auth = generated[0].Clone() + } + } + if auth == nil { + auth = &coreauth.Auth{ + ID: authID, + Provider: provider, + Label: label, + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "path": path, + "source": path, + }, + Metadata: metadata, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + } + } + auth.ID = authID + auth.FileName = filepath.Base(path) + if hasLastRefresh { + auth.LastRefreshedAt = lastRefresh + } + if h != nil && h.authManager != nil { + if existing, ok := h.authManager.GetByID(authID); ok { + auth.CreatedAt = existing.CreatedAt + if !hasLastRefresh { + auth.LastRefreshedAt = existing.LastRefreshedAt + } + auth.NextRefreshAfter = existing.NextRefreshAfter + auth.Runtime = existing.Runtime + } + } + coreauth.ApplyCustomHeadersFromMetadata(auth) + return auth, nil +} + +func (h *Handler) upsertAuthRecord(ctx context.Context, auth *coreauth.Auth) error { + if h == nil || h.authManager == nil || auth == nil { + return nil + } + if existing, ok := h.authManager.GetByID(auth.ID); ok { + auth.CreatedAt = existing.CreatedAt + _, err := h.authManager.Update(ctx, auth) + return err + } + _, err := h.authManager.Register(ctx, auth) + return err +} diff --git a/internal/api/handlers/management/auth_files_fields.go b/internal/api/handlers/management/auth_files_fields.go new file mode 100644 index 000000000..a1d633d6a --- /dev/null +++ b/internal/api/handlers/management/auth_files_fields.go @@ -0,0 +1,684 @@ +package management + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "path/filepath" + "runtime" + "strconv" + "strings" + "time" + + "github.com/gin-gonic/gin" + sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +// PatchAuthFileStatus toggles the disabled state of an auth file +func (h *Handler) PatchAuthFileStatus(c *gin.Context) { + if h.authManager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) + return + } + + var req struct { + Name string `json:"name"` + AuthIndex string `json:"auth_index"` + Disabled *bool `json:"disabled"` + } + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) + return + } + + name := strings.TrimSpace(req.Name) + authIndex := strings.TrimSpace(req.AuthIndex) + if name == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) + return + } + if req.Disabled == nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "disabled is required"}) + return + } + + ctx := c.Request.Context() + + targetAuth, _ := h.lookupAuthFile(name, authIndex) + if targetAuth == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"}) + return + } + if coreauth.IsPluginVirtualAuth(targetAuth) { + // Allow status changes only when targeting the source auth file name, matching delete semantics. + // Expanded virtual project auths still cannot be modified independently. + if !isPluginVirtualSourceDelete(name, targetAuth) { + c.JSON(http.StatusConflict, gin.H{"error": errPluginVirtualAuth.Error()}) + return + } + if errPatch := h.patchPluginVirtualSourceStatus(ctx, targetAuth, *req.Disabled); errPatch != nil { + status := http.StatusInternalServerError + if errors.Is(errPatch, errAuthFileNotFound) || os.IsNotExist(errPatch) { + status = http.StatusNotFound + } + c.JSON(status, gin.H{"error": errPatch.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"status": "ok", "disabled": *req.Disabled}) + return + } + + if coreauth.IsConfigAPIKeyAuth(targetAuth) { + h.mu.Lock() + handled, errToggle := toggleConfigAPIKeyExcludedAll(h.cfg, targetAuth, *req.Disabled) + if errToggle != nil { + h.mu.Unlock() + c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update config api key: %v", errToggle)}) + return + } + if !handled { + h.mu.Unlock() + c.JSON(http.StatusNotFound, gin.H{"error": "config api key entry not found"}) + return + } + cfgSnapshot, okSnapshot := h.saveConfigAndSnapshotLocked(c) + h.mu.Unlock() + if !okSnapshot { + return + } + h.reloadConfigAfterManagementSave(ctx, cfgSnapshot) + if h.tokenStore != nil { + _ = h.tokenStore.Delete(ctx, targetAuth.ID) + } + c.JSON(http.StatusOK, gin.H{ + "status": "ok", + "disabled": *req.Disabled, + "via": "config:excluded-models", + "excluded_pattern": configAPIKeyDisablePattern, + }) + return + } + + applyAuthDisabledState(targetAuth, *req.Disabled) + if _, err := h.authManager.Update(ctx, targetAuth); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update auth: %v", err)}) + return + } + + c.JSON(http.StatusOK, gin.H{"status": "ok", "disabled": *req.Disabled}) +} + +// patchPluginVirtualSourceStatus toggles disabled on a plugin multi-auth source file and all +// runtime auths expanded from it. Virtual project children cannot be toggled independently. +func (h *Handler) patchPluginVirtualSourceStatus(ctx context.Context, targetAuth *coreauth.Auth, disabled bool) error { + if h == nil || h.authManager == nil || targetAuth == nil { + return fmt.Errorf("core auth manager unavailable") + } + sourcePath := strings.TrimSpace(authAttribute(targetAuth, coreauth.AttributeVirtualSource)) + if sourcePath == "" { + sourcePath = strings.TrimSpace(authAttribute(targetAuth, "path")) + } + if sourcePath == "" { + return errPluginVirtualAuth + } + if errWrite := setSourceAuthFileDisabled(sourcePath, disabled); errWrite != nil { + if os.IsNotExist(errWrite) { + return errAuthFileNotFound + } + return fmt.Errorf("failed to update source auth file: %w", errWrite) + } + now := time.Now() + for _, auth := range h.authManager.List() { + if auth == nil { + continue + } + if !sameAuthFilePath(authAttribute(auth, "path"), sourcePath) && + !sameAuthFilePath(authAttribute(auth, coreauth.AttributeVirtualSource), sourcePath) { + continue + } + applyAuthDisabledState(auth, disabled) + auth.UpdatedAt = now + if _, errUpdate := h.authManager.Update(ctx, auth); errUpdate != nil { + return fmt.Errorf("failed to update auth %s: %w", auth.ID, errUpdate) + } + } + return nil +} + +func setSourceAuthFileDisabled(path string, disabled bool) error { + path = strings.TrimSpace(path) + if path == "" { + return fmt.Errorf("source auth path is empty") + } + data, errRead := os.ReadFile(path) + if errRead != nil { + return errRead + } + metadata := make(map[string]any) + if len(bytes.TrimSpace(data)) > 0 { + if errUnmarshal := json.Unmarshal(data, &metadata); errUnmarshal != nil { + return fmt.Errorf("invalid auth file: %w", errUnmarshal) + } + } + if metadata == nil { + metadata = make(map[string]any) + } + metadata["disabled"] = disabled + raw, errMarshal := json.Marshal(metadata) + if errMarshal != nil { + return fmt.Errorf("marshal auth file: %w", errMarshal) + } + if errWrite := os.WriteFile(path, raw, 0o600); errWrite != nil { + return errWrite + } + return nil +} + +func applyAuthDisabledState(auth *coreauth.Auth, disabled bool) { + if auth == nil { + return + } + auth.Disabled = disabled + if disabled { + auth.Status = coreauth.StatusDisabled + auth.StatusMessage = "disabled via management API" + } else { + auth.Status = coreauth.StatusActive + auth.StatusMessage = "" + } + auth.UpdatedAt = time.Now() + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["disabled"] = disabled +} + +// PatchAuthFileFields updates arbitrary metadata fields of an auth file. +func (h *Handler) PatchAuthFileFields(c *gin.Context) { + if h.authManager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "core auth manager unavailable"}) + return + } + + var req map[string]json.RawMessage + decoder := json.NewDecoder(c.Request.Body) + decoder.UseNumber() + if err := decoder.Decode(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"}) + return + } + + nameRaw, ok := req["name"] + if !ok { + c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) + return + } + var nameValue string + if err := json.Unmarshal(nameRaw, &nameValue); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) + return + } + name := strings.TrimSpace(nameValue) + if name == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "name is required"}) + return + } + delete(req, "name") + + ctx := c.Request.Context() + + // Find auth by name or ID + var targetAuth *coreauth.Auth + if auth, ok := h.authManager.GetByID(name); ok { + targetAuth = auth + } else { + auths := h.authManager.List() + for _, auth := range auths { + if auth.FileName == name { + targetAuth = auth + break + } + } + } + + if targetAuth == nil { + c.JSON(http.StatusNotFound, gin.H{"error": "auth file not found"}) + return + } + if coreauth.IsPluginVirtualAuth(targetAuth) { + c.JSON(http.StatusConflict, gin.H{"error": errPluginVirtualAuth.Error()}) + return + } + + changed := false + touchedRoots := make(map[string]struct{}, len(req)) + for key, rawValue := range req { + fieldPath := strings.TrimSpace(key) + if fieldPath == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "field name is required"}) + return + } + value, errDecode := decodeAuthFileFieldValue(rawValue) + if errDecode != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("invalid field %s", fieldPath)}) + return + } + if targetAuth.Metadata == nil { + targetAuth.Metadata = make(map[string]any) + } + + if fieldPath == "headers" { + applyAuthFileHeadersPatch(targetAuth, value) + } else if errSet := setAuthFileMetadataValue(targetAuth.Metadata, fieldPath, value); errSet != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": errSet.Error()}) + return + } + if root := rootAuthFileField(fieldPath); root != "" { + touchedRoots[root] = struct{}{} + } + changed = true + } + if changed { + syncAuthFileMetadataFields(targetAuth, touchedRoots) + } + + if !changed { + c.JSON(http.StatusBadRequest, gin.H{"error": "no fields to update"}) + return + } + + targetAuth.UpdatedAt = time.Now() + + if _, err := h.authManager.Update(ctx, targetAuth); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to update auth: %v", err)}) + return + } + + c.JSON(http.StatusOK, gin.H{"status": "ok"}) +} + +func decodeAuthFileFieldValue(raw json.RawMessage) (any, error) { + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return nil, err + } + return value, nil +} + +func rootAuthFileField(path string) string { + path = strings.TrimSpace(path) + if path == "" { + return "" + } + if idx := strings.Index(path, "."); idx >= 0 { + return strings.TrimSpace(path[:idx]) + } + return path +} + +func setAuthFileMetadataValue(metadata map[string]any, path string, value any) error { + if metadata == nil { + return fmt.Errorf("metadata is nil") + } + parts := strings.Split(path, ".") + current := metadata + for i, rawPart := range parts { + part := strings.TrimSpace(rawPart) + if part == "" { + return fmt.Errorf("invalid field path: %s", path) + } + if i == len(parts)-1 { + current[part] = value + return nil + } + next, ok := current[part].(map[string]any) + if !ok { + next = make(map[string]any) + current[part] = next + } + current = next + } + return nil +} + +func applyAuthFileHeadersPatch(auth *coreauth.Auth, value any) { + if auth == nil { + return + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + headersPatch, ok := authFileHeadersStringMap(value) + if !ok { + auth.Metadata["headers"] = value + return + } + + existingHeaders := coreauth.ExtractCustomHeadersFromMetadata(auth.Metadata) + nextHeaders := make(map[string]string, len(existingHeaders)) + for key, val := range existingHeaders { + nextHeaders[key] = val + } + for key, value := range headersPatch { + name := strings.TrimSpace(key) + if name == "" { + continue + } + val := strings.TrimSpace(value) + if val == "" { + delete(nextHeaders, name) + continue + } + nextHeaders[name] = val + } + + if len(nextHeaders) == 0 { + delete(auth.Metadata, "headers") + return + } + metaHeaders := make(map[string]any, len(nextHeaders)) + for key, value := range nextHeaders { + metaHeaders[key] = value + } + auth.Metadata["headers"] = metaHeaders +} + +func authFileHeadersStringMap(value any) (map[string]string, bool) { + switch typed := value.(type) { + case map[string]string: + return typed, true + case map[string]any: + out := make(map[string]string, len(typed)) + for key, rawValue := range typed { + value, ok := rawValue.(string) + if !ok { + return nil, false + } + out[key] = value + } + return out, true + default: + return nil, false + } +} + +func syncAuthFileMetadataFields(auth *coreauth.Auth, touchedRoots map[string]struct{}) { + if auth == nil || len(touchedRoots) == 0 { + return + } + if _, ok := touchedRoots["prefix"]; ok { + if prefix, okString := auth.Metadata["prefix"].(string); okString { + auth.Prefix = strings.TrimSpace(prefix) + } + } + if _, ok := touchedRoots["proxy_url"]; ok { + if proxyURL, okString := auth.Metadata["proxy_url"].(string); okString { + auth.ProxyURL = strings.TrimSpace(proxyURL) + } + } + if _, ok := touchedRoots["headers"]; ok { + syncAuthFileHeaderAttributes(auth) + } + if _, ok := touchedRoots["priority"]; ok { + syncAuthFilePriorityAttribute(auth) + } + if _, ok := touchedRoots["note"]; ok { + syncAuthFileNoteAttribute(auth) + } + if _, ok := touchedRoots["websockets"]; ok { + syncAuthFileWebsocketsAttribute(auth) + } + if _, ok := touchedRoots["disabled"]; ok { + syncAuthFileDisabledState(auth) + } +} + +func syncAuthFileHeaderAttributes(auth *coreauth.Auth) { + if auth == nil { + return + } + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + for key := range auth.Attributes { + if strings.HasPrefix(key, "header:") { + delete(auth.Attributes, key) + } + } + for name, value := range coreauth.ExtractCustomHeadersFromMetadata(auth.Metadata) { + auth.Attributes["header:"+name] = value + } +} + +func syncAuthFilePriorityAttribute(auth *coreauth.Auth) { + if auth == nil { + return + } + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + priority, ok := authFileIntValue(auth.Metadata["priority"]) + if !ok { + delete(auth.Attributes, "priority") + return + } + if priority == 0 { + delete(auth.Attributes, "priority") + return + } + auth.Attributes["priority"] = strconv.Itoa(priority) +} + +func authFileIntValue(value any) (int, bool) { + switch typed := value.(type) { + case int: + return typed, true + case int64: + return int(typed), true + case float64: + return int(typed), true + case json.Number: + if i, err := typed.Int64(); err == nil { + return int(i), true + } + case string: + if i, err := strconv.Atoi(strings.TrimSpace(typed)); err == nil { + return i, true + } + } + return 0, false +} + +func syncAuthFileNoteAttribute(auth *coreauth.Auth) { + if auth == nil { + return + } + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + note, ok := auth.Metadata["note"].(string) + if !ok { + delete(auth.Attributes, "note") + return + } + note = strings.TrimSpace(note) + if note == "" { + delete(auth.Attributes, "note") + return + } + auth.Attributes["note"] = note +} + +func syncAuthFileWebsocketsAttribute(auth *coreauth.Auth) { + if auth == nil { + return + } + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + websockets, ok := authFileBoolValue(auth.Metadata["websockets"]) + if !ok { + delete(auth.Attributes, "websockets") + return + } + auth.Attributes["websockets"] = strconv.FormatBool(websockets) +} + +func authFileBoolValue(value any) (bool, bool) { + switch typed := value.(type) { + case bool: + return typed, true + case string: + parsed, errParse := strconv.ParseBool(strings.TrimSpace(typed)) + if errParse == nil { + return parsed, true + } + } + return false, false +} + +func syncAuthFileDisabledState(auth *coreauth.Auth) { + if auth == nil { + return + } + disabled, ok := authFileBoolValue(auth.Metadata["disabled"]) + if !ok { + return + } + auth.Disabled = disabled + if disabled { + auth.Status = coreauth.StatusDisabled + if strings.TrimSpace(auth.StatusMessage) == "" { + auth.StatusMessage = "disabled via management API" + } + return + } + auth.Status = coreauth.StatusActive + auth.StatusMessage = "" +} + +func (h *Handler) removeAuth(ctx context.Context, id string) { + if h == nil || h.authManager == nil { + return + } + id = strings.TrimSpace(id) + if id == "" { + return + } + if _, ok := h.authManager.GetByID(id); ok { + h.authManager.Remove(ctx, id) + return + } + authID := h.authIDForPath(id) + if authID == "" { + return + } + h.authManager.Remove(ctx, authID) +} + +func (h *Handler) removeAuthsForPath(ctx context.Context, path string, fallbackID string) { + if h == nil || h.authManager == nil { + return + } + removed := false + for _, auth := range h.authManager.List() { + if auth == nil { + continue + } + if sameAuthFilePath(authAttribute(auth, "path"), path) || sameAuthFilePath(authAttribute(auth, coreauth.AttributeVirtualSource), path) { + h.removeAuth(ctx, auth.ID) + removed = true + } + } + if removed { + return + } + if strings.TrimSpace(fallbackID) != "" { + h.removeAuth(ctx, fallbackID) + return + } + h.removeAuth(ctx, path) +} + +func sameAuthFilePath(left, right string) bool { + left = cleanAuthFilePath(left) + right = cleanAuthFilePath(right) + if left == "" || right == "" { + return false + } + if runtime.GOOS == "windows" { + return strings.EqualFold(left, right) + } + return left == right +} + +func cleanAuthFilePath(path string) string { + path = strings.TrimSpace(path) + if path == "" { + return "" + } + if abs, errAbs := filepath.Abs(path); errAbs == nil && strings.TrimSpace(abs) != "" { + path = abs + } + return filepath.Clean(path) +} + +func (h *Handler) deleteTokenRecord(ctx context.Context, path string) error { + if strings.TrimSpace(path) == "" { + return fmt.Errorf("auth path is empty") + } + store := h.tokenStoreWithBaseDir() + if store == nil { + return fmt.Errorf("token store unavailable") + } + return store.Delete(ctx, path) +} + +func (h *Handler) tokenStoreWithBaseDir() coreauth.Store { + if h == nil { + return nil + } + store := h.tokenStore + if store == nil { + store = sdkAuth.GetTokenStore() + h.tokenStore = store + } + if h.cfg != nil { + if dirSetter, ok := store.(interface{ SetBaseDir(string) }); ok { + dirSetter.SetBaseDir(h.cfg.AuthDir) + } + } + return store +} + +func (h *Handler) saveTokenRecord(ctx context.Context, record *coreauth.Auth) (string, error) { + if record == nil { + return "", fmt.Errorf("token record is nil") + } + store := h.tokenStoreWithBaseDir() + if store == nil { + return "", fmt.Errorf("token store unavailable") + } + if h.postAuthHook != nil { + if err := h.postAuthHook(ctx, record); err != nil { + return "", fmt.Errorf("post-auth hook failed: %w", err) + } + } + savedPath, errSave := store.Save(ctx, record) + if errSave != nil { + return savedPath, errSave + } + if h.postAuthPersistHook != nil { + if errHook := h.postAuthPersistHook(ctx, record); errHook != nil { + return savedPath, fmt.Errorf("post-auth persist hook failed: %w", errHook) + } + } + return savedPath, nil +} diff --git a/internal/api/handlers/management/auth_files_oauth_callback.go b/internal/api/handlers/management/auth_files_oauth_callback.go new file mode 100644 index 000000000..1b9ac82b6 --- /dev/null +++ b/internal/api/handlers/management/auth_files_oauth_callback.go @@ -0,0 +1,220 @@ +package management + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + log "github.com/sirupsen/logrus" +) + +const ( + anthropicCallbackPort = 54545 + codexCallbackPort = 1455 +) + +type callbackForwarder struct { + provider string + server *http.Server + done chan struct{} +} + +func isWebUIRequest(c *gin.Context) bool { + raw := strings.TrimSpace(c.Query("is_webui")) + if raw == "" { + return false + } + switch strings.ToLower(raw) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +func startCallbackForwarder(port int, provider, targetBase string) (*callbackForwarder, error) { + callbackForwardersMu.Lock() + prev := callbackForwarders[port] + if prev != nil { + delete(callbackForwarders, port) + } + callbackForwardersMu.Unlock() + + if prev != nil { + stopForwarderInstance(port, prev) + } + + addr := fmt.Sprintf("0.0.0.0:%d", port) + ln, err := net.Listen("tcp", addr) + if err != nil { + return nil, fmt.Errorf("failed to listen on %s: %w", addr, err) + } + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + target := targetBase + if raw := r.URL.RawQuery; raw != "" { + if strings.Contains(target, "?") { + target = target + "&" + raw + } else { + target = target + "?" + raw + } + } + w.Header().Set("Cache-Control", "no-store") + http.Redirect(w, r, target, http.StatusFound) + }) + + srv := &http.Server{ + Handler: handler, + ReadHeaderTimeout: 5 * time.Second, + WriteTimeout: 5 * time.Second, + } + done := make(chan struct{}) + + go func() { + if errServe := srv.Serve(ln); errServe != nil && !errors.Is(errServe, http.ErrServerClosed) { + log.WithError(errServe).Warnf("callback forwarder for %s stopped unexpectedly", provider) + } + close(done) + }() + + forwarder := &callbackForwarder{ + provider: provider, + server: srv, + done: done, + } + + callbackForwardersMu.Lock() + callbackForwarders[port] = forwarder + callbackForwardersMu.Unlock() + + log.Infof("callback forwarder for %s listening on %s", provider, addr) + + return forwarder, nil +} + +func stopCallbackForwarderInstance(port int, forwarder *callbackForwarder) { + if forwarder == nil { + return + } + callbackForwardersMu.Lock() + if current := callbackForwarders[port]; current == forwarder { + delete(callbackForwarders, port) + } + callbackForwardersMu.Unlock() + + stopForwarderInstance(port, forwarder) +} + +func stopForwarderInstance(port int, forwarder *callbackForwarder) { + if forwarder == nil || forwarder.server == nil { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + if err := forwarder.server.Shutdown(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.WithError(err).Warnf("failed to shut down callback forwarder on port %d", port) + } + + select { + case <-forwarder.done: + case <-time.After(2 * time.Second): + } + + log.Infof("callback forwarder on port %d stopped", port) +} + +func (h *Handler) managementCallbackURL(path string) (string, error) { + if h == nil || h.cfg == nil || h.cfg.Port <= 0 { + return "", fmt.Errorf("server port is not configured") + } + if !strings.HasPrefix(path, "/") { + path = "/" + path + } + scheme := "http" + if h.cfg.TLS.Enable { + scheme = "https" + } + return fmt.Sprintf("%s://127.0.0.1:%d%s", scheme, h.cfg.Port, path), nil +} + +func pluginAuthProviderFromPath(path string) (string, bool) { + path = strings.TrimSpace(path) + const prefix = "/v0/management/" + const suffix = "-auth-url" + if !strings.HasPrefix(path, prefix) || !strings.HasSuffix(path, suffix) { + return "", false + } + provider := strings.TrimSuffix(strings.TrimPrefix(path, prefix), suffix) + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" { + return "", false + } + for _, r := range provider { + switch { + case r >= 'a' && r <= 'z': + case r >= '0' && r <= '9': + case r == '-': + default: + return "", false + } + } + return provider, true +} + +func (h *Handler) ServePluginAuthURL(c *gin.Context) bool { + if h == nil || c == nil || c.Request == nil || c.Request.URL == nil { + return false + } + h.mu.Lock() + host := h.pluginHost + h.mu.Unlock() + if host == nil { + return false + } + provider, ok := pluginAuthProviderFromPath(c.Request.URL.Path) + if !ok || !host.HasAuthProvider(provider) { + return false + } + + ctx := PopulateAuthContext(context.Background(), c) + baseURL, errBaseURL := h.managementCallbackURL("/v0/management/oauth-callback") + if errBaseURL != nil { + log.WithError(errBaseURL).Error("failed to compute plugin auth callback URL") + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) + return true + } + resp, handled, errStart := host.StartLogin(ctx, provider, baseURL) + if !handled { + return false + } + if errStart != nil { + log.WithError(errStart).Error("failed to start plugin auth login") + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) + return true + } + state := strings.TrimSpace(resp.State) + if state == "" { + log.WithField("provider", provider).Error("plugin auth provider returned empty state") + c.JSON(http.StatusBadGateway, gin.H{"error": "invalid oauth state"}) + return true + } + if errState := ValidateOAuthState(state); errState != nil { + log.WithError(errState).WithField("provider", provider).Error("plugin auth provider returned invalid state") + c.JSON(http.StatusBadGateway, gin.H{"error": "invalid oauth state"}) + return true + } + if errRegister := RegisterPluginOAuthSession(state, provider, resp.Metadata); errRegister != nil { + log.WithError(errRegister).WithField("provider", provider).Error("failed to register plugin oauth session") + c.JSON(http.StatusBadGateway, gin.H{"error": "failed to generate authorization url"}) + return true + } + c.JSON(http.StatusOK, gin.H{"status": "ok", "url": resp.URL, "state": state}) + return true +} diff --git a/internal/api/handlers/management/auth_files_provider_oauth.go b/internal/api/handlers/management/auth_files_provider_oauth.go new file mode 100644 index 000000000..1c35ae809 --- /dev/null +++ b/internal/api/handlers/management/auth_files_provider_oauth.go @@ -0,0 +1,875 @@ +package management + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "path/filepath" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/antigravity" + "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude" + "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/kimi" + xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" + "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" + "github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + log "github.com/sirupsen/logrus" +) + +type codexOAuthService interface { + GenerateAuthURL(state string, pkceCodes *codex.PKCECodes) (string, error) + ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *codex.PKCECodes) (*codex.CodexAuthBundle, error) + CreateTokenStorage(bundle *codex.CodexAuthBundle) *codex.CodexTokenStorage +} + +func (h *Handler) RequestAnthropicToken(c *gin.Context) { + ctx := context.Background() + ctx = PopulateAuthContext(ctx, c) + + fmt.Println("Initializing Claude authentication...") + + // Generate PKCE codes + pkceCodes, err := claude.GeneratePKCECodes() + if err != nil { + log.Errorf("Failed to generate PKCE codes: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"}) + return + } + + // Generate random state parameter + state, err := misc.GenerateRandomState() + if err != nil { + log.Errorf("Failed to generate state parameter: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"}) + return + } + + // Initialize Claude auth service + anthropicAuth := claude.NewClaudeAuth(h.cfg) + + // Generate authorization URL (then override redirect_uri to reuse server port) + authURL, state, err := anthropicAuth.GenerateAuthURL(state, pkceCodes) + if err != nil { + log.Errorf("Failed to generate authorization URL: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) + return + } + + RegisterOAuthSession(state, "anthropic") + + isWebUI := isWebUIRequest(c) + var forwarder *callbackForwarder + if isWebUI { + targetURL, errTarget := h.managementCallbackURL("/anthropic/callback") + if errTarget != nil { + log.WithError(errTarget).Error("failed to compute anthropic callback target") + c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"}) + return + } + var errStart error + if forwarder, errStart = startCallbackForwarder(anthropicCallbackPort, "anthropic", targetURL); errStart != nil { + log.WithError(errStart).Error("failed to start anthropic callback forwarder") + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"}) + return + } + } + + go func() { + if isWebUI { + defer stopCallbackForwarderInstance(anthropicCallbackPort, forwarder) + } + + // Helper: wait for callback file + waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-anthropic-%s.oauth", state)) + waitForFile := func(path string, timeout time.Duration) (map[string]string, error) { + deadline := time.Now().Add(timeout) + for { + if !IsOAuthSessionPending(state, "anthropic") { + return nil, errOAuthSessionNotPending + } + if time.Now().After(deadline) { + SetOAuthSessionError(state, "Timeout waiting for OAuth callback") + return nil, fmt.Errorf("timeout waiting for OAuth callback") + } + data, errRead := os.ReadFile(path) + if errRead == nil { + var m map[string]string + _ = json.Unmarshal(data, &m) + _ = os.Remove(path) + return m, nil + } + time.Sleep(500 * time.Millisecond) + } + } + + fmt.Println("Waiting for authentication callback...") + // Wait up to 5 minutes + resultMap, errWait := waitForFile(waitFile, 5*time.Minute) + if errWait != nil { + if errors.Is(errWait, errOAuthSessionNotPending) { + return + } + authErr := claude.NewAuthenticationError(claude.ErrCallbackTimeout, errWait) + log.Error(claude.GetUserFriendlyMessage(authErr)) + return + } + if errStr := resultMap["error"]; errStr != "" { + oauthErr := claude.NewOAuthError(errStr, "", http.StatusBadRequest) + log.Error(claude.GetUserFriendlyMessage(oauthErr)) + SetOAuthSessionError(state, "Bad request") + return + } + if resultMap["state"] != state { + authErr := claude.NewAuthenticationError(claude.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, resultMap["state"])) + log.Error(claude.GetUserFriendlyMessage(authErr)) + SetOAuthSessionError(state, "State code error") + return + } + + // Parse code (Claude may append state after '#') + rawCode := resultMap["code"] + code := strings.Split(rawCode, "#")[0] + + // Exchange code for tokens using internal auth service + bundle, errExchange := anthropicAuth.ExchangeCodeForTokens(ctx, code, state, pkceCodes) + if errExchange != nil { + authErr := claude.NewAuthenticationError(claude.ErrCodeExchangeFailed, errExchange) + log.Errorf("Failed to exchange authorization code for tokens: %v", authErr) + SetOAuthSessionError(state, "Failed to exchange authorization code for tokens") + return + } + + // Create token storage + tokenStorage := anthropicAuth.CreateTokenStorage(bundle) + record := &coreauth.Auth{ + ID: fmt.Sprintf("claude-%s.json", tokenStorage.Email), + Provider: "claude", + FileName: fmt.Sprintf("claude-%s.json", tokenStorage.Email), + Storage: tokenStorage, + Metadata: map[string]any{"email": tokenStorage.Email}, + } + if errGuard := guardOAuthSessionPendingForSave(state, "anthropic"); errGuard != nil { + return + } + savedPath, errSave := h.saveTokenRecord(ctx, record) + if errSave != nil { + log.Errorf("Failed to save authentication tokens: %v", errSave) + SetOAuthSessionError(state, "Failed to save authentication tokens") + return + } + + fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) + if bundle.APIKey != "" { + fmt.Println("API key obtained and saved") + } + fmt.Println("You can now use Claude services through this CLI") + CompleteOAuthSession(state) + }() + + c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) +} + +func (h *Handler) RequestCodexToken(c *gin.Context) { + ctx := context.Background() + ctx = PopulateAuthContext(ctx, c) + + fmt.Println("Initializing Codex authentication...") + + // Generate PKCE codes + pkceCodes, err := codex.GeneratePKCECodes() + if err != nil { + log.Errorf("Failed to generate PKCE codes: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate PKCE codes"}) + return + } + + // Generate random state parameter + state, err := misc.GenerateRandomState() + if err != nil { + log.Errorf("Failed to generate state parameter: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"}) + return + } + + // Initialize Codex auth service + openaiAuth := newCodexOAuthService(h.cfg) + + // Generate authorization URL + authURL, err := openaiAuth.GenerateAuthURL(state, pkceCodes) + if err != nil { + log.Errorf("Failed to generate authorization URL: %v", err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) + return + } + + RegisterOAuthSession(state, "codex") + + isWebUI := isWebUIRequest(c) + var forwarder *callbackForwarder + if isWebUI { + targetURL, errTarget := h.managementCallbackURL("/codex/callback") + if errTarget != nil { + log.WithError(errTarget).Error("failed to compute codex callback target") + c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"}) + return + } + var errStart error + if forwarder, errStart = startCallbackForwarder(codexCallbackPort, "codex", targetURL); errStart != nil { + log.WithError(errStart).Error("failed to start codex callback forwarder") + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"}) + return + } + } + + go func() { + if isWebUI { + defer stopCallbackForwarderInstance(codexCallbackPort, forwarder) + } + + // Wait for callback file + waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-codex-%s.oauth", state)) + deadline := time.Now().Add(5 * time.Minute) + var code string + for { + if !IsOAuthSessionPending(state, "codex") { + return + } + if time.Now().After(deadline) { + authErr := codex.NewAuthenticationError(codex.ErrCallbackTimeout, fmt.Errorf("timeout waiting for OAuth callback")) + log.Error(codex.GetUserFriendlyMessage(authErr)) + SetOAuthSessionError(state, "Timeout waiting for OAuth callback") + return + } + if data, errR := os.ReadFile(waitFile); errR == nil { + var m map[string]string + _ = json.Unmarshal(data, &m) + _ = os.Remove(waitFile) + if errStr := m["error"]; errStr != "" { + oauthErr := codex.NewOAuthError(errStr, "", http.StatusBadRequest) + log.Error(codex.GetUserFriendlyMessage(oauthErr)) + SetOAuthSessionError(state, "Bad Request") + return + } + if m["state"] != state { + authErr := codex.NewAuthenticationError(codex.ErrInvalidState, fmt.Errorf("expected %s, got %s", state, m["state"])) + SetOAuthSessionError(state, "State code error") + log.Error(codex.GetUserFriendlyMessage(authErr)) + return + } + code = m["code"] + break + } + time.Sleep(500 * time.Millisecond) + } + + log.Debug("Authorization code received, exchanging for tokens...") + // Exchange code for tokens using internal auth service + bundle, errExchange := openaiAuth.ExchangeCodeForTokens(ctx, code, pkceCodes) + if errExchange != nil { + authErr := codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, errExchange) + SetOAuthSessionError(state, oauthSessionErrorWithCause("Failed to exchange authorization code for tokens", errExchange)) + log.Errorf("Failed to exchange authorization code for tokens: %v", authErr) + return + } + + // Extract additional info for filename generation + claims, _ := codex.ParseJWTToken(bundle.TokenData.IDToken) + planType := "" + hashAccountID := "" + if claims != nil { + planType = strings.TrimSpace(claims.CodexAuthInfo.ChatgptPlanType) + if accountID := claims.GetAccountID(); accountID != "" { + digest := sha256.Sum256([]byte(accountID)) + hashAccountID = hex.EncodeToString(digest[:])[:8] + } + } + + // Create token storage and persist + tokenStorage := openaiAuth.CreateTokenStorage(bundle) + fileName := codex.CredentialFileName(tokenStorage.Email, planType, hashAccountID, true) + record := &coreauth.Auth{ + ID: fileName, + Provider: "codex", + FileName: fileName, + Storage: tokenStorage, + Metadata: map[string]any{ + "email": tokenStorage.Email, + "account_id": tokenStorage.AccountID, + }, + } + if errGuard := guardOAuthSessionPendingForSave(state, "codex"); errGuard != nil { + return + } + savedPath, errSave := h.saveTokenRecord(ctx, record) + if errSave != nil { + SetOAuthSessionError(state, "Failed to save authentication tokens") + log.Errorf("Failed to save authentication tokens: %v", errSave) + return + } + fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) + if bundle.APIKey != "" { + fmt.Println("API key obtained and saved") + } + fmt.Println("You can now use Codex services through this CLI") + CompleteOAuthSession(state) + }() + + c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) +} + +func (h *Handler) RequestAntigravityToken(c *gin.Context) { + ctx := context.Background() + ctx = PopulateAuthContext(ctx, c) + + fmt.Println("Initializing Antigravity authentication...") + + authSvc := antigravity.NewAntigravityAuth(h.cfg, nil) + + state, errState := misc.GenerateRandomState() + if errState != nil { + log.Errorf("Failed to generate state parameter: %v", errState) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate state parameter"}) + return + } + + redirectURI := fmt.Sprintf("http://localhost:%d/oauth-callback", antigravity.CallbackPort) + authURL := authSvc.BuildAuthURL(state, redirectURI) + + RegisterOAuthSession(state, "antigravity") + + isWebUI := isWebUIRequest(c) + var forwarder *callbackForwarder + if isWebUI { + targetURL, errTarget := h.managementCallbackURL("/antigravity/callback") + if errTarget != nil { + log.WithError(errTarget).Error("failed to compute antigravity callback target") + c.JSON(http.StatusInternalServerError, gin.H{"error": "callback server unavailable"}) + return + } + var errStart error + if forwarder, errStart = startCallbackForwarder(antigravity.CallbackPort, "antigravity", targetURL); errStart != nil { + log.WithError(errStart).Error("failed to start antigravity callback forwarder") + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start callback server"}) + return + } + } + + go func() { + if isWebUI { + defer stopCallbackForwarderInstance(antigravity.CallbackPort, forwarder) + } + + waitFile := filepath.Join(h.cfg.AuthDir, fmt.Sprintf(".oauth-antigravity-%s.oauth", state)) + deadline := time.Now().Add(5 * time.Minute) + var authCode string + for { + if !IsOAuthSessionPending(state, "antigravity") { + return + } + if time.Now().After(deadline) { + log.Error("oauth flow timed out") + SetOAuthSessionError(state, "OAuth flow timed out") + return + } + if data, errReadFile := os.ReadFile(waitFile); errReadFile == nil { + var payload map[string]string + _ = json.Unmarshal(data, &payload) + _ = os.Remove(waitFile) + if errStr := strings.TrimSpace(payload["error"]); errStr != "" { + log.Errorf("Authentication failed: %s", errStr) + SetOAuthSessionError(state, "Authentication failed") + return + } + if payloadState := strings.TrimSpace(payload["state"]); payloadState != "" && payloadState != state { + log.Errorf("Authentication failed: state mismatch") + SetOAuthSessionError(state, "Authentication failed: state mismatch") + return + } + authCode = strings.TrimSpace(payload["code"]) + if authCode == "" { + log.Error("Authentication failed: code not found") + SetOAuthSessionError(state, "Authentication failed: code not found") + return + } + break + } + time.Sleep(500 * time.Millisecond) + } + + tokenResp, errToken := authSvc.ExchangeCodeForTokens(ctx, authCode, redirectURI) + if errToken != nil { + log.Errorf("Failed to exchange token: %v", errToken) + SetOAuthSessionError(state, "Failed to exchange token") + return + } + + accessToken := strings.TrimSpace(tokenResp.AccessToken) + if accessToken == "" { + log.Error("antigravity: token exchange returned empty access token") + SetOAuthSessionError(state, "Failed to exchange token") + return + } + + email, errInfo := authSvc.FetchUserInfo(ctx, accessToken) + if errInfo != nil { + log.Errorf("Failed to fetch user info: %v", errInfo) + SetOAuthSessionError(state, "Failed to fetch user info") + return + } + email = strings.TrimSpace(email) + if email == "" { + log.Error("antigravity: user info returned empty email") + SetOAuthSessionError(state, "Failed to fetch user info") + return + } + + projectID := "" + if accessToken != "" { + fetchedProjectID, errProject := authSvc.FetchProjectID(ctx, accessToken) + if errProject != nil { + log.Warnf("antigravity: failed to fetch project ID: %v", errProject) + } else { + projectID = fetchedProjectID + log.Infof("antigravity: obtained project ID %s", util.HideAPIKey(projectID)) + } + } + + now := time.Now() + metadata := map[string]any{ + "type": "antigravity", + "access_token": tokenResp.AccessToken, + "refresh_token": tokenResp.RefreshToken, + "expires_in": tokenResp.ExpiresIn, + "timestamp": now.UnixMilli(), + "expired": now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339), + } + if email != "" { + metadata["email"] = email + } + if projectID != "" { + metadata["project_id"] = projectID + } + + fileName := antigravity.CredentialFileName(email) + label := strings.TrimSpace(email) + if label == "" { + label = "antigravity" + } + + record := &coreauth.Auth{ + ID: fileName, + Provider: "antigravity", + FileName: fileName, + Label: label, + Metadata: metadata, + } + if errGuard := guardOAuthSessionPendingForSave(state, "antigravity"); errGuard != nil { + return + } + savedPath, errSave := h.saveTokenRecord(ctx, record) + if errSave != nil { + log.Errorf("Failed to save token to file: %v", errSave) + SetOAuthSessionError(state, "Failed to save token to file") + return + } + + CompleteOAuthSession(state) + fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) + if projectID != "" { + fmt.Printf("Using GCP project: %s\n", util.HideAPIKey(projectID)) + } + fmt.Println("You can now use Antigravity services through this CLI") + }() + + c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) +} + +func (h *Handler) RequestXAIToken(c *gin.Context) { + ctx := context.Background() + ctx = PopulateAuthContext(ctx, c) + + fmt.Println("Initializing xAI authentication...") + + state := fmt.Sprintf("xai-%d", time.Now().UnixNano()) + authSvc := xaiauth.NewXAIAuth(h.cfg) + + deviceFlow, errStartDeviceFlow := authSvc.StartDeviceFlow(ctx) + if errStartDeviceFlow != nil { + log.Errorf("Failed to start xAI device flow: %v", errStartDeviceFlow) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to start device authorization flow"}) + return + } + authURL := strings.TrimSpace(deviceFlow.VerificationURIComplete) + if authURL == "" { + authURL = strings.TrimSpace(deviceFlow.VerificationURI) + } + + RegisterOAuthSession(state, "xai") + + go func() { + pollCtx, cancelPoll := context.WithCancel(ctx) + defer cancelPoll() + go watchOAuthSessionCancel(pollCtx, cancelPoll, state, "xai") + + fmt.Println("Waiting for xAI authentication...") + bundle, errWaitForAuthorization := authSvc.WaitForAuthorization(pollCtx, deviceFlow) + if errWaitForAuthorization != nil { + if !IsOAuthSessionPending(state, "xai") { + return + } + log.Errorf("xAI authentication failed: %v", errWaitForAuthorization) + SetOAuthSessionError(state, oauthSessionErrorWithCause("Authentication failed", errWaitForAuthorization)) + return + } + if !IsOAuthSessionPending(state, "xai") { + return + } + + tokenStorage := authSvc.CreateTokenStorage(bundle) + if tokenStorage == nil || strings.TrimSpace(tokenStorage.AccessToken) == "" { + log.Error("xAI token exchange returned empty access token") + SetOAuthSessionError(state, "Failed to exchange token") + return + } + + fileName := xaiauth.CredentialFileName(tokenStorage.Email, tokenStorage.Subject) + label := strings.TrimSpace(tokenStorage.Email) + if label == "" { + label = "xAI" + } + + metadata := map[string]any{ + "type": "xai", + "access_token": tokenStorage.AccessToken, + "refresh_token": tokenStorage.RefreshToken, + "id_token": tokenStorage.IDToken, + "token_type": tokenStorage.TokenType, + "expires_in": tokenStorage.ExpiresIn, + "expired": tokenStorage.Expire, + "last_refresh": tokenStorage.LastRefresh, + "base_url": tokenStorage.BaseURL, + "token_endpoint": tokenStorage.TokenEndpoint, + "auth_kind": "oauth", + } + if tokenStorage.Email != "" { + metadata["email"] = tokenStorage.Email + } + if tokenStorage.Subject != "" { + metadata["sub"] = tokenStorage.Subject + } + + record := &coreauth.Auth{ + ID: fileName, + Provider: "xai", + FileName: fileName, + Label: label, + Storage: tokenStorage, + Metadata: metadata, + Attributes: map[string]string{ + "auth_kind": "oauth", + "base_url": tokenStorage.BaseURL, + }, + } + if errGuard := guardOAuthSessionPendingForSave(state, "xai"); errGuard != nil { + return + } + savedPath, errSave := h.saveTokenRecord(ctx, record) + if errSave != nil { + log.Errorf("Failed to save xAI token to file: %v", errSave) + SetOAuthSessionError(state, "Failed to save token to file") + return + } + + CompleteOAuthSession(state) + fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) + fmt.Println("You can now use xAI services through this CLI") + }() + + response := gin.H{"status": "ok", "url": authURL, "state": state, "flow": "device"} + if userCode := strings.TrimSpace(deviceFlow.UserCode); userCode != "" { + response["user_code"] = userCode + } + if deviceFlow.ExpiresIn > 0 { + response["expires_in"] = deviceFlow.ExpiresIn + } else { + response["expires_in"] = int(xaiauth.MaxPollDuration / time.Second) + } + c.JSON(200, response) +} + +func (h *Handler) RequestKimiToken(c *gin.Context) { + ctx := context.Background() + ctx = PopulateAuthContext(ctx, c) + + fmt.Println("Initializing Kimi authentication...") + + state := fmt.Sprintf("kmi-%d", time.Now().UnixNano()) + // Initialize Kimi auth service + kimiAuth := kimi.NewKimiAuth(h.cfg) + + // Generate authorization URL + deviceFlow, errStartDeviceFlow := kimiAuth.StartDeviceFlow(ctx) + if errStartDeviceFlow != nil { + log.Errorf("Failed to generate authorization URL: %v", errStartDeviceFlow) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate authorization url"}) + return + } + authURL := deviceFlow.VerificationURIComplete + if authURL == "" { + authURL = deviceFlow.VerificationURI + } + + RegisterOAuthSession(state, "kimi") + + go func() { + pollCtx, cancelPoll := context.WithCancel(ctx) + defer cancelPoll() + go watchOAuthSessionCancel(pollCtx, cancelPoll, state, "kimi") + + fmt.Println("Waiting for authentication...") + authBundle, errWaitForAuthorization := kimiAuth.WaitForAuthorization(pollCtx, deviceFlow) + if errWaitForAuthorization != nil { + if !IsOAuthSessionPending(state, "kimi") { + return + } + SetOAuthSessionError(state, oauthSessionErrorWithCause("Authentication failed", errWaitForAuthorization)) + fmt.Printf("Authentication failed: %v\n", errWaitForAuthorization) + return + } + if !IsOAuthSessionPending(state, "kimi") { + return + } + + // Create token storage + tokenStorage := kimiAuth.CreateTokenStorage(authBundle) + + metadata := map[string]any{ + "type": "kimi", + "access_token": authBundle.TokenData.AccessToken, + "refresh_token": authBundle.TokenData.RefreshToken, + "token_type": authBundle.TokenData.TokenType, + "scope": authBundle.TokenData.Scope, + "timestamp": time.Now().UnixMilli(), + } + if authBundle.TokenData.ExpiresAt > 0 { + expired := time.Unix(authBundle.TokenData.ExpiresAt, 0).UTC().Format(time.RFC3339) + metadata["expired"] = expired + } + if strings.TrimSpace(authBundle.DeviceID) != "" { + metadata["device_id"] = strings.TrimSpace(authBundle.DeviceID) + } + + fileName := fmt.Sprintf("kimi-%d.json", time.Now().UnixMilli()) + record := &coreauth.Auth{ + ID: fileName, + Provider: "kimi", + FileName: fileName, + Label: "Kimi User", + Storage: tokenStorage, + Metadata: metadata, + } + if errGuard := guardOAuthSessionPendingForSave(state, "kimi"); errGuard != nil { + return + } + savedPath, errSave := h.saveTokenRecord(ctx, record) + if errSave != nil { + log.Errorf("Failed to save authentication tokens: %v", errSave) + SetOAuthSessionError(state, "Failed to save authentication tokens") + return + } + + fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) + fmt.Println("You can now use Kimi services through this CLI") + CompleteOAuthSession(state) + }() + + response := gin.H{"status": "ok", "url": authURL, "state": state, "flow": "device"} + if userCode := strings.TrimSpace(deviceFlow.UserCode); userCode != "" { + response["user_code"] = userCode + } + if deviceFlow.ExpiresIn > 0 { + response["expires_in"] = deviceFlow.ExpiresIn + } + c.JSON(200, response) +} + +// watchOAuthSessionCancel cancels pollCtx once the OAuth session is no longer pending. +func watchOAuthSessionCancel(pollCtx context.Context, cancel context.CancelFunc, state, provider string) { + if cancel == nil { + return + } + ticker := time.NewTicker(2 * time.Second) + defer ticker.Stop() + for { + select { + case <-pollCtx.Done(): + return + case <-ticker.C: + if !IsOAuthSessionPending(state, provider) { + cancel() + return + } + } + } +} + +// CancelAuthSession cancels a pending OAuth session identified by state. +// Protected by management auth. Safe for both callback and device-code flows: +// waiters check IsOAuthSessionPending and exit without saving credentials. +func (h *Handler) CancelAuthSession(c *gin.Context) { + state := strings.TrimSpace(c.Query("state")) + if state == "" { + c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "missing state"}) + return + } + if err := ValidateOAuthState(state); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid state"}) + return + } + cancelled := CancelOAuthSession(state) + c.JSON(http.StatusOK, gin.H{"status": "ok", "cancelled": cancelled}) +} + +func (h *Handler) GetAuthStatus(c *gin.Context) { + state := strings.TrimSpace(c.Query("state")) + if state == "" { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + return + } + if err := ValidateOAuthState(state); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"status": "error", "error": "invalid state"}) + return + } + + provider, status, isPlugin, metadata, completed, ok := GetOAuthSessionDetails(state) + if !ok { + c.JSON(http.StatusOK, gin.H{"status": "error", "error": "unknown or expired state"}) + return + } + if completed { + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + return + } + if status != "" { + c.JSON(http.StatusOK, gin.H{"status": "error", "error": status}) + return + } + h.mu.Lock() + host := h.pluginHost + h.mu.Unlock() + if isPlugin && host != nil && host.HasAuthProvider(provider) { + ctx := PopulateAuthContext(context.Background(), c) + resp, handled, errPoll := host.PollLogin(ctx, provider, state, metadata) + if handled { + if errPoll != nil { + message := strings.TrimSpace(errPoll.Error()) + if message == "" { + message = "Authentication failed" + } + SetOAuthSessionError(state, message) + c.JSON(http.StatusOK, gin.H{"status": "error", "error": message}) + return + } + switch resp.Status { + case "", pluginapi.AuthLoginStatusPending: + c.JSON(http.StatusOK, gin.H{"status": "wait"}) + return + case pluginapi.AuthLoginStatusError: + message := strings.TrimSpace(resp.Message) + if message == "" { + message = "Authentication failed" + } + SetOAuthSessionError(state, message) + c.JSON(http.StatusOK, gin.H{"status": "error", "error": message}) + return + case pluginapi.AuthLoginStatusSuccess: + records := pluginLoginPollAuths(host, resp) + if len(records) == 0 { + SetOAuthSessionError(state, "Authentication failed") + c.JSON(http.StatusOK, gin.H{"status": "error", "error": "Authentication failed"}) + return + } + if errSave := h.savePluginLoginRecords(ctx, records); errSave != nil { + log.WithError(errSave).WithField("provider", provider).Error("failed to save plugin auth tokens") + SetOAuthSessionError(state, "Failed to save authentication tokens") + c.JSON(http.StatusOK, gin.H{"status": "error", "error": "Failed to save authentication tokens"}) + return + } + CompleteOAuthSession(state) + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + return + default: + c.JSON(http.StatusOK, gin.H{"status": "wait"}) + return + } + } + } + c.JSON(http.StatusOK, gin.H{"status": "wait"}) +} + +func pluginLoginPollAuths(host *pluginhost.Host, resp pluginapi.AuthLoginPollResponse) []*coreauth.Auth { + if host == nil { + return nil + } + authDatas := resp.Auths + if len(authDatas) == 0 { + authDatas = []pluginapi.AuthData{resp.Auth} + } + records := make([]*coreauth.Auth, 0, len(authDatas)) + for _, authData := range authDatas { + record := host.AuthDataToCoreAuth(authData, "", "") + if record == nil { + return nil + } + records = append(records, record) + } + return records +} + +func (h *Handler) savePluginLoginRecords(ctx context.Context, records []*coreauth.Auth) error { + savedPaths := make([]string, 0, len(records)) + for _, record := range records { + savedPath, errSave := h.saveTokenRecord(ctx, record) + if strings.TrimSpace(savedPath) != "" { + savedPaths = append(savedPaths, savedPath) + } + if errSave != nil { + h.rollbackSavedTokenRecords(ctx, savedPaths) + return errSave + } + } + return nil +} + +func (h *Handler) rollbackSavedTokenRecords(ctx context.Context, savedPaths []string) { + for i := len(savedPaths) - 1; i >= 0; i-- { + path := strings.TrimSpace(savedPaths[i]) + if path == "" { + continue + } + if errDelete := h.deleteTokenRecord(ctx, path); errDelete != nil { + log.WithError(errDelete).WithField("path", path).Warn("failed to roll back plugin auth token") + } + h.removeAuthsForPath(ctx, path, path) + } +} + +// PopulateAuthContext extracts request info and adds it to the context +func PopulateAuthContext(ctx context.Context, c *gin.Context) context.Context { + info := &coreauth.RequestInfo{ + Query: c.Request.URL.Query(), + Headers: c.Request.Header, + } + return coreauth.WithRequestInfo(ctx, info) +} diff --git a/internal/api/server.go b/internal/api/server.go index f263e82fb..f12ff4bf4 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -6,199 +6,34 @@ package api import ( "context" - "crypto/subtle" "crypto/tls" - "encoding/json" "errors" "fmt" - "io" "net" "net/http" "os" - "path/filepath" - "sort" - "strconv" "strings" "sync" "sync/atomic" "time" "github.com/gin-gonic/gin" - "github.com/router-for-me/CLIProxyAPI/v7/internal/access" managementHandlers "github.com/router-for-me/CLIProxyAPI/v7/internal/api/handlers/management" "github.com/router-for-me/CLIProxyAPI/v7/internal/api/middleware" - "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" - claudemodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/claude/models" codexlive "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/live" - codexmodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/models" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/managementasset" "github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost" "github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" - "github.com/router-for-me/CLIProxyAPI/v7/internal/safemode" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers/claude" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers/gemini" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers/openai" - sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" - coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" log "github.com/sirupsen/logrus" "golang.org/x/net/http2" "gopkg.in/yaml.v3" ) -const oauthCallbackSuccessHTML = `Authentication successful

Authentication successful!

You can close this window.

This window will close automatically in 5 seconds.

` - -const codexAlphaSearchSourceFormat = "codex-alpha-search" - -var corsExposedResponseHeaders = []string{ - logging.CPATraceIDHeader, - "X-CPA-VERSION", - "X-CPA-COMMIT", - "X-CPA-BUILD-DATE", - "X-CPA-SUPPORT-PLUGIN", - "X-CPA-HOME-VERSION", - "X-CPA-HOME-BUILD-DATE", - "X-SERVER-VERSION", - "X-SERVER-BUILD-DATE", -} - -var corsExposedResponseHeadersJoined = strings.Join(corsExposedResponseHeaders, ", ") - -const ( - exampleAPIKeyManagementPath = "/management.html" - exampleAPIKeyManagementURL = "/management.html?safe-mode=configure" -) - -type serverOptionConfig struct { - extraMiddleware []gin.HandlerFunc - engineConfigurator func(*gin.Engine) - routerConfigurator func(*gin.Engine, *handlers.BaseAPIHandler, *config.Config) - requestLoggerFactory func(*config.Config, string) logging.RequestLogger - localPassword string - keepAliveEnabled bool - keepAliveTimeout time.Duration - keepAliveOnTimeout func() - postAuthHook auth.PostAuthHook - postAuthPersistHook auth.PostAuthHook - pluginHost *pluginhost.Host - configReloadHook func(context.Context, *config.Config) - exampleAPIKeySafeMode bool -} - -// ServerOption customises HTTP server construction. -type ServerOption func(*serverOptionConfig) - -func defaultRequestLoggerFactory(cfg *config.Config, configPath string) logging.RequestLogger { - configDir := filepath.Dir(configPath) - logsDir := logging.ResolveLogDirectory(cfg) - logger := logging.NewFileRequestLogger(cfg.RequestLog, logsDir, configDir, cfg.ErrorLogsMaxFiles) - logger.SetHomeEnabled(cfg != nil && cfg.Home.Enabled) - return logger -} - -func effectiveSDKConfig(cfg *config.Config) *config.SDKConfig { - if cfg == nil { - return nil - } - sdkCfg := cfg.SDKConfig - sdkCfg.CodexOptimizeMultiAgentV2 = cfg.Codex.OptimizeMultiAgentV2 - if cfg.CommercialMode { - sdkCfg.RequestLog = false - } - return &sdkCfg -} - -// WithMiddleware appends additional Gin middleware during server construction. -func WithMiddleware(mw ...gin.HandlerFunc) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.extraMiddleware = append(cfg.extraMiddleware, mw...) - } -} - -// WithEngineConfigurator allows callers to mutate the Gin engine prior to middleware setup. -func WithEngineConfigurator(fn func(*gin.Engine)) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.engineConfigurator = fn - } -} - -// WithRouterConfigurator appends a callback after default routes are registered. -func WithRouterConfigurator(fn func(*gin.Engine, *handlers.BaseAPIHandler, *config.Config)) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.routerConfigurator = fn - } -} - -// WithLocalManagementPassword stores a runtime-only management password accepted for localhost requests. -func WithLocalManagementPassword(password string) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.localPassword = password - } -} - -// WithKeepAliveEndpoint enables a keep-alive endpoint with the provided timeout and callback. -func WithKeepAliveEndpoint(timeout time.Duration, onTimeout func()) ServerOption { - return func(cfg *serverOptionConfig) { - if timeout <= 0 || onTimeout == nil { - return - } - cfg.keepAliveEnabled = true - cfg.keepAliveTimeout = timeout - cfg.keepAliveOnTimeout = onTimeout - } -} - -// WithRequestLoggerFactory customises request logger creation. -func WithRequestLoggerFactory(factory func(*config.Config, string) logging.RequestLogger) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.requestLoggerFactory = factory - } -} - -// WithPostAuthHook registers a hook to be called after auth record creation. -func WithPostAuthHook(hook auth.PostAuthHook) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.postAuthHook = hook - } -} - -// WithPostAuthPersistHook registers a hook to be called after auth persistence. -func WithPostAuthPersistHook(hook auth.PostAuthHook) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.postAuthPersistHook = hook - } -} - -// WithPluginHost registers dynamic plugin HTTP adapters with the server. -func WithPluginHost(host *pluginhost.Host) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.pluginHost = host - } -} - -// WithConfigReloadHook registers a callback used after management saves config changes. -func WithConfigReloadHook(hook func(context.Context, *config.Config)) ServerOption { - return func(cfg *serverOptionConfig) { - cfg.configReloadHook = hook - } -} - -// WithExampleAPIKeySafeMode blocks proxy API endpoints while template API keys remain configured. -func WithExampleAPIKeySafeMode() ServerOption { - return func(cfg *serverOptionConfig) { - cfg.exampleAPIKeySafeMode = true - } -} - // Server represents the main API server. // It encapsulates the Gin engine, HTTP server, handlers, and configuration. type Server struct { @@ -419,1331 +254,6 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk return s } -func (s *Server) homeHeartbeatMiddleware() gin.HandlerFunc { - return func(c *gin.Context) { - if s == nil || s.cfg == nil || !s.cfg.Home.Enabled { - c.Next() - return - } - if c != nil && c.Request != nil { - path := c.Request.URL.Path - if strings.HasPrefix(path, "/v0/management/") || path == "/v0/management" || strings.HasPrefix(path, "/v0/resource/plugins/") || path == "/management.html" { - c.Next() - return - } - } - client := home.Current() - if client == nil || !client.HeartbeatOK() { - c.AbortWithStatus(http.StatusServiceUnavailable) - return - } - c.Next() - } -} - -func (s *Server) exampleAPIKeySafeModeRequired(cfg *config.Config) bool { - return s != nil && s.exampleAPIKeySafeModeEnabled && cfg != nil && safemode.HasExampleAPIKeys(cfg.APIKeys) -} - -func (s *Server) exampleAPIKeySafeModeMiddleware() gin.HandlerFunc { - return func(c *gin.Context) { - if s == nil || !s.exampleAPIKeySafeModeActive.Load() || c == nil || c.Request == nil || c.Request.URL == nil { - c.Next() - return - } - - path := c.Request.URL.Path - if path == exampleAPIKeyManagementPath && c.Query("safe-mode") == "configure" { - c.Next() - return - } - if (path == "/" || path == exampleAPIKeyManagementPath) && (c.Request.Method == http.MethodGet || c.Request.Method == http.MethodHead) { - s.serveExampleAPIKeyWarningPage(c) - return - } - if !isExampleAPIKeySafeModeProxyPath(path) { - c.Next() - return - } - - c.Header("X-CPA-SAFE-MODE", "example-api-key") - c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ - "error": "unsafe_example_api_key", - "message": "Proxy API endpoints are disabled because api-keys contains template values. Open /management.html?safe-mode=configure, update api-keys in Management, then retry.", - }) - } -} - -func (s *Server) serveExampleAPIKeyWarningPage(c *gin.Context) { - cfg := s.cfg - var keys []string - if cfg != nil { - keys = safemode.ExampleAPIKeys(cfg.APIKeys) - } - c.Header("Content-Type", "text/html; charset=utf-8") - c.Header("Cache-Control", "no-store") - if c.Request.Method == http.MethodHead { - c.Status(http.StatusOK) - c.Abort() - return - } - c.String(http.StatusOK, safemode.ExampleAPIKeyWarningPageHTML(keys, exampleAPIKeyManagementURL)) - c.Abort() -} - -func isExampleAPIKeySafeModeProxyPath(path string) bool { - switch { - case path == "/v1" || strings.HasPrefix(path, "/v1/"): - return true - case path == "/v1beta" || strings.HasPrefix(path, "/v1beta/"): - return true - case path == "/openai/v1" || strings.HasPrefix(path, "/openai/v1/"): - return true - case path == "/backend-api/codex" || strings.HasPrefix(path, "/backend-api/codex/"): - return true - default: - return false - } -} - -// setupRoutes configures the API routes for the server. -// It defines the endpoints and associates them with their respective handlers. -func (s *Server) setupRoutes() { - healthzHandler := func(c *gin.Context) { - if c.Request.Method == http.MethodHead { - c.Status(http.StatusOK) - return - } - - c.JSON(http.StatusOK, gin.H{"status": "ok"}) - } - s.engine.GET("/healthz", healthzHandler) - s.engine.HEAD("/healthz", healthzHandler) - - s.engine.GET("/management.html", s.serveManagementControlPanel) - openaiHandlers := openai.NewOpenAIAPIHandler(s.handlers) - geminiHandlers := gemini.NewGeminiAPIHandler(s.handlers) - claudeCodeHandlers := claude.NewClaudeCodeAPIHandler(s.handlers) - openaiResponsesHandlers := openai.NewOpenAIResponsesAPIHandler(s.handlers) - s.codexLiveHandler = codexlive.NewHandler(s.handlers.AuthManager, s.cfg) - - // OpenAI compatible API routes - v1 := s.engine.Group("/v1") - v1.Use(AuthMiddleware(s.accessManager)) - { - v1.GET("/models", s.unifiedModelsHandler(openaiHandlers, claudeCodeHandlers)) - v1.POST("/chat/completions", openaiHandlers.ChatCompletions) - v1.POST("/completions", openaiHandlers.Completions) - v1.POST("/images/generations", openaiHandlers.ImagesGenerations) - v1.POST("/images/edits", openaiHandlers.ImagesEdits) - v1.POST("/videos", openaiHandlers.XAIVideosGenerations) - v1.POST("/videos/generations", openaiHandlers.XAIVideosGenerations) - v1.POST("/videos/edits", openaiHandlers.XAIVideosEdits) - v1.POST("/videos/extensions", openaiHandlers.XAIVideosExtensions) - v1.GET("/videos/:request_id", openaiHandlers.XAIVideosRetrieve) - v1.POST("/messages", claudeCodeHandlers.ClaudeMessages) - v1.POST("/messages/count_tokens", claudeCodeHandlers.ClaudeCountTokens) - v1.GET("/responses", openaiResponsesHandlers.ResponsesWebsocket) - v1.POST("/responses", openaiResponsesHandlers.Responses) - v1.POST("/responses/compact", openaiResponsesHandlers.Compact) - v1.POST("/alpha/search", s.codexAlphaSearch) - v1.POST("/live", s.codexLiveHandler.Handle) - v1.GET("/live/:call_id", s.codexLiveHandler.HandleSideband) - v1.POST("/realtime/calls", s.codexLiveHandler.Handle) - v1.GET("/realtime/calls/:call_id", s.codexLiveHandler.HandleSideband) - v1.GET("/realtime", s.codexLiveHandler.HandleSideband) - } - - openaiV1 := s.engine.Group("/openai/v1") - openaiV1.Use(AuthMiddleware(s.accessManager)) - { - openaiV1.POST("/videos", openaiHandlers.VideosCreate) - openaiV1.GET("/videos/:video_id/content", openaiHandlers.VideosContent) - openaiV1.GET("/videos/:video_id", openaiHandlers.VideosRetrieve) - } - - // Codex CLI direct route aliases (chatgpt_base_url compatible) - codexDirect := s.engine.Group("/backend-api/codex") - codexDirect.Use(AuthMiddleware(s.accessManager)) - { - codexDirect.GET("/responses", openaiResponsesHandlers.ResponsesWebsocket) - codexDirect.POST("/responses", openaiResponsesHandlers.Responses) - codexDirect.POST("/responses/compact", openaiResponsesHandlers.Compact) - codexDirect.POST("/alpha/search", s.codexAlphaSearch) - } - - // Gemini compatible API routes - v1beta := s.engine.Group("/v1beta") - v1beta.Use(AuthMiddleware(s.accessManager)) - { - v1beta.GET("/models", s.geminiModelsHandler(geminiHandlers)) - v1beta.POST("/interactions", geminiHandlers.Interactions) - v1beta.POST("/models/*action", geminiHandlers.GeminiHandler) - v1beta.GET("/models/*action", s.geminiGetHandler(geminiHandlers)) - } - - // Root endpoint - s.engine.GET("/", func(c *gin.Context) { - c.JSON(http.StatusOK, gin.H{ - "message": "CLI Proxy API Server", - "endpoints": []string{ - "POST /v1/chat/completions", - "POST /v1/completions", - "GET /v1/models", - }, - }) - }) - - // OAuth callback endpoints (reuse main server port) - // These endpoints receive provider redirects and persist - // the short-lived code/state for the waiting goroutine. - s.engine.GET("/anthropic/callback", func(c *gin.Context) { - code := c.Query("code") - state := c.Query("state") - errStr := c.Query("error") - if errStr == "" { - errStr = c.Query("error_description") - } - if state != "" { - _, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "anthropic", state, code, errStr) - } - c.Header("Content-Type", "text/html; charset=utf-8") - c.String(http.StatusOK, oauthCallbackSuccessHTML) - }) - - s.engine.GET("/codex/callback", func(c *gin.Context) { - code := c.Query("code") - state := c.Query("state") - errStr := c.Query("error") - if errStr == "" { - errStr = c.Query("error_description") - } - if state != "" { - _, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "codex", state, code, errStr) - } - c.Header("Content-Type", "text/html; charset=utf-8") - c.String(http.StatusOK, oauthCallbackSuccessHTML) - }) - - s.engine.GET("/antigravity/callback", func(c *gin.Context) { - code := c.Query("code") - state := c.Query("state") - errStr := c.Query("error") - if errStr == "" { - errStr = c.Query("error_description") - } - if state != "" { - _, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "antigravity", state, code, errStr) - } - c.Header("Content-Type", "text/html; charset=utf-8") - c.String(http.StatusOK, oauthCallbackSuccessHTML) - }) - - // Management routes are registered lazily by registerManagementRoutes when a secret is configured. -} - -func (s *Server) codexAlphaSearchModelRouterHost() handlers.PluginModelRouterHost { - if s == nil { - return nil - } - if s.pluginHost != nil { - return s.pluginHost - } - if s.handlers != nil && s.handlers.ModelRouterHost != nil { - return s.handlers.ModelRouterHost - } - return nil -} - -func (s *Server) codexAlphaSearchSelectionModel(ctx context.Context, c *gin.Context, body []byte, model string) (string, error) { - host := s.codexAlphaSearchModelRouterHost() - if host == nil { - return model, nil - } - - var headers http.Header - queryValues := make(map[string][]string) - requestPath := "" - if c != nil && c.Request != nil { - headers = c.Request.Header.Clone() - if c.Request.URL != nil { - queryValues = c.Request.URL.Query() - requestPath = c.Request.URL.Path - } - } - metadata := map[string]any{ - coreexecutor.RequestedModelMetadataKey: model, - } - if requestPath != "" { - metadata[coreexecutor.RequestPathMetadataKey] = requestPath - } - resp, handled := host.RouteModel(ctx, pluginapi.ModelRouteRequest{ - SourceFormat: codexAlphaSearchSourceFormat, - RequestedModel: model, - Headers: headers, - Query: queryValues, - Body: body, - Metadata: metadata, - }) - if !handled || !resp.Handled { - return model, nil - } - if resp.TargetKind != pluginapi.ModelRouteTargetProvider || !strings.EqualFold(strings.TrimSpace(resp.Target), "codex") { - return "", fmt.Errorf("unsupported Codex Alpha Search model route target %q (%q)", resp.TargetKind, resp.Target) - } - if targetModel := strings.TrimSpace(resp.TargetModel); targetModel != "" { - return targetModel, nil - } - return model, nil -} - -func sanitizeCodexAlphaSearchBody(body []byte) []byte { - var payload map[string]json.RawMessage - if errUnmarshal := json.Unmarshal(body, &payload); errUnmarshal != nil || payload == nil { - return body - } - - removed := false - for _, field := range []string{"prompt_cache_key", "prompt_cache_retention"} { - if _, exists := payload[field]; exists { - delete(payload, field) - removed = true - } - } - if !removed { - return body - } - - sanitizedBody, errMarshal := json.Marshal(payload) - if errMarshal != nil { - return body - } - return sanitizedBody -} - -func homeSelectionAttemptContext(ctx context.Context, selection *auth.HomeDispatchSelection) (context.Context, func(), error) { - if selection == nil { - return nil, func() {}, errors.New("Home dispatch selection is nil") - } - return selection.AttemptContext(ctx) -} - -// codexAlphaSearch forwards the standalone search endpoint used by current -// Codex clients. Unlike /responses, this payload is already in Codex search -// format and must not pass through a protocol translator. -func (s *Server) codexAlphaSearch(c *gin.Context) { - if s == nil || s.handlers == nil || s.handlers.AuthManager == nil { - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "Codex auth manager unavailable"}) - return - } - - body, err := io.ReadAll(io.LimitReader(c.Request.Body, 16<<20)) - if err != nil { - c.JSON(http.StatusBadRequest, gin.H{"error": "Failed to read search request"}) - return - } - - var routing struct { - ID string `json:"id"` - Model string `json:"model"` - } - _ = json.Unmarshal(body, &routing) - upstreamRequestBody := sanitizeCodexAlphaSearchBody(body) - - selectionHeaders := c.Request.Header.Clone() - if sessionID := strings.TrimSpace(routing.ID); sessionID != "" { - selectionHeaders.Set("X-Session-ID", sessionID) - } - ctx := context.WithValue(c.Request.Context(), "gin", c) - selectionModel, errRoute := s.codexAlphaSearchSelectionModel(ctx, c, body, strings.TrimSpace(routing.Model)) - if errRoute != nil { - log.WithError(errRoute).Warn("codex alpha search: model router returned an unsupported target") - c.JSON(http.StatusServiceUnavailable, gin.H{"error": errRoute.Error()}) - return - } - selectionOpts := coreexecutor.Options{Headers: selectionHeaders, OriginalRequest: body} - var selection *auth.HomeDispatchSelection - var selected *auth.Auth - if s.handlers.AuthManager.HomeEnabled() { - selection, err = s.handlers.AuthManager.SelectHomeAuthByKind(ctx, "codex", selectionModel, auth.AuthKindOAuth, selectionOpts) - if selection != nil { - selected = selection.CloneAuth() - } - } else { - selected, err = s.handlers.AuthManager.SelectAuthByKind(ctx, "codex", selectionModel, auth.AuthKindOAuth, selectionOpts) - } - if err != nil { - status := http.StatusServiceUnavailable - if statusError, ok := err.(interface{ StatusCode() int }); ok && statusError.StatusCode() > 0 { - status = statusError.StatusCode() - } - for _, value := range auth.SafeResponseHeaders(err).Values("Retry-After") { - c.Writer.Header().Add("Retry-After", value) - } - c.JSON(status, gin.H{"error": err.Error()}) - return - } - if selected == nil { - if selection != nil { - selection.End("missing_auth") - } - c.JSON(http.StatusServiceUnavailable, gin.H{"error": "Codex auth unavailable"}) - return - } - var releaseAttempt func() - if selection != nil { - attemptCtx, release, errBind := homeSelectionAttemptContext(ctx, selection) - if errBind != nil { - selection.End("attempt_bind_failed") - c.JSON(http.StatusServiceUnavailable, gin.H{"error": errBind.Error()}) - return - } - ctx = attemptCtx - releaseAttempt = release - defer releaseAttempt() - } - logging.SetGinCPATraceID(c, selected.EnsureIndex()) - - headers := make(http.Header) - headers.Set("Content-Type", "application/json") - headers.Set("Accept", "application/json") - headers.Set("Originator", "codex_cli_rs") - for _, name := range []string{"Version", "User-Agent", "Session_id", "X-Client-Request-Id"} { - if value := strings.TrimSpace(c.GetHeader(name)); value != "" { - headers.Set(name, value) - } - } - if accountID, ok := selected.Metadata["account_id"].(string); ok && strings.TrimSpace(accountID) != "" { - headers.Set("Chatgpt-Account-Id", accountID) - } - - const upstreamURL = "https://chatgpt.com/backend-api/codex/alpha/search" - req, err := s.handlers.AuthManager.NewHttpRequest( - ctx, selected, http.MethodPost, upstreamURL, upstreamRequestBody, headers, - ) - if err != nil { - if selection != nil { - selection.End("request_build_failed") - } - c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) - return - } - - var authID, authLabel, authType, authValue string - if selected != nil { - authID = selected.ID - authLabel = selected.Label - authType, authValue = selected.AccountInfo() - } - helpHeaders := req.Header.Clone() - helps.RecordAPIRequest(ctx, s.cfg, helps.UpstreamRequestLog{ - URL: upstreamURL, - Method: http.MethodPost, - Headers: helpHeaders, - Body: upstreamRequestBody, - Provider: "codex", - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - if errCtx := ctx.Err(); errCtx != nil { - if selection != nil { - selection.End("attempt_canceled") - } - c.JSON(http.StatusRequestTimeout, gin.H{"error": errCtx.Error()}) - return - } - resp, err := s.handlers.AuthManager.HttpRequest(ctx, selected, req) - if err != nil { - if selection != nil { - selection.End("request_failed") - } - helps.RecordAPIResponseError(ctx, s.cfg, err) - c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) - return - } - closeResponseBody := func() error { - errClose := resp.Body.Close() - if errClose != nil { - log.Errorf("codex alpha search: close response body error: %v", errClose) - } - return errClose - } - if selection != nil { - if errBind := selection.Bind(closeResponseBody); errBind != nil { - selection.End("response_bind_failed") - c.JSON(http.StatusServiceUnavailable, gin.H{"error": errBind.Error()}) - return - } - defer selection.End("response_closed") - } else { - defer func() { _ = closeResponseBody() }() - } - helps.RecordAPIResponseMetadata(ctx, s.cfg, resp.StatusCode, resp.Header.Clone()) - upstreamBody, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20)) - if err != nil { - helps.RecordAPIResponseError(ctx, s.cfg, err) - c.JSON(http.StatusBadGateway, gin.H{"error": "Failed to read Codex search response"}) - return - } - helps.AppendAPIResponseChunk(ctx, s.cfg, upstreamBody) - if contentType := resp.Header.Get("Content-Type"); contentType != "" { - c.Header("Content-Type", contentType) - } - c.Status(resp.StatusCode) - _, _ = c.Writer.Write(upstreamBody) -} - -// AttachWebsocketRoute registers a websocket upgrade handler on the primary Gin engine. -// The handler is served as-is without additional middleware beyond the standard stack already configured. -func (s *Server) AttachWebsocketRoute(path string, handler http.Handler) { - if s == nil || s.engine == nil || handler == nil { - return - } - trimmed := strings.TrimSpace(path) - if trimmed == "" { - trimmed = "/v1/ws" - } - if !strings.HasPrefix(trimmed, "/") { - trimmed = "/" + trimmed - } - s.wsRouteMu.Lock() - if _, exists := s.wsRoutes[trimmed]; exists { - s.wsRouteMu.Unlock() - return - } - s.wsRoutes[trimmed] = struct{}{} - s.wsRouteMu.Unlock() - - authMiddleware := AuthMiddleware(s.accessManager) - conditionalAuth := func(c *gin.Context) { - if !s.wsAuthEnabled.Load() { - c.Next() - return - } - authMiddleware(c) - } - finalHandler := func(c *gin.Context) { - handler.ServeHTTP(c.Writer, c.Request) - c.Abort() - } - - s.engine.GET(trimmed, conditionalAuth, finalHandler) -} - -func (s *Server) registerManagementRoutes() { - if s == nil || s.engine == nil || s.mgmt == nil { - return - } - if !s.managementRoutesRegistered.CompareAndSwap(false, true) { - return - } - - log.Info("management routes registered after secret key configuration") - - s.engine.POST("/v0/management/oauth-callback", s.managementAvailabilityMiddleware(), s.mgmt.PostOAuthCallback) - s.engine.GET("/v0/management/oauth-callback", s.managementAvailabilityMiddleware(), s.mgmt.GetOAuthCallback) - - mgmt := s.engine.Group("/v0/management") - mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware()) - { - mgmt.GET("/config", s.mgmt.GetConfig) - mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML) - mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML) - mgmt.GET("/latest-version", s.mgmt.GetLatestVersion) - mgmt.GET("/plugins", s.mgmt.ListPlugins) - mgmt.GET("/plugin-store", s.mgmt.ListPluginStore) - mgmt.POST("/plugin-store/:id/install", s.mgmt.InstallPluginFromStore) - mgmt.DELETE("/plugins/:id", s.mgmt.DeletePlugin) - mgmt.PATCH("/plugins/:id/enabled", s.mgmt.PatchPluginEnabled) - mgmt.GET("/plugins/:id/config", s.mgmt.GetPluginConfig) - mgmt.PUT("/plugins/:id/config", s.mgmt.PutPluginConfig) - mgmt.PATCH("/plugins/:id/config", s.mgmt.PatchPluginConfig) - - mgmt.GET("/debug", s.mgmt.GetDebug) - mgmt.PUT("/debug", s.mgmt.PutDebug) - mgmt.PATCH("/debug", s.mgmt.PutDebug) - - mgmt.GET("/logging-to-file", s.mgmt.GetLoggingToFile) - mgmt.PUT("/logging-to-file", s.mgmt.PutLoggingToFile) - mgmt.PATCH("/logging-to-file", s.mgmt.PutLoggingToFile) - - mgmt.GET("/logs-max-total-size-mb", s.mgmt.GetLogsMaxTotalSizeMB) - mgmt.PUT("/logs-max-total-size-mb", s.mgmt.PutLogsMaxTotalSizeMB) - mgmt.PATCH("/logs-max-total-size-mb", s.mgmt.PutLogsMaxTotalSizeMB) - - mgmt.GET("/error-logs-max-files", s.mgmt.GetErrorLogsMaxFiles) - mgmt.PUT("/error-logs-max-files", s.mgmt.PutErrorLogsMaxFiles) - mgmt.PATCH("/error-logs-max-files", s.mgmt.PutErrorLogsMaxFiles) - - mgmt.GET("/usage-statistics-enabled", s.mgmt.GetUsageStatisticsEnabled) - mgmt.PUT("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled) - mgmt.PATCH("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled) - - mgmt.GET("/proxy-url", s.mgmt.GetProxyURL) - mgmt.PUT("/proxy-url", s.mgmt.PutProxyURL) - mgmt.PATCH("/proxy-url", s.mgmt.PutProxyURL) - mgmt.DELETE("/proxy-url", s.mgmt.DeleteProxyURL) - - mgmt.POST("/api-call", s.mgmt.APICall) - - mgmt.GET("/quota-exceeded/switch-project", s.mgmt.GetSwitchProject) - mgmt.PUT("/quota-exceeded/switch-project", s.mgmt.PutSwitchProject) - mgmt.PATCH("/quota-exceeded/switch-project", s.mgmt.PutSwitchProject) - - mgmt.GET("/quota-exceeded/switch-preview-model", s.mgmt.GetSwitchPreviewModel) - mgmt.PUT("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel) - mgmt.PATCH("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel) - mgmt.POST("/reset-quota", s.mgmt.ResetQuota) - - mgmt.GET("/api-keys", s.mgmt.GetAPIKeys) - mgmt.PUT("/api-keys", s.mgmt.PutAPIKeys) - mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys) - mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys) - mgmt.GET("/api-key-usage", s.mgmt.GetAPIKeyUsage) - mgmt.GET("/usage-queue", s.mgmt.GetUsageQueue) - - mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys) - mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys) - mgmt.PATCH("/gemini-api-key", s.mgmt.PatchGeminiKey) - mgmt.DELETE("/gemini-api-key", s.mgmt.DeleteGeminiKey) - - mgmt.GET("/interactions-api-key", s.mgmt.GetInteractionsKeys) - mgmt.PUT("/interactions-api-key", s.mgmt.PutInteractionsKeys) - mgmt.PATCH("/interactions-api-key", s.mgmt.PatchInteractionsKey) - mgmt.DELETE("/interactions-api-key", s.mgmt.DeleteInteractionsKey) - - mgmt.GET("/logs", s.mgmt.GetLogs) - mgmt.DELETE("/logs", s.mgmt.DeleteLogs) - mgmt.GET("/request-error-logs", s.mgmt.GetRequestErrorLogs) - mgmt.GET("/request-error-logs/:name", s.mgmt.DownloadRequestErrorLog) - mgmt.GET("/request-log-by-id/:id", s.mgmt.GetRequestLogByID) - mgmt.GET("/request-log", s.mgmt.GetRequestLog) - mgmt.PUT("/request-log", s.mgmt.PutRequestLog) - mgmt.PATCH("/request-log", s.mgmt.PutRequestLog) - mgmt.GET("/ws-auth", s.mgmt.GetWebsocketAuth) - mgmt.PUT("/ws-auth", s.mgmt.PutWebsocketAuth) - mgmt.PATCH("/ws-auth", s.mgmt.PutWebsocketAuth) - - mgmt.GET("/request-retry", s.mgmt.GetRequestRetry) - mgmt.PUT("/request-retry", s.mgmt.PutRequestRetry) - mgmt.PATCH("/request-retry", s.mgmt.PutRequestRetry) - mgmt.GET("/max-retry-interval", s.mgmt.GetMaxRetryInterval) - mgmt.PUT("/max-retry-interval", s.mgmt.PutMaxRetryInterval) - mgmt.PATCH("/max-retry-interval", s.mgmt.PutMaxRetryInterval) - - mgmt.GET("/force-model-prefix", s.mgmt.GetForceModelPrefix) - mgmt.PUT("/force-model-prefix", s.mgmt.PutForceModelPrefix) - mgmt.PATCH("/force-model-prefix", s.mgmt.PutForceModelPrefix) - - mgmt.GET("/routing/strategy", s.mgmt.GetRoutingStrategy) - mgmt.PUT("/routing/strategy", s.mgmt.PutRoutingStrategy) - mgmt.PATCH("/routing/strategy", s.mgmt.PutRoutingStrategy) - - mgmt.GET("/claude-api-key", s.mgmt.GetClaudeKeys) - mgmt.PUT("/claude-api-key", s.mgmt.PutClaudeKeys) - mgmt.PATCH("/claude-api-key", s.mgmt.PatchClaudeKey) - mgmt.DELETE("/claude-api-key", s.mgmt.DeleteClaudeKey) - - mgmt.GET("/codex-api-key", s.mgmt.GetCodexKeys) - mgmt.PUT("/codex-api-key", s.mgmt.PutCodexKeys) - mgmt.PATCH("/codex-api-key", s.mgmt.PatchCodexKey) - mgmt.DELETE("/codex-api-key", s.mgmt.DeleteCodexKey) - - mgmt.GET("/xai-api-key", s.mgmt.GetXAIKeys) - mgmt.PUT("/xai-api-key", s.mgmt.PutXAIKeys) - mgmt.PATCH("/xai-api-key", s.mgmt.PatchXAIKey) - mgmt.DELETE("/xai-api-key", s.mgmt.DeleteXAIKey) - - mgmt.GET("/openai-compatibility", s.mgmt.GetOpenAICompat) - mgmt.PUT("/openai-compatibility", s.mgmt.PutOpenAICompat) - mgmt.PATCH("/openai-compatibility", s.mgmt.PatchOpenAICompat) - mgmt.DELETE("/openai-compatibility", s.mgmt.DeleteOpenAICompat) - - mgmt.GET("/vertex-api-key", s.mgmt.GetVertexCompatKeys) - mgmt.PUT("/vertex-api-key", s.mgmt.PutVertexCompatKeys) - mgmt.PATCH("/vertex-api-key", s.mgmt.PatchVertexCompatKey) - mgmt.DELETE("/vertex-api-key", s.mgmt.DeleteVertexCompatKey) - - mgmt.GET("/oauth-excluded-models", s.mgmt.GetOAuthExcludedModels) - mgmt.PUT("/oauth-excluded-models", s.mgmt.PutOAuthExcludedModels) - mgmt.PATCH("/oauth-excluded-models", s.mgmt.PatchOAuthExcludedModels) - mgmt.DELETE("/oauth-excluded-models", s.mgmt.DeleteOAuthExcludedModels) - - mgmt.GET("/oauth-model-alias", s.mgmt.GetOAuthModelAlias) - mgmt.PUT("/oauth-model-alias", s.mgmt.PutOAuthModelAlias) - mgmt.PATCH("/oauth-model-alias", s.mgmt.PatchOAuthModelAlias) - mgmt.DELETE("/oauth-model-alias", s.mgmt.DeleteOAuthModelAlias) - - mgmt.GET("/auth-files", s.mgmt.ListAuthFiles) - mgmt.GET("/auth-files/models", s.mgmt.GetAuthFileModels) - mgmt.GET("/model-definitions/:channel", s.mgmt.GetStaticModelDefinitions) - mgmt.GET("/auth-files/download", s.mgmt.DownloadAuthFile) - mgmt.POST("/auth-files", s.mgmt.UploadAuthFile) - mgmt.DELETE("/auth-files", s.mgmt.DeleteAuthFile) - mgmt.PATCH("/auth-files/status", s.mgmt.PatchAuthFileStatus) - mgmt.PATCH("/auth-files/fields", s.mgmt.PatchAuthFileFields) - mgmt.POST("/vertex/import", s.mgmt.ImportVertexCredential) - - mgmt.GET("/anthropic-auth-url", s.mgmt.RequestAnthropicToken) - mgmt.GET("/codex-auth-url", s.mgmt.RequestCodexToken) - mgmt.GET("/antigravity-auth-url", s.mgmt.RequestAntigravityToken) - mgmt.GET("/kimi-auth-url", s.mgmt.RequestKimiToken) - mgmt.GET("/xai-auth-url", s.mgmt.RequestXAIToken) - mgmt.GET("/get-auth-status", s.mgmt.GetAuthStatus) - mgmt.DELETE("/oauth-session", s.mgmt.CancelAuthSession) - } -} - -func (s *Server) managementAvailabilityMiddleware() gin.HandlerFunc { - return func(c *gin.Context) { - if !s.managementAvailable(c) { - return - } - c.Next() - } -} - -func (s *Server) managementAvailable(c *gin.Context) bool { - if s == nil || s.cfg == nil { - c.AbortWithStatus(http.StatusNotFound) - return false - } - if s.cfg.Home.Enabled { - c.AbortWithStatus(http.StatusNotFound) - return false - } - if !s.managementRoutesEnabled.Load() { - c.AbortWithStatus(http.StatusNotFound) - return false - } - return true -} - -func (s *Server) refreshPluginManagementRoutes() { - if s == nil || s.pluginHost == nil || s.engine == nil { - return - } - s.pluginHost.RegisterManagementRoutes(context.Background(), s.registeredManagementRouteKeys()) -} - -// RefreshPluginManagementRoutes rebuilds plugin-owned Management API routes. -func (s *Server) RefreshPluginManagementRoutes() { - s.refreshPluginManagementRoutes() -} - -func (s *Server) registeredManagementRouteKeys() map[string]struct{} { - out := make(map[string]struct{}) - if s == nil || s.engine == nil { - return out - } - for _, route := range s.engine.Routes() { - if strings.HasPrefix(route.Path, "/v0/management/") || route.Path == "/v0/management" { - out[strings.ToUpper(strings.TrimSpace(route.Method))+" "+route.Path] = struct{}{} - } - } - return out -} - -func (s *Server) pluginManagementNoRoute(c *gin.Context) { - if s == nil || c == nil || c.Request == nil || c.Request.URL == nil { - if c != nil { - c.AbortWithStatus(http.StatusNotFound) - } - return - } - path := c.Request.URL.Path - if strings.HasPrefix(path, "/v0/resource/plugins/") { - s.pluginResourceNoRoute(c) - return - } - if path != "/v0/management" && !strings.HasPrefix(path, "/v0/management/") { - c.AbortWithStatus(http.StatusNotFound) - return - } - if s.pluginHost == nil || s.mgmt == nil { - c.AbortWithStatus(http.StatusNotFound) - return - } - if !s.managementAvailable(c) { - return - } - s.mgmt.Middleware()(c) - if c.IsAborted() { - return - } - if s.mgmt.ServePluginAuthURL(c) { - c.Abort() - return - } - if s.pluginHost.ServeManagementHTTP(c.Writer, c.Request) { - c.Abort() - return - } - c.AbortWithStatus(http.StatusNotFound) -} - -func (s *Server) pluginResourceNoRoute(c *gin.Context) { - if s == nil || c == nil || c.Request == nil || c.Request.URL == nil { - if c != nil { - c.AbortWithStatus(http.StatusNotFound) - } - return - } - if s.cfg == nil || s.cfg.Home.Enabled || s.pluginHost == nil { - c.AbortWithStatus(http.StatusNotFound) - return - } - if s.pluginHost.ServeResourceHTTP(c.Writer, c.Request) { - c.Abort() - return - } - c.AbortWithStatus(http.StatusNotFound) -} - -func (s *Server) serveManagementControlPanel(c *gin.Context) { - cfg := s.cfg - if cfg == nil || cfg.Home.Enabled || cfg.RemoteManagement.DisableControlPanel { - c.AbortWithStatus(http.StatusNotFound) - return - } - filePath := managementasset.FilePath(s.configFilePath) - if strings.TrimSpace(filePath) == "" { - c.AbortWithStatus(http.StatusNotFound) - return - } - - if _, err := os.Stat(filePath); err != nil { - if os.IsNotExist(err) { - // Synchronously ensure management.html is available with a detached context. - // Control panel bootstrap should not be canceled by client disconnects. - if !managementasset.EnsureLatestManagementHTML(context.Background(), managementasset.StaticDir(s.configFilePath), cfg.ProxyURL, cfg.RemoteManagement.PanelGitHubRepository) { - c.AbortWithStatus(http.StatusNotFound) - return - } - } else { - log.WithError(err).Error("failed to stat management control panel asset") - c.AbortWithStatus(http.StatusInternalServerError) - return - } - } - - c.File(filePath) -} - -func (s *Server) enableKeepAlive(timeout time.Duration, onTimeout func()) { - if timeout <= 0 || onTimeout == nil { - return - } - - s.keepAliveEnabled = true - s.keepAliveTimeout = timeout - s.keepAliveOnTimeout = onTimeout - s.keepAliveHeartbeat = make(chan struct{}, 1) - s.keepAliveStop = make(chan struct{}, 1) - - s.engine.GET("/keep-alive", s.handleKeepAlive) - - go s.watchKeepAlive() -} - -func (s *Server) handleKeepAlive(c *gin.Context) { - if s.localPassword != "" { - provided := strings.TrimSpace(c.GetHeader("Authorization")) - if provided != "" { - parts := strings.SplitN(provided, " ", 2) - if len(parts) == 2 && strings.EqualFold(parts[0], "bearer") { - provided = parts[1] - } - } - if provided == "" { - provided = strings.TrimSpace(c.GetHeader("X-Local-Password")) - } - if subtle.ConstantTimeCompare([]byte(provided), []byte(s.localPassword)) != 1 { - c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid password"}) - return - } - } - - s.signalKeepAlive() - c.JSON(http.StatusOK, gin.H{"status": "ok"}) -} - -func (s *Server) signalKeepAlive() { - if !s.keepAliveEnabled { - return - } - select { - case s.keepAliveHeartbeat <- struct{}{}: - default: - } -} - -func (s *Server) watchKeepAlive() { - if !s.keepAliveEnabled { - return - } - - timer := time.NewTimer(s.keepAliveTimeout) - defer timer.Stop() - - for { - select { - case <-timer.C: - log.Warnf("keep-alive endpoint idle for %s, shutting down", s.keepAliveTimeout) - if s.keepAliveOnTimeout != nil { - s.keepAliveOnTimeout() - } - return - case <-s.keepAliveHeartbeat: - if !timer.Stop() { - select { - case <-timer.C: - default: - } - } - timer.Reset(s.keepAliveTimeout) - case <-s.keepAliveStop: - return - } - } -} - -// isAnthropicModelsRequest reports whether a /v1/models request should be served in -// Anthropic format. Anthropic API clients send the Anthropic-Version header; Claude -// Code additionally uses a claude-cli User-Agent. -func isAnthropicModelsRequest(c *gin.Context) bool { - if c.GetHeader("Anthropic-Version") != "" { - return true - } - return strings.HasPrefix(c.GetHeader("User-Agent"), "claude-cli") -} - -// unifiedModelsHandler creates a unified handler for the /v1/models endpoint -// that routes to different handlers based on the request. -// Anthropic API requests (Anthropic-Version header, or a claude-cli User-Agent) -// route to the Claude handler, otherwise they route to the OpenAI handler. -func (s *Server) unifiedModelsHandler(openaiHandler *openai.OpenAIAPIHandler, claudeHandler *claude.ClaudeCodeAPIHandler) gin.HandlerFunc { - return func(c *gin.Context) { - if _, ok := c.Request.URL.Query()["client_version"]; ok { - if s != nil && s.cfg != nil && s.cfg.Home.Enabled { - s.handleHomeCodexClientModels(c) - return - } - openaiHandler.OpenAIModels(c) - return - } - - if s != nil && s.cfg != nil && s.cfg.Home.Enabled { - s.handleHomeModels(c) - return - } - - // Route to Claude handler for Anthropic API requests. - if isAnthropicModelsRequest(c) { - claudeHandler.ClaudeModels(c) - } else { - openaiHandler.OpenAIModels(c) - } - } -} - -// handleHomeCodexClientModels builds the Codex client catalog from Home model IDs. -// Template metadata still comes from the local/remote codex_client_models catalog. -func (s *Server) handleHomeCodexClientModels(c *gin.Context) { - entries, ok := s.loadHomeModelEntries(c) - if !ok { - return - } - - models := make([]map[string]any, 0, len(entries)) - for _, entry := range entries { - model := map[string]any{ - "id": entry.id, - "object": "model", - } - if entry.created > 0 { - model["created"] = entry.created - } - if entry.ownedBy != "" { - model["owned_by"] = entry.ownedBy - } - if entry.displayName != "" { - model["display_name"] = entry.displayName - model["description"] = entry.displayName - } - models = append(models, model) - } - - c.JSON(http.StatusOK, codexmodels.BuildResponse(models, nil, s.cfg.Codex.OptimizeMultiAgentV2)) -} - -func (s *Server) geminiModelsHandler(geminiHandler *gemini.GeminiAPIHandler) gin.HandlerFunc { - return func(c *gin.Context) { - if s != nil && s.cfg != nil && s.cfg.Home.Enabled { - s.handleHomeGeminiModels(c) - return - } - - geminiHandler.GeminiModels(c) - } -} - -func (s *Server) geminiGetHandler(geminiHandler *gemini.GeminiAPIHandler) gin.HandlerFunc { - return func(c *gin.Context) { - if s != nil && s.cfg != nil && s.cfg.Home.Enabled { - s.handleHomeGeminiModel(c) - return - } - - geminiHandler.GeminiGetHandler(c) - } -} - -type homeModelEntry struct { - id string - created int64 - ownedBy string - displayName string - contextLength int - maxCompletionTokens int -} - -func (s *Server) handleHomeModels(c *gin.Context) { - entries, ok := s.loadHomeModelEntries(c) - if !ok { - return - } - - isClaude := isAnthropicModelsRequest(c) - - if isClaude { - c.JSON(http.StatusOK, claudemodels.BuildResponse(formatHomeClaudeModels(entries))) - return - } - - filtered := make([]map[string]any, 0, len(entries)) - for _, entry := range entries { - model := map[string]any{ - "id": entry.id, - "object": "model", - } - if entry.created > 0 { - model["created"] = entry.created - } - if entry.ownedBy != "" { - model["owned_by"] = entry.ownedBy - } - filtered = append(filtered, model) - } - c.JSON(http.StatusOK, gin.H{ - "object": "list", - "data": filtered, - }) -} - -func formatHomeClaudeModels(entries []homeModelEntry) []map[string]any { - out := make([]map[string]any, 0, len(entries)) - for _, entry := range entries { - out = append(out, formatHomeClaudeModel(entry)) - } - return out -} - -func formatHomeClaudeModel(entry homeModelEntry) map[string]any { - displayName := entry.displayName - if displayName == "" { - displayName = entry.id - } - maxInput := entry.contextLength - if maxInput <= 0 { - maxInput = registry.DefaultClaudeMaxInputTokens - } - maxOutput := entry.maxCompletionTokens - if maxOutput <= 0 { - maxOutput = registry.DefaultClaudeMaxOutputTokens - } - model := map[string]any{ - "id": entry.id, - "object": "model", - "owned_by": entry.ownedBy, - "type": "model", - "display_name": displayName, - "max_input_tokens": maxInput, - "max_tokens": maxOutput, - } - if entry.created > 0 { - model["created_at"] = time.Unix(entry.created, 0).UTC().Format(time.RFC3339) - } - return model -} - -func (s *Server) handleHomeGeminiModels(c *gin.Context) { - entries, ok := s.loadHomeModelEntries(c) - if !ok { - return - } - - c.JSON(http.StatusOK, gin.H{ - "models": formatHomeGeminiModels(entries), - }) -} - -func (s *Server) handleHomeGeminiModel(c *gin.Context) { - entries, ok := s.loadHomeModelEntries(c) - if !ok { - return - } - - action := strings.TrimPrefix(c.Param("action"), "/") - action = strings.TrimSpace(action) - for _, entry := range entries { - if homeGeminiModelMatches(entry, action) { - c.JSON(http.StatusOK, formatHomeGeminiModel(entry)) - return - } - } - - c.JSON(http.StatusNotFound, handlers.ErrorResponse{ - Error: handlers.ErrorDetail{ - Message: "Not Found", - Type: "not_found", - }, - }) -} - -func (s *Server) loadHomeModelEntries(c *gin.Context) ([]homeModelEntry, bool) { - if s == nil || c == nil || c.Request == nil { - return nil, false - } - client := home.Current() - if client == nil { - c.JSON(http.StatusServiceUnavailable, handlers.ErrorResponse{ - Error: handlers.ErrorDetail{ - Message: "home control center unavailable", - Type: "server_error", - }, - }) - return nil, false - } - - raw, errGet := client.GetModels(c.Request.Context(), c.Request.Header, c.Request.URL.Query()) - if errGet != nil { - c.JSON(http.StatusBadGateway, handlers.ErrorResponse{ - Error: handlers.ErrorDetail{ - Message: errGet.Error(), - Type: "server_error", - }, - }) - return nil, false - } - - if statusCode, ok := homeModelsAuthStatus(raw); ok { - c.JSON(statusCode, handlers.ErrorResponse{ - Error: handlers.ErrorDetail{ - Message: homeModelsErrorMessage(raw), - Type: "authentication_error", - }, - }) - return nil, false - } - - entries, errDecode := decodeHomeModels(raw) - if errDecode != nil { - c.JSON(http.StatusBadGateway, handlers.ErrorResponse{ - Error: handlers.ErrorDetail{ - Message: errDecode.Error(), - Type: "server_error", - }, - }) - return nil, false - } - - return entries, true -} - -func formatHomeGeminiModels(entries []homeModelEntry) []map[string]any { - out := make([]map[string]any, 0, len(entries)) - for _, entry := range entries { - out = append(out, formatHomeGeminiModel(entry)) - } - return out -} - -func formatHomeGeminiModel(entry homeModelEntry) map[string]any { - name := entry.id - if !strings.HasPrefix(name, "models/") { - name = "models/" + name - } - displayName := entry.displayName - if displayName == "" { - displayName = entry.id - } - return map[string]any{ - "name": name, - "displayName": displayName, - "description": displayName, - "supportedGenerationMethods": []string{"generateContent"}, - } -} - -func homeGeminiModelMatches(entry homeModelEntry, action string) bool { - id := strings.TrimSpace(entry.id) - if id == "" || action == "" { - return false - } - normalizedAction := strings.TrimPrefix(action, "models/") - normalizedID := strings.TrimPrefix(id, "models/") - return action == id || action == "models/"+id || normalizedAction == normalizedID -} - -// homeModelsAuthStatus inspects a home models response for an authentication/error envelope. -// It returns the HTTP status code to surface (401 for credential issues, 502 otherwise) -// and true when the payload is an error response rather than model data. -func homeModelsAuthStatus(raw []byte) (int, bool) { - errType := homeModelsErrorType(raw) - if errType == "" { - return 0, false - } - if errType == "no_credentials" || errType == "invalid_credential" { - return http.StatusUnauthorized, true - } - return http.StatusBadGateway, true -} - -func homeModelsErrorType(raw []byte) string { - top, ok := unmarshalHomeModelsTopLevel(raw) - if !ok { - return "" - } - rawErr, exists := top["error"] - if !exists { - return "" - } - var errObj struct { - Type string `json:"type"` - } - if errUnmarshal := json.Unmarshal(rawErr, &errObj); errUnmarshal != nil { - return "" - } - return strings.TrimSpace(errObj.Type) -} - -func homeModelsErrorMessage(raw []byte) string { - top, ok := unmarshalHomeModelsTopLevel(raw) - if !ok { - return "home models request failed" - } - rawErr, exists := top["error"] - if !exists { - return "home models request failed" - } - var errObj struct { - Message string `json:"message"` - } - if errUnmarshal := json.Unmarshal(rawErr, &errObj); errUnmarshal != nil { - return "home models request failed" - } - if msg := strings.TrimSpace(errObj.Message); msg != "" { - return msg - } - return "home models request failed" -} - -func unmarshalHomeModelsTopLevel(raw []byte) (map[string]json.RawMessage, bool) { - if len(raw) == 0 { - return nil, false - } - var top map[string]json.RawMessage - if errUnmarshal := json.Unmarshal(raw, &top); errUnmarshal != nil { - return nil, false - } - return top, true -} - -func decodeHomeModels(raw []byte) ([]homeModelEntry, error) { - if len(raw) == 0 { - return nil, fmt.Errorf("home models payload is empty") - } - - var bySection map[string][]map[string]any - if err := json.Unmarshal(raw, &bySection); err != nil { - return nil, fmt.Errorf("parse home models payload: %w", err) - } - if len(bySection) == 0 { - return nil, fmt.Errorf("home models payload has no sections") - } - - seen := make(map[string]struct{}) - out := make([]homeModelEntry, 0, 256) - for _, models := range bySection { - for _, model := range models { - id, _ := model["id"].(string) - id = strings.TrimSpace(id) - if id == "" { - name, _ := model["name"].(string) - name = strings.TrimSpace(name) - id = strings.TrimPrefix(name, "models/") - } - if id == "" { - continue - } - if _, ok := seen[id]; ok { - continue - } - seen[id] = struct{}{} - - ownedBy, _ := model["owned_by"].(string) - ownedBy = strings.TrimSpace(ownedBy) - displayName, _ := model["display_name"].(string) - displayName = strings.TrimSpace(displayName) - if displayName == "" { - displayName, _ = model["displayName"].(string) - displayName = strings.TrimSpace(displayName) - } - - out = append(out, homeModelEntry{ - id: id, - created: homeModelInt64Value(model, "created"), - ownedBy: ownedBy, - displayName: displayName, - contextLength: int(homeModelInt64Value(model, "context_length", "contextLength", "inputTokenLimit", "max_input_tokens")), - maxCompletionTokens: int(homeModelInt64Value(model, "max_completion_tokens", "maxCompletionTokens", "outputTokenLimit", "max_tokens")), - }) - } - } - - sort.Slice(out, func(i, j int) bool { return out[i].id < out[j].id }) - if len(out) == 0 { - return nil, fmt.Errorf("home models payload contains no models") - } - return out, nil -} - -func homeModelInt64Value(model map[string]any, keys ...string) int64 { - for _, key := range keys { - switch value := model[key].(type) { - case float64: - return int64(value) - case int64: - return value - case int: - return int64(value) - case json.Number: - if n, errInt := value.Int64(); errInt == nil { - return n - } - case string: - if n, errParse := strconv.ParseInt(strings.TrimSpace(value), 10, 64); errParse == nil { - return n - } - } - } - return 0 -} - // Start begins listening for and serving HTTP or HTTPS requests. // It's a blocking call and will only return on an unrecoverable error. // @@ -1887,314 +397,3 @@ func (s *Server) Stop(ctx context.Context) error { log.Debug("API server stopped") return nil } - -// corsMiddleware returns a Gin middleware handler that adds CORS headers -// to every response, allowing cross-origin requests. -// -// Returns: -// - gin.HandlerFunc: The CORS middleware handler -func corsMiddleware() gin.HandlerFunc { - return func(c *gin.Context) { - c.Header("Access-Control-Allow-Origin", "*") - c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS") - c.Header("Access-Control-Allow-Headers", "*") - c.Header("Access-Control-Expose-Headers", corsExposedResponseHeadersJoined) - - if c.Request.Method == "OPTIONS" { - c.AbortWithStatus(http.StatusNoContent) - return - } - - c.Next() - } -} - -func (s *Server) applyAccessConfig(oldCfg, newCfg *config.Config) bool { - if s == nil || s.accessManager == nil || newCfg == nil { - return false - } - if _, err := access.ApplyAccessProviders(s.accessManager, oldCfg, newCfg); err != nil { - return false - } - return true -} - -// UpdateClients updates the server's client list and configuration. -// This method is called when the configuration or authentication tokens change. -// -// Parameters: -// - clients: The new slice of AI service clients -// - cfg: The new application configuration -func (s *Server) UpdateClients(cfg *config.Config) { - s.UpdateClientsContext(context.Background(), cfg) -} - -// UpdateClientsContext updates runtime clients while honoring cancellation between -// short configuration and filesystem operations. -func (s *Server) UpdateClientsContext(ctx context.Context, cfg *config.Config) bool { - if s == nil || cfg == nil { - return false - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - // Reconstruct old config from YAML snapshot to avoid reference sharing issues - var oldCfg *config.Config - if len(s.oldConfigYaml) > 0 { - _ = yaml.Unmarshal(s.oldConfigYaml, &oldCfg) - } - - // Update request logger enabled state if it has changed - previousRequestLog := false - if oldCfg != nil { - previousRequestLog = oldCfg.RequestLog - } - if s.requestLogger != nil && (oldCfg == nil || previousRequestLog != cfg.RequestLog) { - if s.loggerToggle != nil { - s.loggerToggle(cfg.RequestLog) - } else if toggler, ok := s.requestLogger.(interface{ SetEnabled(bool) }); ok { - toggler.SetEnabled(cfg.RequestLog) - } - } - - if oldCfg == nil || oldCfg.Home.Enabled != cfg.Home.Enabled { - if setter, ok := s.requestLogger.(interface{ SetHomeEnabled(bool) }); ok { - setter.SetHomeEnabled(cfg.Home.Enabled) - } - } - - if oldCfg == nil || oldCfg.LoggingToFile != cfg.LoggingToFile || oldCfg.LogsMaxTotalSizeMB != cfg.LogsMaxTotalSizeMB { - if err := logging.ConfigureLogOutput(cfg); err != nil { - log.Errorf("failed to reconfigure log output: %v", err) - } - if errContext := ctx.Err(); errContext != nil { - return false - } - } - - if oldCfg == nil || oldCfg.UsageStatisticsEnabled != cfg.UsageStatisticsEnabled { - redisqueue.SetUsageStatisticsEnabled(cfg.UsageStatisticsEnabled) - } - - if oldCfg == nil || oldCfg.RedisUsageQueueRetentionSeconds != cfg.RedisUsageQueueRetentionSeconds { - redisqueue.SetRetentionSeconds(cfg.RedisUsageQueueRetentionSeconds) - } - - if s.requestLogger != nil && (oldCfg == nil || oldCfg.ErrorLogsMaxFiles != cfg.ErrorLogsMaxFiles) { - if setter, ok := s.requestLogger.(interface{ SetErrorLogsMaxFiles(int) }); ok { - setter.SetErrorLogsMaxFiles(cfg.ErrorLogsMaxFiles) - } - } - - if oldCfg == nil || oldCfg.DisableCooling != cfg.DisableCooling { - auth.SetQuotaCooldownDisabled(cfg.DisableCooling) - } - if oldCfg == nil || oldCfg.TransientErrorCooldownSeconds != cfg.TransientErrorCooldownSeconds { - auth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) - } - - if oldCfg != nil && oldCfg.DisableImageGeneration != cfg.DisableImageGeneration { - log.Infof("disable-image-generation updated: %v -> %v", oldCfg.DisableImageGeneration, cfg.DisableImageGeneration) - } - - applySignatureCacheConfig(oldCfg, cfg) - - if s.handlers != nil && s.handlers.AuthManager != nil { - s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials) - } - - // Update log level dynamically when debug flag changes - if oldCfg == nil || oldCfg.Debug != cfg.Debug { - util.SetLogLevel(cfg) - } - - prevSecretEmpty := true - if oldCfg != nil { - prevSecretEmpty = oldCfg.RemoteManagement.SecretKey == "" - } - newSecretEmpty := cfg.RemoteManagement.SecretKey == "" - if s.envManagementSecret { - s.registerManagementRoutes() - if s.managementRoutesEnabled.CompareAndSwap(false, true) { - log.Info("management routes enabled via MANAGEMENT_PASSWORD") - } else { - s.managementRoutesEnabled.Store(true) - } - } else { - switch { - case prevSecretEmpty && !newSecretEmpty: - s.registerManagementRoutes() - if s.managementRoutesEnabled.CompareAndSwap(false, true) { - log.Info("management routes enabled after secret key update") - } else { - s.managementRoutesEnabled.Store(true) - } - case !prevSecretEmpty && newSecretEmpty: - if s.managementRoutesEnabled.CompareAndSwap(true, false) { - log.Info("management routes disabled after secret key removal") - } else { - s.managementRoutesEnabled.Store(false) - } - default: - s.managementRoutesEnabled.Store(!newSecretEmpty) - } - } - redisqueue.SetEnabled(s.managementRoutesEnabled.Load() || (cfg != nil && cfg.Home.Enabled)) - - exampleAPIKeySafeModeRequired := s.exampleAPIKeySafeModeRequired(cfg) - if exampleAPIKeySafeModeRequired { - s.exampleAPIKeySafeModeActive.Store(true) - } - accessConfigApplied := s.applyAccessConfig(oldCfg, cfg) - if accessConfigApplied || exampleAPIKeySafeModeRequired { - s.exampleAPIKeySafeModeActive.Store(exampleAPIKeySafeModeRequired) - } - s.cfg = cfg - if s.codexLiveHandler != nil { - if errUpdate := s.codexLiveHandler.UpdateConfig(cfg); errUpdate != nil { - log.WithError(errUpdate).Error("failed to update Codex Live media relay configuration") - } - } - s.wsAuthEnabled.Store(cfg.WebsocketAuth) - if oldCfg != nil && s.wsAuthChanged != nil && oldCfg.WebsocketAuth != cfg.WebsocketAuth { - s.wsAuthChanged(oldCfg.WebsocketAuth, cfg.WebsocketAuth) - } - managementasset.SetCurrentConfig(cfg) - if errContext := ctx.Err(); errContext != nil { - return false - } - // Save YAML snapshot for next comparison - s.oldConfigYaml, _ = yaml.Marshal(cfg) - - s.handlers.UpdateClients(effectiveSDKConfig(cfg)) - s.handlers.SetPluginHost(s.pluginHost) - if s.pluginHost != nil { - s.pluginHost.SetModelExecutor(s.handlers) - s.pluginHost.SetAuthManager(s.handlers.AuthManager) - } - - if s.mgmt != nil { - s.mgmt.SetConfig(cfg) - s.mgmt.SetAuthManager(s.handlers.AuthManager) - s.mgmt.SetPluginHost(s.pluginHost) - } - s.refreshPluginManagementRoutes() - - // Count client sources from configuration and auth store. - authEntries := 0 - if cfg != nil && !cfg.Home.Enabled { - tokenStore := sdkAuth.GetTokenStore() - if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok { - dirSetter.SetBaseDir(cfg.AuthDir) - } - authEntries = util.CountAuthFiles(ctx, tokenStore) - if errContext := ctx.Err(); errContext != nil { - return false - } - } - geminiAPIKeyCount := len(cfg.GeminiKey) - interactionsAPIKeyCount := len(cfg.InteractionsKey) - claudeAPIKeyCount := len(cfg.ClaudeKey) - codexAPIKeyCount := len(cfg.CodexKey) - xaiAPIKeyCount := len(cfg.XAIKey) - vertexAICompatCount := len(cfg.VertexCompatAPIKey) - openAICompatCount := 0 - for i := range cfg.OpenAICompatibility { - entry := cfg.OpenAICompatibility[i] - if entry.Disabled { - continue - } - openAICompatCount += len(entry.APIKeyEntries) - } - - total := authEntries + geminiAPIKeyCount + interactionsAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + xaiAPIKeyCount + vertexAICompatCount + openAICompatCount - fmt.Printf("server clients and configuration updated: %d clients (%d auth entries + %d Gemini API keys + %d Interactions API keys + %d Claude API keys + %d Codex keys + %d xAI keys + %d Vertex-compat + %d OpenAI-compat)\n", - total, - authEntries, - geminiAPIKeyCount, - interactionsAPIKeyCount, - claudeAPIKeyCount, - codexAPIKeyCount, - xaiAPIKeyCount, - vertexAICompatCount, - openAICompatCount, - ) - return ctx.Err() == nil -} - -func (s *Server) SetWebsocketAuthChangeHandler(fn func(bool, bool)) { - if s == nil { - return - } - s.wsAuthChanged = fn -} - -// (management handlers moved to internal/api/handlers/management) - -// AuthMiddleware returns a Gin middleware handler that authenticates requests -// using the configured authentication providers. When no providers are available, -// it allows all requests (legacy behaviour). -func AuthMiddleware(manager *sdkaccess.Manager) gin.HandlerFunc { - return func(c *gin.Context) { - if manager == nil { - c.Next() - return - } - - result, err := manager.Authenticate(c.Request.Context(), c.Request) - if err == nil { - if result != nil { - c.Set("userApiKey", result.Principal) - c.Set("accessProvider", result.Provider) - if len(result.Metadata) > 0 { - c.Set("accessMetadata", result.Metadata) - } - } - c.Next() - return - } - - statusCode := err.HTTPStatusCode() - if statusCode >= http.StatusInternalServerError { - log.Errorf("authentication middleware error: %v", err) - } - c.AbortWithStatusJSON(statusCode, gin.H{"error": err.Message}) - } -} - -func configuredSignatureCacheEnabled(cfg *config.Config) bool { - if cfg != nil && cfg.AntigravitySignatureCacheEnabled != nil { - return *cfg.AntigravitySignatureCacheEnabled - } - return true -} - -func applySignatureCacheConfig(oldCfg, cfg *config.Config) { - newVal := configuredSignatureCacheEnabled(cfg) - newStrict := configuredSignatureBypassStrict(cfg) - if oldCfg == nil { - cache.SetSignatureCacheEnabled(newVal) - cache.SetSignatureBypassStrictMode(newStrict) - return - } - - oldVal := configuredSignatureCacheEnabled(oldCfg) - if oldVal != newVal { - cache.SetSignatureCacheEnabled(newVal) - } - - oldStrict := configuredSignatureBypassStrict(oldCfg) - if oldStrict != newStrict { - cache.SetSignatureBypassStrictMode(newStrict) - } -} - -func configuredSignatureBypassStrict(cfg *config.Config) bool { - if cfg != nil && cfg.AntigravitySignatureBypassStrict != nil { - return *cfg.AntigravitySignatureBypassStrict - } - return false -} diff --git a/internal/api/server_keepalive.go b/internal/api/server_keepalive.go new file mode 100644 index 000000000..28080ae7f --- /dev/null +++ b/internal/api/server_keepalive.go @@ -0,0 +1,89 @@ +package api + +import ( + "crypto/subtle" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + log "github.com/sirupsen/logrus" +) + +func (s *Server) enableKeepAlive(timeout time.Duration, onTimeout func()) { + if timeout <= 0 || onTimeout == nil { + return + } + + s.keepAliveEnabled = true + s.keepAliveTimeout = timeout + s.keepAliveOnTimeout = onTimeout + s.keepAliveHeartbeat = make(chan struct{}, 1) + s.keepAliveStop = make(chan struct{}, 1) + + s.engine.GET("/keep-alive", s.handleKeepAlive) + + go s.watchKeepAlive() +} + +func (s *Server) handleKeepAlive(c *gin.Context) { + if s.localPassword != "" { + provided := strings.TrimSpace(c.GetHeader("Authorization")) + if provided != "" { + parts := strings.SplitN(provided, " ", 2) + if len(parts) == 2 && strings.EqualFold(parts[0], "bearer") { + provided = parts[1] + } + } + if provided == "" { + provided = strings.TrimSpace(c.GetHeader("X-Local-Password")) + } + if subtle.ConstantTimeCompare([]byte(provided), []byte(s.localPassword)) != 1 { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "invalid password"}) + return + } + } + + s.signalKeepAlive() + c.JSON(http.StatusOK, gin.H{"status": "ok"}) +} + +func (s *Server) signalKeepAlive() { + if !s.keepAliveEnabled { + return + } + select { + case s.keepAliveHeartbeat <- struct{}{}: + default: + } +} + +func (s *Server) watchKeepAlive() { + if !s.keepAliveEnabled { + return + } + + timer := time.NewTimer(s.keepAliveTimeout) + defer timer.Stop() + + for { + select { + case <-timer.C: + log.Warnf("keep-alive endpoint idle for %s, shutting down", s.keepAliveTimeout) + if s.keepAliveOnTimeout != nil { + s.keepAliveOnTimeout() + } + return + case <-s.keepAliveHeartbeat: + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timer.Reset(s.keepAliveTimeout) + case <-s.keepAliveStop: + return + } + } +} diff --git a/internal/api/server_management.go b/internal/api/server_management.go new file mode 100644 index 000000000..3a8a53a33 --- /dev/null +++ b/internal/api/server_management.go @@ -0,0 +1,312 @@ +package api + +import ( + "context" + "net/http" + "os" + "strings" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/managementasset" + log "github.com/sirupsen/logrus" +) + +func (s *Server) registerManagementRoutes() { + if s == nil || s.engine == nil || s.mgmt == nil { + return + } + if !s.managementRoutesRegistered.CompareAndSwap(false, true) { + return + } + + log.Info("management routes registered after secret key configuration") + + s.engine.POST("/v0/management/oauth-callback", s.managementAvailabilityMiddleware(), s.mgmt.PostOAuthCallback) + s.engine.GET("/v0/management/oauth-callback", s.managementAvailabilityMiddleware(), s.mgmt.GetOAuthCallback) + + mgmt := s.engine.Group("/v0/management") + mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware()) + { + mgmt.GET("/config", s.mgmt.GetConfig) + mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML) + mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML) + mgmt.GET("/latest-version", s.mgmt.GetLatestVersion) + mgmt.GET("/plugins", s.mgmt.ListPlugins) + mgmt.GET("/plugin-store", s.mgmt.ListPluginStore) + mgmt.POST("/plugin-store/:id/install", s.mgmt.InstallPluginFromStore) + mgmt.DELETE("/plugins/:id", s.mgmt.DeletePlugin) + mgmt.PATCH("/plugins/:id/enabled", s.mgmt.PatchPluginEnabled) + mgmt.GET("/plugins/:id/config", s.mgmt.GetPluginConfig) + mgmt.PUT("/plugins/:id/config", s.mgmt.PutPluginConfig) + mgmt.PATCH("/plugins/:id/config", s.mgmt.PatchPluginConfig) + + mgmt.GET("/debug", s.mgmt.GetDebug) + mgmt.PUT("/debug", s.mgmt.PutDebug) + mgmt.PATCH("/debug", s.mgmt.PutDebug) + + mgmt.GET("/logging-to-file", s.mgmt.GetLoggingToFile) + mgmt.PUT("/logging-to-file", s.mgmt.PutLoggingToFile) + mgmt.PATCH("/logging-to-file", s.mgmt.PutLoggingToFile) + + mgmt.GET("/logs-max-total-size-mb", s.mgmt.GetLogsMaxTotalSizeMB) + mgmt.PUT("/logs-max-total-size-mb", s.mgmt.PutLogsMaxTotalSizeMB) + mgmt.PATCH("/logs-max-total-size-mb", s.mgmt.PutLogsMaxTotalSizeMB) + + mgmt.GET("/error-logs-max-files", s.mgmt.GetErrorLogsMaxFiles) + mgmt.PUT("/error-logs-max-files", s.mgmt.PutErrorLogsMaxFiles) + mgmt.PATCH("/error-logs-max-files", s.mgmt.PutErrorLogsMaxFiles) + + mgmt.GET("/usage-statistics-enabled", s.mgmt.GetUsageStatisticsEnabled) + mgmt.PUT("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled) + mgmt.PATCH("/usage-statistics-enabled", s.mgmt.PutUsageStatisticsEnabled) + + mgmt.GET("/proxy-url", s.mgmt.GetProxyURL) + mgmt.PUT("/proxy-url", s.mgmt.PutProxyURL) + mgmt.PATCH("/proxy-url", s.mgmt.PutProxyURL) + mgmt.DELETE("/proxy-url", s.mgmt.DeleteProxyURL) + + mgmt.POST("/api-call", s.mgmt.APICall) + + mgmt.GET("/quota-exceeded/switch-project", s.mgmt.GetSwitchProject) + mgmt.PUT("/quota-exceeded/switch-project", s.mgmt.PutSwitchProject) + mgmt.PATCH("/quota-exceeded/switch-project", s.mgmt.PutSwitchProject) + + mgmt.GET("/quota-exceeded/switch-preview-model", s.mgmt.GetSwitchPreviewModel) + mgmt.PUT("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel) + mgmt.PATCH("/quota-exceeded/switch-preview-model", s.mgmt.PutSwitchPreviewModel) + mgmt.POST("/reset-quota", s.mgmt.ResetQuota) + + mgmt.GET("/api-keys", s.mgmt.GetAPIKeys) + mgmt.PUT("/api-keys", s.mgmt.PutAPIKeys) + mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys) + mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys) + mgmt.GET("/api-key-usage", s.mgmt.GetAPIKeyUsage) + mgmt.GET("/usage-queue", s.mgmt.GetUsageQueue) + + mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys) + mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys) + mgmt.PATCH("/gemini-api-key", s.mgmt.PatchGeminiKey) + mgmt.DELETE("/gemini-api-key", s.mgmt.DeleteGeminiKey) + + mgmt.GET("/interactions-api-key", s.mgmt.GetInteractionsKeys) + mgmt.PUT("/interactions-api-key", s.mgmt.PutInteractionsKeys) + mgmt.PATCH("/interactions-api-key", s.mgmt.PatchInteractionsKey) + mgmt.DELETE("/interactions-api-key", s.mgmt.DeleteInteractionsKey) + + mgmt.GET("/logs", s.mgmt.GetLogs) + mgmt.DELETE("/logs", s.mgmt.DeleteLogs) + mgmt.GET("/request-error-logs", s.mgmt.GetRequestErrorLogs) + mgmt.GET("/request-error-logs/:name", s.mgmt.DownloadRequestErrorLog) + mgmt.GET("/request-log-by-id/:id", s.mgmt.GetRequestLogByID) + mgmt.GET("/request-log", s.mgmt.GetRequestLog) + mgmt.PUT("/request-log", s.mgmt.PutRequestLog) + mgmt.PATCH("/request-log", s.mgmt.PutRequestLog) + mgmt.GET("/ws-auth", s.mgmt.GetWebsocketAuth) + mgmt.PUT("/ws-auth", s.mgmt.PutWebsocketAuth) + mgmt.PATCH("/ws-auth", s.mgmt.PutWebsocketAuth) + + mgmt.GET("/request-retry", s.mgmt.GetRequestRetry) + mgmt.PUT("/request-retry", s.mgmt.PutRequestRetry) + mgmt.PATCH("/request-retry", s.mgmt.PutRequestRetry) + mgmt.GET("/max-retry-interval", s.mgmt.GetMaxRetryInterval) + mgmt.PUT("/max-retry-interval", s.mgmt.PutMaxRetryInterval) + mgmt.PATCH("/max-retry-interval", s.mgmt.PutMaxRetryInterval) + + mgmt.GET("/force-model-prefix", s.mgmt.GetForceModelPrefix) + mgmt.PUT("/force-model-prefix", s.mgmt.PutForceModelPrefix) + mgmt.PATCH("/force-model-prefix", s.mgmt.PutForceModelPrefix) + + mgmt.GET("/routing/strategy", s.mgmt.GetRoutingStrategy) + mgmt.PUT("/routing/strategy", s.mgmt.PutRoutingStrategy) + mgmt.PATCH("/routing/strategy", s.mgmt.PutRoutingStrategy) + + mgmt.GET("/claude-api-key", s.mgmt.GetClaudeKeys) + mgmt.PUT("/claude-api-key", s.mgmt.PutClaudeKeys) + mgmt.PATCH("/claude-api-key", s.mgmt.PatchClaudeKey) + mgmt.DELETE("/claude-api-key", s.mgmt.DeleteClaudeKey) + + mgmt.GET("/codex-api-key", s.mgmt.GetCodexKeys) + mgmt.PUT("/codex-api-key", s.mgmt.PutCodexKeys) + mgmt.PATCH("/codex-api-key", s.mgmt.PatchCodexKey) + mgmt.DELETE("/codex-api-key", s.mgmt.DeleteCodexKey) + + mgmt.GET("/xai-api-key", s.mgmt.GetXAIKeys) + mgmt.PUT("/xai-api-key", s.mgmt.PutXAIKeys) + mgmt.PATCH("/xai-api-key", s.mgmt.PatchXAIKey) + mgmt.DELETE("/xai-api-key", s.mgmt.DeleteXAIKey) + + mgmt.GET("/openai-compatibility", s.mgmt.GetOpenAICompat) + mgmt.PUT("/openai-compatibility", s.mgmt.PutOpenAICompat) + mgmt.PATCH("/openai-compatibility", s.mgmt.PatchOpenAICompat) + mgmt.DELETE("/openai-compatibility", s.mgmt.DeleteOpenAICompat) + + mgmt.GET("/vertex-api-key", s.mgmt.GetVertexCompatKeys) + mgmt.PUT("/vertex-api-key", s.mgmt.PutVertexCompatKeys) + mgmt.PATCH("/vertex-api-key", s.mgmt.PatchVertexCompatKey) + mgmt.DELETE("/vertex-api-key", s.mgmt.DeleteVertexCompatKey) + + mgmt.GET("/oauth-excluded-models", s.mgmt.GetOAuthExcludedModels) + mgmt.PUT("/oauth-excluded-models", s.mgmt.PutOAuthExcludedModels) + mgmt.PATCH("/oauth-excluded-models", s.mgmt.PatchOAuthExcludedModels) + mgmt.DELETE("/oauth-excluded-models", s.mgmt.DeleteOAuthExcludedModels) + + mgmt.GET("/oauth-model-alias", s.mgmt.GetOAuthModelAlias) + mgmt.PUT("/oauth-model-alias", s.mgmt.PutOAuthModelAlias) + mgmt.PATCH("/oauth-model-alias", s.mgmt.PatchOAuthModelAlias) + mgmt.DELETE("/oauth-model-alias", s.mgmt.DeleteOAuthModelAlias) + + mgmt.GET("/auth-files", s.mgmt.ListAuthFiles) + mgmt.GET("/auth-files/models", s.mgmt.GetAuthFileModels) + mgmt.GET("/model-definitions/:channel", s.mgmt.GetStaticModelDefinitions) + mgmt.GET("/auth-files/download", s.mgmt.DownloadAuthFile) + mgmt.POST("/auth-files", s.mgmt.UploadAuthFile) + mgmt.DELETE("/auth-files", s.mgmt.DeleteAuthFile) + mgmt.PATCH("/auth-files/status", s.mgmt.PatchAuthFileStatus) + mgmt.PATCH("/auth-files/fields", s.mgmt.PatchAuthFileFields) + mgmt.POST("/vertex/import", s.mgmt.ImportVertexCredential) + + mgmt.GET("/anthropic-auth-url", s.mgmt.RequestAnthropicToken) + mgmt.GET("/codex-auth-url", s.mgmt.RequestCodexToken) + mgmt.GET("/antigravity-auth-url", s.mgmt.RequestAntigravityToken) + mgmt.GET("/kimi-auth-url", s.mgmt.RequestKimiToken) + mgmt.GET("/xai-auth-url", s.mgmt.RequestXAIToken) + mgmt.GET("/get-auth-status", s.mgmt.GetAuthStatus) + mgmt.DELETE("/oauth-session", s.mgmt.CancelAuthSession) + } +} + +func (s *Server) managementAvailabilityMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + if !s.managementAvailable(c) { + return + } + c.Next() + } +} + +func (s *Server) managementAvailable(c *gin.Context) bool { + if s == nil || s.cfg == nil { + c.AbortWithStatus(http.StatusNotFound) + return false + } + if s.cfg.Home.Enabled { + c.AbortWithStatus(http.StatusNotFound) + return false + } + if !s.managementRoutesEnabled.Load() { + c.AbortWithStatus(http.StatusNotFound) + return false + } + return true +} + +func (s *Server) refreshPluginManagementRoutes() { + if s == nil || s.pluginHost == nil || s.engine == nil { + return + } + s.pluginHost.RegisterManagementRoutes(context.Background(), s.registeredManagementRouteKeys()) +} + +// RefreshPluginManagementRoutes rebuilds plugin-owned Management API routes. +func (s *Server) RefreshPluginManagementRoutes() { + s.refreshPluginManagementRoutes() +} + +func (s *Server) registeredManagementRouteKeys() map[string]struct{} { + out := make(map[string]struct{}) + if s == nil || s.engine == nil { + return out + } + for _, route := range s.engine.Routes() { + if strings.HasPrefix(route.Path, "/v0/management/") || route.Path == "/v0/management" { + out[strings.ToUpper(strings.TrimSpace(route.Method))+" "+route.Path] = struct{}{} + } + } + return out +} + +func (s *Server) pluginManagementNoRoute(c *gin.Context) { + if s == nil || c == nil || c.Request == nil || c.Request.URL == nil { + if c != nil { + c.AbortWithStatus(http.StatusNotFound) + } + return + } + path := c.Request.URL.Path + if strings.HasPrefix(path, "/v0/resource/plugins/") { + s.pluginResourceNoRoute(c) + return + } + if path != "/v0/management" && !strings.HasPrefix(path, "/v0/management/") { + c.AbortWithStatus(http.StatusNotFound) + return + } + if s.pluginHost == nil || s.mgmt == nil { + c.AbortWithStatus(http.StatusNotFound) + return + } + if !s.managementAvailable(c) { + return + } + s.mgmt.Middleware()(c) + if c.IsAborted() { + return + } + if s.mgmt.ServePluginAuthURL(c) { + c.Abort() + return + } + if s.pluginHost.ServeManagementHTTP(c.Writer, c.Request) { + c.Abort() + return + } + c.AbortWithStatus(http.StatusNotFound) +} + +func (s *Server) pluginResourceNoRoute(c *gin.Context) { + if s == nil || c == nil || c.Request == nil || c.Request.URL == nil { + if c != nil { + c.AbortWithStatus(http.StatusNotFound) + } + return + } + if s.cfg == nil || s.cfg.Home.Enabled || s.pluginHost == nil { + c.AbortWithStatus(http.StatusNotFound) + return + } + if s.pluginHost.ServeResourceHTTP(c.Writer, c.Request) { + c.Abort() + return + } + c.AbortWithStatus(http.StatusNotFound) +} + +func (s *Server) serveManagementControlPanel(c *gin.Context) { + cfg := s.cfg + if cfg == nil || cfg.Home.Enabled || cfg.RemoteManagement.DisableControlPanel { + c.AbortWithStatus(http.StatusNotFound) + return + } + filePath := managementasset.FilePath(s.configFilePath) + if strings.TrimSpace(filePath) == "" { + c.AbortWithStatus(http.StatusNotFound) + return + } + + if _, err := os.Stat(filePath); err != nil { + if os.IsNotExist(err) { + // Synchronously ensure management.html is available with a detached context. + // Control panel bootstrap should not be canceled by client disconnects. + if !managementasset.EnsureLatestManagementHTML(context.Background(), managementasset.StaticDir(s.configFilePath), cfg.ProxyURL, cfg.RemoteManagement.PanelGitHubRepository) { + c.AbortWithStatus(http.StatusNotFound) + return + } + } else { + log.WithError(err).Error("failed to stat management control panel asset") + c.AbortWithStatus(http.StatusInternalServerError) + return + } + } + + c.File(filePath) +} diff --git a/internal/api/server_middleware.go b/internal/api/server_middleware.go new file mode 100644 index 000000000..8511d4238 --- /dev/null +++ b/internal/api/server_middleware.go @@ -0,0 +1,172 @@ +package api + +import ( + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/internal/safemode" + sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" + log "github.com/sirupsen/logrus" +) + +var corsExposedResponseHeaders = []string{ + logging.CPATraceIDHeader, + "X-CPA-VERSION", + "X-CPA-COMMIT", + "X-CPA-BUILD-DATE", + "X-CPA-SUPPORT-PLUGIN", + "X-CPA-HOME-VERSION", + "X-CPA-HOME-BUILD-DATE", + "X-SERVER-VERSION", + "X-SERVER-BUILD-DATE", +} + +var corsExposedResponseHeadersJoined = strings.Join(corsExposedResponseHeaders, ", ") + +const ( + exampleAPIKeyManagementPath = "/management.html" + exampleAPIKeyManagementURL = "/management.html?safe-mode=configure" +) + +func (s *Server) homeHeartbeatMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + if s == nil || s.cfg == nil || !s.cfg.Home.Enabled { + c.Next() + return + } + if c != nil && c.Request != nil { + path := c.Request.URL.Path + if strings.HasPrefix(path, "/v0/management/") || path == "/v0/management" || strings.HasPrefix(path, "/v0/resource/plugins/") || path == "/management.html" { + c.Next() + return + } + } + client := home.Current() + if client == nil || !client.HeartbeatOK() { + c.AbortWithStatus(http.StatusServiceUnavailable) + return + } + c.Next() + } +} + +func (s *Server) exampleAPIKeySafeModeRequired(cfg *config.Config) bool { + return s != nil && s.exampleAPIKeySafeModeEnabled && cfg != nil && safemode.HasExampleAPIKeys(cfg.APIKeys) +} + +func (s *Server) exampleAPIKeySafeModeMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + if s == nil || !s.exampleAPIKeySafeModeActive.Load() || c == nil || c.Request == nil || c.Request.URL == nil { + c.Next() + return + } + + path := c.Request.URL.Path + if path == exampleAPIKeyManagementPath && c.Query("safe-mode") == "configure" { + c.Next() + return + } + if (path == "/" || path == exampleAPIKeyManagementPath) && (c.Request.Method == http.MethodGet || c.Request.Method == http.MethodHead) { + s.serveExampleAPIKeyWarningPage(c) + return + } + if !isExampleAPIKeySafeModeProxyPath(path) { + c.Next() + return + } + + c.Header("X-CPA-SAFE-MODE", "example-api-key") + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "unsafe_example_api_key", + "message": "Proxy API endpoints are disabled because api-keys contains template values. Open /management.html?safe-mode=configure, update api-keys in Management, then retry.", + }) + } +} + +func (s *Server) serveExampleAPIKeyWarningPage(c *gin.Context) { + cfg := s.cfg + var keys []string + if cfg != nil { + keys = safemode.ExampleAPIKeys(cfg.APIKeys) + } + c.Header("Content-Type", "text/html; charset=utf-8") + c.Header("Cache-Control", "no-store") + if c.Request.Method == http.MethodHead { + c.Status(http.StatusOK) + c.Abort() + return + } + c.String(http.StatusOK, safemode.ExampleAPIKeyWarningPageHTML(keys, exampleAPIKeyManagementURL)) + c.Abort() +} + +func isExampleAPIKeySafeModeProxyPath(path string) bool { + switch { + case path == "/v1" || strings.HasPrefix(path, "/v1/"): + return true + case path == "/v1beta" || strings.HasPrefix(path, "/v1beta/"): + return true + case path == "/openai/v1" || strings.HasPrefix(path, "/openai/v1/"): + return true + case path == "/backend-api/codex" || strings.HasPrefix(path, "/backend-api/codex/"): + return true + default: + return false + } +} + +// corsMiddleware returns a Gin middleware handler that adds CORS headers +// to every response, allowing cross-origin requests. +// +// Returns: +// - gin.HandlerFunc: The CORS middleware handler +func corsMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + c.Header("Access-Control-Allow-Origin", "*") + c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS") + c.Header("Access-Control-Allow-Headers", "*") + c.Header("Access-Control-Expose-Headers", corsExposedResponseHeadersJoined) + + if c.Request.Method == "OPTIONS" { + c.AbortWithStatus(http.StatusNoContent) + return + } + + c.Next() + } +} + +// AuthMiddleware returns a Gin middleware handler that authenticates requests +// using the configured authentication providers. When no providers are available, +// it allows all requests (legacy behaviour). +func AuthMiddleware(manager *sdkaccess.Manager) gin.HandlerFunc { + return func(c *gin.Context) { + if manager == nil { + c.Next() + return + } + + result, err := manager.Authenticate(c.Request.Context(), c.Request) + if err == nil { + if result != nil { + c.Set("userApiKey", result.Principal) + c.Set("accessProvider", result.Provider) + if len(result.Metadata) > 0 { + c.Set("accessMetadata", result.Metadata) + } + } + c.Next() + return + } + + statusCode := err.HTTPStatusCode() + if statusCode >= http.StatusInternalServerError { + log.Errorf("authentication middleware error: %v", err) + } + c.AbortWithStatusJSON(statusCode, gin.H{"error": err.Message}) + } +} diff --git a/internal/api/server_options.go b/internal/api/server_options.go new file mode 100644 index 000000000..ef254febc --- /dev/null +++ b/internal/api/server_options.go @@ -0,0 +1,135 @@ +package api + +import ( + "context" + "path/filepath" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +type serverOptionConfig struct { + extraMiddleware []gin.HandlerFunc + engineConfigurator func(*gin.Engine) + routerConfigurator func(*gin.Engine, *handlers.BaseAPIHandler, *config.Config) + requestLoggerFactory func(*config.Config, string) logging.RequestLogger + localPassword string + keepAliveEnabled bool + keepAliveTimeout time.Duration + keepAliveOnTimeout func() + postAuthHook auth.PostAuthHook + postAuthPersistHook auth.PostAuthHook + pluginHost *pluginhost.Host + configReloadHook func(context.Context, *config.Config) + exampleAPIKeySafeMode bool +} + +// ServerOption customises HTTP server construction. +type ServerOption func(*serverOptionConfig) + +func defaultRequestLoggerFactory(cfg *config.Config, configPath string) logging.RequestLogger { + configDir := filepath.Dir(configPath) + logsDir := logging.ResolveLogDirectory(cfg) + logger := logging.NewFileRequestLogger(cfg.RequestLog, logsDir, configDir, cfg.ErrorLogsMaxFiles) + logger.SetHomeEnabled(cfg != nil && cfg.Home.Enabled) + return logger +} + +func effectiveSDKConfig(cfg *config.Config) *config.SDKConfig { + if cfg == nil { + return nil + } + sdkCfg := cfg.SDKConfig + sdkCfg.CodexOptimizeMultiAgentV2 = cfg.Codex.OptimizeMultiAgentV2 + if cfg.CommercialMode { + sdkCfg.RequestLog = false + } + return &sdkCfg +} + +// WithMiddleware appends additional Gin middleware during server construction. +func WithMiddleware(mw ...gin.HandlerFunc) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.extraMiddleware = append(cfg.extraMiddleware, mw...) + } +} + +// WithEngineConfigurator allows callers to mutate the Gin engine prior to middleware setup. +func WithEngineConfigurator(fn func(*gin.Engine)) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.engineConfigurator = fn + } +} + +// WithRouterConfigurator appends a callback after default routes are registered. +func WithRouterConfigurator(fn func(*gin.Engine, *handlers.BaseAPIHandler, *config.Config)) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.routerConfigurator = fn + } +} + +// WithLocalManagementPassword stores a runtime-only management password accepted for localhost requests. +func WithLocalManagementPassword(password string) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.localPassword = password + } +} + +// WithKeepAliveEndpoint enables a keep-alive endpoint with the provided timeout and callback. +func WithKeepAliveEndpoint(timeout time.Duration, onTimeout func()) ServerOption { + return func(cfg *serverOptionConfig) { + if timeout <= 0 || onTimeout == nil { + return + } + cfg.keepAliveEnabled = true + cfg.keepAliveTimeout = timeout + cfg.keepAliveOnTimeout = onTimeout + } +} + +// WithRequestLoggerFactory customises request logger creation. +func WithRequestLoggerFactory(factory func(*config.Config, string) logging.RequestLogger) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.requestLoggerFactory = factory + } +} + +// WithPostAuthHook registers a hook to be called after auth record creation. +func WithPostAuthHook(hook auth.PostAuthHook) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.postAuthHook = hook + } +} + +// WithPostAuthPersistHook registers a hook to be called after auth persistence. +func WithPostAuthPersistHook(hook auth.PostAuthHook) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.postAuthPersistHook = hook + } +} + +// WithPluginHost registers dynamic plugin HTTP adapters with the server. +func WithPluginHost(host *pluginhost.Host) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.pluginHost = host + } +} + +// WithConfigReloadHook registers a callback used after management saves config changes. +func WithConfigReloadHook(hook func(context.Context, *config.Config)) ServerOption { + return func(cfg *serverOptionConfig) { + cfg.configReloadHook = hook + } +} + +// WithExampleAPIKeySafeMode blocks proxy API endpoints while template API keys remain configured. +func WithExampleAPIKeySafeMode() ServerOption { + return func(cfg *serverOptionConfig) { + cfg.exampleAPIKeySafeMode = true + } +} diff --git a/internal/api/server_reload.go b/internal/api/server_reload.go new file mode 100644 index 000000000..95c8c6706 --- /dev/null +++ b/internal/api/server_reload.go @@ -0,0 +1,276 @@ +package api + +import ( + "context" + "fmt" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/access" + "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/internal/managementasset" + "github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" + "gopkg.in/yaml.v3" +) + +func (s *Server) applyAccessConfig(oldCfg, newCfg *config.Config) bool { + if s == nil || s.accessManager == nil || newCfg == nil { + return false + } + if _, err := access.ApplyAccessProviders(s.accessManager, oldCfg, newCfg); err != nil { + return false + } + return true +} + +// UpdateClients updates the server's client list and configuration. +// This method is called when the configuration or authentication tokens change. +// +// Parameters: +// - clients: The new slice of AI service clients +// - cfg: The new application configuration +func (s *Server) UpdateClients(cfg *config.Config) { + s.UpdateClientsContext(context.Background(), cfg) +} + +// UpdateClientsContext updates runtime clients while honoring cancellation between +// short configuration and filesystem operations. +func (s *Server) UpdateClientsContext(ctx context.Context, cfg *config.Config) bool { + if s == nil || cfg == nil { + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + // Reconstruct old config from YAML snapshot to avoid reference sharing issues + var oldCfg *config.Config + if len(s.oldConfigYaml) > 0 { + _ = yaml.Unmarshal(s.oldConfigYaml, &oldCfg) + } + + // Update request logger enabled state if it has changed + previousRequestLog := false + if oldCfg != nil { + previousRequestLog = oldCfg.RequestLog + } + if s.requestLogger != nil && (oldCfg == nil || previousRequestLog != cfg.RequestLog) { + if s.loggerToggle != nil { + s.loggerToggle(cfg.RequestLog) + } else if toggler, ok := s.requestLogger.(interface{ SetEnabled(bool) }); ok { + toggler.SetEnabled(cfg.RequestLog) + } + } + + if oldCfg == nil || oldCfg.Home.Enabled != cfg.Home.Enabled { + if setter, ok := s.requestLogger.(interface{ SetHomeEnabled(bool) }); ok { + setter.SetHomeEnabled(cfg.Home.Enabled) + } + } + + if oldCfg == nil || oldCfg.LoggingToFile != cfg.LoggingToFile || oldCfg.LogsMaxTotalSizeMB != cfg.LogsMaxTotalSizeMB { + if err := logging.ConfigureLogOutput(cfg); err != nil { + log.Errorf("failed to reconfigure log output: %v", err) + } + if errContext := ctx.Err(); errContext != nil { + return false + } + } + + if oldCfg == nil || oldCfg.UsageStatisticsEnabled != cfg.UsageStatisticsEnabled { + redisqueue.SetUsageStatisticsEnabled(cfg.UsageStatisticsEnabled) + } + + if oldCfg == nil || oldCfg.RedisUsageQueueRetentionSeconds != cfg.RedisUsageQueueRetentionSeconds { + redisqueue.SetRetentionSeconds(cfg.RedisUsageQueueRetentionSeconds) + } + + if s.requestLogger != nil && (oldCfg == nil || oldCfg.ErrorLogsMaxFiles != cfg.ErrorLogsMaxFiles) { + if setter, ok := s.requestLogger.(interface{ SetErrorLogsMaxFiles(int) }); ok { + setter.SetErrorLogsMaxFiles(cfg.ErrorLogsMaxFiles) + } + } + + if oldCfg == nil || oldCfg.DisableCooling != cfg.DisableCooling { + auth.SetQuotaCooldownDisabled(cfg.DisableCooling) + } + if oldCfg == nil || oldCfg.TransientErrorCooldownSeconds != cfg.TransientErrorCooldownSeconds { + auth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) + } + + if oldCfg != nil && oldCfg.DisableImageGeneration != cfg.DisableImageGeneration { + log.Infof("disable-image-generation updated: %v -> %v", oldCfg.DisableImageGeneration, cfg.DisableImageGeneration) + } + + applySignatureCacheConfig(oldCfg, cfg) + + if s.handlers != nil && s.handlers.AuthManager != nil { + s.handlers.AuthManager.SetRetryConfig(cfg.RequestRetry, time.Duration(cfg.MaxRetryInterval)*time.Second, cfg.MaxRetryCredentials) + } + + // Update log level dynamically when debug flag changes + if oldCfg == nil || oldCfg.Debug != cfg.Debug { + util.SetLogLevel(cfg) + } + + prevSecretEmpty := true + if oldCfg != nil { + prevSecretEmpty = oldCfg.RemoteManagement.SecretKey == "" + } + newSecretEmpty := cfg.RemoteManagement.SecretKey == "" + if s.envManagementSecret { + s.registerManagementRoutes() + if s.managementRoutesEnabled.CompareAndSwap(false, true) { + log.Info("management routes enabled via MANAGEMENT_PASSWORD") + } else { + s.managementRoutesEnabled.Store(true) + } + } else { + switch { + case prevSecretEmpty && !newSecretEmpty: + s.registerManagementRoutes() + if s.managementRoutesEnabled.CompareAndSwap(false, true) { + log.Info("management routes enabled after secret key update") + } else { + s.managementRoutesEnabled.Store(true) + } + case !prevSecretEmpty && newSecretEmpty: + if s.managementRoutesEnabled.CompareAndSwap(true, false) { + log.Info("management routes disabled after secret key removal") + } else { + s.managementRoutesEnabled.Store(false) + } + default: + s.managementRoutesEnabled.Store(!newSecretEmpty) + } + } + redisqueue.SetEnabled(s.managementRoutesEnabled.Load() || (cfg != nil && cfg.Home.Enabled)) + + exampleAPIKeySafeModeRequired := s.exampleAPIKeySafeModeRequired(cfg) + if exampleAPIKeySafeModeRequired { + s.exampleAPIKeySafeModeActive.Store(true) + } + accessConfigApplied := s.applyAccessConfig(oldCfg, cfg) + if accessConfigApplied || exampleAPIKeySafeModeRequired { + s.exampleAPIKeySafeModeActive.Store(exampleAPIKeySafeModeRequired) + } + s.cfg = cfg + if s.codexLiveHandler != nil { + if errUpdate := s.codexLiveHandler.UpdateConfig(cfg); errUpdate != nil { + log.WithError(errUpdate).Error("failed to update Codex Live media relay configuration") + } + } + s.wsAuthEnabled.Store(cfg.WebsocketAuth) + if oldCfg != nil && s.wsAuthChanged != nil && oldCfg.WebsocketAuth != cfg.WebsocketAuth { + s.wsAuthChanged(oldCfg.WebsocketAuth, cfg.WebsocketAuth) + } + managementasset.SetCurrentConfig(cfg) + if errContext := ctx.Err(); errContext != nil { + return false + } + // Save YAML snapshot for next comparison + s.oldConfigYaml, _ = yaml.Marshal(cfg) + + s.handlers.UpdateClients(effectiveSDKConfig(cfg)) + s.handlers.SetPluginHost(s.pluginHost) + if s.pluginHost != nil { + s.pluginHost.SetModelExecutor(s.handlers) + s.pluginHost.SetAuthManager(s.handlers.AuthManager) + } + + if s.mgmt != nil { + s.mgmt.SetConfig(cfg) + s.mgmt.SetAuthManager(s.handlers.AuthManager) + s.mgmt.SetPluginHost(s.pluginHost) + } + s.refreshPluginManagementRoutes() + + // Count client sources from configuration and auth store. + authEntries := 0 + if cfg != nil && !cfg.Home.Enabled { + tokenStore := sdkAuth.GetTokenStore() + if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok { + dirSetter.SetBaseDir(cfg.AuthDir) + } + authEntries = util.CountAuthFiles(ctx, tokenStore) + if errContext := ctx.Err(); errContext != nil { + return false + } + } + geminiAPIKeyCount := len(cfg.GeminiKey) + interactionsAPIKeyCount := len(cfg.InteractionsKey) + claudeAPIKeyCount := len(cfg.ClaudeKey) + codexAPIKeyCount := len(cfg.CodexKey) + xaiAPIKeyCount := len(cfg.XAIKey) + vertexAICompatCount := len(cfg.VertexCompatAPIKey) + openAICompatCount := 0 + for i := range cfg.OpenAICompatibility { + entry := cfg.OpenAICompatibility[i] + if entry.Disabled { + continue + } + openAICompatCount += len(entry.APIKeyEntries) + } + + total := authEntries + geminiAPIKeyCount + interactionsAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + xaiAPIKeyCount + vertexAICompatCount + openAICompatCount + fmt.Printf("server clients and configuration updated: %d clients (%d auth entries + %d Gemini API keys + %d Interactions API keys + %d Claude API keys + %d Codex keys + %d xAI keys + %d Vertex-compat + %d OpenAI-compat)\n", + total, + authEntries, + geminiAPIKeyCount, + interactionsAPIKeyCount, + claudeAPIKeyCount, + codexAPIKeyCount, + xaiAPIKeyCount, + vertexAICompatCount, + openAICompatCount, + ) + return ctx.Err() == nil +} + +func (s *Server) SetWebsocketAuthChangeHandler(fn func(bool, bool)) { + if s == nil { + return + } + s.wsAuthChanged = fn +} + +func configuredSignatureCacheEnabled(cfg *config.Config) bool { + if cfg != nil && cfg.AntigravitySignatureCacheEnabled != nil { + return *cfg.AntigravitySignatureCacheEnabled + } + return true +} + +func applySignatureCacheConfig(oldCfg, cfg *config.Config) { + newVal := configuredSignatureCacheEnabled(cfg) + newStrict := configuredSignatureBypassStrict(cfg) + if oldCfg == nil { + cache.SetSignatureCacheEnabled(newVal) + cache.SetSignatureBypassStrictMode(newStrict) + return + } + + oldVal := configuredSignatureCacheEnabled(oldCfg) + if oldVal != newVal { + cache.SetSignatureCacheEnabled(newVal) + } + + oldStrict := configuredSignatureBypassStrict(oldCfg) + if oldStrict != newStrict { + cache.SetSignatureBypassStrictMode(newStrict) + } +} + +func configuredSignatureBypassStrict(cfg *config.Config) bool { + if cfg != nil && cfg.AntigravitySignatureBypassStrict != nil { + return *cfg.AntigravitySignatureBypassStrict + } + return false +} diff --git a/internal/api/server_routes.go b/internal/api/server_routes.go new file mode 100644 index 000000000..e79015807 --- /dev/null +++ b/internal/api/server_routes.go @@ -0,0 +1,896 @@ +package api + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "sort" + "strconv" + "strings" + "time" + + "github.com/gin-gonic/gin" + managementHandlers "github.com/router-for-me/CLIProxyAPI/v7/internal/api/handlers/management" + claudemodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/claude/models" + codexlive "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/live" + codexmodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/models" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers/claude" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers/gemini" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers/openai" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + log "github.com/sirupsen/logrus" +) + +const oauthCallbackSuccessHTML = `Authentication successful

Authentication successful!

You can close this window.

This window will close automatically in 5 seconds.

` + +const codexAlphaSearchSourceFormat = "codex-alpha-search" + +// setupRoutes configures the API routes for the server. +// It defines the endpoints and associates them with their respective handlers. +func (s *Server) setupRoutes() { + healthzHandler := func(c *gin.Context) { + if c.Request.Method == http.MethodHead { + c.Status(http.StatusOK) + return + } + + c.JSON(http.StatusOK, gin.H{"status": "ok"}) + } + s.engine.GET("/healthz", healthzHandler) + s.engine.HEAD("/healthz", healthzHandler) + + s.engine.GET("/management.html", s.serveManagementControlPanel) + openaiHandlers := openai.NewOpenAIAPIHandler(s.handlers) + geminiHandlers := gemini.NewGeminiAPIHandler(s.handlers) + claudeCodeHandlers := claude.NewClaudeCodeAPIHandler(s.handlers) + openaiResponsesHandlers := openai.NewOpenAIResponsesAPIHandler(s.handlers) + s.codexLiveHandler = codexlive.NewHandler(s.handlers.AuthManager, s.cfg) + + // OpenAI compatible API routes + v1 := s.engine.Group("/v1") + v1.Use(AuthMiddleware(s.accessManager)) + { + v1.GET("/models", s.unifiedModelsHandler(openaiHandlers, claudeCodeHandlers)) + v1.POST("/chat/completions", openaiHandlers.ChatCompletions) + v1.POST("/completions", openaiHandlers.Completions) + v1.POST("/images/generations", openaiHandlers.ImagesGenerations) + v1.POST("/images/edits", openaiHandlers.ImagesEdits) + v1.POST("/videos", openaiHandlers.XAIVideosGenerations) + v1.POST("/videos/generations", openaiHandlers.XAIVideosGenerations) + v1.POST("/videos/edits", openaiHandlers.XAIVideosEdits) + v1.POST("/videos/extensions", openaiHandlers.XAIVideosExtensions) + v1.GET("/videos/:request_id", openaiHandlers.XAIVideosRetrieve) + v1.POST("/messages", claudeCodeHandlers.ClaudeMessages) + v1.POST("/messages/count_tokens", claudeCodeHandlers.ClaudeCountTokens) + v1.GET("/responses", openaiResponsesHandlers.ResponsesWebsocket) + v1.POST("/responses", openaiResponsesHandlers.Responses) + v1.POST("/responses/compact", openaiResponsesHandlers.Compact) + v1.POST("/alpha/search", s.codexAlphaSearch) + v1.POST("/live", s.codexLiveHandler.Handle) + v1.GET("/live/:call_id", s.codexLiveHandler.HandleSideband) + v1.POST("/realtime/calls", s.codexLiveHandler.Handle) + v1.GET("/realtime/calls/:call_id", s.codexLiveHandler.HandleSideband) + v1.GET("/realtime", s.codexLiveHandler.HandleSideband) + } + + openaiV1 := s.engine.Group("/openai/v1") + openaiV1.Use(AuthMiddleware(s.accessManager)) + { + openaiV1.POST("/videos", openaiHandlers.VideosCreate) + openaiV1.GET("/videos/:video_id/content", openaiHandlers.VideosContent) + openaiV1.GET("/videos/:video_id", openaiHandlers.VideosRetrieve) + } + + // Codex CLI direct route aliases (chatgpt_base_url compatible) + codexDirect := s.engine.Group("/backend-api/codex") + codexDirect.Use(AuthMiddleware(s.accessManager)) + { + codexDirect.GET("/responses", openaiResponsesHandlers.ResponsesWebsocket) + codexDirect.POST("/responses", openaiResponsesHandlers.Responses) + codexDirect.POST("/responses/compact", openaiResponsesHandlers.Compact) + codexDirect.POST("/alpha/search", s.codexAlphaSearch) + } + + // Gemini compatible API routes + v1beta := s.engine.Group("/v1beta") + v1beta.Use(AuthMiddleware(s.accessManager)) + { + v1beta.GET("/models", s.geminiModelsHandler(geminiHandlers)) + v1beta.POST("/interactions", geminiHandlers.Interactions) + v1beta.POST("/models/*action", geminiHandlers.GeminiHandler) + v1beta.GET("/models/*action", s.geminiGetHandler(geminiHandlers)) + } + + // Root endpoint + s.engine.GET("/", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{ + "message": "CLI Proxy API Server", + "endpoints": []string{ + "POST /v1/chat/completions", + "POST /v1/completions", + "GET /v1/models", + }, + }) + }) + + // OAuth callback endpoints (reuse main server port) + // These endpoints receive provider redirects and persist + // the short-lived code/state for the waiting goroutine. + s.engine.GET("/anthropic/callback", func(c *gin.Context) { + code := c.Query("code") + state := c.Query("state") + errStr := c.Query("error") + if errStr == "" { + errStr = c.Query("error_description") + } + if state != "" { + _, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "anthropic", state, code, errStr) + } + c.Header("Content-Type", "text/html; charset=utf-8") + c.String(http.StatusOK, oauthCallbackSuccessHTML) + }) + + s.engine.GET("/codex/callback", func(c *gin.Context) { + code := c.Query("code") + state := c.Query("state") + errStr := c.Query("error") + if errStr == "" { + errStr = c.Query("error_description") + } + if state != "" { + _, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "codex", state, code, errStr) + } + c.Header("Content-Type", "text/html; charset=utf-8") + c.String(http.StatusOK, oauthCallbackSuccessHTML) + }) + + s.engine.GET("/antigravity/callback", func(c *gin.Context) { + code := c.Query("code") + state := c.Query("state") + errStr := c.Query("error") + if errStr == "" { + errStr = c.Query("error_description") + } + if state != "" { + _, _ = managementHandlers.WriteOAuthCallbackFileForPendingSession(s.cfg.AuthDir, "antigravity", state, code, errStr) + } + c.Header("Content-Type", "text/html; charset=utf-8") + c.String(http.StatusOK, oauthCallbackSuccessHTML) + }) + + // Management routes are registered lazily by registerManagementRoutes when a secret is configured. +} + +func (s *Server) codexAlphaSearchModelRouterHost() handlers.PluginModelRouterHost { + if s == nil { + return nil + } + if s.pluginHost != nil { + return s.pluginHost + } + if s.handlers != nil && s.handlers.ModelRouterHost != nil { + return s.handlers.ModelRouterHost + } + return nil +} + +func (s *Server) codexAlphaSearchSelectionModel(ctx context.Context, c *gin.Context, body []byte, model string) (string, error) { + host := s.codexAlphaSearchModelRouterHost() + if host == nil { + return model, nil + } + + var headers http.Header + queryValues := make(map[string][]string) + requestPath := "" + if c != nil && c.Request != nil { + headers = c.Request.Header.Clone() + if c.Request.URL != nil { + queryValues = c.Request.URL.Query() + requestPath = c.Request.URL.Path + } + } + metadata := map[string]any{ + coreexecutor.RequestedModelMetadataKey: model, + } + if requestPath != "" { + metadata[coreexecutor.RequestPathMetadataKey] = requestPath + } + resp, handled := host.RouteModel(ctx, pluginapi.ModelRouteRequest{ + SourceFormat: codexAlphaSearchSourceFormat, + RequestedModel: model, + Headers: headers, + Query: queryValues, + Body: body, + Metadata: metadata, + }) + if !handled || !resp.Handled { + return model, nil + } + if resp.TargetKind != pluginapi.ModelRouteTargetProvider || !strings.EqualFold(strings.TrimSpace(resp.Target), "codex") { + return "", fmt.Errorf("unsupported Codex Alpha Search model route target %q (%q)", resp.TargetKind, resp.Target) + } + if targetModel := strings.TrimSpace(resp.TargetModel); targetModel != "" { + return targetModel, nil + } + return model, nil +} + +func sanitizeCodexAlphaSearchBody(body []byte) []byte { + var payload map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(body, &payload); errUnmarshal != nil || payload == nil { + return body + } + + removed := false + for _, field := range []string{"prompt_cache_key", "prompt_cache_retention"} { + if _, exists := payload[field]; exists { + delete(payload, field) + removed = true + } + } + if !removed { + return body + } + + sanitizedBody, errMarshal := json.Marshal(payload) + if errMarshal != nil { + return body + } + return sanitizedBody +} + +func homeSelectionAttemptContext(ctx context.Context, selection *auth.HomeDispatchSelection) (context.Context, func(), error) { + if selection == nil { + return nil, func() {}, errors.New("Home dispatch selection is nil") + } + return selection.AttemptContext(ctx) +} + +// codexAlphaSearch forwards the standalone search endpoint used by current +// Codex clients. Unlike /responses, this payload is already in Codex search +// format and must not pass through a protocol translator. +func (s *Server) codexAlphaSearch(c *gin.Context) { + if s == nil || s.handlers == nil || s.handlers.AuthManager == nil { + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "Codex auth manager unavailable"}) + return + } + + body, err := io.ReadAll(io.LimitReader(c.Request.Body, 16<<20)) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "Failed to read search request"}) + return + } + + var routing struct { + ID string `json:"id"` + Model string `json:"model"` + } + _ = json.Unmarshal(body, &routing) + upstreamRequestBody := sanitizeCodexAlphaSearchBody(body) + + selectionHeaders := c.Request.Header.Clone() + if sessionID := strings.TrimSpace(routing.ID); sessionID != "" { + selectionHeaders.Set("X-Session-ID", sessionID) + } + ctx := context.WithValue(c.Request.Context(), "gin", c) + selectionModel, errRoute := s.codexAlphaSearchSelectionModel(ctx, c, body, strings.TrimSpace(routing.Model)) + if errRoute != nil { + log.WithError(errRoute).Warn("codex alpha search: model router returned an unsupported target") + c.JSON(http.StatusServiceUnavailable, gin.H{"error": errRoute.Error()}) + return + } + selectionOpts := coreexecutor.Options{Headers: selectionHeaders, OriginalRequest: body} + var selection *auth.HomeDispatchSelection + var selected *auth.Auth + if s.handlers.AuthManager.HomeEnabled() { + selection, err = s.handlers.AuthManager.SelectHomeAuthByKind(ctx, "codex", selectionModel, auth.AuthKindOAuth, selectionOpts) + if selection != nil { + selected = selection.CloneAuth() + } + } else { + selected, err = s.handlers.AuthManager.SelectAuthByKind(ctx, "codex", selectionModel, auth.AuthKindOAuth, selectionOpts) + } + if err != nil { + status := http.StatusServiceUnavailable + if statusError, ok := err.(interface{ StatusCode() int }); ok && statusError.StatusCode() > 0 { + status = statusError.StatusCode() + } + for _, value := range auth.SafeResponseHeaders(err).Values("Retry-After") { + c.Writer.Header().Add("Retry-After", value) + } + c.JSON(status, gin.H{"error": err.Error()}) + return + } + if selected == nil { + if selection != nil { + selection.End("missing_auth") + } + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "Codex auth unavailable"}) + return + } + var releaseAttempt func() + if selection != nil { + attemptCtx, release, errBind := homeSelectionAttemptContext(ctx, selection) + if errBind != nil { + selection.End("attempt_bind_failed") + c.JSON(http.StatusServiceUnavailable, gin.H{"error": errBind.Error()}) + return + } + ctx = attemptCtx + releaseAttempt = release + defer releaseAttempt() + } + logging.SetGinCPATraceID(c, selected.EnsureIndex()) + + headers := make(http.Header) + headers.Set("Content-Type", "application/json") + headers.Set("Accept", "application/json") + headers.Set("Originator", "codex_cli_rs") + for _, name := range []string{"Version", "User-Agent", "Session_id", "X-Client-Request-Id"} { + if value := strings.TrimSpace(c.GetHeader(name)); value != "" { + headers.Set(name, value) + } + } + if accountID, ok := selected.Metadata["account_id"].(string); ok && strings.TrimSpace(accountID) != "" { + headers.Set("Chatgpt-Account-Id", accountID) + } + + const upstreamURL = "https://chatgpt.com/backend-api/codex/alpha/search" + req, err := s.handlers.AuthManager.NewHttpRequest( + ctx, selected, http.MethodPost, upstreamURL, upstreamRequestBody, headers, + ) + if err != nil { + if selection != nil { + selection.End("request_build_failed") + } + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + + var authID, authLabel, authType, authValue string + if selected != nil { + authID = selected.ID + authLabel = selected.Label + authType, authValue = selected.AccountInfo() + } + helpHeaders := req.Header.Clone() + helps.RecordAPIRequest(ctx, s.cfg, helps.UpstreamRequestLog{ + URL: upstreamURL, + Method: http.MethodPost, + Headers: helpHeaders, + Body: upstreamRequestBody, + Provider: "codex", + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + if errCtx := ctx.Err(); errCtx != nil { + if selection != nil { + selection.End("attempt_canceled") + } + c.JSON(http.StatusRequestTimeout, gin.H{"error": errCtx.Error()}) + return + } + resp, err := s.handlers.AuthManager.HttpRequest(ctx, selected, req) + if err != nil { + if selection != nil { + selection.End("request_failed") + } + helps.RecordAPIResponseError(ctx, s.cfg, err) + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + closeResponseBody := func() error { + errClose := resp.Body.Close() + if errClose != nil { + log.Errorf("codex alpha search: close response body error: %v", errClose) + } + return errClose + } + if selection != nil { + if errBind := selection.Bind(closeResponseBody); errBind != nil { + selection.End("response_bind_failed") + c.JSON(http.StatusServiceUnavailable, gin.H{"error": errBind.Error()}) + return + } + defer selection.End("response_closed") + } else { + defer func() { _ = closeResponseBody() }() + } + helps.RecordAPIResponseMetadata(ctx, s.cfg, resp.StatusCode, resp.Header.Clone()) + upstreamBody, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20)) + if err != nil { + helps.RecordAPIResponseError(ctx, s.cfg, err) + c.JSON(http.StatusBadGateway, gin.H{"error": "Failed to read Codex search response"}) + return + } + helps.AppendAPIResponseChunk(ctx, s.cfg, upstreamBody) + if contentType := resp.Header.Get("Content-Type"); contentType != "" { + c.Header("Content-Type", contentType) + } + c.Status(resp.StatusCode) + _, _ = c.Writer.Write(upstreamBody) +} + +// AttachWebsocketRoute registers a websocket upgrade handler on the primary Gin engine. +// The handler is served as-is without additional middleware beyond the standard stack already configured. +func (s *Server) AttachWebsocketRoute(path string, handler http.Handler) { + if s == nil || s.engine == nil || handler == nil { + return + } + trimmed := strings.TrimSpace(path) + if trimmed == "" { + trimmed = "/v1/ws" + } + if !strings.HasPrefix(trimmed, "/") { + trimmed = "/" + trimmed + } + s.wsRouteMu.Lock() + if _, exists := s.wsRoutes[trimmed]; exists { + s.wsRouteMu.Unlock() + return + } + s.wsRoutes[trimmed] = struct{}{} + s.wsRouteMu.Unlock() + + authMiddleware := AuthMiddleware(s.accessManager) + conditionalAuth := func(c *gin.Context) { + if !s.wsAuthEnabled.Load() { + c.Next() + return + } + authMiddleware(c) + } + finalHandler := func(c *gin.Context) { + handler.ServeHTTP(c.Writer, c.Request) + c.Abort() + } + + s.engine.GET(trimmed, conditionalAuth, finalHandler) +} + +// isAnthropicModelsRequest reports whether a /v1/models request should be served in +// Anthropic format. Anthropic API clients send the Anthropic-Version header; Claude +// Code additionally uses a claude-cli User-Agent. +func isAnthropicModelsRequest(c *gin.Context) bool { + if c.GetHeader("Anthropic-Version") != "" { + return true + } + return strings.HasPrefix(c.GetHeader("User-Agent"), "claude-cli") +} + +// unifiedModelsHandler creates a unified handler for the /v1/models endpoint +// that routes to different handlers based on the request. +// Anthropic API requests (Anthropic-Version header, or a claude-cli User-Agent) +// route to the Claude handler, otherwise they route to the OpenAI handler. +func (s *Server) unifiedModelsHandler(openaiHandler *openai.OpenAIAPIHandler, claudeHandler *claude.ClaudeCodeAPIHandler) gin.HandlerFunc { + return func(c *gin.Context) { + if _, ok := c.Request.URL.Query()["client_version"]; ok { + if s != nil && s.cfg != nil && s.cfg.Home.Enabled { + s.handleHomeCodexClientModels(c) + return + } + openaiHandler.OpenAIModels(c) + return + } + + if s != nil && s.cfg != nil && s.cfg.Home.Enabled { + s.handleHomeModels(c) + return + } + + // Route to Claude handler for Anthropic API requests. + if isAnthropicModelsRequest(c) { + claudeHandler.ClaudeModels(c) + } else { + openaiHandler.OpenAIModels(c) + } + } +} + +// handleHomeCodexClientModels builds the Codex client catalog from Home model IDs. +// Template metadata still comes from the local/remote codex_client_models catalog. +func (s *Server) handleHomeCodexClientModels(c *gin.Context) { + entries, ok := s.loadHomeModelEntries(c) + if !ok { + return + } + + models := make([]map[string]any, 0, len(entries)) + for _, entry := range entries { + model := map[string]any{ + "id": entry.id, + "object": "model", + } + if entry.created > 0 { + model["created"] = entry.created + } + if entry.ownedBy != "" { + model["owned_by"] = entry.ownedBy + } + if entry.displayName != "" { + model["display_name"] = entry.displayName + model["description"] = entry.displayName + } + models = append(models, model) + } + + c.JSON(http.StatusOK, codexmodels.BuildResponse(models, nil, s.cfg.Codex.OptimizeMultiAgentV2)) +} + +func (s *Server) geminiModelsHandler(geminiHandler *gemini.GeminiAPIHandler) gin.HandlerFunc { + return func(c *gin.Context) { + if s != nil && s.cfg != nil && s.cfg.Home.Enabled { + s.handleHomeGeminiModels(c) + return + } + + geminiHandler.GeminiModels(c) + } +} + +func (s *Server) geminiGetHandler(geminiHandler *gemini.GeminiAPIHandler) gin.HandlerFunc { + return func(c *gin.Context) { + if s != nil && s.cfg != nil && s.cfg.Home.Enabled { + s.handleHomeGeminiModel(c) + return + } + + geminiHandler.GeminiGetHandler(c) + } +} + +type homeModelEntry struct { + id string + created int64 + ownedBy string + displayName string + contextLength int + maxCompletionTokens int +} + +func (s *Server) handleHomeModels(c *gin.Context) { + entries, ok := s.loadHomeModelEntries(c) + if !ok { + return + } + + isClaude := isAnthropicModelsRequest(c) + + if isClaude { + c.JSON(http.StatusOK, claudemodels.BuildResponse(formatHomeClaudeModels(entries))) + return + } + + filtered := make([]map[string]any, 0, len(entries)) + for _, entry := range entries { + model := map[string]any{ + "id": entry.id, + "object": "model", + } + if entry.created > 0 { + model["created"] = entry.created + } + if entry.ownedBy != "" { + model["owned_by"] = entry.ownedBy + } + filtered = append(filtered, model) + } + c.JSON(http.StatusOK, gin.H{ + "object": "list", + "data": filtered, + }) +} + +func formatHomeClaudeModels(entries []homeModelEntry) []map[string]any { + out := make([]map[string]any, 0, len(entries)) + for _, entry := range entries { + out = append(out, formatHomeClaudeModel(entry)) + } + return out +} + +func formatHomeClaudeModel(entry homeModelEntry) map[string]any { + displayName := entry.displayName + if displayName == "" { + displayName = entry.id + } + maxInput := entry.contextLength + if maxInput <= 0 { + maxInput = registry.DefaultClaudeMaxInputTokens + } + maxOutput := entry.maxCompletionTokens + if maxOutput <= 0 { + maxOutput = registry.DefaultClaudeMaxOutputTokens + } + model := map[string]any{ + "id": entry.id, + "object": "model", + "owned_by": entry.ownedBy, + "type": "model", + "display_name": displayName, + "max_input_tokens": maxInput, + "max_tokens": maxOutput, + } + if entry.created > 0 { + model["created_at"] = time.Unix(entry.created, 0).UTC().Format(time.RFC3339) + } + return model +} + +func (s *Server) handleHomeGeminiModels(c *gin.Context) { + entries, ok := s.loadHomeModelEntries(c) + if !ok { + return + } + + c.JSON(http.StatusOK, gin.H{ + "models": formatHomeGeminiModels(entries), + }) +} + +func (s *Server) handleHomeGeminiModel(c *gin.Context) { + entries, ok := s.loadHomeModelEntries(c) + if !ok { + return + } + + action := strings.TrimPrefix(c.Param("action"), "/") + action = strings.TrimSpace(action) + for _, entry := range entries { + if homeGeminiModelMatches(entry, action) { + c.JSON(http.StatusOK, formatHomeGeminiModel(entry)) + return + } + } + + c.JSON(http.StatusNotFound, handlers.ErrorResponse{ + Error: handlers.ErrorDetail{ + Message: "Not Found", + Type: "not_found", + }, + }) +} + +func (s *Server) loadHomeModelEntries(c *gin.Context) ([]homeModelEntry, bool) { + if s == nil || c == nil || c.Request == nil { + return nil, false + } + client := home.Current() + if client == nil { + c.JSON(http.StatusServiceUnavailable, handlers.ErrorResponse{ + Error: handlers.ErrorDetail{ + Message: "home control center unavailable", + Type: "server_error", + }, + }) + return nil, false + } + + raw, errGet := client.GetModels(c.Request.Context(), c.Request.Header, c.Request.URL.Query()) + if errGet != nil { + c.JSON(http.StatusBadGateway, handlers.ErrorResponse{ + Error: handlers.ErrorDetail{ + Message: errGet.Error(), + Type: "server_error", + }, + }) + return nil, false + } + + if statusCode, ok := homeModelsAuthStatus(raw); ok { + c.JSON(statusCode, handlers.ErrorResponse{ + Error: handlers.ErrorDetail{ + Message: homeModelsErrorMessage(raw), + Type: "authentication_error", + }, + }) + return nil, false + } + + entries, errDecode := decodeHomeModels(raw) + if errDecode != nil { + c.JSON(http.StatusBadGateway, handlers.ErrorResponse{ + Error: handlers.ErrorDetail{ + Message: errDecode.Error(), + Type: "server_error", + }, + }) + return nil, false + } + + return entries, true +} + +func formatHomeGeminiModels(entries []homeModelEntry) []map[string]any { + out := make([]map[string]any, 0, len(entries)) + for _, entry := range entries { + out = append(out, formatHomeGeminiModel(entry)) + } + return out +} + +func formatHomeGeminiModel(entry homeModelEntry) map[string]any { + name := entry.id + if !strings.HasPrefix(name, "models/") { + name = "models/" + name + } + displayName := entry.displayName + if displayName == "" { + displayName = entry.id + } + return map[string]any{ + "name": name, + "displayName": displayName, + "description": displayName, + "supportedGenerationMethods": []string{"generateContent"}, + } +} + +func homeGeminiModelMatches(entry homeModelEntry, action string) bool { + id := strings.TrimSpace(entry.id) + if id == "" || action == "" { + return false + } + normalizedAction := strings.TrimPrefix(action, "models/") + normalizedID := strings.TrimPrefix(id, "models/") + return action == id || action == "models/"+id || normalizedAction == normalizedID +} + +// homeModelsAuthStatus inspects a home models response for an authentication/error envelope. +// It returns the HTTP status code to surface (401 for credential issues, 502 otherwise) +// and true when the payload is an error response rather than model data. +func homeModelsAuthStatus(raw []byte) (int, bool) { + errType := homeModelsErrorType(raw) + if errType == "" { + return 0, false + } + if errType == "no_credentials" || errType == "invalid_credential" { + return http.StatusUnauthorized, true + } + return http.StatusBadGateway, true +} + +func homeModelsErrorType(raw []byte) string { + top, ok := unmarshalHomeModelsTopLevel(raw) + if !ok { + return "" + } + rawErr, exists := top["error"] + if !exists { + return "" + } + var errObj struct { + Type string `json:"type"` + } + if errUnmarshal := json.Unmarshal(rawErr, &errObj); errUnmarshal != nil { + return "" + } + return strings.TrimSpace(errObj.Type) +} + +func homeModelsErrorMessage(raw []byte) string { + top, ok := unmarshalHomeModelsTopLevel(raw) + if !ok { + return "home models request failed" + } + rawErr, exists := top["error"] + if !exists { + return "home models request failed" + } + var errObj struct { + Message string `json:"message"` + } + if errUnmarshal := json.Unmarshal(rawErr, &errObj); errUnmarshal != nil { + return "home models request failed" + } + if msg := strings.TrimSpace(errObj.Message); msg != "" { + return msg + } + return "home models request failed" +} + +func unmarshalHomeModelsTopLevel(raw []byte) (map[string]json.RawMessage, bool) { + if len(raw) == 0 { + return nil, false + } + var top map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(raw, &top); errUnmarshal != nil { + return nil, false + } + return top, true +} + +func decodeHomeModels(raw []byte) ([]homeModelEntry, error) { + if len(raw) == 0 { + return nil, fmt.Errorf("home models payload is empty") + } + + var bySection map[string][]map[string]any + if err := json.Unmarshal(raw, &bySection); err != nil { + return nil, fmt.Errorf("parse home models payload: %w", err) + } + if len(bySection) == 0 { + return nil, fmt.Errorf("home models payload has no sections") + } + + seen := make(map[string]struct{}) + out := make([]homeModelEntry, 0, 256) + for _, models := range bySection { + for _, model := range models { + id, _ := model["id"].(string) + id = strings.TrimSpace(id) + if id == "" { + name, _ := model["name"].(string) + name = strings.TrimSpace(name) + id = strings.TrimPrefix(name, "models/") + } + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + + ownedBy, _ := model["owned_by"].(string) + ownedBy = strings.TrimSpace(ownedBy) + displayName, _ := model["display_name"].(string) + displayName = strings.TrimSpace(displayName) + if displayName == "" { + displayName, _ = model["displayName"].(string) + displayName = strings.TrimSpace(displayName) + } + + out = append(out, homeModelEntry{ + id: id, + created: homeModelInt64Value(model, "created"), + ownedBy: ownedBy, + displayName: displayName, + contextLength: int(homeModelInt64Value(model, "context_length", "contextLength", "inputTokenLimit", "max_input_tokens")), + maxCompletionTokens: int(homeModelInt64Value(model, "max_completion_tokens", "maxCompletionTokens", "outputTokenLimit", "max_tokens")), + }) + } + } + + sort.Slice(out, func(i, j int) bool { return out[i].id < out[j].id }) + if len(out) == 0 { + return nil, fmt.Errorf("home models payload contains no models") + } + return out, nil +} + +func homeModelInt64Value(model map[string]any, keys ...string) int64 { + for _, key := range keys { + switch value := model[key].(type) { + case float64: + return int64(value) + case int64: + return value + case int: + return int64(value) + case json.Number: + if n, errInt := value.Int64(); errInt == nil { + return n + } + case string: + if n, errParse := strconv.ParseInt(strings.TrimSpace(value), 10, 64); errParse == nil { + return n + } + } + } + return 0 +} diff --git a/internal/config/config.go b/internal/config/config.go index b4353f1fe..e8111b73b 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -4,28 +4,6 @@ // debug settings, proxy configuration, and API keys. package config -import ( - "bytes" - "encoding/json" - "errors" - "fmt" - "os" - "strings" - "syscall" - - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - sdkpluginstore "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginstore" - log "github.com/sirupsen/logrus" - "golang.org/x/crypto/bcrypt" - "gopkg.in/yaml.v3" -) - -const ( - DefaultPanelGitHubRepository = "https://github.com/router-for-me/Cli-Proxy-API-Management-Center" - DefaultPprofAddr = "127.0.0.1:8316" - DefaultAuthDir = "~/.cli-proxy-api" -) - // Config represents the application's configuration, loaded from a YAML file. type Config struct { SDKConfig `yaml:",inline"` @@ -177,1887 +155,3 @@ type Config struct { // Payload defines default and override rules for provider payload parameters. Payload PayloadConfig `yaml:"payload" json:"payload"` } - -// PluginsConfig holds dynamic plugin system settings. -type PluginsConfig struct { - // Enabled toggles dynamic plugin loading. - Enabled bool `yaml:"enabled" json:"enabled"` - // Dir is the plugin discovery directory. - Dir string `yaml:"dir" json:"dir"` - // StoreSources appends third-party plugin store registries to the built-in official source. - StoreSources []string `yaml:"store-sources,omitempty" json:"store-sources,omitempty"` - // StoreAuth defines optional auth rules for plugin store registry, metadata, and artifact requests. - StoreAuth []sdkpluginstore.AuthConfig `yaml:"store-auth,omitempty" json:"store-auth,omitempty"` - // AuthRevision changes when Home-managed plugin credentials change. - AuthRevision int64 `yaml:"auth-revision,omitempty" json:"auth-revision,omitempty"` - // Configs stores per-plugin instance configuration by plugin ID. - Configs map[string]PluginInstanceConfig `yaml:"configs" json:"configs"` -} - -// PluginInstanceConfig stores host-owned plugin settings and the original plugin YAML subtree. -type PluginInstanceConfig struct { - // Enabled toggles this plugin instance. Nil is normalized to false during YAML parsing. - Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` - // Priority controls plugin startup and routing order. - Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` - // Raw preserves the full original plugin configuration YAML subtree. - Raw yaml.Node `yaml:"-" json:"-"` -} - -// UnmarshalYAML extracts host-owned fields while preserving the full original YAML node. -func (c *PluginInstanceConfig) UnmarshalYAML(value *yaml.Node) error { - if c == nil { - return nil - } - - c.Priority = 0 - defaultEnabled := false - c.Enabled = &defaultEnabled - - if value == nil || value.Kind == 0 { - c.Raw = *defaultPluginInstanceConfigNode() - return nil - } - - c.Raw = *deepCopyNode(value) - if value.Kind != yaml.MappingNode { - return nil - } - - for i := 0; i+1 < len(value.Content); i += 2 { - key := value.Content[i] - node := value.Content[i+1] - if key == nil { - continue - } - switch key.Value { - case "enabled": - var enabled bool - if errDecodeEnabled := node.Decode(&enabled); errDecodeEnabled != nil { - return fmt.Errorf("parse plugin enabled: %w", errDecodeEnabled) - } - c.Enabled = &enabled - case "priority": - var priority int - if errDecodePriority := node.Decode(&priority); errDecodePriority != nil { - return fmt.Errorf("parse plugin priority: %w", errDecodePriority) - } - c.Priority = priority - } - } - - return nil -} - -// MarshalYAML returns the preserved raw plugin YAML subtree for lossless config output. -func (c PluginInstanceConfig) MarshalYAML() (any, error) { - if c.Raw.Kind == 0 { - return defaultPluginInstanceConfigNode(), nil - } - return deepCopyNode(&c.Raw), nil -} - -func defaultPluginInstanceConfigNode() *yaml.Node { - return &yaml.Node{ - Kind: yaml.MappingNode, - Tag: "!!map", - Content: []*yaml.Node{}, - } -} - -// ClaudeHeaderDefaults configures default header values injected into Claude API requests. -// In legacy mode, UserAgent/PackageVersion/RuntimeVersion/Timeout act as fallbacks when -// the client omits them, while OS/Arch remain runtime-derived. When stabilized device -// profiles are enabled, OS/Arch become the pinned platform baseline, while -// UserAgent/PackageVersion/RuntimeVersion seed the upgradeable software fingerprint. -type ClaudeHeaderDefaults struct { - UserAgent string `yaml:"user-agent" json:"user-agent"` - PackageVersion string `yaml:"package-version" json:"package-version"` - RuntimeVersion string `yaml:"runtime-version" json:"runtime-version"` - OS string `yaml:"os" json:"os"` - Arch string `yaml:"arch" json:"arch"` - Timeout string `yaml:"timeout" json:"timeout"` - StabilizeDeviceProfile *bool `yaml:"stabilize-device-profile,omitempty" json:"stabilize-device-profile,omitempty"` -} - -// CodexHeaderDefaults configures fallback header values injected into Codex -// model requests for OAuth/file-backed auth when the client omits them. -// UserAgent applies to HTTP and websocket requests; BetaFeatures only applies to websockets. -type CodexHeaderDefaults struct { - UserAgent string `yaml:"user-agent" json:"user-agent"` - BetaFeatures string `yaml:"beta-features" json:"beta-features"` -} - -// CodexConfig configures provider-wide Codex request behavior. -type CodexConfig struct { - IdentityConfuse bool `yaml:"identity-confuse" json:"identity-confuse"` - // OptimizeMultiAgentV2 optimizes official Codex multi-agent requests. - OptimizeMultiAgentV2 bool `yaml:"optimize-multi-agent-v2" json:"optimize-multi-agent-v2"` - // LiveMediaRelay terminates and relays Codex Live WebRTC media in this process. - LiveMediaRelay CodexLiveMediaRelayConfig `yaml:"live-media-relay" json:"live-media-relay"` -} - -// CodexLiveMediaRelayConfig configures the in-process Codex Live WebRTC gateway. -type CodexLiveMediaRelayConfig struct { - Enabled bool `yaml:"enabled" json:"enabled"` - MaxSessions int `yaml:"max-sessions" json:"max-sessions"` - DisablePrivateRemoteIPs bool `yaml:"disable-private-remote-ips" json:"disable-private-remote-ips"` - PublicIP string `yaml:"public-ip" json:"public-ip"` - UDPPortMin uint16 `yaml:"udp-port-min" json:"udp-port-min"` - UDPPortMax uint16 `yaml:"udp-port-max" json:"udp-port-max"` - ICEServers []CodexLiveICEServer `yaml:"ice-servers" json:"ice-servers"` -} - -// CodexLiveICEServer configures a STUN or TURN server for the media relay. -type CodexLiveICEServer struct { - URLs []string `yaml:"urls" json:"urls"` - Username string `yaml:"username" json:"-"` - Credential string `yaml:"credential" json:"-"` -} - -// TLSConfig holds HTTPS server settings. -type TLSConfig struct { - // Enable toggles HTTPS server mode. - Enable bool `yaml:"enable" json:"enable"` - // Cert is the path to the TLS certificate file. - Cert string `yaml:"cert" json:"cert"` - // Key is the path to the TLS private key file. - Key string `yaml:"key" json:"key"` -} - -// PprofConfig holds pprof HTTP server settings. -type PprofConfig struct { - // Enable toggles the pprof HTTP debug server. - Enable bool `yaml:"enable" json:"enable"` - // Addr is the host:port address for the pprof HTTP server. - Addr string `yaml:"addr" json:"addr"` -} - -// RemoteManagement holds management API configuration under 'remote-management'. -type RemoteManagement struct { - // AllowRemote toggles remote (non-localhost) access to management API. - AllowRemote bool `yaml:"allow-remote"` - // SecretKey is the management key (plaintext or bcrypt hashed). YAML key intentionally 'secret-key'. - SecretKey string `yaml:"secret-key"` - // DisableControlPanel skips serving and syncing the bundled management UI when true. - DisableControlPanel bool `yaml:"disable-control-panel"` - // DisableAutoUpdatePanel disables automatic periodic background updates of the management panel asset from GitHub. - // When false (the default), the background updater remains enabled; when true, the panel is only downloaded on first access if missing. - DisableAutoUpdatePanel bool `yaml:"disable-auto-update-panel"` - // PanelGitHubRepository overrides the GitHub repository used to fetch the management panel asset. - // Accepts either a repository URL (https://github.com/org/repo) or an API releases endpoint. - PanelGitHubRepository string `yaml:"panel-github-repository"` -} - -// QuotaExceeded defines the behavior when API quota limits are exceeded. -// It provides configuration options for automatic failover mechanisms. -type QuotaExceeded struct { - // SwitchProject indicates whether to automatically switch to another project when a quota is exceeded. - SwitchProject bool `yaml:"switch-project" json:"switch-project"` - - // SwitchPreviewModel indicates whether to automatically switch to a preview model when a quota is exceeded. - SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"` - - // AntigravityCredits enables credits-based last-resort fallback for Claude models. - // When all free-tier auths are exhausted (429/503), the conductor retries with - // an auth that has available Google One AI credits. - AntigravityCredits bool `yaml:"antigravity-credits" json:"antigravity-credits"` -} - -// RoutingConfig configures how credentials are selected for requests. -type RoutingConfig struct { - // Strategy selects the credential selection strategy. - // Supported values: "round-robin" (default), "fill-first". - Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"` - - // SessionAffinity enables universal session-sticky routing for all clients. - // Session IDs are extracted from multiple sources: - // metadata.user_id (Claude Code session format), X-Session-ID, Session_id (Codex), - // X-Client-Request-Id (PI), metadata.user_id, conversation_id, or message hash. - // Automatic failover is always enabled when bound auth becomes unavailable. - SessionAffinity bool `yaml:"session-affinity,omitempty" json:"session-affinity,omitempty"` - - // SessionAffinityTTL specifies how long session-to-auth bindings are retained. - // Default: 1h. Accepts duration strings like "30m", "1h", "2h30m". - SessionAffinityTTL string `yaml:"session-affinity-ttl,omitempty" json:"session-affinity-ttl,omitempty"` -} - -// OAuthModelAlias defines a model ID alias for a specific channel. -// It maps the upstream model name (Name) to the client-visible alias (Alias). -// When Fork is true, the alias is added as an additional model in listings while -// keeping the original model ID available. -type OAuthModelAlias struct { - Name string `yaml:"name" json:"name"` - Alias string `yaml:"alias" json:"alias"` - Fork bool `yaml:"fork,omitempty" json:"fork,omitempty"` - - // DisplayName is the optional human-readable name shown in model catalogs. - DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` - - ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` -} - -// PayloadConfig defines default and override parameter rules applied to provider payloads. -type PayloadConfig struct { - // Default defines rules that only set parameters when they are missing in the payload. - Default []PayloadRule `yaml:"default" json:"default"` - // DefaultRaw defines rules that set raw JSON values only when they are missing. - DefaultRaw []PayloadRule `yaml:"default-raw" json:"default-raw"` - // Override defines rules that always set parameters, overwriting any existing values. - Override []PayloadRule `yaml:"override" json:"override"` - // OverrideRaw defines rules that always set raw JSON values, overwriting any existing values. - OverrideRaw []PayloadRule `yaml:"override-raw" json:"override-raw"` - // Filter defines rules that remove parameters from the payload by JSON path. - Filter []PayloadFilterRule `yaml:"filter" json:"filter"` -} - -// PayloadFilterRule describes a rule to remove specific JSON paths from matching model payloads. -type PayloadFilterRule struct { - // Models lists model entries with name pattern and protocol constraint. - Models []PayloadModelRule `yaml:"models" json:"models"` - // Params lists JSON paths (gjson/sjson syntax) to remove from the payload. - Params []string `yaml:"params" json:"params"` -} - -// PayloadRule describes a single rule targeting a list of models with parameter updates. -type PayloadRule struct { - // Models lists model entries with name pattern and protocol constraint. - Models []PayloadModelRule `yaml:"models" json:"models"` - // Params maps JSON paths (gjson/sjson syntax) to values written into the payload. - // For *-raw rules, values are treated as raw JSON fragments (strings are used as-is). - Params map[string]any `yaml:"params" json:"params"` -} - -// PayloadModelRule ties a model name pattern to a specific translator protocol. -type PayloadModelRule struct { - // Name is the model name or wildcard pattern (e.g., "gpt-*", "*-5", "gemini-*-pro"). - Name string `yaml:"name" json:"name"` - // Protocol restricts the rule to a specific translator format (e.g., "gemini", "responses"). - Protocol string `yaml:"protocol" json:"protocol"` - // Headers restricts the rule to requests whose headers match all configured wildcard patterns. - Headers map[string]string `yaml:"headers" json:"headers"` - // FromProtocol restricts the rule to a specific source protocol (e.g., "gemini", "responses"). - FromProtocol string `yaml:"from-protocol" json:"from-protocol"` - // Match requires payload JSON paths to equal the configured values. - Match []map[string]any `yaml:"match" json:"match"` - // NotMatch requires payload JSON paths to not equal the configured values. - NotMatch []map[string]any `yaml:"not-match" json:"not-match"` - // Exist requires payload JSON paths to exist and not be null. - Exist []string `yaml:"exist" json:"exist"` - // NotExist requires payload JSON paths to be missing or null. - NotExist []string `yaml:"not-exist" json:"not-exist"` -} - -// CloakConfig configures request cloaking for non-Claude-Code clients. -// Cloaking disguises API requests to appear as originating from the official Claude Code CLI. -type CloakConfig struct { - // Mode controls cloaking behavior: "auto" (default), "always", or "never". - // - "auto": cloak only when client is not Claude Code (based on User-Agent) - // - "always": always apply cloaking regardless of client - // - "never": never apply cloaking - Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` - - // StrictMode controls how system prompts are handled when cloaking. - // - false (default): prepend Claude Code prompt to user system messages - // - true: strip all user system messages, keep only Claude Code prompt - StrictMode bool `yaml:"strict-mode,omitempty" json:"strict-mode,omitempty"` - - // SensitiveWords is a list of words to obfuscate with zero-width characters. - // This can help bypass certain content filters. - SensitiveWords []string `yaml:"sensitive-words,omitempty" json:"sensitive-words,omitempty"` - - // CacheUserID controls whether Claude user_id values are cached per API key. - // When false, a fresh random user_id is generated for every request. - CacheUserID *bool `yaml:"cache-user-id,omitempty" json:"cache-user-id,omitempty"` -} - -// ClaudeKey represents the configuration for a Claude API key, -// including the API key itself and an optional base URL for the API endpoint. -type ClaudeKey struct { - // APIKey is the authentication key for accessing Claude API services. - APIKey string `yaml:"api-key" json:"api-key"` - - // Priority controls selection preference when multiple credentials match. - // Higher values are preferred; defaults to 0. - Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` - - // Prefix optionally namespaces models for this credential (e.g., "teamA/claude-sonnet-4"). - Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` - - // BaseURL is the base URL for the Claude API endpoint. - // If empty, the default Claude API URL will be used. - BaseURL string `yaml:"base-url" json:"base-url"` - - // ProxyURL overrides the global proxy setting for this API key if provided. - ProxyURL string `yaml:"proxy-url" json:"proxy-url"` - - // Models defines upstream model names and aliases for request routing. - Models []ClaudeModel `yaml:"models" json:"models"` - - // Headers optionally adds extra HTTP headers for requests sent with this key. - Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` - - // ExcludedModels lists model IDs that should be excluded for this provider. - ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` - - // RebuildMidSystemMessage moves Claude messages with role "system" into the top-level system field. - RebuildMidSystemMessage bool `yaml:"rebuild-mid-system-message,omitempty" json:"rebuild-mid-system-message,omitempty"` - - // DisableCooling disables auth/model cooldown scheduling for this credential when true. - DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` - - // Cloak configures request cloaking for non-Claude-Code clients. - Cloak *CloakConfig `yaml:"cloak,omitempty" json:"cloak,omitempty"` - - // ExperimentalCCHSigning enables opt-in final-body cch signing for cloaked - // Claude /v1/messages requests. It is disabled by default so upstream seed - // changes do not alter the proxy's legacy behavior. - ExperimentalCCHSigning bool `yaml:"experimental-cch-signing,omitempty" json:"experimental-cch-signing,omitempty"` -} - -func (k ClaudeKey) GetAPIKey() string { return k.APIKey } -func (k ClaudeKey) GetBaseURL() string { return k.BaseURL } - -// ClaudeModel describes a mapping between an alias and the actual upstream model name. -type ClaudeModel struct { - // Name is the upstream model identifier used when issuing requests. - Name string `yaml:"name" json:"name"` - - // Alias is the client-facing model name that maps to Name. - Alias string `yaml:"alias" json:"alias"` - - // DisplayName is the optional human-readable name shown in model catalogs. - DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` - - // ForceMapping rewrites upstream response model fields back to Alias. - ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` -} - -func (m ClaudeModel) GetName() string { return m.Name } -func (m ClaudeModel) GetAlias() string { return m.Alias } -func (m ClaudeModel) GetDisplayName() string { return m.DisplayName } -func (m ClaudeModel) GetForceMapping() bool { return m.ForceMapping } - -// CodexKey represents the configuration for a Codex API key, -// including the API key itself and an optional base URL for the API endpoint. -type CodexKey struct { - // APIKey is the authentication key for accessing Codex API services. - APIKey string `yaml:"api-key" json:"api-key"` - - // Priority controls selection preference when multiple credentials match. - // Higher values are preferred; defaults to 0. - Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` - - // Prefix optionally namespaces models for this credential (e.g., "teamA/gpt-5-codex"). - Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` - - // BaseURL is the base URL for the Codex API endpoint. - // If empty, the default Codex API URL will be used. - BaseURL string `yaml:"base-url" json:"base-url"` - - // Websockets enables the Responses API websocket transport for this credential. - Websockets bool `yaml:"websockets,omitempty" json:"websockets,omitempty"` - - // ProxyURL overrides the global proxy setting for this API key if provided. - ProxyURL string `yaml:"proxy-url" json:"proxy-url"` - - // Models defines upstream model names and aliases for request routing. - Models []CodexModel `yaml:"models" json:"models"` - - // Headers optionally adds extra HTTP headers for requests sent with this key. - Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` - - // ExcludedModels lists model IDs that should be excluded for this provider. - ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` - - // DisableCooling disables auth/model cooldown scheduling for this credential when true. - DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` -} - -func (k CodexKey) GetAPIKey() string { return k.APIKey } -func (k CodexKey) GetBaseURL() string { return k.BaseURL } - -// CodexModel describes a mapping between an alias and the actual upstream model name. -type CodexModel struct { - // Name is the upstream model identifier used when issuing requests. - Name string `yaml:"name" json:"name"` - - // Alias is the client-facing model name that maps to Name. - Alias string `yaml:"alias" json:"alias"` - - // DisplayName is the optional human-readable name shown in model catalogs. - DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` - - // ForceMapping rewrites upstream response model fields back to Alias. - ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` -} - -func (m CodexModel) GetName() string { return m.Name } -func (m CodexModel) GetAlias() string { return m.Alias } -func (m CodexModel) GetDisplayName() string { return m.DisplayName } -func (m CodexModel) GetForceMapping() bool { return m.ForceMapping } - -// XAIKey uses the Codex API key structure for native xAI execution. -type XAIKey = CodexKey - -// XAIModel uses the Codex model mapping structure for xAI models. -type XAIModel = CodexModel - -// GeminiKey represents the configuration for a Gemini API key, -// including optional overrides for upstream base URL, proxy routing, and headers. -type GeminiKey struct { - // APIKey is the authentication key for accessing Gemini API services. - APIKey string `yaml:"api-key" json:"api-key"` - - // Priority controls selection preference when multiple credentials match. - // Higher values are preferred; defaults to 0. - Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` - - // Prefix optionally namespaces models for this credential (e.g., "teamA/gemini-3-pro-preview"). - Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` - - // BaseURL optionally overrides the Gemini API endpoint. - BaseURL string `yaml:"base-url,omitempty" json:"base-url,omitempty"` - - // ProxyURL optionally overrides the global proxy for this API key. - ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"` - - // Models defines upstream model names and aliases for request routing. - Models []GeminiModel `yaml:"models,omitempty" json:"models,omitempty"` - - // Headers optionally adds extra HTTP headers for requests sent with this key. - Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` - - // ExcludedModels lists model IDs that should be excluded for this provider. - ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` - - // DisableCooling disables auth/model cooldown scheduling for this credential when true. - DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` -} - -func (k GeminiKey) GetAPIKey() string { return k.APIKey } -func (k GeminiKey) GetBaseURL() string { return k.BaseURL } - -// GeminiModel describes a mapping between an alias and the actual upstream model name. -type GeminiModel struct { - // Name is the upstream model identifier used when issuing requests. - Name string `yaml:"name" json:"name"` - - // Alias is the client-facing model name that maps to Name. - Alias string `yaml:"alias" json:"alias"` - - // DisplayName is the optional human-readable name shown in model catalogs. - DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` - - // ForceMapping rewrites upstream response model fields back to Alias. - ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` -} - -func (m GeminiModel) GetName() string { return m.Name } -func (m GeminiModel) GetAlias() string { return m.Alias } -func (m GeminiModel) GetDisplayName() string { return m.DisplayName } -func (m GeminiModel) GetForceMapping() bool { return m.ForceMapping } - -// OpenAICompatibility represents the configuration for OpenAI API compatibility -// with external providers, allowing model aliases to be routed through OpenAI API format. -type OpenAICompatibility struct { - // Name is the identifier for this OpenAI compatibility configuration. - Name string `yaml:"name" json:"name"` - - // Priority controls selection preference when multiple providers or credentials match. - // Higher values are preferred; defaults to 0. - Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` - - // Disabled prevents this provider from being used for routing. - Disabled bool `yaml:"disabled,omitempty" json:"disabled,omitempty"` - - // Prefix optionally namespaces model aliases for this provider (e.g., "teamA/kimi-k2"). - Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` - - // BaseURL is the base URL for the external OpenAI-compatible API endpoint. - BaseURL string `yaml:"base-url" json:"base-url"` - - // APIKeyEntries defines API keys with optional per-key proxy configuration. - APIKeyEntries []OpenAICompatibilityAPIKey `yaml:"api-key-entries,omitempty" json:"api-key-entries,omitempty"` - - // Models defines the model configurations including aliases for routing. - Models []OpenAICompatibilityModel `yaml:"models" json:"models"` - - // Headers optionally adds extra HTTP headers for requests sent to this provider. - Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` - - // DisableCooling disables auth/model cooldown scheduling for this provider when true. - DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` -} - -// OpenAICompatibilityAPIKey represents an API key configuration with optional proxy setting. -type OpenAICompatibilityAPIKey struct { - // APIKey is the authentication key for accessing the external API services. - APIKey string `yaml:"api-key" json:"api-key"` - - // ProxyURL overrides the global proxy setting for this API key if provided. - ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"` -} - -// OpenAICompatibilityModel represents a model configuration for OpenAI compatibility, -// including the actual model name and its alias for API routing. -type OpenAICompatibilityModel struct { - // Name is the actual model name used by the external provider. - Name string `yaml:"name" json:"name"` - - // Alias is the model name alias that clients will use to reference this model. - Alias string `yaml:"alias" json:"alias"` - - // DisplayName is the optional human-readable name shown in model catalogs. - DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` - - // ForceMapping rewrites upstream response model fields back to Alias. - ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` - - // Image marks this model as callable through /v1/images/generations and /v1/images/edits. - Image bool `yaml:"image,omitempty" json:"image,omitempty"` - - // InputModalities declares chat/responses input capabilities (e.g. text, image) for Codex and other clients. - // This is separate from Image, which only enables /v1/images/* endpoints. - InputModalities []string `yaml:"input-modalities,omitempty" json:"input-modalities,omitempty"` - - // OutputModalities declares supported output modalities when known (e.g. text, image). - OutputModalities []string `yaml:"output-modalities,omitempty" json:"output-modalities,omitempty"` - - // Thinking configures the thinking/reasoning capability for this model. - // If nil, the model defaults to level-based reasoning with levels ["low", "medium", "high"]. - Thinking *registry.ThinkingSupport `yaml:"thinking,omitempty" json:"thinking,omitempty"` -} - -func (m OpenAICompatibilityModel) GetName() string { return m.Name } -func (m OpenAICompatibilityModel) GetAlias() string { return m.Alias } -func (m OpenAICompatibilityModel) GetDisplayName() string { return m.DisplayName } -func (m OpenAICompatibilityModel) GetForceMapping() bool { return m.ForceMapping } - -// LoadConfig reads a YAML configuration file from the given path, -// unmarshals it into a Config struct, applies environment variable overrides, -// and returns it. -// -// Parameters: -// - configFile: The path to the YAML configuration file -// -// Returns: -// - *Config: The loaded configuration -// - error: An error if the configuration could not be loaded -func LoadConfig(configFile string) (*Config, error) { - return LoadConfigOptional(configFile, false) -} - -// LoadConfigOptional reads YAML from configFile. -// If optional is true and the file is missing, it returns an empty Config. -// If optional is true and the file is empty or invalid, it returns an empty Config. -func LoadConfigOptional(configFile string, optional bool) (*Config, error) { - // Read the entire configuration file into memory. - data, err := os.ReadFile(configFile) - if err != nil { - if optional { - if os.IsNotExist(err) || errors.Is(err, syscall.EISDIR) { - // Missing and optional: return empty config (cloud deploy standby). - cfg := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} - cfg.NormalizePluginsConfig() - return cfg, nil - } - } - return nil, fmt.Errorf("failed to read config file: %w", err) - } - - // In cloud deploy mode (optional=true), if file is empty or contains only whitespace, return empty config. - if optional && len(bytes.TrimSpace(data)) == 0 { - cfg := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} - cfg.NormalizePluginsConfig() - return cfg, nil - } - - // Unmarshal the YAML data into the Config struct. - var cfg Config - // Set defaults before unmarshal so that absent keys keep defaults. - cfg.Host = "" // Default empty: binds to all interfaces (IPv4 + IPv6) - cfg.LoggingToFile = false - cfg.LogsMaxTotalSizeMB = 0 - cfg.ErrorLogsMaxFiles = 10 - cfg.UsageStatisticsEnabled = false - cfg.RedisUsageQueueRetentionSeconds = 60 - cfg.DisableCooling = false - cfg.SaveCooldownStatus = false - cfg.TransientErrorCooldownSeconds = 0 - cfg.DisableImageGeneration = DisableImageGenerationOff - cfg.WebsocketAuth = true - cfg.Pprof.Enable = false - cfg.Pprof.Addr = DefaultPprofAddr - cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository - cfg.CredentialInFlight = DefaultCredentialInFlightConfig() - if err = yaml.Unmarshal(data, &cfg); err != nil { - if optional { - // In cloud deploy mode, if YAML parsing fails, return empty config instead of error. - cfgOptional := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} - cfgOptional.NormalizePluginsConfig() - return cfgOptional, nil - } - return nil, fmt.Errorf("failed to parse config file: %w", err) - } - - cfg.CredentialConcurrency = cfg.CredentialConcurrency.WithDefaults() - if errValidate := cfg.CredentialInFlight.Validate(); errValidate != nil { - return nil, errValidate - } - if errValidate := cfg.Codex.LiveMediaRelay.Validate(); errValidate != nil { - return nil, errValidate - } - - // Hash remote management key if plaintext is detected (nested) - // We consider a value to be already hashed if it looks like a bcrypt hash ($2a$, $2b$, or $2y$ prefix). - if cfg.RemoteManagement.SecretKey != "" && !looksLikeBcrypt(cfg.RemoteManagement.SecretKey) { - hashed, errHash := hashSecret(cfg.RemoteManagement.SecretKey) - if errHash != nil { - return nil, fmt.Errorf("failed to hash remote management key: %w", errHash) - } - cfg.RemoteManagement.SecretKey = hashed - - // Persist the hashed value back to the config file to avoid re-hashing on next startup. - // Preserve YAML comments and ordering; update only the nested key. - _ = SaveConfigPreserveCommentsUpdateNestedScalar(configFile, []string{"remote-management", "secret-key"}, hashed) - } - - cfg.RemoteManagement.PanelGitHubRepository = strings.TrimSpace(cfg.RemoteManagement.PanelGitHubRepository) - if cfg.RemoteManagement.PanelGitHubRepository == "" { - cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository - } - - cfg.Pprof.Addr = strings.TrimSpace(cfg.Pprof.Addr) - if cfg.Pprof.Addr == "" { - cfg.Pprof.Addr = DefaultPprofAddr - } - - if cfg.LogsMaxTotalSizeMB < 0 { - cfg.LogsMaxTotalSizeMB = 0 - } - - if cfg.ErrorLogsMaxFiles < 0 { - cfg.ErrorLogsMaxFiles = 10 - } - - if cfg.RedisUsageQueueRetentionSeconds <= 0 { - cfg.RedisUsageQueueRetentionSeconds = 60 - } else if cfg.RedisUsageQueueRetentionSeconds > 3600 { - log.WithField("value", cfg.RedisUsageQueueRetentionSeconds).Warn("redis-usage-queue-retention-seconds too large; clamping to 3600") - cfg.RedisUsageQueueRetentionSeconds = 3600 - } - - if cfg.MaxRetryCredentials < 0 { - cfg.MaxRetryCredentials = 0 - } - - cfg.NormalizePluginsConfig() - if errResolvePluginsDir := cfg.ResolvePluginsDir(); errResolvePluginsDir != nil && cfg.Plugins.Enabled { - return nil, errResolvePluginsDir - } - - // Sanitize Gemini API key configuration and migrate legacy entries. - cfg.SanitizeGeminiKeys() - - // Sanitize native Interactions API key configuration. - cfg.SanitizeInteractionsKeys() - - // Sanitize Vertex-compatible API keys. - cfg.SanitizeVertexCompatKeys() - - // Sanitize Codex keys: drop entries without base-url - cfg.SanitizeCodexKeys() - - // Sanitize xAI keys: drop entries without base-url - cfg.SanitizeXAIKeys() - - // Sanitize Codex header defaults. - cfg.SanitizeCodexHeaderDefaults() - - // Sanitize Claude header defaults. - cfg.SanitizeClaudeHeaderDefaults() - - // Sanitize Claude key headers - cfg.SanitizeClaudeKeys() - - // Sanitize OpenAI compatibility providers: drop entries without base-url - cfg.SanitizeOpenAICompatibility() - - // Normalize OAuth provider model exclusion map. - cfg.OAuthExcludedModels = NormalizeOAuthExcludedModels(cfg.OAuthExcludedModels) - - // Normalize global OAuth model name aliases. - cfg.SanitizeOAuthModelAlias() - - // Validate raw payload rules and drop invalid entries. - cfg.SanitizePayloadRules() - - // Return the populated configuration struct. - return &cfg, nil -} - -// NormalizePluginsConfig applies default plugin configuration values. -func (cfg *Config) NormalizePluginsConfig() { - if cfg == nil { - return - } - cfg.Plugins.Dir = strings.TrimSpace(cfg.Plugins.Dir) - if cfg.Plugins.Dir == "" { - cfg.Plugins.Dir = defaultPluginsDir - } - if len(cfg.Plugins.StoreSources) > 0 { - sources := make([]string, 0, len(cfg.Plugins.StoreSources)) - for _, source := range cfg.Plugins.StoreSources { - source = strings.TrimSpace(source) - if source == "" { - continue - } - sources = append(sources, source) - } - cfg.Plugins.StoreSources = sources - } - cfg.Plugins.StoreAuth = sdkpluginstore.NormalizeAuthConfigs(cfg.Plugins.StoreAuth) - if cfg.Plugins.Configs == nil { - cfg.Plugins.Configs = map[string]PluginInstanceConfig{} - } -} - -// SanitizePayloadRules validates raw JSON payload rule params and drops invalid rules. -func (cfg *Config) SanitizePayloadRules() { - if cfg == nil { - return - } - cfg.Payload.DefaultRaw = sanitizePayloadRawRules(cfg.Payload.DefaultRaw, "default-raw") - cfg.Payload.OverrideRaw = sanitizePayloadRawRules(cfg.Payload.OverrideRaw, "override-raw") -} - -func sanitizePayloadRawRules(rules []PayloadRule, section string) []PayloadRule { - if len(rules) == 0 { - return rules - } - out := make([]PayloadRule, 0, len(rules)) - for i := range rules { - rule := rules[i] - if len(rule.Params) == 0 { - continue - } - invalid := false - for path, value := range rule.Params { - raw, ok := payloadRawString(value) - if !ok { - continue - } - trimmed := bytes.TrimSpace(raw) - if len(trimmed) == 0 || !json.Valid(trimmed) { - log.WithFields(log.Fields{ - "section": section, - "rule_index": i + 1, - "param": path, - }).Warn("payload rule dropped: invalid raw JSON") - invalid = true - break - } - } - if invalid { - continue - } - out = append(out, rule) - } - return out -} - -func payloadRawString(value any) ([]byte, bool) { - switch typed := value.(type) { - case string: - return []byte(typed), true - case []byte: - return typed, true - default: - return nil, false - } -} - -// SanitizeCodexHeaderDefaults trims surrounding whitespace from the -// configured Codex header fallback values. -func (cfg *Config) SanitizeCodexHeaderDefaults() { - if cfg == nil { - return - } - cfg.CodexHeaderDefaults.UserAgent = strings.TrimSpace(cfg.CodexHeaderDefaults.UserAgent) - cfg.CodexHeaderDefaults.BetaFeatures = strings.TrimSpace(cfg.CodexHeaderDefaults.BetaFeatures) -} - -// SanitizeClaudeHeaderDefaults trims surrounding whitespace from the -// configured Claude fingerprint baseline values. -func (cfg *Config) SanitizeClaudeHeaderDefaults() { - if cfg == nil { - return - } - cfg.ClaudeHeaderDefaults.UserAgent = strings.TrimSpace(cfg.ClaudeHeaderDefaults.UserAgent) - cfg.ClaudeHeaderDefaults.PackageVersion = strings.TrimSpace(cfg.ClaudeHeaderDefaults.PackageVersion) - cfg.ClaudeHeaderDefaults.RuntimeVersion = strings.TrimSpace(cfg.ClaudeHeaderDefaults.RuntimeVersion) - cfg.ClaudeHeaderDefaults.OS = strings.TrimSpace(cfg.ClaudeHeaderDefaults.OS) - cfg.ClaudeHeaderDefaults.Arch = strings.TrimSpace(cfg.ClaudeHeaderDefaults.Arch) - cfg.ClaudeHeaderDefaults.Timeout = strings.TrimSpace(cfg.ClaudeHeaderDefaults.Timeout) -} - -// SanitizeOAuthModelAlias normalizes and deduplicates global OAuth model name aliases. -// It trims whitespace, normalizes channel keys to lower-case, drops empty entries, -// allows multiple aliases per upstream name, and ensures aliases are unique within each channel. -func (cfg *Config) SanitizeOAuthModelAlias() { - if cfg == nil || len(cfg.OAuthModelAlias) == 0 { - return - } - out := make(map[string][]OAuthModelAlias, len(cfg.OAuthModelAlias)) - for rawChannel, aliases := range cfg.OAuthModelAlias { - channel := strings.ToLower(strings.TrimSpace(rawChannel)) - if channel == "" || len(aliases) == 0 { - continue - } - seenAlias := make(map[string]struct{}, len(aliases)) - clean := make([]OAuthModelAlias, 0, len(aliases)) - for _, entry := range aliases { - name := strings.TrimSpace(entry.Name) - alias := strings.TrimSpace(entry.Alias) - if name == "" || alias == "" { - continue - } - if strings.EqualFold(name, alias) { - continue - } - aliasKey := strings.ToLower(alias) - if _, ok := seenAlias[aliasKey]; ok { - continue - } - seenAlias[aliasKey] = struct{}{} - clean = append(clean, OAuthModelAlias{ - Name: name, - Alias: alias, - Fork: entry.Fork, - DisplayName: strings.TrimSpace(entry.DisplayName), - ForceMapping: entry.ForceMapping, - }) - } - if len(clean) > 0 { - out[channel] = clean - } - } - cfg.OAuthModelAlias = out -} - -// SanitizeOpenAICompatibility removes OpenAI-compatibility provider entries that are -// not actionable, specifically those missing a BaseURL. It trims whitespace before -// evaluation and preserves the relative order of remaining entries. -func (cfg *Config) SanitizeOpenAICompatibility() { - if cfg == nil || len(cfg.OpenAICompatibility) == 0 { - return - } - out := make([]OpenAICompatibility, 0, len(cfg.OpenAICompatibility)) - for i := range cfg.OpenAICompatibility { - e := cfg.OpenAICompatibility[i] - e.Name = strings.TrimSpace(e.Name) - e.Prefix = normalizeModelPrefix(e.Prefix) - e.BaseURL = strings.TrimSpace(e.BaseURL) - e.Headers = NormalizeHeaders(e.Headers) - if e.BaseURL == "" { - // Skip providers with no base-url; treated as removed - continue - } - out = append(out, e) - } - cfg.OpenAICompatibility = out -} - -// SanitizeCodexKeys removes Codex API key entries missing a BaseURL. -// It trims whitespace and preserves order for remaining entries. -func (cfg *Config) SanitizeCodexKeys() { - if cfg == nil { - return - } - cfg.CodexKey = sanitizeCodexKeyEntries(cfg.CodexKey) -} - -// SanitizeXAIKeys removes xAI API key entries missing a BaseURL. -// It applies the same normalization rules as codex-api-key. -func (cfg *Config) SanitizeXAIKeys() { - if cfg == nil { - return - } - cfg.XAIKey = sanitizeCodexKeyEntries(cfg.XAIKey) -} - -func sanitizeCodexKeyEntries(entries []CodexKey) []CodexKey { - if len(entries) == 0 { - return entries - } - out := make([]CodexKey, 0, len(entries)) - for i := range entries { - e := entries[i] - e.Prefix = normalizeModelPrefix(e.Prefix) - e.BaseURL = strings.TrimSpace(e.BaseURL) - e.Headers = NormalizeHeaders(e.Headers) - e.ExcludedModels = NormalizeExcludedModels(e.ExcludedModels) - if e.BaseURL == "" { - continue - } - out = append(out, e) - } - return out -} - -// SanitizeClaudeKeys normalizes headers for Claude credentials. -func (cfg *Config) SanitizeClaudeKeys() { - if cfg == nil || len(cfg.ClaudeKey) == 0 { - return - } - for i := range cfg.ClaudeKey { - entry := &cfg.ClaudeKey[i] - entry.Prefix = normalizeModelPrefix(entry.Prefix) - entry.Headers = NormalizeHeaders(entry.Headers) - entry.ExcludedModels = NormalizeExcludedModels(entry.ExcludedModels) - } -} - -func sanitizeGeminiKeyEntries(entries []GeminiKey) []GeminiKey { - seen := make(map[string]struct{}, len(entries)) - out := entries[:0] - for i := range entries { - entry := entries[i] - entry.APIKey = strings.TrimSpace(entry.APIKey) - if entry.APIKey == "" { - continue - } - entry.Prefix = normalizeModelPrefix(entry.Prefix) - entry.BaseURL = strings.TrimSpace(entry.BaseURL) - entry.ProxyURL = strings.TrimSpace(entry.ProxyURL) - entry.Headers = NormalizeHeaders(entry.Headers) - entry.ExcludedModels = NormalizeExcludedModels(entry.ExcludedModels) - uniqueKey := entry.APIKey + "|" + entry.BaseURL - if _, exists := seen[uniqueKey]; exists { - continue - } - seen[uniqueKey] = struct{}{} - out = append(out, entry) - } - return out -} - -// SanitizeGeminiKeys deduplicates and normalizes Gemini credentials. -// It uses API key + base URL as the uniqueness key. -func (cfg *Config) SanitizeGeminiKeys() { - if cfg == nil { - return - } - cfg.GeminiKey = sanitizeGeminiKeyEntries(cfg.GeminiKey) -} - -// SanitizeInteractionsKeys deduplicates and normalizes native Interactions credentials. -// It uses API key + base URL as the uniqueness key. -func (cfg *Config) SanitizeInteractionsKeys() { - if cfg == nil { - return - } - cfg.InteractionsKey = sanitizeGeminiKeyEntries(cfg.InteractionsKey) -} - -func normalizeModelPrefix(prefix string) string { - trimmed := strings.TrimSpace(prefix) - trimmed = strings.Trim(trimmed, "/") - if trimmed == "" { - return "" - } - if strings.Contains(trimmed, "/") { - return "" - } - return trimmed -} - -// looksLikeBcrypt returns true if the provided string appears to be a bcrypt hash. -func looksLikeBcrypt(s string) bool { - return len(s) > 4 && (s[:4] == "$2a$" || s[:4] == "$2b$" || s[:4] == "$2y$") -} - -// NormalizeHeaders trims header keys and values and removes empty pairs. -func NormalizeHeaders(headers map[string]string) map[string]string { - if len(headers) == 0 { - return nil - } - clean := make(map[string]string, len(headers)) - for k, v := range headers { - key := strings.TrimSpace(k) - val := strings.TrimSpace(v) - if key == "" || val == "" { - continue - } - clean[key] = val - } - if len(clean) == 0 { - return nil - } - return clean -} - -// NormalizeExcludedModels trims, lowercases, and deduplicates model exclusion patterns. -// It preserves the order of first occurrences and drops empty entries. -func NormalizeExcludedModels(models []string) []string { - if len(models) == 0 { - return nil - } - seen := make(map[string]struct{}, len(models)) - out := make([]string, 0, len(models)) - for _, raw := range models { - trimmed := strings.ToLower(strings.TrimSpace(raw)) - if trimmed == "" { - continue - } - if _, exists := seen[trimmed]; exists { - continue - } - seen[trimmed] = struct{}{} - out = append(out, trimmed) - } - if len(out) == 0 { - return nil - } - return out -} - -// NormalizeOAuthExcludedModels cleans provider -> excluded models mappings by normalizing provider keys -// and applying model exclusion normalization to each entry. -func NormalizeOAuthExcludedModels(entries map[string][]string) map[string][]string { - if len(entries) == 0 { - return nil - } - out := make(map[string][]string, len(entries)) - for provider, models := range entries { - key := strings.ToLower(strings.TrimSpace(provider)) - if key == "" { - continue - } - normalized := NormalizeExcludedModels(models) - if len(normalized) == 0 { - continue - } - out[key] = normalized - } - if len(out) == 0 { - return nil - } - return out -} - -// hashSecret hashes the given secret using bcrypt. -func hashSecret(secret string) (string, error) { - // Use default cost for simplicity. - hashedBytes, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost) - if err != nil { - return "", err - } - return string(hashedBytes), nil -} - -// SaveConfigPreserveComments writes the config back to YAML while preserving existing comments -// and key ordering by loading the original file into a yaml.Node tree and updating values in-place. -func SaveConfigPreserveComments(configFile string, cfg *Config) error { - persistCfg := cfg - // Load original YAML as a node tree to preserve comments and ordering. - data, err := os.ReadFile(configFile) - if err != nil { - return err - } - - var original yaml.Node - if err = yaml.Unmarshal(data, &original); err != nil { - return err - } - if original.Kind != yaml.DocumentNode || len(original.Content) == 0 { - return fmt.Errorf("invalid yaml document structure") - } - if original.Content[0] == nil || original.Content[0].Kind != yaml.MappingNode { - return fmt.Errorf("expected root mapping node") - } - - // Marshal the current cfg to YAML, then unmarshal to a yaml.Node we can merge from. - rendered, err := yaml.Marshal(persistCfg) - if err != nil { - return err - } - var generated yaml.Node - if err = yaml.Unmarshal(rendered, &generated); err != nil { - return err - } - if generated.Kind != yaml.DocumentNode || len(generated.Content) == 0 || generated.Content[0] == nil { - return fmt.Errorf("invalid generated yaml structure") - } - if generated.Content[0].Kind != yaml.MappingNode { - return fmt.Errorf("expected generated root mapping node") - } - - // Remove deprecated sections before merging back the sanitized config. - removeLegacyAuthBlock(original.Content[0]) - removeLegacyOpenAICompatAPIKeys(original.Content[0]) - removeRemovedIntegrationKeys(original.Content[0]) - removeLegacyGenerativeLanguageKeys(original.Content[0]) - - pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-excluded-models") - pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-model-alias") - pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "plugins", "configs") - - // Merge generated into original in-place, preserving comments/order of existing nodes. - mergeMappingPreserve(original.Content[0], generated.Content[0]) - normalizeCollectionNodeStyles(original.Content[0]) - - // Write back. - f, err := os.Create(configFile) - if err != nil { - return err - } - defer func() { _ = f.Close() }() - var buf bytes.Buffer - enc := yaml.NewEncoder(&buf) - enc.SetIndent(2) - if err = enc.Encode(&original); err != nil { - _ = enc.Close() - return err - } - if err = enc.Close(); err != nil { - return err - } - data = NormalizeCommentIndentation(buf.Bytes()) - _, err = f.Write(data) - return err -} - -// SaveConfigPreserveCommentsUpdateNestedScalar updates a nested scalar key path like ["a","b"] -// while preserving comments and positions. -func SaveConfigPreserveCommentsUpdateNestedScalar(configFile string, path []string, value string) error { - data, err := os.ReadFile(configFile) - if err != nil { - return err - } - var root yaml.Node - if err = yaml.Unmarshal(data, &root); err != nil { - return err - } - if root.Kind != yaml.DocumentNode || len(root.Content) == 0 { - return fmt.Errorf("invalid yaml document structure") - } - node := root.Content[0] - // descend mapping nodes following path - for i, key := range path { - if i == len(path)-1 { - // set final scalar - v := getOrCreateMapValue(node, key) - v.Kind = yaml.ScalarNode - v.Tag = "!!str" - v.Value = value - } else { - next := getOrCreateMapValue(node, key) - if next.Kind != yaml.MappingNode { - next.Kind = yaml.MappingNode - next.Tag = "!!map" - } - node = next - } - } - f, err := os.Create(configFile) - if err != nil { - return err - } - defer func() { _ = f.Close() }() - var buf bytes.Buffer - enc := yaml.NewEncoder(&buf) - enc.SetIndent(2) - if err = enc.Encode(&root); err != nil { - _ = enc.Close() - return err - } - if err = enc.Close(); err != nil { - return err - } - data = NormalizeCommentIndentation(buf.Bytes()) - _, err = f.Write(data) - return err -} - -// NormalizeCommentIndentation removes indentation from standalone YAML comment lines to keep them left aligned. -func NormalizeCommentIndentation(data []byte) []byte { - lines := bytes.Split(data, []byte("\n")) - changed := false - for i, line := range lines { - trimmed := bytes.TrimLeft(line, " \t") - if len(trimmed) == 0 || trimmed[0] != '#' { - continue - } - if len(trimmed) == len(line) { - continue - } - lines[i] = append([]byte(nil), trimmed...) - changed = true - } - if !changed { - return data - } - return bytes.Join(lines, []byte("\n")) -} - -// getOrCreateMapValue finds the value node for a given key in a mapping node. -// If not found, it appends a new key/value pair and returns the new value node. -func getOrCreateMapValue(mapNode *yaml.Node, key string) *yaml.Node { - if mapNode.Kind != yaml.MappingNode { - mapNode.Kind = yaml.MappingNode - mapNode.Tag = "!!map" - mapNode.Content = nil - } - for i := 0; i+1 < len(mapNode.Content); i += 2 { - k := mapNode.Content[i] - if k.Value == key { - return mapNode.Content[i+1] - } - } - // append new key/value - mapNode.Content = append(mapNode.Content, &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}) - val := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: ""} - mapNode.Content = append(mapNode.Content, val) - return val -} - -// mergeMappingPreserve merges keys from src into dst mapping node while preserving -// key order and comments of existing keys in dst. New keys are only added if their -// value is non-zero and not a known default to avoid polluting the config with defaults. -func mergeMappingPreserve(dst, src *yaml.Node, path ...[]string) { - var currentPath []string - if len(path) > 0 { - currentPath = path[0] - } - - if dst == nil || src == nil { - return - } - if dst.Kind != yaml.MappingNode || src.Kind != yaml.MappingNode { - // If kinds do not match, prefer replacing dst with src semantics in-place - // but keep dst node object to preserve any attached comments at the parent level. - copyNodeShallow(dst, src) - return - } - for i := 0; i+1 < len(src.Content); i += 2 { - sk := src.Content[i] - sv := src.Content[i+1] - idx := findMapKeyIndex(dst, sk.Value) - childPath := appendPath(currentPath, sk.Value) - if idx >= 0 { - // Merge into existing value node (always update, even to zero values) - dv := dst.Content[idx+1] - mergeNodePreserve(dv, sv, childPath) - } else { - // New key: only add if value is non-zero and not a known default - candidate := deepCopyNode(sv) - pruneKnownDefaultsInNewNode(childPath, candidate) - if isKnownDefaultValue(childPath, candidate) { - continue - } - dst.Content = append(dst.Content, deepCopyNode(sk), candidate) - } - } -} - -// mergeNodePreserve merges src into dst for scalars, mappings and sequences while -// reusing destination nodes to keep comments and anchors. For sequences, it updates -// in-place by index. -func mergeNodePreserve(dst, src *yaml.Node, path ...[]string) { - var currentPath []string - if len(path) > 0 { - currentPath = path[0] - } - - if dst == nil || src == nil { - return - } - switch src.Kind { - case yaml.MappingNode: - if dst.Kind != yaml.MappingNode { - copyNodeShallow(dst, src) - } - mergeMappingPreserve(dst, src, currentPath) - case yaml.SequenceNode: - // Preserve explicit null style if dst was null and src is empty sequence - if dst.Kind == yaml.ScalarNode && dst.Tag == "!!null" && len(src.Content) == 0 { - // Keep as null to preserve original style - return - } - if dst.Kind != yaml.SequenceNode { - dst.Kind = yaml.SequenceNode - dst.Tag = "!!seq" - dst.Content = nil - } - reorderSequenceForMerge(dst, src) - // Update elements in place - minContent := len(dst.Content) - if len(src.Content) < minContent { - minContent = len(src.Content) - } - for i := 0; i < minContent; i++ { - if dst.Content[i] == nil { - dst.Content[i] = deepCopyNode(src.Content[i]) - continue - } - mergeNodePreserve(dst.Content[i], src.Content[i], currentPath) - if dst.Content[i] != nil && src.Content[i] != nil && - dst.Content[i].Kind == yaml.MappingNode && src.Content[i].Kind == yaml.MappingNode { - pruneMissingMapKeys(dst.Content[i], src.Content[i]) - } - } - // Append any extra items from src - for i := len(dst.Content); i < len(src.Content); i++ { - dst.Content = append(dst.Content, deepCopyNode(src.Content[i])) - } - // Truncate if dst has extra items not in src - if len(src.Content) < len(dst.Content) { - dst.Content = dst.Content[:len(src.Content)] - } - case yaml.ScalarNode, yaml.AliasNode: - // For scalars, update Tag and Value but keep Style from dst to preserve quoting - dst.Kind = src.Kind - dst.Tag = src.Tag - dst.Value = src.Value - // Keep dst.Style as-is intentionally - case 0: - // Unknown/empty kind; do nothing - default: - // Fallback: replace shallowly - copyNodeShallow(dst, src) - } -} - -// findMapKeyIndex returns the index of key node in dst mapping (index of key, not value). -// Returns -1 when not found. -func findMapKeyIndex(mapNode *yaml.Node, key string) int { - if mapNode == nil || mapNode.Kind != yaml.MappingNode { - return -1 - } - for i := 0; i+1 < len(mapNode.Content); i += 2 { - if mapNode.Content[i] != nil && mapNode.Content[i].Value == key { - return i - } - } - return -1 -} - -// appendPath appends a key to the path, returning a new slice to avoid modifying the original. -func appendPath(path []string, key string) []string { - if len(path) == 0 { - return []string{key} - } - newPath := make([]string, len(path)+1) - copy(newPath, path) - newPath[len(path)] = key - return newPath -} - -// isKnownDefaultValue returns true if the given node at the specified path -// represents a known default value that should not be written to the config file. -// This prevents non-zero defaults from polluting the config. -func isKnownDefaultValue(path []string, node *yaml.Node) bool { - // First check if it's a zero value - if isZeroValueNode(node) { - return true - } - - // Match known non-zero defaults by exact dotted path. - if len(path) == 0 { - return false - } - - fullPath := strings.Join(path, ".") - - // Check string defaults - if node.Kind == yaml.ScalarNode && node.Tag == "!!str" { - switch fullPath { - case "pprof.addr": - return node.Value == DefaultPprofAddr - case "remote-management.panel-github-repository": - return node.Value == DefaultPanelGitHubRepository - case "plugins.dir": - return node.Value == "plugins" - case "routing.strategy": - return node.Value == "round-robin" - } - } - - // Check integer defaults - if node.Kind == yaml.ScalarNode && node.Tag == "!!int" { - switch fullPath { - case "error-logs-max-files": - return node.Value == "10" - } - } - - return false -} - -// pruneKnownDefaultsInNewNode removes default-valued descendants from a new node -// before it is appended into the destination YAML tree. -func pruneKnownDefaultsInNewNode(path []string, node *yaml.Node) { - if node == nil { - return - } - - switch node.Kind { - case yaml.MappingNode: - filtered := make([]*yaml.Node, 0, len(node.Content)) - for i := 0; i+1 < len(node.Content); i += 2 { - keyNode := node.Content[i] - valueNode := node.Content[i+1] - if keyNode == nil || valueNode == nil { - continue - } - - childPath := appendPath(path, keyNode.Value) - if isKnownDefaultValue(childPath, valueNode) { - continue - } - - pruneKnownDefaultsInNewNode(childPath, valueNode) - if (valueNode.Kind == yaml.MappingNode || valueNode.Kind == yaml.SequenceNode) && - len(valueNode.Content) == 0 { - continue - } - - filtered = append(filtered, keyNode, valueNode) - } - node.Content = filtered - case yaml.SequenceNode: - for _, child := range node.Content { - pruneKnownDefaultsInNewNode(path, child) - } - } -} - -// isZeroValueNode returns true if the YAML node represents a zero/default value -// that should not be written as a new key to preserve config cleanliness. -// For mappings and sequences, recursively checks if all children are zero values. -func isZeroValueNode(node *yaml.Node) bool { - if node == nil { - return true - } - switch node.Kind { - case yaml.ScalarNode: - switch node.Tag { - case "!!bool": - return node.Value == "false" - case "!!int", "!!float": - return node.Value == "0" || node.Value == "0.0" - case "!!str": - return node.Value == "" - case "!!null": - return true - } - case yaml.SequenceNode: - if len(node.Content) == 0 { - return true - } - // Check if all elements are zero values - for _, child := range node.Content { - if !isZeroValueNode(child) { - return false - } - } - return true - case yaml.MappingNode: - if len(node.Content) == 0 { - return true - } - // Check if all values are zero values (values are at odd indices) - for i := 1; i < len(node.Content); i += 2 { - if !isZeroValueNode(node.Content[i]) { - return false - } - } - return true - } - return false -} - -// deepCopyNode creates a deep copy of a yaml.Node graph. -func deepCopyNode(n *yaml.Node) *yaml.Node { - return deepCopyNodeSeen(n, map[*yaml.Node]*yaml.Node{}) -} - -func deepCopyNodeSeen(n *yaml.Node, seen map[*yaml.Node]*yaml.Node) *yaml.Node { - if n == nil { - return nil - } - if cp, ok := seen[n]; ok { - return cp - } - cp := *n - seen[n] = &cp - if n.Alias != nil { - cp.Alias = deepCopyNodeSeen(n.Alias, seen) - } - if len(n.Content) > 0 { - cp.Content = make([]*yaml.Node, len(n.Content)) - for i := range n.Content { - cp.Content[i] = deepCopyNodeSeen(n.Content[i], seen) - } - } - return &cp -} - -// copyNodeShallow copies type/tag/value and resets content to match src, but -// keeps the same destination node pointer to preserve parent relations/comments. -func copyNodeShallow(dst, src *yaml.Node) { - if dst == nil || src == nil { - return - } - dst.Kind = src.Kind - dst.Tag = src.Tag - dst.Value = src.Value - // Replace content with deep copy from src - if len(src.Content) > 0 { - dst.Content = make([]*yaml.Node, len(src.Content)) - for i := range src.Content { - dst.Content[i] = deepCopyNode(src.Content[i]) - } - } else { - dst.Content = nil - } -} - -func reorderSequenceForMerge(dst, src *yaml.Node) { - if dst == nil || src == nil { - return - } - if len(dst.Content) == 0 { - return - } - if len(src.Content) == 0 { - return - } - original := append([]*yaml.Node(nil), dst.Content...) - used := make([]bool, len(original)) - ordered := make([]*yaml.Node, len(src.Content)) - for i := range src.Content { - if idx := matchSequenceElement(original, used, src.Content[i]); idx >= 0 { - ordered[i] = original[idx] - used[idx] = true - } - } - dst.Content = ordered -} - -func matchSequenceElement(original []*yaml.Node, used []bool, target *yaml.Node) int { - if target == nil { - return -1 - } - switch target.Kind { - case yaml.MappingNode: - id := sequenceElementIdentity(target) - if id != "" { - for i := range original { - if used[i] || original[i] == nil || original[i].Kind != yaml.MappingNode { - continue - } - if sequenceElementIdentity(original[i]) == id { - return i - } - } - } - case yaml.ScalarNode: - val := strings.TrimSpace(target.Value) - if val != "" { - for i := range original { - if used[i] || original[i] == nil || original[i].Kind != yaml.ScalarNode { - continue - } - if strings.TrimSpace(original[i].Value) == val { - return i - } - } - } - default: - } - // Fallback to structural equality to preserve nodes lacking explicit identifiers. - for i := range original { - if used[i] || original[i] == nil { - continue - } - if nodesStructurallyEqual(original[i], target) { - return i - } - } - return -1 -} - -func sequenceElementIdentity(node *yaml.Node) string { - if node == nil || node.Kind != yaml.MappingNode { - return "" - } - identityKeys := []string{"id", "name", "alias", "api-key", "api_key", "apikey", "key", "provider", "model"} - for _, k := range identityKeys { - if v := mappingScalarValue(node, k); v != "" { - return k + "=" + v - } - } - for i := 0; i+1 < len(node.Content); i += 2 { - keyNode := node.Content[i] - valNode := node.Content[i+1] - if keyNode == nil || valNode == nil || valNode.Kind != yaml.ScalarNode { - continue - } - val := strings.TrimSpace(valNode.Value) - if val != "" { - return strings.ToLower(strings.TrimSpace(keyNode.Value)) + "=" + val - } - } - return "" -} - -func mappingScalarValue(node *yaml.Node, key string) string { - if node == nil || node.Kind != yaml.MappingNode { - return "" - } - lowerKey := strings.ToLower(key) - for i := 0; i+1 < len(node.Content); i += 2 { - keyNode := node.Content[i] - valNode := node.Content[i+1] - if keyNode == nil || valNode == nil || valNode.Kind != yaml.ScalarNode { - continue - } - if strings.ToLower(strings.TrimSpace(keyNode.Value)) == lowerKey { - return strings.TrimSpace(valNode.Value) - } - } - return "" -} - -func nodesStructurallyEqual(a, b *yaml.Node) bool { - if a == nil || b == nil { - return a == b - } - if a.Kind != b.Kind { - return false - } - switch a.Kind { - case yaml.MappingNode: - if len(a.Content) != len(b.Content) { - return false - } - for i := 0; i+1 < len(a.Content); i += 2 { - if !nodesStructurallyEqual(a.Content[i], b.Content[i]) { - return false - } - if !nodesStructurallyEqual(a.Content[i+1], b.Content[i+1]) { - return false - } - } - return true - case yaml.SequenceNode: - if len(a.Content) != len(b.Content) { - return false - } - for i := range a.Content { - if !nodesStructurallyEqual(a.Content[i], b.Content[i]) { - return false - } - } - return true - case yaml.ScalarNode: - return strings.TrimSpace(a.Value) == strings.TrimSpace(b.Value) - case yaml.AliasNode: - return nodesStructurallyEqual(a.Alias, b.Alias) - default: - return strings.TrimSpace(a.Value) == strings.TrimSpace(b.Value) - } -} - -func removeMapKey(mapNode *yaml.Node, key string) { - if mapNode == nil || mapNode.Kind != yaml.MappingNode || key == "" { - return - } - for i := 0; i+1 < len(mapNode.Content); i += 2 { - if mapNode.Content[i] != nil && mapNode.Content[i].Value == key { - mapNode.Content = append(mapNode.Content[:i], mapNode.Content[i+2:]...) - return - } - } -} - -func pruneMappingToGeneratedKeys(dstRoot, srcRoot *yaml.Node, keyPath ...string) { - if len(keyPath) == 0 || dstRoot == nil || srcRoot == nil { - return - } - if len(keyPath) > 1 { - dstParent := dstRoot - srcParent := srcRoot - for _, key := range keyPath[:len(keyPath)-1] { - if key == "" || dstParent == nil || dstParent.Kind != yaml.MappingNode { - return - } - dstIdx := findMapKeyIndex(dstParent, key) - if dstIdx < 0 || dstIdx+1 >= len(dstParent.Content) { - return - } - dstParent = dstParent.Content[dstIdx+1] - - if srcParent != nil && srcParent.Kind == yaml.MappingNode { - srcIdx := findMapKeyIndex(srcParent, key) - if srcIdx >= 0 && srcIdx+1 < len(srcParent.Content) { - srcParent = srcParent.Content[srcIdx+1] - } else { - srcParent = nil - } - } - } - if srcParent == nil || srcParent.Kind != yaml.MappingNode { - removeMapKey(dstParent, keyPath[len(keyPath)-1]) - return - } - pruneMappingToGeneratedKeys(dstParent, srcParent, keyPath[len(keyPath)-1]) - return - } - key := keyPath[0] - if key == "" { - return - } - if dstRoot.Kind != yaml.MappingNode || srcRoot.Kind != yaml.MappingNode { - return - } - dstIdx := findMapKeyIndex(dstRoot, key) - if dstIdx < 0 || dstIdx+1 >= len(dstRoot.Content) { - return - } - srcIdx := findMapKeyIndex(srcRoot, key) - if srcIdx < 0 { - // Keep an explicit empty mapping for oauth-model-alias when it was previously present. - // When users delete the last channel from oauth-model-alias via the management API, - // we want that deletion to persist across hot reloads and restarts. - if key == "oauth-model-alias" { - dstRoot.Content[dstIdx+1] = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} - return - } - removeMapKey(dstRoot, key) - return - } - if srcIdx+1 >= len(srcRoot.Content) { - return - } - srcVal := srcRoot.Content[srcIdx+1] - dstVal := dstRoot.Content[dstIdx+1] - if srcVal == nil { - dstRoot.Content[dstIdx+1] = nil - return - } - if srcVal.Kind != yaml.MappingNode { - dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal) - return - } - if dstVal == nil || dstVal.Kind != yaml.MappingNode { - dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal) - return - } - pruneMissingMapKeys(dstVal, srcVal) -} - -func pruneMissingMapKeys(dstMap, srcMap *yaml.Node) { - if dstMap == nil || srcMap == nil || dstMap.Kind != yaml.MappingNode || srcMap.Kind != yaml.MappingNode { - return - } - keep := make(map[string]struct{}, len(srcMap.Content)/2) - for i := 0; i+1 < len(srcMap.Content); i += 2 { - keyNode := srcMap.Content[i] - if keyNode == nil { - continue - } - key := strings.TrimSpace(keyNode.Value) - if key == "" { - continue - } - keep[key] = struct{}{} - } - for i := 0; i+1 < len(dstMap.Content); { - keyNode := dstMap.Content[i] - if keyNode == nil { - i += 2 - continue - } - key := strings.TrimSpace(keyNode.Value) - if _, ok := keep[key]; !ok { - dstMap.Content = append(dstMap.Content[:i], dstMap.Content[i+2:]...) - continue - } - i += 2 - } -} - -// normalizeCollectionNodeStyles forces YAML collections to use block notation, keeping -// lists and maps readable. Empty sequences retain flow style ([]) so empty list markers -// remain compact. -func normalizeCollectionNodeStyles(node *yaml.Node) { - if node == nil { - return - } - switch node.Kind { - case yaml.MappingNode: - node.Style = 0 - for i := range node.Content { - normalizeCollectionNodeStyles(node.Content[i]) - } - case yaml.SequenceNode: - if len(node.Content) == 0 { - node.Style = yaml.FlowStyle - } else { - node.Style = 0 - } - for i := range node.Content { - normalizeCollectionNodeStyles(node.Content[i]) - } - default: - // Scalars keep their existing style to preserve quoting - } -} - -func removeLegacyOpenAICompatAPIKeys(root *yaml.Node) { - if root == nil || root.Kind != yaml.MappingNode { - return - } - idx := findMapKeyIndex(root, "openai-compatibility") - if idx < 0 || idx+1 >= len(root.Content) { - return - } - seq := root.Content[idx+1] - if seq == nil || seq.Kind != yaml.SequenceNode { - return - } - for i := range seq.Content { - if seq.Content[i] != nil && seq.Content[i].Kind == yaml.MappingNode { - removeMapKey(seq.Content[i], "api-keys") - } - } -} - -func removeRemovedIntegrationKeys(root *yaml.Node) { - if root == nil || root.Kind != yaml.MappingNode { - return - } - removeMapKey(root, "ampcode") - removeMapKey(root, "amp-upstream-url") - removeMapKey(root, "amp-upstream-api-key") - removeMapKey(root, "amp-restrict-management-to-localhost") - removeMapKey(root, "amp-model-mappings") -} - -func removeLegacyGenerativeLanguageKeys(root *yaml.Node) { - if root == nil || root.Kind != yaml.MappingNode { - return - } - removeMapKey(root, "generative-language-api-key") -} - -func removeLegacyAuthBlock(root *yaml.Node) { - if root == nil || root.Kind != yaml.MappingNode { - return - } - removeMapKey(root, "auth") -} diff --git a/internal/config/config_defaults.go b/internal/config/config_defaults.go new file mode 100644 index 000000000..8e57ab80d --- /dev/null +++ b/internal/config/config_defaults.go @@ -0,0 +1,7 @@ +package config + +const ( + DefaultPanelGitHubRepository = "https://github.com/router-for-me/Cli-Proxy-API-Management-Center" + DefaultPprofAddr = "127.0.0.1:8316" + DefaultAuthDir = "~/.cli-proxy-api" +) diff --git a/internal/config/config_load.go b/internal/config/config_load.go new file mode 100644 index 000000000..61570320a --- /dev/null +++ b/internal/config/config_load.go @@ -0,0 +1,176 @@ +package config + +import ( + "bytes" + "errors" + "fmt" + "os" + "strings" + "syscall" + + log "github.com/sirupsen/logrus" + "gopkg.in/yaml.v3" +) + +// LoadConfig reads a YAML configuration file from the given path, +// unmarshals it into a Config struct, applies environment variable overrides, +// and returns it. +// +// Parameters: +// - configFile: The path to the YAML configuration file +// +// Returns: +// - *Config: The loaded configuration +// - error: An error if the configuration could not be loaded +func LoadConfig(configFile string) (*Config, error) { + return LoadConfigOptional(configFile, false) +} + +// LoadConfigOptional reads YAML from configFile. +// If optional is true and the file is missing, it returns an empty Config. +// If optional is true and the file is empty or invalid, it returns an empty Config. +func LoadConfigOptional(configFile string, optional bool) (*Config, error) { + // Read the entire configuration file into memory. + data, err := os.ReadFile(configFile) + if err != nil { + if optional { + if os.IsNotExist(err) || errors.Is(err, syscall.EISDIR) { + // Missing and optional: return empty config (cloud deploy standby). + cfg := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} + cfg.NormalizePluginsConfig() + return cfg, nil + } + } + return nil, fmt.Errorf("failed to read config file: %w", err) + } + + // In cloud deploy mode (optional=true), if file is empty or contains only whitespace, return empty config. + if optional && len(bytes.TrimSpace(data)) == 0 { + cfg := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} + cfg.NormalizePluginsConfig() + return cfg, nil + } + + // Unmarshal the YAML data into the Config struct. + var cfg Config + // Set defaults before unmarshal so that absent keys keep defaults. + cfg.Host = "" // Default empty: binds to all interfaces (IPv4 + IPv6) + cfg.LoggingToFile = false + cfg.LogsMaxTotalSizeMB = 0 + cfg.ErrorLogsMaxFiles = 10 + cfg.UsageStatisticsEnabled = false + cfg.RedisUsageQueueRetentionSeconds = 60 + cfg.DisableCooling = false + cfg.SaveCooldownStatus = false + cfg.TransientErrorCooldownSeconds = 0 + cfg.DisableImageGeneration = DisableImageGenerationOff + cfg.WebsocketAuth = true + cfg.Pprof.Enable = false + cfg.Pprof.Addr = DefaultPprofAddr + cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository + cfg.CredentialInFlight = DefaultCredentialInFlightConfig() + if err = yaml.Unmarshal(data, &cfg); err != nil { + if optional { + // In cloud deploy mode, if YAML parsing fails, return empty config instead of error. + cfgOptional := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} + cfgOptional.NormalizePluginsConfig() + return cfgOptional, nil + } + return nil, fmt.Errorf("failed to parse config file: %w", err) + } + + cfg.CredentialConcurrency = cfg.CredentialConcurrency.WithDefaults() + if errValidate := cfg.CredentialInFlight.Validate(); errValidate != nil { + return nil, errValidate + } + if errValidate := cfg.Codex.LiveMediaRelay.Validate(); errValidate != nil { + return nil, errValidate + } + + // Hash remote management key if plaintext is detected (nested) + // We consider a value to be already hashed if it looks like a bcrypt hash ($2a$, $2b$, or $2y$ prefix). + if cfg.RemoteManagement.SecretKey != "" && !looksLikeBcrypt(cfg.RemoteManagement.SecretKey) { + hashed, errHash := hashSecret(cfg.RemoteManagement.SecretKey) + if errHash != nil { + return nil, fmt.Errorf("failed to hash remote management key: %w", errHash) + } + cfg.RemoteManagement.SecretKey = hashed + + // Persist the hashed value back to the config file to avoid re-hashing on next startup. + // Preserve YAML comments and ordering; update only the nested key. + _ = SaveConfigPreserveCommentsUpdateNestedScalar(configFile, []string{"remote-management", "secret-key"}, hashed) + } + + cfg.RemoteManagement.PanelGitHubRepository = strings.TrimSpace(cfg.RemoteManagement.PanelGitHubRepository) + if cfg.RemoteManagement.PanelGitHubRepository == "" { + cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository + } + + cfg.Pprof.Addr = strings.TrimSpace(cfg.Pprof.Addr) + if cfg.Pprof.Addr == "" { + cfg.Pprof.Addr = DefaultPprofAddr + } + + if cfg.LogsMaxTotalSizeMB < 0 { + cfg.LogsMaxTotalSizeMB = 0 + } + + if cfg.ErrorLogsMaxFiles < 0 { + cfg.ErrorLogsMaxFiles = 10 + } + + if cfg.RedisUsageQueueRetentionSeconds <= 0 { + cfg.RedisUsageQueueRetentionSeconds = 60 + } else if cfg.RedisUsageQueueRetentionSeconds > 3600 { + log.WithField("value", cfg.RedisUsageQueueRetentionSeconds).Warn("redis-usage-queue-retention-seconds too large; clamping to 3600") + cfg.RedisUsageQueueRetentionSeconds = 3600 + } + + if cfg.MaxRetryCredentials < 0 { + cfg.MaxRetryCredentials = 0 + } + + cfg.NormalizePluginsConfig() + if errResolvePluginsDir := cfg.ResolvePluginsDir(); errResolvePluginsDir != nil && cfg.Plugins.Enabled { + return nil, errResolvePluginsDir + } + + // Sanitize Gemini API key configuration and migrate legacy entries. + cfg.SanitizeGeminiKeys() + + // Sanitize native Interactions API key configuration. + cfg.SanitizeInteractionsKeys() + + // Sanitize Vertex-compatible API keys. + cfg.SanitizeVertexCompatKeys() + + // Sanitize Codex keys: drop entries without base-url + cfg.SanitizeCodexKeys() + + // Sanitize xAI keys: drop entries without base-url + cfg.SanitizeXAIKeys() + + // Sanitize Codex header defaults. + cfg.SanitizeCodexHeaderDefaults() + + // Sanitize Claude header defaults. + cfg.SanitizeClaudeHeaderDefaults() + + // Sanitize Claude key headers + cfg.SanitizeClaudeKeys() + + // Sanitize OpenAI compatibility providers: drop entries without base-url + cfg.SanitizeOpenAICompatibility() + + // Normalize OAuth provider model exclusion map. + cfg.OAuthExcludedModels = NormalizeOAuthExcludedModels(cfg.OAuthExcludedModels) + + // Normalize global OAuth model name aliases. + cfg.SanitizeOAuthModelAlias() + + // Validate raw payload rules and drop invalid entries. + cfg.SanitizePayloadRules() + + // Return the populated configuration struct. + return &cfg, nil +} diff --git a/internal/config/config_normalization.go b/internal/config/config_normalization.go new file mode 100644 index 000000000..697f6c3d0 --- /dev/null +++ b/internal/config/config_normalization.go @@ -0,0 +1,297 @@ +package config + +import ( + "strings" + + sdkpluginstore "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginstore" +) + +// NormalizePluginsConfig applies default plugin configuration values. +func (cfg *Config) NormalizePluginsConfig() { + if cfg == nil { + return + } + cfg.Plugins.Dir = strings.TrimSpace(cfg.Plugins.Dir) + if cfg.Plugins.Dir == "" { + cfg.Plugins.Dir = defaultPluginsDir + } + if len(cfg.Plugins.StoreSources) > 0 { + sources := make([]string, 0, len(cfg.Plugins.StoreSources)) + for _, source := range cfg.Plugins.StoreSources { + source = strings.TrimSpace(source) + if source == "" { + continue + } + sources = append(sources, source) + } + cfg.Plugins.StoreSources = sources + } + cfg.Plugins.StoreAuth = sdkpluginstore.NormalizeAuthConfigs(cfg.Plugins.StoreAuth) + if cfg.Plugins.Configs == nil { + cfg.Plugins.Configs = map[string]PluginInstanceConfig{} + } +} + +// SanitizeCodexHeaderDefaults trims surrounding whitespace from the +// configured Codex header fallback values. +func (cfg *Config) SanitizeCodexHeaderDefaults() { + if cfg == nil { + return + } + cfg.CodexHeaderDefaults.UserAgent = strings.TrimSpace(cfg.CodexHeaderDefaults.UserAgent) + cfg.CodexHeaderDefaults.BetaFeatures = strings.TrimSpace(cfg.CodexHeaderDefaults.BetaFeatures) +} + +// SanitizeClaudeHeaderDefaults trims surrounding whitespace from the +// configured Claude fingerprint baseline values. +func (cfg *Config) SanitizeClaudeHeaderDefaults() { + if cfg == nil { + return + } + cfg.ClaudeHeaderDefaults.UserAgent = strings.TrimSpace(cfg.ClaudeHeaderDefaults.UserAgent) + cfg.ClaudeHeaderDefaults.PackageVersion = strings.TrimSpace(cfg.ClaudeHeaderDefaults.PackageVersion) + cfg.ClaudeHeaderDefaults.RuntimeVersion = strings.TrimSpace(cfg.ClaudeHeaderDefaults.RuntimeVersion) + cfg.ClaudeHeaderDefaults.OS = strings.TrimSpace(cfg.ClaudeHeaderDefaults.OS) + cfg.ClaudeHeaderDefaults.Arch = strings.TrimSpace(cfg.ClaudeHeaderDefaults.Arch) + cfg.ClaudeHeaderDefaults.Timeout = strings.TrimSpace(cfg.ClaudeHeaderDefaults.Timeout) +} + +// SanitizeOAuthModelAlias normalizes and deduplicates global OAuth model name aliases. +// It trims whitespace, normalizes channel keys to lower-case, drops empty entries, +// allows multiple aliases per upstream name, and ensures aliases are unique within each channel. +func (cfg *Config) SanitizeOAuthModelAlias() { + if cfg == nil || len(cfg.OAuthModelAlias) == 0 { + return + } + out := make(map[string][]OAuthModelAlias, len(cfg.OAuthModelAlias)) + for rawChannel, aliases := range cfg.OAuthModelAlias { + channel := strings.ToLower(strings.TrimSpace(rawChannel)) + if channel == "" || len(aliases) == 0 { + continue + } + seenAlias := make(map[string]struct{}, len(aliases)) + clean := make([]OAuthModelAlias, 0, len(aliases)) + for _, entry := range aliases { + name := strings.TrimSpace(entry.Name) + alias := strings.TrimSpace(entry.Alias) + if name == "" || alias == "" { + continue + } + if strings.EqualFold(name, alias) { + continue + } + aliasKey := strings.ToLower(alias) + if _, ok := seenAlias[aliasKey]; ok { + continue + } + seenAlias[aliasKey] = struct{}{} + clean = append(clean, OAuthModelAlias{ + Name: name, + Alias: alias, + Fork: entry.Fork, + DisplayName: strings.TrimSpace(entry.DisplayName), + ForceMapping: entry.ForceMapping, + }) + } + if len(clean) > 0 { + out[channel] = clean + } + } + cfg.OAuthModelAlias = out +} + +// SanitizeOpenAICompatibility removes OpenAI-compatibility provider entries that are +// not actionable, specifically those missing a BaseURL. It trims whitespace before +// evaluation and preserves the relative order of remaining entries. +func (cfg *Config) SanitizeOpenAICompatibility() { + if cfg == nil || len(cfg.OpenAICompatibility) == 0 { + return + } + out := make([]OpenAICompatibility, 0, len(cfg.OpenAICompatibility)) + for i := range cfg.OpenAICompatibility { + e := cfg.OpenAICompatibility[i] + e.Name = strings.TrimSpace(e.Name) + e.Prefix = normalizeModelPrefix(e.Prefix) + e.BaseURL = strings.TrimSpace(e.BaseURL) + e.Headers = NormalizeHeaders(e.Headers) + if e.BaseURL == "" { + // Skip providers with no base-url; treated as removed + continue + } + out = append(out, e) + } + cfg.OpenAICompatibility = out +} + +// SanitizeCodexKeys removes Codex API key entries missing a BaseURL. +// It trims whitespace and preserves order for remaining entries. +func (cfg *Config) SanitizeCodexKeys() { + if cfg == nil { + return + } + cfg.CodexKey = sanitizeCodexKeyEntries(cfg.CodexKey) +} + +// SanitizeXAIKeys removes xAI API key entries missing a BaseURL. +// It applies the same normalization rules as codex-api-key. +func (cfg *Config) SanitizeXAIKeys() { + if cfg == nil { + return + } + cfg.XAIKey = sanitizeCodexKeyEntries(cfg.XAIKey) +} + +func sanitizeCodexKeyEntries(entries []CodexKey) []CodexKey { + if len(entries) == 0 { + return entries + } + out := make([]CodexKey, 0, len(entries)) + for i := range entries { + e := entries[i] + e.Prefix = normalizeModelPrefix(e.Prefix) + e.BaseURL = strings.TrimSpace(e.BaseURL) + e.Headers = NormalizeHeaders(e.Headers) + e.ExcludedModels = NormalizeExcludedModels(e.ExcludedModels) + if e.BaseURL == "" { + continue + } + out = append(out, e) + } + return out +} + +// SanitizeClaudeKeys normalizes headers for Claude credentials. +func (cfg *Config) SanitizeClaudeKeys() { + if cfg == nil || len(cfg.ClaudeKey) == 0 { + return + } + for i := range cfg.ClaudeKey { + entry := &cfg.ClaudeKey[i] + entry.Prefix = normalizeModelPrefix(entry.Prefix) + entry.Headers = NormalizeHeaders(entry.Headers) + entry.ExcludedModels = NormalizeExcludedModels(entry.ExcludedModels) + } +} + +func sanitizeGeminiKeyEntries(entries []GeminiKey) []GeminiKey { + seen := make(map[string]struct{}, len(entries)) + out := entries[:0] + for i := range entries { + entry := entries[i] + entry.APIKey = strings.TrimSpace(entry.APIKey) + if entry.APIKey == "" { + continue + } + entry.Prefix = normalizeModelPrefix(entry.Prefix) + entry.BaseURL = strings.TrimSpace(entry.BaseURL) + entry.ProxyURL = strings.TrimSpace(entry.ProxyURL) + entry.Headers = NormalizeHeaders(entry.Headers) + entry.ExcludedModels = NormalizeExcludedModels(entry.ExcludedModels) + uniqueKey := entry.APIKey + "|" + entry.BaseURL + if _, exists := seen[uniqueKey]; exists { + continue + } + seen[uniqueKey] = struct{}{} + out = append(out, entry) + } + return out +} + +// SanitizeGeminiKeys deduplicates and normalizes Gemini credentials. +// It uses API key + base URL as the uniqueness key. +func (cfg *Config) SanitizeGeminiKeys() { + if cfg == nil { + return + } + cfg.GeminiKey = sanitizeGeminiKeyEntries(cfg.GeminiKey) +} + +// SanitizeInteractionsKeys deduplicates and normalizes native Interactions credentials. +// It uses API key + base URL as the uniqueness key. +func (cfg *Config) SanitizeInteractionsKeys() { + if cfg == nil { + return + } + cfg.InteractionsKey = sanitizeGeminiKeyEntries(cfg.InteractionsKey) +} + +func normalizeModelPrefix(prefix string) string { + trimmed := strings.TrimSpace(prefix) + trimmed = strings.Trim(trimmed, "/") + if trimmed == "" { + return "" + } + if strings.Contains(trimmed, "/") { + return "" + } + return trimmed +} + +// NormalizeHeaders trims header keys and values and removes empty pairs. +func NormalizeHeaders(headers map[string]string) map[string]string { + if len(headers) == 0 { + return nil + } + clean := make(map[string]string, len(headers)) + for k, v := range headers { + key := strings.TrimSpace(k) + val := strings.TrimSpace(v) + if key == "" || val == "" { + continue + } + clean[key] = val + } + if len(clean) == 0 { + return nil + } + return clean +} + +// NormalizeExcludedModels trims, lowercases, and deduplicates model exclusion patterns. +// It preserves the order of first occurrences and drops empty entries. +func NormalizeExcludedModels(models []string) []string { + if len(models) == 0 { + return nil + } + seen := make(map[string]struct{}, len(models)) + out := make([]string, 0, len(models)) + for _, raw := range models { + trimmed := strings.ToLower(strings.TrimSpace(raw)) + if trimmed == "" { + continue + } + if _, exists := seen[trimmed]; exists { + continue + } + seen[trimmed] = struct{}{} + out = append(out, trimmed) + } + if len(out) == 0 { + return nil + } + return out +} + +// NormalizeOAuthExcludedModels cleans provider -> excluded models mappings by normalizing provider keys +// and applying model exclusion normalization to each entry. +func NormalizeOAuthExcludedModels(entries map[string][]string) map[string][]string { + if len(entries) == 0 { + return nil + } + out := make(map[string][]string, len(entries)) + for provider, models := range entries { + key := strings.ToLower(strings.TrimSpace(provider)) + if key == "" { + continue + } + normalized := NormalizeExcludedModels(models) + if len(normalized) == 0 { + continue + } + out[key] = normalized + } + if len(out) == 0 { + return nil + } + return out +} diff --git a/internal/config/config_types.go b/internal/config/config_types.go new file mode 100644 index 000000000..dca5fd03f --- /dev/null +++ b/internal/config/config_types.go @@ -0,0 +1,580 @@ +package config + +import ( + "fmt" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + sdkpluginstore "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginstore" + "gopkg.in/yaml.v3" +) + +// PluginsConfig holds dynamic plugin system settings. +type PluginsConfig struct { + // Enabled toggles dynamic plugin loading. + Enabled bool `yaml:"enabled" json:"enabled"` + // Dir is the plugin discovery directory. + Dir string `yaml:"dir" json:"dir"` + // StoreSources appends third-party plugin store registries to the built-in official source. + StoreSources []string `yaml:"store-sources,omitempty" json:"store-sources,omitempty"` + // StoreAuth defines optional auth rules for plugin store registry, metadata, and artifact requests. + StoreAuth []sdkpluginstore.AuthConfig `yaml:"store-auth,omitempty" json:"store-auth,omitempty"` + // AuthRevision changes when Home-managed plugin credentials change. + AuthRevision int64 `yaml:"auth-revision,omitempty" json:"auth-revision,omitempty"` + // Configs stores per-plugin instance configuration by plugin ID. + Configs map[string]PluginInstanceConfig `yaml:"configs" json:"configs"` +} + +// PluginInstanceConfig stores host-owned plugin settings and the original plugin YAML subtree. +type PluginInstanceConfig struct { + // Enabled toggles this plugin instance. Nil is normalized to false during YAML parsing. + Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` + // Priority controls plugin startup and routing order. + Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + // Raw preserves the full original plugin configuration YAML subtree. + Raw yaml.Node `yaml:"-" json:"-"` +} + +// UnmarshalYAML extracts host-owned fields while preserving the full original YAML node. +func (c *PluginInstanceConfig) UnmarshalYAML(value *yaml.Node) error { + if c == nil { + return nil + } + + c.Priority = 0 + defaultEnabled := false + c.Enabled = &defaultEnabled + + if value == nil || value.Kind == 0 { + c.Raw = *defaultPluginInstanceConfigNode() + return nil + } + + c.Raw = *deepCopyNode(value) + if value.Kind != yaml.MappingNode { + return nil + } + + for i := 0; i+1 < len(value.Content); i += 2 { + key := value.Content[i] + node := value.Content[i+1] + if key == nil { + continue + } + switch key.Value { + case "enabled": + var enabled bool + if errDecodeEnabled := node.Decode(&enabled); errDecodeEnabled != nil { + return fmt.Errorf("parse plugin enabled: %w", errDecodeEnabled) + } + c.Enabled = &enabled + case "priority": + var priority int + if errDecodePriority := node.Decode(&priority); errDecodePriority != nil { + return fmt.Errorf("parse plugin priority: %w", errDecodePriority) + } + c.Priority = priority + } + } + + return nil +} + +// MarshalYAML returns the preserved raw plugin YAML subtree for lossless config output. +func (c PluginInstanceConfig) MarshalYAML() (any, error) { + if c.Raw.Kind == 0 { + return defaultPluginInstanceConfigNode(), nil + } + return deepCopyNode(&c.Raw), nil +} + +func defaultPluginInstanceConfigNode() *yaml.Node { + return &yaml.Node{ + Kind: yaml.MappingNode, + Tag: "!!map", + Content: []*yaml.Node{}, + } +} + +// ClaudeHeaderDefaults configures default header values injected into Claude API requests. +// In legacy mode, UserAgent/PackageVersion/RuntimeVersion/Timeout act as fallbacks when +// the client omits them, while OS/Arch remain runtime-derived. When stabilized device +// profiles are enabled, OS/Arch become the pinned platform baseline, while +// UserAgent/PackageVersion/RuntimeVersion seed the upgradeable software fingerprint. +type ClaudeHeaderDefaults struct { + UserAgent string `yaml:"user-agent" json:"user-agent"` + PackageVersion string `yaml:"package-version" json:"package-version"` + RuntimeVersion string `yaml:"runtime-version" json:"runtime-version"` + OS string `yaml:"os" json:"os"` + Arch string `yaml:"arch" json:"arch"` + Timeout string `yaml:"timeout" json:"timeout"` + StabilizeDeviceProfile *bool `yaml:"stabilize-device-profile,omitempty" json:"stabilize-device-profile,omitempty"` +} + +// CodexHeaderDefaults configures fallback header values injected into Codex +// model requests for OAuth/file-backed auth when the client omits them. +// UserAgent applies to HTTP and websocket requests; BetaFeatures only applies to websockets. +type CodexHeaderDefaults struct { + UserAgent string `yaml:"user-agent" json:"user-agent"` + BetaFeatures string `yaml:"beta-features" json:"beta-features"` +} + +// CodexConfig configures provider-wide Codex request behavior. +type CodexConfig struct { + IdentityConfuse bool `yaml:"identity-confuse" json:"identity-confuse"` + // OptimizeMultiAgentV2 optimizes official Codex multi-agent requests. + OptimizeMultiAgentV2 bool `yaml:"optimize-multi-agent-v2" json:"optimize-multi-agent-v2"` + // LiveMediaRelay terminates and relays Codex Live WebRTC media in this process. + LiveMediaRelay CodexLiveMediaRelayConfig `yaml:"live-media-relay" json:"live-media-relay"` +} + +// CodexLiveMediaRelayConfig configures the in-process Codex Live WebRTC gateway. +type CodexLiveMediaRelayConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + MaxSessions int `yaml:"max-sessions" json:"max-sessions"` + DisablePrivateRemoteIPs bool `yaml:"disable-private-remote-ips" json:"disable-private-remote-ips"` + PublicIP string `yaml:"public-ip" json:"public-ip"` + UDPPortMin uint16 `yaml:"udp-port-min" json:"udp-port-min"` + UDPPortMax uint16 `yaml:"udp-port-max" json:"udp-port-max"` + ICEServers []CodexLiveICEServer `yaml:"ice-servers" json:"ice-servers"` +} + +// CodexLiveICEServer configures a STUN or TURN server for the media relay. +type CodexLiveICEServer struct { + URLs []string `yaml:"urls" json:"urls"` + Username string `yaml:"username" json:"-"` + Credential string `yaml:"credential" json:"-"` +} + +// TLSConfig holds HTTPS server settings. +type TLSConfig struct { + // Enable toggles HTTPS server mode. + Enable bool `yaml:"enable" json:"enable"` + // Cert is the path to the TLS certificate file. + Cert string `yaml:"cert" json:"cert"` + // Key is the path to the TLS private key file. + Key string `yaml:"key" json:"key"` +} + +// PprofConfig holds pprof HTTP server settings. +type PprofConfig struct { + // Enable toggles the pprof HTTP debug server. + Enable bool `yaml:"enable" json:"enable"` + // Addr is the host:port address for the pprof HTTP server. + Addr string `yaml:"addr" json:"addr"` +} + +// RemoteManagement holds management API configuration under 'remote-management'. +type RemoteManagement struct { + // AllowRemote toggles remote (non-localhost) access to management API. + AllowRemote bool `yaml:"allow-remote"` + // SecretKey is the management key (plaintext or bcrypt hashed). YAML key intentionally 'secret-key'. + SecretKey string `yaml:"secret-key"` + // DisableControlPanel skips serving and syncing the bundled management UI when true. + DisableControlPanel bool `yaml:"disable-control-panel"` + // DisableAutoUpdatePanel disables automatic periodic background updates of the management panel asset from GitHub. + // When false (the default), the background updater remains enabled; when true, the panel is only downloaded on first access if missing. + DisableAutoUpdatePanel bool `yaml:"disable-auto-update-panel"` + // PanelGitHubRepository overrides the GitHub repository used to fetch the management panel asset. + // Accepts either a repository URL (https://github.com/org/repo) or an API releases endpoint. + PanelGitHubRepository string `yaml:"panel-github-repository"` +} + +// QuotaExceeded defines the behavior when API quota limits are exceeded. +// It provides configuration options for automatic failover mechanisms. +type QuotaExceeded struct { + // SwitchProject indicates whether to automatically switch to another project when a quota is exceeded. + SwitchProject bool `yaml:"switch-project" json:"switch-project"` + + // SwitchPreviewModel indicates whether to automatically switch to a preview model when a quota is exceeded. + SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"` + + // AntigravityCredits enables credits-based last-resort fallback for Claude models. + // When all free-tier auths are exhausted (429/503), the conductor retries with + // an auth that has available Google One AI credits. + AntigravityCredits bool `yaml:"antigravity-credits" json:"antigravity-credits"` +} + +// RoutingConfig configures how credentials are selected for requests. +type RoutingConfig struct { + // Strategy selects the credential selection strategy. + // Supported values: "round-robin" (default), "fill-first". + Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"` + + // SessionAffinity enables universal session-sticky routing for all clients. + // Session IDs are extracted from multiple sources: + // metadata.user_id (Claude Code session format), X-Session-ID, Session_id (Codex), + // X-Client-Request-Id (PI), metadata.user_id, conversation_id, or message hash. + // Automatic failover is always enabled when bound auth becomes unavailable. + SessionAffinity bool `yaml:"session-affinity,omitempty" json:"session-affinity,omitempty"` + + // SessionAffinityTTL specifies how long session-to-auth bindings are retained. + // Default: 1h. Accepts duration strings like "30m", "1h", "2h30m". + SessionAffinityTTL string `yaml:"session-affinity-ttl,omitempty" json:"session-affinity-ttl,omitempty"` +} + +// OAuthModelAlias defines a model ID alias for a specific channel. +// It maps the upstream model name (Name) to the client-visible alias (Alias). +// When Fork is true, the alias is added as an additional model in listings while +// keeping the original model ID available. +type OAuthModelAlias struct { + Name string `yaml:"name" json:"name"` + Alias string `yaml:"alias" json:"alias"` + Fork bool `yaml:"fork,omitempty" json:"fork,omitempty"` + + // DisplayName is the optional human-readable name shown in model catalogs. + DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` + + ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` +} + +// PayloadConfig defines default and override parameter rules applied to provider payloads. +type PayloadConfig struct { + // Default defines rules that only set parameters when they are missing in the payload. + Default []PayloadRule `yaml:"default" json:"default"` + // DefaultRaw defines rules that set raw JSON values only when they are missing. + DefaultRaw []PayloadRule `yaml:"default-raw" json:"default-raw"` + // Override defines rules that always set parameters, overwriting any existing values. + Override []PayloadRule `yaml:"override" json:"override"` + // OverrideRaw defines rules that always set raw JSON values, overwriting any existing values. + OverrideRaw []PayloadRule `yaml:"override-raw" json:"override-raw"` + // Filter defines rules that remove parameters from the payload by JSON path. + Filter []PayloadFilterRule `yaml:"filter" json:"filter"` +} + +// PayloadFilterRule describes a rule to remove specific JSON paths from matching model payloads. +type PayloadFilterRule struct { + // Models lists model entries with name pattern and protocol constraint. + Models []PayloadModelRule `yaml:"models" json:"models"` + // Params lists JSON paths (gjson/sjson syntax) to remove from the payload. + Params []string `yaml:"params" json:"params"` +} + +// PayloadRule describes a single rule targeting a list of models with parameter updates. +type PayloadRule struct { + // Models lists model entries with name pattern and protocol constraint. + Models []PayloadModelRule `yaml:"models" json:"models"` + // Params maps JSON paths (gjson/sjson syntax) to values written into the payload. + // For *-raw rules, values are treated as raw JSON fragments (strings are used as-is). + Params map[string]any `yaml:"params" json:"params"` +} + +// PayloadModelRule ties a model name pattern to a specific translator protocol. +type PayloadModelRule struct { + // Name is the model name or wildcard pattern (e.g., "gpt-*", "*-5", "gemini-*-pro"). + Name string `yaml:"name" json:"name"` + // Protocol restricts the rule to a specific translator format (e.g., "gemini", "responses"). + Protocol string `yaml:"protocol" json:"protocol"` + // Headers restricts the rule to requests whose headers match all configured wildcard patterns. + Headers map[string]string `yaml:"headers" json:"headers"` + // FromProtocol restricts the rule to a specific source protocol (e.g., "gemini", "responses"). + FromProtocol string `yaml:"from-protocol" json:"from-protocol"` + // Match requires payload JSON paths to equal the configured values. + Match []map[string]any `yaml:"match" json:"match"` + // NotMatch requires payload JSON paths to not equal the configured values. + NotMatch []map[string]any `yaml:"not-match" json:"not-match"` + // Exist requires payload JSON paths to exist and not be null. + Exist []string `yaml:"exist" json:"exist"` + // NotExist requires payload JSON paths to be missing or null. + NotExist []string `yaml:"not-exist" json:"not-exist"` +} + +// CloakConfig configures request cloaking for non-Claude-Code clients. +// Cloaking disguises API requests to appear as originating from the official Claude Code CLI. +type CloakConfig struct { + // Mode controls cloaking behavior: "auto" (default), "always", or "never". + // - "auto": cloak only when client is not Claude Code (based on User-Agent) + // - "always": always apply cloaking regardless of client + // - "never": never apply cloaking + Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` + + // StrictMode controls how system prompts are handled when cloaking. + // - false (default): prepend Claude Code prompt to user system messages + // - true: strip all user system messages, keep only Claude Code prompt + StrictMode bool `yaml:"strict-mode,omitempty" json:"strict-mode,omitempty"` + + // SensitiveWords is a list of words to obfuscate with zero-width characters. + // This can help bypass certain content filters. + SensitiveWords []string `yaml:"sensitive-words,omitempty" json:"sensitive-words,omitempty"` + + // CacheUserID controls whether Claude user_id values are cached per API key. + // When false, a fresh random user_id is generated for every request. + CacheUserID *bool `yaml:"cache-user-id,omitempty" json:"cache-user-id,omitempty"` +} + +// ClaudeKey represents the configuration for a Claude API key, +// including the API key itself and an optional base URL for the API endpoint. +type ClaudeKey struct { + // APIKey is the authentication key for accessing Claude API services. + APIKey string `yaml:"api-key" json:"api-key"` + + // Priority controls selection preference when multiple credentials match. + // Higher values are preferred; defaults to 0. + Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + + // Prefix optionally namespaces models for this credential (e.g., "teamA/claude-sonnet-4"). + Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` + + // BaseURL is the base URL for the Claude API endpoint. + // If empty, the default Claude API URL will be used. + BaseURL string `yaml:"base-url" json:"base-url"` + + // ProxyURL overrides the global proxy setting for this API key if provided. + ProxyURL string `yaml:"proxy-url" json:"proxy-url"` + + // Models defines upstream model names and aliases for request routing. + Models []ClaudeModel `yaml:"models" json:"models"` + + // Headers optionally adds extra HTTP headers for requests sent with this key. + Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` + + // ExcludedModels lists model IDs that should be excluded for this provider. + ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` + + // RebuildMidSystemMessage moves Claude messages with role "system" into the top-level system field. + RebuildMidSystemMessage bool `yaml:"rebuild-mid-system-message,omitempty" json:"rebuild-mid-system-message,omitempty"` + + // DisableCooling disables auth/model cooldown scheduling for this credential when true. + DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` + + // Cloak configures request cloaking for non-Claude-Code clients. + Cloak *CloakConfig `yaml:"cloak,omitempty" json:"cloak,omitempty"` + + // ExperimentalCCHSigning enables opt-in final-body cch signing for cloaked + // Claude /v1/messages requests. It is disabled by default so upstream seed + // changes do not alter the proxy's legacy behavior. + ExperimentalCCHSigning bool `yaml:"experimental-cch-signing,omitempty" json:"experimental-cch-signing,omitempty"` +} + +func (k ClaudeKey) GetAPIKey() string { return k.APIKey } + +func (k ClaudeKey) GetBaseURL() string { return k.BaseURL } + +// ClaudeModel describes a mapping between an alias and the actual upstream model name. +type ClaudeModel struct { + // Name is the upstream model identifier used when issuing requests. + Name string `yaml:"name" json:"name"` + + // Alias is the client-facing model name that maps to Name. + Alias string `yaml:"alias" json:"alias"` + + // DisplayName is the optional human-readable name shown in model catalogs. + DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` + + // ForceMapping rewrites upstream response model fields back to Alias. + ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` +} + +func (m ClaudeModel) GetName() string { return m.Name } + +func (m ClaudeModel) GetAlias() string { return m.Alias } + +func (m ClaudeModel) GetDisplayName() string { return m.DisplayName } + +func (m ClaudeModel) GetForceMapping() bool { return m.ForceMapping } + +// CodexKey represents the configuration for a Codex API key, +// including the API key itself and an optional base URL for the API endpoint. +type CodexKey struct { + // APIKey is the authentication key for accessing Codex API services. + APIKey string `yaml:"api-key" json:"api-key"` + + // Priority controls selection preference when multiple credentials match. + // Higher values are preferred; defaults to 0. + Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + + // Prefix optionally namespaces models for this credential (e.g., "teamA/gpt-5-codex"). + Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` + + // BaseURL is the base URL for the Codex API endpoint. + // If empty, the default Codex API URL will be used. + BaseURL string `yaml:"base-url" json:"base-url"` + + // Websockets enables the Responses API websocket transport for this credential. + Websockets bool `yaml:"websockets,omitempty" json:"websockets,omitempty"` + + // ProxyURL overrides the global proxy setting for this API key if provided. + ProxyURL string `yaml:"proxy-url" json:"proxy-url"` + + // Models defines upstream model names and aliases for request routing. + Models []CodexModel `yaml:"models" json:"models"` + + // Headers optionally adds extra HTTP headers for requests sent with this key. + Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` + + // ExcludedModels lists model IDs that should be excluded for this provider. + ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` + + // DisableCooling disables auth/model cooldown scheduling for this credential when true. + DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` +} + +func (k CodexKey) GetAPIKey() string { return k.APIKey } + +func (k CodexKey) GetBaseURL() string { return k.BaseURL } + +// CodexModel describes a mapping between an alias and the actual upstream model name. +type CodexModel struct { + // Name is the upstream model identifier used when issuing requests. + Name string `yaml:"name" json:"name"` + + // Alias is the client-facing model name that maps to Name. + Alias string `yaml:"alias" json:"alias"` + + // DisplayName is the optional human-readable name shown in model catalogs. + DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` + + // ForceMapping rewrites upstream response model fields back to Alias. + ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` +} + +func (m CodexModel) GetName() string { return m.Name } + +func (m CodexModel) GetAlias() string { return m.Alias } + +func (m CodexModel) GetDisplayName() string { return m.DisplayName } + +func (m CodexModel) GetForceMapping() bool { return m.ForceMapping } + +// XAIKey uses the Codex API key structure for native xAI execution. +type XAIKey = CodexKey + +// XAIModel uses the Codex model mapping structure for xAI models. +type XAIModel = CodexModel + +// GeminiKey represents the configuration for a Gemini API key, +// including optional overrides for upstream base URL, proxy routing, and headers. +type GeminiKey struct { + // APIKey is the authentication key for accessing Gemini API services. + APIKey string `yaml:"api-key" json:"api-key"` + + // Priority controls selection preference when multiple credentials match. + // Higher values are preferred; defaults to 0. + Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + + // Prefix optionally namespaces models for this credential (e.g., "teamA/gemini-3-pro-preview"). + Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` + + // BaseURL optionally overrides the Gemini API endpoint. + BaseURL string `yaml:"base-url,omitempty" json:"base-url,omitempty"` + + // ProxyURL optionally overrides the global proxy for this API key. + ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"` + + // Models defines upstream model names and aliases for request routing. + Models []GeminiModel `yaml:"models,omitempty" json:"models,omitempty"` + + // Headers optionally adds extra HTTP headers for requests sent with this key. + Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` + + // ExcludedModels lists model IDs that should be excluded for this provider. + ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` + + // DisableCooling disables auth/model cooldown scheduling for this credential when true. + DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` +} + +func (k GeminiKey) GetAPIKey() string { return k.APIKey } + +func (k GeminiKey) GetBaseURL() string { return k.BaseURL } + +// GeminiModel describes a mapping between an alias and the actual upstream model name. +type GeminiModel struct { + // Name is the upstream model identifier used when issuing requests. + Name string `yaml:"name" json:"name"` + + // Alias is the client-facing model name that maps to Name. + Alias string `yaml:"alias" json:"alias"` + + // DisplayName is the optional human-readable name shown in model catalogs. + DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` + + // ForceMapping rewrites upstream response model fields back to Alias. + ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` +} + +func (m GeminiModel) GetName() string { return m.Name } + +func (m GeminiModel) GetAlias() string { return m.Alias } + +func (m GeminiModel) GetDisplayName() string { return m.DisplayName } + +func (m GeminiModel) GetForceMapping() bool { return m.ForceMapping } + +// OpenAICompatibility represents the configuration for OpenAI API compatibility +// with external providers, allowing model aliases to be routed through OpenAI API format. +type OpenAICompatibility struct { + // Name is the identifier for this OpenAI compatibility configuration. + Name string `yaml:"name" json:"name"` + + // Priority controls selection preference when multiple providers or credentials match. + // Higher values are preferred; defaults to 0. + Priority int `yaml:"priority,omitempty" json:"priority,omitempty"` + + // Disabled prevents this provider from being used for routing. + Disabled bool `yaml:"disabled,omitempty" json:"disabled,omitempty"` + + // Prefix optionally namespaces model aliases for this provider (e.g., "teamA/kimi-k2"). + Prefix string `yaml:"prefix,omitempty" json:"prefix,omitempty"` + + // BaseURL is the base URL for the external OpenAI-compatible API endpoint. + BaseURL string `yaml:"base-url" json:"base-url"` + + // APIKeyEntries defines API keys with optional per-key proxy configuration. + APIKeyEntries []OpenAICompatibilityAPIKey `yaml:"api-key-entries,omitempty" json:"api-key-entries,omitempty"` + + // Models defines the model configurations including aliases for routing. + Models []OpenAICompatibilityModel `yaml:"models" json:"models"` + + // Headers optionally adds extra HTTP headers for requests sent to this provider. + Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` + + // DisableCooling disables auth/model cooldown scheduling for this provider when true. + DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` +} + +// OpenAICompatibilityAPIKey represents an API key configuration with optional proxy setting. +type OpenAICompatibilityAPIKey struct { + // APIKey is the authentication key for accessing the external API services. + APIKey string `yaml:"api-key" json:"api-key"` + + // ProxyURL overrides the global proxy setting for this API key if provided. + ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"` +} + +// OpenAICompatibilityModel represents a model configuration for OpenAI compatibility, +// including the actual model name and its alias for API routing. +type OpenAICompatibilityModel struct { + // Name is the actual model name used by the external provider. + Name string `yaml:"name" json:"name"` + + // Alias is the model name alias that clients will use to reference this model. + Alias string `yaml:"alias" json:"alias"` + + // DisplayName is the optional human-readable name shown in model catalogs. + DisplayName string `yaml:"display-name,omitempty" json:"display-name,omitempty"` + + // ForceMapping rewrites upstream response model fields back to Alias. + ForceMapping bool `yaml:"force-mapping,omitempty" json:"force-mapping,omitempty"` + + // Image marks this model as callable through /v1/images/generations and /v1/images/edits. + Image bool `yaml:"image,omitempty" json:"image,omitempty"` + + // InputModalities declares chat/responses input capabilities (e.g. text, image) for Codex and other clients. + // This is separate from Image, which only enables /v1/images/* endpoints. + InputModalities []string `yaml:"input-modalities,omitempty" json:"input-modalities,omitempty"` + + // OutputModalities declares supported output modalities when known (e.g. text, image). + OutputModalities []string `yaml:"output-modalities,omitempty" json:"output-modalities,omitempty"` + + // Thinking configures the thinking/reasoning capability for this model. + // If nil, the model defaults to level-based reasoning with levels ["low", "medium", "high"]. + Thinking *registry.ThinkingSupport `yaml:"thinking,omitempty" json:"thinking,omitempty"` +} + +func (m OpenAICompatibilityModel) GetName() string { return m.Name } + +func (m OpenAICompatibilityModel) GetAlias() string { return m.Alias } + +func (m OpenAICompatibilityModel) GetDisplayName() string { return m.DisplayName } + +func (m OpenAICompatibilityModel) GetForceMapping() bool { return m.ForceMapping } diff --git a/internal/config/config_validation.go b/internal/config/config_validation.go new file mode 100644 index 000000000..7961e9ee3 --- /dev/null +++ b/internal/config/config_validation.go @@ -0,0 +1,79 @@ +package config + +import ( + "bytes" + "encoding/json" + + log "github.com/sirupsen/logrus" + "golang.org/x/crypto/bcrypt" +) + +// SanitizePayloadRules validates raw JSON payload rule params and drops invalid rules. +func (cfg *Config) SanitizePayloadRules() { + if cfg == nil { + return + } + cfg.Payload.DefaultRaw = sanitizePayloadRawRules(cfg.Payload.DefaultRaw, "default-raw") + cfg.Payload.OverrideRaw = sanitizePayloadRawRules(cfg.Payload.OverrideRaw, "override-raw") +} + +func sanitizePayloadRawRules(rules []PayloadRule, section string) []PayloadRule { + if len(rules) == 0 { + return rules + } + out := make([]PayloadRule, 0, len(rules)) + for i := range rules { + rule := rules[i] + if len(rule.Params) == 0 { + continue + } + invalid := false + for path, value := range rule.Params { + raw, ok := payloadRawString(value) + if !ok { + continue + } + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || !json.Valid(trimmed) { + log.WithFields(log.Fields{ + "section": section, + "rule_index": i + 1, + "param": path, + }).Warn("payload rule dropped: invalid raw JSON") + invalid = true + break + } + } + if invalid { + continue + } + out = append(out, rule) + } + return out +} + +func payloadRawString(value any) ([]byte, bool) { + switch typed := value.(type) { + case string: + return []byte(typed), true + case []byte: + return typed, true + default: + return nil, false + } +} + +// looksLikeBcrypt returns true if the provided string appears to be a bcrypt hash. +func looksLikeBcrypt(s string) bool { + return len(s) > 4 && (s[:4] == "$2a$" || s[:4] == "$2b$" || s[:4] == "$2y$") +} + +// hashSecret hashes the given secret using bcrypt. +func hashSecret(secret string) (string, error) { + // Use default cost for simplicity. + hashedBytes, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.DefaultCost) + if err != nil { + return "", err + } + return string(hashedBytes), nil +} diff --git a/internal/config/config_yaml.go b/internal/config/config_yaml.go new file mode 100644 index 000000000..69f4490b6 --- /dev/null +++ b/internal/config/config_yaml.go @@ -0,0 +1,815 @@ +package config + +import ( + "bytes" + "fmt" + "os" + "strings" + + "gopkg.in/yaml.v3" +) + +// SaveConfigPreserveComments writes the config back to YAML while preserving existing comments +// and key ordering by loading the original file into a yaml.Node tree and updating values in-place. +func SaveConfigPreserveComments(configFile string, cfg *Config) error { + persistCfg := cfg + // Load original YAML as a node tree to preserve comments and ordering. + data, err := os.ReadFile(configFile) + if err != nil { + return err + } + + var original yaml.Node + if err = yaml.Unmarshal(data, &original); err != nil { + return err + } + if original.Kind != yaml.DocumentNode || len(original.Content) == 0 { + return fmt.Errorf("invalid yaml document structure") + } + if original.Content[0] == nil || original.Content[0].Kind != yaml.MappingNode { + return fmt.Errorf("expected root mapping node") + } + + // Marshal the current cfg to YAML, then unmarshal to a yaml.Node we can merge from. + rendered, err := yaml.Marshal(persistCfg) + if err != nil { + return err + } + var generated yaml.Node + if err = yaml.Unmarshal(rendered, &generated); err != nil { + return err + } + if generated.Kind != yaml.DocumentNode || len(generated.Content) == 0 || generated.Content[0] == nil { + return fmt.Errorf("invalid generated yaml structure") + } + if generated.Content[0].Kind != yaml.MappingNode { + return fmt.Errorf("expected generated root mapping node") + } + + // Remove deprecated sections before merging back the sanitized config. + removeLegacyAuthBlock(original.Content[0]) + removeLegacyOpenAICompatAPIKeys(original.Content[0]) + removeRemovedIntegrationKeys(original.Content[0]) + removeLegacyGenerativeLanguageKeys(original.Content[0]) + + pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-excluded-models") + pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-model-alias") + pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "plugins", "configs") + + // Merge generated into original in-place, preserving comments/order of existing nodes. + mergeMappingPreserve(original.Content[0], generated.Content[0]) + normalizeCollectionNodeStyles(original.Content[0]) + + // Write back. + f, err := os.Create(configFile) + if err != nil { + return err + } + defer func() { _ = f.Close() }() + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(2) + if err = enc.Encode(&original); err != nil { + _ = enc.Close() + return err + } + if err = enc.Close(); err != nil { + return err + } + data = NormalizeCommentIndentation(buf.Bytes()) + _, err = f.Write(data) + return err +} + +// SaveConfigPreserveCommentsUpdateNestedScalar updates a nested scalar key path like ["a","b"] +// while preserving comments and positions. +func SaveConfigPreserveCommentsUpdateNestedScalar(configFile string, path []string, value string) error { + data, err := os.ReadFile(configFile) + if err != nil { + return err + } + var root yaml.Node + if err = yaml.Unmarshal(data, &root); err != nil { + return err + } + if root.Kind != yaml.DocumentNode || len(root.Content) == 0 { + return fmt.Errorf("invalid yaml document structure") + } + node := root.Content[0] + // descend mapping nodes following path + for i, key := range path { + if i == len(path)-1 { + // set final scalar + v := getOrCreateMapValue(node, key) + v.Kind = yaml.ScalarNode + v.Tag = "!!str" + v.Value = value + } else { + next := getOrCreateMapValue(node, key) + if next.Kind != yaml.MappingNode { + next.Kind = yaml.MappingNode + next.Tag = "!!map" + } + node = next + } + } + f, err := os.Create(configFile) + if err != nil { + return err + } + defer func() { _ = f.Close() }() + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(2) + if err = enc.Encode(&root); err != nil { + _ = enc.Close() + return err + } + if err = enc.Close(); err != nil { + return err + } + data = NormalizeCommentIndentation(buf.Bytes()) + _, err = f.Write(data) + return err +} + +// NormalizeCommentIndentation removes indentation from standalone YAML comment lines to keep them left aligned. +func NormalizeCommentIndentation(data []byte) []byte { + lines := bytes.Split(data, []byte("\n")) + changed := false + for i, line := range lines { + trimmed := bytes.TrimLeft(line, " \t") + if len(trimmed) == 0 || trimmed[0] != '#' { + continue + } + if len(trimmed) == len(line) { + continue + } + lines[i] = append([]byte(nil), trimmed...) + changed = true + } + if !changed { + return data + } + return bytes.Join(lines, []byte("\n")) +} + +// getOrCreateMapValue finds the value node for a given key in a mapping node. +// If not found, it appends a new key/value pair and returns the new value node. +func getOrCreateMapValue(mapNode *yaml.Node, key string) *yaml.Node { + if mapNode.Kind != yaml.MappingNode { + mapNode.Kind = yaml.MappingNode + mapNode.Tag = "!!map" + mapNode.Content = nil + } + for i := 0; i+1 < len(mapNode.Content); i += 2 { + k := mapNode.Content[i] + if k.Value == key { + return mapNode.Content[i+1] + } + } + // append new key/value + mapNode.Content = append(mapNode.Content, &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}) + val := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: ""} + mapNode.Content = append(mapNode.Content, val) + return val +} + +// mergeMappingPreserve merges keys from src into dst mapping node while preserving +// key order and comments of existing keys in dst. New keys are only added if their +// value is non-zero and not a known default to avoid polluting the config with defaults. +func mergeMappingPreserve(dst, src *yaml.Node, path ...[]string) { + var currentPath []string + if len(path) > 0 { + currentPath = path[0] + } + + if dst == nil || src == nil { + return + } + if dst.Kind != yaml.MappingNode || src.Kind != yaml.MappingNode { + // If kinds do not match, prefer replacing dst with src semantics in-place + // but keep dst node object to preserve any attached comments at the parent level. + copyNodeShallow(dst, src) + return + } + for i := 0; i+1 < len(src.Content); i += 2 { + sk := src.Content[i] + sv := src.Content[i+1] + idx := findMapKeyIndex(dst, sk.Value) + childPath := appendPath(currentPath, sk.Value) + if idx >= 0 { + // Merge into existing value node (always update, even to zero values) + dv := dst.Content[idx+1] + mergeNodePreserve(dv, sv, childPath) + } else { + // New key: only add if value is non-zero and not a known default + candidate := deepCopyNode(sv) + pruneKnownDefaultsInNewNode(childPath, candidate) + if isKnownDefaultValue(childPath, candidate) { + continue + } + dst.Content = append(dst.Content, deepCopyNode(sk), candidate) + } + } +} + +// mergeNodePreserve merges src into dst for scalars, mappings and sequences while +// reusing destination nodes to keep comments and anchors. For sequences, it updates +// in-place by index. +func mergeNodePreserve(dst, src *yaml.Node, path ...[]string) { + var currentPath []string + if len(path) > 0 { + currentPath = path[0] + } + + if dst == nil || src == nil { + return + } + switch src.Kind { + case yaml.MappingNode: + if dst.Kind != yaml.MappingNode { + copyNodeShallow(dst, src) + } + mergeMappingPreserve(dst, src, currentPath) + case yaml.SequenceNode: + // Preserve explicit null style if dst was null and src is empty sequence + if dst.Kind == yaml.ScalarNode && dst.Tag == "!!null" && len(src.Content) == 0 { + // Keep as null to preserve original style + return + } + if dst.Kind != yaml.SequenceNode { + dst.Kind = yaml.SequenceNode + dst.Tag = "!!seq" + dst.Content = nil + } + reorderSequenceForMerge(dst, src) + // Update elements in place + minContent := len(dst.Content) + if len(src.Content) < minContent { + minContent = len(src.Content) + } + for i := 0; i < minContent; i++ { + if dst.Content[i] == nil { + dst.Content[i] = deepCopyNode(src.Content[i]) + continue + } + mergeNodePreserve(dst.Content[i], src.Content[i], currentPath) + if dst.Content[i] != nil && src.Content[i] != nil && + dst.Content[i].Kind == yaml.MappingNode && src.Content[i].Kind == yaml.MappingNode { + pruneMissingMapKeys(dst.Content[i], src.Content[i]) + } + } + // Append any extra items from src + for i := len(dst.Content); i < len(src.Content); i++ { + dst.Content = append(dst.Content, deepCopyNode(src.Content[i])) + } + // Truncate if dst has extra items not in src + if len(src.Content) < len(dst.Content) { + dst.Content = dst.Content[:len(src.Content)] + } + case yaml.ScalarNode, yaml.AliasNode: + // For scalars, update Tag and Value but keep Style from dst to preserve quoting + dst.Kind = src.Kind + dst.Tag = src.Tag + dst.Value = src.Value + // Keep dst.Style as-is intentionally + case 0: + // Unknown/empty kind; do nothing + default: + // Fallback: replace shallowly + copyNodeShallow(dst, src) + } +} + +// findMapKeyIndex returns the index of key node in dst mapping (index of key, not value). +// Returns -1 when not found. +func findMapKeyIndex(mapNode *yaml.Node, key string) int { + if mapNode == nil || mapNode.Kind != yaml.MappingNode { + return -1 + } + for i := 0; i+1 < len(mapNode.Content); i += 2 { + if mapNode.Content[i] != nil && mapNode.Content[i].Value == key { + return i + } + } + return -1 +} + +// appendPath appends a key to the path, returning a new slice to avoid modifying the original. +func appendPath(path []string, key string) []string { + if len(path) == 0 { + return []string{key} + } + newPath := make([]string, len(path)+1) + copy(newPath, path) + newPath[len(path)] = key + return newPath +} + +// isKnownDefaultValue returns true if the given node at the specified path +// represents a known default value that should not be written to the config file. +// This prevents non-zero defaults from polluting the config. +func isKnownDefaultValue(path []string, node *yaml.Node) bool { + // First check if it's a zero value + if isZeroValueNode(node) { + return true + } + + // Match known non-zero defaults by exact dotted path. + if len(path) == 0 { + return false + } + + fullPath := strings.Join(path, ".") + + // Check string defaults + if node.Kind == yaml.ScalarNode && node.Tag == "!!str" { + switch fullPath { + case "pprof.addr": + return node.Value == DefaultPprofAddr + case "remote-management.panel-github-repository": + return node.Value == DefaultPanelGitHubRepository + case "plugins.dir": + return node.Value == "plugins" + case "routing.strategy": + return node.Value == "round-robin" + } + } + + // Check integer defaults + if node.Kind == yaml.ScalarNode && node.Tag == "!!int" { + switch fullPath { + case "error-logs-max-files": + return node.Value == "10" + } + } + + return false +} + +// pruneKnownDefaultsInNewNode removes default-valued descendants from a new node +// before it is appended into the destination YAML tree. +func pruneKnownDefaultsInNewNode(path []string, node *yaml.Node) { + if node == nil { + return + } + + switch node.Kind { + case yaml.MappingNode: + filtered := make([]*yaml.Node, 0, len(node.Content)) + for i := 0; i+1 < len(node.Content); i += 2 { + keyNode := node.Content[i] + valueNode := node.Content[i+1] + if keyNode == nil || valueNode == nil { + continue + } + + childPath := appendPath(path, keyNode.Value) + if isKnownDefaultValue(childPath, valueNode) { + continue + } + + pruneKnownDefaultsInNewNode(childPath, valueNode) + if (valueNode.Kind == yaml.MappingNode || valueNode.Kind == yaml.SequenceNode) && + len(valueNode.Content) == 0 { + continue + } + + filtered = append(filtered, keyNode, valueNode) + } + node.Content = filtered + case yaml.SequenceNode: + for _, child := range node.Content { + pruneKnownDefaultsInNewNode(path, child) + } + } +} + +// isZeroValueNode returns true if the YAML node represents a zero/default value +// that should not be written as a new key to preserve config cleanliness. +// For mappings and sequences, recursively checks if all children are zero values. +func isZeroValueNode(node *yaml.Node) bool { + if node == nil { + return true + } + switch node.Kind { + case yaml.ScalarNode: + switch node.Tag { + case "!!bool": + return node.Value == "false" + case "!!int", "!!float": + return node.Value == "0" || node.Value == "0.0" + case "!!str": + return node.Value == "" + case "!!null": + return true + } + case yaml.SequenceNode: + if len(node.Content) == 0 { + return true + } + // Check if all elements are zero values + for _, child := range node.Content { + if !isZeroValueNode(child) { + return false + } + } + return true + case yaml.MappingNode: + if len(node.Content) == 0 { + return true + } + // Check if all values are zero values (values are at odd indices) + for i := 1; i < len(node.Content); i += 2 { + if !isZeroValueNode(node.Content[i]) { + return false + } + } + return true + } + return false +} + +// deepCopyNode creates a deep copy of a yaml.Node graph. +func deepCopyNode(n *yaml.Node) *yaml.Node { + return deepCopyNodeSeen(n, map[*yaml.Node]*yaml.Node{}) +} + +func deepCopyNodeSeen(n *yaml.Node, seen map[*yaml.Node]*yaml.Node) *yaml.Node { + if n == nil { + return nil + } + if cp, ok := seen[n]; ok { + return cp + } + cp := *n + seen[n] = &cp + if n.Alias != nil { + cp.Alias = deepCopyNodeSeen(n.Alias, seen) + } + if len(n.Content) > 0 { + cp.Content = make([]*yaml.Node, len(n.Content)) + for i := range n.Content { + cp.Content[i] = deepCopyNodeSeen(n.Content[i], seen) + } + } + return &cp +} + +// copyNodeShallow copies type/tag/value and resets content to match src, but +// keeps the same destination node pointer to preserve parent relations/comments. +func copyNodeShallow(dst, src *yaml.Node) { + if dst == nil || src == nil { + return + } + dst.Kind = src.Kind + dst.Tag = src.Tag + dst.Value = src.Value + // Replace content with deep copy from src + if len(src.Content) > 0 { + dst.Content = make([]*yaml.Node, len(src.Content)) + for i := range src.Content { + dst.Content[i] = deepCopyNode(src.Content[i]) + } + } else { + dst.Content = nil + } +} + +func reorderSequenceForMerge(dst, src *yaml.Node) { + if dst == nil || src == nil { + return + } + if len(dst.Content) == 0 { + return + } + if len(src.Content) == 0 { + return + } + original := append([]*yaml.Node(nil), dst.Content...) + used := make([]bool, len(original)) + ordered := make([]*yaml.Node, len(src.Content)) + for i := range src.Content { + if idx := matchSequenceElement(original, used, src.Content[i]); idx >= 0 { + ordered[i] = original[idx] + used[idx] = true + } + } + dst.Content = ordered +} + +func matchSequenceElement(original []*yaml.Node, used []bool, target *yaml.Node) int { + if target == nil { + return -1 + } + switch target.Kind { + case yaml.MappingNode: + id := sequenceElementIdentity(target) + if id != "" { + for i := range original { + if used[i] || original[i] == nil || original[i].Kind != yaml.MappingNode { + continue + } + if sequenceElementIdentity(original[i]) == id { + return i + } + } + } + case yaml.ScalarNode: + val := strings.TrimSpace(target.Value) + if val != "" { + for i := range original { + if used[i] || original[i] == nil || original[i].Kind != yaml.ScalarNode { + continue + } + if strings.TrimSpace(original[i].Value) == val { + return i + } + } + } + default: + } + // Fallback to structural equality to preserve nodes lacking explicit identifiers. + for i := range original { + if used[i] || original[i] == nil { + continue + } + if nodesStructurallyEqual(original[i], target) { + return i + } + } + return -1 +} + +func sequenceElementIdentity(node *yaml.Node) string { + if node == nil || node.Kind != yaml.MappingNode { + return "" + } + identityKeys := []string{"id", "name", "alias", "api-key", "api_key", "apikey", "key", "provider", "model"} + for _, k := range identityKeys { + if v := mappingScalarValue(node, k); v != "" { + return k + "=" + v + } + } + for i := 0; i+1 < len(node.Content); i += 2 { + keyNode := node.Content[i] + valNode := node.Content[i+1] + if keyNode == nil || valNode == nil || valNode.Kind != yaml.ScalarNode { + continue + } + val := strings.TrimSpace(valNode.Value) + if val != "" { + return strings.ToLower(strings.TrimSpace(keyNode.Value)) + "=" + val + } + } + return "" +} + +func mappingScalarValue(node *yaml.Node, key string) string { + if node == nil || node.Kind != yaml.MappingNode { + return "" + } + lowerKey := strings.ToLower(key) + for i := 0; i+1 < len(node.Content); i += 2 { + keyNode := node.Content[i] + valNode := node.Content[i+1] + if keyNode == nil || valNode == nil || valNode.Kind != yaml.ScalarNode { + continue + } + if strings.ToLower(strings.TrimSpace(keyNode.Value)) == lowerKey { + return strings.TrimSpace(valNode.Value) + } + } + return "" +} + +func nodesStructurallyEqual(a, b *yaml.Node) bool { + if a == nil || b == nil { + return a == b + } + if a.Kind != b.Kind { + return false + } + switch a.Kind { + case yaml.MappingNode: + if len(a.Content) != len(b.Content) { + return false + } + for i := 0; i+1 < len(a.Content); i += 2 { + if !nodesStructurallyEqual(a.Content[i], b.Content[i]) { + return false + } + if !nodesStructurallyEqual(a.Content[i+1], b.Content[i+1]) { + return false + } + } + return true + case yaml.SequenceNode: + if len(a.Content) != len(b.Content) { + return false + } + for i := range a.Content { + if !nodesStructurallyEqual(a.Content[i], b.Content[i]) { + return false + } + } + return true + case yaml.ScalarNode: + return strings.TrimSpace(a.Value) == strings.TrimSpace(b.Value) + case yaml.AliasNode: + return nodesStructurallyEqual(a.Alias, b.Alias) + default: + return strings.TrimSpace(a.Value) == strings.TrimSpace(b.Value) + } +} + +func removeMapKey(mapNode *yaml.Node, key string) { + if mapNode == nil || mapNode.Kind != yaml.MappingNode || key == "" { + return + } + for i := 0; i+1 < len(mapNode.Content); i += 2 { + if mapNode.Content[i] != nil && mapNode.Content[i].Value == key { + mapNode.Content = append(mapNode.Content[:i], mapNode.Content[i+2:]...) + return + } + } +} + +func pruneMappingToGeneratedKeys(dstRoot, srcRoot *yaml.Node, keyPath ...string) { + if len(keyPath) == 0 || dstRoot == nil || srcRoot == nil { + return + } + if len(keyPath) > 1 { + dstParent := dstRoot + srcParent := srcRoot + for _, key := range keyPath[:len(keyPath)-1] { + if key == "" || dstParent == nil || dstParent.Kind != yaml.MappingNode { + return + } + dstIdx := findMapKeyIndex(dstParent, key) + if dstIdx < 0 || dstIdx+1 >= len(dstParent.Content) { + return + } + dstParent = dstParent.Content[dstIdx+1] + + if srcParent != nil && srcParent.Kind == yaml.MappingNode { + srcIdx := findMapKeyIndex(srcParent, key) + if srcIdx >= 0 && srcIdx+1 < len(srcParent.Content) { + srcParent = srcParent.Content[srcIdx+1] + } else { + srcParent = nil + } + } + } + if srcParent == nil || srcParent.Kind != yaml.MappingNode { + removeMapKey(dstParent, keyPath[len(keyPath)-1]) + return + } + pruneMappingToGeneratedKeys(dstParent, srcParent, keyPath[len(keyPath)-1]) + return + } + key := keyPath[0] + if key == "" { + return + } + if dstRoot.Kind != yaml.MappingNode || srcRoot.Kind != yaml.MappingNode { + return + } + dstIdx := findMapKeyIndex(dstRoot, key) + if dstIdx < 0 || dstIdx+1 >= len(dstRoot.Content) { + return + } + srcIdx := findMapKeyIndex(srcRoot, key) + if srcIdx < 0 { + // Keep an explicit empty mapping for oauth-model-alias when it was previously present. + // When users delete the last channel from oauth-model-alias via the management API, + // we want that deletion to persist across hot reloads and restarts. + if key == "oauth-model-alias" { + dstRoot.Content[dstIdx+1] = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + return + } + removeMapKey(dstRoot, key) + return + } + if srcIdx+1 >= len(srcRoot.Content) { + return + } + srcVal := srcRoot.Content[srcIdx+1] + dstVal := dstRoot.Content[dstIdx+1] + if srcVal == nil { + dstRoot.Content[dstIdx+1] = nil + return + } + if srcVal.Kind != yaml.MappingNode { + dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal) + return + } + if dstVal == nil || dstVal.Kind != yaml.MappingNode { + dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal) + return + } + pruneMissingMapKeys(dstVal, srcVal) +} + +func pruneMissingMapKeys(dstMap, srcMap *yaml.Node) { + if dstMap == nil || srcMap == nil || dstMap.Kind != yaml.MappingNode || srcMap.Kind != yaml.MappingNode { + return + } + keep := make(map[string]struct{}, len(srcMap.Content)/2) + for i := 0; i+1 < len(srcMap.Content); i += 2 { + keyNode := srcMap.Content[i] + if keyNode == nil { + continue + } + key := strings.TrimSpace(keyNode.Value) + if key == "" { + continue + } + keep[key] = struct{}{} + } + for i := 0; i+1 < len(dstMap.Content); { + keyNode := dstMap.Content[i] + if keyNode == nil { + i += 2 + continue + } + key := strings.TrimSpace(keyNode.Value) + if _, ok := keep[key]; !ok { + dstMap.Content = append(dstMap.Content[:i], dstMap.Content[i+2:]...) + continue + } + i += 2 + } +} + +// normalizeCollectionNodeStyles forces YAML collections to use block notation, keeping +// lists and maps readable. Empty sequences retain flow style ([]) so empty list markers +// remain compact. +func normalizeCollectionNodeStyles(node *yaml.Node) { + if node == nil { + return + } + switch node.Kind { + case yaml.MappingNode: + node.Style = 0 + for i := range node.Content { + normalizeCollectionNodeStyles(node.Content[i]) + } + case yaml.SequenceNode: + if len(node.Content) == 0 { + node.Style = yaml.FlowStyle + } else { + node.Style = 0 + } + for i := range node.Content { + normalizeCollectionNodeStyles(node.Content[i]) + } + default: + // Scalars keep their existing style to preserve quoting + } +} + +func removeLegacyOpenAICompatAPIKeys(root *yaml.Node) { + if root == nil || root.Kind != yaml.MappingNode { + return + } + idx := findMapKeyIndex(root, "openai-compatibility") + if idx < 0 || idx+1 >= len(root.Content) { + return + } + seq := root.Content[idx+1] + if seq == nil || seq.Kind != yaml.SequenceNode { + return + } + for i := range seq.Content { + if seq.Content[i] != nil && seq.Content[i].Kind == yaml.MappingNode { + removeMapKey(seq.Content[i], "api-keys") + } + } +} + +func removeRemovedIntegrationKeys(root *yaml.Node) { + if root == nil || root.Kind != yaml.MappingNode { + return + } + removeMapKey(root, "ampcode") + removeMapKey(root, "amp-upstream-url") + removeMapKey(root, "amp-upstream-api-key") + removeMapKey(root, "amp-restrict-management-to-localhost") + removeMapKey(root, "amp-model-mappings") +} + +func removeLegacyGenerativeLanguageKeys(root *yaml.Node) { + if root == nil || root.Kind != yaml.MappingNode { + return + } + removeMapKey(root, "generative-language-api-key") +} + +func removeLegacyAuthBlock(root *yaml.Node) { + if root == nil || root.Kind != yaml.MappingNode { + return + } + removeMapKey(root, "auth") +} diff --git a/internal/logging/request_logger.go b/internal/logging/request_logger.go index 46fd220f7..8a51f9455 100644 --- a/internal/logging/request_logger.go +++ b/internal/logging/request_logger.go @@ -4,35 +4,13 @@ package logging import ( - "bufio" - "bytes" - "compress/flate" - "compress/gzip" - "context" - "encoding/json" "fmt" - "io" - "os" "path/filepath" - "regexp" - "sort" - "strings" - "sync" - "sync/atomic" "time" - "github.com/andybalholm/brotli" - "github.com/klauspost/compress/zstd" - log "github.com/sirupsen/logrus" - - "github.com/router-for-me/CLIProxyAPI/v7/internal/buildinfo" - "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" ) -var requestLogID atomic.Uint64 - const ( WebsocketTimelineSourceContextKey = "WEBSOCKET_TIMELINE_SOURCE" APIRequestSourceContextKey = "API_REQUEST_SOURCE" @@ -45,259 +23,6 @@ const ( // DeferredAPIRequest builds an upstream request log only when an error log needs it. type DeferredAPIRequest func() []byte -type homeRequestLogClient interface { - HeartbeatOK() bool - RPushRequestLog(ctx context.Context, payload []byte) error -} - -var currentHomeRequestLogClient = func() homeRequestLogClient { - return home.Current() -} - -// FileBodySource stores large log sections as ordered temp-file parts. -type FileBodySource struct { - mu sync.Mutex - dir string - paths []string - cleaned bool -} - -// NewFileBodySourceInDir creates a temp-backed source under baseDir. -func NewFileBodySourceInDir(baseDir string, prefix string) (*FileBodySource, error) { - prefix = sanitizeTempPrefix(prefix) - baseDir = strings.TrimSpace(baseDir) - if baseDir == "" { - return nil, fmt.Errorf("base directory is required") - } - if errMkdir := os.MkdirAll(baseDir, 0755); errMkdir != nil { - return nil, errMkdir - } - dir, errCreate := os.MkdirTemp(baseDir, "request-log-parts-"+prefix+"-*") - if errCreate != nil { - return nil, errCreate - } - return &FileBodySource{dir: dir}, nil -} - -func sanitizeTempPrefix(prefix string) string { - prefix = strings.TrimSpace(prefix) - if prefix == "" { - return "log" - } - var builder strings.Builder - for _, r := range prefix { - switch { - case r >= 'a' && r <= 'z': - builder.WriteRune(r) - case r >= 'A' && r <= 'Z': - builder.WriteRune(r) - case r >= '0' && r <= '9': - builder.WriteRune(r) - case r == '-' || r == '_': - builder.WriteRune(r) - default: - builder.WriteByte('-') - } - } - out := strings.Trim(builder.String(), "-_") - if out == "" { - return "log" - } - return out -} - -// CreatePart creates one ordered detail log part. -func (s *FileBodySource) CreatePart(prefix string) (*os.File, error) { - if s == nil { - return nil, fmt.Errorf("file body source is nil") - } - s.mu.Lock() - defer s.mu.Unlock() - if s.cleaned { - return nil, fmt.Errorf("file body source has been cleaned") - } - prefix = sanitizeTempPrefix(prefix) - if errMkdir := os.MkdirAll(s.dir, 0755); errMkdir != nil { - return nil, errMkdir - } - file, errCreate := os.CreateTemp(s.dir, prefix+"-*.tmp") - if errCreate != nil { - return nil, errCreate - } - s.paths = append(s.paths, file.Name()) - return file, nil -} - -// AppendPart appends one complete ordered part to the source. -func (s *FileBodySource) AppendPart(data []byte) error { - data = bytes.TrimSpace(data) - if len(data) == 0 { - return nil - } - file, errCreate := s.CreatePart("part") - if errCreate != nil { - return errCreate - } - writeErr := writeLogPart(file, data, false) - if errClose := file.Close(); errClose != nil { - if writeErr == nil { - writeErr = errClose - } - } - return writeErr -} - -// AppendBytes appends raw bytes to a single ordered part. -func (s *FileBodySource) AppendBytes(data []byte) error { - if s == nil { - return fmt.Errorf("file body source is nil") - } - if len(data) == 0 { - return nil - } - s.mu.Lock() - defer s.mu.Unlock() - if s.cleaned { - return fmt.Errorf("file body source has been cleaned") - } - if errMkdir := os.MkdirAll(s.dir, 0755); errMkdir != nil { - return errMkdir - } - - var file *os.File - var errOpen error - if len(s.paths) == 0 { - file, errOpen = os.CreateTemp(s.dir, "part-*.tmp") - if errOpen == nil { - s.paths = append(s.paths, file.Name()) - } - } else { - file, errOpen = os.OpenFile(s.paths[len(s.paths)-1], os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) - } - if errOpen != nil { - return errOpen - } - - _, writeErr := file.Write(data) - if errClose := file.Close(); errClose != nil { - if writeErr == nil { - writeErr = errClose - } - } - return writeErr -} - -// HasPayload reports whether any detail parts were recorded. -func (s *FileBodySource) HasPayload() bool { - if s == nil { - return false - } - s.mu.Lock() - defer s.mu.Unlock() - return len(s.paths) > 0 && !s.cleaned -} - -// Paths returns a copy of the ordered part paths. -func (s *FileBodySource) Paths() []string { - if s == nil { - return nil - } - s.mu.Lock() - defer s.mu.Unlock() - out := make([]string, len(s.paths)) - copy(out, s.paths) - return out -} - -// WriteTo merges all ordered parts into w. -func (s *FileBodySource) WriteTo(w io.Writer) error { - if s == nil || w == nil { - return nil - } - paths := s.Paths() - wrote := false - for _, path := range paths { - file, errOpen := os.Open(path) - if errOpen != nil { - if os.IsNotExist(errOpen) { - continue - } - return errOpen - } - if wrote { - if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { - if errClose := file.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close log part file") - } - return errWrite - } - } - _, errCopy := io.Copy(w, file) - if errClose := file.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close log part file") - if errCopy == nil { - errCopy = errClose - } - } - if errCopy != nil { - return errCopy - } - wrote = true - } - return nil -} - -// Bytes merges all ordered parts into memory. -func (s *FileBodySource) Bytes() ([]byte, error) { - var buf bytes.Buffer - if errWrite := s.WriteTo(&buf); errWrite != nil { - return nil, errWrite - } - return buf.Bytes(), nil -} - -// Cleanup removes all temp detail parts and their directory. -func (s *FileBodySource) Cleanup() error { - if s == nil { - return nil - } - s.mu.Lock() - if s.cleaned { - s.mu.Unlock() - return nil - } - paths := make([]string, len(s.paths)) - copy(paths, s.paths) - dir := s.dir - s.paths = nil - s.cleaned = true - s.mu.Unlock() - - var firstErr error - for _, path := range paths { - if errRemove := os.Remove(path); errRemove != nil && !os.IsNotExist(errRemove) && firstErr == nil { - firstErr = errRemove - } - } - if dir != "" { - if errRemove := os.RemoveAll(dir); errRemove != nil && firstErr == nil { - firstErr = errRemove - } - } - return firstErr -} - -func cleanupFileBodySources(sources ...*FileBodySource) { - for _, source := range sources { - if source == nil { - continue - } - if errCleanup := source.Cleanup(); errCleanup != nil { - log.WithError(errCleanup).Warn("failed to clean up log part files") - } - } -} - // RequestLogger defines the interface for logging HTTP requests and responses. // It provides methods for logging both regular and streaming HTTP request/response cycles. type RequestLogger interface { @@ -421,58 +146,6 @@ type FileRequestLogger struct { homeEnabled bool } -type homeRequestLogPayload struct { - Headers map[string][]string `json:"headers,omitempty"` - RequestID string `json:"request_id,omitempty"` - RequestLog string `json:"request_log,omitempty"` -} - -func cloneHeaders(headers map[string][]string) map[string][]string { - if len(headers) == 0 { - return nil - } - out := make(map[string][]string, len(headers)) - for key, values := range headers { - if strings.TrimSpace(key) == "" { - continue - } - if values == nil { - out[key] = nil - continue - } - copied := make([]string, len(values)) - copy(copied, values) - out[key] = copied - } - if len(out) == 0 { - return nil - } - return out -} - -func (l *FileRequestLogger) forwardRequestLogToHome(ctx context.Context, headers map[string][]string, requestID string, logText string) error { - if l == nil || !l.homeEnabled { - return nil - } - client := currentHomeRequestLogClient() - if client == nil || !client.HeartbeatOK() { - return nil - } - payload := homeRequestLogPayload{ - Headers: cloneHeaders(headers), - RequestID: strings.TrimSpace(requestID), - RequestLog: logText, - } - raw, errMarshal := json.Marshal(&payload) - if errMarshal != nil { - return errMarshal - } - if ctx == nil { - ctx = context.Background() - } - return client.RPushRequestLog(ctx, raw) -} - // NewFileRequestLogger creates a new file-based request logger. // // Parameters: @@ -500,15 +173,6 @@ func NewFileRequestLogger(enabled bool, logsDir string, configDir string, errorL } } -// SetHomeEnabled toggles home request-log forwarding. -// When enabled, request logs are not written to disk and are instead forwarded to home via Redis RESP. -func (l *FileRequestLogger) SetHomeEnabled(enabled bool) { - if l == nil { - return - } - l.homeEnabled = enabled -} - // IsEnabled returns whether request logging is currently enabled. // // Returns: @@ -541,1631 +205,3 @@ func (l *FileRequestLogger) NewFileBodySource(prefix string) (*FileBodySource, e } return NewFileBodySourceInDir(l.logsDir, prefix) } - -// LogRequest logs a complete non-streaming request/response cycle to a file. -// -// Parameters: -// - url: The request URL -// - method: The HTTP method -// - requestHeaders: The request headers -// - body: The request body -// - statusCode: The response status code -// - responseHeaders: The response headers -// - response: The raw response data -// - apiRequest: The API request data -// - apiResponse: The API response data -// - requestID: Optional request ID for log file naming -// - requestTimestamp: When the request was received -// - apiResponseTimestamp: When the API response was received -// -// Returns: -// - error: An error if logging fails, nil otherwise -func (l *FileRequestLogger) LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { - return l.logRequest(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline, apiResponseErrors, false, requestID, requestTimestamp, apiResponseTimestamp) -} - -// LogRequestWithOptions logs a request with optional forced logging behavior. -// The force flag allows writing error logs even when regular request logging is disabled. -func (l *FileRequestLogger) LogRequestWithOptions(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { - return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, nil, apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, nil, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) -} - -func (l *FileRequestLogger) logRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { - return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, nil, apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, nil, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) -} - -// LogRequestWithOptionsAndSources logs a request with optional file-backed large sections. -func (l *FileRequestLogger) LogRequestWithOptionsAndSources(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline []byte, websocketTimelineSource *FileBodySource, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { - return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, websocketTimelineSource, apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, apiWebsocketTimelineSource, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) -} - -// LogRequestWithOptionsAndAllSources logs a request with optional file-backed request and response sections. -func (l *FileRequestLogger) LogRequestWithOptionsAndAllSources(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline []byte, websocketTimelineSource *FileBodySource, apiRequest []byte, apiRequestSource *FileBodySource, apiResponse []byte, apiResponseSource *FileBodySource, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { - return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, websocketTimelineSource, apiRequest, apiRequestSource, apiResponse, apiResponseSource, apiWebsocketTimeline, apiWebsocketTimelineSource, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) -} - -func (l *FileRequestLogger) logRequestWithSources(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline []byte, websocketTimelineSource *FileBodySource, apiRequest []byte, apiRequestSource *FileBodySource, apiResponse []byte, apiResponseSource *FileBodySource, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { - defer cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource, apiWebsocketTimelineSource) - - if !l.enabled && !force { - return nil - } - - if l.homeEnabled && l.enabled { - responseToWrite, decompressErr := l.decompressResponse(responseHeaders, response) - if decompressErr != nil { - responseToWrite = response - } - - var buf bytes.Buffer - writeErr := l.writeNonStreamingLog( - &buf, - url, - method, - requestHeaders, - body, - "", - websocketTimeline, - websocketTimelineSource, - apiRequest, - apiRequestSource, - apiResponse, - apiResponseSource, - apiWebsocketTimeline, - apiWebsocketTimelineSource, - apiResponseErrors, - statusCode, - responseHeaders, - responseToWrite, - decompressErr, - requestTimestamp, - apiResponseTimestamp, - ) - if writeErr != nil { - return fmt.Errorf("failed to build request log content: %w", writeErr) - } - return l.forwardRequestLogToHome(context.Background(), requestHeaders, requestID, buf.String()) - } - - // Ensure logs directory exists - if errEnsure := l.ensureLogsDir(); errEnsure != nil { - return fmt.Errorf("failed to create logs directory: %w", errEnsure) - } - - // Generate filename with request ID - filename := l.generateFilename(url, requestID) - if force && !l.enabled { - filename = l.generateErrorFilename(url, requestID) - } - filePath := filepath.Join(l.logsDir, filename) - - requestBodyPath, errTemp := l.writeRequestBodyTempFile(body) - if errTemp != nil { - log.WithError(errTemp).Warn("failed to create request body temp file, falling back to direct write") - } - if requestBodyPath != "" { - defer func() { - if errRemove := os.Remove(requestBodyPath); errRemove != nil { - log.WithError(errRemove).Warn("failed to remove request body temp file") - } - }() - } - - responseToWrite, decompressErr := l.decompressResponse(responseHeaders, response) - if decompressErr != nil { - // If decompression fails, continue with original response and annotate the log output. - responseToWrite = response - } - - logFile, errOpen := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) - if errOpen != nil { - return fmt.Errorf("failed to create log file: %w", errOpen) - } - - writeErr := l.writeNonStreamingLog( - logFile, - url, - method, - requestHeaders, - body, - requestBodyPath, - websocketTimeline, - websocketTimelineSource, - apiRequest, - apiRequestSource, - apiResponse, - apiResponseSource, - apiWebsocketTimeline, - apiWebsocketTimelineSource, - apiResponseErrors, - statusCode, - responseHeaders, - responseToWrite, - decompressErr, - requestTimestamp, - apiResponseTimestamp, - ) - if errClose := logFile.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close request log file") - if writeErr == nil { - return errClose - } - } - if writeErr != nil { - return fmt.Errorf("failed to write log file: %w", writeErr) - } - - if force && !l.enabled { - if errCleanup := l.cleanupOldErrorLogs(); errCleanup != nil { - log.WithError(errCleanup).Warn("failed to clean up old error logs") - } - } - - return nil -} - -// LogStreamingRequest initiates logging for a streaming request. -// -// Parameters: -// - url: The request URL -// - method: The HTTP method -// - headers: The request headers -// - body: The request body -// - requestID: Optional request ID for log file naming -// -// Returns: -// - StreamingLogWriter: A writer for streaming response chunks -// - error: An error if logging initialization fails, nil otherwise -func (l *FileRequestLogger) LogStreamingRequest(url, method string, headers map[string][]string, body []byte, requestID string) (StreamingLogWriter, error) { - if !l.enabled { - return &NoOpStreamingLogWriter{}, nil - } - - if l.homeEnabled { - client := currentHomeRequestLogClient() - if client == nil || !client.HeartbeatOK() { - return &NoOpStreamingLogWriter{}, nil - } - return newHomeStreamingLogWriter(url, method, headers, body, requestID), nil - } - - // Ensure logs directory exists - if err := l.ensureLogsDir(); err != nil { - return nil, fmt.Errorf("failed to create logs directory: %w", err) - } - - // Generate filename with request ID - filename := l.generateFilename(url, requestID) - filePath := filepath.Join(l.logsDir, filename) - - requestHeaders := make(map[string][]string, len(headers)) - for key, values := range headers { - headerValues := make([]string, len(values)) - copy(headerValues, values) - requestHeaders[key] = headerValues - } - - requestBodyPath, errTemp := l.writeRequestBodyTempFile(body) - if errTemp != nil { - return nil, fmt.Errorf("failed to create request body temp file: %w", errTemp) - } - - responseBodyFile, errCreate := os.CreateTemp(l.logsDir, "response-body-*.tmp") - if errCreate != nil { - _ = os.Remove(requestBodyPath) - return nil, fmt.Errorf("failed to create response body temp file: %w", errCreate) - } - responseBodyPath := responseBodyFile.Name() - - // Create streaming writer - writer := &FileStreamingLogWriter{ - logFilePath: filePath, - url: url, - method: method, - timestamp: time.Now(), - requestHeaders: requestHeaders, - requestBodyPath: requestBodyPath, - responseBodyPath: responseBodyPath, - responseBodyFile: responseBodyFile, - chunkChan: make(chan []byte, 100), // Buffered channel for async writes - closeChan: make(chan struct{}), - errorChan: make(chan error, 1), - } - - // Start async writer goroutine - go writer.asyncWriter() - - return writer, nil -} - -// generateErrorFilename creates a filename with an error prefix to differentiate forced error logs. -func (l *FileRequestLogger) generateErrorFilename(url string, requestID ...string) string { - return fmt.Sprintf("error-%s", l.generateFilename(url, requestID...)) -} - -// ensureLogsDir creates the logs directory if it doesn't exist. -// -// Returns: -// - error: An error if directory creation fails, nil otherwise -func (l *FileRequestLogger) ensureLogsDir() error { - if _, err := os.Stat(l.logsDir); os.IsNotExist(err) { - return os.MkdirAll(l.logsDir, 0755) - } - return nil -} - -// generateFilename creates a sanitized filename from the URL path and current timestamp. -// Format: v1-responses-2025-12-23T195811-a1b2c3d4.log -// -// Parameters: -// - url: The request URL -// - requestID: Optional request ID to include in filename -// -// Returns: -// - string: A sanitized filename for the log file -func (l *FileRequestLogger) generateFilename(url string, requestID ...string) string { - // Extract path from URL - path := url - if strings.Contains(url, "?") { - path = strings.Split(url, "?")[0] - } - - // Remove leading slash - if strings.HasPrefix(path, "/") { - path = path[1:] - } - - // Sanitize path for filename - sanitized := l.sanitizeForFilename(path) - - // Add timestamp - timestamp := time.Now().Format("2006-01-02T150405") - - // Use request ID if provided, otherwise use sequential ID - var idPart string - if len(requestID) > 0 && requestID[0] != "" { - idPart = requestID[0] - } else { - id := requestLogID.Add(1) - idPart = fmt.Sprintf("%d", id) - } - - return fmt.Sprintf("%s-%s-%s.log", sanitized, timestamp, idPart) -} - -// sanitizeForFilename replaces characters that are not safe for filenames. -// -// Parameters: -// - path: The path to sanitize -// -// Returns: -// - string: A sanitized filename -func (l *FileRequestLogger) sanitizeForFilename(path string) string { - // Replace slashes with hyphens - sanitized := strings.ReplaceAll(path, "/", "-") - - // Replace colons with hyphens - sanitized = strings.ReplaceAll(sanitized, ":", "-") - - // Replace other problematic characters with hyphens - reg := regexp.MustCompile(`[<>:"|?*\s]`) - sanitized = reg.ReplaceAllString(sanitized, "-") - - // Remove multiple consecutive hyphens - reg = regexp.MustCompile(`-+`) - sanitized = reg.ReplaceAllString(sanitized, "-") - - // Remove leading/trailing hyphens - sanitized = strings.Trim(sanitized, "-") - - // Handle empty result - if sanitized == "" { - sanitized = "root" - } - - return sanitized -} - -// cleanupOldErrorLogs keeps only the newest errorLogsMaxFiles forced error log files. -func (l *FileRequestLogger) cleanupOldErrorLogs() error { - if l.errorLogsMaxFiles <= 0 { - return nil - } - - entries, errRead := os.ReadDir(l.logsDir) - if errRead != nil { - return errRead - } - - type logFile struct { - name string - modTime time.Time - } - - var files []logFile - for _, entry := range entries { - if entry.IsDir() { - continue - } - name := entry.Name() - if !strings.HasPrefix(name, "error-") || !strings.HasSuffix(name, ".log") { - continue - } - info, errInfo := entry.Info() - if errInfo != nil { - log.WithError(errInfo).Warn("failed to read error log info") - continue - } - files = append(files, logFile{name: name, modTime: info.ModTime()}) - } - - if len(files) <= l.errorLogsMaxFiles { - return nil - } - - sort.Slice(files, func(i, j int) bool { - return files[i].modTime.After(files[j].modTime) - }) - - for _, file := range files[l.errorLogsMaxFiles:] { - if errRemove := os.Remove(filepath.Join(l.logsDir, file.name)); errRemove != nil { - log.WithError(errRemove).Warnf("failed to remove old error log: %s", file.name) - } - } - - return nil -} - -func (l *FileRequestLogger) writeRequestBodyTempFile(body []byte) (string, error) { - tmpFile, errCreate := os.CreateTemp(l.logsDir, "request-body-*.tmp") - if errCreate != nil { - return "", errCreate - } - tmpPath := tmpFile.Name() - - if _, errCopy := io.Copy(tmpFile, bytes.NewReader(body)); errCopy != nil { - _ = tmpFile.Close() - _ = os.Remove(tmpPath) - return "", errCopy - } - if errClose := tmpFile.Close(); errClose != nil { - _ = os.Remove(tmpPath) - return "", errClose - } - return tmpPath, nil -} - -func (l *FileRequestLogger) writeNonStreamingLog( - w io.Writer, - url, method string, - requestHeaders map[string][]string, - requestBody []byte, - requestBodyPath string, - websocketTimeline []byte, - websocketTimelineSource *FileBodySource, - apiRequest []byte, - apiRequestSource *FileBodySource, - apiResponse []byte, - apiResponseSource *FileBodySource, - apiWebsocketTimeline []byte, - apiWebsocketTimelineSource *FileBodySource, - apiResponseErrors []*interfaces.ErrorMessage, - statusCode int, - responseHeaders map[string][]string, - response []byte, - decompressErr error, - requestTimestamp time.Time, - apiResponseTimestamp time.Time, -) error { - if requestTimestamp.IsZero() { - requestTimestamp = time.Now() - } - isWebsocketTranscript := hasSectionPayload(websocketTimeline) || hasFileBodySourcePayload(websocketTimelineSource) - downstreamTransport := inferDownstreamTransport(requestHeaders, websocketTimeline, websocketTimelineSource) - upstreamTransport := inferUpstreamTransport(apiRequest, apiRequestSource, apiResponse, apiResponseSource, apiWebsocketTimeline, apiWebsocketTimelineSource, apiResponseErrors) - if errWrite := writeRequestInfoWithBody(w, url, method, requestHeaders, requestBody, requestBodyPath, requestTimestamp, downstreamTransport, upstreamTransport, !isWebsocketTranscript); errWrite != nil { - return errWrite - } - if errWrite := writeAPISectionWithSource(w, "=== WEBSOCKET TIMELINE ===\n", "=== WEBSOCKET TIMELINE", websocketTimeline, websocketTimelineSource, time.Time{}); errWrite != nil { - return errWrite - } - if errWrite := writeAPISectionWithSource(w, "=== API WEBSOCKET TIMELINE ===\n", "=== API WEBSOCKET TIMELINE", apiWebsocketTimeline, apiWebsocketTimelineSource, time.Time{}); errWrite != nil { - return errWrite - } - if errWrite := writePreformattedAPISectionWithSource(w, "=== API REQUEST ===\n", "=== API REQUEST", apiRequest, apiRequestSource, time.Time{}); errWrite != nil { - return errWrite - } - if errWrite := writeAPIErrorResponses(w, apiResponseErrors); errWrite != nil { - return errWrite - } - if errWrite := writePreformattedAPISectionWithSource(w, "=== API RESPONSE ===\n", "=== API RESPONSE", apiResponse, apiResponseSource, apiResponseTimestamp); errWrite != nil { - return errWrite - } - if isWebsocketTranscript { - // Intentionally omit the generic downstream HTTP response section for websocket - // transcripts. The durable session exchange is captured in WEBSOCKET TIMELINE, - // and appending a one-off upgrade response snapshot would dilute that transcript. - return nil - } - return writeResponseSection(w, statusCode, true, responseHeaders, bytes.NewReader(response), decompressErr, true) -} - -func writeRequestInfoWithBody( - w io.Writer, - url, method string, - headers map[string][]string, - body []byte, - bodyPath string, - timestamp time.Time, - downstreamTransport string, - upstreamTransport string, - includeBody bool, -) error { - if _, errWrite := io.WriteString(w, "=== REQUEST INFO ===\n"); errWrite != nil { - return errWrite - } - if _, errWrite := io.WriteString(w, fmt.Sprintf("Version: %s\n", buildinfo.Version)); errWrite != nil { - return errWrite - } - if _, errWrite := io.WriteString(w, fmt.Sprintf("URL: %s\n", url)); errWrite != nil { - return errWrite - } - if _, errWrite := io.WriteString(w, fmt.Sprintf("Method: %s\n", method)); errWrite != nil { - return errWrite - } - if strings.TrimSpace(downstreamTransport) != "" { - if _, errWrite := io.WriteString(w, fmt.Sprintf("Downstream Transport: %s\n", downstreamTransport)); errWrite != nil { - return errWrite - } - } - if strings.TrimSpace(upstreamTransport) != "" { - if _, errWrite := io.WriteString(w, fmt.Sprintf("Upstream Transport: %s\n", upstreamTransport)); errWrite != nil { - return errWrite - } - } - if _, errWrite := io.WriteString(w, fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano))); errWrite != nil { - return errWrite - } - if errWrite := writeSectionSpacing(w, 1); errWrite != nil { - return errWrite - } - - if _, errWrite := io.WriteString(w, "=== HEADERS ===\n"); errWrite != nil { - return errWrite - } - for key, values := range headers { - for _, value := range values { - masked := util.MaskSensitiveHeaderValue(key, value) - if _, errWrite := io.WriteString(w, fmt.Sprintf("%s: %s\n", key, masked)); errWrite != nil { - return errWrite - } - } - } - if errWrite := writeSectionSpacing(w, 1); errWrite != nil { - return errWrite - } - - if !includeBody { - return nil - } - - if _, errWrite := io.WriteString(w, "=== REQUEST BODY ===\n"); errWrite != nil { - return errWrite - } - - bodyTrailingNewlines := 1 - if bodyPath != "" { - bodyFile, errOpen := os.Open(bodyPath) - if errOpen != nil { - return errOpen - } - tracker := &trailingNewlineTrackingWriter{writer: w} - written, errCopy := io.Copy(tracker, bodyFile) - if errCopy != nil { - _ = bodyFile.Close() - return errCopy - } - if written > 0 { - bodyTrailingNewlines = tracker.trailingNewlines - } - if errClose := bodyFile.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close request body temp file") - } - } else if _, errWrite := w.Write(body); errWrite != nil { - return errWrite - } else if len(body) > 0 { - bodyTrailingNewlines = countTrailingNewlinesBytes(body) - } - if errWrite := writeSectionSpacing(w, bodyTrailingNewlines); errWrite != nil { - return errWrite - } - return nil -} - -func countTrailingNewlinesBytes(payload []byte) int { - count := 0 - for i := len(payload) - 1; i >= 0; i-- { - if payload[i] != '\n' { - break - } - count++ - } - return count -} - -func writeSectionSpacing(w io.Writer, trailingNewlines int) error { - missingNewlines := 3 - trailingNewlines - if missingNewlines <= 0 { - return nil - } - _, errWrite := io.WriteString(w, strings.Repeat("\n", missingNewlines)) - return errWrite -} - -type trailingNewlineTrackingWriter struct { - writer io.Writer - trailingNewlines int -} - -func (t *trailingNewlineTrackingWriter) Write(payload []byte) (int, error) { - written, errWrite := t.writer.Write(payload) - if written > 0 { - writtenPayload := payload[:written] - trailingNewlines := countTrailingNewlinesBytes(writtenPayload) - if trailingNewlines == len(writtenPayload) { - t.trailingNewlines += trailingNewlines - } else { - t.trailingNewlines = trailingNewlines - } - } - return written, errWrite -} - -func hasSectionPayload(payload []byte) bool { - return len(bytes.TrimSpace(payload)) > 0 -} - -func hasFileBodySourcePayload(source *FileBodySource) bool { - return source != nil && source.HasPayload() -} - -func inferDownstreamTransport(headers map[string][]string, websocketTimeline []byte, websocketTimelineSource *FileBodySource) string { - if hasSectionPayload(websocketTimeline) || hasFileBodySourcePayload(websocketTimelineSource) { - return "websocket" - } - for key, values := range headers { - if strings.EqualFold(strings.TrimSpace(key), "Upgrade") { - for _, value := range values { - if strings.EqualFold(strings.TrimSpace(value), "websocket") { - return "websocket" - } - } - } - } - return "http" -} - -func inferUpstreamTransport(apiRequest []byte, apiRequestSource *FileBodySource, apiResponse []byte, apiResponseSource *FileBodySource, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, _ []*interfaces.ErrorMessage) string { - hasHTTP := hasSectionPayload(apiRequest) || hasFileBodySourcePayload(apiRequestSource) || hasSectionPayload(apiResponse) || hasFileBodySourcePayload(apiResponseSource) - hasWS := hasSectionPayload(apiWebsocketTimeline) || hasFileBodySourcePayload(apiWebsocketTimelineSource) - switch { - case hasHTTP && hasWS: - return "websocket+http" - case hasWS: - return "websocket" - case hasHTTP: - return "http" - default: - return "" - } -} - -func writeLogPart(w io.Writer, payload []byte, prependNewline bool) error { - if w == nil { - return nil - } - if prependNewline { - if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { - return errWrite - } - } - if _, errWrite := w.Write(payload); errWrite != nil { - return errWrite - } - if !bytes.HasSuffix(payload, []byte("\n")) { - if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { - return errWrite - } - } - return nil -} - -func writeAPISection(w io.Writer, sectionHeader string, sectionPrefix string, payload []byte, timestamp time.Time) error { - if len(payload) == 0 { - return nil - } - - if bytes.HasPrefix(payload, []byte(sectionPrefix)) { - if _, errWrite := w.Write(payload); errWrite != nil { - return errWrite - } - } else { - if _, errWrite := io.WriteString(w, sectionHeader); errWrite != nil { - return errWrite - } - if !timestamp.IsZero() { - if _, errWrite := io.WriteString(w, fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano))); errWrite != nil { - return errWrite - } - } - if _, errWrite := w.Write(payload); errWrite != nil { - return errWrite - } - } - - if errWrite := writeSectionSpacing(w, countTrailingNewlinesBytes(payload)); errWrite != nil { - return errWrite - } - return nil -} - -func writeAPISectionWithSource(w io.Writer, sectionHeader string, sectionPrefix string, payload []byte, source *FileBodySource, timestamp time.Time) error { - if !hasFileBodySourcePayload(source) { - return writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp) - } - if len(payload) > 0 { - if errWrite := writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp); errWrite != nil { - return errWrite - } - } - if _, errWrite := io.WriteString(w, sectionHeader); errWrite != nil { - return errWrite - } - if !timestamp.IsZero() { - if _, errWrite := io.WriteString(w, fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano))); errWrite != nil { - return errWrite - } - } - tracker := &trailingNewlineTrackingWriter{writer: w} - if errWrite := source.WriteTo(tracker); errWrite != nil { - return errWrite - } - if errWrite := writeSectionSpacing(w, tracker.trailingNewlines); errWrite != nil { - return errWrite - } - return nil -} - -func writePreformattedAPISectionWithSource(w io.Writer, sectionHeader string, sectionPrefix string, payload []byte, source *FileBodySource, timestamp time.Time) error { - if !hasFileBodySourcePayload(source) { - return writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp) - } - if len(payload) > 0 { - if errWrite := writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp); errWrite != nil { - return errWrite - } - } - tracker := &trailingNewlineTrackingWriter{writer: w} - if errWrite := source.WriteTo(tracker); errWrite != nil { - return errWrite - } - if errWrite := writeSectionSpacing(w, tracker.trailingNewlines); errWrite != nil { - return errWrite - } - return nil -} - -func writeAPIErrorResponses(w io.Writer, apiResponseErrors []*interfaces.ErrorMessage) error { - for i := 0; i < len(apiResponseErrors); i++ { - if apiResponseErrors[i] == nil { - continue - } - if _, errWrite := io.WriteString(w, "=== API ERROR RESPONSE ===\n"); errWrite != nil { - return errWrite - } - if _, errWrite := io.WriteString(w, fmt.Sprintf("HTTP Status: %d\n", apiResponseErrors[i].StatusCode)); errWrite != nil { - return errWrite - } - trailingNewlines := 1 - if apiResponseErrors[i].Error != nil { - errText := apiResponseErrors[i].Error.Error() - if _, errWrite := io.WriteString(w, errText); errWrite != nil { - return errWrite - } - if errText != "" { - trailingNewlines = countTrailingNewlinesBytes([]byte(errText)) - } - } - if errWrite := writeSectionSpacing(w, trailingNewlines); errWrite != nil { - return errWrite - } - } - return nil -} - -func writeResponseSection(w io.Writer, statusCode int, statusWritten bool, responseHeaders map[string][]string, responseReader io.Reader, decompressErr error, trailingNewline bool) error { - if _, errWrite := io.WriteString(w, "=== RESPONSE ===\n"); errWrite != nil { - return errWrite - } - if statusWritten { - if _, errWrite := io.WriteString(w, fmt.Sprintf("Status: %d\n", statusCode)); errWrite != nil { - return errWrite - } - } - - if responseHeaders != nil { - for key, values := range responseHeaders { - for _, value := range values { - if _, errWrite := io.WriteString(w, fmt.Sprintf("%s: %s\n", key, value)); errWrite != nil { - return errWrite - } - } - } - } - - var bufferedReader *bufio.Reader - if responseReader != nil { - bufferedReader = bufio.NewReader(responseReader) - } - if !responseBodyStartsWithLeadingNewline(bufferedReader) { - if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { - return errWrite - } - } - - if bufferedReader != nil { - if _, errCopy := io.Copy(w, bufferedReader); errCopy != nil { - return errCopy - } - } - if decompressErr != nil { - if _, errWrite := io.WriteString(w, fmt.Sprintf("\n[DECOMPRESSION ERROR: %v]", decompressErr)); errWrite != nil { - return errWrite - } - } - - if trailingNewline { - if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { - return errWrite - } - } - return nil -} - -func responseBodyStartsWithLeadingNewline(reader *bufio.Reader) bool { - if reader == nil { - return false - } - if peeked, _ := reader.Peek(2); len(peeked) >= 2 && peeked[0] == '\r' && peeked[1] == '\n' { - return true - } - if peeked, _ := reader.Peek(1); len(peeked) >= 1 && peeked[0] == '\n' { - return true - } - return false -} - -// formatLogContent creates the complete log content for non-streaming requests. -// -// Parameters: -// - url: The request URL -// - method: The HTTP method -// - headers: The request headers -// - body: The request body -// - websocketTimeline: The downstream websocket event timeline -// - apiRequest: The API request data -// - apiResponse: The API response data -// - response: The raw response data -// - status: The response status code -// - responseHeaders: The response headers -// -// Returns: -// - string: The formatted log content -func (l *FileRequestLogger) formatLogContent(url, method string, headers map[string][]string, body, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline, response []byte, status int, responseHeaders map[string][]string, apiResponseErrors []*interfaces.ErrorMessage) string { - var content strings.Builder - isWebsocketTranscript := hasSectionPayload(websocketTimeline) - downstreamTransport := inferDownstreamTransport(headers, websocketTimeline, nil) - upstreamTransport := inferUpstreamTransport(apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, nil, apiResponseErrors) - - // Request info - content.WriteString(l.formatRequestInfo(url, method, headers, body, downstreamTransport, upstreamTransport, !isWebsocketTranscript)) - - if len(websocketTimeline) > 0 { - if bytes.HasPrefix(websocketTimeline, []byte("=== WEBSOCKET TIMELINE")) { - content.Write(websocketTimeline) - if !bytes.HasSuffix(websocketTimeline, []byte("\n")) { - content.WriteString("\n") - } - } else { - content.WriteString("=== WEBSOCKET TIMELINE ===\n") - content.Write(websocketTimeline) - content.WriteString("\n") - } - content.WriteString("\n") - } - - if len(apiWebsocketTimeline) > 0 { - if bytes.HasPrefix(apiWebsocketTimeline, []byte("=== API WEBSOCKET TIMELINE")) { - content.Write(apiWebsocketTimeline) - if !bytes.HasSuffix(apiWebsocketTimeline, []byte("\n")) { - content.WriteString("\n") - } - } else { - content.WriteString("=== API WEBSOCKET TIMELINE ===\n") - content.Write(apiWebsocketTimeline) - content.WriteString("\n") - } - content.WriteString("\n") - } - - if len(apiRequest) > 0 { - if bytes.HasPrefix(apiRequest, []byte("=== API REQUEST")) { - content.Write(apiRequest) - if !bytes.HasSuffix(apiRequest, []byte("\n")) { - content.WriteString("\n") - } - } else { - content.WriteString("=== API REQUEST ===\n") - content.Write(apiRequest) - content.WriteString("\n") - } - content.WriteString("\n") - } - - for i := 0; i < len(apiResponseErrors); i++ { - content.WriteString("=== API ERROR RESPONSE ===\n") - content.WriteString(fmt.Sprintf("HTTP Status: %d\n", apiResponseErrors[i].StatusCode)) - content.WriteString(apiResponseErrors[i].Error.Error()) - content.WriteString("\n\n") - } - - if len(apiResponse) > 0 { - if bytes.HasPrefix(apiResponse, []byte("=== API RESPONSE")) { - content.Write(apiResponse) - if !bytes.HasSuffix(apiResponse, []byte("\n")) { - content.WriteString("\n") - } - } else { - content.WriteString("=== API RESPONSE ===\n") - content.Write(apiResponse) - content.WriteString("\n") - } - content.WriteString("\n") - } - - if isWebsocketTranscript { - // Mirror writeNonStreamingLog: websocket transcripts end with the dedicated - // timeline sections instead of a generic downstream HTTP response block. - return content.String() - } - - // Response section - content.WriteString("=== RESPONSE ===\n") - content.WriteString(fmt.Sprintf("Status: %d\n", status)) - - if responseHeaders != nil { - for key, values := range responseHeaders { - for _, value := range values { - content.WriteString(fmt.Sprintf("%s: %s\n", key, value)) - } - } - } - - content.WriteString("\n") - content.Write(response) - content.WriteString("\n") - - return content.String() -} - -// decompressResponse decompresses response data based on Content-Encoding header. -// -// Parameters: -// - responseHeaders: The response headers -// - response: The response data to decompress -// -// Returns: -// - []byte: The decompressed response data -// - error: An error if decompression fails, nil otherwise -func (l *FileRequestLogger) decompressResponse(responseHeaders map[string][]string, response []byte) ([]byte, error) { - if responseHeaders == nil || len(response) == 0 { - return response, nil - } - - // Check Content-Encoding header - var contentEncoding string - for key, values := range responseHeaders { - if strings.ToLower(key) == "content-encoding" && len(values) > 0 { - contentEncoding = strings.ToLower(values[0]) - break - } - } - - switch contentEncoding { - case "gzip": - return l.decompressGzip(response) - case "deflate": - return l.decompressDeflate(response) - case "br": - return l.decompressBrotli(response) - case "zstd": - return l.decompressZstd(response) - default: - // No compression or unsupported compression - return response, nil - } -} - -// decompressGzip decompresses gzip-encoded data. -// -// Parameters: -// - data: The gzip-encoded data to decompress -// -// Returns: -// - []byte: The decompressed data -// - error: An error if decompression fails, nil otherwise -func (l *FileRequestLogger) decompressGzip(data []byte) ([]byte, error) { - reader, err := gzip.NewReader(bytes.NewReader(data)) - if err != nil { - return nil, fmt.Errorf("failed to create gzip reader: %w", err) - } - defer func() { - if errClose := reader.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close gzip reader in request logger") - } - }() - - decompressed, err := io.ReadAll(reader) - if err != nil { - return nil, fmt.Errorf("failed to decompress gzip data: %w", err) - } - - return decompressed, nil -} - -// decompressDeflate decompresses deflate-encoded data. -// -// Parameters: -// - data: The deflate-encoded data to decompress -// -// Returns: -// - []byte: The decompressed data -// - error: An error if decompression fails, nil otherwise -func (l *FileRequestLogger) decompressDeflate(data []byte) ([]byte, error) { - reader := flate.NewReader(bytes.NewReader(data)) - defer func() { - if errClose := reader.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close deflate reader in request logger") - } - }() - - decompressed, err := io.ReadAll(reader) - if err != nil { - return nil, fmt.Errorf("failed to decompress deflate data: %w", err) - } - - return decompressed, nil -} - -// decompressBrotli decompresses brotli-encoded data. -// -// Parameters: -// - data: The brotli-encoded data to decompress -// -// Returns: -// - []byte: The decompressed data -// - error: An error if decompression fails, nil otherwise -func (l *FileRequestLogger) decompressBrotli(data []byte) ([]byte, error) { - reader := brotli.NewReader(bytes.NewReader(data)) - - decompressed, err := io.ReadAll(reader) - if err != nil { - return nil, fmt.Errorf("failed to decompress brotli data: %w", err) - } - - return decompressed, nil -} - -// decompressZstd decompresses zstd-encoded data. -// -// Parameters: -// - data: The zstd-encoded data to decompress -// -// Returns: -// - []byte: The decompressed data -// - error: An error if decompression fails, nil otherwise -func (l *FileRequestLogger) decompressZstd(data []byte) ([]byte, error) { - decoder, err := zstd.NewReader(bytes.NewReader(data)) - if err != nil { - return nil, fmt.Errorf("failed to create zstd reader: %w", err) - } - defer decoder.Close() - - decompressed, err := io.ReadAll(decoder) - if err != nil { - return nil, fmt.Errorf("failed to decompress zstd data: %w", err) - } - - return decompressed, nil -} - -// formatRequestInfo creates the request information section of the log. -// -// Parameters: -// - url: The request URL -// - method: The HTTP method -// - headers: The request headers -// - body: The request body -// -// Returns: -// - string: The formatted request information -func (l *FileRequestLogger) formatRequestInfo(url, method string, headers map[string][]string, body []byte, downstreamTransport string, upstreamTransport string, includeBody bool) string { - var content strings.Builder - - content.WriteString("=== REQUEST INFO ===\n") - content.WriteString(fmt.Sprintf("Version: %s\n", buildinfo.Version)) - content.WriteString(fmt.Sprintf("URL: %s\n", url)) - content.WriteString(fmt.Sprintf("Method: %s\n", method)) - if strings.TrimSpace(downstreamTransport) != "" { - content.WriteString(fmt.Sprintf("Downstream Transport: %s\n", downstreamTransport)) - } - if strings.TrimSpace(upstreamTransport) != "" { - content.WriteString(fmt.Sprintf("Upstream Transport: %s\n", upstreamTransport)) - } - content.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano))) - content.WriteString("\n") - - content.WriteString("=== HEADERS ===\n") - for key, values := range headers { - for _, value := range values { - masked := util.MaskSensitiveHeaderValue(key, value) - content.WriteString(fmt.Sprintf("%s: %s\n", key, masked)) - } - } - content.WriteString("\n") - - if !includeBody { - return content.String() - } - - content.WriteString("=== REQUEST BODY ===\n") - content.Write(body) - content.WriteString("\n\n") - - return content.String() -} - -// FileStreamingLogWriter implements StreamingLogWriter for file-based streaming logs. -// It spools streaming response chunks to a temporary file to avoid retaining large responses in memory. -// The final log file is assembled when Close is called. -type FileStreamingLogWriter struct { - // logFilePath is the final log file path. - logFilePath string - - // url is the request URL (masked upstream in middleware). - url string - - // method is the HTTP method. - method string - - // timestamp is captured when the streaming log is initialized. - timestamp time.Time - - // requestHeaders stores the request headers. - requestHeaders map[string][]string - - // requestBodyPath is a temporary file path holding the request body. - requestBodyPath string - - // responseBodyPath is a temporary file path holding the streaming response body. - responseBodyPath string - - // responseBodyFile is the temp file where chunks are appended by the async writer. - responseBodyFile *os.File - - // chunkChan is a channel for receiving response chunks to spool. - chunkChan chan []byte - - // closeChan is a channel for signaling when the writer is closed. - closeChan chan struct{} - - // errorChan is a channel for reporting errors during writing. - errorChan chan error - - // responseStatus stores the HTTP status code. - responseStatus int - - // statusWritten indicates whether a non-zero status was recorded. - statusWritten bool - - // responseHeaders stores the response headers. - responseHeaders map[string][]string - - // apiRequest stores the upstream API request data. - apiRequest []byte - - // apiRequestSource stores file-backed upstream API request data. - apiRequestSource *FileBodySource - - // apiResponse stores the upstream API response data. - apiResponse []byte - - // apiResponseSource stores file-backed upstream API response data. - apiResponseSource *FileBodySource - - // apiWebsocketTimeline stores the upstream websocket event timeline. - apiWebsocketTimeline []byte - - // apiResponseTimestamp captures when the API response was received. - apiResponseTimestamp time.Time -} - -// WriteChunkAsync writes a response chunk asynchronously (non-blocking). -// -// Parameters: -// - chunk: The response chunk to write -func (w *FileStreamingLogWriter) WriteChunkAsync(chunk []byte) { - if w.chunkChan == nil { - return - } - - // Make a copy of the chunk to avoid data races - chunkCopy := make([]byte, len(chunk)) - copy(chunkCopy, chunk) - - // Non-blocking send - select { - case w.chunkChan <- chunkCopy: - default: - // Channel is full, skip this chunk to avoid blocking - } -} - -// WriteStatus buffers the response status and headers for later writing. -// -// Parameters: -// - status: The response status code -// - headers: The response headers -// -// Returns: -// - error: Always returns nil (buffering cannot fail) -func (w *FileStreamingLogWriter) WriteStatus(status int, headers map[string][]string) error { - if status == 0 { - return nil - } - - w.responseStatus = status - if headers != nil { - w.responseHeaders = make(map[string][]string, len(headers)) - for key, values := range headers { - headerValues := make([]string, len(values)) - copy(headerValues, values) - w.responseHeaders[key] = headerValues - } - } - w.statusWritten = true - return nil -} - -// WriteAPIRequest buffers the upstream API request details for later writing. -// -// Parameters: -// - apiRequest: The API request data (typically includes URL, headers, body sent upstream) -// -// Returns: -// - error: Always returns nil (buffering cannot fail) -func (w *FileStreamingLogWriter) WriteAPIRequest(apiRequest []byte) error { - if len(apiRequest) == 0 { - return nil - } - w.apiRequest = bytes.Clone(apiRequest) - return nil -} - -// WriteAPIRequestSource buffers a file-backed upstream API request for final writing. -func (w *FileStreamingLogWriter) WriteAPIRequestSource(apiRequestSource *FileBodySource) error { - if apiRequestSource == nil || !apiRequestSource.HasPayload() { - return nil - } - w.apiRequestSource = apiRequestSource - return nil -} - -// WriteAPIResponse buffers the upstream API response details for later writing. -// -// Parameters: -// - apiResponse: The API response data -// -// Returns: -// - error: Always returns nil (buffering cannot fail) -func (w *FileStreamingLogWriter) WriteAPIResponse(apiResponse []byte) error { - if len(apiResponse) == 0 { - return nil - } - w.apiResponse = bytes.Clone(apiResponse) - return nil -} - -// WriteAPIResponseSource buffers a file-backed upstream API response for final writing. -func (w *FileStreamingLogWriter) WriteAPIResponseSource(apiResponseSource *FileBodySource) error { - if apiResponseSource == nil || !apiResponseSource.HasPayload() { - return nil - } - w.apiResponseSource = apiResponseSource - return nil -} - -// WriteAPIWebsocketTimeline buffers the upstream websocket timeline for later writing. -// -// Parameters: -// - apiWebsocketTimeline: The upstream websocket event timeline -// -// Returns: -// - error: Always returns nil (buffering cannot fail) -func (w *FileStreamingLogWriter) WriteAPIWebsocketTimeline(apiWebsocketTimeline []byte) error { - if len(apiWebsocketTimeline) == 0 { - return nil - } - w.apiWebsocketTimeline = bytes.Clone(apiWebsocketTimeline) - return nil -} - -func (w *FileStreamingLogWriter) SetFirstChunkTimestamp(timestamp time.Time) { - if !timestamp.IsZero() { - w.apiResponseTimestamp = timestamp - } -} - -// Close finalizes the log file and cleans up resources. -// It writes all buffered data to the file in the correct order: -// API WEBSOCKET TIMELINE -> API REQUEST -> API RESPONSE -> RESPONSE (status, headers, body chunks) -// -// Returns: -// - error: An error if closing fails, nil otherwise -func (w *FileStreamingLogWriter) Close() error { - if w.chunkChan != nil { - close(w.chunkChan) - } - - // Wait for async writer to finish spooling chunks - if w.closeChan != nil { - <-w.closeChan - w.chunkChan = nil - } - - select { - case errWrite := <-w.errorChan: - w.cleanupTempFiles() - return errWrite - default: - } - - if w.logFilePath == "" { - w.cleanupTempFiles() - return nil - } - - logFile, errOpen := os.OpenFile(w.logFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) - if errOpen != nil { - w.cleanupTempFiles() - return fmt.Errorf("failed to create log file: %w", errOpen) - } - - writeErr := w.writeFinalLog(logFile) - if errClose := logFile.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close request log file") - if writeErr == nil { - writeErr = errClose - } - } - - w.cleanupTempFiles() - return writeErr -} - -// asyncWriter runs in a goroutine to buffer chunks from the channel. -// It continuously reads chunks from the channel and appends them to a temp file for later assembly. -func (w *FileStreamingLogWriter) asyncWriter() { - defer close(w.closeChan) - - for chunk := range w.chunkChan { - if w.responseBodyFile == nil { - continue - } - if _, errWrite := w.responseBodyFile.Write(chunk); errWrite != nil { - select { - case w.errorChan <- errWrite: - default: - } - if errClose := w.responseBodyFile.Close(); errClose != nil { - select { - case w.errorChan <- errClose: - default: - } - } - w.responseBodyFile = nil - } - } - - if w.responseBodyFile == nil { - return - } - if errClose := w.responseBodyFile.Close(); errClose != nil { - select { - case w.errorChan <- errClose: - default: - } - } - w.responseBodyFile = nil -} - -func (w *FileStreamingLogWriter) writeFinalLog(logFile *os.File) error { - if errWrite := writeRequestInfoWithBody(logFile, w.url, w.method, w.requestHeaders, nil, w.requestBodyPath, w.timestamp, "http", inferUpstreamTransport(w.apiRequest, w.apiRequestSource, w.apiResponse, w.apiResponseSource, w.apiWebsocketTimeline, nil, nil), true); errWrite != nil { - return errWrite - } - if errWrite := writeAPISection(logFile, "=== API WEBSOCKET TIMELINE ===\n", "=== API WEBSOCKET TIMELINE", w.apiWebsocketTimeline, time.Time{}); errWrite != nil { - return errWrite - } - if errWrite := writePreformattedAPISectionWithSource(logFile, "=== API REQUEST ===\n", "=== API REQUEST", w.apiRequest, w.apiRequestSource, time.Time{}); errWrite != nil { - return errWrite - } - if errWrite := writePreformattedAPISectionWithSource(logFile, "=== API RESPONSE ===\n", "=== API RESPONSE", w.apiResponse, w.apiResponseSource, w.apiResponseTimestamp); errWrite != nil { - return errWrite - } - - responseBodyFile, errOpen := os.Open(w.responseBodyPath) - if errOpen != nil { - return errOpen - } - defer func() { - if errClose := responseBodyFile.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close response body temp file") - } - }() - - return writeResponseSection(logFile, w.responseStatus, w.statusWritten, w.responseHeaders, responseBodyFile, nil, false) -} - -func (w *FileStreamingLogWriter) cleanupTempFiles() { - if w.requestBodyPath != "" { - if errRemove := os.Remove(w.requestBodyPath); errRemove != nil { - log.WithError(errRemove).Warn("failed to remove request body temp file") - } - w.requestBodyPath = "" - } - - if w.responseBodyPath != "" { - if errRemove := os.Remove(w.responseBodyPath); errRemove != nil { - log.WithError(errRemove).Warn("failed to remove response body temp file") - } - w.responseBodyPath = "" - } -} - -// NoOpStreamingLogWriter is a no-operation implementation for when logging is disabled. -// It implements the StreamingLogWriter interface but performs no actual logging operations. -type NoOpStreamingLogWriter struct{} - -// WriteChunkAsync is a no-op implementation that does nothing. -// -// Parameters: -// - chunk: The response chunk (ignored) -func (w *NoOpStreamingLogWriter) WriteChunkAsync(_ []byte) {} - -// WriteStatus is a no-op implementation that does nothing and always returns nil. -// -// Parameters: -// - status: The response status code (ignored) -// - headers: The response headers (ignored) -// -// Returns: -// - error: Always returns nil -func (w *NoOpStreamingLogWriter) WriteStatus(_ int, _ map[string][]string) error { - return nil -} - -// WriteAPIRequest is a no-op implementation that does nothing and always returns nil. -// -// Parameters: -// - apiRequest: The API request data (ignored) -// -// Returns: -// - error: Always returns nil -func (w *NoOpStreamingLogWriter) WriteAPIRequest(_ []byte) error { - return nil -} - -// WriteAPIResponse is a no-op implementation that does nothing and always returns nil. -// -// Parameters: -// - apiResponse: The API response data (ignored) -// -// Returns: -// - error: Always returns nil -func (w *NoOpStreamingLogWriter) WriteAPIResponse(_ []byte) error { - return nil -} - -// WriteAPIWebsocketTimeline is a no-op implementation that does nothing and always returns nil. -// -// Parameters: -// - apiWebsocketTimeline: The upstream websocket event timeline (ignored) -// -// Returns: -// - error: Always returns nil -func (w *NoOpStreamingLogWriter) WriteAPIWebsocketTimeline(_ []byte) error { - return nil -} - -func (w *NoOpStreamingLogWriter) SetFirstChunkTimestamp(_ time.Time) {} - -// Close is a no-op implementation that does nothing and always returns nil. -// -// Returns: -// - error: Always returns nil -func (w *NoOpStreamingLogWriter) Close() error { return nil } - -type homeStreamingLogWriter struct { - url string - method string - timestamp time.Time - - requestHeaders map[string][]string - requestBody []byte - - chunkChan chan []byte - doneChan chan struct{} - - responseStatus int - statusWritten bool - responseHeaders map[string][]string - responseBody bytes.Buffer - apiRequest []byte - apiResponse []byte - apiWebsocketTime []byte - requestID string - apiResponseTS time.Time - firstChunkTS time.Time -} - -func newHomeStreamingLogWriter(url, method string, headers map[string][]string, body []byte, requestID string) *homeStreamingLogWriter { - requestHeaders := make(map[string][]string, len(headers)) - for key, values := range headers { - headerValues := make([]string, len(values)) - copy(headerValues, values) - requestHeaders[key] = headerValues - } - - writer := &homeStreamingLogWriter{ - url: url, - method: method, - timestamp: time.Now(), - requestHeaders: requestHeaders, - requestBody: append([]byte(nil), body...), - requestID: strings.TrimSpace(requestID), - chunkChan: make(chan []byte, 100), - doneChan: make(chan struct{}), - } - - go writer.asyncWriter() - return writer -} - -func (w *homeStreamingLogWriter) asyncWriter() { - defer close(w.doneChan) - for chunk := range w.chunkChan { - if len(chunk) == 0 { - continue - } - _, _ = w.responseBody.Write(chunk) - } -} - -func (w *homeStreamingLogWriter) WriteChunkAsync(chunk []byte) { - if w == nil || w.chunkChan == nil || len(chunk) == 0 { - return - } - select { - case w.chunkChan <- append([]byte(nil), chunk...): - default: - } -} - -func (w *homeStreamingLogWriter) WriteStatus(status int, headers map[string][]string) error { - if w == nil || status == 0 { - return nil - } - w.responseStatus = status - w.statusWritten = true - if headers != nil { - w.responseHeaders = make(map[string][]string, len(headers)) - for key, values := range headers { - copied := make([]string, len(values)) - copy(copied, values) - w.responseHeaders[key] = copied - } - } - return nil -} - -func (w *homeStreamingLogWriter) WriteAPIRequest(apiRequest []byte) error { - if w == nil || len(apiRequest) == 0 { - return nil - } - w.apiRequest = bytes.Clone(apiRequest) - return nil -} - -func (w *homeStreamingLogWriter) WriteAPIResponse(apiResponse []byte) error { - if w == nil || len(apiResponse) == 0 { - return nil - } - w.apiResponse = bytes.Clone(apiResponse) - return nil -} - -func (w *homeStreamingLogWriter) WriteAPIWebsocketTimeline(apiWebsocketTimeline []byte) error { - if w == nil || len(apiWebsocketTimeline) == 0 { - return nil - } - w.apiWebsocketTime = bytes.Clone(apiWebsocketTimeline) - return nil -} - -func (w *homeStreamingLogWriter) SetFirstChunkTimestamp(timestamp time.Time) { - if w == nil { - return - } - if !timestamp.IsZero() { - w.firstChunkTS = timestamp - w.apiResponseTS = timestamp - } -} - -func (w *homeStreamingLogWriter) Close() error { - if w == nil { - return nil - } - - client := currentHomeRequestLogClient() - if client == nil || !client.HeartbeatOK() { - return nil - } - - if w.chunkChan != nil { - close(w.chunkChan) - <-w.doneChan - w.chunkChan = nil - } - - responsePayload := w.responseBody.Bytes() - - var buf bytes.Buffer - upstreamTransport := inferUpstreamTransport(w.apiRequest, nil, w.apiResponse, nil, w.apiWebsocketTime, nil, nil) - if errWrite := writeRequestInfoWithBody(&buf, w.url, w.method, w.requestHeaders, w.requestBody, "", w.timestamp, "http", upstreamTransport, true); errWrite != nil { - return errWrite - } - if errWrite := writeAPISection(&buf, "=== API WEBSOCKET TIMELINE ===\n", "=== API WEBSOCKET TIMELINE", w.apiWebsocketTime, time.Time{}); errWrite != nil { - return errWrite - } - if errWrite := writeAPISection(&buf, "=== API REQUEST ===\n", "=== API REQUEST", w.apiRequest, time.Time{}); errWrite != nil { - return errWrite - } - if errWrite := writeAPISection(&buf, "=== API RESPONSE ===\n", "=== API RESPONSE", w.apiResponse, w.apiResponseTS); errWrite != nil { - return errWrite - } - if errWrite := writeResponseSection(&buf, w.responseStatus, w.statusWritten, w.responseHeaders, bytes.NewReader(responsePayload), nil, false); errWrite != nil { - return errWrite - } - - payload := homeRequestLogPayload{ - Headers: cloneHeaders(w.requestHeaders), - RequestID: w.requestID, - RequestLog: buf.String(), - } - raw, errMarshal := json.Marshal(&payload) - if errMarshal != nil { - return errMarshal - } - return client.RPushRequestLog(context.Background(), raw) -} diff --git a/internal/logging/request_logger_body_source.go b/internal/logging/request_logger_body_source.go new file mode 100644 index 000000000..7589166ed --- /dev/null +++ b/internal/logging/request_logger_body_source.go @@ -0,0 +1,256 @@ +package logging + +import ( + "bytes" + "fmt" + "io" + "os" + "strings" + "sync" + + log "github.com/sirupsen/logrus" +) + +// FileBodySource stores large log sections as ordered temp-file parts. +type FileBodySource struct { + mu sync.Mutex + dir string + paths []string + cleaned bool +} + +// NewFileBodySourceInDir creates a temp-backed source under baseDir. +func NewFileBodySourceInDir(baseDir string, prefix string) (*FileBodySource, error) { + prefix = sanitizeTempPrefix(prefix) + baseDir = strings.TrimSpace(baseDir) + if baseDir == "" { + return nil, fmt.Errorf("base directory is required") + } + if errMkdir := os.MkdirAll(baseDir, 0755); errMkdir != nil { + return nil, errMkdir + } + dir, errCreate := os.MkdirTemp(baseDir, "request-log-parts-"+prefix+"-*") + if errCreate != nil { + return nil, errCreate + } + return &FileBodySource{dir: dir}, nil +} + +func sanitizeTempPrefix(prefix string) string { + prefix = strings.TrimSpace(prefix) + if prefix == "" { + return "log" + } + var builder strings.Builder + for _, r := range prefix { + switch { + case r >= 'a' && r <= 'z': + builder.WriteRune(r) + case r >= 'A' && r <= 'Z': + builder.WriteRune(r) + case r >= '0' && r <= '9': + builder.WriteRune(r) + case r == '-' || r == '_': + builder.WriteRune(r) + default: + builder.WriteByte('-') + } + } + out := strings.Trim(builder.String(), "-_") + if out == "" { + return "log" + } + return out +} + +// CreatePart creates one ordered detail log part. +func (s *FileBodySource) CreatePart(prefix string) (*os.File, error) { + if s == nil { + return nil, fmt.Errorf("file body source is nil") + } + s.mu.Lock() + defer s.mu.Unlock() + if s.cleaned { + return nil, fmt.Errorf("file body source has been cleaned") + } + prefix = sanitizeTempPrefix(prefix) + if errMkdir := os.MkdirAll(s.dir, 0755); errMkdir != nil { + return nil, errMkdir + } + file, errCreate := os.CreateTemp(s.dir, prefix+"-*.tmp") + if errCreate != nil { + return nil, errCreate + } + s.paths = append(s.paths, file.Name()) + return file, nil +} + +// AppendPart appends one complete ordered part to the source. +func (s *FileBodySource) AppendPart(data []byte) error { + data = bytes.TrimSpace(data) + if len(data) == 0 { + return nil + } + file, errCreate := s.CreatePart("part") + if errCreate != nil { + return errCreate + } + writeErr := writeLogPart(file, data, false) + if errClose := file.Close(); errClose != nil { + if writeErr == nil { + writeErr = errClose + } + } + return writeErr +} + +// AppendBytes appends raw bytes to a single ordered part. +func (s *FileBodySource) AppendBytes(data []byte) error { + if s == nil { + return fmt.Errorf("file body source is nil") + } + if len(data) == 0 { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + if s.cleaned { + return fmt.Errorf("file body source has been cleaned") + } + if errMkdir := os.MkdirAll(s.dir, 0755); errMkdir != nil { + return errMkdir + } + + var file *os.File + var errOpen error + if len(s.paths) == 0 { + file, errOpen = os.CreateTemp(s.dir, "part-*.tmp") + if errOpen == nil { + s.paths = append(s.paths, file.Name()) + } + } else { + file, errOpen = os.OpenFile(s.paths[len(s.paths)-1], os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + } + if errOpen != nil { + return errOpen + } + + _, writeErr := file.Write(data) + if errClose := file.Close(); errClose != nil { + if writeErr == nil { + writeErr = errClose + } + } + return writeErr +} + +// HasPayload reports whether any detail parts were recorded. +func (s *FileBodySource) HasPayload() bool { + if s == nil { + return false + } + s.mu.Lock() + defer s.mu.Unlock() + return len(s.paths) > 0 && !s.cleaned +} + +// Paths returns a copy of the ordered part paths. +func (s *FileBodySource) Paths() []string { + if s == nil { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + out := make([]string, len(s.paths)) + copy(out, s.paths) + return out +} + +// WriteTo merges all ordered parts into w. +func (s *FileBodySource) WriteTo(w io.Writer) error { + if s == nil || w == nil { + return nil + } + paths := s.Paths() + wrote := false + for _, path := range paths { + file, errOpen := os.Open(path) + if errOpen != nil { + if os.IsNotExist(errOpen) { + continue + } + return errOpen + } + if wrote { + if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { + if errClose := file.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close log part file") + } + return errWrite + } + } + _, errCopy := io.Copy(w, file) + if errClose := file.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close log part file") + if errCopy == nil { + errCopy = errClose + } + } + if errCopy != nil { + return errCopy + } + wrote = true + } + return nil +} + +// Bytes merges all ordered parts into memory. +func (s *FileBodySource) Bytes() ([]byte, error) { + var buf bytes.Buffer + if errWrite := s.WriteTo(&buf); errWrite != nil { + return nil, errWrite + } + return buf.Bytes(), nil +} + +// Cleanup removes all temp detail parts and their directory. +func (s *FileBodySource) Cleanup() error { + if s == nil { + return nil + } + s.mu.Lock() + if s.cleaned { + s.mu.Unlock() + return nil + } + paths := make([]string, len(s.paths)) + copy(paths, s.paths) + dir := s.dir + s.paths = nil + s.cleaned = true + s.mu.Unlock() + + var firstErr error + for _, path := range paths { + if errRemove := os.Remove(path); errRemove != nil && !os.IsNotExist(errRemove) && firstErr == nil { + firstErr = errRemove + } + } + if dir != "" { + if errRemove := os.RemoveAll(dir); errRemove != nil && firstErr == nil { + firstErr = errRemove + } + } + return firstErr +} + +func cleanupFileBodySources(sources ...*FileBodySource) { + for _, source := range sources { + if source == nil { + continue + } + if errCleanup := source.Cleanup(); errCleanup != nil { + log.WithError(errCleanup).Warn("failed to clean up log part files") + } + } +} diff --git a/internal/logging/request_logger_format.go b/internal/logging/request_logger_format.go new file mode 100644 index 000000000..0f476c750 --- /dev/null +++ b/internal/logging/request_logger_format.go @@ -0,0 +1,720 @@ +package logging + +import ( + "bufio" + "bytes" + "compress/flate" + "compress/gzip" + "fmt" + "io" + "os" + "strings" + "time" + + "github.com/andybalholm/brotli" + "github.com/klauspost/compress/zstd" + "github.com/router-for-me/CLIProxyAPI/v7/internal/buildinfo" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + log "github.com/sirupsen/logrus" +) + +func (l *FileRequestLogger) writeNonStreamingLog( + w io.Writer, + url, method string, + requestHeaders map[string][]string, + requestBody []byte, + requestBodyPath string, + websocketTimeline []byte, + websocketTimelineSource *FileBodySource, + apiRequest []byte, + apiRequestSource *FileBodySource, + apiResponse []byte, + apiResponseSource *FileBodySource, + apiWebsocketTimeline []byte, + apiWebsocketTimelineSource *FileBodySource, + apiResponseErrors []*interfaces.ErrorMessage, + statusCode int, + responseHeaders map[string][]string, + response []byte, + decompressErr error, + requestTimestamp time.Time, + apiResponseTimestamp time.Time, +) error { + if requestTimestamp.IsZero() { + requestTimestamp = time.Now() + } + isWebsocketTranscript := hasSectionPayload(websocketTimeline) || hasFileBodySourcePayload(websocketTimelineSource) + downstreamTransport := inferDownstreamTransport(requestHeaders, websocketTimeline, websocketTimelineSource) + upstreamTransport := inferUpstreamTransport(apiRequest, apiRequestSource, apiResponse, apiResponseSource, apiWebsocketTimeline, apiWebsocketTimelineSource, apiResponseErrors) + if errWrite := writeRequestInfoWithBody(w, url, method, requestHeaders, requestBody, requestBodyPath, requestTimestamp, downstreamTransport, upstreamTransport, !isWebsocketTranscript); errWrite != nil { + return errWrite + } + if errWrite := writeAPISectionWithSource(w, "=== WEBSOCKET TIMELINE ===\n", "=== WEBSOCKET TIMELINE", websocketTimeline, websocketTimelineSource, time.Time{}); errWrite != nil { + return errWrite + } + if errWrite := writeAPISectionWithSource(w, "=== API WEBSOCKET TIMELINE ===\n", "=== API WEBSOCKET TIMELINE", apiWebsocketTimeline, apiWebsocketTimelineSource, time.Time{}); errWrite != nil { + return errWrite + } + if errWrite := writePreformattedAPISectionWithSource(w, "=== API REQUEST ===\n", "=== API REQUEST", apiRequest, apiRequestSource, time.Time{}); errWrite != nil { + return errWrite + } + if errWrite := writeAPIErrorResponses(w, apiResponseErrors); errWrite != nil { + return errWrite + } + if errWrite := writePreformattedAPISectionWithSource(w, "=== API RESPONSE ===\n", "=== API RESPONSE", apiResponse, apiResponseSource, apiResponseTimestamp); errWrite != nil { + return errWrite + } + if isWebsocketTranscript { + // Intentionally omit the generic downstream HTTP response section for websocket + // transcripts. The durable session exchange is captured in WEBSOCKET TIMELINE, + // and appending a one-off upgrade response snapshot would dilute that transcript. + return nil + } + return writeResponseSection(w, statusCode, true, responseHeaders, bytes.NewReader(response), decompressErr, true) +} + +func writeRequestInfoWithBody( + w io.Writer, + url, method string, + headers map[string][]string, + body []byte, + bodyPath string, + timestamp time.Time, + downstreamTransport string, + upstreamTransport string, + includeBody bool, +) error { + if _, errWrite := io.WriteString(w, "=== REQUEST INFO ===\n"); errWrite != nil { + return errWrite + } + if _, errWrite := io.WriteString(w, fmt.Sprintf("Version: %s\n", buildinfo.Version)); errWrite != nil { + return errWrite + } + if _, errWrite := io.WriteString(w, fmt.Sprintf("URL: %s\n", url)); errWrite != nil { + return errWrite + } + if _, errWrite := io.WriteString(w, fmt.Sprintf("Method: %s\n", method)); errWrite != nil { + return errWrite + } + if strings.TrimSpace(downstreamTransport) != "" { + if _, errWrite := io.WriteString(w, fmt.Sprintf("Downstream Transport: %s\n", downstreamTransport)); errWrite != nil { + return errWrite + } + } + if strings.TrimSpace(upstreamTransport) != "" { + if _, errWrite := io.WriteString(w, fmt.Sprintf("Upstream Transport: %s\n", upstreamTransport)); errWrite != nil { + return errWrite + } + } + if _, errWrite := io.WriteString(w, fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano))); errWrite != nil { + return errWrite + } + if errWrite := writeSectionSpacing(w, 1); errWrite != nil { + return errWrite + } + + if _, errWrite := io.WriteString(w, "=== HEADERS ===\n"); errWrite != nil { + return errWrite + } + for key, values := range headers { + for _, value := range values { + masked := util.MaskSensitiveHeaderValue(key, value) + if _, errWrite := io.WriteString(w, fmt.Sprintf("%s: %s\n", key, masked)); errWrite != nil { + return errWrite + } + } + } + if errWrite := writeSectionSpacing(w, 1); errWrite != nil { + return errWrite + } + + if !includeBody { + return nil + } + + if _, errWrite := io.WriteString(w, "=== REQUEST BODY ===\n"); errWrite != nil { + return errWrite + } + + bodyTrailingNewlines := 1 + if bodyPath != "" { + bodyFile, errOpen := os.Open(bodyPath) + if errOpen != nil { + return errOpen + } + tracker := &trailingNewlineTrackingWriter{writer: w} + written, errCopy := io.Copy(tracker, bodyFile) + if errCopy != nil { + _ = bodyFile.Close() + return errCopy + } + if written > 0 { + bodyTrailingNewlines = tracker.trailingNewlines + } + if errClose := bodyFile.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close request body temp file") + } + } else if _, errWrite := w.Write(body); errWrite != nil { + return errWrite + } else if len(body) > 0 { + bodyTrailingNewlines = countTrailingNewlinesBytes(body) + } + if errWrite := writeSectionSpacing(w, bodyTrailingNewlines); errWrite != nil { + return errWrite + } + return nil +} + +func countTrailingNewlinesBytes(payload []byte) int { + count := 0 + for i := len(payload) - 1; i >= 0; i-- { + if payload[i] != '\n' { + break + } + count++ + } + return count +} + +func writeSectionSpacing(w io.Writer, trailingNewlines int) error { + missingNewlines := 3 - trailingNewlines + if missingNewlines <= 0 { + return nil + } + _, errWrite := io.WriteString(w, strings.Repeat("\n", missingNewlines)) + return errWrite +} + +type trailingNewlineTrackingWriter struct { + writer io.Writer + trailingNewlines int +} + +func (t *trailingNewlineTrackingWriter) Write(payload []byte) (int, error) { + written, errWrite := t.writer.Write(payload) + if written > 0 { + writtenPayload := payload[:written] + trailingNewlines := countTrailingNewlinesBytes(writtenPayload) + if trailingNewlines == len(writtenPayload) { + t.trailingNewlines += trailingNewlines + } else { + t.trailingNewlines = trailingNewlines + } + } + return written, errWrite +} + +func hasSectionPayload(payload []byte) bool { + return len(bytes.TrimSpace(payload)) > 0 +} + +func hasFileBodySourcePayload(source *FileBodySource) bool { + return source != nil && source.HasPayload() +} + +func inferDownstreamTransport(headers map[string][]string, websocketTimeline []byte, websocketTimelineSource *FileBodySource) string { + if hasSectionPayload(websocketTimeline) || hasFileBodySourcePayload(websocketTimelineSource) { + return "websocket" + } + for key, values := range headers { + if strings.EqualFold(strings.TrimSpace(key), "Upgrade") { + for _, value := range values { + if strings.EqualFold(strings.TrimSpace(value), "websocket") { + return "websocket" + } + } + } + } + return "http" +} + +func inferUpstreamTransport(apiRequest []byte, apiRequestSource *FileBodySource, apiResponse []byte, apiResponseSource *FileBodySource, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, _ []*interfaces.ErrorMessage) string { + hasHTTP := hasSectionPayload(apiRequest) || hasFileBodySourcePayload(apiRequestSource) || hasSectionPayload(apiResponse) || hasFileBodySourcePayload(apiResponseSource) + hasWS := hasSectionPayload(apiWebsocketTimeline) || hasFileBodySourcePayload(apiWebsocketTimelineSource) + switch { + case hasHTTP && hasWS: + return "websocket+http" + case hasWS: + return "websocket" + case hasHTTP: + return "http" + default: + return "" + } +} + +func writeLogPart(w io.Writer, payload []byte, prependNewline bool) error { + if w == nil { + return nil + } + if prependNewline { + if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { + return errWrite + } + } + if _, errWrite := w.Write(payload); errWrite != nil { + return errWrite + } + if !bytes.HasSuffix(payload, []byte("\n")) { + if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { + return errWrite + } + } + return nil +} + +func writeAPISection(w io.Writer, sectionHeader string, sectionPrefix string, payload []byte, timestamp time.Time) error { + if len(payload) == 0 { + return nil + } + + if bytes.HasPrefix(payload, []byte(sectionPrefix)) { + if _, errWrite := w.Write(payload); errWrite != nil { + return errWrite + } + } else { + if _, errWrite := io.WriteString(w, sectionHeader); errWrite != nil { + return errWrite + } + if !timestamp.IsZero() { + if _, errWrite := io.WriteString(w, fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano))); errWrite != nil { + return errWrite + } + } + if _, errWrite := w.Write(payload); errWrite != nil { + return errWrite + } + } + + if errWrite := writeSectionSpacing(w, countTrailingNewlinesBytes(payload)); errWrite != nil { + return errWrite + } + return nil +} + +func writeAPISectionWithSource(w io.Writer, sectionHeader string, sectionPrefix string, payload []byte, source *FileBodySource, timestamp time.Time) error { + if !hasFileBodySourcePayload(source) { + return writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp) + } + if len(payload) > 0 { + if errWrite := writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp); errWrite != nil { + return errWrite + } + } + if _, errWrite := io.WriteString(w, sectionHeader); errWrite != nil { + return errWrite + } + if !timestamp.IsZero() { + if _, errWrite := io.WriteString(w, fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano))); errWrite != nil { + return errWrite + } + } + tracker := &trailingNewlineTrackingWriter{writer: w} + if errWrite := source.WriteTo(tracker); errWrite != nil { + return errWrite + } + if errWrite := writeSectionSpacing(w, tracker.trailingNewlines); errWrite != nil { + return errWrite + } + return nil +} + +func writePreformattedAPISectionWithSource(w io.Writer, sectionHeader string, sectionPrefix string, payload []byte, source *FileBodySource, timestamp time.Time) error { + if !hasFileBodySourcePayload(source) { + return writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp) + } + if len(payload) > 0 { + if errWrite := writeAPISection(w, sectionHeader, sectionPrefix, payload, timestamp); errWrite != nil { + return errWrite + } + } + tracker := &trailingNewlineTrackingWriter{writer: w} + if errWrite := source.WriteTo(tracker); errWrite != nil { + return errWrite + } + if errWrite := writeSectionSpacing(w, tracker.trailingNewlines); errWrite != nil { + return errWrite + } + return nil +} + +func writeAPIErrorResponses(w io.Writer, apiResponseErrors []*interfaces.ErrorMessage) error { + for i := 0; i < len(apiResponseErrors); i++ { + if apiResponseErrors[i] == nil { + continue + } + if _, errWrite := io.WriteString(w, "=== API ERROR RESPONSE ===\n"); errWrite != nil { + return errWrite + } + if _, errWrite := io.WriteString(w, fmt.Sprintf("HTTP Status: %d\n", apiResponseErrors[i].StatusCode)); errWrite != nil { + return errWrite + } + trailingNewlines := 1 + if apiResponseErrors[i].Error != nil { + errText := apiResponseErrors[i].Error.Error() + if _, errWrite := io.WriteString(w, errText); errWrite != nil { + return errWrite + } + if errText != "" { + trailingNewlines = countTrailingNewlinesBytes([]byte(errText)) + } + } + if errWrite := writeSectionSpacing(w, trailingNewlines); errWrite != nil { + return errWrite + } + } + return nil +} + +func writeResponseSection(w io.Writer, statusCode int, statusWritten bool, responseHeaders map[string][]string, responseReader io.Reader, decompressErr error, trailingNewline bool) error { + if _, errWrite := io.WriteString(w, "=== RESPONSE ===\n"); errWrite != nil { + return errWrite + } + if statusWritten { + if _, errWrite := io.WriteString(w, fmt.Sprintf("Status: %d\n", statusCode)); errWrite != nil { + return errWrite + } + } + + if responseHeaders != nil { + for key, values := range responseHeaders { + for _, value := range values { + if _, errWrite := io.WriteString(w, fmt.Sprintf("%s: %s\n", key, value)); errWrite != nil { + return errWrite + } + } + } + } + + var bufferedReader *bufio.Reader + if responseReader != nil { + bufferedReader = bufio.NewReader(responseReader) + } + if !responseBodyStartsWithLeadingNewline(bufferedReader) { + if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { + return errWrite + } + } + + if bufferedReader != nil { + if _, errCopy := io.Copy(w, bufferedReader); errCopy != nil { + return errCopy + } + } + if decompressErr != nil { + if _, errWrite := io.WriteString(w, fmt.Sprintf("\n[DECOMPRESSION ERROR: %v]", decompressErr)); errWrite != nil { + return errWrite + } + } + + if trailingNewline { + if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { + return errWrite + } + } + return nil +} + +func responseBodyStartsWithLeadingNewline(reader *bufio.Reader) bool { + if reader == nil { + return false + } + if peeked, _ := reader.Peek(2); len(peeked) >= 2 && peeked[0] == '\r' && peeked[1] == '\n' { + return true + } + if peeked, _ := reader.Peek(1); len(peeked) >= 1 && peeked[0] == '\n' { + return true + } + return false +} + +// formatLogContent creates the complete log content for non-streaming requests. +// +// Parameters: +// - url: The request URL +// - method: The HTTP method +// - headers: The request headers +// - body: The request body +// - websocketTimeline: The downstream websocket event timeline +// - apiRequest: The API request data +// - apiResponse: The API response data +// - response: The raw response data +// - status: The response status code +// - responseHeaders: The response headers +// +// Returns: +// - string: The formatted log content +func (l *FileRequestLogger) formatLogContent(url, method string, headers map[string][]string, body, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline, response []byte, status int, responseHeaders map[string][]string, apiResponseErrors []*interfaces.ErrorMessage) string { + var content strings.Builder + isWebsocketTranscript := hasSectionPayload(websocketTimeline) + downstreamTransport := inferDownstreamTransport(headers, websocketTimeline, nil) + upstreamTransport := inferUpstreamTransport(apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, nil, apiResponseErrors) + + // Request info + content.WriteString(l.formatRequestInfo(url, method, headers, body, downstreamTransport, upstreamTransport, !isWebsocketTranscript)) + + if len(websocketTimeline) > 0 { + if bytes.HasPrefix(websocketTimeline, []byte("=== WEBSOCKET TIMELINE")) { + content.Write(websocketTimeline) + if !bytes.HasSuffix(websocketTimeline, []byte("\n")) { + content.WriteString("\n") + } + } else { + content.WriteString("=== WEBSOCKET TIMELINE ===\n") + content.Write(websocketTimeline) + content.WriteString("\n") + } + content.WriteString("\n") + } + + if len(apiWebsocketTimeline) > 0 { + if bytes.HasPrefix(apiWebsocketTimeline, []byte("=== API WEBSOCKET TIMELINE")) { + content.Write(apiWebsocketTimeline) + if !bytes.HasSuffix(apiWebsocketTimeline, []byte("\n")) { + content.WriteString("\n") + } + } else { + content.WriteString("=== API WEBSOCKET TIMELINE ===\n") + content.Write(apiWebsocketTimeline) + content.WriteString("\n") + } + content.WriteString("\n") + } + + if len(apiRequest) > 0 { + if bytes.HasPrefix(apiRequest, []byte("=== API REQUEST")) { + content.Write(apiRequest) + if !bytes.HasSuffix(apiRequest, []byte("\n")) { + content.WriteString("\n") + } + } else { + content.WriteString("=== API REQUEST ===\n") + content.Write(apiRequest) + content.WriteString("\n") + } + content.WriteString("\n") + } + + for i := 0; i < len(apiResponseErrors); i++ { + content.WriteString("=== API ERROR RESPONSE ===\n") + content.WriteString(fmt.Sprintf("HTTP Status: %d\n", apiResponseErrors[i].StatusCode)) + content.WriteString(apiResponseErrors[i].Error.Error()) + content.WriteString("\n\n") + } + + if len(apiResponse) > 0 { + if bytes.HasPrefix(apiResponse, []byte("=== API RESPONSE")) { + content.Write(apiResponse) + if !bytes.HasSuffix(apiResponse, []byte("\n")) { + content.WriteString("\n") + } + } else { + content.WriteString("=== API RESPONSE ===\n") + content.Write(apiResponse) + content.WriteString("\n") + } + content.WriteString("\n") + } + + if isWebsocketTranscript { + // Mirror writeNonStreamingLog: websocket transcripts end with the dedicated + // timeline sections instead of a generic downstream HTTP response block. + return content.String() + } + + // Response section + content.WriteString("=== RESPONSE ===\n") + content.WriteString(fmt.Sprintf("Status: %d\n", status)) + + if responseHeaders != nil { + for key, values := range responseHeaders { + for _, value := range values { + content.WriteString(fmt.Sprintf("%s: %s\n", key, value)) + } + } + } + + content.WriteString("\n") + content.Write(response) + content.WriteString("\n") + + return content.String() +} + +// decompressResponse decompresses response data based on Content-Encoding header. +// +// Parameters: +// - responseHeaders: The response headers +// - response: The response data to decompress +// +// Returns: +// - []byte: The decompressed response data +// - error: An error if decompression fails, nil otherwise +func (l *FileRequestLogger) decompressResponse(responseHeaders map[string][]string, response []byte) ([]byte, error) { + if responseHeaders == nil || len(response) == 0 { + return response, nil + } + + // Check Content-Encoding header + var contentEncoding string + for key, values := range responseHeaders { + if strings.ToLower(key) == "content-encoding" && len(values) > 0 { + contentEncoding = strings.ToLower(values[0]) + break + } + } + + switch contentEncoding { + case "gzip": + return l.decompressGzip(response) + case "deflate": + return l.decompressDeflate(response) + case "br": + return l.decompressBrotli(response) + case "zstd": + return l.decompressZstd(response) + default: + // No compression or unsupported compression + return response, nil + } +} + +// decompressGzip decompresses gzip-encoded data. +// +// Parameters: +// - data: The gzip-encoded data to decompress +// +// Returns: +// - []byte: The decompressed data +// - error: An error if decompression fails, nil otherwise +func (l *FileRequestLogger) decompressGzip(data []byte) ([]byte, error) { + reader, err := gzip.NewReader(bytes.NewReader(data)) + if err != nil { + return nil, fmt.Errorf("failed to create gzip reader: %w", err) + } + defer func() { + if errClose := reader.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close gzip reader in request logger") + } + }() + + decompressed, err := io.ReadAll(reader) + if err != nil { + return nil, fmt.Errorf("failed to decompress gzip data: %w", err) + } + + return decompressed, nil +} + +// decompressDeflate decompresses deflate-encoded data. +// +// Parameters: +// - data: The deflate-encoded data to decompress +// +// Returns: +// - []byte: The decompressed data +// - error: An error if decompression fails, nil otherwise +func (l *FileRequestLogger) decompressDeflate(data []byte) ([]byte, error) { + reader := flate.NewReader(bytes.NewReader(data)) + defer func() { + if errClose := reader.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close deflate reader in request logger") + } + }() + + decompressed, err := io.ReadAll(reader) + if err != nil { + return nil, fmt.Errorf("failed to decompress deflate data: %w", err) + } + + return decompressed, nil +} + +// decompressBrotli decompresses brotli-encoded data. +// +// Parameters: +// - data: The brotli-encoded data to decompress +// +// Returns: +// - []byte: The decompressed data +// - error: An error if decompression fails, nil otherwise +func (l *FileRequestLogger) decompressBrotli(data []byte) ([]byte, error) { + reader := brotli.NewReader(bytes.NewReader(data)) + + decompressed, err := io.ReadAll(reader) + if err != nil { + return nil, fmt.Errorf("failed to decompress brotli data: %w", err) + } + + return decompressed, nil +} + +// decompressZstd decompresses zstd-encoded data. +// +// Parameters: +// - data: The zstd-encoded data to decompress +// +// Returns: +// - []byte: The decompressed data +// - error: An error if decompression fails, nil otherwise +func (l *FileRequestLogger) decompressZstd(data []byte) ([]byte, error) { + decoder, err := zstd.NewReader(bytes.NewReader(data)) + if err != nil { + return nil, fmt.Errorf("failed to create zstd reader: %w", err) + } + defer decoder.Close() + + decompressed, err := io.ReadAll(decoder) + if err != nil { + return nil, fmt.Errorf("failed to decompress zstd data: %w", err) + } + + return decompressed, nil +} + +// formatRequestInfo creates the request information section of the log. +// +// Parameters: +// - url: The request URL +// - method: The HTTP method +// - headers: The request headers +// - body: The request body +// +// Returns: +// - string: The formatted request information +func (l *FileRequestLogger) formatRequestInfo(url, method string, headers map[string][]string, body []byte, downstreamTransport string, upstreamTransport string, includeBody bool) string { + var content strings.Builder + + content.WriteString("=== REQUEST INFO ===\n") + content.WriteString(fmt.Sprintf("Version: %s\n", buildinfo.Version)) + content.WriteString(fmt.Sprintf("URL: %s\n", url)) + content.WriteString(fmt.Sprintf("Method: %s\n", method)) + if strings.TrimSpace(downstreamTransport) != "" { + content.WriteString(fmt.Sprintf("Downstream Transport: %s\n", downstreamTransport)) + } + if strings.TrimSpace(upstreamTransport) != "" { + content.WriteString(fmt.Sprintf("Upstream Transport: %s\n", upstreamTransport)) + } + content.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano))) + content.WriteString("\n") + + content.WriteString("=== HEADERS ===\n") + for key, values := range headers { + for _, value := range values { + masked := util.MaskSensitiveHeaderValue(key, value) + content.WriteString(fmt.Sprintf("%s: %s\n", key, masked)) + } + } + content.WriteString("\n") + + if !includeBody { + return content.String() + } + + content.WriteString("=== REQUEST BODY ===\n") + content.Write(body) + content.WriteString("\n\n") + + return content.String() +} diff --git a/internal/logging/request_logger_home.go b/internal/logging/request_logger_home.go new file mode 100644 index 000000000..939386504 --- /dev/null +++ b/internal/logging/request_logger_home.go @@ -0,0 +1,246 @@ +package logging + +import ( + "bytes" + "context" + "encoding/json" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" +) + +type homeRequestLogClient interface { + HeartbeatOK() bool + RPushRequestLog(ctx context.Context, payload []byte) error +} + +var currentHomeRequestLogClient = func() homeRequestLogClient { + return home.Current() +} + +type homeRequestLogPayload struct { + Headers map[string][]string `json:"headers,omitempty"` + RequestID string `json:"request_id,omitempty"` + RequestLog string `json:"request_log,omitempty"` +} + +func cloneHeaders(headers map[string][]string) map[string][]string { + if len(headers) == 0 { + return nil + } + out := make(map[string][]string, len(headers)) + for key, values := range headers { + if strings.TrimSpace(key) == "" { + continue + } + if values == nil { + out[key] = nil + continue + } + copied := make([]string, len(values)) + copy(copied, values) + out[key] = copied + } + if len(out) == 0 { + return nil + } + return out +} + +func (l *FileRequestLogger) forwardRequestLogToHome(ctx context.Context, headers map[string][]string, requestID string, logText string) error { + if l == nil || !l.homeEnabled { + return nil + } + client := currentHomeRequestLogClient() + if client == nil || !client.HeartbeatOK() { + return nil + } + payload := homeRequestLogPayload{ + Headers: cloneHeaders(headers), + RequestID: strings.TrimSpace(requestID), + RequestLog: logText, + } + raw, errMarshal := json.Marshal(&payload) + if errMarshal != nil { + return errMarshal + } + if ctx == nil { + ctx = context.Background() + } + return client.RPushRequestLog(ctx, raw) +} + +// SetHomeEnabled toggles home request-log forwarding. +// When enabled, request logs are not written to disk and are instead forwarded to home via Redis RESP. +func (l *FileRequestLogger) SetHomeEnabled(enabled bool) { + if l == nil { + return + } + l.homeEnabled = enabled +} + +type homeStreamingLogWriter struct { + url string + method string + timestamp time.Time + + requestHeaders map[string][]string + requestBody []byte + + chunkChan chan []byte + doneChan chan struct{} + + responseStatus int + statusWritten bool + responseHeaders map[string][]string + responseBody bytes.Buffer + apiRequest []byte + apiResponse []byte + apiWebsocketTime []byte + requestID string + apiResponseTS time.Time + firstChunkTS time.Time +} + +func newHomeStreamingLogWriter(url, method string, headers map[string][]string, body []byte, requestID string) *homeStreamingLogWriter { + requestHeaders := make(map[string][]string, len(headers)) + for key, values := range headers { + headerValues := make([]string, len(values)) + copy(headerValues, values) + requestHeaders[key] = headerValues + } + + writer := &homeStreamingLogWriter{ + url: url, + method: method, + timestamp: time.Now(), + requestHeaders: requestHeaders, + requestBody: append([]byte(nil), body...), + requestID: strings.TrimSpace(requestID), + chunkChan: make(chan []byte, 100), + doneChan: make(chan struct{}), + } + + go writer.asyncWriter() + return writer +} + +func (w *homeStreamingLogWriter) asyncWriter() { + defer close(w.doneChan) + for chunk := range w.chunkChan { + if len(chunk) == 0 { + continue + } + _, _ = w.responseBody.Write(chunk) + } +} + +func (w *homeStreamingLogWriter) WriteChunkAsync(chunk []byte) { + if w == nil || w.chunkChan == nil || len(chunk) == 0 { + return + } + select { + case w.chunkChan <- append([]byte(nil), chunk...): + default: + } +} + +func (w *homeStreamingLogWriter) WriteStatus(status int, headers map[string][]string) error { + if w == nil || status == 0 { + return nil + } + w.responseStatus = status + w.statusWritten = true + if headers != nil { + w.responseHeaders = make(map[string][]string, len(headers)) + for key, values := range headers { + copied := make([]string, len(values)) + copy(copied, values) + w.responseHeaders[key] = copied + } + } + return nil +} + +func (w *homeStreamingLogWriter) WriteAPIRequest(apiRequest []byte) error { + if w == nil || len(apiRequest) == 0 { + return nil + } + w.apiRequest = bytes.Clone(apiRequest) + return nil +} + +func (w *homeStreamingLogWriter) WriteAPIResponse(apiResponse []byte) error { + if w == nil || len(apiResponse) == 0 { + return nil + } + w.apiResponse = bytes.Clone(apiResponse) + return nil +} + +func (w *homeStreamingLogWriter) WriteAPIWebsocketTimeline(apiWebsocketTimeline []byte) error { + if w == nil || len(apiWebsocketTimeline) == 0 { + return nil + } + w.apiWebsocketTime = bytes.Clone(apiWebsocketTimeline) + return nil +} + +func (w *homeStreamingLogWriter) SetFirstChunkTimestamp(timestamp time.Time) { + if w == nil { + return + } + if !timestamp.IsZero() { + w.firstChunkTS = timestamp + w.apiResponseTS = timestamp + } +} + +func (w *homeStreamingLogWriter) Close() error { + if w == nil { + return nil + } + + client := currentHomeRequestLogClient() + if client == nil || !client.HeartbeatOK() { + return nil + } + + if w.chunkChan != nil { + close(w.chunkChan) + <-w.doneChan + w.chunkChan = nil + } + + responsePayload := w.responseBody.Bytes() + + var buf bytes.Buffer + upstreamTransport := inferUpstreamTransport(w.apiRequest, nil, w.apiResponse, nil, w.apiWebsocketTime, nil, nil) + if errWrite := writeRequestInfoWithBody(&buf, w.url, w.method, w.requestHeaders, w.requestBody, "", w.timestamp, "http", upstreamTransport, true); errWrite != nil { + return errWrite + } + if errWrite := writeAPISection(&buf, "=== API WEBSOCKET TIMELINE ===\n", "=== API WEBSOCKET TIMELINE", w.apiWebsocketTime, time.Time{}); errWrite != nil { + return errWrite + } + if errWrite := writeAPISection(&buf, "=== API REQUEST ===\n", "=== API REQUEST", w.apiRequest, time.Time{}); errWrite != nil { + return errWrite + } + if errWrite := writeAPISection(&buf, "=== API RESPONSE ===\n", "=== API RESPONSE", w.apiResponse, w.apiResponseTS); errWrite != nil { + return errWrite + } + if errWrite := writeResponseSection(&buf, w.responseStatus, w.statusWritten, w.responseHeaders, bytes.NewReader(responsePayload), nil, false); errWrite != nil { + return errWrite + } + + payload := homeRequestLogPayload{ + Headers: cloneHeaders(w.requestHeaders), + RequestID: w.requestID, + RequestLog: buf.String(), + } + raw, errMarshal := json.Marshal(&payload) + if errMarshal != nil { + return errMarshal + } + return client.RPushRequestLog(context.Background(), raw) +} diff --git a/internal/logging/request_logger_streaming.go b/internal/logging/request_logger_streaming.go new file mode 100644 index 000000000..0462175f7 --- /dev/null +++ b/internal/logging/request_logger_streaming.go @@ -0,0 +1,380 @@ +package logging + +import ( + "bytes" + "fmt" + "os" + "time" + + log "github.com/sirupsen/logrus" +) + +// FileStreamingLogWriter implements StreamingLogWriter for file-based streaming logs. +// It spools streaming response chunks to a temporary file to avoid retaining large responses in memory. +// The final log file is assembled when Close is called. +type FileStreamingLogWriter struct { + // logFilePath is the final log file path. + logFilePath string + + // url is the request URL (masked upstream in middleware). + url string + + // method is the HTTP method. + method string + + // timestamp is captured when the streaming log is initialized. + timestamp time.Time + + // requestHeaders stores the request headers. + requestHeaders map[string][]string + + // requestBodyPath is a temporary file path holding the request body. + requestBodyPath string + + // responseBodyPath is a temporary file path holding the streaming response body. + responseBodyPath string + + // responseBodyFile is the temp file where chunks are appended by the async writer. + responseBodyFile *os.File + + // chunkChan is a channel for receiving response chunks to spool. + chunkChan chan []byte + + // closeChan is a channel for signaling when the writer is closed. + closeChan chan struct{} + + // errorChan is a channel for reporting errors during writing. + errorChan chan error + + // responseStatus stores the HTTP status code. + responseStatus int + + // statusWritten indicates whether a non-zero status was recorded. + statusWritten bool + + // responseHeaders stores the response headers. + responseHeaders map[string][]string + + // apiRequest stores the upstream API request data. + apiRequest []byte + + // apiRequestSource stores file-backed upstream API request data. + apiRequestSource *FileBodySource + + // apiResponse stores the upstream API response data. + apiResponse []byte + + // apiResponseSource stores file-backed upstream API response data. + apiResponseSource *FileBodySource + + // apiWebsocketTimeline stores the upstream websocket event timeline. + apiWebsocketTimeline []byte + + // apiResponseTimestamp captures when the API response was received. + apiResponseTimestamp time.Time +} + +// WriteChunkAsync writes a response chunk asynchronously (non-blocking). +// +// Parameters: +// - chunk: The response chunk to write +func (w *FileStreamingLogWriter) WriteChunkAsync(chunk []byte) { + if w.chunkChan == nil { + return + } + + // Make a copy of the chunk to avoid data races + chunkCopy := make([]byte, len(chunk)) + copy(chunkCopy, chunk) + + // Non-blocking send + select { + case w.chunkChan <- chunkCopy: + default: + // Channel is full, skip this chunk to avoid blocking + } +} + +// WriteStatus buffers the response status and headers for later writing. +// +// Parameters: +// - status: The response status code +// - headers: The response headers +// +// Returns: +// - error: Always returns nil (buffering cannot fail) +func (w *FileStreamingLogWriter) WriteStatus(status int, headers map[string][]string) error { + if status == 0 { + return nil + } + + w.responseStatus = status + if headers != nil { + w.responseHeaders = make(map[string][]string, len(headers)) + for key, values := range headers { + headerValues := make([]string, len(values)) + copy(headerValues, values) + w.responseHeaders[key] = headerValues + } + } + w.statusWritten = true + return nil +} + +// WriteAPIRequest buffers the upstream API request details for later writing. +// +// Parameters: +// - apiRequest: The API request data (typically includes URL, headers, body sent upstream) +// +// Returns: +// - error: Always returns nil (buffering cannot fail) +func (w *FileStreamingLogWriter) WriteAPIRequest(apiRequest []byte) error { + if len(apiRequest) == 0 { + return nil + } + w.apiRequest = bytes.Clone(apiRequest) + return nil +} + +// WriteAPIRequestSource buffers a file-backed upstream API request for final writing. +func (w *FileStreamingLogWriter) WriteAPIRequestSource(apiRequestSource *FileBodySource) error { + if apiRequestSource == nil || !apiRequestSource.HasPayload() { + return nil + } + w.apiRequestSource = apiRequestSource + return nil +} + +// WriteAPIResponse buffers the upstream API response details for later writing. +// +// Parameters: +// - apiResponse: The API response data +// +// Returns: +// - error: Always returns nil (buffering cannot fail) +func (w *FileStreamingLogWriter) WriteAPIResponse(apiResponse []byte) error { + if len(apiResponse) == 0 { + return nil + } + w.apiResponse = bytes.Clone(apiResponse) + return nil +} + +// WriteAPIResponseSource buffers a file-backed upstream API response for final writing. +func (w *FileStreamingLogWriter) WriteAPIResponseSource(apiResponseSource *FileBodySource) error { + if apiResponseSource == nil || !apiResponseSource.HasPayload() { + return nil + } + w.apiResponseSource = apiResponseSource + return nil +} + +// WriteAPIWebsocketTimeline buffers the upstream websocket timeline for later writing. +// +// Parameters: +// - apiWebsocketTimeline: The upstream websocket event timeline +// +// Returns: +// - error: Always returns nil (buffering cannot fail) +func (w *FileStreamingLogWriter) WriteAPIWebsocketTimeline(apiWebsocketTimeline []byte) error { + if len(apiWebsocketTimeline) == 0 { + return nil + } + w.apiWebsocketTimeline = bytes.Clone(apiWebsocketTimeline) + return nil +} + +func (w *FileStreamingLogWriter) SetFirstChunkTimestamp(timestamp time.Time) { + if !timestamp.IsZero() { + w.apiResponseTimestamp = timestamp + } +} + +// Close finalizes the log file and cleans up resources. +// It writes all buffered data to the file in the correct order: +// API WEBSOCKET TIMELINE -> API REQUEST -> API RESPONSE -> RESPONSE (status, headers, body chunks) +// +// Returns: +// - error: An error if closing fails, nil otherwise +func (w *FileStreamingLogWriter) Close() error { + if w.chunkChan != nil { + close(w.chunkChan) + } + + // Wait for async writer to finish spooling chunks + if w.closeChan != nil { + <-w.closeChan + w.chunkChan = nil + } + + select { + case errWrite := <-w.errorChan: + w.cleanupTempFiles() + return errWrite + default: + } + + if w.logFilePath == "" { + w.cleanupTempFiles() + return nil + } + + logFile, errOpen := os.OpenFile(w.logFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) + if errOpen != nil { + w.cleanupTempFiles() + return fmt.Errorf("failed to create log file: %w", errOpen) + } + + writeErr := w.writeFinalLog(logFile) + if errClose := logFile.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close request log file") + if writeErr == nil { + writeErr = errClose + } + } + + w.cleanupTempFiles() + return writeErr +} + +// asyncWriter runs in a goroutine to buffer chunks from the channel. +// It continuously reads chunks from the channel and appends them to a temp file for later assembly. +func (w *FileStreamingLogWriter) asyncWriter() { + defer close(w.closeChan) + + for chunk := range w.chunkChan { + if w.responseBodyFile == nil { + continue + } + if _, errWrite := w.responseBodyFile.Write(chunk); errWrite != nil { + select { + case w.errorChan <- errWrite: + default: + } + if errClose := w.responseBodyFile.Close(); errClose != nil { + select { + case w.errorChan <- errClose: + default: + } + } + w.responseBodyFile = nil + } + } + + if w.responseBodyFile == nil { + return + } + if errClose := w.responseBodyFile.Close(); errClose != nil { + select { + case w.errorChan <- errClose: + default: + } + } + w.responseBodyFile = nil +} + +func (w *FileStreamingLogWriter) writeFinalLog(logFile *os.File) error { + if errWrite := writeRequestInfoWithBody(logFile, w.url, w.method, w.requestHeaders, nil, w.requestBodyPath, w.timestamp, "http", inferUpstreamTransport(w.apiRequest, w.apiRequestSource, w.apiResponse, w.apiResponseSource, w.apiWebsocketTimeline, nil, nil), true); errWrite != nil { + return errWrite + } + if errWrite := writeAPISection(logFile, "=== API WEBSOCKET TIMELINE ===\n", "=== API WEBSOCKET TIMELINE", w.apiWebsocketTimeline, time.Time{}); errWrite != nil { + return errWrite + } + if errWrite := writePreformattedAPISectionWithSource(logFile, "=== API REQUEST ===\n", "=== API REQUEST", w.apiRequest, w.apiRequestSource, time.Time{}); errWrite != nil { + return errWrite + } + if errWrite := writePreformattedAPISectionWithSource(logFile, "=== API RESPONSE ===\n", "=== API RESPONSE", w.apiResponse, w.apiResponseSource, w.apiResponseTimestamp); errWrite != nil { + return errWrite + } + + responseBodyFile, errOpen := os.Open(w.responseBodyPath) + if errOpen != nil { + return errOpen + } + defer func() { + if errClose := responseBodyFile.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close response body temp file") + } + }() + + return writeResponseSection(logFile, w.responseStatus, w.statusWritten, w.responseHeaders, responseBodyFile, nil, false) +} + +func (w *FileStreamingLogWriter) cleanupTempFiles() { + if w.requestBodyPath != "" { + if errRemove := os.Remove(w.requestBodyPath); errRemove != nil { + log.WithError(errRemove).Warn("failed to remove request body temp file") + } + w.requestBodyPath = "" + } + + if w.responseBodyPath != "" { + if errRemove := os.Remove(w.responseBodyPath); errRemove != nil { + log.WithError(errRemove).Warn("failed to remove response body temp file") + } + w.responseBodyPath = "" + } +} + +// NoOpStreamingLogWriter is a no-operation implementation for when logging is disabled. +// It implements the StreamingLogWriter interface but performs no actual logging operations. +type NoOpStreamingLogWriter struct{} + +// WriteChunkAsync is a no-op implementation that does nothing. +// +// Parameters: +// - chunk: The response chunk (ignored) +func (w *NoOpStreamingLogWriter) WriteChunkAsync(_ []byte) {} + +// WriteStatus is a no-op implementation that does nothing and always returns nil. +// +// Parameters: +// - status: The response status code (ignored) +// - headers: The response headers (ignored) +// +// Returns: +// - error: Always returns nil +func (w *NoOpStreamingLogWriter) WriteStatus(_ int, _ map[string][]string) error { + return nil +} + +// WriteAPIRequest is a no-op implementation that does nothing and always returns nil. +// +// Parameters: +// - apiRequest: The API request data (ignored) +// +// Returns: +// - error: Always returns nil +func (w *NoOpStreamingLogWriter) WriteAPIRequest(_ []byte) error { + return nil +} + +// WriteAPIResponse is a no-op implementation that does nothing and always returns nil. +// +// Parameters: +// - apiResponse: The API response data (ignored) +// +// Returns: +// - error: Always returns nil +func (w *NoOpStreamingLogWriter) WriteAPIResponse(_ []byte) error { + return nil +} + +// WriteAPIWebsocketTimeline is a no-op implementation that does nothing and always returns nil. +// +// Parameters: +// - apiWebsocketTimeline: The upstream websocket event timeline (ignored) +// +// Returns: +// - error: Always returns nil +func (w *NoOpStreamingLogWriter) WriteAPIWebsocketTimeline(_ []byte) error { + return nil +} + +func (w *NoOpStreamingLogWriter) SetFirstChunkTimestamp(_ time.Time) {} + +// Close is a no-op implementation that does nothing and always returns nil. +// +// Returns: +// - error: Always returns nil +func (w *NoOpStreamingLogWriter) Close() error { return nil } diff --git a/internal/logging/request_logger_writer.go b/internal/logging/request_logger_writer.go new file mode 100644 index 000000000..e5f80e7d0 --- /dev/null +++ b/internal/logging/request_logger_writer.go @@ -0,0 +1,413 @@ +package logging + +import ( + "bytes" + "context" + "fmt" + "io" + "os" + "path/filepath" + "regexp" + "sort" + "strings" + "sync/atomic" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + log "github.com/sirupsen/logrus" +) + +var requestLogID atomic.Uint64 + +// LogRequest logs a complete non-streaming request/response cycle to a file. +// +// Parameters: +// - url: The request URL +// - method: The HTTP method +// - requestHeaders: The request headers +// - body: The request body +// - statusCode: The response status code +// - responseHeaders: The response headers +// - response: The raw response data +// - apiRequest: The API request data +// - apiResponse: The API response data +// - requestID: Optional request ID for log file naming +// - requestTimestamp: When the request was received +// - apiResponseTimestamp: When the API response was received +// +// Returns: +// - error: An error if logging fails, nil otherwise +func (l *FileRequestLogger) LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { + return l.logRequest(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline, apiResponseErrors, false, requestID, requestTimestamp, apiResponseTimestamp) +} + +// LogRequestWithOptions logs a request with optional forced logging behavior. +// The force flag allows writing error logs even when regular request logging is disabled. +func (l *FileRequestLogger) LogRequestWithOptions(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { + return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, nil, apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, nil, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) +} + +func (l *FileRequestLogger) logRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { + return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, nil, apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, nil, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) +} + +// LogRequestWithOptionsAndSources logs a request with optional file-backed large sections. +func (l *FileRequestLogger) LogRequestWithOptionsAndSources(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline []byte, websocketTimelineSource *FileBodySource, apiRequest, apiResponse, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { + return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, websocketTimelineSource, apiRequest, nil, apiResponse, nil, apiWebsocketTimeline, apiWebsocketTimelineSource, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) +} + +// LogRequestWithOptionsAndAllSources logs a request with optional file-backed request and response sections. +func (l *FileRequestLogger) LogRequestWithOptionsAndAllSources(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline []byte, websocketTimelineSource *FileBodySource, apiRequest []byte, apiRequestSource *FileBodySource, apiResponse []byte, apiResponseSource *FileBodySource, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { + return l.logRequestWithSources(url, method, requestHeaders, body, statusCode, responseHeaders, response, websocketTimeline, websocketTimelineSource, apiRequest, apiRequestSource, apiResponse, apiResponseSource, apiWebsocketTimeline, apiWebsocketTimelineSource, apiResponseErrors, force, requestID, requestTimestamp, apiResponseTimestamp) +} + +func (l *FileRequestLogger) logRequestWithSources(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, websocketTimeline []byte, websocketTimelineSource *FileBodySource, apiRequest []byte, apiRequestSource *FileBodySource, apiResponse []byte, apiResponseSource *FileBodySource, apiWebsocketTimeline []byte, apiWebsocketTimelineSource *FileBodySource, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string, requestTimestamp, apiResponseTimestamp time.Time) error { + defer cleanupFileBodySources(websocketTimelineSource, apiRequestSource, apiResponseSource, apiWebsocketTimelineSource) + + if !l.enabled && !force { + return nil + } + + if l.homeEnabled && l.enabled { + responseToWrite, decompressErr := l.decompressResponse(responseHeaders, response) + if decompressErr != nil { + responseToWrite = response + } + + var buf bytes.Buffer + writeErr := l.writeNonStreamingLog( + &buf, + url, + method, + requestHeaders, + body, + "", + websocketTimeline, + websocketTimelineSource, + apiRequest, + apiRequestSource, + apiResponse, + apiResponseSource, + apiWebsocketTimeline, + apiWebsocketTimelineSource, + apiResponseErrors, + statusCode, + responseHeaders, + responseToWrite, + decompressErr, + requestTimestamp, + apiResponseTimestamp, + ) + if writeErr != nil { + return fmt.Errorf("failed to build request log content: %w", writeErr) + } + return l.forwardRequestLogToHome(context.Background(), requestHeaders, requestID, buf.String()) + } + + // Ensure logs directory exists + if errEnsure := l.ensureLogsDir(); errEnsure != nil { + return fmt.Errorf("failed to create logs directory: %w", errEnsure) + } + + // Generate filename with request ID + filename := l.generateFilename(url, requestID) + if force && !l.enabled { + filename = l.generateErrorFilename(url, requestID) + } + filePath := filepath.Join(l.logsDir, filename) + + requestBodyPath, errTemp := l.writeRequestBodyTempFile(body) + if errTemp != nil { + log.WithError(errTemp).Warn("failed to create request body temp file, falling back to direct write") + } + if requestBodyPath != "" { + defer func() { + if errRemove := os.Remove(requestBodyPath); errRemove != nil { + log.WithError(errRemove).Warn("failed to remove request body temp file") + } + }() + } + + responseToWrite, decompressErr := l.decompressResponse(responseHeaders, response) + if decompressErr != nil { + // If decompression fails, continue with original response and annotate the log output. + responseToWrite = response + } + + logFile, errOpen := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) + if errOpen != nil { + return fmt.Errorf("failed to create log file: %w", errOpen) + } + + writeErr := l.writeNonStreamingLog( + logFile, + url, + method, + requestHeaders, + body, + requestBodyPath, + websocketTimeline, + websocketTimelineSource, + apiRequest, + apiRequestSource, + apiResponse, + apiResponseSource, + apiWebsocketTimeline, + apiWebsocketTimelineSource, + apiResponseErrors, + statusCode, + responseHeaders, + responseToWrite, + decompressErr, + requestTimestamp, + apiResponseTimestamp, + ) + if errClose := logFile.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close request log file") + if writeErr == nil { + return errClose + } + } + if writeErr != nil { + return fmt.Errorf("failed to write log file: %w", writeErr) + } + + if force && !l.enabled { + if errCleanup := l.cleanupOldErrorLogs(); errCleanup != nil { + log.WithError(errCleanup).Warn("failed to clean up old error logs") + } + } + + return nil +} + +// LogStreamingRequest initiates logging for a streaming request. +// +// Parameters: +// - url: The request URL +// - method: The HTTP method +// - headers: The request headers +// - body: The request body +// - requestID: Optional request ID for log file naming +// +// Returns: +// - StreamingLogWriter: A writer for streaming response chunks +// - error: An error if logging initialization fails, nil otherwise +func (l *FileRequestLogger) LogStreamingRequest(url, method string, headers map[string][]string, body []byte, requestID string) (StreamingLogWriter, error) { + if !l.enabled { + return &NoOpStreamingLogWriter{}, nil + } + + if l.homeEnabled { + client := currentHomeRequestLogClient() + if client == nil || !client.HeartbeatOK() { + return &NoOpStreamingLogWriter{}, nil + } + return newHomeStreamingLogWriter(url, method, headers, body, requestID), nil + } + + // Ensure logs directory exists + if err := l.ensureLogsDir(); err != nil { + return nil, fmt.Errorf("failed to create logs directory: %w", err) + } + + // Generate filename with request ID + filename := l.generateFilename(url, requestID) + filePath := filepath.Join(l.logsDir, filename) + + requestHeaders := make(map[string][]string, len(headers)) + for key, values := range headers { + headerValues := make([]string, len(values)) + copy(headerValues, values) + requestHeaders[key] = headerValues + } + + requestBodyPath, errTemp := l.writeRequestBodyTempFile(body) + if errTemp != nil { + return nil, fmt.Errorf("failed to create request body temp file: %w", errTemp) + } + + responseBodyFile, errCreate := os.CreateTemp(l.logsDir, "response-body-*.tmp") + if errCreate != nil { + _ = os.Remove(requestBodyPath) + return nil, fmt.Errorf("failed to create response body temp file: %w", errCreate) + } + responseBodyPath := responseBodyFile.Name() + + // Create streaming writer + writer := &FileStreamingLogWriter{ + logFilePath: filePath, + url: url, + method: method, + timestamp: time.Now(), + requestHeaders: requestHeaders, + requestBodyPath: requestBodyPath, + responseBodyPath: responseBodyPath, + responseBodyFile: responseBodyFile, + chunkChan: make(chan []byte, 100), // Buffered channel for async writes + closeChan: make(chan struct{}), + errorChan: make(chan error, 1), + } + + // Start async writer goroutine + go writer.asyncWriter() + + return writer, nil +} + +// generateErrorFilename creates a filename with an error prefix to differentiate forced error logs. +func (l *FileRequestLogger) generateErrorFilename(url string, requestID ...string) string { + return fmt.Sprintf("error-%s", l.generateFilename(url, requestID...)) +} + +// ensureLogsDir creates the logs directory if it doesn't exist. +// +// Returns: +// - error: An error if directory creation fails, nil otherwise +func (l *FileRequestLogger) ensureLogsDir() error { + if _, err := os.Stat(l.logsDir); os.IsNotExist(err) { + return os.MkdirAll(l.logsDir, 0755) + } + return nil +} + +// generateFilename creates a sanitized filename from the URL path and current timestamp. +// Format: v1-responses-2025-12-23T195811-a1b2c3d4.log +// +// Parameters: +// - url: The request URL +// - requestID: Optional request ID to include in filename +// +// Returns: +// - string: A sanitized filename for the log file +func (l *FileRequestLogger) generateFilename(url string, requestID ...string) string { + // Extract path from URL + path := url + if strings.Contains(url, "?") { + path = strings.Split(url, "?")[0] + } + + // Remove leading slash + if strings.HasPrefix(path, "/") { + path = path[1:] + } + + // Sanitize path for filename + sanitized := l.sanitizeForFilename(path) + + // Add timestamp + timestamp := time.Now().Format("2006-01-02T150405") + + // Use request ID if provided, otherwise use sequential ID + var idPart string + if len(requestID) > 0 && requestID[0] != "" { + idPart = requestID[0] + } else { + id := requestLogID.Add(1) + idPart = fmt.Sprintf("%d", id) + } + + return fmt.Sprintf("%s-%s-%s.log", sanitized, timestamp, idPart) +} + +// sanitizeForFilename replaces characters that are not safe for filenames. +// +// Parameters: +// - path: The path to sanitize +// +// Returns: +// - string: A sanitized filename +func (l *FileRequestLogger) sanitizeForFilename(path string) string { + // Replace slashes with hyphens + sanitized := strings.ReplaceAll(path, "/", "-") + + // Replace colons with hyphens + sanitized = strings.ReplaceAll(sanitized, ":", "-") + + // Replace other problematic characters with hyphens + reg := regexp.MustCompile(`[<>:"|?*\s]`) + sanitized = reg.ReplaceAllString(sanitized, "-") + + // Remove multiple consecutive hyphens + reg = regexp.MustCompile(`-+`) + sanitized = reg.ReplaceAllString(sanitized, "-") + + // Remove leading/trailing hyphens + sanitized = strings.Trim(sanitized, "-") + + // Handle empty result + if sanitized == "" { + sanitized = "root" + } + + return sanitized +} + +// cleanupOldErrorLogs keeps only the newest errorLogsMaxFiles forced error log files. +func (l *FileRequestLogger) cleanupOldErrorLogs() error { + if l.errorLogsMaxFiles <= 0 { + return nil + } + + entries, errRead := os.ReadDir(l.logsDir) + if errRead != nil { + return errRead + } + + type logFile struct { + name string + modTime time.Time + } + + var files []logFile + for _, entry := range entries { + if entry.IsDir() { + continue + } + name := entry.Name() + if !strings.HasPrefix(name, "error-") || !strings.HasSuffix(name, ".log") { + continue + } + info, errInfo := entry.Info() + if errInfo != nil { + log.WithError(errInfo).Warn("failed to read error log info") + continue + } + files = append(files, logFile{name: name, modTime: info.ModTime()}) + } + + if len(files) <= l.errorLogsMaxFiles { + return nil + } + + sort.Slice(files, func(i, j int) bool { + return files[i].modTime.After(files[j].modTime) + }) + + for _, file := range files[l.errorLogsMaxFiles:] { + if errRemove := os.Remove(filepath.Join(l.logsDir, file.name)); errRemove != nil { + log.WithError(errRemove).Warnf("failed to remove old error log: %s", file.name) + } + } + + return nil +} + +func (l *FileRequestLogger) writeRequestBodyTempFile(body []byte) (string, error) { + tmpFile, errCreate := os.CreateTemp(l.logsDir, "request-body-*.tmp") + if errCreate != nil { + return "", errCreate + } + tmpPath := tmpFile.Name() + + if _, errCopy := io.Copy(tmpFile, bytes.NewReader(body)); errCopy != nil { + _ = tmpFile.Close() + _ = os.Remove(tmpPath) + return "", errCopy + } + if errClose := tmpFile.Close(); errClose != nil { + _ = os.Remove(tmpPath) + return "", errClose + } + return tmpPath, nil +} diff --git a/internal/pluginhost/adapters.go b/internal/pluginhost/adapters.go index 385ffa16e..542fc6b14 100644 --- a/internal/pluginhost/adapters.go +++ b/internal/pluginhost/adapters.go @@ -1,24 +1,12 @@ package pluginhost import ( - "bytes" "context" - "encoding/json" "fmt" - "io" - "net/http" - "net/url" - "reflect" - "runtime/debug" - "sort" "strings" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" - sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" - coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" - coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" _ "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator/builtin" @@ -511,1875 +499,3 @@ func (h *Host) callModelsForAuth(ctx context.Context, record capabilityRecord, p HTTPClient: h.newHTTPClient(auth), }) } - -func (h *Host) callRequestInterceptor(ctx context.Context, record capabilityRecord, method string, call func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error), req pluginapi.RequestInterceptRequest) (out pluginapi.RequestInterceptResponse, ok bool) { - if h == nil || call == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { - return pluginapi.RequestInterceptResponse{}, false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, method, recovered) - out = pluginapi.RequestInterceptResponse{} - ok = false - } - }() - resp, errIntercept := call(ctx, req) - if errIntercept != nil { - log.Warnf("pluginhost: request interceptor %s failed: %v", record.id, errIntercept) - return pluginapi.RequestInterceptResponse{}, false - } - return resp, true -} - -func (h *Host) callResponseInterceptor(ctx context.Context, record capabilityRecord, interceptor pluginapi.ResponseInterceptor, req pluginapi.ResponseInterceptRequest) (out pluginapi.ResponseInterceptResponse, ok bool) { - if h == nil || interceptor == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { - return pluginapi.ResponseInterceptResponse{}, false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, "ResponseInterceptor.InterceptResponse", recovered) - out = pluginapi.ResponseInterceptResponse{} - ok = false - } - }() - resp, errIntercept := interceptor.InterceptResponse(ctx, req) - if errIntercept != nil { - log.Warnf("pluginhost: response interceptor %s failed: %v", record.id, errIntercept) - return pluginapi.ResponseInterceptResponse{}, false - } - return resp, true -} - -func (h *Host) callStreamChunkInterceptor(ctx context.Context, record capabilityRecord, interceptor pluginapi.StreamChunkInterceptor, req pluginapi.StreamChunkInterceptRequest) (out pluginapi.StreamChunkInterceptResponse, ok bool) { - if h == nil || interceptor == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { - return pluginapi.StreamChunkInterceptResponse{}, false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, "StreamChunkInterceptor.InterceptStreamChunk", recovered) - out = pluginapi.StreamChunkInterceptResponse{} - ok = false - } - }() - resp, errIntercept := interceptor.InterceptStreamChunk(ctx, req) - if errIntercept != nil { - log.Warnf("pluginhost: stream chunk interceptor %s failed: %v", record.id, errIntercept) - return pluginapi.StreamChunkInterceptResponse{}, false - } - return resp, true -} - -func (h *Host) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { - return h.InterceptRequestBeforeAuthExcept(ctx, req, "") -} - -func (h *Host) InterceptRequestBeforeAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { - return h.interceptRequest(ctx, req, "RequestInterceptor.InterceptRequestBeforeAuth", func(interceptor pluginapi.RequestInterceptor, ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { - return interceptor.InterceptRequestBeforeAuth(ctx, req) - }, skipPluginID) -} - -func (h *Host) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { - return h.InterceptRequestAfterAuthExcept(ctx, req, "") -} - -func (h *Host) InterceptRequestAfterAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { - return h.interceptRequest(ctx, req, "RequestInterceptor.InterceptRequestAfterAuth", func(interceptor pluginapi.RequestInterceptor, ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { - return interceptor.InterceptRequestAfterAuth(ctx, req) - }, skipPluginID) -} - -func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, method string, invoke func(pluginapi.RequestInterceptor, context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error), skipPluginID string) pluginapi.RequestInterceptResponse { - current := pluginapi.RequestInterceptResponse{ - Headers: cloneHeader(req.Headers), - Body: bytes.Clone(req.Body), - } - skipPluginID = strings.TrimSpace(skipPluginID) - for _, record := range h.activeRecords() { - interceptor := record.plugin.Capabilities.RequestInterceptor - if h.isPluginFused(record.id) || interceptor == nil || record.id == skipPluginID { - continue - } - nextReq := req - nextReq.Headers = cloneHeader(current.Headers) - nextReq.Body = bytes.Clone(current.Body) - nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) - if resp, ok := h.callRequestInterceptor(ctx, record, method, func(callCtx context.Context, callReq pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { - return invoke(interceptor, callCtx, callReq) - }, nextReq); ok { - current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) - if len(resp.Body) > 0 { - current.Body = bytes.Clone(resp.Body) - } - } - } - return current -} - -func (h *Host) InterceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { - return h.InterceptResponseExcept(ctx, req, "") -} - -func (h *Host) InterceptResponseExcept(ctx context.Context, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse { - current := pluginapi.ResponseInterceptResponse{ - Headers: cloneHeader(req.ResponseHeaders), - Body: bytes.Clone(req.Body), - } - skipPluginID = strings.TrimSpace(skipPluginID) - for _, record := range h.activeRecords() { - interceptor := record.plugin.Capabilities.ResponseInterceptor - if h.isPluginFused(record.id) || interceptor == nil || record.id == skipPluginID { - continue - } - nextReq := req - nextReq.RequestHeaders = cloneHeader(req.RequestHeaders) - nextReq.ResponseHeaders = cloneHeader(current.Headers) - nextReq.OriginalRequest = bytes.Clone(req.OriginalRequest) - nextReq.RequestBody = bytes.Clone(req.RequestBody) - nextReq.Body = bytes.Clone(current.Body) - nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) - if resp, ok := h.callResponseInterceptor(ctx, record, interceptor, nextReq); ok { - current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) - if len(resp.Body) > 0 { - current.Body = bytes.Clone(resp.Body) - } - } - } - return current -} - -func (h *Host) InterceptStreamChunk(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { - return h.InterceptStreamChunkExcept(ctx, req, "") -} - -func (h *Host) InterceptStreamChunkExcept(ctx context.Context, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse { - current := pluginapi.StreamChunkInterceptResponse{ - Headers: cloneHeader(req.ResponseHeaders), - Body: bytes.Clone(req.Body), - } - skipPluginID = strings.TrimSpace(skipPluginID) - for _, record := range h.activeRecords() { - interceptor := record.plugin.Capabilities.StreamChunkInterceptor - if h.isPluginFused(record.id) || interceptor == nil || current.DropChunk || record.id == skipPluginID { - continue - } - nextReq := req - nextReq.RequestHeaders = cloneHeader(req.RequestHeaders) - nextReq.ResponseHeaders = cloneHeader(current.Headers) - nextReq.OriginalRequest = bytes.Clone(req.OriginalRequest) - nextReq.RequestBody = bytes.Clone(req.RequestBody) - nextReq.Body = bytes.Clone(current.Body) - nextReq.HistoryChunks = cloneByteSlices(req.HistoryChunks) - nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) - if resp, ok := h.callStreamChunkInterceptor(ctx, record, interceptor, nextReq); ok { - current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) - if len(resp.Body) > 0 { - current.Body = bytes.Clone(resp.Body) - } - if resp.DropChunk { - current.DropChunk = true - } - } - } - return current -} - -func (h *Host) HasStreamInterceptors() bool { - if h == nil { - return false - } - for _, record := range h.activeRecords() { - if h.isPluginFused(record.id) { - continue - } - if record.plugin.Capabilities.StreamChunkInterceptor != nil { - return true - } - } - return false -} - -func (h *Host) HasRequestInterceptors() bool { - if h == nil { - return false - } - for _, record := range h.activeRecords() { - if h.isPluginFused(record.id) { - continue - } - if record.plugin.Capabilities.RequestInterceptor != nil { - return true - } - } - return false -} - -func (h *Host) commitModelClients(snap *Snapshot, modelRegistry modelRegistry, registrations []modelClientRegistration, nextClients map[string]struct{}, nextProviders map[string]string, nextModelRegistrations map[string]pluginModelRegistration) { - if h == nil || modelRegistry == nil { - return - } - - staleClients := make([]string, 0) - h.mu.Lock() - if h.Snapshot() != snap { - h.mu.Unlock() - return - } - for clientID := range h.modelClientIDs { - if _, okClient := nextClients[clientID]; !okClient { - staleClients = append(staleClients, clientID) - } - } - h.modelClientIDs = nextClients - h.modelProviders = nextProviders - h.modelRegistrations = nextModelRegistrations - h.mu.Unlock() - - for _, registration := range registrations { - modelRegistry.RegisterClient(registration.clientID, registration.provider, registration.models) - } - for _, clientID := range staleClients { - modelRegistry.UnregisterClient(clientID) - } -} - -type executorManager interface { - Executor(provider string) (coreauth.ProviderExecutor, bool) - RegisterExecutor(coreauth.ProviderExecutor) - UnregisterExecutor(provider string) -} - -type executorRegistration struct { - provider string - adapter *executorAdapter -} - -func (h *Host) RegisterExecutors(manager executorManager, modelRegistry modelProviderRegistry) { - if h == nil || manager == nil { - return - } - - snap := h.Snapshot() - records := h.activeRecordsFromSnapshot(snap) - registrations := h.snapshotModelRegistrations() - selectedModels := make(map[string][]*registry.ModelInfo) - providerModels := make(map[string][]*registry.ModelInfo) - claimedModels := make(map[string]struct{}) - claimedProviders := make(map[string]string) - for _, registration := range registrations { - if !registration.hasExecutor { - appendModelsForProvider(providerModels, registration.provider, registration.models) - } - } - for _, record := range records { - executor := record.plugin.Capabilities.Executor - if executor == nil || h.isPluginFused(record.id) { - continue - } - provider, okProvider := h.executorProvider(record, executor) - if !okProvider { - continue - } - registration := h.modelRegistration(record.id) - if h.providerHasNativeExecutor(manager, provider) { - appendModelsForProvider(providerModels, provider, registration.models) - continue - } - if len(registration.models) == 0 { - continue - } - if owner := claimedProviders[provider]; owner != "" && owner != record.id { - continue - } - for _, model := range registration.models { - modelID := strings.TrimSpace(model.ID) - if modelID == "" { - continue - } - if _, claimed := claimedModels[modelID]; claimed { - continue - } - if h.modelHasNativeExecutor(manager, modelRegistry, modelID) { - continue - } - claimedModels[modelID] = struct{}{} - claimedProviders[provider] = record.id - selectedModels[record.id] = append(selectedModels[record.id], model) - } - } - - seenProviders := make(map[string]struct{}) - nextProviders := make(map[string]struct{}) - nextModelClients := make(map[string]struct{}) - executorRegistrations := make([]executorRegistration, 0) - modelClientRegistrations := make([]modelClientRegistration, 0) - for _, record := range records { - executor := record.plugin.Capabilities.Executor - if executor == nil || h.isPluginFused(record.id) { - continue - } - - provider, okProvider := h.executorProvider(record, executor) - if !okProvider { - continue - } - registration := h.modelRegistration(record.id) - if len(registration.models) > 0 && len(selectedModels[record.id]) == 0 { - continue - } - if _, seenProvider := seenProviders[provider]; seenProvider { - continue - } - seenProviders[provider] = struct{}{} - if h.providerHasNativeExecutor(manager, provider) { - continue - } - - nextProviders[provider] = struct{}{} - executorRegistrations = append(executorRegistrations, newExecutorAdapterRegistration(h, record, provider, executor)) - appendModelsForProvider(providerModels, provider, selectedModels[record.id]) - if len(selectedModels[record.id]) > 0 { - clientID := pluginExecutorModelClientID(record.id, provider) - modelClientRegistrations = append(modelClientRegistrations, modelClientRegistration{ - clientID: clientID, - provider: provider, - models: selectedModels[record.id], - }) - nextModelClients[clientID] = struct{}{} - } - } - h.commitExecutorState(snap, manager, modelRegistry, providerModels, executorRegistrations, nextProviders, modelClientRegistrations, nextModelClients) -} - -func pluginExecutorModelClientID(pluginID, provider string) string { - return "plugin:" + pluginID + ":" + provider + ":executor" -} - -func (h *Host) commitExecutorState(snap *Snapshot, manager executorManager, modelRegistry modelRegistry, providerModels map[string][]*registry.ModelInfo, registrations []executorRegistration, nextProviders map[string]struct{}, modelClientRegistrations []modelClientRegistration, nextModelClients map[string]struct{}) { - if h == nil || manager == nil { - return - } - - h.mu.Lock() - if h.Snapshot() != snap { - h.mu.Unlock() - return - } - - h.providerModels = make(map[string][]*registryModelInfo, len(providerModels)) - for provider, models := range providerModels { - h.providerModels[provider] = cloneRegistryModels(models) - } - - staleProviders := make([]string, 0) - for provider := range h.executorProviders { - if _, okProvider := nextProviders[provider]; !okProvider { - staleProviders = append(staleProviders, provider) - } - } - h.executorProviders = nextProviders - if nextModelClients == nil { - nextModelClients = make(map[string]struct{}) - } - staleModelClients := make([]string, 0) - for clientID := range h.executorModelClientIDs { - if _, okClient := nextModelClients[clientID]; !okClient { - staleModelClients = append(staleModelClients, clientID) - } - } - h.executorModelClientIDs = nextModelClients - - for _, registration := range registrations { - if registration.adapter == nil || registration.provider == "" { - continue - } - manager.RegisterExecutor(registration.adapter) - } - for _, provider := range staleProviders { - existing, okExecutor := manager.Executor(provider) - if !okExecutor || !h.ownsExecutor(existing) { - continue - } - manager.UnregisterExecutor(provider) - } - h.mu.Unlock() - - if modelRegistry == nil { - return - } - for _, registration := range modelClientRegistrations { - modelRegistry.RegisterClient(registration.clientID, registration.provider, registration.models) - } - for _, clientID := range staleModelClients { - modelRegistry.UnregisterClient(clientID) - } -} - -func newExecutorAdapterRegistration(h *Host, record capabilityRecord, provider string, executor pluginapi.ProviderExecutor) executorRegistration { - return executorRegistration{ - provider: provider, - adapter: &executorAdapter{ - host: h, - pluginID: record.id, - path: record.path, - version: record.version, - provider: provider, - executor: executor, - inputFormats: normalizeExecutorFormats(record.plugin.Capabilities.ExecutorInputFormats), - outputFormats: normalizeExecutorFormats(record.plugin.Capabilities.ExecutorOutputFormats), - }, - } -} - -func (h *Host) snapshotModelRegistrations() []pluginModelRegistration { - if h == nil { - return nil - } - h.mu.Lock() - defer h.mu.Unlock() - registrations := make([]pluginModelRegistration, 0, len(h.modelRegistrations)) - for _, registration := range h.modelRegistrations { - registration.models = cloneRegistryModels(registration.models) - registrations = append(registrations, registration) - } - sort.SliceStable(registrations, func(i, j int) bool { - if registrations[i].priority == registrations[j].priority { - return registrations[i].pluginID < registrations[j].pluginID - } - return registrations[i].priority > registrations[j].priority - }) - return registrations -} - -func (h *Host) modelRegistration(pluginID string) pluginModelRegistration { - if h == nil { - return pluginModelRegistration{} - } - h.mu.Lock() - defer h.mu.Unlock() - registration := h.modelRegistrations[pluginID] - registration.models = cloneRegistryModels(registration.models) - return registration -} - -func (h *Host) executorProvider(record capabilityRecord, executor pluginapi.ProviderExecutor) (string, bool) { - if h == nil || !h.recordCurrent(record) { - return "", false - } - provider := h.modelProvider(record.id) - if provider == "" { - identifier, okIdentifier := h.callExecutorIdentifier(record.id, executor) - if !okIdentifier { - return "", false - } - provider = identifier - } - provider = strings.ToLower(strings.TrimSpace(provider)) - return provider, provider != "" -} - -func (h *Host) callExecutorIdentifier(pluginID string, executor pluginapi.ProviderExecutor) (provider string, ok bool) { - if h == nil || executor == nil || h.isPluginFused(pluginID) { - return "", false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(pluginID, "Executor.Identifier", recovered) - provider = "" - ok = false - } - }() - return executor.Identifier(), true -} - -func (h *Host) providerHasNativeExecutor(manager executorManager, provider string) bool { - if h == nil || manager == nil { - return false - } - existing, okExecutor := manager.Executor(provider) - return okExecutor && existing != nil && !h.ownsExecutor(existing) -} - -func (h *Host) modelHasNativeExecutor(manager executorManager, modelRegistry modelProviderRegistry, modelID string) bool { - if h == nil || manager == nil || modelRegistry == nil { - return false - } - for _, provider := range modelRegistry.GetModelProviders(modelID) { - if h.providerHasNativeExecutor(manager, provider) { - return true - } - } - return false -} - -func appendModelsForProvider(out map[string][]*registry.ModelInfo, provider string, models []*registry.ModelInfo) { - provider = strings.ToLower(strings.TrimSpace(provider)) - if provider == "" || len(models) == 0 { - return - } - seen := make(map[string]struct{}, len(out[provider])+len(models)) - for _, model := range out[provider] { - if model != nil && strings.TrimSpace(model.ID) != "" { - seen[strings.TrimSpace(model.ID)] = struct{}{} - } - } - for _, model := range models { - if model == nil { - continue - } - modelID := strings.TrimSpace(model.ID) - if modelID == "" { - continue - } - if _, exists := seen[modelID]; exists { - continue - } - seen[modelID] = struct{}{} - out[provider] = append(out[provider], cloneRegistryModels([]*registry.ModelInfo{model})...) - } -} - -func (h *Host) ModelsForProvider(provider string) []*registry.ModelInfo { - if h == nil { - return nil - } - provider = strings.ToLower(strings.TrimSpace(provider)) - if provider == "" { - return nil - } - h.mu.Lock() - defer h.mu.Unlock() - return cloneRegistryModels(h.providerModels[provider]) -} - -func (h *Host) HasExecutorCandidateProvider(provider string) bool { - if h == nil { - return false - } - provider = strings.ToLower(strings.TrimSpace(provider)) - if provider == "" { - return false - } - for _, record := range h.activeRecords() { - executor := record.plugin.Capabilities.Executor - if executor == nil || h.isPluginFused(record.id) { - continue - } - candidate, okCandidate := h.executorProvider(record, executor) - if okCandidate && candidate == provider { - return true - } - } - return false -} - -func (h *Host) ownsExecutor(executor coreauth.ProviderExecutor) bool { - adapter, okAdapter := executor.(*executorAdapter) - return okAdapter && adapter != nil && adapter.host == h -} - -func (h *Host) modelProvider(pluginID string) string { - if h == nil { - return "" - } - h.mu.Lock() - defer h.mu.Unlock() - return h.modelProviders[pluginID] -} - -func (h *Host) RegisterFrontendAuthProviders() { - if h == nil { - return - } - - type exclusiveFrontendAuthCandidate struct { - key string - pluginID string - priority int - } - - nextKeys := make(map[string]struct{}) - var bestExclusive exclusiveFrontendAuthCandidate - for _, record := range h.activeRecords() { - provider := record.plugin.Capabilities.FrontendAuthProvider - if provider == nil || h.isPluginFused(record.id) { - continue - } - adapter := &accessAdapter{ - host: h, - pluginID: record.id, - path: record.path, - version: record.version, - provider: provider, - } - key := strings.TrimSpace(adapter.Identifier()) - if key == "" { - continue - } - sdkaccess.RegisterProvider(key, adapter) - nextKeys[key] = struct{}{} - if record.plugin.Capabilities.FrontendAuthProviderExclusive { - candidate := exclusiveFrontendAuthCandidate{ - key: key, - pluginID: record.id, - priority: record.priority, - } - if bestExclusive.key == "" || - candidate.priority > bestExclusive.priority || - (candidate.priority == bestExclusive.priority && candidate.pluginID < bestExclusive.pluginID) { - bestExclusive = candidate - } - } - } - - if bestExclusive.key != "" { - sdkaccess.SetExclusiveProvider(bestExclusive.key) - } else { - sdkaccess.ClearExclusiveProvider() - } - h.pruneStaleAccessProviders(nextKeys) -} - -func (h *Host) pruneStaleAccessProviders(nextKeys map[string]struct{}) { - if h == nil { - return - } - - staleKeys := make([]string, 0) - h.mu.Lock() - for key := range h.accessProviderKeys { - if _, okKey := nextKeys[key]; !okKey { - staleKeys = append(staleKeys, key) - } - } - h.accessProviderKeys = nextKeys - h.mu.Unlock() - - for _, key := range staleKeys { - sdkaccess.UnregisterProvider(key) - } -} - -func (h *Host) RegisterUsagePlugins() { - if h == nil { - return - } - - for _, record := range h.activeRecords() { - plugin := record.plugin.Capabilities.UsagePlugin - if plugin == nil || h.isPluginFused(record.id) { - continue - } - coreusage.RegisterNamedPlugin("plugin:"+record.id, &usageAdapter{ - host: h, - pluginID: record.id, - plugin: plugin, - }) - } -} - -func (h *Host) refreshThinkingProviders(records []capabilityRecord) { - thinking.ClearPluginProviders() - if h == nil { - return - } - for _, record := range records { - applier := record.plugin.Capabilities.ThinkingApplier - if applier == nil || h.isPluginFused(record.id) { - continue - } - provider, okProvider := h.callThinkingIdentifier(record, applier) - if !okProvider { - continue - } - thinking.RegisterPluginProvider(record.id, provider, record.priority, &thinkingAdapter{ - host: h, - pluginID: record.id, - path: record.path, - version: record.version, - provider: provider, - applier: applier, - }) - } -} - -func (h *Host) callThinkingIdentifier(record capabilityRecord, applier pluginapi.ThinkingApplier) (provider string, ok bool) { - if h == nil || applier == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { - return "", false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, "ThinkingApplier.Identifier", recovered) - provider = "" - ok = false - } - }() - provider = strings.ToLower(strings.TrimSpace(applier.Identifier())) - if provider == "" { - return "", false - } - return provider, true -} - -func (h *Host) currentUsagePlugin(pluginID string) pluginapi.UsagePlugin { - if h == nil || strings.TrimSpace(pluginID) == "" { - return nil - } - for _, record := range h.activeRecords() { - if record.id != pluginID { - continue - } - if h.isPluginFused(record.id) { - return nil - } - return record.plugin.Capabilities.UsagePlugin - } - return nil -} - -func (h *Host) fusePlugin(id, method string, recovered any) { - if h == nil { - return - } - h.mu.Lock() - h.fused[id] = fmt.Sprintf("%s panic: %v", method, recovered) - h.mu.Unlock() - thinking.UnregisterPluginProviders(id) - log.WithField("plugin_id", id).WithField("method", method).Errorf("pluginhost: plugin panic recovered: %v\n%s", recovered, debug.Stack()) -} - -func (h *Host) isPluginFused(id string) bool { - if h == nil { - return false - } - h.mu.Lock() - _, fused := h.fused[id] - h.mu.Unlock() - return fused -} - -type accessAdapter struct { - host *Host - pluginID string - path string - version string - provider pluginapi.FrontendAuthProvider -} - -func (a *accessAdapter) Identifier() (identifier string) { - if a == nil || a.provider == nil { - return "" - } - defer func() { - if recovered := recover(); recovered != nil { - if a.host != nil { - a.host.fusePlugin(a.pluginID, "FrontendAuthProvider.Identifier", recovered) - } - identifier = "" - } - }() - pluginID := strings.TrimSpace(a.pluginID) - providerID := strings.TrimSpace(a.provider.Identifier()) - if pluginID == "" || providerID == "" { - return "" - } - return "plugin:" + pluginID + ":" + providerID -} - -func (a *accessAdapter) Authenticate(ctx context.Context, r *http.Request) (result *sdkaccess.Result, authErr *sdkaccess.AuthError) { - if a == nil || a.provider == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { - return nil, sdkaccess.NewNotHandledError() - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(a.pluginID, "FrontendAuthProvider.Authenticate", recovered) - result = nil - authErr = sdkaccess.NewNotHandledError() - } - }() - - body, errReadAll := readAndRestoreRequestBody(r) - if errReadAll != nil { - return nil, sdkaccess.NewInternalAuthError("failed to read plugin auth request body", errReadAll) - } - resp, errAuthenticate := a.provider.Authenticate(ctx, pluginapi.FrontendAuthRequest{ - Method: r.Method, - Path: r.URL.Path, - Headers: cloneHeader(r.Header), - Query: cloneValues(r.URL.Query()), - Body: bytes.Clone(body), - }) - if errAuthenticate != nil || !resp.Authenticated { - return nil, sdkaccess.NewNotHandledError() - } - providerID := a.Identifier() - if providerID == "" { - return nil, sdkaccess.NewNotHandledError() - } - return &sdkaccess.Result{ - Provider: providerID, - Principal: resp.Principal, - Metadata: cloneStringMap(resp.Metadata), - }, nil -} - -type executorAdapter struct { - host *Host - pluginID string - path string - version string - provider string - executor pluginapi.ProviderExecutor - inputFormats []sdktranslator.Format - outputFormats []sdktranslator.Format -} - -func (a *executorAdapter) Identifier() string { - if a == nil { - return "" - } - return a.provider -} - -type preparedExecutorCall struct { - req coreexecutor.Request - opts coreexecutor.Options - inputRequested sdktranslator.Format - requestedFormat sdktranslator.Format - inputFormat sdktranslator.Format - outputFormat sdktranslator.Format -} - -func (a *executorAdapter) prepareExecutorCall(req coreexecutor.Request, opts coreexecutor.Options) (preparedExecutorCall, error) { - inputRequested := executorInputFormat(req, opts) - requestedFormat := executorRequestedFormat(req, opts) - inputFormat, errInput := a.selectExecutorInputFormat(inputRequested) - if errInput != nil { - return preparedExecutorCall{}, errInput - } - outputFormat, errOutput := a.selectExecutorOutputFormat(requestedFormat, inputFormat) - if errOutput != nil { - return preparedExecutorCall{}, errOutput - } - - nativeReq := req - nativeOpts := opts - if inputRequested != "" && inputRequested != inputFormat { - nativeReq.Payload = sdktranslator.TranslateRequest(inputRequested, inputFormat, req.Model, req.Payload, opts.Stream) - } - nativeReq.Format = outputFormat - nativeOpts.SourceFormat = inputFormat - nativeOpts.ResponseFormat = outputFormat - - return preparedExecutorCall{ - req: nativeReq, - opts: nativeOpts, - inputRequested: inputRequested, - requestedFormat: requestedFormat, - inputFormat: inputFormat, - outputFormat: outputFormat, - }, nil -} - -func (a *executorAdapter) RequestToFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { - if a == nil { - return "" - } - inputRequested := executorInputFormat(req, opts) - inputFormat, errInput := a.selectExecutorInputFormat(inputRequested) - if errInput != nil { - return "" - } - return inputFormat -} - -func executorInputFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { - if opts.SourceFormat != "" { - return normalizeExecutorFormatName(opts.SourceFormat.String()) - } - if req.Format != "" { - return normalizeExecutorFormatName(req.Format.String()) - } - return sdktranslator.FormatOpenAI -} - -func executorRequestedFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { - if format := coreexecutor.ResponseFormatOrSource(opts); format != "" { - return normalizeExecutorFormatName(format.String()) - } - if req.Format != "" { - return normalizeExecutorFormatName(req.Format.String()) - } - return sdktranslator.FormatOpenAI -} - -func (a *executorAdapter) selectExecutorInputFormat(requested sdktranslator.Format) (sdktranslator.Format, error) { - if len(a.inputFormats) == 0 { - return "", fmt.Errorf("plugin executor %s declares no input formats", a.Identifier()) - } - if executorFormatContains(a.inputFormats, requested) { - return requested, nil - } - for _, format := range a.inputFormats { - if requested == "" || sdktranslator.HasRequestTransformer(requested, format) { - return format, nil - } - } - return "", fmt.Errorf("plugin executor %s does not support input format %q", a.Identifier(), requested) -} - -func (a *executorAdapter) selectExecutorOutputFormat(requested, inputFormat sdktranslator.Format) (sdktranslator.Format, error) { - if len(a.outputFormats) == 0 { - return "", fmt.Errorf("plugin executor %s declares no output formats", a.Identifier()) - } - if executorFormatContains(a.outputFormats, requested) { - return requested, nil - } - if executorFormatContains(a.outputFormats, inputFormat) && a.executorResponseTranslationAvailable(inputFormat, requested) { - return inputFormat, nil - } - for _, format := range a.outputFormats { - if requested == "" || a.executorResponseTranslationAvailable(format, requested) { - return format, nil - } - } - return "", fmt.Errorf("plugin executor %s does not support output format %q", a.Identifier(), requested) -} - -func (a *executorAdapter) executorResponseTranslationAvailable(from, to sdktranslator.Format) bool { - if from == "" || to == "" || from == to { - return true - } - if sdktranslator.HasResponseTransformer(to, from) { - return true - } - return a != nil && a.host.hasResponseTranslator() -} - -func (h *Host) hasResponseTranslator() bool { - for _, record := range h.activeRecords() { - if h.isPluginFused(record.id) || record.plugin.Capabilities.ResponseTranslator == nil { - continue - } - return true - } - return false -} - -func executorNativeStreamResponseTranslatorExists(from, to sdktranslator.Format) bool { - if from == "" || to == "" || from == to { - return true - } - return sdktranslator.HasStreamResponseTransformer(to, from) -} - -func (a *executorAdapter) translateExecutorResponse(ctx context.Context, prepared preparedExecutorCall, payload []byte, stream bool, param *any) []byte { - if prepared.requestedFormat == "" || prepared.outputFormat == prepared.requestedFormat { - return bytes.Clone(payload) - } - originalRequest := prepared.opts.OriginalRequest - if len(originalRequest) == 0 { - originalRequest = prepared.req.Payload - } - if stream { - frames := a.translateExecutorStreamPayload(ctx, prepared, payload, param) - if len(frames) == 0 { - return nil - } - if len(frames) == 1 { - return bytes.Clone(frames[0]) - } - return bytes.Join(frames, nil) - } - return sdktranslator.TranslateNonStream(ctx, prepared.outputFormat, prepared.requestedFormat, prepared.req.Model, originalRequest, prepared.req.Payload, payload, param) -} - -func (a *executorAdapter) translateExecutorStreamChunks(ctx context.Context, prepared preparedExecutorCall, in <-chan pluginapi.ExecutorStreamChunk) <-chan pluginapi.ExecutorStreamChunk { - if prepared.requestedFormat == "" || prepared.outputFormat == prepared.requestedFormat { - return in - } - if in == nil { - return nil - } - if ctx == nil { - ctx = context.Background() - } - out := make(chan pluginapi.ExecutorStreamChunk) - go func() { - defer close(out) - var param any - for { - select { - case <-ctx.Done(): - return - case chunk, ok := <-in: - if !ok { - a.emitTranslatedExecutorStreamTail(ctx, prepared, out, ¶m) - return - } - if chunk.Err != nil { - _ = sendExecutorPluginStreamChunk(ctx, out, chunk) - continue - } - frames := a.translateExecutorStreamPayload(ctx, prepared, chunk.Payload, ¶m) - for _, frame := range frames { - if !sendExecutorPluginStreamChunk(ctx, out, pluginapi.ExecutorStreamChunk{Payload: frame}) { - return - } - } - } - } - }() - return out -} - -func (a *executorAdapter) translateExecutorStreamPayload(ctx context.Context, prepared preparedExecutorCall, payload []byte, param *any) [][]byte { - originalRequest := prepared.opts.OriginalRequest - if len(originalRequest) == 0 { - originalRequest = prepared.req.Payload - } - frames := sdktranslator.TranslateStream(ctx, prepared.outputFormat, prepared.requestedFormat, prepared.req.Model, originalRequest, prepared.req.Payload, payload, param) - if executorStreamTranslationFellBack(prepared, payload, frames) { - return nil - } - return frames -} - -func executorStreamTranslationFellBack(prepared preparedExecutorCall, payload []byte, frames [][]byte) bool { - if prepared.requestedFormat == "" || prepared.outputFormat == "" || prepared.outputFormat == prepared.requestedFormat { - return false - } - if len(frames) != 1 || !bytes.Equal(frames[0], payload) { - return false - } - // A plugin executor only reaches this path after host-side response translation - // has been selected. An unchanged single frame is the SDK registry fallback, - // not a valid translated frame to send to the client. - return executorNativeStreamResponseTranslatorExists(prepared.outputFormat, prepared.requestedFormat) -} - -func (a *executorAdapter) emitTranslatedExecutorStreamTail(ctx context.Context, prepared preparedExecutorCall, out chan<- pluginapi.ExecutorStreamChunk, param *any) { - tail := executorStreamDonePayload(prepared.outputFormat) - if len(tail) == 0 { - return - } - frames := a.translateExecutorStreamPayload(ctx, prepared, tail, param) - for _, frame := range frames { - if !sendExecutorPluginStreamChunk(ctx, out, pluginapi.ExecutorStreamChunk{Payload: frame}) { - return - } - } -} - -func executorStreamDonePayload(format sdktranslator.Format) []byte { - switch format { - case sdktranslator.FormatOpenAI: - return []byte("data: [DONE]") - default: - return nil - } -} - -func sendExecutorPluginStreamChunk(ctx context.Context, out chan<- pluginapi.ExecutorStreamChunk, chunk pluginapi.ExecutorStreamChunk) bool { - select { - case out <- pluginapi.ExecutorStreamChunk{Payload: bytes.Clone(chunk.Payload), Err: chunk.Err}: - return true - case <-ctx.Done(): - return false - } -} - -func (a *executorAdapter) Execute(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (resp coreexecutor.Response, err error) { - if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { - return coreexecutor.Response{}, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(a.pluginID, "Executor.Execute", recovered) - resp = coreexecutor.Response{} - err = fmt.Errorf("plugin executor %s panic: %v", a.Identifier(), recovered) - } - }() - - prepared, errPrepare := a.prepareExecutorCall(req, opts) - if errPrepare != nil { - return coreexecutor.Response{}, errPrepare - } - pluginResp, errExecute := a.executor.Execute(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) - if errExecute != nil { - return coreexecutor.Response{}, errExecute - } - return coreexecutor.Response{ - Payload: a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil), - Metadata: cloneAnyMap(pluginResp.Metadata), - Headers: cloneHeader(pluginResp.Headers), - }, nil -} - -func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (result *coreexecutor.StreamResult, err error) { - if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { - return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(a.pluginID, "Executor.ExecuteStream", recovered) - result = nil - err = fmt.Errorf("plugin executor %s stream panic: %v", a.Identifier(), recovered) - } - }() - - prepared, errPrepare := a.prepareExecutorCall(req, opts) - if errPrepare != nil { - return nil, errPrepare - } - pluginResp, errExecuteStream := a.executor.ExecuteStream(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) - if errExecuteStream != nil { - return nil, errExecuteStream - } - return &coreexecutor.StreamResult{ - Headers: cloneHeader(pluginResp.Headers), - Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, pluginResp.Chunks)), - }, nil -} - -func (a *executorAdapter) Refresh(ctx context.Context, auth *coreauth.Auth) (refreshed *coreauth.Auth, err error) { - if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { - return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) - } - record := a.host.authProviderRecord(authProvider(auth)) - if record == nil || record.plugin.Capabilities.AuthProvider == nil { - return auth.Clone(), nil - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(record.id, "AuthProvider.RefreshAuth", recovered) - refreshed = nil - err = fmt.Errorf("plugin executor %s refresh panic: %v", a.Identifier(), recovered) - } - }() - - pluginResp, errRefresh := record.plugin.Capabilities.AuthProvider.RefreshAuth(ctx, pluginapi.AuthRefreshRequest{ - AuthID: authID(auth), - AuthProvider: authProvider(auth), - StorageJSON: storageJSONFromAuth(auth), - Metadata: cloneAnyMap(authMetadata(auth)), - Attributes: authAttributes(auth), - Host: a.host.hostConfigSummary(), - HTTPClient: a.host.newHTTPClient(auth), - }) - if errRefresh != nil { - return nil, errRefresh - } - data := pluginResp.Auth - if strings.TrimSpace(data.Provider) == "" { - data.Provider = authProvider(auth) - } - if strings.TrimSpace(data.ID) == "" { - data.ID = authID(auth) - } - if strings.TrimSpace(data.FileName) == "" && auth != nil { - data.FileName = auth.FileName - } - if strings.TrimSpace(data.Label) == "" && auth != nil { - data.Label = auth.Label - } - if strings.TrimSpace(data.Prefix) == "" && auth != nil { - data.Prefix = auth.Prefix - } - if strings.TrimSpace(data.ProxyURL) == "" && auth != nil { - data.ProxyURL = auth.ProxyURL - } - if len(data.Metadata) == 0 && auth != nil { - data.Metadata = cloneAnyMap(auth.Metadata) - } - if len(data.Attributes) == 0 && auth != nil { - data.Attributes = cloneStringMap(auth.Attributes) - } - if len(data.StorageJSON) == 0 { - data.StorageJSON = storageJSONFromAuth(auth) - } - if pluginResp.NextRefreshAfter.IsZero() && auth != nil { - data.NextRefreshAfter = auth.NextRefreshAfter - } - if !pluginResp.NextRefreshAfter.IsZero() { - data.NextRefreshAfter = pluginResp.NextRefreshAfter - } - next := a.host.AuthDataToCoreAuth(data, "", data.FileName) - if next == nil { - return nil, fmt.Errorf("plugin executor %s refresh returned invalid auth data", a.Identifier()) - } - if auth != nil { - next.CreatedAt = auth.CreatedAt - next.UpdatedAt = auth.UpdatedAt - } - return next, nil -} - -func (a *executorAdapter) CountTokens(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (resp coreexecutor.Response, err error) { - if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { - return coreexecutor.Response{}, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(a.pluginID, "Executor.CountTokens", recovered) - resp = coreexecutor.Response{} - err = fmt.Errorf("plugin executor %s count tokens panic: %v", a.Identifier(), recovered) - } - }() - - prepared, errPrepare := a.prepareExecutorCall(req, opts) - if errPrepare != nil { - return coreexecutor.Response{}, errPrepare - } - pluginResp, errCountTokens := a.executor.CountTokens(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) - if errCountTokens != nil { - return coreexecutor.Response{}, errCountTokens - } - return coreexecutor.Response{ - Payload: a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil), - Metadata: cloneAnyMap(pluginResp.Metadata), - Headers: cloneHeader(pluginResp.Headers), - }, nil -} - -func (a *executorAdapter) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (resp *http.Response, err error) { - if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { - return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) - } - if req == nil { - return nil, fmt.Errorf("plugin executor %s received nil HTTP request", a.Identifier()) - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(a.pluginID, "Executor.HttpRequest", recovered) - resp = nil - err = fmt.Errorf("plugin executor %s http request panic: %v", a.Identifier(), recovered) - } - }() - body, errReadAll := readAndRestoreRequestBody(req) - if errReadAll != nil { - return nil, fmt.Errorf("read plugin http request body: %w", errReadAll) - } - pluginResp, errHTTPRequest := a.executor.HttpRequest(ctx, pluginapi.ExecutorHTTPRequest{ - AuthID: authID(auth), - AuthProvider: authProvider(auth), - Method: req.Method, - URL: req.URL.String(), - Headers: cloneHeader(req.Header), - Body: bytes.Clone(body), - StorageJSON: storageJSONFromAuth(auth), - Metadata: cloneAnyMap(authMetadata(auth)), - Attributes: authAttributes(auth), - HTTPClient: a.host.newHTTPClient(auth, a.provider), - }) - if errHTTPRequest != nil { - return nil, errHTTPRequest - } - status := pluginResp.StatusCode - if status == 0 { - status = http.StatusOK - } - resp = &http.Response{ - StatusCode: status, - Status: fmt.Sprintf("%d %s", status, http.StatusText(status)), - Header: cloneHeader(pluginResp.Headers), - Body: io.NopCloser(bytes.NewReader(bytes.Clone(pluginResp.Body))), - Request: req, - } - return resp, nil -} - -type usageAdapter struct { - host *Host - pluginID string - plugin pluginapi.UsagePlugin -} - -type thinkingAdapter struct { - host *Host - pluginID string - path string - version string - provider string - applier pluginapi.ThinkingApplier -} - -func (a *usageAdapter) HandleUsage(ctx context.Context, record coreusage.Record) { - if a == nil { - return - } - plugin := a.host.currentUsagePlugin(a.pluginID) - if plugin == nil { - return - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(a.pluginID, "UsagePlugin.HandleUsage", recovered) - } - }() - plugin.HandleUsage(ctx, pluginapi.UsageRecord{ - Provider: record.Provider, - ExecutorType: record.ExecutorType, - Model: record.Model, - Alias: record.Alias, - APIKey: record.APIKey, - AuthID: record.AuthID, - AuthIndex: record.AuthIndex, - AuthType: record.AuthType, - Source: record.Source, - ReasoningEffort: record.ReasoningEffort, - ServiceTier: record.ServiceTier, - Generate: coreusage.GenerateEnabled(record.Generate), - RequestedAt: record.RequestedAt, - Latency: record.Latency, - TTFT: record.TTFT, - Failed: record.Failed, - Failure: pluginapi.UsageFailure{ - StatusCode: record.Fail.StatusCode, - Body: record.Fail.Body, - }, - Detail: pluginapi.UsageDetail{ - InputTokens: record.Detail.InputTokens, - OutputTokens: record.Detail.OutputTokens, - ReasoningTokens: record.Detail.ReasoningTokens, - CachedTokens: record.Detail.CachedTokens, - CacheReadTokens: record.Detail.CacheReadTokens, - CacheCreationTokens: record.Detail.CacheCreationTokens, - TotalTokens: record.Detail.TotalTokens, - }, - ResponseHeaders: cloneHeader(record.ResponseHeaders), - }) -} - -func (a *thinkingAdapter) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) (out []byte, err error) { - if a == nil || a.applier == nil || a.host == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { - return bytes.Clone(body), nil - } - defer func() { - if recovered := recover(); recovered != nil { - a.host.fusePlugin(a.pluginID, "ThinkingApplier.ApplyThinking", recovered) - out = bytes.Clone(body) - err = nil - } - }() - resp, errApply := a.applier.ApplyThinking(context.Background(), pluginapi.ThinkingApplyRequest{ - Provider: a.provider, - Model: registryModelInfoToPluginModelInfo(modelInfo), - Config: pluginapi.ThinkingConfig{ - Mode: config.Mode.String(), - Budget: config.Budget, - Level: string(config.Level), - }, - Body: bytes.Clone(body), - }) - if errApply != nil || len(resp.Body) == 0 { - return bytes.Clone(body), nil - } - return bytes.Clone(resp.Body), nil -} - -func (h *Host) NormalizeRequest(ctx context.Context, from, to sdktranslator.Format, model string, body []byte, stream bool) []byte { - current := bytes.Clone(body) - for _, record := range h.activeRecords() { - if h.isPluginFused(record.id) || record.plugin.Capabilities.RequestNormalizer == nil { - continue - } - if normalized, ok := h.callRequestNormalizer(ctx, record, from, to, model, current, stream); ok { - current = normalized - } - } - return current -} - -func (h *Host) TranslateRequest(ctx context.Context, from, to sdktranslator.Format, model string, body []byte, stream bool) ([]byte, bool) { - for _, record := range h.activeRecords() { - if h.isPluginFused(record.id) || record.plugin.Capabilities.RequestTranslator == nil { - continue - } - if translated, ok := h.callRequestTranslator(ctx, record, from, to, model, body, stream); ok { - return translated, true - } - } - return bytes.Clone(body), false -} - -func (h *Host) NormalizeResponseBefore(ctx context.Context, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) []byte { - current := bytes.Clone(body) - for _, record := range h.activeRecords() { - normalizer := record.plugin.Capabilities.ResponseBeforeTranslator - if h.isPluginFused(record.id) || normalizer == nil { - continue - } - if normalized, ok := h.callResponseNormalizer(ctx, record, "ResponseBeforeTranslator.NormalizeResponse", normalizer, from, to, model, originalRequestRawJSON, requestRawJSON, current, stream); ok { - current = normalized - } - } - return current -} - -func (h *Host) TranslateResponse(ctx context.Context, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) ([]byte, bool) { - for _, record := range h.activeRecords() { - translator := record.plugin.Capabilities.ResponseTranslator - if h.isPluginFused(record.id) || translator == nil { - continue - } - if translated, ok := h.callResponseTranslator(ctx, record, translator, from, to, model, originalRequestRawJSON, requestRawJSON, body, stream); ok { - return translated, true - } - } - return bytes.Clone(body), false -} - -func (h *Host) NormalizeResponseAfter(ctx context.Context, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) []byte { - current := bytes.Clone(body) - for _, record := range h.activeRecords() { - normalizer := record.plugin.Capabilities.ResponseAfterTranslator - if h.isPluginFused(record.id) || normalizer == nil { - continue - } - if normalized, ok := h.callResponseNormalizer(ctx, record, "ResponseAfterTranslator.NormalizeResponse", normalizer, from, to, model, originalRequestRawJSON, requestRawJSON, current, stream); ok { - current = normalized - } - } - return current -} - -func (h *Host) callRequestNormalizer(ctx context.Context, record capabilityRecord, from, to sdktranslator.Format, model string, body []byte, stream bool) (out []byte, ok bool) { - if h == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) || record.plugin.Capabilities.RequestNormalizer == nil { - return nil, false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, "RequestNormalizer.NormalizeRequest", recovered) - out = nil - ok = false - } - }() - resp, errNormalizeRequest := record.plugin.Capabilities.RequestNormalizer.NormalizeRequest(ctx, pluginapi.RequestTransformRequest{ - FromFormat: from.String(), - ToFormat: to.String(), - Model: model, - Stream: stream, - Body: bytes.Clone(body), - }) - if errNormalizeRequest != nil || len(resp.Body) == 0 { - return nil, false - } - return bytes.Clone(resp.Body), true -} - -func (h *Host) callRequestTranslator(ctx context.Context, record capabilityRecord, from, to sdktranslator.Format, model string, body []byte, stream bool) (out []byte, ok bool) { - if h == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) || record.plugin.Capabilities.RequestTranslator == nil { - return nil, false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, "RequestTranslator.TranslateRequest", recovered) - out = nil - ok = false - } - }() - resp, errTranslateRequest := record.plugin.Capabilities.RequestTranslator.TranslateRequest(ctx, pluginapi.RequestTransformRequest{ - FromFormat: from.String(), - ToFormat: to.String(), - Model: model, - Stream: stream, - Body: bytes.Clone(body), - }) - if errTranslateRequest != nil || len(resp.Body) == 0 { - return nil, false - } - return bytes.Clone(resp.Body), true -} - -func (h *Host) callResponseNormalizer(ctx context.Context, record capabilityRecord, method string, normalizer pluginapi.ResponseNormalizer, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) (out []byte, ok bool) { - if h == nil || normalizer == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { - return nil, false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, method, recovered) - out = nil - ok = false - } - }() - resp, errNormalizeResponse := normalizer.NormalizeResponse(ctx, pluginapi.ResponseTransformRequest{ - FromFormat: from.String(), - ToFormat: to.String(), - Model: model, - Stream: stream, - OriginalRequest: bytes.Clone(originalRequestRawJSON), - TranslatedRequest: bytes.Clone(requestRawJSON), - Body: bytes.Clone(body), - }) - if errNormalizeResponse != nil || len(resp.Body) == 0 { - return nil, false - } - return bytes.Clone(resp.Body), true -} - -func (h *Host) callResponseTranslator(ctx context.Context, record capabilityRecord, translator pluginapi.ResponseTranslator, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) (out []byte, ok bool) { - if h == nil || translator == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { - return nil, false - } - defer func() { - if recovered := recover(); recovered != nil { - h.fusePlugin(record.id, "ResponseTranslator.TranslateResponse", recovered) - out = nil - ok = false - } - }() - resp, errTranslateResponse := translator.TranslateResponse(ctx, pluginapi.ResponseTransformRequest{ - FromFormat: from.String(), - ToFormat: to.String(), - Model: model, - Stream: stream, - OriginalRequest: bytes.Clone(originalRequestRawJSON), - TranslatedRequest: bytes.Clone(requestRawJSON), - Body: bytes.Clone(body), - }) - if errTranslateResponse != nil || len(resp.Body) == 0 { - return nil, false - } - return bytes.Clone(resp.Body), true -} - -func buildExecutorRequest(host *Host, provider string, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) pluginapi.ExecutorRequest { - return pluginapi.ExecutorRequest{ - AuthID: authID(auth), - AuthProvider: authProvider(auth), - Model: req.Model, - Format: req.Format.String(), - Stream: opts.Stream, - Alt: opts.Alt, - Headers: cloneHeader(opts.Headers), - Query: cloneValues(opts.Query), - OriginalRequest: bytes.Clone(opts.OriginalRequest), - SourceFormat: opts.SourceFormat.String(), - Payload: bytes.Clone(req.Payload), - Metadata: mergeExecutorMetadata(req.Metadata, opts.Metadata), - StorageJSON: storageJSONFromAuth(auth), - AuthMetadata: cloneAnyMap(authMetadata(auth)), - AuthAttributes: authAttributes(auth), - HTTPClient: host.newHTTPClient(auth, provider), - } -} - -func storageJSONFromAuth(auth *coreauth.Auth) []byte { - if auth == nil { - return nil - } - if rawProvider, okRaw := auth.Storage.(interface{ RawJSON() []byte }); okRaw { - return bytes.Clone(rawProvider.RawJSON()) - } - if len(auth.Metadata) == 0 { - return nil - } - data, errMarshal := json.Marshal(auth.Metadata) - if errMarshal != nil { - return nil - } - return data -} - -func authAttributes(auth *coreauth.Auth) map[string]string { - if auth == nil { - return nil - } - return cloneStringMap(auth.Attributes) -} - -func mergeExecutorMetadata(reqMetadata, optsMetadata map[string]any) map[string]any { - if len(reqMetadata) == 0 && len(optsMetadata) == 0 { - return nil - } - merged := make(map[string]any, len(reqMetadata)+len(optsMetadata)) - for key, value := range reqMetadata { - merged[key] = value - } - for key, value := range optsMetadata { - merged[key] = value - } - return merged -} - -func mapExecutorStreamChunks(ctx context.Context, in <-chan pluginapi.ExecutorStreamChunk) <-chan coreexecutor.StreamChunk { - if ctx == nil { - ctx = context.Background() - } - out := make(chan coreexecutor.StreamChunk) - if in == nil { - close(out) - return out - } - go func() { - defer close(out) - for { - var mapped coreexecutor.StreamChunk - select { - case <-ctx.Done(): - return - case chunk, ok := <-in: - if !ok { - return - } - mapped = coreexecutor.StreamChunk{ - Payload: bytes.Clone(chunk.Payload), - Err: chunk.Err, - } - } - select { - case <-ctx.Done(): - return - case out <- mapped: - } - } - }() - return out -} - -func readAndRestoreRequestBody(r *http.Request) ([]byte, error) { - if r == nil || r.Body == nil { - return nil, nil - } - body, errReadAll := io.ReadAll(r.Body) - if errReadAll != nil { - r.Body = io.NopCloser(bytes.NewReader(body)) - return nil, errReadAll - } - r.Body = io.NopCloser(bytes.NewReader(body)) - return body, nil -} - -func authID(auth *coreauth.Auth) string { - if auth == nil { - return "" - } - return auth.ID -} - -func authProvider(auth *coreauth.Auth) string { - if auth == nil { - return "" - } - return auth.Provider -} - -func authMetadata(auth *coreauth.Auth) map[string]any { - if auth == nil { - return nil - } - return auth.Metadata -} - -func cloneHeader(in http.Header) http.Header { - if len(in) == 0 { - return nil - } - out := make(http.Header, len(in)) - for key, values := range in { - out[key] = append([]string(nil), values...) - } - return out -} - -func mergeHeaders(current, updates http.Header, clear []string) http.Header { - out := cloneHeader(current) - if out == nil { - out = make(http.Header) - } - for _, key := range clear { - out.Del(key) - } - for key, values := range updates { - out.Del(key) - for _, value := range values { - out.Add(key, value) - } - } - return out -} - -func cloneByteSlices(in [][]byte) [][]byte { - if len(in) == 0 { - return nil - } - out := make([][]byte, 0, len(in)) - for _, item := range in { - out = append(out, bytes.Clone(item)) - } - return out -} - -func cloneValues(in url.Values) url.Values { - if len(in) == 0 { - return nil - } - out := make(url.Values, len(in)) - for key, values := range in { - out[key] = append([]string(nil), values...) - } - return out -} - -func cloneAnyMap(in map[string]any) map[string]any { - if len(in) == 0 { - return nil - } - out := make(map[string]any, len(in)) - for key, value := range in { - out[key] = value - } - return out -} - -func cloneInterceptorMetadata(in map[string]any) map[string]any { - if len(in) == 0 { - return nil - } - visited := make(map[metadataCloneVisit]reflect.Value) - out := make(map[string]any, len(in)) - for key, value := range in { - out[key] = cloneInterceptorMetadataAny(reflect.ValueOf(value), visited) - } - return out -} - -type metadataCloneVisit struct { - typ reflect.Type - ptr uintptr -} - -func cloneInterceptorMetadataAny(value reflect.Value, visited map[metadataCloneVisit]reflect.Value) any { - cloned := cloneInterceptorMetadataReflectValue(value, visited) - if !cloned.IsValid() { - return nil - } - return cloned.Interface() -} - -func cloneInterceptorMetadataReflectValue(value reflect.Value, visited map[metadataCloneVisit]reflect.Value) reflect.Value { - if !value.IsValid() { - return reflect.Value{} - } - - switch value.Kind() { - case reflect.Interface: - if value.IsNil() { - return reflect.Zero(value.Type()) - } - return cloneInterceptorMetadataReflectValue(value.Elem(), visited) - case reflect.Pointer: - if value.IsNil() { - return reflect.Zero(value.Type()) - } - visit := metadataCloneVisit{typ: value.Type(), ptr: value.Pointer()} - if existing, okExisting := visited[visit]; okExisting { - return existing - } - out := reflect.New(value.Type().Elem()) - visited[visit] = out - clonedElem := cloneInterceptorMetadataReflectValue(value.Elem(), visited) - if clonedElem.IsValid() { - outElem := out.Elem() - if clonedElem.Type().AssignableTo(outElem.Type()) { - outElem.Set(clonedElem) - } else if clonedElem.Type().ConvertibleTo(outElem.Type()) { - outElem.Set(clonedElem.Convert(outElem.Type())) - } - } - return out - case reflect.Map: - if value.IsNil() { - return reflect.Zero(value.Type()) - } - visit := metadataCloneVisit{typ: value.Type(), ptr: value.Pointer()} - if existing, okExisting := visited[visit]; okExisting { - return existing - } - out := reflect.MakeMapWithSize(value.Type(), value.Len()) - visited[visit] = out - iter := value.MapRange() - for iter.Next() { - keyValue := adaptClonedValue(iter.Key(), cloneInterceptorMetadataReflectValue(iter.Key(), visited)) - valValue := adaptClonedValue(iter.Value(), cloneInterceptorMetadataReflectValue(iter.Value(), visited)) - out.SetMapIndex(keyValue, valValue) - } - return out - case reflect.Slice: - if value.IsNil() { - return reflect.Zero(value.Type()) - } - if value.Type().Elem().Kind() == reflect.Uint8 { - out := reflect.MakeSlice(value.Type(), value.Len(), value.Len()) - reflect.Copy(out, value) - return out - } - visit := metadataCloneVisit{typ: value.Type(), ptr: value.Pointer()} - if existing, okExisting := visited[visit]; okExisting { - return existing - } - out := reflect.MakeSlice(value.Type(), value.Len(), value.Len()) - visited[visit] = out - for i := 0; i < value.Len(); i++ { - clonedItem := cloneInterceptorMetadataReflectValue(value.Index(i), visited) - if !clonedItem.IsValid() { - continue - } - out.Index(i).Set(adaptClonedValue(value.Index(i), clonedItem)) - } - return out - case reflect.Array: - out := reflect.New(value.Type()).Elem() - for i := 0; i < value.Len(); i++ { - clonedItem := cloneInterceptorMetadataReflectValue(value.Index(i), visited) - if !clonedItem.IsValid() { - continue - } - out.Index(i).Set(adaptClonedValue(value.Index(i), clonedItem)) - } - return out - case reflect.Struct: - out := reflect.New(value.Type()).Elem() - // Preserve unexported fields and deep-clone exported fields on a best-effort basis. - out.Set(value) - for i := 0; i < value.NumField(); i++ { - field := value.Field(i) - if !out.Field(i).CanSet() { - continue - } - fieldClone := cloneInterceptorMetadataReflectValue(field, visited) - if !fieldClone.IsValid() { - continue - } - out.Field(i).Set(adaptClonedValue(field, fieldClone)) - } - return out - default: - return value - } -} - -func adaptClonedValue(original, cloned reflect.Value) reflect.Value { - if !cloned.IsValid() { - return original - } - if cloned.Type().AssignableTo(original.Type()) { - return cloned - } - if cloned.Type().ConvertibleTo(original.Type()) { - return cloned.Convert(original.Type()) - } - return original -} - -func cloneStringMap(in map[string]string) map[string]string { - if len(in) == 0 { - return nil - } - out := make(map[string]string, len(in)) - for key, value := range in { - out[key] = value - } - return out -} diff --git a/internal/pluginhost/adapters_auth.go b/internal/pluginhost/adapters_auth.go new file mode 100644 index 000000000..bb4c54a17 --- /dev/null +++ b/internal/pluginhost/adapters_auth.go @@ -0,0 +1,149 @@ +package pluginhost + +import ( + "bytes" + "context" + "net/http" + "strings" + + sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func (h *Host) RegisterFrontendAuthProviders() { + if h == nil { + return + } + + type exclusiveFrontendAuthCandidate struct { + key string + pluginID string + priority int + } + + nextKeys := make(map[string]struct{}) + var bestExclusive exclusiveFrontendAuthCandidate + for _, record := range h.activeRecords() { + provider := record.plugin.Capabilities.FrontendAuthProvider + if provider == nil || h.isPluginFused(record.id) { + continue + } + adapter := &accessAdapter{ + host: h, + pluginID: record.id, + path: record.path, + version: record.version, + provider: provider, + } + key := strings.TrimSpace(adapter.Identifier()) + if key == "" { + continue + } + sdkaccess.RegisterProvider(key, adapter) + nextKeys[key] = struct{}{} + if record.plugin.Capabilities.FrontendAuthProviderExclusive { + candidate := exclusiveFrontendAuthCandidate{ + key: key, + pluginID: record.id, + priority: record.priority, + } + if bestExclusive.key == "" || + candidate.priority > bestExclusive.priority || + (candidate.priority == bestExclusive.priority && candidate.pluginID < bestExclusive.pluginID) { + bestExclusive = candidate + } + } + } + + if bestExclusive.key != "" { + sdkaccess.SetExclusiveProvider(bestExclusive.key) + } else { + sdkaccess.ClearExclusiveProvider() + } + h.pruneStaleAccessProviders(nextKeys) +} + +func (h *Host) pruneStaleAccessProviders(nextKeys map[string]struct{}) { + if h == nil { + return + } + + staleKeys := make([]string, 0) + h.mu.Lock() + for key := range h.accessProviderKeys { + if _, okKey := nextKeys[key]; !okKey { + staleKeys = append(staleKeys, key) + } + } + h.accessProviderKeys = nextKeys + h.mu.Unlock() + + for _, key := range staleKeys { + sdkaccess.UnregisterProvider(key) + } +} + +type accessAdapter struct { + host *Host + pluginID string + path string + version string + provider pluginapi.FrontendAuthProvider +} + +func (a *accessAdapter) Identifier() (identifier string) { + if a == nil || a.provider == nil { + return "" + } + defer func() { + if recovered := recover(); recovered != nil { + if a.host != nil { + a.host.fusePlugin(a.pluginID, "FrontendAuthProvider.Identifier", recovered) + } + identifier = "" + } + }() + pluginID := strings.TrimSpace(a.pluginID) + providerID := strings.TrimSpace(a.provider.Identifier()) + if pluginID == "" || providerID == "" { + return "" + } + return "plugin:" + pluginID + ":" + providerID +} + +func (a *accessAdapter) Authenticate(ctx context.Context, r *http.Request) (result *sdkaccess.Result, authErr *sdkaccess.AuthError) { + if a == nil || a.provider == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { + return nil, sdkaccess.NewNotHandledError() + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(a.pluginID, "FrontendAuthProvider.Authenticate", recovered) + result = nil + authErr = sdkaccess.NewNotHandledError() + } + }() + + body, errReadAll := readAndRestoreRequestBody(r) + if errReadAll != nil { + return nil, sdkaccess.NewInternalAuthError("failed to read plugin auth request body", errReadAll) + } + resp, errAuthenticate := a.provider.Authenticate(ctx, pluginapi.FrontendAuthRequest{ + Method: r.Method, + Path: r.URL.Path, + Headers: cloneHeader(r.Header), + Query: cloneValues(r.URL.Query()), + Body: bytes.Clone(body), + }) + if errAuthenticate != nil || !resp.Authenticated { + return nil, sdkaccess.NewNotHandledError() + } + providerID := a.Identifier() + if providerID == "" { + return nil, sdkaccess.NewNotHandledError() + } + return &sdkaccess.Result{ + Provider: providerID, + Principal: resp.Principal, + Metadata: cloneStringMap(resp.Metadata), + }, nil +} diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go new file mode 100644 index 000000000..3ec4863ea --- /dev/null +++ b/internal/pluginhost/adapters_executors.go @@ -0,0 +1,922 @@ +package pluginhost + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sort" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +type executorManager interface { + Executor(provider string) (coreauth.ProviderExecutor, bool) + RegisterExecutor(coreauth.ProviderExecutor) + UnregisterExecutor(provider string) +} + +type executorRegistration struct { + provider string + adapter *executorAdapter +} + +func (h *Host) RegisterExecutors(manager executorManager, modelRegistry modelProviderRegistry) { + if h == nil || manager == nil { + return + } + + snap := h.Snapshot() + records := h.activeRecordsFromSnapshot(snap) + registrations := h.snapshotModelRegistrations() + selectedModels := make(map[string][]*registry.ModelInfo) + providerModels := make(map[string][]*registry.ModelInfo) + claimedModels := make(map[string]struct{}) + claimedProviders := make(map[string]string) + for _, registration := range registrations { + if !registration.hasExecutor { + appendModelsForProvider(providerModels, registration.provider, registration.models) + } + } + for _, record := range records { + executor := record.plugin.Capabilities.Executor + if executor == nil || h.isPluginFused(record.id) { + continue + } + provider, okProvider := h.executorProvider(record, executor) + if !okProvider { + continue + } + registration := h.modelRegistration(record.id) + if h.providerHasNativeExecutor(manager, provider) { + appendModelsForProvider(providerModels, provider, registration.models) + continue + } + if len(registration.models) == 0 { + continue + } + if owner := claimedProviders[provider]; owner != "" && owner != record.id { + continue + } + for _, model := range registration.models { + modelID := strings.TrimSpace(model.ID) + if modelID == "" { + continue + } + if _, claimed := claimedModels[modelID]; claimed { + continue + } + if h.modelHasNativeExecutor(manager, modelRegistry, modelID) { + continue + } + claimedModels[modelID] = struct{}{} + claimedProviders[provider] = record.id + selectedModels[record.id] = append(selectedModels[record.id], model) + } + } + + seenProviders := make(map[string]struct{}) + nextProviders := make(map[string]struct{}) + nextModelClients := make(map[string]struct{}) + executorRegistrations := make([]executorRegistration, 0) + modelClientRegistrations := make([]modelClientRegistration, 0) + for _, record := range records { + executor := record.plugin.Capabilities.Executor + if executor == nil || h.isPluginFused(record.id) { + continue + } + + provider, okProvider := h.executorProvider(record, executor) + if !okProvider { + continue + } + registration := h.modelRegistration(record.id) + if len(registration.models) > 0 && len(selectedModels[record.id]) == 0 { + continue + } + if _, seenProvider := seenProviders[provider]; seenProvider { + continue + } + seenProviders[provider] = struct{}{} + if h.providerHasNativeExecutor(manager, provider) { + continue + } + + nextProviders[provider] = struct{}{} + executorRegistrations = append(executorRegistrations, newExecutorAdapterRegistration(h, record, provider, executor)) + appendModelsForProvider(providerModels, provider, selectedModels[record.id]) + if len(selectedModels[record.id]) > 0 { + clientID := pluginExecutorModelClientID(record.id, provider) + modelClientRegistrations = append(modelClientRegistrations, modelClientRegistration{ + clientID: clientID, + provider: provider, + models: selectedModels[record.id], + }) + nextModelClients[clientID] = struct{}{} + } + } + h.commitExecutorState(snap, manager, modelRegistry, providerModels, executorRegistrations, nextProviders, modelClientRegistrations, nextModelClients) +} + +func pluginExecutorModelClientID(pluginID, provider string) string { + return "plugin:" + pluginID + ":" + provider + ":executor" +} + +func (h *Host) commitExecutorState(snap *Snapshot, manager executorManager, modelRegistry modelRegistry, providerModels map[string][]*registry.ModelInfo, registrations []executorRegistration, nextProviders map[string]struct{}, modelClientRegistrations []modelClientRegistration, nextModelClients map[string]struct{}) { + if h == nil || manager == nil { + return + } + + h.mu.Lock() + if h.Snapshot() != snap { + h.mu.Unlock() + return + } + + h.providerModels = make(map[string][]*registryModelInfo, len(providerModels)) + for provider, models := range providerModels { + h.providerModels[provider] = cloneRegistryModels(models) + } + + staleProviders := make([]string, 0) + for provider := range h.executorProviders { + if _, okProvider := nextProviders[provider]; !okProvider { + staleProviders = append(staleProviders, provider) + } + } + h.executorProviders = nextProviders + if nextModelClients == nil { + nextModelClients = make(map[string]struct{}) + } + staleModelClients := make([]string, 0) + for clientID := range h.executorModelClientIDs { + if _, okClient := nextModelClients[clientID]; !okClient { + staleModelClients = append(staleModelClients, clientID) + } + } + h.executorModelClientIDs = nextModelClients + + for _, registration := range registrations { + if registration.adapter == nil || registration.provider == "" { + continue + } + manager.RegisterExecutor(registration.adapter) + } + for _, provider := range staleProviders { + existing, okExecutor := manager.Executor(provider) + if !okExecutor || !h.ownsExecutor(existing) { + continue + } + manager.UnregisterExecutor(provider) + } + h.mu.Unlock() + + if modelRegistry == nil { + return + } + for _, registration := range modelClientRegistrations { + modelRegistry.RegisterClient(registration.clientID, registration.provider, registration.models) + } + for _, clientID := range staleModelClients { + modelRegistry.UnregisterClient(clientID) + } +} + +func newExecutorAdapterRegistration(h *Host, record capabilityRecord, provider string, executor pluginapi.ProviderExecutor) executorRegistration { + return executorRegistration{ + provider: provider, + adapter: &executorAdapter{ + host: h, + pluginID: record.id, + path: record.path, + version: record.version, + provider: provider, + executor: executor, + inputFormats: normalizeExecutorFormats(record.plugin.Capabilities.ExecutorInputFormats), + outputFormats: normalizeExecutorFormats(record.plugin.Capabilities.ExecutorOutputFormats), + }, + } +} + +func (h *Host) snapshotModelRegistrations() []pluginModelRegistration { + if h == nil { + return nil + } + h.mu.Lock() + defer h.mu.Unlock() + registrations := make([]pluginModelRegistration, 0, len(h.modelRegistrations)) + for _, registration := range h.modelRegistrations { + registration.models = cloneRegistryModels(registration.models) + registrations = append(registrations, registration) + } + sort.SliceStable(registrations, func(i, j int) bool { + if registrations[i].priority == registrations[j].priority { + return registrations[i].pluginID < registrations[j].pluginID + } + return registrations[i].priority > registrations[j].priority + }) + return registrations +} + +func (h *Host) modelRegistration(pluginID string) pluginModelRegistration { + if h == nil { + return pluginModelRegistration{} + } + h.mu.Lock() + defer h.mu.Unlock() + registration := h.modelRegistrations[pluginID] + registration.models = cloneRegistryModels(registration.models) + return registration +} + +func (h *Host) executorProvider(record capabilityRecord, executor pluginapi.ProviderExecutor) (string, bool) { + if h == nil || !h.recordCurrent(record) { + return "", false + } + provider := h.modelProvider(record.id) + if provider == "" { + identifier, okIdentifier := h.callExecutorIdentifier(record.id, executor) + if !okIdentifier { + return "", false + } + provider = identifier + } + provider = strings.ToLower(strings.TrimSpace(provider)) + return provider, provider != "" +} + +func (h *Host) callExecutorIdentifier(pluginID string, executor pluginapi.ProviderExecutor) (provider string, ok bool) { + if h == nil || executor == nil || h.isPluginFused(pluginID) { + return "", false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(pluginID, "Executor.Identifier", recovered) + provider = "" + ok = false + } + }() + return executor.Identifier(), true +} + +func (h *Host) providerHasNativeExecutor(manager executorManager, provider string) bool { + if h == nil || manager == nil { + return false + } + existing, okExecutor := manager.Executor(provider) + return okExecutor && existing != nil && !h.ownsExecutor(existing) +} + +func (h *Host) modelHasNativeExecutor(manager executorManager, modelRegistry modelProviderRegistry, modelID string) bool { + if h == nil || manager == nil || modelRegistry == nil { + return false + } + for _, provider := range modelRegistry.GetModelProviders(modelID) { + if h.providerHasNativeExecutor(manager, provider) { + return true + } + } + return false +} + +func appendModelsForProvider(out map[string][]*registry.ModelInfo, provider string, models []*registry.ModelInfo) { + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" || len(models) == 0 { + return + } + seen := make(map[string]struct{}, len(out[provider])+len(models)) + for _, model := range out[provider] { + if model != nil && strings.TrimSpace(model.ID) != "" { + seen[strings.TrimSpace(model.ID)] = struct{}{} + } + } + for _, model := range models { + if model == nil { + continue + } + modelID := strings.TrimSpace(model.ID) + if modelID == "" { + continue + } + if _, exists := seen[modelID]; exists { + continue + } + seen[modelID] = struct{}{} + out[provider] = append(out[provider], cloneRegistryModels([]*registry.ModelInfo{model})...) + } +} + +func (h *Host) ModelsForProvider(provider string) []*registry.ModelInfo { + if h == nil { + return nil + } + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" { + return nil + } + h.mu.Lock() + defer h.mu.Unlock() + return cloneRegistryModels(h.providerModels[provider]) +} + +func (h *Host) HasExecutorCandidateProvider(provider string) bool { + if h == nil { + return false + } + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" { + return false + } + for _, record := range h.activeRecords() { + executor := record.plugin.Capabilities.Executor + if executor == nil || h.isPluginFused(record.id) { + continue + } + candidate, okCandidate := h.executorProvider(record, executor) + if okCandidate && candidate == provider { + return true + } + } + return false +} + +func (h *Host) ownsExecutor(executor coreauth.ProviderExecutor) bool { + adapter, okAdapter := executor.(*executorAdapter) + return okAdapter && adapter != nil && adapter.host == h +} + +func (h *Host) modelProvider(pluginID string) string { + if h == nil { + return "" + } + h.mu.Lock() + defer h.mu.Unlock() + return h.modelProviders[pluginID] +} + +type executorAdapter struct { + host *Host + pluginID string + path string + version string + provider string + executor pluginapi.ProviderExecutor + inputFormats []sdktranslator.Format + outputFormats []sdktranslator.Format +} + +func (a *executorAdapter) Identifier() string { + if a == nil { + return "" + } + return a.provider +} + +type preparedExecutorCall struct { + req coreexecutor.Request + opts coreexecutor.Options + inputRequested sdktranslator.Format + requestedFormat sdktranslator.Format + inputFormat sdktranslator.Format + outputFormat sdktranslator.Format +} + +func (a *executorAdapter) prepareExecutorCall(req coreexecutor.Request, opts coreexecutor.Options) (preparedExecutorCall, error) { + inputRequested := executorInputFormat(req, opts) + requestedFormat := executorRequestedFormat(req, opts) + inputFormat, errInput := a.selectExecutorInputFormat(inputRequested) + if errInput != nil { + return preparedExecutorCall{}, errInput + } + outputFormat, errOutput := a.selectExecutorOutputFormat(requestedFormat, inputFormat) + if errOutput != nil { + return preparedExecutorCall{}, errOutput + } + + nativeReq := req + nativeOpts := opts + if inputRequested != "" && inputRequested != inputFormat { + nativeReq.Payload = sdktranslator.TranslateRequest(inputRequested, inputFormat, req.Model, req.Payload, opts.Stream) + } + nativeReq.Format = outputFormat + nativeOpts.SourceFormat = inputFormat + nativeOpts.ResponseFormat = outputFormat + + return preparedExecutorCall{ + req: nativeReq, + opts: nativeOpts, + inputRequested: inputRequested, + requestedFormat: requestedFormat, + inputFormat: inputFormat, + outputFormat: outputFormat, + }, nil +} + +func (a *executorAdapter) RequestToFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { + if a == nil { + return "" + } + inputRequested := executorInputFormat(req, opts) + inputFormat, errInput := a.selectExecutorInputFormat(inputRequested) + if errInput != nil { + return "" + } + return inputFormat +} + +func executorInputFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { + if opts.SourceFormat != "" { + return normalizeExecutorFormatName(opts.SourceFormat.String()) + } + if req.Format != "" { + return normalizeExecutorFormatName(req.Format.String()) + } + return sdktranslator.FormatOpenAI +} + +func executorRequestedFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { + if format := coreexecutor.ResponseFormatOrSource(opts); format != "" { + return normalizeExecutorFormatName(format.String()) + } + if req.Format != "" { + return normalizeExecutorFormatName(req.Format.String()) + } + return sdktranslator.FormatOpenAI +} + +func (a *executorAdapter) selectExecutorInputFormat(requested sdktranslator.Format) (sdktranslator.Format, error) { + if len(a.inputFormats) == 0 { + return "", fmt.Errorf("plugin executor %s declares no input formats", a.Identifier()) + } + if executorFormatContains(a.inputFormats, requested) { + return requested, nil + } + for _, format := range a.inputFormats { + if requested == "" || sdktranslator.HasRequestTransformer(requested, format) { + return format, nil + } + } + return "", fmt.Errorf("plugin executor %s does not support input format %q", a.Identifier(), requested) +} + +func (a *executorAdapter) selectExecutorOutputFormat(requested, inputFormat sdktranslator.Format) (sdktranslator.Format, error) { + if len(a.outputFormats) == 0 { + return "", fmt.Errorf("plugin executor %s declares no output formats", a.Identifier()) + } + if executorFormatContains(a.outputFormats, requested) { + return requested, nil + } + if executorFormatContains(a.outputFormats, inputFormat) && a.executorResponseTranslationAvailable(inputFormat, requested) { + return inputFormat, nil + } + for _, format := range a.outputFormats { + if requested == "" || a.executorResponseTranslationAvailable(format, requested) { + return format, nil + } + } + return "", fmt.Errorf("plugin executor %s does not support output format %q", a.Identifier(), requested) +} + +func (a *executorAdapter) executorResponseTranslationAvailable(from, to sdktranslator.Format) bool { + if from == "" || to == "" || from == to { + return true + } + if sdktranslator.HasResponseTransformer(to, from) { + return true + } + return a != nil && a.host.hasResponseTranslator() +} + +func (h *Host) hasResponseTranslator() bool { + for _, record := range h.activeRecords() { + if h.isPluginFused(record.id) || record.plugin.Capabilities.ResponseTranslator == nil { + continue + } + return true + } + return false +} + +func executorNativeStreamResponseTranslatorExists(from, to sdktranslator.Format) bool { + if from == "" || to == "" || from == to { + return true + } + return sdktranslator.HasStreamResponseTransformer(to, from) +} + +func (a *executorAdapter) translateExecutorResponse(ctx context.Context, prepared preparedExecutorCall, payload []byte, stream bool, param *any) []byte { + if prepared.requestedFormat == "" || prepared.outputFormat == prepared.requestedFormat { + return bytes.Clone(payload) + } + originalRequest := prepared.opts.OriginalRequest + if len(originalRequest) == 0 { + originalRequest = prepared.req.Payload + } + if stream { + frames := a.translateExecutorStreamPayload(ctx, prepared, payload, param) + if len(frames) == 0 { + return nil + } + if len(frames) == 1 { + return bytes.Clone(frames[0]) + } + return bytes.Join(frames, nil) + } + return sdktranslator.TranslateNonStream(ctx, prepared.outputFormat, prepared.requestedFormat, prepared.req.Model, originalRequest, prepared.req.Payload, payload, param) +} + +func (a *executorAdapter) translateExecutorStreamChunks(ctx context.Context, prepared preparedExecutorCall, in <-chan pluginapi.ExecutorStreamChunk) <-chan pluginapi.ExecutorStreamChunk { + if prepared.requestedFormat == "" || prepared.outputFormat == prepared.requestedFormat { + return in + } + if in == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + out := make(chan pluginapi.ExecutorStreamChunk) + go func() { + defer close(out) + var param any + for { + select { + case <-ctx.Done(): + return + case chunk, ok := <-in: + if !ok { + a.emitTranslatedExecutorStreamTail(ctx, prepared, out, ¶m) + return + } + if chunk.Err != nil { + _ = sendExecutorPluginStreamChunk(ctx, out, chunk) + continue + } + frames := a.translateExecutorStreamPayload(ctx, prepared, chunk.Payload, ¶m) + for _, frame := range frames { + if !sendExecutorPluginStreamChunk(ctx, out, pluginapi.ExecutorStreamChunk{Payload: frame}) { + return + } + } + } + } + }() + return out +} + +func (a *executorAdapter) translateExecutorStreamPayload(ctx context.Context, prepared preparedExecutorCall, payload []byte, param *any) [][]byte { + originalRequest := prepared.opts.OriginalRequest + if len(originalRequest) == 0 { + originalRequest = prepared.req.Payload + } + frames := sdktranslator.TranslateStream(ctx, prepared.outputFormat, prepared.requestedFormat, prepared.req.Model, originalRequest, prepared.req.Payload, payload, param) + if executorStreamTranslationFellBack(prepared, payload, frames) { + return nil + } + return frames +} + +func executorStreamTranslationFellBack(prepared preparedExecutorCall, payload []byte, frames [][]byte) bool { + if prepared.requestedFormat == "" || prepared.outputFormat == "" || prepared.outputFormat == prepared.requestedFormat { + return false + } + if len(frames) != 1 || !bytes.Equal(frames[0], payload) { + return false + } + // A plugin executor only reaches this path after host-side response translation + // has been selected. An unchanged single frame is the SDK registry fallback, + // not a valid translated frame to send to the client. + return executorNativeStreamResponseTranslatorExists(prepared.outputFormat, prepared.requestedFormat) +} + +func (a *executorAdapter) emitTranslatedExecutorStreamTail(ctx context.Context, prepared preparedExecutorCall, out chan<- pluginapi.ExecutorStreamChunk, param *any) { + tail := executorStreamDonePayload(prepared.outputFormat) + if len(tail) == 0 { + return + } + frames := a.translateExecutorStreamPayload(ctx, prepared, tail, param) + for _, frame := range frames { + if !sendExecutorPluginStreamChunk(ctx, out, pluginapi.ExecutorStreamChunk{Payload: frame}) { + return + } + } +} + +func executorStreamDonePayload(format sdktranslator.Format) []byte { + switch format { + case sdktranslator.FormatOpenAI: + return []byte("data: [DONE]") + default: + return nil + } +} + +func sendExecutorPluginStreamChunk(ctx context.Context, out chan<- pluginapi.ExecutorStreamChunk, chunk pluginapi.ExecutorStreamChunk) bool { + select { + case out <- pluginapi.ExecutorStreamChunk{Payload: bytes.Clone(chunk.Payload), Err: chunk.Err}: + return true + case <-ctx.Done(): + return false + } +} + +func (a *executorAdapter) Execute(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (resp coreexecutor.Response, err error) { + if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { + return coreexecutor.Response{}, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(a.pluginID, "Executor.Execute", recovered) + resp = coreexecutor.Response{} + err = fmt.Errorf("plugin executor %s panic: %v", a.Identifier(), recovered) + } + }() + + prepared, errPrepare := a.prepareExecutorCall(req, opts) + if errPrepare != nil { + return coreexecutor.Response{}, errPrepare + } + pluginResp, errExecute := a.executor.Execute(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) + if errExecute != nil { + return coreexecutor.Response{}, errExecute + } + return coreexecutor.Response{ + Payload: a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil), + Metadata: cloneAnyMap(pluginResp.Metadata), + Headers: cloneHeader(pluginResp.Headers), + }, nil +} + +func (a *executorAdapter) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (result *coreexecutor.StreamResult, err error) { + if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { + return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(a.pluginID, "Executor.ExecuteStream", recovered) + result = nil + err = fmt.Errorf("plugin executor %s stream panic: %v", a.Identifier(), recovered) + } + }() + + prepared, errPrepare := a.prepareExecutorCall(req, opts) + if errPrepare != nil { + return nil, errPrepare + } + pluginResp, errExecuteStream := a.executor.ExecuteStream(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) + if errExecuteStream != nil { + return nil, errExecuteStream + } + return &coreexecutor.StreamResult{ + Headers: cloneHeader(pluginResp.Headers), + Chunks: mapExecutorStreamChunks(ctx, a.translateExecutorStreamChunks(ctx, prepared, pluginResp.Chunks)), + }, nil +} + +func (a *executorAdapter) Refresh(ctx context.Context, auth *coreauth.Auth) (refreshed *coreauth.Auth, err error) { + if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { + return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) + } + record := a.host.authProviderRecord(authProvider(auth)) + if record == nil || record.plugin.Capabilities.AuthProvider == nil { + return auth.Clone(), nil + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(record.id, "AuthProvider.RefreshAuth", recovered) + refreshed = nil + err = fmt.Errorf("plugin executor %s refresh panic: %v", a.Identifier(), recovered) + } + }() + + pluginResp, errRefresh := record.plugin.Capabilities.AuthProvider.RefreshAuth(ctx, pluginapi.AuthRefreshRequest{ + AuthID: authID(auth), + AuthProvider: authProvider(auth), + StorageJSON: storageJSONFromAuth(auth), + Metadata: cloneAnyMap(authMetadata(auth)), + Attributes: authAttributes(auth), + Host: a.host.hostConfigSummary(), + HTTPClient: a.host.newHTTPClient(auth), + }) + if errRefresh != nil { + return nil, errRefresh + } + data := pluginResp.Auth + if strings.TrimSpace(data.Provider) == "" { + data.Provider = authProvider(auth) + } + if strings.TrimSpace(data.ID) == "" { + data.ID = authID(auth) + } + if strings.TrimSpace(data.FileName) == "" && auth != nil { + data.FileName = auth.FileName + } + if strings.TrimSpace(data.Label) == "" && auth != nil { + data.Label = auth.Label + } + if strings.TrimSpace(data.Prefix) == "" && auth != nil { + data.Prefix = auth.Prefix + } + if strings.TrimSpace(data.ProxyURL) == "" && auth != nil { + data.ProxyURL = auth.ProxyURL + } + if len(data.Metadata) == 0 && auth != nil { + data.Metadata = cloneAnyMap(auth.Metadata) + } + if len(data.Attributes) == 0 && auth != nil { + data.Attributes = cloneStringMap(auth.Attributes) + } + if len(data.StorageJSON) == 0 { + data.StorageJSON = storageJSONFromAuth(auth) + } + if pluginResp.NextRefreshAfter.IsZero() && auth != nil { + data.NextRefreshAfter = auth.NextRefreshAfter + } + if !pluginResp.NextRefreshAfter.IsZero() { + data.NextRefreshAfter = pluginResp.NextRefreshAfter + } + next := a.host.AuthDataToCoreAuth(data, "", data.FileName) + if next == nil { + return nil, fmt.Errorf("plugin executor %s refresh returned invalid auth data", a.Identifier()) + } + if auth != nil { + next.CreatedAt = auth.CreatedAt + next.UpdatedAt = auth.UpdatedAt + } + return next, nil +} + +func (a *executorAdapter) CountTokens(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (resp coreexecutor.Response, err error) { + if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { + return coreexecutor.Response{}, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(a.pluginID, "Executor.CountTokens", recovered) + resp = coreexecutor.Response{} + err = fmt.Errorf("plugin executor %s count tokens panic: %v", a.Identifier(), recovered) + } + }() + + prepared, errPrepare := a.prepareExecutorCall(req, opts) + if errPrepare != nil { + return coreexecutor.Response{}, errPrepare + } + pluginResp, errCountTokens := a.executor.CountTokens(ctx, buildExecutorRequest(a.host, a.provider, auth, prepared.req, prepared.opts)) + if errCountTokens != nil { + return coreexecutor.Response{}, errCountTokens + } + return coreexecutor.Response{ + Payload: a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil), + Metadata: cloneAnyMap(pluginResp.Metadata), + Headers: cloneHeader(pluginResp.Headers), + }, nil +} + +func (a *executorAdapter) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (resp *http.Response, err error) { + if a == nil || a.executor == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { + return nil, fmt.Errorf("plugin executor %s is unavailable", a.Identifier()) + } + if req == nil { + return nil, fmt.Errorf("plugin executor %s received nil HTTP request", a.Identifier()) + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(a.pluginID, "Executor.HttpRequest", recovered) + resp = nil + err = fmt.Errorf("plugin executor %s http request panic: %v", a.Identifier(), recovered) + } + }() + body, errReadAll := readAndRestoreRequestBody(req) + if errReadAll != nil { + return nil, fmt.Errorf("read plugin http request body: %w", errReadAll) + } + pluginResp, errHTTPRequest := a.executor.HttpRequest(ctx, pluginapi.ExecutorHTTPRequest{ + AuthID: authID(auth), + AuthProvider: authProvider(auth), + Method: req.Method, + URL: req.URL.String(), + Headers: cloneHeader(req.Header), + Body: bytes.Clone(body), + StorageJSON: storageJSONFromAuth(auth), + Metadata: cloneAnyMap(authMetadata(auth)), + Attributes: authAttributes(auth), + HTTPClient: a.host.newHTTPClient(auth, a.provider), + }) + if errHTTPRequest != nil { + return nil, errHTTPRequest + } + status := pluginResp.StatusCode + if status == 0 { + status = http.StatusOK + } + resp = &http.Response{ + StatusCode: status, + Status: fmt.Sprintf("%d %s", status, http.StatusText(status)), + Header: cloneHeader(pluginResp.Headers), + Body: io.NopCloser(bytes.NewReader(bytes.Clone(pluginResp.Body))), + Request: req, + } + return resp, nil +} + +func buildExecutorRequest(host *Host, provider string, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) pluginapi.ExecutorRequest { + return pluginapi.ExecutorRequest{ + AuthID: authID(auth), + AuthProvider: authProvider(auth), + Model: req.Model, + Format: req.Format.String(), + Stream: opts.Stream, + Alt: opts.Alt, + Headers: cloneHeader(opts.Headers), + Query: cloneValues(opts.Query), + OriginalRequest: bytes.Clone(opts.OriginalRequest), + SourceFormat: opts.SourceFormat.String(), + Payload: bytes.Clone(req.Payload), + Metadata: mergeExecutorMetadata(req.Metadata, opts.Metadata), + StorageJSON: storageJSONFromAuth(auth), + AuthMetadata: cloneAnyMap(authMetadata(auth)), + AuthAttributes: authAttributes(auth), + HTTPClient: host.newHTTPClient(auth, provider), + } +} + +func storageJSONFromAuth(auth *coreauth.Auth) []byte { + if auth == nil { + return nil + } + if rawProvider, okRaw := auth.Storage.(interface{ RawJSON() []byte }); okRaw { + return bytes.Clone(rawProvider.RawJSON()) + } + if len(auth.Metadata) == 0 { + return nil + } + data, errMarshal := json.Marshal(auth.Metadata) + if errMarshal != nil { + return nil + } + return data +} + +func authAttributes(auth *coreauth.Auth) map[string]string { + if auth == nil { + return nil + } + return cloneStringMap(auth.Attributes) +} + +func mergeExecutorMetadata(reqMetadata, optsMetadata map[string]any) map[string]any { + if len(reqMetadata) == 0 && len(optsMetadata) == 0 { + return nil + } + merged := make(map[string]any, len(reqMetadata)+len(optsMetadata)) + for key, value := range reqMetadata { + merged[key] = value + } + for key, value := range optsMetadata { + merged[key] = value + } + return merged +} + +func mapExecutorStreamChunks(ctx context.Context, in <-chan pluginapi.ExecutorStreamChunk) <-chan coreexecutor.StreamChunk { + if ctx == nil { + ctx = context.Background() + } + out := make(chan coreexecutor.StreamChunk) + if in == nil { + close(out) + return out + } + go func() { + defer close(out) + for { + var mapped coreexecutor.StreamChunk + select { + case <-ctx.Done(): + return + case chunk, ok := <-in: + if !ok { + return + } + mapped = coreexecutor.StreamChunk{ + Payload: bytes.Clone(chunk.Payload), + Err: chunk.Err, + } + } + select { + case <-ctx.Done(): + return + case out <- mapped: + } + } + }() + return out +} diff --git a/internal/pluginhost/adapters_interceptors.go b/internal/pluginhost/adapters_interceptors.go new file mode 100644 index 000000000..228784044 --- /dev/null +++ b/internal/pluginhost/adapters_interceptors.go @@ -0,0 +1,492 @@ +package pluginhost + +import ( + "bytes" + "context" + "io" + "net/http" + "net/url" + "reflect" + "strings" + + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + log "github.com/sirupsen/logrus" +) + +func (h *Host) callRequestInterceptor(ctx context.Context, record capabilityRecord, method string, call func(context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error), req pluginapi.RequestInterceptRequest) (out pluginapi.RequestInterceptResponse, ok bool) { + if h == nil || call == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return pluginapi.RequestInterceptResponse{}, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, method, recovered) + out = pluginapi.RequestInterceptResponse{} + ok = false + } + }() + resp, errIntercept := call(ctx, req) + if errIntercept != nil { + log.Warnf("pluginhost: request interceptor %s failed: %v", record.id, errIntercept) + return pluginapi.RequestInterceptResponse{}, false + } + return resp, true +} + +func (h *Host) callResponseInterceptor(ctx context.Context, record capabilityRecord, interceptor pluginapi.ResponseInterceptor, req pluginapi.ResponseInterceptRequest) (out pluginapi.ResponseInterceptResponse, ok bool) { + if h == nil || interceptor == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return pluginapi.ResponseInterceptResponse{}, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "ResponseInterceptor.InterceptResponse", recovered) + out = pluginapi.ResponseInterceptResponse{} + ok = false + } + }() + resp, errIntercept := interceptor.InterceptResponse(ctx, req) + if errIntercept != nil { + log.Warnf("pluginhost: response interceptor %s failed: %v", record.id, errIntercept) + return pluginapi.ResponseInterceptResponse{}, false + } + return resp, true +} + +func (h *Host) callStreamChunkInterceptor(ctx context.Context, record capabilityRecord, interceptor pluginapi.StreamChunkInterceptor, req pluginapi.StreamChunkInterceptRequest) (out pluginapi.StreamChunkInterceptResponse, ok bool) { + if h == nil || interceptor == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return pluginapi.StreamChunkInterceptResponse{}, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "StreamChunkInterceptor.InterceptStreamChunk", recovered) + out = pluginapi.StreamChunkInterceptResponse{} + ok = false + } + }() + resp, errIntercept := interceptor.InterceptStreamChunk(ctx, req) + if errIntercept != nil { + log.Warnf("pluginhost: stream chunk interceptor %s failed: %v", record.id, errIntercept) + return pluginapi.StreamChunkInterceptResponse{}, false + } + return resp, true +} + +func (h *Host) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + return h.InterceptRequestBeforeAuthExcept(ctx, req, "") +} + +func (h *Host) InterceptRequestBeforeAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + return h.interceptRequest(ctx, req, "RequestInterceptor.InterceptRequestBeforeAuth", func(interceptor pluginapi.RequestInterceptor, ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return interceptor.InterceptRequestBeforeAuth(ctx, req) + }, skipPluginID) +} + +func (h *Host) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse { + return h.InterceptRequestAfterAuthExcept(ctx, req, "") +} + +func (h *Host) InterceptRequestAfterAuthExcept(ctx context.Context, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + return h.interceptRequest(ctx, req, "RequestInterceptor.InterceptRequestAfterAuth", func(interceptor pluginapi.RequestInterceptor, ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return interceptor.InterceptRequestAfterAuth(ctx, req) + }, skipPluginID) +} + +func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, method string, invoke func(pluginapi.RequestInterceptor, context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error), skipPluginID string) pluginapi.RequestInterceptResponse { + current := pluginapi.RequestInterceptResponse{ + Headers: cloneHeader(req.Headers), + Body: bytes.Clone(req.Body), + } + skipPluginID = strings.TrimSpace(skipPluginID) + for _, record := range h.activeRecords() { + interceptor := record.plugin.Capabilities.RequestInterceptor + if h.isPluginFused(record.id) || interceptor == nil || record.id == skipPluginID { + continue + } + nextReq := req + nextReq.Headers = cloneHeader(current.Headers) + nextReq.Body = bytes.Clone(current.Body) + nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) + if resp, ok := h.callRequestInterceptor(ctx, record, method, func(callCtx context.Context, callReq pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { + return invoke(interceptor, callCtx, callReq) + }, nextReq); ok { + current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + current.Body = bytes.Clone(resp.Body) + } + } + } + return current +} + +func (h *Host) InterceptResponse(ctx context.Context, req pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse { + return h.InterceptResponseExcept(ctx, req, "") +} + +func (h *Host) InterceptResponseExcept(ctx context.Context, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse { + current := pluginapi.ResponseInterceptResponse{ + Headers: cloneHeader(req.ResponseHeaders), + Body: bytes.Clone(req.Body), + } + skipPluginID = strings.TrimSpace(skipPluginID) + for _, record := range h.activeRecords() { + interceptor := record.plugin.Capabilities.ResponseInterceptor + if h.isPluginFused(record.id) || interceptor == nil || record.id == skipPluginID { + continue + } + nextReq := req + nextReq.RequestHeaders = cloneHeader(req.RequestHeaders) + nextReq.ResponseHeaders = cloneHeader(current.Headers) + nextReq.OriginalRequest = bytes.Clone(req.OriginalRequest) + nextReq.RequestBody = bytes.Clone(req.RequestBody) + nextReq.Body = bytes.Clone(current.Body) + nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) + if resp, ok := h.callResponseInterceptor(ctx, record, interceptor, nextReq); ok { + current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + current.Body = bytes.Clone(resp.Body) + } + } + } + return current +} + +func (h *Host) InterceptStreamChunk(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + return h.InterceptStreamChunkExcept(ctx, req, "") +} + +func (h *Host) InterceptStreamChunkExcept(ctx context.Context, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse { + current := pluginapi.StreamChunkInterceptResponse{ + Headers: cloneHeader(req.ResponseHeaders), + Body: bytes.Clone(req.Body), + } + skipPluginID = strings.TrimSpace(skipPluginID) + for _, record := range h.activeRecords() { + interceptor := record.plugin.Capabilities.StreamChunkInterceptor + if h.isPluginFused(record.id) || interceptor == nil || current.DropChunk || record.id == skipPluginID { + continue + } + nextReq := req + nextReq.RequestHeaders = cloneHeader(req.RequestHeaders) + nextReq.ResponseHeaders = cloneHeader(current.Headers) + nextReq.OriginalRequest = bytes.Clone(req.OriginalRequest) + nextReq.RequestBody = bytes.Clone(req.RequestBody) + nextReq.Body = bytes.Clone(current.Body) + nextReq.HistoryChunks = cloneByteSlices(req.HistoryChunks) + nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) + if resp, ok := h.callStreamChunkInterceptor(ctx, record, interceptor, nextReq); ok { + current.Headers = mergeHeaders(current.Headers, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + current.Body = bytes.Clone(resp.Body) + } + if resp.DropChunk { + current.DropChunk = true + } + } + } + return current +} + +func (h *Host) HasStreamInterceptors() bool { + if h == nil { + return false + } + for _, record := range h.activeRecords() { + if h.isPluginFused(record.id) { + continue + } + if record.plugin.Capabilities.StreamChunkInterceptor != nil { + return true + } + } + return false +} + +func (h *Host) HasRequestInterceptors() bool { + if h == nil { + return false + } + for _, record := range h.activeRecords() { + if h.isPluginFused(record.id) { + continue + } + if record.plugin.Capabilities.RequestInterceptor != nil { + return true + } + } + return false +} + +func (h *Host) commitModelClients(snap *Snapshot, modelRegistry modelRegistry, registrations []modelClientRegistration, nextClients map[string]struct{}, nextProviders map[string]string, nextModelRegistrations map[string]pluginModelRegistration) { + if h == nil || modelRegistry == nil { + return + } + + staleClients := make([]string, 0) + h.mu.Lock() + if h.Snapshot() != snap { + h.mu.Unlock() + return + } + for clientID := range h.modelClientIDs { + if _, okClient := nextClients[clientID]; !okClient { + staleClients = append(staleClients, clientID) + } + } + h.modelClientIDs = nextClients + h.modelProviders = nextProviders + h.modelRegistrations = nextModelRegistrations + h.mu.Unlock() + + for _, registration := range registrations { + modelRegistry.RegisterClient(registration.clientID, registration.provider, registration.models) + } + for _, clientID := range staleClients { + modelRegistry.UnregisterClient(clientID) + } +} + +func readAndRestoreRequestBody(r *http.Request) ([]byte, error) { + if r == nil || r.Body == nil { + return nil, nil + } + body, errReadAll := io.ReadAll(r.Body) + if errReadAll != nil { + r.Body = io.NopCloser(bytes.NewReader(body)) + return nil, errReadAll + } + r.Body = io.NopCloser(bytes.NewReader(body)) + return body, nil +} + +func authID(auth *coreauth.Auth) string { + if auth == nil { + return "" + } + return auth.ID +} + +func authProvider(auth *coreauth.Auth) string { + if auth == nil { + return "" + } + return auth.Provider +} + +func authMetadata(auth *coreauth.Auth) map[string]any { + if auth == nil { + return nil + } + return auth.Metadata +} + +func cloneHeader(in http.Header) http.Header { + if len(in) == 0 { + return nil + } + out := make(http.Header, len(in)) + for key, values := range in { + out[key] = append([]string(nil), values...) + } + return out +} + +func mergeHeaders(current, updates http.Header, clear []string) http.Header { + out := cloneHeader(current) + if out == nil { + out = make(http.Header) + } + for _, key := range clear { + out.Del(key) + } + for key, values := range updates { + out.Del(key) + for _, value := range values { + out.Add(key, value) + } + } + return out +} + +func cloneByteSlices(in [][]byte) [][]byte { + if len(in) == 0 { + return nil + } + out := make([][]byte, 0, len(in)) + for _, item := range in { + out = append(out, bytes.Clone(item)) + } + return out +} + +func cloneValues(in url.Values) url.Values { + if len(in) == 0 { + return nil + } + out := make(url.Values, len(in)) + for key, values := range in { + out[key] = append([]string(nil), values...) + } + return out +} + +func cloneAnyMap(in map[string]any) map[string]any { + if len(in) == 0 { + return nil + } + out := make(map[string]any, len(in)) + for key, value := range in { + out[key] = value + } + return out +} + +func cloneInterceptorMetadata(in map[string]any) map[string]any { + if len(in) == 0 { + return nil + } + visited := make(map[metadataCloneVisit]reflect.Value) + out := make(map[string]any, len(in)) + for key, value := range in { + out[key] = cloneInterceptorMetadataAny(reflect.ValueOf(value), visited) + } + return out +} + +type metadataCloneVisit struct { + typ reflect.Type + ptr uintptr +} + +func cloneInterceptorMetadataAny(value reflect.Value, visited map[metadataCloneVisit]reflect.Value) any { + cloned := cloneInterceptorMetadataReflectValue(value, visited) + if !cloned.IsValid() { + return nil + } + return cloned.Interface() +} + +func cloneInterceptorMetadataReflectValue(value reflect.Value, visited map[metadataCloneVisit]reflect.Value) reflect.Value { + if !value.IsValid() { + return reflect.Value{} + } + + switch value.Kind() { + case reflect.Interface: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + return cloneInterceptorMetadataReflectValue(value.Elem(), visited) + case reflect.Pointer: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + visit := metadataCloneVisit{typ: value.Type(), ptr: value.Pointer()} + if existing, okExisting := visited[visit]; okExisting { + return existing + } + out := reflect.New(value.Type().Elem()) + visited[visit] = out + clonedElem := cloneInterceptorMetadataReflectValue(value.Elem(), visited) + if clonedElem.IsValid() { + outElem := out.Elem() + if clonedElem.Type().AssignableTo(outElem.Type()) { + outElem.Set(clonedElem) + } else if clonedElem.Type().ConvertibleTo(outElem.Type()) { + outElem.Set(clonedElem.Convert(outElem.Type())) + } + } + return out + case reflect.Map: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + visit := metadataCloneVisit{typ: value.Type(), ptr: value.Pointer()} + if existing, okExisting := visited[visit]; okExisting { + return existing + } + out := reflect.MakeMapWithSize(value.Type(), value.Len()) + visited[visit] = out + iter := value.MapRange() + for iter.Next() { + keyValue := adaptClonedValue(iter.Key(), cloneInterceptorMetadataReflectValue(iter.Key(), visited)) + valValue := adaptClonedValue(iter.Value(), cloneInterceptorMetadataReflectValue(iter.Value(), visited)) + out.SetMapIndex(keyValue, valValue) + } + return out + case reflect.Slice: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + if value.Type().Elem().Kind() == reflect.Uint8 { + out := reflect.MakeSlice(value.Type(), value.Len(), value.Len()) + reflect.Copy(out, value) + return out + } + visit := metadataCloneVisit{typ: value.Type(), ptr: value.Pointer()} + if existing, okExisting := visited[visit]; okExisting { + return existing + } + out := reflect.MakeSlice(value.Type(), value.Len(), value.Len()) + visited[visit] = out + for i := 0; i < value.Len(); i++ { + clonedItem := cloneInterceptorMetadataReflectValue(value.Index(i), visited) + if !clonedItem.IsValid() { + continue + } + out.Index(i).Set(adaptClonedValue(value.Index(i), clonedItem)) + } + return out + case reflect.Array: + out := reflect.New(value.Type()).Elem() + for i := 0; i < value.Len(); i++ { + clonedItem := cloneInterceptorMetadataReflectValue(value.Index(i), visited) + if !clonedItem.IsValid() { + continue + } + out.Index(i).Set(adaptClonedValue(value.Index(i), clonedItem)) + } + return out + case reflect.Struct: + out := reflect.New(value.Type()).Elem() + // Preserve unexported fields and deep-clone exported fields on a best-effort basis. + out.Set(value) + for i := 0; i < value.NumField(); i++ { + field := value.Field(i) + if !out.Field(i).CanSet() { + continue + } + fieldClone := cloneInterceptorMetadataReflectValue(field, visited) + if !fieldClone.IsValid() { + continue + } + out.Field(i).Set(adaptClonedValue(field, fieldClone)) + } + return out + default: + return value + } +} + +func adaptClonedValue(original, cloned reflect.Value) reflect.Value { + if !cloned.IsValid() { + return original + } + if cloned.Type().AssignableTo(original.Type()) { + return cloned + } + if cloned.Type().ConvertibleTo(original.Type()) { + return cloned.Convert(original.Type()) + } + return original +} + +func cloneStringMap(in map[string]string) map[string]string { + if len(in) == 0 { + return nil + } + out := make(map[string]string, len(in)) + for key, value := range in { + out[key] = value + } + return out +} diff --git a/internal/pluginhost/adapters_usage_translation.go b/internal/pluginhost/adapters_usage_translation.go new file mode 100644 index 000000000..2201eb6c8 --- /dev/null +++ b/internal/pluginhost/adapters_usage_translation.go @@ -0,0 +1,369 @@ +package pluginhost + +import ( + "bytes" + "context" + "fmt" + "runtime/debug" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" +) + +func (h *Host) RegisterUsagePlugins() { + if h == nil { + return + } + + for _, record := range h.activeRecords() { + plugin := record.plugin.Capabilities.UsagePlugin + if plugin == nil || h.isPluginFused(record.id) { + continue + } + coreusage.RegisterNamedPlugin("plugin:"+record.id, &usageAdapter{ + host: h, + pluginID: record.id, + plugin: plugin, + }) + } +} + +func (h *Host) refreshThinkingProviders(records []capabilityRecord) { + thinking.ClearPluginProviders() + if h == nil { + return + } + for _, record := range records { + applier := record.plugin.Capabilities.ThinkingApplier + if applier == nil || h.isPluginFused(record.id) { + continue + } + provider, okProvider := h.callThinkingIdentifier(record, applier) + if !okProvider { + continue + } + thinking.RegisterPluginProvider(record.id, provider, record.priority, &thinkingAdapter{ + host: h, + pluginID: record.id, + path: record.path, + version: record.version, + provider: provider, + applier: applier, + }) + } +} + +func (h *Host) callThinkingIdentifier(record capabilityRecord, applier pluginapi.ThinkingApplier) (provider string, ok bool) { + if h == nil || applier == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return "", false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "ThinkingApplier.Identifier", recovered) + provider = "" + ok = false + } + }() + provider = strings.ToLower(strings.TrimSpace(applier.Identifier())) + if provider == "" { + return "", false + } + return provider, true +} + +func (h *Host) currentUsagePlugin(pluginID string) pluginapi.UsagePlugin { + if h == nil || strings.TrimSpace(pluginID) == "" { + return nil + } + for _, record := range h.activeRecords() { + if record.id != pluginID { + continue + } + if h.isPluginFused(record.id) { + return nil + } + return record.plugin.Capabilities.UsagePlugin + } + return nil +} + +func (h *Host) fusePlugin(id, method string, recovered any) { + if h == nil { + return + } + h.mu.Lock() + h.fused[id] = fmt.Sprintf("%s panic: %v", method, recovered) + h.mu.Unlock() + thinking.UnregisterPluginProviders(id) + log.WithField("plugin_id", id).WithField("method", method).Errorf("pluginhost: plugin panic recovered: %v\n%s", recovered, debug.Stack()) +} + +func (h *Host) isPluginFused(id string) bool { + if h == nil { + return false + } + h.mu.Lock() + _, fused := h.fused[id] + h.mu.Unlock() + return fused +} + +type usageAdapter struct { + host *Host + pluginID string + plugin pluginapi.UsagePlugin +} + +type thinkingAdapter struct { + host *Host + pluginID string + path string + version string + provider string + applier pluginapi.ThinkingApplier +} + +func (a *usageAdapter) HandleUsage(ctx context.Context, record coreusage.Record) { + if a == nil { + return + } + plugin := a.host.currentUsagePlugin(a.pluginID) + if plugin == nil { + return + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(a.pluginID, "UsagePlugin.HandleUsage", recovered) + } + }() + plugin.HandleUsage(ctx, pluginapi.UsageRecord{ + Provider: record.Provider, + ExecutorType: record.ExecutorType, + Model: record.Model, + Alias: record.Alias, + APIKey: record.APIKey, + AuthID: record.AuthID, + AuthIndex: record.AuthIndex, + AuthType: record.AuthType, + Source: record.Source, + ReasoningEffort: record.ReasoningEffort, + ServiceTier: record.ServiceTier, + Generate: coreusage.GenerateEnabled(record.Generate), + RequestedAt: record.RequestedAt, + Latency: record.Latency, + TTFT: record.TTFT, + Failed: record.Failed, + Failure: pluginapi.UsageFailure{ + StatusCode: record.Fail.StatusCode, + Body: record.Fail.Body, + }, + Detail: pluginapi.UsageDetail{ + InputTokens: record.Detail.InputTokens, + OutputTokens: record.Detail.OutputTokens, + ReasoningTokens: record.Detail.ReasoningTokens, + CachedTokens: record.Detail.CachedTokens, + CacheReadTokens: record.Detail.CacheReadTokens, + CacheCreationTokens: record.Detail.CacheCreationTokens, + TotalTokens: record.Detail.TotalTokens, + }, + ResponseHeaders: cloneHeader(record.ResponseHeaders), + }) +} + +func (a *thinkingAdapter) Apply(body []byte, config thinking.ThinkingConfig, modelInfo *registry.ModelInfo) (out []byte, err error) { + if a == nil || a.applier == nil || a.host == nil || a.host.isPluginFused(a.pluginID) || !a.host.pluginIdentityCurrent(a.pluginID, a.path, a.version) { + return bytes.Clone(body), nil + } + defer func() { + if recovered := recover(); recovered != nil { + a.host.fusePlugin(a.pluginID, "ThinkingApplier.ApplyThinking", recovered) + out = bytes.Clone(body) + err = nil + } + }() + resp, errApply := a.applier.ApplyThinking(context.Background(), pluginapi.ThinkingApplyRequest{ + Provider: a.provider, + Model: registryModelInfoToPluginModelInfo(modelInfo), + Config: pluginapi.ThinkingConfig{ + Mode: config.Mode.String(), + Budget: config.Budget, + Level: string(config.Level), + }, + Body: bytes.Clone(body), + }) + if errApply != nil || len(resp.Body) == 0 { + return bytes.Clone(body), nil + } + return bytes.Clone(resp.Body), nil +} + +func (h *Host) NormalizeRequest(ctx context.Context, from, to sdktranslator.Format, model string, body []byte, stream bool) []byte { + current := bytes.Clone(body) + for _, record := range h.activeRecords() { + if h.isPluginFused(record.id) || record.plugin.Capabilities.RequestNormalizer == nil { + continue + } + if normalized, ok := h.callRequestNormalizer(ctx, record, from, to, model, current, stream); ok { + current = normalized + } + } + return current +} + +func (h *Host) TranslateRequest(ctx context.Context, from, to sdktranslator.Format, model string, body []byte, stream bool) ([]byte, bool) { + for _, record := range h.activeRecords() { + if h.isPluginFused(record.id) || record.plugin.Capabilities.RequestTranslator == nil { + continue + } + if translated, ok := h.callRequestTranslator(ctx, record, from, to, model, body, stream); ok { + return translated, true + } + } + return bytes.Clone(body), false +} + +func (h *Host) NormalizeResponseBefore(ctx context.Context, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) []byte { + current := bytes.Clone(body) + for _, record := range h.activeRecords() { + normalizer := record.plugin.Capabilities.ResponseBeforeTranslator + if h.isPluginFused(record.id) || normalizer == nil { + continue + } + if normalized, ok := h.callResponseNormalizer(ctx, record, "ResponseBeforeTranslator.NormalizeResponse", normalizer, from, to, model, originalRequestRawJSON, requestRawJSON, current, stream); ok { + current = normalized + } + } + return current +} + +func (h *Host) TranslateResponse(ctx context.Context, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) ([]byte, bool) { + for _, record := range h.activeRecords() { + translator := record.plugin.Capabilities.ResponseTranslator + if h.isPluginFused(record.id) || translator == nil { + continue + } + if translated, ok := h.callResponseTranslator(ctx, record, translator, from, to, model, originalRequestRawJSON, requestRawJSON, body, stream); ok { + return translated, true + } + } + return bytes.Clone(body), false +} + +func (h *Host) NormalizeResponseAfter(ctx context.Context, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) []byte { + current := bytes.Clone(body) + for _, record := range h.activeRecords() { + normalizer := record.plugin.Capabilities.ResponseAfterTranslator + if h.isPluginFused(record.id) || normalizer == nil { + continue + } + if normalized, ok := h.callResponseNormalizer(ctx, record, "ResponseAfterTranslator.NormalizeResponse", normalizer, from, to, model, originalRequestRawJSON, requestRawJSON, current, stream); ok { + current = normalized + } + } + return current +} + +func (h *Host) callRequestNormalizer(ctx context.Context, record capabilityRecord, from, to sdktranslator.Format, model string, body []byte, stream bool) (out []byte, ok bool) { + if h == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) || record.plugin.Capabilities.RequestNormalizer == nil { + return nil, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "RequestNormalizer.NormalizeRequest", recovered) + out = nil + ok = false + } + }() + resp, errNormalizeRequest := record.plugin.Capabilities.RequestNormalizer.NormalizeRequest(ctx, pluginapi.RequestTransformRequest{ + FromFormat: from.String(), + ToFormat: to.String(), + Model: model, + Stream: stream, + Body: bytes.Clone(body), + }) + if errNormalizeRequest != nil || len(resp.Body) == 0 { + return nil, false + } + return bytes.Clone(resp.Body), true +} + +func (h *Host) callRequestTranslator(ctx context.Context, record capabilityRecord, from, to sdktranslator.Format, model string, body []byte, stream bool) (out []byte, ok bool) { + if h == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) || record.plugin.Capabilities.RequestTranslator == nil { + return nil, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "RequestTranslator.TranslateRequest", recovered) + out = nil + ok = false + } + }() + resp, errTranslateRequest := record.plugin.Capabilities.RequestTranslator.TranslateRequest(ctx, pluginapi.RequestTransformRequest{ + FromFormat: from.String(), + ToFormat: to.String(), + Model: model, + Stream: stream, + Body: bytes.Clone(body), + }) + if errTranslateRequest != nil || len(resp.Body) == 0 { + return nil, false + } + return bytes.Clone(resp.Body), true +} + +func (h *Host) callResponseNormalizer(ctx context.Context, record capabilityRecord, method string, normalizer pluginapi.ResponseNormalizer, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) (out []byte, ok bool) { + if h == nil || normalizer == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return nil, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, method, recovered) + out = nil + ok = false + } + }() + resp, errNormalizeResponse := normalizer.NormalizeResponse(ctx, pluginapi.ResponseTransformRequest{ + FromFormat: from.String(), + ToFormat: to.String(), + Model: model, + Stream: stream, + OriginalRequest: bytes.Clone(originalRequestRawJSON), + TranslatedRequest: bytes.Clone(requestRawJSON), + Body: bytes.Clone(body), + }) + if errNormalizeResponse != nil || len(resp.Body) == 0 { + return nil, false + } + return bytes.Clone(resp.Body), true +} + +func (h *Host) callResponseTranslator(ctx context.Context, record capabilityRecord, translator pluginapi.ResponseTranslator, from, to sdktranslator.Format, model string, originalRequestRawJSON, requestRawJSON, body []byte, stream bool) (out []byte, ok bool) { + if h == nil || translator == nil || h.isPluginFused(record.id) || !h.recordCurrent(record) { + return nil, false + } + defer func() { + if recovered := recover(); recovered != nil { + h.fusePlugin(record.id, "ResponseTranslator.TranslateResponse", recovered) + out = nil + ok = false + } + }() + resp, errTranslateResponse := translator.TranslateResponse(ctx, pluginapi.ResponseTransformRequest{ + FromFormat: from.String(), + ToFormat: to.String(), + Model: model, + Stream: stream, + OriginalRequest: bytes.Clone(originalRequestRawJSON), + TranslatedRequest: bytes.Clone(requestRawJSON), + Body: bytes.Clone(body), + }) + if errTranslateResponse != nil || len(resp.Body) == 0 { + return nil, false + } + return bytes.Clone(resp.Body), true +} diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index a804608c6..460a73077 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -4,43 +4,25 @@ package executor import ( - "bufio" - "bytes" "context" - "crypto/sha256" "crypto/tls" - "encoding/binary" "encoding/json" - "errors" "fmt" - "io" - "math/rand" "net/http" - "net/url" - "strconv" "strings" "sync" "time" - "github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" - "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" internalsignature "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" antigravityclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/antigravity/claude" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" - sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" - "golang.org/x/sync/singleflight" ) const ( @@ -61,170 +43,6 @@ const ( // systemInstruction = "You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding.You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question.**Absolute paths only****Proactiveness**" ) -type antigravity429Category string - -type antigravityCreditsFailureState struct { - PermanentlyDisabled bool - ExplicitBalanceExhausted bool -} - -type antigravity429DecisionKind string - -const ( - antigravity429Unknown antigravity429Category = "unknown" - antigravity429RateLimited antigravity429Category = "rate_limited" - antigravity429QuotaExhausted antigravity429Category = "quota_exhausted" - antigravity429SoftRateLimit antigravity429Category = "soft_rate_limit" - antigravity429DecisionSoftRetry antigravity429DecisionKind = "soft_retry" - antigravity429DecisionInstantRetrySameAuth antigravity429DecisionKind = "instant_retry_same_auth" - antigravity429DecisionShortCooldownSwitchAuth antigravity429DecisionKind = "short_cooldown_switch_auth" - antigravity429DecisionFullQuotaExhausted antigravity429DecisionKind = "full_quota_exhausted" -) - -type antigravity429Decision struct { - kind antigravity429DecisionKind - retryAfter *time.Duration - reason string -} - -var ( - randSource = rand.New(rand.NewSource(time.Now().UnixNano())) - randSourceMutex sync.Mutex - antigravityCreditsFailureByAuth sync.Map - antigravityShortCooldownByAuth sync.Map - antigravityCreditsBalanceByAuth sync.Map // auth.ID → antigravityCreditsBalance - antigravityCreditsHintRefreshByID sync.Map // auth.ID → *antigravityCreditsHintRefreshState - antigravityRefreshGroup singleflight.Group - antigravityQuotaExhaustedKeywords = []string{ - "quota_exhausted", - "quota exhausted", - } -) - -type antigravityKVClient interface { - KVGet(ctx context.Context, key string) ([]byte, bool, error) - KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) - KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) - KVDel(ctx context.Context, keys ...string) (int64, error) -} - -var currentAntigravityKVClient = func() (antigravityKVClient, bool, error) { - return homekv.CurrentKVClient() -} - -type antigravityCreditsBalance struct { - CreditAmount float64 - MinCreditAmount float64 - PaidTierID string - Known bool -} - -type antigravityCreditsHintRefreshState struct { - mu sync.Mutex - lastAttempt time.Time -} - -type antigravityTokenRefreshData struct { - AccessToken string `json:"access_token"` - RefreshToken string `json:"refresh_token"` - ExpiresIn int64 `json:"expires_in"` - TokenType string `json:"token_type"` -} - -func antigravityAuthHasCredits(auth *cliproxyauth.Auth) bool { - ok, err := antigravityAuthHasCreditsRequired(context.Background(), auth) - if err != nil { - log.Errorf("antigravity executor: home kv credits check error: %v", err) - return false - } - return ok -} - -func antigravityAuthHasCreditsRequired(ctx context.Context, auth *cliproxyauth.Auth) (bool, error) { - if auth == nil || strings.TrimSpace(auth.ID) == "" { - return false, nil - } - authID := strings.TrimSpace(auth.ID) - if hint, ok, errHint := cliproxyauth.GetAntigravityCreditsHintRequired(ctx, authID); errHint != nil { - return false, errHint - } else if ok && hint.Known { - return hint.Available, nil - } - - client, homeMode, errClient := currentAntigravityKVClient() - if homeMode { - if errClient != nil { - return false, errClient - } - raw, found, errBalance := client.KVGet(ctx, antigravityCreditsBalanceKey(authID)) - if errBalance != nil { - return false, errBalance - } - if !found { - return true, nil - } - var homeBalance antigravityCreditsBalance - if errUnmarshal := json.Unmarshal(raw, &homeBalance); errUnmarshal != nil { - return false, errUnmarshal - } - return antigravityCreditsBalanceAvailable(authID, homeBalance), nil - } - - val, ok := antigravityCreditsBalanceByAuth.Load(authID) - if !ok { - return true, nil // optimistic: assume credits available when balance unknown - } - bal, valid := val.(antigravityCreditsBalance) - if !valid { - antigravityCreditsBalanceByAuth.Delete(authID) - return false, nil - } - return antigravityCreditsBalanceAvailable(authID, bal), nil -} - -func antigravityCreditsBalanceAvailable(authID string, bal antigravityCreditsBalance) bool { - if !bal.Known { - return false - } - available := bal.CreditAmount >= bal.MinCreditAmount - cliproxyauth.SetAntigravityCreditsHint(strings.TrimSpace(authID), cliproxyauth.AntigravityCreditsHint{ - Known: true, - Available: available, - CreditAmount: bal.CreditAmount, - MinCreditAmount: bal.MinCreditAmount, - PaidTierID: bal.PaidTierID, - UpdatedAt: time.Now(), - }) - return available -} - -// parseMetaFloat extracts a float64 from auth.Metadata (handles string and numeric types). -func parseMetaFloat(metadata map[string]any, key string) (float64, bool) { - v, ok := metadata[key] - if !ok { - return 0, false - } - switch typed := v.(type) { - case float64: - return typed, true - case int: - return float64(typed), true - case int64: - return float64(typed), true - case uint64: - return float64(typed), true - case json.Number: - if f, err := typed.Float64(); err == nil { - return f, true - } - case string: - if f, err := strconv.ParseFloat(strings.TrimSpace(typed), 64); err == nil { - return f, true - } - } - return 0, false -} - // AntigravityExecutor proxies requests to the antigravity upstream. type AntigravityExecutor struct { cfg *config.Config @@ -608,2522 +426,3 @@ func (e *AntigravityExecutor) HttpRequest(ctx context.Context, auth *cliproxyaut httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) return httpClient.Do(httpReq) } - -func injectEnabledCreditTypes(payload []byte) []byte { - if len(payload) == 0 { - return nil - } - if !gjson.ValidBytes(payload) { - return nil - } - updated, err := sjson.SetRawBytes(payload, "enabledCreditTypes", []byte(`["GOOGLE_ONE_AI"]`)) - if err != nil { - return nil - } - return updated -} - -func classifyAntigravity429(body []byte) antigravity429Category { - switch decideAntigravity429(body).kind { - case antigravity429DecisionInstantRetrySameAuth, antigravity429DecisionShortCooldownSwitchAuth: - return antigravity429RateLimited - case antigravity429DecisionFullQuotaExhausted: - return antigravity429QuotaExhausted - case antigravity429DecisionSoftRetry: - return antigravity429SoftRateLimit - default: - return antigravity429Unknown - } -} - -func decideAntigravity429(body []byte) antigravity429Decision { - decision := antigravity429Decision{kind: antigravity429DecisionSoftRetry} - if len(body) == 0 { - return decision - } - - if retryAfter, parseErr := helps.ParseRetryDelay(body); parseErr == nil && retryAfter != nil { - decision.retryAfter = retryAfter - } - - status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String()) - if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") { - return decision - } - - details := gjson.GetBytes(body, "error.details") - if details.Exists() && details.IsArray() { - for _, detail := range details.Array() { - if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" { - continue - } - reason := strings.TrimSpace(detail.Get("reason").String()) - decision.reason = reason - switch { - case strings.EqualFold(reason, "QUOTA_EXHAUSTED"): - decision.kind = antigravity429DecisionFullQuotaExhausted - return decision - case strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED"): - if decision.retryAfter == nil { - decision.kind = antigravity429DecisionSoftRetry - return decision - } - switch { - case *decision.retryAfter < antigravityInstantRetryThreshold: - decision.kind = antigravity429DecisionInstantRetrySameAuth - case *decision.retryAfter < antigravityShortQuotaCooldownThreshold: - decision.kind = antigravity429DecisionShortCooldownSwitchAuth - default: - decision.kind = antigravity429DecisionFullQuotaExhausted - } - return decision - } - } - } - - lowerBody := strings.ToLower(string(body)) - for _, keyword := range antigravityQuotaExhaustedKeywords { - if strings.Contains(lowerBody, keyword) { - decision.kind = antigravity429DecisionFullQuotaExhausted - decision.reason = "quota_exhausted" - return decision - } - } - - decision.kind = antigravity429DecisionSoftRetry - return decision -} - -func antigravityCreditsRetryEnabled(cfg *config.Config) bool { - return cfg != nil && cfg.QuotaExceeded.AntigravityCredits -} - -func clearAntigravityCreditsFailureState(auth *cliproxyauth.Auth) { - if auth == nil || strings.TrimSpace(auth.ID) == "" { - return - } - antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID)) -} -func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) { - if auth == nil || strings.TrimSpace(auth.ID) == "" { - return - } - authID := strings.TrimSpace(auth.ID) - state := antigravityCreditsFailureState{ - PermanentlyDisabled: true, - ExplicitBalanceExhausted: true, - } - antigravityCreditsFailureByAuth.Store(authID, state) - bal := antigravityCreditsBalance{ - CreditAmount: 0, - MinCreditAmount: 1, - Known: true, - } - storeAntigravityCreditsBalanceBestEffort(authID, bal) - cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ - Known: true, - Available: false, - CreditAmount: 0, - MinCreditAmount: 1, - UpdatedAt: time.Now(), - }) -} - -func clearAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) { - if auth == nil || strings.TrimSpace(auth.ID) == "" { - return - } - antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID)) -} - -func antigravityHasExplicitCreditsBalanceExhaustedReason(body []byte) bool { - if len(body) == 0 { - return false - } - details := gjson.GetBytes(body, "error.details") - if !details.Exists() || !details.IsArray() { - return false - } - for _, detail := range details.Array() { - if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" { - continue - } - reason := strings.TrimSpace(detail.Get("reason").String()) - if strings.EqualFold(reason, "INSUFFICIENT_G1_CREDITS_BALANCE") { - return true - } - } - return false -} - -func newAntigravityStatusErr(statusCode int, body []byte) statusErr { - err := statusErr{code: statusCode, msg: string(body)} - if statusCode == http.StatusTooManyRequests { - if retryAfter, parseErr := helps.ParseRetryDelay(body); parseErr == nil && retryAfter != nil { - err.retryAfter = retryAfter - } - } - return err -} - -// Execute performs a non-streaming request to the Antigravity API. -func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - if opts.Alt == "responses/compact" { - return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} - } - baseModel := thinking.ParseSuffix(req.Model).ModelName - if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { - return resp, homeKVUnavailableStatusErr(errCooldown) - } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { - log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) - d := remaining - return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} - } - - isClaude := strings.Contains(strings.ToLower(baseModel), "claude") - if isClaude || strings.Contains(baseModel, "gemini-3-pro") || strings.Contains(baseModel, "gemini-3.1-flash-image") { - return e.executeClaudeNonStream(ctx, auth, req, opts) - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("antigravity") - - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) - if errValidate != nil { - return resp, errValidate - } - req.Payload = originalPayload - token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) - if errToken != nil { - return resp, errToken - } - if updatedAuth != nil { - auth = updatedAuth - } - originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, false) - translated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, false) - - translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return resp, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers) - translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated) - reporter.SetTranslatedReasoningEffort(translated, to.String()) - - useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) - - baseURLs := antigravityBaseURLFallbackOrder(auth) - httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - attempts := antigravityRetryAttempts(auth, e.cfg) - -attemptLoop: - for attempt := 0; attempt < attempts; attempt++ { - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - requestPayload := translated - if useCredits { - if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { - requestPayload = cp - helps.MarkCreditsUsed(ctx) - } - } - replayScope := antigravityReasoningReplayScope{} - if antigravityUsesReasoningReplayCache(baseModel) { - var errReplay error - requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) - if errReplay != nil { - err = errReplay - return resp, err - } - } - - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) - if errReq != nil { - err = errReq - return resp, err - } - - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return resp, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errDo - return resp, err - } - - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - err = errRead - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - - if httpResp.StatusCode == http.StatusTooManyRequests { - decision := decideAntigravity429(bodyBytes) - switch decision.kind { - case antigravity429DecisionInstantRetrySameAuth: - if attempt+1 < attempts { - if decision.retryAfter != nil && *decision.retryAfter > 0 { - wait := antigravityInstantRetryDelay(*decision.retryAfter) - log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) - if errWait := antigravityWait(ctx, wait); errWait != nil { - return resp, errWait - } - } - continue attemptLoop - } - case antigravity429DecisionShortCooldownSwitchAuth: - if decision.retryAfter != nil && *decision.retryAfter > 0 { - if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { - err = homeKVUnavailableStatusErr(errMarkCooldown) - return resp, err - } - log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) - } - case antigravity429DecisionFullQuotaExhausted: - if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { - markAntigravityCreditsPermanentlyDisabled(auth) - } - // No credits logic - just fall through to error return below - } - } - - if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), bodyBytes)) - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { - delay := antigravityTransient429RetryDelay(attempt) - log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if attempt+1 < attempts { - delay := antigravityNoCapacityRetryDelay(attempt) - log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { - if attempt+1 < attempts { - delay := antigravitySoftRateLimitDelay(attempt) - log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { - err = errClear - return resp, err - } - err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) - return resp, err - } - - // Success - if useCredits { - clearAntigravityCreditsFailureState(auth) - } - cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, requestPayload, bodyBytes) - bodyBytes = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, bodyBytes) - reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes)) - var param any - converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bodyBytes, ¶m) - resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} - reporter.EnsurePublished(ctx) - return resp, nil - } - - switch { - case lastStatus != 0: - err = newAntigravityStatusErr(lastStatus, lastBody) - case lastErr != nil: - err = lastErr - default: - err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} - } - return resp, err - } - - return resp, err -} - -// executeClaudeNonStream performs a claude non-streaming request to the Antigravity API. -func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - baseModel := thinking.ParseSuffix(req.Model).ModelName - if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { - return resp, homeKVUnavailableStatusErr(errCooldown) - } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { - log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) - d := remaining - return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("antigravity") - - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) - if errValidate != nil { - return resp, errValidate - } - req.Payload = originalPayload - token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) - if errToken != nil { - return resp, errToken - } - if updatedAuth != nil { - auth = updatedAuth - } - originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, true) - translated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, true) - - translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return resp, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers) - translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated) - reporter.SetTranslatedReasoningEffort(translated, to.String()) - - useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) - - baseURLs := antigravityBaseURLFallbackOrder(auth) - httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - - attempts := antigravityRetryAttempts(auth, e.cfg) - -attemptLoop: - for attempt := 0; attempt < attempts; attempt++ { - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - requestPayload := translated - if useCredits { - if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { - requestPayload = cp - helps.MarkCreditsUsed(ctx) - } - } - replayScope := antigravityReasoningReplayScope{} - if antigravityUsesReasoningReplayCache(baseModel) { - var errReplay error - requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) - if errReplay != nil { - err = errReplay - return resp, err - } - } - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) - if errReq != nil { - err = errReq - return resp, err - } - - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return resp, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errDo - return resp, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { - err = errRead - return resp, err - } - if errCtx := ctx.Err(); errCtx != nil { - err = errCtx - return resp, err - } - lastStatus = 0 - lastBody = nil - lastErr = errRead - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errRead - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - if httpResp.StatusCode == http.StatusTooManyRequests { - decision := decideAntigravity429(bodyBytes) - - switch decision.kind { - case antigravity429DecisionInstantRetrySameAuth: - if attempt+1 < attempts { - if decision.retryAfter != nil && *decision.retryAfter > 0 { - wait := antigravityInstantRetryDelay(*decision.retryAfter) - log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) - if errWait := antigravityWait(ctx, wait); errWait != nil { - return resp, errWait - } - } - continue attemptLoop - } - case antigravity429DecisionShortCooldownSwitchAuth: - if decision.retryAfter != nil && *decision.retryAfter > 0 { - if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { - err = homeKVUnavailableStatusErr(errMarkCooldown) - return resp, err - } - log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) - } - case antigravity429DecisionFullQuotaExhausted: - if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { - markAntigravityCreditsPermanentlyDisabled(auth) - } - // No credits logic - just fall through to error return below - } - } - - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { - delay := antigravityTransient429RetryDelay(attempt) - log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if attempt+1 < attempts { - delay := antigravityNoCapacityRetryDelay(attempt) - log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { - if attempt+1 < attempts { - delay := antigravitySoftRateLimitDelay(attempt) - log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return resp, errWait - } - continue attemptLoop - } - } - if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { - err = errClear - return resp, err - } - err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) - return resp, err - } - - // Stream success - if useCredits { - clearAntigravityCreditsFailureState(auth) - } - replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) - out := make(chan cliproxyexecutor.StreamChunk) - go func(resp *http.Response) { - defer close(out) - defer func() { - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - }() - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(nil, streamScannerBuffer) - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if replayAccumulator != nil { - replayAccumulator.ObserveSSELine(line) - } - - // Filter usage metadata for all models - // Only retain usage statistics in the terminal chunk - line = helps.FilterSSEUsageMetadata(line) - - payload := helps.JSONPayload(line) - if payload == nil { - continue - } - - if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { - reporter.Publish(ctx, detail) - } - - out <- cliproxyexecutor.StreamChunk{Payload: payload} - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - out <- cliproxyexecutor.StreamChunk{Err: errScan} - } else { - if replayAccumulator != nil { - replayAccumulator.Commit(ctx) - } - reporter.EnsurePublished(ctx) - } - }(httpResp) - - var buffer bytes.Buffer - for chunk := range out { - if chunk.Err != nil { - return resp, chunk.Err - } - if len(chunk.Payload) > 0 { - _, _ = buffer.Write(chunk.Payload) - _, _ = buffer.Write([]byte("\n")) - } - } - resp = cliproxyexecutor.Response{Payload: e.convertStreamToNonStream(buffer.Bytes())} - - resp.Payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, resp.Payload) - reporter.Publish(ctx, helps.ParseAntigravityUsage(resp.Payload)) - var param any - converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, resp.Payload, ¶m) - resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} - reporter.EnsurePublished(ctx) - - return resp, nil - } - - switch { - case lastStatus != 0: - err = newAntigravityStatusErr(lastStatus, lastBody) - case lastErr != nil: - err = lastErr - default: - err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} - } - return resp, err - } - - return resp, err -} - -func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { - responseTemplate := "" - var traceID string - var finishReason string - var modelVersion string - var responseID string - var role string - var usageRaw string - parts := make([]map[string]interface{}, 0) - var pendingKind string - var pendingText strings.Builder - var pendingThoughtSig string - - flushPending := func() { - if pendingKind == "" { - return - } - text := pendingText.String() - switch pendingKind { - case "text": - if strings.TrimSpace(text) == "" { - pendingKind = "" - pendingText.Reset() - pendingThoughtSig = "" - return - } - parts = append(parts, map[string]interface{}{"text": text}) - case "thought": - if strings.TrimSpace(text) == "" && pendingThoughtSig == "" { - pendingKind = "" - pendingText.Reset() - pendingThoughtSig = "" - return - } - part := map[string]interface{}{"thought": true} - part["text"] = text - if pendingThoughtSig != "" { - part["thoughtSignature"] = pendingThoughtSig - } - parts = append(parts, part) - } - pendingKind = "" - pendingText.Reset() - pendingThoughtSig = "" - } - - normalizePart := func(partResult gjson.Result) map[string]interface{} { - var m map[string]interface{} - _ = json.Unmarshal([]byte(partResult.Raw), &m) - if m == nil { - m = map[string]interface{}{} - } - sig := partResult.Get("thoughtSignature").String() - if sig == "" { - sig = partResult.Get("thought_signature").String() - } - if sig != "" { - m["thoughtSignature"] = sig - delete(m, "thought_signature") - } - if inlineData, ok := m["inline_data"]; ok { - m["inlineData"] = inlineData - delete(m, "inline_data") - } - return m - } - - for _, line := range bytes.Split(stream, []byte("\n")) { - trimmed := bytes.TrimSpace(line) - if len(trimmed) == 0 || !gjson.ValidBytes(trimmed) { - continue - } - - root := gjson.ParseBytes(trimmed) - responseNode := root.Get("response") - if !responseNode.Exists() { - if root.Get("candidates").Exists() { - responseNode = root - } else { - continue - } - } - responseTemplate = responseNode.Raw - - if traceResult := root.Get("traceId"); traceResult.Exists() && traceResult.String() != "" { - traceID = traceResult.String() - } - - if roleResult := responseNode.Get("candidates.0.content.role"); roleResult.Exists() { - role = roleResult.String() - } - - if finishResult := responseNode.Get("candidates.0.finishReason"); finishResult.Exists() && finishResult.String() != "" { - finishReason = finishResult.String() - } - - if modelResult := responseNode.Get("modelVersion"); modelResult.Exists() && modelResult.String() != "" { - modelVersion = modelResult.String() - } - if responseIDResult := responseNode.Get("responseId"); responseIDResult.Exists() && responseIDResult.String() != "" { - responseID = responseIDResult.String() - } - if usageResult := responseNode.Get("usageMetadata"); usageResult.Exists() { - usageRaw = usageResult.Raw - } else if usageMetadataResult := root.Get("usageMetadata"); usageMetadataResult.Exists() { - usageRaw = usageMetadataResult.Raw - } - - if partsResult := responseNode.Get("candidates.0.content.parts"); partsResult.IsArray() { - for _, part := range partsResult.Array() { - hasFunctionCall := part.Get("functionCall").Exists() - hasInlineData := part.Get("inlineData").Exists() || part.Get("inline_data").Exists() - sig := part.Get("thoughtSignature").String() - if sig == "" { - sig = part.Get("thought_signature").String() - } - text := part.Get("text").String() - thought := part.Get("thought").Bool() - - if hasFunctionCall || hasInlineData { - flushPending() - parts = append(parts, normalizePart(part)) - continue - } - - if thought || part.Get("text").Exists() { - kind := "text" - if thought { - kind = "thought" - } - if pendingKind != "" && pendingKind != kind { - flushPending() - } - pendingKind = kind - pendingText.WriteString(text) - if kind == "thought" && sig != "" { - pendingThoughtSig = sig - } - continue - } - - flushPending() - parts = append(parts, normalizePart(part)) - } - } - } - flushPending() - - if responseTemplate == "" { - responseTemplate = `{"candidates":[{"content":{"role":"model","parts":[]}}]}` - } - - partsJSON, _ := json.Marshal(parts) - updatedTemplate, _ := sjson.SetRawBytes([]byte(responseTemplate), "candidates.0.content.parts", partsJSON) - responseTemplate = string(updatedTemplate) - if role != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.content.role", role) - responseTemplate = string(updatedTemplate) - } - if finishReason != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.finishReason", finishReason) - responseTemplate = string(updatedTemplate) - } - if modelVersion != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "modelVersion", modelVersion) - responseTemplate = string(updatedTemplate) - } - if responseID != "" { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "responseId", responseID) - responseTemplate = string(updatedTemplate) - } - if usageRaw != "" { - updatedTemplate, _ = sjson.SetRawBytes([]byte(responseTemplate), "usageMetadata", []byte(usageRaw)) - responseTemplate = string(updatedTemplate) - } else if !gjson.Get(responseTemplate, "usageMetadata").Exists() { - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.promptTokenCount", 0) - responseTemplate = string(updatedTemplate) - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.candidatesTokenCount", 0) - responseTemplate = string(updatedTemplate) - updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.totalTokenCount", 0) - responseTemplate = string(updatedTemplate) - } - - output := `{"response":{},"traceId":""}` - updatedOutput, _ := sjson.SetRawBytes([]byte(output), "response", []byte(responseTemplate)) - output = string(updatedOutput) - if traceID != "" { - updatedOutput, _ = sjson.SetBytes([]byte(output), "traceId", traceID) - output = string(updatedOutput) - } - return []byte(output) -} - -// ExecuteStream performs a streaming request to the Antigravity API. -func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { - if opts.Alt == "responses/compact" { - return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} - } - baseModel := thinking.ParseSuffix(req.Model).ModelName - - ctx = context.WithValue(ctx, "alt", "") - if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { - return nil, homeKVUnavailableStatusErr(errCooldown) - } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { - log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) - d := remaining - return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("antigravity") - - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) - if errValidate != nil { - return nil, errValidate - } - req.Payload = originalPayload - token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) - if errToken != nil { - return nil, errToken - } - if updatedAuth != nil { - auth = updatedAuth - } - - originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, true) - translated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, true) - - translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return nil, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers) - translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated) - translated, _ = sjson.DeleteBytes(translated, "request.stream") - reporter.SetTranslatedReasoningEffort(translated, to.String()) - - useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) - - baseURLs := antigravityBaseURLFallbackOrder(auth) - httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - - attempts := antigravityRetryAttempts(auth, e.cfg) - -attemptLoop: - for attempt := 0; attempt < attempts; attempt++ { - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - requestPayload := translated - if useCredits { - if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { - requestPayload = cp - helps.MarkCreditsUsed(ctx) - } - } - replayScope := antigravityReasoningReplayScope{} - if antigravityUsesReasoningReplayCache(baseModel) { - var errReplay error - requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) - if errReplay != nil { - err = errReplay - return nil, err - } - } - httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) - if errReq != nil { - err = errReq - return nil, err - } - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return nil, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errDo - return nil, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { - err = errRead - return nil, err - } - if errCtx := ctx.Err(); errCtx != nil { - err = errCtx - return nil, err - } - lastStatus = 0 - lastBody = nil - lastErr = errRead - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - err = errRead - return nil, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - if httpResp.StatusCode == http.StatusTooManyRequests { - decision := decideAntigravity429(bodyBytes) - - switch decision.kind { - case antigravity429DecisionInstantRetrySameAuth: - if attempt+1 < attempts { - if decision.retryAfter != nil && *decision.retryAfter > 0 { - wait := antigravityInstantRetryDelay(*decision.retryAfter) - log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) - if errWait := antigravityWait(ctx, wait); errWait != nil { - return nil, errWait - } - } - continue attemptLoop - } - case antigravity429DecisionShortCooldownSwitchAuth: - if decision.retryAfter != nil && *decision.retryAfter > 0 { - if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { - err = homeKVUnavailableStatusErr(errMarkCooldown) - return nil, err - } - log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel) - } - case antigravity429DecisionFullQuotaExhausted: - if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { - markAntigravityCreditsPermanentlyDisabled(auth) - } - // No credits logic - just fall through to error return below - } - } - - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { - delay := antigravityTransient429RetryDelay(attempt) - log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return nil, errWait - } - continue attemptLoop - } - if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - if attempt+1 < attempts { - delay := antigravityNoCapacityRetryDelay(attempt) - log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return nil, errWait - } - continue attemptLoop - } - } - if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { - if attempt+1 < attempts { - delay := antigravitySoftRateLimitDelay(attempt) - log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) - if errWait := antigravityWait(ctx, delay); errWait != nil { - return nil, errWait - } - continue attemptLoop - } - } - if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { - err = errClear - return nil, err - } - err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) - return nil, err - } - - // Stream success - if useCredits { - clearAntigravityCreditsFailureState(auth) - } - replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) - out := make(chan cliproxyexecutor.StreamChunk) - go func(resp *http.Response) { - defer close(out) - defer func() { - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response line error: %v", errClose) - } - }() - scanner := bufio.NewScanner(resp.Body) - scanner.Buffer(nil, streamScannerBuffer) - claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) - var param any - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if replayAccumulator != nil { - replayAccumulator.ObserveSSELine(line) - } - - // Filter usage metadata for all models - // Only retain usage statistics in the terminal chunk - line = helps.FilterSSEUsageMetadata(line) - - payload := helps.JSONPayload(line) - if payload == nil { - continue - } - - if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { - reporter.Publish(ctx, detail) - } - - payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, payload) - chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m, claudeInputTokens) - for i := range chunks { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: - case <-ctx.Done(): - return - } - } - } - tail := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m, claudeInputTokens) - for i := range tail { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: tail[i]}: - case <-ctx.Done(): - return - } - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - select { - case out <- cliproxyexecutor.StreamChunk{Err: errScan}: - case <-ctx.Done(): - } - } else { - if replayAccumulator != nil { - replayAccumulator.Commit(ctx) - } - reporter.EnsurePublished(ctx) - } - }(httpResp) - return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil - } - - switch { - case lastStatus != 0: - err = newAntigravityStatusErr(lastStatus, lastBody) - case lastErr != nil: - err = lastErr - default: - err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} - } - return nil, err - } - - return nil, err -} - -// Refresh refreshes the authentication credentials using the refresh token. -func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { - return refreshed, err - } - if auth == nil { - return auth, nil - } - updated, errRefresh := e.refreshToken(ctx, auth.Clone()) - if errRefresh != nil { - return nil, errRefresh - } - return updated, nil -} - -func (e *AntigravityExecutor) ShouldPrepareRequestAuth(auth *cliproxyauth.Auth) bool { - return antigravityProjectIDFromAuth(auth) == "" -} - -func (e *AntigravityExecutor) PrepareRequestAuth(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - if auth == nil || !e.ShouldPrepareRequestAuth(auth) { - return nil, nil - } - - updated := auth.Clone() - token, refreshedAuth, errToken := e.ensureAccessToken(ctx, updated) - if errToken != nil { - return nil, errToken - } - if refreshedAuth != nil { - updated = refreshedAuth - } - if antigravityProjectIDFromAuth(updated) != "" { - return updated, nil - } - - projectID, errProject := e.fetchAntigravityProjectID(ctx, updated, token) - if errProject != nil { - return nil, missingAntigravityProjectIDError(errProject) - } - if projectID == "" { - return nil, missingAntigravityProjectIDError(nil) - } - if updated.Metadata == nil { - updated.Metadata = make(map[string]any) - } - updated.Metadata["project_id"] = projectID - return updated, nil -} - -// CountTokens counts tokens for the given request using the Antigravity API. -func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - baseModel := thinking.ParseSuffix(req.Model).ModelName - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("antigravity") - respCtx := context.WithValue(ctx, "alt", opts.Alt) - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayloadSource, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayloadSource) - if errValidate != nil { - return cliproxyexecutor.Response{}, errValidate - } - req.Payload = originalPayloadSource - token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) - if errToken != nil { - return cliproxyexecutor.Response{}, errToken - } - if updatedAuth != nil { - auth = updatedAuth - } - if strings.TrimSpace(token) == "" { - return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} - } - - // Prepare payload once (doesn't depend on baseURL) - payload := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, false) - - payload, err := thinking.ApplyThinking(payload, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return cliproxyexecutor.Response{}, err - } - payload = sanitizeAntigravityGeminiRequestSignatures(baseModel, payload) - preparedPayload, _, errReplay := prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, payload) - if errReplay != nil { - return cliproxyexecutor.Response{}, errReplay - } - payload = preparedPayload - - payload = helps.DeleteJSONField(payload, "project") - payload = helps.DeleteJSONField(payload, "model") - payload = helps.DeleteJSONField(payload, "request.safetySettings") - - baseURLs := antigravityBaseURLFallbackOrder(auth) - httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - - var lastStatus int - var lastBody []byte - var lastErr error - - for idx, baseURL := range baseURLs { - base := strings.TrimSuffix(baseURL, "/") - if base == "" { - base = buildBaseURL(auth) - } - - var requestURL strings.Builder - requestURL.WriteString(base) - requestURL.WriteString(antigravityCountTokensPath) - if opts.Alt != "" { - requestURL.WriteString("?$alt=") - requestURL.WriteString(url.QueryEscape(opts.Alt)) - } - - httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload)) - if errReq != nil { - return cliproxyexecutor.Response{}, errReq - } - httpReq.Close = true - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Authorization", "Bearer "+token) - httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) - if host := resolveHost(base); host != "" { - httpReq.Host = host - } - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) - - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: requestURL.String(), - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: payload, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errDo) - if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { - return cliproxyexecutor.Response{}, errDo - } - lastStatus = 0 - lastBody = nil - lastErr = errDo - if idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - return cliproxyexecutor.Response{}, errDo - } - - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - return cliproxyexecutor.Response{}, errRead - } - helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) - - if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { - count := gjson.GetBytes(bodyBytes, "totalTokens").Int() - translated := sdktranslator.TranslateTokenCount(respCtx, to, responseFormat, count, bodyBytes) - return cliproxyexecutor.Response{Payload: translated, Headers: httpResp.Header.Clone()}, nil - } - - lastStatus = httpResp.StatusCode - lastBody = append([]byte(nil), bodyBytes...) - lastErr = nil - if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { - log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) - continue - } - sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} - if httpResp.StatusCode == http.StatusTooManyRequests { - if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - return cliproxyexecutor.Response{}, sErr - } - - switch { - case lastStatus != 0: - sErr := statusErr{code: lastStatus, msg: string(lastBody)} - if lastStatus == http.StatusTooManyRequests { - if retryAfter, parseErr := helps.ParseRetryDelay(lastBody); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - return cliproxyexecutor.Response{}, sErr - case lastErr != nil: - return cliproxyexecutor.Response{}, lastErr - default: - return cliproxyexecutor.Response{}, statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} - } -} - -func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *cliproxyauth.Auth) (string, *cliproxyauth.Auth, error) { - if auth == nil { - return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} - } - accessToken := metaStringValue(auth.Metadata, "access_token") - expiry := tokenExpiry(auth.Metadata) - if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) { - e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken) - return accessToken, nil, nil - } - refreshCtx := context.Background() - if ctx != nil { - if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { - refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) - } - } - if refreshed, handled, err := helps.RefreshAuthViaHome(refreshCtx, e.cfg, auth); handled { - if err != nil { - return "", nil, err - } - token := metaStringValue(refreshed.Metadata, "access_token") - if strings.TrimSpace(token) == "" { - return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} - } - e.maybeRefreshAntigravityCreditsHint(ctx, refreshed, token) - return token, refreshed, nil - } - - updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone()) - if errRefresh != nil { - return "", nil, errRefresh - } - return metaStringValue(updated.Metadata, "access_token"), updated, nil -} - -func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) { - if e == nil || auth == nil || !antigravityCreditsRetryEnabled(e.cfg) { - return - } - if ctx != nil && ctx.Err() != nil { - return - } - authID := strings.TrimSpace(auth.ID) - if authID == "" { - return - } - if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && hint.Known { - return - } - if strings.TrimSpace(accessToken) == "" { - accessToken = metaStringValue(auth.Metadata, "access_token") - } - if strings.TrimSpace(accessToken) == "" { - return - } - - if client, homeMode, errClient := currentAntigravityKVClient(); homeMode { - if errClient != nil { - log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errClient) - return - } - written, errSetNX := client.KVSetNX(context.Background(), antigravityCreditsRefreshLockKey(authID), []byte("1"), antigravityCreditsHintRefreshInterval) - if errSetNX != nil { - log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errSetNX) - return - } - if !written { - return - } - refreshCtx := context.Background() - if ctx != nil { - if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { - refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) - } - } - refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout) - authCopy := auth.Clone() - go func(auth *cliproxyauth.Auth, token string) { - defer cancel() - e.updateAntigravityCreditsBalance(refreshCtx, auth, token) - }(authCopy, accessToken) - return - } - - state := &antigravityCreditsHintRefreshState{} - if existing, loaded := antigravityCreditsHintRefreshByID.LoadOrStore(authID, state); loaded { - if cast, ok := existing.(*antigravityCreditsHintRefreshState); ok && cast != nil { - state = cast - } else { - antigravityCreditsHintRefreshByID.Delete(authID) - antigravityCreditsHintRefreshByID.Store(authID, state) - } - } - - now := time.Now() - if !state.mu.TryLock() { - return - } - if !state.lastAttempt.IsZero() && now.Sub(state.lastAttempt) < antigravityCreditsHintRefreshInterval { - state.mu.Unlock() - return - } - state.lastAttempt = now - - refreshCtx := context.Background() - if ctx != nil { - if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { - refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) - } - } - refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout) - authCopy := auth.Clone() - - go func(state *antigravityCreditsHintRefreshState, auth *cliproxyauth.Auth, token string) { - defer cancel() - defer state.mu.Unlock() - e.updateAntigravityCreditsBalance(refreshCtx, auth, token) - }(state, authCopy, accessToken) -} - -func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - if auth == nil { - return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} - } - refreshToken := metaStringValue(auth.Metadata, "refresh_token") - if refreshToken == "" { - return auth, statusErr{code: http.StatusUnauthorized, msg: "missing refresh token"} - } - if ctx == nil { - ctx = context.Background() - } - refreshToken = strings.TrimSpace(refreshToken) - - result, errRefresh, _ := antigravityRefreshGroup.Do(refreshToken, func() (interface{}, error) { - return e.refreshTokenSingleFlight(context.WithoutCancel(ctx), auth, refreshToken) - }) - if errRefresh != nil { - return auth, errRefresh - } - tokenResp, ok := result.(*antigravityTokenRefreshData) - if !ok || tokenResp == nil { - return auth, fmt.Errorf("antigravity token refresh failed: invalid single-flight result") - } - - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["access_token"] = tokenResp.AccessToken - if tokenResp.RefreshToken != "" { - auth.Metadata["refresh_token"] = tokenResp.RefreshToken - } - auth.Metadata["expires_in"] = tokenResp.ExpiresIn - now := time.Now() - auth.Metadata["timestamp"] = now.UnixMilli() - auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339) - auth.Metadata["type"] = antigravityAuthType - if errProject := e.ensureAntigravityProjectID(ctx, auth, tokenResp.AccessToken); errProject != nil { - log.Warnf("antigravity executor: ensure project id failed: %v", errProject) - } - e.updateAntigravityCreditsBalance(ctx, auth, tokenResp.AccessToken) - return auth, nil -} - -func (e *AntigravityExecutor) refreshTokenSingleFlight(ctx context.Context, auth *cliproxyauth.Auth, refreshToken string) (*antigravityTokenRefreshData, error) { - form := url.Values{} - form.Set("client_id", antigravityClientID) - form.Set("client_secret", antigravityClientSecret) - form.Set("grant_type", "refresh_token") - form.Set("refresh_token", refreshToken) - - httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode())) - if errReq != nil { - return nil, errReq - } - httpReq.Header.Set("Host", "oauth2.googleapis.com") - httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") - // Real Antigravity uses Go's default User-Agent for OAuth token refresh - httpReq.Header.Set("User-Agent", "Go-http-client/2.0") - - httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - return nil, errDo - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close response body error: %v", errClose) - } - }() - - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errRead != nil { - return nil, errRead - } - - if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} - if httpResp.StatusCode == http.StatusTooManyRequests { - if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { - sErr.retryAfter = retryAfter - } - } - return nil, sErr - } - - var tokenResp antigravityTokenRefreshData - if errUnmarshal := json.Unmarshal(bodyBytes, &tokenResp); errUnmarshal != nil { - return nil, errUnmarshal - } - - return &tokenResp, nil -} - -func (e *AntigravityExecutor) ensureAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) error { - if auth == nil { - return nil - } - - if antigravityProjectIDFromAuth(auth) != "" { - return nil - } - - projectID, errFetch := e.fetchAntigravityProjectID(ctx, auth, accessToken) - if errFetch != nil { - return errFetch - } - if projectID == "" { - return nil - } - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["project_id"] = projectID - - return nil -} - -func (e *AntigravityExecutor) fetchAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) (string, error) { - token := strings.TrimSpace(accessToken) - if token == "" { - token = metaStringValue(auth.Metadata, "access_token") - } - if token == "" { - return "", nil - } - - httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - projectID, errFetch := sdkAuth.FetchAntigravityProjectID(ctx, token, httpClient) - if errFetch != nil { - return "", errFetch - } - return strings.TrimSpace(projectID), nil -} - -func (e *AntigravityExecutor) projectIDForRequest(_ context.Context, auth *cliproxyauth.Auth, _ string) (string, error) { - if projectID := antigravityProjectIDFromAuth(auth); projectID != "" { - return projectID, nil - } - return "", missingAntigravityProjectIDError(nil) -} - -func antigravityProjectIDFromAuth(auth *cliproxyauth.Auth) string { - if auth == nil || auth.Metadata == nil { - return "" - } - if pid, ok := auth.Metadata["project_id"].(string); ok { - return strings.TrimSpace(pid) - } - return "" -} - -func missingAntigravityProjectIDError(cause error) statusErr { - msg := "antigravity auth missing project_id" - if cause != nil { - msg = fmt.Sprintf("%s: %v", msg, cause) - } - return statusErr{code: http.StatusBadRequest, msg: msg} -} - -func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) { - if auth == nil || strings.TrimSpace(auth.ID) == "" { - return - } - token := strings.TrimSpace(accessToken) - if token == "" { - token = metaStringValue(auth.Metadata, "access_token") - } - if token == "" { - return - } - - userAgent := resolveUserAgent(auth) - loadReqBody, errMarshal := json.Marshal(map[string]any{ - "metadata": map[string]string{ - "ideType": "ANTIGRAVITY", - }, - }) - if errMarshal != nil { - log.Debugf("antigravity executor: marshal loadCodeAssist request error: %v", errMarshal) - return - } - baseURL := antigravityLoadCodeAssistBaseURL(auth) - endpointURL := strings.TrimSuffix(baseURL, "/") + "/v1internal:loadCodeAssist" - httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(loadReqBody)) - if errReq != nil { - log.Debugf("antigravity executor: create loadCodeAssist request error: %v", errReq) - return - } - httpReq.Header.Set("Authorization", "Bearer "+token) - httpReq.Header.Set("Accept", "*/*") - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("User-Agent", userAgent) - - httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) - httpResp, errDo := httpClient.Do(httpReq) - if errDo != nil { - log.Debugf("antigravity executor: loadCodeAssist request error: %v", errDo) - return - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("antigravity executor: close loadCodeAssist response body error: %v", errClose) - } - }() - - bodyBytes, errRead := io.ReadAll(httpResp.Body) - if errRead != nil || httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { - log.Debugf("antigravity executor: loadCodeAssist returned status %d, err=%v", httpResp.StatusCode, errRead) - return - } - - authID := strings.TrimSpace(auth.ID) - paidTierID := strings.TrimSpace(gjson.GetBytes(bodyBytes, "paidTier.id").String()) - - credits := gjson.GetBytes(bodyBytes, "paidTier.availableCredits") - if !credits.IsArray() { - cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ - Known: true, - Available: false, - PaidTierID: paidTierID, - UpdatedAt: time.Now(), - }) - return - } - for _, credit := range credits.Array() { - if !strings.EqualFold(credit.Get("creditType").String(), "GOOGLE_ONE_AI") { - continue - } - creditAmount, errCA := strconv.ParseFloat(strings.TrimSpace(credit.Get("creditAmount").String()), 64) - if errCA != nil { - continue - } - minAmount, errMA := strconv.ParseFloat(strings.TrimSpace(credit.Get("minimumCreditAmountForUsage").String()), 64) - if errMA != nil { - continue - } - bal := antigravityCreditsBalance{ - CreditAmount: creditAmount, - MinCreditAmount: minAmount, - PaidTierID: paidTierID, - Known: true, - } - storeAntigravityCreditsBalanceBestEffort(authID, bal) - cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ - Known: true, - Available: creditAmount >= minAmount, - CreditAmount: creditAmount, - MinCreditAmount: minAmount, - PaidTierID: paidTierID, - UpdatedAt: time.Now(), - }) - if creditAmount >= minAmount { - clearAntigravityCreditsPermanentlyDisabled(auth) - } - return - } -} - -func (e *AntigravityExecutor) buildRequest(ctx context.Context, auth *cliproxyauth.Auth, token, modelName string, payload []byte, stream bool, alt, baseURL string, derivedSessionIDs ...string) (*http.Request, error) { - if token == "" { - return nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} - } - - base := strings.TrimSuffix(baseURL, "/") - if base == "" { - base = buildBaseURL(auth) - } - path := antigravityGeneratePath - if stream { - path = antigravityStreamPath - } - var requestURL strings.Builder - requestURL.WriteString(base) - requestURL.WriteString(path) - if stream { - if alt != "" { - requestURL.WriteString("?$alt=") - requestURL.WriteString(url.QueryEscape(alt)) - } else { - requestURL.WriteString("?alt=sse") - } - } else if alt != "" { - requestURL.WriteString("?$alt=") - requestURL.WriteString(url.QueryEscape(alt)) - } - - projectID, errProject := e.projectIDForRequest(ctx, auth, token) - if errProject != nil { - return nil, errProject - } - payload = geminiToAntigravity(modelName, payload, projectID, derivedSessionIDs...) - - // Cap maxOutputTokens to model's max_completion_tokens from registry - if maxOut := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxOut.Exists() && maxOut.Type == gjson.Number { - if modelInfo := registry.LookupModelInfo(modelName, "antigravity"); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 { - if int(maxOut.Int()) > modelInfo.MaxCompletionTokens { - payload, _ = sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", modelInfo.MaxCompletionTokens) - } - } - } - - useAntigravitySchema := strings.Contains(modelName, "claude") || strings.Contains(modelName, "gemini-3-pro") || strings.Contains(modelName, "gemini-3.1-pro") - var ( - bodyReader io.Reader - payloadLog []byte - ) - if antigravityRequestNeedsSchemaSanitization(payload) { - payloadStr := sanitizeAntigravityRequestSchemas(string(payload), useAntigravitySchema) - - if strings.Contains(modelName, "claude") { - updated, _ := sjson.SetBytes([]byte(payloadStr), "request.toolConfig.functionCallingConfig.mode", "VALIDATED") - payloadStr = string(updated) - } else { - payloadStr, _ = sjson.Delete(payloadStr, "request.generationConfig.maxOutputTokens") - } - - payloadStrBytes := applyAntigravityNativeSignatureReplayIfNeeded(modelName, []byte(payloadStr)) - bodyReader = bytes.NewReader(payloadStrBytes) - if e.cfg != nil && e.cfg.RequestLog { - payloadLog = append([]byte(nil), payloadStrBytes...) - } - } else { - if strings.Contains(modelName, "claude") { - payload, _ = sjson.SetBytes(payload, "request.toolConfig.functionCallingConfig.mode", "VALIDATED") - } else { - payload, _ = sjson.DeleteBytes(payload, "request.generationConfig.maxOutputTokens") - } - - payload = applyAntigravityNativeSignatureReplayIfNeeded(modelName, payload) - bodyReader = bytes.NewReader(payload) - if e.cfg != nil && e.cfg.RequestLog { - payloadLog = append([]byte(nil), payload...) - } - } - - // if useAntigravitySchema { - // systemInstructionPartsResult := gjson.Get(payloadStr, "request.systemInstruction.parts") - // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.role", "user") - // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.0.text", systemInstruction) - // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.1.text", fmt.Sprintf("Please ignore following [ignore]%s[/ignore]", systemInstruction)) - - // if systemInstructionPartsResult.Exists() && systemInstructionPartsResult.IsArray() { - // for _, partResult := range systemInstructionPartsResult.Array() { - // payloadStr, _ = sjson.SetRawBytes([]byte(payloadStr), "request.systemInstruction.parts.-1", []byte(partResult.Raw)) - // } - // } - // } - - httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bodyReader) - if errReq != nil { - return nil, errReq - } - httpReq.Close = true - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Authorization", "Bearer "+token) - httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) - if host := resolveHost(base); host != "" { - httpReq.Host = host - } - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(httpReq, attrs) - - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: requestURL.String(), - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: payloadLog, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - return httpReq, nil -} - -// sanitizeAntigravityRequestSchemas cleans the JSON schemas carried by an Antigravity request. -// -// Cleaning is applied only to the payload locations that actually hold a JSON schema. The schema -// cleaner rewrites keys such as "title", "format", "default" and "const", which are also ordinary -// data keys inside functionCall arguments replayed from conversation history. Running it over the -// whole document silently mutated that history, so tools lost required argument fields and the -// model imitated the corrupted examples on later turns. -func sanitizeAntigravityRequestSchemas(payloadStr string, useAntigravitySchema bool) string { - for _, base := range antigravityFunctionDeclarationPaths(payloadStr) { - oldPath := base + ".parametersJsonSchema" - if !gjson.Get(payloadStr, oldPath).Exists() { - continue - } - renamed, errRename := util.RenameKey(payloadStr, oldPath, base+".parameters") - if errRename != nil { - log.Debugf("antigravity: failed to rename %s: %v", oldPath, errRename) - continue - } - payloadStr = renamed - } - - clean := util.CleanJSONSchemaForGemini - if useAntigravitySchema { - clean = util.CleanJSONSchemaForAntigravity - } - - for _, schemaPath := range antigravitySchemaPaths(payloadStr) { - schema := gjson.Get(payloadStr, schemaPath) - if !schema.Exists() { - continue - } - updated, errSet := sjson.SetRawBytes([]byte(payloadStr), schemaPath, []byte(cleanNestedSchema(clean, schema.Raw))) - if errSet != nil { - log.Debugf("antigravity: failed to write cleaned schema at %s: %v", schemaPath, errSet) - continue - } - payloadStr = string(updated) - } - - return payloadStr -} - -// antigravitySchemaWrapperKey nests a schema during cleaning. It is never sent upstream. -const antigravitySchemaWrapperKey = "schema" - -// cleanNestedSchema cleans a schema with it nested one level down, then unwraps it. -// -// The cleaner deliberately skips placeholder insertion for a top-level schema, but Claude's -// VALIDATED mode needs every tool schema to declare at least one required property. Whole-payload -// cleaning always saw tool schemas nested inside the request, so nesting is reproduced here to keep -// the emitted schema byte-identical to the previous behaviour. -func cleanNestedSchema(clean func(string) string, schemaRaw string) string { - wrapped, errWrap := sjson.SetRaw("{}", antigravitySchemaWrapperKey, schemaRaw) - if errWrap != nil { - return clean(schemaRaw) - } - if unwrapped := gjson.Get(clean(wrapped), antigravitySchemaWrapperKey); unwrapped.Exists() { - return unwrapped.Raw - } - return clean(schemaRaw) -} - -// antigravityFunctionDeclarationPaths returns the path of every function declaration in the request. -// Both the camelCase and snake_case spellings are accepted because callers reach this executor -// through different translators. -func antigravityFunctionDeclarationPaths(payloadStr string) []string { - tools := gjson.Get(payloadStr, "request.tools") - if !tools.IsArray() { - return nil - } - paths := make([]string, 0, len(tools.Array())) - for i, tool := range tools.Array() { - for _, declKey := range []string{"functionDeclarations", "function_declarations"} { - decls := tool.Get(declKey) - if !decls.IsArray() { - continue - } - for j := range decls.Array() { - paths = append(paths, fmt.Sprintf("request.tools.%d.%s.%d", i, declKey, j)) - } - } - } - return paths -} - -// antigravitySchemaPaths returns every payload path that holds a JSON schema document. -// A function declaration may carry a schema for its parameters and for its result, so all of -// them must be cleaned; anything omitted here reaches the upstream API uncleaned. -func antigravitySchemaPaths(payloadStr string) []string { - paths := make([]string, 0, 12) - for _, base := range antigravityFunctionDeclarationPaths(payloadStr) { - for _, key := range antigravityDeclarationSchemaKeys { - if gjson.Get(payloadStr, base+"."+key).IsObject() { - paths = append(paths, base+"."+key) - } - } - } - for _, container := range antigravityGenerationConfigContainers { - for _, key := range antigravityGenerationSchemaKeys { - p := container + "." + key - if gjson.Get(payloadStr, p).IsObject() { - paths = append(paths, p) - } - } - } - return paths -} - -// The upstream API is proto-JSON and accepts either spelling, and the Gemini translator forwards -// whichever one the client sent. Both are therefore cleaned where they sit rather than renamed: -// renaming would alter the body the client asked for, and only the unsupported keywords inside a -// schema cause upstream errors. The one exception is parametersJsonSchema, renamed onto parameters -// above because whole-payload cleaning did the same. -var ( - antigravityDeclarationSchemaKeys = []string{ - "parameters", "parametersJsonSchema", "parameters_json_schema", - "response", "responseJsonSchema", "response_json_schema", - } - antigravityGenerationConfigContainers = []string{ - "request.generationConfig", "request.generation_config", - } - antigravityGenerationSchemaKeys = []string{ - "responseSchema", "responseJsonSchema", "response_schema", "response_json_schema", - } -) - -func antigravityRequestNeedsSchemaSanitization(payload []byte) bool { - if gjson.GetBytes(payload, "request.tools.0").Exists() { - return true - } - for _, container := range antigravityGenerationConfigContainers { - for _, key := range antigravityGenerationSchemaKeys { - if gjson.GetBytes(payload, container+"."+key).Exists() { - return true - } - } - } - return false -} - -func tokenExpiry(metadata map[string]any) time.Time { - if metadata == nil { - return time.Time{} - } - if expStr, ok := metadata["expired"].(string); ok { - expStr = strings.TrimSpace(expStr) - if expStr != "" { - if parsed, errParse := time.Parse(time.RFC3339, expStr); errParse == nil { - return parsed - } - } - } - expiresIn, hasExpires := int64Value(metadata["expires_in"]) - tsMs, hasTimestamp := int64Value(metadata["timestamp"]) - if hasExpires && hasTimestamp { - return time.Unix(0, tsMs*int64(time.Millisecond)).Add(time.Duration(expiresIn) * time.Second) - } - return time.Time{} -} - -func metaStringValue(metadata map[string]any, key string) string { - if metadata == nil { - return "" - } - if v, ok := metadata[key]; ok { - switch typed := v.(type) { - case string: - return strings.TrimSpace(typed) - case []byte: - return strings.TrimSpace(string(typed)) - } - } - return "" -} - -func int64Value(value any) (int64, bool) { - switch typed := value.(type) { - case int: - return int64(typed), true - case int64: - return typed, true - case float64: - return int64(typed), true - case json.Number: - if i, errParse := typed.Int64(); errParse == nil { - return i, true - } - case string: - if strings.TrimSpace(typed) == "" { - return 0, false - } - if i, errParse := strconv.ParseInt(strings.TrimSpace(typed), 10, 64); errParse == nil { - return i, true - } - } - return 0, false -} - -func buildBaseURL(auth *cliproxyauth.Auth) string { - if baseURLs := antigravityBaseURLFallbackOrder(auth); len(baseURLs) > 0 { - return baseURLs[0] - } - return antigravityBaseURLDaily -} - -func antigravityLoadCodeAssistBaseURL(auth *cliproxyauth.Auth) string { - if base := resolveCustomAntigravityBaseURL(auth); base != "" { - return base - } - return antigravityBaseURLProd -} - -func resolveHost(base string) string { - parsed, errParse := url.Parse(base) - if errParse != nil { - return "" - } - if parsed.Host != "" { - return parsed.Host - } - return strings.TrimPrefix(strings.TrimPrefix(base, "https://"), "http://") -} - -func resolveUserAgent(auth *cliproxyauth.Auth) string { - return misc.AntigravityRequestUserAgent(antigravityConfiguredUserAgent(auth)) -} - -func resolveLoadCodeAssistUserAgent(auth *cliproxyauth.Auth) string { - return misc.AntigravityLoadCodeAssistUserAgent(antigravityConfiguredUserAgent(auth)) -} - -func antigravityConfiguredUserAgent(auth *cliproxyauth.Auth) string { - raw := "" - if auth != nil { - if auth.Attributes != nil { - if ua := strings.TrimSpace(auth.Attributes["user_agent"]); ua != "" { - raw = ua - } - } - if raw == "" && auth.Metadata != nil { - if ua, ok := auth.Metadata["user_agent"].(string); ok && strings.TrimSpace(ua) != "" { - raw = strings.TrimSpace(ua) - } - } - } - return raw -} - -func antigravityRetryAttempts(auth *cliproxyauth.Auth, cfg *config.Config) int { - retry := 0 - if cfg != nil { - retry = cfg.RequestRetry - } - if auth != nil { - if override, ok := auth.RequestRetryOverride(); ok { - retry = override - } - } - if retry < 0 { - retry = 0 - } - attempts := retry + 1 - if attempts < 1 { - return 1 - } - return attempts -} - -func antigravityShouldRetryNoCapacity(statusCode int, body []byte) bool { - if statusCode != http.StatusServiceUnavailable { - return false - } - if len(body) == 0 { - return false - } - msg := strings.ToLower(string(body)) - return strings.Contains(msg, "no capacity available") -} - -func antigravityShouldRetryTransientResourceExhausted429(statusCode int, body []byte) bool { - if statusCode != http.StatusTooManyRequests { - return false - } - if len(body) == 0 { - return false - } - if classifyAntigravity429(body) != antigravity429Unknown { - return false - } - status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String()) - if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") { - return false - } - msg := strings.ToLower(string(body)) - return strings.Contains(msg, "resource has been exhausted") -} - -func antigravityShouldRetrySoftRateLimit(statusCode int, body []byte) bool { - if statusCode != http.StatusTooManyRequests { - return false - } - return decideAntigravity429(body).kind == antigravity429DecisionSoftRetry -} - -func antigravityShouldBypassShortCooldown(ctx context.Context, cfg *config.Config) bool { - return cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(cfg) -} - -func antigravitySoftRateLimitDelay(attempt int) time.Duration { - if attempt < 0 { - attempt = 0 - } - base := time.Duration(attempt+1) * 500 * time.Millisecond - if base > 3*time.Second { - base = 3 * time.Second - } - return base -} - -func antigravityShortCooldownKey(auth *cliproxyauth.Auth, modelName string) string { - if auth == nil { - return "" - } - authID := strings.TrimSpace(auth.ID) - modelName = strings.TrimSpace(modelName) - if authID == "" || modelName == "" { - return "" - } - return authID + "|" + modelName + "|sc" -} - -func antigravityCreditsBalanceKey(authID string) string { - return "cpa:antigravity:credits-balance:" + strings.TrimSpace(authID) -} - -func antigravityCreditsRefreshLockKey(authID string) string { - return "cpa:antigravity:credits-refresh-lock:" + strings.TrimSpace(authID) -} - -func antigravityShortCooldownKVKey(auth *cliproxyauth.Auth, modelName string) string { - if auth == nil { - return "" - } - authID := strings.TrimSpace(auth.ID) - modelName = strings.TrimSpace(modelName) - if authID == "" || modelName == "" { - return "" - } - return "cpa:antigravity:short-cooldown:" + authID + ":" + homekv.HashKeyPart(modelName) -} - -func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration) { - inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(context.Background(), auth, modelName, now) - if errCooldown != nil { - log.Errorf("antigravity executor: home kv cooldown read error: %v", errCooldown) - return false, 0 - } - return inCooldown, remaining -} - -func antigravityIsInShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration, error) { - kvKey := antigravityShortCooldownKVKey(auth, modelName) - client, homeMode, errClient := currentAntigravityKVClient() - if homeMode { - if errClient != nil { - return false, 0, errClient - } - if kvKey == "" { - return false, 0, nil - } - raw, found, errGet := client.KVGet(ctx, kvKey) - if errGet != nil || !found { - return false, 0, errGet - } - untilNano, errParse := strconv.ParseInt(strings.TrimSpace(string(raw)), 10, 64) - if errParse != nil { - return false, 0, errParse - } - remaining := time.Unix(0, untilNano).Sub(now) - if remaining <= 0 { - if _, errDel := client.KVDel(ctx, kvKey); errDel != nil { - return false, 0, errDel - } - return false, 0, nil - } - return true, remaining, nil - } - - key := antigravityShortCooldownKey(auth, modelName) - if key == "" { - return false, 0, nil - } - value, ok := antigravityShortCooldownByAuth.Load(key) - if !ok { - return false, 0, nil - } - until, ok := value.(time.Time) - if !ok || until.IsZero() { - antigravityShortCooldownByAuth.Delete(key) - return false, 0, nil - } - remaining := until.Sub(now) - if remaining <= 0 { - antigravityShortCooldownByAuth.Delete(key) - return false, 0, nil - } - return true, remaining, nil -} - -func markAntigravityShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) { - if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, modelName, now, duration); errMark != nil { - log.Errorf("antigravity executor: home kv cooldown write error: %v", errMark) - } -} - -func markAntigravityShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) error { - kvKey := antigravityShortCooldownKVKey(auth, modelName) - client, homeMode, errClient := currentAntigravityKVClient() - if homeMode { - if errClient != nil { - return errClient - } - if kvKey == "" || duration <= 0 { - return nil - } - until := now.Add(duration) - written, errSet := client.KVSet(ctx, kvKey, []byte(strconv.FormatInt(until.UnixNano(), 10)), homekv.KVSetOptions{EX: duration + 5*time.Second}) - if errSet != nil { - return errSet - } - if !written { - return fmt.Errorf("home kv store unavailable") - } - return nil - } - - key := antigravityShortCooldownKey(auth, modelName) - if key == "" { - return nil - } - antigravityShortCooldownByAuth.Store(key, now.Add(duration)) - return nil -} - -func storeAntigravityCreditsBalanceBestEffort(authID string, bal antigravityCreditsBalance) { - authID = strings.TrimSpace(authID) - if authID == "" { - return - } - if client, homeMode, errClient := currentAntigravityKVClient(); homeMode { - if errClient != nil { - log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errClient) - return - } - raw, errMarshal := json.Marshal(bal) - if errMarshal != nil { - log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errMarshal) - return - } - if _, errSet := client.KVSet(context.Background(), antigravityCreditsBalanceKey(authID), raw, homekv.KVSetOptions{EX: 30 * time.Minute}); errSet != nil { - log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errSet) - } - return - } - antigravityCreditsBalanceByAuth.Store(authID, bal) -} - -func homeKVUnavailableStatusErr(cause error) statusErr { - if cause == nil { - return statusErr{code: http.StatusServiceUnavailable, msg: "home kv store unavailable"} - } - return statusErr{code: http.StatusServiceUnavailable, msg: fmt.Sprintf("home kv store unavailable: %v", cause)} -} - -func antigravityNoCapacityRetryDelay(attempt int) time.Duration { - if attempt < 0 { - attempt = 0 - } - delay := time.Duration(attempt+1) * 250 * time.Millisecond - if delay > 2*time.Second { - delay = 2 * time.Second - } - return delay -} - -func antigravityTransient429RetryDelay(attempt int) time.Duration { - if attempt < 0 { - attempt = 0 - } - delay := time.Duration(attempt+1) * 100 * time.Millisecond - if delay > 500*time.Millisecond { - delay = 500 * time.Millisecond - } - return delay -} - -func antigravityInstantRetryDelay(wait time.Duration) time.Duration { - if wait <= 0 { - return 0 - } - return wait + 800*time.Millisecond -} - -func antigravityWait(ctx context.Context, wait time.Duration) error { - if wait <= 0 { - return nil - } - timer := time.NewTimer(wait) - defer timer.Stop() - select { - case <-ctx.Done(): - return ctx.Err() - case <-timer.C: - return nil - } -} - -var antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string { - if base := resolveCustomAntigravityBaseURL(auth); base != "" { - return []string{base} - } - return []string{ - antigravityBaseURLDaily, - antigravityBaseURLProd, - // antigravitySandboxBaseURLDaily, - } -} - -func resolveCustomAntigravityBaseURL(auth *cliproxyauth.Auth) string { - if auth == nil { - return "" - } - if auth.Attributes != nil { - if v := strings.TrimSpace(auth.Attributes["base_url"]); v != "" { - return strings.TrimSuffix(v, "/") - } - } - if auth.Metadata != nil { - if v, ok := auth.Metadata["base_url"].(string); ok { - v = strings.TrimSpace(v) - if v != "" { - return strings.TrimSuffix(v, "/") - } - } - } - return "" -} - -func geminiToAntigravity(modelName string, payload []byte, projectID string, derivedSessionIDs ...string) []byte { - template := payload - template = helps.SetStringIfDifferent(template, "model", modelName) - template = helps.SetStringIfDifferent(template, "userAgent", "antigravity") - - isImageModel := strings.Contains(modelName, "image") - reqType := strings.TrimSpace(gjson.GetBytes(template, "requestType").String()) - if reqType == "" { - if isImageModel { - reqType = "image_gen" - } else { - reqType = "agent" - } - template, _ = sjson.SetBytes(template, "requestType", reqType) - } - - if projectID != "" { - template = helps.SetStringIfDifferent(template, "project", projectID) - } else { - template, _ = sjson.DeleteBytes(template, "project") - } - - if isImageModel { - template, _ = sjson.SetBytes(template, "requestId", generateImageGenRequestID()) - } else if reqType != "web_search" { - template, _ = sjson.SetBytes(template, "requestId", generateRequestID()) - sessionID := strings.TrimSpace(gjson.GetBytes(template, "request.sessionId").String()) - if sessionID == "" && len(derivedSessionIDs) > 0 { - sessionID = strings.TrimSpace(derivedSessionIDs[0]) - } - if sessionID == "" { - sessionID = generateStableSessionID(payload) - } - template, _ = sjson.SetBytes(template, "request.sessionId", sessionID) - } - - template, _ = sjson.DeleteBytes(template, "request.safetySettings") - if toolConfig := gjson.GetBytes(template, "toolConfig"); toolConfig.Exists() && !gjson.GetBytes(template, "request.toolConfig").Exists() { - template, _ = sjson.SetRawBytes(template, "request.toolConfig", []byte(toolConfig.Raw)) - template, _ = sjson.DeleteBytes(template, "toolConfig") - } - return template -} - -func generateRequestID() string { - return "agent-" + uuid.NewString() -} - -func generateImageGenRequestID() string { - return fmt.Sprintf("image_gen/%d/%s/12", time.Now().UnixMilli(), uuid.NewString()) -} - -func generateSessionID() string { - randSourceMutex.Lock() - n := randSource.Int63n(9_000_000_000_000_000_000) - randSourceMutex.Unlock() - return "-" + strconv.FormatInt(n, 10) -} - -func generateStableSessionID(payload []byte) string { - contents := gjson.GetBytes(payload, "request.contents") - if contents.IsArray() { - for _, content := range contents.Array() { - if content.Get("role").String() == "user" { - text := content.Get("parts.0.text").String() - if text != "" { - h := sha256.Sum256([]byte(text)) - n := int64(binary.BigEndian.Uint64(h[:8])) & 0x7FFFFFFFFFFFFFFF - return "-" + strconv.FormatInt(n, 10) - } - } - } - } - return generateSessionID() -} diff --git a/internal/runtime/executor/antigravity_executor_auth.go b/internal/runtime/executor/antigravity_executor_auth.go new file mode 100644 index 000000000..108eb914e --- /dev/null +++ b/internal/runtime/executor/antigravity_executor_auth.go @@ -0,0 +1,320 @@ +package executor + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" +) + +// Refresh refreshes the authentication credentials using the refresh token. +func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { + if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { + return refreshed, err + } + if auth == nil { + return auth, nil + } + updated, errRefresh := e.refreshToken(ctx, auth.Clone()) + if errRefresh != nil { + return nil, errRefresh + } + return updated, nil +} + +func (e *AntigravityExecutor) ShouldPrepareRequestAuth(auth *cliproxyauth.Auth) bool { + return antigravityProjectIDFromAuth(auth) == "" +} + +func (e *AntigravityExecutor) PrepareRequestAuth(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { + if auth == nil || !e.ShouldPrepareRequestAuth(auth) { + return nil, nil + } + + updated := auth.Clone() + token, refreshedAuth, errToken := e.ensureAccessToken(ctx, updated) + if errToken != nil { + return nil, errToken + } + if refreshedAuth != nil { + updated = refreshedAuth + } + if antigravityProjectIDFromAuth(updated) != "" { + return updated, nil + } + + projectID, errProject := e.fetchAntigravityProjectID(ctx, updated, token) + if errProject != nil { + return nil, missingAntigravityProjectIDError(errProject) + } + if projectID == "" { + return nil, missingAntigravityProjectIDError(nil) + } + if updated.Metadata == nil { + updated.Metadata = make(map[string]any) + } + updated.Metadata["project_id"] = projectID + return updated, nil +} + +func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *cliproxyauth.Auth) (string, *cliproxyauth.Auth, error) { + if auth == nil { + return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} + } + accessToken := metaStringValue(auth.Metadata, "access_token") + expiry := tokenExpiry(auth.Metadata) + if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) { + e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken) + return accessToken, nil, nil + } + refreshCtx := context.Background() + if ctx != nil { + if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { + refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) + } + } + if refreshed, handled, err := helps.RefreshAuthViaHome(refreshCtx, e.cfg, auth); handled { + if err != nil { + return "", nil, err + } + token := metaStringValue(refreshed.Metadata, "access_token") + if strings.TrimSpace(token) == "" { + return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} + } + e.maybeRefreshAntigravityCreditsHint(ctx, refreshed, token) + return token, refreshed, nil + } + + updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone()) + if errRefresh != nil { + return "", nil, errRefresh + } + return metaStringValue(updated.Metadata, "access_token"), updated, nil +} + +func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { + if auth == nil { + return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} + } + refreshToken := metaStringValue(auth.Metadata, "refresh_token") + if refreshToken == "" { + return auth, statusErr{code: http.StatusUnauthorized, msg: "missing refresh token"} + } + if ctx == nil { + ctx = context.Background() + } + refreshToken = strings.TrimSpace(refreshToken) + + result, errRefresh, _ := antigravityRefreshGroup.Do(refreshToken, func() (interface{}, error) { + return e.refreshTokenSingleFlight(context.WithoutCancel(ctx), auth, refreshToken) + }) + if errRefresh != nil { + return auth, errRefresh + } + tokenResp, ok := result.(*antigravityTokenRefreshData) + if !ok || tokenResp == nil { + return auth, fmt.Errorf("antigravity token refresh failed: invalid single-flight result") + } + + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["access_token"] = tokenResp.AccessToken + if tokenResp.RefreshToken != "" { + auth.Metadata["refresh_token"] = tokenResp.RefreshToken + } + auth.Metadata["expires_in"] = tokenResp.ExpiresIn + now := time.Now() + auth.Metadata["timestamp"] = now.UnixMilli() + auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339) + auth.Metadata["type"] = antigravityAuthType + if errProject := e.ensureAntigravityProjectID(ctx, auth, tokenResp.AccessToken); errProject != nil { + log.Warnf("antigravity executor: ensure project id failed: %v", errProject) + } + e.updateAntigravityCreditsBalance(ctx, auth, tokenResp.AccessToken) + return auth, nil +} + +func (e *AntigravityExecutor) refreshTokenSingleFlight(ctx context.Context, auth *cliproxyauth.Auth, refreshToken string) (*antigravityTokenRefreshData, error) { + form := url.Values{} + form.Set("client_id", antigravityClientID) + form.Set("client_secret", antigravityClientSecret) + form.Set("grant_type", "refresh_token") + form.Set("refresh_token", refreshToken) + + httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode())) + if errReq != nil { + return nil, errReq + } + httpReq.Header.Set("Host", "oauth2.googleapis.com") + httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + // Real Antigravity uses Go's default User-Agent for OAuth token refresh + httpReq.Header.Set("User-Agent", "Go-http-client/2.0") + + httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + return nil, errDo + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + }() + + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errRead != nil { + return nil, errRead + } + + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} + if httpResp.StatusCode == http.StatusTooManyRequests { + if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { + sErr.retryAfter = retryAfter + } + } + return nil, sErr + } + + var tokenResp antigravityTokenRefreshData + if errUnmarshal := json.Unmarshal(bodyBytes, &tokenResp); errUnmarshal != nil { + return nil, errUnmarshal + } + + return &tokenResp, nil +} + +func (e *AntigravityExecutor) ensureAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) error { + if auth == nil { + return nil + } + + if antigravityProjectIDFromAuth(auth) != "" { + return nil + } + + projectID, errFetch := e.fetchAntigravityProjectID(ctx, auth, accessToken) + if errFetch != nil { + return errFetch + } + if projectID == "" { + return nil + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["project_id"] = projectID + + return nil +} + +func (e *AntigravityExecutor) fetchAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) (string, error) { + token := strings.TrimSpace(accessToken) + if token == "" { + token = metaStringValue(auth.Metadata, "access_token") + } + if token == "" { + return "", nil + } + + httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) + projectID, errFetch := sdkAuth.FetchAntigravityProjectID(ctx, token, httpClient) + if errFetch != nil { + return "", errFetch + } + return strings.TrimSpace(projectID), nil +} + +func (e *AntigravityExecutor) projectIDForRequest(_ context.Context, auth *cliproxyauth.Auth, _ string) (string, error) { + if projectID := antigravityProjectIDFromAuth(auth); projectID != "" { + return projectID, nil + } + return "", missingAntigravityProjectIDError(nil) +} + +func antigravityProjectIDFromAuth(auth *cliproxyauth.Auth) string { + if auth == nil || auth.Metadata == nil { + return "" + } + if pid, ok := auth.Metadata["project_id"].(string); ok { + return strings.TrimSpace(pid) + } + return "" +} + +func missingAntigravityProjectIDError(cause error) statusErr { + msg := "antigravity auth missing project_id" + if cause != nil { + msg = fmt.Sprintf("%s: %v", msg, cause) + } + return statusErr{code: http.StatusBadRequest, msg: msg} +} + +func tokenExpiry(metadata map[string]any) time.Time { + if metadata == nil { + return time.Time{} + } + if expStr, ok := metadata["expired"].(string); ok { + expStr = strings.TrimSpace(expStr) + if expStr != "" { + if parsed, errParse := time.Parse(time.RFC3339, expStr); errParse == nil { + return parsed + } + } + } + expiresIn, hasExpires := int64Value(metadata["expires_in"]) + tsMs, hasTimestamp := int64Value(metadata["timestamp"]) + if hasExpires && hasTimestamp { + return time.Unix(0, tsMs*int64(time.Millisecond)).Add(time.Duration(expiresIn) * time.Second) + } + return time.Time{} +} + +func metaStringValue(metadata map[string]any, key string) string { + if metadata == nil { + return "" + } + if v, ok := metadata[key]; ok { + switch typed := v.(type) { + case string: + return strings.TrimSpace(typed) + case []byte: + return strings.TrimSpace(string(typed)) + } + } + return "" +} + +func int64Value(value any) (int64, bool) { + switch typed := value.(type) { + case int: + return int64(typed), true + case int64: + return typed, true + case float64: + return int64(typed), true + case json.Number: + if i, errParse := typed.Int64(); errParse == nil { + return i, true + } + case string: + if strings.TrimSpace(typed) == "" { + return 0, false + } + if i, errParse := strconv.ParseInt(strings.TrimSpace(typed), 10, 64); errParse == nil { + return i, true + } + } + return 0, false +} diff --git a/internal/runtime/executor/antigravity_executor_credits.go b/internal/runtime/executor/antigravity_executor_credits.go new file mode 100644 index 000000000..0eff72c71 --- /dev/null +++ b/internal/runtime/executor/antigravity_executor_credits.go @@ -0,0 +1,795 @@ +package executor + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "math/rand" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + "golang.org/x/sync/singleflight" +) + +type antigravity429Category string + +type antigravityCreditsFailureState struct { + PermanentlyDisabled bool + ExplicitBalanceExhausted bool +} + +type antigravity429DecisionKind string + +const ( + antigravity429Unknown antigravity429Category = "unknown" + antigravity429RateLimited antigravity429Category = "rate_limited" + antigravity429QuotaExhausted antigravity429Category = "quota_exhausted" + antigravity429SoftRateLimit antigravity429Category = "soft_rate_limit" + antigravity429DecisionSoftRetry antigravity429DecisionKind = "soft_retry" + antigravity429DecisionInstantRetrySameAuth antigravity429DecisionKind = "instant_retry_same_auth" + antigravity429DecisionShortCooldownSwitchAuth antigravity429DecisionKind = "short_cooldown_switch_auth" + antigravity429DecisionFullQuotaExhausted antigravity429DecisionKind = "full_quota_exhausted" +) + +type antigravity429Decision struct { + kind antigravity429DecisionKind + retryAfter *time.Duration + reason string +} + +var ( + randSource = rand.New(rand.NewSource(time.Now().UnixNano())) + randSourceMutex sync.Mutex + antigravityCreditsFailureByAuth sync.Map + antigravityShortCooldownByAuth sync.Map + antigravityCreditsBalanceByAuth sync.Map // auth.ID → antigravityCreditsBalance + antigravityCreditsHintRefreshByID sync.Map // auth.ID → *antigravityCreditsHintRefreshState + antigravityRefreshGroup singleflight.Group + antigravityQuotaExhaustedKeywords = []string{ + "quota_exhausted", + "quota exhausted", + } +) + +type antigravityKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) + KVDel(ctx context.Context, keys ...string) (int64, error) +} + +var currentAntigravityKVClient = func() (antigravityKVClient, bool, error) { + return homekv.CurrentKVClient() +} + +type antigravityCreditsBalance struct { + CreditAmount float64 + MinCreditAmount float64 + PaidTierID string + Known bool +} + +type antigravityCreditsHintRefreshState struct { + mu sync.Mutex + lastAttempt time.Time +} + +type antigravityTokenRefreshData struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + ExpiresIn int64 `json:"expires_in"` + TokenType string `json:"token_type"` +} + +func antigravityAuthHasCredits(auth *cliproxyauth.Auth) bool { + ok, err := antigravityAuthHasCreditsRequired(context.Background(), auth) + if err != nil { + log.Errorf("antigravity executor: home kv credits check error: %v", err) + return false + } + return ok +} + +func antigravityAuthHasCreditsRequired(ctx context.Context, auth *cliproxyauth.Auth) (bool, error) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return false, nil + } + authID := strings.TrimSpace(auth.ID) + if hint, ok, errHint := cliproxyauth.GetAntigravityCreditsHintRequired(ctx, authID); errHint != nil { + return false, errHint + } else if ok && hint.Known { + return hint.Available, nil + } + + client, homeMode, errClient := currentAntigravityKVClient() + if homeMode { + if errClient != nil { + return false, errClient + } + raw, found, errBalance := client.KVGet(ctx, antigravityCreditsBalanceKey(authID)) + if errBalance != nil { + return false, errBalance + } + if !found { + return true, nil + } + var homeBalance antigravityCreditsBalance + if errUnmarshal := json.Unmarshal(raw, &homeBalance); errUnmarshal != nil { + return false, errUnmarshal + } + return antigravityCreditsBalanceAvailable(authID, homeBalance), nil + } + + val, ok := antigravityCreditsBalanceByAuth.Load(authID) + if !ok { + return true, nil // optimistic: assume credits available when balance unknown + } + bal, valid := val.(antigravityCreditsBalance) + if !valid { + antigravityCreditsBalanceByAuth.Delete(authID) + return false, nil + } + return antigravityCreditsBalanceAvailable(authID, bal), nil +} + +func antigravityCreditsBalanceAvailable(authID string, bal antigravityCreditsBalance) bool { + if !bal.Known { + return false + } + available := bal.CreditAmount >= bal.MinCreditAmount + cliproxyauth.SetAntigravityCreditsHint(strings.TrimSpace(authID), cliproxyauth.AntigravityCreditsHint{ + Known: true, + Available: available, + CreditAmount: bal.CreditAmount, + MinCreditAmount: bal.MinCreditAmount, + PaidTierID: bal.PaidTierID, + UpdatedAt: time.Now(), + }) + return available +} + +// parseMetaFloat extracts a float64 from auth.Metadata (handles string and numeric types). +func parseMetaFloat(metadata map[string]any, key string) (float64, bool) { + v, ok := metadata[key] + if !ok { + return 0, false + } + switch typed := v.(type) { + case float64: + return typed, true + case int: + return float64(typed), true + case int64: + return float64(typed), true + case uint64: + return float64(typed), true + case json.Number: + if f, err := typed.Float64(); err == nil { + return f, true + } + case string: + if f, err := strconv.ParseFloat(strings.TrimSpace(typed), 64); err == nil { + return f, true + } + } + return 0, false +} +func injectEnabledCreditTypes(payload []byte) []byte { + if len(payload) == 0 { + return nil + } + if !gjson.ValidBytes(payload) { + return nil + } + updated, err := sjson.SetRawBytes(payload, "enabledCreditTypes", []byte(`["GOOGLE_ONE_AI"]`)) + if err != nil { + return nil + } + return updated +} + +func classifyAntigravity429(body []byte) antigravity429Category { + switch decideAntigravity429(body).kind { + case antigravity429DecisionInstantRetrySameAuth, antigravity429DecisionShortCooldownSwitchAuth: + return antigravity429RateLimited + case antigravity429DecisionFullQuotaExhausted: + return antigravity429QuotaExhausted + case antigravity429DecisionSoftRetry: + return antigravity429SoftRateLimit + default: + return antigravity429Unknown + } +} + +func decideAntigravity429(body []byte) antigravity429Decision { + decision := antigravity429Decision{kind: antigravity429DecisionSoftRetry} + if len(body) == 0 { + return decision + } + + if retryAfter, parseErr := helps.ParseRetryDelay(body); parseErr == nil && retryAfter != nil { + decision.retryAfter = retryAfter + } + + status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String()) + if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") { + return decision + } + + details := gjson.GetBytes(body, "error.details") + if details.Exists() && details.IsArray() { + for _, detail := range details.Array() { + if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" { + continue + } + reason := strings.TrimSpace(detail.Get("reason").String()) + decision.reason = reason + switch { + case strings.EqualFold(reason, "QUOTA_EXHAUSTED"): + decision.kind = antigravity429DecisionFullQuotaExhausted + return decision + case strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED"): + if decision.retryAfter == nil { + decision.kind = antigravity429DecisionSoftRetry + return decision + } + switch { + case *decision.retryAfter < antigravityInstantRetryThreshold: + decision.kind = antigravity429DecisionInstantRetrySameAuth + case *decision.retryAfter < antigravityShortQuotaCooldownThreshold: + decision.kind = antigravity429DecisionShortCooldownSwitchAuth + default: + decision.kind = antigravity429DecisionFullQuotaExhausted + } + return decision + } + } + } + + lowerBody := strings.ToLower(string(body)) + for _, keyword := range antigravityQuotaExhaustedKeywords { + if strings.Contains(lowerBody, keyword) { + decision.kind = antigravity429DecisionFullQuotaExhausted + decision.reason = "quota_exhausted" + return decision + } + } + + decision.kind = antigravity429DecisionSoftRetry + return decision +} + +func antigravityCreditsRetryEnabled(cfg *config.Config) bool { + return cfg != nil && cfg.QuotaExceeded.AntigravityCredits +} + +func clearAntigravityCreditsFailureState(auth *cliproxyauth.Auth) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return + } + antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID)) +} +func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return + } + authID := strings.TrimSpace(auth.ID) + state := antigravityCreditsFailureState{ + PermanentlyDisabled: true, + ExplicitBalanceExhausted: true, + } + antigravityCreditsFailureByAuth.Store(authID, state) + bal := antigravityCreditsBalance{ + CreditAmount: 0, + MinCreditAmount: 1, + Known: true, + } + storeAntigravityCreditsBalanceBestEffort(authID, bal) + cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ + Known: true, + Available: false, + CreditAmount: 0, + MinCreditAmount: 1, + UpdatedAt: time.Now(), + }) +} + +func clearAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return + } + antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID)) +} + +func antigravityHasExplicitCreditsBalanceExhaustedReason(body []byte) bool { + if len(body) == 0 { + return false + } + details := gjson.GetBytes(body, "error.details") + if !details.Exists() || !details.IsArray() { + return false + } + for _, detail := range details.Array() { + if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" { + continue + } + reason := strings.TrimSpace(detail.Get("reason").String()) + if strings.EqualFold(reason, "INSUFFICIENT_G1_CREDITS_BALANCE") { + return true + } + } + return false +} + +func newAntigravityStatusErr(statusCode int, body []byte) statusErr { + err := statusErr{code: statusCode, msg: string(body)} + if statusCode == http.StatusTooManyRequests { + if retryAfter, parseErr := helps.ParseRetryDelay(body); parseErr == nil && retryAfter != nil { + err.retryAfter = retryAfter + } + } + return err +} +func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) { + if e == nil || auth == nil || !antigravityCreditsRetryEnabled(e.cfg) { + return + } + if ctx != nil && ctx.Err() != nil { + return + } + authID := strings.TrimSpace(auth.ID) + if authID == "" { + return + } + if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && hint.Known { + return + } + if strings.TrimSpace(accessToken) == "" { + accessToken = metaStringValue(auth.Metadata, "access_token") + } + if strings.TrimSpace(accessToken) == "" { + return + } + + if client, homeMode, errClient := currentAntigravityKVClient(); homeMode { + if errClient != nil { + log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errClient) + return + } + written, errSetNX := client.KVSetNX(context.Background(), antigravityCreditsRefreshLockKey(authID), []byte("1"), antigravityCreditsHintRefreshInterval) + if errSetNX != nil { + log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errSetNX) + return + } + if !written { + return + } + refreshCtx := context.Background() + if ctx != nil { + if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { + refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) + } + } + refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout) + authCopy := auth.Clone() + go func(auth *cliproxyauth.Auth, token string) { + defer cancel() + e.updateAntigravityCreditsBalance(refreshCtx, auth, token) + }(authCopy, accessToken) + return + } + + state := &antigravityCreditsHintRefreshState{} + if existing, loaded := antigravityCreditsHintRefreshByID.LoadOrStore(authID, state); loaded { + if cast, ok := existing.(*antigravityCreditsHintRefreshState); ok && cast != nil { + state = cast + } else { + antigravityCreditsHintRefreshByID.Delete(authID) + antigravityCreditsHintRefreshByID.Store(authID, state) + } + } + + now := time.Now() + if !state.mu.TryLock() { + return + } + if !state.lastAttempt.IsZero() && now.Sub(state.lastAttempt) < antigravityCreditsHintRefreshInterval { + state.mu.Unlock() + return + } + state.lastAttempt = now + + refreshCtx := context.Background() + if ctx != nil { + if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { + refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) + } + } + refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout) + authCopy := auth.Clone() + + go func(state *antigravityCreditsHintRefreshState, auth *cliproxyauth.Auth, token string) { + defer cancel() + defer state.mu.Unlock() + e.updateAntigravityCreditsBalance(refreshCtx, auth, token) + }(state, authCopy, accessToken) +} + +func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return + } + token := strings.TrimSpace(accessToken) + if token == "" { + token = metaStringValue(auth.Metadata, "access_token") + } + if token == "" { + return + } + + userAgent := resolveUserAgent(auth) + loadReqBody, errMarshal := json.Marshal(map[string]any{ + "metadata": map[string]string{ + "ideType": "ANTIGRAVITY", + }, + }) + if errMarshal != nil { + log.Debugf("antigravity executor: marshal loadCodeAssist request error: %v", errMarshal) + return + } + baseURL := antigravityLoadCodeAssistBaseURL(auth) + endpointURL := strings.TrimSuffix(baseURL, "/") + "/v1internal:loadCodeAssist" + httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(loadReqBody)) + if errReq != nil { + log.Debugf("antigravity executor: create loadCodeAssist request error: %v", errReq) + return + } + httpReq.Header.Set("Authorization", "Bearer "+token) + httpReq.Header.Set("Accept", "*/*") + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("User-Agent", userAgent) + + httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + log.Debugf("antigravity executor: loadCodeAssist request error: %v", errDo) + return + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close loadCodeAssist response body error: %v", errClose) + } + }() + + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errRead != nil || httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + log.Debugf("antigravity executor: loadCodeAssist returned status %d, err=%v", httpResp.StatusCode, errRead) + return + } + + authID := strings.TrimSpace(auth.ID) + paidTierID := strings.TrimSpace(gjson.GetBytes(bodyBytes, "paidTier.id").String()) + + credits := gjson.GetBytes(bodyBytes, "paidTier.availableCredits") + if !credits.IsArray() { + cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ + Known: true, + Available: false, + PaidTierID: paidTierID, + UpdatedAt: time.Now(), + }) + return + } + for _, credit := range credits.Array() { + if !strings.EqualFold(credit.Get("creditType").String(), "GOOGLE_ONE_AI") { + continue + } + creditAmount, errCA := strconv.ParseFloat(strings.TrimSpace(credit.Get("creditAmount").String()), 64) + if errCA != nil { + continue + } + minAmount, errMA := strconv.ParseFloat(strings.TrimSpace(credit.Get("minimumCreditAmountForUsage").String()), 64) + if errMA != nil { + continue + } + bal := antigravityCreditsBalance{ + CreditAmount: creditAmount, + MinCreditAmount: minAmount, + PaidTierID: paidTierID, + Known: true, + } + storeAntigravityCreditsBalanceBestEffort(authID, bal) + cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{ + Known: true, + Available: creditAmount >= minAmount, + CreditAmount: creditAmount, + MinCreditAmount: minAmount, + PaidTierID: paidTierID, + UpdatedAt: time.Now(), + }) + if creditAmount >= minAmount { + clearAntigravityCreditsPermanentlyDisabled(auth) + } + return + } +} +func antigravityRetryAttempts(auth *cliproxyauth.Auth, cfg *config.Config) int { + retry := 0 + if cfg != nil { + retry = cfg.RequestRetry + } + if auth != nil { + if override, ok := auth.RequestRetryOverride(); ok { + retry = override + } + } + if retry < 0 { + retry = 0 + } + attempts := retry + 1 + if attempts < 1 { + return 1 + } + return attempts +} + +func antigravityShouldRetryNoCapacity(statusCode int, body []byte) bool { + if statusCode != http.StatusServiceUnavailable { + return false + } + if len(body) == 0 { + return false + } + msg := strings.ToLower(string(body)) + return strings.Contains(msg, "no capacity available") +} + +func antigravityShouldRetryTransientResourceExhausted429(statusCode int, body []byte) bool { + if statusCode != http.StatusTooManyRequests { + return false + } + if len(body) == 0 { + return false + } + if classifyAntigravity429(body) != antigravity429Unknown { + return false + } + status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String()) + if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") { + return false + } + msg := strings.ToLower(string(body)) + return strings.Contains(msg, "resource has been exhausted") +} + +func antigravityShouldRetrySoftRateLimit(statusCode int, body []byte) bool { + if statusCode != http.StatusTooManyRequests { + return false + } + return decideAntigravity429(body).kind == antigravity429DecisionSoftRetry +} + +func antigravityShouldBypassShortCooldown(ctx context.Context, cfg *config.Config) bool { + return cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(cfg) +} + +func antigravitySoftRateLimitDelay(attempt int) time.Duration { + if attempt < 0 { + attempt = 0 + } + base := time.Duration(attempt+1) * 500 * time.Millisecond + if base > 3*time.Second { + base = 3 * time.Second + } + return base +} + +func antigravityShortCooldownKey(auth *cliproxyauth.Auth, modelName string) string { + if auth == nil { + return "" + } + authID := strings.TrimSpace(auth.ID) + modelName = strings.TrimSpace(modelName) + if authID == "" || modelName == "" { + return "" + } + return authID + "|" + modelName + "|sc" +} + +func antigravityCreditsBalanceKey(authID string) string { + return "cpa:antigravity:credits-balance:" + strings.TrimSpace(authID) +} + +func antigravityCreditsRefreshLockKey(authID string) string { + return "cpa:antigravity:credits-refresh-lock:" + strings.TrimSpace(authID) +} + +func antigravityShortCooldownKVKey(auth *cliproxyauth.Auth, modelName string) string { + if auth == nil { + return "" + } + authID := strings.TrimSpace(auth.ID) + modelName = strings.TrimSpace(modelName) + if authID == "" || modelName == "" { + return "" + } + return "cpa:antigravity:short-cooldown:" + authID + ":" + homekv.HashKeyPart(modelName) +} + +func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration) { + inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(context.Background(), auth, modelName, now) + if errCooldown != nil { + log.Errorf("antigravity executor: home kv cooldown read error: %v", errCooldown) + return false, 0 + } + return inCooldown, remaining +} + +func antigravityIsInShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration, error) { + kvKey := antigravityShortCooldownKVKey(auth, modelName) + client, homeMode, errClient := currentAntigravityKVClient() + if homeMode { + if errClient != nil { + return false, 0, errClient + } + if kvKey == "" { + return false, 0, nil + } + raw, found, errGet := client.KVGet(ctx, kvKey) + if errGet != nil || !found { + return false, 0, errGet + } + untilNano, errParse := strconv.ParseInt(strings.TrimSpace(string(raw)), 10, 64) + if errParse != nil { + return false, 0, errParse + } + remaining := time.Unix(0, untilNano).Sub(now) + if remaining <= 0 { + if _, errDel := client.KVDel(ctx, kvKey); errDel != nil { + return false, 0, errDel + } + return false, 0, nil + } + return true, remaining, nil + } + + key := antigravityShortCooldownKey(auth, modelName) + if key == "" { + return false, 0, nil + } + value, ok := antigravityShortCooldownByAuth.Load(key) + if !ok { + return false, 0, nil + } + until, ok := value.(time.Time) + if !ok || until.IsZero() { + antigravityShortCooldownByAuth.Delete(key) + return false, 0, nil + } + remaining := until.Sub(now) + if remaining <= 0 { + antigravityShortCooldownByAuth.Delete(key) + return false, 0, nil + } + return true, remaining, nil +} + +func markAntigravityShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) { + if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, modelName, now, duration); errMark != nil { + log.Errorf("antigravity executor: home kv cooldown write error: %v", errMark) + } +} + +func markAntigravityShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) error { + kvKey := antigravityShortCooldownKVKey(auth, modelName) + client, homeMode, errClient := currentAntigravityKVClient() + if homeMode { + if errClient != nil { + return errClient + } + if kvKey == "" || duration <= 0 { + return nil + } + until := now.Add(duration) + written, errSet := client.KVSet(ctx, kvKey, []byte(strconv.FormatInt(until.UnixNano(), 10)), homekv.KVSetOptions{EX: duration + 5*time.Second}) + if errSet != nil { + return errSet + } + if !written { + return fmt.Errorf("home kv store unavailable") + } + return nil + } + + key := antigravityShortCooldownKey(auth, modelName) + if key == "" { + return nil + } + antigravityShortCooldownByAuth.Store(key, now.Add(duration)) + return nil +} + +func storeAntigravityCreditsBalanceBestEffort(authID string, bal antigravityCreditsBalance) { + authID = strings.TrimSpace(authID) + if authID == "" { + return + } + if client, homeMode, errClient := currentAntigravityKVClient(); homeMode { + if errClient != nil { + log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errClient) + return + } + raw, errMarshal := json.Marshal(bal) + if errMarshal != nil { + log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errMarshal) + return + } + if _, errSet := client.KVSet(context.Background(), antigravityCreditsBalanceKey(authID), raw, homekv.KVSetOptions{EX: 30 * time.Minute}); errSet != nil { + log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errSet) + } + return + } + antigravityCreditsBalanceByAuth.Store(authID, bal) +} + +func homeKVUnavailableStatusErr(cause error) statusErr { + if cause == nil { + return statusErr{code: http.StatusServiceUnavailable, msg: "home kv store unavailable"} + } + return statusErr{code: http.StatusServiceUnavailable, msg: fmt.Sprintf("home kv store unavailable: %v", cause)} +} + +func antigravityNoCapacityRetryDelay(attempt int) time.Duration { + if attempt < 0 { + attempt = 0 + } + delay := time.Duration(attempt+1) * 250 * time.Millisecond + if delay > 2*time.Second { + delay = 2 * time.Second + } + return delay +} + +func antigravityTransient429RetryDelay(attempt int) time.Duration { + if attempt < 0 { + attempt = 0 + } + delay := time.Duration(attempt+1) * 100 * time.Millisecond + if delay > 500*time.Millisecond { + delay = 500 * time.Millisecond + } + return delay +} + +func antigravityInstantRetryDelay(wait time.Duration) time.Duration { + if wait <= 0 { + return 0 + } + return wait + 800*time.Millisecond +} + +func antigravityWait(ctx context.Context, wait time.Duration) error { + if wait <= 0 { + return nil + } + timer := time.NewTimer(wait) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} diff --git a/internal/runtime/executor/antigravity_executor_execute.go b/internal/runtime/executor/antigravity_executor_execute.go new file mode 100644 index 000000000..415914ae6 --- /dev/null +++ b/internal/runtime/executor/antigravity_executor_execute.go @@ -0,0 +1,738 @@ +package executor + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +// Execute performs a non-streaming request to the Antigravity API. +func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + if opts.Alt == "responses/compact" { + return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { + return resp, homeKVUnavailableStatusErr(errCooldown) + } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { + log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) + d := remaining + return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} + } + + isClaude := strings.Contains(strings.ToLower(baseModel), "claude") + if isClaude || strings.Contains(baseModel, "gemini-3-pro") || strings.Contains(baseModel, "gemini-3.1-flash-image") { + return e.executeClaudeNonStream(ctx, auth, req, opts) + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("antigravity") + + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) + if errValidate != nil { + return resp, errValidate + } + req.Payload = originalPayload + token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) + if errToken != nil { + return resp, errToken + } + if updatedAuth != nil { + auth = updatedAuth + } + originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, false) + translated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, false) + + translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return resp, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers) + translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated) + reporter.SetTranslatedReasoningEffort(translated, to.String()) + + useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) + + baseURLs := antigravityBaseURLFallbackOrder(auth) + httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + attempts := antigravityRetryAttempts(auth, e.cfg) + +attemptLoop: + for attempt := 0; attempt < attempts; attempt++ { + var lastStatus int + var lastBody []byte + var lastErr error + + for idx, baseURL := range baseURLs { + requestPayload := translated + if useCredits { + if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { + requestPayload = cp + helps.MarkCreditsUsed(ctx) + } + } + replayScope := antigravityReasoningReplayScope{} + if antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) + if errReplay != nil { + err = errReplay + return resp, err + } + } + + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) + if errReq != nil { + err = errReq + return resp, err + } + + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { + return resp, errDo + } + lastStatus = 0 + lastBody = nil + lastErr = errDo + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + err = errDo + return resp, err + } + + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + err = errRead + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + + if httpResp.StatusCode == http.StatusTooManyRequests { + decision := decideAntigravity429(bodyBytes) + switch decision.kind { + case antigravity429DecisionInstantRetrySameAuth: + if attempt+1 < attempts { + if decision.retryAfter != nil && *decision.retryAfter > 0 { + wait := antigravityInstantRetryDelay(*decision.retryAfter) + log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) + if errWait := antigravityWait(ctx, wait); errWait != nil { + return resp, errWait + } + } + continue attemptLoop + } + case antigravity429DecisionShortCooldownSwitchAuth: + if decision.retryAfter != nil && *decision.retryAfter > 0 { + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return resp, err + } + log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) + } + case antigravity429DecisionFullQuotaExhausted: + if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { + markAntigravityCreditsPermanentlyDisabled(auth) + } + // No credits logic - just fall through to error return below + } + } + + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), bodyBytes)) + lastStatus = httpResp.StatusCode + lastBody = append([]byte(nil), bodyBytes...) + lastErr = nil + if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { + delay := antigravityTransient429RetryDelay(attempt) + log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return resp, errWait + } + continue attemptLoop + } + if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + if attempt+1 < attempts { + delay := antigravityNoCapacityRetryDelay(attempt) + log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return resp, errWait + } + continue attemptLoop + } + } + if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { + if attempt+1 < attempts { + delay := antigravitySoftRateLimitDelay(attempt) + log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return resp, errWait + } + continue attemptLoop + } + } + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { + err = errClear + return resp, err + } + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) + return resp, err + } + + // Success + if useCredits { + clearAntigravityCreditsFailureState(auth) + } + cacheAntigravityReasoningReplayFromResponse(ctx, replayScope, requestPayload, bodyBytes) + bodyBytes = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, bodyBytes) + reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes)) + var param any + converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bodyBytes, ¶m) + resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} + reporter.EnsurePublished(ctx) + return resp, nil + } + + switch { + case lastStatus != 0: + err = newAntigravityStatusErr(lastStatus, lastBody) + case lastErr != nil: + err = lastErr + default: + err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + } + return resp, err + } + + return resp, err +} + +// executeClaudeNonStream performs a claude non-streaming request to the Antigravity API. +func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + baseModel := thinking.ParseSuffix(req.Model).ModelName + if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { + return resp, homeKVUnavailableStatusErr(errCooldown) + } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { + log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) + d := remaining + return resp, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("antigravity") + + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) + if errValidate != nil { + return resp, errValidate + } + req.Payload = originalPayload + token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) + if errToken != nil { + return resp, errToken + } + if updatedAuth != nil { + auth = updatedAuth + } + originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, true) + translated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, true) + + translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return resp, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers) + translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated) + reporter.SetTranslatedReasoningEffort(translated, to.String()) + + useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) + + baseURLs := antigravityBaseURLFallbackOrder(auth) + httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + + attempts := antigravityRetryAttempts(auth, e.cfg) + +attemptLoop: + for attempt := 0; attempt < attempts; attempt++ { + var lastStatus int + var lastBody []byte + var lastErr error + + for idx, baseURL := range baseURLs { + requestPayload := translated + if useCredits { + if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { + requestPayload = cp + helps.MarkCreditsUsed(ctx) + } + } + replayScope := antigravityReasoningReplayScope{} + if antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) + if errReplay != nil { + err = errReplay + return resp, err + } + } + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) + if errReq != nil { + err = errReq + return resp, err + } + + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { + return resp, errDo + } + lastStatus = 0 + lastBody = nil + lastErr = errDo + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + err = errDo + return resp, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { + err = errRead + return resp, err + } + if errCtx := ctx.Err(); errCtx != nil { + err = errCtx + return resp, err + } + lastStatus = 0 + lastBody = nil + lastErr = errRead + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + err = errRead + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if httpResp.StatusCode == http.StatusTooManyRequests { + decision := decideAntigravity429(bodyBytes) + + switch decision.kind { + case antigravity429DecisionInstantRetrySameAuth: + if attempt+1 < attempts { + if decision.retryAfter != nil && *decision.retryAfter > 0 { + wait := antigravityInstantRetryDelay(*decision.retryAfter) + log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) + if errWait := antigravityWait(ctx, wait); errWait != nil { + return resp, errWait + } + } + continue attemptLoop + } + case antigravity429DecisionShortCooldownSwitchAuth: + if decision.retryAfter != nil && *decision.retryAfter > 0 { + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return resp, err + } + log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel) + } + case antigravity429DecisionFullQuotaExhausted: + if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { + markAntigravityCreditsPermanentlyDisabled(auth) + } + // No credits logic - just fall through to error return below + } + } + + lastStatus = httpResp.StatusCode + lastBody = append([]byte(nil), bodyBytes...) + lastErr = nil + if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { + delay := antigravityTransient429RetryDelay(attempt) + log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return resp, errWait + } + continue attemptLoop + } + if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + if attempt+1 < attempts { + delay := antigravityNoCapacityRetryDelay(attempt) + log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return resp, errWait + } + continue attemptLoop + } + } + if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { + if attempt+1 < attempts { + delay := antigravitySoftRateLimitDelay(attempt) + log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return resp, errWait + } + continue attemptLoop + } + } + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { + err = errClear + return resp, err + } + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) + return resp, err + } + + // Stream success + if useCredits { + clearAntigravityCreditsFailureState(auth) + } + replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) + out := make(chan cliproxyexecutor.StreamChunk) + go func(resp *http.Response) { + defer close(out) + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + }() + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(nil, streamScannerBuffer) + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } + + // Filter usage metadata for all models + // Only retain usage statistics in the terminal chunk + line = helps.FilterSSEUsageMetadata(line) + + payload := helps.JSONPayload(line) + if payload == nil { + continue + } + + if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { + reporter.Publish(ctx, detail) + } + + out <- cliproxyexecutor.StreamChunk{Payload: payload} + } + if errScan := scanner.Err(); errScan != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + out <- cliproxyexecutor.StreamChunk{Err: errScan} + } else { + if replayAccumulator != nil { + replayAccumulator.Commit(ctx) + } + reporter.EnsurePublished(ctx) + } + }(httpResp) + + var buffer bytes.Buffer + for chunk := range out { + if chunk.Err != nil { + return resp, chunk.Err + } + if len(chunk.Payload) > 0 { + _, _ = buffer.Write(chunk.Payload) + _, _ = buffer.Write([]byte("\n")) + } + } + resp = cliproxyexecutor.Response{Payload: e.convertStreamToNonStream(buffer.Bytes())} + + resp.Payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, resp.Payload) + reporter.Publish(ctx, helps.ParseAntigravityUsage(resp.Payload)) + var param any + converted := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, resp.Payload, ¶m) + resp = cliproxyexecutor.Response{Payload: converted, Headers: httpResp.Header.Clone()} + reporter.EnsurePublished(ctx) + + return resp, nil + } + + switch { + case lastStatus != 0: + err = newAntigravityStatusErr(lastStatus, lastBody) + case lastErr != nil: + err = lastErr + default: + err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + } + return resp, err + } + + return resp, err +} + +func (e *AntigravityExecutor) convertStreamToNonStream(stream []byte) []byte { + responseTemplate := "" + var traceID string + var finishReason string + var modelVersion string + var responseID string + var role string + var usageRaw string + parts := make([]map[string]interface{}, 0) + var pendingKind string + var pendingText strings.Builder + var pendingThoughtSig string + + flushPending := func() { + if pendingKind == "" { + return + } + text := pendingText.String() + switch pendingKind { + case "text": + if strings.TrimSpace(text) == "" { + pendingKind = "" + pendingText.Reset() + pendingThoughtSig = "" + return + } + parts = append(parts, map[string]interface{}{"text": text}) + case "thought": + if strings.TrimSpace(text) == "" && pendingThoughtSig == "" { + pendingKind = "" + pendingText.Reset() + pendingThoughtSig = "" + return + } + part := map[string]interface{}{"thought": true} + part["text"] = text + if pendingThoughtSig != "" { + part["thoughtSignature"] = pendingThoughtSig + } + parts = append(parts, part) + } + pendingKind = "" + pendingText.Reset() + pendingThoughtSig = "" + } + + normalizePart := func(partResult gjson.Result) map[string]interface{} { + var m map[string]interface{} + _ = json.Unmarshal([]byte(partResult.Raw), &m) + if m == nil { + m = map[string]interface{}{} + } + sig := partResult.Get("thoughtSignature").String() + if sig == "" { + sig = partResult.Get("thought_signature").String() + } + if sig != "" { + m["thoughtSignature"] = sig + delete(m, "thought_signature") + } + if inlineData, ok := m["inline_data"]; ok { + m["inlineData"] = inlineData + delete(m, "inline_data") + } + return m + } + + for _, line := range bytes.Split(stream, []byte("\n")) { + trimmed := bytes.TrimSpace(line) + if len(trimmed) == 0 || !gjson.ValidBytes(trimmed) { + continue + } + + root := gjson.ParseBytes(trimmed) + responseNode := root.Get("response") + if !responseNode.Exists() { + if root.Get("candidates").Exists() { + responseNode = root + } else { + continue + } + } + responseTemplate = responseNode.Raw + + if traceResult := root.Get("traceId"); traceResult.Exists() && traceResult.String() != "" { + traceID = traceResult.String() + } + + if roleResult := responseNode.Get("candidates.0.content.role"); roleResult.Exists() { + role = roleResult.String() + } + + if finishResult := responseNode.Get("candidates.0.finishReason"); finishResult.Exists() && finishResult.String() != "" { + finishReason = finishResult.String() + } + + if modelResult := responseNode.Get("modelVersion"); modelResult.Exists() && modelResult.String() != "" { + modelVersion = modelResult.String() + } + if responseIDResult := responseNode.Get("responseId"); responseIDResult.Exists() && responseIDResult.String() != "" { + responseID = responseIDResult.String() + } + if usageResult := responseNode.Get("usageMetadata"); usageResult.Exists() { + usageRaw = usageResult.Raw + } else if usageMetadataResult := root.Get("usageMetadata"); usageMetadataResult.Exists() { + usageRaw = usageMetadataResult.Raw + } + + if partsResult := responseNode.Get("candidates.0.content.parts"); partsResult.IsArray() { + for _, part := range partsResult.Array() { + hasFunctionCall := part.Get("functionCall").Exists() + hasInlineData := part.Get("inlineData").Exists() || part.Get("inline_data").Exists() + sig := part.Get("thoughtSignature").String() + if sig == "" { + sig = part.Get("thought_signature").String() + } + text := part.Get("text").String() + thought := part.Get("thought").Bool() + + if hasFunctionCall || hasInlineData { + flushPending() + parts = append(parts, normalizePart(part)) + continue + } + + if thought || part.Get("text").Exists() { + kind := "text" + if thought { + kind = "thought" + } + if pendingKind != "" && pendingKind != kind { + flushPending() + } + pendingKind = kind + pendingText.WriteString(text) + if kind == "thought" && sig != "" { + pendingThoughtSig = sig + } + continue + } + + flushPending() + parts = append(parts, normalizePart(part)) + } + } + } + flushPending() + + if responseTemplate == "" { + responseTemplate = `{"candidates":[{"content":{"role":"model","parts":[]}}]}` + } + + partsJSON, _ := json.Marshal(parts) + updatedTemplate, _ := sjson.SetRawBytes([]byte(responseTemplate), "candidates.0.content.parts", partsJSON) + responseTemplate = string(updatedTemplate) + if role != "" { + updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.content.role", role) + responseTemplate = string(updatedTemplate) + } + if finishReason != "" { + updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "candidates.0.finishReason", finishReason) + responseTemplate = string(updatedTemplate) + } + if modelVersion != "" { + updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "modelVersion", modelVersion) + responseTemplate = string(updatedTemplate) + } + if responseID != "" { + updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "responseId", responseID) + responseTemplate = string(updatedTemplate) + } + if usageRaw != "" { + updatedTemplate, _ = sjson.SetRawBytes([]byte(responseTemplate), "usageMetadata", []byte(usageRaw)) + responseTemplate = string(updatedTemplate) + } else if !gjson.Get(responseTemplate, "usageMetadata").Exists() { + updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.promptTokenCount", 0) + responseTemplate = string(updatedTemplate) + updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.candidatesTokenCount", 0) + responseTemplate = string(updatedTemplate) + updatedTemplate, _ = sjson.SetBytes([]byte(responseTemplate), "usageMetadata.totalTokenCount", 0) + responseTemplate = string(updatedTemplate) + } + + output := `{"response":{},"traceId":""}` + updatedOutput, _ := sjson.SetRawBytes([]byte(output), "response", []byte(responseTemplate)) + output = string(updatedOutput) + if traceID != "" { + updatedOutput, _ = sjson.SetBytes([]byte(output), "traceId", traceID) + output = string(updatedOutput) + } + return []byte(output) +} diff --git a/internal/runtime/executor/antigravity_executor_request.go b/internal/runtime/executor/antigravity_executor_request.go new file mode 100644 index 000000000..ae0f51a42 --- /dev/null +++ b/internal/runtime/executor/antigravity_executor_request.go @@ -0,0 +1,449 @@ +package executor + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/binary" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (e *AntigravityExecutor) buildRequest(ctx context.Context, auth *cliproxyauth.Auth, token, modelName string, payload []byte, stream bool, alt, baseURL string, derivedSessionIDs ...string) (*http.Request, error) { + if token == "" { + return nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} + } + + base := strings.TrimSuffix(baseURL, "/") + if base == "" { + base = buildBaseURL(auth) + } + path := antigravityGeneratePath + if stream { + path = antigravityStreamPath + } + var requestURL strings.Builder + requestURL.WriteString(base) + requestURL.WriteString(path) + if stream { + if alt != "" { + requestURL.WriteString("?$alt=") + requestURL.WriteString(url.QueryEscape(alt)) + } else { + requestURL.WriteString("?alt=sse") + } + } else if alt != "" { + requestURL.WriteString("?$alt=") + requestURL.WriteString(url.QueryEscape(alt)) + } + + projectID, errProject := e.projectIDForRequest(ctx, auth, token) + if errProject != nil { + return nil, errProject + } + payload = geminiToAntigravity(modelName, payload, projectID, derivedSessionIDs...) + + // Cap maxOutputTokens to model's max_completion_tokens from registry + if maxOut := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxOut.Exists() && maxOut.Type == gjson.Number { + if modelInfo := registry.LookupModelInfo(modelName, "antigravity"); modelInfo != nil && modelInfo.MaxCompletionTokens > 0 { + if int(maxOut.Int()) > modelInfo.MaxCompletionTokens { + payload, _ = sjson.SetBytes(payload, "request.generationConfig.maxOutputTokens", modelInfo.MaxCompletionTokens) + } + } + } + + useAntigravitySchema := strings.Contains(modelName, "claude") || strings.Contains(modelName, "gemini-3-pro") || strings.Contains(modelName, "gemini-3.1-pro") + var ( + bodyReader io.Reader + payloadLog []byte + ) + if antigravityRequestNeedsSchemaSanitization(payload) { + payloadStr := sanitizeAntigravityRequestSchemas(string(payload), useAntigravitySchema) + + if strings.Contains(modelName, "claude") { + updated, _ := sjson.SetBytes([]byte(payloadStr), "request.toolConfig.functionCallingConfig.mode", "VALIDATED") + payloadStr = string(updated) + } else { + payloadStr, _ = sjson.Delete(payloadStr, "request.generationConfig.maxOutputTokens") + } + + payloadStrBytes := applyAntigravityNativeSignatureReplayIfNeeded(modelName, []byte(payloadStr)) + bodyReader = bytes.NewReader(payloadStrBytes) + if e.cfg != nil && e.cfg.RequestLog { + payloadLog = append([]byte(nil), payloadStrBytes...) + } + } else { + if strings.Contains(modelName, "claude") { + payload, _ = sjson.SetBytes(payload, "request.toolConfig.functionCallingConfig.mode", "VALIDATED") + } else { + payload, _ = sjson.DeleteBytes(payload, "request.generationConfig.maxOutputTokens") + } + + payload = applyAntigravityNativeSignatureReplayIfNeeded(modelName, payload) + bodyReader = bytes.NewReader(payload) + if e.cfg != nil && e.cfg.RequestLog { + payloadLog = append([]byte(nil), payload...) + } + } + + // if useAntigravitySchema { + // systemInstructionPartsResult := gjson.Get(payloadStr, "request.systemInstruction.parts") + // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.role", "user") + // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.0.text", systemInstruction) + // payloadStr, _ = sjson.SetBytes([]byte(payloadStr), "request.systemInstruction.parts.1.text", fmt.Sprintf("Please ignore following [ignore]%s[/ignore]", systemInstruction)) + + // if systemInstructionPartsResult.Exists() && systemInstructionPartsResult.IsArray() { + // for _, partResult := range systemInstructionPartsResult.Array() { + // payloadStr, _ = sjson.SetRawBytes([]byte(payloadStr), "request.systemInstruction.parts.-1", []byte(partResult.Raw)) + // } + // } + // } + + httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bodyReader) + if errReq != nil { + return nil, errReq + } + httpReq.Close = true + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Authorization", "Bearer "+token) + httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) + if host := resolveHost(base); host != "" { + httpReq.Host = host + } + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: requestURL.String(), + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: payloadLog, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + return httpReq, nil +} + +// sanitizeAntigravityRequestSchemas cleans the JSON schemas carried by an Antigravity request. +// +// Cleaning is applied only to the payload locations that actually hold a JSON schema. The schema +// cleaner rewrites keys such as "title", "format", "default" and "const", which are also ordinary +// data keys inside functionCall arguments replayed from conversation history. Running it over the +// whole document silently mutated that history, so tools lost required argument fields and the +// model imitated the corrupted examples on later turns. +func sanitizeAntigravityRequestSchemas(payloadStr string, useAntigravitySchema bool) string { + for _, base := range antigravityFunctionDeclarationPaths(payloadStr) { + oldPath := base + ".parametersJsonSchema" + if !gjson.Get(payloadStr, oldPath).Exists() { + continue + } + renamed, errRename := util.RenameKey(payloadStr, oldPath, base+".parameters") + if errRename != nil { + log.Debugf("antigravity: failed to rename %s: %v", oldPath, errRename) + continue + } + payloadStr = renamed + } + + clean := util.CleanJSONSchemaForGemini + if useAntigravitySchema { + clean = util.CleanJSONSchemaForAntigravity + } + + for _, schemaPath := range antigravitySchemaPaths(payloadStr) { + schema := gjson.Get(payloadStr, schemaPath) + if !schema.Exists() { + continue + } + updated, errSet := sjson.SetRawBytes([]byte(payloadStr), schemaPath, []byte(cleanNestedSchema(clean, schema.Raw))) + if errSet != nil { + log.Debugf("antigravity: failed to write cleaned schema at %s: %v", schemaPath, errSet) + continue + } + payloadStr = string(updated) + } + + return payloadStr +} + +// antigravitySchemaWrapperKey nests a schema during cleaning. It is never sent upstream. +const antigravitySchemaWrapperKey = "schema" + +// cleanNestedSchema cleans a schema with it nested one level down, then unwraps it. +// +// The cleaner deliberately skips placeholder insertion for a top-level schema, but Claude's +// VALIDATED mode needs every tool schema to declare at least one required property. Whole-payload +// cleaning always saw tool schemas nested inside the request, so nesting is reproduced here to keep +// the emitted schema byte-identical to the previous behaviour. +func cleanNestedSchema(clean func(string) string, schemaRaw string) string { + wrapped, errWrap := sjson.SetRaw("{}", antigravitySchemaWrapperKey, schemaRaw) + if errWrap != nil { + return clean(schemaRaw) + } + if unwrapped := gjson.Get(clean(wrapped), antigravitySchemaWrapperKey); unwrapped.Exists() { + return unwrapped.Raw + } + return clean(schemaRaw) +} + +// antigravityFunctionDeclarationPaths returns the path of every function declaration in the request. +// Both the camelCase and snake_case spellings are accepted because callers reach this executor +// through different translators. +func antigravityFunctionDeclarationPaths(payloadStr string) []string { + tools := gjson.Get(payloadStr, "request.tools") + if !tools.IsArray() { + return nil + } + paths := make([]string, 0, len(tools.Array())) + for i, tool := range tools.Array() { + for _, declKey := range []string{"functionDeclarations", "function_declarations"} { + decls := tool.Get(declKey) + if !decls.IsArray() { + continue + } + for j := range decls.Array() { + paths = append(paths, fmt.Sprintf("request.tools.%d.%s.%d", i, declKey, j)) + } + } + } + return paths +} + +// antigravitySchemaPaths returns every payload path that holds a JSON schema document. +// A function declaration may carry a schema for its parameters and for its result, so all of +// them must be cleaned; anything omitted here reaches the upstream API uncleaned. +func antigravitySchemaPaths(payloadStr string) []string { + paths := make([]string, 0, 12) + for _, base := range antigravityFunctionDeclarationPaths(payloadStr) { + for _, key := range antigravityDeclarationSchemaKeys { + if gjson.Get(payloadStr, base+"."+key).IsObject() { + paths = append(paths, base+"."+key) + } + } + } + for _, container := range antigravityGenerationConfigContainers { + for _, key := range antigravityGenerationSchemaKeys { + p := container + "." + key + if gjson.Get(payloadStr, p).IsObject() { + paths = append(paths, p) + } + } + } + return paths +} + +// The upstream API is proto-JSON and accepts either spelling, and the Gemini translator forwards +// whichever one the client sent. Both are therefore cleaned where they sit rather than renamed: +// renaming would alter the body the client asked for, and only the unsupported keywords inside a +// schema cause upstream errors. The one exception is parametersJsonSchema, renamed onto parameters +// above because whole-payload cleaning did the same. +var ( + antigravityDeclarationSchemaKeys = []string{ + "parameters", "parametersJsonSchema", "parameters_json_schema", + "response", "responseJsonSchema", "response_json_schema", + } + antigravityGenerationConfigContainers = []string{ + "request.generationConfig", "request.generation_config", + } + antigravityGenerationSchemaKeys = []string{ + "responseSchema", "responseJsonSchema", "response_schema", "response_json_schema", + } +) + +func antigravityRequestNeedsSchemaSanitization(payload []byte) bool { + if gjson.GetBytes(payload, "request.tools.0").Exists() { + return true + } + for _, container := range antigravityGenerationConfigContainers { + for _, key := range antigravityGenerationSchemaKeys { + if gjson.GetBytes(payload, container+"."+key).Exists() { + return true + } + } + } + return false +} +func buildBaseURL(auth *cliproxyauth.Auth) string { + if baseURLs := antigravityBaseURLFallbackOrder(auth); len(baseURLs) > 0 { + return baseURLs[0] + } + return antigravityBaseURLDaily +} + +func antigravityLoadCodeAssistBaseURL(auth *cliproxyauth.Auth) string { + if base := resolveCustomAntigravityBaseURL(auth); base != "" { + return base + } + return antigravityBaseURLProd +} + +func resolveHost(base string) string { + parsed, errParse := url.Parse(base) + if errParse != nil { + return "" + } + if parsed.Host != "" { + return parsed.Host + } + return strings.TrimPrefix(strings.TrimPrefix(base, "https://"), "http://") +} + +func resolveUserAgent(auth *cliproxyauth.Auth) string { + return misc.AntigravityRequestUserAgent(antigravityConfiguredUserAgent(auth)) +} + +func resolveLoadCodeAssistUserAgent(auth *cliproxyauth.Auth) string { + return misc.AntigravityLoadCodeAssistUserAgent(antigravityConfiguredUserAgent(auth)) +} + +func antigravityConfiguredUserAgent(auth *cliproxyauth.Auth) string { + raw := "" + if auth != nil { + if auth.Attributes != nil { + if ua := strings.TrimSpace(auth.Attributes["user_agent"]); ua != "" { + raw = ua + } + } + if raw == "" && auth.Metadata != nil { + if ua, ok := auth.Metadata["user_agent"].(string); ok && strings.TrimSpace(ua) != "" { + raw = strings.TrimSpace(ua) + } + } + } + return raw +} + +var antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string { + if base := resolveCustomAntigravityBaseURL(auth); base != "" { + return []string{base} + } + return []string{ + antigravityBaseURLDaily, + antigravityBaseURLProd, + // antigravitySandboxBaseURLDaily, + } +} + +func resolveCustomAntigravityBaseURL(auth *cliproxyauth.Auth) string { + if auth == nil { + return "" + } + if auth.Attributes != nil { + if v := strings.TrimSpace(auth.Attributes["base_url"]); v != "" { + return strings.TrimSuffix(v, "/") + } + } + if auth.Metadata != nil { + if v, ok := auth.Metadata["base_url"].(string); ok { + v = strings.TrimSpace(v) + if v != "" { + return strings.TrimSuffix(v, "/") + } + } + } + return "" +} + +func geminiToAntigravity(modelName string, payload []byte, projectID string, derivedSessionIDs ...string) []byte { + template := payload + template = helps.SetStringIfDifferent(template, "model", modelName) + template = helps.SetStringIfDifferent(template, "userAgent", "antigravity") + + isImageModel := strings.Contains(modelName, "image") + reqType := strings.TrimSpace(gjson.GetBytes(template, "requestType").String()) + if reqType == "" { + if isImageModel { + reqType = "image_gen" + } else { + reqType = "agent" + } + template, _ = sjson.SetBytes(template, "requestType", reqType) + } + + if projectID != "" { + template = helps.SetStringIfDifferent(template, "project", projectID) + } else { + template, _ = sjson.DeleteBytes(template, "project") + } + + if isImageModel { + template, _ = sjson.SetBytes(template, "requestId", generateImageGenRequestID()) + } else if reqType != "web_search" { + template, _ = sjson.SetBytes(template, "requestId", generateRequestID()) + sessionID := strings.TrimSpace(gjson.GetBytes(template, "request.sessionId").String()) + if sessionID == "" && len(derivedSessionIDs) > 0 { + sessionID = strings.TrimSpace(derivedSessionIDs[0]) + } + if sessionID == "" { + sessionID = generateStableSessionID(payload) + } + template, _ = sjson.SetBytes(template, "request.sessionId", sessionID) + } + + template, _ = sjson.DeleteBytes(template, "request.safetySettings") + if toolConfig := gjson.GetBytes(template, "toolConfig"); toolConfig.Exists() && !gjson.GetBytes(template, "request.toolConfig").Exists() { + template, _ = sjson.SetRawBytes(template, "request.toolConfig", []byte(toolConfig.Raw)) + template, _ = sjson.DeleteBytes(template, "toolConfig") + } + return template +} + +func generateRequestID() string { + return "agent-" + uuid.NewString() +} + +func generateImageGenRequestID() string { + return fmt.Sprintf("image_gen/%d/%s/12", time.Now().UnixMilli(), uuid.NewString()) +} + +func generateSessionID() string { + randSourceMutex.Lock() + n := randSource.Int63n(9_000_000_000_000_000_000) + randSourceMutex.Unlock() + return "-" + strconv.FormatInt(n, 10) +} + +func generateStableSessionID(payload []byte) string { + contents := gjson.GetBytes(payload, "request.contents") + if contents.IsArray() { + for _, content := range contents.Array() { + if content.Get("role").String() == "user" { + text := content.Get("parts.0.text").String() + if text != "" { + h := sha256.Sum256([]byte(text)) + n := int64(binary.BigEndian.Uint64(h[:8])) & 0x7FFFFFFFFFFFFFFF + return "-" + strconv.FormatInt(n, 10) + } + } + } + } + return generateSessionID() +} diff --git a/internal/runtime/executor/antigravity_executor_stream.go b/internal/runtime/executor/antigravity_executor_stream.go new file mode 100644 index 000000000..4990946bb --- /dev/null +++ b/internal/runtime/executor/antigravity_executor_stream.go @@ -0,0 +1,319 @@ +package executor + +import ( + "bufio" + "bytes" + "context" + "errors" + "fmt" + "io" + "net/http" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/sjson" +) + +// ExecuteStream performs a streaming request to the Antigravity API. +func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { + if opts.Alt == "responses/compact" { + return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + + ctx = context.WithValue(ctx, "alt", "") + if inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(ctx, auth, baseModel, time.Now()); errCooldown != nil { + return nil, homeKVUnavailableStatusErr(errCooldown) + } else if inCooldown && !antigravityShouldBypassShortCooldown(ctx, e.cfg) { + log.Debugf("antigravity executor: auth %s in short cooldown for model %s (%s remaining), returning 429 to switch auth", auth.ID, baseModel, remaining) + d := remaining + return nil, statusErr{code: http.StatusTooManyRequests, msg: fmt.Sprintf("auth in short cooldown, %s remaining", remaining), retryAfter: &d} + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("antigravity") + + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalPayload, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayload) + if errValidate != nil { + return nil, errValidate + } + req.Payload = originalPayload + token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) + if errToken != nil { + return nil, errToken + } + if updatedAuth != nil { + auth = updatedAuth + } + + originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, true) + translated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, true) + + translated, err = thinking.ApplyThinking(translated, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return nil, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + translated = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, "antigravity", from.String(), "request", translated, originalTranslated, requestedModel, requestPath, opts.Headers) + translated = sanitizeAntigravityGeminiRequestSignatures(baseModel, translated) + translated, _ = sjson.DeleteBytes(translated, "request.stream") + reporter.SetTranslatedReasoningEffort(translated, to.String()) + + useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg) + + baseURLs := antigravityBaseURLFallbackOrder(auth) + httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + + attempts := antigravityRetryAttempts(auth, e.cfg) + +attemptLoop: + for attempt := 0; attempt < attempts; attempt++ { + var lastStatus int + var lastBody []byte + var lastErr error + + for idx, baseURL := range baseURLs { + requestPayload := translated + if useCredits { + if cp := injectEnabledCreditTypes(translated); len(cp) > 0 { + requestPayload = cp + helps.MarkCreditsUsed(ctx) + } + } + replayScope := antigravityReasoningReplayScope{} + if antigravityUsesReasoningReplayCache(baseModel) { + var errReplay error + requestPayload, replayScope, errReplay = prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, requestPayload) + if errReplay != nil { + err = errReplay + return nil, err + } + } + httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL, helps.DerivedAntigravitySessionID(opts.Metadata, req.Metadata)) + if errReq != nil { + err = errReq + return nil, err + } + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { + return nil, errDo + } + lastStatus = 0 + lastBody = nil + lastErr = errDo + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + err = errDo + return nil, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + if errors.Is(errRead, context.Canceled) || errors.Is(errRead, context.DeadlineExceeded) { + err = errRead + return nil, err + } + if errCtx := ctx.Err(); errCtx != nil { + err = errCtx + return nil, err + } + lastStatus = 0 + lastBody = nil + lastErr = errRead + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + err = errRead + return nil, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + if httpResp.StatusCode == http.StatusTooManyRequests { + decision := decideAntigravity429(bodyBytes) + + switch decision.kind { + case antigravity429DecisionInstantRetrySameAuth: + if attempt+1 < attempts { + if decision.retryAfter != nil && *decision.retryAfter > 0 { + wait := antigravityInstantRetryDelay(*decision.retryAfter) + log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait) + if errWait := antigravityWait(ctx, wait); errWait != nil { + return nil, errWait + } + } + continue attemptLoop + } + case antigravity429DecisionShortCooldownSwitchAuth: + if decision.retryAfter != nil && *decision.retryAfter > 0 { + if errMarkCooldown := markAntigravityShortCooldownRequired(ctx, auth, baseModel, time.Now(), *decision.retryAfter); errMarkCooldown != nil { + err = homeKVUnavailableStatusErr(errMarkCooldown) + return nil, err + } + log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel) + } + case antigravity429DecisionFullQuotaExhausted: + if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) { + markAntigravityCreditsPermanentlyDisabled(auth) + } + // No credits logic - just fall through to error return below + } + } + + lastStatus = httpResp.StatusCode + lastBody = append([]byte(nil), bodyBytes...) + lastErr = nil + if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + if antigravityShouldRetryTransientResourceExhausted429(httpResp.StatusCode, bodyBytes) && attempt+1 < attempts { + delay := antigravityTransient429RetryDelay(attempt) + log.Debugf("antigravity executor: transient 429 resource exhausted for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return nil, errWait + } + continue attemptLoop + } + if antigravityShouldRetryNoCapacity(httpResp.StatusCode, bodyBytes) { + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: no capacity on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + if attempt+1 < attempts { + delay := antigravityNoCapacityRetryDelay(attempt) + log.Debugf("antigravity executor: no capacity for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return nil, errWait + } + continue attemptLoop + } + } + if antigravityShouldRetrySoftRateLimit(httpResp.StatusCode, bodyBytes) { + if attempt+1 < attempts { + delay := antigravitySoftRateLimitDelay(attempt) + log.Debugf("antigravity executor: soft rate limit for model %s, retrying in %s (attempt %d/%d)", baseModel, delay, attempt+1, attempts) + if errWait := antigravityWait(ctx, delay); errWait != nil { + return nil, errWait + } + continue attemptLoop + } + } + if errClear := clearAntigravityReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, bodyBytes); errClear != nil { + err = errClear + return nil, err + } + err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes) + return nil, err + } + + // Stream success + if useCredits { + clearAntigravityCreditsFailureState(auth) + } + replayAccumulator := newAntigravityReasoningReplayAccumulator(replayScope, requestPayload) + out := make(chan cliproxyexecutor.StreamChunk) + go func(resp *http.Response) { + defer close(out) + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response line error: %v", errClose) + } + }() + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(nil, streamScannerBuffer) + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) + var param any + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if replayAccumulator != nil { + replayAccumulator.ObserveSSELine(line) + } + + // Filter usage metadata for all models + // Only retain usage statistics in the terminal chunk + line = helps.FilterSSEUsageMetadata(line) + + payload := helps.JSONPayload(line) + if payload == nil { + continue + } + + if detail, ok := helps.ParseAntigravityStreamUsage(payload); ok { + reporter.Publish(ctx, detail) + } + + payload = e.resolveWebSearchGroundingURLs(ctx, auth, from, originalPayload, translated, payload) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(payload), ¶m, claudeInputTokens) + for i := range chunks { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: + case <-ctx.Done(): + return + } + } + } + tail := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("[DONE]"), ¶m, claudeInputTokens) + for i := range tail { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: tail[i]}: + case <-ctx.Done(): + return + } + } + if errScan := scanner.Err(); errScan != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + select { + case out <- cliproxyexecutor.StreamChunk{Err: errScan}: + case <-ctx.Done(): + } + } else { + if replayAccumulator != nil { + replayAccumulator.Commit(ctx) + } + reporter.EnsurePublished(ctx) + } + }(httpResp) + return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil + } + + switch { + case lastStatus != 0: + err = newAntigravityStatusErr(lastStatus, lastBody) + case lastErr != nil: + err = lastErr + default: + err = statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + } + return nil, err + } + + return nil, err +} diff --git a/internal/runtime/executor/antigravity_executor_tokens.go b/internal/runtime/executor/antigravity_executor_tokens.go new file mode 100644 index 000000000..98d1d561b --- /dev/null +++ b/internal/runtime/executor/antigravity_executor_tokens.go @@ -0,0 +1,188 @@ +package executor + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/url" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +// CountTokens counts tokens for the given request using the Antigravity API. +func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + baseModel := thinking.ParseSuffix(req.Model).ModelName + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("antigravity") + respCtx := context.WithValue(ctx, "alt", opts.Alt) + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayloadSource, errValidate := validateAntigravityRequestSignatures(ctx, baseModel, from, originalPayloadSource) + if errValidate != nil { + return cliproxyexecutor.Response{}, errValidate + } + req.Payload = originalPayloadSource + token, updatedAuth, errToken := e.ensureAccessToken(ctx, auth) + if errToken != nil { + return cliproxyexecutor.Response{}, errToken + } + if updatedAuth != nil { + auth = updatedAuth + } + if strings.TrimSpace(token) == "" { + return cliproxyexecutor.Response{}, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} + } + + // Prepare payload once (doesn't depend on baseURL) + payload := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, false) + + payload, err := thinking.ApplyThinking(payload, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return cliproxyexecutor.Response{}, err + } + payload = sanitizeAntigravityGeminiRequestSignatures(baseModel, payload) + preparedPayload, _, errReplay := prepareAntigravityGeminiReasoningReplayPayload(ctx, baseModel, req, opts, payload) + if errReplay != nil { + return cliproxyexecutor.Response{}, errReplay + } + payload = preparedPayload + + payload = helps.DeleteJSONField(payload, "project") + payload = helps.DeleteJSONField(payload, "model") + payload = helps.DeleteJSONField(payload, "request.safetySettings") + + baseURLs := antigravityBaseURLFallbackOrder(auth) + httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) + + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + + var lastStatus int + var lastBody []byte + var lastErr error + + for idx, baseURL := range baseURLs { + base := strings.TrimSuffix(baseURL, "/") + if base == "" { + base = buildBaseURL(auth) + } + + var requestURL strings.Builder + requestURL.WriteString(base) + requestURL.WriteString(antigravityCountTokensPath) + if opts.Alt != "" { + requestURL.WriteString("?$alt=") + requestURL.WriteString(url.QueryEscape(opts.Alt)) + } + + httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(payload)) + if errReq != nil { + return cliproxyexecutor.Response{}, errReq + } + httpReq.Close = true + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Authorization", "Bearer "+token) + httpReq.Header.Set("User-Agent", resolveUserAgent(auth)) + if host := resolveHost(base); host != "" { + httpReq.Host = host + } + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(httpReq, attrs) + + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: requestURL.String(), + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: payload, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errDo) + if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) { + return cliproxyexecutor.Response{}, errDo + } + lastStatus = 0 + lastBody = nil + lastErr = errDo + if idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + return cliproxyexecutor.Response{}, errDo + } + + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + bodyBytes, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("antigravity executor: close response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + return cliproxyexecutor.Response{}, errRead + } + helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes) + + if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices { + count := gjson.GetBytes(bodyBytes, "totalTokens").Int() + translated := sdktranslator.TranslateTokenCount(respCtx, to, responseFormat, count, bodyBytes) + return cliproxyexecutor.Response{Payload: translated, Headers: httpResp.Header.Clone()}, nil + } + + lastStatus = httpResp.StatusCode + lastBody = append([]byte(nil), bodyBytes...) + lastErr = nil + if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) { + log.Debugf("antigravity executor: rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1]) + continue + } + sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} + if httpResp.StatusCode == http.StatusTooManyRequests { + if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { + sErr.retryAfter = retryAfter + } + } + return cliproxyexecutor.Response{}, sErr + } + + switch { + case lastStatus != 0: + sErr := statusErr{code: lastStatus, msg: string(lastBody)} + if lastStatus == http.StatusTooManyRequests { + if retryAfter, parseErr := helps.ParseRetryDelay(lastBody); parseErr == nil && retryAfter != nil { + sErr.retryAfter = retryAfter + } + } + return cliproxyexecutor.Response{}, sErr + case lastErr != nil: + return cliproxyexecutor.Response{}, lastErr + default: + return cliproxyexecutor.Response{}, statusErr{code: http.StatusServiceUnavailable, msg: "antigravity executor: no base url available"} + } +} diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index 8e10f9893..cd511c157 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -1,38 +1,20 @@ package executor import ( - "bufio" "bytes" - "compress/flate" - "compress/gzip" "context" - "crypto/sha256" - "encoding/hex" "fmt" - "io" "net/http" "strings" - "time" - "github.com/andybalholm/brotli" - "github.com/google/uuid" - "github.com/klauspost/compress/zstd" - claudeauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" sigcompat "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" - sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "github.com/tidwall/sjson" - - "github.com/gin-gonic/gin" ) // ClaudeExecutor is a stateless executor for Anthropic Claude over the messages API. @@ -262,2482 +244,3 @@ func (e *ClaudeExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Aut httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) return httpClient.Do(httpReq) } - -func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - if opts.Alt == "responses/compact" { - return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} - } - baseModel := thinking.ParseSuffix(req.Model).ModelName - upstreamModel := e.upstreamModel(baseModel) - - apiKey, baseURL := claudeCreds(auth) - if baseURL == "" { - baseURL = "https://api.anthropic.com" - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("claude") - // Use streaming translation to preserve function calling, except for claude. - stream := from != to - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, stream) - body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream) - body = helps.SetStringIfDifferent(body, "model", upstreamModel) - - body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return resp, err - } - if rebuildMidSystemMessageEnabled(e.cfg, auth) { - body = rebuildMidSystemMessagesToTopLevel(body) - } - - // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) - // based on client type and configuration. - body, err = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) - if err != nil { - return resp, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body = ensureModelMaxTokens(body, baseModel) - - // Disable thinking if tool_choice forces tool use (Anthropic API constraint) - body = disableThinkingIfToolChoiceForced(body) - body = normalizeClaudeSamplingForUpstream(body) - // Claude OAuth (and this executor's redact-thinking beta) returns signature-only - // thinking blocks unless display is set to "summarized". - body = ensureClaudeThinkingDisplay(body) - - // Auto-inject cache_control if missing (optimization for ClawdBot/clients without caching support) - if countCacheControls(body) == 0 { - body = ensureCacheControl(body) - } - - // Enforce Anthropic's cache_control block limit (max 4 breakpoints per request). - // Cloaking and ensureCacheControl may push the total over 4 when the client - // already sends multiple cache_control blocks. - body = enforceCacheControlLimit(body, 4) - - // Normalize TTL values to prevent ordering violations under prompt-caching-scope-2026-01-05. - // A 1h-TTL block must not appear after a 5m-TTL block in evaluation order (tools→system→messages). - body = normalizeCacheControlTTL(body) - - // Extract betas from body and convert to header - var extraBetas []string - extraBetas, body = extractAndRemoveBetas(body) - bodyForTranslation := body - bodyForUpstream := body - oauthToken := isClaudeOAuthToken(apiKey) - var oauthToolNamesReverseMap map[string]string - if oauthToken { - bodyForUpstream, oauthToolNamesReverseMap = prepareClaudeOAuthToolNamesForUpstream(bodyForUpstream, claudeToolPrefix, auth.ToolPrefixDisabled()) - } - bodyForUpstream = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, bodyForUpstream, baseModel) - // Enable cch signing by default for OAuth tokens (not just experimental flag). - // Claude Code always computes cch; missing or invalid cch is a detectable fingerprint. - if oauthToken || experimentalCCHSigningEnabled(e.cfg, auth) { - bodyForUpstream = signAnthropicMessagesBody(bodyForUpstream) - } - reporter.SetTranslatedReasoningEffort(bodyForUpstream, to.String()) - - url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL) - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyForUpstream)) - if err != nil { - return resp, err - } - if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg, opts.Headers); errHeaders != nil { - return resp, errHeaders - } - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: url, - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: bodyForUpstream, - Provider: e.upstreamRequestLogProvider(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - // Decompress error responses — pass the Content-Encoding value (may be empty) - // and let decodeResponseBody handle both header-declared and magic-byte-detected - // compression. This keeps error-path behaviour consistent with the success path. - errBody, decErr := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) - if decErr != nil { - helps.RecordAPIResponseError(ctx, e.cfg, decErr) - msg := fmt.Sprintf("failed to decode error response body: %v", decErr) - helps.LogWithRequestID(ctx).Warn(msg) - return resp, statusErr{code: httpResp.StatusCode, msg: msg} - } - b, readErr := io.ReadAll(errBody) - if readErr != nil { - helps.RecordAPIResponseError(ctx, e.cfg, readErr) - msg := fmt.Sprintf("failed to read error response body: %v", readErr) - helps.LogWithRequestID(ctx).Warn(msg) - b = []byte(msg) - } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) - err = statusErr{code: httpResp.StatusCode, msg: string(b)} - if errClose := errBody.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - return resp, err - } - decodedBody, err := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - return resp, err - } - defer func() { - if errClose := decodedBody.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - }() - data, err := io.ReadAll(decodedBody) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - if stream { - if errValidate := validateClaudeStreamingResponse(data); errValidate != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errValidate) - return resp, errValidate - } - lines := bytes.Split(data, []byte("\n")) - for _, line := range lines { - if detail, ok := helps.ParseClaudeStreamUsage(line); ok { - reporter.Publish(ctx, detail) - } - } - } else { - reporter.Publish(ctx, helps.ParseClaudeUsage(data)) - } - data = restoreClaudeOAuthToolNamesFromResponse(data, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) - data = e.restoreResponseModel(data, req.Model) - var param any - out := sdktranslator.TranslateNonStream( - ctx, - to, - responseFormat, - req.Model, - opts.OriginalRequest, - bodyForTranslation, - data, - ¶m, - ) - resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} - return resp, nil -} - -func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { - if opts.Alt == "responses/compact" { - return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} - } - baseModel := thinking.ParseSuffix(req.Model).ModelName - upstreamModel := e.upstreamModel(baseModel) - - apiKey, baseURL := claudeCreds(auth) - if baseURL == "" { - baseURL = "https://api.anthropic.com" - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("claude") - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, true) - body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, true) - body = helps.SetStringIfDifferent(body, "model", upstreamModel) - - body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return nil, err - } - if rebuildMidSystemMessageEnabled(e.cfg, auth) { - body = rebuildMidSystemMessagesToTopLevel(body) - } - - // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) - // based on client type and configuration. - body, err = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) - if err != nil { - return nil, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body = ensureModelMaxTokens(body, baseModel) - - // Disable thinking if tool_choice forces tool use (Anthropic API constraint) - body = disableThinkingIfToolChoiceForced(body) - body = normalizeClaudeSamplingForUpstream(body) - // Claude OAuth (and this executor's redact-thinking beta) returns signature-only - // thinking blocks unless display is set to "summarized". - body = ensureClaudeThinkingDisplay(body) - - // Auto-inject cache_control if missing (optimization for ClawdBot/clients without caching support) - if countCacheControls(body) == 0 { - body = ensureCacheControl(body) - } - - // Enforce Anthropic's cache_control block limit (max 4 breakpoints per request). - body = enforceCacheControlLimit(body, 4) - - // Normalize TTL values to prevent ordering violations under prompt-caching-scope-2026-01-05. - body = normalizeCacheControlTTL(body) - - // Extract betas from body and convert to header - var extraBetas []string - extraBetas, body = extractAndRemoveBetas(body) - bodyForTranslation := body - bodyForUpstream := body - oauthToken := isClaudeOAuthToken(apiKey) - var oauthToolNamesReverseMap map[string]string - if oauthToken { - bodyForUpstream, oauthToolNamesReverseMap = prepareClaudeOAuthToolNamesForUpstream(bodyForUpstream, claudeToolPrefix, auth.ToolPrefixDisabled()) - } - bodyForUpstream = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, bodyForUpstream, baseModel) - // Enable cch signing by default for OAuth tokens (not just experimental flag). - if oauthToken || experimentalCCHSigningEnabled(e.cfg, auth) { - bodyForUpstream = signAnthropicMessagesBody(bodyForUpstream) - } - reporter.SetTranslatedReasoningEffort(bodyForUpstream, to.String()) - - url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL) - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyForUpstream)) - if err != nil { - return nil, err - } - if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, true, extraBetas, e.cfg, opts.Headers); errHeaders != nil { - return nil, errHeaders - } - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: url, - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: bodyForUpstream, - Provider: e.upstreamRequestLogProvider(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return nil, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - // Decompress error responses — pass the Content-Encoding value (may be empty) - // and let decodeResponseBody handle both header-declared and magic-byte-detected - // compression. This keeps error-path behaviour consistent with the success path. - errBody, decErr := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) - if decErr != nil { - helps.RecordAPIResponseError(ctx, e.cfg, decErr) - msg := fmt.Sprintf("failed to decode error response body: %v", decErr) - helps.LogWithRequestID(ctx).Warn(msg) - return nil, statusErr{code: httpResp.StatusCode, msg: msg} - } - b, readErr := io.ReadAll(errBody) - if readErr != nil { - helps.RecordAPIResponseError(ctx, e.cfg, readErr) - msg := fmt.Sprintf("failed to read error response body: %v", readErr) - helps.LogWithRequestID(ctx).Warn(msg) - b = []byte(msg) - } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) - if errClose := errBody.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - err = statusErr{code: httpResp.StatusCode, msg: string(b)} - return nil, err - } - decodedBody, err := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - return nil, err - } - out := make(chan cliproxyexecutor.StreamChunk) - go func() { - defer close(out) - defer func() { - if errClose := decodedBody.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - }() - - // If the response target is Claude, directly forward complete SSE events without translation. - if responseFormat == to { - scanner := bufio.NewScanner(decodedBody) - scanner.Buffer(nil, 52_428_800) // 50MB - var event bytes.Buffer - flushEvent := func() bool { - if event.Len() == 0 { - return true - } - cloned := bytes.Clone(event.Bytes()) - event.Reset() - select { - case out <- cliproxyexecutor.StreamChunk{Payload: cloned}: - return true - case <-ctx.Done(): - return false - } - } - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if detail, ok := helps.ParseClaudeStreamUsage(line); ok { - reporter.Publish(ctx, detail) - } - line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) - line = e.restoreResponseModel(line, req.Model) - event.Write(line) - event.WriteByte('\n') - if len(bytes.TrimSpace(line)) == 0 && !flushEvent() { - return - } - } - if !flushEvent() { - return - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - select { - case out <- cliproxyexecutor.StreamChunk{Err: errScan}: - case <-ctx.Done(): - } - } - return - } - - // For other formats, use translation - scanner := bufio.NewScanner(decodedBody) - scanner.Buffer(nil, 52_428_800) // 50MB - var param any - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - if detail, ok := helps.ParseClaudeStreamUsage(line); ok { - reporter.Publish(ctx, detail) - } - line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) - line = e.restoreResponseModel(line, req.Model) - chunks := sdktranslator.TranslateStream( - ctx, - to, - responseFormat, - req.Model, - opts.OriginalRequest, - bodyForTranslation, - bytes.Clone(line), - ¶m, - ) - for i := range chunks { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: - case <-ctx.Done(): - return - } - } - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - select { - case out <- cliproxyexecutor.StreamChunk{Err: errScan}: - case <-ctx.Done(): - } - } - }() - return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil -} - -func validateClaudeStreamingResponse(data []byte) error { - scanner := bufio.NewScanner(bytes.NewReader(data)) - scanner.Buffer(nil, 52_428_800) - - hasData := false - hasMessageStart := false - hasMessageDelta := false - - for scanner.Scan() { - line := bytes.TrimSpace(scanner.Bytes()) - if len(line) == 0 || !bytes.HasPrefix(line, []byte("data:")) { - continue - } - payload := bytes.TrimSpace(line[len("data:"):]) - if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) { - continue - } - hasData = true - if !gjson.ValidBytes(payload) { - return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned malformed stream data"} - } - - root := gjson.ParseBytes(payload) - switch root.Get("type").String() { - case "error": - message := strings.TrimSpace(root.Get("error.message").String()) - if message == "" { - message = strings.TrimSpace(root.Get("error.type").String()) - } - if message == "" { - message = "unknown upstream error" - } - return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned error event: " + message} - case "message_start": - message := root.Get("message") - if strings.TrimSpace(message.Get("id").String()) == "" || strings.TrimSpace(message.Get("model").String()) == "" { - return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream message_start is missing id or model"} - } - hasMessageStart = true - case "message_delta": - hasMessageDelta = true - } - } - if errScan := scanner.Err(); errScan != nil { - return errScan - } - if !hasData { - return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned empty stream response"} - } - if !hasMessageStart { - return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response is missing message_start"} - } - if !hasMessageDelta { - return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response ended before message completion"} - } - return nil -} - -func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - baseModel := thinking.ParseSuffix(req.Model).ModelName - upstreamModel := e.upstreamModel(baseModel) - - apiKey, baseURL := claudeCreds(auth) - if baseURL == "" { - baseURL = "https://api.anthropic.com" - } - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("claude") - // Use streaming translation to preserve function calling, except for claude. - stream := from != to - body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream) - body = helps.SetStringIfDifferent(body, "model", upstreamModel) - if rebuildMidSystemMessageEnabled(e.cfg, auth) { - body = rebuildMidSystemMessagesToTopLevel(body) - } - - if !strings.HasPrefix(baseModel, "claude-3-5-haiku") { - body = checkSystemInstructions(body) - } - - // Keep count_tokens requests compatible with Anthropic cache-control constraints too. - body = enforceCacheControlLimit(body, 4) - body = normalizeCacheControlTTL(body) - - // Extract betas from body and convert to header (for count_tokens too) - var extraBetas []string - extraBetas, body = extractAndRemoveBetas(body) - if isClaudeOAuthToken(apiKey) { - body, _ = prepareClaudeOAuthToolNamesForUpstream(body, claudeToolPrefix, auth.ToolPrefixDisabled()) - } - body = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, body, baseModel) - - url := fmt.Sprintf("%s/v1/messages/count_tokens?beta=true", baseURL) - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) - if err != nil { - return cliproxyexecutor.Response{}, err - } - if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg, opts.Headers); errHeaders != nil { - return cliproxyexecutor.Response{}, errHeaders - } - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: url, - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: body, - Provider: e.upstreamRequestLogProvider(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) - resp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return cliproxyexecutor.Response{}, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, resp.StatusCode, resp.Header.Clone()) - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - // Decompress error responses — pass the Content-Encoding value (may be empty) - // and let decodeResponseBody handle both header-declared and magic-byte-detected - // compression. This keeps error-path behaviour consistent with the success path. - errBody, decErr := decodeResponseBody(resp.Body, resp.Header.Get("Content-Encoding")) - if decErr != nil { - helps.RecordAPIResponseError(ctx, e.cfg, decErr) - msg := fmt.Sprintf("failed to decode error response body: %v", decErr) - helps.LogWithRequestID(ctx).Warn(msg) - return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: msg} - } - b, readErr := io.ReadAll(errBody) - if readErr != nil { - helps.RecordAPIResponseError(ctx, e.cfg, readErr) - msg := fmt.Sprintf("failed to read error response body: %v", readErr) - helps.LogWithRequestID(ctx).Warn(msg) - b = []byte(msg) - } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) - if errClose := errBody.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)} - } - decodedBody, err := decodeResponseBody(resp.Body, resp.Header.Get("Content-Encoding")) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - return cliproxyexecutor.Response{}, err - } - defer func() { - if errClose := decodedBody.Close(); errClose != nil { - log.Errorf("response body close error: %v", errClose) - } - }() - data, err := io.ReadAll(decodedBody) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return cliproxyexecutor.Response{}, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - count := gjson.GetBytes(data, "input_tokens").Int() - out := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, data) - return cliproxyexecutor.Response{Payload: out, Headers: resp.Header.Clone()}, nil -} - -func (e *ClaudeExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - log.Debugf("claude executor: refresh called") - if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { - return refreshed, err - } - if auth == nil { - return nil, fmt.Errorf("claude executor: auth is nil") - } - var refreshToken string - if auth.Metadata != nil { - if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" { - refreshToken = v - } - } - if refreshToken == "" { - return auth, nil - } - svc := claudeauth.NewClaudeAuthWithProxyURL(e.cfg, auth.ProxyURL) - td, err := svc.RefreshTokensWithRetry(ctx, refreshToken, 3) - if err != nil { - return nil, err - } - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["access_token"] = td.AccessToken - if td.RefreshToken != "" { - auth.Metadata["refresh_token"] = td.RefreshToken - } - auth.Metadata["email"] = td.Email - auth.Metadata["expired"] = td.Expire - auth.Metadata["type"] = "claude" - now := time.Now().Format(time.RFC3339) - auth.Metadata["last_refresh"] = now - return auth, nil -} - -// extractAndRemoveBetas extracts the "betas" array from the body and removes it. -// Returns the extracted betas as a string slice and the modified body. -func extractAndRemoveBetas(body []byte) ([]string, []byte) { - betasResult := gjson.GetBytes(body, "betas") - if !betasResult.Exists() { - return nil, body - } - var betas []string - if betasResult.IsArray() { - for _, item := range betasResult.Array() { - if s := strings.TrimSpace(item.String()); s != "" { - betas = append(betas, s) - } - } - } else if s := strings.TrimSpace(betasResult.String()); s != "" { - betas = append(betas, s) - } - body, _ = sjson.DeleteBytes(body, "betas") - return betas, body -} - -// disableThinkingIfToolChoiceForced checks if tool_choice forces tool use and disables thinking. -// Anthropic API does not allow thinking when tool_choice is set to "any" or a specific tool. -// See: https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#important-considerations -func disableThinkingIfToolChoiceForced(body []byte) []byte { - toolChoiceType := gjson.GetBytes(body, "tool_choice.type").String() - // "auto" is allowed with thinking, but "any" or "tool" (specific tool) are not - if toolChoiceType == "any" || toolChoiceType == "tool" { - // Remove thinking configuration entirely to avoid API error - body, _ = sjson.DeleteBytes(body, "thinking") - // Adaptive thinking may also set output_config.effort; remove it to avoid - // leaking thinking controls when tool_choice forces tool use. - body, _ = sjson.DeleteBytes(body, "output_config.effort") - if oc := gjson.GetBytes(body, "output_config"); oc.Exists() && oc.IsObject() && len(oc.Map()) == 0 { - body, _ = sjson.DeleteBytes(body, "output_config") - } - } - return body -} - -// normalizeClaudeSamplingForUpstream keeps Anthropic message requests valid. -func normalizeClaudeSamplingForUpstream(body []byte) []byte { - body, _ = sjson.DeleteBytes(body, "temperature") - body, _ = sjson.DeleteBytes(body, "top_p") - - thinkingType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "thinking.type").String())) - switch thinkingType { - case "enabled", "adaptive", "auto": - body, _ = sjson.DeleteBytes(body, "top_p") - body, _ = sjson.DeleteBytes(body, "top_k") - } - return body -} - -// ensureClaudeThinkingDisplay defaults thinking.display to "summarized" when thinking -// is active and the client did not set display. Without this, Claude backends that -// enable redact-thinking return signature-only thinking blocks (empty thinking text). -// Explicit client values such as "omitted" are preserved. -func ensureClaudeThinkingDisplay(body []byte) []byte { - thinkingType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "thinking.type").String())) - switch thinkingType { - case "enabled", "adaptive", "auto": - default: - return body - } - if display := strings.TrimSpace(gjson.GetBytes(body, "thinking.display").String()); display != "" { - return body - } - out, err := sjson.SetBytes(body, "thinking.display", "summarized") - if err != nil { - return body - } - return out -} - -type compositeReadCloser struct { - io.Reader - closers []func() error -} - -func (c *compositeReadCloser) Close() error { - var firstErr error - for i := range c.closers { - if c.closers[i] == nil { - continue - } - if err := c.closers[i](); err != nil && firstErr == nil { - firstErr = err - } - } - return firstErr -} - -// peekableBody wraps a bufio.Reader around the original ReadCloser so that -// magic bytes can be inspected without consuming them from the stream. -type peekableBody struct { - *bufio.Reader - closer io.Closer -} - -func (p *peekableBody) Close() error { - return p.closer.Close() -} - -func decodeResponseBody(body io.ReadCloser, contentEncoding string) (io.ReadCloser, error) { - if body == nil { - return nil, fmt.Errorf("response body is nil") - } - if contentEncoding == "" { - // No Content-Encoding header. Attempt best-effort magic-byte detection to - // handle misbehaving upstreams that compress without setting the header. - // Only gzip (1f 8b) and zstd (28 b5 2f fd) have reliable magic sequences; - // br and deflate have none and are left as-is. - // The bufio wrapper preserves unread bytes so callers always see the full - // stream regardless of whether decompression was applied. - pb := &peekableBody{Reader: bufio.NewReader(body), closer: body} - magic, peekErr := pb.Peek(4) - if peekErr == nil || (peekErr == io.EOF && len(magic) >= 2) { - switch { - case len(magic) >= 2 && magic[0] == 0x1f && magic[1] == 0x8b: - gzipReader, gzErr := gzip.NewReader(pb) - if gzErr != nil { - _ = pb.Close() - return nil, fmt.Errorf("magic-byte gzip: failed to create reader: %w", gzErr) - } - return &compositeReadCloser{ - Reader: gzipReader, - closers: []func() error{ - gzipReader.Close, - pb.Close, - }, - }, nil - case len(magic) >= 4 && magic[0] == 0x28 && magic[1] == 0xb5 && magic[2] == 0x2f && magic[3] == 0xfd: - decoder, zdErr := zstd.NewReader(pb) - if zdErr != nil { - _ = pb.Close() - return nil, fmt.Errorf("magic-byte zstd: failed to create reader: %w", zdErr) - } - return &compositeReadCloser{ - Reader: decoder, - closers: []func() error{ - func() error { decoder.Close(); return nil }, - pb.Close, - }, - }, nil - } - } - return pb, nil - } - encodings := strings.Split(contentEncoding, ",") - for _, raw := range encodings { - encoding := strings.TrimSpace(strings.ToLower(raw)) - switch encoding { - case "", "identity": - continue - case "gzip": - gzipReader, err := gzip.NewReader(body) - if err != nil { - _ = body.Close() - return nil, fmt.Errorf("failed to create gzip reader: %w", err) - } - return &compositeReadCloser{ - Reader: gzipReader, - closers: []func() error{ - gzipReader.Close, - func() error { return body.Close() }, - }, - }, nil - case "deflate": - deflateReader := flate.NewReader(body) - return &compositeReadCloser{ - Reader: deflateReader, - closers: []func() error{ - deflateReader.Close, - func() error { return body.Close() }, - }, - }, nil - case "br": - return &compositeReadCloser{ - Reader: brotli.NewReader(body), - closers: []func() error{ - func() error { return body.Close() }, - }, - }, nil - case "zstd": - decoder, err := zstd.NewReader(body) - if err != nil { - _ = body.Close() - return nil, fmt.Errorf("failed to create zstd reader: %w", err) - } - return &compositeReadCloser{ - Reader: decoder, - closers: []func() error{ - func() error { decoder.Close(); return nil }, - func() error { return body.Close() }, - }, - }, nil - default: - continue - } - } - return body, nil -} - -func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, stream bool, extraBetas []string, cfg *config.Config, incomingHeaders http.Header) error { - if r == nil { - return nil - } - hdrDefault := func(cfgVal, fallback string) string { - if cfgVal != "" { - return cfgVal - } - return fallback - } - - var hd config.ClaudeHeaderDefaults - if cfg != nil { - hd = cfg.ClaudeHeaderDefaults - } - - useAPIKey := auth != nil && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["api_key"]) != "" - isAnthropicBase := r.URL != nil && strings.EqualFold(r.URL.Scheme, "https") && strings.EqualFold(r.URL.Host, "api.anthropic.com") - if isAnthropicBase && useAPIKey { - r.Header.Del("Authorization") - r.Header.Set("x-api-key", apiKey) - } else { - r.Header.Set("Authorization", "Bearer "+apiKey) - } - r.Header.Set("Content-Type", "application/json") - - if incomingHeaders == nil { - if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - incomingHeaders = ginCtx.Request.Header - } - } - stabilizeDeviceProfile := helps.ClaudeDeviceProfileStabilizationEnabled(cfg) - var deviceProfile helps.ClaudeDeviceProfile - if stabilizeDeviceProfile { - var errDeviceProfile error - deviceProfile, errDeviceProfile = helps.ResolveClaudeDeviceProfileRequired(r.Context(), auth, apiKey, incomingHeaders, cfg) - if errDeviceProfile != nil { - return errDeviceProfile - } - } - - baseBetas := "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27,prompt-caching-scope-2026-01-05,structured-outputs-2025-12-15,fast-mode-2026-02-01,redact-thinking-2026-02-12,token-efficient-tools-2026-03-28" - if val := strings.TrimSpace(strings.Join(incomingHeaders.Values("Anthropic-Beta"), ",")); val != "" { - baseBetas = val - if !strings.Contains(val, "oauth") { - baseBetas += ",oauth-2025-04-20" - } - } - if !strings.Contains(baseBetas, "interleaved-thinking") { - baseBetas += ",interleaved-thinking-2025-05-14" - } - - // Merge extra betas from request body and request flags. - if len(extraBetas) > 0 { - existingSet := make(map[string]bool) - for _, b := range strings.Split(baseBetas, ",") { - betaName := strings.TrimSpace(b) - if betaName != "" { - existingSet[betaName] = true - } - } - for _, beta := range extraBetas { - beta = strings.TrimSpace(beta) - if beta != "" && !existingSet[beta] { - baseBetas += "," + beta - existingSet[beta] = true - } - } - } - r.Header.Set("Anthropic-Beta", baseBetas) - - misc.EnsureHeader(r.Header, incomingHeaders, "Anthropic-Version", "2023-06-01") - // Only set browser access header for API key mode; real Claude Code CLI does not send it. - if useAPIKey { - misc.EnsureHeader(r.Header, incomingHeaders, "Anthropic-Dangerous-Direct-Browser-Access", "true") - } - misc.EnsureHeader(r.Header, incomingHeaders, "X-App", "cli") - // Values below match Claude Code 2.1.63 / @anthropic-ai/sdk 0.74.0 (updated 2026-02-28). - misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Retry-Count", "0") - misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Runtime", "node") - misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Lang", "js") - misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Timeout", hdrDefault(hd.Timeout, "600")) - // Session ID: stable per auth/apiKey, matches Claude Code's X-Claude-Code-Session-Id header. - sessionID, errSessionID := helps.CachedSessionIDRequired(r.Context(), apiKey) - if errSessionID != nil { - return errSessionID - } - misc.EnsureHeader(r.Header, incomingHeaders, "X-Claude-Code-Session-Id", sessionID) - // Per-request UUID, matches Claude Code's x-client-request-id for first-party API. - if isAnthropicBase { - misc.EnsureHeader(r.Header, incomingHeaders, "x-client-request-id", uuid.New().String()) - } - r.Header.Set("Connection", "keep-alive") - if stream { - r.Header.Set("Accept", "text/event-stream") - // SSE streams must not be compressed: the downstream scanner reads - // line-delimited text and cannot parse compressed bytes. Using - // "identity" tells the upstream to send an uncompressed stream. - r.Header.Set("Accept-Encoding", "identity") - } else { - r.Header.Set("Accept", "application/json") - r.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd") - } - // Legacy mode keeps OS/Arch runtime-derived; stabilized mode pins OS/Arch - // to the configured baseline while still allowing newer official - // User-Agent/package/runtime tuples to upgrade the software fingerprint. - if stabilizeDeviceProfile { - helps.ApplyClaudeDeviceProfileHeaders(r, deviceProfile) - } else { - helps.ApplyClaudeLegacyDeviceHeaders(r, incomingHeaders, cfg) - } - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(r, attrs) - // Re-enforce Accept-Encoding: identity after ApplyCustomHeadersFromAttrs, which - // may override it with a user-configured value. Compressed SSE breaks the line - // scanner regardless of user preference, so this is non-negotiable for streams. - if stream { - r.Header.Set("Accept-Encoding", "identity") - } - return nil -} - -func claudeCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) { - if a == nil { - return "", "" - } - if a.Attributes != nil { - apiKey = a.Attributes["api_key"] - baseURL = a.Attributes["base_url"] - } - if apiKey == "" && a.Metadata != nil { - if v, ok := a.Metadata["access_token"].(string); ok { - apiKey = v - } - } - return -} - -func checkSystemInstructions(payload []byte) []byte { - return checkSystemInstructionsWithSigningMode(payload, false, false, false, "2.1.63", "", "") -} - -func rebuildMidSystemMessagesToTopLevel(payload []byte) []byte { - messages := gjson.GetBytes(payload, "messages") - if !messages.IsArray() { - return payload - } - - var movedSystemParts []string - keptMessages := make([]string, 0, int(messages.Get("#").Int())) - messages.ForEach(func(_, message gjson.Result) bool { - if strings.EqualFold(strings.TrimSpace(message.Get("role").String()), "system") { - movedSystemParts = append(movedSystemParts, claudeSystemTextParts(message.Get("content"))...) - return true - } - keptMessages = append(keptMessages, message.Raw) - return true - }) - if len(movedSystemParts) == 0 { - return payload - } - - systemParts := claudeSystemTextParts(gjson.GetBytes(payload, "system")) - systemParts = append(systemParts, movedSystemParts...) - if len(systemParts) > 0 { - if updated, errSetSystem := sjson.SetRawBytes(payload, "system", rawJSONArray(systemParts)); errSetSystem == nil { - payload = updated - } - } - if updated, errSetMessages := sjson.SetRawBytes(payload, "messages", rawJSONArray(keptMessages)); errSetMessages == nil { - payload = updated - } - return payload -} - -func claudeSystemTextParts(content gjson.Result) []string { - if !content.Exists() { - return nil - } - if content.Type == gjson.String { - text := content.String() - if strings.TrimSpace(text) == "" { - return nil - } - block := []byte(`{"type":"text","text":""}`) - block, _ = sjson.SetBytes(block, "text", text) - return []string{string(block)} - } - if !content.IsArray() { - return nil - } - - var parts []string - content.ForEach(func(_, item gjson.Result) bool { - if item.Type == gjson.String { - text := item.String() - if strings.TrimSpace(text) != "" { - block := []byte(`{"type":"text","text":""}`) - block, _ = sjson.SetBytes(block, "text", text) - parts = append(parts, string(block)) - } - return true - } - if item.IsObject() && item.Get("type").String() == "text" && strings.TrimSpace(item.Get("text").String()) != "" { - parts = append(parts, item.Raw) - } - return true - }) - return parts -} - -func rawJSONArray(items []string) []byte { - if len(items) == 0 { - return []byte("[]") - } - var builder strings.Builder - builder.WriteByte('[') - for i, item := range items { - if i > 0 { - builder.WriteByte(',') - } - builder.WriteString(item) - } - builder.WriteByte(']') - return []byte(builder.String()) -} - -func isClaudeOAuthToken(apiKey string) bool { - return strings.Contains(apiKey, "sk-ant-oat") -} - -// prepareClaudeOAuthToolNamesForUpstream applies the Claude OAuth tool-name -// transforms in the same order across request paths. Remap runs before prefixing -// so any future non-empty prefix still composes correctly with the per-request -// reverse map. -func prepareClaudeOAuthToolNamesForUpstream(body []byte, prefix string, prefixDisabled bool) ([]byte, map[string]string) { - body, reverseMap := remapOAuthToolNames(body) - if !prefixDisabled { - body = applyClaudeToolPrefix(body, prefix) - } - return body, reverseMap -} - -// restoreClaudeOAuthToolNamesFromResponse undoes the Claude OAuth tool-name -// transforms for non-stream responses in reverse order. -func restoreClaudeOAuthToolNamesFromResponse(body []byte, prefix string, prefixDisabled bool, reverseMap map[string]string) []byte { - if !prefixDisabled { - body = stripClaudeToolPrefixFromResponse(body, prefix) - } - return reverseRemapOAuthToolNames(body, reverseMap) -} - -// restoreClaudeOAuthToolNamesFromStreamLine undoes the Claude OAuth tool-name -// transforms for SSE lines in reverse order. -func restoreClaudeOAuthToolNamesFromStreamLine(line []byte, prefix string, prefixDisabled bool, reverseMap map[string]string) []byte { - if !prefixDisabled { - line = stripClaudeToolPrefixFromStreamLine(line, prefix) - } - return reverseRemapOAuthToolNamesFromStreamLine(line, reverseMap) -} - -// remapOAuthToolNames renames third-party tool names to Claude Code equivalents -// and removes tools without an official counterpart. This prevents Anthropic from -// fingerprinting the request as a third-party client via tool naming patterns. -// -// It operates on: tools[].name, tool_choice.name, and all tool_use/tool_reference -// references in messages. Removed tools' corresponding tool_result blocks are preserved -// (they just become orphaned, which is safe for Claude). -// -// The returned map is keyed on the upstream (TitleCase) name and maps to the -// client-supplied original name. Callers MUST pass this map to the reverse -// functions so only names the client actually caused us to rewrite are restored -// on the response. A global reverse map (the previous implementation) incorrectly -// rewrote names the client originally sent in TitleCase (e.g. `Bash`) -// when any OTHER tool in the same request triggered a forward rename (e.g. -// `glob` -> `Glob`), because the global reverse map contained `Bash` -> `bash` -// regardless of what the client originally sent. -func remapOAuthToolNames(body []byte) ([]byte, map[string]string) { - reverseMap := make(map[string]string, len(oauthToolRenameMap)) - recordRename := func(original, renamed string) { - // Preserve the first-seen original name if the same upstream name is - // produced from multiple call sites; they all map back identically. - if _, exists := reverseMap[renamed]; !exists { - reverseMap[renamed] = original - } - } - - // 1. Rewrite tools array in a single pass (if present). - // IMPORTANT: do not mutate names first and then rebuild from an older gjson - // snapshot. gjson results are snapshots of the original bytes; rebuilding from a - // stale snapshot will preserve removals but overwrite renamed names back to their - // original lowercase values. - tools := gjson.GetBytes(body, "tools") - toolsNeedRewrite := false - if tools.Exists() && tools.IsArray() { - tools.ForEach(func(_, tool gjson.Result) bool { - if tool.Get("type").Exists() && tool.Get("type").String() != "" { - return true - } - name := tool.Get("name").String() - toolsNeedRewrite = oauthToolsToRemove[name] - if !toolsNeedRewrite { - newName, ok := oauthToolRenameMap[name] - toolsNeedRewrite = ok && newName != name - } - return !toolsNeedRewrite - }) - } - if toolsNeedRewrite { - var toolsJSON strings.Builder - toolsJSON.WriteByte('[') - toolCount := 0 - tools.ForEach(func(_, tool gjson.Result) bool { - // Keep Anthropic built-in tools (web_search, code_execution, etc.) unchanged. - if tool.Get("type").Exists() && tool.Get("type").String() != "" { - if toolCount > 0 { - toolsJSON.WriteByte(',') - } - toolsJSON.WriteString(tool.Raw) - toolCount++ - return true - } - - name := tool.Get("name").String() - if oauthToolsToRemove[name] { - return true - } - - toolJSON := tool.Raw - if newName, ok := oauthToolRenameMap[name]; ok && newName != name { - updatedTool, err := sjson.Set(toolJSON, "name", newName) - if err == nil { - toolJSON = updatedTool - recordRename(name, newName) - } - } - - if toolCount > 0 { - toolsJSON.WriteByte(',') - } - toolsJSON.WriteString(toolJSON) - toolCount++ - return true - }) - toolsJSON.WriteByte(']') - body, _ = sjson.SetRawBytes(body, "tools", []byte(toolsJSON.String())) - } - - // 2. Rename tool_choice if it references a known tool - toolChoiceType := gjson.GetBytes(body, "tool_choice.type").String() - if toolChoiceType == "tool" { - tcName := gjson.GetBytes(body, "tool_choice.name").String() - if oauthToolsToRemove[tcName] { - // The chosen tool was removed from the tools array, so drop tool_choice to - // keep the payload internally consistent and fall back to normal auto tool use. - body, _ = sjson.DeleteBytes(body, "tool_choice") - } else if newName, ok := oauthToolRenameMap[tcName]; ok && newName != tcName { - body, _ = sjson.SetBytes(body, "tool_choice.name", newName) - recordRename(tcName, newName) - } - } - - // 3. Rename tool references in messages - messages := gjson.GetBytes(body, "messages") - if messages.Exists() && messages.IsArray() { - messages.ForEach(func(msgIndex, msg gjson.Result) bool { - content := msg.Get("content") - if !content.Exists() || !content.IsArray() { - return true - } - content.ForEach(func(contentIndex, part gjson.Result) bool { - partType := part.Get("type").String() - switch partType { - case "tool_use": - name := part.Get("name").String() - if newName, ok := oauthToolRenameMap[name]; ok && newName != name { - path := fmt.Sprintf("messages.%d.content.%d.name", msgIndex.Int(), contentIndex.Int()) - body, _ = sjson.SetBytes(body, path, newName) - recordRename(name, newName) - } - case "tool_reference": - toolName := part.Get("tool_name").String() - if newName, ok := oauthToolRenameMap[toolName]; ok && newName != toolName { - path := fmt.Sprintf("messages.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int()) - body, _ = sjson.SetBytes(body, path, newName) - recordRename(toolName, newName) - } - case "tool_result": - // Handle nested tool_reference blocks inside tool_result.content[] - toolID := part.Get("tool_use_id").String() - _ = toolID // tool_use_id stays as-is - nestedContent := part.Get("content") - if nestedContent.Exists() && nestedContent.IsArray() { - nestedContent.ForEach(func(nestedIndex, nestedPart gjson.Result) bool { - if nestedPart.Get("type").String() == "tool_reference" { - nestedToolName := nestedPart.Get("tool_name").String() - if newName, ok := oauthToolRenameMap[nestedToolName]; ok && newName != nestedToolName { - nestedPath := fmt.Sprintf("messages.%d.content.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int(), nestedIndex.Int()) - body, _ = sjson.SetBytes(body, nestedPath, newName) - recordRename(nestedToolName, newName) - } - } - return true - }) - } - } - return true - }) - return true - }) - } - - return body, reverseMap -} - -// reverseRemapOAuthToolNames reverses the tool name mapping for non-stream responses -// using the per-request map produced by remapOAuthToolNames. Names the client sent -// that were NOT forward-renamed are passed through unchanged. -func reverseRemapOAuthToolNames(body []byte, reverseMap map[string]string) []byte { - if len(reverseMap) == 0 { - return body - } - content := gjson.GetBytes(body, "content") - if !content.Exists() || !content.IsArray() { - return body - } - content.ForEach(func(index, part gjson.Result) bool { - partType := part.Get("type").String() - switch partType { - case "tool_use": - name := part.Get("name").String() - if origName, ok := reverseMap[name]; ok { - path := fmt.Sprintf("content.%d.name", index.Int()) - body, _ = sjson.SetBytes(body, path, origName) - } - case "tool_reference": - toolName := part.Get("tool_name").String() - if origName, ok := reverseMap[toolName]; ok { - path := fmt.Sprintf("content.%d.tool_name", index.Int()) - body, _ = sjson.SetBytes(body, path, origName) - } - } - return true - }) - return body -} - -// reverseRemapOAuthToolNamesFromStreamLine reverses the tool name mapping for SSE -// stream lines, using the per-request reverseMap produced by remapOAuthToolNames. -func reverseRemapOAuthToolNamesFromStreamLine(line []byte, reverseMap map[string]string) []byte { - if len(reverseMap) == 0 { - return line - } - payload := helps.JSONPayload(line) - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return line - } - - contentBlock := gjson.GetBytes(payload, "content_block") - if !contentBlock.Exists() { - return line - } - - blockType := contentBlock.Get("type").String() - var updated []byte - var err error - - switch blockType { - case "tool_use": - name := contentBlock.Get("name").String() - if origName, ok := reverseMap[name]; ok { - updated, err = sjson.SetBytes(payload, "content_block.name", origName) - if err != nil { - return line - } - } else { - return line - } - case "tool_reference": - toolName := contentBlock.Get("tool_name").String() - if origName, ok := reverseMap[toolName]; ok { - updated, err = sjson.SetBytes(payload, "content_block.tool_name", origName) - if err != nil { - return line - } - } else { - return line - } - default: - return line - } - - trimmed := bytes.TrimSpace(line) - if bytes.HasPrefix(trimmed, []byte("data:")) { - return append([]byte("data: "), updated...) - } - return updated -} - -func applyClaudeToolPrefix(body []byte, prefix string) []byte { - if prefix == "" { - return body - } - - // Collect built-in tool names from the authoritative fallback seed list and - // augment it with any typed built-ins present in the current request body. - builtinTools := helps.AugmentClaudeBuiltinToolRegistry(body, nil) - - if tools := gjson.GetBytes(body, "tools"); tools.Exists() && tools.IsArray() { - tools.ForEach(func(index, tool gjson.Result) bool { - // Skip built-in tools (web_search, code_execution, etc.) which have - // a "type" field and require their name to remain unchanged. - if tool.Get("type").Exists() && tool.Get("type").String() != "" { - if n := tool.Get("name").String(); n != "" { - builtinTools[n] = true - } - return true - } - name := tool.Get("name").String() - if name == "" || strings.HasPrefix(name, prefix) { - return true - } - path := fmt.Sprintf("tools.%d.name", index.Int()) - body, _ = sjson.SetBytes(body, path, prefix+name) - return true - }) - } - - if gjson.GetBytes(body, "tool_choice.type").String() == "tool" { - name := gjson.GetBytes(body, "tool_choice.name").String() - if name != "" && !strings.HasPrefix(name, prefix) && !builtinTools[name] { - body, _ = sjson.SetBytes(body, "tool_choice.name", prefix+name) - } - } - - if messages := gjson.GetBytes(body, "messages"); messages.Exists() && messages.IsArray() { - messages.ForEach(func(msgIndex, msg gjson.Result) bool { - content := msg.Get("content") - if !content.Exists() || !content.IsArray() { - return true - } - content.ForEach(func(contentIndex, part gjson.Result) bool { - partType := part.Get("type").String() - switch partType { - case "tool_use": - name := part.Get("name").String() - if name == "" || strings.HasPrefix(name, prefix) || builtinTools[name] { - return true - } - path := fmt.Sprintf("messages.%d.content.%d.name", msgIndex.Int(), contentIndex.Int()) - body, _ = sjson.SetBytes(body, path, prefix+name) - case "tool_reference": - toolName := part.Get("tool_name").String() - if toolName == "" || strings.HasPrefix(toolName, prefix) || builtinTools[toolName] { - return true - } - path := fmt.Sprintf("messages.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int()) - body, _ = sjson.SetBytes(body, path, prefix+toolName) - case "tool_result": - // Handle nested tool_reference blocks inside tool_result.content[] - nestedContent := part.Get("content") - if nestedContent.Exists() && nestedContent.IsArray() { - nestedContent.ForEach(func(nestedIndex, nestedPart gjson.Result) bool { - if nestedPart.Get("type").String() == "tool_reference" { - nestedToolName := nestedPart.Get("tool_name").String() - if nestedToolName != "" && !strings.HasPrefix(nestedToolName, prefix) && !builtinTools[nestedToolName] { - nestedPath := fmt.Sprintf("messages.%d.content.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int(), nestedIndex.Int()) - body, _ = sjson.SetBytes(body, nestedPath, prefix+nestedToolName) - } - } - return true - }) - } - } - return true - }) - return true - }) - } - - return body -} - -func stripClaudeToolPrefixFromResponse(body []byte, prefix string) []byte { - if prefix == "" { - return body - } - content := gjson.GetBytes(body, "content") - if !content.Exists() || !content.IsArray() { - return body - } - content.ForEach(func(index, part gjson.Result) bool { - partType := part.Get("type").String() - switch partType { - case "tool_use": - name := part.Get("name").String() - if !strings.HasPrefix(name, prefix) { - return true - } - path := fmt.Sprintf("content.%d.name", index.Int()) - body, _ = sjson.SetBytes(body, path, strings.TrimPrefix(name, prefix)) - case "tool_reference": - toolName := part.Get("tool_name").String() - if !strings.HasPrefix(toolName, prefix) { - return true - } - path := fmt.Sprintf("content.%d.tool_name", index.Int()) - body, _ = sjson.SetBytes(body, path, strings.TrimPrefix(toolName, prefix)) - case "tool_result": - // Handle nested tool_reference blocks inside tool_result.content[] - nestedContent := part.Get("content") - if nestedContent.Exists() && nestedContent.IsArray() { - nestedContent.ForEach(func(nestedIndex, nestedPart gjson.Result) bool { - if nestedPart.Get("type").String() == "tool_reference" { - nestedToolName := nestedPart.Get("tool_name").String() - if strings.HasPrefix(nestedToolName, prefix) { - nestedPath := fmt.Sprintf("content.%d.content.%d.tool_name", index.Int(), nestedIndex.Int()) - body, _ = sjson.SetBytes(body, nestedPath, strings.TrimPrefix(nestedToolName, prefix)) - } - } - return true - }) - } - } - return true - }) - return body -} - -func stripClaudeToolPrefixFromStreamLine(line []byte, prefix string) []byte { - if prefix == "" { - return line - } - payload := helps.JSONPayload(line) - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return line - } - contentBlock := gjson.GetBytes(payload, "content_block") - if !contentBlock.Exists() { - return line - } - - blockType := contentBlock.Get("type").String() - var updated []byte - var err error - - switch blockType { - case "tool_use": - name := contentBlock.Get("name").String() - if !strings.HasPrefix(name, prefix) { - return line - } - updated, err = sjson.SetBytes(payload, "content_block.name", strings.TrimPrefix(name, prefix)) - if err != nil { - return line - } - case "tool_reference": - toolName := contentBlock.Get("tool_name").String() - if !strings.HasPrefix(toolName, prefix) { - return line - } - updated, err = sjson.SetBytes(payload, "content_block.tool_name", strings.TrimPrefix(toolName, prefix)) - if err != nil { - return line - } - default: - return line - } - - trimmed := bytes.TrimSpace(line) - if bytes.HasPrefix(trimmed, []byte("data:")) { - return append([]byte("data: "), updated...) - } - return updated -} - -// getClientUserAgent extracts the client User-Agent from the gin context. -func getClientUserAgent(ctx context.Context) string { - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - return ginCtx.GetHeader("User-Agent") - } - return "" -} - -// parseEntrypointFromUA extracts the entrypoint from a Claude Code User-Agent. -// Format: "claude-cli/x.y.z (external, cli)" → "cli" -// Format: "claude-cli/x.y.z (external, vscode)" → "vscode" -// Returns "cli" if parsing fails or UA is not Claude Code. -func parseEntrypointFromUA(userAgent string) string { - // Find content inside parentheses - start := strings.Index(userAgent, "(") - end := strings.LastIndex(userAgent, ")") - if start < 0 || end <= start { - return "cli" - } - inner := userAgent[start+1 : end] - // Split by comma, take the second part (entrypoint is at index 1, after USER_TYPE) - // Format: "(USER_TYPE, ENTRYPOINT[, extra...])" - parts := strings.Split(inner, ",") - if len(parts) >= 2 { - ep := strings.TrimSpace(parts[1]) - if ep != "" { - return ep - } - } - return "cli" -} - -// getWorkloadFromContext extracts workload identifier from the gin request headers. -func getWorkloadFromContext(ctx context.Context) string { - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - return strings.TrimSpace(ginCtx.GetHeader("X-CPA-Claude-Workload")) - } - return "" -} - -// getCloakConfigFromAuth extracts cloak configuration from the auth's attributes, -// falling back to its stored metadata (the raw OAuth/token JSON). Returns -// (cloakMode, strictMode, sensitiveWords, cacheUserID); an empty cloakMode means -// the credential did not explicitly configure a mode. -func getCloakConfigFromAuth(auth *cliproxyauth.Auth) (cloakMode string, strictMode bool, sensitiveWords []string, cacheUserID bool) { - if auth == nil { - return "", false, nil, false - } - - // lookupCloakAttr prefers the executor-facing Attributes, then falls back to the - // raw metadata blob (e.g. the OAuth/token JSON) so file-based credentials can - // carry cloak settings without a matching claude-api-key config entry. - lookupCloakAttr := func(key string) string { - if auth.Attributes != nil { - if value := strings.TrimSpace(auth.Attributes[key]); value != "" { - return value - } - } - if auth.Metadata != nil { - if value, ok := auth.Metadata[key].(string); ok { - return strings.TrimSpace(value) - } - } - return "" - } - - // An empty cloakMode means this credential did not explicitly configure a mode, - // allowing the caller to fall back to the global/default behavior. - cloakMode = lookupCloakAttr("cloak_mode") - - strictMode = strings.EqualFold(lookupCloakAttr("cloak_strict_mode"), "true") - - if wordsStr := lookupCloakAttr("cloak_sensitive_words"); wordsStr != "" { - sensitiveWords = strings.Split(wordsStr, ",") - for i := range sensitiveWords { - sensitiveWords[i] = strings.TrimSpace(sensitiveWords[i]) - } - } - - cacheUserID = strings.EqualFold(lookupCloakAttr("cloak_cache_user_id"), "true") - - return cloakMode, strictMode, sensitiveWords, cacheUserID -} - -// injectFakeUserID generates and injects a fake user ID into the request metadata. -// When useCache is false, a new user ID is generated for every call. -func injectFakeUserID(ctx context.Context, payload []byte, apiKey string, useCache bool) ([]byte, error) { - generateID := func() (string, error) { - if useCache { - return helps.CachedUserIDRequired(ctx, apiKey) - } - return helps.GenerateFakeUserID(), nil - } - - metadata := gjson.GetBytes(payload, "metadata") - if !metadata.Exists() { - userID, errUserID := generateID() - if errUserID != nil { - return nil, errUserID - } - payload, _ = sjson.SetBytes(payload, "metadata.user_id", userID) - return payload, nil - } - - existingUserID := gjson.GetBytes(payload, "metadata.user_id").String() - if existingUserID == "" || !helps.IsValidUserID(existingUserID) { - userID, errUserID := generateID() - if errUserID != nil { - return nil, errUserID - } - payload, _ = sjson.SetBytes(payload, "metadata.user_id", userID) - } - return payload, nil -} - -// fingerprintSalt is the salt used by Claude Code to compute the 3-char build fingerprint. -const fingerprintSalt = "59cf53e54c78" - -// computeFingerprint computes the 3-char build fingerprint that Claude Code embeds in cc_version. -// Algorithm: SHA256(salt + messageText[4] + messageText[7] + messageText[20] + version)[:3] -func computeFingerprint(messageText, version string) string { - indices := [3]int{4, 7, 20} - runes := []rune(messageText) - var sb strings.Builder - for _, idx := range indices { - if idx < len(runes) { - sb.WriteRune(runes[idx]) - } else { - sb.WriteRune('0') - } - } - input := fingerprintSalt + sb.String() + version - h := sha256.Sum256([]byte(input)) - return hex.EncodeToString(h[:])[:3] -} - -// generateBillingHeader creates the x-anthropic-billing-header text block that -// real Claude Code prepends to every system prompt array. -// Format: x-anthropic-billing-header: cc_version=.; cc_entrypoint=; cch=; [cc_workload=;] -func generateBillingHeader(payload []byte, experimentalCCHSigning bool, version, messageText, entrypoint, workload string) string { - if entrypoint == "" { - entrypoint = "cli" - } - buildHash := computeFingerprint(messageText, version) - workloadPart := "" - if workload != "" { - workloadPart = fmt.Sprintf(" cc_workload=%s;", workload) - } - - if experimentalCCHSigning { - return fmt.Sprintf("x-anthropic-billing-header: cc_version=%s.%s; cc_entrypoint=%s; cch=00000;%s", version, buildHash, entrypoint, workloadPart) - } - - // Generate a deterministic cch hash from the payload content (system + messages + tools). - h := sha256.Sum256(payload) - cch := hex.EncodeToString(h[:])[:5] - return fmt.Sprintf("x-anthropic-billing-header: cc_version=%s.%s; cc_entrypoint=%s; cch=%s;%s", version, buildHash, entrypoint, cch, workloadPart) -} - -func checkSystemInstructionsWithMode(payload []byte, strictMode bool) []byte { - return checkSystemInstructionsWithSigningMode(payload, strictMode, false, false, "2.1.63", "", "") -} - -// checkSystemInstructionsWithSigningMode injects Claude Code-style system blocks: -// -// system[0]: billing header (no cache_control) -// system[1]: agent identifier (cache_control ephemeral, scope=org) -// system[2]: core intro prompt (cache_control ephemeral, scope=global) -// system[3]: system instructions (no cache_control) -// system[4]: doing tasks (no cache_control) -// system[5]: user system messages moved to first user message -func checkSystemInstructionsWithSigningMode(payload []byte, strictMode bool, experimentalCCHSigning bool, oauthMode bool, version, entrypoint, workload string) []byte { - system := gjson.GetBytes(payload, "system") - - // Extract original message text for fingerprint computation (before billing injection). - // Use the first system text block's content as the fingerprint source. - messageText := "" - if system.IsArray() { - system.ForEach(func(_, part gjson.Result) bool { - if part.Get("type").String() == "text" { - messageText = part.Get("text").String() - return false - } - return true - }) - } else if system.Type == gjson.String { - messageText = system.String() - } - - // Skip if already injected - firstText := gjson.GetBytes(payload, "system.0.text").String() - if strings.HasPrefix(firstText, "x-anthropic-billing-header:") { - return payload - } - - billingText := generateBillingHeader(payload, experimentalCCHSigning, version, messageText, entrypoint, workload) - billingBlock := buildTextBlock(billingText, nil) - - // Build system blocks matching real Claude Code structure. - // Important: Claude Code's internal cacheScope='org' does NOT serialize to - // scope='org' in the API request. Only scope='global' is sent explicitly. - // The system prompt prefix block is sent without cache_control. - agentBlock := buildTextBlock("You are Claude Code, Anthropic's official CLI for Claude.", nil) - staticPrompt := strings.Join([]string{ - helps.ClaudeCodeIntro, - helps.ClaudeCodeSystem, - helps.ClaudeCodeDoingTasks, - helps.ClaudeCodeToneAndStyle, - helps.ClaudeCodeOutputEfficiency, - }, "\n\n") - staticBlock := buildTextBlock(staticPrompt, nil) - - systemResult := "[" + billingBlock + "," + agentBlock + "," + staticBlock + "]" - payload, _ = sjson.SetRawBytes(payload, "system", []byte(systemResult)) - - // Collect user system instructions and prepend to first user message - if !strictMode { - var userSystemParts []string - if system.IsArray() { - system.ForEach(func(_, part gjson.Result) bool { - if part.Get("type").String() == "text" { - txt := strings.TrimSpace(part.Get("text").String()) - if txt != "" { - userSystemParts = append(userSystemParts, txt) - } - } - return true - }) - } else if system.Type == gjson.String && strings.TrimSpace(system.String()) != "" { - userSystemParts = append(userSystemParts, strings.TrimSpace(system.String())) - } - - if len(userSystemParts) > 0 { - combined := strings.Join(userSystemParts, "\n\n") - if oauthMode { - combined = sanitizeForwardedSystemPrompt(combined) - } - if strings.TrimSpace(combined) != "" { - payload = prependToFirstUserMessage(payload, combined) - } - } - } - - return payload -} - -// sanitizeForwardedSystemPrompt reduces forwarded third-party system context to a -// tiny neutral reminder for Claude OAuth cloaking. The goal is to preserve only -// the minimum tool/task guidance while removing virtually all client-specific -// prompt structure that Anthropic may classify as third-party agent traffic. -func sanitizeForwardedSystemPrompt(text string) string { - if strings.TrimSpace(text) == "" { - return "" - } - return strings.TrimSpace(`Use the available tools when needed to help with software engineering tasks. -Keep responses concise and focused on the user's request. -Prefer acting on the user's task over describing product-specific workflows.`) -} - -// buildTextBlock constructs a JSON text block object with proper escaping. -// Uses sjson.SetBytes to handle multi-line text, quotes, and control characters. -// cacheControl is optional; pass nil to omit cache_control. -func buildTextBlock(text string, cacheControl map[string]string) string { - block := []byte(`{"type":"text"}`) - block, _ = sjson.SetBytes(block, "text", text) - if cacheControl != nil && len(cacheControl) > 0 { - // Build cache_control JSON manually to avoid sjson map marshaling issues. - // sjson.SetBytes with map[string]string may not produce expected structure. - cc := `{"type":"ephemeral"` - if t, ok := cacheControl["ttl"]; ok { - cc += fmt.Sprintf(`,"ttl":"%s"`, t) - } - cc += "}" - block, _ = sjson.SetRawBytes(block, "cache_control", []byte(cc)) - } - return string(block) -} - -// prependToFirstUserMessage prepends text content to the first user message. -// This avoids putting non-Claude-Code system instructions in system[] which -// triggers Anthropic's extra usage billing for OAuth-proxied requests. -func prependToFirstUserMessage(payload []byte, text string) []byte { - messages := gjson.GetBytes(payload, "messages") - if !messages.Exists() || !messages.IsArray() { - return payload - } - - // Find the first user message index - firstUserIdx := -1 - messages.ForEach(func(idx, msg gjson.Result) bool { - if msg.Get("role").String() == "user" { - firstUserIdx = int(idx.Int()) - return false - } - return true - }) - - if firstUserIdx < 0 { - return payload - } - - prefixBlock := fmt.Sprintf(` -As you answer the user's questions, you can use the following context from the system: -%s - -IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task. - -`, text) - - contentPath := fmt.Sprintf("messages.%d.content", firstUserIdx) - content := gjson.GetBytes(payload, contentPath) - - if content.IsArray() { - newBlock := fmt.Sprintf(`{"type":"text","text":%q}`, prefixBlock) - var newArray string - if content.Raw == "[]" || content.Raw == "" { - newArray = "[" + newBlock + "]" - } else { - newArray = "[" + newBlock + "," + content.Raw[1:] - } - payload, _ = sjson.SetRawBytes(payload, contentPath, []byte(newArray)) - } else if content.Type == gjson.String { - newText := prefixBlock + content.String() - payload, _ = sjson.SetBytes(payload, contentPath, newText) - } - - return payload -} - -// applyCloaking applies cloaking transformations to the payload based on config and client. -// Cloaking includes: system prompt injection, fake user ID, and sensitive word obfuscation. -func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, payload []byte, model string, apiKey string) ([]byte, error) { - clientUserAgent := getClientUserAgent(ctx) - // Enable cch signing for OAuth tokens by default (not just experimental flag). - oauthToken := isClaudeOAuthToken(apiKey) - useCCHSigning := oauthToken || experimentalCCHSigningEnabled(cfg, auth) - - // Get cloak config from ClaudeKey configuration - cloakCfg := resolveClaudeKeyCloakConfig(cfg, auth) - attrMode, attrStrict, attrWords, attrCache := getCloakConfigFromAuth(auth) - - // Determine cloak settings. Precedence (low -> high): - // built-in "auto" default - // -> global disable-claude-cloak-mode switch (forces "never") - // -> per-credential settings from auth attributes/metadata - // -> per claude-api-key cloak config - cloakMode := "auto" - if cfg != nil && cfg.DisableClaudeCloakMode { - cloakMode = "never" - } - strictMode := attrStrict - sensitiveWords := attrWords - cacheUserID := attrCache - - if attrMode != "" { - cloakMode = attrMode - } - - if cloakCfg != nil { - if mode := strings.TrimSpace(cloakCfg.Mode); mode != "" { - cloakMode = mode - } - if cloakCfg.StrictMode { - strictMode = true - } - if len(cloakCfg.SensitiveWords) > 0 { - sensitiveWords = cloakCfg.SensitiveWords - } - if cloakCfg.CacheUserID != nil { - cacheUserID = *cloakCfg.CacheUserID - } - } - - // Determine if cloaking should be applied - if !helps.ShouldCloak(cloakMode, clientUserAgent) { - return payload, nil - } - - // Skip system instructions for claude-3-5-haiku models - if !strings.HasPrefix(model, "claude-3-5-haiku") { - billingVersion := helps.DefaultClaudeVersion(cfg) - entrypoint := parseEntrypointFromUA(clientUserAgent) - workload := getWorkloadFromContext(ctx) - payload = checkSystemInstructionsWithSigningMode(payload, strictMode, useCCHSigning, oauthToken, billingVersion, entrypoint, workload) - } - - // Inject fake user ID - var errFakeUserID error - payload, errFakeUserID = injectFakeUserID(ctx, payload, apiKey, cacheUserID) - if errFakeUserID != nil { - return nil, errFakeUserID - } - - // Apply sensitive word obfuscation - if len(sensitiveWords) > 0 { - matcher := helps.BuildSensitiveWordMatcher(sensitiveWords) - payload = helps.ObfuscateSensitiveWords(payload, matcher) - } - - return payload, nil -} - -// ensureCacheControl injects cache_control breakpoints into the payload for optimal prompt caching. -// According to Anthropic's documentation, cache prefixes are created in order: tools -> system -> messages. -// This function adds cache_control to: -// 1. The LAST non-deferred tool in the tools array (caches all preceding tool definitions) -// 2. The LAST system prompt element -// 3. The SECOND-TO-LAST user turn (caches conversation history for multi-turn) -// -// Up to 4 cache breakpoints are allowed per request. Tools, System, and Messages are INDEPENDENT breakpoints. -// This enables up to 90% cost reduction on cached tokens (cache read = 0.1x base price). -// See: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching -func ensureCacheControl(payload []byte) []byte { - // 1. Inject cache_control into the LAST non-deferred tool - // Tools are cached first in the hierarchy, so this is the most important breakpoint. - payload = injectToolsCacheControl(payload) - - // 2. Inject cache_control into the LAST system prompt element - // System is the second level in the cache hierarchy. - payload = injectSystemCacheControl(payload) - - // 3. Inject cache_control into messages for multi-turn conversation caching - // This caches the conversation history up to the second-to-last user turn. - payload = injectMessagesCacheControl(payload) - - return payload -} - -func countCacheControls(payload []byte) int { - count := 0 - - // Check system - system := gjson.GetBytes(payload, "system") - if system.IsArray() { - system.ForEach(func(_, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - count++ - } - return true - }) - } - - // Check tools - tools := gjson.GetBytes(payload, "tools") - if tools.IsArray() { - tools.ForEach(func(_, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - count++ - } - return true - }) - } - - // Check messages - messages := gjson.GetBytes(payload, "messages") - if messages.IsArray() { - messages.ForEach(func(_, msg gjson.Result) bool { - content := msg.Get("content") - if content.IsArray() { - content.ForEach(func(_, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - count++ - } - return true - }) - } - return true - }) - } - - return count -} - -// normalizeCacheControlTTL ensures cache_control TTL values don't violate the -// prompt-caching-scope-2026-01-05 ordering constraint: a 1h-TTL block must not -// appear after a 5m-TTL block anywhere in the evaluation order. -// -// Anthropic evaluates blocks in order: tools → system (index 0..N) → messages. -// Within each section, blocks are evaluated in array order. A 5m (default) block -// followed by a 1h block at ANY later position is an error — including within -// the same section (e.g. system[1]=5m then system[3]=1h). -// -// Strategy: walk all cache_control blocks in evaluation order. Once a 5m block -// is seen, strip ttl from ALL subsequent 1h blocks (downgrading them to 5m). -func normalizeCacheControlTTL(payload []byte) []byte { - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return payload - } - - original := payload - seen5m := false - modified := false - - processBlock := func(path string, obj gjson.Result) { - cc := obj.Get("cache_control") - if !cc.Exists() { - return - } - if !cc.IsObject() { - seen5m = true - return - } - ttl := cc.Get("ttl") - if ttl.Type != gjson.String || ttl.String() != "1h" { - seen5m = true - return - } - if !seen5m { - return - } - ttlPath := path + ".cache_control.ttl" - updated, errDel := sjson.DeleteBytes(payload, ttlPath) - if errDel != nil { - return - } - payload = updated - modified = true - } - - tools := gjson.GetBytes(payload, "tools") - if tools.IsArray() { - tools.ForEach(func(idx, item gjson.Result) bool { - processBlock(fmt.Sprintf("tools.%d", int(idx.Int())), item) - return true - }) - } - - system := gjson.GetBytes(payload, "system") - if system.IsArray() { - system.ForEach(func(idx, item gjson.Result) bool { - processBlock(fmt.Sprintf("system.%d", int(idx.Int())), item) - return true - }) - } - - messages := gjson.GetBytes(payload, "messages") - if messages.IsArray() { - messages.ForEach(func(msgIdx, msg gjson.Result) bool { - content := msg.Get("content") - if !content.IsArray() { - return true - } - content.ForEach(func(itemIdx, item gjson.Result) bool { - processBlock(fmt.Sprintf("messages.%d.content.%d", int(msgIdx.Int()), int(itemIdx.Int())), item) - return true - }) - return true - }) - } - - if !modified { - return original - } - return payload -} - -// enforceCacheControlLimit removes excess cache_control blocks from a payload -// so the total does not exceed the Anthropic API limit (currently 4). -// -// Anthropic evaluates cache breakpoints in order: tools → system → messages. -// The most valuable breakpoints are: -// 1. Last tool — caches ALL tool definitions -// 2. Last system block — caches ALL system content -// 3. Recent messages — cache conversation context -// -// Removal priority (strip lowest-value first): -// -// Phase 1: system blocks earliest-first, preserving the last one. -// Phase 2: tool blocks earliest-first, preserving the last one. -// Phase 3: message content blocks earliest-first. -// Phase 4: remaining system blocks (last system). -// Phase 5: remaining tool blocks (last tool). -func enforceCacheControlLimit(payload []byte, maxBlocks int) []byte { - if len(payload) == 0 || !gjson.ValidBytes(payload) { - return payload - } - - total := countCacheControls(payload) - if total <= maxBlocks { - return payload - } - - excess := total - maxBlocks - - system := gjson.GetBytes(payload, "system") - if system.IsArray() { - lastIdx := -1 - system.ForEach(func(idx, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - lastIdx = int(idx.Int()) - } - return true - }) - if lastIdx >= 0 { - system.ForEach(func(idx, item gjson.Result) bool { - if excess <= 0 { - return false - } - i := int(idx.Int()) - if i == lastIdx { - return true - } - if !item.Get("cache_control").Exists() { - return true - } - path := fmt.Sprintf("system.%d.cache_control", i) - updated, errDel := sjson.DeleteBytes(payload, path) - if errDel != nil { - return true - } - payload = updated - excess-- - return true - }) - } - } - if excess <= 0 { - return payload - } - - tools := gjson.GetBytes(payload, "tools") - if tools.IsArray() { - lastIdx := -1 - tools.ForEach(func(idx, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - lastIdx = int(idx.Int()) - } - return true - }) - if lastIdx >= 0 { - tools.ForEach(func(idx, item gjson.Result) bool { - if excess <= 0 { - return false - } - i := int(idx.Int()) - if i == lastIdx { - return true - } - if !item.Get("cache_control").Exists() { - return true - } - path := fmt.Sprintf("tools.%d.cache_control", i) - updated, errDel := sjson.DeleteBytes(payload, path) - if errDel != nil { - return true - } - payload = updated - excess-- - return true - }) - } - } - if excess <= 0 { - return payload - } - - messages := gjson.GetBytes(payload, "messages") - if messages.IsArray() { - messages.ForEach(func(msgIdx, msg gjson.Result) bool { - if excess <= 0 { - return false - } - content := msg.Get("content") - if !content.IsArray() { - return true - } - content.ForEach(func(itemIdx, item gjson.Result) bool { - if excess <= 0 { - return false - } - if !item.Get("cache_control").Exists() { - return true - } - path := fmt.Sprintf("messages.%d.content.%d.cache_control", int(msgIdx.Int()), int(itemIdx.Int())) - updated, errDel := sjson.DeleteBytes(payload, path) - if errDel != nil { - return true - } - payload = updated - excess-- - return true - }) - return true - }) - } - if excess <= 0 { - return payload - } - - system = gjson.GetBytes(payload, "system") - if system.IsArray() { - system.ForEach(func(idx, item gjson.Result) bool { - if excess <= 0 { - return false - } - if !item.Get("cache_control").Exists() { - return true - } - path := fmt.Sprintf("system.%d.cache_control", int(idx.Int())) - updated, errDel := sjson.DeleteBytes(payload, path) - if errDel != nil { - return true - } - payload = updated - excess-- - return true - }) - } - if excess <= 0 { - return payload - } - - tools = gjson.GetBytes(payload, "tools") - if tools.IsArray() { - tools.ForEach(func(idx, item gjson.Result) bool { - if excess <= 0 { - return false - } - if !item.Get("cache_control").Exists() { - return true - } - path := fmt.Sprintf("tools.%d.cache_control", int(idx.Int())) - updated, errDel := sjson.DeleteBytes(payload, path) - if errDel != nil { - return true - } - payload = updated - excess-- - return true - }) - } - - return payload -} - -// injectMessagesCacheControl adds cache_control to the second-to-last user turn for multi-turn caching. -// Per Anthropic docs: "Place cache_control on the second-to-last User message to let the model reuse the earlier cache." -// This enables caching of conversation history, which is especially beneficial for long multi-turn conversations. -// Only adds cache_control if: -// - There are at least 2 user turns in the conversation -// - No message content already has cache_control -func injectMessagesCacheControl(payload []byte) []byte { - messages := gjson.GetBytes(payload, "messages") - if !messages.Exists() || !messages.IsArray() { - return payload - } - - // Check if ANY message content already has cache_control - hasCacheControlInMessages := false - messages.ForEach(func(_, msg gjson.Result) bool { - content := msg.Get("content") - if content.IsArray() { - content.ForEach(func(_, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - hasCacheControlInMessages = true - return false - } - return true - }) - } - return !hasCacheControlInMessages - }) - if hasCacheControlInMessages { - return payload - } - - // Find all user message indices - var userMsgIndices []int - messages.ForEach(func(index gjson.Result, msg gjson.Result) bool { - if msg.Get("role").String() == "user" { - userMsgIndices = append(userMsgIndices, int(index.Int())) - } - return true - }) - - // Need at least 2 user turns to cache the second-to-last - if len(userMsgIndices) < 2 { - return payload - } - - // Get the second-to-last user message index - secondToLastUserIdx := userMsgIndices[len(userMsgIndices)-2] - - // Get the content of this message - contentPath := fmt.Sprintf("messages.%d.content", secondToLastUserIdx) - content := gjson.GetBytes(payload, contentPath) - - if content.IsArray() { - // Add cache_control to the last content block of this message - contentCount := int(content.Get("#").Int()) - if contentCount > 0 { - cacheControlPath := fmt.Sprintf("messages.%d.content.%d.cache_control", secondToLastUserIdx, contentCount-1) - result, err := sjson.SetBytes(payload, cacheControlPath, map[string]string{"type": "ephemeral"}) - if err != nil { - log.Warnf("failed to inject cache_control into messages: %v", err) - return payload - } - payload = result - } - } else if content.Type == gjson.String { - // Convert string content to array with cache_control - text := content.String() - newContent := []map[string]interface{}{ - { - "type": "text", - "text": text, - "cache_control": map[string]string{ - "type": "ephemeral", - }, - }, - } - result, err := sjson.SetBytes(payload, contentPath, newContent) - if err != nil { - log.Warnf("failed to inject cache_control into message string content: %v", err) - return payload - } - payload = result - } - - return payload -} - -// injectToolsCacheControl adds cache_control to the last non-deferred tool in the tools array. -// Deferred tools cannot use prompt caching, so trailing deferred tools are skipped. -// This only adds cache_control if NO tool in the array already has it. -func injectToolsCacheControl(payload []byte) []byte { - tools := gjson.GetBytes(payload, "tools") - if !tools.Exists() || !tools.IsArray() { - return payload - } - - // Check if ANY tool already has cache_control and find the last eligible tool. - hasCacheControlInTools := false - lastEligibleToolIndex := -1 - tools.ForEach(func(index, tool gjson.Result) bool { - if tool.Get("cache_control").Exists() { - hasCacheControlInTools = true - return false - } - if !tool.Get("defer_loading").Bool() { - lastEligibleToolIndex = int(index.Int()) - } - return true - }) - if hasCacheControlInTools || lastEligibleToolIndex < 0 { - return payload - } - - lastToolPath := fmt.Sprintf("tools.%d.cache_control", lastEligibleToolIndex) - result, err := sjson.SetBytes(payload, lastToolPath, map[string]string{"type": "ephemeral"}) - if err != nil { - log.Warnf("failed to inject cache_control into tools array: %v", err) - return payload - } - - return result -} - -// injectSystemCacheControl adds cache_control to the last element in the system prompt. -// Converts string system prompts to array format if needed. -// This only adds cache_control if NO system element already has it. -func injectSystemCacheControl(payload []byte) []byte { - system := gjson.GetBytes(payload, "system") - if !system.Exists() { - return payload - } - - if system.IsArray() { - count := int(system.Get("#").Int()) - if count == 0 { - return payload - } - - // Check if ANY system element already has cache_control - hasCacheControlInSystem := false - system.ForEach(func(_, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - hasCacheControlInSystem = true - return false - } - return true - }) - if hasCacheControlInSystem { - return payload - } - - // Add cache_control to the last system element - lastSystemPath := fmt.Sprintf("system.%d.cache_control", count-1) - result, err := sjson.SetBytes(payload, lastSystemPath, map[string]string{"type": "ephemeral"}) - if err != nil { - log.Warnf("failed to inject cache_control into system array: %v", err) - return payload - } - payload = result - } else if system.Type == gjson.String { - // Convert string system prompt to array with cache_control - // "system": "text" -> "system": [{"type": "text", "text": "text", "cache_control": {"type": "ephemeral"}}] - text := system.String() - newSystem := []map[string]interface{}{ - { - "type": "text", - "text": text, - "cache_control": map[string]string{ - "type": "ephemeral", - }, - }, - } - result, err := sjson.SetBytes(payload, "system", newSystem) - if err != nil { - log.Warnf("failed to inject cache_control into system string: %v", err) - return payload - } - payload = result - } - - return payload -} - -func ensureModelMaxTokens(body []byte, modelID string) []byte { - if len(body) == 0 || !gjson.ValidBytes(body) { - return body - } - - if maxTokens := gjson.GetBytes(body, "max_tokens"); maxTokens.Exists() { - return body - } - - for _, provider := range registry.GetGlobalRegistry().GetModelProviders(strings.TrimSpace(modelID)) { - if strings.EqualFold(provider, "claude") { - maxTokens := defaultModelMaxTokens - if info := registry.GetGlobalRegistry().GetModelInfo(strings.TrimSpace(modelID), "claude"); info != nil && info.MaxCompletionTokens > 0 { - maxTokens = info.MaxCompletionTokens - } - body, _ = sjson.SetBytes(body, "max_tokens", maxTokens) - return body - } - } - - return body -} diff --git a/internal/runtime/executor/claude_executor_auth.go b/internal/runtime/executor/claude_executor_auth.go new file mode 100644 index 000000000..679cd4de8 --- /dev/null +++ b/internal/runtime/executor/claude_executor_auth.go @@ -0,0 +1,49 @@ +package executor + +import ( + "context" + "fmt" + "time" + + claudeauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/claude" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" +) + +func (e *ClaudeExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { + log.Debugf("claude executor: refresh called") + if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { + return refreshed, err + } + if auth == nil { + return nil, fmt.Errorf("claude executor: auth is nil") + } + var refreshToken string + if auth.Metadata != nil { + if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" { + refreshToken = v + } + } + if refreshToken == "" { + return auth, nil + } + svc := claudeauth.NewClaudeAuthWithProxyURL(e.cfg, auth.ProxyURL) + td, err := svc.RefreshTokensWithRetry(ctx, refreshToken, 3) + if err != nil { + return nil, err + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["access_token"] = td.AccessToken + if td.RefreshToken != "" { + auth.Metadata["refresh_token"] = td.RefreshToken + } + auth.Metadata["email"] = td.Email + auth.Metadata["expired"] = td.Expire + auth.Metadata["type"] = "claude" + now := time.Now().Format(time.RFC3339) + auth.Metadata["last_refresh"] = now + return auth, nil +} diff --git a/internal/runtime/executor/claude_executor_cloaking.go b/internal/runtime/executor/claude_executor_cloaking.go new file mode 100644 index 000000000..071f069ff --- /dev/null +++ b/internal/runtime/executor/claude_executor_cloaking.go @@ -0,0 +1,960 @@ +package executor + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + + "github.com/gin-gonic/gin" +) + +// getClientUserAgent extracts the client User-Agent from the gin context. +func getClientUserAgent(ctx context.Context) string { + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + return ginCtx.GetHeader("User-Agent") + } + return "" +} + +// parseEntrypointFromUA extracts the entrypoint from a Claude Code User-Agent. +// Format: "claude-cli/x.y.z (external, cli)" → "cli" +// Format: "claude-cli/x.y.z (external, vscode)" → "vscode" +// Returns "cli" if parsing fails or UA is not Claude Code. +func parseEntrypointFromUA(userAgent string) string { + // Find content inside parentheses + start := strings.Index(userAgent, "(") + end := strings.LastIndex(userAgent, ")") + if start < 0 || end <= start { + return "cli" + } + inner := userAgent[start+1 : end] + // Split by comma, take the second part (entrypoint is at index 1, after USER_TYPE) + // Format: "(USER_TYPE, ENTRYPOINT[, extra...])" + parts := strings.Split(inner, ",") + if len(parts) >= 2 { + ep := strings.TrimSpace(parts[1]) + if ep != "" { + return ep + } + } + return "cli" +} + +// getWorkloadFromContext extracts workload identifier from the gin request headers. +func getWorkloadFromContext(ctx context.Context) string { + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + return strings.TrimSpace(ginCtx.GetHeader("X-CPA-Claude-Workload")) + } + return "" +} + +// getCloakConfigFromAuth extracts cloak configuration from the auth's attributes, +// falling back to its stored metadata (the raw OAuth/token JSON). Returns +// (cloakMode, strictMode, sensitiveWords, cacheUserID); an empty cloakMode means +// the credential did not explicitly configure a mode. +func getCloakConfigFromAuth(auth *cliproxyauth.Auth) (cloakMode string, strictMode bool, sensitiveWords []string, cacheUserID bool) { + if auth == nil { + return "", false, nil, false + } + + // lookupCloakAttr prefers the executor-facing Attributes, then falls back to the + // raw metadata blob (e.g. the OAuth/token JSON) so file-based credentials can + // carry cloak settings without a matching claude-api-key config entry. + lookupCloakAttr := func(key string) string { + if auth.Attributes != nil { + if value := strings.TrimSpace(auth.Attributes[key]); value != "" { + return value + } + } + if auth.Metadata != nil { + if value, ok := auth.Metadata[key].(string); ok { + return strings.TrimSpace(value) + } + } + return "" + } + + // An empty cloakMode means this credential did not explicitly configure a mode, + // allowing the caller to fall back to the global/default behavior. + cloakMode = lookupCloakAttr("cloak_mode") + + strictMode = strings.EqualFold(lookupCloakAttr("cloak_strict_mode"), "true") + + if wordsStr := lookupCloakAttr("cloak_sensitive_words"); wordsStr != "" { + sensitiveWords = strings.Split(wordsStr, ",") + for i := range sensitiveWords { + sensitiveWords[i] = strings.TrimSpace(sensitiveWords[i]) + } + } + + cacheUserID = strings.EqualFold(lookupCloakAttr("cloak_cache_user_id"), "true") + + return cloakMode, strictMode, sensitiveWords, cacheUserID +} + +// injectFakeUserID generates and injects a fake user ID into the request metadata. +// When useCache is false, a new user ID is generated for every call. +func injectFakeUserID(ctx context.Context, payload []byte, apiKey string, useCache bool) ([]byte, error) { + generateID := func() (string, error) { + if useCache { + return helps.CachedUserIDRequired(ctx, apiKey) + } + return helps.GenerateFakeUserID(), nil + } + + metadata := gjson.GetBytes(payload, "metadata") + if !metadata.Exists() { + userID, errUserID := generateID() + if errUserID != nil { + return nil, errUserID + } + payload, _ = sjson.SetBytes(payload, "metadata.user_id", userID) + return payload, nil + } + + existingUserID := gjson.GetBytes(payload, "metadata.user_id").String() + if existingUserID == "" || !helps.IsValidUserID(existingUserID) { + userID, errUserID := generateID() + if errUserID != nil { + return nil, errUserID + } + payload, _ = sjson.SetBytes(payload, "metadata.user_id", userID) + } + return payload, nil +} + +// fingerprintSalt is the salt used by Claude Code to compute the 3-char build fingerprint. +const fingerprintSalt = "59cf53e54c78" + +// computeFingerprint computes the 3-char build fingerprint that Claude Code embeds in cc_version. +// Algorithm: SHA256(salt + messageText[4] + messageText[7] + messageText[20] + version)[:3] +func computeFingerprint(messageText, version string) string { + indices := [3]int{4, 7, 20} + runes := []rune(messageText) + var sb strings.Builder + for _, idx := range indices { + if idx < len(runes) { + sb.WriteRune(runes[idx]) + } else { + sb.WriteRune('0') + } + } + input := fingerprintSalt + sb.String() + version + h := sha256.Sum256([]byte(input)) + return hex.EncodeToString(h[:])[:3] +} + +// generateBillingHeader creates the x-anthropic-billing-header text block that +// real Claude Code prepends to every system prompt array. +// Format: x-anthropic-billing-header: cc_version=.; cc_entrypoint=; cch=; [cc_workload=;] +func generateBillingHeader(payload []byte, experimentalCCHSigning bool, version, messageText, entrypoint, workload string) string { + if entrypoint == "" { + entrypoint = "cli" + } + buildHash := computeFingerprint(messageText, version) + workloadPart := "" + if workload != "" { + workloadPart = fmt.Sprintf(" cc_workload=%s;", workload) + } + + if experimentalCCHSigning { + return fmt.Sprintf("x-anthropic-billing-header: cc_version=%s.%s; cc_entrypoint=%s; cch=00000;%s", version, buildHash, entrypoint, workloadPart) + } + + // Generate a deterministic cch hash from the payload content (system + messages + tools). + h := sha256.Sum256(payload) + cch := hex.EncodeToString(h[:])[:5] + return fmt.Sprintf("x-anthropic-billing-header: cc_version=%s.%s; cc_entrypoint=%s; cch=%s;%s", version, buildHash, entrypoint, cch, workloadPart) +} + +func checkSystemInstructionsWithMode(payload []byte, strictMode bool) []byte { + return checkSystemInstructionsWithSigningMode(payload, strictMode, false, false, "2.1.63", "", "") +} + +// checkSystemInstructionsWithSigningMode injects Claude Code-style system blocks: +// +// system[0]: billing header (no cache_control) +// system[1]: agent identifier (cache_control ephemeral, scope=org) +// system[2]: core intro prompt (cache_control ephemeral, scope=global) +// system[3]: system instructions (no cache_control) +// system[4]: doing tasks (no cache_control) +// system[5]: user system messages moved to first user message +func checkSystemInstructionsWithSigningMode(payload []byte, strictMode bool, experimentalCCHSigning bool, oauthMode bool, version, entrypoint, workload string) []byte { + system := gjson.GetBytes(payload, "system") + + // Extract original message text for fingerprint computation (before billing injection). + // Use the first system text block's content as the fingerprint source. + messageText := "" + if system.IsArray() { + system.ForEach(func(_, part gjson.Result) bool { + if part.Get("type").String() == "text" { + messageText = part.Get("text").String() + return false + } + return true + }) + } else if system.Type == gjson.String { + messageText = system.String() + } + + // Skip if already injected + firstText := gjson.GetBytes(payload, "system.0.text").String() + if strings.HasPrefix(firstText, "x-anthropic-billing-header:") { + return payload + } + + billingText := generateBillingHeader(payload, experimentalCCHSigning, version, messageText, entrypoint, workload) + billingBlock := buildTextBlock(billingText, nil) + + // Build system blocks matching real Claude Code structure. + // Important: Claude Code's internal cacheScope='org' does NOT serialize to + // scope='org' in the API request. Only scope='global' is sent explicitly. + // The system prompt prefix block is sent without cache_control. + agentBlock := buildTextBlock("You are Claude Code, Anthropic's official CLI for Claude.", nil) + staticPrompt := strings.Join([]string{ + helps.ClaudeCodeIntro, + helps.ClaudeCodeSystem, + helps.ClaudeCodeDoingTasks, + helps.ClaudeCodeToneAndStyle, + helps.ClaudeCodeOutputEfficiency, + }, "\n\n") + staticBlock := buildTextBlock(staticPrompt, nil) + + systemResult := "[" + billingBlock + "," + agentBlock + "," + staticBlock + "]" + payload, _ = sjson.SetRawBytes(payload, "system", []byte(systemResult)) + + // Collect user system instructions and prepend to first user message + if !strictMode { + var userSystemParts []string + if system.IsArray() { + system.ForEach(func(_, part gjson.Result) bool { + if part.Get("type").String() == "text" { + txt := strings.TrimSpace(part.Get("text").String()) + if txt != "" { + userSystemParts = append(userSystemParts, txt) + } + } + return true + }) + } else if system.Type == gjson.String && strings.TrimSpace(system.String()) != "" { + userSystemParts = append(userSystemParts, strings.TrimSpace(system.String())) + } + + if len(userSystemParts) > 0 { + combined := strings.Join(userSystemParts, "\n\n") + if oauthMode { + combined = sanitizeForwardedSystemPrompt(combined) + } + if strings.TrimSpace(combined) != "" { + payload = prependToFirstUserMessage(payload, combined) + } + } + } + + return payload +} + +// sanitizeForwardedSystemPrompt reduces forwarded third-party system context to a +// tiny neutral reminder for Claude OAuth cloaking. The goal is to preserve only +// the minimum tool/task guidance while removing virtually all client-specific +// prompt structure that Anthropic may classify as third-party agent traffic. +func sanitizeForwardedSystemPrompt(text string) string { + if strings.TrimSpace(text) == "" { + return "" + } + return strings.TrimSpace(`Use the available tools when needed to help with software engineering tasks. +Keep responses concise and focused on the user's request. +Prefer acting on the user's task over describing product-specific workflows.`) +} + +// buildTextBlock constructs a JSON text block object with proper escaping. +// Uses sjson.SetBytes to handle multi-line text, quotes, and control characters. +// cacheControl is optional; pass nil to omit cache_control. +func buildTextBlock(text string, cacheControl map[string]string) string { + block := []byte(`{"type":"text"}`) + block, _ = sjson.SetBytes(block, "text", text) + if cacheControl != nil && len(cacheControl) > 0 { + // Build cache_control JSON manually to avoid sjson map marshaling issues. + // sjson.SetBytes with map[string]string may not produce expected structure. + cc := `{"type":"ephemeral"` + if t, ok := cacheControl["ttl"]; ok { + cc += fmt.Sprintf(`,"ttl":"%s"`, t) + } + cc += "}" + block, _ = sjson.SetRawBytes(block, "cache_control", []byte(cc)) + } + return string(block) +} + +// prependToFirstUserMessage prepends text content to the first user message. +// This avoids putting non-Claude-Code system instructions in system[] which +// triggers Anthropic's extra usage billing for OAuth-proxied requests. +func prependToFirstUserMessage(payload []byte, text string) []byte { + messages := gjson.GetBytes(payload, "messages") + if !messages.Exists() || !messages.IsArray() { + return payload + } + + // Find the first user message index + firstUserIdx := -1 + messages.ForEach(func(idx, msg gjson.Result) bool { + if msg.Get("role").String() == "user" { + firstUserIdx = int(idx.Int()) + return false + } + return true + }) + + if firstUserIdx < 0 { + return payload + } + + prefixBlock := fmt.Sprintf(` +As you answer the user's questions, you can use the following context from the system: +%s + +IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task. + +`, text) + + contentPath := fmt.Sprintf("messages.%d.content", firstUserIdx) + content := gjson.GetBytes(payload, contentPath) + + if content.IsArray() { + newBlock := fmt.Sprintf(`{"type":"text","text":%q}`, prefixBlock) + var newArray string + if content.Raw == "[]" || content.Raw == "" { + newArray = "[" + newBlock + "]" + } else { + newArray = "[" + newBlock + "," + content.Raw[1:] + } + payload, _ = sjson.SetRawBytes(payload, contentPath, []byte(newArray)) + } else if content.Type == gjson.String { + newText := prefixBlock + content.String() + payload, _ = sjson.SetBytes(payload, contentPath, newText) + } + + return payload +} + +// applyCloaking applies cloaking transformations to the payload based on config and client. +// Cloaking includes: system prompt injection, fake user ID, and sensitive word obfuscation. +func applyCloaking(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, payload []byte, model string, apiKey string) ([]byte, error) { + clientUserAgent := getClientUserAgent(ctx) + // Enable cch signing for OAuth tokens by default (not just experimental flag). + oauthToken := isClaudeOAuthToken(apiKey) + useCCHSigning := oauthToken || experimentalCCHSigningEnabled(cfg, auth) + + // Get cloak config from ClaudeKey configuration + cloakCfg := resolveClaudeKeyCloakConfig(cfg, auth) + attrMode, attrStrict, attrWords, attrCache := getCloakConfigFromAuth(auth) + + // Determine cloak settings. Precedence (low -> high): + // built-in "auto" default + // -> global disable-claude-cloak-mode switch (forces "never") + // -> per-credential settings from auth attributes/metadata + // -> per claude-api-key cloak config + cloakMode := "auto" + if cfg != nil && cfg.DisableClaudeCloakMode { + cloakMode = "never" + } + strictMode := attrStrict + sensitiveWords := attrWords + cacheUserID := attrCache + + if attrMode != "" { + cloakMode = attrMode + } + + if cloakCfg != nil { + if mode := strings.TrimSpace(cloakCfg.Mode); mode != "" { + cloakMode = mode + } + if cloakCfg.StrictMode { + strictMode = true + } + if len(cloakCfg.SensitiveWords) > 0 { + sensitiveWords = cloakCfg.SensitiveWords + } + if cloakCfg.CacheUserID != nil { + cacheUserID = *cloakCfg.CacheUserID + } + } + + // Determine if cloaking should be applied + if !helps.ShouldCloak(cloakMode, clientUserAgent) { + return payload, nil + } + + // Skip system instructions for claude-3-5-haiku models + if !strings.HasPrefix(model, "claude-3-5-haiku") { + billingVersion := helps.DefaultClaudeVersion(cfg) + entrypoint := parseEntrypointFromUA(clientUserAgent) + workload := getWorkloadFromContext(ctx) + payload = checkSystemInstructionsWithSigningMode(payload, strictMode, useCCHSigning, oauthToken, billingVersion, entrypoint, workload) + } + + // Inject fake user ID + var errFakeUserID error + payload, errFakeUserID = injectFakeUserID(ctx, payload, apiKey, cacheUserID) + if errFakeUserID != nil { + return nil, errFakeUserID + } + + // Apply sensitive word obfuscation + if len(sensitiveWords) > 0 { + matcher := helps.BuildSensitiveWordMatcher(sensitiveWords) + payload = helps.ObfuscateSensitiveWords(payload, matcher) + } + + return payload, nil +} + +// ensureCacheControl injects cache_control breakpoints into the payload for optimal prompt caching. +// According to Anthropic's documentation, cache prefixes are created in order: tools -> system -> messages. +// This function adds cache_control to: +// 1. The LAST non-deferred tool in the tools array (caches all preceding tool definitions) +// 2. The LAST system prompt element +// 3. The SECOND-TO-LAST user turn (caches conversation history for multi-turn) +// +// Up to 4 cache breakpoints are allowed per request. Tools, System, and Messages are INDEPENDENT breakpoints. +// This enables up to 90% cost reduction on cached tokens (cache read = 0.1x base price). +// See: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching +func ensureCacheControl(payload []byte) []byte { + // 1. Inject cache_control into the LAST non-deferred tool + // Tools are cached first in the hierarchy, so this is the most important breakpoint. + payload = injectToolsCacheControl(payload) + + // 2. Inject cache_control into the LAST system prompt element + // System is the second level in the cache hierarchy. + payload = injectSystemCacheControl(payload) + + // 3. Inject cache_control into messages for multi-turn conversation caching + // This caches the conversation history up to the second-to-last user turn. + payload = injectMessagesCacheControl(payload) + + return payload +} + +func countCacheControls(payload []byte) int { + count := 0 + + // Check system + system := gjson.GetBytes(payload, "system") + if system.IsArray() { + system.ForEach(func(_, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + count++ + } + return true + }) + } + + // Check tools + tools := gjson.GetBytes(payload, "tools") + if tools.IsArray() { + tools.ForEach(func(_, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + count++ + } + return true + }) + } + + // Check messages + messages := gjson.GetBytes(payload, "messages") + if messages.IsArray() { + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if content.IsArray() { + content.ForEach(func(_, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + count++ + } + return true + }) + } + return true + }) + } + + return count +} + +// normalizeCacheControlTTL ensures cache_control TTL values don't violate the +// prompt-caching-scope-2026-01-05 ordering constraint: a 1h-TTL block must not +// appear after a 5m-TTL block anywhere in the evaluation order. +// +// Anthropic evaluates blocks in order: tools → system (index 0..N) → messages. +// Within each section, blocks are evaluated in array order. A 5m (default) block +// followed by a 1h block at ANY later position is an error — including within +// the same section (e.g. system[1]=5m then system[3]=1h). +// +// Strategy: walk all cache_control blocks in evaluation order. Once a 5m block +// is seen, strip ttl from ALL subsequent 1h blocks (downgrading them to 5m). +func normalizeCacheControlTTL(payload []byte) []byte { + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return payload + } + + original := payload + seen5m := false + modified := false + + processBlock := func(path string, obj gjson.Result) { + cc := obj.Get("cache_control") + if !cc.Exists() { + return + } + if !cc.IsObject() { + seen5m = true + return + } + ttl := cc.Get("ttl") + if ttl.Type != gjson.String || ttl.String() != "1h" { + seen5m = true + return + } + if !seen5m { + return + } + ttlPath := path + ".cache_control.ttl" + updated, errDel := sjson.DeleteBytes(payload, ttlPath) + if errDel != nil { + return + } + payload = updated + modified = true + } + + tools := gjson.GetBytes(payload, "tools") + if tools.IsArray() { + tools.ForEach(func(idx, item gjson.Result) bool { + processBlock(fmt.Sprintf("tools.%d", int(idx.Int())), item) + return true + }) + } + + system := gjson.GetBytes(payload, "system") + if system.IsArray() { + system.ForEach(func(idx, item gjson.Result) bool { + processBlock(fmt.Sprintf("system.%d", int(idx.Int())), item) + return true + }) + } + + messages := gjson.GetBytes(payload, "messages") + if messages.IsArray() { + messages.ForEach(func(msgIdx, msg gjson.Result) bool { + content := msg.Get("content") + if !content.IsArray() { + return true + } + content.ForEach(func(itemIdx, item gjson.Result) bool { + processBlock(fmt.Sprintf("messages.%d.content.%d", int(msgIdx.Int()), int(itemIdx.Int())), item) + return true + }) + return true + }) + } + + if !modified { + return original + } + return payload +} + +// enforceCacheControlLimit removes excess cache_control blocks from a payload +// so the total does not exceed the Anthropic API limit (currently 4). +// +// Anthropic evaluates cache breakpoints in order: tools → system → messages. +// The most valuable breakpoints are: +// 1. Last tool — caches ALL tool definitions +// 2. Last system block — caches ALL system content +// 3. Recent messages — cache conversation context +// +// Removal priority (strip lowest-value first): +// +// Phase 1: system blocks earliest-first, preserving the last one. +// Phase 2: tool blocks earliest-first, preserving the last one. +// Phase 3: message content blocks earliest-first. +// Phase 4: remaining system blocks (last system). +// Phase 5: remaining tool blocks (last tool). +func enforceCacheControlLimit(payload []byte, maxBlocks int) []byte { + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return payload + } + + total := countCacheControls(payload) + if total <= maxBlocks { + return payload + } + + excess := total - maxBlocks + + system := gjson.GetBytes(payload, "system") + if system.IsArray() { + lastIdx := -1 + system.ForEach(func(idx, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + lastIdx = int(idx.Int()) + } + return true + }) + if lastIdx >= 0 { + system.ForEach(func(idx, item gjson.Result) bool { + if excess <= 0 { + return false + } + i := int(idx.Int()) + if i == lastIdx { + return true + } + if !item.Get("cache_control").Exists() { + return true + } + path := fmt.Sprintf("system.%d.cache_control", i) + updated, errDel := sjson.DeleteBytes(payload, path) + if errDel != nil { + return true + } + payload = updated + excess-- + return true + }) + } + } + if excess <= 0 { + return payload + } + + tools := gjson.GetBytes(payload, "tools") + if tools.IsArray() { + lastIdx := -1 + tools.ForEach(func(idx, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + lastIdx = int(idx.Int()) + } + return true + }) + if lastIdx >= 0 { + tools.ForEach(func(idx, item gjson.Result) bool { + if excess <= 0 { + return false + } + i := int(idx.Int()) + if i == lastIdx { + return true + } + if !item.Get("cache_control").Exists() { + return true + } + path := fmt.Sprintf("tools.%d.cache_control", i) + updated, errDel := sjson.DeleteBytes(payload, path) + if errDel != nil { + return true + } + payload = updated + excess-- + return true + }) + } + } + if excess <= 0 { + return payload + } + + messages := gjson.GetBytes(payload, "messages") + if messages.IsArray() { + messages.ForEach(func(msgIdx, msg gjson.Result) bool { + if excess <= 0 { + return false + } + content := msg.Get("content") + if !content.IsArray() { + return true + } + content.ForEach(func(itemIdx, item gjson.Result) bool { + if excess <= 0 { + return false + } + if !item.Get("cache_control").Exists() { + return true + } + path := fmt.Sprintf("messages.%d.content.%d.cache_control", int(msgIdx.Int()), int(itemIdx.Int())) + updated, errDel := sjson.DeleteBytes(payload, path) + if errDel != nil { + return true + } + payload = updated + excess-- + return true + }) + return true + }) + } + if excess <= 0 { + return payload + } + + system = gjson.GetBytes(payload, "system") + if system.IsArray() { + system.ForEach(func(idx, item gjson.Result) bool { + if excess <= 0 { + return false + } + if !item.Get("cache_control").Exists() { + return true + } + path := fmt.Sprintf("system.%d.cache_control", int(idx.Int())) + updated, errDel := sjson.DeleteBytes(payload, path) + if errDel != nil { + return true + } + payload = updated + excess-- + return true + }) + } + if excess <= 0 { + return payload + } + + tools = gjson.GetBytes(payload, "tools") + if tools.IsArray() { + tools.ForEach(func(idx, item gjson.Result) bool { + if excess <= 0 { + return false + } + if !item.Get("cache_control").Exists() { + return true + } + path := fmt.Sprintf("tools.%d.cache_control", int(idx.Int())) + updated, errDel := sjson.DeleteBytes(payload, path) + if errDel != nil { + return true + } + payload = updated + excess-- + return true + }) + } + + return payload +} + +// injectMessagesCacheControl adds cache_control to the second-to-last user turn for multi-turn caching. +// Per Anthropic docs: "Place cache_control on the second-to-last User message to let the model reuse the earlier cache." +// This enables caching of conversation history, which is especially beneficial for long multi-turn conversations. +// Only adds cache_control if: +// - There are at least 2 user turns in the conversation +// - No message content already has cache_control +func injectMessagesCacheControl(payload []byte) []byte { + messages := gjson.GetBytes(payload, "messages") + if !messages.Exists() || !messages.IsArray() { + return payload + } + + // Check if ANY message content already has cache_control + hasCacheControlInMessages := false + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if content.IsArray() { + content.ForEach(func(_, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + hasCacheControlInMessages = true + return false + } + return true + }) + } + return !hasCacheControlInMessages + }) + if hasCacheControlInMessages { + return payload + } + + // Find all user message indices + var userMsgIndices []int + messages.ForEach(func(index gjson.Result, msg gjson.Result) bool { + if msg.Get("role").String() == "user" { + userMsgIndices = append(userMsgIndices, int(index.Int())) + } + return true + }) + + // Need at least 2 user turns to cache the second-to-last + if len(userMsgIndices) < 2 { + return payload + } + + // Get the second-to-last user message index + secondToLastUserIdx := userMsgIndices[len(userMsgIndices)-2] + + // Get the content of this message + contentPath := fmt.Sprintf("messages.%d.content", secondToLastUserIdx) + content := gjson.GetBytes(payload, contentPath) + + if content.IsArray() { + // Add cache_control to the last content block of this message + contentCount := int(content.Get("#").Int()) + if contentCount > 0 { + cacheControlPath := fmt.Sprintf("messages.%d.content.%d.cache_control", secondToLastUserIdx, contentCount-1) + result, err := sjson.SetBytes(payload, cacheControlPath, map[string]string{"type": "ephemeral"}) + if err != nil { + log.Warnf("failed to inject cache_control into messages: %v", err) + return payload + } + payload = result + } + } else if content.Type == gjson.String { + // Convert string content to array with cache_control + text := content.String() + newContent := []map[string]interface{}{ + { + "type": "text", + "text": text, + "cache_control": map[string]string{ + "type": "ephemeral", + }, + }, + } + result, err := sjson.SetBytes(payload, contentPath, newContent) + if err != nil { + log.Warnf("failed to inject cache_control into message string content: %v", err) + return payload + } + payload = result + } + + return payload +} + +// injectToolsCacheControl adds cache_control to the last non-deferred tool in the tools array. +// Deferred tools cannot use prompt caching, so trailing deferred tools are skipped. +// This only adds cache_control if NO tool in the array already has it. +func injectToolsCacheControl(payload []byte) []byte { + tools := gjson.GetBytes(payload, "tools") + if !tools.Exists() || !tools.IsArray() { + return payload + } + + // Check if ANY tool already has cache_control and find the last eligible tool. + hasCacheControlInTools := false + lastEligibleToolIndex := -1 + tools.ForEach(func(index, tool gjson.Result) bool { + if tool.Get("cache_control").Exists() { + hasCacheControlInTools = true + return false + } + if !tool.Get("defer_loading").Bool() { + lastEligibleToolIndex = int(index.Int()) + } + return true + }) + if hasCacheControlInTools || lastEligibleToolIndex < 0 { + return payload + } + + lastToolPath := fmt.Sprintf("tools.%d.cache_control", lastEligibleToolIndex) + result, err := sjson.SetBytes(payload, lastToolPath, map[string]string{"type": "ephemeral"}) + if err != nil { + log.Warnf("failed to inject cache_control into tools array: %v", err) + return payload + } + + return result +} + +// injectSystemCacheControl adds cache_control to the last element in the system prompt. +// Converts string system prompts to array format if needed. +// This only adds cache_control if NO system element already has it. +func injectSystemCacheControl(payload []byte) []byte { + system := gjson.GetBytes(payload, "system") + if !system.Exists() { + return payload + } + + if system.IsArray() { + count := int(system.Get("#").Int()) + if count == 0 { + return payload + } + + // Check if ANY system element already has cache_control + hasCacheControlInSystem := false + system.ForEach(func(_, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + hasCacheControlInSystem = true + return false + } + return true + }) + if hasCacheControlInSystem { + return payload + } + + // Add cache_control to the last system element + lastSystemPath := fmt.Sprintf("system.%d.cache_control", count-1) + result, err := sjson.SetBytes(payload, lastSystemPath, map[string]string{"type": "ephemeral"}) + if err != nil { + log.Warnf("failed to inject cache_control into system array: %v", err) + return payload + } + payload = result + } else if system.Type == gjson.String { + // Convert string system prompt to array with cache_control + // "system": "text" -> "system": [{"type": "text", "text": "text", "cache_control": {"type": "ephemeral"}}] + text := system.String() + newSystem := []map[string]interface{}{ + { + "type": "text", + "text": text, + "cache_control": map[string]string{ + "type": "ephemeral", + }, + }, + } + result, err := sjson.SetBytes(payload, "system", newSystem) + if err != nil { + log.Warnf("failed to inject cache_control into system string: %v", err) + return payload + } + payload = result + } + + return payload +} + +func ensureModelMaxTokens(body []byte, modelID string) []byte { + if len(body) == 0 || !gjson.ValidBytes(body) { + return body + } + + if maxTokens := gjson.GetBytes(body, "max_tokens"); maxTokens.Exists() { + return body + } + + for _, provider := range registry.GetGlobalRegistry().GetModelProviders(strings.TrimSpace(modelID)) { + if strings.EqualFold(provider, "claude") { + maxTokens := defaultModelMaxTokens + if info := registry.GetGlobalRegistry().GetModelInfo(strings.TrimSpace(modelID), "claude"); info != nil && info.MaxCompletionTokens > 0 { + maxTokens = info.MaxCompletionTokens + } + body, _ = sjson.SetBytes(body, "max_tokens", maxTokens) + return body + } + } + + return body +} diff --git a/internal/runtime/executor/claude_executor_execute.go b/internal/runtime/executor/claude_executor_execute.go new file mode 100644 index 000000000..191b43a33 --- /dev/null +++ b/internal/runtime/executor/claude_executor_execute.go @@ -0,0 +1,213 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" +) + +func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + if opts.Alt == "responses/compact" { + return resp, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + upstreamModel := e.upstreamModel(baseModel) + + apiKey, baseURL := claudeCreds(auth) + if baseURL == "" { + baseURL = "https://api.anthropic.com" + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("claude") + // Use streaming translation to preserve function calling, except for claude. + stream := from != to + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, stream) + body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream) + body = helps.SetStringIfDifferent(body, "model", upstreamModel) + + body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return resp, err + } + if rebuildMidSystemMessageEnabled(e.cfg, auth) { + body = rebuildMidSystemMessagesToTopLevel(body) + } + + // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) + // based on client type and configuration. + body, err = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) + if err != nil { + return resp, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body = ensureModelMaxTokens(body, baseModel) + + // Disable thinking if tool_choice forces tool use (Anthropic API constraint) + body = disableThinkingIfToolChoiceForced(body) + body = normalizeClaudeSamplingForUpstream(body) + // Claude OAuth (and this executor's redact-thinking beta) returns signature-only + // thinking blocks unless display is set to "summarized". + body = ensureClaudeThinkingDisplay(body) + + // Auto-inject cache_control if missing (optimization for ClawdBot/clients without caching support) + if countCacheControls(body) == 0 { + body = ensureCacheControl(body) + } + + // Enforce Anthropic's cache_control block limit (max 4 breakpoints per request). + // Cloaking and ensureCacheControl may push the total over 4 when the client + // already sends multiple cache_control blocks. + body = enforceCacheControlLimit(body, 4) + + // Normalize TTL values to prevent ordering violations under prompt-caching-scope-2026-01-05. + // A 1h-TTL block must not appear after a 5m-TTL block in evaluation order (tools→system→messages). + body = normalizeCacheControlTTL(body) + + // Extract betas from body and convert to header + var extraBetas []string + extraBetas, body = extractAndRemoveBetas(body) + bodyForTranslation := body + bodyForUpstream := body + oauthToken := isClaudeOAuthToken(apiKey) + var oauthToolNamesReverseMap map[string]string + if oauthToken { + bodyForUpstream, oauthToolNamesReverseMap = prepareClaudeOAuthToolNamesForUpstream(bodyForUpstream, claudeToolPrefix, auth.ToolPrefixDisabled()) + } + bodyForUpstream = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, bodyForUpstream, baseModel) + // Enable cch signing by default for OAuth tokens (not just experimental flag). + // Claude Code always computes cch; missing or invalid cch is a detectable fingerprint. + if oauthToken || experimentalCCHSigningEnabled(e.cfg, auth) { + bodyForUpstream = signAnthropicMessagesBody(bodyForUpstream) + } + reporter.SetTranslatedReasoningEffort(bodyForUpstream, to.String()) + + url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyForUpstream)) + if err != nil { + return resp, err + } + if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg, opts.Headers); errHeaders != nil { + return resp, errHeaders + } + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: url, + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: bodyForUpstream, + Provider: e.upstreamRequestLogProvider(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + // Decompress error responses — pass the Content-Encoding value (may be empty) + // and let decodeResponseBody handle both header-declared and magic-byte-detected + // compression. This keeps error-path behaviour consistent with the success path. + errBody, decErr := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) + if decErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, decErr) + msg := fmt.Sprintf("failed to decode error response body: %v", decErr) + helps.LogWithRequestID(ctx).Warn(msg) + return resp, statusErr{code: httpResp.StatusCode, msg: msg} + } + b, readErr := io.ReadAll(errBody) + if readErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, readErr) + msg := fmt.Sprintf("failed to read error response body: %v", readErr) + helps.LogWithRequestID(ctx).Warn(msg) + b = []byte(msg) + } + helps.AppendAPIResponseChunk(ctx, e.cfg, b) + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) + err = statusErr{code: httpResp.StatusCode, msg: string(b)} + if errClose := errBody.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + return resp, err + } + decodedBody, err := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + return resp, err + } + defer func() { + if errClose := decodedBody.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + }() + data, err := io.ReadAll(decodedBody) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + if stream { + if errValidate := validateClaudeStreamingResponse(data); errValidate != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errValidate) + return resp, errValidate + } + lines := bytes.Split(data, []byte("\n")) + for _, line := range lines { + if detail, ok := helps.ParseClaudeStreamUsage(line); ok { + reporter.Publish(ctx, detail) + } + } + } else { + reporter.Publish(ctx, helps.ParseClaudeUsage(data)) + } + data = restoreClaudeOAuthToolNamesFromResponse(data, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) + data = e.restoreResponseModel(data, req.Model) + var param any + out := sdktranslator.TranslateNonStream( + ctx, + to, + responseFormat, + req.Model, + opts.OriginalRequest, + bodyForTranslation, + data, + ¶m, + ) + resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} + return resp, nil +} diff --git a/internal/runtime/executor/claude_executor_request.go b/internal/runtime/executor/claude_executor_request.go new file mode 100644 index 000000000..847132e94 --- /dev/null +++ b/internal/runtime/executor/claude_executor_request.go @@ -0,0 +1,908 @@ +package executor + +import ( + "bufio" + "bytes" + "compress/flate" + "compress/gzip" + "fmt" + "io" + "net/http" + "strings" + + "github.com/andybalholm/brotli" + "github.com/google/uuid" + "github.com/klauspost/compress/zstd" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + + "github.com/gin-gonic/gin" +) + +// extractAndRemoveBetas extracts the "betas" array from the body and removes it. +// Returns the extracted betas as a string slice and the modified body. +func extractAndRemoveBetas(body []byte) ([]string, []byte) { + betasResult := gjson.GetBytes(body, "betas") + if !betasResult.Exists() { + return nil, body + } + var betas []string + if betasResult.IsArray() { + for _, item := range betasResult.Array() { + if s := strings.TrimSpace(item.String()); s != "" { + betas = append(betas, s) + } + } + } else if s := strings.TrimSpace(betasResult.String()); s != "" { + betas = append(betas, s) + } + body, _ = sjson.DeleteBytes(body, "betas") + return betas, body +} + +// disableThinkingIfToolChoiceForced checks if tool_choice forces tool use and disables thinking. +// Anthropic API does not allow thinking when tool_choice is set to "any" or a specific tool. +// See: https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#important-considerations +func disableThinkingIfToolChoiceForced(body []byte) []byte { + toolChoiceType := gjson.GetBytes(body, "tool_choice.type").String() + // "auto" is allowed with thinking, but "any" or "tool" (specific tool) are not + if toolChoiceType == "any" || toolChoiceType == "tool" { + // Remove thinking configuration entirely to avoid API error + body, _ = sjson.DeleteBytes(body, "thinking") + // Adaptive thinking may also set output_config.effort; remove it to avoid + // leaking thinking controls when tool_choice forces tool use. + body, _ = sjson.DeleteBytes(body, "output_config.effort") + if oc := gjson.GetBytes(body, "output_config"); oc.Exists() && oc.IsObject() && len(oc.Map()) == 0 { + body, _ = sjson.DeleteBytes(body, "output_config") + } + } + return body +} + +// normalizeClaudeSamplingForUpstream keeps Anthropic message requests valid. +func normalizeClaudeSamplingForUpstream(body []byte) []byte { + body, _ = sjson.DeleteBytes(body, "temperature") + body, _ = sjson.DeleteBytes(body, "top_p") + + thinkingType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "thinking.type").String())) + switch thinkingType { + case "enabled", "adaptive", "auto": + body, _ = sjson.DeleteBytes(body, "top_p") + body, _ = sjson.DeleteBytes(body, "top_k") + } + return body +} + +// ensureClaudeThinkingDisplay defaults thinking.display to "summarized" when thinking +// is active and the client did not set display. Without this, Claude backends that +// enable redact-thinking return signature-only thinking blocks (empty thinking text). +// Explicit client values such as "omitted" are preserved. +func ensureClaudeThinkingDisplay(body []byte) []byte { + thinkingType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "thinking.type").String())) + switch thinkingType { + case "enabled", "adaptive", "auto": + default: + return body + } + if display := strings.TrimSpace(gjson.GetBytes(body, "thinking.display").String()); display != "" { + return body + } + out, err := sjson.SetBytes(body, "thinking.display", "summarized") + if err != nil { + return body + } + return out +} + +type compositeReadCloser struct { + io.Reader + closers []func() error +} + +func (c *compositeReadCloser) Close() error { + var firstErr error + for i := range c.closers { + if c.closers[i] == nil { + continue + } + if err := c.closers[i](); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + +// peekableBody wraps a bufio.Reader around the original ReadCloser so that +// magic bytes can be inspected without consuming them from the stream. +type peekableBody struct { + *bufio.Reader + closer io.Closer +} + +func (p *peekableBody) Close() error { + return p.closer.Close() +} + +func decodeResponseBody(body io.ReadCloser, contentEncoding string) (io.ReadCloser, error) { + if body == nil { + return nil, fmt.Errorf("response body is nil") + } + if contentEncoding == "" { + // No Content-Encoding header. Attempt best-effort magic-byte detection to + // handle misbehaving upstreams that compress without setting the header. + // Only gzip (1f 8b) and zstd (28 b5 2f fd) have reliable magic sequences; + // br and deflate have none and are left as-is. + // The bufio wrapper preserves unread bytes so callers always see the full + // stream regardless of whether decompression was applied. + pb := &peekableBody{Reader: bufio.NewReader(body), closer: body} + magic, peekErr := pb.Peek(4) + if peekErr == nil || (peekErr == io.EOF && len(magic) >= 2) { + switch { + case len(magic) >= 2 && magic[0] == 0x1f && magic[1] == 0x8b: + gzipReader, gzErr := gzip.NewReader(pb) + if gzErr != nil { + _ = pb.Close() + return nil, fmt.Errorf("magic-byte gzip: failed to create reader: %w", gzErr) + } + return &compositeReadCloser{ + Reader: gzipReader, + closers: []func() error{ + gzipReader.Close, + pb.Close, + }, + }, nil + case len(magic) >= 4 && magic[0] == 0x28 && magic[1] == 0xb5 && magic[2] == 0x2f && magic[3] == 0xfd: + decoder, zdErr := zstd.NewReader(pb) + if zdErr != nil { + _ = pb.Close() + return nil, fmt.Errorf("magic-byte zstd: failed to create reader: %w", zdErr) + } + return &compositeReadCloser{ + Reader: decoder, + closers: []func() error{ + func() error { decoder.Close(); return nil }, + pb.Close, + }, + }, nil + } + } + return pb, nil + } + encodings := strings.Split(contentEncoding, ",") + for _, raw := range encodings { + encoding := strings.TrimSpace(strings.ToLower(raw)) + switch encoding { + case "", "identity": + continue + case "gzip": + gzipReader, err := gzip.NewReader(body) + if err != nil { + _ = body.Close() + return nil, fmt.Errorf("failed to create gzip reader: %w", err) + } + return &compositeReadCloser{ + Reader: gzipReader, + closers: []func() error{ + gzipReader.Close, + func() error { return body.Close() }, + }, + }, nil + case "deflate": + deflateReader := flate.NewReader(body) + return &compositeReadCloser{ + Reader: deflateReader, + closers: []func() error{ + deflateReader.Close, + func() error { return body.Close() }, + }, + }, nil + case "br": + return &compositeReadCloser{ + Reader: brotli.NewReader(body), + closers: []func() error{ + func() error { return body.Close() }, + }, + }, nil + case "zstd": + decoder, err := zstd.NewReader(body) + if err != nil { + _ = body.Close() + return nil, fmt.Errorf("failed to create zstd reader: %w", err) + } + return &compositeReadCloser{ + Reader: decoder, + closers: []func() error{ + func() error { decoder.Close(); return nil }, + func() error { return body.Close() }, + }, + }, nil + default: + continue + } + } + return body, nil +} + +func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, stream bool, extraBetas []string, cfg *config.Config, incomingHeaders http.Header) error { + if r == nil { + return nil + } + hdrDefault := func(cfgVal, fallback string) string { + if cfgVal != "" { + return cfgVal + } + return fallback + } + + var hd config.ClaudeHeaderDefaults + if cfg != nil { + hd = cfg.ClaudeHeaderDefaults + } + + useAPIKey := auth != nil && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["api_key"]) != "" + isAnthropicBase := r.URL != nil && strings.EqualFold(r.URL.Scheme, "https") && strings.EqualFold(r.URL.Host, "api.anthropic.com") + if isAnthropicBase && useAPIKey { + r.Header.Del("Authorization") + r.Header.Set("x-api-key", apiKey) + } else { + r.Header.Set("Authorization", "Bearer "+apiKey) + } + r.Header.Set("Content-Type", "application/json") + + if incomingHeaders == nil { + if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + incomingHeaders = ginCtx.Request.Header + } + } + stabilizeDeviceProfile := helps.ClaudeDeviceProfileStabilizationEnabled(cfg) + var deviceProfile helps.ClaudeDeviceProfile + if stabilizeDeviceProfile { + var errDeviceProfile error + deviceProfile, errDeviceProfile = helps.ResolveClaudeDeviceProfileRequired(r.Context(), auth, apiKey, incomingHeaders, cfg) + if errDeviceProfile != nil { + return errDeviceProfile + } + } + + baseBetas := "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,context-management-2025-06-27,prompt-caching-scope-2026-01-05,structured-outputs-2025-12-15,fast-mode-2026-02-01,redact-thinking-2026-02-12,token-efficient-tools-2026-03-28" + if val := strings.TrimSpace(strings.Join(incomingHeaders.Values("Anthropic-Beta"), ",")); val != "" { + baseBetas = val + if !strings.Contains(val, "oauth") { + baseBetas += ",oauth-2025-04-20" + } + } + if !strings.Contains(baseBetas, "interleaved-thinking") { + baseBetas += ",interleaved-thinking-2025-05-14" + } + + // Merge extra betas from request body and request flags. + if len(extraBetas) > 0 { + existingSet := make(map[string]bool) + for _, b := range strings.Split(baseBetas, ",") { + betaName := strings.TrimSpace(b) + if betaName != "" { + existingSet[betaName] = true + } + } + for _, beta := range extraBetas { + beta = strings.TrimSpace(beta) + if beta != "" && !existingSet[beta] { + baseBetas += "," + beta + existingSet[beta] = true + } + } + } + r.Header.Set("Anthropic-Beta", baseBetas) + + misc.EnsureHeader(r.Header, incomingHeaders, "Anthropic-Version", "2023-06-01") + // Only set browser access header for API key mode; real Claude Code CLI does not send it. + if useAPIKey { + misc.EnsureHeader(r.Header, incomingHeaders, "Anthropic-Dangerous-Direct-Browser-Access", "true") + } + misc.EnsureHeader(r.Header, incomingHeaders, "X-App", "cli") + // Values below match Claude Code 2.1.63 / @anthropic-ai/sdk 0.74.0 (updated 2026-02-28). + misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Retry-Count", "0") + misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Runtime", "node") + misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Lang", "js") + misc.EnsureHeader(r.Header, incomingHeaders, "X-Stainless-Timeout", hdrDefault(hd.Timeout, "600")) + // Session ID: stable per auth/apiKey, matches Claude Code's X-Claude-Code-Session-Id header. + sessionID, errSessionID := helps.CachedSessionIDRequired(r.Context(), apiKey) + if errSessionID != nil { + return errSessionID + } + misc.EnsureHeader(r.Header, incomingHeaders, "X-Claude-Code-Session-Id", sessionID) + // Per-request UUID, matches Claude Code's x-client-request-id for first-party API. + if isAnthropicBase { + misc.EnsureHeader(r.Header, incomingHeaders, "x-client-request-id", uuid.New().String()) + } + r.Header.Set("Connection", "keep-alive") + if stream { + r.Header.Set("Accept", "text/event-stream") + // SSE streams must not be compressed: the downstream scanner reads + // line-delimited text and cannot parse compressed bytes. Using + // "identity" tells the upstream to send an uncompressed stream. + r.Header.Set("Accept-Encoding", "identity") + } else { + r.Header.Set("Accept", "application/json") + r.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd") + } + // Legacy mode keeps OS/Arch runtime-derived; stabilized mode pins OS/Arch + // to the configured baseline while still allowing newer official + // User-Agent/package/runtime tuples to upgrade the software fingerprint. + if stabilizeDeviceProfile { + helps.ApplyClaudeDeviceProfileHeaders(r, deviceProfile) + } else { + helps.ApplyClaudeLegacyDeviceHeaders(r, incomingHeaders, cfg) + } + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(r, attrs) + // Re-enforce Accept-Encoding: identity after ApplyCustomHeadersFromAttrs, which + // may override it with a user-configured value. Compressed SSE breaks the line + // scanner regardless of user preference, so this is non-negotiable for streams. + if stream { + r.Header.Set("Accept-Encoding", "identity") + } + return nil +} + +func claudeCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) { + if a == nil { + return "", "" + } + if a.Attributes != nil { + apiKey = a.Attributes["api_key"] + baseURL = a.Attributes["base_url"] + } + if apiKey == "" && a.Metadata != nil { + if v, ok := a.Metadata["access_token"].(string); ok { + apiKey = v + } + } + return +} + +func checkSystemInstructions(payload []byte) []byte { + return checkSystemInstructionsWithSigningMode(payload, false, false, false, "2.1.63", "", "") +} + +func rebuildMidSystemMessagesToTopLevel(payload []byte) []byte { + messages := gjson.GetBytes(payload, "messages") + if !messages.IsArray() { + return payload + } + + var movedSystemParts []string + keptMessages := make([]string, 0, int(messages.Get("#").Int())) + messages.ForEach(func(_, message gjson.Result) bool { + if strings.EqualFold(strings.TrimSpace(message.Get("role").String()), "system") { + movedSystemParts = append(movedSystemParts, claudeSystemTextParts(message.Get("content"))...) + return true + } + keptMessages = append(keptMessages, message.Raw) + return true + }) + if len(movedSystemParts) == 0 { + return payload + } + + systemParts := claudeSystemTextParts(gjson.GetBytes(payload, "system")) + systemParts = append(systemParts, movedSystemParts...) + if len(systemParts) > 0 { + if updated, errSetSystem := sjson.SetRawBytes(payload, "system", rawJSONArray(systemParts)); errSetSystem == nil { + payload = updated + } + } + if updated, errSetMessages := sjson.SetRawBytes(payload, "messages", rawJSONArray(keptMessages)); errSetMessages == nil { + payload = updated + } + return payload +} + +func claudeSystemTextParts(content gjson.Result) []string { + if !content.Exists() { + return nil + } + if content.Type == gjson.String { + text := content.String() + if strings.TrimSpace(text) == "" { + return nil + } + block := []byte(`{"type":"text","text":""}`) + block, _ = sjson.SetBytes(block, "text", text) + return []string{string(block)} + } + if !content.IsArray() { + return nil + } + + var parts []string + content.ForEach(func(_, item gjson.Result) bool { + if item.Type == gjson.String { + text := item.String() + if strings.TrimSpace(text) != "" { + block := []byte(`{"type":"text","text":""}`) + block, _ = sjson.SetBytes(block, "text", text) + parts = append(parts, string(block)) + } + return true + } + if item.IsObject() && item.Get("type").String() == "text" && strings.TrimSpace(item.Get("text").String()) != "" { + parts = append(parts, item.Raw) + } + return true + }) + return parts +} + +func rawJSONArray(items []string) []byte { + if len(items) == 0 { + return []byte("[]") + } + var builder strings.Builder + builder.WriteByte('[') + for i, item := range items { + if i > 0 { + builder.WriteByte(',') + } + builder.WriteString(item) + } + builder.WriteByte(']') + return []byte(builder.String()) +} + +func isClaudeOAuthToken(apiKey string) bool { + return strings.Contains(apiKey, "sk-ant-oat") +} + +// prepareClaudeOAuthToolNamesForUpstream applies the Claude OAuth tool-name +// transforms in the same order across request paths. Remap runs before prefixing +// so any future non-empty prefix still composes correctly with the per-request +// reverse map. +func prepareClaudeOAuthToolNamesForUpstream(body []byte, prefix string, prefixDisabled bool) ([]byte, map[string]string) { + body, reverseMap := remapOAuthToolNames(body) + if !prefixDisabled { + body = applyClaudeToolPrefix(body, prefix) + } + return body, reverseMap +} + +// restoreClaudeOAuthToolNamesFromResponse undoes the Claude OAuth tool-name +// transforms for non-stream responses in reverse order. +func restoreClaudeOAuthToolNamesFromResponse(body []byte, prefix string, prefixDisabled bool, reverseMap map[string]string) []byte { + if !prefixDisabled { + body = stripClaudeToolPrefixFromResponse(body, prefix) + } + return reverseRemapOAuthToolNames(body, reverseMap) +} + +// restoreClaudeOAuthToolNamesFromStreamLine undoes the Claude OAuth tool-name +// transforms for SSE lines in reverse order. +func restoreClaudeOAuthToolNamesFromStreamLine(line []byte, prefix string, prefixDisabled bool, reverseMap map[string]string) []byte { + if !prefixDisabled { + line = stripClaudeToolPrefixFromStreamLine(line, prefix) + } + return reverseRemapOAuthToolNamesFromStreamLine(line, reverseMap) +} + +// remapOAuthToolNames renames third-party tool names to Claude Code equivalents +// and removes tools without an official counterpart. This prevents Anthropic from +// fingerprinting the request as a third-party client via tool naming patterns. +// +// It operates on: tools[].name, tool_choice.name, and all tool_use/tool_reference +// references in messages. Removed tools' corresponding tool_result blocks are preserved +// (they just become orphaned, which is safe for Claude). +// +// The returned map is keyed on the upstream (TitleCase) name and maps to the +// client-supplied original name. Callers MUST pass this map to the reverse +// functions so only names the client actually caused us to rewrite are restored +// on the response. A global reverse map (the previous implementation) incorrectly +// rewrote names the client originally sent in TitleCase (e.g. `Bash`) +// when any OTHER tool in the same request triggered a forward rename (e.g. +// `glob` -> `Glob`), because the global reverse map contained `Bash` -> `bash` +// regardless of what the client originally sent. +func remapOAuthToolNames(body []byte) ([]byte, map[string]string) { + reverseMap := make(map[string]string, len(oauthToolRenameMap)) + recordRename := func(original, renamed string) { + // Preserve the first-seen original name if the same upstream name is + // produced from multiple call sites; they all map back identically. + if _, exists := reverseMap[renamed]; !exists { + reverseMap[renamed] = original + } + } + + // 1. Rewrite tools array in a single pass (if present). + // IMPORTANT: do not mutate names first and then rebuild from an older gjson + // snapshot. gjson results are snapshots of the original bytes; rebuilding from a + // stale snapshot will preserve removals but overwrite renamed names back to their + // original lowercase values. + tools := gjson.GetBytes(body, "tools") + toolsNeedRewrite := false + if tools.Exists() && tools.IsArray() { + tools.ForEach(func(_, tool gjson.Result) bool { + if tool.Get("type").Exists() && tool.Get("type").String() != "" { + return true + } + name := tool.Get("name").String() + toolsNeedRewrite = oauthToolsToRemove[name] + if !toolsNeedRewrite { + newName, ok := oauthToolRenameMap[name] + toolsNeedRewrite = ok && newName != name + } + return !toolsNeedRewrite + }) + } + if toolsNeedRewrite { + var toolsJSON strings.Builder + toolsJSON.WriteByte('[') + toolCount := 0 + tools.ForEach(func(_, tool gjson.Result) bool { + // Keep Anthropic built-in tools (web_search, code_execution, etc.) unchanged. + if tool.Get("type").Exists() && tool.Get("type").String() != "" { + if toolCount > 0 { + toolsJSON.WriteByte(',') + } + toolsJSON.WriteString(tool.Raw) + toolCount++ + return true + } + + name := tool.Get("name").String() + if oauthToolsToRemove[name] { + return true + } + + toolJSON := tool.Raw + if newName, ok := oauthToolRenameMap[name]; ok && newName != name { + updatedTool, err := sjson.Set(toolJSON, "name", newName) + if err == nil { + toolJSON = updatedTool + recordRename(name, newName) + } + } + + if toolCount > 0 { + toolsJSON.WriteByte(',') + } + toolsJSON.WriteString(toolJSON) + toolCount++ + return true + }) + toolsJSON.WriteByte(']') + body, _ = sjson.SetRawBytes(body, "tools", []byte(toolsJSON.String())) + } + + // 2. Rename tool_choice if it references a known tool + toolChoiceType := gjson.GetBytes(body, "tool_choice.type").String() + if toolChoiceType == "tool" { + tcName := gjson.GetBytes(body, "tool_choice.name").String() + if oauthToolsToRemove[tcName] { + // The chosen tool was removed from the tools array, so drop tool_choice to + // keep the payload internally consistent and fall back to normal auto tool use. + body, _ = sjson.DeleteBytes(body, "tool_choice") + } else if newName, ok := oauthToolRenameMap[tcName]; ok && newName != tcName { + body, _ = sjson.SetBytes(body, "tool_choice.name", newName) + recordRename(tcName, newName) + } + } + + // 3. Rename tool references in messages + messages := gjson.GetBytes(body, "messages") + if messages.Exists() && messages.IsArray() { + messages.ForEach(func(msgIndex, msg gjson.Result) bool { + content := msg.Get("content") + if !content.Exists() || !content.IsArray() { + return true + } + content.ForEach(func(contentIndex, part gjson.Result) bool { + partType := part.Get("type").String() + switch partType { + case "tool_use": + name := part.Get("name").String() + if newName, ok := oauthToolRenameMap[name]; ok && newName != name { + path := fmt.Sprintf("messages.%d.content.%d.name", msgIndex.Int(), contentIndex.Int()) + body, _ = sjson.SetBytes(body, path, newName) + recordRename(name, newName) + } + case "tool_reference": + toolName := part.Get("tool_name").String() + if newName, ok := oauthToolRenameMap[toolName]; ok && newName != toolName { + path := fmt.Sprintf("messages.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int()) + body, _ = sjson.SetBytes(body, path, newName) + recordRename(toolName, newName) + } + case "tool_result": + // Handle nested tool_reference blocks inside tool_result.content[] + toolID := part.Get("tool_use_id").String() + _ = toolID // tool_use_id stays as-is + nestedContent := part.Get("content") + if nestedContent.Exists() && nestedContent.IsArray() { + nestedContent.ForEach(func(nestedIndex, nestedPart gjson.Result) bool { + if nestedPart.Get("type").String() == "tool_reference" { + nestedToolName := nestedPart.Get("tool_name").String() + if newName, ok := oauthToolRenameMap[nestedToolName]; ok && newName != nestedToolName { + nestedPath := fmt.Sprintf("messages.%d.content.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int(), nestedIndex.Int()) + body, _ = sjson.SetBytes(body, nestedPath, newName) + recordRename(nestedToolName, newName) + } + } + return true + }) + } + } + return true + }) + return true + }) + } + + return body, reverseMap +} + +// reverseRemapOAuthToolNames reverses the tool name mapping for non-stream responses +// using the per-request map produced by remapOAuthToolNames. Names the client sent +// that were NOT forward-renamed are passed through unchanged. +func reverseRemapOAuthToolNames(body []byte, reverseMap map[string]string) []byte { + if len(reverseMap) == 0 { + return body + } + content := gjson.GetBytes(body, "content") + if !content.Exists() || !content.IsArray() { + return body + } + content.ForEach(func(index, part gjson.Result) bool { + partType := part.Get("type").String() + switch partType { + case "tool_use": + name := part.Get("name").String() + if origName, ok := reverseMap[name]; ok { + path := fmt.Sprintf("content.%d.name", index.Int()) + body, _ = sjson.SetBytes(body, path, origName) + } + case "tool_reference": + toolName := part.Get("tool_name").String() + if origName, ok := reverseMap[toolName]; ok { + path := fmt.Sprintf("content.%d.tool_name", index.Int()) + body, _ = sjson.SetBytes(body, path, origName) + } + } + return true + }) + return body +} + +// reverseRemapOAuthToolNamesFromStreamLine reverses the tool name mapping for SSE +// stream lines, using the per-request reverseMap produced by remapOAuthToolNames. +func reverseRemapOAuthToolNamesFromStreamLine(line []byte, reverseMap map[string]string) []byte { + if len(reverseMap) == 0 { + return line + } + payload := helps.JSONPayload(line) + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return line + } + + contentBlock := gjson.GetBytes(payload, "content_block") + if !contentBlock.Exists() { + return line + } + + blockType := contentBlock.Get("type").String() + var updated []byte + var err error + + switch blockType { + case "tool_use": + name := contentBlock.Get("name").String() + if origName, ok := reverseMap[name]; ok { + updated, err = sjson.SetBytes(payload, "content_block.name", origName) + if err != nil { + return line + } + } else { + return line + } + case "tool_reference": + toolName := contentBlock.Get("tool_name").String() + if origName, ok := reverseMap[toolName]; ok { + updated, err = sjson.SetBytes(payload, "content_block.tool_name", origName) + if err != nil { + return line + } + } else { + return line + } + default: + return line + } + + trimmed := bytes.TrimSpace(line) + if bytes.HasPrefix(trimmed, []byte("data:")) { + return append([]byte("data: "), updated...) + } + return updated +} + +func applyClaudeToolPrefix(body []byte, prefix string) []byte { + if prefix == "" { + return body + } + + // Collect built-in tool names from the authoritative fallback seed list and + // augment it with any typed built-ins present in the current request body. + builtinTools := helps.AugmentClaudeBuiltinToolRegistry(body, nil) + + if tools := gjson.GetBytes(body, "tools"); tools.Exists() && tools.IsArray() { + tools.ForEach(func(index, tool gjson.Result) bool { + // Skip built-in tools (web_search, code_execution, etc.) which have + // a "type" field and require their name to remain unchanged. + if tool.Get("type").Exists() && tool.Get("type").String() != "" { + if n := tool.Get("name").String(); n != "" { + builtinTools[n] = true + } + return true + } + name := tool.Get("name").String() + if name == "" || strings.HasPrefix(name, prefix) { + return true + } + path := fmt.Sprintf("tools.%d.name", index.Int()) + body, _ = sjson.SetBytes(body, path, prefix+name) + return true + }) + } + + if gjson.GetBytes(body, "tool_choice.type").String() == "tool" { + name := gjson.GetBytes(body, "tool_choice.name").String() + if name != "" && !strings.HasPrefix(name, prefix) && !builtinTools[name] { + body, _ = sjson.SetBytes(body, "tool_choice.name", prefix+name) + } + } + + if messages := gjson.GetBytes(body, "messages"); messages.Exists() && messages.IsArray() { + messages.ForEach(func(msgIndex, msg gjson.Result) bool { + content := msg.Get("content") + if !content.Exists() || !content.IsArray() { + return true + } + content.ForEach(func(contentIndex, part gjson.Result) bool { + partType := part.Get("type").String() + switch partType { + case "tool_use": + name := part.Get("name").String() + if name == "" || strings.HasPrefix(name, prefix) || builtinTools[name] { + return true + } + path := fmt.Sprintf("messages.%d.content.%d.name", msgIndex.Int(), contentIndex.Int()) + body, _ = sjson.SetBytes(body, path, prefix+name) + case "tool_reference": + toolName := part.Get("tool_name").String() + if toolName == "" || strings.HasPrefix(toolName, prefix) || builtinTools[toolName] { + return true + } + path := fmt.Sprintf("messages.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int()) + body, _ = sjson.SetBytes(body, path, prefix+toolName) + case "tool_result": + // Handle nested tool_reference blocks inside tool_result.content[] + nestedContent := part.Get("content") + if nestedContent.Exists() && nestedContent.IsArray() { + nestedContent.ForEach(func(nestedIndex, nestedPart gjson.Result) bool { + if nestedPart.Get("type").String() == "tool_reference" { + nestedToolName := nestedPart.Get("tool_name").String() + if nestedToolName != "" && !strings.HasPrefix(nestedToolName, prefix) && !builtinTools[nestedToolName] { + nestedPath := fmt.Sprintf("messages.%d.content.%d.content.%d.tool_name", msgIndex.Int(), contentIndex.Int(), nestedIndex.Int()) + body, _ = sjson.SetBytes(body, nestedPath, prefix+nestedToolName) + } + } + return true + }) + } + } + return true + }) + return true + }) + } + + return body +} + +func stripClaudeToolPrefixFromResponse(body []byte, prefix string) []byte { + if prefix == "" { + return body + } + content := gjson.GetBytes(body, "content") + if !content.Exists() || !content.IsArray() { + return body + } + content.ForEach(func(index, part gjson.Result) bool { + partType := part.Get("type").String() + switch partType { + case "tool_use": + name := part.Get("name").String() + if !strings.HasPrefix(name, prefix) { + return true + } + path := fmt.Sprintf("content.%d.name", index.Int()) + body, _ = sjson.SetBytes(body, path, strings.TrimPrefix(name, prefix)) + case "tool_reference": + toolName := part.Get("tool_name").String() + if !strings.HasPrefix(toolName, prefix) { + return true + } + path := fmt.Sprintf("content.%d.tool_name", index.Int()) + body, _ = sjson.SetBytes(body, path, strings.TrimPrefix(toolName, prefix)) + case "tool_result": + // Handle nested tool_reference blocks inside tool_result.content[] + nestedContent := part.Get("content") + if nestedContent.Exists() && nestedContent.IsArray() { + nestedContent.ForEach(func(nestedIndex, nestedPart gjson.Result) bool { + if nestedPart.Get("type").String() == "tool_reference" { + nestedToolName := nestedPart.Get("tool_name").String() + if strings.HasPrefix(nestedToolName, prefix) { + nestedPath := fmt.Sprintf("content.%d.content.%d.tool_name", index.Int(), nestedIndex.Int()) + body, _ = sjson.SetBytes(body, nestedPath, strings.TrimPrefix(nestedToolName, prefix)) + } + } + return true + }) + } + } + return true + }) + return body +} + +func stripClaudeToolPrefixFromStreamLine(line []byte, prefix string) []byte { + if prefix == "" { + return line + } + payload := helps.JSONPayload(line) + if len(payload) == 0 || !gjson.ValidBytes(payload) { + return line + } + contentBlock := gjson.GetBytes(payload, "content_block") + if !contentBlock.Exists() { + return line + } + + blockType := contentBlock.Get("type").String() + var updated []byte + var err error + + switch blockType { + case "tool_use": + name := contentBlock.Get("name").String() + if !strings.HasPrefix(name, prefix) { + return line + } + updated, err = sjson.SetBytes(payload, "content_block.name", strings.TrimPrefix(name, prefix)) + if err != nil { + return line + } + case "tool_reference": + toolName := contentBlock.Get("tool_name").String() + if !strings.HasPrefix(toolName, prefix) { + return line + } + updated, err = sjson.SetBytes(payload, "content_block.tool_name", strings.TrimPrefix(toolName, prefix)) + if err != nil { + return line + } + default: + return line + } + + trimmed := bytes.TrimSpace(line) + if bytes.HasPrefix(trimmed, []byte("data:")) { + return append([]byte("data: "), updated...) + } + return updated +} diff --git a/internal/runtime/executor/claude_executor_stream.go b/internal/runtime/executor/claude_executor_stream.go new file mode 100644 index 000000000..7f5499338 --- /dev/null +++ b/internal/runtime/executor/claude_executor_stream.go @@ -0,0 +1,323 @@ +package executor + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "net/http" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { + if opts.Alt == "responses/compact" { + return nil, statusErr{code: http.StatusNotImplemented, msg: "/responses/compact not supported"} + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + upstreamModel := e.upstreamModel(baseModel) + + apiKey, baseURL := claudeCreds(auth) + if baseURL == "" { + baseURL = "https://api.anthropic.com" + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("claude") + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalTranslated := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, originalPayload, true) + body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, true) + body = helps.SetStringIfDifferent(body, "model", upstreamModel) + + body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return nil, err + } + if rebuildMidSystemMessageEnabled(e.cfg, auth) { + body = rebuildMidSystemMessagesToTopLevel(body) + } + + // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) + // based on client type and configuration. + body, err = applyCloaking(ctx, e.cfg, auth, body, baseModel, apiKey) + if err != nil { + return nil, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body = ensureModelMaxTokens(body, baseModel) + + // Disable thinking if tool_choice forces tool use (Anthropic API constraint) + body = disableThinkingIfToolChoiceForced(body) + body = normalizeClaudeSamplingForUpstream(body) + // Claude OAuth (and this executor's redact-thinking beta) returns signature-only + // thinking blocks unless display is set to "summarized". + body = ensureClaudeThinkingDisplay(body) + + // Auto-inject cache_control if missing (optimization for ClawdBot/clients without caching support) + if countCacheControls(body) == 0 { + body = ensureCacheControl(body) + } + + // Enforce Anthropic's cache_control block limit (max 4 breakpoints per request). + body = enforceCacheControlLimit(body, 4) + + // Normalize TTL values to prevent ordering violations under prompt-caching-scope-2026-01-05. + body = normalizeCacheControlTTL(body) + + // Extract betas from body and convert to header + var extraBetas []string + extraBetas, body = extractAndRemoveBetas(body) + bodyForTranslation := body + bodyForUpstream := body + oauthToken := isClaudeOAuthToken(apiKey) + var oauthToolNamesReverseMap map[string]string + if oauthToken { + bodyForUpstream, oauthToolNamesReverseMap = prepareClaudeOAuthToolNamesForUpstream(bodyForUpstream, claudeToolPrefix, auth.ToolPrefixDisabled()) + } + bodyForUpstream = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, bodyForUpstream, baseModel) + // Enable cch signing by default for OAuth tokens (not just experimental flag). + if oauthToken || experimentalCCHSigningEnabled(e.cfg, auth) { + bodyForUpstream = signAnthropicMessagesBody(bodyForUpstream) + } + reporter.SetTranslatedReasoningEffort(bodyForUpstream, to.String()) + + url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyForUpstream)) + if err != nil { + return nil, err + } + if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, true, extraBetas, e.cfg, opts.Headers); errHeaders != nil { + return nil, errHeaders + } + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: url, + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: bodyForUpstream, + Provider: e.upstreamRequestLogProvider(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return nil, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + // Decompress error responses — pass the Content-Encoding value (may be empty) + // and let decodeResponseBody handle both header-declared and magic-byte-detected + // compression. This keeps error-path behaviour consistent with the success path. + errBody, decErr := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) + if decErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, decErr) + msg := fmt.Sprintf("failed to decode error response body: %v", decErr) + helps.LogWithRequestID(ctx).Warn(msg) + return nil, statusErr{code: httpResp.StatusCode, msg: msg} + } + b, readErr := io.ReadAll(errBody) + if readErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, readErr) + msg := fmt.Sprintf("failed to read error response body: %v", readErr) + helps.LogWithRequestID(ctx).Warn(msg) + b = []byte(msg) + } + helps.AppendAPIResponseChunk(ctx, e.cfg, b) + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) + if errClose := errBody.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + err = statusErr{code: httpResp.StatusCode, msg: string(b)} + return nil, err + } + decodedBody, err := decodeResponseBody(httpResp.Body, httpResp.Header.Get("Content-Encoding")) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + return nil, err + } + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(out) + defer func() { + if errClose := decodedBody.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + }() + + // If the response target is Claude, directly forward complete SSE events without translation. + if responseFormat == to { + scanner := bufio.NewScanner(decodedBody) + scanner.Buffer(nil, 52_428_800) // 50MB + var event bytes.Buffer + flushEvent := func() bool { + if event.Len() == 0 { + return true + } + cloned := bytes.Clone(event.Bytes()) + event.Reset() + select { + case out <- cliproxyexecutor.StreamChunk{Payload: cloned}: + return true + case <-ctx.Done(): + return false + } + } + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if detail, ok := helps.ParseClaudeStreamUsage(line); ok { + reporter.Publish(ctx, detail) + } + line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) + line = e.restoreResponseModel(line, req.Model) + event.Write(line) + event.WriteByte('\n') + if len(bytes.TrimSpace(line)) == 0 && !flushEvent() { + return + } + } + if !flushEvent() { + return + } + if errScan := scanner.Err(); errScan != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + select { + case out <- cliproxyexecutor.StreamChunk{Err: errScan}: + case <-ctx.Done(): + } + } + return + } + + // For other formats, use translation + scanner := bufio.NewScanner(decodedBody) + scanner.Buffer(nil, 52_428_800) // 50MB + var param any + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + if detail, ok := helps.ParseClaudeStreamUsage(line); ok { + reporter.Publish(ctx, detail) + } + line = restoreClaudeOAuthToolNamesFromStreamLine(line, claudeToolPrefix, auth.ToolPrefixDisabled(), oauthToolNamesReverseMap) + line = e.restoreResponseModel(line, req.Model) + chunks := sdktranslator.TranslateStream( + ctx, + to, + responseFormat, + req.Model, + opts.OriginalRequest, + bodyForTranslation, + bytes.Clone(line), + ¶m, + ) + for i := range chunks { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: + case <-ctx.Done(): + return + } + } + } + if errScan := scanner.Err(); errScan != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + select { + case out <- cliproxyexecutor.StreamChunk{Err: errScan}: + case <-ctx.Done(): + } + } + }() + return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil +} + +func validateClaudeStreamingResponse(data []byte) error { + scanner := bufio.NewScanner(bytes.NewReader(data)) + scanner.Buffer(nil, 52_428_800) + + hasData := false + hasMessageStart := false + hasMessageDelta := false + + for scanner.Scan() { + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 || !bytes.HasPrefix(line, []byte("data:")) { + continue + } + payload := bytes.TrimSpace(line[len("data:"):]) + if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) { + continue + } + hasData = true + if !gjson.ValidBytes(payload) { + return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned malformed stream data"} + } + + root := gjson.ParseBytes(payload) + switch root.Get("type").String() { + case "error": + message := strings.TrimSpace(root.Get("error.message").String()) + if message == "" { + message = strings.TrimSpace(root.Get("error.type").String()) + } + if message == "" { + message = "unknown upstream error" + } + return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned error event: " + message} + case "message_start": + message := root.Get("message") + if strings.TrimSpace(message.Get("id").String()) == "" || strings.TrimSpace(message.Get("model").String()) == "" { + return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream message_start is missing id or model"} + } + hasMessageStart = true + case "message_delta": + hasMessageDelta = true + } + } + if errScan := scanner.Err(); errScan != nil { + return errScan + } + if !hasData { + return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream returned empty stream response"} + } + if !hasMessageStart { + return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response is missing message_start"} + } + if !hasMessageDelta { + return statusErr{code: http.StatusBadGateway, msg: "claude executor: upstream stream response ended before message completion"} + } + return nil +} diff --git a/internal/runtime/executor/claude_executor_tokens.go b/internal/runtime/executor/claude_executor_tokens.go new file mode 100644 index 000000000..725f45ef1 --- /dev/null +++ b/internal/runtime/executor/claude_executor_tokens.go @@ -0,0 +1,135 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + baseModel := thinking.ParseSuffix(req.Model).ModelName + upstreamModel := e.upstreamModel(baseModel) + + apiKey, baseURL := claudeCreds(auth) + if baseURL == "" { + baseURL = "https://api.anthropic.com" + } + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("claude") + // Use streaming translation to preserve function calling, except for claude. + stream := from != to + body := helps.TranslateRequestWithCodexMultiAgentV2(ctx, opts.Headers, e.cfg, from, to, baseModel, req.Payload, stream) + body = helps.SetStringIfDifferent(body, "model", upstreamModel) + if rebuildMidSystemMessageEnabled(e.cfg, auth) { + body = rebuildMidSystemMessagesToTopLevel(body) + } + + if !strings.HasPrefix(baseModel, "claude-3-5-haiku") { + body = checkSystemInstructions(body) + } + + // Keep count_tokens requests compatible with Anthropic cache-control constraints too. + body = enforceCacheControlLimit(body, 4) + body = normalizeCacheControlTTL(body) + + // Extract betas from body and convert to header (for count_tokens too) + var extraBetas []string + extraBetas, body = extractAndRemoveBetas(body) + if isClaudeOAuthToken(apiKey) { + body, _ = prepareClaudeOAuthToolNamesForUpstream(body, claudeToolPrefix, auth.ToolPrefixDisabled()) + } + body = sanitizeClaudeMessagesForClaudeUpstreamWithDebug(ctx, body, baseModel) + + url := fmt.Sprintf("%s/v1/messages/count_tokens?beta=true", baseURL) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return cliproxyexecutor.Response{}, err + } + if errHeaders := applyClaudeHeaders(httpReq, auth, apiKey, false, extraBetas, e.cfg, opts.Headers); errHeaders != nil { + return cliproxyexecutor.Response{}, errHeaders + } + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: url, + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: body, + Provider: e.upstreamRequestLogProvider(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) + resp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return cliproxyexecutor.Response{}, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, resp.StatusCode, resp.Header.Clone()) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + // Decompress error responses — pass the Content-Encoding value (may be empty) + // and let decodeResponseBody handle both header-declared and magic-byte-detected + // compression. This keeps error-path behaviour consistent with the success path. + errBody, decErr := decodeResponseBody(resp.Body, resp.Header.Get("Content-Encoding")) + if decErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, decErr) + msg := fmt.Sprintf("failed to decode error response body: %v", decErr) + helps.LogWithRequestID(ctx).Warn(msg) + return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: msg} + } + b, readErr := io.ReadAll(errBody) + if readErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, readErr) + msg := fmt.Sprintf("failed to read error response body: %v", readErr) + helps.LogWithRequestID(ctx).Warn(msg) + b = []byte(msg) + } + helps.AppendAPIResponseChunk(ctx, e.cfg, b) + if errClose := errBody.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + return cliproxyexecutor.Response{}, statusErr{code: resp.StatusCode, msg: string(b)} + } + decodedBody, err := decodeResponseBody(resp.Body, resp.Header.Get("Content-Encoding")) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + return cliproxyexecutor.Response{}, err + } + defer func() { + if errClose := decodedBody.Close(); errClose != nil { + log.Errorf("response body close error: %v", errClose) + } + }() + data, err := io.ReadAll(decodedBody) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return cliproxyexecutor.Response{}, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + count := gjson.GetBytes(data, "input_tokens").Int() + out := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, data) + return cliproxyexecutor.Response{Payload: out, Headers: resp.Header.Clone()}, nil +} diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 3f2eb7e56..82d465918 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -1,285 +1,6 @@ package executor -import ( - "bufio" - "bytes" - "context" - "crypto/sha256" - "encoding/hex" - "fmt" - "io" - "net/http" - "sort" - "strings" - "time" - - codexauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex" - internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" - "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" - "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" - cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" - sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" - log "github.com/sirupsen/logrus" - "github.com/tidwall/gjson" - "github.com/tidwall/sjson" - "github.com/tiktoken-go/tokenizer" - - "github.com/gin-gonic/gin" - "github.com/google/uuid" -) - -const ( - codexUserAgent = "codex-tui/0.135.0 (Mac OS 26.5.0; arm64) iTerm.app/3.6.10 (codex-tui; 0.135.0)" - codexOriginator = "codex-tui" - codexDefaultImageToolModel = "gpt-image-2" - codexResponsesLiteHeader = "X-OpenAI-Internal-Codex-Responses-Lite" - codexResponsesLiteMetadata = "client_metadata.ws_request_header_x_openai_internal_codex_responses_lite" -) - -var dataTag = []byte("data:") - -const codexIncompleteStreamMessage = "stream error: stream disconnected before completion: stream closed before response.completed" - -type codexIncompleteStreamError struct { - statusErr -} - -func newCodexIncompleteStreamError() codexIncompleteStreamError { - return codexIncompleteStreamError{statusErr: statusErr{ - code: http.StatusRequestTimeout, - msg: codexIncompleteStreamMessage, - }} -} - -func (codexIncompleteStreamError) IsRequestScoped() bool { - return true -} - -// Streamed Codex responses may emit response.output_item.done events while leaving -// response.completed.response.output empty. Keep the stream path aligned with the -// already-patched non-stream path by reconstructing response.output from those items. -func collectCodexOutputItemDone(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback *[][]byte) { - itemResult := gjson.GetBytes(eventData, "item") - if !itemResult.Exists() || itemResult.Type != gjson.JSON { - return - } - outputIndexResult := gjson.GetBytes(eventData, "output_index") - if outputIndexResult.Exists() { - outputItemsByIndex[outputIndexResult.Int()] = []byte(itemResult.Raw) - return - } - *outputItemsFallback = append(*outputItemsFallback, []byte(itemResult.Raw)) -} - -func patchCodexCompletedOutput(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { - outputResult := gjson.GetBytes(eventData, "response.output") - shouldPatchOutput := (!outputResult.Exists() || !outputResult.IsArray() || len(outputResult.Array()) == 0) && (len(outputItemsByIndex) > 0 || len(outputItemsFallback) > 0) - if !shouldPatchOutput { - return eventData - } - - indexes := make([]int64, 0, len(outputItemsByIndex)) - for idx := range outputItemsByIndex { - indexes = append(indexes, idx) - } - sort.Slice(indexes, func(i, j int) bool { - return indexes[i] < indexes[j] - }) - - items := make([][]byte, 0, len(outputItemsByIndex)+len(outputItemsFallback)) - for _, idx := range indexes { - items = append(items, outputItemsByIndex[idx]) - } - items = append(items, outputItemsFallback...) - - outputArray := []byte("[]") - if len(items) > 0 { - var buf bytes.Buffer - totalLen := 2 - for _, item := range items { - totalLen += len(item) - } - if len(items) > 1 { - totalLen += len(items) - 1 - } - buf.Grow(totalLen) - buf.WriteByte('[') - for i, item := range items { - if i > 0 { - buf.WriteByte(',') - } - buf.Write(item) - } - buf.WriteByte(']') - outputArray = buf.Bytes() - } - - completedDataPatched, _ := sjson.SetRawBytes(eventData, "response.output", outputArray) - return completedDataPatched -} - -func codexTerminalStreamContextLengthErr(eventData []byte) (statusErr, bool) { - streamErr, body, ok := codexTerminalStreamErr(eventData) - if !ok || !codexTerminalErrorIsContextLength(body) { - return statusErr{}, false - } - return streamErr, true -} - -func codexTerminalStreamErr(eventData []byte) (statusErr, []byte, bool) { - body, ok := codexTerminalFailureBody(eventData) - if !ok || !codexTerminalStreamErrShouldHandle(body) { - return statusErr{}, nil, false - } - return newCodexStatusErr(http.StatusBadRequest, body), body, true -} - -func codexTerminalFailureErr(eventData []byte) (statusErr, []byte, bool) { - if streamErr, body, ok := codexTerminalStreamErr(eventData); ok { - return streamErr, body, true - } - body, ok := codexTerminalFailureBody(eventData) - if !ok { - return statusErr{}, nil, false - } - return newCodexStatusErr(codexTerminalFailureStatus(body), body), body, true -} - -func codexTerminalFailureStatus(body []byte) int { - for _, path := range []string{"error.status_code", "error.status"} { - if status := int(gjson.GetBytes(body, path).Int()); status >= 400 && status <= 599 { - return status - } - } - - errorType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.type").String())) - errorCode := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.code").String())) - switch { - case errorType == "invalid_request_error", errorType == "bad_request_error": - return http.StatusBadRequest - case errorType == "authentication_error", errorCode == "invalid_api_key", errorCode == "unauthorized": - return http.StatusUnauthorized - case errorType == "permission_error", errorCode == "forbidden", errorCode == "permission_denied": - return http.StatusForbidden - case errorType == "not_found_error", errorCode == "not_found", errorCode == "model_not_found": - return http.StatusNotFound - case errorType == "rate_limit_error", errorCode == "rate_limit_exceeded": - return http.StatusTooManyRequests - default: - return http.StatusBadGateway - } -} - -func codexTerminalFailureBody(eventData []byte) ([]byte, bool) { - eventType := gjson.GetBytes(eventData, "type").String() - var body []byte - switch eventType { - case "error": - body = codexTerminalErrorBody(eventData, "error") - if len(body) == 0 { - body = codexTerminalTopLevelErrorBody(eventData) - } - case "response.failed": - body = codexTerminalErrorBody(eventData, "response.error") - if len(body) == 0 { - body = codexTerminalErrorBody(eventData, "error") - } - default: - return nil, false - } - if len(body) == 0 { - body = []byte(`{"error":{"message":"upstream stream failed without error details"}}`) - } - return body, true -} - -func codexTerminalStreamErrShouldHandle(body []byte) bool { - if codexTerminalErrorIsContextLength(body) { - return true - } - if isCodexUsageLimitError(body) || isCodexModelCapacityError(body) { - return true - } - code, _, ok := codexStatusErrorClassification(http.StatusBadRequest, body) - return ok && code == "thinking_signature_invalid" -} - -func codexTerminalErrorBody(eventData []byte, path string) []byte { - errorResult := gjson.GetBytes(eventData, path) - if !errorResult.Exists() { - return nil - } - body := []byte(`{"error":{}}`) - if errorResult.Type == gjson.JSON { - body, _ = sjson.SetRawBytes(body, "error", []byte(errorResult.Raw)) - } else if message := strings.TrimSpace(errorResult.String()); message != "" { - body, _ = sjson.SetBytes(body, "error.message", message) - } - if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { - if message := strings.TrimSpace(gjson.GetBytes(eventData, "response.error.message").String()); message != "" { - body, _ = sjson.SetBytes(body, "error.message", message) - } - } - if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { - if code := strings.TrimSpace(gjson.GetBytes(body, "error.code").String()); code != "" { - body, _ = sjson.SetBytes(body, "error.message", code) - } - } - if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { - if errorType := strings.TrimSpace(gjson.GetBytes(body, "error.type").String()); errorType != "" { - body, _ = sjson.SetBytes(body, "error.message", errorType) - } - } - return body -} - -func codexTerminalTopLevelErrorBody(eventData []byte) []byte { - message := strings.TrimSpace(gjson.GetBytes(eventData, "message").String()) - code := strings.TrimSpace(gjson.GetBytes(eventData, "code").String()) - errorType := strings.TrimSpace(gjson.GetBytes(eventData, "error_type").String()) - param := strings.TrimSpace(gjson.GetBytes(eventData, "param").String()) - if message == "" && code == "" && errorType == "" && param == "" { - return nil - } - - body := []byte(`{"error":{}}`) - if message != "" { - body, _ = sjson.SetBytes(body, "error.message", message) - } - if code != "" { - body, _ = sjson.SetBytes(body, "error.code", code) - } - if errorType != "" { - body, _ = sjson.SetBytes(body, "error.type", errorType) - } - if param != "" { - body, _ = sjson.SetBytes(body, "error.param", param) - } - if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { - if code != "" { - body, _ = sjson.SetBytes(body, "error.message", code) - } else if errorType != "" { - body, _ = sjson.SetBytes(body, "error.message", errorType) - } - } - return body -} - -func codexTerminalErrorIsContextLength(body []byte) bool { - errorCode := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.code").String())) - message := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.message").String())) - return errorCode == "context_length_exceeded" || - errorCode == "context_too_large" || - strings.Contains(message, "context window") || - strings.Contains(message, "context length") || - strings.Contains(message, "too many tokens") -} +import "github.com/router-for-me/CLIProxyAPI/v7/internal/config" // CodexExecutor is a stateless executor for Codex (OpenAI Responses API entrypoint). // If api_key is unavailable on auth, it falls back to legacy via ClientAdapter. @@ -290,2072 +11,3 @@ type CodexExecutor struct { func NewCodexExecutor(cfg *config.Config) *CodexExecutor { return &CodexExecutor{cfg: cfg} } func (e *CodexExecutor) Identifier() string { return "codex" } - -func translateCodexRequestPair(from, to sdktranslator.Format, model string, originalPayload, payload []byte, stream bool) ([]byte, []byte) { - if bytes.Equal(originalPayload, payload) { - body := sdktranslator.TranslateRequest(from, to, model, payload, stream) - return body, body - } - originalTranslated := sdktranslator.TranslateRequest(from, to, model, originalPayload, stream) - body := sdktranslator.TranslateRequest(from, to, model, payload, stream) - return originalTranslated, body -} - -type codexReasoningReplayScope struct { - modelName string - sessionKey string - requestFingerprint string -} - -func (s codexReasoningReplayScope) valid() bool { - return strings.TrimSpace(s.modelName) != "" && strings.TrimSpace(s.sessionKey) != "" -} - -func applyCodexReasoningReplayCache(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) ([]byte, codexReasoningReplayScope) { - updated, scope, _ := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) - return updated, scope -} - -func applyCodexReasoningReplayCacheRequired(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) ([]byte, codexReasoningReplayScope, error) { - scope := codexReasoningReplayScopeFromRequest(ctx, from, req, opts, body) - if !scope.valid() { - return body, scope, nil - } - items, ok, errReplay := internalcache.GetCodexReasoningReplayItemsRequired(ctx, scope.modelName, scope.sessionKey) - if errReplay != nil || !ok { - return body, scope, errReplay - } - updated, ok := insertCodexReasoningReplayTurns(body, items) - if !ok { - return body, scope, nil - } - return updated, scope, nil -} - -func codexReasoningReplayScopeFromRequest(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) codexReasoningReplayScope { - if !codexReasoningReplayEnabledForSource(from) { - return codexReasoningReplayScope{} - } - modelName := strings.TrimSpace(gjson.GetBytes(body, "model").String()) - if modelName == "" { - modelName = thinking.ParseSuffix(req.Model).ModelName - } - inputItems := gjson.GetBytes(body, "input").Array() - return codexReasoningReplayScope{ - modelName: modelName, - sessionKey: codexReasoningReplaySessionKey(ctx, from, req, opts, body), - requestFingerprint: codexReplayInputPrefixFingerprint(inputItems, len(inputItems)), - } -} - -func codexReasoningReplayEnabledForSource(from sdktranslator.Format) bool { - return sourceFormatEqual(from, sdktranslator.FormatClaude) -} - -func sourceFormatEqual(from, want sdktranslator.Format) bool { - return strings.EqualFold(strings.TrimSpace(from.String()), want.String()) -} - -func codexClaudeCodeReplaySessionKey(ctx context.Context, payload []byte, headers http.Header) string { - sessionKey, _ := helps.ClaudeCodeExecutionScope(ctx, payload, headers) - return sessionKey -} - -func codexReasoningReplaySessionKey(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) string { - if ctx == nil { - ctx = context.Background() - } - if sourceFormatEqual(from, sdktranslator.FormatClaude) { - if sessionKey := codexClaudeCodeReplaySessionKey(ctx, req.Payload, opts.Headers); sessionKey != "" { - return sessionKey - } - } - if value := metadataString(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { - return "execution:" + value - } - if value := metadataString(req.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { - return "execution:" + value - } - if value := codexReasoningReplaySessionKeyFromPayload(body); value != "" { - return value - } - if value := codexReasoningReplaySessionKeyFromPayload(req.Payload); value != "" { - return value - } - if value := codexReasoningReplaySessionKeyFromHeaders(opts.Headers); value != "" { - return value - } - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - if value := codexReasoningReplaySessionKeyFromHeaders(ginCtx.Request.Header); value != "" { - return value - } - } - if sourceFormatEqual(from, sdktranslator.FormatOpenAI) { - if apiKey := strings.TrimSpace(helps.APIKeyFromContext(ctx)); apiKey != "" { - return "prompt-cache:" + uuid.NewSHA1(uuid.NameSpaceOID, []byte("cli-proxy-api:codex:prompt-cache:"+apiKey)).String() - } - } - return "" -} - -func metadataString(metadata map[string]any, key string) string { - if len(metadata) == 0 { - return "" - } - raw, ok := metadata[key] - if !ok || raw == nil { - return "" - } - switch v := raw.(type) { - case string: - return strings.TrimSpace(v) - case []byte: - return strings.TrimSpace(string(v)) - default: - return "" - } -} - -func codexReasoningReplaySessionKeyFromPayload(payload []byte) string { - if len(payload) == 0 { - return "" - } - if promptCacheKey := strings.TrimSpace(gjson.GetBytes(payload, "prompt_cache_key").String()); promptCacheKey != "" { - return "prompt-cache:" + promptCacheKey - } - if windowID := strings.TrimSpace(gjson.GetBytes(payload, "client_metadata.x-codex-window-id").String()); windowID != "" { - return "window:" + windowID - } - if turnMetadata := strings.TrimSpace(gjson.GetBytes(payload, "client_metadata.x-codex-turn-metadata").String()); turnMetadata != "" { - return codexReasoningReplaySessionKeyFromTurnMetadata(turnMetadata) - } - return "" -} - -func codexReasoningReplaySessionKeyFromHeaders(headers http.Header) string { - if headers == nil { - return "" - } - if turnMetadata := strings.TrimSpace(headers.Get("X-Codex-Turn-Metadata")); turnMetadata != "" { - if key := codexReasoningReplaySessionKeyFromTurnMetadata(turnMetadata); key != "" { - return key - } - } - if windowID := strings.TrimSpace(headerValueCaseInsensitive(headers, "X-Codex-Window-Id")); windowID != "" { - return "window:" + windowID - } - for _, headerName := range []string{"Session_id", "session_id", "Session-Id"} { - if value := strings.TrimSpace(headerValueCaseInsensitive(headers, headerName)); value != "" { - return "session-id:" + value - } - } - if conversationID := strings.TrimSpace(headerValueCaseInsensitive(headers, "Conversation_id")); conversationID != "" { - return "conversation_id:" + conversationID - } - return "" -} - -func codexReasoningReplaySessionKeyFromTurnMetadata(turnMetadata string) string { - if promptCacheKey := strings.TrimSpace(gjson.Get(turnMetadata, "prompt_cache_key").String()); promptCacheKey != "" { - return "prompt-cache:" + promptCacheKey - } - if windowID := strings.TrimSpace(gjson.Get(turnMetadata, "window_id").String()); windowID != "" { - return "window:" + windowID - } - return "" -} - -func codexInputHasValidReasoningEncryptedContent(body []byte) bool { - input := gjson.GetBytes(body, "input") - if !input.IsArray() { - return false - } - for _, item := range input.Array() { - if strings.TrimSpace(item.Get("type").String()) != "reasoning" { - continue - } - encryptedContent := item.Get("encrypted_content") - if encryptedContent.Type != gjson.String { - continue - } - if _, err := signature.InspectGPTReasoningSignature(encryptedContent.String()); err == nil { - return true - } - } - return false -} - -type codexReasoningReplayTurn struct { - marked bool - assistantFingerprint string - requestFingerprint string - callIDs []string - items [][]byte -} - -func insertCodexReasoningReplayTurns(body []byte, replayItems [][]byte) ([]byte, bool) { - input := gjson.GetBytes(body, "input") - if !input.IsArray() || len(replayItems) == 0 { - return body, false - } - inputItems := input.Array() - turns := splitCodexReasoningReplayTurns(replayItems) - insertions := make(map[int][][]byte) - usedAnchorIndexes := make(map[int]bool) - fallbackAnchorEnd := len(inputItems) - 1 - inserted := false - for turnIndex := len(turns) - 1; turnIndex >= 0; turnIndex-- { - turn := turns[turnIndex] - if len(turn.items) == 0 { - continue - } - if !turn.marked { - items := filterCodexReasoningReplayItemsForInput(body, turn.items) - if len(items) == 0 { - continue - } - index := codexReasoningReplayInsertIndex(inputItems, items) - items = codexAlignReasoningReplayToolCallIDs(inputItems, items) - insertions[index] = append(items, insertions[index]...) - inserted = true - continue - } - - anchorIndex, matched := codexReasoningReplayTurnAnchorIndex(inputItems, turn, fallbackAnchorEnd, usedAnchorIndexes) - if !matched { - continue - } - usedAnchorIndexes[anchorIndex] = true - if turn.requestFingerprint == "" { - fallbackAnchorEnd = anchorIndex - 1 - } - items := filterCodexReasoningReplayTurnItems(inputItems, turn.items) - if len(items) == 0 { - continue - } - items = codexAlignReasoningReplayToolCallIDs(inputItems, items) - insertions[anchorIndex] = append(items, insertions[anchorIndex]...) - inserted = true - } - if !inserted { - return body, false - } - - items := make([]string, 0, len(inputItems)+len(replayItems)) - for index, inputItem := range inputItems { - for _, replayItem := range insertions[index] { - items = append(items, string(replayItem)) - } - items = append(items, inputItem.Raw) - } - for _, replayItem := range insertions[len(inputItems)] { - items = append(items, string(replayItem)) - } - updated, err := sjson.SetRawBytes(body, "input", []byte("["+strings.Join(items, ",")+"]")) - if err != nil { - return body, false - } - return updated, true -} - -func splitCodexReasoningReplayTurns(items [][]byte) []codexReasoningReplayTurn { - turns := make([]codexReasoningReplayTurn, 0) - current := codexReasoningReplayTurn{} - appendCurrent := func() { - if len(current.items) > 0 { - turns = append(turns, current) - } - } - for _, item := range items { - itemResult := gjson.ParseBytes(item) - if strings.TrimSpace(itemResult.Get("type").String()) == internalcache.CodexReasoningReplayTurnType { - appendCurrent() - current = codexReasoningReplayTurn{ - marked: true, - assistantFingerprint: strings.TrimSpace(itemResult.Get("assistant_fingerprint").String()), - requestFingerprint: strings.TrimSpace(itemResult.Get("request_fingerprint").String()), - } - if callIDs := itemResult.Get("call_ids"); callIDs.IsArray() { - for _, callIDResult := range callIDs.Array() { - if callID := strings.TrimSpace(callIDResult.String()); callID != "" { - current.callIDs = append(current.callIDs, callID) - } - } - } - continue - } - current.items = append(current.items, item) - } - appendCurrent() - return turns -} - -func codexReasoningReplayTurnAnchorIndex(inputItems []gjson.Result, turn codexReasoningReplayTurn, fallbackEnd int, used map[int]bool) (int, bool) { - searchEnd := fallbackEnd - if turn.requestFingerprint != "" { - searchEnd = len(inputItems) - 1 - } - if searchEnd >= len(inputItems) { - searchEnd = len(inputItems) - 1 - } - matchesRequestPrefix := func(index int) bool { - return turn.requestFingerprint == "" || codexReplayInputPrefixFingerprint(inputItems, index) == turn.requestFingerprint - } - if len(turn.callIDs) > 0 { - callIDs := make(map[string]bool) - for _, callID := range turn.callIDs { - for _, candidate := range codexReplayComparableCallIDs(callID) { - callIDs[candidate] = true - } - } - for index := searchEnd; index >= 0; index-- { - if used[index] || !matchesRequestPrefix(index) { - continue - } - itemType := strings.TrimSpace(inputItems[index].Get("type").String()) - if itemType != "function_call" && itemType != "custom_tool_call" && itemType != "function_call_output" && itemType != "custom_tool_call_output" { - continue - } - for _, candidate := range codexReplayComparableCallIDs(inputItems[index].Get("call_id").String()) { - if callIDs[candidate] { - return index, true - } - } - } - } - if turn.assistantFingerprint != "" { - for index := searchEnd; index >= 0; index-- { - if used[index] || !matchesRequestPrefix(index) { - continue - } - if codexReplayAssistantMessageFingerprint(inputItems[index]) == turn.assistantFingerprint { - return index, true - } - } - } - if len(turn.callIDs) == 0 && turn.assistantFingerprint == "" { - return codexReasoningReplayInsertIndex(inputItems, turn.items), true - } - return 0, false -} - -func filterCodexReasoningReplayTurnItems(inputItems []gjson.Result, items [][]byte) [][]byte { - existingReasoning := make(map[string]bool) - existingCalls := make(map[string]bool) - existingOutputs := make(map[string]bool) - for _, inputItem := range inputItems { - itemType := strings.TrimSpace(inputItem.Get("type").String()) - switch itemType { - case "reasoning": - if encryptedContent := strings.TrimSpace(inputItem.Get("encrypted_content").String()); encryptedContent != "" { - existingReasoning[encryptedContent] = true - } - case "function_call_output", "custom_tool_call_output": - for _, candidate := range codexReplayComparableCallIDs(inputItem.Get("call_id").String()) { - existingOutputs[candidate] = true - } - } - for _, key := range codexReplayToolCallKeys(inputItem) { - existingCalls[key] = true - } - } - - filtered := make([][]byte, 0, len(items)) - for _, item := range items { - itemResult := gjson.ParseBytes(item) - switch strings.TrimSpace(itemResult.Get("type").String()) { - case "reasoning": - if existingReasoning[strings.TrimSpace(itemResult.Get("encrypted_content").String())] { - continue - } - case "function_call", "custom_tool_call": - keys := codexReplayToolCallKeys(itemResult) - if len(keys) == 0 || codexReplayAnyToolCallKeyExists(existingCalls, keys) { - continue - } - hasMatchingOutput := false - for _, candidate := range codexReplayComparableCallIDs(itemResult.Get("call_id").String()) { - if existingOutputs[candidate] { - hasMatchingOutput = true - break - } - } - if !hasMatchingOutput { - continue - } - for _, key := range keys { - existingCalls[key] = true - } - default: - continue - } - filtered = append(filtered, item) - } - return filtered -} - -func codexReplayAssistantMessageFingerprint(item gjson.Result) string { - itemType := strings.TrimSpace(item.Get("type").String()) - if itemType != "" && itemType != "message" { - return "" - } - if !strings.EqualFold(strings.TrimSpace(item.Get("role").String()), "assistant") { - return "" - } - content := item.Get("content") - var builder strings.Builder - if content.Type == gjson.String { - builder.WriteString(content.String()) - } else if content.IsArray() { - for _, part := range content.Array() { - switch strings.TrimSpace(part.Get("type").String()) { - case "input_text", "output_text": - builder.WriteString(part.Get("text").String()) - case "refusal": - builder.WriteString("\x00refusal\x00") - builder.WriteString(part.Get("refusal").String()) - default: - return "" - } - } - } else { - return "" - } - if builder.Len() == 0 { - return "" - } - sum := sha256.Sum256([]byte(builder.String())) - return hex.EncodeToString(sum[:]) -} - -func codexReplayInputPrefixFingerprint(inputItems []gjson.Result, end int) string { - if end < 0 || end > len(inputItems) { - return "" - } - hasher := sha256.New() - for index := 0; index < end; index++ { - _, _ = hasher.Write([]byte("\x00item\x00")) - _, _ = hasher.Write([]byte(inputItems[index].Raw)) - } - return hex.EncodeToString(hasher.Sum(nil)) -} - -func filterCodexReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte { - input := gjson.GetBytes(body, "input") - if !input.IsArray() { - return nil - } - - hasInputReasoning := codexInputHasValidReasoningEncryptedContent(body) - existingCalls := make(map[string]bool) - existingOutputs := make(map[string]bool) - for _, inputItem := range input.Array() { - itemType := strings.TrimSpace(inputItem.Get("type").String()) - if itemType == "function_call_output" || itemType == "custom_tool_call_output" { - callID := strings.TrimSpace(inputItem.Get("call_id").String()) - if callID != "" { - for _, candidate := range codexReplayComparableCallIDs(callID) { - existingOutputs[candidate] = true - } - } - } - for _, key := range codexReplayToolCallKeys(inputItem) { - existingCalls[key] = true - } - } - - filtered := make([][]byte, 0, len(items)) - for _, item := range items { - itemResult := gjson.ParseBytes(item) - switch strings.TrimSpace(itemResult.Get("type").String()) { - case "reasoning": - if hasInputReasoning { - continue - } - case "function_call", "custom_tool_call": - keys := codexReplayToolCallKeys(itemResult) - if len(keys) == 0 || codexReplayAnyToolCallKeyExists(existingCalls, keys) { - continue - } - // Only inject if there is a matching output in the request - hasMatchingOutput := false - callID := strings.TrimSpace(itemResult.Get("call_id").String()) - if callID != "" { - for _, candidate := range codexReplayComparableCallIDs(callID) { - if existingOutputs[candidate] { - hasMatchingOutput = true - break - } - } - } - if !hasMatchingOutput { - continue - } - for _, key := range keys { - existingCalls[key] = true - } - default: - continue - } - filtered = append(filtered, item) - } - return filtered -} - -func insertCodexReasoningReplayItems(body []byte, replayItems [][]byte) ([]byte, bool) { - input := gjson.GetBytes(body, "input") - if !input.IsArray() || len(replayItems) == 0 { - return body, false - } - inputItems := input.Array() - insertIndex := codexReasoningReplayInsertIndex(inputItems, replayItems) - replayItems = codexAlignReasoningReplayToolCallIDs(inputItems, replayItems) - items := make([]string, 0, len(inputItems)+len(replayItems)) - for i, inputItem := range inputItems { - if i == insertIndex { - for _, replayItem := range replayItems { - items = append(items, string(replayItem)) - } - } - items = append(items, inputItem.Raw) - } - if insertIndex == len(inputItems) { - for _, replayItem := range replayItems { - items = append(items, string(replayItem)) - } - } - updated, err := sjson.SetRawBytes(body, "input", []byte("["+strings.Join(items, ",")+"]")) - if err != nil { - return body, false - } - return updated, true -} - -func codexReasoningReplayInsertIndex(inputItems []gjson.Result, replayItems [][]byte) int { - replayCallIDs := make(map[string]bool) - for _, replayItem := range replayItems { - itemResult := gjson.ParseBytes(replayItem) - itemType := strings.TrimSpace(itemResult.Get("type").String()) - if itemType != "function_call" && itemType != "custom_tool_call" { - continue - } - for _, callID := range codexReplayComparableCallIDs(itemResult.Get("call_id").String()) { - replayCallIDs[callID] = true - } - } - if len(replayCallIDs) > 0 { - for index, inputItem := range inputItems { - itemType := strings.TrimSpace(inputItem.Get("type").String()) - if itemType != "function_call_output" && itemType != "custom_tool_call_output" { - continue - } - callID := strings.TrimSpace(inputItem.Get("call_id").String()) - if callID == "" || replayCallIDs[callID] { - return index - } - } - } - for index := len(inputItems) - 1; index >= 0; index-- { - inputItem := inputItems[index] - if role, ok := codexReplayMessageRole(inputItem); ok && role == "assistant" { - return index - } - } - for index, inputItem := range inputItems { - if shouldInsertCodexReasoningReplayBefore(inputItem) { - return index - } - } - return len(inputItems) -} - -func codexAlignReasoningReplayToolCallIDs(inputItems []gjson.Result, replayItems [][]byte) [][]byte { - outputCallIDs := codexReplayOutputCallIDs(inputItems) - if len(outputCallIDs) == 0 { - return replayItems - } - - aligned := make([][]byte, 0, len(replayItems)) - for _, replayItem := range replayItems { - itemResult := gjson.ParseBytes(replayItem) - itemType := strings.TrimSpace(itemResult.Get("type").String()) - if itemType != "function_call" && itemType != "custom_tool_call" { - aligned = append(aligned, replayItem) - continue - } - - callID := strings.TrimSpace(itemResult.Get("call_id").String()) - outputCallID := "" - for _, candidate := range codexReplayComparableCallIDs(callID) { - if value := outputCallIDs[candidate]; value != "" { - outputCallID = value - break - } - } - if outputCallID == "" || outputCallID == callID { - aligned = append(aligned, replayItem) - continue - } - - updated, err := sjson.SetBytes(replayItem, "call_id", outputCallID) - if err != nil { - aligned = append(aligned, replayItem) - continue - } - aligned = append(aligned, updated) - } - return aligned -} - -func codexReplayOutputCallIDs(inputItems []gjson.Result) map[string]string { - outputCallIDs := make(map[string]string) - for _, inputItem := range inputItems { - itemType := strings.TrimSpace(inputItem.Get("type").String()) - if itemType != "function_call_output" && itemType != "custom_tool_call_output" { - continue - } - callID := strings.TrimSpace(inputItem.Get("call_id").String()) - if callID == "" { - continue - } - for _, candidate := range codexReplayComparableCallIDs(callID) { - outputCallIDs[candidate] = callID - } - } - return outputCallIDs -} - -func shouldInsertCodexReasoningReplayBefore(item gjson.Result) bool { - role, ok := codexReplayMessageRole(item) - if !ok { - return true - } - switch role { - case "developer", "system": - return false - default: - return true - } -} - -func codexReplayMessageRole(item gjson.Result) (string, bool) { - itemType := strings.TrimSpace(item.Get("type").String()) - role := strings.ToLower(strings.TrimSpace(item.Get("role").String())) - if role == "" || (itemType != "" && itemType != "message") { - return "", false - } - return role, true -} - -func codexReplayToolCallKeys(item gjson.Result) []string { - itemType := strings.TrimSpace(item.Get("type").String()) - if itemType != "function_call" && itemType != "custom_tool_call" { - return nil - } - callIDs := codexReplayComparableCallIDs(item.Get("call_id").String()) - if len(callIDs) == 0 { - return nil - } - keys := make([]string, 0, len(callIDs)) - for _, callID := range callIDs { - keys = append(keys, itemType+":"+callID) - } - return keys -} - -func codexReplayAnyToolCallKeyExists(existing map[string]bool, keys []string) bool { - for _, key := range keys { - if existing[key] { - return true - } - } - return false -} - -func codexReplayComparableCallIDs(callID string) []string { - callID = strings.TrimSpace(callID) - if callID == "" { - return nil - } - - claudeVisibleCallID := shortenCodexReplayCallIDIfNeeded(util.SanitizeClaudeToolID(callID)) - if claudeVisibleCallID == "" || claudeVisibleCallID == callID { - return []string{callID} - } - return []string{callID, claudeVisibleCallID} -} - -func shortenCodexReplayCallIDIfNeeded(id string) string { - const limit = 64 - if len(id) <= limit { - return id - } - - sum := sha256.Sum256([]byte(id)) - suffix := "_" + hex.EncodeToString(sum[:8]) - prefixLen := limit - len(suffix) - if prefixLen <= 0 { - return suffix[len(suffix)-limit:] - } - return id[:prefixLen] + suffix -} - -func cacheCodexReasoningReplayFromCompleted(scope codexReasoningReplayScope, completedData []byte) { - if !scope.valid() { - return - } - output := gjson.GetBytes(completedData, "response.output") - if !output.IsArray() { - return - } - replayItems := make([][]byte, 0, len(output.Array())) - callIDs := make([]string, 0) - assistantFingerprint := "" - for _, item := range output.Array() { - switch strings.TrimSpace(item.Get("type").String()) { - case "reasoning": - replayItems = append(replayItems, []byte(item.Raw)) - case "function_call", "custom_tool_call": - replayItems = append(replayItems, []byte(item.Raw)) - if callID := strings.TrimSpace(item.Get("call_id").String()); callID != "" { - callIDs = append(callIDs, callID) - } - case "message": - if fingerprint := codexReplayAssistantMessageFingerprint(item); fingerprint != "" { - assistantFingerprint = fingerprint - } - } - } - if len(replayItems) == 0 { - return - } - - hasher := sha256.New() - _, _ = hasher.Write([]byte(scope.requestFingerprint)) - _, _ = hasher.Write([]byte("\x00assistant\x00" + assistantFingerprint)) - for _, callID := range callIDs { - _, _ = hasher.Write([]byte("\x00call\x00" + callID)) - } - for _, item := range replayItems { - _, _ = hasher.Write([]byte("\x00item\x00")) - _, _ = hasher.Write(item) - } - marker := []byte(`{"type":"` + internalcache.CodexReasoningReplayTurnType + `"}`) - marker, _ = sjson.SetBytes(marker, "id", hex.EncodeToString(hasher.Sum(nil))) - if assistantFingerprint != "" { - marker, _ = sjson.SetBytes(marker, "assistant_fingerprint", assistantFingerprint) - } - if scope.requestFingerprint != "" { - marker, _ = sjson.SetBytes(marker, "request_fingerprint", scope.requestFingerprint) - } - for _, callID := range callIDs { - marker, _ = sjson.SetBytes(marker, "call_ids.-1", callID) - } - items := make([][]byte, 0, len(replayItems)+1) - items = append(items, marker) - items = append(items, replayItems...) - internalcache.AppendCodexReasoningReplayItemsBestEffort(context.Background(), scope.modelName, scope.sessionKey, items) -} - -func clearCodexReasoningReplayOnInvalidSignature(ctx context.Context, scope codexReasoningReplayScope, statusCode int, body []byte) error { - if !scope.valid() { - return nil - } - code, _, ok := codexStatusErrorClassification(statusCode, body) - if ok && code == "thinking_signature_invalid" { - return internalcache.DeleteCodexReasoningReplayItemRequired(ctx, scope.modelName, scope.sessionKey) - } - return nil -} - -// PrepareRequest injects Codex credentials into the outgoing HTTP request. -func (e *CodexExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { - if req == nil { - return nil - } - apiKey, _ := codexCreds(auth) - if strings.TrimSpace(apiKey) != "" { - req.Header.Set("Authorization", "Bearer "+apiKey) - } - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(req, attrs) - return nil -} - -// HttpRequest injects Codex credentials into the request and executes it. -func (e *CodexExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) { - if req == nil { - return nil, fmt.Errorf("codex executor: request is nil") - } - if ctx == nil { - ctx = req.Context() - } - httpReq := req.WithContext(ctx) - if err := e.PrepareRequest(httpReq, auth); err != nil { - return nil, err - } - httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) - return httpClient.Do(httpReq) -} - -func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - if opts.Alt == "responses/compact" { - return e.executeCompact(ctx, auth, req, opts) - } - if isCodexOpenAIImageRequest(opts) { - return e.executeOpenAIImage(ctx, auth, req, opts) - } - baseModel := thinking.ParseSuffix(req.Model).ModelName - - apiKey, baseURL := codexCreds(auth) - if baseURL == "" { - baseURL = "https://chatgpt.com/backend-api/codex" - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("codex") - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, false) - - body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return resp, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body = helps.SetStringIfDifferent(body, "model", baseModel) - body = helps.SetBoolIfDifferent(body, "stream", true) - body, _ = sjson.DeleteBytes(body, "previous_response_id") - body, _ = sjson.DeleteBytes(body, "generate") - body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") - body, _ = sjson.DeleteBytes(body, "safety_identifier") - body, _ = sjson.DeleteBytes(body, "stream_options") - body = normalizeCodexInstructions(body) - if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) - } - body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) - body = normalizeCodexParallelToolCalls(body, opts.Headers) - body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) - body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) - if errReplay != nil { - return resp, errReplay - } - reporter.SetTranslatedReasoningEffort(body, to.String()) - - url := strings.TrimSuffix(baseURL, "/") + "/responses" - var identityState codexIdentityConfuseState - httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) - if err != nil { - return resp, err - } - applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) - applyModelHeaderOverrides(httpReq.Header, baseModel) - applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: url, - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: upstreamBody, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("codex executor: close response body error: %v", errClose) - } - }() - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) - b = applyCodexIdentityConfuseResponsePayload(b, identityState) - if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClearReplay != nil { - return resp, errClearReplay - } - helps.AppendAPIResponseChunk(ctx, e.cfg, b) - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) - err = newCodexStatusErr(httpResp.StatusCode, b) - return resp, err - } - data, errRead := io.ReadAll(httpResp.Body) - upstreamData := applyCodexIdentityConfuseResponsePayload(data, identityState) - helps.AppendAPIResponseChunk(ctx, e.cfg, upstreamData) - - lines := bytes.Split(upstreamData, []byte("\n")) - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte - for _, line := range lines { - if !bytes.HasPrefix(line, dataTag) { - continue - } - - eventData := bytes.TrimSpace(line[5:]) - eventData = helps.RestoreCodexMultiAgentV2Response(eventData, optimizeMultiAgentV2) - eventType := gjson.GetBytes(eventData, "type").String() - - if streamErr, terminalBody, ok := codexTerminalFailureErr(eventData); ok { - if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { - return resp, errClearReplay - } - err = streamErr - return resp, err - } - - if eventType == "response.output_item.done" { - itemResult := gjson.GetBytes(eventData, "item") - if !itemResult.Exists() || itemResult.Type != gjson.JSON { - continue - } - outputIndexResult := gjson.GetBytes(eventData, "output_index") - if outputIndexResult.Exists() { - outputItemsByIndex[outputIndexResult.Int()] = []byte(itemResult.Raw) - } else { - outputItemsFallback = append(outputItemsFallback, []byte(itemResult.Raw)) - } - continue - } - - if eventType != "response.completed" && eventType != "response.incomplete" { - continue - } - - if detail, ok := helps.ParseCodexUsage(eventData); ok { - reporter.Publish(ctx, detail) - } - publishCodexImageToolUsage(ctx, reporter, body, eventData) - - completedData := patchCodexCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) - if eventType == "response.completed" { - cacheCodexReasoningReplayFromCompleted(replayScope, completedData) - } - - var param any - clientCompletedData := applyCodexIdentityExposeResponsePayload(completedData, identityState) - out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, clientCompletedData, ¶m) - resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} - return resp, nil - } - if errRead != nil { - if errCtx := ctx.Err(); errCtx != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errCtx) - err = errCtx - return resp, err - } - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - } - err = newCodexIncompleteStreamError() - return resp, err -} - -func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - baseModel := thinking.ParseSuffix(req.Model).ModelName - - apiKey, baseURL := codexCreds(auth) - if baseURL == "" { - baseURL = "https://chatgpt.com/backend-api/codex" - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("openai-response") - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, false) - - body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return resp, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body = helps.SetStringIfDifferent(body, "model", baseModel) - body, _ = sjson.DeleteBytes(body, "stream") - body = normalizeCodexInstructions(body) - body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) - body = normalizeCodexParallelToolCalls(body, opts.Headers) - body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) - reporter.SetTranslatedReasoningEffort(body, to.String()) - - url := strings.TrimSuffix(baseURL, "/") + "/responses/compact" - var identityState codexIdentityConfuseState - httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) - if err != nil { - return resp, err - } - applyCodexHeaders(httpReq, auth, apiKey, false, e.cfg) - applyModelHeaderOverrides(httpReq.Header, baseModel) - applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: url, - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: upstreamBody, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("codex executor: close response body error: %v", errClose) - } - }() - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - b, _ := io.ReadAll(httpResp.Body) - b = applyCodexIdentityConfuseResponsePayload(b, identityState) - helps.AppendAPIResponseChunk(ctx, e.cfg, b) - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) - err = newCodexStatusErr(httpResp.StatusCode, b) - return resp, err - } - data, err := io.ReadAll(httpResp.Body) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - upstreamData := applyCodexIdentityConfuseResponsePayload(data, identityState) - helps.AppendAPIResponseChunk(ctx, e.cfg, upstreamData) - upstreamData = helps.RestoreCodexMultiAgentV2Response(upstreamData, optimizeMultiAgentV2) - reporter.Publish(ctx, helps.ParseOpenAIUsage(upstreamData)) - reporter.EnsurePublished(ctx) - var param any - clientData := applyCodexIdentityExposeResponsePayload(upstreamData, identityState) - out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, clientData, ¶m) - resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} - return resp, nil -} - -func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { - if opts.Alt == "responses/compact" { - return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} - } - if isCodexOpenAIImageRequest(opts) { - return e.executeOpenAIImageStream(ctx, auth, req, opts) - } - baseModel := thinking.ParseSuffix(req.Model).ModelName - - apiKey, baseURL := codexCreds(auth) - if baseURL == "" { - baseURL = "https://chatgpt.com/backend-api/codex" - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("codex") - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, true) - - body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return nil, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.DeleteBytes(body, "previous_response_id") - body, _ = sjson.DeleteBytes(body, "generate") - body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") - body, _ = sjson.DeleteBytes(body, "safety_identifier") - body, _ = sjson.DeleteBytes(body, "stream_options") - body = helps.SetStringIfDifferent(body, "model", baseModel) - body = normalizeCodexInstructions(body) - if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) - } - body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) - body = normalizeCodexParallelToolCalls(body, opts.Headers) - body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) - body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) - if errReplay != nil { - return nil, errReplay - } - reporter.SetTranslatedReasoningEffort(body, to.String()) - - url := strings.TrimSuffix(baseURL, "/") + "/responses" - var identityState codexIdentityConfuseState - httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) - if err != nil { - return nil, err - } - applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) - applyModelHeaderOverrides(httpReq.Header, baseModel) - applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: url, - Method: http.MethodPost, - Headers: httpReq.Header.Clone(), - Body: upstreamBody, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - - httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return nil, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - data, readErr := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("codex executor: close response body error: %v", errClose) - } - if readErr != nil { - helps.RecordAPIResponseError(ctx, e.cfg, readErr) - return nil, readErr - } - data = applyCodexIdentityConfuseResponsePayload(data, identityState) - if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, data); errClearReplay != nil { - return nil, errClearReplay - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) - err = newCodexStatusErr(httpResp.StatusCode, data) - return nil, err - } - out := make(chan cliproxyexecutor.StreamChunk) - go func() { - defer close(out) - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("codex executor: close response body error: %v", errClose) - } - }() - scanner := bufio.NewScanner(httpResp.Body) - scanner.Buffer(nil, 52_428_800) // 50MB - claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) - var param any - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte - for scanner.Scan() { - line := applyCodexIdentityConfuseResponsePayload(scanner.Bytes(), identityState) - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - translatedLine := bytes.Clone(line) - terminalSuccess := false - - if bytes.HasPrefix(line, dataTag) { - data := bytes.TrimSpace(line[5:]) - data = helps.RestoreCodexMultiAgentV2Response(data, optimizeMultiAgentV2) - translatedLine = append([]byte("data: "), data...) - eventType := gjson.GetBytes(data, "type").String() - if streamErr, terminalBody, ok := codexTerminalFailureErr(data); ok { - if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errClearReplay) - reporter.PublishFailure(ctx, errClearReplay) - select { - case out <- cliproxyexecutor.StreamChunk{Err: errClearReplay}: - case <-ctx.Done(): - } - return - } - helps.RecordAPIResponseError(ctx, e.cfg, streamErr) - reporter.PublishFailure(ctx, streamErr) - select { - case out <- cliproxyexecutor.StreamChunk{Err: streamErr}: - case <-ctx.Done(): - } - return - } - switch eventType { - case "response.output_item.done": - collectCodexOutputItemDone(data, outputItemsByIndex, &outputItemsFallback) - case "response.completed", "response.incomplete": - terminalSuccess = true - if detail, ok := helps.ParseCodexUsage(data); ok { - reporter.Publish(ctx, detail) - } - publishCodexImageToolUsage(ctx, reporter, body, data) - data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback) - if eventType == "response.completed" { - cacheCodexReasoningReplayFromCompleted(replayScope, data) - } - translatedLine = append([]byte("data: "), data...) - } - } - - translatedLine = applyCodexIdentityExposeResponsePayload(translatedLine, identityState) - chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m, claudeInputTokens) - for i := range chunks { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: - case <-ctx.Done(): - return - } - } - if terminalSuccess { - return - } - } - if errScan := scanner.Err(); errScan != nil { - if ctx.Err() != nil { - return - } - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - } - streamErr := newCodexIncompleteStreamError() - helps.RecordAPIResponseError(ctx, e.cfg, streamErr) - reporter.PublishFailure(ctx, streamErr) - select { - case out <- cliproxyexecutor.StreamChunk{Err: streamErr}: - case <-ctx.Done(): - } - }() - return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil -} - -func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - baseModel := thinking.ParseSuffix(req.Model).ModelName - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("codex") - body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false) - - body, err := thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return cliproxyexecutor.Response{}, err - } - - body = helps.SetStringIfDifferent(body, "model", baseModel) - body, _ = sjson.DeleteBytes(body, "previous_response_id") - body, _ = sjson.DeleteBytes(body, "generate") - body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") - body, _ = sjson.DeleteBytes(body, "safety_identifier") - body, _ = sjson.DeleteBytes(body, "stream_options") - body = helps.SetBoolIfDifferent(body, "stream", false) - body = normalizeCodexInstructions(body) - - enc, err := tokenizerForCodexModel(baseModel) - if err != nil { - return cliproxyexecutor.Response{}, fmt.Errorf("codex executor: tokenizer init failed: %w", err) - } - - count, err := countCodexInputTokens(enc, body) - if err != nil { - return cliproxyexecutor.Response{}, fmt.Errorf("codex executor: token counting failed: %w", err) - } - - usageJSON := fmt.Sprintf(`{"response":{"usage":{"input_tokens":%d,"output_tokens":0,"total_tokens":%d}}}`, count, count) - translated := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, []byte(usageJSON)) - return cliproxyexecutor.Response{Payload: translated}, nil -} - -func tokenizerForCodexModel(model string) (tokenizer.Codec, error) { - sanitized := strings.ToLower(strings.TrimSpace(model)) - switch { - case sanitized == "": - return tokenizer.Get(tokenizer.Cl100kBase) - case strings.HasPrefix(sanitized, "gpt-5"): - return tokenizer.ForModel(tokenizer.GPT5) - case strings.HasPrefix(sanitized, "gpt-4.1"): - return tokenizer.ForModel(tokenizer.GPT41) - case strings.HasPrefix(sanitized, "gpt-4o"): - return tokenizer.ForModel(tokenizer.GPT4o) - case strings.HasPrefix(sanitized, "gpt-4"): - return tokenizer.ForModel(tokenizer.GPT4) - case strings.HasPrefix(sanitized, "gpt-3.5"), strings.HasPrefix(sanitized, "gpt-3"): - return tokenizer.ForModel(tokenizer.GPT35Turbo) - default: - return tokenizer.Get(tokenizer.Cl100kBase) - } -} - -func countCodexInputTokens(enc tokenizer.Codec, body []byte) (int64, error) { - if enc == nil { - return 0, fmt.Errorf("encoder is nil") - } - if len(body) == 0 { - return 0, nil - } - - root := gjson.ParseBytes(body) - var segments []string - - if inst := strings.TrimSpace(root.Get("instructions").String()); inst != "" { - segments = append(segments, inst) - } - - inputItems := root.Get("input") - if inputItems.IsArray() { - arr := inputItems.Array() - for i := range arr { - item := arr[i] - switch item.Get("type").String() { - case "message": - content := item.Get("content") - if content.IsArray() { - parts := content.Array() - for j := range parts { - part := parts[j] - if text := strings.TrimSpace(part.Get("text").String()); text != "" { - segments = append(segments, text) - } - } - } - case "function_call": - if name := strings.TrimSpace(item.Get("name").String()); name != "" { - segments = append(segments, name) - } - if args := strings.TrimSpace(item.Get("arguments").String()); args != "" { - segments = append(segments, args) - } - case "function_call_output": - if out := strings.TrimSpace(item.Get("output").String()); out != "" { - segments = append(segments, out) - } - default: - if text := strings.TrimSpace(item.Get("text").String()); text != "" { - segments = append(segments, text) - } - } - } - } - - tools := root.Get("tools") - if tools.IsArray() { - tarr := tools.Array() - for i := range tarr { - tool := tarr[i] - if name := strings.TrimSpace(tool.Get("name").String()); name != "" { - segments = append(segments, name) - } - if desc := strings.TrimSpace(tool.Get("description").String()); desc != "" { - segments = append(segments, desc) - } - if params := tool.Get("parameters"); params.Exists() { - val := params.Raw - if params.Type == gjson.String { - val = params.String() - } - if trimmed := strings.TrimSpace(val); trimmed != "" { - segments = append(segments, trimmed) - } - } - } - } - - textFormat := root.Get("text.format") - if textFormat.Exists() { - if name := strings.TrimSpace(textFormat.Get("name").String()); name != "" { - segments = append(segments, name) - } - if schema := textFormat.Get("schema"); schema.Exists() { - val := schema.Raw - if schema.Type == gjson.String { - val = schema.String() - } - if trimmed := strings.TrimSpace(val); trimmed != "" { - segments = append(segments, trimmed) - } - } - } - - text := strings.Join(segments, "\n") - if text == "" { - return 0, nil - } - - count, err := enc.Count(text) - if err != nil { - return 0, err - } - return int64(count), nil -} - -func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - log.Debugf("codex executor: refresh called") - if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { - return refreshed, err - } - if auth == nil { - return nil, statusErr{code: 500, msg: "codex executor: auth is nil"} - } - var refreshToken string - if auth.Metadata != nil { - if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" { - refreshToken = v - } - } - if refreshToken == "" { - return auth, nil - } - svc := codexauth.NewCodexAuthWithProxyURL(e.cfg, auth.ProxyURL) - td, err := svc.RefreshTokensWithRetry(ctx, refreshToken, 3) - if err != nil { - return nil, err - } - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["id_token"] = td.IDToken - auth.Metadata["access_token"] = td.AccessToken - if td.RefreshToken != "" { - auth.Metadata["refresh_token"] = td.RefreshToken - } - if td.AccountID != "" { - auth.Metadata["account_id"] = td.AccountID - } - auth.Metadata["email"] = td.Email - // Use unified key in files - auth.Metadata["expired"] = td.Expire - auth.Metadata["type"] = "codex" - now := time.Now().Format(time.RFC3339) - auth.Metadata["last_refresh"] = now - return auth, nil -} - -type codexIdentityConfuseState struct { - enabled bool - authID string - originalPromptCacheKey string - promptCacheKey string - turnIDs []codexIdentityReplacement -} - -type codexIdentityReplacement struct { - original string - confused string -} - -func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Format, url string, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, userPayload []byte, rawJSON []byte, headerSets ...http.Header) (*http.Request, []byte, codexIdentityConfuseState, error) { - var headers http.Header - if len(headerSets) > 0 { - headers = headerSets[0] - } - var cache helps.CodexCache - if sourceFormatEqual(from, sdktranslator.FormatClaude) { - modelName := strings.TrimSpace(gjson.GetBytes(rawJSON, "model").String()) - if modelName == "" { - modelName = thinking.ParseSuffix(req.Model).ModelName - } - cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, modelName, req.Payload, headers) - if errCache != nil { - return nil, nil, codexIdentityConfuseState{}, errCache - } - if ok { - cache = cached - } - } else if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) { - promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key") - if promptCacheKey.Exists() { - cache.ID = promptCacheKey.String() - } - } else if sourceFormatEqual(from, sdktranslator.FormatOpenAI) { - if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() { - cache.ID = strings.TrimSpace(promptCacheKey.String()) - } - if cache.ID == "" { - cache.ID = helps.ProviderSessionUUID("codex", req.Metadata) - } - if cache.ID == "" { - if apiKey := strings.TrimSpace(helps.APIKeyFromContext(ctx)); apiKey != "" { - cache.ID = uuid.NewSHA1(uuid.NameSpaceOID, []byte("cli-proxy-api:codex:prompt-cache:"+apiKey)).String() - } - } - } - if cache.ID == "" { - cache.ID = helps.ProviderSessionUUID("codex", req.Metadata) - } - - if cache.ID != "" { - rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", cache.ID) - } - rawJSON = helps.SanitizeCodexInputItemIDs(rawJSON) - var identityState codexIdentityConfuseState - rawJSON, identityState = applyCodexIdentityConfuseBody(e.cfg, auth, userPayload, rawJSON) - if identityState.promptCacheKey != "" { - cache.ID = identityState.promptCacheKey - } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(rawJSON)) - if err != nil { - return nil, nil, codexIdentityConfuseState{}, err - } - if cache.ID != "" { - httpReq.Header.Set("Session_id", cache.ID) - } - return httpReq, rawJSON, identityState, nil -} - -func applyCodexIdentityConfuseBody(cfg *config.Config, auth *cliproxyauth.Auth, userPayload []byte, rawJSON []byte) ([]byte, codexIdentityConfuseState) { - if !codexIdentityConfuseEnabled(cfg) || auth == nil || strings.TrimSpace(auth.ID) == "" || len(rawJSON) == 0 { - return rawJSON, codexIdentityConfuseState{} - } - - state := codexIdentityConfuseState{enabled: true, authID: strings.TrimSpace(auth.ID)} - if promptCacheKey := strings.TrimSpace(gjson.GetBytes(userPayload, "prompt_cache_key").String()); promptCacheKey != "" { - state.originalPromptCacheKey = promptCacheKey - state.promptCacheKey = codexIdentityConfuseUUID(auth.ID, "prompt-cache", promptCacheKey) - rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", state.promptCacheKey) - } - if installationID := strings.TrimSpace(gjson.GetBytes(userPayload, "client_metadata.x-codex-installation-id").String()); installationID != "" { - rawJSON, _ = sjson.SetBytes(rawJSON, "client_metadata.x-codex-installation-id", codexIdentityConfuseUUID(auth.ID, "installation", installationID)) - } - if turnMetadata := strings.TrimSpace(gjson.GetBytes(rawJSON, "client_metadata.x-codex-turn-metadata").String()); turnMetadata != "" { - rawJSON, _ = sjson.SetBytes(rawJSON, "client_metadata.x-codex-turn-metadata", applyCodexTurnMetadataIdentityConfuse(turnMetadata, &state)) - } - if state.promptCacheKey != "" { - if windowID := strings.TrimSpace(gjson.GetBytes(rawJSON, "client_metadata.x-codex-window-id").String()); windowID != "" { - rawJSON, _ = sjson.SetBytes(rawJSON, "client_metadata.x-codex-window-id", state.promptCacheKey+":0") - } - } - - return rawJSON, state -} - -func applyCodexIdentityConfuseHeaders(headers http.Header, state *codexIdentityConfuseState) { - if headers == nil { - return - } - if state == nil || !state.enabled { - return - } - - if rawTurnMetadata := strings.TrimSpace(headers.Get("X-Codex-Turn-Metadata")); rawTurnMetadata != "" { - headers.Set("X-Codex-Turn-Metadata", applyCodexTurnMetadataIdentityConfuse(rawTurnMetadata, state)) - } - if state.promptCacheKey == "" { - return - } - - setCodexSessionHeaderCasePreserved(headers, "Session_id", state.promptCacheKey) - if headerValueCaseInsensitive(headers, "Conversation_id") != "" { - setHeaderCasePreserved(headers, "Conversation_id", state.promptCacheKey) - } - headers.Set("X-Client-Request-Id", state.promptCacheKey) - headers.Set("Thread-Id", state.promptCacheKey) - headers.Set("X-Codex-Window-Id", state.promptCacheKey+":0") -} - -func applyCodexTurnMetadataIdentityConfuse(rawTurnMetadata string, state *codexIdentityConfuseState) string { - updatedTurnMetadata := rawTurnMetadata - if state == nil || !state.enabled { - return updatedTurnMetadata - } - if state.promptCacheKey != "" && gjson.Get(rawTurnMetadata, "prompt_cache_key").Exists() { - updatedTurnMetadata, _ = sjson.Set(updatedTurnMetadata, "prompt_cache_key", state.promptCacheKey) - } else if state.promptCacheKey != "" && state.originalPromptCacheKey != "" { - updatedTurnMetadata = strings.ReplaceAll(updatedTurnMetadata, state.originalPromptCacheKey, state.promptCacheKey) - } - if turnID := strings.TrimSpace(gjson.Get(rawTurnMetadata, "turn_id").String()); turnID != "" { - updatedTurnMetadata, _ = sjson.Set(updatedTurnMetadata, "turn_id", state.confuseTurnID(turnID)) - } - if state.promptCacheKey != "" && gjson.Get(rawTurnMetadata, "window_id").Exists() { - updatedTurnMetadata, _ = sjson.Set(updatedTurnMetadata, "window_id", state.promptCacheKey+":0") - } - return updatedTurnMetadata -} - -func applyCodexIdentityConfuseResponsePayload(payload []byte, state codexIdentityConfuseState) []byte { - payload = replaceCodexIdentityResponsePayload(payload, state.originalPromptCacheKey, state.promptCacheKey) - for _, turnID := range state.turnIDs { - payload = replaceCodexIdentityResponsePayload(payload, turnID.original, turnID.confused) - } - return payload -} - -func applyCodexIdentityExposeResponsePayload(payload []byte, state codexIdentityConfuseState) []byte { - payload = replaceCodexIdentityResponsePayload(payload, state.promptCacheKey, state.originalPromptCacheKey) - for _, turnID := range state.turnIDs { - payload = replaceCodexIdentityResponsePayload(payload, turnID.confused, turnID.original) - } - return payload -} - -func (state *codexIdentityConfuseState) confuseTurnID(turnID string) string { - turnID = strings.TrimSpace(turnID) - if state == nil || !state.enabled || strings.TrimSpace(state.authID) == "" || turnID == "" { - return turnID - } - for _, replacement := range state.turnIDs { - if replacement.original == turnID || replacement.confused == turnID { - return replacement.confused - } - } - confusedTurnID := codexIdentityConfuseUUID(state.authID, "turn", turnID) - state.turnIDs = append(state.turnIDs, codexIdentityReplacement{original: turnID, confused: confusedTurnID}) - return confusedTurnID -} - -func replaceCodexIdentityResponsePayload(payload []byte, from string, to string) []byte { - from = strings.TrimSpace(from) - to = strings.TrimSpace(to) - if len(payload) == 0 || from == "" || to == "" || from == to || !bytes.Contains(payload, []byte(from)) { - return payload - } - return bytes.ReplaceAll(payload, []byte(from), []byte(to)) -} - -func codexIdentityConfuseEnabled(cfg *config.Config) bool { - if cfg == nil || !cfg.Codex.IdentityConfuse { - return false - } - strategy := strings.ToLower(strings.TrimSpace(cfg.Routing.Strategy)) - return cfg.Routing.SessionAffinity || strategy == "fill-first" || strategy == "fillfirst" || strategy == "ff" -} - -func codexIdentityConfuseUUID(authID string, kind string, value string) string { - name := strings.Join([]string{"cli-proxy-api", "codex", "identity-confuse", kind, strings.TrimSpace(authID), strings.TrimSpace(value)}, ":") - return uuid.NewSHA1(uuid.NameSpaceOID, []byte(name)).String() -} - -func applyCodexHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config) { - var ginHeaders http.Header - if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - ginHeaders = ginCtx.Request.Header - } - applyCodexHeadersFromSources(r, auth, token, stream, cfg, ginHeaders) -} - -// applyModelHeaderOverrides forces models.json config.override_header onto upstream headers. -func applyModelHeaderOverrides(headers http.Header, modelName string) { - if headers == nil { - return - } - overrides := registry.ModelOverrideHeaders(modelName) - if len(overrides) == 0 { - return - } - for key, value := range overrides { - headers.Set(key, value) - } - if strings.Contains(headers.Get("User-Agent"), "Mac OS") && codexSessionHeaderValue(headers) == "" { - headers.Set("Session_id", uuid.NewString()) - } -} - -// applyCodexDirectImageHeaders sets Codex upstream headers for direct /images/* calls. -// Downstream client User-Agent values are not forwarded to reduce Cloudflare 1010 blocks. -func applyCodexDirectImageHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config) { - var ginHeaders http.Header - if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - ginHeaders = ginCtx.Request.Header.Clone() - ginHeaders.Del("User-Agent") - } - applyCodexHeadersFromSources(r, auth, token, stream, cfg, ginHeaders) -} - -func applyCodexHeadersFromSources(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config, ginHeaders http.Header) { - r.Header.Set("Content-Type", "application/json") - r.Header.Set("Authorization", "Bearer "+token) - - if ginHeaders != nil && ginHeaders.Get("X-Codex-Beta-Features") != "" { - r.Header.Set("X-Codex-Beta-Features", ginHeaders.Get("X-Codex-Beta-Features")) - } - misc.EnsureHeader(r.Header, ginHeaders, "Version", "") - misc.EnsureHeader(r.Header, ginHeaders, "X-Codex-Turn-Metadata", "") - misc.EnsureHeader(r.Header, ginHeaders, "X-Client-Request-Id", "") - cfgUserAgent, _ := codexHeaderDefaults(cfg, auth) - ensureHeaderWithConfigPrecedence(r.Header, ginHeaders, "User-Agent", cfgUserAgent, codexUserAgent) - - if strings.Contains(r.Header.Get("User-Agent"), "Mac OS") { - misc.EnsureHeader(r.Header, ginHeaders, "Session_id", uuid.NewString()) - } - - if stream { - r.Header.Set("Accept", "text/event-stream") - } else { - r.Header.Set("Accept", "application/json") - } - r.Header.Set("Connection", "Keep-Alive") - - isAPIKey := false - if auth != nil && auth.Attributes != nil { - if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" { - isAPIKey = true - } - } - if originator := strings.TrimSpace(ginHeaders.Get("Originator")); originator != "" { - r.Header.Set("Originator", originator) - } else if !isAPIKey { - r.Header.Set("Originator", codexOriginator) - } - if !isAPIKey { - if auth != nil && auth.Metadata != nil { - if accountID, ok := auth.Metadata["account_id"].(string); ok { - r.Header.Set("Chatgpt-Account-Id", accountID) - } - } - } - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(r, attrs) -} - -func newCodexStatusErr(statusCode int, body []byte) statusErr { - errCode := statusCode - if isCodexModelCapacityError(body) || isCodexUsageLimitError(body) { - errCode = http.StatusTooManyRequests - } - body = classifyCodexStatusError(errCode, body) - err := statusErr{code: errCode, msg: string(body)} - if retryAfter := parseCodexRetryAfter(errCode, body, time.Now()); retryAfter != nil { - err.retryAfter = retryAfter - } - return err -} - -func classifyCodexStatusError(statusCode int, body []byte) []byte { - code, errType, ok := codexStatusErrorClassification(statusCode, body) - if !ok { - return body - } - message := gjson.GetBytes(body, "error.message").String() - if message == "" { - message = gjson.GetBytes(body, "message").String() - } - if message == "" { - message = strings.TrimSpace(string(body)) - } - if message == "" { - message = http.StatusText(statusCode) - } - out := []byte(`{"error":{}}`) - out, _ = sjson.SetBytes(out, "error.message", message) - out, _ = sjson.SetBytes(out, "error.type", errType) - out, _ = sjson.SetBytes(out, "error.code", code) - return out -} - -func codexStatusErrorClassification(statusCode int, body []byte) (code string, errType string, ok bool) { - errorMessage := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.message").String())) - if errorMessage == "" { - errorMessage = strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "message").String())) - } - lower := strings.ToLower(strings.TrimSpace(string(body))) - upstreamCode := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.code").String())) - upstreamType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.type").String())) - isInvalidRequest := upstreamType == "" || upstreamType == "invalid_request_error" - - switch { - case statusCode == http.StatusRequestEntityTooLarge || upstreamCode == "context_length_exceeded" || upstreamCode == "context_too_large" || isInvalidRequest && (strings.Contains(errorMessage, "context length") || strings.Contains(errorMessage, "context_length") || strings.Contains(errorMessage, "maximum context") || strings.Contains(errorMessage, "too many tokens")): - return "context_too_large", "invalid_request_error", true - case strings.Contains(lower, "invalid signature in thinking block") || strings.Contains(lower, "invalid_encrypted_content"): - return "thinking_signature_invalid", "invalid_request_error", true - case upstreamCode == "previous_response_not_found" || strings.Contains(lower, "previous_response_not_found") || strings.Contains(lower, "previous_response_id") && strings.Contains(lower, "not found"): - return "previous_response_not_found", "invalid_request_error", true - case statusCode == http.StatusUnauthorized || upstreamType == "authentication_error" || upstreamCode == "invalid_api_key" || strings.Contains(lower, "invalid or expired token") || strings.Contains(lower, "refresh_token_reused"): - return "auth_unavailable", "authentication_error", true - default: - return "", "", false - } -} - -func normalizeCodexInstructions(body []byte) []byte { - instructions := gjson.GetBytes(body, "instructions") - if !instructions.Exists() || instructions.Type == gjson.Null { - body, _ = sjson.SetBytes(body, "instructions", "") - } - return body -} - -var imageGenToolJSON = []byte(`{"type":"image_generation","output_format":"png"}`) -var imageGenToolArrayJSON = []byte(`[{"type":"image_generation","output_format":"png"}]`) - -func isCodexFreePlanAuth(auth *cliproxyauth.Auth) bool { - if auth == nil || auth.Attributes == nil { - return false - } - if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { - return false - } - return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free") -} - -func isImageGenerationFunctionTool(tool gjson.Result) bool { - switch tool.Get("type").String() { - case "function": - return tool.Get("name").String() == "image_gen.imagegen" - case "namespace": - if tool.Get("name").String() != "image_gen" { - return false - } - tools := tool.Get("tools") - if !tools.IsArray() { - return false - } - for _, nestedTool := range tools.Array() { - if nestedTool.Get("type").String() == "function" && nestedTool.Get("name").String() == "imagegen" { - return true - } - } - } - return false -} - -func isCodexResponsesLiteRequest(body []byte, headers http.Header) bool { - if strings.EqualFold(strings.TrimSpace(headers.Get(codexResponsesLiteHeader)), "true") { - return true - } - // Codex Desktop mirrors websocket-only request headers into client_metadata. - value := gjson.GetBytes(body, codexResponsesLiteMetadata) - if !value.Exists() { - return false - } - return value.Type == gjson.True || value.Type == gjson.String && strings.EqualFold(strings.TrimSpace(value.String()), "true") -} - -func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth.Auth, headers http.Header) []byte { - if isCodexResponsesLiteRequest(body, headers) { - return body - } - if strings.HasSuffix(baseModel, "spark") { - return body - } - if isCodexFreePlanAuth(auth) { - return body - } - - tools := gjson.GetBytes(body, "tools") - if !tools.Exists() || !tools.IsArray() { - body, _ = sjson.SetRawBytes(body, "tools", imageGenToolArrayJSON) - return body - } - for _, t := range tools.Array() { - if t.Get("type").String() == "image_generation" || isImageGenerationFunctionTool(t) { - return body - } - } - body, _ = sjson.SetRawBytes(body, "tools.-1", imageGenToolJSON) - return body -} - -func normalizeCodexParallelToolCalls(body []byte, headers http.Header) []byte { - if isCodexResponsesLiteRequest(body, headers) { - body = helps.SetBoolIfDifferent(body, "parallel_tool_calls", false) - return body - } - return normalizeCodexParallelToolCallsForTools(body) -} - -func normalizeCodexParallelToolCallsForTools(body []byte) []byte { - if !gjson.GetBytes(body, "parallel_tool_calls").Exists() { - return body - } - - tools := gjson.GetBytes(body, "tools") - hasTools := tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 - if hasTools { - return body - } - - body, _ = sjson.DeleteBytes(body, "parallel_tool_calls") - return body -} - -func publishCodexImageToolUsage(ctx context.Context, reporter *helps.UsageReporter, body []byte, completedData []byte) { - detail, ok := helps.ParseCodexImageToolUsage(completedData) - if !ok { - return - } - reporter.EnsurePublished(ctx) - reporter.PublishAdditionalModel(ctx, codexImageGenerationToolModel(body), detail) -} - -func codexImageGenerationToolModel(body []byte) string { - tools := gjson.GetBytes(body, "tools") - if tools.IsArray() { - for _, tool := range tools.Array() { - if tool.Get("type").String() != "image_generation" { - continue - } - if model := strings.TrimSpace(tool.Get("model").String()); model != "" { - return model - } - break - } - } - return codexDefaultImageToolModel -} - -func isCodexModelCapacityError(errorBody []byte) bool { - if len(errorBody) == 0 { - return false - } - candidates := []string{ - gjson.GetBytes(errorBody, "error.message").String(), - gjson.GetBytes(errorBody, "message").String(), - string(errorBody), - } - for _, candidate := range candidates { - lower := strings.ToLower(strings.TrimSpace(candidate)) - if lower == "" { - continue - } - if strings.Contains(lower, "selected model is at capacity") || - strings.Contains(lower, "model is at capacity. please try a different model") { - return true - } - } - return false -} - -// isCodexUsageLimitError reports whether the error body represents a Codex -// quota/plan-limit exhaustion (error.type == "usage_limit_reached"). This is the -// signal Codex emits when a credential's usage quota is depleted, and it carries -// reset timing (resets_at/resets_in_seconds) parsed by parseCodexRetryAfter. -// Transient per-minute rate limits (rate_limit_error/rate_limit_exceeded) are -// intentionally excluded, as they should be retried rather than cooled down. -func isCodexUsageLimitError(errorBody []byte) bool { - if len(errorBody) == 0 { - return false - } - candidates := []string{ - gjson.GetBytes(errorBody, "error.type").String(), - gjson.GetBytes(errorBody, "type").String(), - } - for _, candidate := range candidates { - if strings.EqualFold(strings.TrimSpace(candidate), "usage_limit_reached") { - return true - } - } - return false -} - -func parseCodexRetryAfter(statusCode int, errorBody []byte, now time.Time) *time.Duration { - if statusCode != http.StatusTooManyRequests || len(errorBody) == 0 { - return nil - } - if strings.TrimSpace(gjson.GetBytes(errorBody, "error.type").String()) != "usage_limit_reached" { - return nil - } - if resetsAt := gjson.GetBytes(errorBody, "error.resets_at").Int(); resetsAt > 0 { - resetAtTime := time.Unix(resetsAt, 0) - if resetAtTime.After(now) { - retryAfter := resetAtTime.Sub(now) - return &retryAfter - } - } - if resetsInSeconds := gjson.GetBytes(errorBody, "error.resets_in_seconds").Int(); resetsInSeconds > 0 { - retryAfter := time.Duration(resetsInSeconds) * time.Second - return &retryAfter - } - return nil -} - -func codexCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) { - if a == nil { - return "", "" - } - if a.Attributes != nil { - apiKey = a.Attributes["api_key"] - baseURL = a.Attributes["base_url"] - } - if apiKey == "" && a.Metadata != nil { - if v, ok := a.Metadata["access_token"].(string); ok { - apiKey = v - } - } - return -} - -func (e *CodexExecutor) resolveCodexConfig(auth *cliproxyauth.Auth) *config.CodexKey { - if auth == nil || e.cfg == nil { - return nil - } - var attrKey, attrBase string - if auth.Attributes != nil { - attrKey = strings.TrimSpace(auth.Attributes["api_key"]) - attrBase = strings.TrimSpace(auth.Attributes["base_url"]) - } - for i := range e.cfg.CodexKey { - entry := &e.cfg.CodexKey[i] - cfgKey := strings.TrimSpace(entry.APIKey) - cfgBase := strings.TrimSpace(entry.BaseURL) - if attrKey != "" && attrBase != "" { - if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { - return entry - } - continue - } - if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { - if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - if attrKey != "" { - for i := range e.cfg.CodexKey { - entry := &e.cfg.CodexKey[i] - if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { - return entry - } - } - } - return nil -} diff --git a/internal/runtime/executor/codex_executor_auth.go b/internal/runtime/executor/codex_executor_auth.go new file mode 100644 index 000000000..e200d6902 --- /dev/null +++ b/internal/runtime/executor/codex_executor_auth.go @@ -0,0 +1,110 @@ +package executor + +import ( + "context" + "strings" + "time" + + codexauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" +) + +func (e *CodexExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { + log.Debugf("codex executor: refresh called") + if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { + return refreshed, err + } + if auth == nil { + return nil, statusErr{code: 500, msg: "codex executor: auth is nil"} + } + var refreshToken string + if auth.Metadata != nil { + if v, ok := auth.Metadata["refresh_token"].(string); ok && v != "" { + refreshToken = v + } + } + if refreshToken == "" { + return auth, nil + } + svc := codexauth.NewCodexAuthWithProxyURL(e.cfg, auth.ProxyURL) + td, err := svc.RefreshTokensWithRetry(ctx, refreshToken, 3) + if err != nil { + return nil, err + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["id_token"] = td.IDToken + auth.Metadata["access_token"] = td.AccessToken + if td.RefreshToken != "" { + auth.Metadata["refresh_token"] = td.RefreshToken + } + if td.AccountID != "" { + auth.Metadata["account_id"] = td.AccountID + } + auth.Metadata["email"] = td.Email + // Use unified key in files + auth.Metadata["expired"] = td.Expire + auth.Metadata["type"] = "codex" + now := time.Now().Format(time.RFC3339) + auth.Metadata["last_refresh"] = now + return auth, nil +} + +func codexCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) { + if a == nil { + return "", "" + } + if a.Attributes != nil { + apiKey = a.Attributes["api_key"] + baseURL = a.Attributes["base_url"] + } + if apiKey == "" && a.Metadata != nil { + if v, ok := a.Metadata["access_token"].(string); ok { + apiKey = v + } + } + return +} + +func (e *CodexExecutor) resolveCodexConfig(auth *cliproxyauth.Auth) *config.CodexKey { + if auth == nil || e.cfg == nil { + return nil + } + var attrKey, attrBase string + if auth.Attributes != nil { + attrKey = strings.TrimSpace(auth.Attributes["api_key"]) + attrBase = strings.TrimSpace(auth.Attributes["base_url"]) + } + for i := range e.cfg.CodexKey { + entry := &e.cfg.CodexKey[i] + cfgKey := strings.TrimSpace(entry.APIKey) + cfgBase := strings.TrimSpace(entry.BaseURL) + if attrKey != "" && attrBase != "" { + if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { + return entry + } + continue + } + if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { + if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey != "" { + for i := range e.cfg.CodexKey { + entry := &e.cfg.CodexKey[i] + if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { + return entry + } + } + } + return nil +} diff --git a/internal/runtime/executor/codex_executor_execute.go b/internal/runtime/executor/codex_executor_execute.go new file mode 100644 index 000000000..87279814d --- /dev/null +++ b/internal/runtime/executor/codex_executor_execute.go @@ -0,0 +1,295 @@ +package executor + +import ( + "bytes" + "context" + "io" + "net/http" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + if opts.Alt == "responses/compact" { + return e.executeCompact(ctx, auth, req, opts) + } + if isCodexOpenAIImageRequest(opts) { + return e.executeOpenAIImage(ctx, auth, req, opts) + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + + apiKey, baseURL := codexCreds(auth) + if baseURL == "" { + baseURL = "https://chatgpt.com/backend-api/codex" + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("codex") + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, false) + + body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return resp, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = helps.SetBoolIfDifferent(body, "stream", true) + body, _ = sjson.DeleteBytes(body, "previous_response_id") + body, _ = sjson.DeleteBytes(body, "generate") + body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") + body, _ = sjson.DeleteBytes(body, "safety_identifier") + body, _ = sjson.DeleteBytes(body, "stream_options") + body = normalizeCodexInstructions(body) + if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { + body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) + } + body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) + body = normalizeCodexParallelToolCalls(body, opts.Headers) + body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return resp, errReplay + } + reporter.SetTranslatedReasoningEffort(body, to.String()) + + url := strings.TrimSuffix(baseURL, "/") + "/responses" + var identityState codexIdentityConfuseState + httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) + if err != nil { + return resp, err + } + applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) + applyModelHeaderOverrides(httpReq.Header, baseModel) + applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: url, + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: upstreamBody, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("codex executor: close response body error: %v", errClose) + } + }() + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + b, _ := io.ReadAll(httpResp.Body) + b = applyCodexIdentityConfuseResponsePayload(b, identityState) + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, b); errClearReplay != nil { + return resp, errClearReplay + } + helps.AppendAPIResponseChunk(ctx, e.cfg, b) + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) + err = newCodexStatusErr(httpResp.StatusCode, b) + return resp, err + } + data, errRead := io.ReadAll(httpResp.Body) + upstreamData := applyCodexIdentityConfuseResponsePayload(data, identityState) + helps.AppendAPIResponseChunk(ctx, e.cfg, upstreamData) + + lines := bytes.Split(upstreamData, []byte("\n")) + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + for _, line := range lines { + if !bytes.HasPrefix(line, dataTag) { + continue + } + + eventData := bytes.TrimSpace(line[5:]) + eventData = helps.RestoreCodexMultiAgentV2Response(eventData, optimizeMultiAgentV2) + eventType := gjson.GetBytes(eventData, "type").String() + + if streamErr, terminalBody, ok := codexTerminalFailureErr(eventData); ok { + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + return resp, errClearReplay + } + err = streamErr + return resp, err + } + + if eventType == "response.output_item.done" { + itemResult := gjson.GetBytes(eventData, "item") + if !itemResult.Exists() || itemResult.Type != gjson.JSON { + continue + } + outputIndexResult := gjson.GetBytes(eventData, "output_index") + if outputIndexResult.Exists() { + outputItemsByIndex[outputIndexResult.Int()] = []byte(itemResult.Raw) + } else { + outputItemsFallback = append(outputItemsFallback, []byte(itemResult.Raw)) + } + continue + } + + if eventType != "response.completed" && eventType != "response.incomplete" { + continue + } + + if detail, ok := helps.ParseCodexUsage(eventData); ok { + reporter.Publish(ctx, detail) + } + publishCodexImageToolUsage(ctx, reporter, body, eventData) + + completedData := patchCodexCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) + if eventType == "response.completed" { + cacheCodexReasoningReplayFromCompleted(replayScope, completedData) + } + + var param any + clientCompletedData := applyCodexIdentityExposeResponsePayload(completedData, identityState) + out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, clientCompletedData, ¶m) + resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} + return resp, nil + } + if errRead != nil { + if errCtx := ctx.Err(); errCtx != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errCtx) + err = errCtx + return resp, err + } + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + } + err = newCodexIncompleteStreamError() + return resp, err +} + +func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + baseModel := thinking.ParseSuffix(req.Model).ModelName + + apiKey, baseURL := codexCreds(auth) + if baseURL == "" { + baseURL = "https://chatgpt.com/backend-api/codex" + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("openai-response") + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, false) + + body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return resp, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body, _ = sjson.DeleteBytes(body, "stream") + body = normalizeCodexInstructions(body) + body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) + body = normalizeCodexParallelToolCalls(body, opts.Headers) + body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) + reporter.SetTranslatedReasoningEffort(body, to.String()) + + url := strings.TrimSuffix(baseURL, "/") + "/responses/compact" + var identityState codexIdentityConfuseState + httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) + if err != nil { + return resp, err + } + applyCodexHeaders(httpReq, auth, apiKey, false, e.cfg) + applyModelHeaderOverrides(httpReq.Header, baseModel) + applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: url, + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: upstreamBody, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("codex executor: close response body error: %v", errClose) + } + }() + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + b, _ := io.ReadAll(httpResp.Body) + b = applyCodexIdentityConfuseResponsePayload(b, identityState) + helps.AppendAPIResponseChunk(ctx, e.cfg, b) + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), b)) + err = newCodexStatusErr(httpResp.StatusCode, b) + return resp, err + } + data, err := io.ReadAll(httpResp.Body) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + upstreamData := applyCodexIdentityConfuseResponsePayload(data, identityState) + helps.AppendAPIResponseChunk(ctx, e.cfg, upstreamData) + upstreamData = helps.RestoreCodexMultiAgentV2Response(upstreamData, optimizeMultiAgentV2) + reporter.Publish(ctx, helps.ParseOpenAIUsage(upstreamData)) + reporter.EnsurePublished(ctx) + var param any + clientData := applyCodexIdentityExposeResponsePayload(upstreamData, identityState) + out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, clientData, ¶m) + resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} + return resp, nil +} diff --git a/internal/runtime/executor/codex_executor_reasoning.go b/internal/runtime/executor/codex_executor_reasoning.go new file mode 100644 index 000000000..25093b5ae --- /dev/null +++ b/internal/runtime/executor/codex_executor_reasoning.go @@ -0,0 +1,788 @@ +package executor + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + internalcache "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +type codexReasoningReplayScope struct { + modelName string + sessionKey string + requestFingerprint string +} + +func (s codexReasoningReplayScope) valid() bool { + return strings.TrimSpace(s.modelName) != "" && strings.TrimSpace(s.sessionKey) != "" +} + +func applyCodexReasoningReplayCache(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) ([]byte, codexReasoningReplayScope) { + updated, scope, _ := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + return updated, scope +} + +func applyCodexReasoningReplayCacheRequired(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) ([]byte, codexReasoningReplayScope, error) { + scope := codexReasoningReplayScopeFromRequest(ctx, from, req, opts, body) + if !scope.valid() { + return body, scope, nil + } + items, ok, errReplay := internalcache.GetCodexReasoningReplayItemsRequired(ctx, scope.modelName, scope.sessionKey) + if errReplay != nil || !ok { + return body, scope, errReplay + } + updated, ok := insertCodexReasoningReplayTurns(body, items) + if !ok { + return body, scope, nil + } + return updated, scope, nil +} + +func codexReasoningReplayScopeFromRequest(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) codexReasoningReplayScope { + if !codexReasoningReplayEnabledForSource(from) { + return codexReasoningReplayScope{} + } + modelName := strings.TrimSpace(gjson.GetBytes(body, "model").String()) + if modelName == "" { + modelName = thinking.ParseSuffix(req.Model).ModelName + } + inputItems := gjson.GetBytes(body, "input").Array() + return codexReasoningReplayScope{ + modelName: modelName, + sessionKey: codexReasoningReplaySessionKey(ctx, from, req, opts, body), + requestFingerprint: codexReplayInputPrefixFingerprint(inputItems, len(inputItems)), + } +} + +func codexReasoningReplayEnabledForSource(from sdktranslator.Format) bool { + return sourceFormatEqual(from, sdktranslator.FormatClaude) +} + +func sourceFormatEqual(from, want sdktranslator.Format) bool { + return strings.EqualFold(strings.TrimSpace(from.String()), want.String()) +} + +func codexClaudeCodeReplaySessionKey(ctx context.Context, payload []byte, headers http.Header) string { + sessionKey, _ := helps.ClaudeCodeExecutionScope(ctx, payload, headers) + return sessionKey +} + +func codexReasoningReplaySessionKey(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, body []byte) string { + if ctx == nil { + ctx = context.Background() + } + if sourceFormatEqual(from, sdktranslator.FormatClaude) { + if sessionKey := codexClaudeCodeReplaySessionKey(ctx, req.Payload, opts.Headers); sessionKey != "" { + return sessionKey + } + } + if value := metadataString(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { + return "execution:" + value + } + if value := metadataString(req.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { + return "execution:" + value + } + if value := codexReasoningReplaySessionKeyFromPayload(body); value != "" { + return value + } + if value := codexReasoningReplaySessionKeyFromPayload(req.Payload); value != "" { + return value + } + if value := codexReasoningReplaySessionKeyFromHeaders(opts.Headers); value != "" { + return value + } + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + if value := codexReasoningReplaySessionKeyFromHeaders(ginCtx.Request.Header); value != "" { + return value + } + } + if sourceFormatEqual(from, sdktranslator.FormatOpenAI) { + if apiKey := strings.TrimSpace(helps.APIKeyFromContext(ctx)); apiKey != "" { + return "prompt-cache:" + uuid.NewSHA1(uuid.NameSpaceOID, []byte("cli-proxy-api:codex:prompt-cache:"+apiKey)).String() + } + } + return "" +} + +func metadataString(metadata map[string]any, key string) string { + if len(metadata) == 0 { + return "" + } + raw, ok := metadata[key] + if !ok || raw == nil { + return "" + } + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) + case []byte: + return strings.TrimSpace(string(v)) + default: + return "" + } +} + +func codexReasoningReplaySessionKeyFromPayload(payload []byte) string { + if len(payload) == 0 { + return "" + } + if promptCacheKey := strings.TrimSpace(gjson.GetBytes(payload, "prompt_cache_key").String()); promptCacheKey != "" { + return "prompt-cache:" + promptCacheKey + } + if windowID := strings.TrimSpace(gjson.GetBytes(payload, "client_metadata.x-codex-window-id").String()); windowID != "" { + return "window:" + windowID + } + if turnMetadata := strings.TrimSpace(gjson.GetBytes(payload, "client_metadata.x-codex-turn-metadata").String()); turnMetadata != "" { + return codexReasoningReplaySessionKeyFromTurnMetadata(turnMetadata) + } + return "" +} + +func codexReasoningReplaySessionKeyFromHeaders(headers http.Header) string { + if headers == nil { + return "" + } + if turnMetadata := strings.TrimSpace(headers.Get("X-Codex-Turn-Metadata")); turnMetadata != "" { + if key := codexReasoningReplaySessionKeyFromTurnMetadata(turnMetadata); key != "" { + return key + } + } + if windowID := strings.TrimSpace(headerValueCaseInsensitive(headers, "X-Codex-Window-Id")); windowID != "" { + return "window:" + windowID + } + for _, headerName := range []string{"Session_id", "session_id", "Session-Id"} { + if value := strings.TrimSpace(headerValueCaseInsensitive(headers, headerName)); value != "" { + return "session-id:" + value + } + } + if conversationID := strings.TrimSpace(headerValueCaseInsensitive(headers, "Conversation_id")); conversationID != "" { + return "conversation_id:" + conversationID + } + return "" +} + +func codexReasoningReplaySessionKeyFromTurnMetadata(turnMetadata string) string { + if promptCacheKey := strings.TrimSpace(gjson.Get(turnMetadata, "prompt_cache_key").String()); promptCacheKey != "" { + return "prompt-cache:" + promptCacheKey + } + if windowID := strings.TrimSpace(gjson.Get(turnMetadata, "window_id").String()); windowID != "" { + return "window:" + windowID + } + return "" +} + +func codexInputHasValidReasoningEncryptedContent(body []byte) bool { + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return false + } + for _, item := range input.Array() { + if strings.TrimSpace(item.Get("type").String()) != "reasoning" { + continue + } + encryptedContent := item.Get("encrypted_content") + if encryptedContent.Type != gjson.String { + continue + } + if _, err := signature.InspectGPTReasoningSignature(encryptedContent.String()); err == nil { + return true + } + } + return false +} + +type codexReasoningReplayTurn struct { + marked bool + assistantFingerprint string + requestFingerprint string + callIDs []string + items [][]byte +} + +func insertCodexReasoningReplayTurns(body []byte, replayItems [][]byte) ([]byte, bool) { + input := gjson.GetBytes(body, "input") + if !input.IsArray() || len(replayItems) == 0 { + return body, false + } + inputItems := input.Array() + turns := splitCodexReasoningReplayTurns(replayItems) + insertions := make(map[int][][]byte) + usedAnchorIndexes := make(map[int]bool) + fallbackAnchorEnd := len(inputItems) - 1 + inserted := false + for turnIndex := len(turns) - 1; turnIndex >= 0; turnIndex-- { + turn := turns[turnIndex] + if len(turn.items) == 0 { + continue + } + if !turn.marked { + items := filterCodexReasoningReplayItemsForInput(body, turn.items) + if len(items) == 0 { + continue + } + index := codexReasoningReplayInsertIndex(inputItems, items) + items = codexAlignReasoningReplayToolCallIDs(inputItems, items) + insertions[index] = append(items, insertions[index]...) + inserted = true + continue + } + + anchorIndex, matched := codexReasoningReplayTurnAnchorIndex(inputItems, turn, fallbackAnchorEnd, usedAnchorIndexes) + if !matched { + continue + } + usedAnchorIndexes[anchorIndex] = true + if turn.requestFingerprint == "" { + fallbackAnchorEnd = anchorIndex - 1 + } + items := filterCodexReasoningReplayTurnItems(inputItems, turn.items) + if len(items) == 0 { + continue + } + items = codexAlignReasoningReplayToolCallIDs(inputItems, items) + insertions[anchorIndex] = append(items, insertions[anchorIndex]...) + inserted = true + } + if !inserted { + return body, false + } + + items := make([]string, 0, len(inputItems)+len(replayItems)) + for index, inputItem := range inputItems { + for _, replayItem := range insertions[index] { + items = append(items, string(replayItem)) + } + items = append(items, inputItem.Raw) + } + for _, replayItem := range insertions[len(inputItems)] { + items = append(items, string(replayItem)) + } + updated, err := sjson.SetRawBytes(body, "input", []byte("["+strings.Join(items, ",")+"]")) + if err != nil { + return body, false + } + return updated, true +} + +func splitCodexReasoningReplayTurns(items [][]byte) []codexReasoningReplayTurn { + turns := make([]codexReasoningReplayTurn, 0) + current := codexReasoningReplayTurn{} + appendCurrent := func() { + if len(current.items) > 0 { + turns = append(turns, current) + } + } + for _, item := range items { + itemResult := gjson.ParseBytes(item) + if strings.TrimSpace(itemResult.Get("type").String()) == internalcache.CodexReasoningReplayTurnType { + appendCurrent() + current = codexReasoningReplayTurn{ + marked: true, + assistantFingerprint: strings.TrimSpace(itemResult.Get("assistant_fingerprint").String()), + requestFingerprint: strings.TrimSpace(itemResult.Get("request_fingerprint").String()), + } + if callIDs := itemResult.Get("call_ids"); callIDs.IsArray() { + for _, callIDResult := range callIDs.Array() { + if callID := strings.TrimSpace(callIDResult.String()); callID != "" { + current.callIDs = append(current.callIDs, callID) + } + } + } + continue + } + current.items = append(current.items, item) + } + appendCurrent() + return turns +} + +func codexReasoningReplayTurnAnchorIndex(inputItems []gjson.Result, turn codexReasoningReplayTurn, fallbackEnd int, used map[int]bool) (int, bool) { + searchEnd := fallbackEnd + if turn.requestFingerprint != "" { + searchEnd = len(inputItems) - 1 + } + if searchEnd >= len(inputItems) { + searchEnd = len(inputItems) - 1 + } + matchesRequestPrefix := func(index int) bool { + return turn.requestFingerprint == "" || codexReplayInputPrefixFingerprint(inputItems, index) == turn.requestFingerprint + } + if len(turn.callIDs) > 0 { + callIDs := make(map[string]bool) + for _, callID := range turn.callIDs { + for _, candidate := range codexReplayComparableCallIDs(callID) { + callIDs[candidate] = true + } + } + for index := searchEnd; index >= 0; index-- { + if used[index] || !matchesRequestPrefix(index) { + continue + } + itemType := strings.TrimSpace(inputItems[index].Get("type").String()) + if itemType != "function_call" && itemType != "custom_tool_call" && itemType != "function_call_output" && itemType != "custom_tool_call_output" { + continue + } + for _, candidate := range codexReplayComparableCallIDs(inputItems[index].Get("call_id").String()) { + if callIDs[candidate] { + return index, true + } + } + } + } + if turn.assistantFingerprint != "" { + for index := searchEnd; index >= 0; index-- { + if used[index] || !matchesRequestPrefix(index) { + continue + } + if codexReplayAssistantMessageFingerprint(inputItems[index]) == turn.assistantFingerprint { + return index, true + } + } + } + if len(turn.callIDs) == 0 && turn.assistantFingerprint == "" { + return codexReasoningReplayInsertIndex(inputItems, turn.items), true + } + return 0, false +} + +func filterCodexReasoningReplayTurnItems(inputItems []gjson.Result, items [][]byte) [][]byte { + existingReasoning := make(map[string]bool) + existingCalls := make(map[string]bool) + existingOutputs := make(map[string]bool) + for _, inputItem := range inputItems { + itemType := strings.TrimSpace(inputItem.Get("type").String()) + switch itemType { + case "reasoning": + if encryptedContent := strings.TrimSpace(inputItem.Get("encrypted_content").String()); encryptedContent != "" { + existingReasoning[encryptedContent] = true + } + case "function_call_output", "custom_tool_call_output": + for _, candidate := range codexReplayComparableCallIDs(inputItem.Get("call_id").String()) { + existingOutputs[candidate] = true + } + } + for _, key := range codexReplayToolCallKeys(inputItem) { + existingCalls[key] = true + } + } + + filtered := make([][]byte, 0, len(items)) + for _, item := range items { + itemResult := gjson.ParseBytes(item) + switch strings.TrimSpace(itemResult.Get("type").String()) { + case "reasoning": + if existingReasoning[strings.TrimSpace(itemResult.Get("encrypted_content").String())] { + continue + } + case "function_call", "custom_tool_call": + keys := codexReplayToolCallKeys(itemResult) + if len(keys) == 0 || codexReplayAnyToolCallKeyExists(existingCalls, keys) { + continue + } + hasMatchingOutput := false + for _, candidate := range codexReplayComparableCallIDs(itemResult.Get("call_id").String()) { + if existingOutputs[candidate] { + hasMatchingOutput = true + break + } + } + if !hasMatchingOutput { + continue + } + for _, key := range keys { + existingCalls[key] = true + } + default: + continue + } + filtered = append(filtered, item) + } + return filtered +} + +func codexReplayAssistantMessageFingerprint(item gjson.Result) string { + itemType := strings.TrimSpace(item.Get("type").String()) + if itemType != "" && itemType != "message" { + return "" + } + if !strings.EqualFold(strings.TrimSpace(item.Get("role").String()), "assistant") { + return "" + } + content := item.Get("content") + var builder strings.Builder + if content.Type == gjson.String { + builder.WriteString(content.String()) + } else if content.IsArray() { + for _, part := range content.Array() { + switch strings.TrimSpace(part.Get("type").String()) { + case "input_text", "output_text": + builder.WriteString(part.Get("text").String()) + case "refusal": + builder.WriteString("\x00refusal\x00") + builder.WriteString(part.Get("refusal").String()) + default: + return "" + } + } + } else { + return "" + } + if builder.Len() == 0 { + return "" + } + sum := sha256.Sum256([]byte(builder.String())) + return hex.EncodeToString(sum[:]) +} + +func codexReplayInputPrefixFingerprint(inputItems []gjson.Result, end int) string { + if end < 0 || end > len(inputItems) { + return "" + } + hasher := sha256.New() + for index := 0; index < end; index++ { + _, _ = hasher.Write([]byte("\x00item\x00")) + _, _ = hasher.Write([]byte(inputItems[index].Raw)) + } + return hex.EncodeToString(hasher.Sum(nil)) +} + +func filterCodexReasoningReplayItemsForInput(body []byte, items [][]byte) [][]byte { + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return nil + } + + hasInputReasoning := codexInputHasValidReasoningEncryptedContent(body) + existingCalls := make(map[string]bool) + existingOutputs := make(map[string]bool) + for _, inputItem := range input.Array() { + itemType := strings.TrimSpace(inputItem.Get("type").String()) + if itemType == "function_call_output" || itemType == "custom_tool_call_output" { + callID := strings.TrimSpace(inputItem.Get("call_id").String()) + if callID != "" { + for _, candidate := range codexReplayComparableCallIDs(callID) { + existingOutputs[candidate] = true + } + } + } + for _, key := range codexReplayToolCallKeys(inputItem) { + existingCalls[key] = true + } + } + + filtered := make([][]byte, 0, len(items)) + for _, item := range items { + itemResult := gjson.ParseBytes(item) + switch strings.TrimSpace(itemResult.Get("type").String()) { + case "reasoning": + if hasInputReasoning { + continue + } + case "function_call", "custom_tool_call": + keys := codexReplayToolCallKeys(itemResult) + if len(keys) == 0 || codexReplayAnyToolCallKeyExists(existingCalls, keys) { + continue + } + // Only inject if there is a matching output in the request + hasMatchingOutput := false + callID := strings.TrimSpace(itemResult.Get("call_id").String()) + if callID != "" { + for _, candidate := range codexReplayComparableCallIDs(callID) { + if existingOutputs[candidate] { + hasMatchingOutput = true + break + } + } + } + if !hasMatchingOutput { + continue + } + for _, key := range keys { + existingCalls[key] = true + } + default: + continue + } + filtered = append(filtered, item) + } + return filtered +} + +func insertCodexReasoningReplayItems(body []byte, replayItems [][]byte) ([]byte, bool) { + input := gjson.GetBytes(body, "input") + if !input.IsArray() || len(replayItems) == 0 { + return body, false + } + inputItems := input.Array() + insertIndex := codexReasoningReplayInsertIndex(inputItems, replayItems) + replayItems = codexAlignReasoningReplayToolCallIDs(inputItems, replayItems) + items := make([]string, 0, len(inputItems)+len(replayItems)) + for i, inputItem := range inputItems { + if i == insertIndex { + for _, replayItem := range replayItems { + items = append(items, string(replayItem)) + } + } + items = append(items, inputItem.Raw) + } + if insertIndex == len(inputItems) { + for _, replayItem := range replayItems { + items = append(items, string(replayItem)) + } + } + updated, err := sjson.SetRawBytes(body, "input", []byte("["+strings.Join(items, ",")+"]")) + if err != nil { + return body, false + } + return updated, true +} + +func codexReasoningReplayInsertIndex(inputItems []gjson.Result, replayItems [][]byte) int { + replayCallIDs := make(map[string]bool) + for _, replayItem := range replayItems { + itemResult := gjson.ParseBytes(replayItem) + itemType := strings.TrimSpace(itemResult.Get("type").String()) + if itemType != "function_call" && itemType != "custom_tool_call" { + continue + } + for _, callID := range codexReplayComparableCallIDs(itemResult.Get("call_id").String()) { + replayCallIDs[callID] = true + } + } + if len(replayCallIDs) > 0 { + for index, inputItem := range inputItems { + itemType := strings.TrimSpace(inputItem.Get("type").String()) + if itemType != "function_call_output" && itemType != "custom_tool_call_output" { + continue + } + callID := strings.TrimSpace(inputItem.Get("call_id").String()) + if callID == "" || replayCallIDs[callID] { + return index + } + } + } + for index := len(inputItems) - 1; index >= 0; index-- { + inputItem := inputItems[index] + if role, ok := codexReplayMessageRole(inputItem); ok && role == "assistant" { + return index + } + } + for index, inputItem := range inputItems { + if shouldInsertCodexReasoningReplayBefore(inputItem) { + return index + } + } + return len(inputItems) +} + +func codexAlignReasoningReplayToolCallIDs(inputItems []gjson.Result, replayItems [][]byte) [][]byte { + outputCallIDs := codexReplayOutputCallIDs(inputItems) + if len(outputCallIDs) == 0 { + return replayItems + } + + aligned := make([][]byte, 0, len(replayItems)) + for _, replayItem := range replayItems { + itemResult := gjson.ParseBytes(replayItem) + itemType := strings.TrimSpace(itemResult.Get("type").String()) + if itemType != "function_call" && itemType != "custom_tool_call" { + aligned = append(aligned, replayItem) + continue + } + + callID := strings.TrimSpace(itemResult.Get("call_id").String()) + outputCallID := "" + for _, candidate := range codexReplayComparableCallIDs(callID) { + if value := outputCallIDs[candidate]; value != "" { + outputCallID = value + break + } + } + if outputCallID == "" || outputCallID == callID { + aligned = append(aligned, replayItem) + continue + } + + updated, err := sjson.SetBytes(replayItem, "call_id", outputCallID) + if err != nil { + aligned = append(aligned, replayItem) + continue + } + aligned = append(aligned, updated) + } + return aligned +} + +func codexReplayOutputCallIDs(inputItems []gjson.Result) map[string]string { + outputCallIDs := make(map[string]string) + for _, inputItem := range inputItems { + itemType := strings.TrimSpace(inputItem.Get("type").String()) + if itemType != "function_call_output" && itemType != "custom_tool_call_output" { + continue + } + callID := strings.TrimSpace(inputItem.Get("call_id").String()) + if callID == "" { + continue + } + for _, candidate := range codexReplayComparableCallIDs(callID) { + outputCallIDs[candidate] = callID + } + } + return outputCallIDs +} + +func shouldInsertCodexReasoningReplayBefore(item gjson.Result) bool { + role, ok := codexReplayMessageRole(item) + if !ok { + return true + } + switch role { + case "developer", "system": + return false + default: + return true + } +} + +func codexReplayMessageRole(item gjson.Result) (string, bool) { + itemType := strings.TrimSpace(item.Get("type").String()) + role := strings.ToLower(strings.TrimSpace(item.Get("role").String())) + if role == "" || (itemType != "" && itemType != "message") { + return "", false + } + return role, true +} + +func codexReplayToolCallKeys(item gjson.Result) []string { + itemType := strings.TrimSpace(item.Get("type").String()) + if itemType != "function_call" && itemType != "custom_tool_call" { + return nil + } + callIDs := codexReplayComparableCallIDs(item.Get("call_id").String()) + if len(callIDs) == 0 { + return nil + } + keys := make([]string, 0, len(callIDs)) + for _, callID := range callIDs { + keys = append(keys, itemType+":"+callID) + } + return keys +} + +func codexReplayAnyToolCallKeyExists(existing map[string]bool, keys []string) bool { + for _, key := range keys { + if existing[key] { + return true + } + } + return false +} + +func codexReplayComparableCallIDs(callID string) []string { + callID = strings.TrimSpace(callID) + if callID == "" { + return nil + } + + claudeVisibleCallID := shortenCodexReplayCallIDIfNeeded(util.SanitizeClaudeToolID(callID)) + if claudeVisibleCallID == "" || claudeVisibleCallID == callID { + return []string{callID} + } + return []string{callID, claudeVisibleCallID} +} + +func shortenCodexReplayCallIDIfNeeded(id string) string { + const limit = 64 + if len(id) <= limit { + return id + } + + sum := sha256.Sum256([]byte(id)) + suffix := "_" + hex.EncodeToString(sum[:8]) + prefixLen := limit - len(suffix) + if prefixLen <= 0 { + return suffix[len(suffix)-limit:] + } + return id[:prefixLen] + suffix +} + +func cacheCodexReasoningReplayFromCompleted(scope codexReasoningReplayScope, completedData []byte) { + if !scope.valid() { + return + } + output := gjson.GetBytes(completedData, "response.output") + if !output.IsArray() { + return + } + replayItems := make([][]byte, 0, len(output.Array())) + callIDs := make([]string, 0) + assistantFingerprint := "" + for _, item := range output.Array() { + switch strings.TrimSpace(item.Get("type").String()) { + case "reasoning": + replayItems = append(replayItems, []byte(item.Raw)) + case "function_call", "custom_tool_call": + replayItems = append(replayItems, []byte(item.Raw)) + if callID := strings.TrimSpace(item.Get("call_id").String()); callID != "" { + callIDs = append(callIDs, callID) + } + case "message": + if fingerprint := codexReplayAssistantMessageFingerprint(item); fingerprint != "" { + assistantFingerprint = fingerprint + } + } + } + if len(replayItems) == 0 { + return + } + + hasher := sha256.New() + _, _ = hasher.Write([]byte(scope.requestFingerprint)) + _, _ = hasher.Write([]byte("\x00assistant\x00" + assistantFingerprint)) + for _, callID := range callIDs { + _, _ = hasher.Write([]byte("\x00call\x00" + callID)) + } + for _, item := range replayItems { + _, _ = hasher.Write([]byte("\x00item\x00")) + _, _ = hasher.Write(item) + } + marker := []byte(`{"type":"` + internalcache.CodexReasoningReplayTurnType + `"}`) + marker, _ = sjson.SetBytes(marker, "id", hex.EncodeToString(hasher.Sum(nil))) + if assistantFingerprint != "" { + marker, _ = sjson.SetBytes(marker, "assistant_fingerprint", assistantFingerprint) + } + if scope.requestFingerprint != "" { + marker, _ = sjson.SetBytes(marker, "request_fingerprint", scope.requestFingerprint) + } + for _, callID := range callIDs { + marker, _ = sjson.SetBytes(marker, "call_ids.-1", callID) + } + items := make([][]byte, 0, len(replayItems)+1) + items = append(items, marker) + items = append(items, replayItems...) + internalcache.AppendCodexReasoningReplayItemsBestEffort(context.Background(), scope.modelName, scope.sessionKey, items) +} + +func clearCodexReasoningReplayOnInvalidSignature(ctx context.Context, scope codexReasoningReplayScope, statusCode int, body []byte) error { + if !scope.valid() { + return nil + } + code, _, ok := codexStatusErrorClassification(statusCode, body) + if ok && code == "thinking_signature_invalid" { + return internalcache.DeleteCodexReasoningReplayItemRequired(ctx, scope.modelName, scope.sessionKey) + } + return nil +} diff --git a/internal/runtime/executor/codex_executor_request.go b/internal/runtime/executor/codex_executor_request.go new file mode 100644 index 000000000..7a5f86348 --- /dev/null +++ b/internal/runtime/executor/codex_executor_request.go @@ -0,0 +1,482 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const ( + codexUserAgent = "codex-tui/0.135.0 (Mac OS 26.5.0; arm64) iTerm.app/3.6.10 (codex-tui; 0.135.0)" + codexOriginator = "codex-tui" + codexDefaultImageToolModel = "gpt-image-2" + codexResponsesLiteHeader = "X-OpenAI-Internal-Codex-Responses-Lite" + codexResponsesLiteMetadata = "client_metadata.ws_request_header_x_openai_internal_codex_responses_lite" +) + +var dataTag = []byte("data:") + +func translateCodexRequestPair(from, to sdktranslator.Format, model string, originalPayload, payload []byte, stream bool) ([]byte, []byte) { + if bytes.Equal(originalPayload, payload) { + body := sdktranslator.TranslateRequest(from, to, model, payload, stream) + return body, body + } + originalTranslated := sdktranslator.TranslateRequest(from, to, model, originalPayload, stream) + body := sdktranslator.TranslateRequest(from, to, model, payload, stream) + return originalTranslated, body +} + +// PrepareRequest injects Codex credentials into the outgoing HTTP request. +func (e *CodexExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error { + if req == nil { + return nil + } + apiKey, _ := codexCreds(auth) + if strings.TrimSpace(apiKey) != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(req, attrs) + return nil +} + +// HttpRequest injects Codex credentials into the request and executes it. +func (e *CodexExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) { + if req == nil { + return nil, fmt.Errorf("codex executor: request is nil") + } + if ctx == nil { + ctx = req.Context() + } + httpReq := req.WithContext(ctx) + if err := e.PrepareRequest(httpReq, auth); err != nil { + return nil, err + } + httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) + return httpClient.Do(httpReq) +} + +type codexIdentityConfuseState struct { + enabled bool + authID string + originalPromptCacheKey string + promptCacheKey string + turnIDs []codexIdentityReplacement +} + +type codexIdentityReplacement struct { + original string + confused string +} + +func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Format, url string, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, userPayload []byte, rawJSON []byte, headerSets ...http.Header) (*http.Request, []byte, codexIdentityConfuseState, error) { + var headers http.Header + if len(headerSets) > 0 { + headers = headerSets[0] + } + var cache helps.CodexCache + if sourceFormatEqual(from, sdktranslator.FormatClaude) { + modelName := strings.TrimSpace(gjson.GetBytes(rawJSON, "model").String()) + if modelName == "" { + modelName = thinking.ParseSuffix(req.Model).ModelName + } + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, modelName, req.Payload, headers) + if errCache != nil { + return nil, nil, codexIdentityConfuseState{}, errCache + } + if ok { + cache = cached + } + } else if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) { + promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key") + if promptCacheKey.Exists() { + cache.ID = promptCacheKey.String() + } + } else if sourceFormatEqual(from, sdktranslator.FormatOpenAI) { + if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() { + cache.ID = strings.TrimSpace(promptCacheKey.String()) + } + if cache.ID == "" { + cache.ID = helps.ProviderSessionUUID("codex", req.Metadata) + } + if cache.ID == "" { + if apiKey := strings.TrimSpace(helps.APIKeyFromContext(ctx)); apiKey != "" { + cache.ID = uuid.NewSHA1(uuid.NameSpaceOID, []byte("cli-proxy-api:codex:prompt-cache:"+apiKey)).String() + } + } + } + if cache.ID == "" { + cache.ID = helps.ProviderSessionUUID("codex", req.Metadata) + } + + if cache.ID != "" { + rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", cache.ID) + } + rawJSON = helps.SanitizeCodexInputItemIDs(rawJSON) + var identityState codexIdentityConfuseState + rawJSON, identityState = applyCodexIdentityConfuseBody(e.cfg, auth, userPayload, rawJSON) + if identityState.promptCacheKey != "" { + cache.ID = identityState.promptCacheKey + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(rawJSON)) + if err != nil { + return nil, nil, codexIdentityConfuseState{}, err + } + if cache.ID != "" { + httpReq.Header.Set("Session_id", cache.ID) + } + return httpReq, rawJSON, identityState, nil +} + +func applyCodexIdentityConfuseBody(cfg *config.Config, auth *cliproxyauth.Auth, userPayload []byte, rawJSON []byte) ([]byte, codexIdentityConfuseState) { + if !codexIdentityConfuseEnabled(cfg) || auth == nil || strings.TrimSpace(auth.ID) == "" || len(rawJSON) == 0 { + return rawJSON, codexIdentityConfuseState{} + } + + state := codexIdentityConfuseState{enabled: true, authID: strings.TrimSpace(auth.ID)} + if promptCacheKey := strings.TrimSpace(gjson.GetBytes(userPayload, "prompt_cache_key").String()); promptCacheKey != "" { + state.originalPromptCacheKey = promptCacheKey + state.promptCacheKey = codexIdentityConfuseUUID(auth.ID, "prompt-cache", promptCacheKey) + rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", state.promptCacheKey) + } + if installationID := strings.TrimSpace(gjson.GetBytes(userPayload, "client_metadata.x-codex-installation-id").String()); installationID != "" { + rawJSON, _ = sjson.SetBytes(rawJSON, "client_metadata.x-codex-installation-id", codexIdentityConfuseUUID(auth.ID, "installation", installationID)) + } + if turnMetadata := strings.TrimSpace(gjson.GetBytes(rawJSON, "client_metadata.x-codex-turn-metadata").String()); turnMetadata != "" { + rawJSON, _ = sjson.SetBytes(rawJSON, "client_metadata.x-codex-turn-metadata", applyCodexTurnMetadataIdentityConfuse(turnMetadata, &state)) + } + if state.promptCacheKey != "" { + if windowID := strings.TrimSpace(gjson.GetBytes(rawJSON, "client_metadata.x-codex-window-id").String()); windowID != "" { + rawJSON, _ = sjson.SetBytes(rawJSON, "client_metadata.x-codex-window-id", state.promptCacheKey+":0") + } + } + + return rawJSON, state +} + +func applyCodexIdentityConfuseHeaders(headers http.Header, state *codexIdentityConfuseState) { + if headers == nil { + return + } + if state == nil || !state.enabled { + return + } + + if rawTurnMetadata := strings.TrimSpace(headers.Get("X-Codex-Turn-Metadata")); rawTurnMetadata != "" { + headers.Set("X-Codex-Turn-Metadata", applyCodexTurnMetadataIdentityConfuse(rawTurnMetadata, state)) + } + if state.promptCacheKey == "" { + return + } + + setCodexSessionHeaderCasePreserved(headers, "Session_id", state.promptCacheKey) + if headerValueCaseInsensitive(headers, "Conversation_id") != "" { + setHeaderCasePreserved(headers, "Conversation_id", state.promptCacheKey) + } + headers.Set("X-Client-Request-Id", state.promptCacheKey) + headers.Set("Thread-Id", state.promptCacheKey) + headers.Set("X-Codex-Window-Id", state.promptCacheKey+":0") +} + +func applyCodexTurnMetadataIdentityConfuse(rawTurnMetadata string, state *codexIdentityConfuseState) string { + updatedTurnMetadata := rawTurnMetadata + if state == nil || !state.enabled { + return updatedTurnMetadata + } + if state.promptCacheKey != "" && gjson.Get(rawTurnMetadata, "prompt_cache_key").Exists() { + updatedTurnMetadata, _ = sjson.Set(updatedTurnMetadata, "prompt_cache_key", state.promptCacheKey) + } else if state.promptCacheKey != "" && state.originalPromptCacheKey != "" { + updatedTurnMetadata = strings.ReplaceAll(updatedTurnMetadata, state.originalPromptCacheKey, state.promptCacheKey) + } + if turnID := strings.TrimSpace(gjson.Get(rawTurnMetadata, "turn_id").String()); turnID != "" { + updatedTurnMetadata, _ = sjson.Set(updatedTurnMetadata, "turn_id", state.confuseTurnID(turnID)) + } + if state.promptCacheKey != "" && gjson.Get(rawTurnMetadata, "window_id").Exists() { + updatedTurnMetadata, _ = sjson.Set(updatedTurnMetadata, "window_id", state.promptCacheKey+":0") + } + return updatedTurnMetadata +} + +func applyCodexIdentityConfuseResponsePayload(payload []byte, state codexIdentityConfuseState) []byte { + payload = replaceCodexIdentityResponsePayload(payload, state.originalPromptCacheKey, state.promptCacheKey) + for _, turnID := range state.turnIDs { + payload = replaceCodexIdentityResponsePayload(payload, turnID.original, turnID.confused) + } + return payload +} + +func applyCodexIdentityExposeResponsePayload(payload []byte, state codexIdentityConfuseState) []byte { + payload = replaceCodexIdentityResponsePayload(payload, state.promptCacheKey, state.originalPromptCacheKey) + for _, turnID := range state.turnIDs { + payload = replaceCodexIdentityResponsePayload(payload, turnID.confused, turnID.original) + } + return payload +} + +func (state *codexIdentityConfuseState) confuseTurnID(turnID string) string { + turnID = strings.TrimSpace(turnID) + if state == nil || !state.enabled || strings.TrimSpace(state.authID) == "" || turnID == "" { + return turnID + } + for _, replacement := range state.turnIDs { + if replacement.original == turnID || replacement.confused == turnID { + return replacement.confused + } + } + confusedTurnID := codexIdentityConfuseUUID(state.authID, "turn", turnID) + state.turnIDs = append(state.turnIDs, codexIdentityReplacement{original: turnID, confused: confusedTurnID}) + return confusedTurnID +} + +func replaceCodexIdentityResponsePayload(payload []byte, from string, to string) []byte { + from = strings.TrimSpace(from) + to = strings.TrimSpace(to) + if len(payload) == 0 || from == "" || to == "" || from == to || !bytes.Contains(payload, []byte(from)) { + return payload + } + return bytes.ReplaceAll(payload, []byte(from), []byte(to)) +} + +func codexIdentityConfuseEnabled(cfg *config.Config) bool { + if cfg == nil || !cfg.Codex.IdentityConfuse { + return false + } + strategy := strings.ToLower(strings.TrimSpace(cfg.Routing.Strategy)) + return cfg.Routing.SessionAffinity || strategy == "fill-first" || strategy == "fillfirst" || strategy == "ff" +} + +func codexIdentityConfuseUUID(authID string, kind string, value string) string { + name := strings.Join([]string{"cli-proxy-api", "codex", "identity-confuse", kind, strings.TrimSpace(authID), strings.TrimSpace(value)}, ":") + return uuid.NewSHA1(uuid.NameSpaceOID, []byte(name)).String() +} + +func applyCodexHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config) { + var ginHeaders http.Header + if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + ginHeaders = ginCtx.Request.Header + } + applyCodexHeadersFromSources(r, auth, token, stream, cfg, ginHeaders) +} + +// applyModelHeaderOverrides forces models.json config.override_header onto upstream headers. +func applyModelHeaderOverrides(headers http.Header, modelName string) { + if headers == nil { + return + } + overrides := registry.ModelOverrideHeaders(modelName) + if len(overrides) == 0 { + return + } + for key, value := range overrides { + headers.Set(key, value) + } + if strings.Contains(headers.Get("User-Agent"), "Mac OS") && codexSessionHeaderValue(headers) == "" { + headers.Set("Session_id", uuid.NewString()) + } +} + +// applyCodexDirectImageHeaders sets Codex upstream headers for direct /images/* calls. +// Downstream client User-Agent values are not forwarded to reduce Cloudflare 1010 blocks. +func applyCodexDirectImageHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config) { + var ginHeaders http.Header + if ginCtx, ok := r.Context().Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + ginHeaders = ginCtx.Request.Header.Clone() + ginHeaders.Del("User-Agent") + } + applyCodexHeadersFromSources(r, auth, token, stream, cfg, ginHeaders) +} + +func applyCodexHeadersFromSources(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, cfg *config.Config, ginHeaders http.Header) { + r.Header.Set("Content-Type", "application/json") + r.Header.Set("Authorization", "Bearer "+token) + + if ginHeaders != nil && ginHeaders.Get("X-Codex-Beta-Features") != "" { + r.Header.Set("X-Codex-Beta-Features", ginHeaders.Get("X-Codex-Beta-Features")) + } + misc.EnsureHeader(r.Header, ginHeaders, "Version", "") + misc.EnsureHeader(r.Header, ginHeaders, "X-Codex-Turn-Metadata", "") + misc.EnsureHeader(r.Header, ginHeaders, "X-Client-Request-Id", "") + cfgUserAgent, _ := codexHeaderDefaults(cfg, auth) + ensureHeaderWithConfigPrecedence(r.Header, ginHeaders, "User-Agent", cfgUserAgent, codexUserAgent) + + if strings.Contains(r.Header.Get("User-Agent"), "Mac OS") { + misc.EnsureHeader(r.Header, ginHeaders, "Session_id", uuid.NewString()) + } + + if stream { + r.Header.Set("Accept", "text/event-stream") + } else { + r.Header.Set("Accept", "application/json") + } + r.Header.Set("Connection", "Keep-Alive") + + isAPIKey := false + if auth != nil && auth.Attributes != nil { + if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" { + isAPIKey = true + } + } + if originator := strings.TrimSpace(ginHeaders.Get("Originator")); originator != "" { + r.Header.Set("Originator", originator) + } else if !isAPIKey { + r.Header.Set("Originator", codexOriginator) + } + if !isAPIKey { + if auth != nil && auth.Metadata != nil { + if accountID, ok := auth.Metadata["account_id"].(string); ok { + r.Header.Set("Chatgpt-Account-Id", accountID) + } + } + } + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(r, attrs) +} + +func normalizeCodexInstructions(body []byte) []byte { + instructions := gjson.GetBytes(body, "instructions") + if !instructions.Exists() || instructions.Type == gjson.Null { + body, _ = sjson.SetBytes(body, "instructions", "") + } + return body +} + +var imageGenToolJSON = []byte(`{"type":"image_generation","output_format":"png"}`) +var imageGenToolArrayJSON = []byte(`[{"type":"image_generation","output_format":"png"}]`) + +func isCodexFreePlanAuth(auth *cliproxyauth.Auth) bool { + if auth == nil || auth.Attributes == nil { + return false + } + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + return false + } + return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free") +} + +func isImageGenerationFunctionTool(tool gjson.Result) bool { + switch tool.Get("type").String() { + case "function": + return tool.Get("name").String() == "image_gen.imagegen" + case "namespace": + if tool.Get("name").String() != "image_gen" { + return false + } + tools := tool.Get("tools") + if !tools.IsArray() { + return false + } + for _, nestedTool := range tools.Array() { + if nestedTool.Get("type").String() == "function" && nestedTool.Get("name").String() == "imagegen" { + return true + } + } + } + return false +} + +func isCodexResponsesLiteRequest(body []byte, headers http.Header) bool { + if strings.EqualFold(strings.TrimSpace(headers.Get(codexResponsesLiteHeader)), "true") { + return true + } + // Codex Desktop mirrors websocket-only request headers into client_metadata. + value := gjson.GetBytes(body, codexResponsesLiteMetadata) + if !value.Exists() { + return false + } + return value.Type == gjson.True || value.Type == gjson.String && strings.EqualFold(strings.TrimSpace(value.String()), "true") +} + +func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth.Auth, headers http.Header) []byte { + if isCodexResponsesLiteRequest(body, headers) { + return body + } + if strings.HasSuffix(baseModel, "spark") { + return body + } + if isCodexFreePlanAuth(auth) { + return body + } + + tools := gjson.GetBytes(body, "tools") + if !tools.Exists() || !tools.IsArray() { + body, _ = sjson.SetRawBytes(body, "tools", imageGenToolArrayJSON) + return body + } + for _, t := range tools.Array() { + if t.Get("type").String() == "image_generation" || isImageGenerationFunctionTool(t) { + return body + } + } + body, _ = sjson.SetRawBytes(body, "tools.-1", imageGenToolJSON) + return body +} + +func normalizeCodexParallelToolCalls(body []byte, headers http.Header) []byte { + if isCodexResponsesLiteRequest(body, headers) { + body = helps.SetBoolIfDifferent(body, "parallel_tool_calls", false) + return body + } + return normalizeCodexParallelToolCallsForTools(body) +} + +func normalizeCodexParallelToolCallsForTools(body []byte) []byte { + if !gjson.GetBytes(body, "parallel_tool_calls").Exists() { + return body + } + + tools := gjson.GetBytes(body, "tools") + hasTools := tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 + if hasTools { + return body + } + + body, _ = sjson.DeleteBytes(body, "parallel_tool_calls") + return body +} + +func publishCodexImageToolUsage(ctx context.Context, reporter *helps.UsageReporter, body []byte, completedData []byte) { + detail, ok := helps.ParseCodexImageToolUsage(completedData) + if !ok { + return + } + reporter.EnsurePublished(ctx) + reporter.PublishAdditionalModel(ctx, codexImageGenerationToolModel(body), detail) +} + +func codexImageGenerationToolModel(body []byte) string { + tools := gjson.GetBytes(body, "tools") + if tools.IsArray() { + for _, tool := range tools.Array() { + if tool.Get("type").String() != "image_generation" { + continue + } + if model := strings.TrimSpace(tool.Get("model").String()); model != "" { + return model + } + break + } + } + return codexDefaultImageToolModel +} diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go new file mode 100644 index 000000000..8d5c89930 --- /dev/null +++ b/internal/runtime/executor/codex_executor_stream.go @@ -0,0 +1,217 @@ +package executor + +import ( + "bufio" + "bytes" + "context" + "io" + "net/http" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { + if opts.Alt == "responses/compact" { + return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} + } + if isCodexOpenAIImageRequest(opts) { + return e.executeOpenAIImageStream(ctx, auth, req, opts) + } + baseModel := thinking.ParseSuffix(req.Model).ModelName + + apiKey, baseURL := codexCreds(auth) + if baseURL == "" { + baseURL = "https://chatgpt.com/backend-api/codex" + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("codex") + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, true) + + body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return nil, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body, _ = sjson.DeleteBytes(body, "previous_response_id") + body, _ = sjson.DeleteBytes(body, "generate") + body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") + body, _ = sjson.DeleteBytes(body, "safety_identifier") + body, _ = sjson.DeleteBytes(body, "stream_options") + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = normalizeCodexInstructions(body) + if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { + body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) + } + body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) + body = normalizeCodexParallelToolCalls(body, opts.Headers) + body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return nil, errReplay + } + reporter.SetTranslatedReasoningEffort(body, to.String()) + + url := strings.TrimSuffix(baseURL, "/") + "/responses" + var identityState codexIdentityConfuseState + httpReq, upstreamBody, identityState, err := e.cacheHelper(ctx, from, url, auth, req, originalPayloadSource, body, opts.Headers) + if err != nil { + return nil, err + } + applyCodexHeaders(httpReq, auth, apiKey, true, e.cfg) + applyModelHeaderOverrides(httpReq.Header, baseModel) + applyCodexIdentityConfuseHeaders(httpReq.Header, &identityState) + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: url, + Method: http.MethodPost, + Headers: httpReq.Header.Clone(), + Body: upstreamBody, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + + httpClient := helps.NewUtlsHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return nil, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + data, readErr := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("codex executor: close response body error: %v", errClose) + } + if readErr != nil { + helps.RecordAPIResponseError(ctx, e.cfg, readErr) + return nil, readErr + } + data = applyCodexIdentityConfuseResponsePayload(data, identityState) + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, httpResp.StatusCode, data); errClearReplay != nil { + return nil, errClearReplay + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) + err = newCodexStatusErr(httpResp.StatusCode, data) + return nil, err + } + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(out) + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("codex executor: close response body error: %v", errClose) + } + }() + scanner := bufio.NewScanner(httpResp.Body) + scanner.Buffer(nil, 52_428_800) // 50MB + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) + var param any + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + for scanner.Scan() { + line := applyCodexIdentityConfuseResponsePayload(scanner.Bytes(), identityState) + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + translatedLine := bytes.Clone(line) + terminalSuccess := false + + if bytes.HasPrefix(line, dataTag) { + data := bytes.TrimSpace(line[5:]) + data = helps.RestoreCodexMultiAgentV2Response(data, optimizeMultiAgentV2) + translatedLine = append([]byte("data: "), data...) + eventType := gjson.GetBytes(data, "type").String() + if streamErr, terminalBody, ok := codexTerminalFailureErr(data); ok { + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + select { + case out <- cliproxyexecutor.StreamChunk{Err: errClearReplay}: + case <-ctx.Done(): + } + return + } + helps.RecordAPIResponseError(ctx, e.cfg, streamErr) + reporter.PublishFailure(ctx, streamErr) + select { + case out <- cliproxyexecutor.StreamChunk{Err: streamErr}: + case <-ctx.Done(): + } + return + } + switch eventType { + case "response.output_item.done": + collectCodexOutputItemDone(data, outputItemsByIndex, &outputItemsFallback) + case "response.completed", "response.incomplete": + terminalSuccess = true + if detail, ok := helps.ParseCodexUsage(data); ok { + reporter.Publish(ctx, detail) + } + publishCodexImageToolUsage(ctx, reporter, body, data) + data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback) + if eventType == "response.completed" { + cacheCodexReasoningReplayFromCompleted(replayScope, data) + } + translatedLine = append([]byte("data: "), data...) + } + } + + translatedLine = applyCodexIdentityExposeResponsePayload(translatedLine, identityState) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m, claudeInputTokens) + for i := range chunks { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: + case <-ctx.Done(): + return + } + } + if terminalSuccess { + return + } + } + if errScan := scanner.Err(); errScan != nil { + if ctx.Err() != nil { + return + } + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + } + streamErr := newCodexIncompleteStreamError() + helps.RecordAPIResponseError(ctx, e.cfg, streamErr) + reporter.PublishFailure(ctx, streamErr) + select { + case out <- cliproxyexecutor.StreamChunk{Err: streamErr}: + case <-ctx.Done(): + } + }() + return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil +} diff --git a/internal/runtime/executor/codex_executor_terminal.go b/internal/runtime/executor/codex_executor_terminal.go new file mode 100644 index 000000000..3ebc3d4ec --- /dev/null +++ b/internal/runtime/executor/codex_executor_terminal.go @@ -0,0 +1,373 @@ +package executor + +import ( + "bytes" + "net/http" + "sort" + "strings" + "time" + + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const codexIncompleteStreamMessage = "stream error: stream disconnected before completion: stream closed before response.completed" + +type codexIncompleteStreamError struct { + statusErr +} + +func newCodexIncompleteStreamError() codexIncompleteStreamError { + return codexIncompleteStreamError{statusErr: statusErr{ + code: http.StatusRequestTimeout, + msg: codexIncompleteStreamMessage, + }} +} + +func (codexIncompleteStreamError) IsRequestScoped() bool { + return true +} + +// Streamed Codex responses may emit response.output_item.done events while leaving +// response.completed.response.output empty. Keep the stream path aligned with the +// already-patched non-stream path by reconstructing response.output from those items. +func collectCodexOutputItemDone(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback *[][]byte) { + itemResult := gjson.GetBytes(eventData, "item") + if !itemResult.Exists() || itemResult.Type != gjson.JSON { + return + } + outputIndexResult := gjson.GetBytes(eventData, "output_index") + if outputIndexResult.Exists() { + outputItemsByIndex[outputIndexResult.Int()] = []byte(itemResult.Raw) + return + } + *outputItemsFallback = append(*outputItemsFallback, []byte(itemResult.Raw)) +} + +func patchCodexCompletedOutput(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { + outputResult := gjson.GetBytes(eventData, "response.output") + shouldPatchOutput := (!outputResult.Exists() || !outputResult.IsArray() || len(outputResult.Array()) == 0) && (len(outputItemsByIndex) > 0 || len(outputItemsFallback) > 0) + if !shouldPatchOutput { + return eventData + } + + indexes := make([]int64, 0, len(outputItemsByIndex)) + for idx := range outputItemsByIndex { + indexes = append(indexes, idx) + } + sort.Slice(indexes, func(i, j int) bool { + return indexes[i] < indexes[j] + }) + + items := make([][]byte, 0, len(outputItemsByIndex)+len(outputItemsFallback)) + for _, idx := range indexes { + items = append(items, outputItemsByIndex[idx]) + } + items = append(items, outputItemsFallback...) + + outputArray := []byte("[]") + if len(items) > 0 { + var buf bytes.Buffer + totalLen := 2 + for _, item := range items { + totalLen += len(item) + } + if len(items) > 1 { + totalLen += len(items) - 1 + } + buf.Grow(totalLen) + buf.WriteByte('[') + for i, item := range items { + if i > 0 { + buf.WriteByte(',') + } + buf.Write(item) + } + buf.WriteByte(']') + outputArray = buf.Bytes() + } + + completedDataPatched, _ := sjson.SetRawBytes(eventData, "response.output", outputArray) + return completedDataPatched +} + +func codexTerminalStreamContextLengthErr(eventData []byte) (statusErr, bool) { + streamErr, body, ok := codexTerminalStreamErr(eventData) + if !ok || !codexTerminalErrorIsContextLength(body) { + return statusErr{}, false + } + return streamErr, true +} + +func codexTerminalStreamErr(eventData []byte) (statusErr, []byte, bool) { + body, ok := codexTerminalFailureBody(eventData) + if !ok || !codexTerminalStreamErrShouldHandle(body) { + return statusErr{}, nil, false + } + return newCodexStatusErr(http.StatusBadRequest, body), body, true +} + +func codexTerminalFailureErr(eventData []byte) (statusErr, []byte, bool) { + if streamErr, body, ok := codexTerminalStreamErr(eventData); ok { + return streamErr, body, true + } + body, ok := codexTerminalFailureBody(eventData) + if !ok { + return statusErr{}, nil, false + } + return newCodexStatusErr(codexTerminalFailureStatus(body), body), body, true +} + +func codexTerminalFailureStatus(body []byte) int { + for _, path := range []string{"error.status_code", "error.status"} { + if status := int(gjson.GetBytes(body, path).Int()); status >= 400 && status <= 599 { + return status + } + } + + errorType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.type").String())) + errorCode := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.code").String())) + switch { + case errorType == "invalid_request_error", errorType == "bad_request_error": + return http.StatusBadRequest + case errorType == "authentication_error", errorCode == "invalid_api_key", errorCode == "unauthorized": + return http.StatusUnauthorized + case errorType == "permission_error", errorCode == "forbidden", errorCode == "permission_denied": + return http.StatusForbidden + case errorType == "not_found_error", errorCode == "not_found", errorCode == "model_not_found": + return http.StatusNotFound + case errorType == "rate_limit_error", errorCode == "rate_limit_exceeded": + return http.StatusTooManyRequests + default: + return http.StatusBadGateway + } +} + +func codexTerminalFailureBody(eventData []byte) ([]byte, bool) { + eventType := gjson.GetBytes(eventData, "type").String() + var body []byte + switch eventType { + case "error": + body = codexTerminalErrorBody(eventData, "error") + if len(body) == 0 { + body = codexTerminalTopLevelErrorBody(eventData) + } + case "response.failed": + body = codexTerminalErrorBody(eventData, "response.error") + if len(body) == 0 { + body = codexTerminalErrorBody(eventData, "error") + } + default: + return nil, false + } + if len(body) == 0 { + body = []byte(`{"error":{"message":"upstream stream failed without error details"}}`) + } + return body, true +} + +func codexTerminalStreamErrShouldHandle(body []byte) bool { + if codexTerminalErrorIsContextLength(body) { + return true + } + if isCodexUsageLimitError(body) || isCodexModelCapacityError(body) { + return true + } + code, _, ok := codexStatusErrorClassification(http.StatusBadRequest, body) + return ok && code == "thinking_signature_invalid" +} + +func codexTerminalErrorBody(eventData []byte, path string) []byte { + errorResult := gjson.GetBytes(eventData, path) + if !errorResult.Exists() { + return nil + } + body := []byte(`{"error":{}}`) + if errorResult.Type == gjson.JSON { + body, _ = sjson.SetRawBytes(body, "error", []byte(errorResult.Raw)) + } else if message := strings.TrimSpace(errorResult.String()); message != "" { + body, _ = sjson.SetBytes(body, "error.message", message) + } + if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { + if message := strings.TrimSpace(gjson.GetBytes(eventData, "response.error.message").String()); message != "" { + body, _ = sjson.SetBytes(body, "error.message", message) + } + } + if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { + if code := strings.TrimSpace(gjson.GetBytes(body, "error.code").String()); code != "" { + body, _ = sjson.SetBytes(body, "error.message", code) + } + } + if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { + if errorType := strings.TrimSpace(gjson.GetBytes(body, "error.type").String()); errorType != "" { + body, _ = sjson.SetBytes(body, "error.message", errorType) + } + } + return body +} + +func codexTerminalTopLevelErrorBody(eventData []byte) []byte { + message := strings.TrimSpace(gjson.GetBytes(eventData, "message").String()) + code := strings.TrimSpace(gjson.GetBytes(eventData, "code").String()) + errorType := strings.TrimSpace(gjson.GetBytes(eventData, "error_type").String()) + param := strings.TrimSpace(gjson.GetBytes(eventData, "param").String()) + if message == "" && code == "" && errorType == "" && param == "" { + return nil + } + + body := []byte(`{"error":{}}`) + if message != "" { + body, _ = sjson.SetBytes(body, "error.message", message) + } + if code != "" { + body, _ = sjson.SetBytes(body, "error.code", code) + } + if errorType != "" { + body, _ = sjson.SetBytes(body, "error.type", errorType) + } + if param != "" { + body, _ = sjson.SetBytes(body, "error.param", param) + } + if strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) == "" { + if code != "" { + body, _ = sjson.SetBytes(body, "error.message", code) + } else if errorType != "" { + body, _ = sjson.SetBytes(body, "error.message", errorType) + } + } + return body +} + +func codexTerminalErrorIsContextLength(body []byte) bool { + errorCode := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.code").String())) + message := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.message").String())) + return errorCode == "context_length_exceeded" || + errorCode == "context_too_large" || + strings.Contains(message, "context window") || + strings.Contains(message, "context length") || + strings.Contains(message, "too many tokens") +} + +func newCodexStatusErr(statusCode int, body []byte) statusErr { + errCode := statusCode + if isCodexModelCapacityError(body) || isCodexUsageLimitError(body) { + errCode = http.StatusTooManyRequests + } + body = classifyCodexStatusError(errCode, body) + err := statusErr{code: errCode, msg: string(body)} + if retryAfter := parseCodexRetryAfter(errCode, body, time.Now()); retryAfter != nil { + err.retryAfter = retryAfter + } + return err +} + +func classifyCodexStatusError(statusCode int, body []byte) []byte { + code, errType, ok := codexStatusErrorClassification(statusCode, body) + if !ok { + return body + } + message := gjson.GetBytes(body, "error.message").String() + if message == "" { + message = gjson.GetBytes(body, "message").String() + } + if message == "" { + message = strings.TrimSpace(string(body)) + } + if message == "" { + message = http.StatusText(statusCode) + } + out := []byte(`{"error":{}}`) + out, _ = sjson.SetBytes(out, "error.message", message) + out, _ = sjson.SetBytes(out, "error.type", errType) + out, _ = sjson.SetBytes(out, "error.code", code) + return out +} + +func codexStatusErrorClassification(statusCode int, body []byte) (code string, errType string, ok bool) { + errorMessage := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.message").String())) + if errorMessage == "" { + errorMessage = strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "message").String())) + } + lower := strings.ToLower(strings.TrimSpace(string(body))) + upstreamCode := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.code").String())) + upstreamType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.type").String())) + isInvalidRequest := upstreamType == "" || upstreamType == "invalid_request_error" + + switch { + case statusCode == http.StatusRequestEntityTooLarge || upstreamCode == "context_length_exceeded" || upstreamCode == "context_too_large" || isInvalidRequest && (strings.Contains(errorMessage, "context length") || strings.Contains(errorMessage, "context_length") || strings.Contains(errorMessage, "maximum context") || strings.Contains(errorMessage, "too many tokens")): + return "context_too_large", "invalid_request_error", true + case strings.Contains(lower, "invalid signature in thinking block") || strings.Contains(lower, "invalid_encrypted_content"): + return "thinking_signature_invalid", "invalid_request_error", true + case upstreamCode == "previous_response_not_found" || strings.Contains(lower, "previous_response_not_found") || strings.Contains(lower, "previous_response_id") && strings.Contains(lower, "not found"): + return "previous_response_not_found", "invalid_request_error", true + case statusCode == http.StatusUnauthorized || upstreamType == "authentication_error" || upstreamCode == "invalid_api_key" || strings.Contains(lower, "invalid or expired token") || strings.Contains(lower, "refresh_token_reused"): + return "auth_unavailable", "authentication_error", true + default: + return "", "", false + } +} + +func isCodexModelCapacityError(errorBody []byte) bool { + if len(errorBody) == 0 { + return false + } + candidates := []string{ + gjson.GetBytes(errorBody, "error.message").String(), + gjson.GetBytes(errorBody, "message").String(), + string(errorBody), + } + for _, candidate := range candidates { + lower := strings.ToLower(strings.TrimSpace(candidate)) + if lower == "" { + continue + } + if strings.Contains(lower, "selected model is at capacity") || + strings.Contains(lower, "model is at capacity. please try a different model") { + return true + } + } + return false +} + +// isCodexUsageLimitError reports whether the error body represents a Codex +// quota/plan-limit exhaustion (error.type == "usage_limit_reached"). This is the +// signal Codex emits when a credential's usage quota is depleted, and it carries +// reset timing (resets_at/resets_in_seconds) parsed by parseCodexRetryAfter. +// Transient per-minute rate limits (rate_limit_error/rate_limit_exceeded) are +// intentionally excluded, as they should be retried rather than cooled down. +func isCodexUsageLimitError(errorBody []byte) bool { + if len(errorBody) == 0 { + return false + } + candidates := []string{ + gjson.GetBytes(errorBody, "error.type").String(), + gjson.GetBytes(errorBody, "type").String(), + } + for _, candidate := range candidates { + if strings.EqualFold(strings.TrimSpace(candidate), "usage_limit_reached") { + return true + } + } + return false +} + +func parseCodexRetryAfter(statusCode int, errorBody []byte, now time.Time) *time.Duration { + if statusCode != http.StatusTooManyRequests || len(errorBody) == 0 { + return nil + } + if strings.TrimSpace(gjson.GetBytes(errorBody, "error.type").String()) != "usage_limit_reached" { + return nil + } + if resetsAt := gjson.GetBytes(errorBody, "error.resets_at").Int(); resetsAt > 0 { + resetAtTime := time.Unix(resetsAt, 0) + if resetAtTime.After(now) { + retryAfter := resetAtTime.Sub(now) + return &retryAfter + } + } + if resetsInSeconds := gjson.GetBytes(errorBody, "error.resets_in_seconds").Int(); resetsInSeconds > 0 { + retryAfter := time.Duration(resetsInSeconds) * time.Second + return &retryAfter + } + return nil +} diff --git a/internal/runtime/executor/codex_executor_tokens.go b/internal/runtime/executor/codex_executor_tokens.go new file mode 100644 index 000000000..9a6877801 --- /dev/null +++ b/internal/runtime/executor/codex_executor_tokens.go @@ -0,0 +1,175 @@ +package executor + +import ( + "context" + "fmt" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + "github.com/tiktoken-go/tokenizer" +) + +func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + baseModel := thinking.ParseSuffix(req.Model).ModelName + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("codex") + body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, false) + + body, err := thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return cliproxyexecutor.Response{}, err + } + + body = helps.SetStringIfDifferent(body, "model", baseModel) + body, _ = sjson.DeleteBytes(body, "previous_response_id") + body, _ = sjson.DeleteBytes(body, "generate") + body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") + body, _ = sjson.DeleteBytes(body, "safety_identifier") + body, _ = sjson.DeleteBytes(body, "stream_options") + body = helps.SetBoolIfDifferent(body, "stream", false) + body = normalizeCodexInstructions(body) + + enc, err := tokenizerForCodexModel(baseModel) + if err != nil { + return cliproxyexecutor.Response{}, fmt.Errorf("codex executor: tokenizer init failed: %w", err) + } + + count, err := countCodexInputTokens(enc, body) + if err != nil { + return cliproxyexecutor.Response{}, fmt.Errorf("codex executor: token counting failed: %w", err) + } + + usageJSON := fmt.Sprintf(`{"response":{"usage":{"input_tokens":%d,"output_tokens":0,"total_tokens":%d}}}`, count, count) + translated := sdktranslator.TranslateTokenCount(ctx, to, responseFormat, count, []byte(usageJSON)) + return cliproxyexecutor.Response{Payload: translated}, nil +} + +func tokenizerForCodexModel(model string) (tokenizer.Codec, error) { + sanitized := strings.ToLower(strings.TrimSpace(model)) + switch { + case sanitized == "": + return tokenizer.Get(tokenizer.Cl100kBase) + case strings.HasPrefix(sanitized, "gpt-5"): + return tokenizer.ForModel(tokenizer.GPT5) + case strings.HasPrefix(sanitized, "gpt-4.1"): + return tokenizer.ForModel(tokenizer.GPT41) + case strings.HasPrefix(sanitized, "gpt-4o"): + return tokenizer.ForModel(tokenizer.GPT4o) + case strings.HasPrefix(sanitized, "gpt-4"): + return tokenizer.ForModel(tokenizer.GPT4) + case strings.HasPrefix(sanitized, "gpt-3.5"), strings.HasPrefix(sanitized, "gpt-3"): + return tokenizer.ForModel(tokenizer.GPT35Turbo) + default: + return tokenizer.Get(tokenizer.Cl100kBase) + } +} + +func countCodexInputTokens(enc tokenizer.Codec, body []byte) (int64, error) { + if enc == nil { + return 0, fmt.Errorf("encoder is nil") + } + if len(body) == 0 { + return 0, nil + } + + root := gjson.ParseBytes(body) + var segments []string + + if inst := strings.TrimSpace(root.Get("instructions").String()); inst != "" { + segments = append(segments, inst) + } + + inputItems := root.Get("input") + if inputItems.IsArray() { + arr := inputItems.Array() + for i := range arr { + item := arr[i] + switch item.Get("type").String() { + case "message": + content := item.Get("content") + if content.IsArray() { + parts := content.Array() + for j := range parts { + part := parts[j] + if text := strings.TrimSpace(part.Get("text").String()); text != "" { + segments = append(segments, text) + } + } + } + case "function_call": + if name := strings.TrimSpace(item.Get("name").String()); name != "" { + segments = append(segments, name) + } + if args := strings.TrimSpace(item.Get("arguments").String()); args != "" { + segments = append(segments, args) + } + case "function_call_output": + if out := strings.TrimSpace(item.Get("output").String()); out != "" { + segments = append(segments, out) + } + default: + if text := strings.TrimSpace(item.Get("text").String()); text != "" { + segments = append(segments, text) + } + } + } + } + + tools := root.Get("tools") + if tools.IsArray() { + tarr := tools.Array() + for i := range tarr { + tool := tarr[i] + if name := strings.TrimSpace(tool.Get("name").String()); name != "" { + segments = append(segments, name) + } + if desc := strings.TrimSpace(tool.Get("description").String()); desc != "" { + segments = append(segments, desc) + } + if params := tool.Get("parameters"); params.Exists() { + val := params.Raw + if params.Type == gjson.String { + val = params.String() + } + if trimmed := strings.TrimSpace(val); trimmed != "" { + segments = append(segments, trimmed) + } + } + } + } + + textFormat := root.Get("text.format") + if textFormat.Exists() { + if name := strings.TrimSpace(textFormat.Get("name").String()); name != "" { + segments = append(segments, name) + } + if schema := textFormat.Get("schema"); schema.Exists() { + val := schema.Raw + if schema.Type == gjson.String { + val = schema.String() + } + if trimmed := strings.TrimSpace(val); trimmed != "" { + segments = append(segments, trimmed) + } + } + } + + text := strings.Join(segments, "\n") + if text == "" { + return 0, nil + } + + count, err := enc.Count(text) + if err != nil { + return 0, err + } + return int64(count), nil +} diff --git a/internal/runtime/executor/codex_websockets_connection.go b/internal/runtime/executor/codex_websockets_connection.go new file mode 100644 index 000000000..7411ed624 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_connection.go @@ -0,0 +1,240 @@ +package executor + +import ( + "bytes" + "context" + "errors" + "fmt" + "net" + "net/http" + "net/url" + "strings" + "time" + + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil" + log "github.com/sirupsen/logrus" + "github.com/tidwall/sjson" + "golang.org/x/net/proxy" +) + +const ( + codexResponsesWebsocketBetaHeaderValue = "responses_websockets=2026-02-06" + codexResponsesWebsocketIdleTimeout = 5 * time.Minute + codexResponsesWebsocketHandshakeTO = 30 * time.Second +) + +func (e *CodexWebsocketsExecutor) dialCodexWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { + dialer := newProxyAwareWebsocketDialer(e.cfg, auth) + dialer.HandshakeTimeout = codexResponsesWebsocketHandshakeTO + dialer.EnableCompression = true + if ctx == nil { + ctx = context.Background() + } + conn, resp, err := dialer.DialContext(ctx, wsURL, headers) + closer := newWebsocketConnectionCloser(conn) + if conn != nil { + // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. + // Negotiating permessage-deflate is fine; we just don't compress outbound messages. + conn.EnableWriteCompression(false) + } + return conn, closer, resp, err +} + +func writeCodexWebsocketMessage(sess *codexWebsocketSession, conn *websocket.Conn, payload []byte) error { + if sess != nil { + return sess.writeMessage(conn, websocket.TextMessage, payload) + } + if conn == nil { + return fmt.Errorf("codex websockets executor: websocket conn is nil") + } + return conn.WriteMessage(websocket.TextMessage, payload) +} + +func mapCodexWebsocketWriteError(sess *codexWebsocketSession, conn *websocket.Conn, err error) error { + if err == nil || sess == nil || conn == nil { + return err + } + upstreamErr := sess.upstreamDisconnectError(conn) + var closeErr *websocket.CloseError + if !errors.As(upstreamErr, &closeErr) || closeErr.Code != websocket.CloseMessageTooBig { + return err + } + return mapCodexWebsocketReadError(upstreamErr) +} + +func shouldRetryCodexWebsocketSend(err error) bool { + if err == nil { + return false + } + var requestErr cliproxyexecutor.RequestScopedError + return !errors.As(err, &requestErr) || !requestErr.IsRequestScoped() +} + +type codexWebsocketMessageTooBigError struct { + statusErr +} + +func (codexWebsocketMessageTooBigError) IsRequestScoped() bool { + return true +} + +func mapCodexWebsocketReadError(err error) error { + if err == nil { + return nil + } + var closeErr *websocket.CloseError + if errors.As(err, &closeErr) && closeErr.Code == websocket.CloseMessageTooBig { + return codexWebsocketMessageTooBigError{statusErr: statusErr{ + code: http.StatusRequestEntityTooLarge, + msg: `{"error":{"message":"upstream websocket message too big","type":"invalid_request_error","code":"message_too_big"}}`, + }} + } + return err +} + +func normalizeCodexWebsocketParallelToolCalls(body []byte, headers http.Header) []byte { + if !isCodexResponsesLiteRequest(body, headers) { + return body + } + body = helps.SetBoolIfDifferent(body, "parallel_tool_calls", false) + return body +} + +func buildCodexWebsocketRequestBody(body []byte) []byte { + if len(body) == 0 { + return nil + } + + // Match codex-rs websocket v2 semantics: every request is `response.create`. + // Incremental follow-up turns continue on the same websocket using + // `previous_response_id` + incremental `input`, not `response.append`. + body = helps.SanitizeCodexInputItemIDs(body) + wsReqBody, errSet := sjson.SetBytes(bytes.Clone(body), "type", "response.create") + if errSet == nil && len(wsReqBody) > 0 { + return wsReqBody + } + fallback := bytes.Clone(body) + fallback, _ = sjson.SetBytes(fallback, "type", "response.create") + return fallback +} + +func readCodexWebsocketMessage(ctx context.Context, sess *codexWebsocketSession, conn *websocket.Conn, readCh chan codexWebsocketRead) (int, []byte, error) { + if sess == nil { + if conn == nil { + return 0, nil, fmt.Errorf("codex websockets executor: websocket conn is nil") + } + _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) + msgType, payload, errRead := conn.ReadMessage() + return msgType, payload, errRead + } + if conn == nil { + return 0, nil, fmt.Errorf("codex websockets executor: websocket conn is nil") + } + if readCh == nil { + return 0, nil, fmt.Errorf("codex websockets executor: session read channel is nil") + } + for { + select { + case <-ctx.Done(): + return 0, nil, ctx.Err() + case ev, ok := <-readCh: + if !ok { + return 0, nil, fmt.Errorf("codex websockets executor: session read channel closed") + } + if ev.conn != conn { + continue + } + if ev.err != nil { + return 0, nil, ev.err + } + return ev.msgType, ev.payload, nil + } + } +} + +func newProxyAwareWebsocketDialer(cfg *config.Config, auth *cliproxyauth.Auth) *websocket.Dialer { + dialer := &websocket.Dialer{ + Proxy: http.ProxyFromEnvironment, + HandshakeTimeout: codexResponsesWebsocketHandshakeTO, + EnableCompression: true, + NetDialContext: (&net.Dialer{ + Timeout: 30 * time.Second, + KeepAlive: 30 * time.Second, + }).DialContext, + } + + proxyURL := "" + if auth != nil { + proxyURL = strings.TrimSpace(auth.ProxyURL) + } + if proxyURL == "" && cfg != nil { + proxyURL = strings.TrimSpace(cfg.ProxyURL) + } + if proxyURL == "" { + return dialer + } + + setting, errParse := proxyutil.Parse(proxyURL) + if errParse != nil { + log.Errorf("codex websockets executor: %v", errParse) + return dialer + } + + switch setting.Mode { + case proxyutil.ModeDirect: + dialer.Proxy = nil + return dialer + case proxyutil.ModeProxy: + default: + return dialer + } + + switch setting.URL.Scheme { + case "socks5", "socks5h": + var proxyAuth *proxy.Auth + if setting.URL.User != nil { + username := setting.URL.User.Username() + password, _ := setting.URL.User.Password() + proxyAuth = &proxy.Auth{User: username, Password: password} + } + socksDialer, errSOCKS5 := proxy.SOCKS5("tcp", setting.URL.Host, proxyAuth, proxy.Direct) + if errSOCKS5 != nil { + log.Errorf("codex websockets executor: create SOCKS5 dialer failed: %v", errSOCKS5) + return dialer + } + dialer.Proxy = nil + dialer.NetDialContext = func(_ context.Context, network, addr string) (net.Conn, error) { + return socksDialer.Dial(network, addr) + } + case "http", "https": + dialer.Proxy = http.ProxyURL(setting.URL) + default: + log.Errorf("codex websockets executor: unsupported proxy scheme: %s", setting.URL.Scheme) + } + + return dialer +} + +func buildCodexResponsesWebsocketURL(httpURL string) (string, error) { + parsed, err := url.Parse(strings.TrimSpace(httpURL)) + if err != nil { + return "", err + } + switch strings.ToLower(parsed.Scheme) { + case "http": + parsed.Scheme = "ws" + case "https": + parsed.Scheme = "wss" + default: + return "", fmt.Errorf("codex websockets executor: unsupported responses websocket URL scheme %q", parsed.Scheme) + } + if strings.TrimSpace(parsed.Host) == "" { + return "", fmt.Errorf("codex websockets executor: responses websocket URL host is empty") + } + return parsed.String(), nil +} diff --git a/internal/runtime/executor/codex_websockets_errors.go b/internal/runtime/executor/codex_websockets_errors.go new file mode 100644 index 000000000..eae0706a3 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_errors.go @@ -0,0 +1,199 @@ +package executor + +import ( + "context" + "io" + "net/http" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +type statusErrWithHeaders struct { + statusErr + headers http.Header +} + +func (e statusErrWithHeaders) Headers() http.Header { + if e.headers == nil { + return nil + } + return e.headers.Clone() +} + +func parseCodexWebsocketError(payload []byte) (error, bool) { + if len(payload) == 0 { + return nil, false + } + if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) != "error" { + return nil, false + } + status := int(gjson.GetBytes(payload, "status").Int()) + if status == 0 { + status = int(gjson.GetBytes(payload, "status_code").Int()) + } + if status <= 0 { + return nil, false + } + + out := buildCodexWebsocketErrorPayload(payload, status) + headers := parseCodexWebsocketErrorHeaders(payload) + statusError := statusErr{code: status, msg: string(out)} + if retryAfter := parseCodexRetryAfter(status, out, time.Now()); retryAfter != nil { + statusError.retryAfter = retryAfter + } else if isCodexWebsocketConnectionLimitError(payload) { + retryAfter := time.Duration(0) + statusError.retryAfter = &retryAfter + } + return statusErrWithHeaders{ + statusErr: statusError, + headers: headers, + }, true +} + +func clearCodexReasoningReplayOnWebsocketError(ctx context.Context, scope codexReasoningReplayScope, payload []byte) error { + status := int(gjson.GetBytes(payload, "status").Int()) + if status == 0 { + status = int(gjson.GetBytes(payload, "status_code").Int()) + } + if status <= 0 { + return nil + } + return clearCodexReasoningReplayOnInvalidSignature(ctx, scope, status, buildCodexWebsocketErrorPayload(payload, status)) +} + +func buildCodexWebsocketErrorPayload(payload []byte, status int) []byte { + out := []byte(`{}`) + out, _ = sjson.SetBytes(out, "status", status) + + if bodyNode := gjson.GetBytes(payload, "body"); bodyNode.Exists() { + out, _ = sjson.SetRawBytes(out, "body", []byte(bodyNode.Raw)) + if bodyErrorNode := bodyNode.Get("error"); bodyErrorNode.Exists() { + out, _ = sjson.SetRawBytes(out, "error", []byte(bodyErrorNode.Raw)) + return out + } + } + + if errNode := gjson.GetBytes(payload, "error"); errNode.Exists() { + out, _ = sjson.SetRawBytes(out, "error", []byte(errNode.Raw)) + return out + } + + out, _ = sjson.SetBytes(out, "error.type", "server_error") + out, _ = sjson.SetBytes(out, "error.message", http.StatusText(status)) + return out +} + +func isCodexWebsocketConnectionLimitError(payload []byte) bool { + if len(payload) == 0 { + return false + } + for _, path := range []string{"error.code", "error.type", "body.error.code", "body.error.type", "code", "error"} { + if strings.TrimSpace(gjson.GetBytes(payload, path).String()) == "websocket_connection_limit_reached" { + return true + } + } + return false +} + +func parseCodexWebsocketErrorHeaders(payload []byte) http.Header { + headersNode := gjson.GetBytes(payload, "headers") + if !headersNode.Exists() || !headersNode.IsObject() { + return nil + } + mapped := make(http.Header) + headersNode.ForEach(func(key, value gjson.Result) bool { + name := strings.TrimSpace(key.String()) + if name == "" { + return true + } + switch value.Type { + case gjson.String: + if v := strings.TrimSpace(value.String()); v != "" { + mapped.Set(name, v) + } + case gjson.Number, gjson.True, gjson.False: + if v := strings.TrimSpace(value.Raw); v != "" { + mapped.Set(name, v) + } + default: + } + return true + }) + if len(mapped) == 0 { + return nil + } + return mapped +} + +func normalizeCodexWebsocketCompletion(payload []byte) []byte { + if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.done" { + updated, err := sjson.SetBytes(payload, "type", "response.completed") + if err == nil && len(updated) > 0 { + return updated + } + } + return payload +} + +func encodeCodexWebsocketAsSSE(payload []byte) []byte { + if len(payload) == 0 { + return nil + } + line := make([]byte, 0, len("data: ")+len(payload)) + line = append(line, []byte("data: ")...) + line = append(line, payload...) + return line +} + +func websocketUpgradeRequestLog(info helps.UpstreamRequestLog) helps.UpstreamRequestLog { + upgradeInfo := info + upgradeInfo.URL = helps.WebsocketUpgradeRequestURL(info.URL) + upgradeInfo.Method = http.MethodGet + upgradeInfo.Body = nil + upgradeInfo.Headers = info.Headers.Clone() + if upgradeInfo.Headers == nil { + upgradeInfo.Headers = make(http.Header) + } + if strings.TrimSpace(upgradeInfo.Headers.Get("Connection")) == "" { + upgradeInfo.Headers.Set("Connection", "Upgrade") + } + if strings.TrimSpace(upgradeInfo.Headers.Get("Upgrade")) == "" { + upgradeInfo.Headers.Set("Upgrade", "websocket") + } + return upgradeInfo +} + +func recordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, resp *http.Response) { + if resp == nil { + return + } + helps.RecordAPIWebsocketHandshake(ctx, cfg, resp.StatusCode, resp.Header.Clone()) + closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error") +} + +func websocketHandshakeBody(resp *http.Response) []byte { + if resp == nil || resp.Body == nil { + return nil + } + body, _ := io.ReadAll(resp.Body) + closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error") + if len(body) == 0 { + return nil + } + return body +} + +func closeHTTPResponseBody(resp *http.Response, logPrefix string) { + if resp == nil || resp.Body == nil { + return + } + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("%s: %v", logPrefix, errClose) + } +} diff --git a/internal/runtime/executor/codex_websockets_execute.go b/internal/runtime/executor/codex_websockets_execute.go new file mode 100644 index 000000000..43f86ad80 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_execute.go @@ -0,0 +1,323 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "net/http" + "strings" + + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + if ctx == nil { + ctx = context.Background() + } + if opts.Alt == "responses/compact" { + return e.CodexExecutor.executeCompact(ctx, auth, req, opts) + } + + baseModel := thinking.ParseSuffix(req.Model).ModelName + apiKey, baseURL := codexCreds(auth) + if baseURL == "" { + baseURL = "https://chatgpt.com/backend-api/codex" + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("codex") + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, false) + + body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return resp, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = helps.SetBoolIfDifferent(body, "stream", true) + body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") + body, _ = sjson.DeleteBytes(body, "safety_identifier") + body = normalizeCodexInstructions(body) + if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { + body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) + } + body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) + body = normalizeCodexWebsocketParallelToolCalls(body, opts.Headers) + body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return resp, errReplay + } + + httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" + wsURL, err := buildCodexResponsesWebsocketURL(httpURL) + if err != nil { + return resp, err + } + + body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body, opts.Headers) + if errPromptCache != nil { + return resp, errPromptCache + } + clientBody := body + var identityState codexIdentityConfuseState + upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) + reporter.SetTranslatedReasoningEffort(clientBody, to.String()) + wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) + applyModelHeaderOverrides(wsHeaders, baseModel) + applyCodexIdentityConfuseHeaders(wsHeaders, &identityState) + + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + + executionSessionID := executionSessionIDFromOptions(opts) + var sess *codexWebsocketSession + sessionLocked := false + unlockSession := func() { + if sess != nil && sessionLocked { + sess.reqMu.Unlock() + sessionLocked = false + } + } + if executionSessionID != "" { + sess = e.getOrCreateSession(executionSessionID) + sess.reqMu.Lock() + sessionLocked = true + defer unlockSession() + } + + wsReqBody := buildCodexWebsocketRequestBody(upstreamBody) + wsReqLog := helps.UpstreamRequestLog{ + URL: wsURL, + Method: "WEBSOCKET", + Headers: wsHeaders.Clone(), + Body: wsReqBody, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + } + helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) + + var conn *websocket.Conn + var closer *websocketConnectionCloser + var respHS *http.Response + var errDial error + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + conn, closer = existingWebsocketSessionConn(sess, authID, wsURL) + if conn == nil { + return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } + } else { + conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + } + if errDial != nil { + bodyErr := websocketHandshakeBody(respHS) + if respHS != nil { + helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) + } + if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { + if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) { + return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } + return e.CodexExecutor.Execute(ctx, auth, req, opts) + } + if respHS != nil && respHS.StatusCode > 0 { + return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) + return resp, errDial + } + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + unlockSession() + closeWebsocketAfterBindFailure(sess, conn, closer) + return resp, errBind + } + recordAPIWebsocketHandshake(ctx, e.cfg, respHS) + reporter.StartResponseTTFT() + if sess == nil { + logCodexWebsocketConnected(executionSessionID, authID, wsURL) + defer func() { + reason := "completed" + if err != nil { + reason = "error" + } + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, reason, err) + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + }() + } + + var readCh chan codexWebsocketRead + if sess != nil { + readCh = sess.activate(conn) + defer func() { + sess.clearActive(conn, readCh) + }() + } + + if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { + errSend = mapCodexWebsocketWriteError(sess, conn, errSend) + if sess != nil { + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend) + if !shouldRetryCodexWebsocketSend(errSend) { + helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) + return resp, errSend + } + return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } + e.invalidateUpstreamConn(sess, conn, "send_error", errSend) + if !shouldRetryCodexWebsocketSend(errSend) { + helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) + return resp, errSend + } + + // Retry once with a fresh websocket connection. This is mainly to handle + // upstream closing the socket between sequential requests within the same + // execution session. + connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + if errDialRetry == nil && connRetry != nil { + previousConn, previousReadCh := conn, readCh + conn = connRetry + closer = closerRetry + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + clearRetryActiveState(sess, previousConn, previousReadCh) + unlockSession() + closeWebsocketAfterBindFailure(sess, conn, closer) + return resp, errBind + } + readCh = sess.activate(conn) + wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) + helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: wsURL, + Method: "WEBSOCKET", + Headers: wsHeaders.Clone(), + Body: wsReqBodyRetry, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) + reporter.StartResponseTTFT() + if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry == nil { + wsReqBody = wsReqBodyRetry + } else { + errSendRetry = mapCodexWebsocketWriteError(sess, connRetry, errSendRetry) + e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) + helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) + return resp, errSendRetry + } + } else { + closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") + helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) + return resp, errDialRetry + } + } else { + helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) + return resp, errSend + } + } + + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + for { + if ctx != nil && ctx.Err() != nil { + return resp, ctx.Err() + } + msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh) + if errRead != nil { + mappedErr := mapCodexWebsocketReadError(errRead) + helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr) + return resp, mappedErr + } + if msgType != websocket.TextMessage { + if msgType == websocket.BinaryMessage { + err = fmt.Errorf("codex websockets executor: unexpected binary message") + if sess != nil { + e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err) + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err) + return resp, err + } + continue + } + + payload = bytes.TrimSpace(payload) + if len(payload) == 0 { + continue + } + reporter.MarkFirstResponseByte() + payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) + helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) + payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) + + if wsErr, ok := parseCodexWebsocketError(payload); ok { + if sess != nil { + e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) + } + if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { + return resp, errClearReplay + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) + return resp, wsErr + } + if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { + if sess != nil { + unlockSession() + e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) + } + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + return resp, errClearReplay + } + return resp, streamErr + } + + payload = normalizeCodexWebsocketCompletion(payload) + eventType := gjson.GetBytes(payload, "type").String() + switch eventType { + case "response.output_item.done": + collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) + case "response.completed": + payload = patchCodexCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) + cacheCodexReasoningReplayFromCompleted(replayScope, payload) + if detail, ok := helps.ParseCodexUsage(payload); ok { + reporter.Publish(ctx, detail) + } + var param any + clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) + out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, clientBody, clientPayload, ¶m) + resp = cliproxyexecutor.Response{Payload: out} + return resp, nil + } + } +} diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 31696bd58..84c40698a 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -3,41 +3,15 @@ package executor import ( - "bytes" "context" - "errors" "fmt" - "io" - "net" "net/http" - "net/url" "strconv" "strings" - "sync" - "time" - "github.com/gin-gonic/gin" - "github.com/google/uuid" - "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" - "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil" - sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" - log "github.com/sirupsen/logrus" - "github.com/tidwall/gjson" - "github.com/tidwall/sjson" - "golang.org/x/net/proxy" -) - -const ( - codexResponsesWebsocketBetaHeaderValue = "responses_websockets=2026-02-06" - codexResponsesWebsocketIdleTimeout = 5 * time.Minute - codexResponsesWebsocketHandshakeTO = 30 * time.Second ) // CodexWebsocketsExecutor executes Codex Responses requests using a WebSocket transport. @@ -50,69 +24,6 @@ type CodexWebsocketsExecutor struct { store *codexWebsocketSessionStore } -type codexWebsocketSessionStore struct { - mu sync.Mutex - sessions map[string]*codexWebsocketSession -} - -var globalCodexWebsocketSessionStore = &codexWebsocketSessionStore{ - sessions: make(map[string]*codexWebsocketSession), -} - -type websocketConnectionCloser struct { - conn *websocket.Conn - once sync.Once - err error -} - -func newWebsocketConnectionCloser(conn *websocket.Conn) *websocketConnectionCloser { - if conn == nil { - return nil - } - return &websocketConnectionCloser{conn: conn} -} - -func (c *websocketConnectionCloser) Close() error { - if c == nil || c.conn == nil { - return nil - } - c.once.Do(func() { - c.err = c.conn.Close() - }) - return c.err -} - -type codexWebsocketSession struct { - sessionID string - - reqMu sync.Mutex - - connMu sync.Mutex - conn *websocket.Conn - connCloser *websocketConnectionCloser - wsURL string - authID string - lifecycleBindMu sync.Mutex - lifecycle cliproxyexecutor.ExecutionLifecycle - lifecycleModel string - - writeMu sync.Mutex - - activeMu sync.Mutex - activeConn *websocket.Conn - activeCh chan codexWebsocketRead - activeDone <-chan struct{} - activeCancel context.CancelFunc - - readerConn *websocket.Conn - - upstreamDisconnectOnce sync.Once - upstreamDisconnectCh chan error - upstreamDisconnectErrMu sync.RWMutex - upstreamDisconnectErrConn *websocket.Conn - upstreamDisconnectErr error -} - func NewCodexWebsocketsExecutor(cfg *config.Config) *CodexWebsocketsExecutor { return &CodexWebsocketsExecutor{ CodexExecutor: NewCodexExecutor(cfg), @@ -120,2128 +31,6 @@ func NewCodexWebsocketsExecutor(cfg *config.Config) *CodexWebsocketsExecutor { } } -type codexWebsocketRead struct { - conn *websocket.Conn - msgType int - payload []byte - err error -} - -func (s *codexWebsocketSession) setActive(conn *websocket.Conn, ch chan codexWebsocketRead) { - if s == nil { - return - } - s.activeMu.Lock() - if s.activeCancel != nil { - s.activeCancel() - s.activeCancel = nil - s.activeDone = nil - } - s.activeConn = conn - s.activeCh = ch - if conn != nil && ch != nil { - activeCtx, activeCancel := context.WithCancel(context.Background()) - s.activeDone = activeCtx.Done() - s.activeCancel = activeCancel - } - s.activeMu.Unlock() -} - -func (s *codexWebsocketSession) activate(conn *websocket.Conn) chan codexWebsocketRead { - if s == nil || conn == nil { - return nil - } - ch := make(chan codexWebsocketRead, 4096) - s.setActive(conn, ch) - return ch -} - -func (s *codexWebsocketSession) activeForConn(conn *websocket.Conn) (chan codexWebsocketRead, <-chan struct{}) { - if s == nil || conn == nil { - return nil, nil - } - s.activeMu.Lock() - defer s.activeMu.Unlock() - if s.activeConn != conn { - return nil, nil - } - return s.activeCh, s.activeDone -} - -func clearRetryActiveState(sess *codexWebsocketSession, conn *websocket.Conn, ch chan codexWebsocketRead) bool { - if sess == nil { - return false - } - return sess.clearActive(conn, ch) -} - -func (s *codexWebsocketSession) clearActive(conn *websocket.Conn, ch chan codexWebsocketRead) bool { - if s == nil { - return false - } - s.activeMu.Lock() - defer s.activeMu.Unlock() - if s.activeConn != conn || s.activeCh != ch { - return false - } - s.activeConn = nil - s.activeCh = nil - if s.activeCancel != nil { - s.activeCancel() - } - s.activeCancel = nil - s.activeDone = nil - return true -} - -func (s *codexWebsocketSession) writeMessage(conn *websocket.Conn, msgType int, payload []byte) error { - if s == nil { - return fmt.Errorf("codex websockets executor: session is nil") - } - if conn == nil { - return fmt.Errorf("codex websockets executor: websocket conn is nil") - } - s.writeMu.Lock() - defer s.writeMu.Unlock() - return conn.WriteMessage(msgType, payload) -} - -// sendTerminalWebsocketRead reports whether it invalidated a full channel's connection before waiting. -func sendTerminalWebsocketRead(ch chan<- codexWebsocketRead, done <-chan struct{}, event codexWebsocketRead, invalidate func()) bool { - select { - case ch <- event: - return false - case <-done: - return false - default: - } - - invalidated := invalidate != nil - if invalidated { - invalidate() - } - select { - case ch <- event: - case <-done: - } - return invalidated -} - -func (s *codexWebsocketSession) configureConn(conn *websocket.Conn) { - if s == nil || conn == nil { - return - } - s.resetUpstreamDisconnectError(conn) - conn.SetPingHandler(func(appData string) error { - s.writeMu.Lock() - defer s.writeMu.Unlock() - // Reply pongs from the same write lock to avoid concurrent writes. - return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second)) - }) - defaultCloseHandler := conn.CloseHandler() - conn.SetCloseHandler(func(code int, text string) error { - s.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: code, Text: text}) - return defaultCloseHandler(code, text) - }) -} - -func (s *codexWebsocketSession) bindExecutionLifecycle(opts cliproxyexecutor.Options, conn *websocket.Conn, closer *websocketConnectionCloser, model string) error { - if closer == nil { - return fmt.Errorf("codex websockets executor: websocket connection closer is nil") - } - if s == nil { - return cliproxyexecutor.BindExecutionResource(opts, closer) - } - lifecycle := opts.ExecutionLifecycle - if lifecycle == nil || conn == nil { - return nil - } - - s.lifecycleBindMu.Lock() - defer s.lifecycleBindMu.Unlock() - - s.connMu.Lock() - if s.conn == conn && s.connCloser == nil { - s.connCloser = closer - } - alreadyBound := s.conn == conn && s.connCloser == closer && s.lifecycle == lifecycle - s.connMu.Unlock() - if alreadyBound { - return nil - } - - if errBind := lifecycle.Bind(func() error { - return s.closeBoundConnection(conn, closer, lifecycle) - }); errBind != nil { - return errBind - } - if retained, ok := lifecycle.(interface{ Retain() }); ok { - retained.Retain() - } - - s.connMu.Lock() - if s.conn != conn || s.connCloser != closer { - s.connMu.Unlock() - return fmt.Errorf("codex websockets executor: websocket connection closed during lifecycle bind") - } - previous := s.lifecycle - s.lifecycle = lifecycle - s.lifecycleModel = strings.TrimSpace(model) - s.connMu.Unlock() - if previous != nil && previous != lifecycle { - previous.End("target_replaced") - } - return nil -} - -func (s *codexWebsocketSession) closeBoundConnection(conn *websocket.Conn, closer *websocketConnectionCloser, lifecycle cliproxyexecutor.ExecutionLifecycle) error { - if s == nil || conn == nil { - return nil - } - s.detachConnection(conn, lifecycle) - errClose := closer.Close() - go lifecycle.End("connection_closed") - return errClose -} - -func (s *codexWebsocketSession) detachConnection(conn *websocket.Conn, lifecycle cliproxyexecutor.ExecutionLifecycle) *websocketConnectionCloser { - if s == nil || conn == nil { - return nil - } - s.connMu.Lock() - var closer *websocketConnectionCloser - matched := s.conn == conn - if matched { - closer = s.connCloser - s.conn = nil - s.connCloser = nil - if s.readerConn == conn { - s.readerConn = nil - } - } - if (lifecycle == nil && matched) || (lifecycle != nil && s.lifecycle == lifecycle) { - s.lifecycle = nil - s.lifecycleModel = "" - } - s.connMu.Unlock() - return closer -} - -func closeWebsocketAfterBindFailure(sess *codexWebsocketSession, conn *websocket.Conn, closer *websocketConnectionCloser) { - if conn == nil || closer == nil { - return - } - if sess != nil { - sess.detachConnection(conn, nil) - } - if errClose := closer.Close(); errClose != nil { - log.Errorf("websockets executor: close lifecycle bind failure connection error: %v", errClose) - } -} - -func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, wsURL string) bool { - if sess == nil { - return false - } - - sess.connMu.Lock() - defer sess.connMu.Unlock() - if strings.TrimSpace(sess.authID) == "" && strings.TrimSpace(sess.wsURL) == "" { - return false - } - return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) || strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL) -} - -func existingWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser) { - if sess == nil { - return nil, nil - } - sess.connMu.Lock() - conn := sess.conn - closer := sess.connCloser - matches := conn != nil && closer != nil && - strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && - strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL) - sess.connMu.Unlock() - if !matches || sess.upstreamDisconnectError(conn) != nil { - return nil, nil - } - return conn, closer -} - -func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) { - if sess == nil { - return nil, nil, "", "", nil - } - - sess.connMu.Lock() - defer sess.connMu.Unlock() - conn := sess.conn - if conn == nil || (strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)) { - return nil, nil, "", "", nil - } - - previousAuthID := sess.authID - previousWSURL := sess.wsURL - lifecycle := sess.lifecycle - closer := sess.connCloser - sess.lifecycle = nil - sess.lifecycleModel = "" - sess.conn = nil - sess.connCloser = nil - if sess.readerConn == conn { - sess.readerConn = nil - } - return conn, closer, previousAuthID, previousWSURL, lifecycle -} - -func (s *codexWebsocketSession) resetUpstreamDisconnectError(conn *websocket.Conn) { - if s == nil || conn == nil { - return - } - s.upstreamDisconnectErrMu.Lock() - s.upstreamDisconnectErrConn = conn - s.upstreamDisconnectErr = nil - s.upstreamDisconnectErrMu.Unlock() -} - -func (s *codexWebsocketSession) setUpstreamDisconnectError(conn *websocket.Conn, err error) { - if s == nil || conn == nil || err == nil { - return - } - s.upstreamDisconnectErrMu.Lock() - if s.upstreamDisconnectErrConn == conn && s.upstreamDisconnectErr == nil { - s.upstreamDisconnectErr = err - } - s.upstreamDisconnectErrMu.Unlock() -} - -func (s *codexWebsocketSession) upstreamDisconnectError(conn *websocket.Conn) error { - if s == nil || conn == nil { - return nil - } - s.upstreamDisconnectErrMu.RLock() - defer s.upstreamDisconnectErrMu.RUnlock() - if s.upstreamDisconnectErrConn != conn { - return nil - } - return s.upstreamDisconnectErr -} - -func (s *codexWebsocketSession) notifyUpstreamDisconnect(err error) { - if s == nil { - return - } - s.upstreamDisconnectOnce.Do(func() { - if s.upstreamDisconnectCh == nil { - return - } - select { - case s.upstreamDisconnectCh <- err: - default: - } - close(s.upstreamDisconnectCh) - }) -} - -func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - if ctx == nil { - ctx = context.Background() - } - if opts.Alt == "responses/compact" { - return e.CodexExecutor.executeCompact(ctx, auth, req, opts) - } - - baseModel := thinking.ParseSuffix(req.Model).ModelName - apiKey, baseURL := codexCreds(auth) - if baseURL == "" { - baseURL = "https://chatgpt.com/backend-api/codex" - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("codex") - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, false) - - body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return resp, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body = helps.SetStringIfDifferent(body, "model", baseModel) - body = helps.SetBoolIfDifferent(body, "stream", true) - body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") - body, _ = sjson.DeleteBytes(body, "safety_identifier") - body = normalizeCodexInstructions(body) - if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) - } - body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) - body = normalizeCodexWebsocketParallelToolCalls(body, opts.Headers) - body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) - body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) - if errReplay != nil { - return resp, errReplay - } - - httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" - wsURL, err := buildCodexResponsesWebsocketURL(httpURL) - if err != nil { - return resp, err - } - - body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body, opts.Headers) - if errPromptCache != nil { - return resp, errPromptCache - } - clientBody := body - var identityState codexIdentityConfuseState - upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) - reporter.SetTranslatedReasoningEffort(clientBody, to.String()) - wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) - applyModelHeaderOverrides(wsHeaders, baseModel) - applyCodexIdentityConfuseHeaders(wsHeaders, &identityState) - - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - - executionSessionID := executionSessionIDFromOptions(opts) - var sess *codexWebsocketSession - sessionLocked := false - unlockSession := func() { - if sess != nil && sessionLocked { - sess.reqMu.Unlock() - sessionLocked = false - } - } - if executionSessionID != "" { - sess = e.getOrCreateSession(executionSessionID) - sess.reqMu.Lock() - sessionLocked = true - defer unlockSession() - } - - wsReqBody := buildCodexWebsocketRequestBody(upstreamBody) - wsReqLog := helps.UpstreamRequestLog{ - URL: wsURL, - Method: "WEBSOCKET", - Headers: wsHeaders.Clone(), - Body: wsReqBody, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - } - helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) - - var conn *websocket.Conn - var closer *websocketConnectionCloser - var respHS *http.Response - var errDial error - if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { - conn, closer = existingWebsocketSessionConn(sess, authID, wsURL) - if conn == nil { - return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() - } - } else { - conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) - } - if errDial != nil { - bodyErr := websocketHandshakeBody(respHS) - if respHS != nil { - helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) - } - if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { - if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) { - return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} - } - return e.CodexExecutor.Execute(ctx, auth, req, opts) - } - if respHS != nil && respHS.StatusCode > 0 { - return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} - } - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) - return resp, errDial - } - if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { - unlockSession() - closeWebsocketAfterBindFailure(sess, conn, closer) - return resp, errBind - } - recordAPIWebsocketHandshake(ctx, e.cfg, respHS) - reporter.StartResponseTTFT() - if sess == nil { - logCodexWebsocketConnected(executionSessionID, authID, wsURL) - defer func() { - reason := "completed" - if err != nil { - reason = "error" - } - logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, reason, err) - if errClose := closer.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) - } - }() - } - - var readCh chan codexWebsocketRead - if sess != nil { - readCh = sess.activate(conn) - defer func() { - sess.clearActive(conn, readCh) - }() - } - - if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { - errSend = mapCodexWebsocketWriteError(sess, conn, errSend) - if sess != nil { - if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { - e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend) - if !shouldRetryCodexWebsocketSend(errSend) { - helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) - return resp, errSend - } - return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() - } - e.invalidateUpstreamConn(sess, conn, "send_error", errSend) - if !shouldRetryCodexWebsocketSend(errSend) { - helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) - return resp, errSend - } - - // Retry once with a fresh websocket connection. This is mainly to handle - // upstream closing the socket between sequential requests within the same - // execution session. - connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) - if errDialRetry == nil && connRetry != nil { - previousConn, previousReadCh := conn, readCh - conn = connRetry - closer = closerRetry - if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { - clearRetryActiveState(sess, previousConn, previousReadCh) - unlockSession() - closeWebsocketAfterBindFailure(sess, conn, closer) - return resp, errBind - } - readCh = sess.activate(conn) - wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) - helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: wsURL, - Method: "WEBSOCKET", - Headers: wsHeaders.Clone(), - Body: wsReqBodyRetry, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) - reporter.StartResponseTTFT() - if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry == nil { - wsReqBody = wsReqBodyRetry - } else { - errSendRetry = mapCodexWebsocketWriteError(sess, connRetry, errSendRetry) - e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) - helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) - return resp, errSendRetry - } - } else { - closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) - return resp, errDialRetry - } - } else { - helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) - return resp, errSend - } - } - - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte - for { - if ctx != nil && ctx.Err() != nil { - return resp, ctx.Err() - } - msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh) - if errRead != nil { - mappedErr := mapCodexWebsocketReadError(errRead) - helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr) - return resp, mappedErr - } - if msgType != websocket.TextMessage { - if msgType == websocket.BinaryMessage { - err = fmt.Errorf("codex websockets executor: unexpected binary message") - if sess != nil { - e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err) - } - helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err) - return resp, err - } - continue - } - - payload = bytes.TrimSpace(payload) - if len(payload) == 0 { - continue - } - reporter.MarkFirstResponseByte() - payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) - helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) - payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) - - if wsErr, ok := parseCodexWebsocketError(payload); ok { - if sess != nil { - e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) - } - if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { - return resp, errClearReplay - } - helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) - return resp, wsErr - } - if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { - if sess != nil { - unlockSession() - e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) - } - if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { - return resp, errClearReplay - } - return resp, streamErr - } - - payload = normalizeCodexWebsocketCompletion(payload) - eventType := gjson.GetBytes(payload, "type").String() - switch eventType { - case "response.output_item.done": - collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) - case "response.completed": - payload = patchCodexCompletedOutput(payload, outputItemsByIndex, outputItemsFallback) - cacheCodexReasoningReplayFromCompleted(replayScope, payload) - if detail, ok := helps.ParseCodexUsage(payload); ok { - reporter.Publish(ctx, detail) - } - var param any - clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) - out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, clientBody, clientPayload, ¶m) - resp = cliproxyexecutor.Response{Payload: out} - return resp, nil - } - } -} - -func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { - log.Debugf("Executing Codex Websockets stream request with auth ID: %s, model: %s", auth.ID, req.Model) - if ctx == nil { - ctx = context.Background() - } - if opts.Alt == "responses/compact" { - return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} - } - - baseModel := thinking.ParseSuffix(req.Model).ModelName - apiKey, baseURL := codexCreds(auth) - if baseURL == "" { - baseURL = "https://chatgpt.com/backend-api/codex" - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - to := sdktranslator.FromString("codex") - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := originalPayloadSource - originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, true) - - body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) - if err != nil { - return nil, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body = helps.SetStringIfDifferent(body, "model", baseModel) - body = normalizeCodexInstructions(body) - if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { - body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) - } - body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) - body = normalizeCodexWebsocketParallelToolCalls(body, opts.Headers) - body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) - body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) - if errReplay != nil { - return nil, errReplay - } - - httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" - wsURL, err := buildCodexResponsesWebsocketURL(httpURL) - if err != nil { - return nil, err - } - - body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body, opts.Headers) - if errPromptCache != nil { - return nil, errPromptCache - } - clientBody := body - var identityState codexIdentityConfuseState - upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) - reporter.SetTranslatedReasoningEffort(clientBody, to.String()) - wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) - applyModelHeaderOverrides(wsHeaders, baseModel) - applyCodexIdentityConfuseHeaders(wsHeaders, &identityState) - - var authID, authLabel, authType, authValue string - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - - executionSessionID := executionSessionIDFromOptions(opts) - var sess *codexWebsocketSession - if executionSessionID != "" { - sess = e.getOrCreateSession(executionSessionID) - if sess != nil { - sess.reqMu.Lock() - } - } - streamSessionLocked := sess != nil - unlockStreamSession := func() { - if sess != nil && streamSessionLocked { - sess.reqMu.Unlock() - streamSessionLocked = false - } - } - - wsReqBody := buildCodexWebsocketRequestBody(upstreamBody) - wsReqLog := helps.UpstreamRequestLog{ - URL: wsURL, - Method: "WEBSOCKET", - Headers: wsHeaders.Clone(), - Body: wsReqBody, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - } - helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) - - var conn *websocket.Conn - var closer *websocketConnectionCloser - var respHS *http.Response - var errDial error - if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { - conn, closer = existingWebsocketSessionConn(sess, authID, wsURL) - if conn == nil { - if sess != nil { - sess.reqMu.Unlock() - } - return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() - } - } else { - conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) - } - var upstreamHeaders http.Header - if respHS != nil { - upstreamHeaders = respHS.Header.Clone() - } - if errDial != nil { - bodyErr := websocketHandshakeBody(respHS) - if respHS != nil { - helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) - } - if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { - if sess != nil { - sess.reqMu.Unlock() - } - if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) { - return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} - } - return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts) - } - if respHS != nil && respHS.StatusCode > 0 { - if sess != nil { - sess.reqMu.Unlock() - } - return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} - } - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) - if sess != nil { - sess.reqMu.Unlock() - } - return nil, errDial - } - if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { - if sess != nil { - sess.reqMu.Unlock() - } - closeWebsocketAfterBindFailure(sess, conn, closer) - return nil, errBind - } - recordAPIWebsocketHandshake(ctx, e.cfg, respHS) - reporter.StartResponseTTFT() - - if sess == nil { - logCodexWebsocketConnected(executionSessionID, authID, wsURL) - } - - var readCh chan codexWebsocketRead - if sess != nil { - readCh = sess.activate(conn) - } - - if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { - errSend = mapCodexWebsocketWriteError(sess, conn, errSend) - helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) - if sess != nil { - if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { - e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend) - sess.clearActive(conn, readCh) - sess.reqMu.Unlock() - if !shouldRetryCodexWebsocketSend(errSend) { - return nil, errSend - } - return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() - } - e.invalidateUpstreamConn(sess, conn, "send_error", errSend) - if !shouldRetryCodexWebsocketSend(errSend) { - sess.clearActive(conn, readCh) - sess.reqMu.Unlock() - return nil, errSend - } - - // Retry once with a new websocket connection for the same execution session. - connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) - if errDialRetry != nil || connRetry == nil { - closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") - helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) - sess.clearActive(conn, readCh) - sess.reqMu.Unlock() - return nil, errDialRetry - } - previousConn, previousReadCh := conn, readCh - conn = connRetry - closer = closerRetry - if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { - clearRetryActiveState(sess, previousConn, previousReadCh) - sess.reqMu.Unlock() - closeWebsocketAfterBindFailure(sess, conn, closer) - return nil, errBind - } - readCh = sess.activate(conn) - wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) - helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: wsURL, - Method: "WEBSOCKET", - Headers: wsHeaders.Clone(), - Body: wsReqBodyRetry, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) - recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) - reporter.StartResponseTTFT() - if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { - errSendRetry = mapCodexWebsocketWriteError(sess, conn, errSendRetry) - helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) - e.invalidateUpstreamConn(sess, conn, "send_error", errSendRetry) - sess.clearActive(conn, readCh) - sess.reqMu.Unlock() - return nil, errSendRetry - } - wsReqBody = wsReqBodyRetry - } else { - logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) - if errClose := closer.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) - } - return nil, errSend - } - } - - out := make(chan cliproxyexecutor.StreamChunk) - go func() { - terminateReason := "completed" - var terminateErr error - - defer close(out) - defer func() { - if sess != nil { - sess.clearActive(conn, readCh) - unlockStreamSession() - return - } - logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, terminateReason, terminateErr) - if errClose := closer.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) - } - }() - - send := func(chunk cliproxyexecutor.StreamChunk) bool { - if ctx == nil { - out <- chunk - return true - } - select { - case out <- chunk: - return true - case <-ctx.Done(): - return false - } - } - - claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) - var param any - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte - for { - if ctx != nil && ctx.Err() != nil { - terminateReason = "context_done" - terminateErr = ctx.Err() - _ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()}) - return - } - msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh) - if errRead != nil { - if sess != nil && ctx != nil && ctx.Err() != nil { - terminateReason = "context_done" - terminateErr = ctx.Err() - _ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()}) - return - } - mappedErr := mapCodexWebsocketReadError(errRead) - terminateReason = "read_error" - terminateErr = mappedErr - helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr) - reporter.PublishFailure(ctx, mappedErr) - _ = send(cliproxyexecutor.StreamChunk{Err: mappedErr}) - return - } - if msgType != websocket.TextMessage { - if msgType == websocket.BinaryMessage { - err = fmt.Errorf("codex websockets executor: unexpected binary message") - terminateReason = "unexpected_binary" - terminateErr = err - helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err) - reporter.PublishFailure(ctx, err) - if sess != nil { - e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err) - } - _ = send(cliproxyexecutor.StreamChunk{Err: err}) - return - } - continue - } - - payload = bytes.TrimSpace(payload) - if len(payload) == 0 { - continue - } - reporter.MarkFirstResponseByte() - payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) - helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) - payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) - - if wsErr, ok := parseCodexWebsocketError(payload); ok { - terminateReason = "upstream_error" - terminateErr = wsErr - if sess != nil { - e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) - } - if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { - terminateErr = errClearReplay - helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) - reporter.PublishFailure(ctx, errClearReplay) - _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) - return - } - helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) - reporter.PublishFailure(ctx, wsErr) - _ = send(cliproxyexecutor.StreamChunk{Err: wsErr}) - return - } - if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { - terminateReason = "upstream_error" - terminateErr = streamErr - if sess != nil { - unlockStreamSession() - e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) - } - if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { - terminateErr = errClearReplay - helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) - reporter.PublishFailure(ctx, errClearReplay) - _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) - return - } - helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr) - reporter.PublishFailure(ctx, streamErr) - _ = send(cliproxyexecutor.StreamChunk{Err: streamErr}) - return - } - - eventType := gjson.GetBytes(payload, "type").String() - isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" - if eventType == "response.output_item.done" { - collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) - } - completedPayload := payload - if eventType == "response.completed" || eventType == "response.done" { - completedPayload = normalizeCodexWebsocketCompletion(completedPayload) - completedPayload = patchCodexCompletedOutput(completedPayload, outputItemsByIndex, outputItemsFallback) - cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload) - if detail, ok := helps.ParseCodexUsage(completedPayload); ok { - reporter.Publish(ctx, detail) - } - } - - clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) - if cliproxyexecutor.DownstreamWebsocket(ctx) { - if !send(cliproxyexecutor.StreamChunk{Payload: clientPayload}) { - terminateReason = "context_done" - terminateErr = ctx.Err() - return - } - if isTerminalEvent { - return - } - continue - } - - payload = normalizeCodexWebsocketCompletion(payload) - if eventType == "response.completed" || eventType == "response.done" { - payload = completedPayload - } - eventType = gjson.GetBytes(payload, "type").String() - clientPayload = applyCodexIdentityExposeResponsePayload(payload, identityState) - line := encodeCodexWebsocketAsSSE(clientPayload) - chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m, claudeInputTokens) - for i := range chunks { - if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { - terminateReason = "context_done" - terminateErr = ctx.Err() - return - } - } - if eventType == "response.completed" || eventType == "response.done" { - return - } - } - }() - - return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil -} - -func (e *CodexWebsocketsExecutor) dialCodexWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { - dialer := newProxyAwareWebsocketDialer(e.cfg, auth) - dialer.HandshakeTimeout = codexResponsesWebsocketHandshakeTO - dialer.EnableCompression = true - if ctx == nil { - ctx = context.Background() - } - conn, resp, err := dialer.DialContext(ctx, wsURL, headers) - closer := newWebsocketConnectionCloser(conn) - if conn != nil { - // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. - // Negotiating permessage-deflate is fine; we just don't compress outbound messages. - conn.EnableWriteCompression(false) - } - return conn, closer, resp, err -} - -func writeCodexWebsocketMessage(sess *codexWebsocketSession, conn *websocket.Conn, payload []byte) error { - if sess != nil { - return sess.writeMessage(conn, websocket.TextMessage, payload) - } - if conn == nil { - return fmt.Errorf("codex websockets executor: websocket conn is nil") - } - return conn.WriteMessage(websocket.TextMessage, payload) -} - -func mapCodexWebsocketWriteError(sess *codexWebsocketSession, conn *websocket.Conn, err error) error { - if err == nil || sess == nil || conn == nil { - return err - } - upstreamErr := sess.upstreamDisconnectError(conn) - var closeErr *websocket.CloseError - if !errors.As(upstreamErr, &closeErr) || closeErr.Code != websocket.CloseMessageTooBig { - return err - } - return mapCodexWebsocketReadError(upstreamErr) -} - -func shouldRetryCodexWebsocketSend(err error) bool { - if err == nil { - return false - } - var requestErr cliproxyexecutor.RequestScopedError - return !errors.As(err, &requestErr) || !requestErr.IsRequestScoped() -} - -type codexWebsocketMessageTooBigError struct { - statusErr -} - -func (codexWebsocketMessageTooBigError) IsRequestScoped() bool { - return true -} - -func mapCodexWebsocketReadError(err error) error { - if err == nil { - return nil - } - var closeErr *websocket.CloseError - if errors.As(err, &closeErr) && closeErr.Code == websocket.CloseMessageTooBig { - return codexWebsocketMessageTooBigError{statusErr: statusErr{ - code: http.StatusRequestEntityTooLarge, - msg: `{"error":{"message":"upstream websocket message too big","type":"invalid_request_error","code":"message_too_big"}}`, - }} - } - return err -} - -func normalizeCodexWebsocketParallelToolCalls(body []byte, headers http.Header) []byte { - if !isCodexResponsesLiteRequest(body, headers) { - return body - } - body = helps.SetBoolIfDifferent(body, "parallel_tool_calls", false) - return body -} - -func buildCodexWebsocketRequestBody(body []byte) []byte { - if len(body) == 0 { - return nil - } - - // Match codex-rs websocket v2 semantics: every request is `response.create`. - // Incremental follow-up turns continue on the same websocket using - // `previous_response_id` + incremental `input`, not `response.append`. - body = helps.SanitizeCodexInputItemIDs(body) - wsReqBody, errSet := sjson.SetBytes(bytes.Clone(body), "type", "response.create") - if errSet == nil && len(wsReqBody) > 0 { - return wsReqBody - } - fallback := bytes.Clone(body) - fallback, _ = sjson.SetBytes(fallback, "type", "response.create") - return fallback -} - -func readCodexWebsocketMessage(ctx context.Context, sess *codexWebsocketSession, conn *websocket.Conn, readCh chan codexWebsocketRead) (int, []byte, error) { - if sess == nil { - if conn == nil { - return 0, nil, fmt.Errorf("codex websockets executor: websocket conn is nil") - } - _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) - msgType, payload, errRead := conn.ReadMessage() - return msgType, payload, errRead - } - if conn == nil { - return 0, nil, fmt.Errorf("codex websockets executor: websocket conn is nil") - } - if readCh == nil { - return 0, nil, fmt.Errorf("codex websockets executor: session read channel is nil") - } - for { - select { - case <-ctx.Done(): - return 0, nil, ctx.Err() - case ev, ok := <-readCh: - if !ok { - return 0, nil, fmt.Errorf("codex websockets executor: session read channel closed") - } - if ev.conn != conn { - continue - } - if ev.err != nil { - return 0, nil, ev.err - } - return ev.msgType, ev.payload, nil - } - } -} - -func newProxyAwareWebsocketDialer(cfg *config.Config, auth *cliproxyauth.Auth) *websocket.Dialer { - dialer := &websocket.Dialer{ - Proxy: http.ProxyFromEnvironment, - HandshakeTimeout: codexResponsesWebsocketHandshakeTO, - EnableCompression: true, - NetDialContext: (&net.Dialer{ - Timeout: 30 * time.Second, - KeepAlive: 30 * time.Second, - }).DialContext, - } - - proxyURL := "" - if auth != nil { - proxyURL = strings.TrimSpace(auth.ProxyURL) - } - if proxyURL == "" && cfg != nil { - proxyURL = strings.TrimSpace(cfg.ProxyURL) - } - if proxyURL == "" { - return dialer - } - - setting, errParse := proxyutil.Parse(proxyURL) - if errParse != nil { - log.Errorf("codex websockets executor: %v", errParse) - return dialer - } - - switch setting.Mode { - case proxyutil.ModeDirect: - dialer.Proxy = nil - return dialer - case proxyutil.ModeProxy: - default: - return dialer - } - - switch setting.URL.Scheme { - case "socks5", "socks5h": - var proxyAuth *proxy.Auth - if setting.URL.User != nil { - username := setting.URL.User.Username() - password, _ := setting.URL.User.Password() - proxyAuth = &proxy.Auth{User: username, Password: password} - } - socksDialer, errSOCKS5 := proxy.SOCKS5("tcp", setting.URL.Host, proxyAuth, proxy.Direct) - if errSOCKS5 != nil { - log.Errorf("codex websockets executor: create SOCKS5 dialer failed: %v", errSOCKS5) - return dialer - } - dialer.Proxy = nil - dialer.NetDialContext = func(_ context.Context, network, addr string) (net.Conn, error) { - return socksDialer.Dial(network, addr) - } - case "http", "https": - dialer.Proxy = http.ProxyURL(setting.URL) - default: - log.Errorf("codex websockets executor: unsupported proxy scheme: %s", setting.URL.Scheme) - } - - return dialer -} - -func buildCodexResponsesWebsocketURL(httpURL string) (string, error) { - parsed, err := url.Parse(strings.TrimSpace(httpURL)) - if err != nil { - return "", err - } - switch strings.ToLower(parsed.Scheme) { - case "http": - parsed.Scheme = "ws" - case "https": - parsed.Scheme = "wss" - default: - return "", fmt.Errorf("codex websockets executor: unsupported responses websocket URL scheme %q", parsed.Scheme) - } - if strings.TrimSpace(parsed.Host) == "" { - return "", fmt.Errorf("codex websockets executor: responses websocket URL host is empty") - } - return parsed.String(), nil -} - -func applyCodexPromptCacheHeaders(from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte) ([]byte, http.Header) { - body, headers, _ := applyCodexPromptCacheHeadersWithContext(context.Background(), from, req, rawJSON) - return body, headers -} - -func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte, headerSets ...http.Header) ([]byte, http.Header, error) { - headers := http.Header{} - if len(rawJSON) == 0 { - return rawJSON, headers, nil - } - - var requestHeaders http.Header - if len(headerSets) > 0 { - requestHeaders = headerSets[0] - } - var cache helps.CodexCache - if sourceFormatEqual(from, sdktranslator.FormatClaude) { - modelName := strings.TrimSpace(gjson.GetBytes(rawJSON, "model").String()) - if modelName == "" { - modelName = thinking.ParseSuffix(req.Model).ModelName - } - cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, modelName, req.Payload, requestHeaders) - if errCache != nil { - return nil, nil, errCache - } - if ok { - cache = cached - } - } else if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) { - if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() { - cache.ID = promptCacheKey.String() - } - } - if cache.ID == "" { - cache.ID = helps.ProviderSessionUUID("codex", req.Metadata) - } - - if cache.ID != "" { - rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", cache.ID) - setHeaderCasePreserved(headers, "session_id", cache.ID) - headers.Set("Conversation_id", cache.ID) - } - - return rawJSON, headers, nil -} - -func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth *cliproxyauth.Auth, token string, cfg *config.Config) http.Header { - if headers == nil { - headers = http.Header{} - } - if strings.TrimSpace(token) != "" { - headers.Set("Authorization", "Bearer "+token) - } - - var ginHeaders http.Header - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - ginHeaders = ginCtx.Request.Header.Clone() - } - - isAPIKey := codexAuthUsesAPIKey(auth) - cfgUserAgent, cfgBetaFeatures := codexHeaderDefaults(cfg, auth) - ensureHeaderWithPriority(headers, ginHeaders, "x-codex-beta-features", cfgBetaFeatures, "") - misc.EnsureHeader(headers, ginHeaders, "x-codex-turn-state", "") - misc.EnsureHeader(headers, ginHeaders, "x-codex-turn-metadata", "") - misc.EnsureHeader(headers, ginHeaders, "x-client-request-id", "") - misc.EnsureHeader(headers, ginHeaders, "x-responsesapi-include-timing-metrics", "") - misc.EnsureHeader(headers, ginHeaders, "Version", "") - if isAPIKey { - ensureHeaderWithPriority(headers, ginHeaders, "User-Agent", "", "") - } else { - ensureHeaderWithConfigPrecedence(headers, ginHeaders, "User-Agent", cfgUserAgent, codexUserAgent) - } - - betaHeader := strings.TrimSpace(headers.Get("OpenAI-Beta")) - if betaHeader == "" && ginHeaders != nil { - betaHeader = strings.TrimSpace(ginHeaders.Get("OpenAI-Beta")) - } - if betaHeader == "" || !strings.Contains(betaHeader, "responses_websockets=") { - betaHeader = codexResponsesWebsocketBetaHeaderValue - } - headers.Set("OpenAI-Beta", betaHeader) - sessionFallback := "" - if strings.Contains(headers.Get("User-Agent"), "Mac OS") { - sessionFallback = uuid.NewString() - } - ensureCodexWebsocketSessionHeader(headers, ginHeaders, sessionFallback) - if originator := strings.TrimSpace(ginHeaders.Get("Originator")); originator != "" { - headers.Set("Originator", originator) - } else if !isAPIKey { - headers.Set("Originator", codexOriginator) - } - if !isAPIKey { - if auth != nil && auth.Metadata != nil { - if accountID, ok := auth.Metadata["account_id"].(string); ok { - if trimmed := strings.TrimSpace(accountID); trimmed != "" { - setHeaderCasePreserved(headers, "ChatGPT-Account-ID", trimmed) - } - } - } - } - - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs) - - return headers -} - -func ensureCodexWebsocketSessionHeader(target http.Header, source http.Header, fallbackValue string) { - if target == nil { - return - } - sessionID := codexSessionHeaderValue(target) - if sessionID == "" { - sessionID = codexSessionHeaderValue(source) - } - if sessionID == "" { - sessionID = strings.TrimSpace(fallbackValue) - } - if sessionID != "" { - setHeaderCasePreserved(target, "session_id", sessionID) - } - deleteHeaderCaseInsensitive(target, "Session-Id") -} - -func codexSessionHeaderValue(headers http.Header) string { - for _, key := range []string{"Session-Id", "Session_id", "session_id"} { - if value := strings.TrimSpace(headerValueCaseInsensitive(headers, key)); value != "" { - return value - } - } - return "" -} - -func codexAuthUsesAPIKey(auth *cliproxyauth.Auth) bool { - if auth == nil || auth.Attributes == nil { - return false - } - return strings.TrimSpace(auth.Attributes["api_key"]) != "" -} - -func ensureHeaderCasePreserved(target http.Header, source http.Header, key, configValue, fallbackValue string) { - if target == nil { - return - } - if strings.TrimSpace(headerValueCaseInsensitive(target, key)) != "" { - return - } - if source != nil { - if val := strings.TrimSpace(headerValueCaseInsensitive(source, key)); val != "" { - setHeaderCasePreserved(target, key, val) - return - } - } - if val := strings.TrimSpace(configValue); val != "" { - setHeaderCasePreserved(target, key, val) - return - } - if val := strings.TrimSpace(fallbackValue); val != "" { - setHeaderCasePreserved(target, key, val) - } -} - -func setHeaderCasePreserved(headers http.Header, key string, value string) { - if headers == nil { - return - } - key = strings.TrimSpace(key) - value = strings.TrimSpace(value) - if key == "" || value == "" { - return - } - deleteHeaderCaseInsensitive(headers, key) - headers[key] = []string{value} -} - -func setCodexSessionHeaderCasePreserved(headers http.Header, fallbackKey string, value string) { - if headers == nil { - return - } - fallbackKey = strings.TrimSpace(fallbackKey) - value = strings.TrimSpace(value) - if fallbackKey == "" || value == "" { - return - } - - selectedKey := "" - if _, ok := headers[fallbackKey]; ok && codexSessionHeaderKeyUsesUnderscore(fallbackKey) { - selectedKey = fallbackKey - } else { - for existingKey := range headers { - if codexSessionHeaderKeyUsesUnderscore(existingKey) { - selectedKey = existingKey - break - } - } - } - if selectedKey == "" { - selectedKey = fallbackKey - } - for existingKey := range headers { - if codexSessionHeaderKey(existingKey) && existingKey != selectedKey { - delete(headers, existingKey) - } - } - headers[selectedKey] = []string{value} -} - -func codexSessionHeaderKey(key string) bool { - normalized := strings.ToLower(strings.TrimSpace(key)) - return normalized == "session_id" || normalized == "session-id" -} - -func codexSessionHeaderKeyUsesUnderscore(key string) bool { - return strings.ToLower(strings.TrimSpace(key)) == "session_id" -} - -func headerValueCaseInsensitive(headers http.Header, key string) string { - key = strings.TrimSpace(key) - if headers == nil || key == "" { - return "" - } - if val := strings.TrimSpace(headers.Get(key)); val != "" { - return val - } - for existingKey, values := range headers { - if !strings.EqualFold(existingKey, key) { - continue - } - for _, value := range values { - if trimmed := strings.TrimSpace(value); trimmed != "" { - return trimmed - } - } - } - return "" -} - -func deleteHeaderCaseInsensitive(headers http.Header, key string) { - for existingKey := range headers { - if strings.EqualFold(existingKey, key) { - delete(headers, existingKey) - } - } -} - -func codexHeaderDefaults(cfg *config.Config, auth *cliproxyauth.Auth) (string, string) { - if cfg == nil || auth == nil { - return "", "" - } - if auth.Attributes != nil { - if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" { - return "", "" - } - } - return strings.TrimSpace(cfg.CodexHeaderDefaults.UserAgent), strings.TrimSpace(cfg.CodexHeaderDefaults.BetaFeatures) -} - -func ensureHeaderWithPriority(target http.Header, source http.Header, key, configValue, fallbackValue string) { - if target == nil { - return - } - if strings.TrimSpace(target.Get(key)) != "" { - return - } - if source != nil { - if val := strings.TrimSpace(source.Get(key)); val != "" { - target.Set(key, val) - return - } - } - if val := strings.TrimSpace(configValue); val != "" { - target.Set(key, val) - return - } - if val := strings.TrimSpace(fallbackValue); val != "" { - target.Set(key, val) - } -} - -func ensureHeaderWithConfigPrecedence(target http.Header, source http.Header, key, configValue, fallbackValue string) { - if target == nil { - return - } - if strings.TrimSpace(target.Get(key)) != "" { - return - } - if val := strings.TrimSpace(configValue); val != "" { - target.Set(key, val) - return - } - if source != nil { - if val := strings.TrimSpace(source.Get(key)); val != "" { - target.Set(key, val) - return - } - } - if val := strings.TrimSpace(fallbackValue); val != "" { - target.Set(key, val) - } -} - -type statusErrWithHeaders struct { - statusErr - headers http.Header -} - -func (e statusErrWithHeaders) Headers() http.Header { - if e.headers == nil { - return nil - } - return e.headers.Clone() -} - -func parseCodexWebsocketError(payload []byte) (error, bool) { - if len(payload) == 0 { - return nil, false - } - if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) != "error" { - return nil, false - } - status := int(gjson.GetBytes(payload, "status").Int()) - if status == 0 { - status = int(gjson.GetBytes(payload, "status_code").Int()) - } - if status <= 0 { - return nil, false - } - - out := buildCodexWebsocketErrorPayload(payload, status) - headers := parseCodexWebsocketErrorHeaders(payload) - statusError := statusErr{code: status, msg: string(out)} - if retryAfter := parseCodexRetryAfter(status, out, time.Now()); retryAfter != nil { - statusError.retryAfter = retryAfter - } else if isCodexWebsocketConnectionLimitError(payload) { - retryAfter := time.Duration(0) - statusError.retryAfter = &retryAfter - } - return statusErrWithHeaders{ - statusErr: statusError, - headers: headers, - }, true -} - -func clearCodexReasoningReplayOnWebsocketError(ctx context.Context, scope codexReasoningReplayScope, payload []byte) error { - status := int(gjson.GetBytes(payload, "status").Int()) - if status == 0 { - status = int(gjson.GetBytes(payload, "status_code").Int()) - } - if status <= 0 { - return nil - } - return clearCodexReasoningReplayOnInvalidSignature(ctx, scope, status, buildCodexWebsocketErrorPayload(payload, status)) -} - -func buildCodexWebsocketErrorPayload(payload []byte, status int) []byte { - out := []byte(`{}`) - out, _ = sjson.SetBytes(out, "status", status) - - if bodyNode := gjson.GetBytes(payload, "body"); bodyNode.Exists() { - out, _ = sjson.SetRawBytes(out, "body", []byte(bodyNode.Raw)) - if bodyErrorNode := bodyNode.Get("error"); bodyErrorNode.Exists() { - out, _ = sjson.SetRawBytes(out, "error", []byte(bodyErrorNode.Raw)) - return out - } - } - - if errNode := gjson.GetBytes(payload, "error"); errNode.Exists() { - out, _ = sjson.SetRawBytes(out, "error", []byte(errNode.Raw)) - return out - } - - out, _ = sjson.SetBytes(out, "error.type", "server_error") - out, _ = sjson.SetBytes(out, "error.message", http.StatusText(status)) - return out -} - -func isCodexWebsocketConnectionLimitError(payload []byte) bool { - if len(payload) == 0 { - return false - } - for _, path := range []string{"error.code", "error.type", "body.error.code", "body.error.type", "code", "error"} { - if strings.TrimSpace(gjson.GetBytes(payload, path).String()) == "websocket_connection_limit_reached" { - return true - } - } - return false -} - -func parseCodexWebsocketErrorHeaders(payload []byte) http.Header { - headersNode := gjson.GetBytes(payload, "headers") - if !headersNode.Exists() || !headersNode.IsObject() { - return nil - } - mapped := make(http.Header) - headersNode.ForEach(func(key, value gjson.Result) bool { - name := strings.TrimSpace(key.String()) - if name == "" { - return true - } - switch value.Type { - case gjson.String: - if v := strings.TrimSpace(value.String()); v != "" { - mapped.Set(name, v) - } - case gjson.Number, gjson.True, gjson.False: - if v := strings.TrimSpace(value.Raw); v != "" { - mapped.Set(name, v) - } - default: - } - return true - }) - if len(mapped) == 0 { - return nil - } - return mapped -} - -func normalizeCodexWebsocketCompletion(payload []byte) []byte { - if strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.done" { - updated, err := sjson.SetBytes(payload, "type", "response.completed") - if err == nil && len(updated) > 0 { - return updated - } - } - return payload -} - -func encodeCodexWebsocketAsSSE(payload []byte) []byte { - if len(payload) == 0 { - return nil - } - line := make([]byte, 0, len("data: ")+len(payload)) - line = append(line, []byte("data: ")...) - line = append(line, payload...) - return line -} - -func websocketUpgradeRequestLog(info helps.UpstreamRequestLog) helps.UpstreamRequestLog { - upgradeInfo := info - upgradeInfo.URL = helps.WebsocketUpgradeRequestURL(info.URL) - upgradeInfo.Method = http.MethodGet - upgradeInfo.Body = nil - upgradeInfo.Headers = info.Headers.Clone() - if upgradeInfo.Headers == nil { - upgradeInfo.Headers = make(http.Header) - } - if strings.TrimSpace(upgradeInfo.Headers.Get("Connection")) == "" { - upgradeInfo.Headers.Set("Connection", "Upgrade") - } - if strings.TrimSpace(upgradeInfo.Headers.Get("Upgrade")) == "" { - upgradeInfo.Headers.Set("Upgrade", "websocket") - } - return upgradeInfo -} - -func recordAPIWebsocketHandshake(ctx context.Context, cfg *config.Config, resp *http.Response) { - if resp == nil { - return - } - helps.RecordAPIWebsocketHandshake(ctx, cfg, resp.StatusCode, resp.Header.Clone()) - closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error") -} - -func websocketHandshakeBody(resp *http.Response) []byte { - if resp == nil || resp.Body == nil { - return nil - } - body, _ := io.ReadAll(resp.Body) - closeHTTPResponseBody(resp, "codex websockets executor: close handshake response body error") - if len(body) == 0 { - return nil - } - return body -} - -func closeHTTPResponseBody(resp *http.Response, logPrefix string) { - if resp == nil || resp.Body == nil { - return - } - if errClose := resp.Body.Close(); errClose != nil { - log.Errorf("%s: %v", logPrefix, errClose) - } -} - -func executionSessionIDFromOptions(opts cliproxyexecutor.Options) string { - if len(opts.Metadata) == 0 { - return "" - } - raw, ok := opts.Metadata[cliproxyexecutor.ExecutionSessionMetadataKey] - if !ok || raw == nil { - return "" - } - switch v := raw.(type) { - case string: - return strings.TrimSpace(v) - case []byte: - return strings.TrimSpace(string(v)) - default: - return "" - } -} - -func (e *CodexWebsocketsExecutor) getOrCreateSession(sessionID string) *codexWebsocketSession { - sessionID = strings.TrimSpace(sessionID) - if sessionID == "" { - return nil - } - if e == nil { - return nil - } - store := e.store - if store == nil { - store = globalCodexWebsocketSessionStore - } - store.mu.Lock() - defer store.mu.Unlock() - if store.sessions == nil { - store.sessions = make(map[string]*codexWebsocketSession) - } - if sess, ok := store.sessions[sessionID]; ok && sess != nil { - return sess - } - sess := &codexWebsocketSession{ - sessionID: sessionID, - upstreamDisconnectCh: make(chan error, 1), - } - store.sessions[sessionID] = sess - return sess -} - -func (e *CodexWebsocketsExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { - sess := e.getOrCreateSession(sessionID) - if sess == nil { - return nil - } - return sess.upstreamDisconnectCh -} - -func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { - if sess == nil { - return e.dialCodexWebsocket(ctx, auth, wsURL, headers) - } - - if staleConn, staleCloser, staleAuthID, staleWSURL, staleLifecycle := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { - logCodexWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) - if staleCloser != nil { - if errClose := staleCloser.Close(); errClose != nil { - log.Errorf("codex websockets executor: close stale websocket error: %v", errClose) - } - } - if staleLifecycle != nil { - staleLifecycle.End("target_changed") - } - } - - sess.connMu.Lock() - conn := sess.conn - closer := sess.connCloser - readerConn := sess.readerConn - sess.connMu.Unlock() - if conn != nil { - if readerConn != conn { - sess.connMu.Lock() - sess.readerConn = conn - sess.connMu.Unlock() - sess.configureConn(conn) - go e.readUpstreamLoop(sess, conn) - } - return conn, closer, nil, nil - } - - conn, closer, resp, errDial := e.dialCodexWebsocket(ctx, auth, wsURL, headers) - if errDial != nil { - return nil, closer, resp, errDial - } - - sess.connMu.Lock() - if sess.conn != nil { - previous := sess.conn - previousCloser := sess.connCloser - sess.connMu.Unlock() - if errClose := closer.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) - } - return previous, previousCloser, nil, nil - } - sess.conn = conn - sess.connCloser = closer - sess.wsURL = wsURL - sess.authID = authID - sess.readerConn = conn - sess.connMu.Unlock() - - sess.configureConn(conn) - go e.readUpstreamLoop(sess, conn) - logCodexWebsocketConnected(sess.sessionID, authID, wsURL) - return conn, closer, resp, nil -} - -func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, conn *websocket.Conn) { - if e == nil || sess == nil || conn == nil { - return - } - for { - _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) - msgType, payload, errRead := conn.ReadMessage() - if errRead != nil { - invalidate := func() { - e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) - } - invalidated := false - ch, done := sess.activeForConn(conn) - if ch != nil { - invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errRead}, invalidate) - if sess.clearActive(conn, ch) { - close(ch) - } - } - if !invalidated { - invalidate() - } - return - } - - if msgType != websocket.TextMessage { - if msgType == websocket.BinaryMessage { - errBinary := fmt.Errorf("codex websockets executor: unexpected binary message") - invalidate := func() { - e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) - } - invalidated := false - ch, done := sess.activeForConn(conn) - if ch != nil { - invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errBinary}, invalidate) - if sess.clearActive(conn, ch) { - close(ch) - } - } - if !invalidated { - invalidate() - } - return - } - continue - } - - ch, done := sess.activeForConn(conn) - if ch == nil { - continue - } - select { - case ch <- codexWebsocketRead{conn: conn, msgType: msgType, payload: payload}: - case <-done: - } - } -} - -func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { - e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, true) -} - -func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithoutDisconnectNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { - e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, false) -} - -func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error, notify bool) { - if sess == nil || conn == nil { - return - } - - sess.connMu.Lock() - current := sess.conn - authID := sess.authID - wsURL := sess.wsURL - sessionID := sess.sessionID - if current == nil || current != conn { - sess.connMu.Unlock() - return - } - lifecycle := sess.lifecycle - closer := sess.connCloser - sess.lifecycle = nil - sess.lifecycleModel = "" - sess.conn = nil - sess.connCloser = nil - if sess.readerConn == conn { - sess.readerConn = nil - } - sess.connMu.Unlock() - - logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err) - if notify { - sess.notifyUpstreamDisconnect(err) - } - if closer != nil { - if errClose := closer.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) - } - } - if lifecycle != nil { - lifecycle.End(reason) - } -} - -func (e *CodexWebsocketsExecutor) CloseExecutionSession(sessionID string) { - sessionID = strings.TrimSpace(sessionID) - if e == nil { - return - } - if sessionID == "" { - return - } - if sessionID == cliproxyauth.CloseAllExecutionSessionsID { - e.closeAllExecutionSessions("executor_shutdown") - return - } - - store := e.store - if store == nil { - store = globalCodexWebsocketSessionStore - } - store.mu.Lock() - sess := store.sessions[sessionID] - delete(store.sessions, sessionID) - store.mu.Unlock() - - e.closeExecutionSession(sess, "session_closed") -} - -func (e *CodexWebsocketsExecutor) closeAllExecutionSessions(reason string) { - if e == nil { - return - } - - store := e.store - if store == nil { - store = globalCodexWebsocketSessionStore - } - store.mu.Lock() - sessions := make([]*codexWebsocketSession, 0, len(store.sessions)) - for sessionID, sess := range store.sessions { - delete(store.sessions, sessionID) - if sess != nil { - sessions = append(sessions, sess) - } - } - store.mu.Unlock() - - for i := range sessions { - e.closeExecutionSession(sessions[i], reason) - } -} - -func (e *CodexWebsocketsExecutor) closeExecutionSession(sess *codexWebsocketSession, reason string) { - closeCodexWebsocketSession(sess, reason) -} - -func closeCodexWebsocketSession(sess *codexWebsocketSession, reason string) { - if sess == nil { - return - } - reason = strings.TrimSpace(reason) - if reason == "" { - reason = "session_closed" - } - - sess.connMu.Lock() - conn := sess.conn - authID := sess.authID - wsURL := sess.wsURL - lifecycle := sess.lifecycle - closer := sess.connCloser - sess.lifecycle = nil - sess.lifecycleModel = "" - sess.conn = nil - sess.connCloser = nil - if sess.readerConn == conn { - sess.readerConn = nil - } - sessionID := sess.sessionID - sess.connMu.Unlock() - - if conn != nil { - logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) - if closer != nil { - if errClose := closer.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) - } - } - } - if lifecycle != nil { - lifecycle.End(reason) - } -} - -func logCodexWebsocketConnected(sessionID string, authID string, wsURL string) { - log.Infof("codex websockets: upstream connected session=%s auth=%s url=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL)) -} - -func logCodexWebsocketDisconnected(sessionID string, authID string, wsURL string, reason string, err error) { - if err != nil { - log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s err=%v", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason), err) - return - } - log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason)) -} - -// CloseCodexWebsocketSessionsForAuthID closes all active Codex upstream websocket sessions -// associated with the supplied auth ID. -func CloseCodexWebsocketSessionsForAuthID(authID string, reason string) { - authID = strings.TrimSpace(authID) - if authID == "" { - return - } - reason = strings.TrimSpace(reason) - if reason == "" { - reason = "auth_removed" - } - - store := globalCodexWebsocketSessionStore - if store == nil { - return - } - - type sessionItem struct { - sessionID string - sess *codexWebsocketSession - } - - store.mu.Lock() - items := make([]sessionItem, 0, len(store.sessions)) - for sessionID, sess := range store.sessions { - items = append(items, sessionItem{sessionID: sessionID, sess: sess}) - } - store.mu.Unlock() - - matches := make([]sessionItem, 0) - for i := range items { - sess := items[i].sess - if sess == nil { - continue - } - sess.connMu.Lock() - sessAuthID := strings.TrimSpace(sess.authID) - sess.connMu.Unlock() - if sessAuthID == authID { - matches = append(matches, items[i]) - } - } - if len(matches) == 0 { - return - } - - toClose := make([]*codexWebsocketSession, 0, len(matches)) - store.mu.Lock() - for i := range matches { - current, ok := store.sessions[matches[i].sessionID] - if !ok || current == nil || current != matches[i].sess { - continue - } - delete(store.sessions, matches[i].sessionID) - toClose = append(toClose, current) - } - store.mu.Unlock() - - for i := range toClose { - closeCodexWebsocketSession(toClose[i], reason) - } -} - // CodexAutoExecutor routes Codex requests to the websocket transport only when: // 1. The downstream transport is websocket, and // 2. The selected auth enables websockets. diff --git a/internal/runtime/executor/codex_websockets_request.go b/internal/runtime/executor/codex_websockets_request.go new file mode 100644 index 000000000..fef258336 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_request.go @@ -0,0 +1,323 @@ +package executor + +import ( + "context" + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" +) + +func applyCodexPromptCacheHeaders(from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte) ([]byte, http.Header) { + body, headers, _ := applyCodexPromptCacheHeadersWithContext(context.Background(), from, req, rawJSON) + return body, headers +} + +func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktranslator.Format, req cliproxyexecutor.Request, rawJSON []byte, headerSets ...http.Header) ([]byte, http.Header, error) { + headers := http.Header{} + if len(rawJSON) == 0 { + return rawJSON, headers, nil + } + + var requestHeaders http.Header + if len(headerSets) > 0 { + requestHeaders = headerSets[0] + } + var cache helps.CodexCache + if sourceFormatEqual(from, sdktranslator.FormatClaude) { + modelName := strings.TrimSpace(gjson.GetBytes(rawJSON, "model").String()) + if modelName == "" { + modelName = thinking.ParseSuffix(req.Model).ModelName + } + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, modelName, req.Payload, requestHeaders) + if errCache != nil { + return nil, nil, errCache + } + if ok { + cache = cached + } + } else if sourceFormatEqual(from, sdktranslator.FormatOpenAIResponse) { + if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() { + cache.ID = promptCacheKey.String() + } + } + if cache.ID == "" { + cache.ID = helps.ProviderSessionUUID("codex", req.Metadata) + } + + if cache.ID != "" { + rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", cache.ID) + setHeaderCasePreserved(headers, "session_id", cache.ID) + headers.Set("Conversation_id", cache.ID) + } + + return rawJSON, headers, nil +} + +func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth *cliproxyauth.Auth, token string, cfg *config.Config) http.Header { + if headers == nil { + headers = http.Header{} + } + if strings.TrimSpace(token) != "" { + headers.Set("Authorization", "Bearer "+token) + } + + var ginHeaders http.Header + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + ginHeaders = ginCtx.Request.Header.Clone() + } + + isAPIKey := codexAuthUsesAPIKey(auth) + cfgUserAgent, cfgBetaFeatures := codexHeaderDefaults(cfg, auth) + ensureHeaderWithPriority(headers, ginHeaders, "x-codex-beta-features", cfgBetaFeatures, "") + misc.EnsureHeader(headers, ginHeaders, "x-codex-turn-state", "") + misc.EnsureHeader(headers, ginHeaders, "x-codex-turn-metadata", "") + misc.EnsureHeader(headers, ginHeaders, "x-client-request-id", "") + misc.EnsureHeader(headers, ginHeaders, "x-responsesapi-include-timing-metrics", "") + misc.EnsureHeader(headers, ginHeaders, "Version", "") + if isAPIKey { + ensureHeaderWithPriority(headers, ginHeaders, "User-Agent", "", "") + } else { + ensureHeaderWithConfigPrecedence(headers, ginHeaders, "User-Agent", cfgUserAgent, codexUserAgent) + } + + betaHeader := strings.TrimSpace(headers.Get("OpenAI-Beta")) + if betaHeader == "" && ginHeaders != nil { + betaHeader = strings.TrimSpace(ginHeaders.Get("OpenAI-Beta")) + } + if betaHeader == "" || !strings.Contains(betaHeader, "responses_websockets=") { + betaHeader = codexResponsesWebsocketBetaHeaderValue + } + headers.Set("OpenAI-Beta", betaHeader) + sessionFallback := "" + if strings.Contains(headers.Get("User-Agent"), "Mac OS") { + sessionFallback = uuid.NewString() + } + ensureCodexWebsocketSessionHeader(headers, ginHeaders, sessionFallback) + if originator := strings.TrimSpace(ginHeaders.Get("Originator")); originator != "" { + headers.Set("Originator", originator) + } else if !isAPIKey { + headers.Set("Originator", codexOriginator) + } + if !isAPIKey { + if auth != nil && auth.Metadata != nil { + if accountID, ok := auth.Metadata["account_id"].(string); ok { + if trimmed := strings.TrimSpace(accountID); trimmed != "" { + setHeaderCasePreserved(headers, "ChatGPT-Account-ID", trimmed) + } + } + } + } + + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(&http.Request{Header: headers}, attrs) + + return headers +} + +func ensureCodexWebsocketSessionHeader(target http.Header, source http.Header, fallbackValue string) { + if target == nil { + return + } + sessionID := codexSessionHeaderValue(target) + if sessionID == "" { + sessionID = codexSessionHeaderValue(source) + } + if sessionID == "" { + sessionID = strings.TrimSpace(fallbackValue) + } + if sessionID != "" { + setHeaderCasePreserved(target, "session_id", sessionID) + } + deleteHeaderCaseInsensitive(target, "Session-Id") +} + +func codexSessionHeaderValue(headers http.Header) string { + for _, key := range []string{"Session-Id", "Session_id", "session_id"} { + if value := strings.TrimSpace(headerValueCaseInsensitive(headers, key)); value != "" { + return value + } + } + return "" +} + +func codexAuthUsesAPIKey(auth *cliproxyauth.Auth) bool { + if auth == nil || auth.Attributes == nil { + return false + } + return strings.TrimSpace(auth.Attributes["api_key"]) != "" +} + +func ensureHeaderCasePreserved(target http.Header, source http.Header, key, configValue, fallbackValue string) { + if target == nil { + return + } + if strings.TrimSpace(headerValueCaseInsensitive(target, key)) != "" { + return + } + if source != nil { + if val := strings.TrimSpace(headerValueCaseInsensitive(source, key)); val != "" { + setHeaderCasePreserved(target, key, val) + return + } + } + if val := strings.TrimSpace(configValue); val != "" { + setHeaderCasePreserved(target, key, val) + return + } + if val := strings.TrimSpace(fallbackValue); val != "" { + setHeaderCasePreserved(target, key, val) + } +} + +func setHeaderCasePreserved(headers http.Header, key string, value string) { + if headers == nil { + return + } + key = strings.TrimSpace(key) + value = strings.TrimSpace(value) + if key == "" || value == "" { + return + } + deleteHeaderCaseInsensitive(headers, key) + headers[key] = []string{value} +} + +func setCodexSessionHeaderCasePreserved(headers http.Header, fallbackKey string, value string) { + if headers == nil { + return + } + fallbackKey = strings.TrimSpace(fallbackKey) + value = strings.TrimSpace(value) + if fallbackKey == "" || value == "" { + return + } + + selectedKey := "" + if _, ok := headers[fallbackKey]; ok && codexSessionHeaderKeyUsesUnderscore(fallbackKey) { + selectedKey = fallbackKey + } else { + for existingKey := range headers { + if codexSessionHeaderKeyUsesUnderscore(existingKey) { + selectedKey = existingKey + break + } + } + } + if selectedKey == "" { + selectedKey = fallbackKey + } + for existingKey := range headers { + if codexSessionHeaderKey(existingKey) && existingKey != selectedKey { + delete(headers, existingKey) + } + } + headers[selectedKey] = []string{value} +} + +func codexSessionHeaderKey(key string) bool { + normalized := strings.ToLower(strings.TrimSpace(key)) + return normalized == "session_id" || normalized == "session-id" +} + +func codexSessionHeaderKeyUsesUnderscore(key string) bool { + return strings.ToLower(strings.TrimSpace(key)) == "session_id" +} + +func headerValueCaseInsensitive(headers http.Header, key string) string { + key = strings.TrimSpace(key) + if headers == nil || key == "" { + return "" + } + if val := strings.TrimSpace(headers.Get(key)); val != "" { + return val + } + for existingKey, values := range headers { + if !strings.EqualFold(existingKey, key) { + continue + } + for _, value := range values { + if trimmed := strings.TrimSpace(value); trimmed != "" { + return trimmed + } + } + } + return "" +} + +func deleteHeaderCaseInsensitive(headers http.Header, key string) { + for existingKey := range headers { + if strings.EqualFold(existingKey, key) { + delete(headers, existingKey) + } + } +} + +func codexHeaderDefaults(cfg *config.Config, auth *cliproxyauth.Auth) (string, string) { + if cfg == nil || auth == nil { + return "", "" + } + if auth.Attributes != nil { + if v := strings.TrimSpace(auth.Attributes["api_key"]); v != "" { + return "", "" + } + } + return strings.TrimSpace(cfg.CodexHeaderDefaults.UserAgent), strings.TrimSpace(cfg.CodexHeaderDefaults.BetaFeatures) +} + +func ensureHeaderWithPriority(target http.Header, source http.Header, key, configValue, fallbackValue string) { + if target == nil { + return + } + if strings.TrimSpace(target.Get(key)) != "" { + return + } + if source != nil { + if val := strings.TrimSpace(source.Get(key)); val != "" { + target.Set(key, val) + return + } + } + if val := strings.TrimSpace(configValue); val != "" { + target.Set(key, val) + return + } + if val := strings.TrimSpace(fallbackValue); val != "" { + target.Set(key, val) + } +} + +func ensureHeaderWithConfigPrecedence(target http.Header, source http.Header, key, configValue, fallbackValue string) { + if target == nil { + return + } + if strings.TrimSpace(target.Get(key)) != "" { + return + } + if val := strings.TrimSpace(configValue); val != "" { + target.Set(key, val) + return + } + if source != nil { + if val := strings.TrimSpace(source.Get(key)); val != "" { + target.Set(key, val) + return + } + } + if val := strings.TrimSpace(fallbackValue); val != "" { + target.Set(key, val) + } +} diff --git a/internal/runtime/executor/codex_websockets_session.go b/internal/runtime/executor/codex_websockets_session.go new file mode 100644 index 000000000..76cb29ac8 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_session.go @@ -0,0 +1,788 @@ +package executor + +import ( + "context" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/gorilla/websocket" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + log "github.com/sirupsen/logrus" +) + +type codexWebsocketSessionStore struct { + mu sync.Mutex + sessions map[string]*codexWebsocketSession +} + +var globalCodexWebsocketSessionStore = &codexWebsocketSessionStore{ + sessions: make(map[string]*codexWebsocketSession), +} + +type websocketConnectionCloser struct { + conn *websocket.Conn + once sync.Once + err error +} + +func newWebsocketConnectionCloser(conn *websocket.Conn) *websocketConnectionCloser { + if conn == nil { + return nil + } + return &websocketConnectionCloser{conn: conn} +} + +func (c *websocketConnectionCloser) Close() error { + if c == nil || c.conn == nil { + return nil + } + c.once.Do(func() { + c.err = c.conn.Close() + }) + return c.err +} + +type codexWebsocketSession struct { + sessionID string + + reqMu sync.Mutex + + connMu sync.Mutex + conn *websocket.Conn + connCloser *websocketConnectionCloser + wsURL string + authID string + lifecycleBindMu sync.Mutex + lifecycle cliproxyexecutor.ExecutionLifecycle + lifecycleModel string + + writeMu sync.Mutex + + activeMu sync.Mutex + activeConn *websocket.Conn + activeCh chan codexWebsocketRead + activeDone <-chan struct{} + activeCancel context.CancelFunc + + readerConn *websocket.Conn + + upstreamDisconnectOnce sync.Once + upstreamDisconnectCh chan error + upstreamDisconnectErrMu sync.RWMutex + upstreamDisconnectErrConn *websocket.Conn + upstreamDisconnectErr error +} + +type codexWebsocketRead struct { + conn *websocket.Conn + msgType int + payload []byte + err error +} + +func (s *codexWebsocketSession) setActive(conn *websocket.Conn, ch chan codexWebsocketRead) { + if s == nil { + return + } + s.activeMu.Lock() + if s.activeCancel != nil { + s.activeCancel() + s.activeCancel = nil + s.activeDone = nil + } + s.activeConn = conn + s.activeCh = ch + if conn != nil && ch != nil { + activeCtx, activeCancel := context.WithCancel(context.Background()) + s.activeDone = activeCtx.Done() + s.activeCancel = activeCancel + } + s.activeMu.Unlock() +} + +func (s *codexWebsocketSession) activate(conn *websocket.Conn) chan codexWebsocketRead { + if s == nil || conn == nil { + return nil + } + ch := make(chan codexWebsocketRead, 4096) + s.setActive(conn, ch) + return ch +} + +func (s *codexWebsocketSession) activeForConn(conn *websocket.Conn) (chan codexWebsocketRead, <-chan struct{}) { + if s == nil || conn == nil { + return nil, nil + } + s.activeMu.Lock() + defer s.activeMu.Unlock() + if s.activeConn != conn { + return nil, nil + } + return s.activeCh, s.activeDone +} + +func clearRetryActiveState(sess *codexWebsocketSession, conn *websocket.Conn, ch chan codexWebsocketRead) bool { + if sess == nil { + return false + } + return sess.clearActive(conn, ch) +} + +func (s *codexWebsocketSession) clearActive(conn *websocket.Conn, ch chan codexWebsocketRead) bool { + if s == nil { + return false + } + s.activeMu.Lock() + defer s.activeMu.Unlock() + if s.activeConn != conn || s.activeCh != ch { + return false + } + s.activeConn = nil + s.activeCh = nil + if s.activeCancel != nil { + s.activeCancel() + } + s.activeCancel = nil + s.activeDone = nil + return true +} + +func (s *codexWebsocketSession) writeMessage(conn *websocket.Conn, msgType int, payload []byte) error { + if s == nil { + return fmt.Errorf("codex websockets executor: session is nil") + } + if conn == nil { + return fmt.Errorf("codex websockets executor: websocket conn is nil") + } + s.writeMu.Lock() + defer s.writeMu.Unlock() + return conn.WriteMessage(msgType, payload) +} + +// sendTerminalWebsocketRead reports whether it invalidated a full channel's connection before waiting. +func sendTerminalWebsocketRead(ch chan<- codexWebsocketRead, done <-chan struct{}, event codexWebsocketRead, invalidate func()) bool { + select { + case ch <- event: + return false + case <-done: + return false + default: + } + + invalidated := invalidate != nil + if invalidated { + invalidate() + } + select { + case ch <- event: + case <-done: + } + return invalidated +} + +func (s *codexWebsocketSession) configureConn(conn *websocket.Conn) { + if s == nil || conn == nil { + return + } + s.resetUpstreamDisconnectError(conn) + conn.SetPingHandler(func(appData string) error { + s.writeMu.Lock() + defer s.writeMu.Unlock() + // Reply pongs from the same write lock to avoid concurrent writes. + return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second)) + }) + defaultCloseHandler := conn.CloseHandler() + conn.SetCloseHandler(func(code int, text string) error { + s.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: code, Text: text}) + return defaultCloseHandler(code, text) + }) +} + +func (s *codexWebsocketSession) bindExecutionLifecycle(opts cliproxyexecutor.Options, conn *websocket.Conn, closer *websocketConnectionCloser, model string) error { + if closer == nil { + return fmt.Errorf("codex websockets executor: websocket connection closer is nil") + } + if s == nil { + return cliproxyexecutor.BindExecutionResource(opts, closer) + } + lifecycle := opts.ExecutionLifecycle + if lifecycle == nil || conn == nil { + return nil + } + + s.lifecycleBindMu.Lock() + defer s.lifecycleBindMu.Unlock() + + s.connMu.Lock() + if s.conn == conn && s.connCloser == nil { + s.connCloser = closer + } + alreadyBound := s.conn == conn && s.connCloser == closer && s.lifecycle == lifecycle + s.connMu.Unlock() + if alreadyBound { + return nil + } + + if errBind := lifecycle.Bind(func() error { + return s.closeBoundConnection(conn, closer, lifecycle) + }); errBind != nil { + return errBind + } + if retained, ok := lifecycle.(interface{ Retain() }); ok { + retained.Retain() + } + + s.connMu.Lock() + if s.conn != conn || s.connCloser != closer { + s.connMu.Unlock() + return fmt.Errorf("codex websockets executor: websocket connection closed during lifecycle bind") + } + previous := s.lifecycle + s.lifecycle = lifecycle + s.lifecycleModel = strings.TrimSpace(model) + s.connMu.Unlock() + if previous != nil && previous != lifecycle { + previous.End("target_replaced") + } + return nil +} + +func (s *codexWebsocketSession) closeBoundConnection(conn *websocket.Conn, closer *websocketConnectionCloser, lifecycle cliproxyexecutor.ExecutionLifecycle) error { + if s == nil || conn == nil { + return nil + } + s.detachConnection(conn, lifecycle) + errClose := closer.Close() + go lifecycle.End("connection_closed") + return errClose +} + +func (s *codexWebsocketSession) detachConnection(conn *websocket.Conn, lifecycle cliproxyexecutor.ExecutionLifecycle) *websocketConnectionCloser { + if s == nil || conn == nil { + return nil + } + s.connMu.Lock() + var closer *websocketConnectionCloser + matched := s.conn == conn + if matched { + closer = s.connCloser + s.conn = nil + s.connCloser = nil + if s.readerConn == conn { + s.readerConn = nil + } + } + if (lifecycle == nil && matched) || (lifecycle != nil && s.lifecycle == lifecycle) { + s.lifecycle = nil + s.lifecycleModel = "" + } + s.connMu.Unlock() + return closer +} + +func closeWebsocketAfterBindFailure(sess *codexWebsocketSession, conn *websocket.Conn, closer *websocketConnectionCloser) { + if conn == nil || closer == nil { + return + } + if sess != nil { + sess.detachConnection(conn, nil) + } + if errClose := closer.Close(); errClose != nil { + log.Errorf("websockets executor: close lifecycle bind failure connection error: %v", errClose) + } +} + +func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, wsURL string) bool { + if sess == nil { + return false + } + + sess.connMu.Lock() + defer sess.connMu.Unlock() + if strings.TrimSpace(sess.authID) == "" && strings.TrimSpace(sess.wsURL) == "" { + return false + } + return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) || strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL) +} + +func existingWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser) { + if sess == nil { + return nil, nil + } + sess.connMu.Lock() + conn := sess.conn + closer := sess.connCloser + matches := conn != nil && closer != nil && + strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && + strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL) + sess.connMu.Unlock() + if !matches || sess.upstreamDisconnectError(conn) != nil { + return nil, nil + } + return conn, closer +} + +func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) { + if sess == nil { + return nil, nil, "", "", nil + } + + sess.connMu.Lock() + defer sess.connMu.Unlock() + conn := sess.conn + if conn == nil || (strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)) { + return nil, nil, "", "", nil + } + + previousAuthID := sess.authID + previousWSURL := sess.wsURL + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" + sess.conn = nil + sess.connCloser = nil + if sess.readerConn == conn { + sess.readerConn = nil + } + return conn, closer, previousAuthID, previousWSURL, lifecycle +} + +func (s *codexWebsocketSession) resetUpstreamDisconnectError(conn *websocket.Conn) { + if s == nil || conn == nil { + return + } + s.upstreamDisconnectErrMu.Lock() + s.upstreamDisconnectErrConn = conn + s.upstreamDisconnectErr = nil + s.upstreamDisconnectErrMu.Unlock() +} + +func (s *codexWebsocketSession) setUpstreamDisconnectError(conn *websocket.Conn, err error) { + if s == nil || conn == nil || err == nil { + return + } + s.upstreamDisconnectErrMu.Lock() + if s.upstreamDisconnectErrConn == conn && s.upstreamDisconnectErr == nil { + s.upstreamDisconnectErr = err + } + s.upstreamDisconnectErrMu.Unlock() +} + +func (s *codexWebsocketSession) upstreamDisconnectError(conn *websocket.Conn) error { + if s == nil || conn == nil { + return nil + } + s.upstreamDisconnectErrMu.RLock() + defer s.upstreamDisconnectErrMu.RUnlock() + if s.upstreamDisconnectErrConn != conn { + return nil + } + return s.upstreamDisconnectErr +} + +func (s *codexWebsocketSession) notifyUpstreamDisconnect(err error) { + if s == nil { + return + } + s.upstreamDisconnectOnce.Do(func() { + if s.upstreamDisconnectCh == nil { + return + } + select { + case s.upstreamDisconnectCh <- err: + default: + } + close(s.upstreamDisconnectCh) + }) +} + +func executionSessionIDFromOptions(opts cliproxyexecutor.Options) string { + if len(opts.Metadata) == 0 { + return "" + } + raw, ok := opts.Metadata[cliproxyexecutor.ExecutionSessionMetadataKey] + if !ok || raw == nil { + return "" + } + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) + case []byte: + return strings.TrimSpace(string(v)) + default: + return "" + } +} + +func (e *CodexWebsocketsExecutor) getOrCreateSession(sessionID string) *codexWebsocketSession { + sessionID = strings.TrimSpace(sessionID) + if sessionID == "" { + return nil + } + if e == nil { + return nil + } + store := e.store + if store == nil { + store = globalCodexWebsocketSessionStore + } + store.mu.Lock() + defer store.mu.Unlock() + if store.sessions == nil { + store.sessions = make(map[string]*codexWebsocketSession) + } + if sess, ok := store.sessions[sessionID]; ok && sess != nil { + return sess + } + sess := &codexWebsocketSession{ + sessionID: sessionID, + upstreamDisconnectCh: make(chan error, 1), + } + store.sessions[sessionID] = sess + return sess +} + +func (e *CodexWebsocketsExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { + sess := e.getOrCreateSession(sessionID) + if sess == nil { + return nil + } + return sess.upstreamDisconnectCh +} + +func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { + if sess == nil { + return e.dialCodexWebsocket(ctx, auth, wsURL, headers) + } + + if staleConn, staleCloser, staleAuthID, staleWSURL, staleLifecycle := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { + logCodexWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) + if staleCloser != nil { + if errClose := staleCloser.Close(); errClose != nil { + log.Errorf("codex websockets executor: close stale websocket error: %v", errClose) + } + } + if staleLifecycle != nil { + staleLifecycle.End("target_changed") + } + } + + sess.connMu.Lock() + conn := sess.conn + closer := sess.connCloser + readerConn := sess.readerConn + sess.connMu.Unlock() + if conn != nil { + if readerConn != conn { + sess.connMu.Lock() + sess.readerConn = conn + sess.connMu.Unlock() + sess.configureConn(conn) + go e.readUpstreamLoop(sess, conn) + } + return conn, closer, nil, nil + } + + conn, closer, resp, errDial := e.dialCodexWebsocket(ctx, auth, wsURL, headers) + if errDial != nil { + return nil, closer, resp, errDial + } + + sess.connMu.Lock() + if sess.conn != nil { + previous := sess.conn + previousCloser := sess.connCloser + sess.connMu.Unlock() + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + return previous, previousCloser, nil, nil + } + sess.conn = conn + sess.connCloser = closer + sess.wsURL = wsURL + sess.authID = authID + sess.readerConn = conn + sess.connMu.Unlock() + + sess.configureConn(conn) + go e.readUpstreamLoop(sess, conn) + logCodexWebsocketConnected(sess.sessionID, authID, wsURL) + return conn, closer, resp, nil +} + +func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, conn *websocket.Conn) { + if e == nil || sess == nil || conn == nil { + return + } + for { + _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) + msgType, payload, errRead := conn.ReadMessage() + if errRead != nil { + invalidate := func() { + e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) + } + invalidated := false + ch, done := sess.activeForConn(conn) + if ch != nil { + invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errRead}, invalidate) + if sess.clearActive(conn, ch) { + close(ch) + } + } + if !invalidated { + invalidate() + } + return + } + + if msgType != websocket.TextMessage { + if msgType == websocket.BinaryMessage { + errBinary := fmt.Errorf("codex websockets executor: unexpected binary message") + invalidate := func() { + e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) + } + invalidated := false + ch, done := sess.activeForConn(conn) + if ch != nil { + invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errBinary}, invalidate) + if sess.clearActive(conn, ch) { + close(ch) + } + } + if !invalidated { + invalidate() + } + return + } + continue + } + + ch, done := sess.activeForConn(conn) + if ch == nil { + continue + } + select { + case ch <- codexWebsocketRead{conn: conn, msgType: msgType, payload: payload}: + case <-done: + } + } +} + +func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { + e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, true) +} + +func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithoutDisconnectNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { + e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, false) +} + +func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error, notify bool) { + if sess == nil || conn == nil { + return + } + + sess.connMu.Lock() + current := sess.conn + authID := sess.authID + wsURL := sess.wsURL + sessionID := sess.sessionID + if current == nil || current != conn { + sess.connMu.Unlock() + return + } + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" + sess.conn = nil + sess.connCloser = nil + if sess.readerConn == conn { + sess.readerConn = nil + } + sess.connMu.Unlock() + + logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err) + if notify { + sess.notifyUpstreamDisconnect(err) + } + if closer != nil { + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + } + if lifecycle != nil { + lifecycle.End(reason) + } +} + +func (e *CodexWebsocketsExecutor) CloseExecutionSession(sessionID string) { + sessionID = strings.TrimSpace(sessionID) + if e == nil { + return + } + if sessionID == "" { + return + } + if sessionID == cliproxyauth.CloseAllExecutionSessionsID { + e.closeAllExecutionSessions("executor_shutdown") + return + } + + store := e.store + if store == nil { + store = globalCodexWebsocketSessionStore + } + store.mu.Lock() + sess := store.sessions[sessionID] + delete(store.sessions, sessionID) + store.mu.Unlock() + + e.closeExecutionSession(sess, "session_closed") +} + +func (e *CodexWebsocketsExecutor) closeAllExecutionSessions(reason string) { + if e == nil { + return + } + + store := e.store + if store == nil { + store = globalCodexWebsocketSessionStore + } + store.mu.Lock() + sessions := make([]*codexWebsocketSession, 0, len(store.sessions)) + for sessionID, sess := range store.sessions { + delete(store.sessions, sessionID) + if sess != nil { + sessions = append(sessions, sess) + } + } + store.mu.Unlock() + + for i := range sessions { + e.closeExecutionSession(sessions[i], reason) + } +} + +func (e *CodexWebsocketsExecutor) closeExecutionSession(sess *codexWebsocketSession, reason string) { + closeCodexWebsocketSession(sess, reason) +} + +func closeCodexWebsocketSession(sess *codexWebsocketSession, reason string) { + if sess == nil { + return + } + reason = strings.TrimSpace(reason) + if reason == "" { + reason = "session_closed" + } + + sess.connMu.Lock() + conn := sess.conn + authID := sess.authID + wsURL := sess.wsURL + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" + sess.conn = nil + sess.connCloser = nil + if sess.readerConn == conn { + sess.readerConn = nil + } + sessionID := sess.sessionID + sess.connMu.Unlock() + + if conn != nil { + logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) + if closer != nil { + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + } + } + if lifecycle != nil { + lifecycle.End(reason) + } +} + +func logCodexWebsocketConnected(sessionID string, authID string, wsURL string) { + log.Infof("codex websockets: upstream connected session=%s auth=%s url=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL)) +} + +func logCodexWebsocketDisconnected(sessionID string, authID string, wsURL string, reason string, err error) { + if err != nil { + log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s err=%v", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason), err) + return + } + log.Infof("codex websockets: upstream disconnected session=%s auth=%s url=%s reason=%s", strings.TrimSpace(sessionID), strings.TrimSpace(authID), strings.TrimSpace(wsURL), strings.TrimSpace(reason)) +} + +// CloseCodexWebsocketSessionsForAuthID closes all active Codex upstream websocket sessions +// associated with the supplied auth ID. +func CloseCodexWebsocketSessionsForAuthID(authID string, reason string) { + authID = strings.TrimSpace(authID) + if authID == "" { + return + } + reason = strings.TrimSpace(reason) + if reason == "" { + reason = "auth_removed" + } + + store := globalCodexWebsocketSessionStore + if store == nil { + return + } + + type sessionItem struct { + sessionID string + sess *codexWebsocketSession + } + + store.mu.Lock() + items := make([]sessionItem, 0, len(store.sessions)) + for sessionID, sess := range store.sessions { + items = append(items, sessionItem{sessionID: sessionID, sess: sess}) + } + store.mu.Unlock() + + matches := make([]sessionItem, 0) + for i := range items { + sess := items[i].sess + if sess == nil { + continue + } + sess.connMu.Lock() + sessAuthID := strings.TrimSpace(sess.authID) + sess.connMu.Unlock() + if sessAuthID == authID { + matches = append(matches, items[i]) + } + } + if len(matches) == 0 { + return + } + + toClose := make([]*codexWebsocketSession, 0, len(matches)) + store.mu.Lock() + for i := range matches { + current, ok := store.sessions[matches[i].sessionID] + if !ok || current == nil || current != matches[i].sess { + continue + } + delete(store.sessions, matches[i].sessionID) + toClose = append(toClose, current) + } + store.mu.Unlock() + + for i := range toClose { + closeCodexWebsocketSession(toClose[i], reason) + } +} diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go new file mode 100644 index 000000000..719e1a363 --- /dev/null +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -0,0 +1,429 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "net/http" + "strings" + + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { + log.Debugf("Executing Codex Websockets stream request with auth ID: %s, model: %s", auth.ID, req.Model) + if ctx == nil { + ctx = context.Background() + } + if opts.Alt == "responses/compact" { + return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} + } + + baseModel := thinking.ParseSuffix(req.Model).ModelName + apiKey, baseURL := codexCreds(auth) + if baseURL == "" { + baseURL = "https://chatgpt.com/backend-api/codex" + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + to := sdktranslator.FromString("codex") + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := originalPayloadSource + originalTranslated, body := translateCodexRequestPair(from, to, baseModel, originalPayload, req.Payload, true) + + body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) + if err != nil { + return nil, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = normalizeCodexInstructions(body) + if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { + body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) + } + body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex websockets executor", body) + body = normalizeCodexWebsocketParallelToolCalls(body, opts.Headers) + body, optimizeMultiAgentV2 := helps.OptimizeCodexMultiAgentV2Request(ctx, opts.Headers, body, e.cfg) + body, replayScope, errReplay := applyCodexReasoningReplayCacheRequired(ctx, from, req, opts, body) + if errReplay != nil { + return nil, errReplay + } + + httpURL := strings.TrimSuffix(baseURL, "/") + "/responses" + wsURL, err := buildCodexResponsesWebsocketURL(httpURL) + if err != nil { + return nil, err + } + + body, wsHeaders, errPromptCache := applyCodexPromptCacheHeadersWithContext(ctx, from, req, body, opts.Headers) + if errPromptCache != nil { + return nil, errPromptCache + } + clientBody := body + var identityState codexIdentityConfuseState + upstreamBody, identityState := applyCodexIdentityConfuseBody(e.cfg, auth, originalPayloadSource, body) + reporter.SetTranslatedReasoningEffort(clientBody, to.String()) + wsHeaders = applyCodexWebsocketHeaders(ctx, wsHeaders, auth, apiKey, e.cfg) + applyModelHeaderOverrides(wsHeaders, baseModel) + applyCodexIdentityConfuseHeaders(wsHeaders, &identityState) + + var authID, authLabel, authType, authValue string + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + + executionSessionID := executionSessionIDFromOptions(opts) + var sess *codexWebsocketSession + if executionSessionID != "" { + sess = e.getOrCreateSession(executionSessionID) + if sess != nil { + sess.reqMu.Lock() + } + } + streamSessionLocked := sess != nil + unlockStreamSession := func() { + if sess != nil && streamSessionLocked { + sess.reqMu.Unlock() + streamSessionLocked = false + } + } + + wsReqBody := buildCodexWebsocketRequestBody(upstreamBody) + wsReqLog := helps.UpstreamRequestLog{ + URL: wsURL, + Method: "WEBSOCKET", + Headers: wsHeaders.Clone(), + Body: wsReqBody, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + } + helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) + + var conn *websocket.Conn + var closer *websocketConnectionCloser + var respHS *http.Response + var errDial error + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + conn, closer = existingWebsocketSessionConn(sess, authID, wsURL) + if conn == nil { + if sess != nil { + sess.reqMu.Unlock() + } + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } + } else { + conn, closer, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + } + var upstreamHeaders http.Header + if respHS != nil { + upstreamHeaders = respHS.Header.Clone() + } + if errDial != nil { + bodyErr := websocketHandshakeBody(respHS) + if respHS != nil { + helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) + } + if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { + if sess != nil { + sess.reqMu.Unlock() + } + if opts.ExecutionLifecycle != nil || cliproxyexecutor.DownstreamWebsocket(ctx) { + return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } + return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts) + } + if respHS != nil && respHS.StatusCode > 0 { + if sess != nil { + sess.reqMu.Unlock() + } + return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) + if sess != nil { + sess.reqMu.Unlock() + } + return nil, errDial + } + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + if sess != nil { + sess.reqMu.Unlock() + } + closeWebsocketAfterBindFailure(sess, conn, closer) + return nil, errBind + } + recordAPIWebsocketHandshake(ctx, e.cfg, respHS) + reporter.StartResponseTTFT() + + if sess == nil { + logCodexWebsocketConnected(executionSessionID, authID, wsURL) + } + + var readCh chan codexWebsocketRead + if sess != nil { + readCh = sess.activate(conn) + } + + if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { + errSend = mapCodexWebsocketWriteError(sess, conn, errSend) + helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) + if sess != nil { + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend) + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + if !shouldRetryCodexWebsocketSend(errSend) { + return nil, errSend + } + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } + e.invalidateUpstreamConn(sess, conn, "send_error", errSend) + if !shouldRetryCodexWebsocketSend(errSend) { + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + return nil, errSend + } + + // Retry once with a new websocket connection for the same execution session. + connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + if errDialRetry != nil || connRetry == nil { + closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") + helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + return nil, errDialRetry + } + previousConn, previousReadCh := conn, readCh + conn = connRetry + closer = closerRetry + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + clearRetryActiveState(sess, previousConn, previousReadCh) + sess.reqMu.Unlock() + closeWebsocketAfterBindFailure(sess, conn, closer) + return nil, errBind + } + readCh = sess.activate(conn) + wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) + helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: wsURL, + Method: "WEBSOCKET", + Headers: wsHeaders.Clone(), + Body: wsReqBodyRetry, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) + recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) + reporter.StartResponseTTFT() + if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { + errSendRetry = mapCodexWebsocketWriteError(sess, conn, errSendRetry) + helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) + e.invalidateUpstreamConn(sess, conn, "send_error", errSendRetry) + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + return nil, errSendRetry + } + wsReqBody = wsReqBodyRetry + } else { + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + return nil, errSend + } + } + + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + terminateReason := "completed" + var terminateErr error + + defer close(out) + defer func() { + if sess != nil { + sess.clearActive(conn, readCh) + unlockStreamSession() + return + } + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, terminateReason, terminateErr) + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + }() + + send := func(chunk cliproxyexecutor.StreamChunk) bool { + if ctx == nil { + out <- chunk + return true + } + select { + case out <- chunk: + return true + case <-ctx.Done(): + return false + } + } + + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) + var param any + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + for { + if ctx != nil && ctx.Err() != nil { + terminateReason = "context_done" + terminateErr = ctx.Err() + _ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()}) + return + } + msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh) + if errRead != nil { + if sess != nil && ctx != nil && ctx.Err() != nil { + terminateReason = "context_done" + terminateErr = ctx.Err() + _ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()}) + return + } + mappedErr := mapCodexWebsocketReadError(errRead) + terminateReason = "read_error" + terminateErr = mappedErr + helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr) + reporter.PublishFailure(ctx, mappedErr) + _ = send(cliproxyexecutor.StreamChunk{Err: mappedErr}) + return + } + if msgType != websocket.TextMessage { + if msgType == websocket.BinaryMessage { + err = fmt.Errorf("codex websockets executor: unexpected binary message") + terminateReason = "unexpected_binary" + terminateErr = err + helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", err) + reporter.PublishFailure(ctx, err) + if sess != nil { + e.invalidateUpstreamConn(sess, conn, "unexpected_binary", err) + } + _ = send(cliproxyexecutor.StreamChunk{Err: err}) + return + } + continue + } + + payload = bytes.TrimSpace(payload) + if len(payload) == 0 { + continue + } + reporter.MarkFirstResponseByte() + payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) + helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) + payload = helps.RestoreCodexMultiAgentV2Response(payload, optimizeMultiAgentV2) + + if wsErr, ok := parseCodexWebsocketError(payload); ok { + terminateReason = "upstream_error" + terminateErr = wsErr + if sess != nil { + e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) + } + if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { + terminateErr = errClearReplay + helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) + return + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) + reporter.PublishFailure(ctx, wsErr) + _ = send(cliproxyexecutor.StreamChunk{Err: wsErr}) + return + } + if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { + terminateReason = "upstream_error" + terminateErr = streamErr + if sess != nil { + unlockStreamSession() + e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) + } + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + terminateErr = errClearReplay + helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + _ = send(cliproxyexecutor.StreamChunk{Err: errClearReplay}) + return + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr) + reporter.PublishFailure(ctx, streamErr) + _ = send(cliproxyexecutor.StreamChunk{Err: streamErr}) + return + } + + eventType := gjson.GetBytes(payload, "type").String() + isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" + if eventType == "response.output_item.done" { + collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) + } + completedPayload := payload + if eventType == "response.completed" || eventType == "response.done" { + completedPayload = normalizeCodexWebsocketCompletion(completedPayload) + completedPayload = patchCodexCompletedOutput(completedPayload, outputItemsByIndex, outputItemsFallback) + cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload) + if detail, ok := helps.ParseCodexUsage(completedPayload); ok { + reporter.Publish(ctx, detail) + } + } + + clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) + if cliproxyexecutor.DownstreamWebsocket(ctx) { + if !send(cliproxyexecutor.StreamChunk{Payload: clientPayload}) { + terminateReason = "context_done" + terminateErr = ctx.Err() + return + } + if isTerminalEvent { + return + } + continue + } + + payload = normalizeCodexWebsocketCompletion(payload) + if eventType == "response.completed" || eventType == "response.done" { + payload = completedPayload + } + eventType = gjson.GetBytes(payload, "type").String() + clientPayload = applyCodexIdentityExposeResponsePayload(payload, identityState) + line := encodeCodexWebsocketAsSSE(clientPayload) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m, claudeInputTokens) + for i := range chunks { + if !send(cliproxyexecutor.StreamChunk{Payload: chunks[i]}) { + terminateReason = "context_done" + terminateErr = ctx.Err() + return + } + } + if eventType == "response.completed" || eventType == "response.done" { + return + } + } + }() + + return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil +} diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 0ff3ea491..e4f0ae486 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -1,34 +1,15 @@ package executor import ( - "bufio" - "bytes" "context" - "encoding/json" "fmt" - "io" "net/http" - "net/url" - "sort" - "strconv" "strings" - "time" - "github.com/google/uuid" - xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" - "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" - cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" - sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" - log "github.com/sirupsen/logrus" - "github.com/tidwall/gjson" - "github.com/tidwall/sjson" - "github.com/tiktoken-go/tokenizer" ) var ( @@ -123,2832 +104,3 @@ func (e *XAIExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) return httpClient.Do(httpReq) } - -func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - if opts.Alt == "responses/compact" { - return e.executeCompact(ctx, auth, req, opts) - } - if endpointPath := xaiImageEndpointPath(opts); endpointPath != "" { - return e.executeImages(ctx, auth, req, endpointPath) - } - if xaiIsVideoRequest(opts) { - return e.executeVideos(ctx, auth, req, opts) - } - - token, _ := xaiCreds(auth) - baseURL := xaiChatBaseURL(auth) - logXAIResolvedBaseURL(ctx, baseURL) - - prepared, err := e.prepareResponsesRequest(ctx, req, opts, true) - if err != nil { - return resp, err - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) - - url := strings.TrimSuffix(baseURL, "/") + "/responses" - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(prepared.body)) - if err != nil { - return resp, err - } - applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID) - e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) - - httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("xai executor: close response body error: %v", errClose) - } - }() - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - data, errRead := io.ReadAll(httpResp.Body) - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - return resp, errRead - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) - return resp, xaiStatusErr(httpResp.StatusCode, data) - } - - data, err := io.ReadAll(httpResp.Body) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte - responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) - for _, line := range bytes.Split(data, []byte("\n")) { - if !bytes.HasPrefix(line, xaiDataTag) { - continue - } - eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):])) - eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) - eventData = responseFilter.apply(eventData) - if len(eventData) == 0 { - continue - } - switch gjson.GetBytes(eventData, "type").String() { - case "response.output_item.done": - xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) - case "response.completed": - if detail, ok := helps.ParseCodexUsage(eventData); ok { - reporter.Publish(ctx, detail) - } - completedData := xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) - completedData = xaiNormalizeReasoningSummaryData(completedData) - cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, completedData) - var param any - out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, completedData, ¶m) - return cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}, nil - } - } - - return resp, statusErr{code: http.StatusRequestTimeout, msg: "xai stream error: stream disconnected before response.completed"} -} - -func (e *XAIExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - prepared, data, headers, errCompact := e.executeCompactRequest(ctx, auth, req, opts) - if errCompact != nil { - return resp, errCompact - } - - var param any - out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, data, ¶m) - return cliproxyexecutor.Response{Payload: out, Headers: headers}, nil -} - -func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*xaiPreparedRequest, []byte, http.Header, error) { - token, _ := xaiCreds(auth) - // Compact must not use xaiChatBaseURL: CLI chat-proxy returns 404 for - // /responses/compact and a 404 cools down the whole xAI auth pool. - baseURL := xaiCompactBaseURL(auth) - logXAIResolvedBaseURL(ctx, baseURL) - - prepared, err := e.prepareResponsesRequestTo(ctx, req, opts, false, sdktranslator.FormatOpenAIResponse) - if err != nil { - return nil, nil, nil, err - } - prepared.body, _ = sjson.DeleteBytes(prepared.body, "stream") - prepared.body, _ = sjson.DeleteBytes(prepared.body, "tools") - for _, field := range []string{"max_output_tokens", "temperature", "top_p", "top_k", "stop"} { - prepared.body, _ = sjson.DeleteBytes(prepared.body, field) - } - prepared.body = xaiRemoveInputItemsByType(prepared.body, "compaction_trigger") - - reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) - - requestURL := strings.TrimSuffix(baseURL, "/") + "/responses/compact" - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL, bytes.NewReader(prepared.body)) - if err != nil { - return nil, nil, nil, err - } - // Official API / custom compact endpoints use standard API headers, not CLI - // chat-proxy identity headers (which applyXAIChatHeaders may still attach for OAuth chat). - applyXAIHeaders(httpReq, auth, token, false, prepared.sessionID) - e.recordXAIRequest(ctx, auth, requestURL, httpReq.Header.Clone(), prepared.body) - - httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return nil, nil, nil, err - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("xai executor: close response body error: %v", errClose) - } - }() - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - - data, err := io.ReadAll(httpResp.Body) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return nil, nil, nil, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) - err = xaiStatusErr(httpResp.StatusCode, data) - return nil, nil, nil, err - } - - reporter.Publish(ctx, helps.ParseOpenAIUsage(data)) - reporter.EnsurePublished(ctx) - clearXAIReasoningReplayAfterCompaction(ctx, prepared.replayScope) - return prepared, data, httpResp.Header.Clone(), nil -} - -func (e *XAIExecutor) executeCompactionTriggerStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { - prepared, data, headers, err := e.executeCompactRequest(ctx, auth, req, opts) - if err != nil { - return nil, err - } - - headers = headers.Clone() - if headers == nil { - headers = make(http.Header) - } - headers.Set("Content-Type", "text/event-stream") - - chunks := xaiBuildCompactionTriggerStreamChunks(prepared, data) - out := make(chan cliproxyexecutor.StreamChunk, len(chunks)) - for _, chunk := range chunks { - out <- cliproxyexecutor.StreamChunk{Payload: chunk} - } - close(out) - return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil -} - -func xaiInputHasItemType(body []byte, itemType string) bool { - input := gjson.GetBytes(body, "input") - if !input.IsArray() { - return false - } - for _, item := range input.Array() { - if item.Get("type").String() == itemType { - return true - } - } - return false -} - -func xaiRemoveInputItemsByType(body []byte, itemType string) []byte { - input := gjson.GetBytes(body, "input") - if !input.IsArray() { - return body - } - - var buf bytes.Buffer - buf.WriteByte('[') - kept := 0 - for _, item := range input.Array() { - if item.Get("type").String() == itemType { - continue - } - if kept > 0 { - buf.WriteByte(',') - } - buf.WriteString(item.Raw) - kept++ - } - buf.WriteByte(']') - - updated, err := sjson.SetRawBytes(body, "input", buf.Bytes()) - if err != nil { - return body - } - return updated -} - -func xaiBuildCompactionTriggerStreamChunks(prepared *xaiPreparedRequest, compactData []byte) [][]byte { - responseID := xaiCompactionResponseID(compactData) - now := time.Now().Unix() - createdAt := gjson.GetBytes(compactData, "created_at").Int() - if createdAt == 0 { - createdAt = now - } - completedAt := gjson.GetBytes(compactData, "completed_at").Int() - if completedAt == 0 { - completedAt = now - } - - item := xaiCompactionOutputItem(compactData, responseID) - output := make([]byte, 0, len(item)+2) - output = append(output, '[') - output = append(output, item...) - output = append(output, ']') - - createdResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "in_progress") - inProgressResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "in_progress") - completedResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "completed") - completedResponse, _ = sjson.SetBytes(completedResponse, "completed_at", completedAt) - completedResponse, _ = sjson.SetRawBytes(completedResponse, "output", output) - if usage := gjson.GetBytes(compactData, "usage"); usage.Exists() { - completedResponse, _ = sjson.SetRawBytes(completedResponse, "usage", []byte(usage.Raw)) - } - - createdPayload := []byte(`{"type":"response.created","sequence_number":0}`) - createdPayload, _ = sjson.SetRawBytes(createdPayload, "response", createdResponse) - inProgressPayload := []byte(`{"type":"response.in_progress","sequence_number":1}`) - inProgressPayload, _ = sjson.SetRawBytes(inProgressPayload, "response", inProgressResponse) - addedPayload := []byte(`{"type":"response.output_item.added","sequence_number":2,"output_index":0}`) - addedPayload, _ = sjson.SetRawBytes(addedPayload, "item", item) - keepalivePayload := []byte(`{"type":"keepalive","sequence_number":3}`) - donePayload := []byte(`{"type":"response.output_item.done","sequence_number":4,"output_index":0}`) - donePayload, _ = sjson.SetRawBytes(donePayload, "item", item) - completedPayload := []byte(`{"type":"response.completed","sequence_number":5}`) - completedPayload, _ = sjson.SetRawBytes(completedPayload, "response", completedResponse) - - return [][]byte{ - xaiBuildSSEFrame("response.created", createdPayload), - xaiBuildSSEFrame("response.in_progress", inProgressPayload), - xaiBuildSSEFrame("response.output_item.added", addedPayload), - xaiBuildSSEFrame("keepalive", keepalivePayload), - xaiBuildSSEFrame("response.output_item.done", donePayload), - xaiBuildSSEFrame("response.completed", completedPayload), - } -} - -func xaiBuildCompactionBaseResponse(prepared *xaiPreparedRequest, compactData []byte, responseID string, createdAt int64, status string) []byte { - response := []byte(`{"id":"","object":"response","created_at":0,"status":"","background":false,"error":null,"incomplete_details":null,"output":[]}`) - response, _ = sjson.SetBytes(response, "id", responseID) - response, _ = sjson.SetBytes(response, "created_at", createdAt) - response, _ = sjson.SetBytes(response, "status", status) - if model := gjson.GetBytes(compactData, "model").String(); model != "" { - response, _ = sjson.SetBytes(response, "model", model) - } else if prepared != nil && prepared.baseModel != "" { - response, _ = sjson.SetBytes(response, "model", prepared.baseModel) - } - - if prepared == nil { - return response - } - for _, field := range []string{ - "instructions", - "max_output_tokens", - "max_tool_calls", - "parallel_tool_calls", - "previous_response_id", - "prompt_cache_key", - "reasoning", - "text", - "tool_choice", - "tools", - "top_logprobs", - "top_p", - "truncation", - "user", - "metadata", - } { - if value := gjson.GetBytes(prepared.body, field); value.Exists() { - response, _ = sjson.SetRawBytes(response, field, []byte(value.Raw)) - } - } - return response -} - -func xaiCompactionOutputItem(compactData []byte, responseID string) []byte { - itemResult := gjson.GetBytes(compactData, "output.0") - item := []byte(`{"type":"compaction"}`) - if itemResult.Exists() && itemResult.Type == gjson.JSON { - item = []byte(itemResult.Raw) - } - if !gjson.GetBytes(item, "type").Exists() { - item, _ = sjson.SetBytes(item, "type", "compaction") - } - if !gjson.GetBytes(item, "id").Exists() { - item, _ = sjson.SetBytes(item, "id", xaiCompactionItemID(responseID)) - } - return item -} - -func xaiCompactionResponseID(compactData []byte) string { - if responseID := strings.TrimSpace(gjson.GetBytes(compactData, "id").String()); responseID != "" { - if strings.HasPrefix(responseID, "resp_") { - return responseID - } - return "resp_" + strings.TrimPrefix(responseID, "cmp_") - } - return fmt.Sprintf("resp_xai_compaction_%d", time.Now().UnixNano()) -} - -func xaiCompactionItemID(responseID string) string { - if suffix := strings.TrimPrefix(responseID, "resp_"); suffix != "" && suffix != responseID { - return "cmp_" + suffix - } - return "cmp_" + responseID -} - -func xaiBuildSSEFrame(eventName string, data []byte) []byte { - out := make([]byte, 0, len(eventName)+len(data)+16) - out = append(out, "event: "...) - out = append(out, eventName...) - out = append(out, '\n') - out = append(out, "data: "...) - out = append(out, data...) - out = append(out, '\n', '\n') - return out -} - -func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, endpointPath string) (resp cliproxyexecutor.Response, err error) { - model := strings.TrimSpace(gjson.GetBytes(req.Payload, "model").String()) - if model == "" { - model = strings.TrimSpace(req.Model) - } - reporter := helps.NewExecutorUsageReporter(ctx, e, model, auth) - defer reporter.TrackFailure(ctx, &err) - - token, baseURL := xaiCreds(auth) - if baseURL == "" { - baseURL = xaiauth.DefaultAPIBaseURL - } - logXAIResolvedBaseURL(ctx, baseURL) - if endpointPath == "" { - endpointPath = xaiDefaultImageEndpointPath - } - - payload := normalizeXAIImageRefs(req.Payload) - url := strings.TrimSuffix(baseURL, "/") + endpointPath - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload)) - if err != nil { - return resp, err - } - applyXAIHeaders(httpReq, auth, token, false, "") - e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), payload) - - httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("xai executor: close response body error: %v", errClose) - } - }() - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - - data, err := io.ReadAll(httpResp.Body) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) - err = xaiStatusErr(httpResp.StatusCode, data) - return resp, err - } - - reporter.EnsurePublished(ctx) - return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil -} - -func (e *XAIExecutor) executeVideos(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { - token, baseURL := xaiCreds(auth) - if baseURL == "" { - baseURL = xaiauth.DefaultAPIBaseURL - } - logXAIResolvedBaseURL(ctx, baseURL) - - payload := normalizeXAIImageRefs(req.Payload) - method := http.MethodPost - endpointPath := xaiVideosGenerationsPath - var body io.Reader = bytes.NewReader(payload) - - switch path := xaiVideoEndpointPath(opts); path { - case xaiVideosGenerationsPath, xaiVideosEditsPath, xaiVideosExtensionsPath: - endpointPath = path - default: - if requestID := strings.TrimSpace(gjson.GetBytes(payload, "request_id").String()); requestID != "" { - method = http.MethodGet - endpointPath = xaiVideosPath + "/" + url.PathEscape(requestID) - body = nil - } - } - requestURL := strings.TrimSuffix(baseURL, "/") + endpointPath - httpReq, err := http.NewRequestWithContext(ctx, method, requestURL, body) - if err != nil { - return resp, err - } - applyXAIHeaders(httpReq, auth, token, false, "") - if method == http.MethodPost { - key := xaiMetadataString(opts.Metadata, xaiIdempotencyKeyMetaKey) - if key == "" && opts.Headers != nil { - key = strings.TrimSpace(opts.Headers.Get("x-idempotency-key")) - } - if key != "" { - httpReq.Header.Set("x-idempotency-key", key) - } - } - e.recordXAIRequest(ctx, auth, requestURL, httpReq.Header.Clone(), payload) - - httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("xai executor: close response body error: %v", errClose) - } - }() - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - - data, err := io.ReadAll(httpResp.Body) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return resp, err - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) - return resp, xaiStatusErr(httpResp.StatusCode, data) - } - - return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil -} - -func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { - if opts.Alt == "responses/compact" { - return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} - } - if xaiInputHasItemType(req.Payload, "compaction_trigger") { - return e.executeCompactionTriggerStream(ctx, auth, req, opts) - } - - token, _ := xaiCreds(auth) - baseURL := xaiChatBaseURL(auth) - logXAIResolvedBaseURL(ctx, baseURL) - - prepared, err := e.prepareResponsesRequest(ctx, req, opts, true) - if err != nil { - return nil, err - } - - reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) - defer reporter.TrackFailure(ctx, &err) - reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) - - url := strings.TrimSuffix(baseURL, "/") + "/responses" - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(prepared.body)) - if err != nil { - return nil, err - } - applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID) - e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) - - httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) - httpClient = reporter.TrackHTTPClient(httpClient) - httpResp, err := httpClient.Do(httpReq) - if err != nil { - helps.RecordAPIResponseError(ctx, e.cfg, err) - return nil, err - } - helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) - if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { - data, errRead := io.ReadAll(httpResp.Body) - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("xai executor: close response body error: %v", errClose) - } - if errRead != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errRead) - return nil, errRead - } - helps.AppendAPIResponseChunk(ctx, e.cfg, data) - helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) - return nil, xaiStatusErr(httpResp.StatusCode, data) - } - - out := make(chan cliproxyexecutor.StreamChunk) - go func() { - defer close(out) - defer func() { - if errClose := httpResp.Body.Close(); errClose != nil { - log.Errorf("xai executor: close response body error: %v", errClose) - } - }() - scanner := bufio.NewScanner(httpResp.Body) - scanner.Buffer(nil, 52_428_800) - claudeInputTokens := helps.NewClaudeInputTokenState(prepared.from, prepared.to, prepared.responseFormat, prepared.originalPayload) - var param any - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte - responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) - var pendingEventLine []byte - emitTranslatedLine := func(translatedLine []byte) bool { - chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, translatedLine, ¶m, claudeInputTokens) - for i := range chunks { - select { - case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: - case <-ctx.Done(): - return false - } - } - return true - } - for scanner.Scan() { - line := scanner.Bytes() - helps.AppendAPIResponseChunk(ctx, e.cfg, line) - - if bytes.HasPrefix(line, xaiEventTag) { - if pendingEventLine != nil && !emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) { - return - } - pendingEventLine = bytes.Clone(line) - continue - } - - if bytes.HasPrefix(line, xaiDataTag) { - eventDataList := xaiNormalizeReasoningSummaryDataEvents(bytes.TrimSpace(line[len(xaiDataTag):])) - hasPendingEventLine := pendingEventLine != nil - for i, eventData := range eventDataList { - eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) - eventData = responseFilter.apply(eventData) - if len(eventData) == 0 { - if hasPendingEventLine && i == 0 { - pendingEventLine = nil - } - continue - } - normalizedEventName := gjson.GetBytes(eventData, "type").String() - switch normalizedEventName { - case "response.output_item.done": - xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) - case "response.completed": - if detail, ok := helps.ParseCodexUsage(eventData); ok { - reporter.Publish(ctx, detail) - } - eventData = xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) - eventData = xaiNormalizeReasoningSummaryData(eventData) - cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, eventData) - normalizedEventName = gjson.GetBytes(eventData, "type").String() - } - - if hasPendingEventLine { - eventLine := []byte("event: " + normalizedEventName) - if i == 0 { - eventLine = xaiNormalizeReasoningSummaryEventLine(pendingEventLine, normalizedEventName) - pendingEventLine = nil - } - if !emitTranslatedLine(eventLine) { - return - } - } - if !emitTranslatedLine(append([]byte("data: "), eventData...)) { - return - } - } - continue - } - - if pendingEventLine != nil { - if !emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) { - return - } - pendingEventLine = nil - } - if !emitTranslatedLine(bytes.Clone(line)) { - return - } - } - if pendingEventLine != nil { - emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) - } - if errScan := scanner.Err(); errScan != nil { - helps.RecordAPIResponseError(ctx, e.cfg, errScan) - reporter.PublishFailure(ctx, errScan) - select { - case out <- cliproxyexecutor.StreamChunk{Err: errScan}: - case <-ctx.Done(): - } - } - }() - return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil -} - -// CountTokens estimates token count for xAI Responses requests. -func (e *XAIExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - prepared, err := e.prepareResponsesRequest(ctx, req, opts, false) - if err != nil { - return cliproxyexecutor.Response{}, err - } - enc, err := tokenizer.Get(tokenizer.O200kBase) - if err != nil { - return cliproxyexecutor.Response{}, fmt.Errorf("xai executor: tokenizer init failed: %w", err) - } - count, err := countXAIInputTokens(enc, prepared.body) - if err != nil { - return cliproxyexecutor.Response{}, fmt.Errorf("xai executor: token counting failed: %w", err) - } - usageJSON := fmt.Sprintf(`{"response":{"usage":{"input_tokens":%d,"output_tokens":0,"total_tokens":%d}}}`, count, count) - translated := sdktranslator.TranslateTokenCount(ctx, prepared.to, prepared.responseFormat, count, []byte(usageJSON)) - return cliproxyexecutor.Response{Payload: translated}, nil -} - -func countXAIInputTokens(enc tokenizer.Codec, body []byte) (int64, error) { - if enc == nil { - return 0, fmt.Errorf("encoder is nil") - } - if len(body) == 0 { - return 0, nil - } - - root := gjson.ParseBytes(body) - segments := make([]string, 0, 32) - xaiAppendTokenString(&segments, root.Get("instructions")) - xaiCollectInputTokenSegments(root.Get("input"), &segments) - xaiCollectToolTokenSegments(root.Get("tools"), &segments) - - textFormat := root.Get("text.format") - if textFormat.Exists() { - xaiAppendTokenString(&segments, textFormat.Get("name")) - xaiAppendTokenJSON(&segments, textFormat.Get("schema")) - } - - if len(segments) == 0 { - return 0, nil - } - count, err := enc.Count(strings.Join(segments, "\n")) - if err != nil { - return 0, err - } - return int64(count), nil -} - -func xaiCollectInputTokenSegments(input gjson.Result, segments *[]string) { - if input.Type == gjson.String { - xaiAppendTokenString(segments, input) - return - } - if !input.IsArray() { - return - } - for _, item := range input.Array() { - switch item.Get("type").String() { - case "message": - xaiCollectContentTokenSegments(item.Get("content"), segments) - case "function_call": - xaiAppendTokenString(segments, item.Get("name")) - xaiAppendTokenJSON(segments, item.Get("arguments")) - case "function_call_output": - xaiAppendTokenJSON(segments, item.Get("output")) - case "reasoning": - for _, part := range item.Get("summary").Array() { - xaiAppendTokenString(segments, part.Get("text")) - } - } - } -} - -func xaiCollectContentTokenSegments(content gjson.Result, segments *[]string) { - if content.Type == gjson.String { - xaiAppendTokenString(segments, content) - return - } - if !content.IsArray() { - return - } - for _, part := range content.Array() { - switch part.Get("type").String() { - case "text", "input_text", "output_text": - xaiAppendTokenString(segments, part.Get("text")) - case "refusal": - xaiAppendTokenString(segments, part.Get("refusal")) - case "input_image": - xaiAppendTokenString(segments, part.Get("image_url")) - xaiAppendTokenString(segments, part.Get("file_id")) - case "input_file": - xaiAppendTokenString(segments, part.Get("file_data")) - xaiAppendTokenString(segments, part.Get("file_url")) - xaiAppendTokenString(segments, part.Get("file_id")) - xaiAppendTokenString(segments, part.Get("filename")) - case "input_audio": - xaiAppendTokenString(segments, part.Get("data")) - xaiAppendTokenString(segments, part.Get("input_audio.data")) - } - } -} - -func xaiCollectToolTokenSegments(tools gjson.Result, segments *[]string) { - if !tools.IsArray() { - return - } - for _, tool := range tools.Array() { - if tool.Get("type").String() != xaiFunctionToolType { - continue - } - xaiAppendTokenString(segments, tool.Get("name")) - xaiAppendTokenString(segments, tool.Get("description")) - xaiAppendTokenJSON(segments, tool.Get("parameters")) - } -} - -func xaiAppendTokenString(segments *[]string, value gjson.Result) { - if text := strings.TrimSpace(value.String()); text != "" { - *segments = append(*segments, text) - } -} - -func xaiAppendTokenJSON(segments *[]string, value gjson.Result) { - if !value.Exists() { - return - } - if value.Type == gjson.String { - xaiAppendTokenString(segments, value) - return - } - if text := strings.TrimSpace(value.Raw); text != "" { - *segments = append(*segments, text) - } -} - -// Refresh refreshes xAI OAuth credentials using the stored refresh token. -func (e *XAIExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { - log.Debugf("xai executor: refresh called") - if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { - return refreshed, err - } - if auth == nil { - return nil, statusErr{code: http.StatusInternalServerError, msg: "xai executor: auth is nil"} - } - refreshToken := xaiMetadataString(auth.Metadata, "refresh_token") - if refreshToken == "" { - return auth, nil - } - tokenEndpoint := xaiMetadataString(auth.Metadata, "token_endpoint") - svc := xaiauth.NewXAIAuthWithProxyURL(e.cfg, auth.ProxyURL) - td, err := svc.RefreshTokens(ctx, refreshToken, tokenEndpoint) - if err != nil { - return nil, err - } - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["type"] = "xai" - auth.Metadata["auth_kind"] = "oauth" - auth.Metadata["access_token"] = td.AccessToken - if td.RefreshToken != "" { - auth.Metadata["refresh_token"] = td.RefreshToken - } - if td.IDToken != "" { - auth.Metadata["id_token"] = td.IDToken - } - if td.TokenType != "" { - auth.Metadata["token_type"] = td.TokenType - } - if td.ExpiresIn > 0 { - auth.Metadata["expires_in"] = td.ExpiresIn - } - if td.Expire != "" { - auth.Metadata["expired"] = td.Expire - } - if td.Email != "" { - auth.Metadata["email"] = td.Email - } - if td.Subject != "" { - auth.Metadata["sub"] = td.Subject - } - if tokenEndpoint != "" { - auth.Metadata["token_endpoint"] = tokenEndpoint - } - if xaiMetadataString(auth.Metadata, "base_url") == "" { - auth.Metadata["base_url"] = xaiauth.DefaultAPIBaseURL - } - auth.Metadata["last_refresh"] = time.Now().UTC().Format(time.RFC3339) - if auth.Attributes == nil { - auth.Attributes = make(map[string]string) - } - auth.Attributes["auth_kind"] = "oauth" - if strings.TrimSpace(auth.Attributes["base_url"]) == "" { - auth.Attributes["base_url"] = xaiauth.DefaultAPIBaseURL - } - return auth, nil -} - -type xaiPreparedRequest struct { - baseModel string - from sdktranslator.Format - responseFormat sdktranslator.Format - to sdktranslator.Format - originalPayload []byte - body []byte - namespaceTools map[string]xaiNamespaceToolRef - clientDeclaredTools map[xaiClientToolKey]struct{} - sessionID string - replayScope xaiReasoningReplayScope - filterInternalXSearch bool -} - -type xaiNamespaceToolRef struct { - namespace string - name string -} - -// xaiClientToolKey identifies a client-declared callable tool using the -// post-restore Responses shape (short name + optional namespace) and the -// effective upstream tool type after normalizeXAITool (client custom tools are -// sent as function). Response call types are matched against this effective -// kind so internal custom_tool_call traces are not exempted merely because a -// client declared an ordinary function/custom tool with the same short name, -// while legitimate function_call responses for normalized custom tools are kept. -type xaiClientToolKey struct { - namespace string - name string - toolType string -} - -func (e *XAIExecutor) prepareResponsesRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) (*xaiPreparedRequest, error) { - return e.prepareResponsesRequestTo(ctx, req, opts, stream, sdktranslator.FormatCodex) -} - -func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool, to sdktranslator.Format) (*xaiPreparedRequest, error) { - baseModel := thinking.ParseSuffix(req.Model).ModelName - from := opts.SourceFormat - responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) - originalPayloadSource := req.Payload - if len(opts.OriginalRequest) > 0 { - originalPayloadSource = opts.OriginalRequest - } - originalPayload := bytes.Clone(originalPayloadSource) - originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, stream) - originalTranslated = preserveXAIResponsesOutputControls(originalTranslated, originalPayload, from) - body := sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), stream) - body = preserveXAIResponsesOutputControls(body, req.Payload, from) - - var err error - body, err = thinking.ApplyThinking(body, req.Model, from.String(), e.Identifier(), e.Identifier()) - if err != nil { - return nil, err - } - - requestedModel := helps.PayloadRequestedModel(opts, req.Model) - requestPath := helps.PayloadRequestPath(opts) - body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body = helps.SetStringIfDifferent(body, "model", baseModel) - body = helps.SetBoolIfDifferent(body, "stream", stream) - body, _ = sjson.DeleteBytes(body, "previous_response_id") - body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") - body, _ = sjson.DeleteBytes(body, "safety_identifier") - body, _ = sjson.DeleteBytes(body, "stream_options") - body = helps.RewriteCodexMultiAgentV2Input(ctx, opts.Headers, body, e.cfg) - namespaceTools := collectXAINamespaceToolRefs(body) - // Collect before normalizeXAITools flattens namespace wrappers so keys match - // the post-restore (namespace, short-name) shape used by the response filter. - clientDeclaredTools := collectXAIClientDeclaredToolKeys(body) - body = normalizeXAITools(body) - body = promoteXAIAdditionalTools(body) - // Drop choices that point at tools removed by normalizeXAITools before we - // inject native x_search, so a surviving allowed_tools / forced choice is not - // left pointing at a deleted tool once only x_search remains. - body = normalizeXAINamespaceToolChoice(body) - body = pruneXAIOrphanedToolChoice(body) - body = normalizeXAIToolChoiceForTools(body) - body = ensureXAINativeXSearchTool(body) - var replayScope xaiReasoningReplayScope - body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body) - if err != nil { - return nil, err - } - body = normalizeXAIInputCustomToolCalls(body) - body = normalizeXAIInputNamespaceToolCalls(body) - body = normalizeXAIInputReasoningItems(body) - body = sanitizeXAIInputEncryptedContent(body) - body = normalizeCodexInstructions(body) - body = sanitizeXAIResponsesBody(body, baseModel) - body = normalizeXAIImageRefs(body) - - sessionID, errSession := xaiResolveComposerSessionID(ctx, req, opts, baseModel) - if errSession != nil { - return nil, errSession - } - if sessionID != "" { - body = helps.SetStringIfDifferent(body, "prompt_cache_key", sessionID) - } - - return &xaiPreparedRequest{ - baseModel: baseModel, - from: from, - responseFormat: responseFormat, - to: to, - originalPayload: originalPayload, - body: body, - namespaceTools: namespaceTools, - clientDeclaredTools: clientDeclaredTools, - sessionID: sessionID, - replayScope: replayScope, - filterInternalXSearch: xaiRequestHasNativeXSearch(body), - }, nil -} - -func (e *XAIExecutor) recordXAIRequest(ctx context.Context, auth *cliproxyauth.Auth, url string, headers http.Header, body []byte) { - var authID, authLabel, authType, authValue string - if auth != nil { - authID = auth.ID - authLabel = auth.Label - authType, authValue = auth.AccountInfo() - } - helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ - URL: url, - Method: http.MethodPost, - Headers: headers, - Body: body, - Provider: e.Identifier(), - AuthID: authID, - AuthLabel: authLabel, - AuthType: authType, - AuthValue: authValue, - }) -} - -func xaiCreds(auth *cliproxyauth.Auth) (token, baseURL string) { - if auth == nil { - return "", "" - } - if auth.Attributes != nil { - token = strings.TrimSpace(auth.Attributes["api_key"]) - baseURL = strings.TrimSpace(auth.Attributes["base_url"]) - } - if auth.Metadata != nil { - if token == "" { - token = xaiMetadataString(auth.Metadata, "access_token") - } - if baseURL == "" { - baseURL = xaiMetadataString(auth.Metadata, "base_url") - } - } - return token, baseURL -} - -// xaiUsingAPI reports whether this xAI auth should use the official API path -// for non-media HTTP chat. OAuth defaults to false to use Grok Build. -func xaiUsingAPI(auth *cliproxyauth.Auth) bool { - if auth == nil { - return true - } - if len(auth.Attributes) > 0 { - if raw := strings.TrimSpace(auth.Attributes[xaiUsingAPIAttr]); raw != "" { - parsed, errParse := strconv.ParseBool(raw) - if errParse == nil { - return parsed - } - } - } - if len(auth.Metadata) > 0 { - raw, ok := auth.Metadata[xaiUsingAPIAttr] - if ok && raw != nil { - switch v := raw.(type) { - case bool: - return v - case string: - parsed, errParse := strconv.ParseBool(strings.TrimSpace(v)) - if errParse == nil { - return parsed - } - default: - } - } - } - if raw := strings.TrimSpace(auth.Attributes["auth_kind"]); raw != "" { - return !strings.EqualFold(raw, "oauth") - } - return !strings.EqualFold(xaiMetadataString(auth.Metadata, "auth_kind"), "oauth") -} - -// xaiChatBaseURL returns the base URL for non-image/video xAI HTTP chat requests. -// When auth using_api is true, the official API base URL logic is used. When it -// is false (including its OAuth default), empty or official default base_url is -// rewritten to the CLI chat-proxy endpoint; an explicit non-default base_url is -// still honored. -// Websocket and compact transports intentionally do not use this helper: -// cli-chat-proxy only accepts HTTP POST chat and does not implement -// /responses/compact (404) or websocket upgrades (405). -func xaiChatBaseURL(auth *cliproxyauth.Auth) string { - _, baseURL := xaiCreds(auth) - if xaiUsingAPI(auth) { - if baseURL == "" { - return xaiauth.DefaultAPIBaseURL - } - return baseURL - } - if baseURL != "" && !xaiIsDefaultAPIBaseURL(baseURL) { - return baseURL - } - return xaiauth.CLIChatProxyBaseURL -} - -// xaiCompactBaseURL returns the base URL for xAI /responses/compact requests. -// Compact must stay on the official API (or an explicit non-CLI-proxy base_url). -// Reusing xaiChatBaseURL would pin OAuth traffic to cli-chat-proxy, which returns -// 404 for /responses/compact and then cools down the auth pool as not_found. -func xaiCompactBaseURL(auth *cliproxyauth.Auth) string { - _, baseURL := xaiCreds(auth) - if baseURL == "" || xaiIsCLIChatProxyBaseURL(baseURL) { - return xaiauth.DefaultAPIBaseURL - } - return baseURL -} - -func xaiNormalizeBaseURL(baseURL string) string { - return strings.TrimRight(strings.TrimSpace(baseURL), "/") -} - -func xaiIsDefaultAPIBaseURL(baseURL string) bool { - return xaiNormalizeBaseURL(baseURL) == xaiNormalizeBaseURL(xaiauth.DefaultAPIBaseURL) -} - -func xaiIsCLIChatProxyBaseURL(baseURL string) bool { - return xaiNormalizeBaseURL(baseURL) == xaiNormalizeBaseURL(xaiauth.CLIChatProxyBaseURL) -} - -// xaiBaseURLSource classifies a resolved xAI base URL for logging. -func xaiBaseURLSource(baseURL string) string { - switch { - case xaiIsDefaultAPIBaseURL(baseURL): - return "DefaultAPIBaseURL" - case xaiIsCLIChatProxyBaseURL(baseURL): - return "CLIChatProxyBaseURL" - default: - return "custom" - } -} - -// logXAIResolvedBaseURL emits a console log for the resolved upstream base URL. -func logXAIResolvedBaseURL(ctx context.Context, baseURL string) { - helps.LogWithRequestID(ctx).Infof("xai: using base_url=%s source=%s", baseURL, xaiBaseURLSource(baseURL)) -} - -func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) { - applyXAIDefaultHeaders(r, token, stream, sessionID) - applyXAICustomHeaders(r, auth) -} - -func applyXAIDefaultHeaders(r *http.Request, token string, stream bool, sessionID string) { - r.Header.Set("Content-Type", "application/json") - if strings.TrimSpace(token) != "" { - r.Header.Set("Authorization", "Bearer "+token) - } - if stream { - r.Header.Set("Accept", "text/event-stream") - } else { - r.Header.Set("Accept", "application/json") - } - r.Header.Set("Connection", "Keep-Alive") - if sessionID != "" { - r.Header.Set("x-grok-conv-id", sessionID) - } -} - -func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth) { - var attrs map[string]string - if auth != nil { - attrs = auth.Attributes - } - util.ApplyCustomHeadersFromAttrs(r, attrs) -} - -// applyXAIChatHeaders applies standard xAI headers for non-image/video chat -// requests. When using_api is true, this matches the standard -// applyXAIHeaders behavior. CLI chat-proxy identity headers are only attached -// when using_api is false and the resolved chat base URL is the official CLI -// chat-proxy endpoint. -func applyXAIChatHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) { - if xaiUsingAPI(auth) { - applyXAIHeaders(r, auth, token, stream, sessionID) - return - } - applyXAIDefaultHeaders(r, token, stream, sessionID) - if xaiIsCLIChatProxyBaseURL(xaiChatBaseURL(auth)) { - r.Header.Set(xaiTokenAuthHeader, xaiTokenAuthValue) - r.Header.Set(xaiClientVersionHeader, xaiClientVersionValue) - r.Header.Set("User-Agent", "xai-grok-workspace/"+xaiClientVersionValue) - } - applyXAICustomHeaders(r, auth) -} - -func xaiResolveComposerSessionID(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, baseModel string) (string, error) { - if sessionID := xaiExecutionSessionID(req, opts); sessionID != "" { - return sessionID, nil - } - if !xaiRequiresIsolatedConversation(baseModel) { - return "", nil - } - cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, baseModel, req.Payload, opts.Headers) - if errCache != nil { - return "", errCache - } - if ok { - return cached.ID, nil - } - return uuid.NewString(), nil -} - -func xaiExecutionSessionID(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) string { - if value := xaiMetadataString(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { - return value - } - if value := xaiMetadataString(req.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { - return value - } - if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() { - if value := strings.TrimSpace(promptCacheKey.String()); value != "" { - return value - } - } - return helps.DerivedSessionUUID("xai", opts.Metadata, req.Metadata) -} - -func xaiRequiresIsolatedConversation(model string) bool { - return strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), xaiComposerModelPrefix) -} - -func xaiImageEndpointPath(opts cliproxyexecutor.Options) string { - if opts.SourceFormat.String() != xaiImageHandlerType { - return "" - } - - path := xaiMetadataString(opts.Metadata, cliproxyexecutor.RequestPathMetadataKey) - if strings.HasSuffix(path, "/images/edits") { - return xaiImagesEditsPath - } - if strings.HasSuffix(path, "/images/generations") { - return xaiImagesGenerationsPath - } - return xaiDefaultImageEndpointPath -} - -// normalizeXAIImageRefs rewrites OpenAI-style image object fields to the xAI -// image API shape before the payload is sent upstream: -// -// {"image":{"image_url":"https://..."}} → {"image":{"url":"https://..."}} -// -// Applies to image / images / reference_images anywhere in the JSON tree, -// including nested objects and array items. Does not rewrite chat content -// parts shaped as {"type":"image_url","image_url":{...}}. -func normalizeXAIImageRefs(body []byte) []byte { - if !gjson.ValidBytes(body) { - return body - } - - decoder := json.NewDecoder(bytes.NewReader(body)) - decoder.UseNumber() - var payload any - if errDecode := decoder.Decode(&payload); errDecode != nil { - return body - } - - if !normalizeXAIImageRefsValue(payload) { - return body - } - normalized, errMarshal := json.Marshal(payload) - if errMarshal != nil { - return body - } - return normalized -} - -func normalizeXAIImageRefsValue(value any) bool { - changed := false - switch node := value.(type) { - case map[string]any: - for key, child := range node { - switch key { - case "image": - changed = normalizeXAIImageRef(child) || changed - case "images", "reference_images": - if refs, ok := child.([]any); ok { - for _, ref := range refs { - changed = normalizeXAIImageRef(ref) || changed - } - } - } - changed = normalizeXAIImageRefsValue(child) || changed - } - case []any: - for _, child := range node { - changed = normalizeXAIImageRefsValue(child) || changed - } - } - return changed -} - -func normalizeXAIImageRef(value any) bool { - ref, ok := value.(map[string]any) - if !ok { - return false - } - - originalURL, _ := ref["url"].(string) - url := strings.TrimSpace(originalURL) - imageURL, hasImageURL := ref["image_url"] - if url == "" { - switch imageURL := imageURL.(type) { - case string: - url = strings.TrimSpace(imageURL) - case map[string]any: - url, _ = imageURL["url"].(string) - url = strings.TrimSpace(url) - } - } - if url == "" { - return false - } - if url == originalURL && !hasImageURL { - return false - } - - // Always emit the xAI field name and drop the OpenAI alias. - ref["url"] = url - delete(ref, "image_url") - return true -} - -func xaiIsVideoRequest(opts cliproxyexecutor.Options) bool { - return opts.SourceFormat.String() == xaiVideoHandlerType -} - -func xaiVideoEndpointPath(opts cliproxyexecutor.Options) string { - if !xaiIsVideoRequest(opts) { - return "" - } - path := xaiMetadataString(opts.Metadata, cliproxyexecutor.RequestPathMetadataKey) - if strings.HasSuffix(path, "/videos/edits") { - return xaiVideosEditsPath - } - if strings.HasSuffix(path, "/videos/extensions") { - return xaiVideosExtensionsPath - } - if strings.HasSuffix(path, "/videos/generations") { - return xaiVideosGenerationsPath - } - return "" -} - -func xaiMetadataString(meta map[string]any, key string) string { - if len(meta) == 0 || key == "" { - return "" - } - value, ok := meta[key] - if !ok || value == nil { - return "" - } - switch typed := value.(type) { - case string: - return strings.TrimSpace(typed) - case fmt.Stringer: - return strings.TrimSpace(typed.String()) - default: - return strings.TrimSpace(fmt.Sprint(typed)) - } -} - -func preserveXAIResponsesOutputControls(body, source []byte, from sdktranslator.Format) []byte { - var maxOutputTokens gjson.Result - switch from { - case sdktranslator.FormatOpenAI: - maxOutputTokens = gjson.GetBytes(source, "max_completion_tokens") - if !maxOutputTokens.Exists() || maxOutputTokens.Type == gjson.Null { - maxOutputTokens = gjson.GetBytes(source, "max_tokens") - } - case sdktranslator.FormatOpenAIResponse: - maxOutputTokens = gjson.GetBytes(source, "max_output_tokens") - default: - return body - } - - if maxOutputTokens.Exists() && maxOutputTokens.Type != gjson.Null { - body, _ = sjson.SetRawBytes(body, "max_output_tokens", []byte(maxOutputTokens.Raw)) - } - for _, field := range []string{"temperature", "top_p", "top_k"} { - value := gjson.GetBytes(source, field) - if value.Exists() && value.Type != gjson.Null { - body, _ = sjson.SetRawBytes(body, field, []byte(value.Raw)) - } - } - return body -} - -func sanitizeXAIResponsesBody(body []byte, model string) []byte { - // stop is supported by Chat Completions but not by xAI's Responses API. - body, _ = sjson.DeleteBytes(body, "stop") - if !xaiSupportsReasoningEffort(model) { - if gjson.GetBytes(body, "reasoning.effort").Exists() { - log.Debugf("xai: stripping reasoning.effort for model %s (no thinking levels in model registry)", model) - } - body, _ = sjson.DeleteBytes(body, "reasoning.effort") - if reasoning := gjson.GetBytes(body, "reasoning"); reasoning.Exists() && reasoning.IsObject() && len(reasoning.Map()) == 0 { - body, _ = sjson.DeleteBytes(body, "reasoning") - } - } - return body -} - -// ensureXAINativeXSearchTool appends {"type":"x_search"} when the final tools -// list does not already include native X Search. When tool_choice restricts the -// model to allowed_tools, x_search is also added there (without duplicates) so -// Grok can select the injected tool. HTTP and websocket executors both prepare -// payloads through prepareResponsesRequestTo, so this runs once before the body -// is submitted upstream. -func ensureXAINativeXSearchTool(body []byte) []byte { - if !gjson.ValidBytes(body) { - return body - } - if !xaiRequestHasNativeXSearch(body) { - tools := gjson.GetBytes(body, "tools") - if !tools.Exists() || !tools.IsArray() { - body, _ = sjson.SetRawBytes(body, "tools", []byte(`[{"type":"x_search"}]`)) - } else { - body, _ = sjson.SetRawBytes(body, "tools.-1", xaiXSearchToolJSON) - } - } - return ensureXAINativeXSearchAllowedTools(body) -} - -// ensureXAINativeXSearchAllowedTools appends x_search to tool_choice.tools when -// the choice mode is allowed_tools and x_search is not already listed. -func ensureXAINativeXSearchAllowedTools(body []byte) []byte { - choice := gjson.GetBytes(body, "tool_choice") - if !choice.IsObject() || choice.Get("type").String() != "allowed_tools" { - return body - } - allowed := choice.Get("tools") - if !allowed.Exists() || !allowed.IsArray() { - body, _ = sjson.SetRawBytes(body, "tool_choice.tools", []byte(`[{"type":"x_search"}]`)) - return body - } - for _, tool := range allowed.Array() { - if strings.TrimSpace(tool.Get("type").String()) == xaiXSearchToolType { - return body - } - } - body, _ = sjson.SetRawBytes(body, "tool_choice.tools.-1", xaiXSearchToolJSON) - return body -} - -// pruneXAIOrphanedToolChoice removes tool_choice entries that no longer match -// any remaining tool after normalizeXAITools filtering. Forced choices that -// reference a deleted tool are dropped entirely; allowed_tools lists keep only -// choices that still resolve against the post-normalization tools set. -func pruneXAIOrphanedToolChoice(body []byte) []byte { - if !gjson.ValidBytes(body) { - return body - } - choice := gjson.GetBytes(body, "tool_choice") - if !choice.Exists() { - return body - } - available := collectXAIAvailableToolChoiceKeys(body) - if choice.Type == gjson.String { - // auto / none / required are not tool references. - return body - } - if !choice.IsObject() { - return body - } - choiceType := strings.TrimSpace(choice.Get("type").String()) - switch choiceType { - case "allowed_tools": - return pruneXAIAllowedToolsChoice(body, available) - default: - if choiceType == "" { - return body - } - if xaiToolChoiceMatchesAvailable(choice, available) { - return body - } - body, _ = sjson.DeleteBytes(body, "tool_choice") - return body - } -} - -func pruneXAIAllowedToolsChoice(body []byte, available map[xaiToolChoiceKey]struct{}) []byte { - allowed := gjson.GetBytes(body, "tool_choice.tools") - if !allowed.Exists() || !allowed.IsArray() { - body, _ = sjson.DeleteBytes(body, "tool_choice") - return body - } - allowedItems := allowed.Array() - filtered := make([][]byte, 0, len(allowedItems)) - changed := false - for _, tool := range allowedItems { - if !xaiToolChoiceMatchesAvailable(tool, available) { - changed = true - continue - } - filtered = append(filtered, []byte(tool.Raw)) - } - if !changed { - return body - } - if len(filtered) == 0 { - body, _ = sjson.DeleteBytes(body, "tool_choice") - return body - } - body, _ = sjson.SetRawBytes(body, "tool_choice.tools", helps.JoinRawJSONArray(filtered)) - return body -} - -// xaiToolChoiceKey identifies a selectable tool the way xAI tool_choice entries -// reference it after namespace qualification: type alone for host tools, or -// type+name for function tools. -type xaiToolChoiceKey struct { - toolType string - name string -} - -func collectXAIAvailableToolChoiceKeys(body []byte) map[xaiToolChoiceKey]struct{} { - keys := make(map[xaiToolChoiceKey]struct{}) - collect := func(tools gjson.Result) { - if !tools.IsArray() { - return - } - for _, tool := range tools.Array() { - toolType := strings.TrimSpace(tool.Get("type").String()) - if toolType == "" { - continue - } - key := xaiToolChoiceKey{toolType: toolType} - if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { - key.name = strings.TrimSpace(tool.Get("name").String()) - if key.name == "" { - continue - } - } - keys[key] = struct{}{} - } - } - collect(gjson.GetBytes(body, "tools")) - input := gjson.GetBytes(body, "input") - if input.IsArray() { - for _, item := range input.Array() { - if item.Get("type").String() == "additional_tools" { - collect(item.Get("tools")) - } - } - } - return keys -} - -func xaiToolChoiceMatchesAvailable(choice gjson.Result, available map[xaiToolChoiceKey]struct{}) bool { - toolType := strings.TrimSpace(choice.Get("type").String()) - if toolType == "" { - return false - } - key := xaiToolChoiceKey{toolType: toolType} - if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { - key.name = strings.TrimSpace(choice.Get("name").String()) - if key.name == "" { - return false - } - } - _, ok := available[key] - return ok -} - -func normalizeXAITools(body []byte) []byte { - if !gjson.ValidBytes(body) { - return body - } - original := body - normalizeAtPath := func(path string) bool { - tools := gjson.GetBytes(body, path) - if !tools.Exists() || !tools.IsArray() { - return true - } - filtered, changed, ok := normalizeXAIToolArray(tools) - if !ok { - return false - } - if !changed { - return true - } - updated, errSet := sjson.SetRawBytes(body, path, filtered) - if errSet != nil { - return false - } - body = updated - return true - } - - if !normalizeAtPath("tools") { - return original - } - input := gjson.GetBytes(body, "input") - if input.Exists() && input.IsArray() { - for index, item := range input.Array() { - if item.Get("type").String() != "additional_tools" { - continue - } - if !normalizeAtPath(fmt.Sprintf("input.%d.tools", index)) { - return original - } - } - } - return body -} - -// promoteXAIAdditionalTools moves Responses Lite tool declarations to the -// top-level tools array because xAI does not accept additional_tools input items. -func promoteXAIAdditionalTools(body []byte) []byte { - if !gjson.ValidBytes(body) { - return body - } - input := gjson.GetBytes(body, "input") - if !input.IsArray() { - return body - } - - inputItems := input.Array() - remainingInput := make([]json.RawMessage, 0, len(inputItems)) - promotedTools := make([]json.RawMessage, 0) - for _, item := range inputItems { - if item.Get("type").String() != "additional_tools" { - remainingInput = append(remainingInput, json.RawMessage(item.Raw)) - continue - } - for _, tool := range item.Get("tools").Array() { - promotedTools = append(promotedTools, json.RawMessage(tool.Raw)) - } - } - if len(remainingInput) == len(inputItems) { - return body - } - - rawInput, errMarshalInput := json.Marshal(remainingInput) - if errMarshalInput != nil { - return body - } - updated, errSetInput := sjson.SetRawBytes(body, "input", rawInput) - if errSetInput != nil { - return body - } - if len(promotedTools) == 0 { - return updated - } - - topLevelTools := gjson.GetBytes(updated, "tools") - tools := make([]json.RawMessage, 0, len(topLevelTools.Array())+len(promotedTools)) - if topLevelTools.IsArray() { - for _, tool := range topLevelTools.Array() { - tools = append(tools, json.RawMessage(tool.Raw)) - } - } - tools = append(tools, promotedTools...) - rawTools, errMarshalTools := json.Marshal(tools) - if errMarshalTools != nil { - return body - } - updated, errSetTools := sjson.SetRawBytes(updated, "tools", rawTools) - if errSetTools != nil { - return body - } - return updated -} - -func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) { - toolItems := tools.Array() - filtered := make([][]byte, 0, len(toolItems)) - changed := false - for _, tool := range toolItems { - toolType := tool.Get("type").String() - if toolType == xaiNamespaceToolType { - changed = true - namespaceName := tool.Get("name").String() - if namespaceTools := tool.Get("tools"); namespaceTools.IsArray() { - for _, nestedTool := range namespaceTools.Array() { - nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName) - if !ok { - return nil, false, false - } - changed = changed || nestedChanged - if len(nestedRaw) > 0 { - filtered = append(filtered, nestedRaw) - } - } - } - continue - } - raw, toolChanged, ok := normalizeXAITool(tool, "") - if !ok { - return nil, false, false - } - changed = changed || toolChanged - if len(raw) > 0 { - filtered = append(filtered, raw) - } - } - if !changed { - return nil, false, true - } - return helps.JoinRawJSONArray(filtered), true, true -} - -// normalizeXAIToolChoiceForTools drops tool_choice and parallel_tool_calls -// when tools are absent or empty (including after normalizeXAITools filtering). -// xAI rejects payloads that include tool_choice without any tools defined. -// Existence checks avoid unnecessary sjson parse/copy passes. -func normalizeXAIToolChoiceForTools(body []byte) []byte { - tools := gjson.GetBytes(body, "tools") - hasTools := tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 - if !hasTools { - input := gjson.GetBytes(body, "input") - if input.Exists() && input.IsArray() { - for _, item := range input.Array() { - additionalTools := item.Get("tools") - if item.Get("type").String() == "additional_tools" && additionalTools.IsArray() && len(additionalTools.Array()) > 0 { - hasTools = true - break - } - } - } - } - if hasTools { - return body - } - if tools.Exists() { - body, _ = sjson.DeleteBytes(body, "tools") - } - if gjson.GetBytes(body, "tool_choice").Exists() { - body, _ = sjson.DeleteBytes(body, "tool_choice") - } - if gjson.GetBytes(body, "parallel_tool_calls").Exists() { - body, _ = sjson.DeleteBytes(body, "parallel_tool_calls") - } - return body -} - -// normalizeXAINamespaceToolChoice qualifies namespaced function choices using -// the same names sent in the flattened tools list. xAI does not accept the -// Responses namespace field on tool choices. -func normalizeXAINamespaceToolChoice(body []byte) []byte { - if !gjson.ValidBytes(body) { - return body - } - original := body - normalizeAtPath := func(path string) bool { - toolChoice := gjson.GetBytes(body, path) - if !toolChoice.IsObject() || toolChoice.Get("type").String() != xaiFunctionToolType { - return true - } - namespaceName := strings.TrimSpace(toolChoice.Get("namespace").String()) - toolName := strings.TrimSpace(toolChoice.Get("name").String()) - qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) - if namespaceName == "" || qualifiedName == "" { - return true - } - updated, errSet := sjson.SetBytes(body, path+".name", qualifiedName) - if errSet != nil { - return false - } - updated, errDelete := sjson.DeleteBytes(updated, path+".namespace") - if errDelete != nil { - return false - } - body = updated - return true - } - - if !normalizeAtPath("tool_choice") { - return original - } - tools := gjson.GetBytes(body, "tool_choice.tools") - if tools.IsArray() { - for index := range tools.Array() { - if !normalizeAtPath(fmt.Sprintf("tool_choice.tools.%d", index)) { - return original - } - } - } - return body -} - -func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) { - toolType := tool.Get("type").String() - changed := false - if toolType == xaiToolSearchType || toolType == xaiImageGenerationToolType { - return nil, true, true - } - if toolType == xaiCustomToolType && tool.Get("name").String() == "apply_patch" { - return nil, true, true - } - - raw := []byte(tool.Raw) - schemaTool := tool - if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { - updatedTool, schemaChanged, ok := normalizeXAIObjectRootUnionBranchTypes(raw) - if !ok { - return nil, false, false - } - raw = updatedTool - if schemaChanged { - schemaTool = gjson.ParseBytes(raw) - changed = true - log.Debugf("xai: added object types to root union branches for tool %s.%s", namespaceName, tool.Get("name").String()) - } - } - if toolType == xaiCustomToolType { - updatedTool, errSet := sjson.SetBytes(raw, "type", xaiFunctionToolType) - if errSet != nil { - return nil, false, false - } - raw = updatedTool - toolType = xaiFunctionToolType - changed = true - } - if toolType == xaiWebSearchToolType && tool.Get("external_web_access").Exists() { - updatedTool, errDel := sjson.DeleteBytes(raw, "external_web_access") - if errDel != nil { - return nil, false, false - } - raw = updatedTool - changed = true - } - if toolType == xaiFunctionToolType && !schemaTool.Get("parameters").Exists() { - updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(`{"type":"object","properties":{}}`)) - if errSet != nil { - return nil, false, false - } - raw = updatedTool - changed = true - } - // Simplify the Codex Desktop automation schema and root unions that xAI - // rejects because function parameters must resolve exclusively to objects. - if toolType == xaiFunctionToolType && xaiFunctionParametersNeedSimplification(schemaTool, namespaceName) { - updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(xaiSafeFunctionParameters)) - if errSet != nil { - return nil, false, false - } - raw = updatedTool - if strict := tool.Get("strict"); strict.Exists() && strict.Bool() { - updatedTool, errSet = sjson.SetBytes(raw, "strict", false) - if errSet != nil { - return nil, false, false - } - raw = updatedTool - } - changed = true - log.Debugf("xai: simplified parameters for tool %s.%s to avoid upstream schema rejection or hang", namespaceName, tool.Get("name").String()) - } - if toolType == xaiFunctionToolType && strings.TrimSpace(namespaceName) != "" { - qualifiedName := qualifyXAINamespaceToolName(namespaceName, tool.Get("name").String()) - if qualifiedName == "" { - return nil, false, false - } - updatedTool, errSet := sjson.SetBytes(raw, "name", qualifiedName) - if errSet != nil { - return nil, false, false - } - raw = updatedTool - changed = true - } - return raw, changed, true -} - -func qualifyXAINamespaceToolName(namespaceName, toolName string) string { - namespaceName = strings.TrimSpace(namespaceName) - toolName = strings.TrimSpace(toolName) - if namespaceName == "" || toolName == "" || strings.HasPrefix(toolName, "mcp__") { - return toolName - } - prefix := namespaceName - if !strings.HasSuffix(prefix, "__") { - prefix += "__" - } - if strings.HasPrefix(toolName, prefix) { - return toolName - } - return prefix + toolName -} - -func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { - refs := make(map[string]xaiNamespaceToolRef) - collect := func(tools gjson.Result) { - if !tools.Exists() || !tools.IsArray() { - return - } - for _, tool := range tools.Array() { - if tool.Get("type").String() != xaiNamespaceToolType { - continue - } - namespaceName := strings.TrimSpace(tool.Get("name").String()) - if namespaceName == "" { - continue - } - for _, nestedTool := range tool.Get("tools").Array() { - toolName := strings.TrimSpace(nestedTool.Get("name").String()) - qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) - if qualifiedName == "" { - continue - } - refs[qualifiedName] = xaiNamespaceToolRef{namespace: namespaceName, name: toolName} - } - } - } - collect(gjson.GetBytes(body, "tools")) - input := gjson.GetBytes(body, "input") - if input.Exists() && input.IsArray() { - for _, item := range input.Array() { - if item.Get("type").String() == "additional_tools" { - collect(item.Get("tools")) - } - } - } - return refs -} - -func normalizeXAIInputCustomToolCalls(body []byte) []byte { - input := gjson.GetBytes(body, "input") - if !input.Exists() || !input.IsArray() { - return body - } - - changed := false - inputArray := input.Array() - items := make([]json.RawMessage, 0, len(inputArray)) - for _, item := range inputArray { - var normalized []byte - switch item.Get("type").String() { - case "custom_tool_call": - callID := strings.TrimSpace(item.Get("call_id").String()) - name := strings.TrimSpace(item.Get("name").String()) - if callID == "" || name == "" { - changed = true - continue - } - normalized = []byte(`{"type":"function_call"}`) - normalized, _ = sjson.SetBytes(normalized, "call_id", callID) - normalized, _ = sjson.SetBytes(normalized, "name", name) - normalized, _ = sjson.SetBytes(normalized, "arguments", xaiCustomToolCallArguments(item.Get("input"))) - case "custom_tool_call_output": - callID := strings.TrimSpace(item.Get("call_id").String()) - if callID == "" { - changed = true - continue - } - normalized = []byte(`{"type":"function_call_output"}`) - normalized, _ = sjson.SetBytes(normalized, "call_id", callID) - normalized, _ = sjson.SetBytes(normalized, "output", xaiCustomToolCallOutput(item.Get("output"))) - default: - items = append(items, json.RawMessage(item.Raw)) - continue - } - items = append(items, json.RawMessage(normalized)) - changed = true - } - if !changed { - return body - } - - rawInput, errMarshal := json.Marshal(items) - if errMarshal != nil { - return body - } - updated, errSet := sjson.SetRawBytes(body, "input", rawInput) - if errSet != nil { - return body - } - return updated -} - -func xaiCustomToolCallArguments(input gjson.Result) string { - if !input.Exists() { - return "{}" - } - if input.Type == gjson.String { - text := input.String() - trimmed := strings.TrimSpace(text) - if gjson.Valid(trimmed) { - parsed := gjson.Parse(trimmed) - if parsed.IsObject() { - return parsed.Raw - } - } - encoded, errMarshal := json.Marshal(text) - if errMarshal != nil { - return "{}" - } - return `{"input":` + string(encoded) + `}` - } - if input.IsObject() { - return input.Raw - } - if input.Raw != "" { - return `{"input":` + input.Raw + `}` - } - return "{}" -} - -func xaiCustomToolCallOutput(output gjson.Result) string { - if !output.Exists() { - return "" - } - if output.Type == gjson.String { - return output.String() - } - return output.Raw -} - -// xAI executes these x_search subtools server-side but exposes their trace as -// client-style tool calls. Hide the trace so Responses clients do not execute it again. -type xaiInternalXSearchResponseFilter struct { - enabled bool - clientDeclaredTools map[xaiClientToolKey]struct{} - droppedOutputIndexes map[int64]struct{} - droppedItemIDs map[string]struct{} -} - -func newXAIInternalXSearchResponseFilter(enabled bool, clientDeclaredTools map[xaiClientToolKey]struct{}) *xaiInternalXSearchResponseFilter { - filter := &xaiInternalXSearchResponseFilter{ - enabled: enabled, - clientDeclaredTools: clientDeclaredTools, - } - if enabled { - filter.droppedOutputIndexes = make(map[int64]struct{}) - filter.droppedItemIDs = make(map[string]struct{}) - } - return filter -} - -func xaiRequestHasNativeXSearch(body []byte) bool { - if gjson.GetBytes(body, `tools.#(type=="x_search")`).Exists() { - return true - } - // Multipath queries return an array of matches; an empty array still Exists(). - // Check the match count instead of Exists() for additional_tools injection. - return len(gjson.GetBytes(body, `input.#(type=="additional_tools")#.tools.#(type=="x_search")`).Array()) > 0 -} - -// collectXAIClientDeclaredToolKeys records client-declared function/custom tools -// using the Responses post-restore identity (short name + optional namespace) and -// the effective upstream tool type after normalizeXAITool. Client custom tools -// are normalized to function before being sent to xAI, so keys use function for -// both declaration kinds. Must run before normalizeXAITools flattens namespace wrappers. -func collectXAIClientDeclaredToolKeys(body []byte) map[xaiClientToolKey]struct{} { - keys := make(map[xaiClientToolKey]struct{}) - collect := func(tools gjson.Result) { - if !tools.Exists() || !tools.IsArray() { - return - } - for _, tool := range tools.Array() { - switch toolType := strings.TrimSpace(tool.Get("type").String()); toolType { - case xaiNamespaceToolType: - namespaceName := strings.TrimSpace(tool.Get("name").String()) - if namespaceName == "" { - continue - } - for _, nestedTool := range tool.Get("tools").Array() { - nestedType := strings.TrimSpace(nestedTool.Get("type").String()) - if nestedType != xaiFunctionToolType && nestedType != xaiCustomToolType { - continue - } - toolName := strings.TrimSpace(nestedTool.Get("name").String()) - if toolName == "" { - continue - } - // normalizeXAITool converts custom → function before upstream send. - keys[xaiClientToolKey{namespace: namespaceName, name: toolName, toolType: xaiEffectiveDeclaredToolType(nestedType)}] = struct{}{} - } - case xaiFunctionToolType, xaiCustomToolType: - toolName := strings.TrimSpace(tool.Get("name").String()) - if toolName == "" { - continue - } - // normalizeXAITool converts custom → function before upstream send. - keys[xaiClientToolKey{namespace: "", name: toolName, toolType: xaiEffectiveDeclaredToolType(toolType)}] = struct{}{} - } - } - } - collect(gjson.GetBytes(body, "tools")) - input := gjson.GetBytes(body, "input") - if input.Exists() && input.IsArray() { - for _, item := range input.Array() { - if item.Get("type").String() == "additional_tools" { - collect(item.Get("tools")) - } - } - } - return keys -} - -// xaiEffectiveDeclaredToolType returns the tool type actually sent upstream -// after normalizeXAITool. Client custom tools are rewritten to function. -func xaiEffectiveDeclaredToolType(toolType string) string { - if strings.TrimSpace(toolType) == xaiCustomToolType { - return xaiFunctionToolType - } - return strings.TrimSpace(toolType) -} - -func xaiIsInternalXSearchToolName(name string) bool { - switch strings.TrimSpace(name) { - case "x_user_search", "x_semantic_search", "x_keyword_search", "x_thread_fetch": - return true - default: - return false - } -} - -// xaiResponseCallDeclaredType maps a Responses output call type to the effective -// upstream tool declaration kind used when matching client-declared tools. -// Client custom tools are normalized to function before upstream send, so only -// function_call can match a client-declared same-name tool; custom_tool_call -// remains the internal X Search trace shape. -func xaiResponseCallDeclaredType(itemType string) string { - switch strings.TrimSpace(itemType) { - case "function_call": - return xaiFunctionToolType - case "custom_tool_call": - return xaiCustomToolType - default: - return "" - } -} - -// xaiIsInternalXSearchCallID reports whether call_id matches the evidenced xAI -// X Search server-side trace prefix (xs_call...), as observed in Responses traffic -// for native x_search subtools (see issue #4282 / PR #4284 fixtures). -func xaiIsInternalXSearchCallID(callID string) bool { - return strings.HasPrefix(strings.TrimSpace(callID), "xs_call") -} - -// xaiIsInternalXSearchCall reports whether an output item is an xAI server-side -// X Search subtool trace that should be hidden from Responses clients. -// -// Evidence from xAI Responses traffic (issue #4282 / PR #4284): -// - native x_search subtools are emitted as custom_tool_call items named -// x_user_search / x_semantic_search / x_keyword_search / x_thread_fetch -// - those traces commonly use call_id values prefixed with "xs_call" -// -// Client tools that share a short name are preserved only when the response call -// kind matches the effective upstream declaration type. Because normalizeXAITool -// rewrites client custom → function, a client custom x_keyword_search is keyed as -// function and therefore preserves function_call while still filtering genuine -// internal custom_tool_call / xs_call* traces. Namespaced restored client tools -// are never treated as internal. -func xaiIsInternalXSearchCall(item gjson.Result, clientDeclaredTools map[xaiClientToolKey]struct{}) bool { - itemType := strings.TrimSpace(item.Get("type").String()) - declaredType := xaiResponseCallDeclaredType(itemType) - if declaredType == "" { - return false - } - name := strings.TrimSpace(item.Get("name").String()) - if !xaiIsInternalXSearchToolName(name) { - return false - } - namespace := strings.TrimSpace(item.Get("namespace").String()) - // Namespaced calls are restored client tools, never xAI internal X Search traces. - if namespace != "" { - return false - } - // Evidenced internal call_id prefix always identifies server-side X Search traces, - // even when a client tool reuses the same short name. - if xaiIsInternalXSearchCallID(item.Get("call_id").String()) { - return true - } - // Preserve only client tools whose effective upstream declaration kind matches - // this call type (function_call ↔ function after custom normalization). - if _, declared := clientDeclaredTools[xaiClientToolKey{namespace: namespace, name: name, toolType: declaredType}]; declared { - return false - } - return true -} - -func (f *xaiInternalXSearchResponseFilter) apply(eventData []byte) []byte { - if f == nil || !f.enabled || len(eventData) == 0 || !gjson.ValidBytes(eventData) { - return eventData - } - - if item := gjson.GetBytes(eventData, "item"); xaiIsInternalXSearchCall(item, f.clientDeclaredTools) { - f.recordDroppedItem(eventData, item) - return nil - } - - eventData = f.filterCompletedOutput(eventData) - if f.referencesDroppedItem(eventData) { - return nil - } - return f.compactOutputIndex(eventData) -} - -func (f *xaiInternalXSearchResponseFilter) recordDroppedItem(eventData []byte, item gjson.Result) { - if outputIndex := gjson.GetBytes(eventData, "output_index"); outputIndex.Exists() { - f.droppedOutputIndexes[outputIndex.Int()] = struct{}{} - } - for _, path := range []string{"id", "call_id"} { - if id := strings.TrimSpace(item.Get(path).String()); id != "" { - f.droppedItemIDs[id] = struct{}{} - } - } -} - -func (f *xaiInternalXSearchResponseFilter) referencesDroppedItem(eventData []byte) bool { - if outputIndex := gjson.GetBytes(eventData, "output_index"); outputIndex.Exists() { - if _, dropped := f.droppedOutputIndexes[outputIndex.Int()]; dropped { - return true - } - } - for _, path := range []string{"item_id", "call_id"} { - id := strings.TrimSpace(gjson.GetBytes(eventData, path).String()) - if _, dropped := f.droppedItemIDs[id]; id != "" && dropped { - return true - } - } - return false -} - -func (f *xaiInternalXSearchResponseFilter) compactOutputIndex(eventData []byte) []byte { - outputIndex := gjson.GetBytes(eventData, "output_index") - if !outputIndex.Exists() { - return eventData - } - original := outputIndex.Int() - removedBefore := int64(0) - for dropped := range f.droppedOutputIndexes { - if dropped < original { - removedBefore++ - } - } - if removedBefore == 0 { - return eventData - } - updated, errSet := sjson.SetBytes(eventData, "output_index", original-removedBefore) - if errSet != nil { - return eventData - } - return updated -} - -func (f *xaiInternalXSearchResponseFilter) filterCompletedOutput(eventData []byte) []byte { - output := gjson.GetBytes(eventData, "response.output") - if !output.IsArray() { - return eventData - } - var clientDeclaredTools map[xaiClientToolKey]struct{} - if f != nil { - clientDeclaredTools = f.clientDeclaredTools - } - items := make([]json.RawMessage, 0, len(output.Array())) - changed := false - for _, item := range output.Array() { - if xaiIsInternalXSearchCall(item, clientDeclaredTools) { - changed = true - continue - } - items = append(items, json.RawMessage(item.Raw)) - } - if !changed { - return eventData - } - rawOutput, errMarshal := json.Marshal(items) - if errMarshal != nil { - return eventData - } - updated, errSet := sjson.SetRawBytes(eventData, "response.output", rawOutput) - if errSet != nil { - return eventData - } - return updated -} - -func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { - if !gjson.ValidBytes(body) { - return body - } - input := gjson.GetBytes(body, "input") - if !input.Exists() || !input.IsArray() { - return body - } - for index, item := range input.Array() { - if item.Get("type").String() != "function_call" { - continue - } - namespaceName := strings.TrimSpace(item.Get("namespace").String()) - toolName := strings.TrimSpace(item.Get("name").String()) - qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) - if namespaceName == "" || qualifiedName == "" { - continue - } - namePath := fmt.Sprintf("input.%d.name", index) - namespacePath := fmt.Sprintf("input.%d.namespace", index) - updated, errSet := sjson.SetBytes(body, namePath, qualifiedName) - if errSet != nil { - continue - } - updated, errDelete := sjson.DeleteBytes(updated, namespacePath) - if errDelete != nil { - continue - } - body = updated - } - return body -} - -func restoreXAINamespaceToolCalls(data []byte, refs map[string]xaiNamespaceToolRef) []byte { - if len(refs) == 0 || len(data) == 0 || !gjson.ValidBytes(data) { - return data - } - data = restoreXAINamespaceToolCallAtPath(data, "item", refs) - output := gjson.GetBytes(data, "response.output") - if output.Exists() && output.IsArray() { - for index := range output.Array() { - data = restoreXAINamespaceToolCallAtPath(data, fmt.Sprintf("response.output.%d", index), refs) - } - } - return data -} - -func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string]xaiNamespaceToolRef) []byte { - if gjson.GetBytes(data, path+".type").String() != "function_call" { - return data - } - qualifiedName := strings.TrimSpace(gjson.GetBytes(data, path+".name").String()) - ref, ok := refs[qualifiedName] - if !ok { - return data - } - updated, errSet := sjson.SetBytes(data, path+".name", ref.name) - if errSet != nil { - return data - } - updated, errSet = sjson.SetBytes(updated, path+".namespace", ref.namespace) - if errSet != nil { - return data - } - return updated -} - -// normalizeXAIObjectRootUnionBranchTypes makes untyped root union branches -// explicitly object-only when the parameter root already permits only objects. -// This preserves the original schema semantics while satisfying xAI validation. -func normalizeXAIObjectRootUnionBranchTypes(tool []byte) ([]byte, bool, bool) { - parameters := gjson.GetBytes(tool, "parameters") - rootType := parameters.Get("type") - if rootType.Type != gjson.String || rootType.String() != "object" { - return tool, false, true - } - - original := tool - changed := false - for _, unionName := range []string{"anyOf", "oneOf"} { - union := parameters.Get(unionName) - if !union.IsArray() { - continue - } - for index, branch := range union.Array() { - if !branch.IsObject() || branch.Get("type").Exists() { - continue - } - updated, errSet := sjson.SetBytes(tool, fmt.Sprintf("parameters.%s.%d.type", unionName, index), "object") - if errSet != nil { - return original, false, false - } - tool = updated - changed = true - } - } - return tool, changed, true -} - -func xaiSchemaTypeIsObjectOnly(schemaType gjson.Result) bool { - if schemaType.Type == gjson.String { - return strings.EqualFold(strings.TrimSpace(schemaType.String()), "object") - } - if !schemaType.IsArray() { - return false - } - types := schemaType.Array() - if len(types) == 0 { - return false - } - for _, schemaTypeItem := range types { - if schemaTypeItem.Type != gjson.String || !strings.EqualFold(strings.TrimSpace(schemaTypeItem.String()), "object") { - return false - } - } - return true -} - -// xaiFunctionParametersNeedSimplification reports whether a function tool, or -// a custom tool normalized to a function, has a schema that xAI cannot accept. -func xaiFunctionParametersNeedSimplification(tool gjson.Result, namespaceName string) bool { - toolType := strings.TrimSpace(tool.Get("type").String()) - isFunction := strings.EqualFold(toolType, xaiFunctionToolType) - isNormalizedCustom := strings.EqualFold(toolType, xaiCustomToolType) - if !isFunction && !isNormalizedCustom { - return false - } - - toolName := strings.TrimSpace(tool.Get("name").String()) - qualifiedAutomationName := xaiCodexAppNamespaceName + "__" + xaiAutomationUpdateToolName - if isFunction && (strings.EqualFold(toolName, qualifiedAutomationName) || - (strings.EqualFold(strings.TrimSpace(namespaceName), xaiCodexAppNamespaceName) && - strings.EqualFold(toolName, xaiAutomationUpdateToolName))) { - return true - } - - parameters := tool.Get("parameters") - for _, unionName := range []string{"anyOf", "oneOf"} { - union := parameters.Get(unionName) - if !union.IsArray() { - continue - } - for _, branch := range union.Array() { - if !xaiSchemaTypeIsObjectOnly(branch.Get("type")) { - return true - } - } - } - return false -} - -func sanitizeXAIInputEncryptedContent(body []byte) []byte { - input := gjson.GetBytes(body, "input") - if !input.Exists() || !input.IsArray() { - return body - } - items := make([]json.RawMessage, 0, len(input.Array())) - changed := false - dropCount := 0 - firstReason := "" - firstItemType := "" - for _, item := range input.Array() { - itemType := strings.TrimSpace(item.Get("type").String()) - if itemType != "reasoning" && itemType != "compaction" { - items = append(items, json.RawMessage(item.Raw)) - continue - } - encryptedContent := item.Get("encrypted_content") - if !encryptedContent.Exists() { - items = append(items, json.RawMessage(item.Raw)) - continue - } - reason := "" - switch encryptedContent.Type { - case gjson.String: - if _, err := signature.InspectGrokEncryptedContent(encryptedContent.String()); err != nil { - reason = err.Error() - } - case gjson.Null: - reason = "encrypted_content is null" - default: - reason = fmt.Sprintf("encrypted_content must be a string, got %s", encryptedContent.Type.String()) - } - if reason == "" { - items = append(items, json.RawMessage(item.Raw)) - continue - } - - if itemType == "compaction" { - changed = true - dropCount++ - if firstReason == "" { - firstReason = reason - firstItemType = itemType - } - continue - } - - next, err := sjson.DeleteBytes([]byte(item.Raw), "encrypted_content") - if err != nil { - items = append(items, json.RawMessage(item.Raw)) - continue - } - items = append(items, json.RawMessage(next)) - changed = true - dropCount++ - if firstReason == "" { - firstReason = reason - firstItemType = itemType - } - } - if !changed { - return body - } - rawInput, err := json.Marshal(items) - if err != nil { - return body - } - updated, err := sjson.SetRawBytes(body, "input", rawInput) - if err != nil { - return body - } - if dropCount > 0 { - log.WithFields(log.Fields{ - "component": "xai_encrypted_content_sanitizer", - "dropped": dropCount, - "first_item_type": firstItemType, - "first_reason": firstReason, - }).Debug("xai executor: removed invalid encrypted_content before upstream") - } - return mergeAdjacentXAIInputReasoningSummaries(updated) -} - -func normalizeXAIInputReasoningItems(body []byte) []byte { - input := gjson.GetBytes(body, "input") - if !input.Exists() || !input.IsArray() { - return body - } - - updated := body - for i, item := range input.Array() { - if item.Get("type").String() != "reasoning" { - continue - } - contentPath := fmt.Sprintf("input.%d.content", i) - if content := gjson.GetBytes(updated, contentPath); content.Exists() && content.Type == gjson.Null { - updatedBody, errDel := sjson.DeleteBytes(updated, contentPath) - if errDel != nil { - return body - } - updated = updatedBody - } - encryptedContentPath := fmt.Sprintf("input.%d.encrypted_content", i) - if encryptedContent := gjson.GetBytes(updated, encryptedContentPath); encryptedContent.Exists() && encryptedContent.Type == gjson.Null { - updatedBody, errDel := sjson.DeleteBytes(updated, encryptedContentPath) - if errDel != nil { - return body - } - updated = updatedBody - } - } - return mergeAdjacentXAIInputReasoningSummaries(updated) -} - -func mergeAdjacentXAIInputReasoningSummaries(body []byte) []byte { - input := gjson.GetBytes(body, "input") - if !input.Exists() || !input.IsArray() { - return body - } - - changed := false - items := make([]json.RawMessage, 0, len(input.Array())) - for _, item := range input.Array() { - if len(items) > 0 && canMergeXAIReasoningSummary(items[len(items)-1], item) { - merged, ok := appendXAIReasoningSummary(items[len(items)-1], item.Get("summary").Array()) - if ok { - items[len(items)-1] = json.RawMessage(merged) - changed = true - continue - } - } - items = append(items, json.RawMessage(item.Raw)) - } - if !changed { - return body - } - - rawInput, errMarshal := json.Marshal(items) - if errMarshal != nil { - return body - } - updated, errSet := sjson.SetRawBytes(body, "input", rawInput) - if errSet != nil { - return body - } - return updated -} - -func canMergeXAIReasoningSummary(previous json.RawMessage, current gjson.Result) bool { - previousItem := gjson.ParseBytes(previous) - if previousItem.Get("type").String() != "reasoning" || current.Get("type").String() != "reasoning" { - return false - } - if !previousItem.Get("summary").IsArray() || !current.Get("summary").IsArray() { - return false - } - if len(current.Get("summary").Array()) == 0 { - return false - } - for name := range current.Map() { - if name != "type" && name != "summary" { - return false - } - } - return true -} - -func appendXAIReasoningSummary(previous json.RawMessage, currentSummary []gjson.Result) ([]byte, bool) { - updated := []byte(previous) - summary := gjson.GetBytes(updated, "summary") - if !summary.IsArray() { - return previous, false - } - nextIndex := len(summary.Array()) - for i, item := range currentSummary { - updatedItem, errSet := sjson.SetRawBytes(updated, fmt.Sprintf("summary.%d", nextIndex+i), []byte(item.Raw)) - if errSet != nil { - return previous, false - } - updated = updatedItem - } - return updated, true -} - -// xaiSupportsReasoningEffort reports whether the model accepts Responses API -// reasoning.effort. Capability comes from model registry thinking metadata -// (static models.json and dynamic registrations), not a hard-coded name allowlist. -func xaiSupportsReasoningEffort(model string) bool { - name := strings.ToLower(strings.TrimSpace(thinking.ParseSuffix(model).ModelName)) - if idx := strings.LastIndex(name, "/"); idx >= 0 { - name = name[idx+1:] - } - if name == "" { - return false - } - info := registry.LookupModelInfo(name, "xai") - if info == nil || info.Thinking == nil { - return false - } - return len(info.Thinking.Levels) > 0 -} - -func xaiNormalizeReasoningSummaryEventLine(line []byte, eventName string) []byte { - if eventName == "" && bytes.HasPrefix(line, xaiEventTag) { - eventName = strings.TrimSpace(string(line[len(xaiEventTag):])) - } - eventName = xaiNormalizeReasoningSummaryEventName(eventName) - if eventName == "" { - return bytes.Clone(line) - } - return []byte("event: " + eventName) -} - -func xaiNormalizeReasoningSummaryEventName(eventName string) string { - switch eventName { - case "response.reasoning_text.delta": - return "response.reasoning_summary_text.delta" - case "response.reasoning_text.done": - return "response.reasoning_summary_part.done" - default: - return eventName - } -} - -func xaiNormalizeReasoningSummaryData(eventData []byte) []byte { - if len(eventData) == 0 || !gjson.ValidBytes(eventData) { - return eventData - } - - normalized := eventData - switch gjson.GetBytes(normalized, "type").String() { - case "response.reasoning_text.delta": - normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_text.delta") - normalized = xaiNormalizeReasoningSummaryIndex(normalized) - case "response.reasoning_text.done": - normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_part.done") - normalized, _ = sjson.SetBytes(normalized, "part.type", "summary_text") - if text := gjson.GetBytes(normalized, "text"); text.Exists() { - normalized, _ = sjson.SetBytes(normalized, "part.text", text.String()) - } - normalized, _ = sjson.DeleteBytes(normalized, "text") - normalized = xaiNormalizeReasoningSummaryIndex(normalized) - case "response.content_part.added": - if gjson.GetBytes(normalized, "part.type").String() == "reasoning_text" { - normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_part.added") - normalized, _ = sjson.SetBytes(normalized, "part.type", "summary_text") - normalized = xaiNormalizeReasoningSummaryIndex(normalized) - } - case "response.content_part.done": - if gjson.GetBytes(normalized, "part.type").String() == "reasoning_text" { - normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_part.done") - normalized, _ = sjson.SetBytes(normalized, "part.type", "summary_text") - normalized = xaiNormalizeReasoningSummaryIndex(normalized) - } - } - - if item := gjson.GetBytes(normalized, "item"); item.Exists() && item.Type == gjson.JSON { - updatedItem := xaiNormalizeReasoningOutputItem([]byte(item.Raw)) - if !bytes.Equal(updatedItem, []byte(item.Raw)) { - normalized, _ = sjson.SetRawBytes(normalized, "item", updatedItem) - } - } - if output := gjson.GetBytes(normalized, "response.output"); output.IsArray() { - updatedOutput, changed := xaiNormalizeReasoningOutputItems(output.Array()) - if changed { - normalized, _ = sjson.SetRawBytes(normalized, "response.output", updatedOutput) - } - } - - return normalized -} - -func xaiNormalizeReasoningSummaryDataEvents(eventData []byte) [][]byte { - if len(eventData) == 0 || !gjson.ValidBytes(eventData) { - return [][]byte{eventData} - } - if gjson.GetBytes(eventData, "type").String() != "response.reasoning_text.done" { - return [][]byte{xaiNormalizeReasoningSummaryData(eventData)} - } - - textDone, _ := sjson.SetBytes(eventData, "type", "response.reasoning_summary_text.done") - textDone = xaiNormalizeReasoningSummaryIndex(textDone) - partDone := xaiNormalizeReasoningSummaryData(eventData) - return [][]byte{textDone, partDone} -} - -func xaiNormalizeReasoningSummaryIndex(eventData []byte) []byte { - contentIndex := gjson.GetBytes(eventData, "content_index") - if contentIndex.Exists() && contentIndex.Raw != "" && !gjson.GetBytes(eventData, "summary_index").Exists() { - eventData, _ = sjson.SetRawBytes(eventData, "summary_index", []byte(contentIndex.Raw)) - } - eventData, _ = sjson.DeleteBytes(eventData, "content_index") - return eventData -} - -func xaiNormalizeReasoningOutputItems(items []gjson.Result) ([]byte, bool) { - var buf bytes.Buffer - buf.WriteByte('[') - changed := false - for i, item := range items { - if i > 0 { - buf.WriteByte(',') - } - updatedItem := xaiNormalizeReasoningOutputItem([]byte(item.Raw)) - if !bytes.Equal(updatedItem, []byte(item.Raw)) { - changed = true - } - buf.Write(updatedItem) - } - buf.WriteByte(']') - return buf.Bytes(), changed -} - -func xaiNormalizeReasoningOutputItem(item []byte) []byte { - if !gjson.ValidBytes(item) || gjson.GetBytes(item, "type").String() != "reasoning" { - return item - } - - normalized := item - if summary := gjson.GetBytes(normalized, "summary"); summary.IsArray() { - updatedSummary, changed := xaiNormalizeReasoningSummaryItems(summary.Array()) - if changed { - normalized, _ = sjson.SetRawBytes(normalized, "summary", updatedSummary) - } - } - - content := gjson.GetBytes(normalized, "content") - if !content.IsArray() { - return normalized - } - - summaryItems := make([]gjson.Result, 0, len(content.Array())) - for _, part := range content.Array() { - if part.Get("type").String() == "reasoning_text" { - summaryItems = append(summaryItems, part) - } - } - if len(summaryItems) == 0 { - return normalized - } - - updatedSummary, _ := xaiNormalizeReasoningSummaryItems(summaryItems) - normalized, _ = sjson.SetRawBytes(normalized, "summary", updatedSummary) - normalized, _ = sjson.DeleteBytes(normalized, "content") - return normalized -} - -func xaiNormalizeReasoningSummaryItems(items []gjson.Result) ([]byte, bool) { - var buf bytes.Buffer - buf.WriteByte('[') - changed := false - for i, item := range items { - if i > 0 { - buf.WriteByte(',') - } - itemRaw := []byte(item.Raw) - if item.Get("type").String() == "reasoning_text" { - var errSet error - itemRaw, errSet = sjson.SetBytes(itemRaw, "type", "summary_text") - if errSet == nil { - changed = true - } - } - buf.Write(itemRaw) - } - buf.WriteByte(']') - return buf.Bytes(), changed -} - -func xaiCollectOutputItemDone(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback *[][]byte) { - itemResult := gjson.GetBytes(eventData, "item") - if !itemResult.Exists() || itemResult.Type != gjson.JSON { - return - } - outputIndexResult := gjson.GetBytes(eventData, "output_index") - if outputIndexResult.Exists() { - outputItemsByIndex[outputIndexResult.Int()] = []byte(itemResult.Raw) - return - } - *outputItemsFallback = append(*outputItemsFallback, []byte(itemResult.Raw)) -} - -func xaiPatchCompletedOutput(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { - outputResult := gjson.GetBytes(eventData, "response.output") - shouldPatchOutput := (!outputResult.Exists() || !outputResult.IsArray() || len(outputResult.Array()) == 0) && (len(outputItemsByIndex) > 0 || len(outputItemsFallback) > 0) - if !shouldPatchOutput { - return eventData - } - - indexes := make([]int64, 0, len(outputItemsByIndex)) - for idx := range outputItemsByIndex { - indexes = append(indexes, idx) - } - sort.Slice(indexes, func(i, j int) bool { - return indexes[i] < indexes[j] - }) - - outputArray := []byte("[]") - var buf bytes.Buffer - buf.WriteByte('[') - wrote := false - for _, idx := range indexes { - if wrote { - buf.WriteByte(',') - } - buf.Write(outputItemsByIndex[idx]) - wrote = true - } - for _, item := range outputItemsFallback { - if wrote { - buf.WriteByte(',') - } - buf.Write(item) - wrote = true - } - buf.WriteByte(']') - if wrote { - outputArray = buf.Bytes() - } - - patched, _ := sjson.SetRawBytes(eventData, "response.output", outputArray) - return patched -} - -// xaiFreeUsageExhaustedCooldown is the free-tier rolling window advertised by -// cli-chat-proxy ("Usage resets over a rolling 24-hour window"). -const xaiFreeUsageExhaustedCooldown = 24 * time.Hour - -// xaiStatusErr wraps upstream error bodies so free-tier exhaustion -// (subscription:free-usage-exhausted) carries a 24h RetryAfter hint for -// auth cooldown / account rotation. Generic 429s stay without an explicit -// retry hint so conductor backoff still applies. -func xaiStatusErr(code int, body []byte) statusErr { - err := statusErr{code: code, msg: string(body)} - if code != http.StatusTooManyRequests || len(body) == 0 { - return err - } - codeStr := strings.ToLower(gjson.GetBytes(body, "code").String()) - msg := strings.ToLower(gjson.GetBytes(body, "error").String()) - if msg == "" { - msg = strings.ToLower(string(body)) - } - if strings.Contains(codeStr, "free-usage-exhausted") || - strings.Contains(msg, "free-usage-exhausted") || - strings.Contains(msg, "included free usage") { - d := xaiFreeUsageExhaustedCooldown - err.retryAfter = &d - } - return err -} diff --git a/internal/runtime/executor/xai_executor_auth.go b/internal/runtime/executor/xai_executor_auth.go new file mode 100644 index 000000000..97074d1e3 --- /dev/null +++ b/internal/runtime/executor/xai_executor_auth.go @@ -0,0 +1,76 @@ +package executor + +import ( + "context" + "net/http" + "strings" + "time" + + xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + log "github.com/sirupsen/logrus" +) + +// Refresh refreshes xAI OAuth credentials using the stored refresh token. +func (e *XAIExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { + log.Debugf("xai executor: refresh called") + if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { + return refreshed, err + } + if auth == nil { + return nil, statusErr{code: http.StatusInternalServerError, msg: "xai executor: auth is nil"} + } + refreshToken := xaiMetadataString(auth.Metadata, "refresh_token") + if refreshToken == "" { + return auth, nil + } + tokenEndpoint := xaiMetadataString(auth.Metadata, "token_endpoint") + svc := xaiauth.NewXAIAuthWithProxyURL(e.cfg, auth.ProxyURL) + td, err := svc.RefreshTokens(ctx, refreshToken, tokenEndpoint) + if err != nil { + return nil, err + } + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["type"] = "xai" + auth.Metadata["auth_kind"] = "oauth" + auth.Metadata["access_token"] = td.AccessToken + if td.RefreshToken != "" { + auth.Metadata["refresh_token"] = td.RefreshToken + } + if td.IDToken != "" { + auth.Metadata["id_token"] = td.IDToken + } + if td.TokenType != "" { + auth.Metadata["token_type"] = td.TokenType + } + if td.ExpiresIn > 0 { + auth.Metadata["expires_in"] = td.ExpiresIn + } + if td.Expire != "" { + auth.Metadata["expired"] = td.Expire + } + if td.Email != "" { + auth.Metadata["email"] = td.Email + } + if td.Subject != "" { + auth.Metadata["sub"] = td.Subject + } + if tokenEndpoint != "" { + auth.Metadata["token_endpoint"] = tokenEndpoint + } + if xaiMetadataString(auth.Metadata, "base_url") == "" { + auth.Metadata["base_url"] = xaiauth.DefaultAPIBaseURL + } + auth.Metadata["last_refresh"] = time.Now().UTC().Format(time.RFC3339) + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + auth.Attributes["auth_kind"] = "oauth" + if strings.TrimSpace(auth.Attributes["base_url"]) == "" { + auth.Attributes["base_url"] = xaiauth.DefaultAPIBaseURL + } + return auth, nil +} diff --git a/internal/runtime/executor/xai_executor_execute.go b/internal/runtime/executor/xai_executor_execute.go new file mode 100644 index 000000000..d5a2b74a2 --- /dev/null +++ b/internal/runtime/executor/xai_executor_execute.go @@ -0,0 +1,382 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + if opts.Alt == "responses/compact" { + return e.executeCompact(ctx, auth, req, opts) + } + if endpointPath := xaiImageEndpointPath(opts); endpointPath != "" { + return e.executeImages(ctx, auth, req, endpointPath) + } + if xaiIsVideoRequest(opts) { + return e.executeVideos(ctx, auth, req, opts) + } + + token, _ := xaiCreds(auth) + baseURL := xaiChatBaseURL(auth) + logXAIResolvedBaseURL(ctx, baseURL) + + prepared, err := e.prepareResponsesRequest(ctx, req, opts, true) + if err != nil { + return resp, err + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) + + url := strings.TrimSuffix(baseURL, "/") + "/responses" + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(prepared.body)) + if err != nil { + return resp, err + } + applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID) + e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) + + httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("xai executor: close response body error: %v", errClose) + } + }() + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + data, errRead := io.ReadAll(httpResp.Body) + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + return resp, errRead + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) + return resp, xaiStatusErr(httpResp.StatusCode, data) + } + + data, err := io.ReadAll(httpResp.Body) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) + for _, line := range bytes.Split(data, []byte("\n")) { + if !bytes.HasPrefix(line, xaiDataTag) { + continue + } + eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):])) + eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) + eventData = responseFilter.apply(eventData) + if len(eventData) == 0 { + continue + } + switch gjson.GetBytes(eventData, "type").String() { + case "response.output_item.done": + xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) + case "response.completed": + if detail, ok := helps.ParseCodexUsage(eventData); ok { + reporter.Publish(ctx, detail) + } + completedData := xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) + completedData = xaiNormalizeReasoningSummaryData(completedData) + cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, completedData) + var param any + out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, completedData, ¶m) + return cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()}, nil + } + } + + return resp, statusErr{code: http.StatusRequestTimeout, msg: "xai stream error: stream disconnected before response.completed"} +} + +func (e *XAIExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + prepared, data, headers, errCompact := e.executeCompactRequest(ctx, auth, req, opts) + if errCompact != nil { + return resp, errCompact + } + + var param any + out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, data, ¶m) + return cliproxyexecutor.Response{Payload: out, Headers: headers}, nil +} + +func (e *XAIExecutor) executeCompactRequest(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*xaiPreparedRequest, []byte, http.Header, error) { + token, _ := xaiCreds(auth) + // Compact must not use xaiChatBaseURL: CLI chat-proxy returns 404 for + // /responses/compact and a 404 cools down the whole xAI auth pool. + baseURL := xaiCompactBaseURL(auth) + logXAIResolvedBaseURL(ctx, baseURL) + + prepared, err := e.prepareResponsesRequestTo(ctx, req, opts, false, sdktranslator.FormatOpenAIResponse) + if err != nil { + return nil, nil, nil, err + } + prepared.body, _ = sjson.DeleteBytes(prepared.body, "stream") + prepared.body, _ = sjson.DeleteBytes(prepared.body, "tools") + for _, field := range []string{"max_output_tokens", "temperature", "top_p", "top_k", "stop"} { + prepared.body, _ = sjson.DeleteBytes(prepared.body, field) + } + prepared.body = xaiRemoveInputItemsByType(prepared.body, "compaction_trigger") + + reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) + + requestURL := strings.TrimSuffix(baseURL, "/") + "/responses/compact" + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL, bytes.NewReader(prepared.body)) + if err != nil { + return nil, nil, nil, err + } + // Official API / custom compact endpoints use standard API headers, not CLI + // chat-proxy identity headers (which applyXAIChatHeaders may still attach for OAuth chat). + applyXAIHeaders(httpReq, auth, token, false, prepared.sessionID) + e.recordXAIRequest(ctx, auth, requestURL, httpReq.Header.Clone(), prepared.body) + + httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return nil, nil, nil, err + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("xai executor: close response body error: %v", errClose) + } + }() + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + + data, err := io.ReadAll(httpResp.Body) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return nil, nil, nil, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) + err = xaiStatusErr(httpResp.StatusCode, data) + return nil, nil, nil, err + } + + reporter.Publish(ctx, helps.ParseOpenAIUsage(data)) + reporter.EnsurePublished(ctx) + clearXAIReasoningReplayAfterCompaction(ctx, prepared.replayScope) + return prepared, data, httpResp.Header.Clone(), nil +} + +func (e *XAIExecutor) executeCompactionTriggerStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + prepared, data, headers, err := e.executeCompactRequest(ctx, auth, req, opts) + if err != nil { + return nil, err + } + + headers = headers.Clone() + if headers == nil { + headers = make(http.Header) + } + headers.Set("Content-Type", "text/event-stream") + + chunks := xaiBuildCompactionTriggerStreamChunks(prepared, data) + out := make(chan cliproxyexecutor.StreamChunk, len(chunks)) + for _, chunk := range chunks { + out <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + close(out) + return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil +} + +func xaiInputHasItemType(body []byte, itemType string) bool { + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return false + } + for _, item := range input.Array() { + if item.Get("type").String() == itemType { + return true + } + } + return false +} + +func xaiRemoveInputItemsByType(body []byte, itemType string) []byte { + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return body + } + + var buf bytes.Buffer + buf.WriteByte('[') + kept := 0 + for _, item := range input.Array() { + if item.Get("type").String() == itemType { + continue + } + if kept > 0 { + buf.WriteByte(',') + } + buf.WriteString(item.Raw) + kept++ + } + buf.WriteByte(']') + + updated, err := sjson.SetRawBytes(body, "input", buf.Bytes()) + if err != nil { + return body + } + return updated +} + +func xaiBuildCompactionTriggerStreamChunks(prepared *xaiPreparedRequest, compactData []byte) [][]byte { + responseID := xaiCompactionResponseID(compactData) + now := time.Now().Unix() + createdAt := gjson.GetBytes(compactData, "created_at").Int() + if createdAt == 0 { + createdAt = now + } + completedAt := gjson.GetBytes(compactData, "completed_at").Int() + if completedAt == 0 { + completedAt = now + } + + item := xaiCompactionOutputItem(compactData, responseID) + output := make([]byte, 0, len(item)+2) + output = append(output, '[') + output = append(output, item...) + output = append(output, ']') + + createdResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "in_progress") + inProgressResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "in_progress") + completedResponse := xaiBuildCompactionBaseResponse(prepared, compactData, responseID, createdAt, "completed") + completedResponse, _ = sjson.SetBytes(completedResponse, "completed_at", completedAt) + completedResponse, _ = sjson.SetRawBytes(completedResponse, "output", output) + if usage := gjson.GetBytes(compactData, "usage"); usage.Exists() { + completedResponse, _ = sjson.SetRawBytes(completedResponse, "usage", []byte(usage.Raw)) + } + + createdPayload := []byte(`{"type":"response.created","sequence_number":0}`) + createdPayload, _ = sjson.SetRawBytes(createdPayload, "response", createdResponse) + inProgressPayload := []byte(`{"type":"response.in_progress","sequence_number":1}`) + inProgressPayload, _ = sjson.SetRawBytes(inProgressPayload, "response", inProgressResponse) + addedPayload := []byte(`{"type":"response.output_item.added","sequence_number":2,"output_index":0}`) + addedPayload, _ = sjson.SetRawBytes(addedPayload, "item", item) + keepalivePayload := []byte(`{"type":"keepalive","sequence_number":3}`) + donePayload := []byte(`{"type":"response.output_item.done","sequence_number":4,"output_index":0}`) + donePayload, _ = sjson.SetRawBytes(donePayload, "item", item) + completedPayload := []byte(`{"type":"response.completed","sequence_number":5}`) + completedPayload, _ = sjson.SetRawBytes(completedPayload, "response", completedResponse) + + return [][]byte{ + xaiBuildSSEFrame("response.created", createdPayload), + xaiBuildSSEFrame("response.in_progress", inProgressPayload), + xaiBuildSSEFrame("response.output_item.added", addedPayload), + xaiBuildSSEFrame("keepalive", keepalivePayload), + xaiBuildSSEFrame("response.output_item.done", donePayload), + xaiBuildSSEFrame("response.completed", completedPayload), + } +} + +func xaiBuildCompactionBaseResponse(prepared *xaiPreparedRequest, compactData []byte, responseID string, createdAt int64, status string) []byte { + response := []byte(`{"id":"","object":"response","created_at":0,"status":"","background":false,"error":null,"incomplete_details":null,"output":[]}`) + response, _ = sjson.SetBytes(response, "id", responseID) + response, _ = sjson.SetBytes(response, "created_at", createdAt) + response, _ = sjson.SetBytes(response, "status", status) + if model := gjson.GetBytes(compactData, "model").String(); model != "" { + response, _ = sjson.SetBytes(response, "model", model) + } else if prepared != nil && prepared.baseModel != "" { + response, _ = sjson.SetBytes(response, "model", prepared.baseModel) + } + + if prepared == nil { + return response + } + for _, field := range []string{ + "instructions", + "max_output_tokens", + "max_tool_calls", + "parallel_tool_calls", + "previous_response_id", + "prompt_cache_key", + "reasoning", + "text", + "tool_choice", + "tools", + "top_logprobs", + "top_p", + "truncation", + "user", + "metadata", + } { + if value := gjson.GetBytes(prepared.body, field); value.Exists() { + response, _ = sjson.SetRawBytes(response, field, []byte(value.Raw)) + } + } + return response +} + +func xaiCompactionOutputItem(compactData []byte, responseID string) []byte { + itemResult := gjson.GetBytes(compactData, "output.0") + item := []byte(`{"type":"compaction"}`) + if itemResult.Exists() && itemResult.Type == gjson.JSON { + item = []byte(itemResult.Raw) + } + if !gjson.GetBytes(item, "type").Exists() { + item, _ = sjson.SetBytes(item, "type", "compaction") + } + if !gjson.GetBytes(item, "id").Exists() { + item, _ = sjson.SetBytes(item, "id", xaiCompactionItemID(responseID)) + } + return item +} + +func xaiCompactionResponseID(compactData []byte) string { + if responseID := strings.TrimSpace(gjson.GetBytes(compactData, "id").String()); responseID != "" { + if strings.HasPrefix(responseID, "resp_") { + return responseID + } + return "resp_" + strings.TrimPrefix(responseID, "cmp_") + } + return fmt.Sprintf("resp_xai_compaction_%d", time.Now().UnixNano()) +} + +func xaiCompactionItemID(responseID string) string { + if suffix := strings.TrimPrefix(responseID, "resp_"); suffix != "" && suffix != responseID { + return "cmp_" + suffix + } + return "cmp_" + responseID +} + +func xaiBuildSSEFrame(eventName string, data []byte) []byte { + out := make([]byte, 0, len(eventName)+len(data)+16) + out = append(out, "event: "...) + out = append(out, eventName...) + out = append(out, '\n') + out = append(out, "data: "...) + out = append(out, data...) + out = append(out, '\n', '\n') + return out +} diff --git a/internal/runtime/executor/xai_executor_media.go b/internal/runtime/executor/xai_executor_media.go new file mode 100644 index 000000000..847448e8a --- /dev/null +++ b/internal/runtime/executor/xai_executor_media.go @@ -0,0 +1,141 @@ +package executor + +import ( + "bytes" + "context" + "io" + "net/http" + "net/url" + "strings" + + xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, endpointPath string) (resp cliproxyexecutor.Response, err error) { + model := strings.TrimSpace(gjson.GetBytes(req.Payload, "model").String()) + if model == "" { + model = strings.TrimSpace(req.Model) + } + reporter := helps.NewExecutorUsageReporter(ctx, e, model, auth) + defer reporter.TrackFailure(ctx, &err) + + token, baseURL := xaiCreds(auth) + if baseURL == "" { + baseURL = xaiauth.DefaultAPIBaseURL + } + logXAIResolvedBaseURL(ctx, baseURL) + if endpointPath == "" { + endpointPath = xaiDefaultImageEndpointPath + } + + payload := normalizeXAIImageRefs(req.Payload) + url := strings.TrimSuffix(baseURL, "/") + endpointPath + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + return resp, err + } + applyXAIHeaders(httpReq, auth, token, false, "") + e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), payload) + + httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("xai executor: close response body error: %v", errClose) + } + }() + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + + data, err := io.ReadAll(httpResp.Body) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) + err = xaiStatusErr(httpResp.StatusCode, data) + return resp, err + } + + reporter.EnsurePublished(ctx) + return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil +} + +func (e *XAIExecutor) executeVideos(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) { + token, baseURL := xaiCreds(auth) + if baseURL == "" { + baseURL = xaiauth.DefaultAPIBaseURL + } + logXAIResolvedBaseURL(ctx, baseURL) + + payload := normalizeXAIImageRefs(req.Payload) + method := http.MethodPost + endpointPath := xaiVideosGenerationsPath + var body io.Reader = bytes.NewReader(payload) + + switch path := xaiVideoEndpointPath(opts); path { + case xaiVideosGenerationsPath, xaiVideosEditsPath, xaiVideosExtensionsPath: + endpointPath = path + default: + if requestID := strings.TrimSpace(gjson.GetBytes(payload, "request_id").String()); requestID != "" { + method = http.MethodGet + endpointPath = xaiVideosPath + "/" + url.PathEscape(requestID) + body = nil + } + } + requestURL := strings.TrimSuffix(baseURL, "/") + endpointPath + httpReq, err := http.NewRequestWithContext(ctx, method, requestURL, body) + if err != nil { + return resp, err + } + applyXAIHeaders(httpReq, auth, token, false, "") + if method == http.MethodPost { + key := xaiMetadataString(opts.Metadata, xaiIdempotencyKeyMetaKey) + if key == "" && opts.Headers != nil { + key = strings.TrimSpace(opts.Headers.Get("x-idempotency-key")) + } + if key != "" { + httpReq.Header.Set("x-idempotency-key", key) + } + } + e.recordXAIRequest(ctx, auth, requestURL, httpReq.Header.Clone(), payload) + + httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("xai executor: close response body error: %v", errClose) + } + }() + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + + data, err := io.ReadAll(httpResp.Body) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return resp, err + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) + return resp, xaiStatusErr(httpResp.StatusCode, data) + } + + return cliproxyexecutor.Response{Payload: data, Headers: httpResp.Header.Clone()}, nil +} diff --git a/internal/runtime/executor/xai_executor_request.go b/internal/runtime/executor/xai_executor_request.go new file mode 100644 index 000000000..623c9ba32 --- /dev/null +++ b/internal/runtime/executor/xai_executor_request.go @@ -0,0 +1,1145 @@ +package executor + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strconv" + "strings" + + "github.com/google/uuid" + xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +type xaiPreparedRequest struct { + baseModel string + from sdktranslator.Format + responseFormat sdktranslator.Format + to sdktranslator.Format + originalPayload []byte + body []byte + namespaceTools map[string]xaiNamespaceToolRef + clientDeclaredTools map[xaiClientToolKey]struct{} + sessionID string + replayScope xaiReasoningReplayScope + filterInternalXSearch bool +} + +type xaiNamespaceToolRef struct { + namespace string + name string +} + +// xaiClientToolKey identifies a client-declared callable tool using the +// post-restore Responses shape (short name + optional namespace) and the +// effective upstream tool type after normalizeXAITool (client custom tools are +// sent as function). Response call types are matched against this effective +// kind so internal custom_tool_call traces are not exempted merely because a +// client declared an ordinary function/custom tool with the same short name, +// while legitimate function_call responses for normalized custom tools are kept. +type xaiClientToolKey struct { + namespace string + name string + toolType string +} + +func (e *XAIExecutor) prepareResponsesRequest(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool) (*xaiPreparedRequest, error) { + return e.prepareResponsesRequestTo(ctx, req, opts, stream, sdktranslator.FormatCodex) +} + +func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, stream bool, to sdktranslator.Format) (*xaiPreparedRequest, error) { + baseModel := thinking.ParseSuffix(req.Model).ModelName + from := opts.SourceFormat + responseFormat := cliproxyexecutor.ResponseFormatOrSource(opts) + originalPayloadSource := req.Payload + if len(opts.OriginalRequest) > 0 { + originalPayloadSource = opts.OriginalRequest + } + originalPayload := bytes.Clone(originalPayloadSource) + originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, stream) + originalTranslated = preserveXAIResponsesOutputControls(originalTranslated, originalPayload, from) + body := sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), stream) + body = preserveXAIResponsesOutputControls(body, req.Payload, from) + + var err error + body, err = thinking.ApplyThinking(body, req.Model, from.String(), e.Identifier(), e.Identifier()) + if err != nil { + return nil, err + } + + requestedModel := helps.PayloadRequestedModel(opts, req.Model) + requestPath := helps.PayloadRequestPath(opts) + body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = helps.SetBoolIfDifferent(body, "stream", stream) + body, _ = sjson.DeleteBytes(body, "previous_response_id") + body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") + body, _ = sjson.DeleteBytes(body, "safety_identifier") + body, _ = sjson.DeleteBytes(body, "stream_options") + body = helps.RewriteCodexMultiAgentV2Input(ctx, opts.Headers, body, e.cfg) + namespaceTools := collectXAINamespaceToolRefs(body) + // Collect before normalizeXAITools flattens namespace wrappers so keys match + // the post-restore (namespace, short-name) shape used by the response filter. + clientDeclaredTools := collectXAIClientDeclaredToolKeys(body) + body = normalizeXAITools(body) + body = promoteXAIAdditionalTools(body) + // Drop choices that point at tools removed by normalizeXAITools before we + // inject native x_search, so a surviving allowed_tools / forced choice is not + // left pointing at a deleted tool once only x_search remains. + body = normalizeXAINamespaceToolChoice(body) + body = pruneXAIOrphanedToolChoice(body) + body = normalizeXAIToolChoiceForTools(body) + body = ensureXAINativeXSearchTool(body) + var replayScope xaiReasoningReplayScope + body, replayScope, err = applyXAIReasoningReplayCacheRequired(ctx, from, req, opts, body) + if err != nil { + return nil, err + } + body = normalizeXAIInputCustomToolCalls(body) + body = normalizeXAIInputNamespaceToolCalls(body) + body = normalizeXAIInputReasoningItems(body) + body = sanitizeXAIInputEncryptedContent(body) + body = normalizeCodexInstructions(body) + body = sanitizeXAIResponsesBody(body, baseModel) + body = normalizeXAIImageRefs(body) + + sessionID, errSession := xaiResolveComposerSessionID(ctx, req, opts, baseModel) + if errSession != nil { + return nil, errSession + } + if sessionID != "" { + body = helps.SetStringIfDifferent(body, "prompt_cache_key", sessionID) + } + + return &xaiPreparedRequest{ + baseModel: baseModel, + from: from, + responseFormat: responseFormat, + to: to, + originalPayload: originalPayload, + body: body, + namespaceTools: namespaceTools, + clientDeclaredTools: clientDeclaredTools, + sessionID: sessionID, + replayScope: replayScope, + filterInternalXSearch: xaiRequestHasNativeXSearch(body), + }, nil +} + +func (e *XAIExecutor) recordXAIRequest(ctx context.Context, auth *cliproxyauth.Auth, url string, headers http.Header, body []byte) { + var authID, authLabel, authType, authValue string + if auth != nil { + authID = auth.ID + authLabel = auth.Label + authType, authValue = auth.AccountInfo() + } + helps.RecordAPIRequest(ctx, e.cfg, helps.UpstreamRequestLog{ + URL: url, + Method: http.MethodPost, + Headers: headers, + Body: body, + Provider: e.Identifier(), + AuthID: authID, + AuthLabel: authLabel, + AuthType: authType, + AuthValue: authValue, + }) +} + +func xaiCreds(auth *cliproxyauth.Auth) (token, baseURL string) { + if auth == nil { + return "", "" + } + if auth.Attributes != nil { + token = strings.TrimSpace(auth.Attributes["api_key"]) + baseURL = strings.TrimSpace(auth.Attributes["base_url"]) + } + if auth.Metadata != nil { + if token == "" { + token = xaiMetadataString(auth.Metadata, "access_token") + } + if baseURL == "" { + baseURL = xaiMetadataString(auth.Metadata, "base_url") + } + } + return token, baseURL +} + +// xaiUsingAPI reports whether this xAI auth should use the official API path +// for non-media HTTP chat. OAuth defaults to false to use Grok Build. +func xaiUsingAPI(auth *cliproxyauth.Auth) bool { + if auth == nil { + return true + } + if len(auth.Attributes) > 0 { + if raw := strings.TrimSpace(auth.Attributes[xaiUsingAPIAttr]); raw != "" { + parsed, errParse := strconv.ParseBool(raw) + if errParse == nil { + return parsed + } + } + } + if len(auth.Metadata) > 0 { + raw, ok := auth.Metadata[xaiUsingAPIAttr] + if ok && raw != nil { + switch v := raw.(type) { + case bool: + return v + case string: + parsed, errParse := strconv.ParseBool(strings.TrimSpace(v)) + if errParse == nil { + return parsed + } + default: + } + } + } + if raw := strings.TrimSpace(auth.Attributes["auth_kind"]); raw != "" { + return !strings.EqualFold(raw, "oauth") + } + return !strings.EqualFold(xaiMetadataString(auth.Metadata, "auth_kind"), "oauth") +} + +// xaiChatBaseURL returns the base URL for non-image/video xAI HTTP chat requests. +// When auth using_api is true, the official API base URL logic is used. When it +// is false (including its OAuth default), empty or official default base_url is +// rewritten to the CLI chat-proxy endpoint; an explicit non-default base_url is +// still honored. +// Websocket and compact transports intentionally do not use this helper: +// cli-chat-proxy only accepts HTTP POST chat and does not implement +// /responses/compact (404) or websocket upgrades (405). +func xaiChatBaseURL(auth *cliproxyauth.Auth) string { + _, baseURL := xaiCreds(auth) + if xaiUsingAPI(auth) { + if baseURL == "" { + return xaiauth.DefaultAPIBaseURL + } + return baseURL + } + if baseURL != "" && !xaiIsDefaultAPIBaseURL(baseURL) { + return baseURL + } + return xaiauth.CLIChatProxyBaseURL +} + +// xaiCompactBaseURL returns the base URL for xAI /responses/compact requests. +// Compact must stay on the official API (or an explicit non-CLI-proxy base_url). +// Reusing xaiChatBaseURL would pin OAuth traffic to cli-chat-proxy, which returns +// 404 for /responses/compact and then cools down the auth pool as not_found. +func xaiCompactBaseURL(auth *cliproxyauth.Auth) string { + _, baseURL := xaiCreds(auth) + if baseURL == "" || xaiIsCLIChatProxyBaseURL(baseURL) { + return xaiauth.DefaultAPIBaseURL + } + return baseURL +} + +func xaiNormalizeBaseURL(baseURL string) string { + return strings.TrimRight(strings.TrimSpace(baseURL), "/") +} + +func xaiIsDefaultAPIBaseURL(baseURL string) bool { + return xaiNormalizeBaseURL(baseURL) == xaiNormalizeBaseURL(xaiauth.DefaultAPIBaseURL) +} + +func xaiIsCLIChatProxyBaseURL(baseURL string) bool { + return xaiNormalizeBaseURL(baseURL) == xaiNormalizeBaseURL(xaiauth.CLIChatProxyBaseURL) +} + +// xaiBaseURLSource classifies a resolved xAI base URL for logging. +func xaiBaseURLSource(baseURL string) string { + switch { + case xaiIsDefaultAPIBaseURL(baseURL): + return "DefaultAPIBaseURL" + case xaiIsCLIChatProxyBaseURL(baseURL): + return "CLIChatProxyBaseURL" + default: + return "custom" + } +} + +// logXAIResolvedBaseURL emits a console log for the resolved upstream base URL. +func logXAIResolvedBaseURL(ctx context.Context, baseURL string) { + helps.LogWithRequestID(ctx).Infof("xai: using base_url=%s source=%s", baseURL, xaiBaseURLSource(baseURL)) +} + +func applyXAIHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) { + applyXAIDefaultHeaders(r, token, stream, sessionID) + applyXAICustomHeaders(r, auth) +} + +func applyXAIDefaultHeaders(r *http.Request, token string, stream bool, sessionID string) { + r.Header.Set("Content-Type", "application/json") + if strings.TrimSpace(token) != "" { + r.Header.Set("Authorization", "Bearer "+token) + } + if stream { + r.Header.Set("Accept", "text/event-stream") + } else { + r.Header.Set("Accept", "application/json") + } + r.Header.Set("Connection", "Keep-Alive") + if sessionID != "" { + r.Header.Set("x-grok-conv-id", sessionID) + } +} + +func applyXAICustomHeaders(r *http.Request, auth *cliproxyauth.Auth) { + var attrs map[string]string + if auth != nil { + attrs = auth.Attributes + } + util.ApplyCustomHeadersFromAttrs(r, attrs) +} + +// applyXAIChatHeaders applies standard xAI headers for non-image/video chat +// requests. When using_api is true, this matches the standard +// applyXAIHeaders behavior. CLI chat-proxy identity headers are only attached +// when using_api is false and the resolved chat base URL is the official CLI +// chat-proxy endpoint. +func applyXAIChatHeaders(r *http.Request, auth *cliproxyauth.Auth, token string, stream bool, sessionID string) { + if xaiUsingAPI(auth) { + applyXAIHeaders(r, auth, token, stream, sessionID) + return + } + applyXAIDefaultHeaders(r, token, stream, sessionID) + if xaiIsCLIChatProxyBaseURL(xaiChatBaseURL(auth)) { + r.Header.Set(xaiTokenAuthHeader, xaiTokenAuthValue) + r.Header.Set(xaiClientVersionHeader, xaiClientVersionValue) + r.Header.Set("User-Agent", "xai-grok-workspace/"+xaiClientVersionValue) + } + applyXAICustomHeaders(r, auth) +} + +func xaiResolveComposerSessionID(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, baseModel string) (string, error) { + if sessionID := xaiExecutionSessionID(req, opts); sessionID != "" { + return sessionID, nil + } + if !xaiRequiresIsolatedConversation(baseModel) { + return "", nil + } + cached, ok, errCache := helps.ClaudeCodePromptCache(ctx, baseModel, req.Payload, opts.Headers) + if errCache != nil { + return "", errCache + } + if ok { + return cached.ID, nil + } + return uuid.NewString(), nil +} + +func xaiExecutionSessionID(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) string { + if value := xaiMetadataString(opts.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { + return value + } + if value := xaiMetadataString(req.Metadata, cliproxyexecutor.ExecutionSessionMetadataKey); value != "" { + return value + } + if promptCacheKey := gjson.GetBytes(req.Payload, "prompt_cache_key"); promptCacheKey.Exists() { + if value := strings.TrimSpace(promptCacheKey.String()); value != "" { + return value + } + } + return helps.DerivedSessionUUID("xai", opts.Metadata, req.Metadata) +} + +func xaiRequiresIsolatedConversation(model string) bool { + return strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), xaiComposerModelPrefix) +} + +func xaiImageEndpointPath(opts cliproxyexecutor.Options) string { + if opts.SourceFormat.String() != xaiImageHandlerType { + return "" + } + + path := xaiMetadataString(opts.Metadata, cliproxyexecutor.RequestPathMetadataKey) + if strings.HasSuffix(path, "/images/edits") { + return xaiImagesEditsPath + } + if strings.HasSuffix(path, "/images/generations") { + return xaiImagesGenerationsPath + } + return xaiDefaultImageEndpointPath +} + +// normalizeXAIImageRefs rewrites OpenAI-style image object fields to the xAI +// image API shape before the payload is sent upstream: +// +// {"image":{"image_url":"https://..."}} → {"image":{"url":"https://..."}} +// +// Applies to image / images / reference_images anywhere in the JSON tree, +// including nested objects and array items. Does not rewrite chat content +// parts shaped as {"type":"image_url","image_url":{...}}. +func normalizeXAIImageRefs(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var payload any + if errDecode := decoder.Decode(&payload); errDecode != nil { + return body + } + + if !normalizeXAIImageRefsValue(payload) { + return body + } + normalized, errMarshal := json.Marshal(payload) + if errMarshal != nil { + return body + } + return normalized +} + +func normalizeXAIImageRefsValue(value any) bool { + changed := false + switch node := value.(type) { + case map[string]any: + for key, child := range node { + switch key { + case "image": + changed = normalizeXAIImageRef(child) || changed + case "images", "reference_images": + if refs, ok := child.([]any); ok { + for _, ref := range refs { + changed = normalizeXAIImageRef(ref) || changed + } + } + } + changed = normalizeXAIImageRefsValue(child) || changed + } + case []any: + for _, child := range node { + changed = normalizeXAIImageRefsValue(child) || changed + } + } + return changed +} + +func normalizeXAIImageRef(value any) bool { + ref, ok := value.(map[string]any) + if !ok { + return false + } + + originalURL, _ := ref["url"].(string) + url := strings.TrimSpace(originalURL) + imageURL, hasImageURL := ref["image_url"] + if url == "" { + switch imageURL := imageURL.(type) { + case string: + url = strings.TrimSpace(imageURL) + case map[string]any: + url, _ = imageURL["url"].(string) + url = strings.TrimSpace(url) + } + } + if url == "" { + return false + } + if url == originalURL && !hasImageURL { + return false + } + + // Always emit the xAI field name and drop the OpenAI alias. + ref["url"] = url + delete(ref, "image_url") + return true +} + +func xaiIsVideoRequest(opts cliproxyexecutor.Options) bool { + return opts.SourceFormat.String() == xaiVideoHandlerType +} + +func xaiVideoEndpointPath(opts cliproxyexecutor.Options) string { + if !xaiIsVideoRequest(opts) { + return "" + } + path := xaiMetadataString(opts.Metadata, cliproxyexecutor.RequestPathMetadataKey) + if strings.HasSuffix(path, "/videos/edits") { + return xaiVideosEditsPath + } + if strings.HasSuffix(path, "/videos/extensions") { + return xaiVideosExtensionsPath + } + if strings.HasSuffix(path, "/videos/generations") { + return xaiVideosGenerationsPath + } + return "" +} + +func xaiMetadataString(meta map[string]any, key string) string { + if len(meta) == 0 || key == "" { + return "" + } + value, ok := meta[key] + if !ok || value == nil { + return "" + } + switch typed := value.(type) { + case string: + return strings.TrimSpace(typed) + case fmt.Stringer: + return strings.TrimSpace(typed.String()) + default: + return strings.TrimSpace(fmt.Sprint(typed)) + } +} + +func preserveXAIResponsesOutputControls(body, source []byte, from sdktranslator.Format) []byte { + var maxOutputTokens gjson.Result + switch from { + case sdktranslator.FormatOpenAI: + maxOutputTokens = gjson.GetBytes(source, "max_completion_tokens") + if !maxOutputTokens.Exists() || maxOutputTokens.Type == gjson.Null { + maxOutputTokens = gjson.GetBytes(source, "max_tokens") + } + case sdktranslator.FormatOpenAIResponse: + maxOutputTokens = gjson.GetBytes(source, "max_output_tokens") + default: + return body + } + + if maxOutputTokens.Exists() && maxOutputTokens.Type != gjson.Null { + body, _ = sjson.SetRawBytes(body, "max_output_tokens", []byte(maxOutputTokens.Raw)) + } + for _, field := range []string{"temperature", "top_p", "top_k"} { + value := gjson.GetBytes(source, field) + if value.Exists() && value.Type != gjson.Null { + body, _ = sjson.SetRawBytes(body, field, []byte(value.Raw)) + } + } + return body +} + +func sanitizeXAIResponsesBody(body []byte, model string) []byte { + // stop is supported by Chat Completions but not by xAI's Responses API. + body, _ = sjson.DeleteBytes(body, "stop") + if !xaiSupportsReasoningEffort(model) { + if gjson.GetBytes(body, "reasoning.effort").Exists() { + log.Debugf("xai: stripping reasoning.effort for model %s (no thinking levels in model registry)", model) + } + body, _ = sjson.DeleteBytes(body, "reasoning.effort") + if reasoning := gjson.GetBytes(body, "reasoning"); reasoning.Exists() && reasoning.IsObject() && len(reasoning.Map()) == 0 { + body, _ = sjson.DeleteBytes(body, "reasoning") + } + } + return body +} + +// ensureXAINativeXSearchTool appends {"type":"x_search"} when the final tools +// list does not already include native X Search. When tool_choice restricts the +// model to allowed_tools, x_search is also added there (without duplicates) so +// Grok can select the injected tool. HTTP and websocket executors both prepare +// payloads through prepareResponsesRequestTo, so this runs once before the body +// is submitted upstream. +func ensureXAINativeXSearchTool(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + if !xaiRequestHasNativeXSearch(body) { + tools := gjson.GetBytes(body, "tools") + if !tools.Exists() || !tools.IsArray() { + body, _ = sjson.SetRawBytes(body, "tools", []byte(`[{"type":"x_search"}]`)) + } else { + body, _ = sjson.SetRawBytes(body, "tools.-1", xaiXSearchToolJSON) + } + } + return ensureXAINativeXSearchAllowedTools(body) +} + +// ensureXAINativeXSearchAllowedTools appends x_search to tool_choice.tools when +// the choice mode is allowed_tools and x_search is not already listed. +func ensureXAINativeXSearchAllowedTools(body []byte) []byte { + choice := gjson.GetBytes(body, "tool_choice") + if !choice.IsObject() || choice.Get("type").String() != "allowed_tools" { + return body + } + allowed := choice.Get("tools") + if !allowed.Exists() || !allowed.IsArray() { + body, _ = sjson.SetRawBytes(body, "tool_choice.tools", []byte(`[{"type":"x_search"}]`)) + return body + } + for _, tool := range allowed.Array() { + if strings.TrimSpace(tool.Get("type").String()) == xaiXSearchToolType { + return body + } + } + body, _ = sjson.SetRawBytes(body, "tool_choice.tools.-1", xaiXSearchToolJSON) + return body +} + +// pruneXAIOrphanedToolChoice removes tool_choice entries that no longer match +// any remaining tool after normalizeXAITools filtering. Forced choices that +// reference a deleted tool are dropped entirely; allowed_tools lists keep only +// choices that still resolve against the post-normalization tools set. +func pruneXAIOrphanedToolChoice(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + choice := gjson.GetBytes(body, "tool_choice") + if !choice.Exists() { + return body + } + available := collectXAIAvailableToolChoiceKeys(body) + if choice.Type == gjson.String { + // auto / none / required are not tool references. + return body + } + if !choice.IsObject() { + return body + } + choiceType := strings.TrimSpace(choice.Get("type").String()) + switch choiceType { + case "allowed_tools": + return pruneXAIAllowedToolsChoice(body, available) + default: + if choiceType == "" { + return body + } + if xaiToolChoiceMatchesAvailable(choice, available) { + return body + } + body, _ = sjson.DeleteBytes(body, "tool_choice") + return body + } +} + +func pruneXAIAllowedToolsChoice(body []byte, available map[xaiToolChoiceKey]struct{}) []byte { + allowed := gjson.GetBytes(body, "tool_choice.tools") + if !allowed.Exists() || !allowed.IsArray() { + body, _ = sjson.DeleteBytes(body, "tool_choice") + return body + } + allowedItems := allowed.Array() + filtered := make([][]byte, 0, len(allowedItems)) + changed := false + for _, tool := range allowedItems { + if !xaiToolChoiceMatchesAvailable(tool, available) { + changed = true + continue + } + filtered = append(filtered, []byte(tool.Raw)) + } + if !changed { + return body + } + if len(filtered) == 0 { + body, _ = sjson.DeleteBytes(body, "tool_choice") + return body + } + body, _ = sjson.SetRawBytes(body, "tool_choice.tools", helps.JoinRawJSONArray(filtered)) + return body +} + +// xaiToolChoiceKey identifies a selectable tool the way xAI tool_choice entries +// reference it after namespace qualification: type alone for host tools, or +// type+name for function tools. +type xaiToolChoiceKey struct { + toolType string + name string +} + +func collectXAIAvailableToolChoiceKeys(body []byte) map[xaiToolChoiceKey]struct{} { + keys := make(map[xaiToolChoiceKey]struct{}) + collect := func(tools gjson.Result) { + if !tools.IsArray() { + return + } + for _, tool := range tools.Array() { + toolType := strings.TrimSpace(tool.Get("type").String()) + if toolType == "" { + continue + } + key := xaiToolChoiceKey{toolType: toolType} + if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { + key.name = strings.TrimSpace(tool.Get("name").String()) + if key.name == "" { + continue + } + } + keys[key] = struct{}{} + } + } + collect(gjson.GetBytes(body, "tools")) + input := gjson.GetBytes(body, "input") + if input.IsArray() { + for _, item := range input.Array() { + if item.Get("type").String() == "additional_tools" { + collect(item.Get("tools")) + } + } + } + return keys +} + +func xaiToolChoiceMatchesAvailable(choice gjson.Result, available map[xaiToolChoiceKey]struct{}) bool { + toolType := strings.TrimSpace(choice.Get("type").String()) + if toolType == "" { + return false + } + key := xaiToolChoiceKey{toolType: toolType} + if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { + key.name = strings.TrimSpace(choice.Get("name").String()) + if key.name == "" { + return false + } + } + _, ok := available[key] + return ok +} + +func normalizeXAITools(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + original := body + normalizeAtPath := func(path string) bool { + tools := gjson.GetBytes(body, path) + if !tools.Exists() || !tools.IsArray() { + return true + } + filtered, changed, ok := normalizeXAIToolArray(tools) + if !ok { + return false + } + if !changed { + return true + } + updated, errSet := sjson.SetRawBytes(body, path, filtered) + if errSet != nil { + return false + } + body = updated + return true + } + + if !normalizeAtPath("tools") { + return original + } + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for index, item := range input.Array() { + if item.Get("type").String() != "additional_tools" { + continue + } + if !normalizeAtPath(fmt.Sprintf("input.%d.tools", index)) { + return original + } + } + } + return body +} + +// promoteXAIAdditionalTools moves Responses Lite tool declarations to the +// top-level tools array because xAI does not accept additional_tools input items. +func promoteXAIAdditionalTools(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + input := gjson.GetBytes(body, "input") + if !input.IsArray() { + return body + } + + inputItems := input.Array() + remainingInput := make([]json.RawMessage, 0, len(inputItems)) + promotedTools := make([]json.RawMessage, 0) + for _, item := range inputItems { + if item.Get("type").String() != "additional_tools" { + remainingInput = append(remainingInput, json.RawMessage(item.Raw)) + continue + } + for _, tool := range item.Get("tools").Array() { + promotedTools = append(promotedTools, json.RawMessage(tool.Raw)) + } + } + if len(remainingInput) == len(inputItems) { + return body + } + + rawInput, errMarshalInput := json.Marshal(remainingInput) + if errMarshalInput != nil { + return body + } + updated, errSetInput := sjson.SetRawBytes(body, "input", rawInput) + if errSetInput != nil { + return body + } + if len(promotedTools) == 0 { + return updated + } + + topLevelTools := gjson.GetBytes(updated, "tools") + tools := make([]json.RawMessage, 0, len(topLevelTools.Array())+len(promotedTools)) + if topLevelTools.IsArray() { + for _, tool := range topLevelTools.Array() { + tools = append(tools, json.RawMessage(tool.Raw)) + } + } + tools = append(tools, promotedTools...) + rawTools, errMarshalTools := json.Marshal(tools) + if errMarshalTools != nil { + return body + } + updated, errSetTools := sjson.SetRawBytes(updated, "tools", rawTools) + if errSetTools != nil { + return body + } + return updated +} + +func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) { + toolItems := tools.Array() + filtered := make([][]byte, 0, len(toolItems)) + changed := false + for _, tool := range toolItems { + toolType := tool.Get("type").String() + if toolType == xaiNamespaceToolType { + changed = true + namespaceName := tool.Get("name").String() + if namespaceTools := tool.Get("tools"); namespaceTools.IsArray() { + for _, nestedTool := range namespaceTools.Array() { + nestedRaw, nestedChanged, ok := normalizeXAITool(nestedTool, namespaceName) + if !ok { + return nil, false, false + } + changed = changed || nestedChanged + if len(nestedRaw) > 0 { + filtered = append(filtered, nestedRaw) + } + } + } + continue + } + raw, toolChanged, ok := normalizeXAITool(tool, "") + if !ok { + return nil, false, false + } + changed = changed || toolChanged + if len(raw) > 0 { + filtered = append(filtered, raw) + } + } + if !changed { + return nil, false, true + } + return helps.JoinRawJSONArray(filtered), true, true +} + +// normalizeXAIToolChoiceForTools drops tool_choice and parallel_tool_calls +// when tools are absent or empty (including after normalizeXAITools filtering). +// xAI rejects payloads that include tool_choice without any tools defined. +// Existence checks avoid unnecessary sjson parse/copy passes. +func normalizeXAIToolChoiceForTools(body []byte) []byte { + tools := gjson.GetBytes(body, "tools") + hasTools := tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 + if !hasTools { + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for _, item := range input.Array() { + additionalTools := item.Get("tools") + if item.Get("type").String() == "additional_tools" && additionalTools.IsArray() && len(additionalTools.Array()) > 0 { + hasTools = true + break + } + } + } + } + if hasTools { + return body + } + if tools.Exists() { + body, _ = sjson.DeleteBytes(body, "tools") + } + if gjson.GetBytes(body, "tool_choice").Exists() { + body, _ = sjson.DeleteBytes(body, "tool_choice") + } + if gjson.GetBytes(body, "parallel_tool_calls").Exists() { + body, _ = sjson.DeleteBytes(body, "parallel_tool_calls") + } + return body +} + +// normalizeXAINamespaceToolChoice qualifies namespaced function choices using +// the same names sent in the flattened tools list. xAI does not accept the +// Responses namespace field on tool choices. +func normalizeXAINamespaceToolChoice(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + original := body + normalizeAtPath := func(path string) bool { + toolChoice := gjson.GetBytes(body, path) + if !toolChoice.IsObject() || toolChoice.Get("type").String() != xaiFunctionToolType { + return true + } + namespaceName := strings.TrimSpace(toolChoice.Get("namespace").String()) + toolName := strings.TrimSpace(toolChoice.Get("name").String()) + qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) + if namespaceName == "" || qualifiedName == "" { + return true + } + updated, errSet := sjson.SetBytes(body, path+".name", qualifiedName) + if errSet != nil { + return false + } + updated, errDelete := sjson.DeleteBytes(updated, path+".namespace") + if errDelete != nil { + return false + } + body = updated + return true + } + + if !normalizeAtPath("tool_choice") { + return original + } + tools := gjson.GetBytes(body, "tool_choice.tools") + if tools.IsArray() { + for index := range tools.Array() { + if !normalizeAtPath(fmt.Sprintf("tool_choice.tools.%d", index)) { + return original + } + } + } + return body +} + +func normalizeXAITool(tool gjson.Result, namespaceName string) ([]byte, bool, bool) { + toolType := tool.Get("type").String() + changed := false + if toolType == xaiToolSearchType || toolType == xaiImageGenerationToolType { + return nil, true, true + } + if toolType == xaiCustomToolType && tool.Get("name").String() == "apply_patch" { + return nil, true, true + } + + raw := []byte(tool.Raw) + schemaTool := tool + if toolType == xaiFunctionToolType || toolType == xaiCustomToolType { + updatedTool, schemaChanged, ok := normalizeXAIObjectRootUnionBranchTypes(raw) + if !ok { + return nil, false, false + } + raw = updatedTool + if schemaChanged { + schemaTool = gjson.ParseBytes(raw) + changed = true + log.Debugf("xai: added object types to root union branches for tool %s.%s", namespaceName, tool.Get("name").String()) + } + } + if toolType == xaiCustomToolType { + updatedTool, errSet := sjson.SetBytes(raw, "type", xaiFunctionToolType) + if errSet != nil { + return nil, false, false + } + raw = updatedTool + toolType = xaiFunctionToolType + changed = true + } + if toolType == xaiWebSearchToolType && tool.Get("external_web_access").Exists() { + updatedTool, errDel := sjson.DeleteBytes(raw, "external_web_access") + if errDel != nil { + return nil, false, false + } + raw = updatedTool + changed = true + } + if toolType == xaiFunctionToolType && !schemaTool.Get("parameters").Exists() { + updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(`{"type":"object","properties":{}}`)) + if errSet != nil { + return nil, false, false + } + raw = updatedTool + changed = true + } + // Simplify the Codex Desktop automation schema and root unions that xAI + // rejects because function parameters must resolve exclusively to objects. + if toolType == xaiFunctionToolType && xaiFunctionParametersNeedSimplification(schemaTool, namespaceName) { + updatedTool, errSet := sjson.SetRawBytes(raw, "parameters", []byte(xaiSafeFunctionParameters)) + if errSet != nil { + return nil, false, false + } + raw = updatedTool + if strict := tool.Get("strict"); strict.Exists() && strict.Bool() { + updatedTool, errSet = sjson.SetBytes(raw, "strict", false) + if errSet != nil { + return nil, false, false + } + raw = updatedTool + } + changed = true + log.Debugf("xai: simplified parameters for tool %s.%s to avoid upstream schema rejection or hang", namespaceName, tool.Get("name").String()) + } + if toolType == xaiFunctionToolType && strings.TrimSpace(namespaceName) != "" { + qualifiedName := qualifyXAINamespaceToolName(namespaceName, tool.Get("name").String()) + if qualifiedName == "" { + return nil, false, false + } + updatedTool, errSet := sjson.SetBytes(raw, "name", qualifiedName) + if errSet != nil { + return nil, false, false + } + raw = updatedTool + changed = true + } + return raw, changed, true +} + +func qualifyXAINamespaceToolName(namespaceName, toolName string) string { + namespaceName = strings.TrimSpace(namespaceName) + toolName = strings.TrimSpace(toolName) + if namespaceName == "" || toolName == "" || strings.HasPrefix(toolName, "mcp__") { + return toolName + } + prefix := namespaceName + if !strings.HasSuffix(prefix, "__") { + prefix += "__" + } + if strings.HasPrefix(toolName, prefix) { + return toolName + } + return prefix + toolName +} + +func collectXAINamespaceToolRefs(body []byte) map[string]xaiNamespaceToolRef { + refs := make(map[string]xaiNamespaceToolRef) + collect := func(tools gjson.Result) { + if !tools.Exists() || !tools.IsArray() { + return + } + for _, tool := range tools.Array() { + if tool.Get("type").String() != xaiNamespaceToolType { + continue + } + namespaceName := strings.TrimSpace(tool.Get("name").String()) + if namespaceName == "" { + continue + } + for _, nestedTool := range tool.Get("tools").Array() { + toolName := strings.TrimSpace(nestedTool.Get("name").String()) + qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) + if qualifiedName == "" { + continue + } + refs[qualifiedName] = xaiNamespaceToolRef{namespace: namespaceName, name: toolName} + } + } + } + collect(gjson.GetBytes(body, "tools")) + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for _, item := range input.Array() { + if item.Get("type").String() == "additional_tools" { + collect(item.Get("tools")) + } + } + } + return refs +} + +func normalizeXAIInputCustomToolCalls(body []byte) []byte { + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body + } + + changed := false + inputArray := input.Array() + items := make([]json.RawMessage, 0, len(inputArray)) + for _, item := range inputArray { + var normalized []byte + switch item.Get("type").String() { + case "custom_tool_call": + callID := strings.TrimSpace(item.Get("call_id").String()) + name := strings.TrimSpace(item.Get("name").String()) + if callID == "" || name == "" { + changed = true + continue + } + normalized = []byte(`{"type":"function_call"}`) + normalized, _ = sjson.SetBytes(normalized, "call_id", callID) + normalized, _ = sjson.SetBytes(normalized, "name", name) + normalized, _ = sjson.SetBytes(normalized, "arguments", xaiCustomToolCallArguments(item.Get("input"))) + case "custom_tool_call_output": + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID == "" { + changed = true + continue + } + normalized = []byte(`{"type":"function_call_output"}`) + normalized, _ = sjson.SetBytes(normalized, "call_id", callID) + normalized, _ = sjson.SetBytes(normalized, "output", xaiCustomToolCallOutput(item.Get("output"))) + default: + items = append(items, json.RawMessage(item.Raw)) + continue + } + items = append(items, json.RawMessage(normalized)) + changed = true + } + if !changed { + return body + } + + rawInput, errMarshal := json.Marshal(items) + if errMarshal != nil { + return body + } + updated, errSet := sjson.SetRawBytes(body, "input", rawInput) + if errSet != nil { + return body + } + return updated +} + +func xaiCustomToolCallArguments(input gjson.Result) string { + if !input.Exists() { + return "{}" + } + if input.Type == gjson.String { + text := input.String() + trimmed := strings.TrimSpace(text) + if gjson.Valid(trimmed) { + parsed := gjson.Parse(trimmed) + if parsed.IsObject() { + return parsed.Raw + } + } + encoded, errMarshal := json.Marshal(text) + if errMarshal != nil { + return "{}" + } + return `{"input":` + string(encoded) + `}` + } + if input.IsObject() { + return input.Raw + } + if input.Raw != "" { + return `{"input":` + input.Raw + `}` + } + return "{}" +} + +func xaiCustomToolCallOutput(output gjson.Result) string { + if !output.Exists() { + return "" + } + if output.Type == gjson.String { + return output.String() + } + return output.Raw +} diff --git a/internal/runtime/executor/xai_executor_response.go b/internal/runtime/executor/xai_executor_response.go new file mode 100644 index 000000000..a0bdbae08 --- /dev/null +++ b/internal/runtime/executor/xai_executor_response.go @@ -0,0 +1,881 @@ +package executor + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "sort" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/signature" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +// xAI executes these x_search subtools server-side but exposes their trace as +// client-style tool calls. Hide the trace so Responses clients do not execute it again. +type xaiInternalXSearchResponseFilter struct { + enabled bool + clientDeclaredTools map[xaiClientToolKey]struct{} + droppedOutputIndexes map[int64]struct{} + droppedItemIDs map[string]struct{} +} + +func newXAIInternalXSearchResponseFilter(enabled bool, clientDeclaredTools map[xaiClientToolKey]struct{}) *xaiInternalXSearchResponseFilter { + filter := &xaiInternalXSearchResponseFilter{ + enabled: enabled, + clientDeclaredTools: clientDeclaredTools, + } + if enabled { + filter.droppedOutputIndexes = make(map[int64]struct{}) + filter.droppedItemIDs = make(map[string]struct{}) + } + return filter +} + +func xaiRequestHasNativeXSearch(body []byte) bool { + if gjson.GetBytes(body, `tools.#(type=="x_search")`).Exists() { + return true + } + // Multipath queries return an array of matches; an empty array still Exists(). + // Check the match count instead of Exists() for additional_tools injection. + return len(gjson.GetBytes(body, `input.#(type=="additional_tools")#.tools.#(type=="x_search")`).Array()) > 0 +} + +// collectXAIClientDeclaredToolKeys records client-declared function/custom tools +// using the Responses post-restore identity (short name + optional namespace) and +// the effective upstream tool type after normalizeXAITool. Client custom tools +// are normalized to function before being sent to xAI, so keys use function for +// both declaration kinds. Must run before normalizeXAITools flattens namespace wrappers. +func collectXAIClientDeclaredToolKeys(body []byte) map[xaiClientToolKey]struct{} { + keys := make(map[xaiClientToolKey]struct{}) + collect := func(tools gjson.Result) { + if !tools.Exists() || !tools.IsArray() { + return + } + for _, tool := range tools.Array() { + switch toolType := strings.TrimSpace(tool.Get("type").String()); toolType { + case xaiNamespaceToolType: + namespaceName := strings.TrimSpace(tool.Get("name").String()) + if namespaceName == "" { + continue + } + for _, nestedTool := range tool.Get("tools").Array() { + nestedType := strings.TrimSpace(nestedTool.Get("type").String()) + if nestedType != xaiFunctionToolType && nestedType != xaiCustomToolType { + continue + } + toolName := strings.TrimSpace(nestedTool.Get("name").String()) + if toolName == "" { + continue + } + // normalizeXAITool converts custom → function before upstream send. + keys[xaiClientToolKey{namespace: namespaceName, name: toolName, toolType: xaiEffectiveDeclaredToolType(nestedType)}] = struct{}{} + } + case xaiFunctionToolType, xaiCustomToolType: + toolName := strings.TrimSpace(tool.Get("name").String()) + if toolName == "" { + continue + } + // normalizeXAITool converts custom → function before upstream send. + keys[xaiClientToolKey{namespace: "", name: toolName, toolType: xaiEffectiveDeclaredToolType(toolType)}] = struct{}{} + } + } + } + collect(gjson.GetBytes(body, "tools")) + input := gjson.GetBytes(body, "input") + if input.Exists() && input.IsArray() { + for _, item := range input.Array() { + if item.Get("type").String() == "additional_tools" { + collect(item.Get("tools")) + } + } + } + return keys +} + +// xaiEffectiveDeclaredToolType returns the tool type actually sent upstream +// after normalizeXAITool. Client custom tools are rewritten to function. +func xaiEffectiveDeclaredToolType(toolType string) string { + if strings.TrimSpace(toolType) == xaiCustomToolType { + return xaiFunctionToolType + } + return strings.TrimSpace(toolType) +} + +func xaiIsInternalXSearchToolName(name string) bool { + switch strings.TrimSpace(name) { + case "x_user_search", "x_semantic_search", "x_keyword_search", "x_thread_fetch": + return true + default: + return false + } +} + +// xaiResponseCallDeclaredType maps a Responses output call type to the effective +// upstream tool declaration kind used when matching client-declared tools. +// Client custom tools are normalized to function before upstream send, so only +// function_call can match a client-declared same-name tool; custom_tool_call +// remains the internal X Search trace shape. +func xaiResponseCallDeclaredType(itemType string) string { + switch strings.TrimSpace(itemType) { + case "function_call": + return xaiFunctionToolType + case "custom_tool_call": + return xaiCustomToolType + default: + return "" + } +} + +// xaiIsInternalXSearchCallID reports whether call_id matches the evidenced xAI +// X Search server-side trace prefix (xs_call...), as observed in Responses traffic +// for native x_search subtools (see issue #4282 / PR #4284 fixtures). +func xaiIsInternalXSearchCallID(callID string) bool { + return strings.HasPrefix(strings.TrimSpace(callID), "xs_call") +} + +// xaiIsInternalXSearchCall reports whether an output item is an xAI server-side +// X Search subtool trace that should be hidden from Responses clients. +// +// Evidence from xAI Responses traffic (issue #4282 / PR #4284): +// - native x_search subtools are emitted as custom_tool_call items named +// x_user_search / x_semantic_search / x_keyword_search / x_thread_fetch +// - those traces commonly use call_id values prefixed with "xs_call" +// +// Client tools that share a short name are preserved only when the response call +// kind matches the effective upstream declaration type. Because normalizeXAITool +// rewrites client custom → function, a client custom x_keyword_search is keyed as +// function and therefore preserves function_call while still filtering genuine +// internal custom_tool_call / xs_call* traces. Namespaced restored client tools +// are never treated as internal. +func xaiIsInternalXSearchCall(item gjson.Result, clientDeclaredTools map[xaiClientToolKey]struct{}) bool { + itemType := strings.TrimSpace(item.Get("type").String()) + declaredType := xaiResponseCallDeclaredType(itemType) + if declaredType == "" { + return false + } + name := strings.TrimSpace(item.Get("name").String()) + if !xaiIsInternalXSearchToolName(name) { + return false + } + namespace := strings.TrimSpace(item.Get("namespace").String()) + // Namespaced calls are restored client tools, never xAI internal X Search traces. + if namespace != "" { + return false + } + // Evidenced internal call_id prefix always identifies server-side X Search traces, + // even when a client tool reuses the same short name. + if xaiIsInternalXSearchCallID(item.Get("call_id").String()) { + return true + } + // Preserve only client tools whose effective upstream declaration kind matches + // this call type (function_call ↔ function after custom normalization). + if _, declared := clientDeclaredTools[xaiClientToolKey{namespace: namespace, name: name, toolType: declaredType}]; declared { + return false + } + return true +} + +func (f *xaiInternalXSearchResponseFilter) apply(eventData []byte) []byte { + if f == nil || !f.enabled || len(eventData) == 0 || !gjson.ValidBytes(eventData) { + return eventData + } + + if item := gjson.GetBytes(eventData, "item"); xaiIsInternalXSearchCall(item, f.clientDeclaredTools) { + f.recordDroppedItem(eventData, item) + return nil + } + + eventData = f.filterCompletedOutput(eventData) + if f.referencesDroppedItem(eventData) { + return nil + } + return f.compactOutputIndex(eventData) +} + +func (f *xaiInternalXSearchResponseFilter) recordDroppedItem(eventData []byte, item gjson.Result) { + if outputIndex := gjson.GetBytes(eventData, "output_index"); outputIndex.Exists() { + f.droppedOutputIndexes[outputIndex.Int()] = struct{}{} + } + for _, path := range []string{"id", "call_id"} { + if id := strings.TrimSpace(item.Get(path).String()); id != "" { + f.droppedItemIDs[id] = struct{}{} + } + } +} + +func (f *xaiInternalXSearchResponseFilter) referencesDroppedItem(eventData []byte) bool { + if outputIndex := gjson.GetBytes(eventData, "output_index"); outputIndex.Exists() { + if _, dropped := f.droppedOutputIndexes[outputIndex.Int()]; dropped { + return true + } + } + for _, path := range []string{"item_id", "call_id"} { + id := strings.TrimSpace(gjson.GetBytes(eventData, path).String()) + if _, dropped := f.droppedItemIDs[id]; id != "" && dropped { + return true + } + } + return false +} + +func (f *xaiInternalXSearchResponseFilter) compactOutputIndex(eventData []byte) []byte { + outputIndex := gjson.GetBytes(eventData, "output_index") + if !outputIndex.Exists() { + return eventData + } + original := outputIndex.Int() + removedBefore := int64(0) + for dropped := range f.droppedOutputIndexes { + if dropped < original { + removedBefore++ + } + } + if removedBefore == 0 { + return eventData + } + updated, errSet := sjson.SetBytes(eventData, "output_index", original-removedBefore) + if errSet != nil { + return eventData + } + return updated +} + +func (f *xaiInternalXSearchResponseFilter) filterCompletedOutput(eventData []byte) []byte { + output := gjson.GetBytes(eventData, "response.output") + if !output.IsArray() { + return eventData + } + var clientDeclaredTools map[xaiClientToolKey]struct{} + if f != nil { + clientDeclaredTools = f.clientDeclaredTools + } + items := make([]json.RawMessage, 0, len(output.Array())) + changed := false + for _, item := range output.Array() { + if xaiIsInternalXSearchCall(item, clientDeclaredTools) { + changed = true + continue + } + items = append(items, json.RawMessage(item.Raw)) + } + if !changed { + return eventData + } + rawOutput, errMarshal := json.Marshal(items) + if errMarshal != nil { + return eventData + } + updated, errSet := sjson.SetRawBytes(eventData, "response.output", rawOutput) + if errSet != nil { + return eventData + } + return updated +} + +func normalizeXAIInputNamespaceToolCalls(body []byte) []byte { + if !gjson.ValidBytes(body) { + return body + } + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body + } + for index, item := range input.Array() { + if item.Get("type").String() != "function_call" { + continue + } + namespaceName := strings.TrimSpace(item.Get("namespace").String()) + toolName := strings.TrimSpace(item.Get("name").String()) + qualifiedName := qualifyXAINamespaceToolName(namespaceName, toolName) + if namespaceName == "" || qualifiedName == "" { + continue + } + namePath := fmt.Sprintf("input.%d.name", index) + namespacePath := fmt.Sprintf("input.%d.namespace", index) + updated, errSet := sjson.SetBytes(body, namePath, qualifiedName) + if errSet != nil { + continue + } + updated, errDelete := sjson.DeleteBytes(updated, namespacePath) + if errDelete != nil { + continue + } + body = updated + } + return body +} + +func restoreXAINamespaceToolCalls(data []byte, refs map[string]xaiNamespaceToolRef) []byte { + if len(refs) == 0 || len(data) == 0 || !gjson.ValidBytes(data) { + return data + } + data = restoreXAINamespaceToolCallAtPath(data, "item", refs) + output := gjson.GetBytes(data, "response.output") + if output.Exists() && output.IsArray() { + for index := range output.Array() { + data = restoreXAINamespaceToolCallAtPath(data, fmt.Sprintf("response.output.%d", index), refs) + } + } + return data +} + +func restoreXAINamespaceToolCallAtPath(data []byte, path string, refs map[string]xaiNamespaceToolRef) []byte { + if gjson.GetBytes(data, path+".type").String() != "function_call" { + return data + } + qualifiedName := strings.TrimSpace(gjson.GetBytes(data, path+".name").String()) + ref, ok := refs[qualifiedName] + if !ok { + return data + } + updated, errSet := sjson.SetBytes(data, path+".name", ref.name) + if errSet != nil { + return data + } + updated, errSet = sjson.SetBytes(updated, path+".namespace", ref.namespace) + if errSet != nil { + return data + } + return updated +} + +// normalizeXAIObjectRootUnionBranchTypes makes untyped root union branches +// explicitly object-only when the parameter root already permits only objects. +// This preserves the original schema semantics while satisfying xAI validation. +func normalizeXAIObjectRootUnionBranchTypes(tool []byte) ([]byte, bool, bool) { + parameters := gjson.GetBytes(tool, "parameters") + rootType := parameters.Get("type") + if rootType.Type != gjson.String || rootType.String() != "object" { + return tool, false, true + } + + original := tool + changed := false + for _, unionName := range []string{"anyOf", "oneOf"} { + union := parameters.Get(unionName) + if !union.IsArray() { + continue + } + for index, branch := range union.Array() { + if !branch.IsObject() || branch.Get("type").Exists() { + continue + } + updated, errSet := sjson.SetBytes(tool, fmt.Sprintf("parameters.%s.%d.type", unionName, index), "object") + if errSet != nil { + return original, false, false + } + tool = updated + changed = true + } + } + return tool, changed, true +} + +func xaiSchemaTypeIsObjectOnly(schemaType gjson.Result) bool { + if schemaType.Type == gjson.String { + return strings.EqualFold(strings.TrimSpace(schemaType.String()), "object") + } + if !schemaType.IsArray() { + return false + } + types := schemaType.Array() + if len(types) == 0 { + return false + } + for _, schemaTypeItem := range types { + if schemaTypeItem.Type != gjson.String || !strings.EqualFold(strings.TrimSpace(schemaTypeItem.String()), "object") { + return false + } + } + return true +} + +// xaiFunctionParametersNeedSimplification reports whether a function tool, or +// a custom tool normalized to a function, has a schema that xAI cannot accept. +func xaiFunctionParametersNeedSimplification(tool gjson.Result, namespaceName string) bool { + toolType := strings.TrimSpace(tool.Get("type").String()) + isFunction := strings.EqualFold(toolType, xaiFunctionToolType) + isNormalizedCustom := strings.EqualFold(toolType, xaiCustomToolType) + if !isFunction && !isNormalizedCustom { + return false + } + + toolName := strings.TrimSpace(tool.Get("name").String()) + qualifiedAutomationName := xaiCodexAppNamespaceName + "__" + xaiAutomationUpdateToolName + if isFunction && (strings.EqualFold(toolName, qualifiedAutomationName) || + (strings.EqualFold(strings.TrimSpace(namespaceName), xaiCodexAppNamespaceName) && + strings.EqualFold(toolName, xaiAutomationUpdateToolName))) { + return true + } + + parameters := tool.Get("parameters") + for _, unionName := range []string{"anyOf", "oneOf"} { + union := parameters.Get(unionName) + if !union.IsArray() { + continue + } + for _, branch := range union.Array() { + if !xaiSchemaTypeIsObjectOnly(branch.Get("type")) { + return true + } + } + } + return false +} + +func sanitizeXAIInputEncryptedContent(body []byte) []byte { + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body + } + items := make([]json.RawMessage, 0, len(input.Array())) + changed := false + dropCount := 0 + firstReason := "" + firstItemType := "" + for _, item := range input.Array() { + itemType := strings.TrimSpace(item.Get("type").String()) + if itemType != "reasoning" && itemType != "compaction" { + items = append(items, json.RawMessage(item.Raw)) + continue + } + encryptedContent := item.Get("encrypted_content") + if !encryptedContent.Exists() { + items = append(items, json.RawMessage(item.Raw)) + continue + } + reason := "" + switch encryptedContent.Type { + case gjson.String: + if _, err := signature.InspectGrokEncryptedContent(encryptedContent.String()); err != nil { + reason = err.Error() + } + case gjson.Null: + reason = "encrypted_content is null" + default: + reason = fmt.Sprintf("encrypted_content must be a string, got %s", encryptedContent.Type.String()) + } + if reason == "" { + items = append(items, json.RawMessage(item.Raw)) + continue + } + + if itemType == "compaction" { + changed = true + dropCount++ + if firstReason == "" { + firstReason = reason + firstItemType = itemType + } + continue + } + + next, err := sjson.DeleteBytes([]byte(item.Raw), "encrypted_content") + if err != nil { + items = append(items, json.RawMessage(item.Raw)) + continue + } + items = append(items, json.RawMessage(next)) + changed = true + dropCount++ + if firstReason == "" { + firstReason = reason + firstItemType = itemType + } + } + if !changed { + return body + } + rawInput, err := json.Marshal(items) + if err != nil { + return body + } + updated, err := sjson.SetRawBytes(body, "input", rawInput) + if err != nil { + return body + } + if dropCount > 0 { + log.WithFields(log.Fields{ + "component": "xai_encrypted_content_sanitizer", + "dropped": dropCount, + "first_item_type": firstItemType, + "first_reason": firstReason, + }).Debug("xai executor: removed invalid encrypted_content before upstream") + } + return mergeAdjacentXAIInputReasoningSummaries(updated) +} + +func normalizeXAIInputReasoningItems(body []byte) []byte { + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body + } + + updated := body + for i, item := range input.Array() { + if item.Get("type").String() != "reasoning" { + continue + } + contentPath := fmt.Sprintf("input.%d.content", i) + if content := gjson.GetBytes(updated, contentPath); content.Exists() && content.Type == gjson.Null { + updatedBody, errDel := sjson.DeleteBytes(updated, contentPath) + if errDel != nil { + return body + } + updated = updatedBody + } + encryptedContentPath := fmt.Sprintf("input.%d.encrypted_content", i) + if encryptedContent := gjson.GetBytes(updated, encryptedContentPath); encryptedContent.Exists() && encryptedContent.Type == gjson.Null { + updatedBody, errDel := sjson.DeleteBytes(updated, encryptedContentPath) + if errDel != nil { + return body + } + updated = updatedBody + } + } + return mergeAdjacentXAIInputReasoningSummaries(updated) +} + +func mergeAdjacentXAIInputReasoningSummaries(body []byte) []byte { + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body + } + + changed := false + items := make([]json.RawMessage, 0, len(input.Array())) + for _, item := range input.Array() { + if len(items) > 0 && canMergeXAIReasoningSummary(items[len(items)-1], item) { + merged, ok := appendXAIReasoningSummary(items[len(items)-1], item.Get("summary").Array()) + if ok { + items[len(items)-1] = json.RawMessage(merged) + changed = true + continue + } + } + items = append(items, json.RawMessage(item.Raw)) + } + if !changed { + return body + } + + rawInput, errMarshal := json.Marshal(items) + if errMarshal != nil { + return body + } + updated, errSet := sjson.SetRawBytes(body, "input", rawInput) + if errSet != nil { + return body + } + return updated +} + +func canMergeXAIReasoningSummary(previous json.RawMessage, current gjson.Result) bool { + previousItem := gjson.ParseBytes(previous) + if previousItem.Get("type").String() != "reasoning" || current.Get("type").String() != "reasoning" { + return false + } + if !previousItem.Get("summary").IsArray() || !current.Get("summary").IsArray() { + return false + } + if len(current.Get("summary").Array()) == 0 { + return false + } + for name := range current.Map() { + if name != "type" && name != "summary" { + return false + } + } + return true +} + +func appendXAIReasoningSummary(previous json.RawMessage, currentSummary []gjson.Result) ([]byte, bool) { + updated := []byte(previous) + summary := gjson.GetBytes(updated, "summary") + if !summary.IsArray() { + return previous, false + } + nextIndex := len(summary.Array()) + for i, item := range currentSummary { + updatedItem, errSet := sjson.SetRawBytes(updated, fmt.Sprintf("summary.%d", nextIndex+i), []byte(item.Raw)) + if errSet != nil { + return previous, false + } + updated = updatedItem + } + return updated, true +} + +// xaiSupportsReasoningEffort reports whether the model accepts Responses API +// reasoning.effort. Capability comes from model registry thinking metadata +// (static models.json and dynamic registrations), not a hard-coded name allowlist. +func xaiSupportsReasoningEffort(model string) bool { + name := strings.ToLower(strings.TrimSpace(thinking.ParseSuffix(model).ModelName)) + if idx := strings.LastIndex(name, "/"); idx >= 0 { + name = name[idx+1:] + } + if name == "" { + return false + } + info := registry.LookupModelInfo(name, "xai") + if info == nil || info.Thinking == nil { + return false + } + return len(info.Thinking.Levels) > 0 +} + +func xaiNormalizeReasoningSummaryEventLine(line []byte, eventName string) []byte { + if eventName == "" && bytes.HasPrefix(line, xaiEventTag) { + eventName = strings.TrimSpace(string(line[len(xaiEventTag):])) + } + eventName = xaiNormalizeReasoningSummaryEventName(eventName) + if eventName == "" { + return bytes.Clone(line) + } + return []byte("event: " + eventName) +} + +func xaiNormalizeReasoningSummaryEventName(eventName string) string { + switch eventName { + case "response.reasoning_text.delta": + return "response.reasoning_summary_text.delta" + case "response.reasoning_text.done": + return "response.reasoning_summary_part.done" + default: + return eventName + } +} + +func xaiNormalizeReasoningSummaryData(eventData []byte) []byte { + if len(eventData) == 0 || !gjson.ValidBytes(eventData) { + return eventData + } + + normalized := eventData + switch gjson.GetBytes(normalized, "type").String() { + case "response.reasoning_text.delta": + normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_text.delta") + normalized = xaiNormalizeReasoningSummaryIndex(normalized) + case "response.reasoning_text.done": + normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_part.done") + normalized, _ = sjson.SetBytes(normalized, "part.type", "summary_text") + if text := gjson.GetBytes(normalized, "text"); text.Exists() { + normalized, _ = sjson.SetBytes(normalized, "part.text", text.String()) + } + normalized, _ = sjson.DeleteBytes(normalized, "text") + normalized = xaiNormalizeReasoningSummaryIndex(normalized) + case "response.content_part.added": + if gjson.GetBytes(normalized, "part.type").String() == "reasoning_text" { + normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_part.added") + normalized, _ = sjson.SetBytes(normalized, "part.type", "summary_text") + normalized = xaiNormalizeReasoningSummaryIndex(normalized) + } + case "response.content_part.done": + if gjson.GetBytes(normalized, "part.type").String() == "reasoning_text" { + normalized, _ = sjson.SetBytes(normalized, "type", "response.reasoning_summary_part.done") + normalized, _ = sjson.SetBytes(normalized, "part.type", "summary_text") + normalized = xaiNormalizeReasoningSummaryIndex(normalized) + } + } + + if item := gjson.GetBytes(normalized, "item"); item.Exists() && item.Type == gjson.JSON { + updatedItem := xaiNormalizeReasoningOutputItem([]byte(item.Raw)) + if !bytes.Equal(updatedItem, []byte(item.Raw)) { + normalized, _ = sjson.SetRawBytes(normalized, "item", updatedItem) + } + } + if output := gjson.GetBytes(normalized, "response.output"); output.IsArray() { + updatedOutput, changed := xaiNormalizeReasoningOutputItems(output.Array()) + if changed { + normalized, _ = sjson.SetRawBytes(normalized, "response.output", updatedOutput) + } + } + + return normalized +} + +func xaiNormalizeReasoningSummaryDataEvents(eventData []byte) [][]byte { + if len(eventData) == 0 || !gjson.ValidBytes(eventData) { + return [][]byte{eventData} + } + if gjson.GetBytes(eventData, "type").String() != "response.reasoning_text.done" { + return [][]byte{xaiNormalizeReasoningSummaryData(eventData)} + } + + textDone, _ := sjson.SetBytes(eventData, "type", "response.reasoning_summary_text.done") + textDone = xaiNormalizeReasoningSummaryIndex(textDone) + partDone := xaiNormalizeReasoningSummaryData(eventData) + return [][]byte{textDone, partDone} +} + +func xaiNormalizeReasoningSummaryIndex(eventData []byte) []byte { + contentIndex := gjson.GetBytes(eventData, "content_index") + if contentIndex.Exists() && contentIndex.Raw != "" && !gjson.GetBytes(eventData, "summary_index").Exists() { + eventData, _ = sjson.SetRawBytes(eventData, "summary_index", []byte(contentIndex.Raw)) + } + eventData, _ = sjson.DeleteBytes(eventData, "content_index") + return eventData +} + +func xaiNormalizeReasoningOutputItems(items []gjson.Result) ([]byte, bool) { + var buf bytes.Buffer + buf.WriteByte('[') + changed := false + for i, item := range items { + if i > 0 { + buf.WriteByte(',') + } + updatedItem := xaiNormalizeReasoningOutputItem([]byte(item.Raw)) + if !bytes.Equal(updatedItem, []byte(item.Raw)) { + changed = true + } + buf.Write(updatedItem) + } + buf.WriteByte(']') + return buf.Bytes(), changed +} + +func xaiNormalizeReasoningOutputItem(item []byte) []byte { + if !gjson.ValidBytes(item) || gjson.GetBytes(item, "type").String() != "reasoning" { + return item + } + + normalized := item + if summary := gjson.GetBytes(normalized, "summary"); summary.IsArray() { + updatedSummary, changed := xaiNormalizeReasoningSummaryItems(summary.Array()) + if changed { + normalized, _ = sjson.SetRawBytes(normalized, "summary", updatedSummary) + } + } + + content := gjson.GetBytes(normalized, "content") + if !content.IsArray() { + return normalized + } + + summaryItems := make([]gjson.Result, 0, len(content.Array())) + for _, part := range content.Array() { + if part.Get("type").String() == "reasoning_text" { + summaryItems = append(summaryItems, part) + } + } + if len(summaryItems) == 0 { + return normalized + } + + updatedSummary, _ := xaiNormalizeReasoningSummaryItems(summaryItems) + normalized, _ = sjson.SetRawBytes(normalized, "summary", updatedSummary) + normalized, _ = sjson.DeleteBytes(normalized, "content") + return normalized +} + +func xaiNormalizeReasoningSummaryItems(items []gjson.Result) ([]byte, bool) { + var buf bytes.Buffer + buf.WriteByte('[') + changed := false + for i, item := range items { + if i > 0 { + buf.WriteByte(',') + } + itemRaw := []byte(item.Raw) + if item.Get("type").String() == "reasoning_text" { + var errSet error + itemRaw, errSet = sjson.SetBytes(itemRaw, "type", "summary_text") + if errSet == nil { + changed = true + } + } + buf.Write(itemRaw) + } + buf.WriteByte(']') + return buf.Bytes(), changed +} + +func xaiCollectOutputItemDone(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback *[][]byte) { + itemResult := gjson.GetBytes(eventData, "item") + if !itemResult.Exists() || itemResult.Type != gjson.JSON { + return + } + outputIndexResult := gjson.GetBytes(eventData, "output_index") + if outputIndexResult.Exists() { + outputItemsByIndex[outputIndexResult.Int()] = []byte(itemResult.Raw) + return + } + *outputItemsFallback = append(*outputItemsFallback, []byte(itemResult.Raw)) +} + +func xaiPatchCompletedOutput(eventData []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { + outputResult := gjson.GetBytes(eventData, "response.output") + shouldPatchOutput := (!outputResult.Exists() || !outputResult.IsArray() || len(outputResult.Array()) == 0) && (len(outputItemsByIndex) > 0 || len(outputItemsFallback) > 0) + if !shouldPatchOutput { + return eventData + } + + indexes := make([]int64, 0, len(outputItemsByIndex)) + for idx := range outputItemsByIndex { + indexes = append(indexes, idx) + } + sort.Slice(indexes, func(i, j int) bool { + return indexes[i] < indexes[j] + }) + + outputArray := []byte("[]") + var buf bytes.Buffer + buf.WriteByte('[') + wrote := false + for _, idx := range indexes { + if wrote { + buf.WriteByte(',') + } + buf.Write(outputItemsByIndex[idx]) + wrote = true + } + for _, item := range outputItemsFallback { + if wrote { + buf.WriteByte(',') + } + buf.Write(item) + wrote = true + } + buf.WriteByte(']') + if wrote { + outputArray = buf.Bytes() + } + + patched, _ := sjson.SetRawBytes(eventData, "response.output", outputArray) + return patched +} + +// xaiFreeUsageExhaustedCooldown is the free-tier rolling window advertised by +// cli-chat-proxy ("Usage resets over a rolling 24-hour window"). +const xaiFreeUsageExhaustedCooldown = 24 * time.Hour + +// xaiStatusErr wraps upstream error bodies so free-tier exhaustion +// (subscription:free-usage-exhausted) carries a 24h RetryAfter hint for +// auth cooldown / account rotation. Generic 429s stay without an explicit +// retry hint so conductor backoff still applies. +func xaiStatusErr(code int, body []byte) statusErr { + err := statusErr{code: code, msg: string(body)} + if code != http.StatusTooManyRequests || len(body) == 0 { + return err + } + codeStr := strings.ToLower(gjson.GetBytes(body, "code").String()) + msg := strings.ToLower(gjson.GetBytes(body, "error").String()) + if msg == "" { + msg = strings.ToLower(string(body)) + } + if strings.Contains(codeStr, "free-usage-exhausted") || + strings.Contains(msg, "free-usage-exhausted") || + strings.Contains(msg, "included free usage") { + d := xaiFreeUsageExhaustedCooldown + err.retryAfter = &d + } + return err +} diff --git a/internal/runtime/executor/xai_executor_stream.go b/internal/runtime/executor/xai_executor_stream.go new file mode 100644 index 000000000..5ccbe292c --- /dev/null +++ b/internal/runtime/executor/xai_executor_stream.go @@ -0,0 +1,174 @@ +package executor + +import ( + "bufio" + "bytes" + "context" + "io" + "net/http" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (_ *cliproxyexecutor.StreamResult, err error) { + if opts.Alt == "responses/compact" { + return nil, statusErr{code: http.StatusBadRequest, msg: "streaming not supported for /responses/compact"} + } + if xaiInputHasItemType(req.Payload, "compaction_trigger") { + return e.executeCompactionTriggerStream(ctx, auth, req, opts) + } + + token, _ := xaiCreds(auth) + baseURL := xaiChatBaseURL(auth) + logXAIResolvedBaseURL(ctx, baseURL) + + prepared, err := e.prepareResponsesRequest(ctx, req, opts, true) + if err != nil { + return nil, err + } + + reporter := helps.NewExecutorUsageReporter(ctx, e, prepared.baseModel, auth) + defer reporter.TrackFailure(ctx, &err) + reporter.SetTranslatedReasoningEffort(prepared.body, e.Identifier()) + + url := strings.TrimSuffix(baseURL, "/") + "/responses" + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(prepared.body)) + if err != nil { + return nil, err + } + applyXAIChatHeaders(httpReq, auth, token, true, prepared.sessionID) + e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), prepared.body) + + httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0) + httpClient = reporter.TrackHTTPClient(httpClient) + httpResp, err := httpClient.Do(httpReq) + if err != nil { + helps.RecordAPIResponseError(ctx, e.cfg, err) + return nil, err + } + helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone()) + if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 { + data, errRead := io.ReadAll(httpResp.Body) + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("xai executor: close response body error: %v", errClose) + } + if errRead != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errRead) + return nil, errRead + } + helps.AppendAPIResponseChunk(ctx, e.cfg, data) + helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data)) + return nil, xaiStatusErr(httpResp.StatusCode, data) + } + + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(out) + defer func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("xai executor: close response body error: %v", errClose) + } + }() + scanner := bufio.NewScanner(httpResp.Body) + scanner.Buffer(nil, 52_428_800) + claudeInputTokens := helps.NewClaudeInputTokenState(prepared.from, prepared.to, prepared.responseFormat, prepared.originalPayload) + var param any + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools) + var pendingEventLine []byte + emitTranslatedLine := func(translatedLine []byte) bool { + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, translatedLine, ¶m, claudeInputTokens) + for i := range chunks { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: + case <-ctx.Done(): + return false + } + } + return true + } + for scanner.Scan() { + line := scanner.Bytes() + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + + if bytes.HasPrefix(line, xaiEventTag) { + if pendingEventLine != nil && !emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) { + return + } + pendingEventLine = bytes.Clone(line) + continue + } + + if bytes.HasPrefix(line, xaiDataTag) { + eventDataList := xaiNormalizeReasoningSummaryDataEvents(bytes.TrimSpace(line[len(xaiDataTag):])) + hasPendingEventLine := pendingEventLine != nil + for i, eventData := range eventDataList { + eventData = restoreXAINamespaceToolCalls(eventData, prepared.namespaceTools) + eventData = responseFilter.apply(eventData) + if len(eventData) == 0 { + if hasPendingEventLine && i == 0 { + pendingEventLine = nil + } + continue + } + normalizedEventName := gjson.GetBytes(eventData, "type").String() + switch normalizedEventName { + case "response.output_item.done": + xaiCollectOutputItemDone(eventData, outputItemsByIndex, &outputItemsFallback) + case "response.completed": + if detail, ok := helps.ParseCodexUsage(eventData); ok { + reporter.Publish(ctx, detail) + } + eventData = xaiPatchCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) + eventData = xaiNormalizeReasoningSummaryData(eventData) + cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, eventData) + normalizedEventName = gjson.GetBytes(eventData, "type").String() + } + + if hasPendingEventLine { + eventLine := []byte("event: " + normalizedEventName) + if i == 0 { + eventLine = xaiNormalizeReasoningSummaryEventLine(pendingEventLine, normalizedEventName) + pendingEventLine = nil + } + if !emitTranslatedLine(eventLine) { + return + } + } + if !emitTranslatedLine(append([]byte("data: "), eventData...)) { + return + } + } + continue + } + + if pendingEventLine != nil { + if !emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) { + return + } + pendingEventLine = nil + } + if !emitTranslatedLine(bytes.Clone(line)) { + return + } + } + if pendingEventLine != nil { + emitTranslatedLine(xaiNormalizeReasoningSummaryEventLine(pendingEventLine, "")) + } + if errScan := scanner.Err(); errScan != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + select { + case out <- cliproxyexecutor.StreamChunk{Err: errScan}: + case <-ctx.Done(): + } + } + }() + return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil +} diff --git a/internal/runtime/executor/xai_executor_tokens.go b/internal/runtime/executor/xai_executor_tokens.go new file mode 100644 index 000000000..0eebb6f54 --- /dev/null +++ b/internal/runtime/executor/xai_executor_tokens.go @@ -0,0 +1,149 @@ +package executor + +import ( + "context" + "fmt" + "strings" + + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "github.com/tidwall/gjson" + "github.com/tiktoken-go/tokenizer" +) + +// CountTokens estimates token count for xAI Responses requests. +func (e *XAIExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + prepared, err := e.prepareResponsesRequest(ctx, req, opts, false) + if err != nil { + return cliproxyexecutor.Response{}, err + } + enc, err := tokenizer.Get(tokenizer.O200kBase) + if err != nil { + return cliproxyexecutor.Response{}, fmt.Errorf("xai executor: tokenizer init failed: %w", err) + } + count, err := countXAIInputTokens(enc, prepared.body) + if err != nil { + return cliproxyexecutor.Response{}, fmt.Errorf("xai executor: token counting failed: %w", err) + } + usageJSON := fmt.Sprintf(`{"response":{"usage":{"input_tokens":%d,"output_tokens":0,"total_tokens":%d}}}`, count, count) + translated := sdktranslator.TranslateTokenCount(ctx, prepared.to, prepared.responseFormat, count, []byte(usageJSON)) + return cliproxyexecutor.Response{Payload: translated}, nil +} + +func countXAIInputTokens(enc tokenizer.Codec, body []byte) (int64, error) { + if enc == nil { + return 0, fmt.Errorf("encoder is nil") + } + if len(body) == 0 { + return 0, nil + } + + root := gjson.ParseBytes(body) + segments := make([]string, 0, 32) + xaiAppendTokenString(&segments, root.Get("instructions")) + xaiCollectInputTokenSegments(root.Get("input"), &segments) + xaiCollectToolTokenSegments(root.Get("tools"), &segments) + + textFormat := root.Get("text.format") + if textFormat.Exists() { + xaiAppendTokenString(&segments, textFormat.Get("name")) + xaiAppendTokenJSON(&segments, textFormat.Get("schema")) + } + + if len(segments) == 0 { + return 0, nil + } + count, err := enc.Count(strings.Join(segments, "\n")) + if err != nil { + return 0, err + } + return int64(count), nil +} + +func xaiCollectInputTokenSegments(input gjson.Result, segments *[]string) { + if input.Type == gjson.String { + xaiAppendTokenString(segments, input) + return + } + if !input.IsArray() { + return + } + for _, item := range input.Array() { + switch item.Get("type").String() { + case "message": + xaiCollectContentTokenSegments(item.Get("content"), segments) + case "function_call": + xaiAppendTokenString(segments, item.Get("name")) + xaiAppendTokenJSON(segments, item.Get("arguments")) + case "function_call_output": + xaiAppendTokenJSON(segments, item.Get("output")) + case "reasoning": + for _, part := range item.Get("summary").Array() { + xaiAppendTokenString(segments, part.Get("text")) + } + } + } +} + +func xaiCollectContentTokenSegments(content gjson.Result, segments *[]string) { + if content.Type == gjson.String { + xaiAppendTokenString(segments, content) + return + } + if !content.IsArray() { + return + } + for _, part := range content.Array() { + switch part.Get("type").String() { + case "text", "input_text", "output_text": + xaiAppendTokenString(segments, part.Get("text")) + case "refusal": + xaiAppendTokenString(segments, part.Get("refusal")) + case "input_image": + xaiAppendTokenString(segments, part.Get("image_url")) + xaiAppendTokenString(segments, part.Get("file_id")) + case "input_file": + xaiAppendTokenString(segments, part.Get("file_data")) + xaiAppendTokenString(segments, part.Get("file_url")) + xaiAppendTokenString(segments, part.Get("file_id")) + xaiAppendTokenString(segments, part.Get("filename")) + case "input_audio": + xaiAppendTokenString(segments, part.Get("data")) + xaiAppendTokenString(segments, part.Get("input_audio.data")) + } + } +} + +func xaiCollectToolTokenSegments(tools gjson.Result, segments *[]string) { + if !tools.IsArray() { + return + } + for _, tool := range tools.Array() { + if tool.Get("type").String() != xaiFunctionToolType { + continue + } + xaiAppendTokenString(segments, tool.Get("name")) + xaiAppendTokenString(segments, tool.Get("description")) + xaiAppendTokenJSON(segments, tool.Get("parameters")) + } +} + +func xaiAppendTokenString(segments *[]string, value gjson.Result) { + if text := strings.TrimSpace(value.String()); text != "" { + *segments = append(*segments, text) + } +} + +func xaiAppendTokenJSON(segments *[]string, value gjson.Result) { + if !value.Exists() { + return + } + if value.Type == gjson.String { + xaiAppendTokenString(segments, value) + return + } + if text := strings.TrimSpace(value.Raw); text != "" { + *segments = append(*segments, text) + } +} diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index f7a168480..3379c0ccf 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -6,10 +6,8 @@ package handlers import ( "bytes" "encoding/json" - "errors" "fmt" "net/http" - "net/url" "reflect" "strings" "sync" @@ -17,18 +15,14 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" - . "github.com/router-for-me/CLIProxyAPI/v7/internal/constant" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" coresession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" - sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson" "golang.org/x/net/context" ) @@ -63,127 +57,6 @@ const ( maxStreamInterceptorHistoryBytes = 1 << 20 ) -type pinnedAuthContextKey struct{} -type selectedAuthCallbackContextKey struct{} -type preparedModelRouteContextKey struct{} -type executionSessionContextKey struct{} -type disallowFreeAuthContextKey struct{} - -// PluginInterceptorHost applies plugin interceptors around handler execution. -type PluginInterceptorHost interface { - InterceptRequestBeforeAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse - InterceptRequestAfterAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse - InterceptResponse(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse - InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse -} - -type pluginInterceptorSkipHost interface { - InterceptRequestBeforeAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse - InterceptRequestAfterAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse - InterceptResponseExcept(context.Context, pluginapi.ResponseInterceptRequest, string) pluginapi.ResponseInterceptResponse - InterceptStreamChunkExcept(context.Context, pluginapi.StreamChunkInterceptRequest, string) pluginapi.StreamChunkInterceptResponse -} - -type streamInterceptorDetector interface { - HasStreamInterceptors() bool -} - -type requestInterceptorDetector interface { - HasRequestInterceptors() bool -} - -// PluginModelRouterHost routes matching requests to a plugin executor, the router's own executor, -// or a built-in provider before model-to-provider resolution and auth selection. -type PluginModelRouterHost interface { - RouteModel(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) -} - -// PluginExecutorHost executes a routed request with a specific plugin executor. -type PluginExecutorHost interface { - ExecutePluginExecutor(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) - ExecutePluginExecutorStream(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) - CountPluginExecutor(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) -} - -type pluginExecutorFormatResolver interface { - PluginExecutorRequestToFormat(string, coreexecutor.Request, coreexecutor.Options) sdktranslator.Format -} - -type pluginModelRouterSkipHost interface { - RouteModelExcept(context.Context, pluginapi.ModelRouteRequest, string) (pluginapi.ModelRouteResponse, bool) -} - -type modelRouterDetector interface { - HasModelRouters() bool -} - -type modelRouterSkipDetector interface { - HasModelRoutersExcept(string) bool -} - -// WithPinnedAuthID returns a child context that requests execution on a specific auth ID. -func WithPinnedAuthID(ctx context.Context, authID string) context.Context { - authID = strings.TrimSpace(authID) - if authID == "" { - return ctx - } - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, pinnedAuthContextKey{}, authID) -} - -// WithSelectedAuthIDCallback returns a child context that receives the selected auth ID. -func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) context.Context { - if callback == nil { - return ctx - } - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, selectedAuthCallbackContextKey{}, callback) -} - -// PrepareStreamModelRoute resolves a stream route once and stores it on the returned context for execution. -// The boolean reports whether the route overrides normal model-to-provider resolution. -func (h *BaseAPIHandler) PrepareStreamModelRoute(ctx context.Context, handlerType string, modelName string, rawJSON []byte) (context.Context, bool) { - if ctx == nil { - ctx = context.Background() - } - decision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, true, modelExecutionOptions{}) - ctx = context.WithValue(ctx, preparedModelRouteContextKey{}, decision) - hasOverride := strings.TrimSpace(decision.ExecutorPluginID) != "" || strings.TrimSpace(decision.Provider) != "" - return ctx, hasOverride -} - -func preparedModelRouteFromContext(ctx context.Context) (modelRouteDecision, bool) { - if ctx == nil { - return modelRouteDecision{}, false - } - decision, ok := ctx.Value(preparedModelRouteContextKey{}).(modelRouteDecision) - return decision, ok -} - -// WithExecutionSessionID returns a child context tagged with a long-lived execution session ID. -func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context { - sessionID = strings.TrimSpace(sessionID) - if sessionID == "" { - return ctx - } - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, executionSessionContextKey{}, sessionID) -} - -// WithDisallowFreeAuth returns a child context that requests skipping known free-tier credentials. -func WithDisallowFreeAuth(ctx context.Context) context.Context { - if ctx == nil { - ctx = context.Background() - } - return context.WithValue(ctx, disallowFreeAuthContextKey{}, true) -} - // BuildErrorResponseBody builds an OpenAI-compatible JSON error response body. // If errText is already valid JSON, it is returned as-is to preserve upstream error payloads. func BuildErrorResponseBody(status int, errText string) []byte { @@ -386,82 +259,6 @@ func setGenerateMetadata(meta map[string]any, rawJSON []byte) { meta[coreexecutor.GenerateMetadataKey] = generate } -// headersFromContext extracts the original HTTP request headers from the gin context -// embedded in the provided context. This allows session affinity selectors to read -// client-provided session headers. -func headersFromContext(ctx context.Context) http.Header { - if ctx == nil { - return nil - } - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { - return ginCtx.Request.Header.Clone() - } - return nil -} - -// queryFromContext extracts the original HTTP request query parameters from the -// gin context embedded in the provided context. Mirrors headersFromContext so -// model routers can observe inbound query parameters for plain HTTP requests, -// where execOptions.Query is not populated by callers. -func queryFromContext(ctx context.Context) url.Values { - if ctx == nil { - return nil - } - if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil && ginCtx.Request.URL != nil { - return ginCtx.Request.URL.Query() - } - return nil -} - -func pinnedAuthIDFromContext(ctx context.Context) string { - if ctx == nil { - return "" - } - raw := ctx.Value(pinnedAuthContextKey{}) - switch v := raw.(type) { - case string: - return strings.TrimSpace(v) - case []byte: - return strings.TrimSpace(string(v)) - default: - return "" - } -} - -func selectedAuthIDCallbackFromContext(ctx context.Context) func(string) { - if ctx == nil { - return nil - } - raw := ctx.Value(selectedAuthCallbackContextKey{}) - if callback, ok := raw.(func(string)); ok && callback != nil { - return callback - } - return nil -} - -func executionSessionIDFromContext(ctx context.Context) string { - if ctx == nil { - return "" - } - raw := ctx.Value(executionSessionContextKey{}) - switch v := raw.(type) { - case string: - return strings.TrimSpace(v) - case []byte: - return strings.TrimSpace(string(v)) - default: - return "" - } -} - -func disallowFreeAuthFromContext(ctx context.Context) bool { - if ctx == nil { - return false - } - raw, ok := ctx.Value(disallowFreeAuthContextKey{}).(bool) - return ok && raw -} - // BaseAPIHandler contains the handlers for API endpoints. // It holds a pool of clients to interact with the backend service and manages // load balancing, client selection, and configuration. @@ -760,1586 +557,6 @@ func appendAPIResponse(c *gin.Context, data []byte) { c.Set("API_RESPONSE", bytes.Clone(data)) } -// ExecuteWithAuthManager executes a non-streaming request via the core auth manager. -// This path is the only supported execution route. -func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) { - return h.executeWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, false) -} - -// ExecuteImageWithAuthManager executes an OpenAI-compatible image endpoint request. -func (h *BaseAPIHandler) ExecuteImageWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) { - return h.executeWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, true) -} - -func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, allowImageModel bool) ([]byte, http.Header, *interfaces.ErrorMessage) { - return h.executeWithAuthManagerFormats(ctx, handlerType, handlerType, modelName, rawJSON, alt, allowImageModel, modelExecutionOptions{}) -} - -func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { - originalRequestedModel := modelName - routeDecision := h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, false, execOptions) - responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol) - if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil { - return nil, nil, errMsg - } - if routeDecision.ExecutorPluginID != "" { - return h.executeWithPluginExecutor(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions) - } - providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, allowImageModel, routeDecision, execOptions) - if errMsg != nil { - return nil, nil, errMsg - } - providers = adjustExecutionProvidersForEntryProtocol(entryProtocol, providers) - reqMeta := requestExecutionMetadata(ctx) - reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel - addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) - addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) - setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON) - setServiceTierMetadata(reqMeta, rawJSON) - setGenerateMetadata(reqMeta, rawJSON) - payload := rawJSON - if len(payload) == 0 { - payload = nil - } - req := coreexecutor.Request{ - Model: normalizedModel, - Payload: payload, - } - afterAuthCapture := &requestAfterAuthCapture{} - opts := coreexecutor.Options{ - Stream: false, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(entryProtocol), - ResponseFormat: sdktranslator.FromString(responseProtocol), - Headers: modelExecutionHeaders(ctx, execOptions.Headers), - Query: modelExecutionQuery(ctx, execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), - } - opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - resp, err := h.AuthManager.Execute(ctx, providers, req, opts) - if err != nil { - err = enrichAuthSelectionError(err, providers, normalizedModel) - status := http.StatusInternalServerError - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} - } - executedReq, executedOpts := afterAuthCapture.apply(req, opts) - rawResponseHeaders := cloneHeader(resp.Headers) - responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) - return body, responseHeaders, nil -} - -// ExecuteCountWithAuthManager executes a non-streaming request via the core auth manager. -// This path is the only supported execution route. -func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) { - return h.executeCountWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, modelExecutionOptions{}) -} - -func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { - originalRequestedModel := modelName - routeDecision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, false, execOptions) - if routeDecision.ExecutorPluginID != "" { - return h.countWithPluginExecutor(ctx, handlerType, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions) - } - providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, false, routeDecision, execOptions) - if errMsg != nil { - return nil, nil, errMsg - } - providers = adjustExecutionProvidersForEntryProtocol(handlerType, providers) - reqMeta := requestExecutionMetadata(ctx) - reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel - addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) - setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON) - setServiceTierMetadata(reqMeta, rawJSON) - setGenerateMetadata(reqMeta, rawJSON) - payload := rawJSON - if len(payload) == 0 { - payload = nil - } - req := coreexecutor.Request{ - Model: normalizedModel, - Payload: payload, - } - afterAuthCapture := &requestAfterAuthCapture{} - opts := coreexecutor.Options{ - Stream: false, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(handlerType), - Headers: modelExecutionHeaders(ctx, execOptions.Headers), - Query: modelExecutionQuery(ctx, execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), - } - opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts) - if err != nil { - err = enrichAuthSelectionError(err, providers, normalizedModel) - status := http.StatusInternalServerError - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} - } - executedReq, executedOpts := afterAuthCapture.apply(req, opts) - rawResponseHeaders := cloneHeader(resp.Headers) - responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) - return body, responseHeaders, nil -} - -func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { - if h.AuthManager != nil && h.AuthManager.HomeEnabled() { - return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} - } - host := h.pluginExecutorHost() - if host == nil { - return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} - } - req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - resp, errExecute := host.ExecutePluginExecutor(ctx, executorPluginID, req, opts) - if errExecute != nil { - return nil, nil, executionErrorMessage(errExecute) - } - rawResponseHeaders := cloneHeader(resp.Headers) - responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) - return body, responseHeaders, nil -} - -func (h *BaseAPIHandler) countWithPluginExecutor(ctx context.Context, handlerType, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { - if h.AuthManager != nil && h.AuthManager.HomeEnabled() { - return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} - } - host := h.pluginExecutorHost() - if host == nil { - return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} - } - req, opts := h.pluginExecutorRequest(ctx, handlerType, handlerType, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - resp, errCount := host.CountPluginExecutor(ctx, executorPluginID, req, opts) - if errCount != nil { - return nil, nil, executionErrorMessage(errCount) - } - rawResponseHeaders := cloneHeader(resp.Headers) - responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) - body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) - return body, responseHeaders, nil -} - -func (h *BaseAPIHandler) pluginExecutorRequest(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt string, stream bool, execOptions modelExecutionOptions) (coreexecutor.Request, coreexecutor.Options) { - reqMeta := requestExecutionMetadata(ctx) - reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel - addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) - addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) - setReasoningEffortMetadata(reqMeta, entryProtocol, modelName, rawJSON) - setServiceTierMetadata(reqMeta, rawJSON) - setGenerateMetadata(reqMeta, rawJSON) - payload := rawJSON - if len(payload) == 0 { - payload = nil - } - req := coreexecutor.Request{Model: modelName, Payload: payload} - opts := coreexecutor.Options{ - Stream: stream, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(entryProtocol), - ResponseFormat: sdktranslator.FromString(responseProtocol), - Headers: modelExecutionHeaders(ctx, execOptions.Headers), - Query: modelExecutionQuery(ctx, execOptions.Query), - Metadata: reqMeta, - } - return req, opts -} - -func (h *BaseAPIHandler) applyRequestInterceptorsAfterPluginExecutorRoute(ctx context.Context, host PluginExecutorHost, executorPluginID, entryProtocol, originalRequestedModel string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options) { - if !requestInterceptorsEnabled(h.interceptorHost()) { - return req, opts - } - toFormat := sdktranslator.FromString(entryProtocol) - if resolver, ok := host.(pluginExecutorFormatResolver); ok && resolver != nil { - if resolved := resolver.PluginExecutorRequestToFormat(executorPluginID, req, opts); resolved != "" { - toFormat = resolved - } - } - resp := h.applyRequestInterceptorsAfterAuth(ctx, coreexecutor.RequestAfterAuthInterceptRequest{ - SourceFormat: opts.SourceFormat, - ToFormat: toFormat, - Model: req.Model, - RequestedModel: originalRequestedModel, - Stream: opts.Stream, - Headers: cloneHeader(opts.Headers), - Body: cloneBytes(req.Payload), - Metadata: opts.Metadata, - }, skipPluginID) - opts.Headers = mergeRequestInterceptorHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) - if len(resp.Body) > 0 { - req.Payload = cloneBytes(resp.Body) - opts.OriginalRequest = cloneBytes(resp.Body) - } - return req, opts -} - -func executionErrorMessage(err error) *interfaces.ErrorMessage { - status := http.StatusInternalServerError - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - return &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} -} - -// ExecuteStreamWithAuthManager executes a streaming request via the core auth manager. -// This path is the only supported execution route. -// The returned http.Header carries upstream response headers captured before streaming begins. -func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { - return h.executeStreamWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, false) -} - -// ExecuteImageStreamWithAuthManager executes a streaming OpenAI-compatible image endpoint request. -func (h *BaseAPIHandler) ExecuteImageStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { - return h.executeStreamWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, true) -} - -func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { - if h.AuthManager != nil && h.AuthManager.HomeEnabled() { - errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} - close(errChan) - return nil, nil, errChan - } - host := h.pluginExecutorHost() - if host == nil { - errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} - close(errChan) - return nil, nil, errChan - } - req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, true, execOptions) - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - streamResult, errStream := host.ExecutePluginExecutorStream(ctx, executorPluginID, req, opts) - if errStream != nil { - errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- executionErrorMessage(errStream) - close(errChan) - return nil, nil, errChan - } - if streamResult == nil { - errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor returned nil stream")} - close(errChan) - return nil, nil, errChan - } - - passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg) - interceptorHost := h.interceptorHost() - streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost) - rawStreamHeaders := cloneHeader(streamResult.Headers) - baseStreamHeaders := cloneHeader(streamResult.Headers) - applyStreamHeaders := func(headers http.Header) { - rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers) - } - if streamInterceptorsActive { - intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ - SourceFormat: responseProtocol, - Model: modelName, - RequestedModel: originalRequestedModel, - RequestHeaders: cloneHeader(opts.Headers), - ResponseHeaders: cloneHeader(rawStreamHeaders), - OriginalRequest: cloneBytes(opts.OriginalRequest), - RequestBody: cloneBytes(req.Payload), - ChunkIndex: pluginapi.StreamChunkHeaderInitIndex, - Metadata: opts.Metadata, - }, execOptions.SkipInterceptorPluginID) - applyStreamHeaders(intercepted.Headers) - } - upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) - if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { - upstreamHeaders = make(http.Header) - } - - dataChan := make(chan []byte) - errChan := make(chan *interfaces.ErrorMessage, 1) - var done <-chan struct{} - if ctx != nil { - done = ctx.Done() - } - chunks := streamResult.Chunks - if chunks == nil { - closed := make(chan coreexecutor.StreamChunk) - close(closed) - chunks = closed - } - go func() { - defer close(dataChan) - defer close(errChan) - chunkIndex := 0 - var historyChunks [][]byte - for { - chunk, ok, canceled := nextStreamChunk(ctx, nil, nil, chunks) - if canceled { - return - } - if !ok { - return - } - if chunk.Err != nil { - select { - case errChan <- executionErrorMessage(chunk.Err): - case <-done: - } - return - } - if len(chunk.Payload) == 0 { - continue - } - payload := cloneBytes(chunk.Payload) - if streamInterceptorsActive { - intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ - SourceFormat: responseProtocol, - Model: modelName, - RequestedModel: originalRequestedModel, - RequestHeaders: cloneHeader(opts.Headers), - ResponseHeaders: cloneHeader(rawStreamHeaders), - OriginalRequest: cloneBytes(opts.OriginalRequest), - RequestBody: cloneBytes(req.Payload), - Body: payload, - HistoryChunks: cloneByteSlices(historyChunks), - ChunkIndex: chunkIndex, - Metadata: opts.Metadata, - }, execOptions.SkipInterceptorPluginID) - applyStreamHeaders(intercepted.Headers) - if len(intercepted.Body) > 0 { - payload = cloneBytes(intercepted.Body) - } - chunkIndex++ - if intercepted.DropChunk { - continue - } - } else { - chunkIndex++ - } - if responseProtocol == "openai-response" { - if errValidate := validateSSEDataJSON(payload); errValidate != nil { - select { - case errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}: - case <-done: - } - return - } - } - select { - case dataChan <- payload: - if streamInterceptorsActive { - historyChunks = appendStreamInterceptorHistory(historyChunks, payload) - } - case <-done: - return - } - } - }() - return dataChan, upstreamHeaders, errChan -} - -func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, allowImageModel bool) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { - return h.executeStreamWithAuthManagerFormats(ctx, handlerType, handlerType, modelName, rawJSON, alt, allowImageModel, modelExecutionOptions{}) -} - -func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { - originalRequestedModel := modelName - routeDecision, preparedRoute := preparedModelRouteFromContext(ctx) - if !preparedRoute { - routeDecision = h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions) - } - responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol) - if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil { - errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- errMsg - close(errChan) - return nil, nil, errChan - } - if routeDecision.ExecutorPluginID != "" { - return h.streamWithPluginExecutor(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions) - } - providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, allowImageModel, routeDecision, execOptions) - if errMsg != nil { - errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- errMsg - close(errChan) - return nil, nil, errChan - } - providers = adjustExecutionProvidersForEntryProtocol(entryProtocol, providers) - reqMeta := requestExecutionMetadata(ctx) - reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel - addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) - addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) - setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON) - setServiceTierMetadata(reqMeta, rawJSON) - setGenerateMetadata(reqMeta, rawJSON) - payload := rawJSON - if len(payload) == 0 { - payload = nil - } - req := coreexecutor.Request{ - Model: normalizedModel, - Payload: payload, - } - afterAuthCapture := &requestAfterAuthCapture{} - opts := coreexecutor.Options{ - Stream: true, - Alt: alt, - OriginalRequest: rawJSON, - SourceFormat: sdktranslator.FromString(entryProtocol), - ResponseFormat: sdktranslator.FromString(responseProtocol), - Headers: modelExecutionHeaders(ctx, execOptions.Headers), - Query: modelExecutionQuery(ctx, execOptions.Query), - RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), - } - opts.Metadata = reqMeta - req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) - streamResult, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts) - if err != nil { - err = enrichAuthSelectionError(err, providers, normalizedModel) - errChan := make(chan *interfaces.ErrorMessage, 1) - status := http.StatusInternalServerError - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - errChan <- &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} - close(errChan) - return nil, nil, errChan - } - if streamResult == nil { - errChan := make(chan *interfaces.ErrorMessage, 1) - errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("auth manager returned nil stream")} - close(errChan) - return nil, nil, errChan - } - executedRequest := func() (coreexecutor.Request, coreexecutor.Options) { - return afterAuthCapture.apply(req, opts) - } - passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg) - interceptorHost := h.interceptorHost() - streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost) - // Resolve bootstrap retries and header initialization before returning so the - // returned header snapshot is never modified by the stream goroutine. - rawStreamHeaders := cloneHeader(streamResult.Headers) - baseStreamHeaders := cloneHeader(streamResult.Headers) - chunks := streamResult.Chunks - if chunks == nil { - closed := make(chan coreexecutor.StreamChunk) - close(closed) - chunks = closed - } - streamClosedBeforeRead := false - streamCanceledBeforeRead := false - streamHeaderInitialized := false - - applyStreamHeaders := func(headers http.Header) { - rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers) - } - - applyStreamHeaderInit := func() { - if !streamInterceptorsActive || streamHeaderInitialized { - return - } - executedReq, executedOpts := executedRequest() - intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ - SourceFormat: responseProtocol, - Model: normalizedModel, - RequestedModel: originalRequestedModel, - RequestHeaders: cloneHeader(executedOpts.Headers), - ResponseHeaders: cloneHeader(rawStreamHeaders), - OriginalRequest: cloneBytes(executedOpts.OriginalRequest), - RequestBody: cloneBytes(executedReq.Payload), - ChunkIndex: pluginapi.StreamChunkHeaderInitIndex, - Metadata: executedOpts.Metadata, - }, execOptions.SkipInterceptorPluginID) - applyStreamHeaders(intercepted.Headers) - streamHeaderInitialized = true - } - - transformStreamPayload := func(payload []byte, chunkIndex *int, historyChunks [][]byte) ([]byte, bool, *interfaces.ErrorMessage) { - applyStreamHeaderInit() - payload = cloneBytes(payload) - if streamInterceptorsActive { - executedReq, executedOpts := executedRequest() - intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ - SourceFormat: responseProtocol, - Model: normalizedModel, - RequestedModel: originalRequestedModel, - RequestHeaders: cloneHeader(executedOpts.Headers), - ResponseHeaders: cloneHeader(rawStreamHeaders), - OriginalRequest: cloneBytes(executedOpts.OriginalRequest), - RequestBody: cloneBytes(executedReq.Payload), - Body: payload, - HistoryChunks: cloneByteSlices(historyChunks), - ChunkIndex: *chunkIndex, - Metadata: executedOpts.Metadata, - }, execOptions.SkipInterceptorPluginID) - applyStreamHeaders(intercepted.Headers) - if len(intercepted.Body) > 0 { - payload = cloneBytes(intercepted.Body) - } - (*chunkIndex)++ - if intercepted.DropChunk { - return nil, false, nil - } - } else { - (*chunkIndex)++ - } - if responseProtocol == "openai-response" { - if errValidate := validateSSEDataJSON(payload); errValidate != nil { - return nil, false, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate} - } - } - return payload, true, nil - } - - var bootstrapPayload []byte - bootstrapChunkIndex := 0 - var bootstrapHistoryChunks [][]byte - var bootstrapStreamErr error - var bootstrapErr *interfaces.ErrorMessage - readInitialStreamChunks := func() { - for { - var chunk coreexecutor.StreamChunk - var ok bool - if ctx != nil { - select { - case <-ctx.Done(): - streamCanceledBeforeRead = true - return - case chunk, ok = <-chunks: - } - } else { - chunk, ok = <-chunks - } - if !ok { - streamClosedBeforeRead = true - applyStreamHeaderInit() - return - } - if chunk.Err != nil { - bootstrapStreamErr = chunk.Err - return - } - if len(chunk.Payload) == 0 { - continue - } - payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &bootstrapChunkIndex, bootstrapHistoryChunks) - if errMsg != nil { - bootstrapErr = errMsg - return - } - if !deliverable { - continue - } - bootstrapPayload = payload - return - } - } - - bootstrapEligible := func(err error) bool { - status := statusFromError(err) - if status == 0 { - return true - } - switch status { - case http.StatusUnauthorized, http.StatusForbidden, http.StatusPaymentRequired, - http.StatusRequestTimeout, http.StatusTooManyRequests: - return true - default: - return status >= http.StatusInternalServerError - } - } - - maxBootstrapRetries := StreamingBootstrapRetries(h.Cfg) - if h.AuthManager.HomeEnabled() { - maxBootstrapRetries = 0 - } - for bootstrapRetries := 0; !streamCanceledBeforeRead; { - readInitialStreamChunks() - if streamCanceledBeforeRead || bootstrapErr != nil || bootstrapStreamErr == nil { - break - } - if bootstrapRetries >= maxBootstrapRetries || !bootstrapEligible(bootstrapStreamErr) { - bootstrapErr = executionErrorMessage(bootstrapStreamErr) - break - } - bootstrapRetries++ - retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) - if retryErr != nil { - bootstrapErr = executionErrorMessage(enrichAuthSelectionError(retryErr, providers, normalizedModel)) - break - } - if retryResult == nil { - bootstrapErr = executionErrorMessage(fmt.Errorf("auth manager returned nil stream")) - break - } - rawStreamHeaders = cloneHeader(retryResult.Headers) - baseStreamHeaders = cloneHeader(retryResult.Headers) - streamHeaderInitialized = false - streamClosedBeforeRead = false - bootstrapStreamErr = nil - bootstrapPayload = nil - bootstrapChunkIndex = 0 - bootstrapHistoryChunks = nil - chunks = retryResult.Chunks - if chunks == nil { - closed := make(chan coreexecutor.StreamChunk) - close(closed) - chunks = closed - } - } - - upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) - if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { - upstreamHeaders = make(http.Header) - } - dataChan := make(chan []byte) - errChan := make(chan *interfaces.ErrorMessage, 1) - - go func() { - defer close(dataChan) - defer close(errChan) - if streamCanceledBeforeRead { - return - } - - sendErr := func(msg *interfaces.ErrorMessage) bool { - if ctx == nil { - errChan <- msg - return true - } - select { - case <-ctx.Done(): - return false - case errChan <- msg: - return true - } - } - - sendData := func(chunk []byte) bool { - if ctx == nil { - dataChan <- chunk - return true - } - select { - case <-ctx.Done(): - return false - case dataChan <- chunk: - return true - } - } - - if bootstrapErr != nil { - _ = sendErr(bootstrapErr) - return - } - - chunkIndex := bootstrapChunkIndex - historyChunks := bootstrapHistoryChunks - if bootstrapPayload != nil { - if okSendData := sendData(bootstrapPayload); !okSendData { - return - } - if streamInterceptorsActive { - historyChunks = appendStreamInterceptorHistory(historyChunks, bootstrapPayload) - } - } - for { - chunk, ok, canceled := nextStreamChunk(ctx, nil, &streamClosedBeforeRead, chunks) - if canceled || !ok { - return - } - if chunk.Err != nil { - _ = sendErr(executionErrorMessage(chunk.Err)) - return - } - if len(chunk.Payload) == 0 { - continue - } - payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &chunkIndex, historyChunks) - if errMsg != nil { - _ = sendErr(errMsg) - return - } - if !deliverable { - continue - } - if okSendData := sendData(payload); !okSendData { - return - } - if streamInterceptorsActive { - historyChunks = appendStreamInterceptorHistory(historyChunks, payload) - } - } - }() - return dataChan, upstreamHeaders, errChan -} - -func validateSSEDataJSON(chunk []byte) error { - for _, line := range bytes.Split(chunk, []byte("\n")) { - line = bytes.TrimSpace(line) - if len(line) == 0 { - continue - } - if !bytes.HasPrefix(line, []byte("data:")) { - continue - } - data := bytes.TrimSpace(line[5:]) - if len(data) == 0 { - continue - } - if bytes.Equal(data, []byte("[DONE]")) { - continue - } - if json.Valid(data) { - continue - } - const max = 512 - preview := data - if len(preview) > max { - preview = preview[:max] - } - return fmt.Errorf("invalid SSE data JSON (len=%d): %q", len(data), preview) - } - return nil -} - -func preferExecutionProvider(providers []string, preferred string) []string { - preferred = strings.ToLower(strings.TrimSpace(preferred)) - if preferred == "" || len(providers) < 2 { - return providers - } - preferredIndex := -1 - for i := range providers { - if strings.ToLower(strings.TrimSpace(providers[i])) == preferred { - preferredIndex = i - break - } - } - if preferredIndex <= 0 { - return providers - } - out := make([]string, 0, len(providers)) - out = append(out, providers[preferredIndex]) - out = append(out, providers[:preferredIndex]...) - out = append(out, providers[preferredIndex+1:]...) - return out -} - -func adjustExecutionProvidersForEntryProtocol(entryProtocol string, providers []string) []string { - if entryProtocol == Interactions { - return preferExecutionProvider(providers, GeminiInteractions) - } - if supportsNativeInteractionsEntryProtocol(entryProtocol) { - return providers - } - return excludeExecutionProvider(providers, GeminiInteractions) -} - -func supportsNativeInteractionsEntryProtocol(entryProtocol string) bool { - switch entryProtocol { - case Interactions, OpenAI, OpenaiResponse, Claude, Gemini: - return true - default: - return false - } -} - -func excludeExecutionProvider(providers []string, excluded string) []string { - excluded = strings.ToLower(strings.TrimSpace(excluded)) - if excluded == "" || len(providers) == 0 { - return providers - } - excludedIndex := -1 - for i := range providers { - if strings.ToLower(strings.TrimSpace(providers[i])) == excluded { - excludedIndex = i - break - } - } - if excludedIndex == -1 { - return providers - } - out := make([]string, 0, len(providers)-1) - out = append(out, providers[:excludedIndex]...) - out = append(out, providers[excludedIndex+1:]...) - return out -} - -func statusFromError(err error) int { - if err == nil { - return 0 - } - if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - return code - } - } - return 0 -} - -func (h *BaseAPIHandler) getRequestDetails(modelName string) (providers []string, normalizedModel string, err *interfaces.ErrorMessage) { - return h.getRequestDetailsWithOptions(modelName, false) -} - -func validateNativeInteractionsExecution(entryProtocol string, execOptions modelExecutionOptions, routeDecision modelRouteDecision) *interfaces.ErrorMessage { - forcedProvider := strings.ToLower(strings.TrimSpace(execOptions.ForcedProvider)) - if forcedProvider == "" || entryProtocol != Interactions { - return nil - } - if routeDecision.ExecutorPluginID != "" { - return nativeInteractionsExecutionError() - } - if routeProvider := strings.ToLower(strings.TrimSpace(routeDecision.Provider)); routeProvider != "" && routeProvider != forcedProvider { - return nativeInteractionsExecutionError() - } - return nil -} - -func nativeInteractionsExecutionError() *interfaces.ErrorMessage { - return &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("agent is only supported for native interactions execution"), - } -} - -// providersForExecution resolves the providers and normalized model for a request. When a model -// router selected a built-in provider, it skips model->provider resolution and uses the router's -// provider (with an optional target model); otherwise it falls back to the registry-based path. -func (h *BaseAPIHandler) providersForExecution(modelName, originalRequestedModel string, allowImageModel bool, routeDecision modelRouteDecision, execOptions modelExecutionOptions) ([]string, string, *interfaces.ErrorMessage) { - forcedProvider := strings.ToLower(strings.TrimSpace(execOptions.ForcedProvider)) - if forcedProvider != "" { - if routeDecision.ExecutorPluginID != "" { - return nil, "", nativeInteractionsExecutionError() - } - if routeProvider := strings.ToLower(strings.TrimSpace(routeDecision.Provider)); routeProvider != "" && routeProvider != forcedProvider { - return nil, "", nativeInteractionsExecutionError() - } - normalizedModel := strings.TrimSpace(modelName) - if normalizedModel == "" { - normalizedModel = strings.TrimSpace(originalRequestedModel) - } - if errMsg := h.validateImageOnlyModel(normalizedModel, allowImageModel); errMsg != nil { - return nil, "", errMsg - } - return []string{forcedProvider}, normalizedModel, nil - } - if routeDecision.Provider != "" { - normalizedModel := originalRequestedModel - if routeDecision.Model != "" { - normalizedModel = routeDecision.Model - } - if errMsg := h.validateImageOnlyModel(normalizedModel, allowImageModel); errMsg != nil { - return nil, "", errMsg - } - return []string{routeDecision.Provider}, normalizedModel, nil - } - return h.getRequestDetailsWithOptions(modelName, allowImageModel) -} - -func (h *BaseAPIHandler) getRequestDetailsWithOptions(modelName string, allowImageModel bool) (providers []string, normalizedModel string, err *interfaces.ErrorMessage) { - resolvedModelName := modelName - initialSuffix := thinking.ParseSuffix(modelName) - if initialSuffix.ModelName == "auto" { - if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() { - resolvedModelName = modelName - } else { - resolvedBase := util.ResolveAutoModel(initialSuffix.ModelName) - if initialSuffix.HasSuffix { - resolvedModelName = fmt.Sprintf("%s(%s)", resolvedBase, initialSuffix.RawSuffix) - } else { - resolvedModelName = resolvedBase - } - } - } else { - if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() { - resolvedModelName = modelName - } else { - resolvedModelName = util.ResolveAutoModel(modelName) - } - } - - parsed := thinking.ParseSuffix(resolvedModelName) - baseModel := strings.TrimSpace(parsed.ModelName) - - if errMsg := h.validateImageOnlyModel(baseModel, allowImageModel); errMsg != nil { - return nil, "", errMsg - } - - if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() { - return []string{"home"}, resolvedModelName, nil - } - - providers = util.GetProviderName(baseModel) - // Fallback: if baseModel has no provider but differs from resolvedModelName, - // try using the full model name. This handles edge cases where custom models - // may be registered with their full suffixed name (e.g., "my-model(8192)"). - // Evaluated in Story 11.8: This fallback is intentionally preserved to support - // custom model registrations that include thinking suffixes. - if len(providers) == 0 && baseModel != resolvedModelName { - providers = util.GetProviderName(resolvedModelName) - } - - if len(providers) == 0 { - return nil, "", &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("unknown provider for model %s", modelName)} - } - - // The thinking suffix is preserved in the model name itself, so no - // metadata-based configuration passing is needed. - return providers, resolvedModelName, nil -} - -func (h *BaseAPIHandler) validateImageOnlyModel(modelName string, allowImageModel bool) *interfaces.ErrorMessage { - baseModel := strings.TrimSpace(thinking.ParseSuffix(modelName).ModelName) - if baseModel == "" { - baseModel = strings.TrimSpace(modelName) - } - if isOpenAIImageOnlyModel(baseModel) && !allowImageModel { - return &interfaces.ErrorMessage{ - StatusCode: http.StatusServiceUnavailable, - Error: fmt.Errorf("model %s is only supported on /v1/images/generations and /v1/images/edits", routeModelBaseName(baseModel)), - } - } - return nil -} - -func isOpenAIImageOnlyModel(model string) bool { - switch strings.ToLower(strings.TrimSpace(routeModelBaseName(model))) { - case "gpt-image-1.5", "gpt-image-2", "grok-imagine-image", "grok-imagine-image-quality": - return true - default: - return false - } -} - -func routeModelBaseName(model string) string { - model = strings.TrimSpace(model) - if idx := strings.LastIndex(model, "/"); idx >= 0 && idx < len(model)-1 { - return strings.TrimSpace(model[idx+1:]) - } - return model -} - -func cloneBytes(src []byte) []byte { - if len(src) == 0 { - return nil - } - dst := make([]byte, len(src)) - copy(dst, src) - return dst -} - -func cloneHeader(src http.Header) http.Header { - if src == nil { - return nil - } - dst := make(http.Header, len(src)) - for key, values := range src { - dst[key] = append([]string(nil), values...) - } - return dst -} - -func cloneByteSlices(src [][]byte) [][]byte { - if len(src) == 0 { - return nil - } - dst := make([][]byte, 0, len(src)) - for _, item := range src { - dst = append(dst, cloneBytes(item)) - } - return dst -} - -func nextStreamChunk(ctx context.Context, pending *[]coreexecutor.StreamChunk, closed *bool, chunks <-chan coreexecutor.StreamChunk) (coreexecutor.StreamChunk, bool, bool) { - if pending != nil && len(*pending) > 0 { - chunk := (*pending)[0] - (*pending)[0] = coreexecutor.StreamChunk{} - *pending = (*pending)[1:] - return chunk, true, false - } - if closed != nil && *closed { - return coreexecutor.StreamChunk{}, false, false - } - var chunk coreexecutor.StreamChunk - var ok bool - if ctx != nil { - select { - case <-ctx.Done(): - return coreexecutor.StreamChunk{}, false, true - case chunk, ok = <-chunks: - } - } else { - chunk, ok = <-chunks - } - if !ok && closed != nil { - *closed = true - } - return chunk, ok, false -} - -func appendStreamInterceptorHistory(history [][]byte, chunk []byte) [][]byte { - if len(chunk) == 0 { - return history - } - history = append(history, cloneBytes(chunk)) - for len(history) > maxStreamInterceptorHistoryChunks || byteSlicesSize(history) > maxStreamInterceptorHistoryBytes { - history[0] = nil - history = history[1:] - } - if len(history) == 0 { - return nil - } - return history -} - -func byteSlicesSize(items [][]byte) int { - total := 0 - for _, item := range items { - total += len(item) - } - return total -} - -func finalInterceptorHeaders(current, intercepted http.Header) http.Header { - if intercepted == nil { - return current - } - if len(intercepted) == 0 { - return nil - } - return cloneHeader(intercepted) -} - -func downstreamHeadersFromExecutor(headers http.Header, passthrough bool) http.Header { - if !passthrough { - return nil - } - return FilterUpstreamHeaders(headers) -} - -func downstreamHeadersAfterInterceptors(baseRaw, finalRaw http.Header, passthrough bool) http.Header { - if passthrough { - return FilterUpstreamHeaders(finalRaw) - } - return FilterUpstreamHeaders(diffHeaders(baseRaw, finalRaw)) -} - -func diffHeaders(base, next http.Header) http.Header { - if len(next) == 0 { - return nil - } - baseValues := make(map[string][]string, len(base)) - for key, values := range base { - baseValues[http.CanonicalHeaderKey(key)] = values - } - out := make(http.Header) - for key, values := range next { - canonicalKey := http.CanonicalHeaderKey(key) - if stringSlicesEqual(baseValues[canonicalKey], values) { - continue - } - out[canonicalKey] = append([]string(nil), values...) - } - if len(out) == 0 { - return nil - } - return out -} - -func stringSlicesEqual(left, right []string) bool { - if len(left) != len(right) { - return false - } - for i := range left { - if left[i] != right[i] { - return false - } - } - return true -} - -func (h *BaseAPIHandler) interceptorHost() PluginInterceptorHost { - if h == nil { - return nil - } - return h.PluginHost -} - -func (h *BaseAPIHandler) modelRouterHost() PluginModelRouterHost { - if h == nil { - return nil - } - if !isNilPluginModelRouterHost(h.ModelRouterHost) { - return h.ModelRouterHost - } - host := h.interceptorHost() - if host == nil { - return nil - } - router, ok := host.(PluginModelRouterHost) - if !ok { - return nil - } - return router -} - -func (h *BaseAPIHandler) pluginExecutorHost() PluginExecutorHost { - if h == nil { - return nil - } - if executorHost, ok := h.ModelRouterHost.(PluginExecutorHost); ok && executorHost != nil { - return executorHost - } - if executorHost, ok := h.PluginHost.(PluginExecutorHost); ok && executorHost != nil { - return executorHost - } - return nil -} - -type modelRouteDecision struct { - ExecutorPluginID string - Provider string - Model string -} - -func routeModel(ctx context.Context, host PluginModelRouterHost, req pluginapi.ModelRouteRequest, skipPluginID string) (pluginapi.ModelRouteResponse, bool) { - if host == nil { - return pluginapi.ModelRouteResponse{}, false - } - skipPluginID = strings.TrimSpace(skipPluginID) - if skipPluginID != "" { - if skipper, ok := host.(pluginModelRouterSkipHost); ok { - return skipper.RouteModelExcept(ctx, req, skipPluginID) - } - return pluginapi.ModelRouteResponse{}, false - } - return host.RouteModel(ctx, req) -} - -func modelRoutersEnabled(host PluginModelRouterHost, skipPluginID string) bool { - if host == nil { - return false - } - skipPluginID = strings.TrimSpace(skipPluginID) - if skipPluginID != "" { - if _, ok := host.(pluginModelRouterSkipHost); !ok { - return false - } - if detector, ok := host.(modelRouterSkipDetector); ok { - return detector.HasModelRoutersExcept(skipPluginID) - } - } - if detector, ok := host.(modelRouterDetector); ok { - return detector.HasModelRouters() - } - // No detector: treat routing as disabled (same conservative default as before any - // ModelRouter existed). Hosts that route must implement HasModelRouters (pluginhost.Host does). - return false -} - -func (h *BaseAPIHandler) applyModelRouter(ctx context.Context, handlerType, modelName string, rawJSON []byte, stream bool, execOptions modelExecutionOptions) modelRouteDecision { - var decision modelRouteDecision - host := h.modelRouterHost() - if host == nil || !modelRoutersEnabled(host, execOptions.SkipRouterPluginID) { - return decision - } - meta := requestExecutionMetadata(ctx) - meta[coreexecutor.RequestedModelMetadataKey] = modelName - addModelExecutionSourceMetadata(meta, execOptions.InternalSource) - resp, ok := routeModel(ctx, host, pluginapi.ModelRouteRequest{ - SourceFormat: handlerType, - RequestedModel: modelName, - Stream: stream, - Headers: modelExecutionHeaders(ctx, execOptions.Headers), - Query: modelExecutionQuery(ctx, execOptions.Query), - Body: cloneBytes(rawJSON), - Metadata: meta, - }, execOptions.SkipRouterPluginID) - if !ok || !resp.Handled { - return decision - } - switch resp.TargetKind { - case pluginapi.ModelRouteTargetSelf, pluginapi.ModelRouteTargetExecutor: - decision.ExecutorPluginID = strings.TrimSpace(resp.Target) - case pluginapi.ModelRouteTargetProvider: - decision.Provider = strings.ToLower(strings.TrimSpace(resp.Target)) - decision.Model = strings.TrimSpace(resp.TargetModel) - } - return decision -} - -func streamInterceptorsEnabled(host PluginInterceptorHost) bool { - if host == nil { - return false - } - if detector, ok := host.(streamInterceptorDetector); ok { - return detector.HasStreamInterceptors() - } - return true -} - -func requestInterceptorsEnabled(host PluginInterceptorHost) bool { - if host == nil { - return false - } - if detector, ok := host.(requestInterceptorDetector); ok { - return detector.HasRequestInterceptors() - } - return true -} - -type requestAfterAuthCapture struct { - mu sync.Mutex - set bool - headers http.Header - body []byte - originalRequest []byte - originalRequestReplaced bool -} - -func (c *requestAfterAuthCapture) record(req coreexecutor.RequestAfterAuthInterceptRequest, resp coreexecutor.RequestAfterAuthInterceptResponse) { - if c == nil { - return - } - headers := mergeRequestInterceptorHeaders(req.Headers, resp.Headers, resp.ClearHeaders) - body := cloneBytes(req.Body) - var originalRequest []byte - originalRequestReplaced := false - if len(resp.Body) > 0 { - body = cloneBytes(resp.Body) - originalRequest = cloneBytes(resp.Body) - originalRequestReplaced = true - } - - c.mu.Lock() - defer c.mu.Unlock() - c.set = true - c.headers = headers - c.body = body - c.originalRequest = originalRequest - c.originalRequestReplaced = originalRequestReplaced -} - -func (c *requestAfterAuthCapture) apply(req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Request, coreexecutor.Options) { - if c == nil { - return req, opts - } - c.mu.Lock() - defer c.mu.Unlock() - if !c.set { - return req, opts - } - req.Payload = cloneBytes(c.body) - opts.Headers = cloneHeader(c.headers) - if c.originalRequestReplaced { - opts.OriginalRequest = cloneBytes(c.originalRequest) - } - return req, opts -} - -func mergeRequestInterceptorHeaders(current, updates http.Header, clear []string) http.Header { - if updates == nil && len(clear) == 0 { - return cloneHeader(current) - } - out := cloneHeader(current) - if out == nil && (len(updates) > 0 || len(clear) > 0) { - out = make(http.Header) - } - for _, key := range clear { - out.Del(key) - } - for key, values := range updates { - out.Del(key) - for _, value := range values { - out.Add(key, value) - } - } - return out -} - -func interceptRequestBeforeAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { - if skipPluginID != "" { - if skipper, ok := host.(pluginInterceptorSkipHost); ok { - return skipper.InterceptRequestBeforeAuthExcept(ctx, req, skipPluginID) - } - } - return host.InterceptRequestBeforeAuth(ctx, req) -} - -func interceptRequestAfterAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { - if skipPluginID != "" { - if skipper, ok := host.(pluginInterceptorSkipHost); ok { - return skipper.InterceptRequestAfterAuthExcept(ctx, req, skipPluginID) - } - } - return host.InterceptRequestAfterAuth(ctx, req) -} - -func interceptResponse(ctx context.Context, host PluginInterceptorHost, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse { - if skipPluginID != "" { - if skipper, ok := host.(pluginInterceptorSkipHost); ok { - return skipper.InterceptResponseExcept(ctx, req, skipPluginID) - } - } - return host.InterceptResponse(ctx, req) -} - -func interceptStreamChunk(ctx context.Context, host PluginInterceptorHost, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse { - if skipPluginID != "" { - if skipper, ok := host.(pluginInterceptorSkipHost); ok { - return skipper.InterceptStreamChunkExcept(ctx, req, skipPluginID) - } - } - return host.InterceptStreamChunk(ctx, req) -} - -func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, handlerType, requestedModel string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options) { - host := h.interceptorHost() - if host == nil { - return req, opts - } - resp := interceptRequestBeforeAuth(ctx, host, pluginapi.RequestInterceptRequest{ - SourceFormat: handlerType, - Model: req.Model, - RequestedModel: requestedModel, - Stream: opts.Stream, - Headers: cloneHeader(opts.Headers), - Body: cloneBytes(req.Payload), - Metadata: opts.Metadata, - }, skipPluginID) - opts.Headers = finalInterceptorHeaders(opts.Headers, resp.Headers) - if len(resp.Body) > 0 { - req.Payload = cloneBytes(resp.Body) - opts.OriginalRequest = cloneBytes(resp.Body) - } - return req, opts -} - -func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture, skipPluginID string) coreexecutor.RequestAfterAuthInterceptor { - if !requestInterceptorsEnabled(h.interceptorHost()) { - return nil - } - return func(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse { - resp := h.applyRequestInterceptorsAfterAuth(ctx, req, skipPluginID) - if capture != nil { - capture.record(req, resp) - } - return resp - } -} - -func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest, skipPluginID string) coreexecutor.RequestAfterAuthInterceptResponse { - host := h.interceptorHost() - if !requestInterceptorsEnabled(host) { - return coreexecutor.RequestAfterAuthInterceptResponse{} - } - resp := interceptRequestAfterAuth(ctx, host, pluginapi.RequestInterceptRequest{ - SourceFormat: req.SourceFormat.String(), - ToFormat: req.ToFormat.String(), - Model: req.Model, - RequestedModel: req.RequestedModel, - Stream: req.Stream, - Headers: cloneHeader(req.Headers), - Body: cloneBytes(req.Body), - Metadata: req.Metadata, - }, skipPluginID) - return coreexecutor.RequestAfterAuthInterceptResponse{ - Headers: resp.Headers, - Body: resp.Body, - ClearHeaders: resp.ClearHeaders, - } -} - -func (h *BaseAPIHandler) applyResponseInterceptors(ctx context.Context, handlerType, normalizedModel, requestedModel string, opts coreexecutor.Options, rawResponseHeaders, responseHeaders http.Header, originalRequest, requestBody, body []byte, statusCode int, skipPluginID string) ([]byte, http.Header) { - host := h.interceptorHost() - if host == nil { - return body, responseHeaders - } - resp := interceptResponse(ctx, host, pluginapi.ResponseInterceptRequest{ - SourceFormat: handlerType, - Model: normalizedModel, - RequestedModel: requestedModel, - Stream: false, - RequestHeaders: cloneHeader(opts.Headers), - ResponseHeaders: cloneHeader(rawResponseHeaders), - OriginalRequest: cloneBytes(originalRequest), - RequestBody: cloneBytes(requestBody), - Body: cloneBytes(body), - StatusCode: statusCode, - Metadata: opts.Metadata, - }, skipPluginID) - responseHeaders = downstreamHeadersAfterInterceptors(rawResponseHeaders, finalInterceptorHeaders(rawResponseHeaders, resp.Headers), PassthroughHeadersEnabled(h.Cfg)) - if len(resp.Body) > 0 { - body = cloneBytes(resp.Body) - } - return body, responseHeaders -} - -func enrichAuthSelectionError(err error, providers []string, model string) error { - if err == nil { - return nil - } - - var authErr *coreauth.Error - if !errors.As(err, &authErr) || authErr == nil { - return err - } - - code := strings.TrimSpace(authErr.Code) - if code != "auth_not_found" && code != "auth_unavailable" { - return err - } - - providerText := strings.Join(providers, ",") - if providerText == "" { - providerText = "unknown" - } - modelText := strings.TrimSpace(model) - if modelText == "" { - modelText = "unknown" - } - - baseMessage := strings.TrimSpace(authErr.Message) - if baseMessage == "" { - baseMessage = "no auth available" - } - detail := fmt.Sprintf("%s (providers=%s, model=%s)", baseMessage, providerText, modelText) - - // Clarify the most common alias confusion between Anthropic route names and internal provider keys. - if strings.Contains(","+providerText+",", ",claude,") { - detail += "; check Claude auth/key session and cooldown state via /v0/management/auth-files" - } - - status := authErr.HTTPStatus - if status <= 0 { - status = http.StatusServiceUnavailable - } - - return &coreauth.Error{ - Code: authErr.Code, - Message: detail, - Retryable: authErr.Retryable, - HTTPStatus: status, - } -} - -// WriteErrorResponse writes an error message to the response writer using the HTTP status embedded in the message. -func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.ErrorMessage) { - status := http.StatusInternalServerError - if msg != nil && msg.StatusCode > 0 { - status = msg.StatusCode - } - if msg != nil && msg.Error != nil { - for _, value := range coreauth.SafeResponseHeaders(msg.Error).Values("Retry-After") { - c.Writer.Header().Add("Retry-After", value) - } - } - if msg != nil && msg.Addon != nil && PassthroughHeadersEnabled(h.Cfg) { - for key, values := range msg.Addon { - if len(values) == 0 || IsCPAReservedResponseHeader(key) { - continue - } - c.Writer.Header().Del(key) - for _, value := range values { - c.Writer.Header().Add(key, value) - } - } - } - - errText := http.StatusText(status) - if msg != nil && msg.Error != nil { - if v := strings.TrimSpace(msg.Error.Error()); v != "" { - errText = v - } - } - - body := BuildErrorResponseBody(status, errText) - // Append first to preserve upstream response logs, then drop duplicate payloads if already recorded. - var previous []byte - if existing, exists := c.Get("API_RESPONSE"); exists { - if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 { - previous = existingBytes - } - } - appendAPIResponse(c, body) - trimmedErrText := strings.TrimSpace(errText) - trimmedBody := bytes.TrimSpace(body) - if len(previous) > 0 { - if (trimmedErrText != "" && bytes.Contains(previous, []byte(trimmedErrText))) || - (len(trimmedBody) > 0 && bytes.Contains(previous, trimmedBody)) { - c.Set("API_RESPONSE", previous) - } - } - - if !c.Writer.Written() { - c.Writer.Header().Set("Content-Type", "application/json") - } - c.Status(status) - _, _ = c.Writer.Write(body) -} - -func (h *BaseAPIHandler) LoggingAPIResponseError(ctx context.Context, err *interfaces.ErrorMessage) { - if h.Cfg.RequestLog { - if ginContext, ok := ctx.Value("gin").(*gin.Context); ok { - if apiResponseErrors, isExist := ginContext.Get("API_RESPONSE_ERROR"); isExist { - if slicesAPIResponseError, isOk := apiResponseErrors.([]*interfaces.ErrorMessage); isOk { - slicesAPIResponseError = append(slicesAPIResponseError, err) - ginContext.Set("API_RESPONSE_ERROR", slicesAPIResponseError) - } - } else { - // Create new response data entry - ginContext.Set("API_RESPONSE_ERROR", []*interfaces.ErrorMessage{err}) - } - } - } -} - // APIHandlerCancelFunc is a function type for canceling an API handler's context. // It can optionally accept parameters, which are used for logging the response. type APIHandlerCancelFunc func(params ...interface{}) diff --git a/sdk/api/handlers/handlers_context.go b/sdk/api/handlers/handlers_context.go new file mode 100644 index 000000000..7926be0f3 --- /dev/null +++ b/sdk/api/handlers/handlers_context.go @@ -0,0 +1,159 @@ +package handlers + +import ( + "net/http" + "net/url" + "strings" + + "github.com/gin-gonic/gin" + "golang.org/x/net/context" +) + +type pinnedAuthContextKey struct{} + +type selectedAuthCallbackContextKey struct{} + +type preparedModelRouteContextKey struct{} + +type executionSessionContextKey struct{} + +type disallowFreeAuthContextKey struct{} + +// WithPinnedAuthID returns a child context that requests execution on a specific auth ID. +func WithPinnedAuthID(ctx context.Context, authID string) context.Context { + authID = strings.TrimSpace(authID) + if authID == "" { + return ctx + } + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, pinnedAuthContextKey{}, authID) +} + +// WithSelectedAuthIDCallback returns a child context that receives the selected auth ID. +func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) context.Context { + if callback == nil { + return ctx + } + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, selectedAuthCallbackContextKey{}, callback) +} + +// PrepareStreamModelRoute resolves a stream route once and stores it on the returned context for execution. +// The boolean reports whether the route overrides normal model-to-provider resolution. +func (h *BaseAPIHandler) PrepareStreamModelRoute(ctx context.Context, handlerType string, modelName string, rawJSON []byte) (context.Context, bool) { + if ctx == nil { + ctx = context.Background() + } + decision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, true, modelExecutionOptions{}) + ctx = context.WithValue(ctx, preparedModelRouteContextKey{}, decision) + hasOverride := strings.TrimSpace(decision.ExecutorPluginID) != "" || strings.TrimSpace(decision.Provider) != "" + return ctx, hasOverride +} + +func preparedModelRouteFromContext(ctx context.Context) (modelRouteDecision, bool) { + if ctx == nil { + return modelRouteDecision{}, false + } + decision, ok := ctx.Value(preparedModelRouteContextKey{}).(modelRouteDecision) + return decision, ok +} + +// WithExecutionSessionID returns a child context tagged with a long-lived execution session ID. +func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context { + sessionID = strings.TrimSpace(sessionID) + if sessionID == "" { + return ctx + } + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, executionSessionContextKey{}, sessionID) +} + +// WithDisallowFreeAuth returns a child context that requests skipping known free-tier credentials. +func WithDisallowFreeAuth(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, disallowFreeAuthContextKey{}, true) +} + +// headersFromContext extracts the original HTTP request headers from the gin context +// embedded in the provided context. This allows session affinity selectors to read +// client-provided session headers. +func headersFromContext(ctx context.Context) http.Header { + if ctx == nil { + return nil + } + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + return ginCtx.Request.Header.Clone() + } + return nil +} + +// queryFromContext extracts the original HTTP request query parameters from the +// gin context embedded in the provided context. Mirrors headersFromContext so +// model routers can observe inbound query parameters for plain HTTP requests, +// where execOptions.Query is not populated by callers. +func queryFromContext(ctx context.Context) url.Values { + if ctx == nil { + return nil + } + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil && ginCtx.Request.URL != nil { + return ginCtx.Request.URL.Query() + } + return nil +} + +func pinnedAuthIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + raw := ctx.Value(pinnedAuthContextKey{}) + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) + case []byte: + return strings.TrimSpace(string(v)) + default: + return "" + } +} + +func selectedAuthIDCallbackFromContext(ctx context.Context) func(string) { + if ctx == nil { + return nil + } + raw := ctx.Value(selectedAuthCallbackContextKey{}) + if callback, ok := raw.(func(string)); ok && callback != nil { + return callback + } + return nil +} + +func executionSessionIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + raw := ctx.Value(executionSessionContextKey{}) + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) + case []byte: + return strings.TrimSpace(string(v)) + default: + return "" + } +} + +func disallowFreeAuthFromContext(ctx context.Context) bool { + if ctx == nil { + return false + } + raw, ok := ctx.Value(disallowFreeAuthContextKey{}).(bool) + return ok && raw +} diff --git a/sdk/api/handlers/handlers_errors.go b/sdk/api/handlers/handlers_errors.go new file mode 100644 index 000000000..960442bd4 --- /dev/null +++ b/sdk/api/handlers/handlers_errors.go @@ -0,0 +1,145 @@ +package handlers + +import ( + "bytes" + "errors" + "fmt" + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "golang.org/x/net/context" +) + +func statusFromError(err error) int { + if err == nil { + return 0 + } + if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { + if code := se.StatusCode(); code > 0 { + return code + } + } + return 0 +} + +func enrichAuthSelectionError(err error, providers []string, model string) error { + if err == nil { + return nil + } + + var authErr *coreauth.Error + if !errors.As(err, &authErr) || authErr == nil { + return err + } + + code := strings.TrimSpace(authErr.Code) + if code != "auth_not_found" && code != "auth_unavailable" { + return err + } + + providerText := strings.Join(providers, ",") + if providerText == "" { + providerText = "unknown" + } + modelText := strings.TrimSpace(model) + if modelText == "" { + modelText = "unknown" + } + + baseMessage := strings.TrimSpace(authErr.Message) + if baseMessage == "" { + baseMessage = "no auth available" + } + detail := fmt.Sprintf("%s (providers=%s, model=%s)", baseMessage, providerText, modelText) + + // Clarify the most common alias confusion between Anthropic route names and internal provider keys. + if strings.Contains(","+providerText+",", ",claude,") { + detail += "; check Claude auth/key session and cooldown state via /v0/management/auth-files" + } + + status := authErr.HTTPStatus + if status <= 0 { + status = http.StatusServiceUnavailable + } + + return &coreauth.Error{ + Code: authErr.Code, + Message: detail, + Retryable: authErr.Retryable, + HTTPStatus: status, + } +} + +// WriteErrorResponse writes an error message to the response writer using the HTTP status embedded in the message. +func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.ErrorMessage) { + status := http.StatusInternalServerError + if msg != nil && msg.StatusCode > 0 { + status = msg.StatusCode + } + if msg != nil && msg.Error != nil { + for _, value := range coreauth.SafeResponseHeaders(msg.Error).Values("Retry-After") { + c.Writer.Header().Add("Retry-After", value) + } + } + if msg != nil && msg.Addon != nil && PassthroughHeadersEnabled(h.Cfg) { + for key, values := range msg.Addon { + if len(values) == 0 || IsCPAReservedResponseHeader(key) { + continue + } + c.Writer.Header().Del(key) + for _, value := range values { + c.Writer.Header().Add(key, value) + } + } + } + + errText := http.StatusText(status) + if msg != nil && msg.Error != nil { + if v := strings.TrimSpace(msg.Error.Error()); v != "" { + errText = v + } + } + + body := BuildErrorResponseBody(status, errText) + // Append first to preserve upstream response logs, then drop duplicate payloads if already recorded. + var previous []byte + if existing, exists := c.Get("API_RESPONSE"); exists { + if existingBytes, ok := existing.([]byte); ok && len(existingBytes) > 0 { + previous = existingBytes + } + } + appendAPIResponse(c, body) + trimmedErrText := strings.TrimSpace(errText) + trimmedBody := bytes.TrimSpace(body) + if len(previous) > 0 { + if (trimmedErrText != "" && bytes.Contains(previous, []byte(trimmedErrText))) || + (len(trimmedBody) > 0 && bytes.Contains(previous, trimmedBody)) { + c.Set("API_RESPONSE", previous) + } + } + + if !c.Writer.Written() { + c.Writer.Header().Set("Content-Type", "application/json") + } + c.Status(status) + _, _ = c.Writer.Write(body) +} + +func (h *BaseAPIHandler) LoggingAPIResponseError(ctx context.Context, err *interfaces.ErrorMessage) { + if h.Cfg.RequestLog { + if ginContext, ok := ctx.Value("gin").(*gin.Context); ok { + if apiResponseErrors, isExist := ginContext.Get("API_RESPONSE_ERROR"); isExist { + if slicesAPIResponseError, isOk := apiResponseErrors.([]*interfaces.ErrorMessage); isOk { + slicesAPIResponseError = append(slicesAPIResponseError, err) + ginContext.Set("API_RESPONSE_ERROR", slicesAPIResponseError) + } + } else { + // Create new response data entry + ginContext.Set("API_RESPONSE_ERROR", []*interfaces.ErrorMessage{err}) + } + } + } +} diff --git a/sdk/api/handlers/handlers_execution.go b/sdk/api/handlers/handlers_execution.go new file mode 100644 index 000000000..18508b781 --- /dev/null +++ b/sdk/api/handlers/handlers_execution.go @@ -0,0 +1,296 @@ +package handlers + +import ( + "fmt" + "net/http" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "golang.org/x/net/context" +) + +// PluginExecutorHost executes a routed request with a specific plugin executor. +type PluginExecutorHost interface { + ExecutePluginExecutor(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) + ExecutePluginExecutorStream(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) + CountPluginExecutor(context.Context, string, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) +} + +type pluginExecutorFormatResolver interface { + PluginExecutorRequestToFormat(string, coreexecutor.Request, coreexecutor.Options) sdktranslator.Format +} + +// ExecuteWithAuthManager executes a non-streaming request via the core auth manager. +// This path is the only supported execution route. +func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) { + return h.executeWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, false) +} + +// ExecuteImageWithAuthManager executes an OpenAI-compatible image endpoint request. +func (h *BaseAPIHandler) ExecuteImageWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) { + return h.executeWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, true) +} + +func (h *BaseAPIHandler) executeWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, allowImageModel bool) ([]byte, http.Header, *interfaces.ErrorMessage) { + return h.executeWithAuthManagerFormats(ctx, handlerType, handlerType, modelName, rawJSON, alt, allowImageModel, modelExecutionOptions{}) +} + +func (h *BaseAPIHandler) executeWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { + originalRequestedModel := modelName + routeDecision := h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, false, execOptions) + responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol) + if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil { + return nil, nil, errMsg + } + if routeDecision.ExecutorPluginID != "" { + return h.executeWithPluginExecutor(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions) + } + providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, allowImageModel, routeDecision, execOptions) + if errMsg != nil { + return nil, nil, errMsg + } + providers = adjustExecutionProvidersForEntryProtocol(entryProtocol, providers) + reqMeta := requestExecutionMetadata(ctx) + reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel + addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) + addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) + setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON) + setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) + payload := rawJSON + if len(payload) == 0 { + payload = nil + } + req := coreexecutor.Request{ + Model: normalizedModel, + Payload: payload, + } + afterAuthCapture := &requestAfterAuthCapture{} + opts := coreexecutor.Options{ + Stream: false, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(entryProtocol), + ResponseFormat: sdktranslator.FromString(responseProtocol), + Headers: modelExecutionHeaders(ctx, execOptions.Headers), + Query: modelExecutionQuery(ctx, execOptions.Query), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), + } + opts.Metadata = reqMeta + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + resp, err := h.AuthManager.Execute(ctx, providers, req, opts) + if err != nil { + err = enrichAuthSelectionError(err, providers, normalizedModel) + status := http.StatusInternalServerError + if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { + if code := se.StatusCode(); code > 0 { + status = code + } + } + var addon http.Header + if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { + if hdr := he.Headers(); hdr != nil { + addon = hdr.Clone() + } + } + return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} + } + executedReq, executedOpts := afterAuthCapture.apply(req, opts) + rawResponseHeaders := cloneHeader(resp.Headers) + responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) + body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + return body, responseHeaders, nil +} + +// ExecuteCountWithAuthManager executes a non-streaming request via the core auth manager. +// This path is the only supported execution route. +func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) ([]byte, http.Header, *interfaces.ErrorMessage) { + return h.executeCountWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, modelExecutionOptions{}) +} + +func (h *BaseAPIHandler) executeCountWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { + originalRequestedModel := modelName + routeDecision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, false, execOptions) + if routeDecision.ExecutorPluginID != "" { + return h.countWithPluginExecutor(ctx, handlerType, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions) + } + providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, false, routeDecision, execOptions) + if errMsg != nil { + return nil, nil, errMsg + } + providers = adjustExecutionProvidersForEntryProtocol(handlerType, providers) + reqMeta := requestExecutionMetadata(ctx) + reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel + addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) + setReasoningEffortMetadata(reqMeta, handlerType, normalizedModel, rawJSON) + setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) + payload := rawJSON + if len(payload) == 0 { + payload = nil + } + req := coreexecutor.Request{ + Model: normalizedModel, + Payload: payload, + } + afterAuthCapture := &requestAfterAuthCapture{} + opts := coreexecutor.Options{ + Stream: false, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(handlerType), + Headers: modelExecutionHeaders(ctx, execOptions.Headers), + Query: modelExecutionQuery(ctx, execOptions.Query), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), + } + opts.Metadata = reqMeta + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts) + if err != nil { + err = enrichAuthSelectionError(err, providers, normalizedModel) + status := http.StatusInternalServerError + if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { + if code := se.StatusCode(); code > 0 { + status = code + } + } + var addon http.Header + if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { + if hdr := he.Headers(); hdr != nil { + addon = hdr.Clone() + } + } + return nil, nil, &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} + } + executedReq, executedOpts := afterAuthCapture.apply(req, opts) + rawResponseHeaders := cloneHeader(resp.Headers) + responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) + body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, normalizedModel, originalRequestedModel, executedOpts, rawResponseHeaders, responseHeaders, executedOpts.OriginalRequest, executedReq.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + return body, responseHeaders, nil +} + +func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { + if h.AuthManager != nil && h.AuthManager.HomeEnabled() { + return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} + } + host := h.pluginExecutorHost() + if host == nil { + return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} + } + req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + resp, errExecute := host.ExecutePluginExecutor(ctx, executorPluginID, req, opts) + if errExecute != nil { + return nil, nil, executionErrorMessage(errExecute) + } + rawResponseHeaders := cloneHeader(resp.Headers) + responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) + body, responseHeaders := h.applyResponseInterceptors(ctx, responseProtocol, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + return body, responseHeaders, nil +} + +func (h *BaseAPIHandler) countWithPluginExecutor(ctx context.Context, handlerType, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { + if h.AuthManager != nil && h.AuthManager.HomeEnabled() { + return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} + } + host := h.pluginExecutorHost() + if host == nil { + return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} + } + req, opts := h.pluginExecutorRequest(ctx, handlerType, handlerType, modelName, originalRequestedModel, rawJSON, alt, false, execOptions) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, handlerType, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + resp, errCount := host.CountPluginExecutor(ctx, executorPluginID, req, opts) + if errCount != nil { + return nil, nil, executionErrorMessage(errCount) + } + rawResponseHeaders := cloneHeader(resp.Headers) + responseHeaders := downstreamHeadersFromExecutor(rawResponseHeaders, PassthroughHeadersEnabled(h.Cfg)) + body, responseHeaders := h.applyResponseInterceptors(ctx, handlerType, modelName, originalRequestedModel, opts, rawResponseHeaders, responseHeaders, opts.OriginalRequest, req.Payload, resp.Payload, http.StatusOK, execOptions.SkipInterceptorPluginID) + return body, responseHeaders, nil +} + +func (h *BaseAPIHandler) pluginExecutorRequest(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt string, stream bool, execOptions modelExecutionOptions) (coreexecutor.Request, coreexecutor.Options) { + reqMeta := requestExecutionMetadata(ctx) + reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel + addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) + addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) + setReasoningEffortMetadata(reqMeta, entryProtocol, modelName, rawJSON) + setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) + payload := rawJSON + if len(payload) == 0 { + payload = nil + } + req := coreexecutor.Request{Model: modelName, Payload: payload} + opts := coreexecutor.Options{ + Stream: stream, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(entryProtocol), + ResponseFormat: sdktranslator.FromString(responseProtocol), + Headers: modelExecutionHeaders(ctx, execOptions.Headers), + Query: modelExecutionQuery(ctx, execOptions.Query), + Metadata: reqMeta, + } + return req, opts +} + +func (h *BaseAPIHandler) applyRequestInterceptorsAfterPluginExecutorRoute(ctx context.Context, host PluginExecutorHost, executorPluginID, entryProtocol, originalRequestedModel string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options) { + if !requestInterceptorsEnabled(h.interceptorHost()) { + return req, opts + } + toFormat := sdktranslator.FromString(entryProtocol) + if resolver, ok := host.(pluginExecutorFormatResolver); ok && resolver != nil { + if resolved := resolver.PluginExecutorRequestToFormat(executorPluginID, req, opts); resolved != "" { + toFormat = resolved + } + } + resp := h.applyRequestInterceptorsAfterAuth(ctx, coreexecutor.RequestAfterAuthInterceptRequest{ + SourceFormat: opts.SourceFormat, + ToFormat: toFormat, + Model: req.Model, + RequestedModel: originalRequestedModel, + Stream: opts.Stream, + Headers: cloneHeader(opts.Headers), + Body: cloneBytes(req.Payload), + Metadata: opts.Metadata, + }, skipPluginID) + opts.Headers = mergeRequestInterceptorHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + req.Payload = cloneBytes(resp.Body) + opts.OriginalRequest = cloneBytes(resp.Body) + } + return req, opts +} + +func executionErrorMessage(err error) *interfaces.ErrorMessage { + status := http.StatusInternalServerError + if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { + if code := se.StatusCode(); code > 0 { + status = code + } + } + var addon http.Header + if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { + if hdr := he.Headers(); hdr != nil { + addon = hdr.Clone() + } + } + return &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} +} + +func (h *BaseAPIHandler) pluginExecutorHost() PluginExecutorHost { + if h == nil { + return nil + } + if executorHost, ok := h.ModelRouterHost.(PluginExecutorHost); ok && executorHost != nil { + return executorHost + } + if executorHost, ok := h.PluginHost.(PluginExecutorHost); ok && executorHost != nil { + return executorHost + } + return nil +} diff --git a/sdk/api/handlers/handlers_interceptors.go b/sdk/api/handlers/handlers_interceptors.go new file mode 100644 index 000000000..eb90d7289 --- /dev/null +++ b/sdk/api/handlers/handlers_interceptors.go @@ -0,0 +1,377 @@ +package handlers + +import ( + "net/http" + "sync" + + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + "golang.org/x/net/context" +) + +// PluginInterceptorHost applies plugin interceptors around handler execution. +type PluginInterceptorHost interface { + InterceptRequestBeforeAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + InterceptRequestAfterAuth(context.Context, pluginapi.RequestInterceptRequest) pluginapi.RequestInterceptResponse + InterceptResponse(context.Context, pluginapi.ResponseInterceptRequest) pluginapi.ResponseInterceptResponse + InterceptStreamChunk(context.Context, pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse +} + +type pluginInterceptorSkipHost interface { + InterceptRequestBeforeAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse + InterceptRequestAfterAuthExcept(context.Context, pluginapi.RequestInterceptRequest, string) pluginapi.RequestInterceptResponse + InterceptResponseExcept(context.Context, pluginapi.ResponseInterceptRequest, string) pluginapi.ResponseInterceptResponse + InterceptStreamChunkExcept(context.Context, pluginapi.StreamChunkInterceptRequest, string) pluginapi.StreamChunkInterceptResponse +} + +type streamInterceptorDetector interface { + HasStreamInterceptors() bool +} + +type requestInterceptorDetector interface { + HasRequestInterceptors() bool +} + +func cloneHeader(src http.Header) http.Header { + if src == nil { + return nil + } + dst := make(http.Header, len(src)) + for key, values := range src { + dst[key] = append([]string(nil), values...) + } + return dst +} + +func cloneByteSlices(src [][]byte) [][]byte { + if len(src) == 0 { + return nil + } + dst := make([][]byte, 0, len(src)) + for _, item := range src { + dst = append(dst, cloneBytes(item)) + } + return dst +} + +func nextStreamChunk(ctx context.Context, pending *[]coreexecutor.StreamChunk, closed *bool, chunks <-chan coreexecutor.StreamChunk) (coreexecutor.StreamChunk, bool, bool) { + if pending != nil && len(*pending) > 0 { + chunk := (*pending)[0] + (*pending)[0] = coreexecutor.StreamChunk{} + *pending = (*pending)[1:] + return chunk, true, false + } + if closed != nil && *closed { + return coreexecutor.StreamChunk{}, false, false + } + var chunk coreexecutor.StreamChunk + var ok bool + if ctx != nil { + select { + case <-ctx.Done(): + return coreexecutor.StreamChunk{}, false, true + case chunk, ok = <-chunks: + } + } else { + chunk, ok = <-chunks + } + if !ok && closed != nil { + *closed = true + } + return chunk, ok, false +} + +func appendStreamInterceptorHistory(history [][]byte, chunk []byte) [][]byte { + if len(chunk) == 0 { + return history + } + history = append(history, cloneBytes(chunk)) + for len(history) > maxStreamInterceptorHistoryChunks || byteSlicesSize(history) > maxStreamInterceptorHistoryBytes { + history[0] = nil + history = history[1:] + } + if len(history) == 0 { + return nil + } + return history +} + +func byteSlicesSize(items [][]byte) int { + total := 0 + for _, item := range items { + total += len(item) + } + return total +} + +func finalInterceptorHeaders(current, intercepted http.Header) http.Header { + if intercepted == nil { + return current + } + if len(intercepted) == 0 { + return nil + } + return cloneHeader(intercepted) +} + +func downstreamHeadersFromExecutor(headers http.Header, passthrough bool) http.Header { + if !passthrough { + return nil + } + return FilterUpstreamHeaders(headers) +} + +func downstreamHeadersAfterInterceptors(baseRaw, finalRaw http.Header, passthrough bool) http.Header { + if passthrough { + return FilterUpstreamHeaders(finalRaw) + } + return FilterUpstreamHeaders(diffHeaders(baseRaw, finalRaw)) +} + +func diffHeaders(base, next http.Header) http.Header { + if len(next) == 0 { + return nil + } + baseValues := make(map[string][]string, len(base)) + for key, values := range base { + baseValues[http.CanonicalHeaderKey(key)] = values + } + out := make(http.Header) + for key, values := range next { + canonicalKey := http.CanonicalHeaderKey(key) + if stringSlicesEqual(baseValues[canonicalKey], values) { + continue + } + out[canonicalKey] = append([]string(nil), values...) + } + if len(out) == 0 { + return nil + } + return out +} + +func stringSlicesEqual(left, right []string) bool { + if len(left) != len(right) { + return false + } + for i := range left { + if left[i] != right[i] { + return false + } + } + return true +} + +func (h *BaseAPIHandler) interceptorHost() PluginInterceptorHost { + if h == nil { + return nil + } + return h.PluginHost +} + +func streamInterceptorsEnabled(host PluginInterceptorHost) bool { + if host == nil { + return false + } + if detector, ok := host.(streamInterceptorDetector); ok { + return detector.HasStreamInterceptors() + } + return true +} + +func requestInterceptorsEnabled(host PluginInterceptorHost) bool { + if host == nil { + return false + } + if detector, ok := host.(requestInterceptorDetector); ok { + return detector.HasRequestInterceptors() + } + return true +} + +type requestAfterAuthCapture struct { + mu sync.Mutex + set bool + headers http.Header + body []byte + originalRequest []byte + originalRequestReplaced bool +} + +func (c *requestAfterAuthCapture) record(req coreexecutor.RequestAfterAuthInterceptRequest, resp coreexecutor.RequestAfterAuthInterceptResponse) { + if c == nil { + return + } + headers := mergeRequestInterceptorHeaders(req.Headers, resp.Headers, resp.ClearHeaders) + body := cloneBytes(req.Body) + var originalRequest []byte + originalRequestReplaced := false + if len(resp.Body) > 0 { + body = cloneBytes(resp.Body) + originalRequest = cloneBytes(resp.Body) + originalRequestReplaced = true + } + + c.mu.Lock() + defer c.mu.Unlock() + c.set = true + c.headers = headers + c.body = body + c.originalRequest = originalRequest + c.originalRequestReplaced = originalRequestReplaced +} + +func (c *requestAfterAuthCapture) apply(req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Request, coreexecutor.Options) { + if c == nil { + return req, opts + } + c.mu.Lock() + defer c.mu.Unlock() + if !c.set { + return req, opts + } + req.Payload = cloneBytes(c.body) + opts.Headers = cloneHeader(c.headers) + if c.originalRequestReplaced { + opts.OriginalRequest = cloneBytes(c.originalRequest) + } + return req, opts +} + +func mergeRequestInterceptorHeaders(current, updates http.Header, clear []string) http.Header { + if updates == nil && len(clear) == 0 { + return cloneHeader(current) + } + out := cloneHeader(current) + if out == nil && (len(updates) > 0 || len(clear) > 0) { + out = make(http.Header) + } + for _, key := range clear { + out.Del(key) + } + for key, values := range updates { + out.Del(key) + for _, value := range values { + out.Add(key, value) + } + } + return out +} + +func interceptRequestBeforeAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptRequestBeforeAuthExcept(ctx, req, skipPluginID) + } + } + return host.InterceptRequestBeforeAuth(ctx, req) +} + +func interceptRequestAfterAuth(ctx context.Context, host PluginInterceptorHost, req pluginapi.RequestInterceptRequest, skipPluginID string) pluginapi.RequestInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptRequestAfterAuthExcept(ctx, req, skipPluginID) + } + } + return host.InterceptRequestAfterAuth(ctx, req) +} + +func interceptResponse(ctx context.Context, host PluginInterceptorHost, req pluginapi.ResponseInterceptRequest, skipPluginID string) pluginapi.ResponseInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptResponseExcept(ctx, req, skipPluginID) + } + } + return host.InterceptResponse(ctx, req) +} + +func interceptStreamChunk(ctx context.Context, host PluginInterceptorHost, req pluginapi.StreamChunkInterceptRequest, skipPluginID string) pluginapi.StreamChunkInterceptResponse { + if skipPluginID != "" { + if skipper, ok := host.(pluginInterceptorSkipHost); ok { + return skipper.InterceptStreamChunkExcept(ctx, req, skipPluginID) + } + } + return host.InterceptStreamChunk(ctx, req) +} + +func (h *BaseAPIHandler) applyRequestInterceptorsBeforeAuth(ctx context.Context, handlerType, requestedModel string, req coreexecutor.Request, opts coreexecutor.Options, skipPluginID string) (coreexecutor.Request, coreexecutor.Options) { + host := h.interceptorHost() + if host == nil { + return req, opts + } + resp := interceptRequestBeforeAuth(ctx, host, pluginapi.RequestInterceptRequest{ + SourceFormat: handlerType, + Model: req.Model, + RequestedModel: requestedModel, + Stream: opts.Stream, + Headers: cloneHeader(opts.Headers), + Body: cloneBytes(req.Payload), + Metadata: opts.Metadata, + }, skipPluginID) + opts.Headers = finalInterceptorHeaders(opts.Headers, resp.Headers) + if len(resp.Body) > 0 { + req.Payload = cloneBytes(resp.Body) + opts.OriginalRequest = cloneBytes(resp.Body) + } + return req, opts +} + +func (h *BaseAPIHandler) requestAfterAuthInterceptor(capture *requestAfterAuthCapture, skipPluginID string) coreexecutor.RequestAfterAuthInterceptor { + if !requestInterceptorsEnabled(h.interceptorHost()) { + return nil + } + return func(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest) coreexecutor.RequestAfterAuthInterceptResponse { + resp := h.applyRequestInterceptorsAfterAuth(ctx, req, skipPluginID) + if capture != nil { + capture.record(req, resp) + } + return resp + } +} + +func (h *BaseAPIHandler) applyRequestInterceptorsAfterAuth(ctx context.Context, req coreexecutor.RequestAfterAuthInterceptRequest, skipPluginID string) coreexecutor.RequestAfterAuthInterceptResponse { + host := h.interceptorHost() + if !requestInterceptorsEnabled(host) { + return coreexecutor.RequestAfterAuthInterceptResponse{} + } + resp := interceptRequestAfterAuth(ctx, host, pluginapi.RequestInterceptRequest{ + SourceFormat: req.SourceFormat.String(), + ToFormat: req.ToFormat.String(), + Model: req.Model, + RequestedModel: req.RequestedModel, + Stream: req.Stream, + Headers: cloneHeader(req.Headers), + Body: cloneBytes(req.Body), + Metadata: req.Metadata, + }, skipPluginID) + return coreexecutor.RequestAfterAuthInterceptResponse{ + Headers: resp.Headers, + Body: resp.Body, + ClearHeaders: resp.ClearHeaders, + } +} + +func (h *BaseAPIHandler) applyResponseInterceptors(ctx context.Context, handlerType, normalizedModel, requestedModel string, opts coreexecutor.Options, rawResponseHeaders, responseHeaders http.Header, originalRequest, requestBody, body []byte, statusCode int, skipPluginID string) ([]byte, http.Header) { + host := h.interceptorHost() + if host == nil { + return body, responseHeaders + } + resp := interceptResponse(ctx, host, pluginapi.ResponseInterceptRequest{ + SourceFormat: handlerType, + Model: normalizedModel, + RequestedModel: requestedModel, + Stream: false, + RequestHeaders: cloneHeader(opts.Headers), + ResponseHeaders: cloneHeader(rawResponseHeaders), + OriginalRequest: cloneBytes(originalRequest), + RequestBody: cloneBytes(requestBody), + Body: cloneBytes(body), + StatusCode: statusCode, + Metadata: opts.Metadata, + }, skipPluginID) + responseHeaders = downstreamHeadersAfterInterceptors(rawResponseHeaders, finalInterceptorHeaders(rawResponseHeaders, resp.Headers), PassthroughHeadersEnabled(h.Cfg)) + if len(resp.Body) > 0 { + body = cloneBytes(resp.Body) + } + return body, responseHeaders +} diff --git a/sdk/api/handlers/handlers_routing.go b/sdk/api/handlers/handlers_routing.go new file mode 100644 index 000000000..c590f415b --- /dev/null +++ b/sdk/api/handlers/handlers_routing.go @@ -0,0 +1,336 @@ +package handlers + +import ( + "fmt" + "net/http" + "strings" + + . "github.com/router-for-me/CLIProxyAPI/v7/internal/constant" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + "golang.org/x/net/context" +) + +// PluginModelRouterHost routes matching requests to a plugin executor, the router's own executor, +// or a built-in provider before model-to-provider resolution and auth selection. +type PluginModelRouterHost interface { + RouteModel(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) +} + +type pluginModelRouterSkipHost interface { + RouteModelExcept(context.Context, pluginapi.ModelRouteRequest, string) (pluginapi.ModelRouteResponse, bool) +} + +type modelRouterDetector interface { + HasModelRouters() bool +} + +type modelRouterSkipDetector interface { + HasModelRoutersExcept(string) bool +} + +func preferExecutionProvider(providers []string, preferred string) []string { + preferred = strings.ToLower(strings.TrimSpace(preferred)) + if preferred == "" || len(providers) < 2 { + return providers + } + preferredIndex := -1 + for i := range providers { + if strings.ToLower(strings.TrimSpace(providers[i])) == preferred { + preferredIndex = i + break + } + } + if preferredIndex <= 0 { + return providers + } + out := make([]string, 0, len(providers)) + out = append(out, providers[preferredIndex]) + out = append(out, providers[:preferredIndex]...) + out = append(out, providers[preferredIndex+1:]...) + return out +} + +func adjustExecutionProvidersForEntryProtocol(entryProtocol string, providers []string) []string { + if entryProtocol == Interactions { + return preferExecutionProvider(providers, GeminiInteractions) + } + if supportsNativeInteractionsEntryProtocol(entryProtocol) { + return providers + } + return excludeExecutionProvider(providers, GeminiInteractions) +} + +func supportsNativeInteractionsEntryProtocol(entryProtocol string) bool { + switch entryProtocol { + case Interactions, OpenAI, OpenaiResponse, Claude, Gemini: + return true + default: + return false + } +} + +func excludeExecutionProvider(providers []string, excluded string) []string { + excluded = strings.ToLower(strings.TrimSpace(excluded)) + if excluded == "" || len(providers) == 0 { + return providers + } + excludedIndex := -1 + for i := range providers { + if strings.ToLower(strings.TrimSpace(providers[i])) == excluded { + excludedIndex = i + break + } + } + if excludedIndex == -1 { + return providers + } + out := make([]string, 0, len(providers)-1) + out = append(out, providers[:excludedIndex]...) + out = append(out, providers[excludedIndex+1:]...) + return out +} + +func (h *BaseAPIHandler) getRequestDetails(modelName string) (providers []string, normalizedModel string, err *interfaces.ErrorMessage) { + return h.getRequestDetailsWithOptions(modelName, false) +} + +func validateNativeInteractionsExecution(entryProtocol string, execOptions modelExecutionOptions, routeDecision modelRouteDecision) *interfaces.ErrorMessage { + forcedProvider := strings.ToLower(strings.TrimSpace(execOptions.ForcedProvider)) + if forcedProvider == "" || entryProtocol != Interactions { + return nil + } + if routeDecision.ExecutorPluginID != "" { + return nativeInteractionsExecutionError() + } + if routeProvider := strings.ToLower(strings.TrimSpace(routeDecision.Provider)); routeProvider != "" && routeProvider != forcedProvider { + return nativeInteractionsExecutionError() + } + return nil +} + +func nativeInteractionsExecutionError() *interfaces.ErrorMessage { + return &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("agent is only supported for native interactions execution"), + } +} + +// providersForExecution resolves the providers and normalized model for a request. When a model +// router selected a built-in provider, it skips model->provider resolution and uses the router's +// provider (with an optional target model); otherwise it falls back to the registry-based path. +func (h *BaseAPIHandler) providersForExecution(modelName, originalRequestedModel string, allowImageModel bool, routeDecision modelRouteDecision, execOptions modelExecutionOptions) ([]string, string, *interfaces.ErrorMessage) { + forcedProvider := strings.ToLower(strings.TrimSpace(execOptions.ForcedProvider)) + if forcedProvider != "" { + if routeDecision.ExecutorPluginID != "" { + return nil, "", nativeInteractionsExecutionError() + } + if routeProvider := strings.ToLower(strings.TrimSpace(routeDecision.Provider)); routeProvider != "" && routeProvider != forcedProvider { + return nil, "", nativeInteractionsExecutionError() + } + normalizedModel := strings.TrimSpace(modelName) + if normalizedModel == "" { + normalizedModel = strings.TrimSpace(originalRequestedModel) + } + if errMsg := h.validateImageOnlyModel(normalizedModel, allowImageModel); errMsg != nil { + return nil, "", errMsg + } + return []string{forcedProvider}, normalizedModel, nil + } + if routeDecision.Provider != "" { + normalizedModel := originalRequestedModel + if routeDecision.Model != "" { + normalizedModel = routeDecision.Model + } + if errMsg := h.validateImageOnlyModel(normalizedModel, allowImageModel); errMsg != nil { + return nil, "", errMsg + } + return []string{routeDecision.Provider}, normalizedModel, nil + } + return h.getRequestDetailsWithOptions(modelName, allowImageModel) +} + +func (h *BaseAPIHandler) getRequestDetailsWithOptions(modelName string, allowImageModel bool) (providers []string, normalizedModel string, err *interfaces.ErrorMessage) { + resolvedModelName := modelName + initialSuffix := thinking.ParseSuffix(modelName) + if initialSuffix.ModelName == "auto" { + if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() { + resolvedModelName = modelName + } else { + resolvedBase := util.ResolveAutoModel(initialSuffix.ModelName) + if initialSuffix.HasSuffix { + resolvedModelName = fmt.Sprintf("%s(%s)", resolvedBase, initialSuffix.RawSuffix) + } else { + resolvedModelName = resolvedBase + } + } + } else { + if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() { + resolvedModelName = modelName + } else { + resolvedModelName = util.ResolveAutoModel(modelName) + } + } + + parsed := thinking.ParseSuffix(resolvedModelName) + baseModel := strings.TrimSpace(parsed.ModelName) + + if errMsg := h.validateImageOnlyModel(baseModel, allowImageModel); errMsg != nil { + return nil, "", errMsg + } + + if h != nil && h.AuthManager != nil && h.AuthManager.HomeEnabled() { + return []string{"home"}, resolvedModelName, nil + } + + providers = util.GetProviderName(baseModel) + // Fallback: if baseModel has no provider but differs from resolvedModelName, + // try using the full model name. This handles edge cases where custom models + // may be registered with their full suffixed name (e.g., "my-model(8192)"). + // Evaluated in Story 11.8: This fallback is intentionally preserved to support + // custom model registrations that include thinking suffixes. + if len(providers) == 0 && baseModel != resolvedModelName { + providers = util.GetProviderName(resolvedModelName) + } + + if len(providers) == 0 { + return nil, "", &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("unknown provider for model %s", modelName)} + } + + // The thinking suffix is preserved in the model name itself, so no + // metadata-based configuration passing is needed. + return providers, resolvedModelName, nil +} + +func (h *BaseAPIHandler) validateImageOnlyModel(modelName string, allowImageModel bool) *interfaces.ErrorMessage { + baseModel := strings.TrimSpace(thinking.ParseSuffix(modelName).ModelName) + if baseModel == "" { + baseModel = strings.TrimSpace(modelName) + } + if isOpenAIImageOnlyModel(baseModel) && !allowImageModel { + return &interfaces.ErrorMessage{ + StatusCode: http.StatusServiceUnavailable, + Error: fmt.Errorf("model %s is only supported on /v1/images/generations and /v1/images/edits", routeModelBaseName(baseModel)), + } + } + return nil +} + +func isOpenAIImageOnlyModel(model string) bool { + switch strings.ToLower(strings.TrimSpace(routeModelBaseName(model))) { + case "gpt-image-1.5", "gpt-image-2", "grok-imagine-image", "grok-imagine-image-quality": + return true + default: + return false + } +} + +func routeModelBaseName(model string) string { + model = strings.TrimSpace(model) + if idx := strings.LastIndex(model, "/"); idx >= 0 && idx < len(model)-1 { + return strings.TrimSpace(model[idx+1:]) + } + return model +} + +func cloneBytes(src []byte) []byte { + if len(src) == 0 { + return nil + } + dst := make([]byte, len(src)) + copy(dst, src) + return dst +} + +func (h *BaseAPIHandler) modelRouterHost() PluginModelRouterHost { + if h == nil { + return nil + } + if !isNilPluginModelRouterHost(h.ModelRouterHost) { + return h.ModelRouterHost + } + host := h.interceptorHost() + if host == nil { + return nil + } + router, ok := host.(PluginModelRouterHost) + if !ok { + return nil + } + return router +} + +type modelRouteDecision struct { + ExecutorPluginID string + Provider string + Model string +} + +func routeModel(ctx context.Context, host PluginModelRouterHost, req pluginapi.ModelRouteRequest, skipPluginID string) (pluginapi.ModelRouteResponse, bool) { + if host == nil { + return pluginapi.ModelRouteResponse{}, false + } + skipPluginID = strings.TrimSpace(skipPluginID) + if skipPluginID != "" { + if skipper, ok := host.(pluginModelRouterSkipHost); ok { + return skipper.RouteModelExcept(ctx, req, skipPluginID) + } + return pluginapi.ModelRouteResponse{}, false + } + return host.RouteModel(ctx, req) +} + +func modelRoutersEnabled(host PluginModelRouterHost, skipPluginID string) bool { + if host == nil { + return false + } + skipPluginID = strings.TrimSpace(skipPluginID) + if skipPluginID != "" { + if _, ok := host.(pluginModelRouterSkipHost); !ok { + return false + } + if detector, ok := host.(modelRouterSkipDetector); ok { + return detector.HasModelRoutersExcept(skipPluginID) + } + } + if detector, ok := host.(modelRouterDetector); ok { + return detector.HasModelRouters() + } + // No detector: treat routing as disabled (same conservative default as before any + // ModelRouter existed). Hosts that route must implement HasModelRouters (pluginhost.Host does). + return false +} + +func (h *BaseAPIHandler) applyModelRouter(ctx context.Context, handlerType, modelName string, rawJSON []byte, stream bool, execOptions modelExecutionOptions) modelRouteDecision { + var decision modelRouteDecision + host := h.modelRouterHost() + if host == nil || !modelRoutersEnabled(host, execOptions.SkipRouterPluginID) { + return decision + } + meta := requestExecutionMetadata(ctx) + meta[coreexecutor.RequestedModelMetadataKey] = modelName + addModelExecutionSourceMetadata(meta, execOptions.InternalSource) + resp, ok := routeModel(ctx, host, pluginapi.ModelRouteRequest{ + SourceFormat: handlerType, + RequestedModel: modelName, + Stream: stream, + Headers: modelExecutionHeaders(ctx, execOptions.Headers), + Query: modelExecutionQuery(ctx, execOptions.Query), + Body: cloneBytes(rawJSON), + Metadata: meta, + }, execOptions.SkipRouterPluginID) + if !ok || !resp.Handled { + return decision + } + switch resp.TargetKind { + case pluginapi.ModelRouteTargetSelf, pluginapi.ModelRouteTargetExecutor: + decision.ExecutorPluginID = strings.TrimSpace(resp.Target) + case pluginapi.ModelRouteTargetProvider: + decision.Provider = strings.ToLower(strings.TrimSpace(resp.Target)) + decision.Model = strings.TrimSpace(resp.TargetModel) + } + return decision +} diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go new file mode 100644 index 000000000..51caffaee --- /dev/null +++ b/sdk/api/handlers/handlers_stream.go @@ -0,0 +1,542 @@ +package handlers + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + "golang.org/x/net/context" +) + +// ExecuteStreamWithAuthManager executes a streaming request via the core auth manager. +// This path is the only supported execution route. +// The returned http.Header carries upstream response headers captured before streaming begins. +func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { + return h.executeStreamWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, false) +} + +// ExecuteImageStreamWithAuthManager executes a streaming OpenAI-compatible image endpoint request. +func (h *BaseAPIHandler) ExecuteImageStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { + return h.executeStreamWithAuthManager(ctx, handlerType, modelName, rawJSON, alt, true) +} + +func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { + if h.AuthManager != nil && h.AuthManager.HomeEnabled() { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} + close(errChan) + return nil, nil, errChan + } + host := h.pluginExecutorHost() + if host == nil { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} + close(errChan) + return nil, nil, errChan + } + req, opts := h.pluginExecutorRequest(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, true, execOptions) + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + req, opts = h.applyRequestInterceptorsAfterPluginExecutorRoute(ctx, host, executorPluginID, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + streamResult, errStream := host.ExecutePluginExecutorStream(ctx, executorPluginID, req, opts) + if errStream != nil { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- executionErrorMessage(errStream) + close(errChan) + return nil, nil, errChan + } + if streamResult == nil { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor returned nil stream")} + close(errChan) + return nil, nil, errChan + } + + passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg) + interceptorHost := h.interceptorHost() + streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost) + rawStreamHeaders := cloneHeader(streamResult.Headers) + baseStreamHeaders := cloneHeader(streamResult.Headers) + applyStreamHeaders := func(headers http.Header) { + rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers) + } + if streamInterceptorsActive { + intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + SourceFormat: responseProtocol, + Model: modelName, + RequestedModel: originalRequestedModel, + RequestHeaders: cloneHeader(opts.Headers), + ResponseHeaders: cloneHeader(rawStreamHeaders), + OriginalRequest: cloneBytes(opts.OriginalRequest), + RequestBody: cloneBytes(req.Payload), + ChunkIndex: pluginapi.StreamChunkHeaderInitIndex, + Metadata: opts.Metadata, + }, execOptions.SkipInterceptorPluginID) + applyStreamHeaders(intercepted.Headers) + } + upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) + if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { + upstreamHeaders = make(http.Header) + } + + dataChan := make(chan []byte) + errChan := make(chan *interfaces.ErrorMessage, 1) + var done <-chan struct{} + if ctx != nil { + done = ctx.Done() + } + chunks := streamResult.Chunks + if chunks == nil { + closed := make(chan coreexecutor.StreamChunk) + close(closed) + chunks = closed + } + go func() { + defer close(dataChan) + defer close(errChan) + chunkIndex := 0 + var historyChunks [][]byte + for { + chunk, ok, canceled := nextStreamChunk(ctx, nil, nil, chunks) + if canceled { + return + } + if !ok { + return + } + if chunk.Err != nil { + select { + case errChan <- executionErrorMessage(chunk.Err): + case <-done: + } + return + } + if len(chunk.Payload) == 0 { + continue + } + payload := cloneBytes(chunk.Payload) + if streamInterceptorsActive { + intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + SourceFormat: responseProtocol, + Model: modelName, + RequestedModel: originalRequestedModel, + RequestHeaders: cloneHeader(opts.Headers), + ResponseHeaders: cloneHeader(rawStreamHeaders), + OriginalRequest: cloneBytes(opts.OriginalRequest), + RequestBody: cloneBytes(req.Payload), + Body: payload, + HistoryChunks: cloneByteSlices(historyChunks), + ChunkIndex: chunkIndex, + Metadata: opts.Metadata, + }, execOptions.SkipInterceptorPluginID) + applyStreamHeaders(intercepted.Headers) + if len(intercepted.Body) > 0 { + payload = cloneBytes(intercepted.Body) + } + chunkIndex++ + if intercepted.DropChunk { + continue + } + } else { + chunkIndex++ + } + if responseProtocol == "openai-response" { + if errValidate := validateSSEDataJSON(payload); errValidate != nil { + select { + case errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}: + case <-done: + } + return + } + } + select { + case dataChan <- payload: + if streamInterceptorsActive { + historyChunks = appendStreamInterceptorHistory(historyChunks, payload) + } + case <-done: + return + } + } + }() + return dataChan, upstreamHeaders, errChan +} + +func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handlerType, modelName string, rawJSON []byte, alt string, allowImageModel bool) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { + return h.executeStreamWithAuthManagerFormats(ctx, handlerType, handlerType, modelName, rawJSON, alt, allowImageModel, modelExecutionOptions{}) +} + +func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { + originalRequestedModel := modelName + routeDecision, preparedRoute := preparedModelRouteFromContext(ctx) + if !preparedRoute { + routeDecision = h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions) + } + responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol) + if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- errMsg + close(errChan) + return nil, nil, errChan + } + if routeDecision.ExecutorPluginID != "" { + return h.streamWithPluginExecutor(ctx, entryProtocol, responseProtocol, modelName, originalRequestedModel, rawJSON, alt, routeDecision.ExecutorPluginID, execOptions) + } + providers, normalizedModel, errMsg := h.providersForExecution(modelName, originalRequestedModel, allowImageModel, routeDecision, execOptions) + if errMsg != nil { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- errMsg + close(errChan) + return nil, nil, errChan + } + providers = adjustExecutionProvidersForEntryProtocol(entryProtocol, providers) + reqMeta := requestExecutionMetadata(ctx) + reqMeta[coreexecutor.RequestedModelMetadataKey] = originalRequestedModel + addAuthSelectionModelMetadata(reqMeta, execOptions.AuthSelectionModel) + addModelExecutionSourceMetadata(reqMeta, execOptions.InternalSource) + setReasoningEffortMetadata(reqMeta, entryProtocol, normalizedModel, rawJSON) + setServiceTierMetadata(reqMeta, rawJSON) + setGenerateMetadata(reqMeta, rawJSON) + payload := rawJSON + if len(payload) == 0 { + payload = nil + } + req := coreexecutor.Request{ + Model: normalizedModel, + Payload: payload, + } + afterAuthCapture := &requestAfterAuthCapture{} + opts := coreexecutor.Options{ + Stream: true, + Alt: alt, + OriginalRequest: rawJSON, + SourceFormat: sdktranslator.FromString(entryProtocol), + ResponseFormat: sdktranslator.FromString(responseProtocol), + Headers: modelExecutionHeaders(ctx, execOptions.Headers), + Query: modelExecutionQuery(ctx, execOptions.Query), + RequestAfterAuthInterceptor: h.requestAfterAuthInterceptor(afterAuthCapture, execOptions.SkipInterceptorPluginID), + } + opts.Metadata = reqMeta + req, opts = h.applyRequestInterceptorsBeforeAuth(ctx, entryProtocol, originalRequestedModel, req, opts, execOptions.SkipInterceptorPluginID) + streamResult, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts) + if err != nil { + err = enrichAuthSelectionError(err, providers, normalizedModel) + errChan := make(chan *interfaces.ErrorMessage, 1) + status := http.StatusInternalServerError + if se, ok := err.(interface{ StatusCode() int }); ok && se != nil { + if code := se.StatusCode(); code > 0 { + status = code + } + } + var addon http.Header + if he, ok := err.(interface{ Headers() http.Header }); ok && he != nil { + if hdr := he.Headers(); hdr != nil { + addon = hdr.Clone() + } + } + errChan <- &interfaces.ErrorMessage{StatusCode: status, Error: err, Addon: addon} + close(errChan) + return nil, nil, errChan + } + if streamResult == nil { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("auth manager returned nil stream")} + close(errChan) + return nil, nil, errChan + } + executedRequest := func() (coreexecutor.Request, coreexecutor.Options) { + return afterAuthCapture.apply(req, opts) + } + passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg) + interceptorHost := h.interceptorHost() + streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost) + // Resolve bootstrap retries and header initialization before returning so the + // returned header snapshot is never modified by the stream goroutine. + rawStreamHeaders := cloneHeader(streamResult.Headers) + baseStreamHeaders := cloneHeader(streamResult.Headers) + chunks := streamResult.Chunks + if chunks == nil { + closed := make(chan coreexecutor.StreamChunk) + close(closed) + chunks = closed + } + streamClosedBeforeRead := false + streamCanceledBeforeRead := false + streamHeaderInitialized := false + + applyStreamHeaders := func(headers http.Header) { + rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers) + } + + applyStreamHeaderInit := func() { + if !streamInterceptorsActive || streamHeaderInitialized { + return + } + executedReq, executedOpts := executedRequest() + intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + SourceFormat: responseProtocol, + Model: normalizedModel, + RequestedModel: originalRequestedModel, + RequestHeaders: cloneHeader(executedOpts.Headers), + ResponseHeaders: cloneHeader(rawStreamHeaders), + OriginalRequest: cloneBytes(executedOpts.OriginalRequest), + RequestBody: cloneBytes(executedReq.Payload), + ChunkIndex: pluginapi.StreamChunkHeaderInitIndex, + Metadata: executedOpts.Metadata, + }, execOptions.SkipInterceptorPluginID) + applyStreamHeaders(intercepted.Headers) + streamHeaderInitialized = true + } + + transformStreamPayload := func(payload []byte, chunkIndex *int, historyChunks [][]byte) ([]byte, bool, *interfaces.ErrorMessage) { + applyStreamHeaderInit() + payload = cloneBytes(payload) + if streamInterceptorsActive { + executedReq, executedOpts := executedRequest() + intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + SourceFormat: responseProtocol, + Model: normalizedModel, + RequestedModel: originalRequestedModel, + RequestHeaders: cloneHeader(executedOpts.Headers), + ResponseHeaders: cloneHeader(rawStreamHeaders), + OriginalRequest: cloneBytes(executedOpts.OriginalRequest), + RequestBody: cloneBytes(executedReq.Payload), + Body: payload, + HistoryChunks: cloneByteSlices(historyChunks), + ChunkIndex: *chunkIndex, + Metadata: executedOpts.Metadata, + }, execOptions.SkipInterceptorPluginID) + applyStreamHeaders(intercepted.Headers) + if len(intercepted.Body) > 0 { + payload = cloneBytes(intercepted.Body) + } + (*chunkIndex)++ + if intercepted.DropChunk { + return nil, false, nil + } + } else { + (*chunkIndex)++ + } + if responseProtocol == "openai-response" { + if errValidate := validateSSEDataJSON(payload); errValidate != nil { + return nil, false, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate} + } + } + return payload, true, nil + } + + var bootstrapPayload []byte + bootstrapChunkIndex := 0 + var bootstrapHistoryChunks [][]byte + var bootstrapStreamErr error + var bootstrapErr *interfaces.ErrorMessage + readInitialStreamChunks := func() { + for { + var chunk coreexecutor.StreamChunk + var ok bool + if ctx != nil { + select { + case <-ctx.Done(): + streamCanceledBeforeRead = true + return + case chunk, ok = <-chunks: + } + } else { + chunk, ok = <-chunks + } + if !ok { + streamClosedBeforeRead = true + applyStreamHeaderInit() + return + } + if chunk.Err != nil { + bootstrapStreamErr = chunk.Err + return + } + if len(chunk.Payload) == 0 { + continue + } + payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &bootstrapChunkIndex, bootstrapHistoryChunks) + if errMsg != nil { + bootstrapErr = errMsg + return + } + if !deliverable { + continue + } + bootstrapPayload = payload + return + } + } + + bootstrapEligible := func(err error) bool { + status := statusFromError(err) + if status == 0 { + return true + } + switch status { + case http.StatusUnauthorized, http.StatusForbidden, http.StatusPaymentRequired, + http.StatusRequestTimeout, http.StatusTooManyRequests: + return true + default: + return status >= http.StatusInternalServerError + } + } + + maxBootstrapRetries := StreamingBootstrapRetries(h.Cfg) + if h.AuthManager.HomeEnabled() { + maxBootstrapRetries = 0 + } + for bootstrapRetries := 0; !streamCanceledBeforeRead; { + readInitialStreamChunks() + if streamCanceledBeforeRead || bootstrapErr != nil || bootstrapStreamErr == nil { + break + } + if bootstrapRetries >= maxBootstrapRetries || !bootstrapEligible(bootstrapStreamErr) { + bootstrapErr = executionErrorMessage(bootstrapStreamErr) + break + } + bootstrapRetries++ + retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) + if retryErr != nil { + bootstrapErr = executionErrorMessage(enrichAuthSelectionError(retryErr, providers, normalizedModel)) + break + } + if retryResult == nil { + bootstrapErr = executionErrorMessage(fmt.Errorf("auth manager returned nil stream")) + break + } + rawStreamHeaders = cloneHeader(retryResult.Headers) + baseStreamHeaders = cloneHeader(retryResult.Headers) + streamHeaderInitialized = false + streamClosedBeforeRead = false + bootstrapStreamErr = nil + bootstrapPayload = nil + bootstrapChunkIndex = 0 + bootstrapHistoryChunks = nil + chunks = retryResult.Chunks + if chunks == nil { + closed := make(chan coreexecutor.StreamChunk) + close(closed) + chunks = closed + } + } + + upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) + if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { + upstreamHeaders = make(http.Header) + } + dataChan := make(chan []byte) + errChan := make(chan *interfaces.ErrorMessage, 1) + + go func() { + defer close(dataChan) + defer close(errChan) + if streamCanceledBeforeRead { + return + } + + sendErr := func(msg *interfaces.ErrorMessage) bool { + if ctx == nil { + errChan <- msg + return true + } + select { + case <-ctx.Done(): + return false + case errChan <- msg: + return true + } + } + + sendData := func(chunk []byte) bool { + if ctx == nil { + dataChan <- chunk + return true + } + select { + case <-ctx.Done(): + return false + case dataChan <- chunk: + return true + } + } + + if bootstrapErr != nil { + _ = sendErr(bootstrapErr) + return + } + + chunkIndex := bootstrapChunkIndex + historyChunks := bootstrapHistoryChunks + if bootstrapPayload != nil { + if okSendData := sendData(bootstrapPayload); !okSendData { + return + } + if streamInterceptorsActive { + historyChunks = appendStreamInterceptorHistory(historyChunks, bootstrapPayload) + } + } + for { + chunk, ok, canceled := nextStreamChunk(ctx, nil, &streamClosedBeforeRead, chunks) + if canceled || !ok { + return + } + if chunk.Err != nil { + _ = sendErr(executionErrorMessage(chunk.Err)) + return + } + if len(chunk.Payload) == 0 { + continue + } + payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &chunkIndex, historyChunks) + if errMsg != nil { + _ = sendErr(errMsg) + return + } + if !deliverable { + continue + } + if okSendData := sendData(payload); !okSendData { + return + } + if streamInterceptorsActive { + historyChunks = appendStreamInterceptorHistory(historyChunks, payload) + } + } + }() + return dataChan, upstreamHeaders, errChan +} + +func validateSSEDataJSON(chunk []byte) error { + for _, line := range bytes.Split(chunk, []byte("\n")) { + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + if !bytes.HasPrefix(line, []byte("data:")) { + continue + } + data := bytes.TrimSpace(line[5:]) + if len(data) == 0 { + continue + } + if bytes.Equal(data, []byte("[DONE]")) { + continue + } + if json.Valid(data) { + continue + } + const max = 512 + preview := data + if len(preview) > max { + preview = preview[:max] + } + return fmt.Errorf("invalid SSE data JSON (len=%d): %q", len(data), preview) + } + return nil +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index 2d017d36a..afcd7b8e6 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -3,13 +3,8 @@ package openai import ( "bytes" "context" - "encoding/json" "errors" - "fmt" - "io" "net/http" - "sort" - "strconv" "strings" "sync" "sync/atomic" @@ -20,10 +15,6 @@ import ( "github.com/google/uuid" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" - requestlogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" @@ -192,171 +183,6 @@ func truncateWebsocketCloseReason(reason string, maxBytes int) string { return truncated.String() } -type websocketTimelineAppender interface { - Append(eventType string, payload []byte, timestamp time.Time) -} - -type responsesWebsocketPinnedAuthState struct { - authID string - modelKey string -} - -type websocketTimelineLog struct { - enabled bool - source *requestlogging.FileBodySource - builder *strings.Builder - - currentPart io.WriteCloser - currentPartHasLog bool -} - -func newWebsocketTimelineLog(enabled bool, source *requestlogging.FileBodySource) *websocketTimelineLog { - if !enabled { - return &websocketTimelineLog{} - } - if source == nil { - return newInMemoryWebsocketTimelineLog() - } - return &websocketTimelineLog{ - enabled: true, - source: source, - } -} - -func newInMemoryWebsocketTimelineLog() *websocketTimelineLog { - return &websocketTimelineLog{ - enabled: true, - builder: &strings.Builder{}, - } -} - -func websocketTimelineSourceFromContext(c *gin.Context) *requestlogging.FileBodySource { - if c == nil { - return nil - } - value, exists := c.Get(requestlogging.WebsocketTimelineSourceContextKey) - if !exists { - return nil - } - source, ok := value.(*requestlogging.FileBodySource) - if !ok { - return nil - } - return source -} - -func (l *websocketTimelineLog) BeginRequest() { - if l == nil || !l.enabled || l.source == nil { - return - } - l.closeCurrentPart() - part, errCreate := l.source.CreatePart("request") - if errCreate != nil { - log.WithError(errCreate).Warn("failed to create websocket request detail log") - return - } - l.currentPart = part - l.currentPartHasLog = false -} - -func (l *websocketTimelineLog) Append(eventType string, payload []byte, timestamp time.Time) { - if l == nil || !l.enabled { - return - } - data := formatWebsocketTimelineEvent(eventType, payload, timestamp) - if len(data) == 0 { - return - } - if l.source != nil { - if l.currentPart == nil { - l.BeginRequest() - } - if l.currentPart == nil { - return - } - if errWrite := writeWebsocketTimelinePart(l.currentPart, data, l.currentPartHasLog); errWrite != nil { - log.WithError(errWrite).Warn("failed to write websocket request detail log") - return - } - l.currentPartHasLog = true - return - } - if l.builder != nil { - writeWebsocketTimelineBuilder(l.builder, data) - } -} - -func (l *websocketTimelineLog) SetContext(c *gin.Context) { - if l == nil || !l.enabled { - return - } - l.closeCurrentPart() - if l.source != nil { - if l.source.HasPayload() { - c.Set(requestlogging.WebsocketTimelineSourceContextKey, l.source) - return - } - if errCleanup := l.source.Cleanup(); errCleanup != nil { - log.WithError(errCleanup).Warn("failed to clean up empty websocket timeline log parts") - } - } - if l.builder != nil { - setWebsocketTimelineBody(c, l.builder.String()) - } -} - -func (l *websocketTimelineLog) String() string { - if l == nil || !l.enabled { - return "" - } - l.closeCurrentPart() - if l.source != nil { - data, errRead := l.source.Bytes() - if errRead != nil { - return "" - } - return string(data) - } - if l.builder == nil { - return "" - } - return l.builder.String() -} - -func (l *websocketTimelineLog) closeCurrentPart() { - if l == nil || l.currentPart == nil { - return - } - if errClose := l.currentPart.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close websocket request detail log") - } - l.currentPart = nil - l.currentPartHasLog = false -} - -func writeWebsocketTimelinePart(w io.Writer, data []byte, prependNewline bool) error { - if w == nil || len(data) == 0 { - return nil - } - if prependNewline { - if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { - return errWrite - } - } - _, errWrite := w.Write(data) - return errWrite -} - -func writeWebsocketTimelineBuilder(builder *strings.Builder, data []byte) { - if builder == nil || len(data) == 0 { - return - } - if builder.Len() > 0 { - builder.WriteString("\n") - } - builder.Write(data) -} - // ResponsesWebsocket handles websocket requests for /v1/responses. // It accepts `response.create` and `response.append` requests and streams // response events back as JSON websocket text messages. @@ -811,1546 +637,3 @@ func responsesWebsocketPreviousResponseNotFoundError() *interfaces.ErrorMessage ), } } - -func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) { - return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true) -} - -func normalizeResponsesWebsocketRequestWithMode(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { - return normalizeResponsesWebsocketRequestWithLastResponseID(rawJSON, lastRequest, lastResponseOutput, "", allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) -} - -func normalizeResponsesWebsocketRequestWithLastResponseID(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { - return normalizeResponsesWebsocketRequestWithIncrementalState(rawJSON, lastRequest, lastResponseOutput, lastResponseID, nil, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) -} - -func normalizeResponsesWebsocketRequestWithIncrementalState(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, lastResponsePendingToolCallIDs []string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { - requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) - switch requestType { - case wsRequestTypeCreate: - // log.Infof("responses websocket: response.create request") - if len(lastRequest) == 0 { - return normalizeResponseCreateRequest(rawJSON) - } - return normalizeResponseSubsequentRequest(rawJSON, lastRequest, lastResponseOutput, lastResponseID, lastResponsePendingToolCallIDs, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) - case wsRequestTypeAppend: - // log.Infof("responses websocket: response.append request") - return normalizeResponseSubsequentRequest(rawJSON, lastRequest, lastResponseOutput, lastResponseID, lastResponsePendingToolCallIDs, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) - default: - return nil, lastRequest, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("unsupported websocket request type: %s", requestType), - } - } -} - -func normalizeResponseCreateRequest(rawJSON []byte) ([]byte, []byte, *interfaces.ErrorMessage) { - normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") - if errDelete != nil { - normalized = bytes.Clone(rawJSON) - } - normalized, _ = sjson.SetBytes(normalized, "stream", true) - if !gjson.GetBytes(normalized, "input").Exists() { - normalized, _ = sjson.SetRawBytes(normalized, "input", []byte("[]")) - } - - modelName := strings.TrimSpace(gjson.GetBytes(normalized, "model").String()) - if modelName == "" { - return nil, nil, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("missing model in response.create request"), - } - } - return normalized, bytes.Clone(normalized), nil -} - -func normalizeResponseSubsequentRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, lastResponsePendingToolCallIDs []string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { - if len(lastRequest) == 0 { - return nil, lastRequest, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("websocket request received before response.create"), - } - } - - nextInput := gjson.GetBytes(rawJSON, "input") - if !nextInput.Exists() || !nextInput.IsArray() { - return nil, lastRequest, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("websocket request requires array field: input"), - } - } - - // Compaction can cause clients to replace local websocket history with a new - // compact transcript on the next `response.create`. When the input already - // contains historical model output items, treating it as an incremental append - // duplicates stale turn-state and can leave late orphaned function_call items. - if shouldReplaceWebsocketTranscript(rawJSON, nextInput) { - normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest) - return normalized, bytes.Clone(normalized), nil - } - - // Websocket v2 mode uses response.create with previous_response_id + incremental input. - // Do not expand it into a full input transcript; upstream expects the incremental payload. - if allowIncrementalInputWithPreviousResponseID { - prev := strings.TrimSpace(gjson.GetBytes(rawJSON, "previous_response_id").String()) - if prev == "" { - if !inputSatisfiesPendingToolCalls(nextInput, lastResponsePendingToolCallIDs) { - normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest) - return normalized, bytes.Clone(normalized), nil - } - prev = strings.TrimSpace(lastResponseID) - } - if prev != "" { - normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") - if errDelete != nil { - normalized = bytes.Clone(rawJSON) - } - normalized, _ = sjson.SetBytes(normalized, "previous_response_id", prev) - if !gjson.GetBytes(normalized, "model").Exists() { - modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) - if modelName != "" { - normalized, _ = sjson.SetBytes(normalized, "model", modelName) - } - } - if !gjson.GetBytes(normalized, "instructions").Exists() { - instructions := gjson.GetBytes(lastRequest, "instructions") - if instructions.Exists() { - normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw)) - } - } - normalized, _ = sjson.SetBytes(normalized, "stream", true) - return normalized, bytes.Clone(normalized), nil - } - } - - // When the client sends a compact replay for a downstream that can consume it - // directly, the input already carries the canonical history. In that case, - // skip merging with stale lastRequest/lastResponseOutput to avoid breaking - // function_call / function_call_output pairings. - // See: https://github.com/router-for-me/CLIProxyAPI/issues/2207 - var mergedInput string - if allowCompactionReplayBypass && inputContainsFullTranscript(nextInput) { - log.Infof("responses websocket: full transcript detected, skipping stale merge (input items=%d)", len(nextInput.Array())) - mergedInput = nextInput.Raw - } else { - appendInputRaw := nextInput.Raw - if inputContainsFullTranscript(nextInput) { - appendInputRaw = inputWithoutCompactionItems(nextInput) - } - - existingInput := gjson.GetBytes(lastRequest, "input") - var errMerge error - mergedInput, errMerge = mergeJSONArrayRaw(existingInput.Raw, normalizeJSONArrayRaw(lastResponseOutput)) - if errMerge != nil { - return nil, lastRequest, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("invalid previous response output: %w", errMerge), - } - } - - mergedInput, errMerge = mergeJSONArrayRaw(mergedInput, appendInputRaw) - if errMerge != nil { - return nil, lastRequest, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("invalid request input: %w", errMerge), - } - } - } - dedupedInput, errDedupeFunctionCalls := dedupeFunctionCallsByCallID(mergedInput) - if errDedupeFunctionCalls == nil { - mergedInput = dedupedInput - } - dedupedInput, errDedupeItemIDs := dedupeInputItemsByID(mergedInput) - if errDedupeItemIDs == nil { - mergedInput = dedupedInput - } - - normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") - if errDelete != nil { - normalized = bytes.Clone(rawJSON) - } - normalized, _ = sjson.DeleteBytes(normalized, "previous_response_id") - var errSet error - normalized, errSet = sjson.SetRawBytes(normalized, "input", []byte(mergedInput)) - if errSet != nil { - return nil, lastRequest, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("failed to merge websocket input: %w", errSet), - } - } - if !gjson.GetBytes(normalized, "model").Exists() { - modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) - if modelName != "" { - normalized, _ = sjson.SetBytes(normalized, "model", modelName) - } - } - if !gjson.GetBytes(normalized, "instructions").Exists() { - instructions := gjson.GetBytes(lastRequest, "instructions") - if instructions.Exists() { - normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw)) - } - } - normalized, _ = sjson.SetBytes(normalized, "stream", true) - return normalized, bytes.Clone(normalized), nil -} - -func shouldReplaceWebsocketTranscript(rawJSON []byte, nextInput gjson.Result) bool { - requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) - if requestType != wsRequestTypeCreate && requestType != wsRequestTypeAppend { - return false - } - previousResponseID := gjson.GetBytes(rawJSON, "previous_response_id") - if strings.TrimSpace(previousResponseID.String()) != "" { - return false - } - if !nextInput.Exists() || !nextInput.IsArray() { - return false - } - if requestType == wsRequestTypeCreate && !previousResponseID.Exists() && inputHasCodexLocalCompactionSummary(nextInput) { - return true - } - - for _, item := range nextInput.Array() { - switch strings.TrimSpace(item.Get("type").String()) { - case "function_call", "custom_tool_call": - return true - case "message": - if strings.TrimSpace(item.Get("role").String()) == "assistant" { - return true - } - } - } - - return false -} - -func inputHasCodexLocalCompactionSummary(input gjson.Result) bool { - if !input.IsArray() { - return false - } - - hasSummary := false - for index, item := range input.Array() { - itemType := strings.TrimSpace(item.Get("type").String()) - if itemType == "additional_tools" { - tools := item.Get("tools") - if index != 0 || strings.TrimSpace(item.Get("role").String()) != "developer" || !tools.IsArray() { - return false - } - for _, tool := range tools.Array() { - if !tool.IsObject() || strings.TrimSpace(tool.Get("type").String()) == "" { - return false - } - } - continue - } - if itemType != "" && itemType != "message" { - return false - } - - role := strings.TrimSpace(item.Get("role").String()) - if role != "user" && role != "developer" { - return false - } - if role == "user" && strings.HasPrefix(codexLocalCompactionMessageText(item), codexLocalCompactionSummaryPrefix+"\n") { - hasSummary = true - } - } - return hasSummary -} - -func codexLocalCompactionMessageText(message gjson.Result) string { - content := message.Get("content") - if content.Type == gjson.String { - return content.String() - } - if !content.IsArray() { - return "" - } - - var text strings.Builder - for _, part := range content.Array() { - if strings.TrimSpace(part.Get("type").String()) == "input_text" { - text.WriteString(part.Get("text").String()) - } - } - return text.String() -} - -func inputSatisfiesPendingToolCalls(input gjson.Result, pendingCallIDs []string) bool { - if len(pendingCallIDs) == 0 { - return true - } - if !input.IsArray() { - return false - } - outputs := make(map[string]struct{}, len(pendingCallIDs)) - for _, item := range input.Array() { - switch strings.TrimSpace(item.Get("type").String()) { - case "function_call_output", "custom_tool_call_output": - callID := strings.TrimSpace(item.Get("call_id").String()) - if callID != "" { - outputs[callID] = struct{}{} - } - } - } - for _, callID := range pendingCallIDs { - callID = strings.TrimSpace(callID) - if callID == "" { - continue - } - if _, ok := outputs[callID]; !ok { - return false - } - } - return true -} - -func normalizeResponseTranscriptReplacement(rawJSON []byte, lastRequest []byte) []byte { - normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") - if errDelete != nil { - normalized = bytes.Clone(rawJSON) - } - normalized, _ = sjson.DeleteBytes(normalized, "previous_response_id") - if !gjson.GetBytes(normalized, "model").Exists() { - modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) - if modelName != "" { - normalized, _ = sjson.SetBytes(normalized, "model", modelName) - } - } - if !gjson.GetBytes(normalized, "instructions").Exists() { - instructions := gjson.GetBytes(lastRequest, "instructions") - if instructions.Exists() { - normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw)) - } - } - normalized, _ = sjson.SetBytes(normalized, "stream", true) - return bytes.Clone(normalized) -} - -func dedupeFunctionCallsByCallID(rawArray string) (string, error) { - rawArray = strings.TrimSpace(rawArray) - if rawArray == "" { - return "[]", nil - } - var items []json.RawMessage - if errUnmarshal := json.Unmarshal([]byte(rawArray), &items); errUnmarshal != nil { - return "", errUnmarshal - } - - seenCallIDs := make(map[string]struct{}, len(items)) - filtered := make([]json.RawMessage, 0, len(items)) - for _, item := range items { - if len(item) == 0 { - continue - } - itemType := strings.TrimSpace(gjson.GetBytes(item, "type").String()) - if isResponsesToolCallType(itemType) { - callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String()) - if callID != "" { - if _, ok := seenCallIDs[callID]; ok { - continue - } - seenCallIDs[callID] = struct{}{} - } - } - filtered = append(filtered, item) - } - - out, errMarshal := json.Marshal(filtered) - if errMarshal != nil { - return "", errMarshal - } - return string(out), nil -} - -func dedupeResponsesWebsocketInputItemsByID(payload []byte) []byte { - input := gjson.GetBytes(payload, "input") - if !input.Exists() || !input.IsArray() { - return payload - } - dedupedInput, errDedupe := dedupeInputItemsByID(input.Raw) - if errDedupe != nil || dedupedInput == input.Raw { - return payload - } - updated, errSet := sjson.SetRawBytes(payload, "input", []byte(dedupedInput)) - if errSet != nil { - return payload - } - return updated -} - -func dedupeInputItemsByID(rawArray string) (string, error) { - rawArray = strings.TrimSpace(rawArray) - if rawArray == "" { - return "[]", nil - } - var items []json.RawMessage - if errUnmarshal := json.Unmarshal([]byte(rawArray), &items); errUnmarshal != nil { - return "", errUnmarshal - } - - // Parse each item's type, id and call_id once; gjson is a scan-based - // parser, so reusing this metadata avoids rescanning every item in each of - // the loops below as the conversation history grows. - type itemMetadata struct { - itemType string - id string - callID string - } - meta := make([]itemMetadata, len(items)) - for i, item := range items { - if len(item) == 0 { - continue - } - res := gjson.GetManyBytes(item, "type", "id", "call_id") - meta[i] = itemMetadata{ - itemType: strings.TrimSpace(res[0].String()), - id: strings.TrimSpace(res[1].String()), - callID: strings.TrimSpace(res[2].String()), - } - } - - // Collect the call_ids that are still referenced by tool-call output - // items. When several input items share the same id, the one we keep must - // preserve any call_id that has a matching output; otherwise the upstream - // rejects the request with "No tool call found for function call output". - referencedCallIDs := make(map[string]struct{}, len(items)) - for i := range items { - switch meta[i].itemType { - case "function_call_output", "custom_tool_call_output": - if meta[i].callID != "" { - referencedCallIDs[meta[i].callID] = struct{}{} - } - } - } - - // For each id, choose the index to keep. The default is the last - // occurrence (matching the original dedupe behavior), but we never replace - // an item whose call_id still has a matching output with one that does not. - // This keeps a single item per id while ensuring retained tool calls stay - // paired with their outputs. - keepIndexByID := make(map[string]int, len(items)) - keepReferencedByID := make(map[string]bool, len(items)) - for i := range items { - itemID := meta[i].id - if itemID == "" { - continue - } - _, referenced := referencedCallIDs[meta[i].callID] - referenced = referenced && meta[i].callID != "" - if _, seen := keepIndexByID[itemID]; !seen { - keepIndexByID[itemID] = i - keepReferencedByID[itemID] = referenced - continue - } - if referenced || !keepReferencedByID[itemID] { - keepIndexByID[itemID] = i - keepReferencedByID[itemID] = referenced - } - } - - filtered := make([]json.RawMessage, 0, len(items)) - for i, item := range items { - if len(item) == 0 { - continue - } - itemID := meta[i].id - if itemID != "" { - if keepIndexByID[itemID] != i { - continue - } - } - filtered = append(filtered, item) - } - - out, errMarshal := json.Marshal(filtered) - if errMarshal != nil { - return "", errMarshal - } - return string(out), nil -} - -func websocketUpstreamSupportsIncrementalInput(attributes map[string]string, metadata map[string]any) bool { - if len(attributes) > 0 { - if raw := strings.TrimSpace(attributes["websockets"]); raw != "" { - parsed, errParse := strconv.ParseBool(raw) - if errParse == nil { - return parsed - } - } - } - if len(metadata) == 0 { - return false - } - raw, ok := metadata["websockets"] - if !ok || raw == nil { - return false - } - switch value := raw.(type) { - case bool: - return value - case string: - parsed, errParse := strconv.ParseBool(strings.TrimSpace(value)) - if errParse == nil { - return parsed - } - default: - } - return false -} - -func (h *OpenAIResponsesAPIHandler) websocketUpstreamSupportsIncrementalInputForModel(modelName string) bool { - auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName) - for _, auth := range auths { - if responsesWebsocketAuthSupportsIncrementalInput(auth) { - return true - } - } - return false -} - -func (h *OpenAIResponsesAPIHandler) websocketUpstreamSupportsCompactionReplayForModel(modelName string) bool { - auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName) - if len(auths) == 0 { - return false - } - for _, auth := range auths { - if !responsesWebsocketAuthSupportsCompactionReplay(auth) { - return false - } - } - return true -} - -func (h *OpenAIResponsesAPIHandler) responsesWebsocketAvailableAuthsForModel(modelName string) ([]*coreauth.Auth, string) { - if h == nil || h.AuthManager == nil { - return nil, "" - } - resolvedModelName := responsesWebsocketResolvedModelName(modelName) - providerSet, modelKey := responsesWebsocketProviderSetForModel(resolvedModelName) - if len(providerSet) == 0 { - return nil, modelKey - } - - registryRef := registry.GetGlobalRegistry() - now := time.Now() - auths := h.AuthManager.List() - available := make([]*coreauth.Auth, 0, len(auths)) - for _, auth := range auths { - if !responsesWebsocketAuthMatchesModel(auth, providerSet, modelKey, registryRef, now) { - continue - } - available = append(available, auth) - } - return available, modelKey -} - -func (h *OpenAIResponsesAPIHandler) responsesWebsocketUsesCodexWebsocketPassthrough(modelName string) bool { - return h.responsesWebsocketUsesUpstreamWebsocketPassthrough(modelName) -} - -func (h *OpenAIResponsesAPIHandler) responsesWebsocketUsesUpstreamWebsocketPassthrough(modelName string) bool { - modelName = strings.TrimSpace(modelName) - if h == nil || h.AuthManager == nil || modelName == "" { - return false - } - auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName) - if len(auths) == 0 { - return false - } - provider := "" - for _, auth := range auths { - if auth == nil { - return false - } - authProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) - if authProvider != "codex" && authProvider != "xai" { - return false - } - if provider == "" { - provider = authProvider - if _, ok := h.AuthManager.Executor(provider); !ok { - return false - } - } else if authProvider != provider { - return false - } - if !websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) { - return false - } - } - return provider != "" -} - -func responsesWebsocketAuthSupportsIncrementalInput(auth *coreauth.Auth) bool { - if auth == nil { - return false - } - return websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) -} - -func responsesWebsocketPinnedAuthMatchesModel(auth *coreauth.Auth, modelName string, pinnedModelKey string, homeRuntime bool) bool { - if auth == nil { - return false - } - providerSet, modelKey := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(modelName)) - providerKey := strings.ToLower(strings.TrimSpace(auth.Provider)) - if _, ok := providerSet[providerKey]; !ok { - return false - } - if !responsesWebsocketAuthAvailableForModel(auth, modelKey, time.Now()) { - return false - } - - if homeRuntime { - return strings.EqualFold(strings.TrimSpace(pinnedModelKey), strings.TrimSpace(modelKey)) - } - return registry.GetGlobalRegistry().ClientSupportsModel(auth.ID, modelKey) -} - -func normalizeResponsesWebsocketPassthroughRequest(rawJSON []byte, modelName string) ([]byte, *interfaces.ErrorMessage) { - if !json.Valid(rawJSON) { - return nil, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("invalid websocket request JSON"), - } - } - - requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) - switch requestType { - case wsRequestTypeCreate, wsRequestTypeAppend: - default: - return nil, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("unsupported websocket request type: %s", requestType), - } - } - - normalized := bytes.Clone(rawJSON) - if strings.TrimSpace(gjson.GetBytes(normalized, "model").String()) == "" { - modelName = strings.TrimSpace(modelName) - if modelName == "" { - return nil, &interfaces.ErrorMessage{ - StatusCode: http.StatusBadRequest, - Error: fmt.Errorf("missing model in response.create request"), - } - } - normalized, _ = sjson.SetBytes(normalized, "model", modelName) - } - normalized, _ = sjson.SetBytes(normalized, "stream", true) - return normalized, nil -} - -func responsesWebsocketResolvedModelName(modelName string) string { - initialSuffix := thinking.ParseSuffix(modelName) - if initialSuffix.ModelName == "auto" { - resolvedBase := util.ResolveAutoModel(initialSuffix.ModelName) - if initialSuffix.HasSuffix { - return fmt.Sprintf("%s(%s)", resolvedBase, initialSuffix.RawSuffix) - } - return resolvedBase - } - return util.ResolveAutoModel(modelName) -} - -func responsesWebsocketProviderSetForModel(resolvedModelName string) (map[string]struct{}, string) { - parsed := thinking.ParseSuffix(resolvedModelName) - baseModel := strings.TrimSpace(parsed.ModelName) - providers := util.GetProviderName(baseModel) - if len(providers) == 0 && baseModel != resolvedModelName { - providers = util.GetProviderName(resolvedModelName) - } - providerSet := make(map[string]struct{}, len(providers)) - for _, provider := range providers { - providerKey := strings.TrimSpace(strings.ToLower(provider)) - if providerKey == "" { - continue - } - providerSet[providerKey] = struct{}{} - } - modelKey := baseModel - if modelKey == "" { - modelKey = strings.TrimSpace(resolvedModelName) - } - return providerSet, modelKey -} - -func responsesWebsocketAuthMatchesModel(auth *coreauth.Auth, providerSet map[string]struct{}, modelKey string, registryRef *registry.ModelRegistry, now time.Time) bool { - if auth == nil { - return false - } - providerKey := strings.TrimSpace(strings.ToLower(auth.Provider)) - if _, ok := providerSet[providerKey]; !ok { - return false - } - if modelKey != "" && registryRef != nil && !registryRef.ClientSupportsModel(auth.ID, modelKey) { - return false - } - return responsesWebsocketAuthAvailableForModel(auth, modelKey, now) -} - -func responsesWebsocketAuthSupportsCompactionReplay(auth *coreauth.Auth) bool { - if auth == nil { - return false - } - return strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") -} - -func responsesWebsocketAuthAvailableForModel(auth *coreauth.Auth, modelName string, now time.Time) bool { - if auth == nil { - return false - } - if auth.Disabled || auth.Status == coreauth.StatusDisabled { - return false - } - if modelName != "" && len(auth.ModelStates) > 0 { - state, ok := auth.ModelStates[modelName] - if (!ok || state == nil) && modelName != "" { - baseModel := strings.TrimSpace(thinking.ParseSuffix(modelName).ModelName) - if baseModel != "" && baseModel != modelName { - state, ok = auth.ModelStates[baseModel] - } - } - if ok && state != nil { - if state.Status == coreauth.StatusDisabled { - return false - } - if state.Unavailable && !state.NextRetryAfter.IsZero() && state.NextRetryAfter.After(now) { - return false - } - return true - } - } - if auth.Unavailable && !auth.NextRetryAfter.IsZero() && auth.NextRetryAfter.After(now) { - return false - } - return true -} - -func shouldHandleResponsesWebsocketPrewarmLocally(rawJSON []byte, lastRequest []byte, allowIncrementalInputWithPreviousResponseID bool) bool { - if allowIncrementalInputWithPreviousResponseID || len(lastRequest) != 0 { - return false - } - if strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) != wsRequestTypeCreate { - return false - } - generateResult := gjson.GetBytes(rawJSON, "generate") - return generateResult.Exists() && !generateResult.Bool() -} - -func writeResponsesWebsocketSyntheticPrewarm( - c *gin.Context, - writer *responsesWebsocketWriter, - requestJSON []byte, - wsTimelineLog websocketTimelineAppender, - sessionID string, -) error { - payloads, errPayloads := syntheticResponsesWebsocketPrewarmPayloads(requestJSON) - if errPayloads != nil { - return errPayloads - } - for i := 0; i < len(payloads); i++ { - markAPIResponseTimestamp(c) - // log.Infof( - // "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", - // sessionID, - // websocket.TextMessage, - // websocketPayloadEventType(payloads[i]), - // websocketPayloadPreview(payloads[i]), - // ) - if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { - log.Warnf( - "responses websocket: downstream_out write failed id=%s event=%s error=%v", - sessionID, - websocketPayloadEventType(payloads[i]), - errWrite, - ) - return errWrite - } - } - return nil -} - -func syntheticResponsesWebsocketPrewarmPayloads(requestJSON []byte) ([][]byte, error) { - responseID := "resp_prewarm_" + uuid.NewString() - createdAt := time.Now().Unix() - modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()) - - createdPayload := []byte(`{"type":"response.created","sequence_number":0,"response":{"id":"","object":"response","created_at":0,"status":"in_progress","background":false,"error":null,"output":[]}}`) - var errSet error - createdPayload, errSet = sjson.SetBytes(createdPayload, "response.id", responseID) - if errSet != nil { - return nil, errSet - } - createdPayload, errSet = sjson.SetBytes(createdPayload, "response.created_at", createdAt) - if errSet != nil { - return nil, errSet - } - if modelName != "" { - createdPayload, errSet = sjson.SetBytes(createdPayload, "response.model", modelName) - if errSet != nil { - return nil, errSet - } - } - - completedPayload := []byte(`{"type":"response.completed","sequence_number":1,"response":{"id":"","object":"response","created_at":0,"status":"completed","background":false,"error":null,"output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) - completedPayload, errSet = sjson.SetBytes(completedPayload, "response.id", responseID) - if errSet != nil { - return nil, errSet - } - completedPayload, errSet = sjson.SetBytes(completedPayload, "response.created_at", createdAt) - if errSet != nil { - return nil, errSet - } - if modelName != "" { - completedPayload, errSet = sjson.SetBytes(completedPayload, "response.model", modelName) - if errSet != nil { - return nil, errSet - } - } - - return [][]byte{createdPayload, completedPayload}, nil -} - -func mergeJSONArrayRaw(existingRaw, appendRaw string) (string, error) { - existingRaw = strings.TrimSpace(existingRaw) - appendRaw = strings.TrimSpace(appendRaw) - if existingRaw == "" { - existingRaw = "[]" - } - if appendRaw == "" { - appendRaw = "[]" - } - - var existing []json.RawMessage - if err := json.Unmarshal([]byte(existingRaw), &existing); err != nil { - return "", err - } - var appendItems []json.RawMessage - if err := json.Unmarshal([]byte(appendRaw), &appendItems); err != nil { - return "", err - } - - merged := append(existing, appendItems...) - out, err := json.Marshal(merged) - if err != nil { - return "", err - } - return string(out), nil -} - -// inputContainsFullTranscript returns true when the input array carries compact -// replay markers that indicate the client already sent the full conversation -// transcript. Merging that input with stale lastRequest/lastResponseOutput -// would duplicate or break function_call/function_call_output pairings, so the -// caller should use the input as-is. -// -// Assistant messages alone are not enough to classify the payload as a replay: -// incremental websocket requests may legitimately append assistant items. -func inputContainsFullTranscript(input gjson.Result) bool { - if !input.IsArray() { - return false - } - for _, item := range input.Array() { - t := item.Get("type").String() - if t == "compaction" || t == "compaction_summary" { - return true - } - } - return false -} - -func inputWithoutCompactionItems(input gjson.Result) string { - if !input.IsArray() { - return normalizeJSONArrayRaw([]byte(input.Raw)) - } - filtered := make([]string, 0, len(input.Array())) - for _, item := range input.Array() { - t := item.Get("type").String() - if t == "compaction" || t == "compaction_summary" { - continue - } - filtered = append(filtered, item.Raw) - } - return "[" + strings.Join(filtered, ",") + "]" -} - -func normalizeJSONArrayRaw(raw []byte) string { - trimmed := strings.TrimSpace(string(raw)) - if trimmed == "" { - return "[]" - } - result := gjson.Parse(trimmed) - if result.Type == gjson.JSON && result.IsArray() { - return trimmed - } - return "[]" -} - -type responsesWebsocketForwardOptions struct { - toolCacheTurn *responsesWebsocketToolCacheTurn - suppressError func(*interfaces.ErrorMessage) bool -} - -func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( - c *gin.Context, - writer *responsesWebsocketWriter, - cancel handlers.APIHandlerCancelFunc, - data <-chan []byte, - errs <-chan *interfaces.ErrorMessage, - wsTimelineLog websocketTimelineAppender, - sessionID string, - options ...responsesWebsocketForwardOptions, -) ([]byte, string, []string, *interfaces.ErrorMessage, error) { - var opts responsesWebsocketForwardOptions - if len(options) > 0 { - opts = options[0] - } - toolCacheTurn := opts.toolCacheTurn - completed := false - completedOutput := []byte("[]") - completedResponseID := "" - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte - pendingToolCallIDs := make(map[string]struct{}) - downstreamSessionKey := "" - if c != nil && c.Request != nil { - downstreamSessionKey = websocketDownstreamSessionKey(c.Request) - } - - for { - select { - case <-c.Request.Context().Done(): - cancel(c.Request.Context().Err()) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, c.Request.Context().Err() - case errMsg, ok := <-errs: - if !ok { - errs = nil - continue - } - if errMsg != nil { - h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) - if opts.suppressError != nil && opts.suppressError(errMsg) { - cancel(errMsg.Error) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil - } - markAPIResponseTimestamp(c) - if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched { - cancel(errMsg.Error) - if errClose != nil { - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errClose - } - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, websocket.ErrCloseSent - } - errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg) - log.Infof( - "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", - sessionID, - websocket.TextMessage, - websocketPayloadEventType(errorPayload), - websocketPayloadPreview(errorPayload), - ) - if errWrite != nil { - // log.Warnf( - // "responses websocket: downstream_out write failed id=%s event=%s error=%v", - // sessionID, - // websocketPayloadEventType(errorPayload), - // errWrite, - // ) - cancel(errMsg.Error) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errWrite - } - } - if errMsg != nil { - cancel(errMsg.Error) - } else { - cancel(nil) - } - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil - case chunk, ok := <-data: - if !ok { - if !completed { - errMsg := &interfaces.ErrorMessage{ - StatusCode: http.StatusRequestTimeout, - Error: fmt.Errorf("stream closed before response.completed"), - } - h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) - markAPIResponseTimestamp(c) - errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg) - log.Infof( - "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", - sessionID, - websocket.TextMessage, - websocketPayloadEventType(errorPayload), - websocketPayloadPreview(errorPayload), - ) - if errWrite != nil { - log.Warnf( - "responses websocket: downstream_out write failed id=%s event=%s error=%v", - sessionID, - websocketPayloadEventType(errorPayload), - errWrite, - ) - cancel(errMsg.Error) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errWrite - } - cancel(errMsg.Error) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil - } - cancel(nil) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, nil - } - - payloads := websocketJSONPayloadsFromChunk(chunk) - for i := range payloads { - collectResponsesWebsocketOutputItem(payloads[i], outputItemsByIndex, &outputItemsFallback) - eventType := gjson.GetBytes(payloads[i], "type").String() - if isResponsesWebsocketCompletionEvent(eventType) { - payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback) - } - if toolCacheTurn != nil { - toolCacheTurn.recordResponse(payloads[i]) - } else { - recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i]) - } - recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i]) - var payloadErrMsg *interfaces.ErrorMessage - if eventType == wsEventTypeError { - payloadErrMsg = responsesWebsocketErrorMessageFromPayload(payloads[i]) - if h != nil { - h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg) - } - if opts.suppressError != nil && opts.suppressError(payloadErrMsg) { - cancel(payloadErrMsg.Error) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil - } - } else if isResponsesWebsocketCompletionEvent(eventType) { - completed = true - completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback) - completedResponseID = responseCompletedIDFromPayload(payloads[i]) - } - markAPIResponseTimestamp(c) - // log.Infof( - // "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", - // sessionID, - // websocket.TextMessage, - // websocketPayloadEventType(payloads[i]), - // websocketPayloadPreview(payloads[i]), - // ) - if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { - log.Warnf( - "responses websocket: downstream_out write failed id=%s event=%s error=%v", - sessionID, - websocketPayloadEventType(payloads[i]), - errWrite, - ) - cancel(errWrite) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, errWrite - } - if payloadErrMsg != nil { - cancel(payloadErrMsg.Error) - return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil - } - } - } - } -} - -func responsesWebsocketErrorStatus(errMsg *interfaces.ErrorMessage) int { - if errMsg == nil { - return 0 - } - status := errMsg.StatusCode - if status <= 0 && errMsg.Error != nil { - if se, ok := errMsg.Error.(interface{ StatusCode() int }); ok && se != nil { - status = se.StatusCode() - } - } - return status -} - -func shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg *interfaces.ErrorMessage) bool { - switch responsesWebsocketErrorStatus(errMsg) { - case http.StatusUnauthorized, http.StatusTooManyRequests: - return true - default: - return false - } -} - -func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool { - if errMsg == nil { - return false - } - switch responsesWebsocketErrorStatus(errMsg) { - case http.StatusUnauthorized, - http.StatusPaymentRequired, - http.StatusForbidden, - http.StatusTooManyRequests, - http.StatusRequestTimeout, - http.StatusBadGateway, - http.StatusServiceUnavailable, - http.StatusGatewayTimeout: - return true - default: - } - if errMsg.Error != nil { - msg := strings.ToLower(errMsg.Error.Error()) - switch { - case strings.Contains(msg, "stream closed before response.completed"), - strings.Contains(msg, "previous_response_not_found"), - strings.Contains(msg, "ws_failed"), - strings.Contains(msg, "upstream stream closed before first payload"), - strings.Contains(msg, "empty_stream"): - return true - } - } - return false -} - -func collectResponsesWebsocketOutputItem(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback *[][]byte) { - if gjson.GetBytes(payload, "type").String() != "response.output_item.done" { - return - } - item := gjson.GetBytes(payload, "item") - if !item.Exists() || !item.IsObject() { - return - } - outputIndex := gjson.GetBytes(payload, "output_index") - if outputIndex.Exists() { - outputItemsByIndex[outputIndex.Int()] = bytes.Clone([]byte(item.Raw)) - return - } - *outputItemsFallback = append(*outputItemsFallback, bytes.Clone([]byte(item.Raw))) -} - -func restoreResponsesWebsocketCompletionOutput(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { - output := gjson.GetBytes(payload, "response.output") - if output.Exists() && output.IsArray() && len(output.Array()) > 0 { - reconciledOutput, changed := reconcileResponsesWebsocketCompletionToolCalls(output, outputItemsByIndex, outputItemsFallback) - if !changed { - return payload - } - restored, errSet := sjson.SetRawBytes(payload, "response.output", reconciledOutput) - if errSet != nil { - return payload - } - return restored - } - if len(outputItemsByIndex) == 0 && len(outputItemsFallback) == 0 { - return payload - } - - restored, errSet := sjson.SetRawBytes(payload, "response.output", responseCompletedOutputFromPayload(payload, outputItemsByIndex, outputItemsFallback)) - if errSet != nil { - return payload - } - return restored -} - -func reconcileResponsesWebsocketCompletionToolCalls(output gjson.Result, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) ([]byte, bool) { - collectedToolCalls := make(map[string]json.RawMessage) - recordCollectedToolCall := func(raw []byte) { - item := gjson.ParseBytes(raw) - if !isCompleteResponsesWebsocketToolCall(item) { - return - } - callID := strings.TrimSpace(item.Get("call_id").String()) - collectedToolCalls[callID] = append(json.RawMessage(nil), raw...) - } - - indexes := make([]int64, 0, len(outputItemsByIndex)) - for index := range outputItemsByIndex { - indexes = append(indexes, index) - } - sort.Slice(indexes, func(i, j int) bool { - return indexes[i] < indexes[j] - }) - for _, index := range indexes { - recordCollectedToolCall(outputItemsByIndex[index]) - } - for _, item := range outputItemsFallback { - recordCollectedToolCall(item) - } - if len(collectedToolCalls) == 0 { - return nil, false - } - - items := output.Array() - reconciled := make([]json.RawMessage, 0, len(items)) - changed := false - for _, item := range items { - raw := json.RawMessage(item.Raw) - if isResponsesToolCallType(item.Get("type").String()) { - callID := strings.TrimSpace(item.Get("call_id").String()) - if collected, ok := collectedToolCalls[callID]; ok && !bytes.Equal(raw, collected) { - raw = collected - changed = true - } - } - reconciled = append(reconciled, raw) - } - if !changed { - return nil, false - } - - marshaledOutput, errMarshal := json.Marshal(reconciled) - if errMarshal != nil { - return nil, false - } - return marshaledOutput, true -} - -func isCompleteResponsesWebsocketToolCall(item gjson.Result) bool { - if !item.Exists() || !item.IsObject() { - return false - } - callID := item.Get("call_id") - name := item.Get("name") - if callID.Type != gjson.String || strings.TrimSpace(callID.String()) == "" || name.Type != gjson.String || strings.TrimSpace(name.String()) == "" { - return false - } - - switch strings.TrimSpace(item.Get("type").String()) { - case "function_call": - arguments := item.Get("arguments") - return arguments.Exists() && arguments.Type == gjson.String - case "custom_tool_call": - input := item.Get("input") - return input.Exists() && input.Type == gjson.String - default: - return false - } -} - -func responseCompletedOutputFromPayload(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { - output := gjson.GetBytes(payload, "response.output") - if output.Exists() && output.IsArray() && len(output.Array()) > 0 { - return bytes.Clone([]byte(output.Raw)) - } - if len(outputItemsByIndex) == 0 && len(outputItemsFallback) == 0 { - return []byte("[]") - } - - indexes := make([]int64, 0, len(outputItemsByIndex)) - for index := range outputItemsByIndex { - indexes = append(indexes, index) - } - sort.Slice(indexes, func(i, j int) bool { - return indexes[i] < indexes[j] - }) - - items := make([]json.RawMessage, 0, len(outputItemsByIndex)+len(outputItemsFallback)) - appendCollectedItem := func(raw []byte) { - item := gjson.ParseBytes(raw) - if isResponsesToolCallType(item.Get("type").String()) && !isCompleteResponsesWebsocketToolCall(item) { - return - } - items = append(items, append(json.RawMessage(nil), raw...)) - } - for _, index := range indexes { - appendCollectedItem(outputItemsByIndex[index]) - } - for _, item := range outputItemsFallback { - appendCollectedItem(item) - } - - marshaledOutput, errMarshal := json.Marshal(items) - if errMarshal != nil { - return []byte("[]") - } - return marshaledOutput -} - -func responseCompletedIDFromPayload(payload []byte) string { - return strings.TrimSpace(gjson.GetBytes(payload, "response.id").String()) -} - -func recordPendingToolCallIDsFromPayload(pending map[string]struct{}, payload []byte) { - if pending == nil || len(payload) == 0 { - return - } - updatePendingToolCallIDsFromItem(pending, gjson.GetBytes(payload, "item")) - output := gjson.GetBytes(payload, "response.output") - if output.IsArray() { - for _, item := range output.Array() { - updatePendingToolCallIDsFromItem(pending, item) - } - } -} - -func updatePendingToolCallIDsFromItem(pending map[string]struct{}, item gjson.Result) { - if pending == nil || !item.Exists() { - return - } - switch strings.TrimSpace(item.Get("type").String()) { - case "function_call", "custom_tool_call": - if !isCompleteResponsesWebsocketToolCall(item) { - return - } - callID := strings.TrimSpace(item.Get("call_id").String()) - pending[callID] = struct{}{} - case "function_call_output", "custom_tool_call_output": - callID := strings.TrimSpace(item.Get("call_id").String()) - if callID != "" { - delete(pending, callID) - } - } -} - -func sortedStringSet(values map[string]struct{}) []string { - if len(values) == 0 { - return nil - } - out := make([]string, 0, len(values)) - for value := range values { - value = strings.TrimSpace(value) - if value != "" { - out = append(out, value) - } - } - sort.Strings(out) - return out -} - -func websocketJSONPayloadsFromChunk(chunk []byte) [][]byte { - payloads := make([][]byte, 0, 2) - lines := bytes.Split(chunk, []byte("\n")) - for i := range lines { - line := bytes.TrimSpace(lines[i]) - if len(line) == 0 || bytes.HasPrefix(line, []byte("event:")) { - continue - } - if bytes.HasPrefix(line, []byte("data:")) { - line = bytes.TrimSpace(line[len("data:"):]) - } - if len(line) == 0 || bytes.Equal(line, []byte(wsDoneMarker)) { - continue - } - if json.Valid(line) { - payloads = append(payloads, bytes.Clone(line)) - } - } - - if len(payloads) > 0 { - return payloads - } - - trimmed := bytes.TrimSpace(chunk) - if bytes.HasPrefix(trimmed, []byte("data:")) { - trimmed = bytes.TrimSpace(trimmed[len("data:"):]) - } - if len(trimmed) > 0 && !bytes.Equal(trimmed, []byte(wsDoneMarker)) && json.Valid(trimmed) { - payloads = append(payloads, bytes.Clone(trimmed)) - } - return payloads -} - -func writeResponsesWebsocketError(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, errMsg *interfaces.ErrorMessage) ([]byte, error) { - status := http.StatusInternalServerError - errText := http.StatusText(status) - if errMsg != nil { - if errMsg.StatusCode > 0 { - status = errMsg.StatusCode - errText = http.StatusText(status) - } - if errMsg.Error != nil && strings.TrimSpace(errMsg.Error.Error()) != "" { - errText = errMsg.Error.Error() - } - } - - body := handlers.BuildErrorResponseBody(status, errText) - payload := []byte(`{}`) - var errSet error - payload, errSet = sjson.SetBytes(payload, "type", wsEventTypeError) - if errSet != nil { - return nil, errSet - } - payload, errSet = sjson.SetBytes(payload, "status", status) - if errSet != nil { - return nil, errSet - } - - if errMsg != nil && errMsg.Addon != nil { - headers := []byte(`{}`) - hasHeaders := false - for key, values := range errMsg.Addon { - if len(values) == 0 { - continue - } - headerPath := strings.ReplaceAll(strings.ReplaceAll(key, `\\`, `\\\\`), ".", `\\.`) - headers, errSet = sjson.SetBytes(headers, headerPath, values[0]) - if errSet != nil { - return nil, errSet - } - hasHeaders = true - } - if hasHeaders { - payload, errSet = sjson.SetRawBytes(payload, "headers", headers) - if errSet != nil { - return nil, errSet - } - } - } - - if len(body) > 0 && json.Valid(body) { - errorNode := gjson.GetBytes(body, "error") - if errorNode.Exists() { - payload, errSet = sjson.SetRawBytes(payload, "error", []byte(errorNode.Raw)) - } else { - payload, errSet = sjson.SetRawBytes(payload, "error", body) - } - if errSet != nil { - return nil, errSet - } - } - - if !gjson.GetBytes(payload, "error").Exists() { - payload, errSet = sjson.SetBytes(payload, "error.type", "server_error") - if errSet != nil { - return nil, errSet - } - payload, errSet = sjson.SetBytes(payload, "error.message", errText) - if errSet != nil { - return nil, errSet - } - } - - return payload, writeResponsesWebsocketPayload(writer, wsTimelineLog, payload, time.Now()) -} - -func appendWebsocketEvent(builder *strings.Builder, eventType string, payload []byte) { - if builder == nil { - return - } - trimmedPayload := bytes.TrimSpace(payload) - if len(trimmedPayload) == 0 { - return - } - if builder.Len() > 0 { - builder.WriteString("\n") - } - builder.WriteString("websocket.") - builder.WriteString(eventType) - builder.WriteString("\n") - builder.Write(trimmedPayload) - builder.WriteString("\n") -} - -func websocketPayloadEventType(payload []byte) string { - eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) - if eventType == "" { - return "-" - } - return eventType -} - -func websocketPayloadPreview(payload []byte) string { - trimmedPayload := bytes.TrimSpace(payload) - if len(trimmedPayload) == 0 { - return "" - } - previewText := strings.ReplaceAll(string(trimmedPayload), "\n", "\\n") - previewText = strings.ReplaceAll(previewText, "\r", "\\r") - return previewText -} - -func isResponsesWebsocketCompletionEvent(eventType string) bool { - return eventType == wsEventTypeCompleted || eventType == wsEventTypeDone -} - -func responsesWebsocketErrorMessageFromPayload(payload []byte) *interfaces.ErrorMessage { - status := int(gjson.GetBytes(payload, "status").Int()) - if status <= 0 { - status = int(gjson.GetBytes(payload, "status_code").Int()) - } - if status <= 0 { - status = http.StatusInternalServerError - } - - errText := strings.TrimSpace(gjson.GetBytes(payload, "error.message").String()) - if errText == "" { - errText = strings.TrimSpace(gjson.GetBytes(payload, "message").String()) - } - if errText == "" { - errText = strings.TrimSpace(string(payload)) - } - if errText == "" { - errText = http.StatusText(status) - } - return &interfaces.ErrorMessage{StatusCode: status, Error: fmt.Errorf("%s", errText)} -} - -func setWebsocketTimelineBody(c *gin.Context, body string) { - setWebsocketBody(c, wsTimelineBodyKey, body) -} - -func setWebsocketBody(c *gin.Context, key string, body string) { - if c == nil { - return - } - trimmedBody := strings.TrimSpace(body) - if trimmedBody == "" { - return - } - c.Set(key, []byte(trimmedBody)) -} - -func writeResponsesWebsocketPayload(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, payload []byte, timestamp time.Time) error { - if wsTimelineLog != nil { - wsTimelineLog.Append("response", payload, timestamp) - } - if writer == nil || writer.conn == nil { - return fmt.Errorf("responses websocket: writer is nil") - } - writer.writeMu.Lock() - defer writer.writeMu.Unlock() - if writer.closing.Load() { - return websocket.ErrCloseSent - } - return writer.conn.WriteMessage(websocket.TextMessage, payload) -} - -func appendWebsocketTimelineDisconnect(timeline websocketTimelineAppender, err error, timestamp time.Time) { - if err == nil { - return - } - if timeline != nil { - timeline.Append("disconnect", []byte(err.Error()), timestamp) - } -} - -func appendWebsocketTimelineEvent(builder *strings.Builder, eventType string, payload []byte, timestamp time.Time) { - if builder == nil { - return - } - writeWebsocketTimelineBuilder(builder, formatWebsocketTimelineEvent(eventType, payload, timestamp)) -} - -func formatWebsocketTimelineEvent(eventType string, payload []byte, timestamp time.Time) []byte { - trimmedPayload := bytes.TrimSpace(payload) - if len(trimmedPayload) == 0 { - return nil - } - var builder strings.Builder - builder.WriteString("Timestamp: ") - builder.WriteString(timestamp.Format(time.RFC3339Nano)) - builder.WriteString("\n") - builder.WriteString("Event: websocket.") - builder.WriteString(eventType) - builder.WriteString("\n") - builder.Write(trimmedPayload) - builder.WriteString("\n") - return []byte(builder.String()) -} - -func markAPIResponseTimestamp(c *gin.Context) { - if c == nil { - return - } - if _, exists := c.Get("API_RESPONSE_TIMESTAMP"); exists { - return - } - c.Set("API_RESPONSE_TIMESTAMP", time.Now()) -} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_forward.go b/sdk/api/handlers/openai/openai_responses_websocket_forward.go new file mode 100644 index 000000000..b465201bc --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_forward.go @@ -0,0 +1,552 @@ +package openai + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "sort" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +type responsesWebsocketForwardOptions struct { + toolCacheTurn *responsesWebsocketToolCacheTurn + suppressError func(*interfaces.ErrorMessage) bool +} + +func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( + c *gin.Context, + writer *responsesWebsocketWriter, + cancel handlers.APIHandlerCancelFunc, + data <-chan []byte, + errs <-chan *interfaces.ErrorMessage, + wsTimelineLog websocketTimelineAppender, + sessionID string, + options ...responsesWebsocketForwardOptions, +) ([]byte, string, []string, *interfaces.ErrorMessage, error) { + var opts responsesWebsocketForwardOptions + if len(options) > 0 { + opts = options[0] + } + toolCacheTurn := opts.toolCacheTurn + completed := false + completedOutput := []byte("[]") + completedResponseID := "" + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + pendingToolCallIDs := make(map[string]struct{}) + downstreamSessionKey := "" + if c != nil && c.Request != nil { + downstreamSessionKey = websocketDownstreamSessionKey(c.Request) + } + + for { + select { + case <-c.Request.Context().Done(): + cancel(c.Request.Context().Err()) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, c.Request.Context().Err() + case errMsg, ok := <-errs: + if !ok { + errs = nil + continue + } + if errMsg != nil { + h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) + if opts.suppressError != nil && opts.suppressError(errMsg) { + cancel(errMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil + } + markAPIResponseTimestamp(c) + if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched { + cancel(errMsg.Error) + if errClose != nil { + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errClose + } + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, websocket.ErrCloseSent + } + errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg) + log.Infof( + "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", + sessionID, + websocket.TextMessage, + websocketPayloadEventType(errorPayload), + websocketPayloadPreview(errorPayload), + ) + if errWrite != nil { + // log.Warnf( + // "responses websocket: downstream_out write failed id=%s event=%s error=%v", + // sessionID, + // websocketPayloadEventType(errorPayload), + // errWrite, + // ) + cancel(errMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errWrite + } + } + if errMsg != nil { + cancel(errMsg.Error) + } else { + cancel(nil) + } + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil + case chunk, ok := <-data: + if !ok { + if !completed { + errMsg := &interfaces.ErrorMessage{ + StatusCode: http.StatusRequestTimeout, + Error: fmt.Errorf("stream closed before response.completed"), + } + h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) + markAPIResponseTimestamp(c) + errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg) + log.Infof( + "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", + sessionID, + websocket.TextMessage, + websocketPayloadEventType(errorPayload), + websocketPayloadPreview(errorPayload), + ) + if errWrite != nil { + log.Warnf( + "responses websocket: downstream_out write failed id=%s event=%s error=%v", + sessionID, + websocketPayloadEventType(errorPayload), + errWrite, + ) + cancel(errMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errWrite + } + cancel(errMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil + } + cancel(nil) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, nil + } + + payloads := websocketJSONPayloadsFromChunk(chunk) + for i := range payloads { + collectResponsesWebsocketOutputItem(payloads[i], outputItemsByIndex, &outputItemsFallback) + eventType := gjson.GetBytes(payloads[i], "type").String() + if isResponsesWebsocketCompletionEvent(eventType) { + payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback) + } + if toolCacheTurn != nil { + toolCacheTurn.recordResponse(payloads[i]) + } else { + recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i]) + } + recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i]) + var payloadErrMsg *interfaces.ErrorMessage + if eventType == wsEventTypeError { + payloadErrMsg = responsesWebsocketErrorMessageFromPayload(payloads[i]) + if h != nil { + h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg) + } + if opts.suppressError != nil && opts.suppressError(payloadErrMsg) { + cancel(payloadErrMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil + } + } else if isResponsesWebsocketCompletionEvent(eventType) { + completed = true + completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback) + completedResponseID = responseCompletedIDFromPayload(payloads[i]) + } + markAPIResponseTimestamp(c) + // log.Infof( + // "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", + // sessionID, + // websocket.TextMessage, + // websocketPayloadEventType(payloads[i]), + // websocketPayloadPreview(payloads[i]), + // ) + if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { + log.Warnf( + "responses websocket: downstream_out write failed id=%s event=%s error=%v", + sessionID, + websocketPayloadEventType(payloads[i]), + errWrite, + ) + cancel(errWrite) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), nil, errWrite + } + if payloadErrMsg != nil { + cancel(payloadErrMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil + } + } + } + } +} + +func responsesWebsocketErrorStatus(errMsg *interfaces.ErrorMessage) int { + if errMsg == nil { + return 0 + } + status := errMsg.StatusCode + if status <= 0 && errMsg.Error != nil { + if se, ok := errMsg.Error.(interface{ StatusCode() int }); ok && se != nil { + status = se.StatusCode() + } + } + return status +} + +func shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg *interfaces.ErrorMessage) bool { + switch responsesWebsocketErrorStatus(errMsg) { + case http.StatusUnauthorized, http.StatusTooManyRequests: + return true + default: + return false + } +} + +func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool { + if errMsg == nil { + return false + } + switch responsesWebsocketErrorStatus(errMsg) { + case http.StatusUnauthorized, + http.StatusPaymentRequired, + http.StatusForbidden, + http.StatusTooManyRequests, + http.StatusRequestTimeout, + http.StatusBadGateway, + http.StatusServiceUnavailable, + http.StatusGatewayTimeout: + return true + default: + } + if errMsg.Error != nil { + msg := strings.ToLower(errMsg.Error.Error()) + switch { + case strings.Contains(msg, "stream closed before response.completed"), + strings.Contains(msg, "previous_response_not_found"), + strings.Contains(msg, "ws_failed"), + strings.Contains(msg, "upstream stream closed before first payload"), + strings.Contains(msg, "empty_stream"): + return true + } + } + return false +} + +func collectResponsesWebsocketOutputItem(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback *[][]byte) { + if gjson.GetBytes(payload, "type").String() != "response.output_item.done" { + return + } + item := gjson.GetBytes(payload, "item") + if !item.Exists() || !item.IsObject() { + return + } + outputIndex := gjson.GetBytes(payload, "output_index") + if outputIndex.Exists() { + outputItemsByIndex[outputIndex.Int()] = bytes.Clone([]byte(item.Raw)) + return + } + *outputItemsFallback = append(*outputItemsFallback, bytes.Clone([]byte(item.Raw))) +} + +func restoreResponsesWebsocketCompletionOutput(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { + output := gjson.GetBytes(payload, "response.output") + if output.Exists() && output.IsArray() && len(output.Array()) > 0 { + reconciledOutput, changed := reconcileResponsesWebsocketCompletionToolCalls(output, outputItemsByIndex, outputItemsFallback) + if !changed { + return payload + } + restored, errSet := sjson.SetRawBytes(payload, "response.output", reconciledOutput) + if errSet != nil { + return payload + } + return restored + } + if len(outputItemsByIndex) == 0 && len(outputItemsFallback) == 0 { + return payload + } + + restored, errSet := sjson.SetRawBytes(payload, "response.output", responseCompletedOutputFromPayload(payload, outputItemsByIndex, outputItemsFallback)) + if errSet != nil { + return payload + } + return restored +} + +func reconcileResponsesWebsocketCompletionToolCalls(output gjson.Result, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) ([]byte, bool) { + collectedToolCalls := make(map[string]json.RawMessage) + recordCollectedToolCall := func(raw []byte) { + item := gjson.ParseBytes(raw) + if !isCompleteResponsesWebsocketToolCall(item) { + return + } + callID := strings.TrimSpace(item.Get("call_id").String()) + collectedToolCalls[callID] = append(json.RawMessage(nil), raw...) + } + + indexes := make([]int64, 0, len(outputItemsByIndex)) + for index := range outputItemsByIndex { + indexes = append(indexes, index) + } + sort.Slice(indexes, func(i, j int) bool { + return indexes[i] < indexes[j] + }) + for _, index := range indexes { + recordCollectedToolCall(outputItemsByIndex[index]) + } + for _, item := range outputItemsFallback { + recordCollectedToolCall(item) + } + if len(collectedToolCalls) == 0 { + return nil, false + } + + items := output.Array() + reconciled := make([]json.RawMessage, 0, len(items)) + changed := false + for _, item := range items { + raw := json.RawMessage(item.Raw) + if isResponsesToolCallType(item.Get("type").String()) { + callID := strings.TrimSpace(item.Get("call_id").String()) + if collected, ok := collectedToolCalls[callID]; ok && !bytes.Equal(raw, collected) { + raw = collected + changed = true + } + } + reconciled = append(reconciled, raw) + } + if !changed { + return nil, false + } + + marshaledOutput, errMarshal := json.Marshal(reconciled) + if errMarshal != nil { + return nil, false + } + return marshaledOutput, true +} + +func isCompleteResponsesWebsocketToolCall(item gjson.Result) bool { + if !item.Exists() || !item.IsObject() { + return false + } + callID := item.Get("call_id") + name := item.Get("name") + if callID.Type != gjson.String || strings.TrimSpace(callID.String()) == "" || name.Type != gjson.String || strings.TrimSpace(name.String()) == "" { + return false + } + + switch strings.TrimSpace(item.Get("type").String()) { + case "function_call": + arguments := item.Get("arguments") + return arguments.Exists() && arguments.Type == gjson.String + case "custom_tool_call": + input := item.Get("input") + return input.Exists() && input.Type == gjson.String + default: + return false + } +} + +func responseCompletedOutputFromPayload(payload []byte, outputItemsByIndex map[int64][]byte, outputItemsFallback [][]byte) []byte { + output := gjson.GetBytes(payload, "response.output") + if output.Exists() && output.IsArray() && len(output.Array()) > 0 { + return bytes.Clone([]byte(output.Raw)) + } + if len(outputItemsByIndex) == 0 && len(outputItemsFallback) == 0 { + return []byte("[]") + } + + indexes := make([]int64, 0, len(outputItemsByIndex)) + for index := range outputItemsByIndex { + indexes = append(indexes, index) + } + sort.Slice(indexes, func(i, j int) bool { + return indexes[i] < indexes[j] + }) + + items := make([]json.RawMessage, 0, len(outputItemsByIndex)+len(outputItemsFallback)) + appendCollectedItem := func(raw []byte) { + item := gjson.ParseBytes(raw) + if isResponsesToolCallType(item.Get("type").String()) && !isCompleteResponsesWebsocketToolCall(item) { + return + } + items = append(items, append(json.RawMessage(nil), raw...)) + } + for _, index := range indexes { + appendCollectedItem(outputItemsByIndex[index]) + } + for _, item := range outputItemsFallback { + appendCollectedItem(item) + } + + marshaledOutput, errMarshal := json.Marshal(items) + if errMarshal != nil { + return []byte("[]") + } + return marshaledOutput +} + +func responseCompletedIDFromPayload(payload []byte) string { + return strings.TrimSpace(gjson.GetBytes(payload, "response.id").String()) +} + +func recordPendingToolCallIDsFromPayload(pending map[string]struct{}, payload []byte) { + if pending == nil || len(payload) == 0 { + return + } + updatePendingToolCallIDsFromItem(pending, gjson.GetBytes(payload, "item")) + output := gjson.GetBytes(payload, "response.output") + if output.IsArray() { + for _, item := range output.Array() { + updatePendingToolCallIDsFromItem(pending, item) + } + } +} + +func updatePendingToolCallIDsFromItem(pending map[string]struct{}, item gjson.Result) { + if pending == nil || !item.Exists() { + return + } + switch strings.TrimSpace(item.Get("type").String()) { + case "function_call", "custom_tool_call": + if !isCompleteResponsesWebsocketToolCall(item) { + return + } + callID := strings.TrimSpace(item.Get("call_id").String()) + pending[callID] = struct{}{} + case "function_call_output", "custom_tool_call_output": + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID != "" { + delete(pending, callID) + } + } +} + +func sortedStringSet(values map[string]struct{}) []string { + if len(values) == 0 { + return nil + } + out := make([]string, 0, len(values)) + for value := range values { + value = strings.TrimSpace(value) + if value != "" { + out = append(out, value) + } + } + sort.Strings(out) + return out +} + +func websocketJSONPayloadsFromChunk(chunk []byte) [][]byte { + payloads := make([][]byte, 0, 2) + lines := bytes.Split(chunk, []byte("\n")) + for i := range lines { + line := bytes.TrimSpace(lines[i]) + if len(line) == 0 || bytes.HasPrefix(line, []byte("event:")) { + continue + } + if bytes.HasPrefix(line, []byte("data:")) { + line = bytes.TrimSpace(line[len("data:"):]) + } + if len(line) == 0 || bytes.Equal(line, []byte(wsDoneMarker)) { + continue + } + if json.Valid(line) { + payloads = append(payloads, bytes.Clone(line)) + } + } + + if len(payloads) > 0 { + return payloads + } + + trimmed := bytes.TrimSpace(chunk) + if bytes.HasPrefix(trimmed, []byte("data:")) { + trimmed = bytes.TrimSpace(trimmed[len("data:"):]) + } + if len(trimmed) > 0 && !bytes.Equal(trimmed, []byte(wsDoneMarker)) && json.Valid(trimmed) { + payloads = append(payloads, bytes.Clone(trimmed)) + } + return payloads +} + +func writeResponsesWebsocketError(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, errMsg *interfaces.ErrorMessage) ([]byte, error) { + status := http.StatusInternalServerError + errText := http.StatusText(status) + if errMsg != nil { + if errMsg.StatusCode > 0 { + status = errMsg.StatusCode + errText = http.StatusText(status) + } + if errMsg.Error != nil && strings.TrimSpace(errMsg.Error.Error()) != "" { + errText = errMsg.Error.Error() + } + } + + body := handlers.BuildErrorResponseBody(status, errText) + payload := []byte(`{}`) + var errSet error + payload, errSet = sjson.SetBytes(payload, "type", wsEventTypeError) + if errSet != nil { + return nil, errSet + } + payload, errSet = sjson.SetBytes(payload, "status", status) + if errSet != nil { + return nil, errSet + } + + if errMsg != nil && errMsg.Addon != nil { + headers := []byte(`{}`) + hasHeaders := false + for key, values := range errMsg.Addon { + if len(values) == 0 { + continue + } + headerPath := strings.ReplaceAll(strings.ReplaceAll(key, `\\`, `\\\\`), ".", `\\.`) + headers, errSet = sjson.SetBytes(headers, headerPath, values[0]) + if errSet != nil { + return nil, errSet + } + hasHeaders = true + } + if hasHeaders { + payload, errSet = sjson.SetRawBytes(payload, "headers", headers) + if errSet != nil { + return nil, errSet + } + } + } + + if len(body) > 0 && json.Valid(body) { + errorNode := gjson.GetBytes(body, "error") + if errorNode.Exists() { + payload, errSet = sjson.SetRawBytes(payload, "error", []byte(errorNode.Raw)) + } else { + payload, errSet = sjson.SetRawBytes(payload, "error", body) + } + if errSet != nil { + return nil, errSet + } + } + + if !gjson.GetBytes(payload, "error").Exists() { + payload, errSet = sjson.SetBytes(payload, "error.type", "server_error") + if errSet != nil { + return nil, errSet + } + payload, errSet = sjson.SetBytes(payload, "error.message", errText) + if errSet != nil { + return nil, errSet + } + } + + return payload, writeResponsesWebsocketPayload(writer, wsTimelineLog, payload, time.Now()) +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_prewarm.go b/sdk/api/handlers/openai/openai_responses_websocket_prewarm.go new file mode 100644 index 000000000..6d2ab7d1b --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_prewarm.go @@ -0,0 +1,173 @@ +package openai + +import ( + "encoding/json" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func shouldHandleResponsesWebsocketPrewarmLocally(rawJSON []byte, lastRequest []byte, allowIncrementalInputWithPreviousResponseID bool) bool { + if allowIncrementalInputWithPreviousResponseID || len(lastRequest) != 0 { + return false + } + if strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) != wsRequestTypeCreate { + return false + } + generateResult := gjson.GetBytes(rawJSON, "generate") + return generateResult.Exists() && !generateResult.Bool() +} + +func writeResponsesWebsocketSyntheticPrewarm( + c *gin.Context, + writer *responsesWebsocketWriter, + requestJSON []byte, + wsTimelineLog websocketTimelineAppender, + sessionID string, +) error { + payloads, errPayloads := syntheticResponsesWebsocketPrewarmPayloads(requestJSON) + if errPayloads != nil { + return errPayloads + } + for i := 0; i < len(payloads); i++ { + markAPIResponseTimestamp(c) + // log.Infof( + // "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", + // sessionID, + // websocket.TextMessage, + // websocketPayloadEventType(payloads[i]), + // websocketPayloadPreview(payloads[i]), + // ) + if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { + log.Warnf( + "responses websocket: downstream_out write failed id=%s event=%s error=%v", + sessionID, + websocketPayloadEventType(payloads[i]), + errWrite, + ) + return errWrite + } + } + return nil +} + +func syntheticResponsesWebsocketPrewarmPayloads(requestJSON []byte) ([][]byte, error) { + responseID := "resp_prewarm_" + uuid.NewString() + createdAt := time.Now().Unix() + modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()) + + createdPayload := []byte(`{"type":"response.created","sequence_number":0,"response":{"id":"","object":"response","created_at":0,"status":"in_progress","background":false,"error":null,"output":[]}}`) + var errSet error + createdPayload, errSet = sjson.SetBytes(createdPayload, "response.id", responseID) + if errSet != nil { + return nil, errSet + } + createdPayload, errSet = sjson.SetBytes(createdPayload, "response.created_at", createdAt) + if errSet != nil { + return nil, errSet + } + if modelName != "" { + createdPayload, errSet = sjson.SetBytes(createdPayload, "response.model", modelName) + if errSet != nil { + return nil, errSet + } + } + + completedPayload := []byte(`{"type":"response.completed","sequence_number":1,"response":{"id":"","object":"response","created_at":0,"status":"completed","background":false,"error":null,"output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + completedPayload, errSet = sjson.SetBytes(completedPayload, "response.id", responseID) + if errSet != nil { + return nil, errSet + } + completedPayload, errSet = sjson.SetBytes(completedPayload, "response.created_at", createdAt) + if errSet != nil { + return nil, errSet + } + if modelName != "" { + completedPayload, errSet = sjson.SetBytes(completedPayload, "response.model", modelName) + if errSet != nil { + return nil, errSet + } + } + + return [][]byte{createdPayload, completedPayload}, nil +} + +func mergeJSONArrayRaw(existingRaw, appendRaw string) (string, error) { + existingRaw = strings.TrimSpace(existingRaw) + appendRaw = strings.TrimSpace(appendRaw) + if existingRaw == "" { + existingRaw = "[]" + } + if appendRaw == "" { + appendRaw = "[]" + } + + var existing []json.RawMessage + if err := json.Unmarshal([]byte(existingRaw), &existing); err != nil { + return "", err + } + var appendItems []json.RawMessage + if err := json.Unmarshal([]byte(appendRaw), &appendItems); err != nil { + return "", err + } + + merged := append(existing, appendItems...) + out, err := json.Marshal(merged) + if err != nil { + return "", err + } + return string(out), nil +} + +// inputContainsFullTranscript returns true when the input array carries compact +// replay markers that indicate the client already sent the full conversation +// transcript. Merging that input with stale lastRequest/lastResponseOutput +// would duplicate or break function_call/function_call_output pairings, so the +// caller should use the input as-is. +// +// Assistant messages alone are not enough to classify the payload as a replay: +// incremental websocket requests may legitimately append assistant items. +func inputContainsFullTranscript(input gjson.Result) bool { + if !input.IsArray() { + return false + } + for _, item := range input.Array() { + t := item.Get("type").String() + if t == "compaction" || t == "compaction_summary" { + return true + } + } + return false +} + +func inputWithoutCompactionItems(input gjson.Result) string { + if !input.IsArray() { + return normalizeJSONArrayRaw([]byte(input.Raw)) + } + filtered := make([]string, 0, len(input.Array())) + for _, item := range input.Array() { + t := item.Get("type").String() + if t == "compaction" || t == "compaction_summary" { + continue + } + filtered = append(filtered, item.Raw) + } + return "[" + strings.Join(filtered, ",") + "]" +} + +func normalizeJSONArrayRaw(raw []byte) string { + trimmed := strings.TrimSpace(string(raw)) + if trimmed == "" { + return "[]" + } + result := gjson.Parse(trimmed) + if result.Type == gjson.JSON && result.IsArray() { + return trimmed + } + return "[]" +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_requests.go b/sdk/api/handlers/openai/openai_responses_websocket_requests.go new file mode 100644 index 000000000..606b075e3 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_requests.go @@ -0,0 +1,506 @@ +package openai + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) { + return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true) +} + +func normalizeResponsesWebsocketRequestWithMode(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { + return normalizeResponsesWebsocketRequestWithLastResponseID(rawJSON, lastRequest, lastResponseOutput, "", allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) +} + +func normalizeResponsesWebsocketRequestWithLastResponseID(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { + return normalizeResponsesWebsocketRequestWithIncrementalState(rawJSON, lastRequest, lastResponseOutput, lastResponseID, nil, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) +} + +func normalizeResponsesWebsocketRequestWithIncrementalState(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, lastResponsePendingToolCallIDs []string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { + requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) + switch requestType { + case wsRequestTypeCreate: + // log.Infof("responses websocket: response.create request") + if len(lastRequest) == 0 { + return normalizeResponseCreateRequest(rawJSON) + } + return normalizeResponseSubsequentRequest(rawJSON, lastRequest, lastResponseOutput, lastResponseID, lastResponsePendingToolCallIDs, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) + case wsRequestTypeAppend: + // log.Infof("responses websocket: response.append request") + return normalizeResponseSubsequentRequest(rawJSON, lastRequest, lastResponseOutput, lastResponseID, lastResponsePendingToolCallIDs, allowIncrementalInputWithPreviousResponseID, allowCompactionReplayBypass) + default: + return nil, lastRequest, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("unsupported websocket request type: %s", requestType), + } + } +} + +func normalizeResponseCreateRequest(rawJSON []byte) ([]byte, []byte, *interfaces.ErrorMessage) { + normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") + if errDelete != nil { + normalized = bytes.Clone(rawJSON) + } + normalized, _ = sjson.SetBytes(normalized, "stream", true) + if !gjson.GetBytes(normalized, "input").Exists() { + normalized, _ = sjson.SetRawBytes(normalized, "input", []byte("[]")) + } + + modelName := strings.TrimSpace(gjson.GetBytes(normalized, "model").String()) + if modelName == "" { + return nil, nil, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("missing model in response.create request"), + } + } + return normalized, bytes.Clone(normalized), nil +} + +func normalizeResponseSubsequentRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte, lastResponseID string, lastResponsePendingToolCallIDs []string, allowIncrementalInputWithPreviousResponseID bool, allowCompactionReplayBypass bool) ([]byte, []byte, *interfaces.ErrorMessage) { + if len(lastRequest) == 0 { + return nil, lastRequest, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("websocket request received before response.create"), + } + } + + nextInput := gjson.GetBytes(rawJSON, "input") + if !nextInput.Exists() || !nextInput.IsArray() { + return nil, lastRequest, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("websocket request requires array field: input"), + } + } + + // Compaction can cause clients to replace local websocket history with a new + // compact transcript on the next `response.create`. When the input already + // contains historical model output items, treating it as an incremental append + // duplicates stale turn-state and can leave late orphaned function_call items. + if shouldReplaceWebsocketTranscript(rawJSON, nextInput) { + normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest) + return normalized, bytes.Clone(normalized), nil + } + + // Websocket v2 mode uses response.create with previous_response_id + incremental input. + // Do not expand it into a full input transcript; upstream expects the incremental payload. + if allowIncrementalInputWithPreviousResponseID { + prev := strings.TrimSpace(gjson.GetBytes(rawJSON, "previous_response_id").String()) + if prev == "" { + if !inputSatisfiesPendingToolCalls(nextInput, lastResponsePendingToolCallIDs) { + normalized := normalizeResponseTranscriptReplacement(rawJSON, lastRequest) + return normalized, bytes.Clone(normalized), nil + } + prev = strings.TrimSpace(lastResponseID) + } + if prev != "" { + normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") + if errDelete != nil { + normalized = bytes.Clone(rawJSON) + } + normalized, _ = sjson.SetBytes(normalized, "previous_response_id", prev) + if !gjson.GetBytes(normalized, "model").Exists() { + modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) + if modelName != "" { + normalized, _ = sjson.SetBytes(normalized, "model", modelName) + } + } + if !gjson.GetBytes(normalized, "instructions").Exists() { + instructions := gjson.GetBytes(lastRequest, "instructions") + if instructions.Exists() { + normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw)) + } + } + normalized, _ = sjson.SetBytes(normalized, "stream", true) + return normalized, bytes.Clone(normalized), nil + } + } + + // When the client sends a compact replay for a downstream that can consume it + // directly, the input already carries the canonical history. In that case, + // skip merging with stale lastRequest/lastResponseOutput to avoid breaking + // function_call / function_call_output pairings. + // See: https://github.com/router-for-me/CLIProxyAPI/issues/2207 + var mergedInput string + if allowCompactionReplayBypass && inputContainsFullTranscript(nextInput) { + log.Infof("responses websocket: full transcript detected, skipping stale merge (input items=%d)", len(nextInput.Array())) + mergedInput = nextInput.Raw + } else { + appendInputRaw := nextInput.Raw + if inputContainsFullTranscript(nextInput) { + appendInputRaw = inputWithoutCompactionItems(nextInput) + } + + existingInput := gjson.GetBytes(lastRequest, "input") + var errMerge error + mergedInput, errMerge = mergeJSONArrayRaw(existingInput.Raw, normalizeJSONArrayRaw(lastResponseOutput)) + if errMerge != nil { + return nil, lastRequest, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("invalid previous response output: %w", errMerge), + } + } + + mergedInput, errMerge = mergeJSONArrayRaw(mergedInput, appendInputRaw) + if errMerge != nil { + return nil, lastRequest, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("invalid request input: %w", errMerge), + } + } + } + dedupedInput, errDedupeFunctionCalls := dedupeFunctionCallsByCallID(mergedInput) + if errDedupeFunctionCalls == nil { + mergedInput = dedupedInput + } + dedupedInput, errDedupeItemIDs := dedupeInputItemsByID(mergedInput) + if errDedupeItemIDs == nil { + mergedInput = dedupedInput + } + + normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") + if errDelete != nil { + normalized = bytes.Clone(rawJSON) + } + normalized, _ = sjson.DeleteBytes(normalized, "previous_response_id") + var errSet error + normalized, errSet = sjson.SetRawBytes(normalized, "input", []byte(mergedInput)) + if errSet != nil { + return nil, lastRequest, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("failed to merge websocket input: %w", errSet), + } + } + if !gjson.GetBytes(normalized, "model").Exists() { + modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) + if modelName != "" { + normalized, _ = sjson.SetBytes(normalized, "model", modelName) + } + } + if !gjson.GetBytes(normalized, "instructions").Exists() { + instructions := gjson.GetBytes(lastRequest, "instructions") + if instructions.Exists() { + normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw)) + } + } + normalized, _ = sjson.SetBytes(normalized, "stream", true) + return normalized, bytes.Clone(normalized), nil +} + +func shouldReplaceWebsocketTranscript(rawJSON []byte, nextInput gjson.Result) bool { + requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) + if requestType != wsRequestTypeCreate && requestType != wsRequestTypeAppend { + return false + } + previousResponseID := gjson.GetBytes(rawJSON, "previous_response_id") + if strings.TrimSpace(previousResponseID.String()) != "" { + return false + } + if !nextInput.Exists() || !nextInput.IsArray() { + return false + } + if requestType == wsRequestTypeCreate && !previousResponseID.Exists() && inputHasCodexLocalCompactionSummary(nextInput) { + return true + } + + for _, item := range nextInput.Array() { + switch strings.TrimSpace(item.Get("type").String()) { + case "function_call", "custom_tool_call": + return true + case "message": + if strings.TrimSpace(item.Get("role").String()) == "assistant" { + return true + } + } + } + + return false +} + +func inputHasCodexLocalCompactionSummary(input gjson.Result) bool { + if !input.IsArray() { + return false + } + + hasSummary := false + for index, item := range input.Array() { + itemType := strings.TrimSpace(item.Get("type").String()) + if itemType == "additional_tools" { + tools := item.Get("tools") + if index != 0 || strings.TrimSpace(item.Get("role").String()) != "developer" || !tools.IsArray() { + return false + } + for _, tool := range tools.Array() { + if !tool.IsObject() || strings.TrimSpace(tool.Get("type").String()) == "" { + return false + } + } + continue + } + if itemType != "" && itemType != "message" { + return false + } + + role := strings.TrimSpace(item.Get("role").String()) + if role != "user" && role != "developer" { + return false + } + if role == "user" && strings.HasPrefix(codexLocalCompactionMessageText(item), codexLocalCompactionSummaryPrefix+"\n") { + hasSummary = true + } + } + return hasSummary +} + +func codexLocalCompactionMessageText(message gjson.Result) string { + content := message.Get("content") + if content.Type == gjson.String { + return content.String() + } + if !content.IsArray() { + return "" + } + + var text strings.Builder + for _, part := range content.Array() { + if strings.TrimSpace(part.Get("type").String()) == "input_text" { + text.WriteString(part.Get("text").String()) + } + } + return text.String() +} + +func inputSatisfiesPendingToolCalls(input gjson.Result, pendingCallIDs []string) bool { + if len(pendingCallIDs) == 0 { + return true + } + if !input.IsArray() { + return false + } + outputs := make(map[string]struct{}, len(pendingCallIDs)) + for _, item := range input.Array() { + switch strings.TrimSpace(item.Get("type").String()) { + case "function_call_output", "custom_tool_call_output": + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID != "" { + outputs[callID] = struct{}{} + } + } + } + for _, callID := range pendingCallIDs { + callID = strings.TrimSpace(callID) + if callID == "" { + continue + } + if _, ok := outputs[callID]; !ok { + return false + } + } + return true +} + +func normalizeResponseTranscriptReplacement(rawJSON []byte, lastRequest []byte) []byte { + normalized, errDelete := sjson.DeleteBytes(rawJSON, "type") + if errDelete != nil { + normalized = bytes.Clone(rawJSON) + } + normalized, _ = sjson.DeleteBytes(normalized, "previous_response_id") + if !gjson.GetBytes(normalized, "model").Exists() { + modelName := strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) + if modelName != "" { + normalized, _ = sjson.SetBytes(normalized, "model", modelName) + } + } + if !gjson.GetBytes(normalized, "instructions").Exists() { + instructions := gjson.GetBytes(lastRequest, "instructions") + if instructions.Exists() { + normalized, _ = sjson.SetRawBytes(normalized, "instructions", []byte(instructions.Raw)) + } + } + normalized, _ = sjson.SetBytes(normalized, "stream", true) + return bytes.Clone(normalized) +} + +func dedupeFunctionCallsByCallID(rawArray string) (string, error) { + rawArray = strings.TrimSpace(rawArray) + if rawArray == "" { + return "[]", nil + } + var items []json.RawMessage + if errUnmarshal := json.Unmarshal([]byte(rawArray), &items); errUnmarshal != nil { + return "", errUnmarshal + } + + seenCallIDs := make(map[string]struct{}, len(items)) + filtered := make([]json.RawMessage, 0, len(items)) + for _, item := range items { + if len(item) == 0 { + continue + } + itemType := strings.TrimSpace(gjson.GetBytes(item, "type").String()) + if isResponsesToolCallType(itemType) { + callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String()) + if callID != "" { + if _, ok := seenCallIDs[callID]; ok { + continue + } + seenCallIDs[callID] = struct{}{} + } + } + filtered = append(filtered, item) + } + + out, errMarshal := json.Marshal(filtered) + if errMarshal != nil { + return "", errMarshal + } + return string(out), nil +} + +func dedupeResponsesWebsocketInputItemsByID(payload []byte) []byte { + input := gjson.GetBytes(payload, "input") + if !input.Exists() || !input.IsArray() { + return payload + } + dedupedInput, errDedupe := dedupeInputItemsByID(input.Raw) + if errDedupe != nil || dedupedInput == input.Raw { + return payload + } + updated, errSet := sjson.SetRawBytes(payload, "input", []byte(dedupedInput)) + if errSet != nil { + return payload + } + return updated +} + +func dedupeInputItemsByID(rawArray string) (string, error) { + rawArray = strings.TrimSpace(rawArray) + if rawArray == "" { + return "[]", nil + } + var items []json.RawMessage + if errUnmarshal := json.Unmarshal([]byte(rawArray), &items); errUnmarshal != nil { + return "", errUnmarshal + } + + // Parse each item's type, id and call_id once; gjson is a scan-based + // parser, so reusing this metadata avoids rescanning every item in each of + // the loops below as the conversation history grows. + type itemMetadata struct { + itemType string + id string + callID string + } + meta := make([]itemMetadata, len(items)) + for i, item := range items { + if len(item) == 0 { + continue + } + res := gjson.GetManyBytes(item, "type", "id", "call_id") + meta[i] = itemMetadata{ + itemType: strings.TrimSpace(res[0].String()), + id: strings.TrimSpace(res[1].String()), + callID: strings.TrimSpace(res[2].String()), + } + } + + // Collect the call_ids that are still referenced by tool-call output + // items. When several input items share the same id, the one we keep must + // preserve any call_id that has a matching output; otherwise the upstream + // rejects the request with "No tool call found for function call output". + referencedCallIDs := make(map[string]struct{}, len(items)) + for i := range items { + switch meta[i].itemType { + case "function_call_output", "custom_tool_call_output": + if meta[i].callID != "" { + referencedCallIDs[meta[i].callID] = struct{}{} + } + } + } + + // For each id, choose the index to keep. The default is the last + // occurrence (matching the original dedupe behavior), but we never replace + // an item whose call_id still has a matching output with one that does not. + // This keeps a single item per id while ensuring retained tool calls stay + // paired with their outputs. + keepIndexByID := make(map[string]int, len(items)) + keepReferencedByID := make(map[string]bool, len(items)) + for i := range items { + itemID := meta[i].id + if itemID == "" { + continue + } + _, referenced := referencedCallIDs[meta[i].callID] + referenced = referenced && meta[i].callID != "" + if _, seen := keepIndexByID[itemID]; !seen { + keepIndexByID[itemID] = i + keepReferencedByID[itemID] = referenced + continue + } + if referenced || !keepReferencedByID[itemID] { + keepIndexByID[itemID] = i + keepReferencedByID[itemID] = referenced + } + } + + filtered := make([]json.RawMessage, 0, len(items)) + for i, item := range items { + if len(item) == 0 { + continue + } + itemID := meta[i].id + if itemID != "" { + if keepIndexByID[itemID] != i { + continue + } + } + filtered = append(filtered, item) + } + + out, errMarshal := json.Marshal(filtered) + if errMarshal != nil { + return "", errMarshal + } + return string(out), nil +} + +func normalizeResponsesWebsocketPassthroughRequest(rawJSON []byte, modelName string) ([]byte, *interfaces.ErrorMessage) { + if !json.Valid(rawJSON) { + return nil, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("invalid websocket request JSON"), + } + } + + requestType := strings.TrimSpace(gjson.GetBytes(rawJSON, "type").String()) + switch requestType { + case wsRequestTypeCreate, wsRequestTypeAppend: + default: + return nil, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("unsupported websocket request type: %s", requestType), + } + } + + normalized := bytes.Clone(rawJSON) + if strings.TrimSpace(gjson.GetBytes(normalized, "model").String()) == "" { + modelName = strings.TrimSpace(modelName) + if modelName == "" { + return nil, &interfaces.ErrorMessage{ + StatusCode: http.StatusBadRequest, + Error: fmt.Errorf("missing model in response.create request"), + } + } + normalized, _ = sjson.SetBytes(normalized, "model", modelName) + } + normalized, _ = sjson.SetBytes(normalized, "stream", true) + return normalized, nil +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_session.go b/sdk/api/handlers/openai/openai_responses_websocket_session.go new file mode 100644 index 000000000..5786da357 --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_session.go @@ -0,0 +1,237 @@ +package openai + +import ( + "fmt" + "strconv" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func websocketUpstreamSupportsIncrementalInput(attributes map[string]string, metadata map[string]any) bool { + if len(attributes) > 0 { + if raw := strings.TrimSpace(attributes["websockets"]); raw != "" { + parsed, errParse := strconv.ParseBool(raw) + if errParse == nil { + return parsed + } + } + } + if len(metadata) == 0 { + return false + } + raw, ok := metadata["websockets"] + if !ok || raw == nil { + return false + } + switch value := raw.(type) { + case bool: + return value + case string: + parsed, errParse := strconv.ParseBool(strings.TrimSpace(value)) + if errParse == nil { + return parsed + } + default: + } + return false +} + +func (h *OpenAIResponsesAPIHandler) websocketUpstreamSupportsIncrementalInputForModel(modelName string) bool { + auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName) + for _, auth := range auths { + if responsesWebsocketAuthSupportsIncrementalInput(auth) { + return true + } + } + return false +} + +func (h *OpenAIResponsesAPIHandler) websocketUpstreamSupportsCompactionReplayForModel(modelName string) bool { + auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName) + if len(auths) == 0 { + return false + } + for _, auth := range auths { + if !responsesWebsocketAuthSupportsCompactionReplay(auth) { + return false + } + } + return true +} + +func (h *OpenAIResponsesAPIHandler) responsesWebsocketAvailableAuthsForModel(modelName string) ([]*coreauth.Auth, string) { + if h == nil || h.AuthManager == nil { + return nil, "" + } + resolvedModelName := responsesWebsocketResolvedModelName(modelName) + providerSet, modelKey := responsesWebsocketProviderSetForModel(resolvedModelName) + if len(providerSet) == 0 { + return nil, modelKey + } + + registryRef := registry.GetGlobalRegistry() + now := time.Now() + auths := h.AuthManager.List() + available := make([]*coreauth.Auth, 0, len(auths)) + for _, auth := range auths { + if !responsesWebsocketAuthMatchesModel(auth, providerSet, modelKey, registryRef, now) { + continue + } + available = append(available, auth) + } + return available, modelKey +} + +func (h *OpenAIResponsesAPIHandler) responsesWebsocketUsesCodexWebsocketPassthrough(modelName string) bool { + return h.responsesWebsocketUsesUpstreamWebsocketPassthrough(modelName) +} + +func (h *OpenAIResponsesAPIHandler) responsesWebsocketUsesUpstreamWebsocketPassthrough(modelName string) bool { + modelName = strings.TrimSpace(modelName) + if h == nil || h.AuthManager == nil || modelName == "" { + return false + } + auths, _ := h.responsesWebsocketAvailableAuthsForModel(modelName) + if len(auths) == 0 { + return false + } + provider := "" + for _, auth := range auths { + if auth == nil { + return false + } + authProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if authProvider != "codex" && authProvider != "xai" { + return false + } + if provider == "" { + provider = authProvider + if _, ok := h.AuthManager.Executor(provider); !ok { + return false + } + } else if authProvider != provider { + return false + } + if !websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) { + return false + } + } + return provider != "" +} + +func responsesWebsocketAuthSupportsIncrementalInput(auth *coreauth.Auth) bool { + if auth == nil { + return false + } + return websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) +} + +func responsesWebsocketPinnedAuthMatchesModel(auth *coreauth.Auth, modelName string, pinnedModelKey string, homeRuntime bool) bool { + if auth == nil { + return false + } + providerSet, modelKey := responsesWebsocketProviderSetForModel(responsesWebsocketResolvedModelName(modelName)) + providerKey := strings.ToLower(strings.TrimSpace(auth.Provider)) + if _, ok := providerSet[providerKey]; !ok { + return false + } + if !responsesWebsocketAuthAvailableForModel(auth, modelKey, time.Now()) { + return false + } + + if homeRuntime { + return strings.EqualFold(strings.TrimSpace(pinnedModelKey), strings.TrimSpace(modelKey)) + } + return registry.GetGlobalRegistry().ClientSupportsModel(auth.ID, modelKey) +} + +func responsesWebsocketResolvedModelName(modelName string) string { + initialSuffix := thinking.ParseSuffix(modelName) + if initialSuffix.ModelName == "auto" { + resolvedBase := util.ResolveAutoModel(initialSuffix.ModelName) + if initialSuffix.HasSuffix { + return fmt.Sprintf("%s(%s)", resolvedBase, initialSuffix.RawSuffix) + } + return resolvedBase + } + return util.ResolveAutoModel(modelName) +} + +func responsesWebsocketProviderSetForModel(resolvedModelName string) (map[string]struct{}, string) { + parsed := thinking.ParseSuffix(resolvedModelName) + baseModel := strings.TrimSpace(parsed.ModelName) + providers := util.GetProviderName(baseModel) + if len(providers) == 0 && baseModel != resolvedModelName { + providers = util.GetProviderName(resolvedModelName) + } + providerSet := make(map[string]struct{}, len(providers)) + for _, provider := range providers { + providerKey := strings.TrimSpace(strings.ToLower(provider)) + if providerKey == "" { + continue + } + providerSet[providerKey] = struct{}{} + } + modelKey := baseModel + if modelKey == "" { + modelKey = strings.TrimSpace(resolvedModelName) + } + return providerSet, modelKey +} + +func responsesWebsocketAuthMatchesModel(auth *coreauth.Auth, providerSet map[string]struct{}, modelKey string, registryRef *registry.ModelRegistry, now time.Time) bool { + if auth == nil { + return false + } + providerKey := strings.TrimSpace(strings.ToLower(auth.Provider)) + if _, ok := providerSet[providerKey]; !ok { + return false + } + if modelKey != "" && registryRef != nil && !registryRef.ClientSupportsModel(auth.ID, modelKey) { + return false + } + return responsesWebsocketAuthAvailableForModel(auth, modelKey, now) +} + +func responsesWebsocketAuthSupportsCompactionReplay(auth *coreauth.Auth) bool { + if auth == nil { + return false + } + return strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") +} + +func responsesWebsocketAuthAvailableForModel(auth *coreauth.Auth, modelName string, now time.Time) bool { + if auth == nil { + return false + } + if auth.Disabled || auth.Status == coreauth.StatusDisabled { + return false + } + if modelName != "" && len(auth.ModelStates) > 0 { + state, ok := auth.ModelStates[modelName] + if (!ok || state == nil) && modelName != "" { + baseModel := strings.TrimSpace(thinking.ParseSuffix(modelName).ModelName) + if baseModel != "" && baseModel != modelName { + state, ok = auth.ModelStates[baseModel] + } + } + if ok && state != nil { + if state.Status == coreauth.StatusDisabled { + return false + } + if state.Unavailable && !state.NextRetryAfter.IsZero() && state.NextRetryAfter.After(now) { + return false + } + return true + } + } + if auth.Unavailable && !auth.NextRetryAfter.IsZero() && auth.NextRetryAfter.After(now) { + return false + } + return true +} diff --git a/sdk/api/handlers/openai/openai_responses_websocket_timeline.go b/sdk/api/handlers/openai/openai_responses_websocket_timeline.go new file mode 100644 index 000000000..8126857ef --- /dev/null +++ b/sdk/api/handlers/openai/openai_responses_websocket_timeline.go @@ -0,0 +1,317 @@ +package openai + +import ( + "bytes" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" + requestlogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +type websocketTimelineAppender interface { + Append(eventType string, payload []byte, timestamp time.Time) +} + +type responsesWebsocketPinnedAuthState struct { + authID string + modelKey string +} + +type websocketTimelineLog struct { + enabled bool + source *requestlogging.FileBodySource + builder *strings.Builder + + currentPart io.WriteCloser + currentPartHasLog bool +} + +func newWebsocketTimelineLog(enabled bool, source *requestlogging.FileBodySource) *websocketTimelineLog { + if !enabled { + return &websocketTimelineLog{} + } + if source == nil { + return newInMemoryWebsocketTimelineLog() + } + return &websocketTimelineLog{ + enabled: true, + source: source, + } +} + +func newInMemoryWebsocketTimelineLog() *websocketTimelineLog { + return &websocketTimelineLog{ + enabled: true, + builder: &strings.Builder{}, + } +} + +func websocketTimelineSourceFromContext(c *gin.Context) *requestlogging.FileBodySource { + if c == nil { + return nil + } + value, exists := c.Get(requestlogging.WebsocketTimelineSourceContextKey) + if !exists { + return nil + } + source, ok := value.(*requestlogging.FileBodySource) + if !ok { + return nil + } + return source +} + +func (l *websocketTimelineLog) BeginRequest() { + if l == nil || !l.enabled || l.source == nil { + return + } + l.closeCurrentPart() + part, errCreate := l.source.CreatePart("request") + if errCreate != nil { + log.WithError(errCreate).Warn("failed to create websocket request detail log") + return + } + l.currentPart = part + l.currentPartHasLog = false +} + +func (l *websocketTimelineLog) Append(eventType string, payload []byte, timestamp time.Time) { + if l == nil || !l.enabled { + return + } + data := formatWebsocketTimelineEvent(eventType, payload, timestamp) + if len(data) == 0 { + return + } + if l.source != nil { + if l.currentPart == nil { + l.BeginRequest() + } + if l.currentPart == nil { + return + } + if errWrite := writeWebsocketTimelinePart(l.currentPart, data, l.currentPartHasLog); errWrite != nil { + log.WithError(errWrite).Warn("failed to write websocket request detail log") + return + } + l.currentPartHasLog = true + return + } + if l.builder != nil { + writeWebsocketTimelineBuilder(l.builder, data) + } +} + +func (l *websocketTimelineLog) SetContext(c *gin.Context) { + if l == nil || !l.enabled { + return + } + l.closeCurrentPart() + if l.source != nil { + if l.source.HasPayload() { + c.Set(requestlogging.WebsocketTimelineSourceContextKey, l.source) + return + } + if errCleanup := l.source.Cleanup(); errCleanup != nil { + log.WithError(errCleanup).Warn("failed to clean up empty websocket timeline log parts") + } + } + if l.builder != nil { + setWebsocketTimelineBody(c, l.builder.String()) + } +} + +func (l *websocketTimelineLog) String() string { + if l == nil || !l.enabled { + return "" + } + l.closeCurrentPart() + if l.source != nil { + data, errRead := l.source.Bytes() + if errRead != nil { + return "" + } + return string(data) + } + if l.builder == nil { + return "" + } + return l.builder.String() +} + +func (l *websocketTimelineLog) closeCurrentPart() { + if l == nil || l.currentPart == nil { + return + } + if errClose := l.currentPart.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close websocket request detail log") + } + l.currentPart = nil + l.currentPartHasLog = false +} + +func writeWebsocketTimelinePart(w io.Writer, data []byte, prependNewline bool) error { + if w == nil || len(data) == 0 { + return nil + } + if prependNewline { + if _, errWrite := io.WriteString(w, "\n"); errWrite != nil { + return errWrite + } + } + _, errWrite := w.Write(data) + return errWrite +} + +func writeWebsocketTimelineBuilder(builder *strings.Builder, data []byte) { + if builder == nil || len(data) == 0 { + return + } + if builder.Len() > 0 { + builder.WriteString("\n") + } + builder.Write(data) +} + +func appendWebsocketEvent(builder *strings.Builder, eventType string, payload []byte) { + if builder == nil { + return + } + trimmedPayload := bytes.TrimSpace(payload) + if len(trimmedPayload) == 0 { + return + } + if builder.Len() > 0 { + builder.WriteString("\n") + } + builder.WriteString("websocket.") + builder.WriteString(eventType) + builder.WriteString("\n") + builder.Write(trimmedPayload) + builder.WriteString("\n") +} + +func websocketPayloadEventType(payload []byte) string { + eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + if eventType == "" { + return "-" + } + return eventType +} + +func websocketPayloadPreview(payload []byte) string { + trimmedPayload := bytes.TrimSpace(payload) + if len(trimmedPayload) == 0 { + return "" + } + previewText := strings.ReplaceAll(string(trimmedPayload), "\n", "\\n") + previewText = strings.ReplaceAll(previewText, "\r", "\\r") + return previewText +} + +func isResponsesWebsocketCompletionEvent(eventType string) bool { + return eventType == wsEventTypeCompleted || eventType == wsEventTypeDone +} + +func responsesWebsocketErrorMessageFromPayload(payload []byte) *interfaces.ErrorMessage { + status := int(gjson.GetBytes(payload, "status").Int()) + if status <= 0 { + status = int(gjson.GetBytes(payload, "status_code").Int()) + } + if status <= 0 { + status = http.StatusInternalServerError + } + + errText := strings.TrimSpace(gjson.GetBytes(payload, "error.message").String()) + if errText == "" { + errText = strings.TrimSpace(gjson.GetBytes(payload, "message").String()) + } + if errText == "" { + errText = strings.TrimSpace(string(payload)) + } + if errText == "" { + errText = http.StatusText(status) + } + return &interfaces.ErrorMessage{StatusCode: status, Error: fmt.Errorf("%s", errText)} +} + +func setWebsocketTimelineBody(c *gin.Context, body string) { + setWebsocketBody(c, wsTimelineBodyKey, body) +} + +func setWebsocketBody(c *gin.Context, key string, body string) { + if c == nil { + return + } + trimmedBody := strings.TrimSpace(body) + if trimmedBody == "" { + return + } + c.Set(key, []byte(trimmedBody)) +} + +func writeResponsesWebsocketPayload(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, payload []byte, timestamp time.Time) error { + if wsTimelineLog != nil { + wsTimelineLog.Append("response", payload, timestamp) + } + if writer == nil || writer.conn == nil { + return fmt.Errorf("responses websocket: writer is nil") + } + writer.writeMu.Lock() + defer writer.writeMu.Unlock() + if writer.closing.Load() { + return websocket.ErrCloseSent + } + return writer.conn.WriteMessage(websocket.TextMessage, payload) +} + +func appendWebsocketTimelineDisconnect(timeline websocketTimelineAppender, err error, timestamp time.Time) { + if err == nil { + return + } + if timeline != nil { + timeline.Append("disconnect", []byte(err.Error()), timestamp) + } +} + +func appendWebsocketTimelineEvent(builder *strings.Builder, eventType string, payload []byte, timestamp time.Time) { + if builder == nil { + return + } + writeWebsocketTimelineBuilder(builder, formatWebsocketTimelineEvent(eventType, payload, timestamp)) +} + +func formatWebsocketTimelineEvent(eventType string, payload []byte, timestamp time.Time) []byte { + trimmedPayload := bytes.TrimSpace(payload) + if len(trimmedPayload) == 0 { + return nil + } + var builder strings.Builder + builder.WriteString("Timestamp: ") + builder.WriteString(timestamp.Format(time.RFC3339Nano)) + builder.WriteString("\n") + builder.WriteString("Event: websocket.") + builder.WriteString(eventType) + builder.WriteString("\n") + builder.Write(trimmedPayload) + builder.WriteString("\n") + return []byte(builder.String()) +} + +func markAPIResponseTimestamp(c *gin.Context) { + if c == nil { + return + } + if _, exists := c.Get("API_RESPONSE_TIMESTAMP"); exists { + return + } + c.Set("API_RESPONSE_TIMESTAMP", time.Now()) +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index a2c5e05bd..cf537b6df 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -1,37 +1,15 @@ package auth import ( - "bytes" "context" - "encoding/json" - "errors" - "fmt" - "io" - "math/rand/v2" "net/http" - "path/filepath" - "sort" - "strconv" - "strings" "sync" "sync/atomic" "time" - "github.com/google/uuid" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/home" - "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" - cliproxysession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" - coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" - sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" - log "github.com/sirupsen/logrus" - "github.com/tidwall/sjson" ) // ProviderExecutor defines the contract required by Manager to execute provider calls. @@ -64,101 +42,6 @@ type ExecutionSessionCloser interface { CloseExecutionSession(sessionID string) } -const ( - homeAuthCountMetadataKey = "__cliproxy_home_auth_count" - // CloseAllExecutionSessionsID asks an executor to release all active execution sessions. - // Executors that do not support this marker may ignore it. - CloseAllExecutionSessionsID = "__all_execution_sessions__" -) - -// RefreshEvaluator allows runtime state to override refresh decisions. -type RefreshEvaluator interface { - ShouldRefresh(now time.Time, auth *Auth) bool -} - -const ( - refreshCheckInterval = 5 * time.Second - refreshMaxConcurrency = 16 - refreshPendingBackoff = time.Minute - refreshFailureBackoff = 5 * time.Minute - // refreshIneffectiveBackoff throttles refresh attempts when an executor returns - // success but the auth still evaluates as needing refresh (e.g. token expiry - // wasn't updated). Without this guard, the auto-refresh loop can tight-loop and - // burn CPU at idle. - refreshIneffectiveBackoff = 30 * time.Second - quotaBackoffBase = time.Second - quotaBackoffMax = 30 * time.Minute - transientErrorCooldown = time.Minute -) - -var quotaCooldownDisabled atomic.Bool -var transientErrorCooldownSeconds atomic.Int64 - -// SetQuotaCooldownDisabled toggles quota cooldown scheduling globally. -func SetQuotaCooldownDisabled(disable bool) { - quotaCooldownDisabled.Store(disable) -} - -// SetTransientErrorCooldownSeconds configures cooldowns for 408/500/502/503/504. -// 0 keeps the legacy default; negative values disable transient error cooldowns. -func SetTransientErrorCooldownSeconds(seconds int) { - transientErrorCooldownSeconds.Store(int64(seconds)) -} - -func quotaCooldownDisabledForAuth(auth *Auth) bool { - return quotaCooldownDisabledForAuthWithConfig(auth, nil) -} - -func quotaCooldownDisabledForAuthWithConfig(auth *Auth, cfg *internalconfig.Config) bool { - if auth != nil { - if override, ok := auth.DisableCoolingOverride(); ok { - return override - } - if providerCoolingDisabledForAuth(auth, cfg) { - return true - } - } - if cfg != nil && cfg.DisableCooling { - return true - } - return quotaCooldownDisabled.Load() -} - -func providerCoolingDisabledForAuth(auth *Auth, cfg *internalconfig.Config) bool { - if auth == nil || cfg == nil { - return false - } - provider := strings.ToLower(strings.TrimSpace(auth.Provider)) - if provider == "" { - return false - } - providerKey := "" - compatName := "" - if auth.Attributes != nil { - providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) - compatName = strings.TrimSpace(auth.Attributes["compat_name"]) - } - if providerKey == "" && compatName == "" && provider != "openai-compatibility" { - return false - } - if providerKey == "" { - providerKey = provider - } - entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, provider) - return entry != nil && entry.DisableCooling -} - -func nextTransientErrorRetryAfter(now time.Time) time.Time { - seconds := transientErrorCooldownSeconds.Load() - if seconds < 0 { - return time.Time{} - } - if seconds == 0 { - return now.Add(transientErrorCooldown) - } - return now.Add(time.Duration(seconds) * time.Second) -} - // Result captures execution outcome used to adjust auth state. type Result struct { // AuthID references the auth that produced this result. @@ -305,7265 +188,3 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager { manager.scheduler = newAuthScheduler(selector) return manager } - -// HomeDispatchBundle is the immutable client and registry pair for one Home lifetime. -type HomeDispatchBundle struct { - client homeAuthDispatcher - registry *executionregistry.Registry - generation uint64 -} - -// PublishHomeDispatch publishes the selectable Home lifetime as one atomic bundle. -func (m *Manager) PublishHomeDispatch(client homeAuthDispatcher, registry *executionregistry.Registry, generation uint64) *HomeDispatchBundle { - if m == nil || client == nil || registry == nil { - return nil - } - bundle := &HomeDispatchBundle{client: client, registry: registry, generation: generation} - m.homeDispatchBundle.Store(bundle) - return bundle -} - -// ClearHomeDispatchBundle removes bundle only when it still belongs to the active lifetime. -func (m *Manager) ClearHomeDispatchBundle(bundle *HomeDispatchBundle) bool { - if m == nil || bundle == nil { - return false - } - return m.homeDispatchBundle.CompareAndSwap(bundle, nil) -} - -// HomeDispatchBundle returns the active Home lifetime bundle. -func (m *Manager) HomeDispatchBundle() *HomeDispatchBundle { - if m == nil { - return nil - } - return m.homeDispatchBundle.Load() -} - -// SetHomeExecutionRegistry preserves the legacy registry API for callers that also install the current dispatcher. -func (m *Manager) SetHomeExecutionRegistry(registry *executionregistry.Registry) { - if m == nil { - return - } - m.PublishHomeDispatch(currentHomeDispatcher(), registry, 0) -} - -// ClearHomeExecutionRegistry removes a matching legacy registry bundle. -func (m *Manager) ClearHomeExecutionRegistry(registry *executionregistry.Registry) bool { - bundle := m.HomeDispatchBundle() - if bundle == nil || bundle.registry != registry { - return false - } - return m.ClearHomeDispatchBundle(bundle) -} - -// HomeExecutionRegistry returns the registry from the active Home lifetime bundle. -func (m *Manager) HomeExecutionRegistry() *executionregistry.Registry { - bundle := m.HomeDispatchBundle() - if bundle == nil { - return nil - } - return bundle.registry -} - -func (m *Manager) SetPluginScheduler(scheduler PluginScheduler) { - if m == nil { - return - } - m.mu.Lock() - m.pluginScheduler = scheduler - m.mu.Unlock() -} - -func (m *Manager) hasPluginScheduler() bool { - if m == nil { - return false - } - m.mu.RLock() - scheduler := m.pluginScheduler - m.mu.RUnlock() - if scheduler == nil { - return false - } - if state, ok := scheduler.(pluginSchedulerState); ok { - return state.HasScheduler() - } - return true -} - -func isBuiltInSelector(selector Selector) bool { - switch selector.(type) { - case *RoundRobinSelector, *FillFirstSelector: - return true - default: - return false - } -} - -func (m *Manager) syncSchedulerFromSnapshot(auths []*Auth) { - if m == nil || m.scheduler == nil { - return - } - m.scheduler.rebuild(auths) -} - -func (m *Manager) syncScheduler() { - if m == nil || m.scheduler == nil { - return - } - m.syncSchedulerFromSnapshot(m.snapshotAuths()) -} - -func (m *Manager) snapshotAuths() []*Auth { - m.mu.RLock() - defer m.mu.RUnlock() - out := make([]*Auth, 0, len(m.auths)) - for _, a := range m.auths { - out = append(out, a.Clone()) - } - return out -} - -// RefreshSchedulerEntry re-upserts a single auth into the scheduler so that its -// supportedModelSet is rebuilt from the current global model registry state. -// This must be called after models have been registered for a newly added auth, -// because the initial scheduler.upsertAuth during Register/Update runs before -// registerModelsForAuth and therefore snapshots an empty model set. -func (m *Manager) RefreshSchedulerEntry(authID string) { - if m == nil || m.scheduler == nil || authID == "" { - return - } - m.mu.RLock() - auth, ok := m.auths[authID] - if !ok || auth == nil { - m.mu.RUnlock() - return - } - snapshot := auth.Clone() - m.mu.RUnlock() - m.scheduler.upsertAuth(snapshot) -} - -// RefreshSchedulerAll rebuilds scheduler entries for every known auth. -func (m *Manager) RefreshSchedulerAll() { - if m == nil { - return - } - m.mu.RLock() - ids := make([]string, 0, len(m.auths)) - for id := range m.auths { - ids = append(ids, id) - } - m.mu.RUnlock() - for _, id := range ids { - m.RefreshSchedulerEntry(id) - } -} - -// ReconcileRegistryModelStates aligns per-model runtime state with the current -// registry snapshot for one auth. -// -// Supported models are reset to a clean state because re-registration already -// cleared the registry-side cooldown/suspension snapshot. ModelStates for -// models that are no longer present in the registry are pruned entirely so -// renamed/removed models cannot keep auth-level status stale. -func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID string) { - if m == nil || authID == "" { - return - } - - supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID) - supported := make(map[string]struct{}, len(supportedModels)) - for _, model := range supportedModels { - if model == nil { - continue - } - modelKey := canonicalModelKey(model.ID) - if modelKey == "" { - continue - } - supported[modelKey] = struct{}{} - } - - var snapshot *Auth - now := time.Now() - - m.mu.Lock() - auth, ok := m.auths[authID] - if ok && auth != nil && len(auth.ModelStates) > 0 { - changed := false - for modelKey, state := range auth.ModelStates { - baseModel := canonicalModelKey(modelKey) - if baseModel == "" { - baseModel = strings.TrimSpace(modelKey) - } - if _, supportedModel := supported[baseModel]; !supportedModel { - // Drop state for models that disappeared from the current registry - // snapshot. Keeping them around leaks stale errors into auth-level - // status, management output, and websocket fallback checks. - delete(auth.ModelStates, modelKey) - changed = true - continue - } - if state == nil { - continue - } - if modelStateIsClean(state) { - continue - } - resetModelState(state, now) - changed = true - } - if len(auth.ModelStates) == 0 { - auth.ModelStates = nil - } - if changed { - updateAggregatedAvailability(auth, now) - if !hasModelError(auth, now) { - auth.LastError = nil - auth.StatusMessage = "" - auth.Status = StatusActive - } - auth.UpdatedAt = now - if errPersist := m.persist(ctx, auth); errPersist != nil { - logEntryWithRequestID(ctx).WithField("auth_id", auth.ID).Warnf("failed to persist auth changes during model state reconciliation: %v", errPersist) - } - snapshot = auth.Clone() - } - } - m.mu.Unlock() - - if m.scheduler != nil && snapshot != nil { - m.scheduler.upsertAuth(snapshot) - } -} - -func (m *Manager) SetSelector(selector Selector) { - if m == nil { - return - } - if selector == nil { - selector = &RoundRobinSelector{} - } - m.mu.Lock() - m.selector = selector - m.mu.Unlock() - if m.scheduler != nil { - m.scheduler.setSelector(selector) - m.syncScheduler() - } -} - -// Selector returns the current credential selector. -func (m *Manager) Selector() Selector { - if m == nil { - return nil - } - m.mu.RLock() - defer m.mu.RUnlock() - return m.selector -} - -// SetStore swaps the underlying persistence store. -func (m *Manager) SetStore(store Store) { - m.mu.Lock() - defer m.mu.Unlock() - m.store = store -} - -// SetCooldownStateStore swaps the independent runtime cooldown state store. -func (m *Manager) SetCooldownStateStore(store CooldownStateStore) { - if m == nil { - return - } - m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - m.mu.Lock() - defer m.mu.Unlock() - m.cooldownStore = store -} - -// SetRoundTripperProvider register a provider that returns a per-auth RoundTripper. -func (m *Manager) SetRoundTripperProvider(p RoundTripperProvider) { - m.mu.Lock() - m.rtProvider = p - m.mu.Unlock() -} - -// SetConfig updates the runtime config snapshot used by request-time helpers. -// Callers should provide the latest config on reload so per-credential alias mapping stays in sync. -func (m *Manager) SetConfig(cfg *internalconfig.Config) { - if m == nil { - return - } - m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - if m.setConfigSnapshotLocked(cfg) { - m.persistCooldownStatesLocked(context.Background()) - } -} - -// SetConfigSnapshot updates only in-memory configuration state. It reports whether -// a caller must persist cleared cooldown state after its commit critical section. -func (m *Manager) SetConfigSnapshot(cfg *internalconfig.Config) bool { - if m == nil { - return false - } - m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - return m.setConfigSnapshotLocked(cfg) -} - -func (m *Manager) setConfigSnapshotLocked(cfg *internalconfig.Config) bool { - if cfg == nil { - cfg = &internalconfig.Config{} - } - m.mu.RLock() - oldCooldownStore := m.cooldownStore - m.mu.RUnlock() - m.runtimeConfig.Store(cfg) - clearedCooldowns := m.clearDisabledCooldownStates(cfg) - if clearedCooldowns && oldCooldownStore != nil { - m.mu.Lock() - if m.cooldownStore == oldCooldownStore { - m.pendingCooldownStateStore = oldCooldownStore - } - m.mu.Unlock() - } - if !cfg.Home.Enabled { - m.clearHomeRuntimeAuths() - } - m.rebuildAPIKeyModelAliasFromRuntimeConfig() - return clearedCooldowns -} - -// ApplyConfigWithCooldownStateStore serializes a config update with its cooldown -// store transition. It persists the resulting state to the captured old store before -// exposing the resolved replacement store. -func (m *Manager) ApplyConfigWithCooldownStateStore(ctx context.Context, cfg *internalconfig.Config, store CooldownStateStore) bool { - if m == nil { - return false - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - - m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - m.mu.RLock() - oldStore := m.cooldownStore - m.mu.RUnlock() - m.setConfigSnapshotLocked(cfg) - if oldStore != nil && !m.persistCooldownStatesToLocked(ctx, oldStore) { - return false - } - if errContext := ctx.Err(); errContext != nil { - return false - } - m.mu.Lock() - defer m.mu.Unlock() - if m.cooldownStore != oldStore { - return false - } - if m.pendingCooldownStateStore == oldStore { - m.pendingCooldownStateStore = nil - } - m.cooldownStore = store - return true -} - -// PersistCooldownStates writes the current cooldown snapshot using ctx. -func (m *Manager) PersistCooldownStates(ctx context.Context) { - m.persistCooldownStates(ctx) -} - -// SwapCooldownStateStore persists cleared state to the old store before replacing it. -// Persistence is deliberately performed without holding the manager lock. -func (m *Manager) SwapCooldownStateStore(ctx context.Context, store CooldownStateStore, persistOld bool) bool { - if m == nil { - return false - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - m.mu.RLock() - oldStore := m.cooldownStore - pendingStore := m.pendingCooldownStateStore - m.mu.RUnlock() - storeToPersist := pendingStore - if storeToPersist == nil && persistOld { - storeToPersist = oldStore - } - if storeToPersist != nil && !m.persistCooldownStatesToLocked(ctx, storeToPersist) { - return false - } - if errContext := ctx.Err(); errContext != nil { - return false - } - m.mu.Lock() - defer m.mu.Unlock() - if m.cooldownStore != oldStore { - return false - } - if m.pendingCooldownStateStore == storeToPersist { - m.pendingCooldownStateStore = nil - } - m.cooldownStore = store - return true -} - -func (m *Manager) cooldownDisabledForAuth(auth *Auth) bool { - if m == nil { - return quotaCooldownDisabledForAuth(auth) - } - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - return quotaCooldownDisabledForAuthWithConfig(auth, cfg) -} - -func (m *Manager) clearDisabledCooldownStates(cfg *internalconfig.Config) bool { - if m == nil { - return false - } - now := time.Now() - snapshots := make([]*Auth, 0) - m.mu.Lock() - for _, auth := range m.auths { - if auth == nil { - continue - } - if !quotaCooldownDisabledForAuthWithConfig(auth, cfg) && !auth.Disabled && auth.Status != StatusDisabled { - continue - } - if clearCooldownStateForAuth(auth, now) { - snapshots = append(snapshots, auth.Clone()) - } - } - m.mu.Unlock() - - if m.scheduler != nil { - for _, snapshot := range snapshots { - m.scheduler.upsertAuth(snapshot) - } - } - return len(snapshots) > 0 -} - -// RestoreCooldownStates restores unexpired persisted cooldown records into registered auths. -func (m *Manager) RestoreCooldownStates(ctx context.Context) error { - if m == nil { - return nil - } - if ctx == nil { - ctx = context.Background() - } - m.mu.RLock() - store := m.cooldownStore - m.mu.RUnlock() - if store == nil { - return nil - } - records, errLoad := store.Load(ctx) - if errLoad != nil { - return errLoad - } - if len(records) == 0 { - return nil - } - - now := time.Now() - authLevelRecords := make([]CooldownStateRecord, 0) - snapshotsByID := make(map[string]*Auth) - - m.mu.Lock() - for _, record := range records { - if strings.TrimSpace(record.Model) == "" { - authLevelRecords = append(authLevelRecords, record) - continue - } - if m.restoreCooldownRecordLocked(record, now) { - if auth := m.auths[strings.TrimSpace(record.AuthID)]; auth != nil { - snapshotsByID[auth.ID] = auth.Clone() - } - } - } - for _, record := range authLevelRecords { - if m.restoreCooldownRecordLocked(record, now) { - if auth := m.auths[strings.TrimSpace(record.AuthID)]; auth != nil { - snapshotsByID[auth.ID] = auth.Clone() - } - } - } - m.mu.Unlock() - - if m.scheduler != nil { - for _, snapshot := range snapshotsByID { - m.scheduler.upsertAuth(snapshot) - } - } - m.persistCooldownStates(ctx) - return nil -} - -func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now time.Time) bool { - authID := strings.TrimSpace(record.AuthID) - if authID == "" || record.NextRetryAfter.IsZero() || !record.NextRetryAfter.After(now) { - return false - } - auth := m.auths[authID] - if auth == nil || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) { - return false - } - updatedAt := record.UpdatedAt - if updatedAt.IsZero() { - updatedAt = now - } - reason := strings.TrimSpace(record.Reason) - model := strings.TrimSpace(record.Model) - quota := record.Quota - if quota.Exceeded && quota.NextRecoverAt.IsZero() { - quota.NextRecoverAt = record.NextRetryAfter - } - - if model == "" { - auth.Unavailable = true - auth.Status = StatusError - auth.NextRetryAfter = record.NextRetryAfter - auth.Quota = quota - auth.UpdatedAt = updatedAt - if reason != "" { - auth.StatusMessage = reason - } - auth.LastError = cloneError(record.LastError) - return true - } - - state := ensureModelState(auth, model) - state.Unavailable = true - state.Status = StatusError - state.NextRetryAfter = record.NextRetryAfter - state.Quota = quota - state.UpdatedAt = updatedAt - if reason != "" { - state.StatusMessage = reason - } - state.LastError = cloneError(record.LastError) - updateAggregatedAvailability(auth, now) - return true -} - -func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { - if auth == nil { - return false - } - changed := false - if auth.Unavailable || !auth.NextRetryAfter.IsZero() || auth.Quota.Exceeded || !auth.Quota.NextRecoverAt.IsZero() { - auth.Unavailable = false - auth.NextRetryAfter = time.Time{} - auth.Quota = QuotaState{} - auth.UpdatedAt = now - changed = true - } - for _, state := range auth.ModelStates { - if state == nil { - continue - } - if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() { - state.Unavailable = false - state.NextRetryAfter = time.Time{} - state.Quota = QuotaState{} - state.UpdatedAt = now - changed = true - } - } - if len(auth.ModelStates) > 0 { - updateAggregatedAvailability(auth, now) - } - return changed -} - -func dedupeStrings(values []string) []string { - if len(values) < 2 { - return values - } - seen := make(map[string]struct{}, len(values)) - out := values[:0] - for _, value := range values { - value = strings.TrimSpace(value) - if value == "" { - continue - } - if _, ok := seen[value]; ok { - continue - } - seen[value] = struct{}{} - out = append(out, value) - } - return out -} - -// ResetQuota clears quota/cooldown state for an auth and resumes registry routing. -func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []string, error) { - if m == nil { - return nil, nil, nil - } - authID = strings.TrimSpace(authID) - if authID == "" { - return nil, nil, fmt.Errorf("auth id is required") - } - - now := time.Now() - var snapshot *Auth - models := make([]string, 0) - registeredModels := modelsForRegisteredAuth(authID) - cooldownStateChanged := false - - m.mu.Lock() - auth, ok := m.auths[authID] - if !ok || auth == nil { - m.mu.Unlock() - return nil, nil, nil - } - - var cooldownRecordsBefore []CooldownStateRecord - trackCooldownState := m.cooldownStore != nil - if trackCooldownState { - cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now) - } - - for modelKey, state := range auth.ModelStates { - if strings.TrimSpace(modelKey) == "" { - continue - } - models = append(models, modelKey) - if state != nil { - resetModelState(state, now) - } - } - if clearCooldownStateForAuth(auth, now) { - if len(models) == 0 { - models = append(models, registeredModels...) - } - } else if len(auth.ModelStates) > 0 { - updateAggregatedAvailability(auth, now) - } - - if len(models) == 0 { - models = append(models, registeredModels...) - } - models = dedupeStrings(models) - - if !auth.Disabled && auth.Status != StatusDisabled && !hasModelError(auth, now) { - auth.LastError = nil - auth.StatusMessage = "" - auth.Status = StatusActive - } - auth.UpdatedAt = now - if errPersist := m.persist(ctx, auth); errPersist != nil { - m.mu.Unlock() - return nil, nil, errPersist - } - snapshot = auth.Clone() - if trackCooldownState { - cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) - cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) - } - m.mu.Unlock() - - for _, modelKey := range models { - registry.GetGlobalRegistry().ClearModelQuotaExceeded(authID, modelKey) - registry.GetGlobalRegistry().ResumeClientModel(authID, modelKey) - } - if m.scheduler != nil && snapshot != nil { - m.scheduler.upsertAuth(snapshot) - } - if snapshot != nil && cooldownStateChanged { - m.persistCooldownStates(ctx) - } - return snapshot, models, nil -} - -func modelsForRegisteredAuth(authID string) []string { - supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID) - models := make([]string, 0, len(supportedModels)) - for _, supportedModel := range supportedModels { - if supportedModel == nil || strings.TrimSpace(supportedModel.ID) == "" { - continue - } - models = append(models, supportedModel.ID) - } - return models -} - -func (m *Manager) persistCooldownStates(ctx context.Context) { - if m == nil { - return - } - m.configCooldownMu.Lock() - defer m.configCooldownMu.Unlock() - m.persistCooldownStatesLocked(ctx) -} - -func (m *Manager) persistCooldownStatesLocked(ctx context.Context) { - m.mu.RLock() - store := m.cooldownStore - m.mu.RUnlock() - if m.persistCooldownStatesToLocked(ctx, store) { - m.mu.Lock() - if m.pendingCooldownStateStore == store { - m.pendingCooldownStateStore = nil - } - m.mu.Unlock() - } -} - -func (m *Manager) persistCooldownStatesToLocked(ctx context.Context, store CooldownStateStore) bool { - if m == nil || store == nil { - return true - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - records := m.cooldownStateRecordsSnapshot() - if errSave := store.Save(ctx, records); errSave != nil { - logEntryWithRequestID(ctx).Warnf("failed to persist cooldown state: %v", errSave) - return false - } - return ctx.Err() == nil -} - -func (m *Manager) cooldownStateRecordsSnapshot() []CooldownStateRecord { - now := time.Now() - records := make([]CooldownStateRecord, 0) - - m.mu.RLock() - for _, auth := range m.auths { - records = append(records, m.cooldownStateRecordsForAuthLocked(auth, now)...) - } - m.mu.RUnlock() - - sort.Slice(records, func(i, j int) bool { - if records[i].Provider != records[j].Provider { - return records[i].Provider < records[j].Provider - } - if records[i].AuthID != records[j].AuthID { - return records[i].AuthID < records[j].AuthID - } - return records[i].Model < records[j].Model - }) - return records -} - -func (m *Manager) cooldownStateRecordsForAuthLocked(auth *Auth, now time.Time) []CooldownStateRecord { - if auth == nil || auth.ID == "" || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) { - return nil - } - records := make([]CooldownStateRecord, 0, 1+len(auth.ModelStates)) - if record, ok := authCooldownStateRecord(auth, now); ok { - records = append(records, record) - } - for model, state := range auth.ModelStates { - if record, ok := modelCooldownStateRecord(auth, model, state, now); ok { - records = append(records, record) - } - } - sort.Slice(records, func(i, j int) bool { - return records[i].Model < records[j].Model - }) - return records -} - -func cooldownStateRecordsEqual(a, b []CooldownStateRecord) bool { - if len(a) != len(b) { - return false - } - for i := range a { - if !cooldownStateRecordEqual(a[i], b[i]) { - return false - } - } - return true -} - -func cooldownStateRecordEqual(a, b CooldownStateRecord) bool { - if a.Provider != b.Provider || - a.AuthID != b.AuthID || - a.AuthFile != b.AuthFile || - a.Model != b.Model || - a.Status != b.Status || - a.Reason != b.Reason || - !a.NextRetryAfter.Equal(b.NextRetryAfter) || - !a.UpdatedAt.Equal(b.UpdatedAt) || - !cooldownQuotaEqual(a.Quota, b.Quota) { - return false - } - return cooldownErrorEqual(a.LastError, b.LastError) -} - -func cooldownQuotaEqual(a, b QuotaState) bool { - return a.Exceeded == b.Exceeded && - a.Reason == b.Reason && - a.BackoffLevel == b.BackoffLevel && - a.NextRecoverAt.Equal(b.NextRecoverAt) -} - -func cooldownErrorEqual(a, b *Error) bool { - if a == nil || b == nil { - return a == b - } - return a.Code == b.Code && - a.Message == b.Message && - a.Retryable == b.Retryable && - a.HTTPStatus == b.HTTPStatus -} - -func authCooldownStateRecord(auth *Auth, now time.Time) (CooldownStateRecord, bool) { - if auth == nil || !auth.Unavailable || auth.NextRetryAfter.IsZero() || !auth.NextRetryAfter.After(now) { - return CooldownStateRecord{}, false - } - return CooldownStateRecord{ - Provider: strings.TrimSpace(auth.Provider), - AuthID: auth.ID, - AuthFile: cooldownAuthFile(auth), - Status: "cooling", - NextRetryAfter: auth.NextRetryAfter, - Reason: cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError), - Quota: auth.Quota, - LastError: cloneError(auth.LastError), - UpdatedAt: auth.UpdatedAt, - }, true -} - -func modelCooldownStateRecord(auth *Auth, model string, state *ModelState, now time.Time) (CooldownStateRecord, bool) { - model = strings.TrimSpace(model) - if auth == nil || state == nil || model == "" || !state.Unavailable || state.NextRetryAfter.IsZero() || !state.NextRetryAfter.After(now) { - return CooldownStateRecord{}, false - } - return CooldownStateRecord{ - Provider: strings.TrimSpace(auth.Provider), - AuthID: auth.ID, - AuthFile: cooldownAuthFile(auth), - Model: model, - Status: "cooling", - NextRetryAfter: state.NextRetryAfter, - Reason: cooldownReason(state.StatusMessage, state.Quota, state.LastError), - Quota: state.Quota, - LastError: cloneError(state.LastError), - UpdatedAt: state.UpdatedAt, - }, true -} - -func cooldownReason(statusMessage string, quota QuotaState, lastErr *Error) string { - if reason := strings.TrimSpace(quota.Reason); reason != "" { - return reason - } - if statusMessage = strings.TrimSpace(statusMessage); statusMessage != "" { - return statusMessage - } - if lastErr != nil { - if code := strings.TrimSpace(lastErr.Code); code != "" { - return code - } - if message := strings.TrimSpace(lastErr.Message); message != "" { - return message - } - } - return "" -} - -// HomeEnabled reports whether the home control plane integration is enabled in the runtime config. -func (m *Manager) HomeEnabled() bool { - if m == nil { - return false - } - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - return cfg != nil && cfg.Home.Enabled -} - -func (m *Manager) localExecutionAllowed() bool { - return m != nil && !m.HomeEnabled() -} - -func (m *Manager) localFallbackAuth(authID string) *Auth { - if !m.localExecutionAllowed() { - return nil - } - m.mu.RLock() - auth := m.auths[strings.TrimSpace(authID)] - m.mu.RUnlock() - if auth == nil { - return nil - } - return auth.Clone() -} - -func (m *Manager) lookupAPIKeyUpstreamModel(authID, requestedModel string) string { - if m == nil { - return "" - } - authID = strings.TrimSpace(authID) - if authID == "" { - return "" - } - requestedModel = strings.TrimSpace(requestedModel) - if requestedModel == "" { - return "" - } - table, _ := m.apiKeyModelAlias.Load().(apiKeyModelAliasTable) - if table == nil { - return "" - } - byAlias := table[authID] - if len(byAlias) == 0 { - return "" - } - key := strings.ToLower(thinking.ParseSuffix(requestedModel).ModelName) - if key == "" { - key = strings.ToLower(requestedModel) - } - resolved := strings.TrimSpace(byAlias[key]) - if resolved == "" { - return "" - } - return preserveRequestedModelSuffix(requestedModel, resolved) -} - -func isAPIKeyAuth(auth *Auth) bool { - if auth == nil { - return false - } - return auth.AuthKind() == AuthKindAPIKey -} - -func isOpenAICompatAPIKeyAuth(auth *Auth) bool { - if !isAPIKeyAuth(auth) { - return false - } - if strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { - return true - } - if auth.Attributes == nil { - return false - } - return strings.TrimSpace(auth.Attributes["compat_name"]) != "" -} - -func openAICompatProviderKey(auth *Auth) string { - if auth == nil { - return "" - } - if auth.Attributes != nil { - if providerKey := strings.TrimSpace(auth.Attributes["provider_key"]); providerKey != "" { - return util.OpenAICompatibleProviderKey(providerKey) - } - if compatName := strings.TrimSpace(auth.Attributes["compat_name"]); compatName != "" { - return util.OpenAICompatibleProviderKey(compatName) - } - } - return util.OpenAICompatibleProviderKey(auth.Provider) -} - -func openAICompatModelPoolKey(auth *Auth, requestedModel string) string { - base := strings.TrimSpace(thinking.ParseSuffix(requestedModel).ModelName) - if base == "" { - base = strings.TrimSpace(requestedModel) - } - return strings.ToLower(strings.TrimSpace(auth.ID)) + "|" + openAICompatProviderKey(auth) + "|" + strings.ToLower(base) -} - -func (m *Manager) nextModelPoolOffset(key string, size int) int { - if m == nil || size <= 1 { - return 0 - } - key = strings.TrimSpace(key) - if key == "" { - return 0 - } - m.mu.Lock() - defer m.mu.Unlock() - if m.modelPoolOffsets == nil { - m.modelPoolOffsets = make(map[string]int) - } - offset := m.modelPoolOffsets[key] - if offset >= 2_147_483_640 { - offset = 0 - } - m.modelPoolOffsets[key] = offset + 1 - if size <= 0 { - return 0 - } - return offset % size -} - -func rotateStrings(values []string, offset int) []string { - if len(values) <= 1 { - return values - } - if offset <= 0 { - out := make([]string, len(values)) - copy(out, values) - return out - } - offset = offset % len(values) - out := make([]string, 0, len(values)) - out = append(out, values[offset:]...) - out = append(out, values[:offset]...) - return out -} - -func (m *Manager) resolveOpenAICompatUpstreamModelPool(auth *Auth, requestedModel string) []string { - if m == nil || !isOpenAICompatAPIKeyAuth(auth) { - return nil - } - requestedModel = strings.TrimSpace(requestedModel) - if requestedModel == "" { - return nil - } - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - if cfg == nil { - cfg = &internalconfig.Config{} - } - providerKey := "" - compatName := "" - if auth.Attributes != nil { - providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) - compatName = strings.TrimSpace(auth.Attributes["compat_name"]) - } - entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider) - if entry == nil { - return nil - } - return resolveModelAliasPoolFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -func preserveRequestedModelSuffix(requestedModel, resolved string) string { - return preserveResolvedModelSuffix(resolved, thinking.ParseSuffix(requestedModel)) -} - -func (m *Manager) executionModelCandidates(auth *Auth, routeModel string) []string { - if auth != nil && auth.Attributes != nil { - if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" { - return []string{homeModel} - } - } - requestedModel := rewriteModelForAuth(routeModel, auth) - requestedModel = m.applyOAuthModelAlias(auth, requestedModel) - if pool := m.resolveOpenAICompatUpstreamModelPool(auth, requestedModel); len(pool) > 0 { - if len(pool) == 1 { - return pool - } - offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, requestedModel), len(pool)) - return rotateStrings(pool, offset) - } - resolved := m.applyAPIKeyModelAlias(auth, requestedModel) - if strings.TrimSpace(resolved) == "" { - resolved = requestedModel - } - return []string{resolved} -} - -func (m *Manager) selectionModelForAuth(auth *Auth, routeModel string) string { - requestedModel := rewriteModelForAuth(routeModel, auth) - if strings.TrimSpace(requestedModel) == "" { - requestedModel = strings.TrimSpace(routeModel) - } - resolvedModel := m.applyOAuthModelAlias(auth, requestedModel) - if strings.TrimSpace(resolvedModel) == "" { - resolvedModel = requestedModel - } - return resolvedModel -} - -func (m *Manager) selectionModelKeyForAuth(auth *Auth, routeModel string) string { - return canonicalModelKey(m.selectionModelForAuth(auth, routeModel)) -} - -func (m *Manager) stateModelForExecution(auth *Auth, routeModel, upstreamModel string, pooled bool) string { - if auth != nil && auth.Attributes != nil { - if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" { - if resolved := strings.TrimSpace(upstreamModel); resolved != "" { - return resolved - } - return homeModel - } - } - stateModel := executionResultModel(routeModel, upstreamModel, pooled) - selectionModel := m.selectionModelForAuth(auth, routeModel) - if canonicalModelKey(selectionModel) == canonicalModelKey(upstreamModel) && strings.TrimSpace(selectionModel) != "" { - return strings.TrimSpace(upstreamModel) - } - return stateModel -} - -func executionResultModel(routeModel, upstreamModel string, pooled bool) string { - if pooled { - if resolved := strings.TrimSpace(upstreamModel); resolved != "" { - return resolved - } - } - if requested := strings.TrimSpace(routeModel); requested != "" { - return requested - } - return strings.TrimSpace(upstreamModel) -} - -func (m *Manager) filterExecutionModels(auth *Auth, routeModel string, candidates []string, pooled bool) []string { - if len(candidates) == 0 { - return nil - } - now := time.Now() - out := make([]string, 0, len(candidates)) - for _, upstreamModel := range candidates { - stateModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) - blocked, _, _ := isAuthBlockedForModel(auth, stateModel, now) - if blocked { - continue - } - out = append(out, upstreamModel) - } - return out -} - -func (m *Manager) preparedExecutionModels(auth *Auth, routeModel string) ([]string, bool) { - candidates := m.executionModelCandidates(auth, routeModel) - pooled := len(candidates) > 1 - return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled -} - -func (m *Manager) preparedExecutionModelsWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) { - candidates, pooled, aliasResult := m.executionModelCandidatesWithAlias(auth, routeModel) - return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled, aliasResult -} - -func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) { - requestedModel := rewriteModelForAuth(routeModel, auth) - aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel) - if aliasResult.ForceMapping && auth != nil && auth.Attributes != nil && strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") { - aliasResult.OriginalAlias = strings.TrimSpace(routeModel) - } - upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult) - - var candidates []string - if auth != nil && auth.Attributes != nil { - if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" { - candidates = []string{homeModel} - } - } - if len(candidates) == 0 { - if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) > 0 { - if len(pool) == 1 { - candidates = pool - } else { - offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, upstreamModel), len(pool)) - candidates = rotateStrings(pool, offset) - } - } else { - resolved := m.applyAPIKeyModelAlias(auth, upstreamModel) - if strings.TrimSpace(resolved) == "" { - resolved = upstreamModel - } - candidates = []string{resolved} - } - } - pooled := len(candidates) > 1 - return candidates, pooled, aliasResult -} - -func (m *Manager) resolveExecutionAliasResult(auth *Auth, routeModel string) OAuthModelAliasResult { - requestedModel := rewriteModelForAuth(routeModel, auth) - return m.resolveExecutionAliasResultForRequested(auth, requestedModel) -} - -func (m *Manager) resolveExecutionAliasResultForRequested(auth *Auth, requestedModel string) OAuthModelAliasResult { - if result := homeForceMappingAliasResult(auth, requestedModel); result.ForceMapping { - return result - } - if auth != nil && auth.AuthKind() == AuthKindAPIKey { - return m.resolveAPIKeyModelAliasWithResult(auth, requestedModel) - } - return m.applyOAuthModelAliasWithResult(auth, requestedModel) -} - -func homeForceMappingAliasResult(auth *Auth, requestedModel string) OAuthModelAliasResult { - if auth == nil || auth.Attributes == nil || !strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") { - return OAuthModelAliasResult{} - } - originalAlias := strings.TrimSpace(auth.Attributes[homeOriginalAliasAttributeKey]) - canonicalOriginalAlias := canonicalHomeConcurrencyModelKey(auth.Attributes[homeOriginalAliasAttributeKey]) - canonicalRequestedModel := canonicalHomeConcurrencyModelKey(requestedModel) - if canonicalOriginalAlias == "" || canonicalOriginalAlias != canonicalRequestedModel { - return OAuthModelAliasResult{} - } - upstreamModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]) - if upstreamModel == "" { - upstreamModel = strings.TrimSpace(requestedModel) - } - return OAuthModelAliasResult{ - UpstreamModel: upstreamModel, - ForceMapping: true, - OriginalAlias: originalAlias, - } -} - -func executionAliasPoolModel(auth *Auth, requestedModel string, aliasResult OAuthModelAliasResult) string { - if auth != nil && auth.AuthKind() == AuthKindAPIKey { - if strings.TrimSpace(requestedModel) != "" { - return requestedModel - } - } - if strings.TrimSpace(aliasResult.UpstreamModel) != "" { - return aliasResult.UpstreamModel - } - return requestedModel -} - -func (m *Manager) resolveAPIKeyModelAliasWithResult(auth *Auth, requestedModel string) OAuthModelAliasResult { - if m == nil || auth == nil { - return OAuthModelAliasResult{} - } - requestedModel = strings.TrimSpace(requestedModel) - if requestedModel == "" { - return OAuthModelAliasResult{} - } - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - if cfg == nil { - cfg = &internalconfig.Config{} - } - provider := strings.ToLower(strings.TrimSpace(auth.Provider)) - var models []modelAliasEntry - switch provider { - case "gemini": - if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil { - models = asModelAliasEntries(entry.Models) - } - case "gemini-interactions": - if entry := resolveInteractionsAPIKeyConfig(cfg, auth); entry != nil { - models = asModelAliasEntries(entry.Models) - } - case "claude": - if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil { - models = asModelAliasEntries(entry.Models) - } - case "codex": - if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil { - models = asModelAliasEntries(entry.Models) - } - case "xai": - if entry := resolveXAIAPIKeyConfig(cfg, auth); entry != nil { - models = asModelAliasEntries(entry.Models) - } - case "vertex": - if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil { - models = asModelAliasEntries(entry.Models) - } - default: - providerKey := "" - compatName := "" - if auth.Attributes != nil { - providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) - compatName = strings.TrimSpace(auth.Attributes["compat_name"]) - } - if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { - if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil { - models = asModelAliasEntries(entry.Models) - } - } - } - if len(models) == 0 { - return OAuthModelAliasResult{UpstreamModel: requestedModel} - } - result := resolveModelAliasResultFromConfigModels(requestedModel, models) - if strings.TrimSpace(result.UpstreamModel) == "" { - return OAuthModelAliasResult{UpstreamModel: requestedModel} - } - return result -} - -func (m *Manager) prepareExecutionModels(auth *Auth, routeModel string) []string { - models, _ := m.preparedExecutionModels(auth, routeModel) - return models -} - -func rewriteForceMappedResponse(resp *cliproxyexecutor.Response, aliasResult OAuthModelAliasResult) { - if resp == nil || !aliasResult.ForceMapping || strings.TrimSpace(aliasResult.OriginalAlias) == "" { - return - } - resp.Payload = rewriteModelInResponse(resp.Payload, aliasResult.OriginalAlias) -} - -func rewriteForceMappedStreamChunk(rewriter *StreamRewriter, payload []byte) []byte { - if rewriter == nil || len(payload) == 0 { - return payload - } - rewritten := rewriter.RewriteChunk(payload) - if len(rewritten) > 0 { - return rewritten - } - if bytes.Contains(payload, []byte("data:")) { - if lineWise := rewriteSSEPayloadLines(payload, rewriter.options.RewriteModel); len(lineWise) > 0 { - return lineWise - } - } - if len(rewriter.pendingBuf) > 0 { - return nil - } - return nil -} - -func finishForceMappedStreamChunks(rewriter *StreamRewriter) []byte { - if rewriter == nil { - return nil - } - return rewriter.Finish() -} - -func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeModel string, now time.Time) ([]*Auth, error) { - if len(auths) == 0 { - return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"} - } - - availableByPriority := make(map[int][]*Auth) - cooldownCount := 0 - var earliest time.Time - for _, candidate := range auths { - checkModel := m.selectionModelForAuth(candidate, routeModel) - blocked, reason, next := isAuthBlockedForModel(candidate, checkModel, now) - if !blocked { - priority := authPriority(candidate) - availableByPriority[priority] = append(availableByPriority[priority], candidate) - continue - } - if reason == blockReasonCooldown { - cooldownCount++ - if !next.IsZero() && (earliest.IsZero() || next.Before(earliest)) { - earliest = next - } - } - } - - if len(availableByPriority) == 0 { - if cooldownCount == len(auths) && !earliest.IsZero() { - providerForError := provider - if providerForError == "mixed" { - providerForError = "" - } - resetIn := earliest.Sub(now) - if resetIn < 0 { - resetIn = 0 - } - return nil, newModelCooldownError(routeModel, providerForError, resetIn) - } - return nil, &Error{Code: "auth_unavailable", Message: "no auth available"} - } - - bestPriority := 0 - found := false - for priority := range availableByPriority { - if !found || priority > bestPriority { - bestPriority = priority - found = true - } - } - - available := availableByPriority[bestPriority] - if len(available) > 1 { - sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID }) - } - return available, nil -} - -func selectionArgForSelector(selector Selector, routeModel string) string { - if isBuiltInSelector(selector) { - return "" - } - return routeModel -} - -func schedulerAttributeSensitive(key string) bool { - key = strings.ToLower(strings.TrimSpace(key)) - normalized := strings.NewReplacer("-", "_", ".", "_", " ", "_").Replace(key) - compact := strings.NewReplacer("_", "", "-", "", ".", "", " ", "").Replace(key) - for _, fragment := range []string{ - "api_key", - "apikey", - "token", - "secret", - "cookie", - "credential", - "password", - "storage", - "authorization", - "auth_header", - "proxy_url", - } { - if strings.Contains(key, fragment) || strings.Contains(normalized, fragment) || strings.Contains(compact, fragment) { - return true - } - } - return false -} - -func schedulerSafeAttributes(src map[string]string) map[string]string { - if len(src) == 0 { - return nil - } - out := make(map[string]string, len(src)) - for key, value := range src { - if schedulerAttributeSensitive(key) { - continue - } - out[key] = value - } - if len(out) == 0 { - return nil - } - return out -} - -func cloneSchedulerAnyMap(src map[string]any) map[string]any { - if len(src) == 0 { - return nil - } - out := make(map[string]any, len(src)) - for key, value := range src { - out[key] = value - } - return out -} - -func cloneAuthSlice(auths []*Auth) []*Auth { - if len(auths) == 0 { - return nil - } - out := make([]*Auth, 0, len(auths)) - for _, auth := range auths { - if auth == nil { - continue - } - out = append(out, auth.Clone()) - } - return out -} - -func schedulerAuthCandidates(auths []*Auth) []pluginapi.SchedulerAuthCandidate { - if len(auths) == 0 { - return nil - } - out := make([]pluginapi.SchedulerAuthCandidate, 0, len(auths)) - for _, auth := range auths { - if auth == nil { - continue - } - out = append(out, pluginapi.SchedulerAuthCandidate{ - ID: auth.ID, - Provider: strings.ToLower(strings.TrimSpace(auth.Provider)), - Priority: authPriority(auth), - Status: string(auth.Status), - Attributes: schedulerSafeAttributes(auth.Attributes), - }) - } - return out -} - -func schedulerProviders(provider string, providers []string) []string { - out := make([]string, 0, len(providers)+1) - seen := make(map[string]struct{}, len(providers)+1) - addProvider := func(value string) { - value = strings.ToLower(strings.TrimSpace(value)) - if value == "" || value == "mixed" { - return - } - if _, ok := seen[value]; ok { - return - } - seen[value] = struct{}{} - out = append(out, value) - } - addProvider(provider) - for _, value := range providers { - addProvider(value) - } - return out -} - -func schedulerOptions(opts cliproxyexecutor.Options) pluginapi.SchedulerOptions { - return pluginapi.SchedulerOptions{ - Headers: cloneHTTPHeader(opts.Headers), - Metadata: cloneSchedulerAnyMap(opts.Metadata), - } -} - -func pickSchedulerAuthByID(candidates []*Auth, authID string) *Auth { - authID = strings.TrimSpace(authID) - if authID == "" { - return nil - } - for _, candidate := range candidates { - if candidate != nil && candidate.ID == authID { - return candidate - } - } - return nil -} - -func builtinSchedulerStrategy(delegate string) (schedulerStrategy, bool) { - switch strings.TrimSpace(delegate) { - case pluginapi.SchedulerBuiltinRoundRobin: - return schedulerStrategyRoundRobin, true - case pluginapi.SchedulerBuiltinFillFirst: - return schedulerStrategyFillFirst, true - default: - return schedulerStrategyCustom, false - } -} - -func (m *Manager) pickViaBuiltinScheduler(ctx context.Context, strategy schedulerStrategy, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, bool, error) { - if m == nil || m.scheduler == nil { - return nil, false, nil - } - providerKey := strings.ToLower(strings.TrimSpace(provider)) - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - for { - var selected *Auth - var errPick error - if providerKey == "mixed" { - selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy) - } - } else { - selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) - } - } - if errPick != nil { - return nil, true, errPick - } - if selected == nil { - return nil, true, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if disallowFreeAuth && isFreeCodexAuth(selected) { - if tried == nil { - tried = make(map[string]struct{}) - } - tried[selected.ID] = struct{}{} - continue - } - return selected, true, nil - } -} - -func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginScheduler, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}, candidates []*Auth) (*Auth, bool, error) { - if scheduler == nil || len(candidates) == 0 { - return nil, false, nil - } - providerKey := strings.ToLower(strings.TrimSpace(provider)) - requestProvider := providerKey - if providerKey == "mixed" { - requestProvider = "" - } - req := pluginapi.SchedulerPickRequest{ - Provider: requestProvider, - Providers: schedulerProviders(providerKey, providers), - Model: model, - Stream: opts.Stream, - Options: schedulerOptions(opts), - Candidates: schedulerAuthCandidates(candidates), - } - resp, handled, errPick := scheduler.PickAuth(ctx, req) - if errPick != nil { - return nil, true, errPick - } - if !handled || !resp.Handled { - return nil, false, nil - } - if selected := pickSchedulerAuthByID(candidates, resp.AuthID); selected != nil { - return selected, true, nil - } - - strategy, okStrategy := builtinSchedulerStrategy(resp.DelegateBuiltin) - if !okStrategy { - return nil, false, nil - } - return m.pickViaBuiltinScheduler(ctx, strategy, providerKey, providers, model, opts, tried) -} - -func (m *Manager) authSupportsRouteModel(registryRef *registry.ModelRegistry, auth *Auth, routeModel string) bool { - if registryRef == nil || auth == nil { - return true - } - routeKey := canonicalModelKey(routeModel) - if routeKey == "" { - return true - } - if registryRef.ClientSupportsModel(auth.ID, routeKey) { - return true - } - selectionKey := m.selectionModelKeyForAuth(auth, routeModel) - return selectionKey != "" && selectionKey != routeKey && registryRef.ClientSupportsModel(auth.ID, selectionKey) -} - -func discardStreamChunks(ch <-chan cliproxyexecutor.StreamChunk) { - if ch == nil { - return - } - go func() { - for range ch { - } - }() -} - -type streamBootstrapError struct { - cause error - headers http.Header -} - -func cloneHTTPHeader(headers http.Header) http.Header { - if headers == nil { - return nil - } - return headers.Clone() -} - -func newStreamBootstrapError(err error, headers http.Header) error { - if err == nil { - return nil - } - return &streamBootstrapError{ - cause: err, - headers: cloneHTTPHeader(headers), - } -} - -func (e *streamBootstrapError) Error() string { - if e == nil || e.cause == nil { - return "" - } - return e.cause.Error() -} - -func (e *streamBootstrapError) Unwrap() error { - if e == nil { - return nil - } - return e.cause -} - -func (e *streamBootstrapError) Headers() http.Header { - if e == nil { - return nil - } - return cloneHTTPHeader(e.headers) -} - -func streamErrorResult(headers http.Header, err error) *cliproxyexecutor.StreamResult { - ch := make(chan cliproxyexecutor.StreamChunk, 1) - ch <- cliproxyexecutor.StreamChunk{Err: err} - close(ch) - return &cliproxyexecutor.StreamResult{ - Headers: cloneHTTPHeader(headers), - Chunks: ch, - } -} - -func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamChunk) ([]cliproxyexecutor.StreamChunk, bool, error) { - if ch == nil { - return nil, true, nil - } - buffered := make([]cliproxyexecutor.StreamChunk, 0, 1) - for { - var ( - chunk cliproxyexecutor.StreamChunk - ok bool - ) - if ctx != nil { - select { - case <-ctx.Done(): - return nil, false, ctx.Err() - case chunk, ok = <-ch: - } - } else { - chunk, ok = <-ch - } - if !ok { - return buffered, true, nil - } - if chunk.Err != nil { - return nil, false, chunk.Err - } - buffered = append(buffered, chunk) - if len(chunk.Payload) > 0 { - return buffered, false, nil - } - } -} - -func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool) *cliproxyexecutor.StreamResult { - out := make(chan cliproxyexecutor.StreamChunk) - go func() { - defer close(out) - var failed bool - forward := true - var rewriter *StreamRewriter - if aliasResult.ForceMapping && strings.TrimSpace(aliasResult.OriginalAlias) != "" { - rewriter = NewStreamRewriter(StreamRewriteOptions{RewriteModel: aliasResult.OriginalAlias}) - } - emit := func(chunk cliproxyexecutor.StreamChunk) bool { - if chunk.Err != nil && !failed { - failed = true - rerr := resultErrorFromError(chunk.Err) - m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}, auth, ephemeralResult) - } - if !forward { - return false - } - if chunk.Err != nil { - if ctx == nil { - out <- chunk - return true - } - select { - case <-ctx.Done(): - forward = false - return false - case out <- chunk: - return true - } - } - if len(chunk.Payload) == 0 { - return true - } - payload := rewriteForceMappedStreamChunk(rewriter, chunk.Payload) - if len(payload) == 0 { - return true - } - chunk.Payload = payload - if ctx == nil { - out <- chunk - return true - } - select { - case <-ctx.Done(): - forward = false - return false - case out <- chunk: - return true - } - } - for _, chunk := range buffered { - if ok := emit(chunk); !ok { - discardStreamChunks(remaining) - return - } - } - for chunk := range remaining { - if ok := emit(chunk); !ok { - discardStreamChunks(remaining) - return - } - } - if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 { - tailChunk := cliproxyexecutor.StreamChunk{Payload: tail} - if !emit(tailChunk) { - return - } - } - if !failed { - m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}, auth, ephemeralResult) - } - }() - return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} -} - -func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult, allowRetry bool, ephemeralResult bool) (*cliproxyexecutor.StreamResult, error) { - if executor == nil { - return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} - } - ctx = contextWithRequestedModelAlias(ctx, opts, routeModel) - var lastErr error - didRefreshOnUnauthorized := false - for idx, execModel := range execModels { - resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) - execReq := req - execReq.Model = execModel - if executionModel != "" { - execReq.Model = executionModel - } - execOpts := opts - execReq, execOpts = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) - if errCtx := ctx.Err(); errCtx != nil { - return nil, errCtx - } - streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) - if errStream != nil { - if errCtx := ctx.Err(); errCtx != nil { - return nil, errCtx - } - if allowRetry { - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, errStream, didRefreshOnUnauthorized); okRefresh { - auth = refreshed - didRefreshOnUnauthorized = true - streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) - if errStream != nil { - if errCtx := ctx.Err(); errCtx != nil { - return nil, errCtx - } - } - } - } - } - if errStream == nil && (streamResult == nil || streamResult.Chunks == nil) { - errStream = &Error{Code: "empty_stream", Message: "upstream stream has no source", Retryable: true} - } - if errStream != nil { - rerr := resultErrorFromError(errStream) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} - result.RetryAfter = retryAfterFromError(errStream) - m.recordExecutionResult(ctx, result, auth, ephemeralResult) - if isRequestInvalidError(errStream) { - return nil, errStream - } - lastErr = errStream - continue - } - - buffered, closed, bootstrapErr := readStreamBootstrap(ctx, streamResult.Chunks) - if bootstrapErr != nil { - if errCtx := ctx.Err(); errCtx != nil { - discardStreamChunks(streamResult.Chunks) - return nil, errCtx - } - if allowRetry { - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, bootstrapErr, didRefreshOnUnauthorized); okRefresh { - discardStreamChunks(streamResult.Chunks) - auth = refreshed - didRefreshOnUnauthorized = true - retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) - if retryErr != nil { - if errCtx := ctx.Err(); errCtx != nil { - return nil, errCtx - } - bootstrapErr = retryErr - streamResult = &cliproxyexecutor.StreamResult{} - } else { - streamResult = retryStream - buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks) - } - } - } - } - if bootstrapErr != nil { - if isRequestInvalidError(bootstrapErr) { - rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} - result.RetryAfter = retryAfterFromError(bootstrapErr) - m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) - return nil, bootstrapErr - } - if idx < len(execModels)-1 { - rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} - result.RetryAfter = retryAfterFromError(bootstrapErr) - m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) - lastErr = bootstrapErr - continue - } - rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} - result.RetryAfter = retryAfterFromError(bootstrapErr) - m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) - return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) - } - - if closed && len(buffered) == 0 { - emptyErr := &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true} - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: emptyErr} - m.recordExecutionResult(ctx, result, auth, ephemeralResult) - if idx < len(execModels)-1 { - lastErr = emptyErr - continue - } - return nil, newStreamBootstrapError(emptyErr, streamResult.Headers) - } - - remaining := streamResult.Chunks - if closed { - closedCh := make(chan cliproxyexecutor.StreamChunk) - close(closedCh) - remaining = closedCh - } - return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, aliasResult, ephemeralResult), nil - } - if lastErr == nil { - lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"} - } - return nil, lastErr -} - -func (m *Manager) rebuildAPIKeyModelAliasFromRuntimeConfig() { - if m == nil { - return - } - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - if cfg == nil { - cfg = &internalconfig.Config{} - } - m.mu.Lock() - defer m.mu.Unlock() - m.rebuildAPIKeyModelAliasLocked(cfg) -} - -// RefreshAPIKeyModelAlias rebuilds the API-key model alias table from the current runtime config. -func (m *Manager) RefreshAPIKeyModelAlias() { - m.rebuildAPIKeyModelAliasFromRuntimeConfig() -} - -func (m *Manager) rebuildAPIKeyModelAliasLocked(cfg *internalconfig.Config) { - if m == nil { - return - } - if cfg == nil { - cfg = &internalconfig.Config{} - } - - out := make(apiKeyModelAliasTable) - for _, auth := range m.auths { - if auth == nil { - continue - } - if strings.TrimSpace(auth.ID) == "" { - continue - } - if auth.AuthKind() != AuthKindAPIKey { - continue - } - - byAlias := make(map[string]string) - provider := strings.ToLower(strings.TrimSpace(auth.Provider)) - switch provider { - case "gemini": - if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil { - compileAPIKeyModelAliasForModels(byAlias, entry.Models) - } - case "gemini-interactions": - if entry := resolveInteractionsAPIKeyConfig(cfg, auth); entry != nil { - compileAPIKeyModelAliasForModels(byAlias, entry.Models) - } - case "claude": - if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil { - compileAPIKeyModelAliasForModels(byAlias, entry.Models) - } - case "codex": - if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil { - compileAPIKeyModelAliasForModels(byAlias, entry.Models) - } - case "xai": - if entry := resolveXAIAPIKeyConfig(cfg, auth); entry != nil { - compileAPIKeyModelAliasForModels(byAlias, entry.Models) - } - case "vertex": - if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil { - compileAPIKeyModelAliasForModels(byAlias, entry.Models) - } - default: - // OpenAI-compat uses config selection from auth.Attributes. - providerKey := "" - compatName := "" - if auth.Attributes != nil { - providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) - compatName = strings.TrimSpace(auth.Attributes["compat_name"]) - } - if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { - if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil { - compileAPIKeyModelAliasForModels(byAlias, entry.Models) - } - } - } - - if len(byAlias) > 0 { - out[auth.ID] = byAlias - } - } - - m.apiKeyModelAlias.Store(out) -} - -func compileAPIKeyModelAliasForModels[T interface { - GetName() string - GetAlias() string -}](out map[string]string, models []T) { - if out == nil { - return - } - for i := range models { - alias := strings.TrimSpace(models[i].GetAlias()) - name := strings.TrimSpace(models[i].GetName()) - if alias == "" || name == "" { - continue - } - aliasKey := strings.ToLower(thinking.ParseSuffix(alias).ModelName) - if aliasKey == "" { - aliasKey = strings.ToLower(alias) - } - // Config priority: first alias wins. - if _, exists := out[aliasKey]; exists { - continue - } - out[aliasKey] = name - // Also allow direct lookup by upstream name (case-insensitive), so lookups on already-upstream - // models remain a cheap no-op. - nameKey := strings.ToLower(thinking.ParseSuffix(name).ModelName) - if nameKey == "" { - nameKey = strings.ToLower(name) - } - if nameKey != "" { - if _, exists := out[nameKey]; !exists { - out[nameKey] = name - } - } - // Preserve config suffix priority by seeding a base-name lookup when name already has suffix. - nameResult := thinking.ParseSuffix(name) - if nameResult.HasSuffix { - baseKey := strings.ToLower(strings.TrimSpace(nameResult.ModelName)) - if baseKey != "" { - if _, exists := out[baseKey]; !exists { - out[baseKey] = name - } - } - } - } -} - -// SetRetryConfig updates retry attempts, credential retry limit and cooldown wait interval. -func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration, maxRetryCredentials int) { - if m == nil { - return - } - if retry < 0 { - retry = 0 - } - if maxRetryCredentials < 0 { - maxRetryCredentials = 0 - } - if maxRetryInterval < 0 { - maxRetryInterval = 0 - } - m.requestRetry.Store(int32(retry)) - m.maxRetryCredentials.Store(int32(maxRetryCredentials)) - m.maxRetryInterval.Store(maxRetryInterval.Nanoseconds()) -} - -// RegisterExecutor registers a provider executor with the manager. -func (m *Manager) RegisterExecutor(executor ProviderExecutor) { - if executor == nil { - return - } - provider := strings.TrimSpace(executor.Identifier()) - if provider == "" { - return - } - - var replaced ProviderExecutor - m.mu.Lock() - replaced = m.executors[provider] - m.executors[provider] = executor - m.mu.Unlock() - - if replaced == nil || replaced == executor { - return - } - if closer, ok := replaced.(ExecutionSessionCloser); ok && closer != nil { - closer.CloseExecutionSession(CloseAllExecutionSessionsID) - } -} - -// UnregisterExecutor removes the executor associated with the provider key. -func (m *Manager) UnregisterExecutor(provider string) { - provider = strings.ToLower(strings.TrimSpace(provider)) - if provider == "" { - return - } - m.mu.Lock() - delete(m.executors, provider) - m.mu.Unlock() -} - -// Register inserts a new auth entry into the manager. -func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { - if auth == nil { - return nil, nil - } - if auth.ID == "" { - auth.ID = uuid.NewString() - } - now := time.Now() - clearedCooldown := false - if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { - clearedCooldown = clearCooldownStateForAuth(auth, now) - } - auth.EnsureIndex() - authClone := auth.Clone() - m.mu.Lock() - m.auths[auth.ID] = authClone - m.mu.Unlock() - if !shouldDeferAPIKeyModelAliasRebuild(ctx) { - m.rebuildAPIKeyModelAliasFromRuntimeConfig() - } - if m.scheduler != nil { - m.scheduler.upsertAuth(authClone) - } - m.queueRefreshReschedule(auth.ID) - _ = m.persist(ctx, auth) - m.hook.OnAuthRegistered(ctx, auth.Clone()) - if clearedCooldown { - m.persistCooldownStates(ctx) - } - return auth.Clone(), nil -} - -// Update replaces an existing auth entry and notifies hooks. -func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { - if auth == nil || auth.ID == "" { - return nil, nil - } - m.mu.Lock() - existing, ok := m.auths[auth.ID] - if !ok || existing == nil { - m.mu.Unlock() - return nil, nil - } - if !auth.indexAssigned && auth.Index == "" { - auth.Index = existing.Index - auth.indexAssigned = existing.indexAssigned - } - auth.Success = existing.Success - auth.Failed = existing.Failed - auth.recentRequests = existing.recentRequests - if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { - if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { - auth.ModelStates = existing.ModelStates - } - } - now := time.Now() - clearedCooldown := false - if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { - clearedCooldown = clearCooldownStateForAuth(auth, now) - } - auth.EnsureIndex() - authClone := auth.Clone() - m.auths[auth.ID] = authClone - m.mu.Unlock() - if !shouldDeferAPIKeyModelAliasRebuild(ctx) { - m.rebuildAPIKeyModelAliasFromRuntimeConfig() - } - if m.scheduler != nil { - m.scheduler.upsertAuth(authClone) - } - m.queueRefreshReschedule(auth.ID) - _ = m.persist(ctx, auth) - m.hook.OnAuthUpdated(ctx, auth.Clone()) - if clearedCooldown { - m.persistCooldownStates(ctx) - } - return auth.Clone(), nil -} - -// Remove deletes an auth from runtime state without persisting. -// Disk and token-store deletion must be handled by the caller. -func (m *Manager) Remove(ctx context.Context, id string) { - if m == nil { - return - } - id = strings.TrimSpace(id) - if id == "" { - return - } - _ = ctx - - m.mu.Lock() - existing := m.auths[id] - if existing == nil { - m.mu.Unlock() - return - } - provider := strings.TrimSpace(existing.Provider) - delete(m.auths, id) - if m.modelPoolOffsets != nil { - delete(m.modelPoolOffsets, id) - } - for sessionID, sessionAuths := range m.homeRuntimeAuths { - if sessionAuths == nil { - continue - } - delete(sessionAuths, id) - if len(sessionAuths) == 0 { - delete(m.homeRuntimeAuths, sessionID) - } - } - m.mu.Unlock() - - if !shouldDeferAPIKeyModelAliasRebuild(ctx) { - m.rebuildAPIKeyModelAliasFromRuntimeConfig() - } - if m.scheduler != nil { - m.scheduler.removeAuth(id) - } - m.queueRefreshUnschedule(id) - m.invalidateSessionAffinity(id) - - if provider != "" { - if exec, ok := m.Executor(provider); ok && exec != nil { - if closer, okCloser := exec.(ExecutionSessionCloser); okCloser { - closer.CloseExecutionSession(CloseAllExecutionSessionsID) - } - } - } - m.persistCooldownStates(ctx) -} - -func (m *Manager) invalidateSessionAffinity(authID string) { - if m == nil || authID == "" { - return - } - if invalidator, ok := m.selector.(interface{ InvalidateAuth(string) }); ok && invalidator != nil { - invalidator.InvalidateAuth(authID) - } -} - -// Load resets manager state from the backing store. -func (m *Manager) Load(ctx context.Context) error { - m.mu.Lock() - if m.store == nil { - m.mu.Unlock() - return nil - } - items, err := m.store.List(ctx) - if err != nil { - m.mu.Unlock() - return err - } - m.auths = make(map[string]*Auth, len(items)) - for _, auth := range items { - if auth == nil || auth.ID == "" { - continue - } - auth.EnsureIndex() - m.auths[auth.ID] = auth.Clone() - } - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - if cfg == nil { - cfg = &internalconfig.Config{} - } - m.rebuildAPIKeyModelAliasLocked(cfg) - m.mu.Unlock() - m.syncScheduler() - return nil -} - -// Execute performs a non-streaming execution using the configured selector and executor. -// It supports multiple providers for the same model and round-robins the starting provider per model. -func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - req, opts = cliproxysession.Enrich(req, opts) - normalized := m.normalizeProviders(providers) - if len(normalized) == 0 { - return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} - } - if m.HomeEnabled() { - return m.executeHome(ctx, normalized, req, opts, false) - } - - _, maxRetryCredentials, maxWait := m.retrySettings() - - var lastErr error - retryModel := authSelectionModelFromOptions(opts, req.Model) - for attempt := 0; ; attempt++ { - resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) - if errExec == nil { - return resp, nil - } - lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) - if !shouldRetry { - break - } - if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { - return cliproxyexecutor.Response{}, errWait - } - } - if lastErr != nil { - if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { - if resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil { - return cliproxyexecutor.Response{}, errCredits - } else if ok { - return resp, nil - } - } - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} -} - -// It supports multiple providers for the same model and round-robins the starting provider per model. -func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - req, opts = cliproxysession.Enrich(req, opts) - normalized := m.normalizeProviders(providers) - if len(normalized) == 0 { - return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} - } - if m.HomeEnabled() { - return m.executeHome(ctx, normalized, req, opts, true) - } - - _, maxRetryCredentials, maxWait := m.retrySettings() - - var lastErr error - retryModel := authSelectionModelFromOptions(opts, req.Model) - for attempt := 0; ; attempt++ { - resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) - if errExec == nil { - return resp, nil - } - lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) - if !shouldRetry { - break - } - if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { - return cliproxyexecutor.Response{}, errWait - } - } - if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} -} - -// ExecuteStream performs a streaming execution using the configured selector and executor. -// It supports multiple providers for the same model and round-robins the starting provider per model. -func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { - req, opts = cliproxysession.Enrich(req, opts) - if m.HomeEnabled() { - if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { - defer unlockSession() - } - } - normalized := m.normalizeProviders(providers) - if len(normalized) == 0 { - return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} - } - - _, maxRetryCredentials, maxWait := m.retrySettings() - - var lastErr error - retryModel := authSelectionModelFromOptions(opts, req.Model) - for attempt := 0; ; attempt++ { - result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) - if errStream == nil { - return result, nil - } - lastErr = errStream - wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, retryModel, maxWait) - if !shouldRetry { - break - } - if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { - return nil, errWait - } - } - if lastErr != nil { - if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { - if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil { - return nil, errCredits - } else if ok { - return result, nil - } - } - var bootstrapErr *streamBootstrapError - if errors.As(lastErr, &bootstrapErr) && bootstrapErr != nil { - return streamErrorResult(bootstrapErr.Headers(), bootstrapErr.cause), nil - } - return nil, lastErr - } - return nil, &Error{Code: "auth_not_found", Message: "no auth available"} -} - -func (m *Manager) executeHome(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool) (cliproxyexecutor.Response, error) { - if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { - defer unlockSession() - } - routeModel := authSelectionModelFromOptions(opts, req.Model) - responseAlias := requestedModelAliasFromOptions(opts, routeModel) - executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) - opts = ensureRequestedModelMetadata(opts, routeModel) - tried := make(map[string]struct{}) - var lastErr error - for homeAuthCount := 1; ; homeAuthCount++ { - selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, withHomeAuthCount(opts, homeAuthCount)) - if errSelection != nil { - if lastErr != nil && isHomeRequestRetryExceededError(errSelection) { - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, errSelection - } - auth := selection.CloneAuthForRoute(routeModel) - if auth == nil || selection.Executor == nil { - selection.End("missing_execution_target") - return cliproxyexecutor.Response{}, &Error{Code: "executor_not_found", Message: "executor not registered"} - } - if _, seen := tried[auth.ID]; seen { - selection.End("repeated_auth") - if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, repeatedHomeAuthError() - } - entry := logEntryWithRequestID(ctx) - debugLogAuthSelection(entry, auth, selection.Provider, routeModel) - if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil { - selection.End("runtime_auth_bind_failed") - return cliproxyexecutor.Response{}, errRuntimeAuth - } - publishSelectedAuthMetadata(opts.Metadata, auth) - tried[auth.ID] = struct{}{} - execCtx, releaseAttempt, errBind := homeExecutionAttemptContext(ctx, selection) - if errBind != nil { - selection.End("attempt_bind_failed") - return cliproxyexecutor.Response{}, errBind - } - if rt := m.roundTripperFor(auth); rt != nil { - execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) - execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) - } - models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) - if aliasResult.ForceMapping && responseAlias != "" { - aliasResult.OriginalAlias = responseAlias - } - if len(models) > 1 { - models = models[:1] - pooled = false - } - if len(models) == 0 { - releaseAttempt() - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil { - return cliproxyexecutor.Response{}, errEnd - } - lastErr = &Error{Code: "auth_not_found", Message: "no execution models available"} - continue - } - preparedAuth, errPrepare := m.prepareHomeRequestAuth(execCtx, selection.Executor, selection) - if errPrepare != nil { - m.reportHomeResult(execCtx, Result{AuthID: auth.ID, Provider: selection.Provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}, auth) - releaseAttempt() - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { - return cliproxyexecutor.Response{}, errEnd - } - lastErr = errPrepare - continue - } - for _, upstreamModel := range models { - resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled) - execReq := req - execReq.Model = upstreamModel - if restoreExecutionModel { - execReq.Model = executionModel - } - execOpts := opts - execOpts.ExecutionLifecycle = selection - execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, selection.Executor, selection.Provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) - if errCtx := execCtx.Err(); errCtx != nil { - releaseAttempt() - selection.End("attempt_canceled") - return cliproxyexecutor.Response{}, errCtx - } - var response cliproxyexecutor.Response - var errExecute error - if countTokens { - response, errExecute = selection.Executor.CountTokens(execCtx, preparedAuth, execReq, execOpts) - } else { - response, errExecute = selection.Executor.Execute(execCtx, preparedAuth, execReq, execOpts) - } - result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil} - if errExecute == nil { - m.reportHomeResult(execCtx, result, preparedAuth) - releaseAttempt() - rewriteForceMappedResponse(&response, aliasResult) - if !m.retainHomeWebsocketSelection(ctx, opts, routeModel, selection) { - selection.End("completed") - } - return response, nil - } - result.Error = resultErrorFromError(errExecute) - result.RetryAfter = retryAfterFromError(errExecute) - m.reportHomeResult(execCtx, result, preparedAuth) - lastErr = errExecute - if isRequestInvalidError(errExecute) { - releaseAttempt() - selection.End("request_invalid") - return cliproxyexecutor.Response{}, errExecute - } - } - releaseAttempt() - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "execution_failed"); errEnd != nil { - return cliproxyexecutor.Response{}, errEnd - } - if errCtx := execCtx.Err(); errCtx != nil && ctx != nil && ctx.Err() != nil { - return cliproxyexecutor.Response{}, errCtx - } - } -} - -type requestToFormatResolver interface { - RequestToFormat(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format -} - -func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExecutor, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, requestedModel string) (cliproxyexecutor.Request, cliproxyexecutor.Options) { - if opts.RequestAfterAuthInterceptor == nil { - return req, opts - } - toFormat := requestToFormat(provider, executor, req, opts) - resp := opts.RequestAfterAuthInterceptor(ctx, cliproxyexecutor.RequestAfterAuthInterceptRequest{ - SourceFormat: opts.SourceFormat, - ToFormat: toFormat, - Model: req.Model, - RequestedModel: requestedModel, - Stream: opts.Stream, - Headers: cloneRequestHeaders(opts.Headers), - Body: bytes.Clone(req.Payload), - Metadata: opts.Metadata, - }) - opts.Headers = mergeRequestHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) - if len(resp.Body) > 0 { - req.Payload = bytes.Clone(resp.Body) - opts.OriginalRequest = bytes.Clone(resp.Body) - } - return req, opts -} - -func requestToFormat(provider string, executor ProviderExecutor, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format { - resolver, ok := executor.(requestToFormatResolver) - if ok && resolver != nil { - formatRequestTo := resolver.RequestToFormat(req, opts) - if formatRequestTo != "" { - return formatRequestTo - } - } - source := opts.SourceFormat.String() - if source == "openai-image" || source == "openai-video" { - return opts.SourceFormat - } - if opts.Alt == "responses/compact" && !opts.Stream { - return sdktranslator.FormatOpenAIResponse - } - switch strings.ToLower(strings.TrimSpace(provider)) { - case "codex": - return sdktranslator.FormatCodex - case "xai": - return sdktranslator.FormatCodex - case "claude": - return sdktranslator.FormatClaude - case "gemini", "vertex", "aistudio": - return sdktranslator.FormatGemini - case "kimi": - return sdktranslator.FormatOpenAI - case "antigravity": - return sdktranslator.FormatAntigravity - default: - return sdktranslator.FormatOpenAI - } -} - -func cloneRequestHeaders(src http.Header) http.Header { - if src == nil { - return nil - } - dst := make(http.Header, len(src)) - for key, values := range src { - dst[key] = append([]string(nil), values...) - } - return dst -} - -func mergeRequestHeaders(current, updates http.Header, clear []string) http.Header { - if updates == nil && len(clear) == 0 { - return current - } - out := cloneRequestHeaders(current) - if out == nil && (len(updates) > 0 || len(clear) > 0) { - out = make(http.Header) - } - for _, key := range clear { - out.Del(key) - } - for key, values := range updates { - out.Del(key) - for _, value := range values { - out.Add(key, value) - } - } - return out -} - -func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { - if len(providers) == 0 { - return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} - } - routeModel := authSelectionModelFromOptions(opts, req.Model) - executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) - opts = ensureRequestedModelMetadata(opts, routeModel) - homeMode := m.HomeEnabled() - homeAuthCount := 1 - tried := make(map[string]struct{}) - attempted := make(map[string]struct{}) - var lastErr error - for { - if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { - if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} - } - pickOpts := opts - if homeMode { - pickOpts = withHomeAuthCount(opts, homeAuthCount) - } - auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) - if errPick != nil { - if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, errPick - } - - entry := logEntryWithRequestID(ctx) - debugLogAuthSelection(entry, auth, provider, routeModel) - publishSelectedAuthMetadata(opts.Metadata, auth) - - tried[auth.ID] = struct{}{} - execCtx := ctx - if rt := m.roundTripperFor(auth); rt != nil { - execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) - execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) - } - execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) - - models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) - if len(models) == 0 { - continue - } - attempted[auth.ID] = struct{}{} - var errPrepare error - auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) - if errPrepare != nil { - result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} - m.MarkResult(execCtx, result) - lastErr = errPrepare - continue - } - var authErr error - didRefreshOnUnauthorized := false - for _, upstreamModel := range models { - resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) - execReq := req - execReq.Model = upstreamModel - if restoreExecutionModel { - execReq.Model = executionModel - } - execOpts := opts - execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) - resp, errExec := executor.Execute(execCtx, auth, execReq, execOpts) - if errExec != nil { - if errCtx := execCtx.Err(); errCtx != nil { - return cliproxyexecutor.Response{}, errCtx - } - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { - auth = refreshed - didRefreshOnUnauthorized = true - resp, errExec = executor.Execute(execCtx, auth, execReq, execOpts) - if errExec != nil { - if errCtx := execCtx.Err(); errCtx != nil { - return cliproxyexecutor.Response{}, errCtx - } - } - } - } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} - if errExec != nil { - result.Error = resultErrorFromError(errExec) - if ra := retryAfterFromError(errExec); ra != nil { - result.RetryAfter = ra - } - m.MarkResult(execCtx, result) - if isRequestInvalidError(errExec) { - return cliproxyexecutor.Response{}, errExec - } - authErr = errExec - continue - } - m.MarkResult(execCtx, result) - rewriteForceMappedResponse(&resp, aliasResult) - return resp, nil - } - if authErr != nil { - if isRequestInvalidError(authErr) { - return cliproxyexecutor.Response{}, authErr - } - lastErr = authErr - if homeMode { - homeAuthCount++ - } - continue - } - } -} - -func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { - if len(providers) == 0 { - return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} - } - routeModel := authSelectionModelFromOptions(opts, req.Model) - executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) - opts = ensureRequestedModelMetadata(opts, routeModel) - homeMode := m.HomeEnabled() - homeAuthCount := 1 - tried := make(map[string]struct{}) - attempted := make(map[string]struct{}) - var lastErr error - for { - if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { - if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} - } - pickOpts := opts - if homeMode { - pickOpts = withHomeAuthCount(opts, homeAuthCount) - } - auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) - if errPick != nil { - if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { - return cliproxyexecutor.Response{}, lastErr - } - return cliproxyexecutor.Response{}, errPick - } - - entry := logEntryWithRequestID(ctx) - debugLogAuthSelection(entry, auth, provider, routeModel) - publishSelectedAuthMetadata(opts.Metadata, auth) - - tried[auth.ID] = struct{}{} - execCtx := ctx - if rt := m.roundTripperFor(auth); rt != nil { - execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) - execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) - } - execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) - - models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) - if len(models) == 0 { - continue - } - attempted[auth.ID] = struct{}{} - var errPrepare error - auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) - if errPrepare != nil { - result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} - m.MarkResult(execCtx, result) - lastErr = errPrepare - continue - } - var authErr error - didRefreshOnUnauthorized := false - for _, upstreamModel := range models { - resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) - execReq := req - execReq.Model = upstreamModel - if restoreExecutionModel { - execReq.Model = executionModel - } - execOpts := opts - execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) - resp, errExec := executor.CountTokens(execCtx, auth, execReq, execOpts) - if errExec != nil { - if errCtx := execCtx.Err(); errCtx != nil { - return cliproxyexecutor.Response{}, errCtx - } - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { - auth = refreshed - didRefreshOnUnauthorized = true - resp, errExec = executor.CountTokens(execCtx, auth, execReq, execOpts) - if errExec != nil { - if errCtx := execCtx.Err(); errCtx != nil { - return cliproxyexecutor.Response{}, errCtx - } - } - } - } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} - if errExec != nil { - result.Error = resultErrorFromError(errExec) - if ra := retryAfterFromError(errExec); ra != nil { - result.RetryAfter = ra - } - // Some Anthropic-compatible upstreams do not implement the - // count_tokens route and return a generic endpoint 404. Record - // the failure for hooks and metrics without suspending a model - // that remains usable through the messages endpoint. - if isCountTokensEndpointNotFoundError(errExec, execReq.Model) { - m.recordAvailabilityNeutralResult(execCtx, result) - } else { - m.MarkResult(execCtx, result) - } - if isRequestInvalidError(errExec) { - return cliproxyexecutor.Response{}, errExec - } - authErr = errExec - continue - } - m.MarkResult(execCtx, result) - rewriteForceMappedResponse(&resp, aliasResult) - return resp, nil - } - if authErr != nil { - if isRequestInvalidError(authErr) { - return cliproxyexecutor.Response{}, authErr - } - lastErr = authErr - if homeMode { - homeAuthCount++ - } - continue - } - } -} - -func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (*cliproxyexecutor.StreamResult, error) { - if len(providers) == 0 { - return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} - } - routeModel := authSelectionModelFromOptions(opts, req.Model) - responseAlias := requestedModelAliasFromOptions(opts, routeModel) - executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) - opts = ensureRequestedModelMetadata(opts, routeModel) - homeMode := m.HomeEnabled() - homeAuthCount := 1 - tried := make(map[string]struct{}) - attempted := make(map[string]struct{}) - var lastErr error - for { - if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { - if lastErr != nil { - return nil, lastErr - } - return nil, &Error{Code: "auth_not_found", Message: "no auth available"} - } - pickOpts := opts - if homeMode { - pickOpts = withHomeAuthCount(opts, homeAuthCount) - } - - var selection *HomeDispatchSelection - var auth *Auth - var executor ProviderExecutor - var provider string - var errPick error - if homeMode { - selection, errPick = m.pickHomeDispatchSelection(ctx, routeModel, pickOpts) - if selection != nil { - auth = selection.CloneAuthForRoute(routeModel) - executor = selection.Executor - provider = selection.Provider - } - } else { - auth, executor, provider, errPick = m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) - } - if errPick != nil { - if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { - return nil, lastErr - } - return nil, errPick - } - if auth == nil || executor == nil { - if selection != nil { - selection.End("missing_execution_target") - } - return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} - } - - entry := logEntryWithRequestID(ctx) - debugLogAuthSelection(entry, auth, provider, routeModel) - if selection != nil { - if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil { - selection.End("runtime_auth_bind_failed") - return nil, errRuntimeAuth - } - } - publishSelectedAuthMetadata(opts.Metadata, auth) - - tried[auth.ID] = struct{}{} - execCtx := ctx - releaseAttempt := func() {} - if selection != nil { - var errBind error - execCtx, releaseAttempt, errBind = homeExecutionAttemptContext(ctx, selection) - if errBind != nil { - selection.End("attempt_bind_failed") - return nil, errBind - } - } - if rt := m.roundTripperFor(auth); rt != nil { - execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) - execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) - } - models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) - if selection != nil && aliasResult.ForceMapping && responseAlias != "" { - aliasResult.OriginalAlias = responseAlias - } - if len(models) == 0 { - if selection != nil { - releaseAttempt() - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil { - return nil, errEnd - } - } - continue - } - attempted[auth.ID] = struct{}{} - var errPrepare error - if selection != nil { - auth, errPrepare = m.prepareHomeRequestAuth(execCtx, executor, selection) - } else { - auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) - } - if errPrepare != nil { - result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} - if selection != nil { - m.reportHomeResult(execCtx, result, auth) - releaseAttempt() - } else { - m.MarkResult(execCtx, result) - } - lastErr = errPrepare - if selection != nil { - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { - return nil, errEnd - } - } - continue - } - execReq := sanitizeDownstreamWebsocketFallbackRequest(execCtx, auth, req) - streamExecutionModel := "" - if restoreExecutionModel { - streamExecutionModel = executionModel - } - execOpts := opts - if selection != nil { - execOpts.ExecutionLifecycle = selection - } - if homeMode && len(models) > 1 { - models = models[:1] - pooled = false - } - streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, !homeMode, selection != nil) - if errStream != nil { - if selection != nil { - releaseAttempt() - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "stream_start_failed"); errEnd != nil { - return nil, errEnd - } - } - if errCtx := execCtx.Err(); errCtx != nil && ctx != nil && ctx.Err() != nil { - return nil, errCtx - } - if isRequestInvalidError(errStream) { - return nil, errStream - } - lastErr = errStream - if homeMode { - homeAuthCount++ - } - continue - } - if selection != nil { - if m.retainHomeWebsocketSelection(ctx, opts, routeModel, selection) { - return wrapHomeStream(ctx, streamResult, nil, releaseAttempt), nil - } - return wrapHomeStream(ctx, streamResult, selection, releaseAttempt), nil - } - return streamResult, nil - } -} - -func homeExecutionAttemptContext(ctx context.Context, selection *HomeDispatchSelection) (context.Context, func(), error) { - if selection == nil { - return nil, func() {}, fmt.Errorf("Home dispatch selection is nil") - } - return selection.AttemptContext(ctx) -} - -func wrapHomeStream(ctx context.Context, result *cliproxyexecutor.StreamResult, selection *HomeDispatchSelection, releaseAttempt func()) *cliproxyexecutor.StreamResult { - if result == nil || result.Chunks == nil { - if releaseAttempt != nil { - releaseAttempt() - } - return result - } - out := make(chan cliproxyexecutor.StreamChunk) - go func() { - defer close(out) - if releaseAttempt != nil { - defer releaseAttempt() - } - if selection != nil { - defer selection.End("stream_closed") - } - forward := true - for { - select { - case <-ctx.Done(): - return - case chunk, ok := <-result.Chunks: - if !ok { - return - } - if !forward { - continue - } - select { - case <-ctx.Done(): - return - case out <- chunk: - } - if chunk.Err != nil && selection != nil { - forward = false - } - } - } - }() - return &cliproxyexecutor.StreamResult{Headers: result.Headers, Chunks: out} -} - -func sanitizeDownstreamWebsocketFallbackRequest(ctx context.Context, auth *Auth, req cliproxyexecutor.Request) cliproxyexecutor.Request { - if !cliproxyexecutor.DownstreamWebsocket(ctx) || authWebsocketsEnabled(auth) || len(req.Payload) == 0 { - return req - } - updated, errDelete := sjson.DeleteBytes(req.Payload, "generate") - if errDelete != nil { - return req - } - req.Payload = updated - return req -} - -func ensureRequestedModelMetadata(opts cliproxyexecutor.Options, requestedModel string) cliproxyexecutor.Options { - requestedModel = strings.TrimSpace(requestedModel) - if requestedModel == "" { - return opts - } - if hasRequestedModelMetadata(opts.Metadata) { - return opts - } - if len(opts.Metadata) == 0 { - opts.Metadata = map[string]any{cliproxyexecutor.RequestedModelMetadataKey: requestedModel} - return opts - } - meta := make(map[string]any, len(opts.Metadata)+1) - for k, v := range opts.Metadata { - meta[k] = v - } - meta[cliproxyexecutor.RequestedModelMetadataKey] = requestedModel - opts.Metadata = meta - return opts -} - -func authSelectionModelFromOptions(opts cliproxyexecutor.Options, fallback string) string { - fallback = strings.TrimSpace(fallback) - if len(opts.Metadata) == 0 { - return fallback - } - raw, ok := opts.Metadata[cliproxyexecutor.AuthSelectionModelMetadataKey] - if !ok || raw == nil { - return fallback - } - switch value := raw.(type) { - case string: - if strings.TrimSpace(value) != "" { - return strings.TrimSpace(value) - } - case []byte: - if strings.TrimSpace(string(value)) != "" { - return strings.TrimSpace(string(value)) - } - } - return fallback -} - -func executionModelForAuthSelection(opts cliproxyexecutor.Options, model string) (string, bool) { - model = strings.TrimSpace(model) - if model == "" { - return "", false - } - selectionModel := authSelectionModelFromOptions(opts, model) - if selectionModel == model { - return "", false - } - return model, true -} - -func withHomeAuthCount(opts cliproxyexecutor.Options, count int) cliproxyexecutor.Options { - if count <= 0 { - count = 1 - } - meta := make(map[string]any, len(opts.Metadata)+1) - for k, v := range opts.Metadata { - meta[k] = v - } - meta[homeAuthCountMetadataKey] = count - opts.Metadata = meta - return opts -} - -func homeAuthCountFromMetadata(meta map[string]any) int { - if len(meta) == 0 { - return 1 - } - switch value := meta[homeAuthCountMetadataKey].(type) { - case int: - if value > 0 { - return value - } - case int64: - if value > 0 { - return int(value) - } - case float64: - if value > 0 { - return int(value) - } - } - return 1 -} - -func hasRequestedModelMetadata(meta map[string]any) bool { - if len(meta) == 0 { - return false - } - raw, ok := meta[cliproxyexecutor.RequestedModelMetadataKey] - if !ok || raw == nil { - return false - } - switch v := raw.(type) { - case string: - return strings.TrimSpace(v) != "" - case []byte: - return strings.TrimSpace(string(v)) != "" - default: - return false - } -} - -type requestAuthPrepareLock struct { - mu sync.Mutex -} - -// prepareHomeRequestAuth prepares a dispatch auth without reading or updating local auth state. -func (m *Manager) prepareHomeRequestAuth(ctx context.Context, executor ProviderExecutor, selection *HomeDispatchSelection) (*Auth, error) { - if m == nil || executor == nil || selection == nil { - return nil, nil - } - auth := selection.CloneAuth() - if auth == nil { - return nil, nil - } - preparer, ok := executor.(RequestAuthPreparer) - if !ok || preparer == nil || !preparer.ShouldPrepareRequestAuth(auth) { - return auth, nil - } - - prepare := func() (*Auth, error) { - target := auth.Clone() - if !preparer.ShouldPrepareRequestAuth(target) { - return target, nil - } - updated, errPrepare := preparer.PrepareRequestAuth(ctx, target) - if errPrepare != nil { - return auth, errPrepare - } - if updated == nil { - return target, nil - } - return updated, nil - } - - id := strings.TrimSpace(auth.ID) - if id == "" { - return prepare() - } - lockValue, _ := m.requestPrepareLocks.LoadOrStore(id, &requestAuthPrepareLock{}) - lock, ok := lockValue.(*requestAuthPrepareLock) - if !ok || lock == nil { - return prepare() - } - lock.mu.Lock() - defer lock.mu.Unlock() - return prepare() -} - -func (m *Manager) prepareRequestAuth(ctx context.Context, executor ProviderExecutor, auth *Auth) (*Auth, error) { - if m == nil || executor == nil || auth == nil { - return auth, nil - } - preparer, ok := executor.(RequestAuthPreparer) - if !ok || preparer == nil || !preparer.ShouldPrepareRequestAuth(auth) { - return auth, nil - } - - id := strings.TrimSpace(auth.ID) - if id == "" { - return preparer.PrepareRequestAuth(ctx, auth.Clone()) - } - - lockValue, _ := m.requestPrepareLocks.LoadOrStore(id, &requestAuthPrepareLock{}) - lock, ok := lockValue.(*requestAuthPrepareLock) - if !ok || lock == nil { - return preparer.PrepareRequestAuth(ctx, auth.Clone()) - } - - lock.mu.Lock() - defer lock.mu.Unlock() - - target := auth.Clone() - m.mu.RLock() - if current := m.auths[id]; current != nil { - target = current.Clone() - } - m.mu.RUnlock() - - if !preparer.ShouldPrepareRequestAuth(target) { - return target, nil - } - - updated, errPrepare := preparer.PrepareRequestAuth(ctx, target) - if errPrepare != nil { - return auth, errPrepare - } - if updated == nil { - return target, nil - } - - saved, errUpdate := m.Update(ctx, updated) - if errUpdate != nil { - return updated, errUpdate - } - if saved != nil { - return saved, nil - } - return updated, nil -} - -func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.Options, fallback string) context.Context { - alias := requestedModelAliasFromOptions(opts, fallback) - ctx = coreusage.WithRequestedModelAlias(ctx, alias) - effort := reasoningEffortFromOptions(opts) - if effort != "" { - ctx = coreusage.WithReasoningEffort(ctx, effort) - } - serviceTier := serviceTierFromOptions(opts) - if serviceTier != "" { - ctx = coreusage.WithServiceTier(ctx, serviceTier) - } - if generate, ok := generateFromOptions(opts); ok { - ctx = coreusage.WithGenerate(ctx, generate) - } - return ctx -} - -func requestedModelAliasFromOptions(opts cliproxyexecutor.Options, fallback string) string { - fallback = strings.TrimSpace(fallback) - if len(opts.Metadata) == 0 { - return fallback - } - raw, ok := opts.Metadata[cliproxyexecutor.RequestedModelMetadataKey] - if !ok || raw == nil { - return fallback - } - switch value := raw.(type) { - case string: - if strings.TrimSpace(value) == "" { - return fallback - } - return strings.TrimSpace(value) - case []byte: - if len(value) == 0 { - return fallback - } - return strings.TrimSpace(string(value)) - default: - return fallback - } -} - -func reasoningEffortFromOptions(opts cliproxyexecutor.Options) string { - if len(opts.Metadata) == 0 { - return "" - } - raw, ok := opts.Metadata[cliproxyexecutor.ReasoningEffortMetadataKey] - if !ok || raw == nil { - return "" - } - switch value := raw.(type) { - case string: - return strings.TrimSpace(value) - case []byte: - return strings.TrimSpace(string(value)) - default: - return "" - } -} - -func serviceTierFromOptions(opts cliproxyexecutor.Options) string { - return stringMetadataValue(opts.Metadata, cliproxyexecutor.ServiceTierMetadataKey) -} - -func generateFromOptions(opts cliproxyexecutor.Options) (bool, bool) { - if len(opts.Metadata) == 0 { - return false, false - } - raw, ok := opts.Metadata[cliproxyexecutor.GenerateMetadataKey] - if !ok || raw == nil { - return false, false - } - switch value := raw.(type) { - case bool: - return value, true - default: - return false, false - } -} - -func stringMetadataValue(metadata map[string]any, key string) string { - if len(metadata) == 0 { - return "" - } - raw, ok := metadata[key] - if !ok || raw == nil { - return "" - } - switch value := raw.(type) { - case string: - return strings.TrimSpace(value) - case []byte: - return strings.TrimSpace(string(value)) - default: - return "" - } -} - -func pinnedAuthIDFromMetadata(meta map[string]any) string { - if len(meta) == 0 { - return "" - } - raw, ok := meta[cliproxyexecutor.PinnedAuthMetadataKey] - if !ok || raw == nil { - return "" - } - switch val := raw.(type) { - case string: - return strings.TrimSpace(val) - case []byte: - return strings.TrimSpace(string(val)) - default: - return "" - } -} - -func disallowFreeAuthFromMetadata(meta map[string]any) bool { - if len(meta) == 0 { - return false - } - raw, ok := meta[cliproxyexecutor.DisallowFreeAuthMetadataKey] - if !ok || raw == nil { - return false - } - switch val := raw.(type) { - case bool: - return val - case string: - parsed, err := strconv.ParseBool(strings.TrimSpace(val)) - return err == nil && parsed - case []byte: - parsed, err := strconv.ParseBool(strings.TrimSpace(string(val))) - return err == nil && parsed - default: - return false - } -} - -func isFreeCodexAuth(auth *Auth) bool { - if auth == nil || auth.Attributes == nil { - return false - } - if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { - return false - } - return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free") -} - -func publishSelectedAuthMetadata(meta map[string]any, auth *Auth) { - if len(meta) == 0 || auth == nil { - return - } - if authID := strings.TrimSpace(auth.ID); authID != "" { - meta[cliproxyexecutor.SelectedAuthMetadataKey] = authID - if callback, ok := meta[cliproxyexecutor.SelectedAuthCallbackMetadataKey].(func(string)); ok && callback != nil { - callback(authID) - } - } - if authIndex := strings.TrimSpace(auth.EnsureIndex()); authIndex != "" { - meta[cliproxyexecutor.SelectedAuthIndexMetadataKey] = authIndex - if callback, ok := meta[cliproxyexecutor.SelectedAuthIndexCallbackMetadataKey].(func(string)); ok && callback != nil { - callback(authIndex) - } - } -} - -func rewriteModelForAuth(model string, auth *Auth) string { - if auth == nil || model == "" { - return model - } - prefix := strings.TrimSpace(auth.Prefix) - if prefix == "" { - return model - } - needle := prefix + "/" - if !strings.HasPrefix(model, needle) { - return model - } - return strings.TrimPrefix(model, needle) -} - -func (m *Manager) applyAPIKeyModelAlias(auth *Auth, requestedModel string) string { - if m == nil || auth == nil { - return requestedModel - } - - if auth.AuthKind() != AuthKindAPIKey { - return requestedModel - } - - requestedModel = strings.TrimSpace(requestedModel) - if requestedModel == "" { - return requestedModel - } - - // Fast path: lookup per-auth mapping table (keyed by auth.ID). - if resolved := m.lookupAPIKeyUpstreamModel(auth.ID, requestedModel); resolved != "" { - return resolved - } - - // Slow path: scan config for the matching credential entry and resolve alias. - // This acts as a safety net if mappings are stale or auth.ID is missing. - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - if cfg == nil { - cfg = &internalconfig.Config{} - } - - provider := strings.ToLower(strings.TrimSpace(auth.Provider)) - upstreamModel := "" - switch provider { - case "gemini": - upstreamModel = resolveUpstreamModelForGeminiAPIKey(cfg, auth, requestedModel) - case "gemini-interactions": - upstreamModel = resolveUpstreamModelForInteractionsAPIKey(cfg, auth, requestedModel) - case "claude": - upstreamModel = resolveUpstreamModelForClaudeAPIKey(cfg, auth, requestedModel) - case "codex": - upstreamModel = resolveUpstreamModelForCodexAPIKey(cfg, auth, requestedModel) - case "xai": - upstreamModel = resolveUpstreamModelForXAIAPIKey(cfg, auth, requestedModel) - case "vertex": - upstreamModel = resolveUpstreamModelForVertexAPIKey(cfg, auth, requestedModel) - default: - upstreamModel = resolveUpstreamModelForOpenAICompatAPIKey(cfg, auth, requestedModel) - } - - // Return upstream model if found, otherwise return requested model. - if upstreamModel != "" { - return upstreamModel - } - return requestedModel -} - -// APIKeyConfigEntry is a generic interface for API key configurations. -type APIKeyConfigEntry interface { - GetAPIKey() string - GetBaseURL() string -} - -func resolveAPIKeyConfig[T APIKeyConfigEntry](entries []T, auth *Auth) *T { - if auth == nil || len(entries) == 0 { - return nil - } - attrKey, attrBase := "", "" - if auth.Attributes != nil { - attrKey = strings.TrimSpace(auth.Attributes["api_key"]) - attrBase = strings.TrimSpace(auth.Attributes["base_url"]) - } - for i := range entries { - entry := &entries[i] - cfgKey := strings.TrimSpace((*entry).GetAPIKey()) - cfgBase := strings.TrimSpace((*entry).GetBaseURL()) - if attrKey != "" && attrBase != "" { - if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { - return entry - } - continue - } - if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { - if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - if attrKey != "" { - for i := range entries { - entry := &entries[i] - if strings.EqualFold(strings.TrimSpace((*entry).GetAPIKey()), attrKey) { - return entry - } - } - } - return nil -} - -func resolveGeminiAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.GeminiKey { - if cfg == nil { - return nil - } - return resolveAPIKeyConfig(cfg.GeminiKey, auth) -} - -func resolveInteractionsAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.GeminiKey { - if cfg == nil { - return nil - } - return resolveAPIKeyConfig(cfg.InteractionsKey, auth) -} - -func resolveClaudeAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.ClaudeKey { - if cfg == nil { - return nil - } - return resolveAPIKeyConfig(cfg.ClaudeKey, auth) -} - -func resolveCodexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.CodexKey { - if cfg == nil { - return nil - } - return resolveAPIKeyConfig(cfg.CodexKey, auth) -} - -func resolveXAIAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.XAIKey { - if cfg == nil { - return nil - } - return resolveAPIKeyConfig(cfg.XAIKey, auth) -} - -func resolveVertexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.VertexCompatKey { - if cfg == nil { - return nil - } - return resolveAPIKeyConfig(cfg.VertexCompatAPIKey, auth) -} - -func resolveUpstreamModelForGeminiAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { - entry := resolveGeminiAPIKeyConfig(cfg, auth) - if entry == nil { - return "" - } - return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -func resolveUpstreamModelForInteractionsAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { - entry := resolveInteractionsAPIKeyConfig(cfg, auth) - if entry == nil { - return "" - } - return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -func resolveUpstreamModelForClaudeAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { - entry := resolveClaudeAPIKeyConfig(cfg, auth) - if entry == nil { - return "" - } - return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -func resolveUpstreamModelForCodexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { - entry := resolveCodexAPIKeyConfig(cfg, auth) - if entry == nil { - return "" - } - return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -func resolveUpstreamModelForXAIAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { - entry := resolveXAIAPIKeyConfig(cfg, auth) - if entry == nil { - return "" - } - return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -func resolveUpstreamModelForVertexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { - entry := resolveVertexAPIKeyConfig(cfg, auth) - if entry == nil { - return "" - } - return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -func resolveUpstreamModelForOpenAICompatAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { - providerKey := "" - compatName := "" - if auth != nil && len(auth.Attributes) > 0 { - providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) - compatName = strings.TrimSpace(auth.Attributes["compat_name"]) - } - if compatName == "" && !strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { - return "" - } - entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider) - if entry == nil { - return "" - } - return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) -} - -type apiKeyModelAliasTable map[string]map[string]string - -func resolveOpenAICompatConfig(cfg *internalconfig.Config, providerKey, compatName, authProvider string) *internalconfig.OpenAICompatibility { - if cfg == nil { - return nil - } - candidates := make([]string, 0, 3) - if v := strings.TrimSpace(compatName); v != "" { - candidates = append(candidates, v) - } - if v := strings.TrimSpace(providerKey); v != "" { - candidates = append(candidates, v) - } - if v := strings.TrimSpace(authProvider); v != "" { - candidates = append(candidates, v) - } - for i := range cfg.OpenAICompatibility { - compat := &cfg.OpenAICompatibility[i] - if compat.Disabled { - continue - } - for _, candidate := range candidates { - if candidate != "" && strings.EqualFold(strings.TrimSpace(candidate), compat.Name) { - return compat - } - } - } - return nil -} - -func asModelAliasEntries[T interface { - GetName() string - GetAlias() string - GetForceMapping() bool -}](models []T) []modelAliasEntry { - if len(models) == 0 { - return nil - } - out := make([]modelAliasEntry, 0, len(models)) - for i := range models { - out = append(out, models[i]) - } - return out -} - -func (m *Manager) normalizeProviders(providers []string) []string { - if len(providers) == 0 { - return nil - } - result := make([]string, 0, len(providers)) - seen := make(map[string]struct{}, len(providers)) - for _, provider := range providers { - p := strings.TrimSpace(strings.ToLower(provider)) - if p == "" { - continue - } - if _, ok := seen[p]; ok { - continue - } - seen[p] = struct{}{} - result = append(result, p) - } - return result -} - -// AvailableProviders returns the set of provider keys that currently have at least one -// registered auth record that is not disabled. It is a best-effort snapshot for routing -// decisions and does not account for per-model cooldowns or transient runtime availability. -// Disabled auths (Disabled flag or StatusDisabled) are excluded so routing does not target -// providers that auth selection would refuse to use, which would otherwise cause execution -// failures instead of falling back to lower-priority routers. -func (m *Manager) AvailableProviders() []string { - if m == nil { - return nil - } - m.mu.RLock() - defer m.mu.RUnlock() - seen := make(map[string]struct{}, len(m.auths)) - out := make([]string, 0, len(m.auths)) - for _, auth := range m.auths { - if auth == nil || auth.Disabled || auth.Status == StatusDisabled { - continue - } - provider := strings.ToLower(strings.TrimSpace(auth.Provider)) - if provider == "" { - continue - } - if _, ok := seen[provider]; ok { - continue - } - seen[provider] = struct{}{} - out = append(out, provider) - } - sort.Strings(out) - return out -} - -// HasProviderAuth reports whether at least one non-disabled auth record is registered for -// the provider. Disabled auths (Disabled flag or StatusDisabled) are excluded to match the -// behavior of auth selection, which refuses to pick disabled credentials. -func (m *Manager) HasProviderAuth(provider string) bool { - if m == nil { - return false - } - provider = strings.ToLower(strings.TrimSpace(provider)) - if provider == "" { - return false - } - m.mu.RLock() - defer m.mu.RUnlock() - for _, auth := range m.auths { - if auth == nil || auth.Disabled || auth.Status == StatusDisabled { - continue - } - if strings.ToLower(strings.TrimSpace(auth.Provider)) == provider { - return true - } - } - return false -} - -func (m *Manager) retrySettings() (int, int, time.Duration) { - if m == nil { - return 0, 0, 0 - } - return int(m.requestRetry.Load()), int(m.maxRetryCredentials.Load()), time.Duration(m.maxRetryInterval.Load()) -} - -func (m *Manager) closestCooldownWait(providers []string, model string, attempt int) (time.Duration, bool) { - if m == nil || len(providers) == 0 { - return 0, false - } - now := time.Now() - defaultRetry := int(m.requestRetry.Load()) - if defaultRetry < 0 { - defaultRetry = 0 - } - providerSet := make(map[string]struct{}, len(providers)) - for i := range providers { - key := strings.TrimSpace(strings.ToLower(providers[i])) - if key == "" { - continue - } - providerSet[key] = struct{}{} - } - m.mu.RLock() - defer m.mu.RUnlock() - var ( - found bool - minWait time.Duration - ) - for _, auth := range m.auths { - if auth == nil { - continue - } - providerKey := executorKeyFromAuth(auth) - if _, ok := providerSet[providerKey]; !ok { - continue - } - effectiveRetry := defaultRetry - if override, ok := auth.RequestRetryOverride(); ok { - effectiveRetry = override - } - if effectiveRetry < 0 { - effectiveRetry = 0 - } - if attempt >= effectiveRetry { - continue - } - checkModel := model - if strings.TrimSpace(model) != "" { - checkModel = m.selectionModelForAuth(auth, model) - } - blocked, reason, next := isAuthBlockedForModel(auth, checkModel, now) - if !blocked || next.IsZero() || reason == blockReasonDisabled { - continue - } - wait := next.Sub(now) - if wait < 0 { - continue - } - if !found || wait < minWait { - minWait = wait - found = true - } - } - return minWait, found -} - -func (m *Manager) retryAllowed(attempt int, providers []string) bool { - if m == nil || attempt < 0 || len(providers) == 0 { - return false - } - defaultRetry := int(m.requestRetry.Load()) - if defaultRetry < 0 { - defaultRetry = 0 - } - providerSet := make(map[string]struct{}, len(providers)) - for i := range providers { - key := strings.TrimSpace(strings.ToLower(providers[i])) - if key == "" { - continue - } - providerSet[key] = struct{}{} - } - if len(providerSet) == 0 { - return false - } - - m.mu.RLock() - defer m.mu.RUnlock() - for _, auth := range m.auths { - if auth == nil { - continue - } - providerKey := executorKeyFromAuth(auth) - if _, ok := providerSet[providerKey]; !ok { - continue - } - effectiveRetry := defaultRetry - if override, ok := auth.RequestRetryOverride(); ok { - effectiveRetry = override - } - if effectiveRetry < 0 { - effectiveRetry = 0 - } - if attempt < effectiveRetry { - return true - } - } - return false -} - -func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []string, model string, maxWait time.Duration) (time.Duration, bool) { - if err == nil { - return 0, false - } - var homeBusy *HomeConcurrencyBusyError - if errors.As(err, &homeBusy) && homeBusy != nil { - return 0, false - } - if maxWait <= 0 { - return 0, false - } - status := statusCodeFromError(err) - if status == http.StatusOK { - return 0, false - } - if isRequestInvalidError(err) { - return 0, false - } - wait, found := m.closestCooldownWait(providers, model, attempt) - if found { - if wait > maxWait { - return 0, false - } - return wait, true - } - if status != http.StatusTooManyRequests { - return 0, false - } - if !m.retryAllowed(attempt, providers) { - return 0, false - } - retryAfter := retryAfterFromError(err) - if retryAfter == nil || *retryAfter <= 0 || *retryAfter > maxWait { - return 0, false - } - return *retryAfter, true -} - -// cooldownWaitJitterCap bounds the random jitter added to cooldown waits so a -// long wait is never extended by more than this amount. -const cooldownWaitJitterCap = 2 * time.Second - -// jitteredCooldownWait adds a small random delay to a cooldown wait so -// concurrent requests waiting on the same recovery deadline do not wake in -// lockstep and stampede the first credential that recovers. The jitter never -// pushes the total wait past maxWait, which callers have already enforced as -// the retry ceiling; maxWait <= 0 means no ceiling. -func jitteredCooldownWait(wait, maxWait time.Duration) time.Duration { - if wait <= 0 { - return wait - } - jitterRange := wait / 4 - if jitterRange > cooldownWaitJitterCap { - jitterRange = cooldownWaitJitterCap - } - if maxWait > 0 && jitterRange > maxWait-wait { - jitterRange = maxWait - wait - } - if jitterRange <= 0 { - return wait - } - return wait + rand.N(jitterRange) -} - -func waitForCooldown(ctx context.Context, wait, maxWait time.Duration) error { - if wait <= 0 { - return nil - } - timer := time.NewTimer(jitteredCooldownWait(wait, maxWait)) - defer timer.Stop() - select { - case <-ctx.Done(): - return ctx.Err() - case <-timer.C: - return nil - } -} - -// MarkResult records an execution result and notifies hooks. -func (m *Manager) MarkResult(ctx context.Context, result Result) { - if result.AuthID == "" { - return - } - - shouldResumeModel := false - shouldSuspendModel := false - suspendReason := "" - clearModelQuota := false - setModelQuota := false - var authSnapshot *Auth - cooldownStateChanged := false - - m.mu.Lock() - if auth, ok := m.auths[result.AuthID]; ok && auth != nil { - now := time.Now() - var cooldownRecordsBefore []CooldownStateRecord - trackCooldownState := m.cooldownStore != nil - if trackCooldownState { - cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now) - } - auth.recordRecentRequest(now, result.Success) - if result.Success { - auth.Success++ - } else { - auth.Failed++ - } - - if result.Success { - if result.Model != "" { - state := ensureModelState(auth, result.Model) - resetModelState(state, now) - updateAggregatedAvailability(auth, now) - if !hasModelError(auth, now) { - auth.LastError = nil - auth.StatusMessage = "" - auth.Status = StatusActive - } - auth.UpdatedAt = now - shouldResumeModel = true - clearModelQuota = true - } else { - clearAuthStateOnSuccess(auth, now) - } - } else { - if result.Model != "" { - if !isRequestScopedResultError(result.Error) { - disableCooling := m.cooldownDisabledForAuth(auth) - state := ensureModelState(auth, result.Model) - state.Unavailable = true - state.Status = StatusError - state.UpdatedAt = now - if result.Error != nil { - state.LastError = cloneError(result.Error) - state.StatusMessage = result.Error.Message - auth.LastError = cloneError(result.Error) - auth.StatusMessage = result.Error.Message - } - - statusCode := statusCodeFromResult(result.Error) - if isModelSupportResultError(result.Error) { - next := now.Add(12 * time.Hour) - state.NextRetryAfter = next - suspendReason = "model_not_supported" - shouldSuspendModel = true - } else if isCloudflareChallengeResultError(result.Error) { - next, backoffLevel := nextCloudflareCooldown(state.Quota.BackoffLevel, disableCooling, now) - state.NextRetryAfter = next - state.StatusMessage = "cloudflare challenge" - if auth.LastError != nil { - auth.StatusMessage = "cloudflare challenge" - } - state.Quota = QuotaState{ - Exceeded: true, - Reason: "cloudflare challenge", - NextRecoverAt: next, - BackoffLevel: backoffLevel, - } - } else if isInvalidGrantResultError(result.Error) { - if disableCooling { - state.NextRetryAfter = time.Time{} - } else { - state.NextRetryAfter = now.Add(30 * time.Minute) - suspendReason = "invalid_grant" - shouldSuspendModel = true - } - } else { - switch statusCode { - case 401: - if disableCooling { - state.NextRetryAfter = time.Time{} - } else { - next := now.Add(30 * time.Minute) - state.NextRetryAfter = next - suspendReason = "unauthorized" - shouldSuspendModel = true - } - case 402, 403: - if disableCooling { - state.NextRetryAfter = time.Time{} - } else { - next := now.Add(30 * time.Minute) - state.NextRetryAfter = next - suspendReason = "payment_required" - shouldSuspendModel = true - } - case 404: - if disableCooling { - state.NextRetryAfter = time.Time{} - } else { - next := now.Add(12 * time.Hour) - state.NextRetryAfter = next - suspendReason = "not_found" - shouldSuspendModel = true - } - case 429: - var next time.Time - backoffLevel := state.Quota.BackoffLevel - if !disableCooling { - if result.RetryAfter != nil { - next = now.Add(*result.RetryAfter) - } else { - next, backoffLevel = quotaCooldownAfterFailure(state.Quota, now) - } - } - state.NextRetryAfter = next - state.Quota = QuotaState{ - Exceeded: true, - Reason: "quota", - NextRecoverAt: next, - BackoffLevel: backoffLevel, - } - if !disableCooling { - suspendReason = "quota" - shouldSuspendModel = true - setModelQuota = true - } - case 408, 500, 502, 503, 504: - if disableCooling { - state.NextRetryAfter = time.Time{} - } else { - state.NextRetryAfter = nextTransientErrorRetryAfter(now) - } - default: - state.NextRetryAfter = time.Time{} - } - } - - auth.Status = StatusError - auth.UpdatedAt = now - updateAggregatedAvailability(auth, now) - } - } else { - disableCooling := m.cooldownDisabledForAuth(auth) - applyAuthFailureState(auth, result.Error, result.RetryAfter, now, disableCooling) - } - } - - _ = m.persist(ctx, auth) - authSnapshot = auth.Clone() - if trackCooldownState { - cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) - cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) - } - } - m.mu.Unlock() - if m.scheduler != nil && authSnapshot != nil { - m.scheduler.upsertAuth(authSnapshot) - } - if authSnapshot != nil && cooldownStateChanged { - m.persistCooldownStates(context.Background()) - } - - if clearModelQuota && result.Model != "" { - registry.GetGlobalRegistry().ClearModelQuotaExceeded(result.AuthID, result.Model) - } - if setModelQuota && result.Model != "" { - registry.GetGlobalRegistry().SetModelQuotaExceeded(result.AuthID, result.Model) - } - if shouldResumeModel { - registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, result.Model) - } else if shouldSuspendModel { - registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, result.Model, suspendReason) - } - - m.hook.OnResult(ctx, result) - m.publishErrorEvent(result, authSnapshot) -} - -func (m *Manager) recordExecutionResult(ctx context.Context, result Result, auth *Auth, ephemeral bool) { - if !ephemeral { - m.MarkResult(ctx, result) - return - } - m.reportHomeResult(ctx, result, auth) -} - -// reportHomeResult only observes a Home dispatch result and never updates local auth state. -func (m *Manager) reportHomeResult(ctx context.Context, result Result, auth *Auth) { - if m == nil || result.AuthID == "" { - return - } - var snapshot *Auth - if auth != nil { - snapshot = auth.Clone() - } - m.hook.OnResult(ctx, result) - m.publishErrorEvent(result, snapshot) -} - -func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Result) { - if result.AuthID == "" { - return - } - - var authSnapshot *Auth - m.mu.Lock() - if auth, ok := m.auths[result.AuthID]; ok && auth != nil { - now := time.Now() - auth.recordRecentRequest(now, result.Success) - if result.Success { - auth.Success++ - } else { - auth.Failed++ - } - _ = m.persist(ctx, auth) - authSnapshot = auth.Clone() - } - m.mu.Unlock() - - m.hook.OnResult(ctx, result) - m.publishErrorEvent(result, authSnapshot) -} - -func ensureModelState(auth *Auth, model string) *ModelState { - if auth == nil || model == "" { - return nil - } - if auth.ModelStates == nil { - auth.ModelStates = make(map[string]*ModelState) - } - if state, ok := auth.ModelStates[model]; ok && state != nil { - return state - } - state := &ModelState{Status: StatusActive} - auth.ModelStates[model] = state - return state -} - -func resetModelState(state *ModelState, now time.Time) { - if state == nil { - return - } - state.Unavailable = false - state.Status = StatusActive - state.StatusMessage = "" - state.NextRetryAfter = time.Time{} - state.LastError = nil - state.Quota = QuotaState{} - state.UpdatedAt = now -} - -func modelStateIsClean(state *ModelState) bool { - if state == nil { - return true - } - if state.Status != StatusActive { - return false - } - if state.Unavailable || state.StatusMessage != "" || !state.NextRetryAfter.IsZero() || state.LastError != nil { - return false - } - if state.Quota.Exceeded || state.Quota.Reason != "" || !state.Quota.NextRecoverAt.IsZero() || state.Quota.BackoffLevel != 0 { - return false - } - return true -} - -func updateAggregatedAvailability(auth *Auth, now time.Time) { - if auth == nil { - return - } - if len(auth.ModelStates) == 0 { - clearAggregatedAvailability(auth) - return - } - allUnavailable := true - earliestRetry := time.Time{} - quotaExceeded := false - quotaRecover := time.Time{} - maxBackoffLevel := 0 - hasState := false - for _, state := range auth.ModelStates { - if state == nil { - continue - } - hasState = true - stateUnavailable := false - if state.Status == StatusDisabled { - stateUnavailable = true - } else if state.Unavailable { - if state.NextRetryAfter.IsZero() { - stateUnavailable = false - } else if state.NextRetryAfter.After(now) { - stateUnavailable = true - if earliestRetry.IsZero() || state.NextRetryAfter.Before(earliestRetry) { - earliestRetry = state.NextRetryAfter - } - } else { - state.Unavailable = false - state.NextRetryAfter = time.Time{} - } - } - if !stateUnavailable { - allUnavailable = false - } - if state.Quota.Exceeded { - quotaExceeded = true - if quotaRecover.IsZero() || (!state.Quota.NextRecoverAt.IsZero() && state.Quota.NextRecoverAt.Before(quotaRecover)) { - quotaRecover = state.Quota.NextRecoverAt - } - if state.Quota.BackoffLevel > maxBackoffLevel { - maxBackoffLevel = state.Quota.BackoffLevel - } - } - } - if !hasState { - clearAggregatedAvailability(auth) - return - } - auth.Unavailable = allUnavailable - if allUnavailable { - auth.NextRetryAfter = earliestRetry - } else { - auth.NextRetryAfter = time.Time{} - } - if quotaExceeded { - auth.Quota.Exceeded = true - auth.Quota.Reason = "quota" - auth.Quota.NextRecoverAt = quotaRecover - auth.Quota.BackoffLevel = maxBackoffLevel - } else { - auth.Quota.Exceeded = false - auth.Quota.Reason = "" - auth.Quota.NextRecoverAt = time.Time{} - auth.Quota.BackoffLevel = 0 - } -} - -func clearAggregatedAvailability(auth *Auth) { - if auth == nil { - return - } - auth.Unavailable = false - auth.NextRetryAfter = time.Time{} - auth.Quota = QuotaState{} -} - -func hasModelError(auth *Auth, now time.Time) bool { - if auth == nil || len(auth.ModelStates) == 0 { - return false - } - for _, state := range auth.ModelStates { - if state == nil { - continue - } - if state.LastError != nil { - return true - } - if state.Status == StatusError { - if state.Unavailable && (state.NextRetryAfter.IsZero() || state.NextRetryAfter.After(now)) { - return true - } - } - } - return false -} - -func clearAuthStateOnSuccess(auth *Auth, now time.Time) { - if auth == nil { - return - } - auth.Unavailable = false - auth.Status = StatusActive - auth.StatusMessage = "" - auth.Quota.Exceeded = false - auth.Quota.Reason = "" - auth.Quota.NextRecoverAt = time.Time{} - auth.Quota.BackoffLevel = 0 - auth.LastError = nil - auth.NextRetryAfter = time.Time{} - auth.UpdatedAt = now -} - -func cloneError(err *Error) *Error { - if err == nil { - return nil - } - return &Error{ - Code: err.Code, - Message: err.Message, - Retryable: err.Retryable, - HTTPStatus: err.HTTPStatus, - } -} - -func errorString(err error) string { - if err == nil { - return "" - } - return err.Error() -} - -func statusCodeFromError(err error) int { - if err == nil { - return 0 - } - type statusCoder interface { - StatusCode() int - } - var sc statusCoder - if errors.As(err, &sc) && sc != nil { - return sc.StatusCode() - } - return 0 -} - -func isRequestScopedError(err error) bool { - if err == nil { - return false - } - requestErr, ok := errors.AsType[cliproxyexecutor.RequestScopedError](err) - return ok && requestErr != nil && requestErr.IsRequestScoped() -} - -func resultErrorFromError(err error) *Error { - if err == nil { - return nil - } - var sourceErr *Error - var resultErr *Error - if errors.As(err, &sourceErr) && sourceErr != nil { - resultErr = cloneError(sourceErr) - } else { - resultErr = &Error{Message: err.Error()} - } - if resultErr.HTTPStatus == 0 { - resultErr.HTTPStatus = statusCodeFromError(err) - } - if isRequestScopedError(err) || isRequestInvalidError(err) { - resultErr.Code = requestScopedErrorCode - } - return resultErr -} - -func isUnauthorizedError(err error) bool { - if err == nil { - return false - } - if statusCodeFromError(err) == http.StatusUnauthorized { - return true - } - raw := strings.ToLower(err.Error()) - return strings.Contains(raw, "status 401") || strings.Contains(raw, "401 unauthorized") -} - -func hasUnauthorizedAuthFailure(auth *Auth) bool { - if auth == nil || auth.LastError == nil { - return false - } - return auth.LastError.StatusCode() == http.StatusUnauthorized || strings.EqualFold(auth.LastError.Code, "unauthorized") -} - -func refreshErrorFromError(err error) *Error { - if err == nil { - return nil - } - statusCode := statusCodeFromError(err) - if statusCode == 0 && isUnauthorizedError(err) { - statusCode = http.StatusUnauthorized - } - authErr := &Error{Message: err.Error(), HTTPStatus: statusCode} - if statusCode == http.StatusUnauthorized { - authErr.Code = "unauthorized" - authErr.Retryable = false - } - return authErr -} - -func retryAfterFromError(err error) *time.Duration { - if err == nil { - return nil - } - type retryAfterProvider interface { - RetryAfter() *time.Duration - } - var rap retryAfterProvider - if !errors.As(err, &rap) || rap == nil { - return nil - } - retryAfter := rap.RetryAfter() - if retryAfter == nil { - return nil - } - value := *retryAfter - return &value -} - -func statusCodeFromResult(err *Error) int { - if err == nil { - return 0 - } - return err.StatusCode() -} - -func isModelSupportErrorMessage(message string) bool { - lower := strings.ToLower(strings.TrimSpace(message)) - if lower == "" { - return false - } - patterns := [...]string{ - "model_not_supported", - "requested model is not supported", - "requested model is unsupported", - "requested model is unavailable", - "model is not supported", - "model not supported", - "unsupported model", - "model unavailable", - "not available for your plan", - "not available for your account", - } - for _, pattern := range patterns { - if strings.Contains(lower, pattern) { - return true - } - } - return false -} - -func isModelSupportError(err error) bool { - if err == nil { - return false - } - status := statusCodeFromError(err) - if status != http.StatusBadRequest && status != http.StatusUnprocessableEntity { - return false - } - return isModelSupportErrorMessage(err.Error()) -} - -func isInvalidGrantErrorMessage(message string) bool { - return strings.Contains(strings.ToLower(message), "invalid_grant") -} - -func isInvalidGrantError(err error) bool { - if err == nil { - return false - } - status := statusCodeFromError(err) - if status != http.StatusBadRequest && status != http.StatusUnauthorized { - return false - } - return isInvalidGrantErrorMessage(err.Error()) -} - -func isInvalidGrantResultError(err *Error) bool { - if err == nil { - return false - } - status := statusCodeFromResult(err) - if status != http.StatusBadRequest && status != http.StatusUnauthorized { - return false - } - return isInvalidGrantErrorMessage(err.Code) || isInvalidGrantErrorMessage(err.Message) -} - -func isModelSupportResultError(err *Error) bool { - if err == nil { - return false - } - status := statusCodeFromResult(err) - if status != http.StatusBadRequest && status != http.StatusUnprocessableEntity { - return false - } - return isModelSupportErrorMessage(err.Message) -} - -func isCloudflareChallengeErrorMessage(message string) bool { - lower := strings.ToLower(strings.TrimSpace(message)) - return strings.Contains(lower, "challenge-platform") || - strings.Contains(lower, "cf-mitigated") || - strings.Contains(lower, "cloudflare challenge") || - (strings.Contains(lower, "cloudflare") && strings.Contains(lower, " 0 { - next = now.Add(cooldown) - } - backoffLevel = nextLevel - } - return next, backoffLevel -} -func isRequestScopedNotFoundMessage(message string) bool { - if message == "" { - return false - } - lower := strings.ToLower(message) - return strings.Contains(lower, "item with id") && - strings.Contains(lower, "not found") && - strings.Contains(lower, "items are not persisted when `store` is set to false") -} - -func isRequestScopedNotFoundResultError(err *Error) bool { - if err == nil || statusCodeFromResult(err) != http.StatusNotFound { - return false - } - return isRequestScopedNotFoundMessage(err.Message) -} - -func isRequestScopedResultError(err *Error) bool { - return err != nil && (err.IsRequestScoped() || isRequestScopedNotFoundResultError(err)) -} - -func isCountTokensEndpointNotFoundError(err error, requestedModel string) bool { - if err == nil || statusCodeFromError(err) != http.StatusNotFound { - return false - } - baseModel := thinking.ParseSuffix(requestedModel).ModelName - return !isExplicitModelNotFoundError(err, baseModel) -} - -func isExplicitModelNotFoundError(err error, requestedModel string) bool { - if err == nil { - return false - } - if authErr, ok := err.(*Error); ok && authErr != nil { - if isModelNotFoundIdentifier(authErr.Code) || isStructuredModelNotFoundError(authErr.Message, requestedModel) { - return true - } - } else if isStructuredModelNotFoundError(err.Error(), requestedModel) { - return true - } - - switch wrapped := err.(type) { - case interface{ Unwrap() []error }: - for _, nested := range wrapped.Unwrap() { - if isExplicitModelNotFoundError(nested, requestedModel) { - return true - } - } - case interface{ Unwrap() error }: - return isExplicitModelNotFoundError(wrapped.Unwrap(), requestedModel) - } - return false -} - -func isStructuredModelNotFoundError(message, requestedModel string) bool { - var payload any - if errJSON := json.Unmarshal([]byte(strings.TrimSpace(message)), &payload); errJSON != nil { - return false - } - return containsStructuredModelNotFound(payload, requestedModel) -} - -func containsStructuredModelNotFound(value any, requestedModel string) bool { - switch typed := value.(type) { - case map[string]any: - notFoundType := false - exactModelReference := false - for key, item := range typed { - text, isString := item.(string) - if isString { - switch strings.ToLower(strings.TrimSpace(key)) { - case "code": - if isModelNotFoundIdentifier(text) { - return true - } - case "type": - if isModelNotFoundIdentifier(text) { - return true - } - notFoundType = notFoundType || isNotFoundErrorIdentifier(text) - case "error", "message", "detail", "error_description", "title": - if isExplicitModelNotFoundMessage(text, requestedModel) { - return true - } - exactModelReference = exactModelReference || isExactRequestedModelReference(text, requestedModel) - } - } - switch item.(type) { - case map[string]any, []any: - if containsStructuredModelNotFound(item, requestedModel) { - return true - } - } - } - return notFoundType && exactModelReference - case []any: - for _, item := range typed { - if text, isString := item.(string); isString && isExplicitModelNotFoundMessage(text, requestedModel) { - return true - } - if containsStructuredModelNotFound(item, requestedModel) { - return true - } - } - } - return false -} - -func isModelNotFoundIdentifier(value string) bool { - candidate := strings.ToLower(strings.TrimSpace(value)) - if fragment := strings.LastIndex(candidate, "#"); fragment >= 0 && fragment+1 < len(candidate) { - candidate = candidate[fragment+1:] - } else { - if query := strings.Index(candidate, "?"); query >= 0 { - candidate = candidate[:query] - } - candidate = strings.TrimRight(candidate, "/") - if separator := strings.LastIndexAny(candidate, "/:"); separator >= 0 { - candidate = candidate[separator+1:] - } - } - normalized := strings.NewReplacer("-", "_", " ", "_").Replace(candidate) - switch normalized { - case "model_not_found", "model_not_found_error", "unknown_model", "model_does_not_exist", "model_not_exist": - return true - default: - return false - } -} - -func isNotFoundErrorIdentifier(value string) bool { - normalized := strings.NewReplacer("-", "_", " ", "_").Replace(strings.ToLower(strings.TrimSpace(value))) - return normalized == "not_found" || normalized == "not_found_error" -} - -func isExplicitModelNotFoundMessage(message, requestedModel string) bool { - lower := strings.Trim(strings.ToLower(strings.TrimSpace(message)), " .!;\t\r\n") - if lower == "" { - return false - } - normalized := strings.NewReplacer("-", "_", " ", "_").Replace(lower) - if strings.Contains(normalized, "model_not_found") || strings.Contains(normalized, "unknown_model") { - return true - } - for _, prefix := range []string{"no such model", "unknown model"} { - if lower != prefix && !strings.HasPrefix(lower, prefix+" ") && !strings.HasPrefix(lower, prefix+":") { - continue - } - remainder := strings.TrimSpace(strings.TrimPrefix(lower, prefix)) - remainder = strings.TrimSpace(strings.TrimPrefix(remainder, ":")) - if remainder == "" { - return true - } - missingSuffix, matches := trimRequestedModelReference(remainder, requestedModel) - return matches && missingSuffix == "" - } - for _, prefix := range []string{"the requested model", "requested model", "the model", "model"} { - if lower != prefix && !strings.HasPrefix(lower, prefix+" ") && !strings.HasPrefix(lower, prefix+":") { - continue - } - remainder := strings.TrimSpace(strings.TrimPrefix(lower, prefix)) - remainder = strings.TrimSpace(strings.TrimPrefix(remainder, ":")) - if isMissingModelPhrase(remainder) { - return true - } - missingSuffix, matches := trimRequestedModelReference(remainder, requestedModel) - return matches && isMissingModelPhrase(missingSuffix) - } - return false -} - -func isExactRequestedModelReference(message, requestedModel string) bool { - lower := strings.Trim(strings.ToLower(strings.TrimSpace(message)), " .!;\t\r\n") - for _, prefix := range []string{"the requested model", "requested model", "the model", "model"} { - if lower != prefix && !strings.HasPrefix(lower, prefix+" ") && !strings.HasPrefix(lower, prefix+":") { - continue - } - remainder := strings.TrimSpace(strings.TrimPrefix(lower, prefix)) - remainder = strings.TrimSpace(strings.TrimPrefix(remainder, ":")) - suffix, matches := trimRequestedModelReference(remainder, requestedModel) - return matches && suffix == "" - } - return false -} - -func trimRequestedModelReference(value, requestedModel string) (string, bool) { - model := strings.ToLower(strings.TrimSpace(requestedModel)) - if model == "" { - return "", false - } - for _, candidate := range []string{model, "'" + model + "'", `"` + model + `"`, "`" + model + "`"} { - if value == candidate { - return "", true - } - if !strings.HasPrefix(value, candidate) { - continue - } - remainder := value[len(candidate):] - if remainder == "" || strings.ContainsRune(" :,", rune(remainder[0])) { - return strings.TrimLeft(remainder, " :,"), true - } - } - return "", false -} - -func isMissingModelPhrase(value string) bool { - switch strings.Trim(value, " .!;\t\r\n") { - case "not found", "was not found", "could not be found", "does not exist", "doesn't exist", "not exist", "is unknown": - return true - default: - return false - } -} - -// isRequestInvalidError returns true if the error represents a client request -// error that should not be retried. Specifically, it treats 400 responses with -// "invalid_request_error", request-scoped 404 item misses caused by `store=false`, -// and all 422 responses as request-shape failures, where switching auths or -// pooled upstream models will not help. Model-support errors are excluded so -// routing can fall through to another auth or upstream. -func isRequestInvalidError(err error) bool { - if err == nil { - return false - } - if isRequestScopedError(err) { - return true - } - if isCloudflareChallengeError(err) { - return false - } - if isInvalidGrantError(err) { - return false - } - if isModelSupportError(err) { - return false - } - status := statusCodeFromError(err) - switch status { - case http.StatusBadRequest: - msg := err.Error() - return strings.Contains(msg, "invalid_request_error") || - strings.Contains(msg, "bad_request_error") || - strings.Contains(msg, "INVALID_ARGUMENT") || - strings.Contains(msg, "FAILED_PRECONDITION") - case http.StatusNotFound: - return isRequestScopedNotFoundMessage(err.Error()) - case http.StatusUnprocessableEntity: - return true - case http.StatusInternalServerError: - msg := err.Error() - return strings.Contains(msg, "\"status\":\"UNKNOWN\"") || - strings.Contains(msg, "\"status\": \"UNKNOWN\"") - default: - return false - } -} - -func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Duration, now time.Time, disableCooling bool) { - if auth == nil { - return - } - if isRequestScopedResultError(resultErr) { - return - } - auth.Unavailable = true - auth.Status = StatusError - auth.UpdatedAt = now - if resultErr != nil { - auth.LastError = cloneError(resultErr) - if resultErr.Message != "" { - auth.StatusMessage = resultErr.Message - } - } - statusCode := statusCodeFromResult(resultErr) - if isCloudflareChallengeResultError(resultErr) { - auth.StatusMessage = "cloudflare challenge" - next, backoffLevel := nextCloudflareCooldown(auth.Quota.BackoffLevel, disableCooling, now) - auth.Quota = QuotaState{ - Exceeded: true, - Reason: "cloudflare challenge", - NextRecoverAt: next, - BackoffLevel: backoffLevel, - } - auth.NextRetryAfter = next - return - } - if isInvalidGrantResultError(resultErr) { - auth.StatusMessage = "invalid_grant" - if disableCooling { - auth.NextRetryAfter = time.Time{} - } else { - auth.NextRetryAfter = now.Add(30 * time.Minute) - } - return - } - switch statusCode { - case 401: - auth.StatusMessage = "unauthorized" - if disableCooling { - auth.NextRetryAfter = time.Time{} - } else { - auth.NextRetryAfter = now.Add(30 * time.Minute) - } - case 402, 403: - auth.StatusMessage = "payment_required" - if disableCooling { - auth.NextRetryAfter = time.Time{} - } else { - auth.NextRetryAfter = now.Add(30 * time.Minute) - } - case 404: - auth.StatusMessage = "not_found" - if disableCooling { - auth.NextRetryAfter = time.Time{} - } else { - auth.NextRetryAfter = now.Add(12 * time.Hour) - } - case 429: - auth.StatusMessage = "quota exhausted" - auth.Quota.Exceeded = true - auth.Quota.Reason = "quota" - var next time.Time - if !disableCooling { - if retryAfter != nil { - next = now.Add(*retryAfter) - } else { - next, auth.Quota.BackoffLevel = quotaCooldownAfterFailure(auth.Quota, now) - } - } - auth.Quota.NextRecoverAt = next - auth.NextRetryAfter = next - case 408, 500, 502, 503, 504: - auth.StatusMessage = "transient upstream error" - if disableCooling { - auth.NextRetryAfter = time.Time{} - } else { - auth.NextRetryAfter = nextTransientErrorRetryAfter(now) - } - default: - if auth.StatusMessage == "" { - auth.StatusMessage = "request failed" - } - } -} - -// quotaCooldownAfterFailure returns the recovery deadline and backoff level for -// a quota failure observed at now. Failures that land while a previous quota -// window is still open reuse that window instead of escalating, so a burst of -// concurrent in-flight failures advances the backoff ladder at most once per -// window. -func quotaCooldownAfterFailure(quota QuotaState, now time.Time) (time.Time, int) { - if quota.NextRecoverAt.After(now) { - return quota.NextRecoverAt, quota.BackoffLevel - } - cooldown, nextLevel := nextQuotaCooldown(quota.BackoffLevel, false) - var next time.Time - if cooldown > 0 { - next = now.Add(cooldown) - } - return next, nextLevel -} - -// nextQuotaCooldown returns the next cooldown duration and updated backoff level for repeated quota errors. -func nextQuotaCooldown(prevLevel int, disableCooling bool) (time.Duration, int) { - if prevLevel < 0 { - prevLevel = 0 - } - if disableCooling { - return 0, prevLevel - } - cooldown := quotaBackoffBase * time.Duration(1<= quotaBackoffMax { - return quotaBackoffMax, prevLevel - } - return cooldown, prevLevel + 1 -} - -// List returns all auth entries currently known by the manager. -func (m *Manager) List() []*Auth { - m.mu.RLock() - defer m.mu.RUnlock() - list := make([]*Auth, 0, len(m.auths)) - for _, auth := range m.auths { - list = append(list, auth.Clone()) - } - return list -} - -// GetByID retrieves an auth entry by its ID. - -func (m *Manager) GetByID(id string) (*Auth, bool) { - if id == "" { - return nil, false - } - m.mu.RLock() - defer m.mu.RUnlock() - auth, ok := m.auths[id] - if !ok { - return nil, false - } - return auth.Clone(), true -} - -// GetExecutionSessionAuthByID retrieves a Home runtime auth scoped to an execution session. -func (m *Manager) GetExecutionSessionAuthByID(sessionID string, authID string) (*Auth, bool) { - sessionID = strings.TrimSpace(sessionID) - authID = strings.TrimSpace(authID) - if m == nil || sessionID == "" || authID == "" { - return nil, false - } - m.mu.RLock() - defer m.mu.RUnlock() - sessionAuths := m.homeRuntimeAuths[sessionID] - auth := sessionAuths[authID] - if auth == nil { - return nil, false - } - return auth.Clone(), true -} - -// Executor returns the registered provider executor for a provider key. -func (m *Manager) Executor(provider string) (ProviderExecutor, bool) { - if m == nil { - return nil, false - } - provider = strings.TrimSpace(provider) - if provider == "" { - return nil, false - } - - m.mu.RLock() - executor, okExecutor := m.executors[provider] - if !okExecutor { - lowerProvider := strings.ToLower(provider) - if lowerProvider != provider { - executor, okExecutor = m.executors[lowerProvider] - } - } - m.mu.RUnlock() - - if !okExecutor || executor == nil { - return nil, false - } - return executor, true -} - -// CloseExecutionSession asks all registered executors to release the supplied execution session. -func (m *Manager) CloseExecutionSession(sessionID string) { - sessionID = strings.TrimSpace(sessionID) - if m == nil || sessionID == "" { - return - } - - m.mu.Lock() - var selections []*HomeDispatchSelection - if sessionID == CloseAllExecutionSessionsID { - m.clearHomeRuntimeAuthsLocked() - selections = m.takeAllHomeSessionSelectionsLocked() - m.clearHomeSessionLocks() - } else { - m.clearHomeRuntimeAuthsForSessionLocked(sessionID) - selections = m.takeHomeSessionSelectionsLocked(sessionID) - m.homeSessionLocks.Delete(sessionID) - } - executors := make([]ProviderExecutor, 0, len(m.executors)) - for _, exec := range m.executors { - executors = append(executors, exec) - } - m.mu.Unlock() - - for _, selection := range selections { - selection.End("session_closed") - } - for i := range executors { - if closer, ok := executors[i].(ExecutionSessionCloser); ok && closer != nil { - closer.CloseExecutionSession(sessionID) - } - } -} - -func (m *Manager) useSchedulerFastPath() bool { - if m == nil || m.scheduler == nil { - return false - } - return isBuiltInSelector(m.selector) -} - -func shouldRetrySchedulerPick(err error) bool { - if err == nil { - return false - } - var cooldownErr *modelCooldownError - if errors.As(err, &cooldownErr) { - return true - } - var authErr *Error - if !errors.As(err, &authErr) || authErr == nil { - return false - } - return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" -} - -func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) bool { - if auth == nil || strings.TrimSpace(routeModel) == "" { - return false - } - return m.selectionModelKeyForAuth(auth, routeModel) != canonicalModelKey(routeModel) -} - -func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { - if m.HomeEnabled() { - auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) - return auth, exec, err - } - - pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - - m.mu.RLock() - selector := m.selector - pluginScheduler := m.pluginScheduler - executor, okExecutor := m.executors[provider] - if !okExecutor { - m.mu.RUnlock() - return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"} - } - candidates := make([]*Auth, 0, len(m.auths)) - modelKey := strings.TrimSpace(model) - // Always use base model name (without thinking suffix) for auth matching. - if modelKey != "" { - parsed := thinking.ParseSuffix(modelKey) - if parsed.ModelName != "" { - modelKey = strings.TrimSpace(parsed.ModelName) - } - } - registryRef := registry.GetGlobalRegistry() - for _, candidate := range m.auths { - if candidate == nil || executorKeyFromAuth(candidate) != provider || candidate.Disabled { - continue - } - if pinnedAuthID != "" && candidate.ID != pinnedAuthID { - continue - } - if disallowFreeAuth && isFreeCodexAuth(candidate) { - continue - } - if _, used := tried[candidate.ID]; used { - continue - } - if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { - continue - } - candidates = append(candidates, candidate) - } - if len(candidates) == 0 { - m.mu.RUnlock() - return nil, nil, &Error{Code: "auth_not_found", Message: "no auth available"} - } - available, errAvailable := m.availableAuthsForRouteModel(candidates, provider, model, time.Now()) - if errAvailable != nil { - m.mu.RUnlock() - return nil, nil, errAvailable - } - available = cloneAuthSlice(available) - m.mu.RUnlock() - - selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, provider, []string{provider}, model, opts, tried, available) - if errPick != nil { - return nil, nil, errPick - } - if !handled { - selected, errPick = selector.Pick(ctx, provider, selectionArgForSelector(selector, model), opts, available) - if errPick != nil { - return nil, nil, errPick - } - } - if selected == nil { - return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() - } - m.mu.Unlock() - } - return authCopy, executor, nil -} - -// SelectAuth selects one credential through the configured scheduling strategy. -// It does not execute or alter the selected credential's result state. -func (m *Manager) SelectAuth(ctx context.Context, provider, model string, opts cliproxyexecutor.Options) (*Auth, error) { - if m != nil && m.HomeEnabled() { - return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} - } - selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, nil) - if errPick != nil { - return nil, errPick - } - if m.HomeEnabled() { - return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} - } - return selected, nil -} - -// SelectAuthByKind selects one credential of the required kind through the -// configured scheduling strategy. Credentials of other kinds are skipped. -func (m *Manager) SelectAuthByKind(ctx context.Context, provider, model, requiredKind string, opts cliproxyexecutor.Options) (*Auth, error) { - if m != nil && m.HomeEnabled() { - return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} - } - requiredKind = normalizeAuthKind(requiredKind) - if requiredKind == "" { - return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} - } - - tried := make(map[string]struct{}) - for { - selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, tried) - if errPick != nil { - return nil, errPick - } - if selected == nil { - return nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if selected.AuthKind() == requiredKind { - if m.HomeEnabled() { - return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} - } - return selected, nil - } - authID := strings.TrimSpace(selected.ID) - if authID == "" { - return nil, &Error{Code: "auth_not_found", Message: "selected auth has no ID"} - } - if _, alreadyTried := tried[authID]; alreadyTried { - return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} - } - tried[authID] = struct{}{} - } -} - -// SelectHomeAuthByKind selects a Home dispatch while retaining its execution scope. -func (m *Manager) SelectHomeAuthByKind(ctx context.Context, provider string, model string, requiredKind string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { - requiredKind = normalizeAuthKind(requiredKind) - if requiredKind == "" { - return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} - } - if m == nil || !m.HomeEnabled() { - return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} - } - - homeAuthCount := homeAuthCountFromMetadata(opts.Metadata) - tried := make(map[string]struct{}) - for { - selectionOpts := withHomeAuthCount(opts, homeAuthCount) - selection, errSelection := m.pickHomeDispatchSelection(ctx, model, selectionOpts) - if errSelection != nil { - return nil, errSelection - } - providerMatches := strings.TrimSpace(provider) == "" || strings.EqualFold(strings.TrimSpace(selection.Provider), strings.TrimSpace(provider)) - kindMatches := selection.Auth != nil && selection.Auth.AuthKind() == requiredKind - if providerMatches && kindMatches { - return selection, nil - } - - authID := "" - if selection.Auth != nil { - authID = strings.TrimSpace(selection.Auth.ID) - } - reason := "auth_kind_mismatch" - if !providerMatches { - reason = "provider_mismatch" - } - if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, reason); errEnd != nil { - return nil, errEnd - } - if authID == "" { - return nil, &Error{Code: "auth_not_found", Message: "selected auth has no ID"} - } - if _, alreadyTried := tried[authID]; alreadyTried { - return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} - } - tried[authID] = struct{}{} - homeAuthCount++ - } -} - -func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { - if m.HomeEnabled() { - auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) - return auth, exec, err - } - - if m.hasPluginScheduler() || !m.useSchedulerFastPath() { - return m.pickNextLegacy(ctx, provider, model, opts, tried) - } - if strings.TrimSpace(model) != "" { - m.mu.RLock() - for _, candidate := range m.auths { - if candidate == nil || executorKeyFromAuth(candidate) != provider || candidate.Disabled { - continue - } - if _, used := tried[candidate.ID]; used { - continue - } - if m.routeAwareSelectionRequired(candidate, model) { - m.mu.RUnlock() - return m.pickNextLegacy(ctx, provider, model, opts, tried) - } - } - m.mu.RUnlock() - } - executor, okExecutor := m.Executor(provider) - if !okExecutor { - return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"} - } - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - for { - selected, errPick := m.scheduler.pickSingle(ctx, provider, model, opts, tried) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, errPick = m.scheduler.pickSingle(ctx, provider, model, opts, tried) - } - if errPick != nil { - return nil, nil, errPick - } - if selected == nil { - return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if disallowFreeAuth && isFreeCodexAuth(selected) { - if tried == nil { - tried = make(map[string]struct{}) - } - tried[selected.ID] = struct{}{} - continue - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() - } - m.mu.Unlock() - } - return authCopy, executor, nil - } -} - -func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { - if m.HomeEnabled() { - return m.pickNextViaHome(ctx, model, opts, tried) - } - - pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - - providerSet := make(map[string]struct{}, len(providers)) - for _, provider := range providers { - p := strings.TrimSpace(strings.ToLower(provider)) - if p == "" { - continue - } - providerSet[p] = struct{}{} - } - if len(providerSet) == 0 { - return nil, nil, "", &Error{Code: "provider_not_found", Message: "no provider supplied"} - } - - m.mu.RLock() - selector := m.selector - pluginScheduler := m.pluginScheduler - candidates := make([]*Auth, 0, len(m.auths)) - modelKey := strings.TrimSpace(model) - // Always use base model name (without thinking suffix) for auth matching. - if modelKey != "" { - parsed := thinking.ParseSuffix(modelKey) - if parsed.ModelName != "" { - modelKey = strings.TrimSpace(parsed.ModelName) - } - } - registryRef := registry.GetGlobalRegistry() - for _, candidate := range m.auths { - if candidate == nil || candidate.Disabled { - continue - } - if pinnedAuthID != "" && candidate.ID != pinnedAuthID { - continue - } - if disallowFreeAuth && isFreeCodexAuth(candidate) { - continue - } - providerKey := executorKeyFromAuth(candidate) - if providerKey == "" { - continue - } - if _, ok := providerSet[providerKey]; !ok { - continue - } - if _, used := tried[candidate.ID]; used { - continue - } - if _, ok := m.executors[providerKey]; !ok { - continue - } - if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { - continue - } - candidates = append(candidates, candidate) - } - if len(candidates) == 0 { - m.mu.RUnlock() - return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} - } - available, errAvailable := m.availableAuthsForRouteModel(candidates, "mixed", model, time.Now()) - if errAvailable != nil { - m.mu.RUnlock() - return nil, nil, "", errAvailable - } - available = cloneAuthSlice(available) - m.mu.RUnlock() - - selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, "mixed", providers, model, opts, tried, available) - if errPick != nil { - return nil, nil, "", errPick - } - if !handled { - selected, errPick = selector.Pick(ctx, "mixed", selectionArgForSelector(selector, model), opts, available) - if errPick != nil { - return nil, nil, "", errPick - } - } - if selected == nil { - return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - providerKey := executorKeyFromAuth(selected) - executor, okExecutor := m.Executor(providerKey) - if !okExecutor { - return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() - } - m.mu.Unlock() - } - return authCopy, executor, providerKey, nil -} - -func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { - if m.HomeEnabled() { - return m.pickNextViaHome(ctx, model, opts, tried) - } - - if m.hasPluginScheduler() || !m.useSchedulerFastPath() { - return m.pickNextMixedLegacy(ctx, providers, model, opts, tried) - } - - eligibleProviders := make([]string, 0, len(providers)) - seenProviders := make(map[string]struct{}, len(providers)) - for _, provider := range providers { - providerKey := strings.TrimSpace(strings.ToLower(provider)) - if providerKey == "" { - continue - } - if _, seen := seenProviders[providerKey]; seen { - continue - } - if _, okExecutor := m.Executor(providerKey); !okExecutor { - continue - } - seenProviders[providerKey] = struct{}{} - eligibleProviders = append(eligibleProviders, providerKey) - } - if len(eligibleProviders) == 0 { - return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} - } - if strings.TrimSpace(model) != "" { - providerSet := make(map[string]struct{}, len(eligibleProviders)) - for _, providerKey := range eligibleProviders { - providerSet[providerKey] = struct{}{} - } - m.mu.RLock() - for _, candidate := range m.auths { - if candidate == nil || candidate.Disabled { - continue - } - if _, ok := providerSet[executorKeyFromAuth(candidate)]; !ok { - continue - } - if _, used := tried[candidate.ID]; used { - continue - } - if m.routeAwareSelectionRequired(candidate, model) { - m.mu.RUnlock() - return m.pickNextMixedLegacy(ctx, providers, model, opts, tried) - } - } - m.mu.RUnlock() - } - - disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) - for { - selected, providerKey, errPick := m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) - if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { - m.syncScheduler() - selected, providerKey, errPick = m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) - } - if errPick != nil { - return nil, nil, "", errPick - } - if selected == nil { - return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} - } - if disallowFreeAuth && isFreeCodexAuth(selected) { - if tried == nil { - tried = make(map[string]struct{}) - } - tried[selected.ID] = struct{}{} - continue - } - executor, okExecutor := m.Executor(providerKey) - if !okExecutor { - return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} - } - authCopy := selected.Clone() - if !selected.indexAssigned { - m.mu.Lock() - if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { - current.EnsureIndex() - authCopy = current.Clone() - } - m.mu.Unlock() - } - return authCopy, executor, providerKey, nil - } -} - -type homeErrorEnvelope struct { - Error *homeErrorDetail `json:"error"` -} - -type homeErrorDetail struct { - Type string `json:"type"` - Message string `json:"message"` - Code string `json:"code,omitempty"` - Retryable bool `json:"retryable,omitempty"` - RetryAfterMS int64 `json:"retry_after_ms,omitempty"` -} - -const ( - homeUpstreamModelAttributeKey = "home_upstream_model" - homeForceMappingAttributeKey = "home_force_mapping" - homeOriginalAliasAttributeKey = "home_original_alias" - homeRequestRetryExceededErrorCode = "request_retry_exceeded" -) - -func isHomeRequestRetryExceededError(err error) bool { - var authErr *Error - if !errors.As(err, &authErr) || authErr == nil { - return false - } - return strings.EqualFold(strings.TrimSpace(authErr.Code), homeRequestRetryExceededErrorCode) -} - -func shouldReturnLastErrorOnPickFailure(homeMode bool, lastErr error, errPick error) bool { - if lastErr == nil { - return false - } - if !homeMode { - return true - } - return isHomeRequestRetryExceededError(errPick) -} - -func homeAuthAlreadyTried(tried map[string]struct{}, authID string) bool { - authID = strings.TrimSpace(authID) - if authID == "" || len(tried) == 0 { - return false - } - _, ok := tried[authID] - return ok -} - -func repeatedHomeAuthError() *Error { - return &Error{ - Code: homeRequestRetryExceededErrorCode, - Message: "home returned a previously tried auth", - HTTPStatus: http.StatusServiceUnavailable, - } -} - -type homeAuthDispatchResponse struct { - Model string `json:"model"` - Provider string `json:"provider"` - AuthIndex string `json:"auth_index"` - UserAPIKey string `json:"user_api_key"` - ForceMapping bool `json:"force_mapping"` - OriginalAlias string `json:"original_alias"` - Auth Auth `json:"auth"` -} - -type homeAuthDispatcher interface { - HeartbeatOK() bool - RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error) - AbortAmbiguousDispatch() -} - -var currentHomeDispatcher = func() homeAuthDispatcher { - return home.Current() -} - -func setHomeUserAPIKeyOnGinContext(ctx context.Context, apiKey string) { - apiKey = strings.TrimSpace(apiKey) - if apiKey == "" || ctx == nil { - return - } - ginCtx, ok := ctx.Value("gin").(interface{ Set(string, any) }) - if !ok || ginCtx == nil { - return - } - ginCtx.Set("userApiKey", apiKey) -} - -func homeDispatchHeaders(ctx context.Context, headers http.Header) http.Header { - apiKey, ok := homeQueryCredentialFromContext(ctx) - if !ok { - return headers - } - out := headers.Clone() - if out == nil { - out = http.Header{} - } - if out.Get("Authorization") != "" || out.Get("X-Goog-Api-Key") != "" || out.Get("X-Api-Key") != "" { - return out - } - out.Set("X-Goog-Api-Key", apiKey) - return out -} - -func homeQueryCredentialFromContext(ctx context.Context) (string, bool) { - if ctx == nil { - return "", false - } - if queryCtx, ok := ctx.Value("gin").(interface{ Query(string) string }); ok && queryCtx != nil { - if apiKey := strings.TrimSpace(queryCtx.Query("key")); apiKey != "" { - return apiKey, true - } - if apiKey := strings.TrimSpace(queryCtx.Query("auth_token")); apiKey != "" { - return apiKey, true - } - } - ginCtx, ok := ctx.Value("gin").(interface{ Get(string) (any, bool) }) - if !ok || ginCtx == nil { - return "", false - } - rawMetadata, ok := ginCtx.Get("accessMetadata") - if !ok { - return "", false - } - source := accessMetadataSource(rawMetadata) - if source != "query-key" && source != "query-auth-token" { - return "", false - } - rawAPIKey, ok := ginCtx.Get("userApiKey") - if !ok { - return "", false - } - apiKey := contextStringValue(rawAPIKey) - if apiKey == "" { - return "", false - } - return apiKey, true -} - -func accessMetadataSource(raw any) string { - switch v := raw.(type) { - case map[string]string: - return strings.TrimSpace(v["source"]) - case map[string]any: - return contextStringValue(v["source"]) - default: - return "" - } -} - -func contextStringValue(raw any) string { - switch v := raw.(type) { - case string: - return strings.TrimSpace(v) - case []byte: - return strings.TrimSpace(string(v)) - default: - return "" - } -} - -func homeExecutionSessionIDFromMetadata(meta map[string]any) string { - if len(meta) == 0 { - return "" - } - raw, ok := meta[cliproxyexecutor.ExecutionSessionMetadataKey] - if !ok || raw == nil { - return "" - } - switch value := raw.(type) { - case string: - return strings.TrimSpace(value) - case []byte: - return strings.TrimSpace(string(value)) - default: - return "" - } -} - -type homeSessionSelectionKey struct { - credentialID string - routeModel string -} - -func (m *Manager) lockHomeWebsocketSession(ctx context.Context, opts cliproxyexecutor.Options) func() { - if m == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { - return nil - } - sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) - if sessionID == "" { - return nil - } - lock, _ := m.homeSessionLocks.LoadOrStore(sessionID, &sync.Mutex{}) - mutex, ok := lock.(*sync.Mutex) - if !ok || mutex == nil { - return nil - } - mutex.Lock() - return mutex.Unlock -} - -func (m *Manager) retainedHomeSessionSelection(ctx context.Context, opts cliproxyexecutor.Options, model string) (*HomeDispatchSelection, bool, error) { - if m == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { - return nil, false, nil - } - sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) - credentialID := pinnedAuthIDFromMetadata(opts.Metadata) - if sessionID == "" { - return nil, false, nil - } - - routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) - var retained *HomeDispatchSelection - var ended []*HomeDispatchSelection - fallbackAttempt := homeAuthCountFromMetadata(opts.Metadata) > 1 - m.mu.Lock() - selections := m.homeSessionSelections[sessionID] - for key, selection := range selections { - if selection == nil { - delete(selections, key) - continue - } - matchesCredential := credentialID == "" || key.credentialID == credentialID - matchesRoute := validRouteModel && key.routeModel == routeModel - if !fallbackAttempt && matchesCredential && selection.Active() && matchesRoute && retained == nil { - retained = selection - continue - } - delete(selections, key) - ended = append(ended, selection) - } - if len(selections) == 0 { - delete(m.homeSessionSelections, sessionID) - } - m.mu.Unlock() - - for _, selection := range ended { - if errWait := m.endHomeSelectionBeforeRedispatch(ctx, selection, "target_changed"); errWait != nil { - return nil, false, errWait - } - } - return retained, retained != nil, nil -} - -func (m *Manager) predictedHomeConcurrencyModel(auth *Auth, routeModel string) (string, bool) { - requestedModel := rewriteModelForAuth(routeModel, auth) - aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel) - upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult) - if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) != 0 { - if len(pool) != 1 { - return "", false - } - upstreamModel = pool[0] - } else { - upstreamModel = m.applyAPIKeyModelAlias(auth, upstreamModel) - } - return validCanonicalHomeConcurrencyModelKey(upstreamModel) -} - -func (m *Manager) endMismatchedHomeSessionSelections(ctx context.Context, sessionID, credentialID, model string, waitForAck bool) error { - if m == nil || sessionID == "" { - return nil - } - routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) - var ended []*HomeDispatchSelection - m.mu.Lock() - selections := m.homeSessionSelections[sessionID] - for key, selection := range selections { - if selection == nil { - delete(selections, key) - continue - } - matchesRoute := validRouteModel && key.routeModel == routeModel - if key.credentialID == credentialID && matchesRoute { - continue - } - delete(selections, key) - ended = append(ended, selection) - } - if len(selections) == 0 { - delete(m.homeSessionSelections, sessionID) - } - m.mu.Unlock() - for _, selection := range ended { - if !waitForAck { - selection.End("target_changed") - continue - } - if errWait := m.endHomeSelectionBeforeRedispatch(ctx, selection, "target_changed"); errWait != nil { - return errWait - } - } - return nil -} - -func (m *Manager) endHomeSelectionBeforeRedispatch(ctx context.Context, selection *HomeDispatchSelection, reason string) error { - if selection == nil { - return nil - } - ticket := selection.EndWithRelease(reason) - if ticket == nil { - return nil - } - - bound := internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound - if m != nil { - if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil { - bound = cfg.CredentialConcurrency.WithDefaults().CPACancelBound - } - } - waitCtx := ctx - if waitCtx == nil { - waitCtx = context.Background() - } - waitCtx, cancelWait := context.WithTimeout(waitCtx, bound) - defer cancelWait() - if errWait := ticket.Wait(waitCtx); errWait != nil { - return &Error{Code: "home_unavailable", Message: "Home did not acknowledge credential release: " + errWait.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} - } - return nil -} - -func (m *Manager) retainHomeWebsocketSelection(ctx context.Context, opts cliproxyexecutor.Options, model string, selection *HomeDispatchSelection) bool { - if m == nil || selection == nil || !selection.Retained() || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil { - return false - } - sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) - credentialID := strings.TrimSpace(selection.Auth.ID) - routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) - if selection.accountedModel == "" { - selection.accountedModel, _ = m.predictedHomeConcurrencyModel(selection.Auth, model) - } - if sessionID == "" || credentialID == "" || !validRouteModel || selection.accountedModel == "" { - return false - } - _ = m.endMismatchedHomeSessionSelections(ctx, sessionID, credentialID, routeModel, false) - key := homeSessionSelectionKey{credentialID: credentialID, routeModel: routeModel} - m.mu.Lock() - if m.homeSessionSelections == nil { - m.homeSessionSelections = make(map[string]map[homeSessionSelectionKey]*HomeDispatchSelection) - } - selections := m.homeSessionSelections[sessionID] - if selections == nil { - selections = make(map[homeSessionSelectionKey]*HomeDispatchSelection) - m.homeSessionSelections[sessionID] = selections - } - previous := selections[key] - selections[key] = selection - m.mu.Unlock() - m.rememberHomeRuntimeAuth(sessionID, selection.Auth) - if previous != nil && previous != selection { - previous.End("target_replaced") - } - return true -} - -func (m *Manager) clearHomeSessionLocks() { - if m == nil { - return - } - m.homeSessionLocks.Range(func(key, _ any) bool { - m.homeSessionLocks.Delete(key) - return true - }) -} - -func (m *Manager) takeHomeSessionSelectionsLocked(sessionID string) []*HomeDispatchSelection { - if m == nil { - return nil - } - selections := m.homeSessionSelections[sessionID] - delete(m.homeSessionSelections, sessionID) - result := make([]*HomeDispatchSelection, 0, len(selections)) - for _, selection := range selections { - result = append(result, selection) - } - return result -} - -func (m *Manager) takeAllHomeSessionSelectionsLocked() []*HomeDispatchSelection { - if m == nil { - return nil - } - result := make([]*HomeDispatchSelection, 0) - for sessionID, selections := range m.homeSessionSelections { - delete(m.homeSessionSelections, sessionID) - for _, selection := range selections { - result = append(result, selection) - } - } - return result -} - -func (m *Manager) clearHomeRuntimeAuths() { - if m == nil { - return - } - m.mu.Lock() - m.clearHomeRuntimeAuthsLocked() - selections := m.takeAllHomeSessionSelectionsLocked() - m.mu.Unlock() - for _, selection := range selections { - selection.End("home_disabled") - } -} - -func (m *Manager) clearHomeRuntimeAuthsLocked() { - if m == nil { - return - } - m.homeRuntimeAuths = make(map[string]map[string]*Auth) - m.homeRuntimeAuthOwners = make(map[string]map[string]*HomeDispatchSelection) -} - -func (m *Manager) clearHomeRuntimeAuthsForSessionLocked(sessionID string) { - sessionID = strings.TrimSpace(sessionID) - if m == nil || sessionID == "" { - return - } - delete(m.homeRuntimeAuths, sessionID) - delete(m.homeRuntimeAuthOwners, sessionID) -} - -func (m *Manager) bindHomeSelectionRuntimeAuth(ctx context.Context, opts cliproxyexecutor.Options, selection *HomeDispatchSelection) error { - if m == nil || selection == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil || !authWebsocketsEnabled(selection.Auth) { - return nil - } - sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) - authID := strings.TrimSpace(selection.Auth.ID) - if sessionID == "" || authID == "" || !selection.runtimeAuthBound.CompareAndSwap(false, true) { - return nil - } - m.rememberHomeSelectionRuntimeAuth(sessionID, selection) - if errBind := selection.Bind(func() error { - m.forgetHomeRuntimeAuth(sessionID, authID, selection) - return nil - }); errBind != nil { - selection.runtimeAuthBound.Store(false) - m.forgetHomeRuntimeAuth(sessionID, authID, selection) - return errBind - } - return nil -} - -func (m *Manager) rememberHomeSelectionRuntimeAuth(sessionID string, selection *HomeDispatchSelection) { - if m == nil || selection == nil || selection.Auth == nil { - return - } - sessionID = strings.TrimSpace(sessionID) - authID := strings.TrimSpace(selection.Auth.ID) - if sessionID == "" || authID == "" { - return - } - m.mu.Lock() - if m.homeRuntimeAuths == nil { - m.homeRuntimeAuths = make(map[string]map[string]*Auth) - } - if m.homeRuntimeAuthOwners == nil { - m.homeRuntimeAuthOwners = make(map[string]map[string]*HomeDispatchSelection) - } - if m.homeRuntimeAuths[sessionID] == nil { - m.homeRuntimeAuths[sessionID] = make(map[string]*Auth) - } - if m.homeRuntimeAuthOwners[sessionID] == nil { - m.homeRuntimeAuthOwners[sessionID] = make(map[string]*HomeDispatchSelection) - } - m.homeRuntimeAuths[sessionID][authID] = selection.Auth.Clone() - m.homeRuntimeAuthOwners[sessionID][authID] = selection - m.mu.Unlock() -} - -func (m *Manager) forgetHomeRuntimeAuth(sessionID string, authID string, owner *HomeDispatchSelection) { - sessionID = strings.TrimSpace(sessionID) - authID = strings.TrimSpace(authID) - if m == nil || sessionID == "" || authID == "" { - return - } - m.mu.Lock() - owners := m.homeRuntimeAuthOwners[sessionID] - if owner != nil && owners[authID] != owner { - m.mu.Unlock() - return - } - sessionAuths := m.homeRuntimeAuths[sessionID] - delete(sessionAuths, authID) - delete(owners, authID) - if len(sessionAuths) == 0 { - delete(m.homeRuntimeAuths, sessionID) - } - if len(owners) == 0 { - delete(m.homeRuntimeAuthOwners, sessionID) - } - m.mu.Unlock() -} - -func (m *Manager) rememberHomeRuntimeAuth(sessionID string, auth *Auth) { - sessionID = strings.TrimSpace(sessionID) - authID := "" - if auth != nil { - authID = strings.TrimSpace(auth.ID) - } - if m == nil || auth == nil || sessionID == "" || authID == "" || !authWebsocketsEnabled(auth) { - return - } - m.mu.Lock() - if m.homeRuntimeAuths == nil { - m.homeRuntimeAuths = make(map[string]map[string]*Auth) - } - sessionAuths := m.homeRuntimeAuths[sessionID] - if sessionAuths == nil { - sessionAuths = make(map[string]*Auth) - m.homeRuntimeAuths[sessionID] = sessionAuths - } - sessionAuths[authID] = auth.Clone() - m.mu.Unlock() -} - -func (m *Manager) homeRuntimeAuthByID(sessionID string, authID string) (*Auth, ProviderExecutor, string, bool) { - sessionID = strings.TrimSpace(sessionID) - authID = strings.TrimSpace(authID) - if m == nil || sessionID == "" || authID == "" { - return nil, nil, "", false - } - m.mu.RLock() - sessionAuths := m.homeRuntimeAuths[sessionID] - auth := sessionAuths[authID] - m.mu.RUnlock() - if auth == nil || !authWebsocketsEnabled(auth) { - return nil, nil, "", false - } - logicalProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) - executorKey := executorKeyFromAuth(auth) - if logicalProvider == "" || executorKey == "" { - return nil, nil, "", false - } - executor, ok := m.Executor(executorKey) - if !ok && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" { - executor, ok = m.Executor("openai-compatibility") - } - if !ok { - return nil, nil, "", false - } - return auth.Clone(), executor, logicalProvider, true -} - -func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { - if m == nil { - return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} - } - if ctx == nil { - ctx = context.Background() - } - selection, errSelection := m.pickHomeDispatchSelection(ctx, model, opts) - if errSelection != nil { - return nil, nil, "", errSelection - } - if selection.Auth == nil || homeAuthAlreadyTried(tried, selection.Auth.ID) { - selection.End("repeated_auth") - return nil, nil, "", repeatedHomeAuthError() - } - auth := selection.CloneAuthForRoute(model) - executor := selection.Executor - provider := selection.Provider - selection.End("legacy_selection_unbound") - return auth, executor, provider, nil -} - -func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { - if m == nil { - return nil, &Error{Code: "auth_not_found", Message: "no auth available"} - } - if ctx == nil { - ctx = context.Background() - } - - requestedModel := strings.TrimSpace(model) - if requestedModel == "" { - requestedModel = requestedModelFromMetadata(opts.Metadata, model) - } - retained, retainedOK, errRetained := m.retainedHomeSessionSelection(ctx, opts, requestedModel) - if errRetained != nil { - return nil, errRetained - } - if retainedOK { - return retained, nil - } - if sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata); sessionID != "" { - if credentialID := pinnedAuthIDFromMetadata(opts.Metadata); credentialID != "" { - if errEnd := m.endMismatchedHomeSessionSelections(ctx, sessionID, credentialID, requestedModel, true); errEnd != nil { - return nil, errEnd - } - } - } - - bundle := m.HomeDispatchBundle() - if bundle == nil || bundle.client == nil || bundle.registry == nil { - return nil, &Error{Code: "home_unavailable", Message: "home dispatch bundle unavailable", HTTPStatus: http.StatusServiceUnavailable} - } - client := bundle.client - registry := bundle.registry - if !client.HeartbeatOK() { - return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} - } - pending, errBegin := registry.BeginDispatch() - if errBegin != nil { - return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} - } - - sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata) - dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers) - raw, errRPop := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) - if errRPop != nil { - if home.IsAmbiguousDispatchError(errRPop) { - client.AbortAmbiguousDispatch() - } - pending.End() - if errors.Is(errRPop, home.ErrAuthNotFound) { - return nil, &Error{Code: "auth_not_found", Message: errRPop.Error(), HTTPStatus: http.StatusServiceUnavailable} - } - return nil, &Error{Code: "home_unavailable", Message: errRPop.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} - } - - envelope, errEnvelope := decodeHomeDispatchConcurrencyEnvelope(raw) - if errEnvelope != nil { - if envelope.Present { - client.AbortAmbiguousDispatch() - } - pending.End() - if envelope.Present { - return nil, invalidHomeConcurrencyResponse("Home returned malformed concurrency tuple") - } - return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} - } - - kind := "http" - if cliproxyexecutor.DownstreamWebsocket(ctx) { - kind = "websocket" - } else if opts.Stream { - kind = "stream" - } - baseScope := executionregistry.ScopeSpec{ - RequestID: logging.GetRequestID(ctx), - Model: requestedModel, - Kind: kind, - StartedAt: time.Now(), - } - var scope *executionregistry.Scope - if envelope.Present { - var errInstall error - scope, errInstall = installHomeConcurrencyScope(registry, pending, envelope.Tuple, baseScope) - if errInstall != nil { - client.AbortAmbiguousDispatch() - pending.End() - return nil, homeConcurrencyInstallError(errInstall) - } - } - endScope := func() { - if scope != nil { - scope.End("local_validation_failed") - return - } - pending.End() - } - if errHome := decodeHomeDispatchError(raw); errHome != nil { - if envelope.Present { - client.AbortAmbiguousDispatch() - endScope() - return nil, invalidHomeConcurrencyResponse("Home returned both accounted concurrency and an error") - } - pending.End() - return nil, errHome - } - - var dispatch homeAuthDispatchResponse - if errUnmarshal := json.Unmarshal(raw, &dispatch); errUnmarshal != nil { - endScope() - return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} - } - auth := dispatch.Auth - if strings.TrimSpace(auth.ID) == "" { - // Backward compatibility: older Home instances returned the auth directly. - if errUnmarshal := json.Unmarshal(raw, &auth); errUnmarshal != nil { - endScope() - return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} - } - } - observedModel := canonicalHomeDispatchModel(dispatch.Model, requestedModel) - if envelope.Present { - observedConcurrencyModel, validModel := validCanonicalHomeConcurrencyModelKey(observedModel) - if !validModel || envelope.Tuple.Model != observedConcurrencyModel { - client.AbortAmbiguousDispatch() - endScope() - return nil, invalidHomeConcurrencyResponse("Home concurrency model does not match dispatched model") - } - } - if !envelope.Present { - baseScope.Model = observedModel - } - - setHomeUserAPIKeyOnGinContext(ctx, dispatch.UserAPIKey) - if upstreamModel := strings.TrimSpace(dispatch.Model); upstreamModel != "" { - if auth.Attributes == nil { - auth.Attributes = make(map[string]string, 3) - } - auth.Attributes[homeUpstreamModelAttributeKey] = upstreamModel - } - if originalAlias := strings.TrimSpace(dispatch.OriginalAlias); dispatch.ForceMapping && originalAlias != "" { - if auth.Attributes == nil { - auth.Attributes = make(map[string]string, 2) - } - auth.Attributes[homeForceMappingAttributeKey] = "true" - auth.Attributes[homeOriginalAliasAttributeKey] = originalAlias - } - if strings.TrimSpace(auth.ID) == "" { - endScope() - return nil, &Error{Code: "invalid_auth", Message: "home returned auth without id", HTTPStatus: http.StatusBadGateway} - } - if errIdentity := verifyAccountedHomeConcurrencyIdentity(envelope.Tuple, &auth, dispatch.AuthIndex); errIdentity != nil { - endScope() - return nil, errIdentity - } - logicalProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) - executorKey := executorKeyFromAuth(&auth) - if logicalProvider == "" || executorKey == "" { - endScope() - return nil, &Error{Code: "invalid_auth", Message: "home returned auth without provider", HTTPStatus: http.StatusBadGateway} - } - - homeAuthIndex := strings.TrimSpace(dispatch.AuthIndex) - if homeAuthIndex != "" { - auth.Index = homeAuthIndex - auth.indexAssigned = true - } else { - auth.EnsureIndex() - } - - executor, okExecutor := m.Executor(executorKey) - if !okExecutor && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" { - executor, okExecutor = m.Executor("openai-compatibility") - } - if !okExecutor { - endScope() - return nil, &Error{Code: "executor_not_found", Message: "executor not registered", HTTPStatus: http.StatusBadGateway} - } - if scope == nil { - var errInstall error - scope, errInstall = installHomeConcurrencyScope(registry, pending, homeConcurrencyTuple{}, executionregistry.ScopeSpec{ - RequestID: baseScope.RequestID, - CredentialID: strings.TrimSpace(auth.ID), - Model: baseScope.Model, - Kind: baseScope.Kind, - StartedAt: baseScope.StartedAt, - }) - if errInstall != nil { - client.AbortAmbiguousDispatch() - pending.End() - return nil, homeConcurrencyInstallError(errInstall) - } - } - - selection, errSelection := newHomeDispatchSelection(auth.Clone(), executor, logicalProvider, scope) - if errSelection != nil { - endScope() - return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} - } - if envelope.Present { - selection.accountedModel = envelope.Tuple.Model - } - if executionSessionID := homeExecutionSessionIDFromMetadata(opts.Metadata); executionSessionID != "" && cliproxyexecutor.DownstreamWebsocket(ctx) { - if errEnd := m.endMismatchedHomeSessionSelections(ctx, executionSessionID, strings.TrimSpace(auth.ID), requestedModel, true); errEnd != nil { - selection.End("target_change_release_failed") - return nil, errEnd - } - } - return selection, nil -} - -func requestedModelFromMetadata(metadata map[string]any, fallback string) string { - if metadata != nil { - if v, ok := metadata[cliproxyexecutor.RequestedModelMetadataKey]; ok { - switch typed := v.(type) { - case string: - if trimmed := strings.TrimSpace(typed); trimmed != "" { - return trimmed - } - case []byte: - if trimmed := strings.TrimSpace(string(typed)); trimmed != "" { - return trimmed - } - } - } - } - fallback = strings.TrimSpace(fallback) - if fallback == "" { - return "unknown" - } - return fallback -} - -func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, routeModel string, opts cliproxyexecutor.Options) ([]creditsCandidateEntry, error) { - if m == nil || !m.localExecutionAllowed() { - return nil, nil - } - pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - var candidates []creditsCandidateEntry - m.mu.RLock() - for _, auth := range m.auths { - if auth == nil || auth.Disabled || auth.Status == StatusDisabled { - continue - } - if pinnedAuthID != "" && auth.ID != pinnedAuthID { - continue - } - if !strings.EqualFold(strings.TrimSpace(auth.Provider), "antigravity") { - continue - } - if !strings.Contains(strings.ToLower(strings.TrimSpace(routeModel)), "claude") { - continue - } - providerKey := executorKeyFromAuth(auth) - executor, ok := m.executors[providerKey] - if !ok { - continue - } - candidates = append(candidates, creditsCandidateEntry{ - auth: auth.Clone(), - executor: executor, - provider: providerKey, - }) - } - m.mu.RUnlock() - - var known []creditsCandidateEntry - var unknown []creditsCandidateEntry - for _, candidate := range candidates { - hint, okHint, errHint := GetAntigravityCreditsHintRequired(ctx, candidate.auth.ID) - if errHint != nil { - return nil, antigravityCreditsKVUnavailableError(errHint) - } - if okHint && hint.Known { - if !hint.Available { - continue - } - known = append(known, candidate) - continue - } - unknown = append(unknown, candidate) - } - sort.Slice(known, func(i, j int) bool { - return known[i].auth.ID < known[j].auth.ID - }) - sort.Slice(unknown, func(i, j int) bool { - return unknown[i].auth.ID < unknown[j].auth.ID - }) - return append(known, unknown...), nil -} - -type creditsCandidateEntry struct { - auth *Auth - executor ProviderExecutor - provider string -} - -func hasAntigravityProvider(providers []string) bool { - for _, p := range providers { - if strings.EqualFold(strings.TrimSpace(p), "antigravity") { - return true - } - } - return false -} - -func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, providers []string) bool { - status := statusCodeFromError(lastErr) - log.WithFields(log.Fields{ - "lastErr": errorString(lastErr), - "status": status, - "providers": providers, - }).Debug("shouldAttemptAntigravityCreditsFallback") - if m == nil || lastErr == nil || m.HomeEnabled() { - return false - } - cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - if cfg == nil || !cfg.QuotaExceeded.AntigravityCredits { - return false - } - switch status { - case http.StatusTooManyRequests, http.StatusServiceUnavailable: - return true - case 0: - var authErr *Error - if errors.As(lastErr, &authErr) && authErr != nil { - return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" || authErr.Code == "model_cooldown" - } - var cooldownErr *modelCooldownError - if errors.As(lastErr, &cooldownErr) { - return true - } - return false - default: - return false - } -} - -func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, bool, error) { - if m != nil && m.HomeEnabled() { - return cliproxyexecutor.Response{}, false, &Error{Code: "home_fallback_unsupported", Message: "Home does not support Antigravity credits fallback", HTTPStatus: http.StatusServiceUnavailable} - } - if !m.localExecutionAllowed() { - return cliproxyexecutor.Response{}, false, nil - } - routeModel := req.Model - candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) - if errCandidates != nil { - return cliproxyexecutor.Response{}, false, errCandidates - } - for _, c := range candidates { - if ctx.Err() != nil { - return cliproxyexecutor.Response{}, false, nil - } - creditsCtx := WithAntigravityCredits(ctx) - if rt := m.roundTripperFor(c.auth); rt != nil { - creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt) - creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt) - } - creditsOpts := ensureRequestedModelMetadata(opts, routeModel) - creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel) - preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) - if errPrepare != nil { - continue - } - c.auth = preparedAuth - publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth) - models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) - if len(models) == 0 { - continue - } - for _, upstreamModel := range models { - resultModel := m.stateModelForExecution(c.auth, routeModel, upstreamModel, pooled) - execReq := req - execReq.Model = upstreamModel - resp, errExec := c.executor.Execute(creditsCtx, c.auth, execReq, creditsOpts) - result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, Success: errExec == nil} - if errExec != nil { - result.Error = resultErrorFromError(errExec) - if ra := retryAfterFromError(errExec); ra != nil { - result.RetryAfter = ra - } - m.MarkResult(creditsCtx, result) - continue - } - m.MarkResult(creditsCtx, result) - rewriteForceMappedResponse(&resp, aliasResult) - return resp, true, nil - } - } - return cliproxyexecutor.Response{}, false, nil -} - -func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, bool, error) { - if m != nil && m.HomeEnabled() { - return nil, false, &Error{Code: "home_fallback_unsupported", Message: "Home does not support Antigravity credits fallback", HTTPStatus: http.StatusServiceUnavailable} - } - if !m.localExecutionAllowed() { - return nil, false, nil - } - routeModel := req.Model - candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) - if errCandidates != nil { - return nil, false, errCandidates - } - for _, c := range candidates { - if ctx.Err() != nil { - return nil, false, nil - } - creditsCtx := WithAntigravityCredits(ctx) - if rt := m.roundTripperFor(c.auth); rt != nil { - creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt) - creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt) - } - creditsOpts := ensureRequestedModelMetadata(opts, routeModel) - preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) - if errPrepare != nil { - continue - } - c.auth = preparedAuth - publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth) - models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) - if len(models) == 0 { - continue - } - result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, "", models, pooled, aliasResult, true, false) - if errStream != nil { - continue - } - return result, true, nil - } - return nil, false, nil -} - -func antigravityCreditsKVUnavailableError(cause error) error { - if cause == nil { - return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable", HTTPStatus: http.StatusServiceUnavailable} - } - return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable: " + cause.Error(), HTTPStatus: http.StatusServiceUnavailable} -} - -func (m *Manager) persist(ctx context.Context, auth *Auth) error { - if m.store == nil || auth == nil { - return nil - } - if shouldSkipPersist(ctx) { - return nil - } - if IsConfigAPIKeyAuth(auth) { - return nil - } - if auth.Attributes != nil { - if v := strings.ToLower(strings.TrimSpace(auth.Attributes["runtime_only"])); v == "true" { - return nil - } - } - if IsPluginVirtualAuth(auth) { - return nil - } - // Skip persistence when metadata is absent (e.g., runtime-only auths). - if auth.Metadata == nil { - return nil - } - _, err := m.store.Save(ctx, auth) - return err -} - -// StartAutoRefresh launches a background loop that evaluates auth freshness -// every few seconds and triggers refresh operations when required. -// Only one loop is kept alive; starting a new one cancels the previous run. -func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duration) { - if interval <= 0 { - interval = refreshCheckInterval - } - - m.mu.Lock() - cancelPrev := m.refreshCancel - m.refreshCancel = nil - m.refreshLoop = nil - m.mu.Unlock() - if cancelPrev != nil { - cancelPrev() - } - - ctx, cancelCtx := context.WithCancel(parent) - workers := refreshMaxConcurrency - if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil && cfg.AuthAutoRefreshWorkers > 0 { - workers = cfg.AuthAutoRefreshWorkers - } - loop := newAuthAutoRefreshLoop(m, interval, workers) - - m.mu.Lock() - m.refreshCancel = cancelCtx - m.refreshLoop = loop - m.mu.Unlock() - - loop.rebuild(time.Now()) - go loop.run(ctx) -} - -// StopAutoRefresh cancels the background refresh loop, if running. -// It also stops the selector if it implements StoppableSelector. -func (m *Manager) StopAutoRefresh() { - m.mu.Lock() - cancel := m.refreshCancel - m.refreshCancel = nil - m.refreshLoop = nil - m.mu.Unlock() - if cancel != nil { - cancel() - } - // Stop selector if it implements StoppableSelector (e.g., SessionAffinitySelector) - if stoppable, ok := m.selector.(StoppableSelector); ok { - stoppable.Stop() - } -} - -func (m *Manager) queueRefreshReschedule(authID string) { - if m == nil || authID == "" { - return - } - m.mu.RLock() - loop := m.refreshLoop - m.mu.RUnlock() - if loop == nil { - return - } - loop.queueReschedule(authID) -} - -func (m *Manager) queueRefreshUnschedule(authID string) { - if m == nil || authID == "" { - return - } - m.mu.RLock() - loop := m.refreshLoop - m.mu.RUnlock() - if loop == nil { - return - } - loop.remove(authID) -} - -func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool { - if a == nil { - return false - } - if hasUnauthorizedAuthFailure(a) { - return false - } - if !a.NextRefreshAfter.IsZero() && now.Before(a.NextRefreshAfter) { - return false - } - if evaluator, ok := a.Runtime.(RefreshEvaluator); ok && evaluator != nil { - return evaluator.ShouldRefresh(now, a) - } - - lastRefresh := a.LastRefreshedAt - if lastRefresh.IsZero() { - if ts, ok := authLastRefreshTimestamp(a); ok { - lastRefresh = ts - } - } - - expiry, hasExpiry := a.ExpirationTime() - - if interval := authPreferredInterval(a); interval > 0 { - if hasExpiry && !expiry.IsZero() { - if !expiry.After(now) { - return true - } - if expiry.Sub(now) <= interval { - return true - } - } - if lastRefresh.IsZero() { - return true - } - return now.Sub(lastRefresh) >= interval - } - - provider := strings.ToLower(a.Provider) - lead := ProviderRefreshLead(provider, a.Runtime) - if lead == nil { - return false - } - if *lead <= 0 { - if hasExpiry && !expiry.IsZero() { - return now.After(expiry) - } - return false - } - if hasExpiry && !expiry.IsZero() { - return time.Until(expiry) <= *lead - } - if !lastRefresh.IsZero() { - return now.Sub(lastRefresh) >= *lead - } - return true -} - -func authPreferredInterval(a *Auth) time.Duration { - if a == nil { - return 0 - } - if d := durationFromMetadata(a.Metadata, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 { - return d - } - if d := durationFromAttributes(a.Attributes, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 { - return d - } - return 0 -} - -func durationFromMetadata(meta map[string]any, keys ...string) time.Duration { - if len(meta) == 0 { - return 0 - } - for _, key := range keys { - if val, ok := meta[key]; ok { - if dur := parseDurationValue(val); dur > 0 { - return dur - } - } - } - return 0 -} - -func durationFromAttributes(attrs map[string]string, keys ...string) time.Duration { - if len(attrs) == 0 { - return 0 - } - for _, key := range keys { - if val, ok := attrs[key]; ok { - if dur := parseDurationString(val); dur > 0 { - return dur - } - } - } - return 0 -} - -func parseDurationValue(val any) time.Duration { - switch v := val.(type) { - case time.Duration: - if v <= 0 { - return 0 - } - return v - case int: - if v <= 0 { - return 0 - } - return time.Duration(v) * time.Second - case int32: - if v <= 0 { - return 0 - } - return time.Duration(v) * time.Second - case int64: - if v <= 0 { - return 0 - } - return time.Duration(v) * time.Second - case uint: - if v == 0 { - return 0 - } - return time.Duration(v) * time.Second - case uint32: - if v == 0 { - return 0 - } - return time.Duration(v) * time.Second - case uint64: - if v == 0 { - return 0 - } - return time.Duration(v) * time.Second - case float32: - if v <= 0 { - return 0 - } - return time.Duration(float64(v) * float64(time.Second)) - case float64: - if v <= 0 { - return 0 - } - return time.Duration(v * float64(time.Second)) - case json.Number: - if i, err := v.Int64(); err == nil { - if i <= 0 { - return 0 - } - return time.Duration(i) * time.Second - } - if f, err := v.Float64(); err == nil && f > 0 { - return time.Duration(f * float64(time.Second)) - } - case string: - return parseDurationString(v) - } - return 0 -} - -func parseDurationString(raw string) time.Duration { - s := strings.TrimSpace(raw) - if s == "" { - return 0 - } - if dur, err := time.ParseDuration(s); err == nil && dur > 0 { - return dur - } - if secs, err := strconv.ParseFloat(s, 64); err == nil && secs > 0 { - return time.Duration(secs * float64(time.Second)) - } - return 0 -} - -func authLastRefreshTimestamp(a *Auth) (time.Time, bool) { - if a == nil { - return time.Time{}, false - } - if a.Metadata != nil { - if ts, ok := lookupMetadataTime(a.Metadata, "last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"); ok { - return ts, true - } - } - if a.Attributes != nil { - for _, key := range []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"} { - if val := strings.TrimSpace(a.Attributes[key]); val != "" { - if ts, ok := parseTimeValue(val); ok { - return ts, true - } - } - } - } - return time.Time{}, false -} - -func lookupMetadataTime(meta map[string]any, keys ...string) (time.Time, bool) { - for _, key := range keys { - if val, ok := meta[key]; ok { - if ts, ok1 := parseTimeValue(val); ok1 { - return ts, true - } - } - } - return time.Time{}, false -} - -func (m *Manager) markRefreshPending(id string, now time.Time) bool { - m.mu.Lock() - auth, ok := m.auths[id] - if !ok || auth == nil { - m.mu.Unlock() - return false - } - if !auth.NextRefreshAfter.IsZero() && now.Before(auth.NextRefreshAfter) { - m.mu.Unlock() - return false - } - auth.NextRefreshAfter = now.Add(refreshPendingBackoff) - m.auths[id] = auth - m.mu.Unlock() - - m.queueRefreshReschedule(id) - return true -} - -type authRefreshLock struct { - mu sync.Mutex -} - -func authAccessToken(auth *Auth) string { - if token := authMetadataString(auth, "access_token"); token != "" { - return token - } - return authMetadataString(auth, "accessToken") -} - -func authHasRefreshCredential(auth *Auth) bool { - if authMetadataString(auth, "refresh_token") != "" { - return true - } - return authMetadataString(auth, "refreshToken") != "" -} - -func clearUnauthorizedModelStates(auth *Auth, now time.Time) []string { - if auth == nil || len(auth.ModelStates) == 0 { - return nil - } - var resumed []string - for model, state := range auth.ModelStates { - if state == nil || state.LastError == nil { - continue - } - if state.LastError.StatusCode() != http.StatusUnauthorized && !strings.EqualFold(state.LastError.Code, "unauthorized") { - continue - } - resetModelState(state, now) - resumed = append(resumed, model) - } - if len(resumed) > 0 { - updateAggregatedAvailability(auth, now) - } - return resumed -} - -// tryRefreshAfterUnauthorized refreshes OAuth credentials once after a 401 so the -// current auth can be retried before fallback/suspend. -func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, execErr error, alreadyTried bool) (*Auth, bool) { - if m == nil || auth == nil || alreadyTried || execErr == nil { - return auth, false - } - if !isUnauthorizedError(execErr) || !authHasRefreshCredential(auth) { - return auth, false - } - log.Debugf("unauthorized response for %s (%s), refreshing credentials before fallback", auth.Provider, auth.ID) - refreshed, errRefresh := m.refreshAuthForRequest(ctx, auth.ID, authAccessToken(auth)) - if errRefresh != nil || refreshed == nil { - log.Debugf("credential refresh before fallback failed for %s (%s): %v", auth.Provider, auth.ID, errRefresh) - return auth, false - } - return refreshed, true -} - -func (m *Manager) refreshAuth(ctx context.Context, id string) { - _, _ = m.refreshAuthForRequest(ctx, id, "") -} - -// refreshAuthForRequest performs a synchronous credential refresh for the given auth. -// failedAccessToken lets concurrent callers reuse a refresh that already replaced the -// access token that produced the unauthorized response. -func (m *Manager) refreshAuthForRequest(ctx context.Context, id, failedAccessToken string) (*Auth, error) { - if m == nil { - return nil, errors.New("auth manager is nil") - } - if ctx == nil { - ctx = context.Background() - } - id = strings.TrimSpace(id) - if id == "" { - return nil, errors.New("auth id is empty") - } - - lockValue, _ := m.refreshLocks.LoadOrStore(id, &authRefreshLock{}) - lock, _ := lockValue.(*authRefreshLock) - if lock == nil { - lock = &authRefreshLock{} - m.refreshLocks.Store(id, lock) - } - lock.mu.Lock() - defer lock.mu.Unlock() - - m.mu.RLock() - auth := m.auths[id] - var exec ProviderExecutor - if auth != nil { - exec = m.executors[auth.Provider] - } - m.mu.RUnlock() - if auth == nil || exec == nil { - return nil, errors.New("auth or executor not found") - } - - // Another request may already have refreshed this credential. - if failedAccessToken != "" { - if currentToken := authAccessToken(auth); currentToken != "" && currentToken != failedAccessToken { - return auth.Clone(), nil - } - } - - cloned := auth.Clone() - updated, err := exec.Refresh(ctx, cloned) - if err != nil && errors.Is(err, context.Canceled) { - log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID) - return nil, err - } - log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err) - now := time.Now() - if err != nil { - unauthorized := isUnauthorizedError(err) - shouldReschedule := false - m.mu.Lock() - if current := m.auths[id]; current != nil { - current.LastError = refreshErrorFromError(err) - if unauthorized { - current.NextRefreshAfter = time.Time{} - current.Unavailable = true - current.Status = StatusError - current.StatusMessage = "unauthorized" - } else { - current.NextRefreshAfter = now.Add(refreshFailureBackoff) - } - m.auths[id] = current - shouldReschedule = true - if m.scheduler != nil { - m.scheduler.upsertAuth(current.Clone()) - } - } - m.mu.Unlock() - if shouldReschedule { - m.queueRefreshReschedule(id) - } - return nil, err - } - if updated == nil { - updated = cloned - } - // Preserve runtime created by the executor during Refresh. - // If executor didn't set one, fall back to the previous runtime. - if updated.Runtime == nil { - updated.Runtime = auth.Runtime - } - updated.LastRefreshedAt = now - updated.NextRefreshAfter = time.Time{} - updated.LastError = nil - updated.StatusMessage = "" - updated.Unavailable = false - if updated.Status == StatusError { - updated.Status = StatusActive - } - updated.UpdatedAt = now - modelsToResume := clearUnauthorizedModelStates(updated, now) - if m.shouldRefresh(updated, now) { - updated.NextRefreshAfter = now.Add(refreshIneffectiveBackoff) - } - saved, errUpdate := m.Update(ctx, updated) - for _, model := range modelsToResume { - registry.GetGlobalRegistry().ResumeClientModel(id, model) - } - if errUpdate != nil { - log.Debugf("persist refreshed auth %s (%s) failed: %v", auth.Provider, auth.ID, errUpdate) - } - if saved != nil { - return saved, nil - } - return updated.Clone(), nil -} - -func (m *Manager) executorFor(provider string) ProviderExecutor { - m.mu.RLock() - defer m.mu.RUnlock() - return m.executors[provider] -} - -// roundTripperContextKey is an unexported context key type to avoid collisions. -type roundTripperContextKey struct{} - -// roundTripperFor retrieves an HTTP RoundTripper for the given auth if a provider is registered. -func (m *Manager) roundTripperFor(auth *Auth) http.RoundTripper { - m.mu.RLock() - p := m.rtProvider - m.mu.RUnlock() - if p == nil || auth == nil { - return nil - } - return p.RoundTripperFor(auth) -} - -// RoundTripperProvider defines a minimal provider of per-auth HTTP transports. -type RoundTripperProvider interface { - RoundTripperFor(auth *Auth) http.RoundTripper -} - -// RequestPreparer is an optional interface that provider executors can implement -// to mutate outbound HTTP requests with provider credentials. -type RequestPreparer interface { - PrepareRequest(req *http.Request, auth *Auth) error -} - -func executorKeyFromAuth(auth *Auth) string { - if auth == nil { - return "" - } - if auth.Attributes != nil { - providerKey := strings.TrimSpace(auth.Attributes["provider_key"]) - compatName := strings.TrimSpace(auth.Attributes["compat_name"]) - if compatName != "" { - if providerKey == "" { - providerKey = compatName - } - return util.OpenAICompatibleProviderKey(providerKey) - } - } - if strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { - providerKey := strings.TrimSpace(auth.Label) - if providerKey == "" { - providerKey = "openai-compatibility" - } - return util.OpenAICompatibleProviderKey(providerKey) - } - return strings.ToLower(strings.TrimSpace(auth.Provider)) -} - -// logEntryWithRequestID returns a logrus entry with request_id field if available in context. -func logEntryWithRequestID(ctx context.Context) *log.Entry { - if ctx == nil { - return log.NewEntry(log.StandardLogger()) - } - if reqID := logging.GetRequestID(ctx); reqID != "" { - return log.WithField("request_id", reqID) - } - return log.NewEntry(log.StandardLogger()) -} - -func debugLogAuthSelection(entry *log.Entry, auth *Auth, provider string, model string) { - if !log.IsLevelEnabled(log.DebugLevel) { - return - } - if entry == nil || auth == nil { - return - } - accountType, accountInfo := auth.AccountInfo() - proxyInfo := auth.ProxyInfo() - suffix := "" - if proxyInfo != "" { - suffix = " " + proxyInfo - } - switch accountType { - case "api_key": - entry.Debugf("Use API key %s for model %s%s", util.HideAPIKey(accountInfo), model, suffix) - case "oauth": - ident := formatOauthIdentity(auth, provider, accountInfo) - entry.Debugf("Use OAuth %s for model %s%s", ident, model, suffix) - } -} - -func formatOauthIdentity(auth *Auth, provider string, accountInfo string) string { - if auth == nil { - return "" - } - // Prefer the auth's provider when available. - providerName := strings.TrimSpace(auth.Provider) - if providerName == "" { - providerName = strings.TrimSpace(provider) - } - // Only log the basename to avoid leaking host paths. - // FileName may be unset for some auth backends; fall back to ID. - authFile := strings.TrimSpace(auth.FileName) - if authFile == "" { - authFile = strings.TrimSpace(auth.ID) - } - if authFile != "" { - authFile = filepath.Base(authFile) - } - parts := make([]string, 0, 3) - if providerName != "" { - parts = append(parts, "provider="+providerName) - } - if authFile != "" { - parts = append(parts, "auth_file="+authFile) - } - if len(parts) == 0 { - return accountInfo - } - return strings.Join(parts, " ") -} - -// InjectCredentials delegates per-provider HTTP request preparation when supported. -// If the registered executor for the auth provider implements RequestPreparer, -// it will be invoked to modify the request (e.g., add headers). -func (m *Manager) InjectCredentials(req *http.Request, authID string) error { - if req == nil || authID == "" { - return nil - } - m.mu.RLock() - a := m.auths[authID] - var exec ProviderExecutor - if a != nil { - exec = m.executors[executorKeyFromAuth(a)] - } - m.mu.RUnlock() - if a == nil || exec == nil { - return nil - } - if p, ok := exec.(RequestPreparer); ok && p != nil { - return p.PrepareRequest(req, a) - } - return nil -} - -// PrepareHttpRequest injects provider credentials into the supplied HTTP request. -func (m *Manager) PrepareHttpRequest(ctx context.Context, auth *Auth, req *http.Request) error { - if m == nil { - return &Error{Code: "provider_not_found", Message: "manager is nil"} - } - if auth == nil { - return &Error{Code: "auth_not_found", Message: "auth is nil"} - } - if req == nil { - return &Error{Code: "invalid_request", Message: "http request is nil"} - } - if ctx != nil { - *req = *req.WithContext(ctx) - } - providerKey := executorKeyFromAuth(auth) - if providerKey == "" { - return &Error{Code: "provider_not_found", Message: "auth provider is empty"} - } - exec := m.executorFor(providerKey) - if exec == nil { - return &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey} - } - preparer, ok := exec.(RequestPreparer) - if !ok || preparer == nil { - return &Error{Code: "not_supported", Message: "executor does not support http request preparation"} - } - return preparer.PrepareRequest(req, auth) -} - -// NewHttpRequest constructs a new HTTP request and injects provider credentials into it. -func (m *Manager) NewHttpRequest(ctx context.Context, auth *Auth, method, targetURL string, body []byte, headers http.Header) (*http.Request, error) { - if ctx == nil { - ctx = context.Background() - } - method = strings.TrimSpace(method) - if method == "" { - method = http.MethodGet - } - var reader io.Reader - if body != nil { - reader = bytes.NewReader(body) - } - httpReq, err := http.NewRequestWithContext(ctx, method, targetURL, reader) - if err != nil { - return nil, err - } - if headers != nil { - httpReq.Header = headers.Clone() - } - if errPrepare := m.PrepareHttpRequest(ctx, auth, httpReq); errPrepare != nil { - return nil, errPrepare - } - return httpReq, nil -} - -// HttpRequest injects provider credentials into the supplied HTTP request and executes it. -func (m *Manager) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) { - if m == nil { - return nil, &Error{Code: "provider_not_found", Message: "manager is nil"} - } - if auth == nil { - return nil, &Error{Code: "auth_not_found", Message: "auth is nil"} - } - if req == nil { - return nil, &Error{Code: "invalid_request", Message: "http request is nil"} - } - providerKey := executorKeyFromAuth(auth) - if providerKey == "" { - return nil, &Error{Code: "provider_not_found", Message: "auth provider is empty"} - } - exec := m.executorFor(providerKey) - if exec == nil { - return nil, &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey} - } - return exec.HttpRequest(ctx, auth, req) -} diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go new file mode 100644 index 000000000..d21ab818a --- /dev/null +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -0,0 +1,1684 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "sort" + "strings" + "sync/atomic" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +var quotaCooldownDisabled atomic.Bool + +var transientErrorCooldownSeconds atomic.Int64 + +// SetQuotaCooldownDisabled toggles quota cooldown scheduling globally. +func SetQuotaCooldownDisabled(disable bool) { + quotaCooldownDisabled.Store(disable) +} + +// SetTransientErrorCooldownSeconds configures cooldowns for 408/500/502/503/504. +// 0 keeps the legacy default; negative values disable transient error cooldowns. +func SetTransientErrorCooldownSeconds(seconds int) { + transientErrorCooldownSeconds.Store(int64(seconds)) +} + +func quotaCooldownDisabledForAuth(auth *Auth) bool { + return quotaCooldownDisabledForAuthWithConfig(auth, nil) +} + +func quotaCooldownDisabledForAuthWithConfig(auth *Auth, cfg *internalconfig.Config) bool { + if auth != nil { + if override, ok := auth.DisableCoolingOverride(); ok { + return override + } + if providerCoolingDisabledForAuth(auth, cfg) { + return true + } + } + if cfg != nil && cfg.DisableCooling { + return true + } + return quotaCooldownDisabled.Load() +} + +func providerCoolingDisabledForAuth(auth *Auth, cfg *internalconfig.Config) bool { + if auth == nil || cfg == nil { + return false + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if provider == "" { + return false + } + providerKey := "" + compatName := "" + if auth.Attributes != nil { + providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) + compatName = strings.TrimSpace(auth.Attributes["compat_name"]) + } + if providerKey == "" && compatName == "" && provider != "openai-compatibility" { + return false + } + if providerKey == "" { + providerKey = provider + } + entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, provider) + return entry != nil && entry.DisableCooling +} + +func nextTransientErrorRetryAfter(now time.Time) time.Time { + seconds := transientErrorCooldownSeconds.Load() + if seconds < 0 { + return time.Time{} + } + if seconds == 0 { + return now.Add(transientErrorCooldown) + } + return now.Add(time.Duration(seconds) * time.Second) +} + +// SetConfig updates the runtime config snapshot used by request-time helpers. +// Callers should provide the latest config on reload so per-credential alias mapping stays in sync. +func (m *Manager) SetConfig(cfg *internalconfig.Config) { + if m == nil { + return + } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + if m.setConfigSnapshotLocked(cfg) { + m.persistCooldownStatesLocked(context.Background()) + } +} + +// SetConfigSnapshot updates only in-memory configuration state. It reports whether +// a caller must persist cleared cooldown state after its commit critical section. +func (m *Manager) SetConfigSnapshot(cfg *internalconfig.Config) bool { + if m == nil { + return false + } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + return m.setConfigSnapshotLocked(cfg) +} + +func (m *Manager) setConfigSnapshotLocked(cfg *internalconfig.Config) bool { + if cfg == nil { + cfg = &internalconfig.Config{} + } + m.mu.RLock() + oldCooldownStore := m.cooldownStore + m.mu.RUnlock() + m.runtimeConfig.Store(cfg) + clearedCooldowns := m.clearDisabledCooldownStates(cfg) + if clearedCooldowns && oldCooldownStore != nil { + m.mu.Lock() + if m.cooldownStore == oldCooldownStore { + m.pendingCooldownStateStore = oldCooldownStore + } + m.mu.Unlock() + } + if !cfg.Home.Enabled { + m.clearHomeRuntimeAuths() + } + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + return clearedCooldowns +} + +// ApplyConfigWithCooldownStateStore serializes a config update with its cooldown +// store transition. It persists the resulting state to the captured old store before +// exposing the resolved replacement store. +func (m *Manager) ApplyConfigWithCooldownStateStore(ctx context.Context, cfg *internalconfig.Config, store CooldownStateStore) bool { + if m == nil { + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + m.mu.RLock() + oldStore := m.cooldownStore + m.mu.RUnlock() + m.setConfigSnapshotLocked(cfg) + if oldStore != nil && !m.persistCooldownStatesToLocked(ctx, oldStore) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if m.cooldownStore != oldStore { + return false + } + if m.pendingCooldownStateStore == oldStore { + m.pendingCooldownStateStore = nil + } + m.cooldownStore = store + return true +} + +// PersistCooldownStates writes the current cooldown snapshot using ctx. +func (m *Manager) PersistCooldownStates(ctx context.Context) { + m.persistCooldownStates(ctx) +} + +// SwapCooldownStateStore persists cleared state to the old store before replacing it. +// Persistence is deliberately performed without holding the manager lock. +func (m *Manager) SwapCooldownStateStore(ctx context.Context, store CooldownStateStore, persistOld bool) bool { + if m == nil { + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + m.mu.RLock() + oldStore := m.cooldownStore + pendingStore := m.pendingCooldownStateStore + m.mu.RUnlock() + storeToPersist := pendingStore + if storeToPersist == nil && persistOld { + storeToPersist = oldStore + } + if storeToPersist != nil && !m.persistCooldownStatesToLocked(ctx, storeToPersist) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if m.cooldownStore != oldStore { + return false + } + if m.pendingCooldownStateStore == storeToPersist { + m.pendingCooldownStateStore = nil + } + m.cooldownStore = store + return true +} + +func (m *Manager) cooldownDisabledForAuth(auth *Auth) bool { + if m == nil { + return quotaCooldownDisabledForAuth(auth) + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + return quotaCooldownDisabledForAuthWithConfig(auth, cfg) +} + +func (m *Manager) clearDisabledCooldownStates(cfg *internalconfig.Config) bool { + if m == nil { + return false + } + now := time.Now() + snapshots := make([]*Auth, 0) + m.mu.Lock() + for _, auth := range m.auths { + if auth == nil { + continue + } + if !quotaCooldownDisabledForAuthWithConfig(auth, cfg) && !auth.Disabled && auth.Status != StatusDisabled { + continue + } + if clearCooldownStateForAuth(auth, now) { + snapshots = append(snapshots, auth.Clone()) + } + } + m.mu.Unlock() + + if m.scheduler != nil { + for _, snapshot := range snapshots { + m.scheduler.upsertAuth(snapshot) + } + } + return len(snapshots) > 0 +} + +// RestoreCooldownStates restores unexpired persisted cooldown records into registered auths. +func (m *Manager) RestoreCooldownStates(ctx context.Context) error { + if m == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + m.mu.RLock() + store := m.cooldownStore + m.mu.RUnlock() + if store == nil { + return nil + } + records, errLoad := store.Load(ctx) + if errLoad != nil { + return errLoad + } + if len(records) == 0 { + return nil + } + + now := time.Now() + authLevelRecords := make([]CooldownStateRecord, 0) + snapshotsByID := make(map[string]*Auth) + + m.mu.Lock() + for _, record := range records { + if strings.TrimSpace(record.Model) == "" { + authLevelRecords = append(authLevelRecords, record) + continue + } + if m.restoreCooldownRecordLocked(record, now) { + if auth := m.auths[strings.TrimSpace(record.AuthID)]; auth != nil { + snapshotsByID[auth.ID] = auth.Clone() + } + } + } + for _, record := range authLevelRecords { + if m.restoreCooldownRecordLocked(record, now) { + if auth := m.auths[strings.TrimSpace(record.AuthID)]; auth != nil { + snapshotsByID[auth.ID] = auth.Clone() + } + } + } + m.mu.Unlock() + + if m.scheduler != nil { + for _, snapshot := range snapshotsByID { + m.scheduler.upsertAuth(snapshot) + } + } + m.persistCooldownStates(ctx) + return nil +} + +func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now time.Time) bool { + authID := strings.TrimSpace(record.AuthID) + if authID == "" || record.NextRetryAfter.IsZero() || !record.NextRetryAfter.After(now) { + return false + } + auth := m.auths[authID] + if auth == nil || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) { + return false + } + updatedAt := record.UpdatedAt + if updatedAt.IsZero() { + updatedAt = now + } + reason := strings.TrimSpace(record.Reason) + model := strings.TrimSpace(record.Model) + quota := record.Quota + if quota.Exceeded && quota.NextRecoverAt.IsZero() { + quota.NextRecoverAt = record.NextRetryAfter + } + + if model == "" { + auth.Unavailable = true + auth.Status = StatusError + auth.NextRetryAfter = record.NextRetryAfter + auth.Quota = quota + auth.UpdatedAt = updatedAt + if reason != "" { + auth.StatusMessage = reason + } + auth.LastError = cloneError(record.LastError) + return true + } + + state := ensureModelState(auth, model) + state.Unavailable = true + state.Status = StatusError + state.NextRetryAfter = record.NextRetryAfter + state.Quota = quota + state.UpdatedAt = updatedAt + if reason != "" { + state.StatusMessage = reason + } + state.LastError = cloneError(record.LastError) + updateAggregatedAvailability(auth, now) + return true +} + +func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { + if auth == nil { + return false + } + changed := false + if auth.Unavailable || !auth.NextRetryAfter.IsZero() || auth.Quota.Exceeded || !auth.Quota.NextRecoverAt.IsZero() { + auth.Unavailable = false + auth.NextRetryAfter = time.Time{} + auth.Quota = QuotaState{} + auth.UpdatedAt = now + changed = true + } + for _, state := range auth.ModelStates { + if state == nil { + continue + } + if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() { + state.Unavailable = false + state.NextRetryAfter = time.Time{} + state.Quota = QuotaState{} + state.UpdatedAt = now + changed = true + } + } + if len(auth.ModelStates) > 0 { + updateAggregatedAvailability(auth, now) + } + return changed +} + +func dedupeStrings(values []string) []string { + if len(values) < 2 { + return values + } + seen := make(map[string]struct{}, len(values)) + out := values[:0] + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" { + continue + } + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + out = append(out, value) + } + return out +} + +// ResetQuota clears quota/cooldown state for an auth and resumes registry routing. +func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []string, error) { + if m == nil { + return nil, nil, nil + } + authID = strings.TrimSpace(authID) + if authID == "" { + return nil, nil, fmt.Errorf("auth id is required") + } + + now := time.Now() + var snapshot *Auth + models := make([]string, 0) + registeredModels := modelsForRegisteredAuth(authID) + cooldownStateChanged := false + + m.mu.Lock() + auth, ok := m.auths[authID] + if !ok || auth == nil { + m.mu.Unlock() + return nil, nil, nil + } + + var cooldownRecordsBefore []CooldownStateRecord + trackCooldownState := m.cooldownStore != nil + if trackCooldownState { + cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now) + } + + for modelKey, state := range auth.ModelStates { + if strings.TrimSpace(modelKey) == "" { + continue + } + models = append(models, modelKey) + if state != nil { + resetModelState(state, now) + } + } + if clearCooldownStateForAuth(auth, now) { + if len(models) == 0 { + models = append(models, registeredModels...) + } + } else if len(auth.ModelStates) > 0 { + updateAggregatedAvailability(auth, now) + } + + if len(models) == 0 { + models = append(models, registeredModels...) + } + models = dedupeStrings(models) + + if !auth.Disabled && auth.Status != StatusDisabled && !hasModelError(auth, now) { + auth.LastError = nil + auth.StatusMessage = "" + auth.Status = StatusActive + } + auth.UpdatedAt = now + if errPersist := m.persist(ctx, auth); errPersist != nil { + m.mu.Unlock() + return nil, nil, errPersist + } + snapshot = auth.Clone() + if trackCooldownState { + cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) + cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) + } + m.mu.Unlock() + + for _, modelKey := range models { + registry.GetGlobalRegistry().ClearModelQuotaExceeded(authID, modelKey) + registry.GetGlobalRegistry().ResumeClientModel(authID, modelKey) + } + if m.scheduler != nil && snapshot != nil { + m.scheduler.upsertAuth(snapshot) + } + if snapshot != nil && cooldownStateChanged { + m.persistCooldownStates(ctx) + } + return snapshot, models, nil +} + +func modelsForRegisteredAuth(authID string) []string { + supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID) + models := make([]string, 0, len(supportedModels)) + for _, supportedModel := range supportedModels { + if supportedModel == nil || strings.TrimSpace(supportedModel.ID) == "" { + continue + } + models = append(models, supportedModel.ID) + } + return models +} + +func (m *Manager) persistCooldownStates(ctx context.Context) { + if m == nil { + return + } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + m.persistCooldownStatesLocked(ctx) +} + +func (m *Manager) persistCooldownStatesLocked(ctx context.Context) { + m.mu.RLock() + store := m.cooldownStore + m.mu.RUnlock() + if m.persistCooldownStatesToLocked(ctx, store) { + m.mu.Lock() + if m.pendingCooldownStateStore == store { + m.pendingCooldownStateStore = nil + } + m.mu.Unlock() + } +} + +func (m *Manager) persistCooldownStatesToLocked(ctx context.Context, store CooldownStateStore) bool { + if m == nil || store == nil { + return true + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + records := m.cooldownStateRecordsSnapshot() + if errSave := store.Save(ctx, records); errSave != nil { + logEntryWithRequestID(ctx).Warnf("failed to persist cooldown state: %v", errSave) + return false + } + return ctx.Err() == nil +} + +func (m *Manager) cooldownStateRecordsSnapshot() []CooldownStateRecord { + now := time.Now() + records := make([]CooldownStateRecord, 0) + + m.mu.RLock() + for _, auth := range m.auths { + records = append(records, m.cooldownStateRecordsForAuthLocked(auth, now)...) + } + m.mu.RUnlock() + + sort.Slice(records, func(i, j int) bool { + if records[i].Provider != records[j].Provider { + return records[i].Provider < records[j].Provider + } + if records[i].AuthID != records[j].AuthID { + return records[i].AuthID < records[j].AuthID + } + return records[i].Model < records[j].Model + }) + return records +} + +func (m *Manager) cooldownStateRecordsForAuthLocked(auth *Auth, now time.Time) []CooldownStateRecord { + if auth == nil || auth.ID == "" || auth.Disabled || auth.Status == StatusDisabled || m.cooldownDisabledForAuth(auth) { + return nil + } + records := make([]CooldownStateRecord, 0, 1+len(auth.ModelStates)) + if record, ok := authCooldownStateRecord(auth, now); ok { + records = append(records, record) + } + for model, state := range auth.ModelStates { + if record, ok := modelCooldownStateRecord(auth, model, state, now); ok { + records = append(records, record) + } + } + sort.Slice(records, func(i, j int) bool { + return records[i].Model < records[j].Model + }) + return records +} + +func cooldownStateRecordsEqual(a, b []CooldownStateRecord) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if !cooldownStateRecordEqual(a[i], b[i]) { + return false + } + } + return true +} + +func cooldownStateRecordEqual(a, b CooldownStateRecord) bool { + if a.Provider != b.Provider || + a.AuthID != b.AuthID || + a.AuthFile != b.AuthFile || + a.Model != b.Model || + a.Status != b.Status || + a.Reason != b.Reason || + !a.NextRetryAfter.Equal(b.NextRetryAfter) || + !a.UpdatedAt.Equal(b.UpdatedAt) || + !cooldownQuotaEqual(a.Quota, b.Quota) { + return false + } + return cooldownErrorEqual(a.LastError, b.LastError) +} + +func cooldownQuotaEqual(a, b QuotaState) bool { + return a.Exceeded == b.Exceeded && + a.Reason == b.Reason && + a.BackoffLevel == b.BackoffLevel && + a.NextRecoverAt.Equal(b.NextRecoverAt) +} + +func cooldownErrorEqual(a, b *Error) bool { + if a == nil || b == nil { + return a == b + } + return a.Code == b.Code && + a.Message == b.Message && + a.Retryable == b.Retryable && + a.HTTPStatus == b.HTTPStatus +} + +func authCooldownStateRecord(auth *Auth, now time.Time) (CooldownStateRecord, bool) { + if auth == nil || !auth.Unavailable || auth.NextRetryAfter.IsZero() || !auth.NextRetryAfter.After(now) { + return CooldownStateRecord{}, false + } + return CooldownStateRecord{ + Provider: strings.TrimSpace(auth.Provider), + AuthID: auth.ID, + AuthFile: cooldownAuthFile(auth), + Status: "cooling", + NextRetryAfter: auth.NextRetryAfter, + Reason: cooldownReason(auth.StatusMessage, auth.Quota, auth.LastError), + Quota: auth.Quota, + LastError: cloneError(auth.LastError), + UpdatedAt: auth.UpdatedAt, + }, true +} + +func modelCooldownStateRecord(auth *Auth, model string, state *ModelState, now time.Time) (CooldownStateRecord, bool) { + model = strings.TrimSpace(model) + if auth == nil || state == nil || model == "" || !state.Unavailable || state.NextRetryAfter.IsZero() || !state.NextRetryAfter.After(now) { + return CooldownStateRecord{}, false + } + return CooldownStateRecord{ + Provider: strings.TrimSpace(auth.Provider), + AuthID: auth.ID, + AuthFile: cooldownAuthFile(auth), + Model: model, + Status: "cooling", + NextRetryAfter: state.NextRetryAfter, + Reason: cooldownReason(state.StatusMessage, state.Quota, state.LastError), + Quota: state.Quota, + LastError: cloneError(state.LastError), + UpdatedAt: state.UpdatedAt, + }, true +} + +func cooldownReason(statusMessage string, quota QuotaState, lastErr *Error) string { + if reason := strings.TrimSpace(quota.Reason); reason != "" { + return reason + } + if statusMessage = strings.TrimSpace(statusMessage); statusMessage != "" { + return statusMessage + } + if lastErr != nil { + if code := strings.TrimSpace(lastErr.Code); code != "" { + return code + } + if message := strings.TrimSpace(lastErr.Message); message != "" { + return message + } + } + return "" +} + +// MarkResult records an execution result and notifies hooks. +func (m *Manager) MarkResult(ctx context.Context, result Result) { + if result.AuthID == "" { + return + } + + shouldResumeModel := false + shouldSuspendModel := false + suspendReason := "" + clearModelQuota := false + setModelQuota := false + var authSnapshot *Auth + cooldownStateChanged := false + + m.mu.Lock() + if auth, ok := m.auths[result.AuthID]; ok && auth != nil { + now := time.Now() + var cooldownRecordsBefore []CooldownStateRecord + trackCooldownState := m.cooldownStore != nil + if trackCooldownState { + cooldownRecordsBefore = m.cooldownStateRecordsForAuthLocked(auth, now) + } + auth.recordRecentRequest(now, result.Success) + if result.Success { + auth.Success++ + } else { + auth.Failed++ + } + + if result.Success { + if result.Model != "" { + state := ensureModelState(auth, result.Model) + resetModelState(state, now) + updateAggregatedAvailability(auth, now) + if !hasModelError(auth, now) { + auth.LastError = nil + auth.StatusMessage = "" + auth.Status = StatusActive + } + auth.UpdatedAt = now + shouldResumeModel = true + clearModelQuota = true + } else { + clearAuthStateOnSuccess(auth, now) + } + } else { + if result.Model != "" { + if !isRequestScopedResultError(result.Error) { + disableCooling := m.cooldownDisabledForAuth(auth) + state := ensureModelState(auth, result.Model) + state.Unavailable = true + state.Status = StatusError + state.UpdatedAt = now + if result.Error != nil { + state.LastError = cloneError(result.Error) + state.StatusMessage = result.Error.Message + auth.LastError = cloneError(result.Error) + auth.StatusMessage = result.Error.Message + } + + statusCode := statusCodeFromResult(result.Error) + if isModelSupportResultError(result.Error) { + next := now.Add(12 * time.Hour) + state.NextRetryAfter = next + suspendReason = "model_not_supported" + shouldSuspendModel = true + } else if isCloudflareChallengeResultError(result.Error) { + next, backoffLevel := nextCloudflareCooldown(state.Quota.BackoffLevel, disableCooling, now) + state.NextRetryAfter = next + state.StatusMessage = "cloudflare challenge" + if auth.LastError != nil { + auth.StatusMessage = "cloudflare challenge" + } + state.Quota = QuotaState{ + Exceeded: true, + Reason: "cloudflare challenge", + NextRecoverAt: next, + BackoffLevel: backoffLevel, + } + } else if isInvalidGrantResultError(result.Error) { + if disableCooling { + state.NextRetryAfter = time.Time{} + } else { + state.NextRetryAfter = now.Add(30 * time.Minute) + suspendReason = "invalid_grant" + shouldSuspendModel = true + } + } else { + switch statusCode { + case 401: + if disableCooling { + state.NextRetryAfter = time.Time{} + } else { + next := now.Add(30 * time.Minute) + state.NextRetryAfter = next + suspendReason = "unauthorized" + shouldSuspendModel = true + } + case 402, 403: + if disableCooling { + state.NextRetryAfter = time.Time{} + } else { + next := now.Add(30 * time.Minute) + state.NextRetryAfter = next + suspendReason = "payment_required" + shouldSuspendModel = true + } + case 404: + if disableCooling { + state.NextRetryAfter = time.Time{} + } else { + next := now.Add(12 * time.Hour) + state.NextRetryAfter = next + suspendReason = "not_found" + shouldSuspendModel = true + } + case 429: + var next time.Time + backoffLevel := state.Quota.BackoffLevel + if !disableCooling { + if result.RetryAfter != nil { + next = now.Add(*result.RetryAfter) + } else { + next, backoffLevel = quotaCooldownAfterFailure(state.Quota, now) + } + } + state.NextRetryAfter = next + state.Quota = QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: next, + BackoffLevel: backoffLevel, + } + if !disableCooling { + suspendReason = "quota" + shouldSuspendModel = true + setModelQuota = true + } + case 408, 500, 502, 503, 504: + if disableCooling { + state.NextRetryAfter = time.Time{} + } else { + state.NextRetryAfter = nextTransientErrorRetryAfter(now) + } + default: + state.NextRetryAfter = time.Time{} + } + } + + auth.Status = StatusError + auth.UpdatedAt = now + updateAggregatedAvailability(auth, now) + } + } else { + disableCooling := m.cooldownDisabledForAuth(auth) + applyAuthFailureState(auth, result.Error, result.RetryAfter, now, disableCooling) + } + } + + _ = m.persist(ctx, auth) + authSnapshot = auth.Clone() + if trackCooldownState { + cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) + cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) + } + } + m.mu.Unlock() + if m.scheduler != nil && authSnapshot != nil { + m.scheduler.upsertAuth(authSnapshot) + } + if authSnapshot != nil && cooldownStateChanged { + m.persistCooldownStates(context.Background()) + } + + if clearModelQuota && result.Model != "" { + registry.GetGlobalRegistry().ClearModelQuotaExceeded(result.AuthID, result.Model) + } + if setModelQuota && result.Model != "" { + registry.GetGlobalRegistry().SetModelQuotaExceeded(result.AuthID, result.Model) + } + if shouldResumeModel { + registry.GetGlobalRegistry().ResumeClientModel(result.AuthID, result.Model) + } else if shouldSuspendModel { + registry.GetGlobalRegistry().SuspendClientModel(result.AuthID, result.Model, suspendReason) + } + + m.hook.OnResult(ctx, result) + m.publishErrorEvent(result, authSnapshot) +} + +func (m *Manager) recordExecutionResult(ctx context.Context, result Result, auth *Auth, ephemeral bool) { + if !ephemeral { + m.MarkResult(ctx, result) + return + } + m.reportHomeResult(ctx, result, auth) +} + +// reportHomeResult only observes a Home dispatch result and never updates local auth state. +func (m *Manager) reportHomeResult(ctx context.Context, result Result, auth *Auth) { + if m == nil || result.AuthID == "" { + return + } + var snapshot *Auth + if auth != nil { + snapshot = auth.Clone() + } + m.hook.OnResult(ctx, result) + m.publishErrorEvent(result, snapshot) +} + +func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Result) { + if result.AuthID == "" { + return + } + + var authSnapshot *Auth + m.mu.Lock() + if auth, ok := m.auths[result.AuthID]; ok && auth != nil { + now := time.Now() + auth.recordRecentRequest(now, result.Success) + if result.Success { + auth.Success++ + } else { + auth.Failed++ + } + _ = m.persist(ctx, auth) + authSnapshot = auth.Clone() + } + m.mu.Unlock() + + m.hook.OnResult(ctx, result) + m.publishErrorEvent(result, authSnapshot) +} + +func ensureModelState(auth *Auth, model string) *ModelState { + if auth == nil || model == "" { + return nil + } + if auth.ModelStates == nil { + auth.ModelStates = make(map[string]*ModelState) + } + if state, ok := auth.ModelStates[model]; ok && state != nil { + return state + } + state := &ModelState{Status: StatusActive} + auth.ModelStates[model] = state + return state +} + +func resetModelState(state *ModelState, now time.Time) { + if state == nil { + return + } + state.Unavailable = false + state.Status = StatusActive + state.StatusMessage = "" + state.NextRetryAfter = time.Time{} + state.LastError = nil + state.Quota = QuotaState{} + state.UpdatedAt = now +} + +func modelStateIsClean(state *ModelState) bool { + if state == nil { + return true + } + if state.Status != StatusActive { + return false + } + if state.Unavailable || state.StatusMessage != "" || !state.NextRetryAfter.IsZero() || state.LastError != nil { + return false + } + if state.Quota.Exceeded || state.Quota.Reason != "" || !state.Quota.NextRecoverAt.IsZero() || state.Quota.BackoffLevel != 0 { + return false + } + return true +} + +func updateAggregatedAvailability(auth *Auth, now time.Time) { + if auth == nil { + return + } + if len(auth.ModelStates) == 0 { + clearAggregatedAvailability(auth) + return + } + allUnavailable := true + earliestRetry := time.Time{} + quotaExceeded := false + quotaRecover := time.Time{} + maxBackoffLevel := 0 + hasState := false + for _, state := range auth.ModelStates { + if state == nil { + continue + } + hasState = true + stateUnavailable := false + if state.Status == StatusDisabled { + stateUnavailable = true + } else if state.Unavailable { + if state.NextRetryAfter.IsZero() { + stateUnavailable = false + } else if state.NextRetryAfter.After(now) { + stateUnavailable = true + if earliestRetry.IsZero() || state.NextRetryAfter.Before(earliestRetry) { + earliestRetry = state.NextRetryAfter + } + } else { + state.Unavailable = false + state.NextRetryAfter = time.Time{} + } + } + if !stateUnavailable { + allUnavailable = false + } + if state.Quota.Exceeded { + quotaExceeded = true + if quotaRecover.IsZero() || (!state.Quota.NextRecoverAt.IsZero() && state.Quota.NextRecoverAt.Before(quotaRecover)) { + quotaRecover = state.Quota.NextRecoverAt + } + if state.Quota.BackoffLevel > maxBackoffLevel { + maxBackoffLevel = state.Quota.BackoffLevel + } + } + } + if !hasState { + clearAggregatedAvailability(auth) + return + } + auth.Unavailable = allUnavailable + if allUnavailable { + auth.NextRetryAfter = earliestRetry + } else { + auth.NextRetryAfter = time.Time{} + } + if quotaExceeded { + auth.Quota.Exceeded = true + auth.Quota.Reason = "quota" + auth.Quota.NextRecoverAt = quotaRecover + auth.Quota.BackoffLevel = maxBackoffLevel + } else { + auth.Quota.Exceeded = false + auth.Quota.Reason = "" + auth.Quota.NextRecoverAt = time.Time{} + auth.Quota.BackoffLevel = 0 + } +} + +func clearAggregatedAvailability(auth *Auth) { + if auth == nil { + return + } + auth.Unavailable = false + auth.NextRetryAfter = time.Time{} + auth.Quota = QuotaState{} +} + +func hasModelError(auth *Auth, now time.Time) bool { + if auth == nil || len(auth.ModelStates) == 0 { + return false + } + for _, state := range auth.ModelStates { + if state == nil { + continue + } + if state.LastError != nil { + return true + } + if state.Status == StatusError { + if state.Unavailable && (state.NextRetryAfter.IsZero() || state.NextRetryAfter.After(now)) { + return true + } + } + } + return false +} + +func clearAuthStateOnSuccess(auth *Auth, now time.Time) { + if auth == nil { + return + } + auth.Unavailable = false + auth.Status = StatusActive + auth.StatusMessage = "" + auth.Quota.Exceeded = false + auth.Quota.Reason = "" + auth.Quota.NextRecoverAt = time.Time{} + auth.Quota.BackoffLevel = 0 + auth.LastError = nil + auth.NextRetryAfter = time.Time{} + auth.UpdatedAt = now +} + +func cloneError(err *Error) *Error { + if err == nil { + return nil + } + return &Error{ + Code: err.Code, + Message: err.Message, + Retryable: err.Retryable, + HTTPStatus: err.HTTPStatus, + } +} + +func errorString(err error) string { + if err == nil { + return "" + } + return err.Error() +} + +func statusCodeFromError(err error) int { + if err == nil { + return 0 + } + type statusCoder interface { + StatusCode() int + } + var sc statusCoder + if errors.As(err, &sc) && sc != nil { + return sc.StatusCode() + } + return 0 +} + +func isRequestScopedError(err error) bool { + if err == nil { + return false + } + requestErr, ok := errors.AsType[cliproxyexecutor.RequestScopedError](err) + return ok && requestErr != nil && requestErr.IsRequestScoped() +} + +func resultErrorFromError(err error) *Error { + if err == nil { + return nil + } + var sourceErr *Error + var resultErr *Error + if errors.As(err, &sourceErr) && sourceErr != nil { + resultErr = cloneError(sourceErr) + } else { + resultErr = &Error{Message: err.Error()} + } + if resultErr.HTTPStatus == 0 { + resultErr.HTTPStatus = statusCodeFromError(err) + } + if isRequestScopedError(err) || isRequestInvalidError(err) { + resultErr.Code = requestScopedErrorCode + } + return resultErr +} + +func isUnauthorizedError(err error) bool { + if err == nil { + return false + } + if statusCodeFromError(err) == http.StatusUnauthorized { + return true + } + raw := strings.ToLower(err.Error()) + return strings.Contains(raw, "status 401") || strings.Contains(raw, "401 unauthorized") +} + +func hasUnauthorizedAuthFailure(auth *Auth) bool { + if auth == nil || auth.LastError == nil { + return false + } + return auth.LastError.StatusCode() == http.StatusUnauthorized || strings.EqualFold(auth.LastError.Code, "unauthorized") +} + +func refreshErrorFromError(err error) *Error { + if err == nil { + return nil + } + statusCode := statusCodeFromError(err) + if statusCode == 0 && isUnauthorizedError(err) { + statusCode = http.StatusUnauthorized + } + authErr := &Error{Message: err.Error(), HTTPStatus: statusCode} + if statusCode == http.StatusUnauthorized { + authErr.Code = "unauthorized" + authErr.Retryable = false + } + return authErr +} + +func retryAfterFromError(err error) *time.Duration { + if err == nil { + return nil + } + type retryAfterProvider interface { + RetryAfter() *time.Duration + } + var rap retryAfterProvider + if !errors.As(err, &rap) || rap == nil { + return nil + } + retryAfter := rap.RetryAfter() + if retryAfter == nil { + return nil + } + value := *retryAfter + return &value +} + +func statusCodeFromResult(err *Error) int { + if err == nil { + return 0 + } + return err.StatusCode() +} + +func isModelSupportErrorMessage(message string) bool { + lower := strings.ToLower(strings.TrimSpace(message)) + if lower == "" { + return false + } + patterns := [...]string{ + "model_not_supported", + "requested model is not supported", + "requested model is unsupported", + "requested model is unavailable", + "model is not supported", + "model not supported", + "unsupported model", + "model unavailable", + "not available for your plan", + "not available for your account", + } + for _, pattern := range patterns { + if strings.Contains(lower, pattern) { + return true + } + } + return false +} + +func isModelSupportError(err error) bool { + if err == nil { + return false + } + status := statusCodeFromError(err) + if status != http.StatusBadRequest && status != http.StatusUnprocessableEntity { + return false + } + return isModelSupportErrorMessage(err.Error()) +} + +func isInvalidGrantErrorMessage(message string) bool { + return strings.Contains(strings.ToLower(message), "invalid_grant") +} + +func isInvalidGrantError(err error) bool { + if err == nil { + return false + } + status := statusCodeFromError(err) + if status != http.StatusBadRequest && status != http.StatusUnauthorized { + return false + } + return isInvalidGrantErrorMessage(err.Error()) +} + +func isInvalidGrantResultError(err *Error) bool { + if err == nil { + return false + } + status := statusCodeFromResult(err) + if status != http.StatusBadRequest && status != http.StatusUnauthorized { + return false + } + return isInvalidGrantErrorMessage(err.Code) || isInvalidGrantErrorMessage(err.Message) +} + +func isModelSupportResultError(err *Error) bool { + if err == nil { + return false + } + status := statusCodeFromResult(err) + if status != http.StatusBadRequest && status != http.StatusUnprocessableEntity { + return false + } + return isModelSupportErrorMessage(err.Message) +} + +func isCloudflareChallengeErrorMessage(message string) bool { + lower := strings.ToLower(strings.TrimSpace(message)) + return strings.Contains(lower, "challenge-platform") || + strings.Contains(lower, "cf-mitigated") || + strings.Contains(lower, "cloudflare challenge") || + (strings.Contains(lower, "cloudflare") && strings.Contains(lower, " 0 { + next = now.Add(cooldown) + } + backoffLevel = nextLevel + } + return next, backoffLevel +} + +func isRequestScopedNotFoundMessage(message string) bool { + if message == "" { + return false + } + lower := strings.ToLower(message) + return strings.Contains(lower, "item with id") && + strings.Contains(lower, "not found") && + strings.Contains(lower, "items are not persisted when `store` is set to false") +} + +func isRequestScopedNotFoundResultError(err *Error) bool { + if err == nil || statusCodeFromResult(err) != http.StatusNotFound { + return false + } + return isRequestScopedNotFoundMessage(err.Message) +} + +func isRequestScopedResultError(err *Error) bool { + return err != nil && (err.IsRequestScoped() || isRequestScopedNotFoundResultError(err)) +} + +func isCountTokensEndpointNotFoundError(err error, requestedModel string) bool { + if err == nil || statusCodeFromError(err) != http.StatusNotFound { + return false + } + baseModel := thinking.ParseSuffix(requestedModel).ModelName + return !isExplicitModelNotFoundError(err, baseModel) +} + +func isExplicitModelNotFoundError(err error, requestedModel string) bool { + if err == nil { + return false + } + if authErr, ok := err.(*Error); ok && authErr != nil { + if isModelNotFoundIdentifier(authErr.Code) || isStructuredModelNotFoundError(authErr.Message, requestedModel) { + return true + } + } else if isStructuredModelNotFoundError(err.Error(), requestedModel) { + return true + } + + switch wrapped := err.(type) { + case interface{ Unwrap() []error }: + for _, nested := range wrapped.Unwrap() { + if isExplicitModelNotFoundError(nested, requestedModel) { + return true + } + } + case interface{ Unwrap() error }: + return isExplicitModelNotFoundError(wrapped.Unwrap(), requestedModel) + } + return false +} + +func isStructuredModelNotFoundError(message, requestedModel string) bool { + var payload any + if errJSON := json.Unmarshal([]byte(strings.TrimSpace(message)), &payload); errJSON != nil { + return false + } + return containsStructuredModelNotFound(payload, requestedModel) +} + +func containsStructuredModelNotFound(value any, requestedModel string) bool { + switch typed := value.(type) { + case map[string]any: + notFoundType := false + exactModelReference := false + for key, item := range typed { + text, isString := item.(string) + if isString { + switch strings.ToLower(strings.TrimSpace(key)) { + case "code": + if isModelNotFoundIdentifier(text) { + return true + } + case "type": + if isModelNotFoundIdentifier(text) { + return true + } + notFoundType = notFoundType || isNotFoundErrorIdentifier(text) + case "error", "message", "detail", "error_description", "title": + if isExplicitModelNotFoundMessage(text, requestedModel) { + return true + } + exactModelReference = exactModelReference || isExactRequestedModelReference(text, requestedModel) + } + } + switch item.(type) { + case map[string]any, []any: + if containsStructuredModelNotFound(item, requestedModel) { + return true + } + } + } + return notFoundType && exactModelReference + case []any: + for _, item := range typed { + if text, isString := item.(string); isString && isExplicitModelNotFoundMessage(text, requestedModel) { + return true + } + if containsStructuredModelNotFound(item, requestedModel) { + return true + } + } + } + return false +} + +func isModelNotFoundIdentifier(value string) bool { + candidate := strings.ToLower(strings.TrimSpace(value)) + if fragment := strings.LastIndex(candidate, "#"); fragment >= 0 && fragment+1 < len(candidate) { + candidate = candidate[fragment+1:] + } else { + if query := strings.Index(candidate, "?"); query >= 0 { + candidate = candidate[:query] + } + candidate = strings.TrimRight(candidate, "/") + if separator := strings.LastIndexAny(candidate, "/:"); separator >= 0 { + candidate = candidate[separator+1:] + } + } + normalized := strings.NewReplacer("-", "_", " ", "_").Replace(candidate) + switch normalized { + case "model_not_found", "model_not_found_error", "unknown_model", "model_does_not_exist", "model_not_exist": + return true + default: + return false + } +} + +func isNotFoundErrorIdentifier(value string) bool { + normalized := strings.NewReplacer("-", "_", " ", "_").Replace(strings.ToLower(strings.TrimSpace(value))) + return normalized == "not_found" || normalized == "not_found_error" +} + +func isExplicitModelNotFoundMessage(message, requestedModel string) bool { + lower := strings.Trim(strings.ToLower(strings.TrimSpace(message)), " .!;\t\r\n") + if lower == "" { + return false + } + normalized := strings.NewReplacer("-", "_", " ", "_").Replace(lower) + if strings.Contains(normalized, "model_not_found") || strings.Contains(normalized, "unknown_model") { + return true + } + for _, prefix := range []string{"no such model", "unknown model"} { + if lower != prefix && !strings.HasPrefix(lower, prefix+" ") && !strings.HasPrefix(lower, prefix+":") { + continue + } + remainder := strings.TrimSpace(strings.TrimPrefix(lower, prefix)) + remainder = strings.TrimSpace(strings.TrimPrefix(remainder, ":")) + if remainder == "" { + return true + } + missingSuffix, matches := trimRequestedModelReference(remainder, requestedModel) + return matches && missingSuffix == "" + } + for _, prefix := range []string{"the requested model", "requested model", "the model", "model"} { + if lower != prefix && !strings.HasPrefix(lower, prefix+" ") && !strings.HasPrefix(lower, prefix+":") { + continue + } + remainder := strings.TrimSpace(strings.TrimPrefix(lower, prefix)) + remainder = strings.TrimSpace(strings.TrimPrefix(remainder, ":")) + if isMissingModelPhrase(remainder) { + return true + } + missingSuffix, matches := trimRequestedModelReference(remainder, requestedModel) + return matches && isMissingModelPhrase(missingSuffix) + } + return false +} + +func isExactRequestedModelReference(message, requestedModel string) bool { + lower := strings.Trim(strings.ToLower(strings.TrimSpace(message)), " .!;\t\r\n") + for _, prefix := range []string{"the requested model", "requested model", "the model", "model"} { + if lower != prefix && !strings.HasPrefix(lower, prefix+" ") && !strings.HasPrefix(lower, prefix+":") { + continue + } + remainder := strings.TrimSpace(strings.TrimPrefix(lower, prefix)) + remainder = strings.TrimSpace(strings.TrimPrefix(remainder, ":")) + suffix, matches := trimRequestedModelReference(remainder, requestedModel) + return matches && suffix == "" + } + return false +} + +func trimRequestedModelReference(value, requestedModel string) (string, bool) { + model := strings.ToLower(strings.TrimSpace(requestedModel)) + if model == "" { + return "", false + } + for _, candidate := range []string{model, "'" + model + "'", `"` + model + `"`, "`" + model + "`"} { + if value == candidate { + return "", true + } + if !strings.HasPrefix(value, candidate) { + continue + } + remainder := value[len(candidate):] + if remainder == "" || strings.ContainsRune(" :,", rune(remainder[0])) { + return strings.TrimLeft(remainder, " :,"), true + } + } + return "", false +} + +func isMissingModelPhrase(value string) bool { + switch strings.Trim(value, " .!;\t\r\n") { + case "not found", "was not found", "could not be found", "does not exist", "doesn't exist", "not exist", "is unknown": + return true + default: + return false + } +} + +// isRequestInvalidError returns true if the error represents a client request +// error that should not be retried. Specifically, it treats 400 responses with +// "invalid_request_error", request-scoped 404 item misses caused by `store=false`, +// and all 422 responses as request-shape failures, where switching auths or +// pooled upstream models will not help. Model-support errors are excluded so +// routing can fall through to another auth or upstream. +func isRequestInvalidError(err error) bool { + if err == nil { + return false + } + if isRequestScopedError(err) { + return true + } + if isCloudflareChallengeError(err) { + return false + } + if isInvalidGrantError(err) { + return false + } + if isModelSupportError(err) { + return false + } + status := statusCodeFromError(err) + switch status { + case http.StatusBadRequest: + msg := err.Error() + return strings.Contains(msg, "invalid_request_error") || + strings.Contains(msg, "bad_request_error") || + strings.Contains(msg, "INVALID_ARGUMENT") || + strings.Contains(msg, "FAILED_PRECONDITION") + case http.StatusNotFound: + return isRequestScopedNotFoundMessage(err.Error()) + case http.StatusUnprocessableEntity: + return true + case http.StatusInternalServerError: + msg := err.Error() + return strings.Contains(msg, "\"status\":\"UNKNOWN\"") || + strings.Contains(msg, "\"status\": \"UNKNOWN\"") + default: + return false + } +} + +func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Duration, now time.Time, disableCooling bool) { + if auth == nil { + return + } + if isRequestScopedResultError(resultErr) { + return + } + auth.Unavailable = true + auth.Status = StatusError + auth.UpdatedAt = now + if resultErr != nil { + auth.LastError = cloneError(resultErr) + if resultErr.Message != "" { + auth.StatusMessage = resultErr.Message + } + } + statusCode := statusCodeFromResult(resultErr) + if isCloudflareChallengeResultError(resultErr) { + auth.StatusMessage = "cloudflare challenge" + next, backoffLevel := nextCloudflareCooldown(auth.Quota.BackoffLevel, disableCooling, now) + auth.Quota = QuotaState{ + Exceeded: true, + Reason: "cloudflare challenge", + NextRecoverAt: next, + BackoffLevel: backoffLevel, + } + auth.NextRetryAfter = next + return + } + if isInvalidGrantResultError(resultErr) { + auth.StatusMessage = "invalid_grant" + if disableCooling { + auth.NextRetryAfter = time.Time{} + } else { + auth.NextRetryAfter = now.Add(30 * time.Minute) + } + return + } + switch statusCode { + case 401: + auth.StatusMessage = "unauthorized" + if disableCooling { + auth.NextRetryAfter = time.Time{} + } else { + auth.NextRetryAfter = now.Add(30 * time.Minute) + } + case 402, 403: + auth.StatusMessage = "payment_required" + if disableCooling { + auth.NextRetryAfter = time.Time{} + } else { + auth.NextRetryAfter = now.Add(30 * time.Minute) + } + case 404: + auth.StatusMessage = "not_found" + if disableCooling { + auth.NextRetryAfter = time.Time{} + } else { + auth.NextRetryAfter = now.Add(12 * time.Hour) + } + case 429: + auth.StatusMessage = "quota exhausted" + auth.Quota.Exceeded = true + auth.Quota.Reason = "quota" + var next time.Time + if !disableCooling { + if retryAfter != nil { + next = now.Add(*retryAfter) + } else { + next, auth.Quota.BackoffLevel = quotaCooldownAfterFailure(auth.Quota, now) + } + } + auth.Quota.NextRecoverAt = next + auth.NextRetryAfter = next + case 408, 500, 502, 503, 504: + auth.StatusMessage = "transient upstream error" + if disableCooling { + auth.NextRetryAfter = time.Time{} + } else { + auth.NextRetryAfter = nextTransientErrorRetryAfter(now) + } + default: + if auth.StatusMessage == "" { + auth.StatusMessage = "request failed" + } + } +} + +// quotaCooldownAfterFailure returns the recovery deadline and backoff level for +// a quota failure observed at now. Failures that land while a previous quota +// window is still open reuse that window instead of escalating, so a burst of +// concurrent in-flight failures advances the backoff ladder at most once per +// window. +func quotaCooldownAfterFailure(quota QuotaState, now time.Time) (time.Time, int) { + if quota.NextRecoverAt.After(now) { + return quota.NextRecoverAt, quota.BackoffLevel + } + cooldown, nextLevel := nextQuotaCooldown(quota.BackoffLevel, false) + var next time.Time + if cooldown > 0 { + next = now.Add(cooldown) + } + return next, nextLevel +} + +// nextQuotaCooldown returns the next cooldown duration and updated backoff level for repeated quota errors. +func nextQuotaCooldown(prevLevel int, disableCooling bool) (time.Duration, int) { + if prevLevel < 0 { + prevLevel = 0 + } + if disableCooling { + return 0, prevLevel + } + cooldown := quotaBackoffBase * time.Duration(1<= quotaBackoffMax { + return quotaBackoffMax, prevLevel + } + return cooldown, prevLevel + 1 +} diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go new file mode 100644 index 000000000..9f31ab793 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -0,0 +1,1221 @@ +package auth + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "path/filepath" + "strconv" + "strings" + "sync" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + cliproxysession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" +) + +// Execute performs a non-streaming execution using the configured selector and executor. +// It supports multiple providers for the same model and round-robins the starting provider per model. +func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + req, opts = cliproxysession.Enrich(req, opts) + normalized := m.normalizeProviders(providers) + if len(normalized) == 0 { + return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} + } + if m.HomeEnabled() { + return m.executeHome(ctx, normalized, req, opts, false) + } + + _, maxRetryCredentials, maxWait := m.retrySettings() + + var lastErr error + retryModel := authSelectionModelFromOptions(opts, req.Model) + for attempt := 0; ; attempt++ { + resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + if errExec == nil { + return resp, nil + } + lastErr = errExec + wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) + if !shouldRetry { + break + } + if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { + return cliproxyexecutor.Response{}, errWait + } + } + if lastErr != nil { + if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { + if resp, ok, errCredits := m.tryAntigravityCreditsExecute(ctx, req, opts); errCredits != nil { + return cliproxyexecutor.Response{}, errCredits + } else if ok { + return resp, nil + } + } + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} +} + +// It supports multiple providers for the same model and round-robins the starting provider per model. +func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + req, opts = cliproxysession.Enrich(req, opts) + normalized := m.normalizeProviders(providers) + if len(normalized) == 0 { + return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} + } + if m.HomeEnabled() { + return m.executeHome(ctx, normalized, req, opts, true) + } + + _, maxRetryCredentials, maxWait := m.retrySettings() + + var lastErr error + retryModel := authSelectionModelFromOptions(opts, req.Model) + for attempt := 0; ; attempt++ { + resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + if errExec == nil { + return resp, nil + } + lastErr = errExec + wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) + if !shouldRetry { + break + } + if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { + return cliproxyexecutor.Response{}, errWait + } + } + if lastErr != nil { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} +} + +// ExecuteStream performs a streaming execution using the configured selector and executor. +// It supports multiple providers for the same model and round-robins the starting provider per model. +func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + req, opts = cliproxysession.Enrich(req, opts) + if m.HomeEnabled() { + if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { + defer unlockSession() + } + } + normalized := m.normalizeProviders(providers) + if len(normalized) == 0 { + return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} + } + + _, maxRetryCredentials, maxWait := m.retrySettings() + + var lastErr error + retryModel := authSelectionModelFromOptions(opts, req.Model) + for attempt := 0; ; attempt++ { + result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + if errStream == nil { + return result, nil + } + lastErr = errStream + wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, retryModel, maxWait) + if !shouldRetry { + break + } + if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { + return nil, errWait + } + } + if lastErr != nil { + if hasAntigravityProvider(normalized) && shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) { + if result, ok, errCredits := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); errCredits != nil { + return nil, errCredits + } else if ok { + return result, nil + } + } + var bootstrapErr *streamBootstrapError + if errors.As(lastErr, &bootstrapErr) && bootstrapErr != nil { + return streamErrorResult(bootstrapErr.Headers(), bootstrapErr.cause), nil + } + return nil, lastErr + } + return nil, &Error{Code: "auth_not_found", Message: "no auth available"} +} + +type requestToFormatResolver interface { + RequestToFormat(req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format +} + +func applyRequestAfterAuthInterceptor(ctx context.Context, executor ProviderExecutor, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, requestedModel string) (cliproxyexecutor.Request, cliproxyexecutor.Options) { + if opts.RequestAfterAuthInterceptor == nil { + return req, opts + } + toFormat := requestToFormat(provider, executor, req, opts) + resp := opts.RequestAfterAuthInterceptor(ctx, cliproxyexecutor.RequestAfterAuthInterceptRequest{ + SourceFormat: opts.SourceFormat, + ToFormat: toFormat, + Model: req.Model, + RequestedModel: requestedModel, + Stream: opts.Stream, + Headers: cloneRequestHeaders(opts.Headers), + Body: bytes.Clone(req.Payload), + Metadata: opts.Metadata, + }) + opts.Headers = mergeRequestHeaders(opts.Headers, resp.Headers, resp.ClearHeaders) + if len(resp.Body) > 0 { + req.Payload = bytes.Clone(resp.Body) + opts.OriginalRequest = bytes.Clone(resp.Body) + } + return req, opts +} + +func requestToFormat(provider string, executor ProviderExecutor, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) sdktranslator.Format { + resolver, ok := executor.(requestToFormatResolver) + if ok && resolver != nil { + formatRequestTo := resolver.RequestToFormat(req, opts) + if formatRequestTo != "" { + return formatRequestTo + } + } + source := opts.SourceFormat.String() + if source == "openai-image" || source == "openai-video" { + return opts.SourceFormat + } + if opts.Alt == "responses/compact" && !opts.Stream { + return sdktranslator.FormatOpenAIResponse + } + switch strings.ToLower(strings.TrimSpace(provider)) { + case "codex": + return sdktranslator.FormatCodex + case "xai": + return sdktranslator.FormatCodex + case "claude": + return sdktranslator.FormatClaude + case "gemini", "vertex", "aistudio": + return sdktranslator.FormatGemini + case "kimi": + return sdktranslator.FormatOpenAI + case "antigravity": + return sdktranslator.FormatAntigravity + default: + return sdktranslator.FormatOpenAI + } +} + +func cloneRequestHeaders(src http.Header) http.Header { + if src == nil { + return nil + } + dst := make(http.Header, len(src)) + for key, values := range src { + dst[key] = append([]string(nil), values...) + } + return dst +} + +func mergeRequestHeaders(current, updates http.Header, clear []string) http.Header { + if updates == nil && len(clear) == 0 { + return current + } + out := cloneRequestHeaders(current) + if out == nil && (len(updates) > 0 || len(clear) > 0) { + out = make(http.Header) + } + for _, key := range clear { + out.Del(key) + } + for key, values := range updates { + out.Del(key) + for _, value := range values { + out.Add(key, value) + } + } + return out +} + +func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { + if len(providers) == 0 { + return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} + } + routeModel := authSelectionModelFromOptions(opts, req.Model) + executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) + opts = ensureRequestedModelMetadata(opts, routeModel) + homeMode := m.HomeEnabled() + homeAuthCount := 1 + tried := make(map[string]struct{}) + attempted := make(map[string]struct{}) + var lastErr error + for { + if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { + if lastErr != nil { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} + } + pickOpts := opts + if homeMode { + pickOpts = withHomeAuthCount(opts, homeAuthCount) + } + auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + if errPick != nil { + if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, errPick + } + + entry := logEntryWithRequestID(ctx) + debugLogAuthSelection(entry, auth, provider, routeModel) + publishSelectedAuthMetadata(opts.Metadata, auth) + + tried[auth.ID] = struct{}{} + execCtx := ctx + if rt := m.roundTripperFor(auth); rt != nil { + execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) + execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) + } + execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) + if len(models) == 0 { + continue + } + attempted[auth.ID] = struct{}{} + var errPrepare error + auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) + if errPrepare != nil { + result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} + m.MarkResult(execCtx, result) + lastErr = errPrepare + continue + } + var authErr error + didRefreshOnUnauthorized := false + for _, upstreamModel := range models { + resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) + execReq := req + execReq.Model = upstreamModel + if restoreExecutionModel { + execReq.Model = executionModel + } + execOpts := opts + execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + resp, errExec := executor.Execute(execCtx, auth, execReq, execOpts) + if errExec != nil { + if errCtx := execCtx.Err(); errCtx != nil { + return cliproxyexecutor.Response{}, errCtx + } + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { + auth = refreshed + didRefreshOnUnauthorized = true + resp, errExec = executor.Execute(execCtx, auth, execReq, execOpts) + if errExec != nil { + if errCtx := execCtx.Err(); errCtx != nil { + return cliproxyexecutor.Response{}, errCtx + } + } + } + } + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} + if errExec != nil { + result.Error = resultErrorFromError(errExec) + if ra := retryAfterFromError(errExec); ra != nil { + result.RetryAfter = ra + } + m.MarkResult(execCtx, result) + if isRequestInvalidError(errExec) { + return cliproxyexecutor.Response{}, errExec + } + authErr = errExec + continue + } + m.MarkResult(execCtx, result) + rewriteForceMappedResponse(&resp, aliasResult) + return resp, nil + } + if authErr != nil { + if isRequestInvalidError(authErr) { + return cliproxyexecutor.Response{}, authErr + } + lastErr = authErr + if homeMode { + homeAuthCount++ + } + continue + } + } +} + +func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { + if len(providers) == 0 { + return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} + } + routeModel := authSelectionModelFromOptions(opts, req.Model) + executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) + opts = ensureRequestedModelMetadata(opts, routeModel) + homeMode := m.HomeEnabled() + homeAuthCount := 1 + tried := make(map[string]struct{}) + attempted := make(map[string]struct{}) + var lastErr error + for { + if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { + if lastErr != nil { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} + } + pickOpts := opts + if homeMode { + pickOpts = withHomeAuthCount(opts, homeAuthCount) + } + auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + if errPick != nil { + if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, errPick + } + + entry := logEntryWithRequestID(ctx) + debugLogAuthSelection(entry, auth, provider, routeModel) + publishSelectedAuthMetadata(opts.Metadata, auth) + + tried[auth.ID] = struct{}{} + execCtx := ctx + if rt := m.roundTripperFor(auth); rt != nil { + execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) + execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) + } + execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) + if len(models) == 0 { + continue + } + attempted[auth.ID] = struct{}{} + var errPrepare error + auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) + if errPrepare != nil { + result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} + m.MarkResult(execCtx, result) + lastErr = errPrepare + continue + } + var authErr error + didRefreshOnUnauthorized := false + for _, upstreamModel := range models { + resultModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) + execReq := req + execReq.Model = upstreamModel + if restoreExecutionModel { + execReq.Model = executionModel + } + execOpts := opts + execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + resp, errExec := executor.CountTokens(execCtx, auth, execReq, execOpts) + if errExec != nil { + if errCtx := execCtx.Err(); errCtx != nil { + return cliproxyexecutor.Response{}, errCtx + } + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(execCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { + auth = refreshed + didRefreshOnUnauthorized = true + resp, errExec = executor.CountTokens(execCtx, auth, execReq, execOpts) + if errExec != nil { + if errCtx := execCtx.Err(); errCtx != nil { + return cliproxyexecutor.Response{}, errCtx + } + } + } + } + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} + if errExec != nil { + result.Error = resultErrorFromError(errExec) + if ra := retryAfterFromError(errExec); ra != nil { + result.RetryAfter = ra + } + // Some Anthropic-compatible upstreams do not implement the + // count_tokens route and return a generic endpoint 404. Record + // the failure for hooks and metrics without suspending a model + // that remains usable through the messages endpoint. + if isCountTokensEndpointNotFoundError(errExec, execReq.Model) { + m.recordAvailabilityNeutralResult(execCtx, result) + } else { + m.MarkResult(execCtx, result) + } + if isRequestInvalidError(errExec) { + return cliproxyexecutor.Response{}, errExec + } + authErr = errExec + continue + } + m.MarkResult(execCtx, result) + rewriteForceMappedResponse(&resp, aliasResult) + return resp, nil + } + if authErr != nil { + if isRequestInvalidError(authErr) { + return cliproxyexecutor.Response{}, authErr + } + lastErr = authErr + if homeMode { + homeAuthCount++ + } + continue + } + } +} + +func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (*cliproxyexecutor.StreamResult, error) { + if len(providers) == 0 { + return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} + } + routeModel := authSelectionModelFromOptions(opts, req.Model) + responseAlias := requestedModelAliasFromOptions(opts, routeModel) + executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) + opts = ensureRequestedModelMetadata(opts, routeModel) + homeMode := m.HomeEnabled() + homeAuthCount := 1 + tried := make(map[string]struct{}) + attempted := make(map[string]struct{}) + var lastErr error + for { + if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { + if lastErr != nil { + return nil, lastErr + } + return nil, &Error{Code: "auth_not_found", Message: "no auth available"} + } + pickOpts := opts + if homeMode { + pickOpts = withHomeAuthCount(opts, homeAuthCount) + } + + var selection *HomeDispatchSelection + var auth *Auth + var executor ProviderExecutor + var provider string + var errPick error + if homeMode { + selection, errPick = m.pickHomeDispatchSelection(ctx, routeModel, pickOpts) + if selection != nil { + auth = selection.CloneAuthForRoute(routeModel) + executor = selection.Executor + provider = selection.Provider + } + } else { + auth, executor, provider, errPick = m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + } + if errPick != nil { + if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { + return nil, lastErr + } + return nil, errPick + } + if auth == nil || executor == nil { + if selection != nil { + selection.End("missing_execution_target") + } + return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} + } + + entry := logEntryWithRequestID(ctx) + debugLogAuthSelection(entry, auth, provider, routeModel) + if selection != nil { + if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil { + selection.End("runtime_auth_bind_failed") + return nil, errRuntimeAuth + } + } + publishSelectedAuthMetadata(opts.Metadata, auth) + + tried[auth.ID] = struct{}{} + execCtx := ctx + releaseAttempt := func() {} + if selection != nil { + var errBind error + execCtx, releaseAttempt, errBind = homeExecutionAttemptContext(ctx, selection) + if errBind != nil { + selection.End("attempt_bind_failed") + return nil, errBind + } + } + if rt := m.roundTripperFor(auth); rt != nil { + execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) + execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) + } + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) + if selection != nil && aliasResult.ForceMapping && responseAlias != "" { + aliasResult.OriginalAlias = responseAlias + } + if len(models) == 0 { + if selection != nil { + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil { + return nil, errEnd + } + } + continue + } + attempted[auth.ID] = struct{}{} + var errPrepare error + if selection != nil { + auth, errPrepare = m.prepareHomeRequestAuth(execCtx, executor, selection) + } else { + auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) + } + if errPrepare != nil { + result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} + if selection != nil { + m.reportHomeResult(execCtx, result, auth) + releaseAttempt() + } else { + m.MarkResult(execCtx, result) + } + lastErr = errPrepare + if selection != nil { + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { + return nil, errEnd + } + } + continue + } + execReq := sanitizeDownstreamWebsocketFallbackRequest(execCtx, auth, req) + streamExecutionModel := "" + if restoreExecutionModel { + streamExecutionModel = executionModel + } + execOpts := opts + if selection != nil { + execOpts.ExecutionLifecycle = selection + } + if homeMode && len(models) > 1 { + models = models[:1] + pooled = false + } + streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, !homeMode, selection != nil) + if errStream != nil { + if selection != nil { + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "stream_start_failed"); errEnd != nil { + return nil, errEnd + } + } + if errCtx := execCtx.Err(); errCtx != nil && ctx != nil && ctx.Err() != nil { + return nil, errCtx + } + if isRequestInvalidError(errStream) { + return nil, errStream + } + lastErr = errStream + if homeMode { + homeAuthCount++ + } + continue + } + if selection != nil { + if m.retainHomeWebsocketSelection(ctx, opts, routeModel, selection) { + return wrapHomeStream(ctx, streamResult, nil, releaseAttempt), nil + } + return wrapHomeStream(ctx, streamResult, selection, releaseAttempt), nil + } + return streamResult, nil + } +} + +func ensureRequestedModelMetadata(opts cliproxyexecutor.Options, requestedModel string) cliproxyexecutor.Options { + requestedModel = strings.TrimSpace(requestedModel) + if requestedModel == "" { + return opts + } + if hasRequestedModelMetadata(opts.Metadata) { + return opts + } + if len(opts.Metadata) == 0 { + opts.Metadata = map[string]any{cliproxyexecutor.RequestedModelMetadataKey: requestedModel} + return opts + } + meta := make(map[string]any, len(opts.Metadata)+1) + for k, v := range opts.Metadata { + meta[k] = v + } + meta[cliproxyexecutor.RequestedModelMetadataKey] = requestedModel + opts.Metadata = meta + return opts +} + +func authSelectionModelFromOptions(opts cliproxyexecutor.Options, fallback string) string { + fallback = strings.TrimSpace(fallback) + if len(opts.Metadata) == 0 { + return fallback + } + raw, ok := opts.Metadata[cliproxyexecutor.AuthSelectionModelMetadataKey] + if !ok || raw == nil { + return fallback + } + switch value := raw.(type) { + case string: + if strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + case []byte: + if strings.TrimSpace(string(value)) != "" { + return strings.TrimSpace(string(value)) + } + } + return fallback +} + +func executionModelForAuthSelection(opts cliproxyexecutor.Options, model string) (string, bool) { + model = strings.TrimSpace(model) + if model == "" { + return "", false + } + selectionModel := authSelectionModelFromOptions(opts, model) + if selectionModel == model { + return "", false + } + return model, true +} + +func withHomeAuthCount(opts cliproxyexecutor.Options, count int) cliproxyexecutor.Options { + if count <= 0 { + count = 1 + } + meta := make(map[string]any, len(opts.Metadata)+1) + for k, v := range opts.Metadata { + meta[k] = v + } + meta[homeAuthCountMetadataKey] = count + opts.Metadata = meta + return opts +} + +func homeAuthCountFromMetadata(meta map[string]any) int { + if len(meta) == 0 { + return 1 + } + switch value := meta[homeAuthCountMetadataKey].(type) { + case int: + if value > 0 { + return value + } + case int64: + if value > 0 { + return int(value) + } + case float64: + if value > 0 { + return int(value) + } + } + return 1 +} + +func hasRequestedModelMetadata(meta map[string]any) bool { + if len(meta) == 0 { + return false + } + raw, ok := meta[cliproxyexecutor.RequestedModelMetadataKey] + if !ok || raw == nil { + return false + } + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) != "" + case []byte: + return strings.TrimSpace(string(v)) != "" + default: + return false + } +} + +type requestAuthPrepareLock struct { + mu sync.Mutex +} + +// prepareHomeRequestAuth prepares a dispatch auth without reading or updating local auth state. +func (m *Manager) prepareHomeRequestAuth(ctx context.Context, executor ProviderExecutor, selection *HomeDispatchSelection) (*Auth, error) { + if m == nil || executor == nil || selection == nil { + return nil, nil + } + auth := selection.CloneAuth() + if auth == nil { + return nil, nil + } + preparer, ok := executor.(RequestAuthPreparer) + if !ok || preparer == nil || !preparer.ShouldPrepareRequestAuth(auth) { + return auth, nil + } + + prepare := func() (*Auth, error) { + target := auth.Clone() + if !preparer.ShouldPrepareRequestAuth(target) { + return target, nil + } + updated, errPrepare := preparer.PrepareRequestAuth(ctx, target) + if errPrepare != nil { + return auth, errPrepare + } + if updated == nil { + return target, nil + } + return updated, nil + } + + id := strings.TrimSpace(auth.ID) + if id == "" { + return prepare() + } + lockValue, _ := m.requestPrepareLocks.LoadOrStore(id, &requestAuthPrepareLock{}) + lock, ok := lockValue.(*requestAuthPrepareLock) + if !ok || lock == nil { + return prepare() + } + lock.mu.Lock() + defer lock.mu.Unlock() + return prepare() +} + +func (m *Manager) prepareRequestAuth(ctx context.Context, executor ProviderExecutor, auth *Auth) (*Auth, error) { + if m == nil || executor == nil || auth == nil { + return auth, nil + } + preparer, ok := executor.(RequestAuthPreparer) + if !ok || preparer == nil || !preparer.ShouldPrepareRequestAuth(auth) { + return auth, nil + } + + id := strings.TrimSpace(auth.ID) + if id == "" { + return preparer.PrepareRequestAuth(ctx, auth.Clone()) + } + + lockValue, _ := m.requestPrepareLocks.LoadOrStore(id, &requestAuthPrepareLock{}) + lock, ok := lockValue.(*requestAuthPrepareLock) + if !ok || lock == nil { + return preparer.PrepareRequestAuth(ctx, auth.Clone()) + } + + lock.mu.Lock() + defer lock.mu.Unlock() + + target := auth.Clone() + m.mu.RLock() + if current := m.auths[id]; current != nil { + target = current.Clone() + } + m.mu.RUnlock() + + if !preparer.ShouldPrepareRequestAuth(target) { + return target, nil + } + + updated, errPrepare := preparer.PrepareRequestAuth(ctx, target) + if errPrepare != nil { + return auth, errPrepare + } + if updated == nil { + return target, nil + } + + saved, errUpdate := m.Update(ctx, updated) + if errUpdate != nil { + return updated, errUpdate + } + if saved != nil { + return saved, nil + } + return updated, nil +} + +func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.Options, fallback string) context.Context { + alias := requestedModelAliasFromOptions(opts, fallback) + ctx = coreusage.WithRequestedModelAlias(ctx, alias) + effort := reasoningEffortFromOptions(opts) + if effort != "" { + ctx = coreusage.WithReasoningEffort(ctx, effort) + } + serviceTier := serviceTierFromOptions(opts) + if serviceTier != "" { + ctx = coreusage.WithServiceTier(ctx, serviceTier) + } + if generate, ok := generateFromOptions(opts); ok { + ctx = coreusage.WithGenerate(ctx, generate) + } + return ctx +} + +func requestedModelAliasFromOptions(opts cliproxyexecutor.Options, fallback string) string { + fallback = strings.TrimSpace(fallback) + if len(opts.Metadata) == 0 { + return fallback + } + raw, ok := opts.Metadata[cliproxyexecutor.RequestedModelMetadataKey] + if !ok || raw == nil { + return fallback + } + switch value := raw.(type) { + case string: + if strings.TrimSpace(value) == "" { + return fallback + } + return strings.TrimSpace(value) + case []byte: + if len(value) == 0 { + return fallback + } + return strings.TrimSpace(string(value)) + default: + return fallback + } +} + +func reasoningEffortFromOptions(opts cliproxyexecutor.Options) string { + if len(opts.Metadata) == 0 { + return "" + } + raw, ok := opts.Metadata[cliproxyexecutor.ReasoningEffortMetadataKey] + if !ok || raw == nil { + return "" + } + switch value := raw.(type) { + case string: + return strings.TrimSpace(value) + case []byte: + return strings.TrimSpace(string(value)) + default: + return "" + } +} + +func serviceTierFromOptions(opts cliproxyexecutor.Options) string { + return stringMetadataValue(opts.Metadata, cliproxyexecutor.ServiceTierMetadataKey) +} + +func generateFromOptions(opts cliproxyexecutor.Options) (bool, bool) { + if len(opts.Metadata) == 0 { + return false, false + } + raw, ok := opts.Metadata[cliproxyexecutor.GenerateMetadataKey] + if !ok || raw == nil { + return false, false + } + switch value := raw.(type) { + case bool: + return value, true + default: + return false, false + } +} + +func stringMetadataValue(metadata map[string]any, key string) string { + if len(metadata) == 0 { + return "" + } + raw, ok := metadata[key] + if !ok || raw == nil { + return "" + } + switch value := raw.(type) { + case string: + return strings.TrimSpace(value) + case []byte: + return strings.TrimSpace(string(value)) + default: + return "" + } +} + +func pinnedAuthIDFromMetadata(meta map[string]any) string { + if len(meta) == 0 { + return "" + } + raw, ok := meta[cliproxyexecutor.PinnedAuthMetadataKey] + if !ok || raw == nil { + return "" + } + switch val := raw.(type) { + case string: + return strings.TrimSpace(val) + case []byte: + return strings.TrimSpace(string(val)) + default: + return "" + } +} + +func disallowFreeAuthFromMetadata(meta map[string]any) bool { + if len(meta) == 0 { + return false + } + raw, ok := meta[cliproxyexecutor.DisallowFreeAuthMetadataKey] + if !ok || raw == nil { + return false + } + switch val := raw.(type) { + case bool: + return val + case string: + parsed, err := strconv.ParseBool(strings.TrimSpace(val)) + return err == nil && parsed + case []byte: + parsed, err := strconv.ParseBool(strings.TrimSpace(string(val))) + return err == nil && parsed + default: + return false + } +} + +func isFreeCodexAuth(auth *Auth) bool { + if auth == nil || auth.Attributes == nil { + return false + } + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + return false + } + return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free") +} + +func publishSelectedAuthMetadata(meta map[string]any, auth *Auth) { + if len(meta) == 0 || auth == nil { + return + } + if authID := strings.TrimSpace(auth.ID); authID != "" { + meta[cliproxyexecutor.SelectedAuthMetadataKey] = authID + if callback, ok := meta[cliproxyexecutor.SelectedAuthCallbackMetadataKey].(func(string)); ok && callback != nil { + callback(authID) + } + } + if authIndex := strings.TrimSpace(auth.EnsureIndex()); authIndex != "" { + meta[cliproxyexecutor.SelectedAuthIndexMetadataKey] = authIndex + if callback, ok := meta[cliproxyexecutor.SelectedAuthIndexCallbackMetadataKey].(func(string)); ok && callback != nil { + callback(authIndex) + } + } +} + +func (m *Manager) executorFor(provider string) ProviderExecutor { + m.mu.RLock() + defer m.mu.RUnlock() + return m.executors[provider] +} + +// roundTripperContextKey is an unexported context key type to avoid collisions. +type roundTripperContextKey struct{} + +// roundTripperFor retrieves an HTTP RoundTripper for the given auth if a provider is registered. +func (m *Manager) roundTripperFor(auth *Auth) http.RoundTripper { + m.mu.RLock() + p := m.rtProvider + m.mu.RUnlock() + if p == nil || auth == nil { + return nil + } + return p.RoundTripperFor(auth) +} + +// RoundTripperProvider defines a minimal provider of per-auth HTTP transports. +type RoundTripperProvider interface { + RoundTripperFor(auth *Auth) http.RoundTripper +} + +// RequestPreparer is an optional interface that provider executors can implement +// to mutate outbound HTTP requests with provider credentials. +type RequestPreparer interface { + PrepareRequest(req *http.Request, auth *Auth) error +} + +func executorKeyFromAuth(auth *Auth) string { + if auth == nil { + return "" + } + if auth.Attributes != nil { + providerKey := strings.TrimSpace(auth.Attributes["provider_key"]) + compatName := strings.TrimSpace(auth.Attributes["compat_name"]) + if compatName != "" { + if providerKey == "" { + providerKey = compatName + } + return util.OpenAICompatibleProviderKey(providerKey) + } + } + if strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { + providerKey := strings.TrimSpace(auth.Label) + if providerKey == "" { + providerKey = "openai-compatibility" + } + return util.OpenAICompatibleProviderKey(providerKey) + } + return strings.ToLower(strings.TrimSpace(auth.Provider)) +} + +// logEntryWithRequestID returns a logrus entry with request_id field if available in context. +func logEntryWithRequestID(ctx context.Context) *log.Entry { + if ctx == nil { + return log.NewEntry(log.StandardLogger()) + } + if reqID := logging.GetRequestID(ctx); reqID != "" { + return log.WithField("request_id", reqID) + } + return log.NewEntry(log.StandardLogger()) +} + +func debugLogAuthSelection(entry *log.Entry, auth *Auth, provider string, model string) { + if !log.IsLevelEnabled(log.DebugLevel) { + return + } + if entry == nil || auth == nil { + return + } + accountType, accountInfo := auth.AccountInfo() + proxyInfo := auth.ProxyInfo() + suffix := "" + if proxyInfo != "" { + suffix = " " + proxyInfo + } + switch accountType { + case "api_key": + entry.Debugf("Use API key %s for model %s%s", util.HideAPIKey(accountInfo), model, suffix) + case "oauth": + ident := formatOauthIdentity(auth, provider, accountInfo) + entry.Debugf("Use OAuth %s for model %s%s", ident, model, suffix) + } +} + +func formatOauthIdentity(auth *Auth, provider string, accountInfo string) string { + if auth == nil { + return "" + } + // Prefer the auth's provider when available. + providerName := strings.TrimSpace(auth.Provider) + if providerName == "" { + providerName = strings.TrimSpace(provider) + } + // Only log the basename to avoid leaking host paths. + // FileName may be unset for some auth backends; fall back to ID. + authFile := strings.TrimSpace(auth.FileName) + if authFile == "" { + authFile = strings.TrimSpace(auth.ID) + } + if authFile != "" { + authFile = filepath.Base(authFile) + } + parts := make([]string, 0, 3) + if providerName != "" { + parts = append(parts, "provider="+providerName) + } + if authFile != "" { + parts = append(parts, "auth_file="+authFile) + } + if len(parts) == 0 { + return accountInfo + } + return strings.Join(parts, " ") +} + +// InjectCredentials delegates per-provider HTTP request preparation when supported. +// If the registered executor for the auth provider implements RequestPreparer, +// it will be invoked to modify the request (e.g., add headers). +func (m *Manager) InjectCredentials(req *http.Request, authID string) error { + if req == nil || authID == "" { + return nil + } + m.mu.RLock() + a := m.auths[authID] + var exec ProviderExecutor + if a != nil { + exec = m.executors[executorKeyFromAuth(a)] + } + m.mu.RUnlock() + if a == nil || exec == nil { + return nil + } + if p, ok := exec.(RequestPreparer); ok && p != nil { + return p.PrepareRequest(req, a) + } + return nil +} + +// PrepareHttpRequest injects provider credentials into the supplied HTTP request. +func (m *Manager) PrepareHttpRequest(ctx context.Context, auth *Auth, req *http.Request) error { + if m == nil { + return &Error{Code: "provider_not_found", Message: "manager is nil"} + } + if auth == nil { + return &Error{Code: "auth_not_found", Message: "auth is nil"} + } + if req == nil { + return &Error{Code: "invalid_request", Message: "http request is nil"} + } + if ctx != nil { + *req = *req.WithContext(ctx) + } + providerKey := executorKeyFromAuth(auth) + if providerKey == "" { + return &Error{Code: "provider_not_found", Message: "auth provider is empty"} + } + exec := m.executorFor(providerKey) + if exec == nil { + return &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey} + } + preparer, ok := exec.(RequestPreparer) + if !ok || preparer == nil { + return &Error{Code: "not_supported", Message: "executor does not support http request preparation"} + } + return preparer.PrepareRequest(req, auth) +} + +// NewHttpRequest constructs a new HTTP request and injects provider credentials into it. +func (m *Manager) NewHttpRequest(ctx context.Context, auth *Auth, method, targetURL string, body []byte, headers http.Header) (*http.Request, error) { + if ctx == nil { + ctx = context.Background() + } + method = strings.TrimSpace(method) + if method == "" { + method = http.MethodGet + } + var reader io.Reader + if body != nil { + reader = bytes.NewReader(body) + } + httpReq, err := http.NewRequestWithContext(ctx, method, targetURL, reader) + if err != nil { + return nil, err + } + if headers != nil { + httpReq.Header = headers.Clone() + } + if errPrepare := m.PrepareHttpRequest(ctx, auth, httpReq); errPrepare != nil { + return nil, errPrepare + } + return httpReq, nil +} + +// HttpRequest injects provider credentials into the supplied HTTP request and executes it. +func (m *Manager) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) { + if m == nil { + return nil, &Error{Code: "provider_not_found", Message: "manager is nil"} + } + if auth == nil { + return nil, &Error{Code: "auth_not_found", Message: "auth is nil"} + } + if req == nil { + return nil, &Error{Code: "invalid_request", Message: "http request is nil"} + } + providerKey := executorKeyFromAuth(auth) + if providerKey == "" { + return nil, &Error{Code: "provider_not_found", Message: "auth provider is empty"} + } + exec := m.executorFor(providerKey) + if exec == nil { + return nil, &Error{Code: "provider_not_found", Message: "executor not registered for provider: " + providerKey} + } + return exec.HttpRequest(ctx, auth, req) +} diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go new file mode 100644 index 000000000..0c85f17a3 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_home.go @@ -0,0 +1,1116 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "sort" + "strings" + "sync" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + log "github.com/sirupsen/logrus" +) + +const ( + homeAuthCountMetadataKey = "__cliproxy_home_auth_count" + // CloseAllExecutionSessionsID asks an executor to release all active execution sessions. + // Executors that do not support this marker may ignore it. + CloseAllExecutionSessionsID = "__all_execution_sessions__" +) + +// HomeDispatchBundle is the immutable client and registry pair for one Home lifetime. +type HomeDispatchBundle struct { + client homeAuthDispatcher + registry *executionregistry.Registry + generation uint64 +} + +// PublishHomeDispatch publishes the selectable Home lifetime as one atomic bundle. +func (m *Manager) PublishHomeDispatch(client homeAuthDispatcher, registry *executionregistry.Registry, generation uint64) *HomeDispatchBundle { + if m == nil || client == nil || registry == nil { + return nil + } + bundle := &HomeDispatchBundle{client: client, registry: registry, generation: generation} + m.homeDispatchBundle.Store(bundle) + return bundle +} + +// ClearHomeDispatchBundle removes bundle only when it still belongs to the active lifetime. +func (m *Manager) ClearHomeDispatchBundle(bundle *HomeDispatchBundle) bool { + if m == nil || bundle == nil { + return false + } + return m.homeDispatchBundle.CompareAndSwap(bundle, nil) +} + +// HomeDispatchBundle returns the active Home lifetime bundle. +func (m *Manager) HomeDispatchBundle() *HomeDispatchBundle { + if m == nil { + return nil + } + return m.homeDispatchBundle.Load() +} + +// SetHomeExecutionRegistry preserves the legacy registry API for callers that also install the current dispatcher. +func (m *Manager) SetHomeExecutionRegistry(registry *executionregistry.Registry) { + if m == nil { + return + } + m.PublishHomeDispatch(currentHomeDispatcher(), registry, 0) +} + +// ClearHomeExecutionRegistry removes a matching legacy registry bundle. +func (m *Manager) ClearHomeExecutionRegistry(registry *executionregistry.Registry) bool { + bundle := m.HomeDispatchBundle() + if bundle == nil || bundle.registry != registry { + return false + } + return m.ClearHomeDispatchBundle(bundle) +} + +// HomeExecutionRegistry returns the registry from the active Home lifetime bundle. +func (m *Manager) HomeExecutionRegistry() *executionregistry.Registry { + bundle := m.HomeDispatchBundle() + if bundle == nil { + return nil + } + return bundle.registry +} + +// HomeEnabled reports whether the home control plane integration is enabled in the runtime config. +func (m *Manager) HomeEnabled() bool { + if m == nil { + return false + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + return cfg != nil && cfg.Home.Enabled +} + +func (m *Manager) localExecutionAllowed() bool { + return m != nil && !m.HomeEnabled() +} + +func (m *Manager) localFallbackAuth(authID string) *Auth { + if !m.localExecutionAllowed() { + return nil + } + m.mu.RLock() + auth := m.auths[strings.TrimSpace(authID)] + m.mu.RUnlock() + if auth == nil { + return nil + } + return auth.Clone() +} + +type homeErrorEnvelope struct { + Error *homeErrorDetail `json:"error"` +} + +type homeErrorDetail struct { + Type string `json:"type"` + Message string `json:"message"` + Code string `json:"code,omitempty"` + Retryable bool `json:"retryable,omitempty"` + RetryAfterMS int64 `json:"retry_after_ms,omitempty"` +} + +const ( + homeUpstreamModelAttributeKey = "home_upstream_model" + homeForceMappingAttributeKey = "home_force_mapping" + homeOriginalAliasAttributeKey = "home_original_alias" + homeRequestRetryExceededErrorCode = "request_retry_exceeded" +) + +func isHomeRequestRetryExceededError(err error) bool { + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + return false + } + return strings.EqualFold(strings.TrimSpace(authErr.Code), homeRequestRetryExceededErrorCode) +} + +func shouldReturnLastErrorOnPickFailure(homeMode bool, lastErr error, errPick error) bool { + if lastErr == nil { + return false + } + if !homeMode { + return true + } + return isHomeRequestRetryExceededError(errPick) +} + +func homeAuthAlreadyTried(tried map[string]struct{}, authID string) bool { + authID = strings.TrimSpace(authID) + if authID == "" || len(tried) == 0 { + return false + } + _, ok := tried[authID] + return ok +} + +func repeatedHomeAuthError() *Error { + return &Error{ + Code: homeRequestRetryExceededErrorCode, + Message: "home returned a previously tried auth", + HTTPStatus: http.StatusServiceUnavailable, + } +} + +type homeAuthDispatchResponse struct { + Model string `json:"model"` + Provider string `json:"provider"` + AuthIndex string `json:"auth_index"` + UserAPIKey string `json:"user_api_key"` + ForceMapping bool `json:"force_mapping"` + OriginalAlias string `json:"original_alias"` + Auth Auth `json:"auth"` +} + +type homeAuthDispatcher interface { + HeartbeatOK() bool + RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error) + AbortAmbiguousDispatch() +} + +var currentHomeDispatcher = func() homeAuthDispatcher { + return home.Current() +} + +func setHomeUserAPIKeyOnGinContext(ctx context.Context, apiKey string) { + apiKey = strings.TrimSpace(apiKey) + if apiKey == "" || ctx == nil { + return + } + ginCtx, ok := ctx.Value("gin").(interface{ Set(string, any) }) + if !ok || ginCtx == nil { + return + } + ginCtx.Set("userApiKey", apiKey) +} + +func homeDispatchHeaders(ctx context.Context, headers http.Header) http.Header { + apiKey, ok := homeQueryCredentialFromContext(ctx) + if !ok { + return headers + } + out := headers.Clone() + if out == nil { + out = http.Header{} + } + if out.Get("Authorization") != "" || out.Get("X-Goog-Api-Key") != "" || out.Get("X-Api-Key") != "" { + return out + } + out.Set("X-Goog-Api-Key", apiKey) + return out +} + +func homeQueryCredentialFromContext(ctx context.Context) (string, bool) { + if ctx == nil { + return "", false + } + if queryCtx, ok := ctx.Value("gin").(interface{ Query(string) string }); ok && queryCtx != nil { + if apiKey := strings.TrimSpace(queryCtx.Query("key")); apiKey != "" { + return apiKey, true + } + if apiKey := strings.TrimSpace(queryCtx.Query("auth_token")); apiKey != "" { + return apiKey, true + } + } + ginCtx, ok := ctx.Value("gin").(interface{ Get(string) (any, bool) }) + if !ok || ginCtx == nil { + return "", false + } + rawMetadata, ok := ginCtx.Get("accessMetadata") + if !ok { + return "", false + } + source := accessMetadataSource(rawMetadata) + if source != "query-key" && source != "query-auth-token" { + return "", false + } + rawAPIKey, ok := ginCtx.Get("userApiKey") + if !ok { + return "", false + } + apiKey := contextStringValue(rawAPIKey) + if apiKey == "" { + return "", false + } + return apiKey, true +} + +func accessMetadataSource(raw any) string { + switch v := raw.(type) { + case map[string]string: + return strings.TrimSpace(v["source"]) + case map[string]any: + return contextStringValue(v["source"]) + default: + return "" + } +} + +func contextStringValue(raw any) string { + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) + case []byte: + return strings.TrimSpace(string(v)) + default: + return "" + } +} + +func homeExecutionSessionIDFromMetadata(meta map[string]any) string { + if len(meta) == 0 { + return "" + } + raw, ok := meta[cliproxyexecutor.ExecutionSessionMetadataKey] + if !ok || raw == nil { + return "" + } + switch value := raw.(type) { + case string: + return strings.TrimSpace(value) + case []byte: + return strings.TrimSpace(string(value)) + default: + return "" + } +} + +type homeSessionSelectionKey struct { + credentialID string + routeModel string +} + +func (m *Manager) lockHomeWebsocketSession(ctx context.Context, opts cliproxyexecutor.Options) func() { + if m == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { + return nil + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + if sessionID == "" { + return nil + } + lock, _ := m.homeSessionLocks.LoadOrStore(sessionID, &sync.Mutex{}) + mutex, ok := lock.(*sync.Mutex) + if !ok || mutex == nil { + return nil + } + mutex.Lock() + return mutex.Unlock +} + +func (m *Manager) retainedHomeSessionSelection(ctx context.Context, opts cliproxyexecutor.Options, model string) (*HomeDispatchSelection, bool, error) { + if m == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { + return nil, false, nil + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + credentialID := pinnedAuthIDFromMetadata(opts.Metadata) + if sessionID == "" { + return nil, false, nil + } + + routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) + var retained *HomeDispatchSelection + var ended []*HomeDispatchSelection + fallbackAttempt := homeAuthCountFromMetadata(opts.Metadata) > 1 + m.mu.Lock() + selections := m.homeSessionSelections[sessionID] + for key, selection := range selections { + if selection == nil { + delete(selections, key) + continue + } + matchesCredential := credentialID == "" || key.credentialID == credentialID + matchesRoute := validRouteModel && key.routeModel == routeModel + if !fallbackAttempt && matchesCredential && selection.Active() && matchesRoute && retained == nil { + retained = selection + continue + } + delete(selections, key) + ended = append(ended, selection) + } + if len(selections) == 0 { + delete(m.homeSessionSelections, sessionID) + } + m.mu.Unlock() + + for _, selection := range ended { + if errWait := m.endHomeSelectionBeforeRedispatch(ctx, selection, "target_changed"); errWait != nil { + return nil, false, errWait + } + } + return retained, retained != nil, nil +} + +func (m *Manager) predictedHomeConcurrencyModel(auth *Auth, routeModel string) (string, bool) { + requestedModel := rewriteModelForAuth(routeModel, auth) + aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel) + upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult) + if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) != 0 { + if len(pool) != 1 { + return "", false + } + upstreamModel = pool[0] + } else { + upstreamModel = m.applyAPIKeyModelAlias(auth, upstreamModel) + } + return validCanonicalHomeConcurrencyModelKey(upstreamModel) +} + +func (m *Manager) endMismatchedHomeSessionSelections(ctx context.Context, sessionID, credentialID, model string, waitForAck bool) error { + if m == nil || sessionID == "" { + return nil + } + routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) + var ended []*HomeDispatchSelection + m.mu.Lock() + selections := m.homeSessionSelections[sessionID] + for key, selection := range selections { + if selection == nil { + delete(selections, key) + continue + } + matchesRoute := validRouteModel && key.routeModel == routeModel + if key.credentialID == credentialID && matchesRoute { + continue + } + delete(selections, key) + ended = append(ended, selection) + } + if len(selections) == 0 { + delete(m.homeSessionSelections, sessionID) + } + m.mu.Unlock() + for _, selection := range ended { + if !waitForAck { + selection.End("target_changed") + continue + } + if errWait := m.endHomeSelectionBeforeRedispatch(ctx, selection, "target_changed"); errWait != nil { + return errWait + } + } + return nil +} + +func (m *Manager) endHomeSelectionBeforeRedispatch(ctx context.Context, selection *HomeDispatchSelection, reason string) error { + if selection == nil { + return nil + } + ticket := selection.EndWithRelease(reason) + if ticket == nil { + return nil + } + + bound := internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound + if m != nil { + if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil { + bound = cfg.CredentialConcurrency.WithDefaults().CPACancelBound + } + } + waitCtx := ctx + if waitCtx == nil { + waitCtx = context.Background() + } + waitCtx, cancelWait := context.WithTimeout(waitCtx, bound) + defer cancelWait() + if errWait := ticket.Wait(waitCtx); errWait != nil { + return &Error{Code: "home_unavailable", Message: "Home did not acknowledge credential release: " + errWait.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} + } + return nil +} + +func (m *Manager) retainHomeWebsocketSelection(ctx context.Context, opts cliproxyexecutor.Options, model string, selection *HomeDispatchSelection) bool { + if m == nil || selection == nil || !selection.Retained() || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil { + return false + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + credentialID := strings.TrimSpace(selection.Auth.ID) + routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) + if selection.accountedModel == "" { + selection.accountedModel, _ = m.predictedHomeConcurrencyModel(selection.Auth, model) + } + if sessionID == "" || credentialID == "" || !validRouteModel || selection.accountedModel == "" { + return false + } + _ = m.endMismatchedHomeSessionSelections(ctx, sessionID, credentialID, routeModel, false) + key := homeSessionSelectionKey{credentialID: credentialID, routeModel: routeModel} + m.mu.Lock() + if m.homeSessionSelections == nil { + m.homeSessionSelections = make(map[string]map[homeSessionSelectionKey]*HomeDispatchSelection) + } + selections := m.homeSessionSelections[sessionID] + if selections == nil { + selections = make(map[homeSessionSelectionKey]*HomeDispatchSelection) + m.homeSessionSelections[sessionID] = selections + } + previous := selections[key] + selections[key] = selection + m.mu.Unlock() + m.rememberHomeRuntimeAuth(sessionID, selection.Auth) + if previous != nil && previous != selection { + previous.End("target_replaced") + } + return true +} + +func (m *Manager) clearHomeSessionLocks() { + if m == nil { + return + } + m.homeSessionLocks.Range(func(key, _ any) bool { + m.homeSessionLocks.Delete(key) + return true + }) +} + +func (m *Manager) takeHomeSessionSelectionsLocked(sessionID string) []*HomeDispatchSelection { + if m == nil { + return nil + } + selections := m.homeSessionSelections[sessionID] + delete(m.homeSessionSelections, sessionID) + result := make([]*HomeDispatchSelection, 0, len(selections)) + for _, selection := range selections { + result = append(result, selection) + } + return result +} + +func (m *Manager) takeAllHomeSessionSelectionsLocked() []*HomeDispatchSelection { + if m == nil { + return nil + } + result := make([]*HomeDispatchSelection, 0) + for sessionID, selections := range m.homeSessionSelections { + delete(m.homeSessionSelections, sessionID) + for _, selection := range selections { + result = append(result, selection) + } + } + return result +} + +func (m *Manager) clearHomeRuntimeAuths() { + if m == nil { + return + } + m.mu.Lock() + m.clearHomeRuntimeAuthsLocked() + selections := m.takeAllHomeSessionSelectionsLocked() + m.mu.Unlock() + for _, selection := range selections { + selection.End("home_disabled") + } +} + +func (m *Manager) clearHomeRuntimeAuthsLocked() { + if m == nil { + return + } + m.homeRuntimeAuths = make(map[string]map[string]*Auth) + m.homeRuntimeAuthOwners = make(map[string]map[string]*HomeDispatchSelection) +} + +func (m *Manager) clearHomeRuntimeAuthsForSessionLocked(sessionID string) { + sessionID = strings.TrimSpace(sessionID) + if m == nil || sessionID == "" { + return + } + delete(m.homeRuntimeAuths, sessionID) + delete(m.homeRuntimeAuthOwners, sessionID) +} + +func (m *Manager) bindHomeSelectionRuntimeAuth(ctx context.Context, opts cliproxyexecutor.Options, selection *HomeDispatchSelection) error { + if m == nil || selection == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil || !authWebsocketsEnabled(selection.Auth) { + return nil + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + authID := strings.TrimSpace(selection.Auth.ID) + if sessionID == "" || authID == "" || !selection.runtimeAuthBound.CompareAndSwap(false, true) { + return nil + } + m.rememberHomeSelectionRuntimeAuth(sessionID, selection) + if errBind := selection.Bind(func() error { + m.forgetHomeRuntimeAuth(sessionID, authID, selection) + return nil + }); errBind != nil { + selection.runtimeAuthBound.Store(false) + m.forgetHomeRuntimeAuth(sessionID, authID, selection) + return errBind + } + return nil +} + +func (m *Manager) rememberHomeSelectionRuntimeAuth(sessionID string, selection *HomeDispatchSelection) { + if m == nil || selection == nil || selection.Auth == nil { + return + } + sessionID = strings.TrimSpace(sessionID) + authID := strings.TrimSpace(selection.Auth.ID) + if sessionID == "" || authID == "" { + return + } + m.mu.Lock() + if m.homeRuntimeAuths == nil { + m.homeRuntimeAuths = make(map[string]map[string]*Auth) + } + if m.homeRuntimeAuthOwners == nil { + m.homeRuntimeAuthOwners = make(map[string]map[string]*HomeDispatchSelection) + } + if m.homeRuntimeAuths[sessionID] == nil { + m.homeRuntimeAuths[sessionID] = make(map[string]*Auth) + } + if m.homeRuntimeAuthOwners[sessionID] == nil { + m.homeRuntimeAuthOwners[sessionID] = make(map[string]*HomeDispatchSelection) + } + m.homeRuntimeAuths[sessionID][authID] = selection.Auth.Clone() + m.homeRuntimeAuthOwners[sessionID][authID] = selection + m.mu.Unlock() +} + +func (m *Manager) forgetHomeRuntimeAuth(sessionID string, authID string, owner *HomeDispatchSelection) { + sessionID = strings.TrimSpace(sessionID) + authID = strings.TrimSpace(authID) + if m == nil || sessionID == "" || authID == "" { + return + } + m.mu.Lock() + owners := m.homeRuntimeAuthOwners[sessionID] + if owner != nil && owners[authID] != owner { + m.mu.Unlock() + return + } + sessionAuths := m.homeRuntimeAuths[sessionID] + delete(sessionAuths, authID) + delete(owners, authID) + if len(sessionAuths) == 0 { + delete(m.homeRuntimeAuths, sessionID) + } + if len(owners) == 0 { + delete(m.homeRuntimeAuthOwners, sessionID) + } + m.mu.Unlock() +} + +func (m *Manager) rememberHomeRuntimeAuth(sessionID string, auth *Auth) { + sessionID = strings.TrimSpace(sessionID) + authID := "" + if auth != nil { + authID = strings.TrimSpace(auth.ID) + } + if m == nil || auth == nil || sessionID == "" || authID == "" || !authWebsocketsEnabled(auth) { + return + } + m.mu.Lock() + if m.homeRuntimeAuths == nil { + m.homeRuntimeAuths = make(map[string]map[string]*Auth) + } + sessionAuths := m.homeRuntimeAuths[sessionID] + if sessionAuths == nil { + sessionAuths = make(map[string]*Auth) + m.homeRuntimeAuths[sessionID] = sessionAuths + } + sessionAuths[authID] = auth.Clone() + m.mu.Unlock() +} + +func (m *Manager) homeRuntimeAuthByID(sessionID string, authID string) (*Auth, ProviderExecutor, string, bool) { + sessionID = strings.TrimSpace(sessionID) + authID = strings.TrimSpace(authID) + if m == nil || sessionID == "" || authID == "" { + return nil, nil, "", false + } + m.mu.RLock() + sessionAuths := m.homeRuntimeAuths[sessionID] + auth := sessionAuths[authID] + m.mu.RUnlock() + if auth == nil || !authWebsocketsEnabled(auth) { + return nil, nil, "", false + } + logicalProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) + executorKey := executorKeyFromAuth(auth) + if logicalProvider == "" || executorKey == "" { + return nil, nil, "", false + } + executor, ok := m.Executor(executorKey) + if !ok && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" { + executor, ok = m.Executor("openai-compatibility") + } + if !ok { + return nil, nil, "", false + } + return auth.Clone(), executor, logicalProvider, true +} + +func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { + if m == nil { + return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} + } + if ctx == nil { + ctx = context.Background() + } + selection, errSelection := m.pickHomeDispatchSelection(ctx, model, opts) + if errSelection != nil { + return nil, nil, "", errSelection + } + if selection.Auth == nil || homeAuthAlreadyTried(tried, selection.Auth.ID) { + selection.End("repeated_auth") + return nil, nil, "", repeatedHomeAuthError() + } + auth := selection.CloneAuthForRoute(model) + executor := selection.Executor + provider := selection.Provider + selection.End("legacy_selection_unbound") + return auth, executor, provider, nil +} + +func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { + if m == nil { + return nil, &Error{Code: "auth_not_found", Message: "no auth available"} + } + if ctx == nil { + ctx = context.Background() + } + + requestedModel := strings.TrimSpace(model) + if requestedModel == "" { + requestedModel = requestedModelFromMetadata(opts.Metadata, model) + } + retained, retainedOK, errRetained := m.retainedHomeSessionSelection(ctx, opts, requestedModel) + if errRetained != nil { + return nil, errRetained + } + if retainedOK { + return retained, nil + } + if sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata); sessionID != "" { + if credentialID := pinnedAuthIDFromMetadata(opts.Metadata); credentialID != "" { + if errEnd := m.endMismatchedHomeSessionSelections(ctx, sessionID, credentialID, requestedModel, true); errEnd != nil { + return nil, errEnd + } + } + } + + bundle := m.HomeDispatchBundle() + if bundle == nil || bundle.client == nil || bundle.registry == nil { + return nil, &Error{Code: "home_unavailable", Message: "home dispatch bundle unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + client := bundle.client + registry := bundle.registry + if !client.HeartbeatOK() { + return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} + } + + sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata) + dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers) + raw, errRPop := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) + if errRPop != nil { + if home.IsAmbiguousDispatchError(errRPop) { + client.AbortAmbiguousDispatch() + } + pending.End() + if errors.Is(errRPop, home.ErrAuthNotFound) { + return nil, &Error{Code: "auth_not_found", Message: errRPop.Error(), HTTPStatus: http.StatusServiceUnavailable} + } + return nil, &Error{Code: "home_unavailable", Message: errRPop.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} + } + + envelope, errEnvelope := decodeHomeDispatchConcurrencyEnvelope(raw) + if errEnvelope != nil { + if envelope.Present { + client.AbortAmbiguousDispatch() + } + pending.End() + if envelope.Present { + return nil, invalidHomeConcurrencyResponse("Home returned malformed concurrency tuple") + } + return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} + } + + kind := "http" + if cliproxyexecutor.DownstreamWebsocket(ctx) { + kind = "websocket" + } else if opts.Stream { + kind = "stream" + } + baseScope := executionregistry.ScopeSpec{ + RequestID: logging.GetRequestID(ctx), + Model: requestedModel, + Kind: kind, + StartedAt: time.Now(), + } + var scope *executionregistry.Scope + if envelope.Present { + var errInstall error + scope, errInstall = installHomeConcurrencyScope(registry, pending, envelope.Tuple, baseScope) + if errInstall != nil { + client.AbortAmbiguousDispatch() + pending.End() + return nil, homeConcurrencyInstallError(errInstall) + } + } + endScope := func() { + if scope != nil { + scope.End("local_validation_failed") + return + } + pending.End() + } + if errHome := decodeHomeDispatchError(raw); errHome != nil { + if envelope.Present { + client.AbortAmbiguousDispatch() + endScope() + return nil, invalidHomeConcurrencyResponse("Home returned both accounted concurrency and an error") + } + pending.End() + return nil, errHome + } + + var dispatch homeAuthDispatchResponse + if errUnmarshal := json.Unmarshal(raw, &dispatch); errUnmarshal != nil { + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} + } + auth := dispatch.Auth + if strings.TrimSpace(auth.ID) == "" { + // Backward compatibility: older Home instances returned the auth directly. + if errUnmarshal := json.Unmarshal(raw, &auth); errUnmarshal != nil { + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} + } + } + observedModel := canonicalHomeDispatchModel(dispatch.Model, requestedModel) + if envelope.Present { + observedConcurrencyModel, validModel := validCanonicalHomeConcurrencyModelKey(observedModel) + if !validModel || envelope.Tuple.Model != observedConcurrencyModel { + client.AbortAmbiguousDispatch() + endScope() + return nil, invalidHomeConcurrencyResponse("Home concurrency model does not match dispatched model") + } + } + if !envelope.Present { + baseScope.Model = observedModel + } + + setHomeUserAPIKeyOnGinContext(ctx, dispatch.UserAPIKey) + if upstreamModel := strings.TrimSpace(dispatch.Model); upstreamModel != "" { + if auth.Attributes == nil { + auth.Attributes = make(map[string]string, 3) + } + auth.Attributes[homeUpstreamModelAttributeKey] = upstreamModel + } + if originalAlias := strings.TrimSpace(dispatch.OriginalAlias); dispatch.ForceMapping && originalAlias != "" { + if auth.Attributes == nil { + auth.Attributes = make(map[string]string, 2) + } + auth.Attributes[homeForceMappingAttributeKey] = "true" + auth.Attributes[homeOriginalAliasAttributeKey] = originalAlias + } + if strings.TrimSpace(auth.ID) == "" { + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned auth without id", HTTPStatus: http.StatusBadGateway} + } + if errIdentity := verifyAccountedHomeConcurrencyIdentity(envelope.Tuple, &auth, dispatch.AuthIndex); errIdentity != nil { + endScope() + return nil, errIdentity + } + logicalProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) + executorKey := executorKeyFromAuth(&auth) + if logicalProvider == "" || executorKey == "" { + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned auth without provider", HTTPStatus: http.StatusBadGateway} + } + + homeAuthIndex := strings.TrimSpace(dispatch.AuthIndex) + if homeAuthIndex != "" { + auth.Index = homeAuthIndex + auth.indexAssigned = true + } else { + auth.EnsureIndex() + } + + executor, okExecutor := m.Executor(executorKey) + if !okExecutor && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" { + executor, okExecutor = m.Executor("openai-compatibility") + } + if !okExecutor { + endScope() + return nil, &Error{Code: "executor_not_found", Message: "executor not registered", HTTPStatus: http.StatusBadGateway} + } + if scope == nil { + var errInstall error + scope, errInstall = installHomeConcurrencyScope(registry, pending, homeConcurrencyTuple{}, executionregistry.ScopeSpec{ + RequestID: baseScope.RequestID, + CredentialID: strings.TrimSpace(auth.ID), + Model: baseScope.Model, + Kind: baseScope.Kind, + StartedAt: baseScope.StartedAt, + }) + if errInstall != nil { + client.AbortAmbiguousDispatch() + pending.End() + return nil, homeConcurrencyInstallError(errInstall) + } + } + + selection, errSelection := newHomeDispatchSelection(auth.Clone(), executor, logicalProvider, scope) + if errSelection != nil { + endScope() + return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} + } + if envelope.Present { + selection.accountedModel = envelope.Tuple.Model + } + if executionSessionID := homeExecutionSessionIDFromMetadata(opts.Metadata); executionSessionID != "" && cliproxyexecutor.DownstreamWebsocket(ctx) { + if errEnd := m.endMismatchedHomeSessionSelections(ctx, executionSessionID, strings.TrimSpace(auth.ID), requestedModel, true); errEnd != nil { + selection.End("target_change_release_failed") + return nil, errEnd + } + } + return selection, nil +} + +func requestedModelFromMetadata(metadata map[string]any, fallback string) string { + if metadata != nil { + if v, ok := metadata[cliproxyexecutor.RequestedModelMetadataKey]; ok { + switch typed := v.(type) { + case string: + if trimmed := strings.TrimSpace(typed); trimmed != "" { + return trimmed + } + case []byte: + if trimmed := strings.TrimSpace(string(typed)); trimmed != "" { + return trimmed + } + } + } + } + fallback = strings.TrimSpace(fallback) + if fallback == "" { + return "unknown" + } + return fallback +} + +func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, routeModel string, opts cliproxyexecutor.Options) ([]creditsCandidateEntry, error) { + if m == nil || !m.localExecutionAllowed() { + return nil, nil + } + pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + var candidates []creditsCandidateEntry + m.mu.RLock() + for _, auth := range m.auths { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled { + continue + } + if pinnedAuthID != "" && auth.ID != pinnedAuthID { + continue + } + if !strings.EqualFold(strings.TrimSpace(auth.Provider), "antigravity") { + continue + } + if !strings.Contains(strings.ToLower(strings.TrimSpace(routeModel)), "claude") { + continue + } + providerKey := executorKeyFromAuth(auth) + executor, ok := m.executors[providerKey] + if !ok { + continue + } + candidates = append(candidates, creditsCandidateEntry{ + auth: auth.Clone(), + executor: executor, + provider: providerKey, + }) + } + m.mu.RUnlock() + + var known []creditsCandidateEntry + var unknown []creditsCandidateEntry + for _, candidate := range candidates { + hint, okHint, errHint := GetAntigravityCreditsHintRequired(ctx, candidate.auth.ID) + if errHint != nil { + return nil, antigravityCreditsKVUnavailableError(errHint) + } + if okHint && hint.Known { + if !hint.Available { + continue + } + known = append(known, candidate) + continue + } + unknown = append(unknown, candidate) + } + sort.Slice(known, func(i, j int) bool { + return known[i].auth.ID < known[j].auth.ID + }) + sort.Slice(unknown, func(i, j int) bool { + return unknown[i].auth.ID < unknown[j].auth.ID + }) + return append(known, unknown...), nil +} + +type creditsCandidateEntry struct { + auth *Auth + executor ProviderExecutor + provider string +} + +func hasAntigravityProvider(providers []string) bool { + for _, p := range providers { + if strings.EqualFold(strings.TrimSpace(p), "antigravity") { + return true + } + } + return false +} + +func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, providers []string) bool { + status := statusCodeFromError(lastErr) + log.WithFields(log.Fields{ + "lastErr": errorString(lastErr), + "status": status, + "providers": providers, + }).Debug("shouldAttemptAntigravityCreditsFallback") + if m == nil || lastErr == nil || m.HomeEnabled() { + return false + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if cfg == nil || !cfg.QuotaExceeded.AntigravityCredits { + return false + } + switch status { + case http.StatusTooManyRequests, http.StatusServiceUnavailable: + return true + case 0: + var authErr *Error + if errors.As(lastErr, &authErr) && authErr != nil { + return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" || authErr.Code == "model_cooldown" + } + var cooldownErr *modelCooldownError + if errors.As(lastErr, &cooldownErr) { + return true + } + return false + default: + return false + } +} + +func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, bool, error) { + if m != nil && m.HomeEnabled() { + return cliproxyexecutor.Response{}, false, &Error{Code: "home_fallback_unsupported", Message: "Home does not support Antigravity credits fallback", HTTPStatus: http.StatusServiceUnavailable} + } + if !m.localExecutionAllowed() { + return cliproxyexecutor.Response{}, false, nil + } + routeModel := req.Model + candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) + if errCandidates != nil { + return cliproxyexecutor.Response{}, false, errCandidates + } + for _, c := range candidates { + if ctx.Err() != nil { + return cliproxyexecutor.Response{}, false, nil + } + creditsCtx := WithAntigravityCredits(ctx) + if rt := m.roundTripperFor(c.auth); rt != nil { + creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt) + creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt) + } + creditsOpts := ensureRequestedModelMetadata(opts, routeModel) + creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel) + preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) + if errPrepare != nil { + continue + } + c.auth = preparedAuth + publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth) + models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) + if len(models) == 0 { + continue + } + for _, upstreamModel := range models { + resultModel := m.stateModelForExecution(c.auth, routeModel, upstreamModel, pooled) + execReq := req + execReq.Model = upstreamModel + resp, errExec := c.executor.Execute(creditsCtx, c.auth, execReq, creditsOpts) + result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, Success: errExec == nil} + if errExec != nil { + result.Error = resultErrorFromError(errExec) + if ra := retryAfterFromError(errExec); ra != nil { + result.RetryAfter = ra + } + m.MarkResult(creditsCtx, result) + continue + } + m.MarkResult(creditsCtx, result) + rewriteForceMappedResponse(&resp, aliasResult) + return resp, true, nil + } + } + return cliproxyexecutor.Response{}, false, nil +} + +func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, bool, error) { + if m != nil && m.HomeEnabled() { + return nil, false, &Error{Code: "home_fallback_unsupported", Message: "Home does not support Antigravity credits fallback", HTTPStatus: http.StatusServiceUnavailable} + } + if !m.localExecutionAllowed() { + return nil, false, nil + } + routeModel := req.Model + candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) + if errCandidates != nil { + return nil, false, errCandidates + } + for _, c := range candidates { + if ctx.Err() != nil { + return nil, false, nil + } + creditsCtx := WithAntigravityCredits(ctx) + if rt := m.roundTripperFor(c.auth); rt != nil { + creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt) + creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt) + } + creditsOpts := ensureRequestedModelMetadata(opts, routeModel) + preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) + if errPrepare != nil { + continue + } + c.auth = preparedAuth + publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth) + models, pooled, aliasResult := m.executionModelCandidatesWithAlias(c.auth, routeModel) + if len(models) == 0 { + continue + } + result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, "", models, pooled, aliasResult, true, false) + if errStream != nil { + continue + } + return result, true, nil + } + return nil, false, nil +} + +func antigravityCreditsKVUnavailableError(cause error) error { + if cause == nil { + return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + return &Error{Code: "home_kv_unavailable", Message: "home kv store unavailable: " + cause.Error(), HTTPStatus: http.StatusServiceUnavailable} +} diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go new file mode 100644 index 000000000..8b655cbbc --- /dev/null +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -0,0 +1,195 @@ +package auth + +import ( + "context" + "fmt" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/tidwall/sjson" +) + +func (m *Manager) executeHome(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool) (cliproxyexecutor.Response, error) { + if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { + defer unlockSession() + } + routeModel := authSelectionModelFromOptions(opts, req.Model) + responseAlias := requestedModelAliasFromOptions(opts, routeModel) + executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) + opts = ensureRequestedModelMetadata(opts, routeModel) + tried := make(map[string]struct{}) + var lastErr error + for homeAuthCount := 1; ; homeAuthCount++ { + selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, withHomeAuthCount(opts, homeAuthCount)) + if errSelection != nil { + if lastErr != nil && isHomeRequestRetryExceededError(errSelection) { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, errSelection + } + auth := selection.CloneAuthForRoute(routeModel) + if auth == nil || selection.Executor == nil { + selection.End("missing_execution_target") + return cliproxyexecutor.Response{}, &Error{Code: "executor_not_found", Message: "executor not registered"} + } + if _, seen := tried[auth.ID]; seen { + selection.End("repeated_auth") + if lastErr != nil { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, repeatedHomeAuthError() + } + entry := logEntryWithRequestID(ctx) + debugLogAuthSelection(entry, auth, selection.Provider, routeModel) + if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil { + selection.End("runtime_auth_bind_failed") + return cliproxyexecutor.Response{}, errRuntimeAuth + } + publishSelectedAuthMetadata(opts.Metadata, auth) + tried[auth.ID] = struct{}{} + execCtx, releaseAttempt, errBind := homeExecutionAttemptContext(ctx, selection) + if errBind != nil { + selection.End("attempt_bind_failed") + return cliproxyexecutor.Response{}, errBind + } + if rt := m.roundTripperFor(auth); rt != nil { + execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) + execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) + } + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) + if aliasResult.ForceMapping && responseAlias != "" { + aliasResult.OriginalAlias = responseAlias + } + if len(models) > 1 { + models = models[:1] + pooled = false + } + if len(models) == 0 { + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil { + return cliproxyexecutor.Response{}, errEnd + } + lastErr = &Error{Code: "auth_not_found", Message: "no execution models available"} + continue + } + preparedAuth, errPrepare := m.prepareHomeRequestAuth(execCtx, selection.Executor, selection) + if errPrepare != nil { + m.reportHomeResult(execCtx, Result{AuthID: auth.ID, Provider: selection.Provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}, auth) + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { + return cliproxyexecutor.Response{}, errEnd + } + lastErr = errPrepare + continue + } + for _, upstreamModel := range models { + resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled) + execReq := req + execReq.Model = upstreamModel + if restoreExecutionModel { + execReq.Model = executionModel + } + execOpts := opts + execOpts.ExecutionLifecycle = selection + execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, selection.Executor, selection.Provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errCtx := execCtx.Err(); errCtx != nil { + releaseAttempt() + selection.End("attempt_canceled") + return cliproxyexecutor.Response{}, errCtx + } + var response cliproxyexecutor.Response + var errExecute error + if countTokens { + response, errExecute = selection.Executor.CountTokens(execCtx, preparedAuth, execReq, execOpts) + } else { + response, errExecute = selection.Executor.Execute(execCtx, preparedAuth, execReq, execOpts) + } + result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil} + if errExecute == nil { + m.reportHomeResult(execCtx, result, preparedAuth) + releaseAttempt() + rewriteForceMappedResponse(&response, aliasResult) + if !m.retainHomeWebsocketSelection(ctx, opts, routeModel, selection) { + selection.End("completed") + } + return response, nil + } + result.Error = resultErrorFromError(errExecute) + result.RetryAfter = retryAfterFromError(errExecute) + m.reportHomeResult(execCtx, result, preparedAuth) + lastErr = errExecute + if isRequestInvalidError(errExecute) { + releaseAttempt() + selection.End("request_invalid") + return cliproxyexecutor.Response{}, errExecute + } + } + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "execution_failed"); errEnd != nil { + return cliproxyexecutor.Response{}, errEnd + } + if errCtx := execCtx.Err(); errCtx != nil && ctx != nil && ctx.Err() != nil { + return cliproxyexecutor.Response{}, errCtx + } + } +} + +func homeExecutionAttemptContext(ctx context.Context, selection *HomeDispatchSelection) (context.Context, func(), error) { + if selection == nil { + return nil, func() {}, fmt.Errorf("Home dispatch selection is nil") + } + return selection.AttemptContext(ctx) +} + +func wrapHomeStream(ctx context.Context, result *cliproxyexecutor.StreamResult, selection *HomeDispatchSelection, releaseAttempt func()) *cliproxyexecutor.StreamResult { + if result == nil || result.Chunks == nil { + if releaseAttempt != nil { + releaseAttempt() + } + return result + } + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(out) + if releaseAttempt != nil { + defer releaseAttempt() + } + if selection != nil { + defer selection.End("stream_closed") + } + forward := true + for { + select { + case <-ctx.Done(): + return + case chunk, ok := <-result.Chunks: + if !ok { + return + } + if !forward { + continue + } + select { + case <-ctx.Done(): + return + case out <- chunk: + } + if chunk.Err != nil && selection != nil { + forward = false + } + } + } + }() + return &cliproxyexecutor.StreamResult{Headers: result.Headers, Chunks: out} +} + +func sanitizeDownstreamWebsocketFallbackRequest(ctx context.Context, auth *Auth, req cliproxyexecutor.Request) cliproxyexecutor.Request { + if !cliproxyexecutor.DownstreamWebsocket(ctx) || authWebsocketsEnabled(auth) || len(req.Payload) == 0 { + return req + } + updated, errDelete := sjson.DeleteBytes(req.Payload, "generate") + if errDelete != nil { + return req + } + req.Payload = updated + return req +} diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go new file mode 100644 index 000000000..109f2f878 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -0,0 +1,262 @@ +package auth + +import ( + "context" + "strings" + "time" + + "github.com/google/uuid" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +// SetRetryConfig updates retry attempts, credential retry limit and cooldown wait interval. +func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration, maxRetryCredentials int) { + if m == nil { + return + } + if retry < 0 { + retry = 0 + } + if maxRetryCredentials < 0 { + maxRetryCredentials = 0 + } + if maxRetryInterval < 0 { + maxRetryInterval = 0 + } + m.requestRetry.Store(int32(retry)) + m.maxRetryCredentials.Store(int32(maxRetryCredentials)) + m.maxRetryInterval.Store(maxRetryInterval.Nanoseconds()) +} + +// RegisterExecutor registers a provider executor with the manager. +func (m *Manager) RegisterExecutor(executor ProviderExecutor) { + if executor == nil { + return + } + provider := strings.TrimSpace(executor.Identifier()) + if provider == "" { + return + } + + var replaced ProviderExecutor + m.mu.Lock() + replaced = m.executors[provider] + m.executors[provider] = executor + m.mu.Unlock() + + if replaced == nil || replaced == executor { + return + } + if closer, ok := replaced.(ExecutionSessionCloser); ok && closer != nil { + closer.CloseExecutionSession(CloseAllExecutionSessionsID) + } +} + +// UnregisterExecutor removes the executor associated with the provider key. +func (m *Manager) UnregisterExecutor(provider string) { + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" { + return + } + m.mu.Lock() + delete(m.executors, provider) + m.mu.Unlock() +} + +// Register inserts a new auth entry into the manager. +func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { + if auth == nil { + return nil, nil + } + if auth.ID == "" { + auth.ID = uuid.NewString() + } + now := time.Now() + clearedCooldown := false + if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { + clearedCooldown = clearCooldownStateForAuth(auth, now) + } + auth.EnsureIndex() + authClone := auth.Clone() + m.mu.Lock() + m.auths[auth.ID] = authClone + m.mu.Unlock() + if !shouldDeferAPIKeyModelAliasRebuild(ctx) { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + } + if m.scheduler != nil { + m.scheduler.upsertAuth(authClone) + } + m.queueRefreshReschedule(auth.ID) + _ = m.persist(ctx, auth) + m.hook.OnAuthRegistered(ctx, auth.Clone()) + if clearedCooldown { + m.persistCooldownStates(ctx) + } + return auth.Clone(), nil +} + +// Update replaces an existing auth entry and notifies hooks. +func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { + if auth == nil || auth.ID == "" { + return nil, nil + } + m.mu.Lock() + existing, ok := m.auths[auth.ID] + if !ok || existing == nil { + m.mu.Unlock() + return nil, nil + } + if !auth.indexAssigned && auth.Index == "" { + auth.Index = existing.Index + auth.indexAssigned = existing.indexAssigned + } + auth.Success = existing.Success + auth.Failed = existing.Failed + auth.recentRequests = existing.recentRequests + if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { + if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { + auth.ModelStates = existing.ModelStates + } + } + now := time.Now() + clearedCooldown := false + if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { + clearedCooldown = clearCooldownStateForAuth(auth, now) + } + auth.EnsureIndex() + authClone := auth.Clone() + m.auths[auth.ID] = authClone + m.mu.Unlock() + if !shouldDeferAPIKeyModelAliasRebuild(ctx) { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + } + if m.scheduler != nil { + m.scheduler.upsertAuth(authClone) + } + m.queueRefreshReschedule(auth.ID) + _ = m.persist(ctx, auth) + m.hook.OnAuthUpdated(ctx, auth.Clone()) + if clearedCooldown { + m.persistCooldownStates(ctx) + } + return auth.Clone(), nil +} + +// Remove deletes an auth from runtime state without persisting. +// Disk and token-store deletion must be handled by the caller. +func (m *Manager) Remove(ctx context.Context, id string) { + if m == nil { + return + } + id = strings.TrimSpace(id) + if id == "" { + return + } + _ = ctx + + m.mu.Lock() + existing := m.auths[id] + if existing == nil { + m.mu.Unlock() + return + } + provider := strings.TrimSpace(existing.Provider) + delete(m.auths, id) + if m.modelPoolOffsets != nil { + delete(m.modelPoolOffsets, id) + } + for sessionID, sessionAuths := range m.homeRuntimeAuths { + if sessionAuths == nil { + continue + } + delete(sessionAuths, id) + if len(sessionAuths) == 0 { + delete(m.homeRuntimeAuths, sessionID) + } + } + m.mu.Unlock() + + if !shouldDeferAPIKeyModelAliasRebuild(ctx) { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + } + if m.scheduler != nil { + m.scheduler.removeAuth(id) + } + m.queueRefreshUnschedule(id) + m.invalidateSessionAffinity(id) + + if provider != "" { + if exec, ok := m.Executor(provider); ok && exec != nil { + if closer, okCloser := exec.(ExecutionSessionCloser); okCloser { + closer.CloseExecutionSession(CloseAllExecutionSessionsID) + } + } + } + m.persistCooldownStates(ctx) +} + +func (m *Manager) invalidateSessionAffinity(authID string) { + if m == nil || authID == "" { + return + } + if invalidator, ok := m.selector.(interface{ InvalidateAuth(string) }); ok && invalidator != nil { + invalidator.InvalidateAuth(authID) + } +} + +// Load resets manager state from the backing store. +func (m *Manager) Load(ctx context.Context) error { + m.mu.Lock() + if m.store == nil { + m.mu.Unlock() + return nil + } + items, err := m.store.List(ctx) + if err != nil { + m.mu.Unlock() + return err + } + m.auths = make(map[string]*Auth, len(items)) + for _, auth := range items { + if auth == nil || auth.ID == "" { + continue + } + auth.EnsureIndex() + m.auths[auth.ID] = auth.Clone() + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if cfg == nil { + cfg = &internalconfig.Config{} + } + m.rebuildAPIKeyModelAliasLocked(cfg) + m.mu.Unlock() + m.syncScheduler() + return nil +} + +func (m *Manager) persist(ctx context.Context, auth *Auth) error { + if m.store == nil || auth == nil { + return nil + } + if shouldSkipPersist(ctx) { + return nil + } + if IsConfigAPIKeyAuth(auth) { + return nil + } + if auth.Attributes != nil { + if v := strings.ToLower(strings.TrimSpace(auth.Attributes["runtime_only"])); v == "true" { + return nil + } + } + if IsPluginVirtualAuth(auth) { + return nil + } + // Skip persistence when metadata is absent (e.g., runtime-only auths). + if auth.Metadata == nil { + return nil + } + _, err := m.store.Save(ctx, auth) + return err +} diff --git a/sdk/cliproxy/auth/conductor_models.go b/sdk/cliproxy/auth/conductor_models.go new file mode 100644 index 000000000..900c7343e --- /dev/null +++ b/sdk/cliproxy/auth/conductor_models.go @@ -0,0 +1,830 @@ +package auth + +import ( + "bytes" + "strings" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func (m *Manager) lookupAPIKeyUpstreamModel(authID, requestedModel string) string { + if m == nil { + return "" + } + authID = strings.TrimSpace(authID) + if authID == "" { + return "" + } + requestedModel = strings.TrimSpace(requestedModel) + if requestedModel == "" { + return "" + } + table, _ := m.apiKeyModelAlias.Load().(apiKeyModelAliasTable) + if table == nil { + return "" + } + byAlias := table[authID] + if len(byAlias) == 0 { + return "" + } + key := strings.ToLower(thinking.ParseSuffix(requestedModel).ModelName) + if key == "" { + key = strings.ToLower(requestedModel) + } + resolved := strings.TrimSpace(byAlias[key]) + if resolved == "" { + return "" + } + return preserveRequestedModelSuffix(requestedModel, resolved) +} + +func isAPIKeyAuth(auth *Auth) bool { + if auth == nil { + return false + } + return auth.AuthKind() == AuthKindAPIKey +} + +func isOpenAICompatAPIKeyAuth(auth *Auth) bool { + if !isAPIKeyAuth(auth) { + return false + } + if strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { + return true + } + if auth.Attributes == nil { + return false + } + return strings.TrimSpace(auth.Attributes["compat_name"]) != "" +} + +func openAICompatProviderKey(auth *Auth) string { + if auth == nil { + return "" + } + if auth.Attributes != nil { + if providerKey := strings.TrimSpace(auth.Attributes["provider_key"]); providerKey != "" { + return util.OpenAICompatibleProviderKey(providerKey) + } + if compatName := strings.TrimSpace(auth.Attributes["compat_name"]); compatName != "" { + return util.OpenAICompatibleProviderKey(compatName) + } + } + return util.OpenAICompatibleProviderKey(auth.Provider) +} + +func openAICompatModelPoolKey(auth *Auth, requestedModel string) string { + base := strings.TrimSpace(thinking.ParseSuffix(requestedModel).ModelName) + if base == "" { + base = strings.TrimSpace(requestedModel) + } + return strings.ToLower(strings.TrimSpace(auth.ID)) + "|" + openAICompatProviderKey(auth) + "|" + strings.ToLower(base) +} + +func (m *Manager) nextModelPoolOffset(key string, size int) int { + if m == nil || size <= 1 { + return 0 + } + key = strings.TrimSpace(key) + if key == "" { + return 0 + } + m.mu.Lock() + defer m.mu.Unlock() + if m.modelPoolOffsets == nil { + m.modelPoolOffsets = make(map[string]int) + } + offset := m.modelPoolOffsets[key] + if offset >= 2_147_483_640 { + offset = 0 + } + m.modelPoolOffsets[key] = offset + 1 + if size <= 0 { + return 0 + } + return offset % size +} + +func rotateStrings(values []string, offset int) []string { + if len(values) <= 1 { + return values + } + if offset <= 0 { + out := make([]string, len(values)) + copy(out, values) + return out + } + offset = offset % len(values) + out := make([]string, 0, len(values)) + out = append(out, values[offset:]...) + out = append(out, values[:offset]...) + return out +} + +func (m *Manager) resolveOpenAICompatUpstreamModelPool(auth *Auth, requestedModel string) []string { + if m == nil || !isOpenAICompatAPIKeyAuth(auth) { + return nil + } + requestedModel = strings.TrimSpace(requestedModel) + if requestedModel == "" { + return nil + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if cfg == nil { + cfg = &internalconfig.Config{} + } + providerKey := "" + compatName := "" + if auth.Attributes != nil { + providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) + compatName = strings.TrimSpace(auth.Attributes["compat_name"]) + } + entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider) + if entry == nil { + return nil + } + return resolveModelAliasPoolFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +func preserveRequestedModelSuffix(requestedModel, resolved string) string { + return preserveResolvedModelSuffix(resolved, thinking.ParseSuffix(requestedModel)) +} + +func (m *Manager) executionModelCandidates(auth *Auth, routeModel string) []string { + if auth != nil && auth.Attributes != nil { + if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" { + return []string{homeModel} + } + } + requestedModel := rewriteModelForAuth(routeModel, auth) + requestedModel = m.applyOAuthModelAlias(auth, requestedModel) + if pool := m.resolveOpenAICompatUpstreamModelPool(auth, requestedModel); len(pool) > 0 { + if len(pool) == 1 { + return pool + } + offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, requestedModel), len(pool)) + return rotateStrings(pool, offset) + } + resolved := m.applyAPIKeyModelAlias(auth, requestedModel) + if strings.TrimSpace(resolved) == "" { + resolved = requestedModel + } + return []string{resolved} +} + +func (m *Manager) selectionModelForAuth(auth *Auth, routeModel string) string { + requestedModel := rewriteModelForAuth(routeModel, auth) + if strings.TrimSpace(requestedModel) == "" { + requestedModel = strings.TrimSpace(routeModel) + } + resolvedModel := m.applyOAuthModelAlias(auth, requestedModel) + if strings.TrimSpace(resolvedModel) == "" { + resolvedModel = requestedModel + } + return resolvedModel +} + +func (m *Manager) selectionModelKeyForAuth(auth *Auth, routeModel string) string { + return canonicalModelKey(m.selectionModelForAuth(auth, routeModel)) +} + +func (m *Manager) stateModelForExecution(auth *Auth, routeModel, upstreamModel string, pooled bool) string { + if auth != nil && auth.Attributes != nil { + if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" { + if resolved := strings.TrimSpace(upstreamModel); resolved != "" { + return resolved + } + return homeModel + } + } + stateModel := executionResultModel(routeModel, upstreamModel, pooled) + selectionModel := m.selectionModelForAuth(auth, routeModel) + if canonicalModelKey(selectionModel) == canonicalModelKey(upstreamModel) && strings.TrimSpace(selectionModel) != "" { + return strings.TrimSpace(upstreamModel) + } + return stateModel +} + +func executionResultModel(routeModel, upstreamModel string, pooled bool) string { + if pooled { + if resolved := strings.TrimSpace(upstreamModel); resolved != "" { + return resolved + } + } + if requested := strings.TrimSpace(routeModel); requested != "" { + return requested + } + return strings.TrimSpace(upstreamModel) +} + +func (m *Manager) filterExecutionModels(auth *Auth, routeModel string, candidates []string, pooled bool) []string { + if len(candidates) == 0 { + return nil + } + now := time.Now() + out := make([]string, 0, len(candidates)) + for _, upstreamModel := range candidates { + stateModel := m.stateModelForExecution(auth, routeModel, upstreamModel, pooled) + blocked, _, _ := isAuthBlockedForModel(auth, stateModel, now) + if blocked { + continue + } + out = append(out, upstreamModel) + } + return out +} + +func (m *Manager) preparedExecutionModels(auth *Auth, routeModel string) ([]string, bool) { + candidates := m.executionModelCandidates(auth, routeModel) + pooled := len(candidates) > 1 + return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled +} + +func (m *Manager) preparedExecutionModelsWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) { + candidates, pooled, aliasResult := m.executionModelCandidatesWithAlias(auth, routeModel) + return m.filterExecutionModels(auth, routeModel, candidates, pooled), pooled, aliasResult +} + +func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) { + requestedModel := rewriteModelForAuth(routeModel, auth) + aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel) + if aliasResult.ForceMapping && auth != nil && auth.Attributes != nil && strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") { + aliasResult.OriginalAlias = strings.TrimSpace(routeModel) + } + upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult) + + var candidates []string + if auth != nil && auth.Attributes != nil { + if homeModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]); homeModel != "" { + candidates = []string{homeModel} + } + } + if len(candidates) == 0 { + if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) > 0 { + if len(pool) == 1 { + candidates = pool + } else { + offset := m.nextModelPoolOffset(openAICompatModelPoolKey(auth, upstreamModel), len(pool)) + candidates = rotateStrings(pool, offset) + } + } else { + resolved := m.applyAPIKeyModelAlias(auth, upstreamModel) + if strings.TrimSpace(resolved) == "" { + resolved = upstreamModel + } + candidates = []string{resolved} + } + } + pooled := len(candidates) > 1 + return candidates, pooled, aliasResult +} + +func (m *Manager) resolveExecutionAliasResult(auth *Auth, routeModel string) OAuthModelAliasResult { + requestedModel := rewriteModelForAuth(routeModel, auth) + return m.resolveExecutionAliasResultForRequested(auth, requestedModel) +} + +func (m *Manager) resolveExecutionAliasResultForRequested(auth *Auth, requestedModel string) OAuthModelAliasResult { + if result := homeForceMappingAliasResult(auth, requestedModel); result.ForceMapping { + return result + } + if auth != nil && auth.AuthKind() == AuthKindAPIKey { + return m.resolveAPIKeyModelAliasWithResult(auth, requestedModel) + } + return m.applyOAuthModelAliasWithResult(auth, requestedModel) +} + +func homeForceMappingAliasResult(auth *Auth, requestedModel string) OAuthModelAliasResult { + if auth == nil || auth.Attributes == nil || !strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") { + return OAuthModelAliasResult{} + } + originalAlias := strings.TrimSpace(auth.Attributes[homeOriginalAliasAttributeKey]) + canonicalOriginalAlias := canonicalHomeConcurrencyModelKey(auth.Attributes[homeOriginalAliasAttributeKey]) + canonicalRequestedModel := canonicalHomeConcurrencyModelKey(requestedModel) + if canonicalOriginalAlias == "" || canonicalOriginalAlias != canonicalRequestedModel { + return OAuthModelAliasResult{} + } + upstreamModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]) + if upstreamModel == "" { + upstreamModel = strings.TrimSpace(requestedModel) + } + return OAuthModelAliasResult{ + UpstreamModel: upstreamModel, + ForceMapping: true, + OriginalAlias: originalAlias, + } +} + +func executionAliasPoolModel(auth *Auth, requestedModel string, aliasResult OAuthModelAliasResult) string { + if auth != nil && auth.AuthKind() == AuthKindAPIKey { + if strings.TrimSpace(requestedModel) != "" { + return requestedModel + } + } + if strings.TrimSpace(aliasResult.UpstreamModel) != "" { + return aliasResult.UpstreamModel + } + return requestedModel +} + +func (m *Manager) resolveAPIKeyModelAliasWithResult(auth *Auth, requestedModel string) OAuthModelAliasResult { + if m == nil || auth == nil { + return OAuthModelAliasResult{} + } + requestedModel = strings.TrimSpace(requestedModel) + if requestedModel == "" { + return OAuthModelAliasResult{} + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if cfg == nil { + cfg = &internalconfig.Config{} + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + var models []modelAliasEntry + switch provider { + case "gemini": + if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil { + models = asModelAliasEntries(entry.Models) + } + case "gemini-interactions": + if entry := resolveInteractionsAPIKeyConfig(cfg, auth); entry != nil { + models = asModelAliasEntries(entry.Models) + } + case "claude": + if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil { + models = asModelAliasEntries(entry.Models) + } + case "codex": + if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil { + models = asModelAliasEntries(entry.Models) + } + case "xai": + if entry := resolveXAIAPIKeyConfig(cfg, auth); entry != nil { + models = asModelAliasEntries(entry.Models) + } + case "vertex": + if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil { + models = asModelAliasEntries(entry.Models) + } + default: + providerKey := "" + compatName := "" + if auth.Attributes != nil { + providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) + compatName = strings.TrimSpace(auth.Attributes["compat_name"]) + } + if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { + if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil { + models = asModelAliasEntries(entry.Models) + } + } + } + if len(models) == 0 { + return OAuthModelAliasResult{UpstreamModel: requestedModel} + } + result := resolveModelAliasResultFromConfigModels(requestedModel, models) + if strings.TrimSpace(result.UpstreamModel) == "" { + return OAuthModelAliasResult{UpstreamModel: requestedModel} + } + return result +} + +func (m *Manager) prepareExecutionModels(auth *Auth, routeModel string) []string { + models, _ := m.preparedExecutionModels(auth, routeModel) + return models +} + +func rewriteForceMappedResponse(resp *cliproxyexecutor.Response, aliasResult OAuthModelAliasResult) { + if resp == nil || !aliasResult.ForceMapping || strings.TrimSpace(aliasResult.OriginalAlias) == "" { + return + } + resp.Payload = rewriteModelInResponse(resp.Payload, aliasResult.OriginalAlias) +} + +func rewriteForceMappedStreamChunk(rewriter *StreamRewriter, payload []byte) []byte { + if rewriter == nil || len(payload) == 0 { + return payload + } + rewritten := rewriter.RewriteChunk(payload) + if len(rewritten) > 0 { + return rewritten + } + if bytes.Contains(payload, []byte("data:")) { + if lineWise := rewriteSSEPayloadLines(payload, rewriter.options.RewriteModel); len(lineWise) > 0 { + return lineWise + } + } + if len(rewriter.pendingBuf) > 0 { + return nil + } + return nil +} + +func finishForceMappedStreamChunks(rewriter *StreamRewriter) []byte { + if rewriter == nil { + return nil + } + return rewriter.Finish() +} + +func (m *Manager) rebuildAPIKeyModelAliasFromRuntimeConfig() { + if m == nil { + return + } + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if cfg == nil { + cfg = &internalconfig.Config{} + } + m.mu.Lock() + defer m.mu.Unlock() + m.rebuildAPIKeyModelAliasLocked(cfg) +} + +// RefreshAPIKeyModelAlias rebuilds the API-key model alias table from the current runtime config. +func (m *Manager) RefreshAPIKeyModelAlias() { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() +} + +func (m *Manager) rebuildAPIKeyModelAliasLocked(cfg *internalconfig.Config) { + if m == nil { + return + } + if cfg == nil { + cfg = &internalconfig.Config{} + } + + out := make(apiKeyModelAliasTable) + for _, auth := range m.auths { + if auth == nil { + continue + } + if strings.TrimSpace(auth.ID) == "" { + continue + } + if auth.AuthKind() != AuthKindAPIKey { + continue + } + + byAlias := make(map[string]string) + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + switch provider { + case "gemini": + if entry := resolveGeminiAPIKeyConfig(cfg, auth); entry != nil { + compileAPIKeyModelAliasForModels(byAlias, entry.Models) + } + case "gemini-interactions": + if entry := resolveInteractionsAPIKeyConfig(cfg, auth); entry != nil { + compileAPIKeyModelAliasForModels(byAlias, entry.Models) + } + case "claude": + if entry := resolveClaudeAPIKeyConfig(cfg, auth); entry != nil { + compileAPIKeyModelAliasForModels(byAlias, entry.Models) + } + case "codex": + if entry := resolveCodexAPIKeyConfig(cfg, auth); entry != nil { + compileAPIKeyModelAliasForModels(byAlias, entry.Models) + } + case "xai": + if entry := resolveXAIAPIKeyConfig(cfg, auth); entry != nil { + compileAPIKeyModelAliasForModels(byAlias, entry.Models) + } + case "vertex": + if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil { + compileAPIKeyModelAliasForModels(byAlias, entry.Models) + } + default: + // OpenAI-compat uses config selection from auth.Attributes. + providerKey := "" + compatName := "" + if auth.Attributes != nil { + providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) + compatName = strings.TrimSpace(auth.Attributes["compat_name"]) + } + if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { + if entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider); entry != nil { + compileAPIKeyModelAliasForModels(byAlias, entry.Models) + } + } + } + + if len(byAlias) > 0 { + out[auth.ID] = byAlias + } + } + + m.apiKeyModelAlias.Store(out) +} + +func compileAPIKeyModelAliasForModels[T interface { + GetName() string + GetAlias() string +}](out map[string]string, models []T) { + if out == nil { + return + } + for i := range models { + alias := strings.TrimSpace(models[i].GetAlias()) + name := strings.TrimSpace(models[i].GetName()) + if alias == "" || name == "" { + continue + } + aliasKey := strings.ToLower(thinking.ParseSuffix(alias).ModelName) + if aliasKey == "" { + aliasKey = strings.ToLower(alias) + } + // Config priority: first alias wins. + if _, exists := out[aliasKey]; exists { + continue + } + out[aliasKey] = name + // Also allow direct lookup by upstream name (case-insensitive), so lookups on already-upstream + // models remain a cheap no-op. + nameKey := strings.ToLower(thinking.ParseSuffix(name).ModelName) + if nameKey == "" { + nameKey = strings.ToLower(name) + } + if nameKey != "" { + if _, exists := out[nameKey]; !exists { + out[nameKey] = name + } + } + // Preserve config suffix priority by seeding a base-name lookup when name already has suffix. + nameResult := thinking.ParseSuffix(name) + if nameResult.HasSuffix { + baseKey := strings.ToLower(strings.TrimSpace(nameResult.ModelName)) + if baseKey != "" { + if _, exists := out[baseKey]; !exists { + out[baseKey] = name + } + } + } + } +} + +func rewriteModelForAuth(model string, auth *Auth) string { + if auth == nil || model == "" { + return model + } + prefix := strings.TrimSpace(auth.Prefix) + if prefix == "" { + return model + } + needle := prefix + "/" + if !strings.HasPrefix(model, needle) { + return model + } + return strings.TrimPrefix(model, needle) +} + +func (m *Manager) applyAPIKeyModelAlias(auth *Auth, requestedModel string) string { + if m == nil || auth == nil { + return requestedModel + } + + if auth.AuthKind() != AuthKindAPIKey { + return requestedModel + } + + requestedModel = strings.TrimSpace(requestedModel) + if requestedModel == "" { + return requestedModel + } + + // Fast path: lookup per-auth mapping table (keyed by auth.ID). + if resolved := m.lookupAPIKeyUpstreamModel(auth.ID, requestedModel); resolved != "" { + return resolved + } + + // Slow path: scan config for the matching credential entry and resolve alias. + // This acts as a safety net if mappings are stale or auth.ID is missing. + cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) + if cfg == nil { + cfg = &internalconfig.Config{} + } + + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + upstreamModel := "" + switch provider { + case "gemini": + upstreamModel = resolveUpstreamModelForGeminiAPIKey(cfg, auth, requestedModel) + case "gemini-interactions": + upstreamModel = resolveUpstreamModelForInteractionsAPIKey(cfg, auth, requestedModel) + case "claude": + upstreamModel = resolveUpstreamModelForClaudeAPIKey(cfg, auth, requestedModel) + case "codex": + upstreamModel = resolveUpstreamModelForCodexAPIKey(cfg, auth, requestedModel) + case "xai": + upstreamModel = resolveUpstreamModelForXAIAPIKey(cfg, auth, requestedModel) + case "vertex": + upstreamModel = resolveUpstreamModelForVertexAPIKey(cfg, auth, requestedModel) + default: + upstreamModel = resolveUpstreamModelForOpenAICompatAPIKey(cfg, auth, requestedModel) + } + + // Return upstream model if found, otherwise return requested model. + if upstreamModel != "" { + return upstreamModel + } + return requestedModel +} + +// APIKeyConfigEntry is a generic interface for API key configurations. +type APIKeyConfigEntry interface { + GetAPIKey() string + GetBaseURL() string +} + +func resolveAPIKeyConfig[T APIKeyConfigEntry](entries []T, auth *Auth) *T { + if auth == nil || len(entries) == 0 { + return nil + } + attrKey, attrBase := "", "" + if auth.Attributes != nil { + attrKey = strings.TrimSpace(auth.Attributes["api_key"]) + attrBase = strings.TrimSpace(auth.Attributes["base_url"]) + } + for i := range entries { + entry := &entries[i] + cfgKey := strings.TrimSpace((*entry).GetAPIKey()) + cfgBase := strings.TrimSpace((*entry).GetBaseURL()) + if attrKey != "" && attrBase != "" { + if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { + return entry + } + continue + } + if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { + if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey != "" { + for i := range entries { + entry := &entries[i] + if strings.EqualFold(strings.TrimSpace((*entry).GetAPIKey()), attrKey) { + return entry + } + } + } + return nil +} + +func resolveGeminiAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.GeminiKey { + if cfg == nil { + return nil + } + return resolveAPIKeyConfig(cfg.GeminiKey, auth) +} + +func resolveInteractionsAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.GeminiKey { + if cfg == nil { + return nil + } + return resolveAPIKeyConfig(cfg.InteractionsKey, auth) +} + +func resolveClaudeAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.ClaudeKey { + if cfg == nil { + return nil + } + return resolveAPIKeyConfig(cfg.ClaudeKey, auth) +} + +func resolveCodexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.CodexKey { + if cfg == nil { + return nil + } + return resolveAPIKeyConfig(cfg.CodexKey, auth) +} + +func resolveXAIAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.XAIKey { + if cfg == nil { + return nil + } + return resolveAPIKeyConfig(cfg.XAIKey, auth) +} + +func resolveVertexAPIKeyConfig(cfg *internalconfig.Config, auth *Auth) *internalconfig.VertexCompatKey { + if cfg == nil { + return nil + } + return resolveAPIKeyConfig(cfg.VertexCompatAPIKey, auth) +} + +func resolveUpstreamModelForGeminiAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { + entry := resolveGeminiAPIKeyConfig(cfg, auth) + if entry == nil { + return "" + } + return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +func resolveUpstreamModelForInteractionsAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { + entry := resolveInteractionsAPIKeyConfig(cfg, auth) + if entry == nil { + return "" + } + return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +func resolveUpstreamModelForClaudeAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { + entry := resolveClaudeAPIKeyConfig(cfg, auth) + if entry == nil { + return "" + } + return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +func resolveUpstreamModelForCodexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { + entry := resolveCodexAPIKeyConfig(cfg, auth) + if entry == nil { + return "" + } + return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +func resolveUpstreamModelForXAIAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { + entry := resolveXAIAPIKeyConfig(cfg, auth) + if entry == nil { + return "" + } + return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +func resolveUpstreamModelForVertexAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { + entry := resolveVertexAPIKeyConfig(cfg, auth) + if entry == nil { + return "" + } + return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +func resolveUpstreamModelForOpenAICompatAPIKey(cfg *internalconfig.Config, auth *Auth, requestedModel string) string { + providerKey := "" + compatName := "" + if auth != nil && len(auth.Attributes) > 0 { + providerKey = strings.TrimSpace(auth.Attributes["provider_key"]) + compatName = strings.TrimSpace(auth.Attributes["compat_name"]) + } + if compatName == "" && !strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") { + return "" + } + entry := resolveOpenAICompatConfig(cfg, providerKey, compatName, auth.Provider) + if entry == nil { + return "" + } + return resolveModelAliasFromConfigModels(requestedModel, asModelAliasEntries(entry.Models)) +} + +type apiKeyModelAliasTable map[string]map[string]string + +func resolveOpenAICompatConfig(cfg *internalconfig.Config, providerKey, compatName, authProvider string) *internalconfig.OpenAICompatibility { + if cfg == nil { + return nil + } + candidates := make([]string, 0, 3) + if v := strings.TrimSpace(compatName); v != "" { + candidates = append(candidates, v) + } + if v := strings.TrimSpace(providerKey); v != "" { + candidates = append(candidates, v) + } + if v := strings.TrimSpace(authProvider); v != "" { + candidates = append(candidates, v) + } + for i := range cfg.OpenAICompatibility { + compat := &cfg.OpenAICompatibility[i] + if compat.Disabled { + continue + } + for _, candidate := range candidates { + if candidate != "" && strings.EqualFold(strings.TrimSpace(candidate), compat.Name) { + return compat + } + } + } + return nil +} + +func asModelAliasEntries[T interface { + GetName() string + GetAlias() string + GetForceMapping() bool +}](models []T) []modelAliasEntry { + if len(models) == 0 { + return nil + } + out := make([]modelAliasEntry, 0, len(models)) + for i := range models { + out = append(out, models[i]) + } + return out +} diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go new file mode 100644 index 000000000..4d9385d4b --- /dev/null +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -0,0 +1,510 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "strconv" + "strings" + "sync" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + log "github.com/sirupsen/logrus" +) + +// RefreshEvaluator allows runtime state to override refresh decisions. +type RefreshEvaluator interface { + ShouldRefresh(now time.Time, auth *Auth) bool +} + +const ( + refreshCheckInterval = 5 * time.Second + refreshMaxConcurrency = 16 + refreshPendingBackoff = time.Minute + refreshFailureBackoff = 5 * time.Minute + // refreshIneffectiveBackoff throttles refresh attempts when an executor returns + // success but the auth still evaluates as needing refresh (e.g. token expiry + // wasn't updated). Without this guard, the auto-refresh loop can tight-loop and + // burn CPU at idle. + refreshIneffectiveBackoff = 30 * time.Second + quotaBackoffBase = time.Second + quotaBackoffMax = 30 * time.Minute + transientErrorCooldown = time.Minute +) + +// StartAutoRefresh launches a background loop that evaluates auth freshness +// every few seconds and triggers refresh operations when required. +// Only one loop is kept alive; starting a new one cancels the previous run. +func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duration) { + if interval <= 0 { + interval = refreshCheckInterval + } + + m.mu.Lock() + cancelPrev := m.refreshCancel + m.refreshCancel = nil + m.refreshLoop = nil + m.mu.Unlock() + if cancelPrev != nil { + cancelPrev() + } + + ctx, cancelCtx := context.WithCancel(parent) + workers := refreshMaxConcurrency + if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil && cfg.AuthAutoRefreshWorkers > 0 { + workers = cfg.AuthAutoRefreshWorkers + } + loop := newAuthAutoRefreshLoop(m, interval, workers) + + m.mu.Lock() + m.refreshCancel = cancelCtx + m.refreshLoop = loop + m.mu.Unlock() + + loop.rebuild(time.Now()) + go loop.run(ctx) +} + +// StopAutoRefresh cancels the background refresh loop, if running. +// It also stops the selector if it implements StoppableSelector. +func (m *Manager) StopAutoRefresh() { + m.mu.Lock() + cancel := m.refreshCancel + m.refreshCancel = nil + m.refreshLoop = nil + m.mu.Unlock() + if cancel != nil { + cancel() + } + // Stop selector if it implements StoppableSelector (e.g., SessionAffinitySelector) + if stoppable, ok := m.selector.(StoppableSelector); ok { + stoppable.Stop() + } +} + +func (m *Manager) queueRefreshReschedule(authID string) { + if m == nil || authID == "" { + return + } + m.mu.RLock() + loop := m.refreshLoop + m.mu.RUnlock() + if loop == nil { + return + } + loop.queueReschedule(authID) +} + +func (m *Manager) queueRefreshUnschedule(authID string) { + if m == nil || authID == "" { + return + } + m.mu.RLock() + loop := m.refreshLoop + m.mu.RUnlock() + if loop == nil { + return + } + loop.remove(authID) +} + +func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool { + if a == nil { + return false + } + if hasUnauthorizedAuthFailure(a) { + return false + } + if !a.NextRefreshAfter.IsZero() && now.Before(a.NextRefreshAfter) { + return false + } + if evaluator, ok := a.Runtime.(RefreshEvaluator); ok && evaluator != nil { + return evaluator.ShouldRefresh(now, a) + } + + lastRefresh := a.LastRefreshedAt + if lastRefresh.IsZero() { + if ts, ok := authLastRefreshTimestamp(a); ok { + lastRefresh = ts + } + } + + expiry, hasExpiry := a.ExpirationTime() + + if interval := authPreferredInterval(a); interval > 0 { + if hasExpiry && !expiry.IsZero() { + if !expiry.After(now) { + return true + } + if expiry.Sub(now) <= interval { + return true + } + } + if lastRefresh.IsZero() { + return true + } + return now.Sub(lastRefresh) >= interval + } + + provider := strings.ToLower(a.Provider) + lead := ProviderRefreshLead(provider, a.Runtime) + if lead == nil { + return false + } + if *lead <= 0 { + if hasExpiry && !expiry.IsZero() { + return now.After(expiry) + } + return false + } + if hasExpiry && !expiry.IsZero() { + return time.Until(expiry) <= *lead + } + if !lastRefresh.IsZero() { + return now.Sub(lastRefresh) >= *lead + } + return true +} + +func authPreferredInterval(a *Auth) time.Duration { + if a == nil { + return 0 + } + if d := durationFromMetadata(a.Metadata, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 { + return d + } + if d := durationFromAttributes(a.Attributes, "refresh_interval_seconds", "refreshIntervalSeconds", "refresh_interval", "refreshInterval"); d > 0 { + return d + } + return 0 +} + +func durationFromMetadata(meta map[string]any, keys ...string) time.Duration { + if len(meta) == 0 { + return 0 + } + for _, key := range keys { + if val, ok := meta[key]; ok { + if dur := parseDurationValue(val); dur > 0 { + return dur + } + } + } + return 0 +} + +func durationFromAttributes(attrs map[string]string, keys ...string) time.Duration { + if len(attrs) == 0 { + return 0 + } + for _, key := range keys { + if val, ok := attrs[key]; ok { + if dur := parseDurationString(val); dur > 0 { + return dur + } + } + } + return 0 +} + +func parseDurationValue(val any) time.Duration { + switch v := val.(type) { + case time.Duration: + if v <= 0 { + return 0 + } + return v + case int: + if v <= 0 { + return 0 + } + return time.Duration(v) * time.Second + case int32: + if v <= 0 { + return 0 + } + return time.Duration(v) * time.Second + case int64: + if v <= 0 { + return 0 + } + return time.Duration(v) * time.Second + case uint: + if v == 0 { + return 0 + } + return time.Duration(v) * time.Second + case uint32: + if v == 0 { + return 0 + } + return time.Duration(v) * time.Second + case uint64: + if v == 0 { + return 0 + } + return time.Duration(v) * time.Second + case float32: + if v <= 0 { + return 0 + } + return time.Duration(float64(v) * float64(time.Second)) + case float64: + if v <= 0 { + return 0 + } + return time.Duration(v * float64(time.Second)) + case json.Number: + if i, err := v.Int64(); err == nil { + if i <= 0 { + return 0 + } + return time.Duration(i) * time.Second + } + if f, err := v.Float64(); err == nil && f > 0 { + return time.Duration(f * float64(time.Second)) + } + case string: + return parseDurationString(v) + } + return 0 +} + +func parseDurationString(raw string) time.Duration { + s := strings.TrimSpace(raw) + if s == "" { + return 0 + } + if dur, err := time.ParseDuration(s); err == nil && dur > 0 { + return dur + } + if secs, err := strconv.ParseFloat(s, 64); err == nil && secs > 0 { + return time.Duration(secs * float64(time.Second)) + } + return 0 +} + +func authLastRefreshTimestamp(a *Auth) (time.Time, bool) { + if a == nil { + return time.Time{}, false + } + if a.Metadata != nil { + if ts, ok := lookupMetadataTime(a.Metadata, "last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"); ok { + return ts, true + } + } + if a.Attributes != nil { + for _, key := range []string{"last_refresh", "lastRefresh", "last_refreshed_at", "lastRefreshedAt"} { + if val := strings.TrimSpace(a.Attributes[key]); val != "" { + if ts, ok := parseTimeValue(val); ok { + return ts, true + } + } + } + } + return time.Time{}, false +} + +func lookupMetadataTime(meta map[string]any, keys ...string) (time.Time, bool) { + for _, key := range keys { + if val, ok := meta[key]; ok { + if ts, ok1 := parseTimeValue(val); ok1 { + return ts, true + } + } + } + return time.Time{}, false +} + +func (m *Manager) markRefreshPending(id string, now time.Time) bool { + m.mu.Lock() + auth, ok := m.auths[id] + if !ok || auth == nil { + m.mu.Unlock() + return false + } + if !auth.NextRefreshAfter.IsZero() && now.Before(auth.NextRefreshAfter) { + m.mu.Unlock() + return false + } + auth.NextRefreshAfter = now.Add(refreshPendingBackoff) + m.auths[id] = auth + m.mu.Unlock() + + m.queueRefreshReschedule(id) + return true +} + +type authRefreshLock struct { + mu sync.Mutex +} + +func authAccessToken(auth *Auth) string { + if token := authMetadataString(auth, "access_token"); token != "" { + return token + } + return authMetadataString(auth, "accessToken") +} + +func authHasRefreshCredential(auth *Auth) bool { + if authMetadataString(auth, "refresh_token") != "" { + return true + } + return authMetadataString(auth, "refreshToken") != "" +} + +func clearUnauthorizedModelStates(auth *Auth, now time.Time) []string { + if auth == nil || len(auth.ModelStates) == 0 { + return nil + } + var resumed []string + for model, state := range auth.ModelStates { + if state == nil || state.LastError == nil { + continue + } + if state.LastError.StatusCode() != http.StatusUnauthorized && !strings.EqualFold(state.LastError.Code, "unauthorized") { + continue + } + resetModelState(state, now) + resumed = append(resumed, model) + } + if len(resumed) > 0 { + updateAggregatedAvailability(auth, now) + } + return resumed +} + +// tryRefreshAfterUnauthorized refreshes OAuth credentials once after a 401 so the +// current auth can be retried before fallback/suspend. +func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, execErr error, alreadyTried bool) (*Auth, bool) { + if m == nil || auth == nil || alreadyTried || execErr == nil { + return auth, false + } + if !isUnauthorizedError(execErr) || !authHasRefreshCredential(auth) { + return auth, false + } + log.Debugf("unauthorized response for %s (%s), refreshing credentials before fallback", auth.Provider, auth.ID) + refreshed, errRefresh := m.refreshAuthForRequest(ctx, auth.ID, authAccessToken(auth)) + if errRefresh != nil || refreshed == nil { + log.Debugf("credential refresh before fallback failed for %s (%s): %v", auth.Provider, auth.ID, errRefresh) + return auth, false + } + return refreshed, true +} + +func (m *Manager) refreshAuth(ctx context.Context, id string) { + _, _ = m.refreshAuthForRequest(ctx, id, "") +} + +// refreshAuthForRequest performs a synchronous credential refresh for the given auth. +// failedAccessToken lets concurrent callers reuse a refresh that already replaced the +// access token that produced the unauthorized response. +func (m *Manager) refreshAuthForRequest(ctx context.Context, id, failedAccessToken string) (*Auth, error) { + if m == nil { + return nil, errors.New("auth manager is nil") + } + if ctx == nil { + ctx = context.Background() + } + id = strings.TrimSpace(id) + if id == "" { + return nil, errors.New("auth id is empty") + } + + lockValue, _ := m.refreshLocks.LoadOrStore(id, &authRefreshLock{}) + lock, _ := lockValue.(*authRefreshLock) + if lock == nil { + lock = &authRefreshLock{} + m.refreshLocks.Store(id, lock) + } + lock.mu.Lock() + defer lock.mu.Unlock() + + m.mu.RLock() + auth := m.auths[id] + var exec ProviderExecutor + if auth != nil { + exec = m.executors[auth.Provider] + } + m.mu.RUnlock() + if auth == nil || exec == nil { + return nil, errors.New("auth or executor not found") + } + + // Another request may already have refreshed this credential. + if failedAccessToken != "" { + if currentToken := authAccessToken(auth); currentToken != "" && currentToken != failedAccessToken { + return auth.Clone(), nil + } + } + + cloned := auth.Clone() + updated, err := exec.Refresh(ctx, cloned) + if err != nil && errors.Is(err, context.Canceled) { + log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID) + return nil, err + } + log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err) + now := time.Now() + if err != nil { + unauthorized := isUnauthorizedError(err) + shouldReschedule := false + m.mu.Lock() + if current := m.auths[id]; current != nil { + current.LastError = refreshErrorFromError(err) + if unauthorized { + current.NextRefreshAfter = time.Time{} + current.Unavailable = true + current.Status = StatusError + current.StatusMessage = "unauthorized" + } else { + current.NextRefreshAfter = now.Add(refreshFailureBackoff) + } + m.auths[id] = current + shouldReschedule = true + if m.scheduler != nil { + m.scheduler.upsertAuth(current.Clone()) + } + } + m.mu.Unlock() + if shouldReschedule { + m.queueRefreshReschedule(id) + } + return nil, err + } + if updated == nil { + updated = cloned + } + // Preserve runtime created by the executor during Refresh. + // If executor didn't set one, fall back to the previous runtime. + if updated.Runtime == nil { + updated.Runtime = auth.Runtime + } + updated.LastRefreshedAt = now + updated.NextRefreshAfter = time.Time{} + updated.LastError = nil + updated.StatusMessage = "" + updated.Unavailable = false + if updated.Status == StatusError { + updated.Status = StatusActive + } + updated.UpdatedAt = now + modelsToResume := clearUnauthorizedModelStates(updated, now) + if m.shouldRefresh(updated, now) { + updated.NextRefreshAfter = now.Add(refreshIneffectiveBackoff) + } + saved, errUpdate := m.Update(ctx, updated) + for _, model := range modelsToResume { + registry.GetGlobalRegistry().ResumeClientModel(id, model) + } + if errUpdate != nil { + log.Debugf("persist refreshed auth %s (%s) failed: %v", auth.Provider, auth.ID, errUpdate) + } + if saved != nil { + return saved, nil + } + return updated.Clone(), nil +} diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go new file mode 100644 index 000000000..97f1a35af --- /dev/null +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -0,0 +1,1361 @@ +package auth + +import ( + "context" + "errors" + "math/rand/v2" + "net/http" + "sort" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" +) + +func (m *Manager) SetPluginScheduler(scheduler PluginScheduler) { + if m == nil { + return + } + m.mu.Lock() + m.pluginScheduler = scheduler + m.mu.Unlock() +} + +func (m *Manager) hasPluginScheduler() bool { + if m == nil { + return false + } + m.mu.RLock() + scheduler := m.pluginScheduler + m.mu.RUnlock() + if scheduler == nil { + return false + } + if state, ok := scheduler.(pluginSchedulerState); ok { + return state.HasScheduler() + } + return true +} + +func isBuiltInSelector(selector Selector) bool { + switch selector.(type) { + case *RoundRobinSelector, *FillFirstSelector: + return true + default: + return false + } +} + +func (m *Manager) syncSchedulerFromSnapshot(auths []*Auth) { + if m == nil || m.scheduler == nil { + return + } + m.scheduler.rebuild(auths) +} + +func (m *Manager) syncScheduler() { + if m == nil || m.scheduler == nil { + return + } + m.syncSchedulerFromSnapshot(m.snapshotAuths()) +} + +func (m *Manager) snapshotAuths() []*Auth { + m.mu.RLock() + defer m.mu.RUnlock() + out := make([]*Auth, 0, len(m.auths)) + for _, a := range m.auths { + out = append(out, a.Clone()) + } + return out +} + +// RefreshSchedulerEntry re-upserts a single auth into the scheduler so that its +// supportedModelSet is rebuilt from the current global model registry state. +// This must be called after models have been registered for a newly added auth, +// because the initial scheduler.upsertAuth during Register/Update runs before +// registerModelsForAuth and therefore snapshots an empty model set. +func (m *Manager) RefreshSchedulerEntry(authID string) { + if m == nil || m.scheduler == nil || authID == "" { + return + } + m.mu.RLock() + auth, ok := m.auths[authID] + if !ok || auth == nil { + m.mu.RUnlock() + return + } + snapshot := auth.Clone() + m.mu.RUnlock() + m.scheduler.upsertAuth(snapshot) +} + +// RefreshSchedulerAll rebuilds scheduler entries for every known auth. +func (m *Manager) RefreshSchedulerAll() { + if m == nil { + return + } + m.mu.RLock() + ids := make([]string, 0, len(m.auths)) + for id := range m.auths { + ids = append(ids, id) + } + m.mu.RUnlock() + for _, id := range ids { + m.RefreshSchedulerEntry(id) + } +} + +// ReconcileRegistryModelStates aligns per-model runtime state with the current +// registry snapshot for one auth. +// +// Supported models are reset to a clean state because re-registration already +// cleared the registry-side cooldown/suspension snapshot. ModelStates for +// models that are no longer present in the registry are pruned entirely so +// renamed/removed models cannot keep auth-level status stale. +func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID string) { + if m == nil || authID == "" { + return + } + + supportedModels := registry.GetGlobalRegistry().GetModelsForClient(authID) + supported := make(map[string]struct{}, len(supportedModels)) + for _, model := range supportedModels { + if model == nil { + continue + } + modelKey := canonicalModelKey(model.ID) + if modelKey == "" { + continue + } + supported[modelKey] = struct{}{} + } + + var snapshot *Auth + now := time.Now() + + m.mu.Lock() + auth, ok := m.auths[authID] + if ok && auth != nil && len(auth.ModelStates) > 0 { + changed := false + for modelKey, state := range auth.ModelStates { + baseModel := canonicalModelKey(modelKey) + if baseModel == "" { + baseModel = strings.TrimSpace(modelKey) + } + if _, supportedModel := supported[baseModel]; !supportedModel { + // Drop state for models that disappeared from the current registry + // snapshot. Keeping them around leaks stale errors into auth-level + // status, management output, and websocket fallback checks. + delete(auth.ModelStates, modelKey) + changed = true + continue + } + if state == nil { + continue + } + if modelStateIsClean(state) { + continue + } + resetModelState(state, now) + changed = true + } + if len(auth.ModelStates) == 0 { + auth.ModelStates = nil + } + if changed { + updateAggregatedAvailability(auth, now) + if !hasModelError(auth, now) { + auth.LastError = nil + auth.StatusMessage = "" + auth.Status = StatusActive + } + auth.UpdatedAt = now + if errPersist := m.persist(ctx, auth); errPersist != nil { + logEntryWithRequestID(ctx).WithField("auth_id", auth.ID).Warnf("failed to persist auth changes during model state reconciliation: %v", errPersist) + } + snapshot = auth.Clone() + } + } + m.mu.Unlock() + + if m.scheduler != nil && snapshot != nil { + m.scheduler.upsertAuth(snapshot) + } +} + +func (m *Manager) SetSelector(selector Selector) { + if m == nil { + return + } + if selector == nil { + selector = &RoundRobinSelector{} + } + m.mu.Lock() + m.selector = selector + m.mu.Unlock() + if m.scheduler != nil { + m.scheduler.setSelector(selector) + m.syncScheduler() + } +} + +// Selector returns the current credential selector. +func (m *Manager) Selector() Selector { + if m == nil { + return nil + } + m.mu.RLock() + defer m.mu.RUnlock() + return m.selector +} + +// SetStore swaps the underlying persistence store. +func (m *Manager) SetStore(store Store) { + m.mu.Lock() + defer m.mu.Unlock() + m.store = store +} + +// SetCooldownStateStore swaps the independent runtime cooldown state store. +func (m *Manager) SetCooldownStateStore(store CooldownStateStore) { + if m == nil { + return + } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + m.mu.Lock() + defer m.mu.Unlock() + m.cooldownStore = store +} + +// SetRoundTripperProvider register a provider that returns a per-auth RoundTripper. +func (m *Manager) SetRoundTripperProvider(p RoundTripperProvider) { + m.mu.Lock() + m.rtProvider = p + m.mu.Unlock() +} + +func (m *Manager) availableAuthsForRouteModel(auths []*Auth, provider, routeModel string, now time.Time) ([]*Auth, error) { + if len(auths) == 0 { + return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"} + } + + availableByPriority := make(map[int][]*Auth) + cooldownCount := 0 + var earliest time.Time + for _, candidate := range auths { + checkModel := m.selectionModelForAuth(candidate, routeModel) + blocked, reason, next := isAuthBlockedForModel(candidate, checkModel, now) + if !blocked { + priority := authPriority(candidate) + availableByPriority[priority] = append(availableByPriority[priority], candidate) + continue + } + if reason == blockReasonCooldown { + cooldownCount++ + if !next.IsZero() && (earliest.IsZero() || next.Before(earliest)) { + earliest = next + } + } + } + + if len(availableByPriority) == 0 { + if cooldownCount == len(auths) && !earliest.IsZero() { + providerForError := provider + if providerForError == "mixed" { + providerForError = "" + } + resetIn := earliest.Sub(now) + if resetIn < 0 { + resetIn = 0 + } + return nil, newModelCooldownError(routeModel, providerForError, resetIn) + } + return nil, &Error{Code: "auth_unavailable", Message: "no auth available"} + } + + bestPriority := 0 + found := false + for priority := range availableByPriority { + if !found || priority > bestPriority { + bestPriority = priority + found = true + } + } + + available := availableByPriority[bestPriority] + if len(available) > 1 { + sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID }) + } + return available, nil +} + +func selectionArgForSelector(selector Selector, routeModel string) string { + if isBuiltInSelector(selector) { + return "" + } + return routeModel +} + +func schedulerAttributeSensitive(key string) bool { + key = strings.ToLower(strings.TrimSpace(key)) + normalized := strings.NewReplacer("-", "_", ".", "_", " ", "_").Replace(key) + compact := strings.NewReplacer("_", "", "-", "", ".", "", " ", "").Replace(key) + for _, fragment := range []string{ + "api_key", + "apikey", + "token", + "secret", + "cookie", + "credential", + "password", + "storage", + "authorization", + "auth_header", + "proxy_url", + } { + if strings.Contains(key, fragment) || strings.Contains(normalized, fragment) || strings.Contains(compact, fragment) { + return true + } + } + return false +} + +func schedulerSafeAttributes(src map[string]string) map[string]string { + if len(src) == 0 { + return nil + } + out := make(map[string]string, len(src)) + for key, value := range src { + if schedulerAttributeSensitive(key) { + continue + } + out[key] = value + } + if len(out) == 0 { + return nil + } + return out +} + +func cloneSchedulerAnyMap(src map[string]any) map[string]any { + if len(src) == 0 { + return nil + } + out := make(map[string]any, len(src)) + for key, value := range src { + out[key] = value + } + return out +} + +func cloneAuthSlice(auths []*Auth) []*Auth { + if len(auths) == 0 { + return nil + } + out := make([]*Auth, 0, len(auths)) + for _, auth := range auths { + if auth == nil { + continue + } + out = append(out, auth.Clone()) + } + return out +} + +func schedulerAuthCandidates(auths []*Auth) []pluginapi.SchedulerAuthCandidate { + if len(auths) == 0 { + return nil + } + out := make([]pluginapi.SchedulerAuthCandidate, 0, len(auths)) + for _, auth := range auths { + if auth == nil { + continue + } + out = append(out, pluginapi.SchedulerAuthCandidate{ + ID: auth.ID, + Provider: strings.ToLower(strings.TrimSpace(auth.Provider)), + Priority: authPriority(auth), + Status: string(auth.Status), + Attributes: schedulerSafeAttributes(auth.Attributes), + }) + } + return out +} + +func schedulerProviders(provider string, providers []string) []string { + out := make([]string, 0, len(providers)+1) + seen := make(map[string]struct{}, len(providers)+1) + addProvider := func(value string) { + value = strings.ToLower(strings.TrimSpace(value)) + if value == "" || value == "mixed" { + return + } + if _, ok := seen[value]; ok { + return + } + seen[value] = struct{}{} + out = append(out, value) + } + addProvider(provider) + for _, value := range providers { + addProvider(value) + } + return out +} + +func schedulerOptions(opts cliproxyexecutor.Options) pluginapi.SchedulerOptions { + return pluginapi.SchedulerOptions{ + Headers: cloneHTTPHeader(opts.Headers), + Metadata: cloneSchedulerAnyMap(opts.Metadata), + } +} + +func pickSchedulerAuthByID(candidates []*Auth, authID string) *Auth { + authID = strings.TrimSpace(authID) + if authID == "" { + return nil + } + for _, candidate := range candidates { + if candidate != nil && candidate.ID == authID { + return candidate + } + } + return nil +} + +func builtinSchedulerStrategy(delegate string) (schedulerStrategy, bool) { + switch strings.TrimSpace(delegate) { + case pluginapi.SchedulerBuiltinRoundRobin: + return schedulerStrategyRoundRobin, true + case pluginapi.SchedulerBuiltinFillFirst: + return schedulerStrategyFillFirst, true + default: + return schedulerStrategyCustom, false + } +} + +func (m *Manager) pickViaBuiltinScheduler(ctx context.Context, strategy schedulerStrategy, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, bool, error) { + if m == nil || m.scheduler == nil { + return nil, false, nil + } + providerKey := strings.ToLower(strings.TrimSpace(provider)) + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + for { + var selected *Auth + var errPick error + if providerKey == "mixed" { + selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, _, errPick = m.scheduler.pickMixedWithStrategy(ctx, providers, model, opts, tried, strategy) + } + } else { + selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, errPick = m.scheduler.pickSingleWithStrategy(ctx, providerKey, model, opts, tried, strategy) + } + } + if errPick != nil { + return nil, true, errPick + } + if selected == nil { + return nil, true, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + if disallowFreeAuth && isFreeCodexAuth(selected) { + if tried == nil { + tried = make(map[string]struct{}) + } + tried[selected.ID] = struct{}{} + continue + } + return selected, true, nil + } +} + +func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginScheduler, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}, candidates []*Auth) (*Auth, bool, error) { + if scheduler == nil || len(candidates) == 0 { + return nil, false, nil + } + providerKey := strings.ToLower(strings.TrimSpace(provider)) + requestProvider := providerKey + if providerKey == "mixed" { + requestProvider = "" + } + req := pluginapi.SchedulerPickRequest{ + Provider: requestProvider, + Providers: schedulerProviders(providerKey, providers), + Model: model, + Stream: opts.Stream, + Options: schedulerOptions(opts), + Candidates: schedulerAuthCandidates(candidates), + } + resp, handled, errPick := scheduler.PickAuth(ctx, req) + if errPick != nil { + return nil, true, errPick + } + if !handled || !resp.Handled { + return nil, false, nil + } + if selected := pickSchedulerAuthByID(candidates, resp.AuthID); selected != nil { + return selected, true, nil + } + + strategy, okStrategy := builtinSchedulerStrategy(resp.DelegateBuiltin) + if !okStrategy { + return nil, false, nil + } + return m.pickViaBuiltinScheduler(ctx, strategy, providerKey, providers, model, opts, tried) +} + +func (m *Manager) authSupportsRouteModel(registryRef *registry.ModelRegistry, auth *Auth, routeModel string) bool { + if registryRef == nil || auth == nil { + return true + } + routeKey := canonicalModelKey(routeModel) + if routeKey == "" { + return true + } + if registryRef.ClientSupportsModel(auth.ID, routeKey) { + return true + } + selectionKey := m.selectionModelKeyForAuth(auth, routeModel) + return selectionKey != "" && selectionKey != routeKey && registryRef.ClientSupportsModel(auth.ID, selectionKey) +} + +func (m *Manager) normalizeProviders(providers []string) []string { + if len(providers) == 0 { + return nil + } + result := make([]string, 0, len(providers)) + seen := make(map[string]struct{}, len(providers)) + for _, provider := range providers { + p := strings.TrimSpace(strings.ToLower(provider)) + if p == "" { + continue + } + if _, ok := seen[p]; ok { + continue + } + seen[p] = struct{}{} + result = append(result, p) + } + return result +} + +// AvailableProviders returns the set of provider keys that currently have at least one +// registered auth record that is not disabled. It is a best-effort snapshot for routing +// decisions and does not account for per-model cooldowns or transient runtime availability. +// Disabled auths (Disabled flag or StatusDisabled) are excluded so routing does not target +// providers that auth selection would refuse to use, which would otherwise cause execution +// failures instead of falling back to lower-priority routers. +func (m *Manager) AvailableProviders() []string { + if m == nil { + return nil + } + m.mu.RLock() + defer m.mu.RUnlock() + seen := make(map[string]struct{}, len(m.auths)) + out := make([]string, 0, len(m.auths)) + for _, auth := range m.auths { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled { + continue + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if provider == "" { + continue + } + if _, ok := seen[provider]; ok { + continue + } + seen[provider] = struct{}{} + out = append(out, provider) + } + sort.Strings(out) + return out +} + +// HasProviderAuth reports whether at least one non-disabled auth record is registered for +// the provider. Disabled auths (Disabled flag or StatusDisabled) are excluded to match the +// behavior of auth selection, which refuses to pick disabled credentials. +func (m *Manager) HasProviderAuth(provider string) bool { + if m == nil { + return false + } + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" { + return false + } + m.mu.RLock() + defer m.mu.RUnlock() + for _, auth := range m.auths { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled { + continue + } + if strings.ToLower(strings.TrimSpace(auth.Provider)) == provider { + return true + } + } + return false +} + +func (m *Manager) retrySettings() (int, int, time.Duration) { + if m == nil { + return 0, 0, 0 + } + return int(m.requestRetry.Load()), int(m.maxRetryCredentials.Load()), time.Duration(m.maxRetryInterval.Load()) +} + +func (m *Manager) closestCooldownWait(providers []string, model string, attempt int) (time.Duration, bool) { + if m == nil || len(providers) == 0 { + return 0, false + } + now := time.Now() + defaultRetry := int(m.requestRetry.Load()) + if defaultRetry < 0 { + defaultRetry = 0 + } + providerSet := make(map[string]struct{}, len(providers)) + for i := range providers { + key := strings.TrimSpace(strings.ToLower(providers[i])) + if key == "" { + continue + } + providerSet[key] = struct{}{} + } + m.mu.RLock() + defer m.mu.RUnlock() + var ( + found bool + minWait time.Duration + ) + for _, auth := range m.auths { + if auth == nil { + continue + } + providerKey := executorKeyFromAuth(auth) + if _, ok := providerSet[providerKey]; !ok { + continue + } + effectiveRetry := defaultRetry + if override, ok := auth.RequestRetryOverride(); ok { + effectiveRetry = override + } + if effectiveRetry < 0 { + effectiveRetry = 0 + } + if attempt >= effectiveRetry { + continue + } + checkModel := model + if strings.TrimSpace(model) != "" { + checkModel = m.selectionModelForAuth(auth, model) + } + blocked, reason, next := isAuthBlockedForModel(auth, checkModel, now) + if !blocked || next.IsZero() || reason == blockReasonDisabled { + continue + } + wait := next.Sub(now) + if wait < 0 { + continue + } + if !found || wait < minWait { + minWait = wait + found = true + } + } + return minWait, found +} + +func (m *Manager) retryAllowed(attempt int, providers []string) bool { + if m == nil || attempt < 0 || len(providers) == 0 { + return false + } + defaultRetry := int(m.requestRetry.Load()) + if defaultRetry < 0 { + defaultRetry = 0 + } + providerSet := make(map[string]struct{}, len(providers)) + for i := range providers { + key := strings.TrimSpace(strings.ToLower(providers[i])) + if key == "" { + continue + } + providerSet[key] = struct{}{} + } + if len(providerSet) == 0 { + return false + } + + m.mu.RLock() + defer m.mu.RUnlock() + for _, auth := range m.auths { + if auth == nil { + continue + } + providerKey := executorKeyFromAuth(auth) + if _, ok := providerSet[providerKey]; !ok { + continue + } + effectiveRetry := defaultRetry + if override, ok := auth.RequestRetryOverride(); ok { + effectiveRetry = override + } + if effectiveRetry < 0 { + effectiveRetry = 0 + } + if attempt < effectiveRetry { + return true + } + } + return false +} + +func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []string, model string, maxWait time.Duration) (time.Duration, bool) { + if err == nil { + return 0, false + } + var homeBusy *HomeConcurrencyBusyError + if errors.As(err, &homeBusy) && homeBusy != nil { + return 0, false + } + if maxWait <= 0 { + return 0, false + } + status := statusCodeFromError(err) + if status == http.StatusOK { + return 0, false + } + if isRequestInvalidError(err) { + return 0, false + } + wait, found := m.closestCooldownWait(providers, model, attempt) + if found { + if wait > maxWait { + return 0, false + } + return wait, true + } + if status != http.StatusTooManyRequests { + return 0, false + } + if !m.retryAllowed(attempt, providers) { + return 0, false + } + retryAfter := retryAfterFromError(err) + if retryAfter == nil || *retryAfter <= 0 || *retryAfter > maxWait { + return 0, false + } + return *retryAfter, true +} + +// cooldownWaitJitterCap bounds the random jitter added to cooldown waits so a +// long wait is never extended by more than this amount. +const cooldownWaitJitterCap = 2 * time.Second + +// jitteredCooldownWait adds a small random delay to a cooldown wait so +// concurrent requests waiting on the same recovery deadline do not wake in +// lockstep and stampede the first credential that recovers. The jitter never +// pushes the total wait past maxWait, which callers have already enforced as +// the retry ceiling; maxWait <= 0 means no ceiling. +func jitteredCooldownWait(wait, maxWait time.Duration) time.Duration { + if wait <= 0 { + return wait + } + jitterRange := wait / 4 + if jitterRange > cooldownWaitJitterCap { + jitterRange = cooldownWaitJitterCap + } + if maxWait > 0 && jitterRange > maxWait-wait { + jitterRange = maxWait - wait + } + if jitterRange <= 0 { + return wait + } + return wait + rand.N(jitterRange) +} + +func waitForCooldown(ctx context.Context, wait, maxWait time.Duration) error { + if wait <= 0 { + return nil + } + timer := time.NewTimer(jitteredCooldownWait(wait, maxWait)) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// List returns all auth entries currently known by the manager. +func (m *Manager) List() []*Auth { + m.mu.RLock() + defer m.mu.RUnlock() + list := make([]*Auth, 0, len(m.auths)) + for _, auth := range m.auths { + list = append(list, auth.Clone()) + } + return list +} + +// GetByID retrieves an auth entry by its ID. +func (m *Manager) GetByID(id string) (*Auth, bool) { + if id == "" { + return nil, false + } + m.mu.RLock() + defer m.mu.RUnlock() + auth, ok := m.auths[id] + if !ok { + return nil, false + } + return auth.Clone(), true +} + +// GetExecutionSessionAuthByID retrieves a Home runtime auth scoped to an execution session. +func (m *Manager) GetExecutionSessionAuthByID(sessionID string, authID string) (*Auth, bool) { + sessionID = strings.TrimSpace(sessionID) + authID = strings.TrimSpace(authID) + if m == nil || sessionID == "" || authID == "" { + return nil, false + } + m.mu.RLock() + defer m.mu.RUnlock() + sessionAuths := m.homeRuntimeAuths[sessionID] + auth := sessionAuths[authID] + if auth == nil { + return nil, false + } + return auth.Clone(), true +} + +// Executor returns the registered provider executor for a provider key. +func (m *Manager) Executor(provider string) (ProviderExecutor, bool) { + if m == nil { + return nil, false + } + provider = strings.TrimSpace(provider) + if provider == "" { + return nil, false + } + + m.mu.RLock() + executor, okExecutor := m.executors[provider] + if !okExecutor { + lowerProvider := strings.ToLower(provider) + if lowerProvider != provider { + executor, okExecutor = m.executors[lowerProvider] + } + } + m.mu.RUnlock() + + if !okExecutor || executor == nil { + return nil, false + } + return executor, true +} + +// CloseExecutionSession asks all registered executors to release the supplied execution session. +func (m *Manager) CloseExecutionSession(sessionID string) { + sessionID = strings.TrimSpace(sessionID) + if m == nil || sessionID == "" { + return + } + + m.mu.Lock() + var selections []*HomeDispatchSelection + if sessionID == CloseAllExecutionSessionsID { + m.clearHomeRuntimeAuthsLocked() + selections = m.takeAllHomeSessionSelectionsLocked() + m.clearHomeSessionLocks() + } else { + m.clearHomeRuntimeAuthsForSessionLocked(sessionID) + selections = m.takeHomeSessionSelectionsLocked(sessionID) + m.homeSessionLocks.Delete(sessionID) + } + executors := make([]ProviderExecutor, 0, len(m.executors)) + for _, exec := range m.executors { + executors = append(executors, exec) + } + m.mu.Unlock() + + for _, selection := range selections { + selection.End("session_closed") + } + for i := range executors { + if closer, ok := executors[i].(ExecutionSessionCloser); ok && closer != nil { + closer.CloseExecutionSession(sessionID) + } + } +} + +func (m *Manager) useSchedulerFastPath() bool { + if m == nil || m.scheduler == nil { + return false + } + return isBuiltInSelector(m.selector) +} + +func shouldRetrySchedulerPick(err error) bool { + if err == nil { + return false + } + var cooldownErr *modelCooldownError + if errors.As(err, &cooldownErr) { + return true + } + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + return false + } + return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" +} + +func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) bool { + if auth == nil || strings.TrimSpace(routeModel) == "" { + return false + } + return m.selectionModelKeyForAuth(auth, routeModel) != canonicalModelKey(routeModel) +} + +func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { + if m.HomeEnabled() { + auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) + return auth, exec, err + } + + pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + + m.mu.RLock() + selector := m.selector + pluginScheduler := m.pluginScheduler + executor, okExecutor := m.executors[provider] + if !okExecutor { + m.mu.RUnlock() + return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"} + } + candidates := make([]*Auth, 0, len(m.auths)) + modelKey := strings.TrimSpace(model) + // Always use base model name (without thinking suffix) for auth matching. + if modelKey != "" { + parsed := thinking.ParseSuffix(modelKey) + if parsed.ModelName != "" { + modelKey = strings.TrimSpace(parsed.ModelName) + } + } + registryRef := registry.GetGlobalRegistry() + for _, candidate := range m.auths { + if candidate == nil || executorKeyFromAuth(candidate) != provider || candidate.Disabled { + continue + } + if pinnedAuthID != "" && candidate.ID != pinnedAuthID { + continue + } + if disallowFreeAuth && isFreeCodexAuth(candidate) { + continue + } + if _, used := tried[candidate.ID]; used { + continue + } + if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { + continue + } + candidates = append(candidates, candidate) + } + if len(candidates) == 0 { + m.mu.RUnlock() + return nil, nil, &Error{Code: "auth_not_found", Message: "no auth available"} + } + available, errAvailable := m.availableAuthsForRouteModel(candidates, provider, model, time.Now()) + if errAvailable != nil { + m.mu.RUnlock() + return nil, nil, errAvailable + } + available = cloneAuthSlice(available) + m.mu.RUnlock() + + selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, provider, []string{provider}, model, opts, tried, available) + if errPick != nil { + return nil, nil, errPick + } + if !handled { + selected, errPick = selector.Pick(ctx, provider, selectionArgForSelector(selector, model), opts, available) + if errPick != nil { + return nil, nil, errPick + } + } + if selected == nil { + return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() + } + m.mu.Unlock() + } + return authCopy, executor, nil +} + +// SelectAuth selects one credential through the configured scheduling strategy. +// It does not execute or alter the selected credential's result state. +func (m *Manager) SelectAuth(ctx context.Context, provider, model string, opts cliproxyexecutor.Options) (*Auth, error) { + if m != nil && m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, nil) + if errPick != nil { + return nil, errPick + } + if m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + return selected, nil +} + +// SelectAuthByKind selects one credential of the required kind through the +// configured scheduling strategy. Credentials of other kinds are skipped. +func (m *Manager) SelectAuthByKind(ctx context.Context, provider, model, requiredKind string, opts cliproxyexecutor.Options) (*Auth, error) { + if m != nil && m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + requiredKind = normalizeAuthKind(requiredKind) + if requiredKind == "" { + return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} + } + + tried := make(map[string]struct{}) + for { + selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, tried) + if errPick != nil { + return nil, errPick + } + if selected == nil { + return nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + if selected.AuthKind() == requiredKind { + if m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + return selected, nil + } + authID := strings.TrimSpace(selected.ID) + if authID == "" { + return nil, &Error{Code: "auth_not_found", Message: "selected auth has no ID"} + } + if _, alreadyTried := tried[authID]; alreadyTried { + return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} + } + tried[authID] = struct{}{} + } +} + +// SelectHomeAuthByKind selects a Home dispatch while retaining its execution scope. +func (m *Manager) SelectHomeAuthByKind(ctx context.Context, provider string, model string, requiredKind string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { + requiredKind = normalizeAuthKind(requiredKind) + if requiredKind == "" { + return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} + } + if m == nil || !m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + + homeAuthCount := homeAuthCountFromMetadata(opts.Metadata) + tried := make(map[string]struct{}) + for { + selectionOpts := withHomeAuthCount(opts, homeAuthCount) + selection, errSelection := m.pickHomeDispatchSelection(ctx, model, selectionOpts) + if errSelection != nil { + return nil, errSelection + } + providerMatches := strings.TrimSpace(provider) == "" || strings.EqualFold(strings.TrimSpace(selection.Provider), strings.TrimSpace(provider)) + kindMatches := selection.Auth != nil && selection.Auth.AuthKind() == requiredKind + if providerMatches && kindMatches { + return selection, nil + } + + authID := "" + if selection.Auth != nil { + authID = strings.TrimSpace(selection.Auth.ID) + } + reason := "auth_kind_mismatch" + if !providerMatches { + reason = "provider_mismatch" + } + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, reason); errEnd != nil { + return nil, errEnd + } + if authID == "" { + return nil, &Error{Code: "auth_not_found", Message: "selected auth has no ID"} + } + if _, alreadyTried := tried[authID]; alreadyTried { + return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} + } + tried[authID] = struct{}{} + homeAuthCount++ + } +} + +func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { + if m.HomeEnabled() { + auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) + return auth, exec, err + } + + if m.hasPluginScheduler() || !m.useSchedulerFastPath() { + return m.pickNextLegacy(ctx, provider, model, opts, tried) + } + if strings.TrimSpace(model) != "" { + m.mu.RLock() + for _, candidate := range m.auths { + if candidate == nil || executorKeyFromAuth(candidate) != provider || candidate.Disabled { + continue + } + if _, used := tried[candidate.ID]; used { + continue + } + if m.routeAwareSelectionRequired(candidate, model) { + m.mu.RUnlock() + return m.pickNextLegacy(ctx, provider, model, opts, tried) + } + } + m.mu.RUnlock() + } + executor, okExecutor := m.Executor(provider) + if !okExecutor { + return nil, nil, &Error{Code: "executor_not_found", Message: "executor not registered"} + } + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + for { + selected, errPick := m.scheduler.pickSingle(ctx, provider, model, opts, tried) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, errPick = m.scheduler.pickSingle(ctx, provider, model, opts, tried) + } + if errPick != nil { + return nil, nil, errPick + } + if selected == nil { + return nil, nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + if disallowFreeAuth && isFreeCodexAuth(selected) { + if tried == nil { + tried = make(map[string]struct{}) + } + tried[selected.ID] = struct{}{} + continue + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() + } + m.mu.Unlock() + } + return authCopy, executor, nil + } +} + +func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { + if m.HomeEnabled() { + return m.pickNextViaHome(ctx, model, opts, tried) + } + + pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + + providerSet := make(map[string]struct{}, len(providers)) + for _, provider := range providers { + p := strings.TrimSpace(strings.ToLower(provider)) + if p == "" { + continue + } + providerSet[p] = struct{}{} + } + if len(providerSet) == 0 { + return nil, nil, "", &Error{Code: "provider_not_found", Message: "no provider supplied"} + } + + m.mu.RLock() + selector := m.selector + pluginScheduler := m.pluginScheduler + candidates := make([]*Auth, 0, len(m.auths)) + modelKey := strings.TrimSpace(model) + // Always use base model name (without thinking suffix) for auth matching. + if modelKey != "" { + parsed := thinking.ParseSuffix(modelKey) + if parsed.ModelName != "" { + modelKey = strings.TrimSpace(parsed.ModelName) + } + } + registryRef := registry.GetGlobalRegistry() + for _, candidate := range m.auths { + if candidate == nil || candidate.Disabled { + continue + } + if pinnedAuthID != "" && candidate.ID != pinnedAuthID { + continue + } + if disallowFreeAuth && isFreeCodexAuth(candidate) { + continue + } + providerKey := executorKeyFromAuth(candidate) + if providerKey == "" { + continue + } + if _, ok := providerSet[providerKey]; !ok { + continue + } + if _, used := tried[candidate.ID]; used { + continue + } + if _, ok := m.executors[providerKey]; !ok { + continue + } + if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { + continue + } + candidates = append(candidates, candidate) + } + if len(candidates) == 0 { + m.mu.RUnlock() + return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} + } + available, errAvailable := m.availableAuthsForRouteModel(candidates, "mixed", model, time.Now()) + if errAvailable != nil { + m.mu.RUnlock() + return nil, nil, "", errAvailable + } + available = cloneAuthSlice(available) + m.mu.RUnlock() + + selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, "mixed", providers, model, opts, tried, available) + if errPick != nil { + return nil, nil, "", errPick + } + if !handled { + selected, errPick = selector.Pick(ctx, "mixed", selectionArgForSelector(selector, model), opts, available) + if errPick != nil { + return nil, nil, "", errPick + } + } + if selected == nil { + return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + providerKey := executorKeyFromAuth(selected) + executor, okExecutor := m.Executor(providerKey) + if !okExecutor { + return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() + } + m.mu.Unlock() + } + return authCopy, executor, providerKey, nil +} + +func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { + if m.HomeEnabled() { + return m.pickNextViaHome(ctx, model, opts, tried) + } + + if m.hasPluginScheduler() || !m.useSchedulerFastPath() { + return m.pickNextMixedLegacy(ctx, providers, model, opts, tried) + } + + eligibleProviders := make([]string, 0, len(providers)) + seenProviders := make(map[string]struct{}, len(providers)) + for _, provider := range providers { + providerKey := strings.TrimSpace(strings.ToLower(provider)) + if providerKey == "" { + continue + } + if _, seen := seenProviders[providerKey]; seen { + continue + } + if _, okExecutor := m.Executor(providerKey); !okExecutor { + continue + } + seenProviders[providerKey] = struct{}{} + eligibleProviders = append(eligibleProviders, providerKey) + } + if len(eligibleProviders) == 0 { + return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} + } + if strings.TrimSpace(model) != "" { + providerSet := make(map[string]struct{}, len(eligibleProviders)) + for _, providerKey := range eligibleProviders { + providerSet[providerKey] = struct{}{} + } + m.mu.RLock() + for _, candidate := range m.auths { + if candidate == nil || candidate.Disabled { + continue + } + if _, ok := providerSet[executorKeyFromAuth(candidate)]; !ok { + continue + } + if _, used := tried[candidate.ID]; used { + continue + } + if m.routeAwareSelectionRequired(candidate, model) { + m.mu.RUnlock() + return m.pickNextMixedLegacy(ctx, providers, model, opts, tried) + } + } + m.mu.RUnlock() + } + + disallowFreeAuth := disallowFreeAuthFromMetadata(opts.Metadata) + for { + selected, providerKey, errPick := m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) + if errPick != nil && model != "" && shouldRetrySchedulerPick(errPick) { + m.syncScheduler() + selected, providerKey, errPick = m.scheduler.pickMixed(ctx, eligibleProviders, model, opts, tried) + } + if errPick != nil { + return nil, nil, "", errPick + } + if selected == nil { + return nil, nil, "", &Error{Code: "auth_not_found", Message: "selector returned no auth"} + } + if disallowFreeAuth && isFreeCodexAuth(selected) { + if tried == nil { + tried = make(map[string]struct{}) + } + tried[selected.ID] = struct{}{} + continue + } + executor, okExecutor := m.Executor(providerKey) + if !okExecutor { + return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered"} + } + authCopy := selected.Clone() + if !selected.indexAssigned { + m.mu.Lock() + if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned { + current.EnsureIndex() + authCopy = current.Clone() + } + m.mu.Unlock() + } + return authCopy, executor, providerKey, nil + } +} diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go new file mode 100644 index 000000000..b81920002 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -0,0 +1,309 @@ +package auth + +import ( + "context" + "net/http" + "strings" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func discardStreamChunks(ch <-chan cliproxyexecutor.StreamChunk) { + if ch == nil { + return + } + go func() { + for range ch { + } + }() +} + +type streamBootstrapError struct { + cause error + headers http.Header +} + +func cloneHTTPHeader(headers http.Header) http.Header { + if headers == nil { + return nil + } + return headers.Clone() +} + +func newStreamBootstrapError(err error, headers http.Header) error { + if err == nil { + return nil + } + return &streamBootstrapError{ + cause: err, + headers: cloneHTTPHeader(headers), + } +} + +func (e *streamBootstrapError) Error() string { + if e == nil || e.cause == nil { + return "" + } + return e.cause.Error() +} + +func (e *streamBootstrapError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func (e *streamBootstrapError) Headers() http.Header { + if e == nil { + return nil + } + return cloneHTTPHeader(e.headers) +} + +func streamErrorResult(headers http.Header, err error) *cliproxyexecutor.StreamResult { + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Err: err} + close(ch) + return &cliproxyexecutor.StreamResult{ + Headers: cloneHTTPHeader(headers), + Chunks: ch, + } +} + +func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamChunk) ([]cliproxyexecutor.StreamChunk, bool, error) { + if ch == nil { + return nil, true, nil + } + buffered := make([]cliproxyexecutor.StreamChunk, 0, 1) + for { + var ( + chunk cliproxyexecutor.StreamChunk + ok bool + ) + if ctx != nil { + select { + case <-ctx.Done(): + return nil, false, ctx.Err() + case chunk, ok = <-ch: + } + } else { + chunk, ok = <-ch + } + if !ok { + return buffered, true, nil + } + if chunk.Err != nil { + return nil, false, chunk.Err + } + buffered = append(buffered, chunk) + if len(chunk.Payload) > 0 { + return buffered, false, nil + } + } +} + +func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool) *cliproxyexecutor.StreamResult { + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(out) + var failed bool + forward := true + var rewriter *StreamRewriter + if aliasResult.ForceMapping && strings.TrimSpace(aliasResult.OriginalAlias) != "" { + rewriter = NewStreamRewriter(StreamRewriteOptions{RewriteModel: aliasResult.OriginalAlias}) + } + emit := func(chunk cliproxyexecutor.StreamChunk) bool { + if chunk.Err != nil && !failed { + failed = true + rerr := resultErrorFromError(chunk.Err) + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}, auth, ephemeralResult) + } + if !forward { + return false + } + if chunk.Err != nil { + if ctx == nil { + out <- chunk + return true + } + select { + case <-ctx.Done(): + forward = false + return false + case out <- chunk: + return true + } + } + if len(chunk.Payload) == 0 { + return true + } + payload := rewriteForceMappedStreamChunk(rewriter, chunk.Payload) + if len(payload) == 0 { + return true + } + chunk.Payload = payload + if ctx == nil { + out <- chunk + return true + } + select { + case <-ctx.Done(): + forward = false + return false + case out <- chunk: + return true + } + } + for _, chunk := range buffered { + if ok := emit(chunk); !ok { + discardStreamChunks(remaining) + return + } + } + for chunk := range remaining { + if ok := emit(chunk); !ok { + discardStreamChunks(remaining) + return + } + } + if tail := finishForceMappedStreamChunks(rewriter); len(tail) > 0 { + tailChunk := cliproxyexecutor.StreamChunk{Payload: tail} + if !emit(tailChunk) { + return + } + } + if !failed { + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}, auth, ephemeralResult) + } + }() + return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} +} + +func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult, allowRetry bool, ephemeralResult bool) (*cliproxyexecutor.StreamResult, error) { + if executor == nil { + return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} + } + ctx = contextWithRequestedModelAlias(ctx, opts, routeModel) + var lastErr error + didRefreshOnUnauthorized := false + for idx, execModel := range execModels { + resultModel := m.stateModelForExecution(auth, routeModel, execModel, pooled) + execReq := req + execReq.Model = execModel + if executionModel != "" { + execReq.Model = executionModel + } + execOpts := opts + execReq, execOpts = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } + streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) + if errStream != nil { + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } + if allowRetry { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, errStream, didRefreshOnUnauthorized); okRefresh { + auth = refreshed + didRefreshOnUnauthorized = true + streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) + if errStream != nil { + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } + } + } + } + } + if errStream == nil && (streamResult == nil || streamResult.Chunks == nil) { + errStream = &Error{Code: "empty_stream", Message: "upstream stream has no source", Retryable: true} + } + if errStream != nil { + rerr := resultErrorFromError(errStream) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result.RetryAfter = retryAfterFromError(errStream) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + if isRequestInvalidError(errStream) { + return nil, errStream + } + lastErr = errStream + continue + } + + buffered, closed, bootstrapErr := readStreamBootstrap(ctx, streamResult.Chunks) + if bootstrapErr != nil { + if errCtx := ctx.Err(); errCtx != nil { + discardStreamChunks(streamResult.Chunks) + return nil, errCtx + } + if allowRetry { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, bootstrapErr, didRefreshOnUnauthorized); okRefresh { + discardStreamChunks(streamResult.Chunks) + auth = refreshed + didRefreshOnUnauthorized = true + retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) + if retryErr != nil { + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } + bootstrapErr = retryErr + streamResult = &cliproxyexecutor.StreamResult{} + } else { + streamResult = retryStream + buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks) + } + } + } + } + if bootstrapErr != nil { + if isRequestInvalidError(bootstrapErr) { + rerr := resultErrorFromError(bootstrapErr) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result.RetryAfter = retryAfterFromError(bootstrapErr) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + discardStreamChunks(streamResult.Chunks) + return nil, bootstrapErr + } + if idx < len(execModels)-1 { + rerr := resultErrorFromError(bootstrapErr) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result.RetryAfter = retryAfterFromError(bootstrapErr) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + discardStreamChunks(streamResult.Chunks) + lastErr = bootstrapErr + continue + } + rerr := resultErrorFromError(bootstrapErr) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} + result.RetryAfter = retryAfterFromError(bootstrapErr) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + discardStreamChunks(streamResult.Chunks) + return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) + } + + if closed && len(buffered) == 0 { + emptyErr := &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: emptyErr} + m.recordExecutionResult(ctx, result, auth, ephemeralResult) + if idx < len(execModels)-1 { + lastErr = emptyErr + continue + } + return nil, newStreamBootstrapError(emptyErr, streamResult.Headers) + } + + remaining := streamResult.Chunks + if closed { + closedCh := make(chan cliproxyexecutor.StreamChunk) + close(closedCh) + remaining = closedCh + } + return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, aliasResult, ephemeralResult), nil + } + if lastErr == nil { + lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"} + } + return nil, lastErr +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 231ffb68e..0077344f5 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -5,38 +5,21 @@ package cliproxy import ( "context" - "errors" - "fmt" - "os" - "strings" "sync" - "sync/atomic" "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/api" - internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" - "github.com/router-for-me/CLIProxyAPI/v7/internal/constant" "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/homeplugins" - "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost" - "github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue" - "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" - "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" - "github.com/router-for-me/CLIProxyAPI/v7/internal/util" "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher" - "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/diff" - "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer" "github.com/router-for-me/CLIProxyAPI/v7/internal/wsrelay" sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" - "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" sdkpluginstore "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginstore" - sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" - log "github.com/sirupsen/logrus" ) // Service wraps the proxy server lifecycle so external programs can embed the CLI proxy. @@ -139,3522 +122,3 @@ type Service struct { homePluginSyncFetch func(context.Context, sdkpluginstore.PluginSyncRequest) (sdkpluginstore.PluginSyncResponse, error) homePluginDeleteTask func(context.Context, *config.Config, home.PluginTask) homeplugins.SyncReport } - -type homeSubscriberSupervisor struct { - cancel context.CancelFunc - done chan struct{} - - publisherMu sync.Mutex - publisherDone <-chan struct{} -} - -func (s *homeSubscriberSupervisor) setPublisherCompletion(done <-chan struct{}) { - if s == nil { - return - } - s.publisherMu.Lock() - s.publisherDone = done - s.publisherMu.Unlock() -} - -func (s *homeSubscriberSupervisor) publisherCompletion() <-chan struct{} { - if s == nil { - return nil - } - s.publisherMu.Lock() - defer s.publisherMu.Unlock() - return s.publisherDone -} - -type homeConfigWorkQueue struct { - mu sync.Mutex - items [][]byte - wake chan struct{} -} - -func newHomeConfigWorkQueue() *homeConfigWorkQueue { - return &homeConfigWorkQueue{wake: make(chan struct{}, 1)} -} - -func (q *homeConfigWorkQueue) enqueue(raw []byte) { - if q == nil { - return - } - item := append([]byte(nil), raw...) - q.mu.Lock() - q.items = append(q.items, item) - q.mu.Unlock() - select { - case q.wake <- struct{}{}: - default: - } -} - -func (q *homeConfigWorkQueue) dequeue(ctx context.Context) ([]byte, bool) { - if q == nil || ctx == nil { - return nil, false - } - for { - if ctx.Err() != nil { - return nil, false - } - q.mu.Lock() - if ctx.Err() != nil { - q.mu.Unlock() - return nil, false - } - if len(q.items) > 0 { - item := q.items[0] - q.items[0] = nil - q.items = q.items[1:] - q.mu.Unlock() - return item, true - } - q.mu.Unlock() - select { - case <-ctx.Done(): - return nil, false - case <-q.wake: - } - } -} - -type homeLogForwarder interface { - Bind(*home.Client) - Deactivate(*home.Client) - Stop() -} - -var startHomeLogForwarder = func(queueSize int) homeLogForwarder { - return logging.StartHomeAppLogForwarder(queueSize) -} - -const ( - modelRegistrationMaxWorkersPerCategory = 5 - modelRegistrationMaxWorkersOpenAICompatibility = 20 - homeSubscriberPreAckRetryBackoff = 100 * time.Millisecond -) - -const ( - modelRegistrationPhaseConfigAPIKey = iota - modelRegistrationPhaseOther -) - -type modelRegistrationTask struct { - phase int - category string - run func(*openAICompatibilityRegistrationCache) -} - -type executorRegistrationOptions struct { - includeBaseline bool - includePlugins bool - forceReplaceAuths bool - auths []*coreauth.Auth -} - -var registerPluginExecutors = func(host *pluginhost.Host, manager *coreauth.Manager) { - if host == nil || manager == nil { - return - } - host.RegisterExecutors(manager, registry.GetGlobalRegistry()) -} - -// RegisterUsagePlugin registers a usage plugin on the global usage manager. -// This allows external code to monitor API usage and token consumption. -// -// Parameters: -// - plugin: The usage plugin to register -func (s *Service) RegisterUsagePlugin(plugin usage.Plugin) { - usage.RegisterPlugin(plugin) -} - -func (s *Service) registerPluginAuthParser() { - var parser PluginAuthParser - if s != nil && s.pluginHost != nil { - parser = s.pluginHost - } - sdkAuth.RegisterPluginAuthParser(parser) - if s != nil && s.watcher != nil { - s.watcher.SetPluginAuthParser(parser) - } -} - -func (s *Service) syncPluginRuntime(ctx context.Context) { - if !s.syncPluginRuntimeConfig(ctx) { - return - } - s.syncPluginModelRuntime(ctx) -} - -func (s *Service) syncPluginRuntimeConfig(ctx context.Context) bool { - if s == nil { - sdkAuth.RegisterPluginAuthParser(nil) - return false - } - s.cfgMu.RLock() - cfg := s.cfg - s.cfgMu.RUnlock() - return s.syncPluginRuntimeConfigForConfig(ctx, cfg) -} - -func (s *Service) syncPluginRuntimeConfigForConfig(ctx context.Context, cfg *config.Config) bool { - if s == nil { - sdkAuth.RegisterPluginAuthParser(nil) - return false - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - - if s.pluginHost != nil { - s.pluginHost.ApplyConfig(ctx, cfg) - } - if errContext := ctx.Err(); errContext != nil { - return false - } - if s.coreManager != nil { - s.coreManager.SetPluginScheduler(s.pluginHost) - } - s.registerPluginAuthParser() - if s.pluginHost == nil { - return false - } - s.pluginHost.RegisterFrontendAuthProviders() - if errContext := ctx.Err(); errContext != nil { - return false - } - if s.accessManager != nil { - s.accessManager.SetProviders(sdkaccess.RegisteredProviders()) - } - s.pluginHost.RegisterUsagePlugins() - sdktranslator.SetPluginHooks(s.pluginHost) - if s.server != nil { - s.server.RefreshPluginManagementRoutes() - } - return ctx.Err() == nil -} - -func (s *Service) syncPluginModelRuntime(ctx context.Context) { - if s == nil || s.pluginHost == nil || s.coreManager == nil { - return - } - if ctx == nil { - ctx = context.Background() - } - s.pluginHost.RegisterModels(ctx, registry.GetGlobalRegistry()) - if ctx.Err() != nil { - return - } - s.cfgMu.RLock() - homeEnabled := s.cfg != nil && s.cfg.Home.Enabled - s.cfgMu.RUnlock() - s.registerAvailableExecutors(ctx, executorRegistrationOptions{ - includeBaseline: homeEnabled, - includePlugins: true, - forceReplaceAuths: false, - auths: s.coreManager.List(), - }) - s.refreshPluginModelRegistrations(ctx) - if ctx.Err() != nil { - return - } - s.coreManager.RefreshSchedulerAll() -} - -func (s *Service) refreshPluginModelRegistrations(ctx context.Context) { - if s == nil || s.pluginHost == nil || s.coreManager == nil { - return - } - s.registerModelsForAuthBatch(ctx, s.coreManager.List()) -} - -func (s *Service) registerModelsForAuthBatch(ctx context.Context, auths []*coreauth.Auth) { - if s == nil || s.coreManager == nil || len(auths) == 0 { - return - } - tasks := make([]modelRegistrationTask, 0, len(auths)) - for _, auth := range auths { - if auth == nil { - continue - } - authForRegistration := auth.Clone() - tasks = append(tasks, modelRegistrationTask{ - phase: modelRegistrationPhase(authForRegistration), - category: modelRegistrationCategory(authForRegistration), - run: func(compatCache *openAICompatibilityRegistrationCache) { - s.completeModelRegistrationForAuthWithCache(ctx, authForRegistration, compatCache) - }, - }) - } - s.runModelRegistrationTasks(ctx, tasks) -} - -func (s *Service) runModelRegistrationTasks(ctx context.Context, tasks []modelRegistrationTask) { - if len(tasks) == 0 { - return - } - if ctx == nil { - ctx = context.Background() - } - - configAPIKeyTasks := make([]modelRegistrationTask, 0) - otherTasks := make([]modelRegistrationTask, 0) - for _, task := range tasks { - if task.phase == modelRegistrationPhaseConfigAPIKey { - configAPIKeyTasks = append(configAPIKeyTasks, task) - continue - } - otherTasks = append(otherTasks, task) - } - - compatCache := s.newOpenAICompatibilityRegistrationCache() - s.runModelRegistrationTaskPhase(ctx, configAPIKeyTasks, compatCache) - s.runModelRegistrationTaskPhase(ctx, otherTasks, compatCache) -} - -func (s *Service) runModelRegistrationTaskPhase(ctx context.Context, tasks []modelRegistrationTask, compatCache *openAICompatibilityRegistrationCache) { - if len(tasks) == 0 { - return - } - - grouped := make(map[string][]modelRegistrationTask) - order := make([]string, 0) - for _, task := range tasks { - if task.run == nil { - continue - } - category := strings.ToLower(strings.TrimSpace(task.category)) - if category == "" { - category = "unknown" - } - if _, exists := grouped[category]; !exists { - order = append(order, category) - } - grouped[category] = append(grouped[category], task) - } - - var wg sync.WaitGroup - for _, category := range order { - group := grouped[category] - workers := len(group) - maxWorkers := modelRegistrationMaxWorkersForCategory(category) - if workers > maxWorkers { - workers = maxWorkers - } - if workers <= 0 { - continue - } - - taskCh := make(chan modelRegistrationTask) - for i := 0; i < workers; i++ { - wg.Add(1) - go func() { - defer wg.Done() - for task := range taskCh { - select { - case <-ctx.Done(): - return - default: - } - task.run(compatCache) - } - }() - } - go func(group []modelRegistrationTask) { - defer close(taskCh) - for _, task := range group { - select { - case <-ctx.Done(): - return - case taskCh <- task: - } - } - }(group) - } - wg.Wait() -} - -func modelRegistrationPhase(auth *coreauth.Auth) int { - if coreauth.IsConfigAPIKeyAuth(auth) { - return modelRegistrationPhaseConfigAPIKey - } - return modelRegistrationPhaseOther -} - -func modelRegistrationCategory(auth *coreauth.Auth) string { - if auth == nil { - return "unknown" - } - provider := strings.ToLower(strings.TrimSpace(auth.Provider)) - if compatProviderKey, _, compatDetected := openAICompatInfoFromAuth(auth); compatDetected { - if compatProviderKey != "" { - provider = compatProviderKey - } else { - provider = "openai-compatibility" - } - } - if provider == "" { - provider = "unknown" - } - - authKind := auth.AuthKind() - if authKind == "" { - return provider - } - return provider + ":" + authKind -} - -func modelRegistrationMaxWorkersForCategory(category string) int { - category = strings.ToLower(strings.TrimSpace(category)) - if strings.HasPrefix(category, "openai-compatible-") || strings.HasPrefix(category, "openai-compatibility") { - return modelRegistrationMaxWorkersOpenAICompatibility - } - return modelRegistrationMaxWorkersPerCategory -} - -func (s *Service) registerModelRefreshCallback() { - // Register callback for startup and periodic model catalog refresh. - // When remote model definitions change, re-register models for affected providers. - // This intentionally rebuilds per-auth model availability from the latest catalog - // snapshot instead of preserving prior registry suppression state. - registry.SetModelRefreshCallback(func(changedProviders []string) { - if s == nil || s.coreManager == nil || len(changedProviders) == 0 { - return - } - - providerSet := make(map[string]bool, len(changedProviders)) - for _, p := range changedProviders { - providerSet[strings.ToLower(strings.TrimSpace(p))] = true - } - - auths := s.coreManager.List() - refreshed := 0 - var refreshedMu sync.Mutex - tasks := make([]modelRegistrationTask, 0, len(auths)) - for _, item := range auths { - if item == nil || item.ID == "" { - continue - } - auth, ok := s.coreManager.GetByID(item.ID) - if !ok || auth == nil || auth.Disabled { - continue - } - provider := strings.ToLower(strings.TrimSpace(auth.Provider)) - if !providerSet[provider] { - continue - } - authForRefresh := auth - tasks = append(tasks, modelRegistrationTask{ - phase: modelRegistrationPhase(authForRefresh), - category: modelRegistrationCategory(authForRefresh), - run: func(compatCache *openAICompatibilityRegistrationCache) { - if s.refreshModelRegistrationForAuthWithCache(authForRefresh, compatCache) { - refreshedMu.Lock() - refreshed++ - refreshedMu.Unlock() - } - }, - }) - } - s.runModelRegistrationTasks(context.Background(), tasks) - - if refreshed > 0 { - log.Infof("re-registered models for %d auth(s) due to model catalog changes: %v", refreshed, changedProviders) - } - }) -} - -// newDefaultAuthManager creates a default authentication manager with supported OAuth providers. -func newDefaultAuthManager() *sdkAuth.Manager { - return sdkAuth.NewManager( - sdkAuth.GetTokenStore(), - sdkAuth.NewCodexAuthenticator(), - sdkAuth.NewClaudeAuthenticator(), - sdkAuth.NewXAIAuthenticator(), - ) -} - -func (s *Service) ensureAuthUpdateQueue(ctx context.Context) { - if s == nil { - return - } - if s.authUpdates == nil { - s.authUpdates = make(chan watcher.AuthUpdate, 256) - } - if s.authQueueStop != nil { - return - } - queueCtx, cancel := context.WithCancel(ctx) - s.authQueueStop = cancel - go s.consumeAuthUpdates(queueCtx) -} - -func (s *Service) consumeAuthUpdates(ctx context.Context) { - ctx = coreauth.WithSkipPersist(ctx) - for { - select { - case <-ctx.Done(): - return - case update, ok := <-s.authUpdates: - if !ok { - return - } - updates := []watcher.AuthUpdate{update} - labelDrain: - for { - select { - case nextUpdate := <-s.authUpdates: - updates = append(updates, nextUpdate) - default: - break labelDrain - } - } - s.handleAuthUpdates(ctx, updates) - } - } -} - -func (s *Service) emitAuthUpdate(ctx context.Context, update watcher.AuthUpdate) { - if s == nil { - return - } - if ctx == nil { - ctx = context.Background() - } - if s.watcher != nil && s.watcher.DispatchRuntimeAuthUpdate(update) { - return - } - if s.authUpdates != nil { - select { - case s.authUpdates <- update: - return - default: - log.Debugf("auth update queue saturated, applying inline action=%v id=%s", update.Action, update.ID) - } - } - s.handleAuthUpdate(ctx, update) -} - -func (s *Service) handleAuthUpdate(ctx context.Context, update watcher.AuthUpdate) { - s.handleAuthUpdates(ctx, []watcher.AuthUpdate{update}) -} - -func (s *Service) handleAuthUpdates(ctx context.Context, updates []watcher.AuthUpdate) { - if s == nil { - return - } - updates = coalesceAuthUpdates(updates) - s.cfgMu.RLock() - cfg := s.cfg - s.cfgMu.RUnlock() - if cfg == nil || s.coreManager == nil { - return - } - - registrationCtx := coreauth.WithDeferredAPIKeyModelAliasRebuild(ctx) - tasks := make([]modelRegistrationTask, 0, len(updates)) - needsPluginSync := false - needsAliasRebuild := false - for _, update := range updates { - switch update.Action { - case watcher.AuthUpdateActionAdd, watcher.AuthUpdateActionModify: - if update.Auth == nil || update.Auth.ID == "" { - continue - } - auth := s.prepareCoreAuthForModelRegistration(registrationCtx, update.Auth) - if auth == nil { - continue - } - needsAliasRebuild = true - authForRegistration := auth - tasks = append(tasks, modelRegistrationTask{ - phase: modelRegistrationPhase(authForRegistration), - category: modelRegistrationCategory(authForRegistration), - run: func(compatCache *openAICompatibilityRegistrationCache) { - s.completeModelRegistrationForAuthWithCache(registrationCtx, authForRegistration, compatCache) - }, - }) - needsPluginSync = true - case watcher.AuthUpdateActionDelete: - id := update.ID - if id == "" && update.Auth != nil { - id = update.Auth.ID - } - if id == "" { - continue - } - s.applyCoreAuthRemoval(registrationCtx, id) - needsAliasRebuild = true - default: - log.Debugf("received unknown auth update action: %v", update.Action) - } - } - - if needsAliasRebuild { - s.coreManager.RefreshAPIKeyModelAlias() - } - s.runModelRegistrationTasks(registrationCtx, tasks) - if needsPluginSync { - s.syncPluginRuntime(registrationCtx) - } -} - -func coalesceAuthUpdates(updates []watcher.AuthUpdate) []watcher.AuthUpdate { - if len(updates) <= 1 { - return updates - } - order := make([]string, 0, len(updates)) - byID := make(map[string]watcher.AuthUpdate, len(updates)) - unkeyed := make([]watcher.AuthUpdate, 0) - for _, update := range updates { - id := authUpdateID(update) - if id == "" { - unkeyed = append(unkeyed, update) - continue - } - if _, exists := byID[id]; !exists { - order = append(order, id) - } - byID[id] = update - } - if len(byID) == 0 { - return unkeyed - } - out := make([]watcher.AuthUpdate, 0, len(byID)+len(unkeyed)) - for _, id := range order { - out = append(out, byID[id]) - } - out = append(out, unkeyed...) - return out -} - -func authUpdateID(update watcher.AuthUpdate) string { - if strings.TrimSpace(update.ID) != "" { - return strings.TrimSpace(update.ID) - } - if update.Auth != nil { - return strings.TrimSpace(update.Auth.ID) - } - return "" -} - -func (s *Service) ensureWebsocketGateway() { - if s == nil { - return - } - if s.wsGateway != nil { - return - } - opts := wsrelay.Options{ - Path: "/v1/ws", - OnConnected: s.wsOnConnected, - OnDisconnected: s.wsOnDisconnected, - LogDebugf: log.Debugf, - LogInfof: log.Infof, - LogWarnf: log.Warnf, - } - s.wsGateway = wsrelay.NewManager(opts) -} - -func (s *Service) wsOnConnected(channelID string) { - if s == nil || channelID == "" { - return - } - if !strings.HasPrefix(strings.ToLower(channelID), "aistudio-") { - return - } - if s.coreManager != nil { - if existing, ok := s.coreManager.GetByID(channelID); ok && existing != nil { - if !existing.Disabled && existing.Status == coreauth.StatusActive { - return - } - } - } - now := time.Now().UTC() - auth := &coreauth.Auth{ - ID: channelID, // keep channel identifier as ID - Provider: "aistudio", // logical provider for switch routing - Label: channelID, // display original channel id - Status: coreauth.StatusActive, - CreatedAt: now, - UpdatedAt: now, - Attributes: map[string]string{"runtime_only": "true"}, - Metadata: map[string]any{"email": channelID}, // metadata drives logging and usage tracking - } - log.Infof("websocket provider connected: %s", channelID) - s.emitAuthUpdate(context.Background(), watcher.AuthUpdate{ - Action: watcher.AuthUpdateActionAdd, - ID: auth.ID, - Auth: auth, - }) -} - -func (s *Service) wsOnDisconnected(channelID string, reason error) { - if s == nil || channelID == "" { - return - } - if reason != nil { - if strings.Contains(reason.Error(), "replaced by new connection") { - log.Infof("websocket provider replaced: %s", channelID) - return - } - log.Warnf("websocket provider disconnected: %s (%v)", channelID, reason) - } else { - log.Infof("websocket provider disconnected: %s", channelID) - } - ctx := context.Background() - s.emitAuthUpdate(ctx, watcher.AuthUpdate{ - Action: watcher.AuthUpdateActionDelete, - ID: channelID, - }) -} - -func (s *Service) applyCoreAuthAddOrUpdate(ctx context.Context, auth *coreauth.Auth) { - auth = s.prepareCoreAuthForModelRegistration(ctx, auth) - if auth == nil { - return - } - s.completeModelRegistrationForAuth(ctx, auth) - s.syncPluginRuntime(ctx) -} - -func (s *Service) prepareCoreAuthForModelRegistration(ctx context.Context, auth *coreauth.Auth) *coreauth.Auth { - if s == nil || s.coreManager == nil || auth == nil || auth.ID == "" { - return nil - } - auth = auth.Clone() - s.ensureExecutorsForAuthWithContext(ctx, auth, false) - - // IMPORTANT: Update coreManager FIRST, before model registration. - // This ensures that configuration changes (proxy_url, prefix, etc.) take effect - // immediately for API calls, rather than waiting for model registration to complete. - op := "register" - var err error - if existing, ok := s.coreManager.GetByID(auth.ID); ok { - auth.CreatedAt = existing.CreatedAt - if !existing.Disabled && existing.Status != coreauth.StatusDisabled && !auth.Disabled && auth.Status != coreauth.StatusDisabled { - auth.LastRefreshedAt = existing.LastRefreshedAt - auth.NextRefreshAfter = existing.NextRefreshAfter - if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { - auth.ModelStates = existing.ModelStates - } - } - op = "update" - _, err = s.coreManager.Update(ctx, auth) - } else { - _, err = s.coreManager.Register(ctx, auth) - } - if err != nil { - log.Errorf("failed to %s auth %s: %v", op, auth.ID, err) - current, ok := s.coreManager.GetByID(auth.ID) - if !ok || current.Disabled { - GlobalModelRegistry().UnregisterClient(auth.ID) - return nil - } - auth = current - } - return auth -} - -func (s *Service) completeModelRegistrationForAuth(ctx context.Context, auth *coreauth.Auth) { - s.completeModelRegistrationForAuthWithCache(ctx, auth, nil) -} - -func (s *Service) completeModelRegistrationForAuthWithCache(ctx context.Context, auth *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) { - if s == nil || s.coreManager == nil || auth == nil || auth.ID == "" { - return - } - if ctx != nil && ctx.Err() != nil { - return - } - s.registerModelsForAuthWithCache(ctx, auth, compatCache) - if ctx != nil && ctx.Err() != nil { - return - } - s.coreManager.ReconcileRegistryModelStates(ctx, auth.ID) - - // Refresh the scheduler entry so that the auth's supportedModelSet is rebuilt - // from the now-populated global model registry. Without this, newly added auths - // have an empty supportedModelSet (because Register/Update upserts into the - // scheduler before registerModelsForAuth runs) and are invisible to the scheduler. - s.coreManager.RefreshSchedulerEntry(auth.ID) -} - -func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) { - if s == nil || id == "" { - return - } - if s.coreManager == nil { - return - } - id = strings.TrimSpace(id) - var provider string - if existing, ok := s.coreManager.GetByID(id); ok && existing != nil { - provider = strings.TrimSpace(existing.Provider) - } - GlobalModelRegistry().UnregisterClient(id) - s.coreManager.Remove(ctx, id) - if strings.EqualFold(provider, "codex") { - executor.CloseCodexWebsocketSessionsForAuthID(id, "auth_removed") - } - if strings.EqualFold(provider, "xai") { - executor.CloseXAIWebsocketSessionsForAuthID(id, "auth_removed") - } - s.syncPluginRuntime(ctx) -} - -func (s *Service) applyRetryConfig(cfg *config.Config) { - if s == nil || s.coreManager == nil || cfg == nil { - return - } - maxInterval := time.Duration(cfg.MaxRetryInterval) * time.Second - s.coreManager.SetRetryConfig(cfg.RequestRetry, maxInterval, cfg.MaxRetryCredentials) - coreauth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) -} - -func (s *Service) configureCooldownStateStore(cfg *config.Config) { - _ = s.configureCooldownStateStoreContext(context.Background(), cfg, false) -} - -func (s *Service) configureCooldownStateStoreContext(ctx context.Context, cfg *config.Config, persistOld bool) bool { - if s == nil || s.coreManager == nil { - return true - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - return s.coreManager.SwapCooldownStateStore(ctx, s.resolveCooldownStateStore(cfg), persistOld) -} - -func (s *Service) resolveCooldownStateStore(cfg *config.Config) coreauth.CooldownStateStore { - if cfg == nil || !cfg.SaveCooldownStatus || cfg.Home.Enabled { - return nil - } - authDir, errResolve := resolveCooldownStateAuthDir(cfg) - if errResolve != nil { - log.Warnf("failed to resolve cooldown state directory: %v", errResolve) - return nil - } - if authDir == "" { - return nil - } - return coreauth.NewFileCooldownStateStoreWithAuthDir(authDir, authDir) -} - -func resolveCooldownStateAuthDir(cfg *config.Config) (string, error) { - if cfg == nil { - return "", nil - } - authDir, errAuthDir := util.ResolveAuthDir(cfg.AuthDir) - if errAuthDir != nil { - return "", errAuthDir - } - return authDir, nil -} - -func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName string, ok bool) { - if a == nil { - return "", "", false - } - if len(a.Attributes) > 0 { - providerKey = strings.TrimSpace(a.Attributes["provider_key"]) - compatName = strings.TrimSpace(a.Attributes["compat_name"]) - if compatName != "" { - if providerKey == "" { - providerKey = compatName - } - return util.OpenAICompatibleProviderKey(providerKey), compatName, true - } - } - if strings.EqualFold(strings.TrimSpace(a.Provider), "openai-compatibility") { - compatName = strings.TrimSpace(a.Label) - providerKey = compatName - if providerKey == "" { - providerKey = "openai-compatibility" - } - return util.OpenAICompatibleProviderKey(providerKey), compatName, true - } - return "", "", false -} - -type openAICompatibilityRegistrationCache struct { - byName map[string]*openAICompatibilityRegistrationEntry -} - -type openAICompatibilityRegistrationEntry struct { - providerKey string - models []*ModelInfo -} - -func (s *Service) newOpenAICompatibilityRegistrationCache() *openAICompatibilityRegistrationCache { - if s == nil { - return nil - } - s.cfgMu.RLock() - cfg := s.cfg - s.cfgMu.RUnlock() - if cfg == nil || len(cfg.OpenAICompatibility) == 0 { - return nil - } - - cache := &openAICompatibilityRegistrationCache{ - byName: make(map[string]*openAICompatibilityRegistrationEntry, len(cfg.OpenAICompatibility)), - } - for i := range cfg.OpenAICompatibility { - compat := &cfg.OpenAICompatibility[i] - if compat.Disabled { - continue - } - compatName := strings.TrimSpace(compat.Name) - key := strings.ToLower(compatName) - if _, exists := cache.byName[key]; exists { - continue - } - providerName := strings.ToLower(compatName) - if providerName == "" { - providerName = "openai-compatibility" - } - cache.byName[key] = &openAICompatibilityRegistrationEntry{ - providerKey: util.OpenAICompatibleProviderKey(providerName), - models: buildOpenAICompatibilityConfigModels(compat), - } - } - if len(cache.byName) == 0 { - return nil - } - return cache -} - -func (c *openAICompatibilityRegistrationCache) lookup(compatName string) (*openAICompatibilityRegistrationEntry, bool) { - if c == nil || len(c.byName) == 0 { - return nil, false - } - entry, ok := c.byName[strings.ToLower(strings.TrimSpace(compatName))] - return entry, ok -} - -func (s *Service) hasNativeOpenAICompatExecutorConfig(a *coreauth.Auth, providerKey string, cfg *config.Config) bool { - if a == nil { - return false - } - providerKey = strings.ToLower(strings.TrimSpace(providerKey)) - if a.Attributes != nil { - if strings.TrimSpace(a.Attributes["base_url"]) != "" { - return true - } - if strings.TrimSpace(a.Attributes["compat_name"]) != "" { - return true - } - } - if strings.EqualFold(strings.TrimSpace(a.Provider), "openai-compatibility") { - return true - } - if s == nil || cfg == nil { - return false - } - - candidates := make([]string, 0, 3) - if providerKey != "" { - candidates = append(candidates, providerKey) - } - if a.Attributes != nil { - if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" { - candidates = append(candidates, strings.ToLower(v)) - } - } - if provider := strings.TrimSpace(a.Provider); provider != "" { - candidates = append(candidates, strings.ToLower(provider)) - } - - for i := range cfg.OpenAICompatibility { - compat := &cfg.OpenAICompatibility[i] - if compat.Disabled { - continue - } - name := strings.ToLower(strings.TrimSpace(compat.Name)) - if name == "" { - continue - } - for _, candidate := range candidates { - if candidate != "" && candidate == name { - return true - } - } - } - return false -} - -func (s *Service) unregisterOpenAICompatExecutor(providerKey string) { - if s == nil || s.coreManager == nil { - return - } - providerKey = strings.ToLower(strings.TrimSpace(providerKey)) - if providerKey == "" { - return - } - existing, okExecutor := s.coreManager.Executor(providerKey) - if !okExecutor || existing == nil { - return - } - if _, okOpenAICompat := existing.(*executor.OpenAICompatExecutor); !okOpenAICompat { - return - } - s.coreManager.UnregisterExecutor(providerKey) -} - -func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) { - s.ensureExecutorsForAuthWithContext(context.Background(), a, false) -} - -func (s *Service) ensureExecutorsForAuthWithMode(a *coreauth.Auth, forceReplace bool) { - s.ensureExecutorsForAuthWithContext(context.Background(), a, forceReplace) -} - -func (s *Service) ensureExecutorsForAuthWithContext(ctx context.Context, a *coreauth.Auth, forceReplace bool) { - if a == nil || (ctx != nil && ctx.Err() != nil) { - return - } - s.registerAvailableExecutors(ctx, executorRegistrationOptions{ - auths: []*coreauth.Auth{a}, - forceReplaceAuths: forceReplace, - }) -} - -func (s *Service) registerAvailableExecutors(ctx context.Context, opts executorRegistrationOptions) { - if s == nil || s.coreManager == nil { - return - } - if ctx == nil { - ctx = context.Background() - } - s.executorRegistrationMu.Lock() - defer s.executorRegistrationMu.Unlock() - if ctx.Err() != nil { - return - } - // Keep all Service-owned executor registration paths here so native, Home, - // auth-derived, and plugin executors stay in the same binding order. - if opts.includeBaseline { - s.registerExecutorsForAuths(baselineExecutorAuths(), opts.forceReplaceAuths) - } - if len(opts.auths) > 0 { - s.registerExecutorsForAuths(opts.auths, opts.forceReplaceAuths) - } - if opts.includePlugins && s.pluginHost != nil { - registerPluginExecutors(s.pluginHost, s.coreManager) - } -} - -func baselineExecutorAuths() []*coreauth.Auth { - providers := []string{ - "codex", - "claude", - constant.Gemini, - constant.GeminiInteractions, - "vertex", - "aistudio", - "antigravity", - "kimi", - "xai", - "openai-compatibility", - } - auths := make([]*coreauth.Auth, 0, len(providers)) - for _, provider := range providers { - auth := &coreauth.Auth{ - ID: provider, - Provider: provider, - } - if provider == "openai-compatibility" { - auth.Attributes = map[string]string{"compat_name": "openai-compatibility"} - } - auths = append(auths, auth) - } - return auths -} - -func (s *Service) registerExecutorsForAuths(auths []*coreauth.Auth, forceReplace bool) { - reboundCodex := false - for _, auth := range auths { - if auth != nil && strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { - if reboundCodex && forceReplace { - continue - } - reboundCodex = true - } - s.registerExecutorForAuth(auth, forceReplace) - } -} - -func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { - if s == nil || s.coreManager == nil || a == nil { - return - } - s.cfgMu.RLock() - cfg := s.cfg - s.cfgMu.RUnlock() - if strings.EqualFold(strings.TrimSpace(a.Provider), "codex") { - if !forceReplace { - existingExecutor, hasExecutor := s.coreManager.Executor("codex") - if hasExecutor { - _, isCodexAutoExecutor := existingExecutor.(*executor.CodexAutoExecutor) - if isCodexAutoExecutor { - return - } - } - } - s.coreManager.RegisterExecutor(executor.NewCodexAutoExecutor(cfg)) - return - } - // Skip disabled auth entries when (re)binding executors. - // Disabled auths can linger during config reloads (e.g., removed OpenAI-compat entries) - // and must not override active provider executors. - if a.Disabled { - return - } - if compatProviderKey, _, isCompat := openAICompatInfoFromAuth(a); isCompat { - if compatProviderKey == "" { - compatProviderKey = strings.ToLower(strings.TrimSpace(a.Provider)) - } - if compatProviderKey == "" { - compatProviderKey = "openai-compatibility" - } - if !forceReplace { - if existingExecutor, hasExecutor := s.coreManager.Executor(compatProviderKey); hasExecutor { - if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { - return - } - } - } - s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(compatProviderKey, cfg)) - return - } - switch strings.ToLower(a.Provider) { - case constant.Gemini: - s.coreManager.RegisterExecutor(executor.NewGeminiExecutor(cfg)) - case constant.GeminiInteractions: - s.coreManager.RegisterExecutor(executor.NewGeminiInteractionsExecutor(cfg)) - case "vertex": - s.coreManager.RegisterExecutor(executor.NewGeminiVertexExecutor(cfg)) - case "aistudio": - if s.wsGateway != nil { - s.coreManager.RegisterExecutor(executor.NewAIStudioExecutor(cfg, a.ID, s.wsGateway)) - } - return - case "antigravity": - s.coreManager.RegisterExecutor(executor.NewAntigravityExecutor(cfg)) - case "claude": - s.coreManager.RegisterExecutor(executor.NewClaudeExecutor(cfg)) - case "kimi": - s.coreManager.RegisterExecutor(executor.NewKimiExecutor(cfg)) - case "xai": - if !forceReplace { - existingExecutor, hasExecutor := s.coreManager.Executor("xai") - if hasExecutor { - existingXAIAutoExecutor, isXAIAutoExecutor := existingExecutor.(*executor.XAIAutoExecutor) - if isXAIAutoExecutor && existingXAIAutoExecutor.UsesConfig(cfg) { - return - } - } - } - s.coreManager.RegisterExecutor(executor.NewXAIAutoExecutor(cfg)) - default: - providerKey := strings.ToLower(strings.TrimSpace(a.Provider)) - if providerKey == "" { - providerKey = "openai-compatibility" - } - if s.pluginHost != nil && - s.pluginHost.HasExecutorCandidateProvider(providerKey) && - !s.hasNativeOpenAICompatExecutorConfig(a, providerKey, cfg) { - s.unregisterOpenAICompatExecutor(providerKey) - return - } - if !forceReplace { - if existingExecutor, hasExecutor := s.coreManager.Executor(providerKey); hasExecutor { - if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { - return - } - } - } - s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(providerKey, cfg)) - } -} - -func (s *Service) registerResolvedModelsForAuth(a *coreauth.Auth, providerKey string, models []*ModelInfo) { - if a == nil || a.ID == "" { - return - } - providerKey = strings.ToLower(strings.TrimSpace(providerKey)) - if providerKey == "" { - GlobalModelRegistry().UnregisterClient(a.ID) - return - } - normalizedModels := make([]*ModelInfo, 0, len(models)) - for _, model := range models { - if model == nil { - continue - } - modelID := strings.TrimSpace(model.ID) - if modelID == "" { - continue - } - clone := *model - clone.ID = modelID - normalizedModels = append(normalizedModels, &clone) - } - if len(normalizedModels) == 0 { - GlobalModelRegistry().UnregisterClient(a.ID) - return - } - GlobalModelRegistry().RegisterClient(a.ID, providerKey, normalizedModels) -} - -func (s *Service) pluginModelsForProvider(providerKey string) []*ModelInfo { - if s == nil || s.pluginHost == nil { - return nil - } - return s.pluginHost.ModelsForProvider(providerKey) -} - -func (s *Service) appendPluginModels(providerKey string, models []*ModelInfo) []*ModelInfo { - pluginModels := s.pluginModelsForProvider(providerKey) - if len(pluginModels) == 0 { - return models - } - out := make([]*ModelInfo, 0, len(models)+len(pluginModels)) - seen := make(map[string]struct{}, len(models)+len(pluginModels)) - for _, model := range models { - if model == nil { - continue - } - modelID := strings.TrimSpace(model.ID) - if modelID != "" { - seen[modelID] = struct{}{} - } - out = append(out, model) - } - for _, model := range pluginModels { - if model == nil { - continue - } - modelID := strings.TrimSpace(model.ID) - if modelID == "" { - continue - } - if _, exists := seen[modelID]; exists { - continue - } - seen[modelID] = struct{}{} - out = append(out, model) - } - return out -} - -func (s *Service) tryRegisterPluginModelsForAuth(ctx context.Context, a *coreauth.Auth, provider, authKind string, excluded []string) bool { - if s == nil || s.pluginHost == nil || a == nil { - return false - } - if ctx != nil && ctx.Err() != nil { - return true - } - result := s.pluginHost.ModelsForAuth(ctx, a) - if ctx != nil && ctx.Err() != nil { - return true - } - if !result.Handled { - return false - } - if result.Err != nil { - return true - } - activeAuth := a - providerKey := strings.ToLower(strings.TrimSpace(result.Provider)) - if providerKey == "" { - providerKey = strings.ToLower(strings.TrimSpace(provider)) - } - if result.Auth != nil && s.coreManager != nil { - result.Auth.ID = a.ID - if result.Auth.Provider == "" { - result.Auth.Provider = a.Provider - } - if result.Auth.FileName == "" { - result.Auth.FileName = a.FileName - } - if result.Auth.Attributes == nil { - result.Auth.Attributes = make(map[string]string) - } - for key, value := range a.Attributes { - if _, exists := result.Auth.Attributes[key]; !exists { - result.Auth.Attributes[key] = value - } - } - if updated, errUpdate := s.coreManager.Update(ctx, result.Auth); errUpdate == nil && updated != nil { - activeAuth = updated.Clone() - } - } - if activeAuth == nil { - activeAuth = a - } - if activeProvider := strings.ToLower(strings.TrimSpace(activeAuth.Provider)); activeProvider != "" { - providerKey = activeProvider - } - if providerKey == "" { - providerKey = strings.ToLower(strings.TrimSpace(provider)) - } - activeAuthKind := activeAuth.AuthKind() - activeExcluded := s.oauthExcludedModels(providerKey, activeAuthKind) - if a == activeAuth && len(activeExcluded) == 0 { - activeExcluded = excluded - } - if activeAuth.Attributes != nil { - if val, ok := activeAuth.Attributes["excluded_models"]; ok && strings.TrimSpace(val) != "" { - activeExcluded = strings.Split(val, ",") - } - } - if ctx != nil && ctx.Err() != nil { - return true - } - models := applyExcludedModels(result.Models, activeExcluded) - models = applyOAuthModelAliasForAuth(s.cfg, providerKey, activeAuthKind, activeAuth.Attributes, models) - if len(models) > 0 { - s.registerResolvedModelsForAuth(activeAuth, providerKey, applyModelPrefixes(models, activeAuth.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix)) - return true - } - GlobalModelRegistry().UnregisterClient(activeAuth.ID) - return true -} - -func (s *Service) applyConfigUpdate(newCfg *config.Config) { - s.applyConfigUpdateWithAuthSynthesis(context.Background(), newCfg, true) -} - -func (s *Service) applyWatcherConfigUpdate(newCfg *config.Config) { - s.applyConfigUpdateWithAuthSynthesis(context.Background(), newCfg, false) -} - -type configCommit struct { - cfg *config.Config - sequence uint64 -} - -type routingRuntimeState struct { - strategy string - sessionAffinity bool - sessionAffinityTTL time.Duration -} - -func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { - state := routingRuntimeState{ - strategy: "round-robin", - sessionAffinityTTL: time.Hour, - } - if cfg == nil { - return state - } - - switch strings.ToLower(strings.TrimSpace(cfg.Routing.Strategy)) { - case "fill-first", "fillfirst", "ff": - state.strategy = "fill-first" - } - state.sessionAffinity = cfg.Routing.SessionAffinity - if ttl := strings.TrimSpace(cfg.Routing.SessionAffinityTTL); ttl != "" { - if parsed, errParse := time.ParseDuration(ttl); errParse == nil && parsed > 0 { - state.sessionAffinityTTL = parsed - } - } - return state -} - -func newRoutingSelector(state routingRuntimeState) coreauth.Selector { - var selector coreauth.Selector - if state.strategy == "fill-first" { - selector = &coreauth.FillFirstSelector{} - } else { - selector = &coreauth.RoundRobinSelector{} - } - if state.sessionAffinity { - selector = coreauth.NewSessionAffinitySelectorWithConfig(coreauth.SessionAffinityConfig{ - Fallback: selector, - TTL: state.sessionAffinityTTL, - }) - } - return selector -} - -func (s *Service) applyConfigUpdateWithAuthSynthesis(ctx context.Context, newCfg *config.Config, synthesizeConfigAuths bool) bool { - commit := s.commitConfigUpdate(newCfg) - if commit.cfg == nil { - return false - } - return s.applyConfigRuntime(ctx, commit, synthesizeConfigAuths) -} - -// commitConfigUpdate applies only in-memory configuration state. Runtime work that -// may block on plugins, models, storage, or networking is deliberately deferred. -func (s *Service) commitConfigUpdate(newCfg *config.Config) configCommit { - if s == nil { - return configCommit{} - } - - s.configUpdateMu.Lock() - defer s.configUpdateMu.Unlock() - - if newCfg == nil { - s.cfgMu.RLock() - newCfg = s.cfg - s.cfgMu.RUnlock() - } - if newCfg == nil { - return configCommit{} - } - - s.cfgMu.Lock() - s.cfg = newCfg - s.cfgMu.Unlock() - s.configSequence++ - return configCommit{cfg: newCfg, sequence: s.configSequence} -} - -func (s *Service) configCommitCurrent(commit configCommit) bool { - if s == nil || commit.sequence == 0 { - return false - } - s.configUpdateMu.Lock() - current := s.configSequence == commit.sequence - s.configUpdateMu.Unlock() - return current -} - -func (s *Service) applyConfigRuntime(ctx context.Context, commit configCommit, synthesizeConfigAuths bool) bool { - cfg := commit.cfg - if s == nil || cfg == nil { - return false - } - s.configRuntimeMu.Lock() - defer s.configRuntimeMu.Unlock() - if !s.configCommitCurrent(commit) { - return false - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - - if !s.applyManagerConfig(ctx, commit) { - return false - } - if errContext := ctx.Err(); errContext != nil { - return false - } - if !s.applyPprofConfigContext(ctx, cfg) { - return false - } - if errContext := ctx.Err(); errContext != nil { - return false - } - if !s.updateServerClientsContext(ctx, cfg) { - return false - } - if errContext := ctx.Err(); errContext != nil { - return false - } - - registrationCtx := coreauth.WithSkipPersist(ctx) - s.syncPluginRuntimeConfigForConfig(registrationCtx, cfg) - if errContext := ctx.Err(); errContext != nil { - return false - } - var auths []*coreauth.Auth - if s.coreManager != nil { - auths = s.coreManager.List() - } - s.registerAvailableExecutors(registrationCtx, executorRegistrationOptions{ - includeBaseline: cfg.Home.Enabled, - forceReplaceAuths: true, - auths: auths, - }) - if errContext := ctx.Err(); errContext != nil { - return false - } - if synthesizeConfigAuths { - s.registerConfigAPIKeyAuths(registrationCtx, cfg) - } - if errContext := ctx.Err(); errContext != nil { - return false - } - if s.coreManager != nil && !cfg.Home.Enabled && cfg.SaveCooldownStatus { - if errRestoreCooldown := s.coreManager.RestoreCooldownStates(registrationCtx); errRestoreCooldown != nil && ctx.Err() == nil { - log.Warnf("failed to restore cooldown state after config update: %v", errRestoreCooldown) - } - } - if errContext := ctx.Err(); errContext != nil { - return false - } - s.syncPluginModelRuntime(registrationCtx) - return ctx.Err() == nil -} - -func (s *Service) applyManagerConfig(ctx context.Context, commit configCommit) bool { - if s == nil || s.coreManager == nil || commit.cfg == nil { - return s != nil && commit.cfg != nil - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return false - } - routingState := normalizedRoutingRuntimeState(commit.cfg) - if s.appliedRoutingState == nil || *s.appliedRoutingState != routingState { - s.coreManager.SetSelector(newRoutingSelector(routingState)) - s.appliedRoutingState = &routingState - } - s.applyRetryConfig(commit.cfg) - store := s.resolveCooldownStateStore(commit.cfg) - if !s.coreManager.ApplyConfigWithCooldownStateStore(ctx, commit.cfg, store) { - return false - } - s.coreManager.SetOAuthModelAlias(commit.cfg.OAuthModelAlias) - return true -} - -func (s *Service) updateServerClientsContext(ctx context.Context, cfg *config.Config) bool { - if s == nil || cfg == nil || (ctx != nil && ctx.Err() != nil) { - return false - } - if s.updateServerClientsContextFn != nil { - return s.updateServerClientsContextFn(ctx, cfg) - } - if s.server == nil { - return true - } - return s.server.UpdateClientsContext(ctx, cfg) -} - -func (s *Service) reloadConfigFromWatcher() bool { - if s == nil || s.watcher == nil { - return false - } - return s.watcher.ReloadConfigIfChanged() -} - -func (s *Service) registerConfigAPIKeyAuths(ctx context.Context, cfg *config.Config) { - if s == nil || s.coreManager == nil || cfg == nil { - return - } - if ctx == nil { - ctx = context.Background() - } - configSynth := synthesizer.NewConfigSynthesizer() - auths, errSynthesize := configSynth.Synthesize(&synthesizer.SynthesisContext{ - Config: cfg, - Now: time.Now(), - IDGenerator: synthesizer.NewStableIDGenerator(), - }) - if errSynthesize != nil { - log.Warnf("failed to synthesize config API key auths: %v", errSynthesize) - return - } - - registrationCtx := coreauth.WithDeferredAPIKeyModelAliasRebuild(ctx) - tasks := make([]modelRegistrationTask, 0, len(auths)) - needsAliasRebuild := false - for _, auth := range auths { - if !coreauth.IsConfigAPIKeyAuth(auth) { - continue - } - prepared := s.prepareCoreAuthForModelRegistration(registrationCtx, auth) - if prepared == nil { - continue - } - needsAliasRebuild = true - authForRegistration := prepared - tasks = append(tasks, modelRegistrationTask{ - phase: modelRegistrationPhaseConfigAPIKey, - category: modelRegistrationCategory(authForRegistration), - run: func(compatCache *openAICompatibilityRegistrationCache) { - s.completeModelRegistrationForAuthWithCache(registrationCtx, authForRegistration, compatCache) - }, - }) - } - if needsAliasRebuild { - s.coreManager.RefreshAPIKeyModelAlias() - } - s.runModelRegistrationTasks(registrationCtx, tasks) -} - -func forceHomeRuntimeConfig(cfg *config.Config) { - if cfg == nil { - return - } - cfg.APIKeys = nil - cfg.UsageStatisticsEnabled = true - cfg.DisableCooling = true - cfg.SaveCooldownStatus = false - cfg.WebsocketAuth = false - cfg.RemoteManagement.AllowRemote = false - cfg.RemoteManagement.DisableControlPanel = true - cfg.Plugins.StoreAuth = nil -} - -func (s *Service) applyHomeOverlay(remoteCfg *config.Config) { - if errApply := s.applyHomeOverlayContext(context.Background(), remoteCfg); errApply != nil { - log.Warnf("failed to apply home config payload: %v", errApply) - } -} - -func (s *Service) applyHomeOverlayContext(ctx context.Context, remoteCfg *config.Config) error { - return s.applyHomeOverlayWithClient(ctx, remoteCfg, nil) -} - -func (s *Service) applyHomeOverlayWithClient(ctx context.Context, remoteCfg *config.Config, client *home.Client) error { - work, errStage := s.stageHomeOverlayWithClient(ctx, remoteCfg, client) - if errStage != nil { - return errStage - } - if ctx != nil { - if errContext := ctx.Err(); errContext != nil { - return errContext - } - } - if work.config != nil { - if !s.applyConfigUpdateWithAuthSynthesis(ctx, work.config, true) { - return context.Canceled - } - work.committed = true - } - if errFinalize := s.finalizeHomePluginWork(ctx, client, work); errFinalize != nil { - return errFinalize - } - return nil -} - -func (s *Service) stageHomeOverlayWithClient(ctx context.Context, remoteCfg *config.Config, client *home.Client) (*homePluginFinalization, error) { - work := &homePluginFinalization{} - if s == nil || remoteCfg == nil { - return work, nil - } - if ctx == nil { - ctx = context.Background() - } - if errContext := ctx.Err(); errContext != nil { - return nil, errContext - } - - s.cfgMu.RLock() - baseCfg := s.cfg - s.cfgMu.RUnlock() - if baseCfg == nil { - return work, nil - } - - merged := *remoteCfg - merged.Host = baseCfg.Host - merged.Port = baseCfg.Port - merged.TLS = baseCfg.TLS - merged.Home = baseCfg.Home - storeAuth := merged.Plugins.StoreAuth - forceHomeRuntimeConfig(&merged) - syncCfg := merged - syncCfg.Plugins.StoreAuth = storeAuth - - logHomeConfigChanges(baseCfg, &merged) - report, syncKey, didSync, errSync := s.syncHomePluginsWithClient(ctx, &syncCfg, client) - if errSync != nil { - return nil, fmt.Errorf("sync home plugins: %w", errSync) - } - if errContext := ctx.Err(); errContext != nil { - return nil, errContext - } - if didSync { - if errLoad := homeplugins.MarkLoadResults(&report, s.pluginHost); errLoad != nil { - return nil, fmt.Errorf("load home plugins: %w", errLoad) - } - } - if strings.TrimSpace(report.Task) != "" { - work.syncKey = syncKey - work.markSynced = true - if strings.TrimSpace(merged.Home.NodeID) != "" { - work.statusWork = append(work.statusWork, homePluginStatusWork{cfg: &merged, report: report}) - } - } - taskWork, errTasks := s.stageHomePluginTasksWithClient(ctx, &merged, client) - if errTasks != nil { - return nil, fmt.Errorf("stage home plugin tasks: %w", errTasks) - } - work.taskWork = append(work.taskWork, taskWork...) - if errContext := ctx.Err(); errContext != nil { - return nil, errContext - } - work.config = &merged - return work, nil -} - -func (s *Service) commitHomeConfig(lifetimeCtx, homeCtx context.Context, generation uint64, work *homePluginFinalization) bool { - if s == nil || work == nil || work.config == nil { - return false - } - - s.homeConfigCommitMu.Lock() - defer s.homeConfigCommitMu.Unlock() - if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) { - return false - } - if s.homeConfigCommitHook != nil { - s.homeConfigCommitHook() - } - if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) { - return false - } - commit := s.commitConfigUpdate(work.config) - if commit.cfg == nil { - return false - } - work.config = commit.cfg - work.configCommit = commit - work.committed = true - return true -} - -func (s *Service) homeLifetimeActive(homeCtx, lifetimeCtx context.Context, generation uint64) bool { - if s == nil || homeCtx.Err() != nil || lifetimeCtx.Err() != nil { - return false - } - s.homeMu.Lock() - active := s.homeGeneration == generation - s.homeMu.Unlock() - return active -} - -func (s *Service) finalizeHomePluginWorkUntilDone(ctx, homeCtx context.Context, generation uint64, client *home.Client, work *homePluginFinalization, publish func() bool) error { - stopClose := closeHomeClientOnCancellation(ctx, client) - defer stopClose() - for { - if errContext := ctx.Err(); errContext != nil { - return errContext - } - - s.homeOwnershipMu.Lock() - if !s.homeLifetimeActive(homeCtx, ctx, generation) { - s.homeOwnershipMu.Unlock() - return context.Canceled - } - errFinalize := s.finalizeHomePluginWork(ctx, client, work) - if errFinalize == nil && (publish == nil || publish()) { - s.homeOwnershipMu.Unlock() - return nil - } - s.homeOwnershipMu.Unlock() - if errFinalize == nil { - return context.Canceled - } - - log.WithError(errFinalize).Warn("failed to finalize home plugins; retrying") - timer := time.NewTimer(homeSubscriberPreAckRetryBackoff) - select { - case <-ctx.Done(): - timer.Stop() - return ctx.Err() - case <-timer.C: - } - } -} - -func closeHomeClientOnCancellation(ctx context.Context, client *home.Client) func() { - if ctx == nil || client == nil { - return func() {} - } - stop := make(chan struct{}) - go func() { - select { - case <-ctx.Done(): - client.Close() - case <-stop: - } - }() - return func() { close(stop) } -} - -func logHomeConfigChanges(oldCfg, newCfg *config.Config) { - if oldCfg == nil || newCfg == nil || !newCfg.Home.Enabled || (!oldCfg.Debug && !newCfg.Debug) { - return - } - - details := diff.BuildConfigChangeDetails(oldCfg, newCfg) - if len(details) == 0 { - return - } - - if newCfg.Debug && !log.IsLevelEnabled(log.DebugLevel) { - util.SetLogLevel(newCfg) - } - - log.Debugf("home config changes detected:") - for _, detail := range details { - log.Debugf(" %s", detail) - } -} - -func (s *Service) startHomeUsageForwarder(ctx context.Context, client *home.Client) { - if s == nil || client == nil { - return - } - if ctx == nil { - ctx = context.Background() - } - - sleep := func(d time.Duration) bool { - if d <= 0 { - return true - } - timer := time.NewTimer(d) - defer timer.Stop() - select { - case <-ctx.Done(): - return false - case <-timer.C: - return true - } - } - - go func() { - for { - select { - case <-ctx.Done(): - return - default: - } - - if !client.HeartbeatOK() { - if !sleep(time.Second) { - return - } - continue - } - - items := redisqueue.PopOldest(64) - if len(items) == 0 { - if !sleep(500 * time.Millisecond) { - return - } - continue - } - - for i := range items { - if errPush := client.LPushUsage(ctx, items[i]); errPush != nil { - for j := i; j < len(items); j++ { - redisqueue.Enqueue(items[j]) - } - if !sleep(time.Second) { - return - } - break - } - } - } - }() -} - -func applyHomeObservationBarrier(registry *executionregistry.Registry, revision int64) { - if registry != nil { - registry.ObserveBarrier(revision) - } -} - -func applyHomeInFlightPublisherConfig(manager *coreauth.Manager, cfg internalconfig.CredentialInFlightConfig) error { - publisherCfg, errConfig := coreauth.HomeInFlightPublisherConfigFromConfig(cfg) - if errConfig != nil { - return errConfig - } - if manager != nil { - manager.ApplyHomeInFlightPublisherConfig(publisherCfg) - } - return nil -} - -func (s *Service) startHomeSubscriber(ctx context.Context) { - if s == nil { - return - } - s.cfgMu.RLock() - cfg := s.cfg - s.cfgMu.RUnlock() - if cfg == nil || !cfg.Home.Enabled { - return - } - - parentCtx := ctx - if parentCtx == nil { - parentCtx = context.Background() - } - - s.homeLifecycleMu.Lock() - defer s.homeLifecycleMu.Unlock() - - if previousSupervisor := s.homeSupervisor; previousSupervisor != nil { - s.homeConfigCommitMu.Lock() - previousSupervisor.cancel() - s.homeConfigCommitMu.Unlock() - <-previousSupervisor.done - } - if !s.drainDetachedHomeLifetime(parentCtx) { - return - } - if parentCtx.Err() != nil { - return - } - - homeCtx, cancel := context.WithCancel(parentCtx) - done := make(chan struct{}) - s.homeMu.Lock() - s.homeGeneration++ - generation := s.homeGeneration - s.homeCancel = cancel - s.homeMu.Unlock() - supervisor := &homeSubscriberSupervisor{cancel: cancel, done: done} - s.homeSupervisor = supervisor - go s.runHomeSubscriber(homeCtx, parentCtx, cfg.Home, generation, supervisor) -} - -func (s *Service) drainDetachedHomeLifetime(parentCtx context.Context) bool { - s.homeMu.Lock() - previousCancel := s.homeCancel - previousClient := s.homeClient - previousRegistry := s.homeRegistry - previousBundle := s.homeDispatchBundle - previousDrainBound := s.homeDrainBound - previousForwarder := s.homeLogForwarder - previousForwarderClient := s.homeLogForwarderClient - s.homeCancel = nil - s.homeClient = nil - s.homeRegistry = nil - s.homeDispatchBundle = nil - s.homeDrainBound = 0 - s.homeLogForwarderClient = nil - s.homeMu.Unlock() - - if s.coreManager != nil { - s.coreManager.ClearHomeDispatchBundle(previousBundle) - } - home.ClearCurrentIf(previousClient) - if previousCancel != nil { - previousCancel() - } - if previousForwarder != nil && previousForwarderClient == previousClient { - previousForwarder.Deactivate(previousClient) - } - if previousRegistry != nil { - if previousDrainBound <= 0 { - previousDrainBound = internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound - } - drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), previousDrainBound) - errDrain := previousRegistry.Drain(drainCtx) - cancelDrain() - if errDrain != nil { - if previousClient != nil { - previousClient.Close() - } - if parentCtx.Err() == nil { - log.WithError(errDrain).Error("failed to drain replaced Home execution registry") - s.cancelServiceRun() - } - return false - } - } - if previousClient != nil { - previousClient.Close() - } - return true -} - -func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.Context, homeCfg internalconfig.HomeConfig, generation uint64, supervisor *homeSubscriberSupervisor) { - defer func() { - s.homeMu.Lock() - if s.homeGeneration == generation { - s.homeCancel = nil - } - s.homeMu.Unlock() - close(supervisor.done) - }() - - for homeCtx.Err() == nil { - supervisor.setPublisherCompletion(nil) - client := home.New(homeCfg) - client.SetManagedLifetime(true) - registry := executionregistry.New() - releaseCtx, releaseCancel := context.WithCancel(context.WithoutCancel(homeCtx)) - releaseFlusher := home.NewReleaseFlusher(client.LimiterConfig, client.PushConcurrencyRelease) - registry.SetReleaseSink(releaseFlusher.MarkDirty) - releaseDone := make(chan struct{}) - go func() { - defer close(releaseDone) - releaseFlusher.Run(releaseCtx) - }() - lifetimeCtx, lifetimeCancel := context.WithCancel(homeCtx) - cancelBound := atomic.Int64{} - cancelBound.Store(int64(internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound)) - queue := newHomeConfigWorkQueue() - ready := make(chan struct{}) - var readyOnce sync.Once - var published atomic.Bool - workerDone := make(chan struct{}) - - go func() { - defer close(workerDone) - s.runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx, generation, client, registry, queue, ready, &published, &cancelBound, supervisor) - }() - - errRun := client.RunConfigSubscriberLifetime(lifetimeCtx, func(raw []byte) error { - parsed, errParse := config.ParseConfigBytes(raw) - if errParse != nil { - log.Warnf("failed to parse home config payload: %v", errParse) - return errParse - } - if errSetLifecycle := client.SetLifecycleConfig(parsed.CredentialConcurrency); errSetLifecycle != nil { - log.Warnf("failed to apply Home lifecycle config: %v", errSetLifecycle) - return errSetLifecycle - } - if errPublisherConfig := applyHomeInFlightPublisherConfig(s.coreManager, parsed.CredentialInFlight); errPublisherConfig != nil { - log.Warnf("failed to apply Home in-flight publisher config: %v", errPublisherConfig) - return errPublisherConfig - } - applyHomeObservationBarrier(registry, parsed.CredentialConcurrency.ObservationBarrierRevision) - cancelBound.Store(int64(parsed.CredentialConcurrency.WithDefaults().CPACancelBound)) - queue.enqueue(raw) - return nil - }, func() { - readyOnce.Do(func() { close(ready) }) - }) - lifetimeCancel() - <-workerDone - if publisherDone := supervisor.publisherCompletion(); publisherDone != nil { - <-publisherDone - } - - s.detachHomeSubscriberLifetime(client, registry) - drainBound := time.Duration(cancelBound.Load()) - drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), drainBound) - errDrain := registry.Drain(drainCtx) - var errFlush error - if errDrain == nil { - errFlush = releaseFlusher.Flush(drainCtx) - } - cancelDrain() - releaseCancel() - <-releaseDone - client.Close() - if errDrain != nil { - if parentCtx.Err() == nil { - log.WithError(errDrain).Error("failed to drain Home execution registry") - s.cancelServiceRun() - } - return - } - if errFlush != nil { - if parentCtx.Err() == nil { - log.WithError(errFlush).Error("failed to flush Home concurrency releases") - s.cancelServiceRun() - } - return - } - if errRun != nil && homeCtx.Err() == nil { - log.WithError(errRun).Warn("home config subscription lifetime ended") - } - if !published.Load() && errRun != nil && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) { - return - } - } -} - -func (s *Service) runHomeConfigWorker(lifetimeCtx, homeCtx context.Context, generation uint64, client *home.Client, registry *executionregistry.Registry, queue *homeConfigWorkQueue, ready <-chan struct{}, published *atomic.Bool, cancelBound *atomic.Int64) { - s.runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx, generation, client, registry, queue, ready, published, cancelBound, nil) -} - -func (s *Service) runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx context.Context, generation uint64, client *home.Client, registry *executionregistry.Registry, queue *homeConfigWorkQueue, ready <-chan struct{}, published *atomic.Bool, cancelBound *atomic.Int64, supervisor *homeSubscriberSupervisor) { - select { - case <-lifetimeCtx.Done(): - return - case <-ready: - } - - for { - if lifetimeCtx.Err() != nil { - return - } - raw, ok := queue.dequeue(lifetimeCtx) - if !ok { - return - } - if lifetimeCtx.Err() != nil { - return - } - - var work *homePluginFinalization - for { - if lifetimeCtx.Err() != nil { - return - } - parsed, errParse := config.ParseConfigBytes(raw) - if errParse == nil { - work, errParse = s.stageHomeOverlayWithClient(lifetimeCtx, parsed, client) - } - if errParse == nil { - break - } - if lifetimeCtx.Err() != nil { - return - } - log.WithError(errParse).Warn("failed to stage home config; retrying") - if !waitForHomeSubscriberRetry(lifetimeCtx, homeSubscriberPreAckRetryBackoff) { - return - } - } - - var publish func() bool - if !published.Load() { - publish = func() bool { - s.homeMu.Lock() - defer s.homeMu.Unlock() - if homeCtx.Err() != nil || lifetimeCtx.Err() != nil || s.homeGeneration != generation { - return false - } - s.homeClient = client - s.homeRegistry = registry - s.homeDrainBound = time.Duration(cancelBound.Load()) - if s.coreManager != nil { - s.homeDispatchBundle = s.coreManager.PublishHomeDispatch(client, registry, generation) - } - home.SetCurrent(client) - if s.homeLogForwarder == nil { - s.homeLogForwarder = startHomeLogForwarder(0) - } - s.homeLogForwarder.Bind(client) - s.homeLogForwarderClient = client - published.Store(true) - return true - } - } - if s.homeConfigStageHook != nil { - s.homeConfigStageHook() - } - if !s.commitHomeConfig(lifetimeCtx, homeCtx, generation, work) { - return - } - if s.homeConfigRuntimeHook != nil { - s.homeConfigRuntimeHook() - } - if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) || !s.applyConfigRuntime(lifetimeCtx, work.configCommit, true) { - return - } - if errFinalize := s.finalizeHomePluginWorkUntilDone(lifetimeCtx, homeCtx, generation, client, work, publish); errFinalize != nil { - if !errors.Is(errFinalize, context.Canceled) { - log.WithError(errFinalize).Warn("home plugin finalization ended") - } - return - } - if publish != nil { - s.startHomeInFlightPublisher(lifetimeCtx, client, registry, supervisor) - s.startHomeUsageForwarder(lifetimeCtx, client) - } - } -} - -func (s *Service) startHomeInFlightPublisher(ctx context.Context, client *home.Client, registry *executionregistry.Registry, supervisor *homeSubscriberSupervisor) { - if s == nil || s.coreManager == nil { - return - } - done := make(chan struct{}) - if supervisor != nil { - supervisor.setPublisherCompletion(done) - } - go func() { - defer close(done) - s.coreManager.StartHomeInFlightPublisher(ctx, client, registry) - }() -} - -func waitForHomeSubscriberRetry(ctx context.Context, delay time.Duration) bool { - timer := time.NewTimer(delay) - defer timer.Stop() - select { - case <-ctx.Done(): - return false - case <-timer.C: - return true - } -} - -func (s *Service) detachHomeSubscriberLifetime(client *home.Client, registry *executionregistry.Registry) { - if s == nil { - return - } - s.homeMu.Lock() - var bundle *coreauth.HomeDispatchBundle - if s.homeClient == client && s.homeRegistry == registry { - bundle = s.homeDispatchBundle - s.homeClient = nil - s.homeRegistry = nil - s.homeDispatchBundle = nil - s.homeDrainBound = 0 - } - forwarder := s.homeLogForwarder - if s.homeLogForwarderClient == client { - s.homeLogForwarderClient = nil - } else { - forwarder = nil - } - s.homeMu.Unlock() - if s.coreManager != nil { - s.coreManager.ClearHomeDispatchBundle(bundle) - } - home.ClearCurrentIf(client) - if forwarder != nil { - forwarder.Deactivate(client) - } -} - -func (s *Service) cancelServiceRun() { - if s == nil { - return - } - s.homeMu.Lock() - cancel := s.runCancel - if cancel == nil { - cancel = s.homeCancel - } - s.homeMu.Unlock() - if cancel != nil { - cancel() - } -} - -// Run starts the service and blocks until the context is cancelled or the server stops. -// It initializes all components including authentication, file watching, HTTP server, -// and starts processing requests. The method blocks until the context is cancelled. -// -// Parameters: -// - ctx: The context for controlling the service lifecycle -// -// Returns: -// - error: An error if the service fails to start or run -func (s *Service) Run(ctx context.Context) error { - if s == nil { - return fmt.Errorf("cliproxy: service is nil") - } - if ctx == nil { - ctx = context.Background() - } - ctx, runCancel := context.WithCancel(ctx) - s.homeMu.Lock() - s.runCancel = runCancel - s.homeMu.Unlock() - defer func() { - runCancel() - s.homeMu.Lock() - if s.runCancel != nil { - s.runCancel = nil - } - s.homeMu.Unlock() - }() - - usage.StartDefault(ctx) - homeEnabled := s.cfg != nil && s.cfg.Home.Enabled - if homeEnabled { - forceHomeRuntimeConfig(s.cfg) - redisqueue.SetUsageStatisticsEnabled(true) - } - - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) - defer shutdownCancel() - defer func() { - if err := s.Shutdown(shutdownCtx); err != nil { - log.Errorf("service shutdown returned error: %v", err) - } - }() - - if !homeEnabled { - if errEnsureAuthDir := s.ensureAuthDir(); errEnsureAuthDir != nil { - return errEnsureAuthDir - } - } - - s.applyRetryConfig(s.cfg) - s.configureCooldownStateStore(s.cfg) - - s.registerPluginAuthParser() - if s.coreManager != nil && !homeEnabled { - if errLoad := s.coreManager.Load(ctx); errLoad != nil { - log.Warnf("failed to load auth store: %v", errLoad) - } - s.registerConfigAPIKeyAuths(coreauth.WithSkipPersist(ctx), s.cfg) - if s.cfg.SaveCooldownStatus { - if errRestoreCooldown := s.coreManager.RestoreCooldownStates(ctx); errRestoreCooldown != nil { - log.Warnf("failed to restore cooldown state: %v", errRestoreCooldown) - } - } - } - - if !homeEnabled { - tokenResult, err := s.tokenProvider.Load(ctx, s.cfg) - if err != nil && !errors.Is(err, context.Canceled) { - return err - } - if tokenResult == nil { - tokenResult = &TokenClientResult{} - } - - apiKeyResult, err := s.apiKeyProvider.Load(ctx, s.cfg) - if err != nil && !errors.Is(err, context.Canceled) { - return err - } - if apiKeyResult == nil { - apiKeyResult = &APIKeyClientResult{} - } - } - - // legacy clients removed; no caches to refresh - - s.ensureWebsocketGateway() - if homeEnabled { - s.registerAvailableExecutors(ctx, executorRegistrationOptions{ - includeBaseline: true, - }) - // Home mode does not expose in-process Redis RESP usage output; usage is forwarded to home instead. - redisqueue.SetEnabled(true) - } - - // handlers no longer depend on legacy clients; pass nil slice initially - s.server = api.NewServer(s.cfg, s.coreManager, s.accessManager, s.configPath, s.serverOptions...) - s.syncPluginRuntimeConfig(ctx) - if homeEnabled { - s.syncPluginModelRuntime(ctx) - } - - if s.authManager == nil { - s.authManager = newDefaultAuthManager() - } - - if homeEnabled { - s.startHomeSubscriber(ctx) - } - - if s.server != nil && s.wsGateway != nil { - s.server.AttachWebsocketRoute(s.wsGateway.Path(), s.wsGateway.Handler()) - s.server.SetWebsocketAuthChangeHandler(func(oldEnabled, newEnabled bool) { - if oldEnabled == newEnabled { - return - } - if !oldEnabled && newEnabled { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if errStop := s.wsGateway.Stop(ctx); errStop != nil { - log.Warnf("failed to reset websocket connections after ws-auth change %t -> %t: %v", oldEnabled, newEnabled, errStop) - return - } - log.Debugf("ws-auth enabled; existing websocket sessions terminated to enforce authentication") - return - } - log.Debugf("ws-auth disabled; existing websocket sessions remain connected") - }) - } - - if s.hooks.OnBeforeStart != nil { - s.hooks.OnBeforeStart(s.cfg) - } - - s.serverErr = make(chan error, 1) - go func() { - if errStart := s.server.Start(); errStart != nil { - s.serverErr <- errStart - } else { - s.serverErr <- nil - } - }() - - time.Sleep(100 * time.Millisecond) - fmt.Printf("API server started successfully on: %s:%d\n", s.cfg.Host, s.cfg.Port) - - s.applyPprofConfig(s.cfg) - - if s.hooks.OnAfterStart != nil { - s.hooks.OnAfterStart(s) - } - - if !homeEnabled { - var watcherWrapper *WatcherWrapper - reloadCallback := func(newCfg *config.Config) { s.applyWatcherConfigUpdate(newCfg) } - - watcherWrapper, errCreate := s.watcherFactory(s.configPath, s.cfg.AuthDir, reloadCallback) - if errCreate != nil { - return fmt.Errorf("cliproxy: failed to create watcher: %w", errCreate) - } - s.watcher = watcherWrapper - s.ensureAuthUpdateQueue(ctx) - if s.authUpdates != nil { - watcherWrapper.SetAuthUpdateQueue(s.authUpdates) - } - watcherWrapper.SetConfig(s.cfg) - s.registerPluginAuthParser() - - watcherCtx, watcherCancel := context.WithCancel(context.Background()) - s.watcherCancel = watcherCancel - if errStart := watcherWrapper.Start(watcherCtx); errStart != nil { - return fmt.Errorf("cliproxy: failed to start watcher: %w", errStart) - } - log.Info("file watcher started for config and auth directory changes") - s.syncPluginModelRuntime(ctx) - } - - s.registerModelRefreshCallback() - - // Prefer core auth manager auto refresh if available. - if s.coreManager != nil && !homeEnabled { - interval := 15 * time.Minute - s.coreManager.StartAutoRefresh(context.Background(), interval) - log.Infof("core auth auto-refresh started (interval=%s)", interval) - } - - select { - case <-ctx.Done(): - log.Debug("service context cancelled, shutting down...") - return ctx.Err() - case errServer := <-s.serverErr: - return errServer - } -} - -// Shutdown gracefully stops background workers and the HTTP server. -// It ensures all resources are properly cleaned up and connections are closed. -// The shutdown is idempotent and can be called multiple times safely. -// -// Parameters: -// - ctx: The context for controlling the shutdown timeout -// -// Returns: -// - error: An error if shutdown fails -func (s *Service) Shutdown(ctx context.Context) error { - if s == nil { - return nil - } - var shutdownErr error - s.shutdownOnce.Do(func() { - if ctx == nil { - ctx = context.Background() - } - - s.homeLifecycleMu.Lock() - if supervisor := s.homeSupervisor; supervisor != nil { - s.homeConfigCommitMu.Lock() - supervisor.cancel() - s.homeConfigCommitMu.Unlock() - <-supervisor.done - } - s.homeMu.Lock() - homeCancel := s.homeCancel - homeClient := s.homeClient - homeRegistry := s.homeRegistry - homeDispatchBundle := s.homeDispatchBundle - homeForwarder := s.homeLogForwarder - homeForwarderClient := s.homeLogForwarderClient - s.homeGeneration++ - s.homeCancel = nil - s.homeClient = nil - s.homeRegistry = nil - s.homeDispatchBundle = nil - s.homeDrainBound = 0 - s.homeLogForwarder = nil - s.homeLogForwarderClient = nil - s.homeMu.Unlock() - if s.coreManager != nil { - s.coreManager.ClearHomeDispatchBundle(homeDispatchBundle) - } - home.ClearCurrentIf(homeClient) - if homeCancel != nil { - homeCancel() - } - if homeRegistry != nil { - if errClose := homeRegistry.Close(); errClose != nil { - log.WithError(errClose).Warn("failed to close Home execution registry during shutdown") - } - } - if homeClient != nil { - homeClient.Close() - } - if homeForwarder != nil { - if homeForwarderClient == homeClient { - homeForwarder.Deactivate(homeClient) - } - homeForwarder.Stop() - } - s.homeLifecycleMu.Unlock() - - // legacy refresh loop removed; only stopping core auth manager below - - if s.watcherCancel != nil { - s.watcherCancel() - } - if s.coreManager != nil { - s.coreManager.StopAutoRefresh() - } - if s.watcher != nil { - if err := s.watcher.Stop(); err != nil { - log.Errorf("failed to stop file watcher: %v", err) - shutdownErr = err - } - } - if s.wsGateway != nil { - if err := s.wsGateway.Stop(ctx); err != nil { - log.Errorf("failed to stop websocket gateway: %v", err) - if shutdownErr == nil { - shutdownErr = err - } - } - } - if s.authQueueStop != nil { - s.authQueueStop() - s.authQueueStop = nil - } - - if errShutdownPprof := s.shutdownPprof(ctx); errShutdownPprof != nil { - log.Errorf("failed to stop pprof server: %v", errShutdownPprof) - if shutdownErr == nil { - shutdownErr = errShutdownPprof - } - } - - // no legacy clients to persist - - if s.server != nil { - shutdownCtx, cancel := context.WithTimeout(ctx, 30*time.Second) - defer cancel() - if err := s.server.Stop(shutdownCtx); err != nil { - log.Errorf("error stopping API server: %v", err) - if shutdownErr == nil { - shutdownErr = err - } - } - } - - if s.pluginHost != nil { - sdktranslator.SetPluginHooks(nil) - sdkAuth.RegisterPluginAuthParser(nil) - if s.watcher != nil { - s.watcher.SetPluginAuthParser(nil) - } - s.pluginHost.ApplyConfig(ctx, &config.Config{}) - s.pluginHost.RegisterModels(ctx, registry.GetGlobalRegistry()) - s.registerAvailableExecutors(ctx, executorRegistrationOptions{ - includePlugins: true, - }) - s.pluginHost.RegisterFrontendAuthProviders() - s.pluginHost.ShutdownAllContext(ctx) - if s.accessManager != nil { - s.accessManager.SetProviders(sdkaccess.RegisteredProviders()) - } - } - - usage.StopDefault() - }) - return shutdownErr -} - -func (s *Service) ensureAuthDir() error { - info, err := os.Stat(s.cfg.AuthDir) - if err != nil { - if os.IsNotExist(err) { - if mkErr := os.MkdirAll(s.cfg.AuthDir, 0o755); mkErr != nil { - return fmt.Errorf("cliproxy: failed to create auth directory %s: %w", s.cfg.AuthDir, mkErr) - } - log.Infof("created missing auth directory: %s", s.cfg.AuthDir) - return nil - } - return fmt.Errorf("cliproxy: error checking auth directory %s: %w", s.cfg.AuthDir, err) - } - if !info.IsDir() { - return fmt.Errorf("cliproxy: auth path exists but is not a directory: %s", s.cfg.AuthDir) - } - return nil -} - -// registerModelsForAuth (re)binds provider models in the global registry using the core auth ID as client identifier. -func (s *Service) registerModelsForAuth(ctx context.Context, a *coreauth.Auth) { - s.registerModelsForAuthWithCache(ctx, a, nil) -} - -func (s *Service) registerModelsForAuthWithCache(ctx context.Context, a *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) { - if a == nil || a.ID == "" { - return - } - if ctx == nil { - ctx = context.Background() - } - if ctx.Err() != nil { - return - } - if a.Disabled { - GlobalModelRegistry().UnregisterClient(a.ID) - return - } - authKind := a.AuthKind() - // Unregister legacy client ID (if present) to avoid double counting - if a.Runtime != nil { - if idGetter, ok := a.Runtime.(interface{ GetClientID() string }); ok { - if rid := idGetter.GetClientID(); rid != "" && rid != a.ID { - GlobalModelRegistry().UnregisterClient(rid) - } - } - } - provider := strings.ToLower(strings.TrimSpace(a.Provider)) - compatProviderKey, compatDisplayName, compatDetected := openAICompatInfoFromAuth(a) - if compatDetected { - provider = "openai-compatibility" - } - excluded := s.oauthExcludedModels(provider, authKind) - // The synthesizer pre-merges per-account and global exclusions into the "excluded_models" attribute. - // If this attribute is present, it represents the complete list of exclusions and overrides the global config. - if a.Attributes != nil { - if val, ok := a.Attributes["excluded_models"]; ok && strings.TrimSpace(val) != "" { - excluded = strings.Split(val, ",") - } - } - if s.tryRegisterPluginModelsForAuth(ctx, a, provider, authKind, excluded) { - return - } - if ctx.Err() != nil { - return - } - var models []*ModelInfo - switch provider { - case constant.Gemini: - models = registry.GetGeminiModels() - if entry := s.resolveConfigGeminiKey(a); entry != nil { - if len(entry.Models) > 0 { - models = buildGeminiConfigModels(entry) - } - if authKind == "apikey" { - excluded = entry.ExcludedModels - } - } - models = applyExcludedModels(models, excluded) - case constant.GeminiInteractions: - models = registry.GetGeminiModels() - if entry := s.resolveConfigInteractionsKey(a); entry != nil { - if len(entry.Models) > 0 { - models = buildGeminiConfigModels(entry) - } - if authKind == "apikey" { - excluded = entry.ExcludedModels - } - } - models = applyExcludedModels(models, excluded) - case "vertex": - // Vertex AI Gemini supports the same model identifiers as Gemini. - models = registry.GetGeminiVertexModels() - if entry := s.resolveConfigVertexCompatKey(a); entry != nil { - if len(entry.Models) > 0 { - models = buildVertexCompatConfigModels(entry) - } - if authKind == "apikey" { - excluded = entry.ExcludedModels - } - } - models = applyExcludedModels(models, excluded) - case "aistudio": - models = registry.GetAIStudioModels() - models = applyExcludedModels(models, excluded) - case "antigravity": - models = registry.GetAntigravityModels() - models = applyAntigravityFetchedModelCapabilities(models, s.fetchAntigravityModelCapabilityHintsForAuth(ctx, a)) - models = applyExcludedModels(models, excluded) - case "claude": - models = registry.GetClaudeModels() - if entry := s.resolveConfigClaudeKey(a); entry != nil { - if len(entry.Models) > 0 { - models = buildClaudeConfigModels(entry) - } - if authKind == "apikey" { - excluded = entry.ExcludedModels - } - } - models = applyExcludedModels(models, excluded) - case "codex": - codexPlanType := "" - if a.Attributes != nil { - codexPlanType = strings.TrimSpace(a.Attributes["plan_type"]) - } - switch strings.ToLower(codexPlanType) { - case "pro": - models = registry.GetCodexProModels() - case "plus": - models = registry.GetCodexPlusModels() - case "team", "business", "go": - models = registry.GetCodexTeamModels() - case "free": - models = registry.GetCodexFreeModels() - default: - models = registry.GetCodexProModels() - } - if entry := s.resolveConfigCodexKey(a); entry != nil { - if len(entry.Models) > 0 { - models = buildCodexConfigModels(entry) - } - if authKind == "apikey" { - excluded = entry.ExcludedModels - } - } - models = applyExcludedModels(models, excluded) - case "kimi": - models = registry.GetKimiModels() - models = applyExcludedModels(models, excluded) - case "xai": - models = registry.GetXAIModels() - if entry := s.resolveConfigXAIKey(a); entry != nil { - if len(entry.Models) > 0 { - models = buildXAIConfigModels(entry) - } - if authKind == "apikey" { - excluded = entry.ExcludedModels - } - } - models = applyExcludedModels(models, excluded) - default: - // Handle OpenAI-compatibility providers by name using config - if s.cfg != nil { - providerKey := provider - compatName := strings.TrimSpace(a.Provider) - isCompatAuth := false - if compatDetected { - if compatProviderKey != "" { - providerKey = compatProviderKey - } - if compatDisplayName != "" { - compatName = compatDisplayName - } - isCompatAuth = true - } - if strings.EqualFold(providerKey, "openai-compatibility") { - isCompatAuth = true - if a.Attributes != nil { - if v := strings.TrimSpace(a.Attributes["compat_name"]); v != "" { - compatName = v - } - if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" { - providerKey = strings.ToLower(v) - isCompatAuth = true - } - } - if providerKey == "openai-compatibility" && compatName != "" { - providerKey = strings.ToLower(compatName) - } - } else if a.Attributes != nil { - if v := strings.TrimSpace(a.Attributes["compat_name"]); v != "" { - compatName = v - isCompatAuth = true - } - if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" { - providerKey = strings.ToLower(v) - isCompatAuth = true - } - } - if cached, ok := compatCache.lookup(compatName); ok { - isCompatAuth = true - if providerKey == "" { - providerKey = cached.providerKey - } - if providerKey == "" { - providerKey = "openai-compatibility" - } - ms := cached.models - if len(ms) > 0 { - ms = s.appendPluginModels(providerKey, ms) - s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) - } else { - ms = s.appendPluginModels(providerKey, nil) - if len(ms) > 0 { - s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) - } else { - GlobalModelRegistry().UnregisterClient(a.ID) - } - } - return - } - for i := range s.cfg.OpenAICompatibility { - compat := &s.cfg.OpenAICompatibility[i] - if compat.Disabled { - continue - } - if strings.EqualFold(compat.Name, compatName) { - isCompatAuth = true - ms := buildOpenAICompatibilityConfigModels(compat) - // Register and return - if len(ms) > 0 { - if providerKey == "" { - providerKey = "openai-compatibility" - } - ms = s.appendPluginModels(providerKey, ms) - s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) - } else { - // Ensure stale registrations are cleared when model list becomes empty. - ms = s.appendPluginModels(providerKey, nil) - if len(ms) > 0 { - s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) - } else { - GlobalModelRegistry().UnregisterClient(a.ID) - } - } - return - } - } - if isCompatAuth { - models = s.appendPluginModels(providerKey, nil) - if len(models) > 0 { - s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(models, a.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix)) - } else { - // No matching provider found or models removed entirely; drop any prior registration. - GlobalModelRegistry().UnregisterClient(a.ID) - } - return - } - } - } - if ctx.Err() != nil { - return - } - models = applyOAuthModelAliasForAuth(s.cfg, provider, authKind, a.Attributes, models) - if ctx.Err() != nil { - return - } - key := provider - if key == "" { - key = strings.ToLower(strings.TrimSpace(a.Provider)) - } - models = s.appendPluginModels(key, models) - if len(models) > 0 { - s.registerResolvedModelsForAuth(a, key, applyModelPrefixes(models, a.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix)) - return - } - - GlobalModelRegistry().UnregisterClient(a.ID) -} - -// refreshModelRegistrationForAuth re-applies the latest model registration for -// one auth and reconciles any concurrent auth changes that race with the -// refresh. Callers are expected to pre-filter provider membership. -// -// Re-registration is deliberate: registry cooldown/suspension state is treated -// as part of the previous registration snapshot and is cleared when the auth is -// rebound to the refreshed model catalog. -func (s *Service) refreshModelRegistrationForAuth(current *coreauth.Auth) bool { - return s.refreshModelRegistrationForAuthWithContext(context.Background(), current, nil) -} - -func (s *Service) refreshModelRegistrationForAuthWithCache(current *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) bool { - return s.refreshModelRegistrationForAuthWithContext(context.Background(), current, compatCache) -} - -func (s *Service) refreshModelRegistrationForAuthWithContext(ctx context.Context, current *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) bool { - if s == nil || s.coreManager == nil || current == nil || current.ID == "" { - return false - } - if ctx == nil { - ctx = context.Background() - } - if ctx.Err() != nil { - return false - } - if !current.Disabled { - s.ensureExecutorsForAuthWithContext(ctx, current, false) - } - s.registerModelsForAuthWithCache(ctx, current, compatCache) - s.coreManager.ReconcileRegistryModelStates(ctx, current.ID) - if ctx.Err() != nil { - return false - } - - latest, ok := s.latestAuthForModelRegistration(current.ID) - if !ok || latest.Disabled { - GlobalModelRegistry().UnregisterClient(current.ID) - s.coreManager.RefreshSchedulerEntry(current.ID) - return false - } - - // Re-apply the latest auth snapshot so concurrent auth updates cannot leave - // stale model registrations behind. This may duplicate registration work when - // no auth fields changed, but keeps the refresh path simple and correct. - s.ensureExecutorsForAuthWithContext(ctx, latest, false) - s.registerModelsForAuthWithCache(ctx, latest, compatCache) - if ctx.Err() != nil { - return false - } - s.coreManager.ReconcileRegistryModelStates(ctx, latest.ID) - s.coreManager.RefreshSchedulerEntry(current.ID) - return true -} - -// latestAuthForModelRegistration returns the latest auth snapshot regardless of -// provider membership. Callers use this after a registration attempt to restore -// whichever state currently owns the client ID in the global registry. -func (s *Service) latestAuthForModelRegistration(authID string) (*coreauth.Auth, bool) { - if s == nil || s.coreManager == nil || authID == "" { - return nil, false - } - auth, ok := s.coreManager.GetByID(authID) - if !ok || auth == nil || auth.ID == "" { - return nil, false - } - return auth, true -} - -func (s *Service) resolveConfigClaudeKey(auth *coreauth.Auth) *config.ClaudeKey { - if auth == nil || s.cfg == nil { - return nil - } - var attrKey, attrBase string - if auth.Attributes != nil { - attrKey = strings.TrimSpace(auth.Attributes["api_key"]) - attrBase = strings.TrimSpace(auth.Attributes["base_url"]) - } - for i := range s.cfg.ClaudeKey { - entry := &s.cfg.ClaudeKey[i] - cfgKey := strings.TrimSpace(entry.APIKey) - cfgBase := strings.TrimSpace(entry.BaseURL) - if attrKey != "" && attrBase != "" { - if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { - return entry - } - continue - } - if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { - if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - if attrKey != "" { - for i := range s.cfg.ClaudeKey { - entry := &s.cfg.ClaudeKey[i] - if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { - return entry - } - } - } - return nil -} - -func (s *Service) resolveConfigGeminiKey(auth *coreauth.Auth) *config.GeminiKey { - if s == nil || s.cfg == nil { - return nil - } - return s.resolveConfigGeminiKeyEntry(auth, s.cfg.GeminiKey) -} - -func (s *Service) resolveConfigInteractionsKey(auth *coreauth.Auth) *config.GeminiKey { - if s == nil || s.cfg == nil { - return nil - } - return s.resolveConfigGeminiKeyEntry(auth, s.cfg.InteractionsKey) -} - -func (s *Service) resolveConfigGeminiKeyEntry(auth *coreauth.Auth, entries []config.GeminiKey) *config.GeminiKey { - if auth == nil || s.cfg == nil { - return nil - } - var attrKey, attrBase string - if auth.Attributes != nil { - attrKey = strings.TrimSpace(auth.Attributes["api_key"]) - attrBase = strings.TrimSpace(auth.Attributes["base_url"]) - } - for i := range entries { - entry := &entries[i] - cfgKey := strings.TrimSpace(entry.APIKey) - cfgBase := strings.TrimSpace(entry.BaseURL) - if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { - if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { - return entry - } - continue - } - if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - return nil -} - -func (s *Service) resolveConfigVertexCompatKey(auth *coreauth.Auth) *config.VertexCompatKey { - if auth == nil || s.cfg == nil { - return nil - } - var attrKey, attrBase string - if auth.Attributes != nil { - attrKey = strings.TrimSpace(auth.Attributes["api_key"]) - attrBase = strings.TrimSpace(auth.Attributes["base_url"]) - } - for i := range s.cfg.VertexCompatAPIKey { - entry := &s.cfg.VertexCompatAPIKey[i] - cfgKey := strings.TrimSpace(entry.APIKey) - cfgBase := strings.TrimSpace(entry.BaseURL) - if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { - if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { - return entry - } - continue - } - if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - if attrKey != "" { - for i := range s.cfg.VertexCompatAPIKey { - entry := &s.cfg.VertexCompatAPIKey[i] - if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { - return entry - } - } - } - return nil -} - -func (s *Service) resolveConfigCodexKey(auth *coreauth.Auth) *config.CodexKey { - if s == nil || s.cfg == nil { - return nil - } - return resolveConfigCodexStyleKey(auth, s.cfg.CodexKey) -} - -func (s *Service) resolveConfigXAIKey(auth *coreauth.Auth) *config.XAIKey { - if s == nil || s.cfg == nil { - return nil - } - return resolveConfigCodexStyleKey(auth, s.cfg.XAIKey) -} - -func resolveConfigCodexStyleKey(auth *coreauth.Auth, entries []config.CodexKey) *config.CodexKey { - if auth == nil { - return nil - } - var attrKey, attrBase string - if auth.Attributes != nil { - attrKey = strings.TrimSpace(auth.Attributes["api_key"]) - attrBase = strings.TrimSpace(auth.Attributes["base_url"]) - } - for i := range entries { - entry := &entries[i] - cfgKey := strings.TrimSpace(entry.APIKey) - cfgBase := strings.TrimSpace(entry.BaseURL) - if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { - if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { - return entry - } - continue - } - if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { - return entry - } - } - return nil -} - -func (s *Service) oauthExcludedModels(provider, authKind string) []string { - cfg := s.cfg - if cfg == nil { - return nil - } - authKindKey := strings.ToLower(strings.TrimSpace(authKind)) - providerKey := strings.ToLower(strings.TrimSpace(provider)) - if authKindKey == "apikey" { - return nil - } - return cfg.OAuthExcludedModels[providerKey] -} - -func applyExcludedModels(models []*ModelInfo, excluded []string) []*ModelInfo { - if len(models) == 0 || len(excluded) == 0 { - return models - } - - patterns := make([]string, 0, len(excluded)) - for _, item := range excluded { - if trimmed := strings.TrimSpace(item); trimmed != "" { - patterns = append(patterns, strings.ToLower(trimmed)) - } - } - if len(patterns) == 0 { - return models - } - - filtered := make([]*ModelInfo, 0, len(models)) - for _, model := range models { - if model == nil { - continue - } - modelID := strings.ToLower(strings.TrimSpace(model.ID)) - blocked := false - for _, pattern := range patterns { - if matchWildcard(pattern, modelID) { - blocked = true - break - } - } - if !blocked { - filtered = append(filtered, model) - } - } - return filtered -} - -func applyModelPrefixes(models []*ModelInfo, prefix string, forceModelPrefix bool) []*ModelInfo { - trimmedPrefix := strings.TrimSpace(prefix) - if trimmedPrefix == "" || len(models) == 0 { - return models - } - - out := make([]*ModelInfo, 0, len(models)*2) - seen := make(map[string]struct{}, len(models)*2) - - addModel := func(model *ModelInfo) { - if model == nil { - return - } - id := strings.TrimSpace(model.ID) - if id == "" { - return - } - if _, exists := seen[id]; exists { - return - } - seen[id] = struct{}{} - out = append(out, model) - } - - for _, model := range models { - if model == nil { - continue - } - baseID := strings.TrimSpace(model.ID) - if baseID == "" { - continue - } - if !forceModelPrefix || trimmedPrefix == baseID { - addModel(model) - } - clone := *model - clone.ID = trimmedPrefix + "/" + baseID - addModel(&clone) - } - return out -} - -// matchWildcard performs case-insensitive wildcard matching where '*' matches any substring. -func matchWildcard(pattern, value string) bool { - if pattern == "" { - return false - } - - // Fast path for exact match (no wildcard present). - if !strings.Contains(pattern, "*") { - return pattern == value - } - - parts := strings.Split(pattern, "*") - // Handle prefix. - if prefix := parts[0]; prefix != "" { - if !strings.HasPrefix(value, prefix) { - return false - } - value = value[len(prefix):] - } - - // Handle suffix. - if suffix := parts[len(parts)-1]; suffix != "" { - if !strings.HasSuffix(value, suffix) { - return false - } - value = value[:len(value)-len(suffix)] - } - - // Handle middle segments in order. - for i := 1; i < len(parts)-1; i++ { - segment := parts[i] - if segment == "" { - continue - } - idx := strings.Index(value, segment) - if idx < 0 { - return false - } - value = value[idx+len(segment):] - } - - return true -} - -type modelEntry interface { - GetName() string - GetAlias() string - GetDisplayName() string -} - -func buildConfiguredModelInfo(model modelEntry, ownedBy, modelType string, created int64, fallbackDisplayName string, userDefined bool) *ModelInfo { - name := strings.TrimSpace(model.GetName()) - alias := strings.TrimSpace(model.GetAlias()) - if alias == "" { - alias = name - } - if alias == "" { - return nil - } - displayName := strings.TrimSpace(model.GetDisplayName()) - if displayName == "" { - displayName = fallbackDisplayName - } - if displayName == "" { - displayName = alias - } - return &ModelInfo{ - ID: alias, - Object: "model", - Created: created, - OwnedBy: ownedBy, - Type: modelType, - DisplayName: displayName, - UserDefined: userDefined, - } -} - -func buildOpenAICompatibilityConfigModels(compat *config.OpenAICompatibility) []*ModelInfo { - if compat == nil || len(compat.Models) == 0 { - return nil - } - now := time.Now().Unix() - models := make([]*ModelInfo, 0, len(compat.Models)) - for i := range compat.Models { - model := compat.Models[i] - modelType := "openai-compatibility" - if model.Image { - modelType = registry.OpenAIImageModelType - } - info := buildConfiguredModelInfo(model, compat.Name, modelType, now, strings.TrimSpace(model.Alias), false) - if info == nil { - continue - } - thinking := model.Thinking - if thinking == nil && !model.Image { - thinking = ®istry.ThinkingSupport{Levels: []string{"low", "medium", "high"}} - } - info.Thinking = thinking - info.SupportedInputModalities = normalizeCompatConfigModalities(model.InputModalities) - info.SupportedOutputModalities = normalizeCompatConfigModalities(model.OutputModalities) - models = append(models, info) - } - return models -} - -func normalizeCompatConfigModalities(raw []string) []string { - if len(raw) == 0 { - return nil - } - out := make([]string, 0, len(raw)) - seen := make(map[string]struct{}, len(raw)) - for _, item := range raw { - modality := strings.ToLower(strings.TrimSpace(item)) - if modality == "" { - continue - } - if _, exists := seen[modality]; exists { - continue - } - seen[modality] = struct{}{} - out = append(out, modality) - } - if len(out) == 0 { - return nil - } - return out -} - -func buildConfigModels[T modelEntry](models []T, ownedBy, modelType string) []*ModelInfo { - if len(models) == 0 { - return nil - } - now := time.Now().Unix() - out := make([]*ModelInfo, 0, len(models)) - seen := make(map[string]struct{}, len(models)) - for i := range models { - model := models[i] - name := strings.TrimSpace(model.GetName()) - info := buildConfiguredModelInfo(model, ownedBy, modelType, now, name, true) - if info == nil { - continue - } - alias := info.ID - key := strings.ToLower(alias) - if _, exists := seen[key]; exists { - continue - } - seen[key] = struct{}{} - if name != "" { - if upstream := registry.LookupStaticModelInfo(name); upstream != nil && upstream.Thinking != nil { - info.Thinking = upstream.Thinking - } - } - out = append(out, info) - } - return out -} - -func buildVertexCompatConfigModels(entry *config.VertexCompatKey) []*ModelInfo { - if entry == nil { - return nil - } - return buildConfigModels(entry.Models, "google", "vertex") -} - -func buildGeminiConfigModels(entry *config.GeminiKey) []*ModelInfo { - if entry == nil { - return nil - } - return buildConfigModels(entry.Models, "google", "gemini") -} - -func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo { - if entry == nil { - return nil - } - return buildConfigModels(entry.Models, "anthropic", "claude") -} - -func buildXAIConfigModels(entry *config.XAIKey) []*ModelInfo { - if entry == nil { - return nil - } - return buildConfigModels(entry.Models, "xai", "xai") -} - -func buildCodexConfigModels(entry *config.CodexKey) []*ModelInfo { - if entry == nil { - return nil - } - - models := registry.WithCodexBuiltins(buildConfigModels(entry.Models, "openai", "openai")) - configuredDisplayNames := make(map[string]string, len(entry.Models)) - seenConfiguredModels := make(map[string]struct{}, len(entry.Models)) - for i := range entry.Models { - model := entry.Models[i] - alias := strings.TrimSpace(model.Alias) - if alias == "" { - alias = strings.TrimSpace(model.Name) - } - if alias == "" { - continue - } - key := strings.ToLower(alias) - if _, exists := seenConfiguredModels[key]; exists { - continue - } - seenConfiguredModels[key] = struct{}{} - - displayName := strings.TrimSpace(model.DisplayName) - if displayName != "" { - configuredDisplayNames[key] = displayName - } - } - for _, model := range models { - if model == nil { - continue - } - if displayName, ok := configuredDisplayNames[strings.ToLower(model.ID)]; ok { - model.DisplayName = displayName - } - } - return models -} - -func rewriteModelInfoName(name, oldID, newID string) string { - trimmed := strings.TrimSpace(name) - if trimmed == "" { - return name - } - oldID = strings.TrimSpace(oldID) - newID = strings.TrimSpace(newID) - if oldID == "" || newID == "" { - return name - } - if strings.EqualFold(oldID, newID) { - return name - } - if strings.EqualFold(trimmed, oldID) { - return newID - } - if strings.HasSuffix(trimmed, "/"+oldID) { - prefix := strings.TrimSuffix(trimmed, oldID) - return prefix + newID - } - if trimmed == "models/"+oldID { - return "models/" + newID - } - return name -} - -func applyOAuthModelAlias(cfg *config.Config, provider, authKind string, models []*ModelInfo) []*ModelInfo { - return applyOAuthModelAliasForAuth(cfg, provider, authKind, nil, models) -} - -func applyOAuthModelAliasForAuth(cfg *config.Config, provider, authKind string, attributes map[string]string, models []*ModelInfo) []*ModelInfo { - if len(models) == 0 { - return models - } - channel := coreauth.OAuthModelAliasChannel(provider, authKind) - if channel == "" { - return models - } - aliases := oauthModelAliasesForAuth(cfg, channel, attributes) - if len(aliases) == 0 { - return models - } - return applyOAuthModelAliasEntries(aliases, models) -} - -func oauthModelAliasesForAuth(cfg *config.Config, channel string, attributes map[string]string) []config.OAuthModelAlias { - perAuthAliases := coreauth.OAuthModelAliasesFromAttributes(attributes) - if cfg == nil || len(cfg.OAuthModelAlias) == 0 { - return perAuthAliases - } - globalAliases := cfg.OAuthModelAlias[channel] - if len(perAuthAliases) == 0 { - return globalAliases - } - if len(globalAliases) == 0 { - return perAuthAliases - } - out := make([]config.OAuthModelAlias, 0, len(perAuthAliases)+len(globalAliases)) - seenAlias := make(map[string]struct{}, len(perAuthAliases)+len(globalAliases)) - add := func(aliases []config.OAuthModelAlias) { - for _, entry := range aliases { - alias := strings.TrimSpace(entry.Alias) - if alias == "" { - continue - } - key := strings.ToLower(alias) - if _, exists := seenAlias[key]; exists { - continue - } - seenAlias[key] = struct{}{} - out = append(out, entry) - } - } - add(perAuthAliases) - add(globalAliases) - return out -} - -func applyOAuthModelAliasEntries(aliases []config.OAuthModelAlias, models []*ModelInfo) []*ModelInfo { - type aliasEntry struct { - alias string - displayName string - fork bool - } - - forward := make(map[string][]aliasEntry, len(aliases)) - for i := range aliases { - name := strings.TrimSpace(aliases[i].Name) - alias := strings.TrimSpace(aliases[i].Alias) - if name == "" || alias == "" { - continue - } - if strings.EqualFold(name, alias) { - continue - } - key := strings.ToLower(name) - forward[key] = append(forward[key], aliasEntry{ - alias: alias, - displayName: strings.TrimSpace(aliases[i].DisplayName), - fork: aliases[i].Fork, - }) - } - if len(forward) == 0 { - return models - } - - out := make([]*ModelInfo, 0, len(models)) - seen := make(map[string]struct{}, len(models)) - for _, model := range models { - if model == nil { - continue - } - id := strings.TrimSpace(model.ID) - if id == "" { - continue - } - key := strings.ToLower(id) - entries := forward[key] - if len(entries) == 0 { - if _, exists := seen[key]; exists { - continue - } - seen[key] = struct{}{} - out = append(out, model) - continue - } - - keepOriginal := false - for _, entry := range entries { - if entry.fork { - keepOriginal = true - break - } - } - if keepOriginal { - if _, exists := seen[key]; !exists { - seen[key] = struct{}{} - out = append(out, model) - } - } - - addedAlias := false - for _, entry := range entries { - mappedID := strings.TrimSpace(entry.alias) - if mappedID == "" { - continue - } - if strings.EqualFold(mappedID, id) { - continue - } - aliasKey := strings.ToLower(mappedID) - if _, exists := seen[aliasKey]; exists { - continue - } - seen[aliasKey] = struct{}{} - clone := *model - clone.ID = mappedID - if entry.displayName != "" { - clone.DisplayName = entry.displayName - } - if clone.Name != "" { - clone.Name = rewriteModelInfoName(clone.Name, id, mappedID) - } - out = append(out, &clone) - addedAlias = true - } - - if !keepOriginal && !addedAlias { - if _, exists := seen[key]; exists { - continue - } - seen[key] = struct{}{} - out = append(out, model) - } - } - return out -} diff --git a/sdk/cliproxy/service_auth.go b/sdk/cliproxy/service_auth.go new file mode 100644 index 000000000..0b1990c21 --- /dev/null +++ b/sdk/cliproxy/service_auth.go @@ -0,0 +1,432 @@ +package cliproxy + +import ( + "context" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher" + "github.com/router-for-me/CLIProxyAPI/v7/internal/wsrelay" + sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + log "github.com/sirupsen/logrus" +) + +// newDefaultAuthManager creates a default authentication manager with supported OAuth providers. +func newDefaultAuthManager() *sdkAuth.Manager { + return sdkAuth.NewManager( + sdkAuth.GetTokenStore(), + sdkAuth.NewCodexAuthenticator(), + sdkAuth.NewClaudeAuthenticator(), + sdkAuth.NewXAIAuthenticator(), + ) +} + +func (s *Service) ensureAuthUpdateQueue(ctx context.Context) { + if s == nil { + return + } + if s.authUpdates == nil { + s.authUpdates = make(chan watcher.AuthUpdate, 256) + } + if s.authQueueStop != nil { + return + } + queueCtx, cancel := context.WithCancel(ctx) + s.authQueueStop = cancel + go s.consumeAuthUpdates(queueCtx) +} + +func (s *Service) consumeAuthUpdates(ctx context.Context) { + ctx = coreauth.WithSkipPersist(ctx) + for { + select { + case <-ctx.Done(): + return + case update, ok := <-s.authUpdates: + if !ok { + return + } + updates := []watcher.AuthUpdate{update} + labelDrain: + for { + select { + case nextUpdate := <-s.authUpdates: + updates = append(updates, nextUpdate) + default: + break labelDrain + } + } + s.handleAuthUpdates(ctx, updates) + } + } +} + +func (s *Service) emitAuthUpdate(ctx context.Context, update watcher.AuthUpdate) { + if s == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + if s.watcher != nil && s.watcher.DispatchRuntimeAuthUpdate(update) { + return + } + if s.authUpdates != nil { + select { + case s.authUpdates <- update: + return + default: + log.Debugf("auth update queue saturated, applying inline action=%v id=%s", update.Action, update.ID) + } + } + s.handleAuthUpdate(ctx, update) +} + +func (s *Service) handleAuthUpdate(ctx context.Context, update watcher.AuthUpdate) { + s.handleAuthUpdates(ctx, []watcher.AuthUpdate{update}) +} + +func (s *Service) handleAuthUpdates(ctx context.Context, updates []watcher.AuthUpdate) { + if s == nil { + return + } + updates = coalesceAuthUpdates(updates) + s.cfgMu.RLock() + cfg := s.cfg + s.cfgMu.RUnlock() + if cfg == nil || s.coreManager == nil { + return + } + + registrationCtx := coreauth.WithDeferredAPIKeyModelAliasRebuild(ctx) + tasks := make([]modelRegistrationTask, 0, len(updates)) + needsPluginSync := false + needsAliasRebuild := false + for _, update := range updates { + switch update.Action { + case watcher.AuthUpdateActionAdd, watcher.AuthUpdateActionModify: + if update.Auth == nil || update.Auth.ID == "" { + continue + } + auth := s.prepareCoreAuthForModelRegistration(registrationCtx, update.Auth) + if auth == nil { + continue + } + needsAliasRebuild = true + authForRegistration := auth + tasks = append(tasks, modelRegistrationTask{ + phase: modelRegistrationPhase(authForRegistration), + category: modelRegistrationCategory(authForRegistration), + run: func(compatCache *openAICompatibilityRegistrationCache) { + s.completeModelRegistrationForAuthWithCache(registrationCtx, authForRegistration, compatCache) + }, + }) + needsPluginSync = true + case watcher.AuthUpdateActionDelete: + id := update.ID + if id == "" && update.Auth != nil { + id = update.Auth.ID + } + if id == "" { + continue + } + s.applyCoreAuthRemoval(registrationCtx, id) + needsAliasRebuild = true + default: + log.Debugf("received unknown auth update action: %v", update.Action) + } + } + + if needsAliasRebuild { + s.coreManager.RefreshAPIKeyModelAlias() + } + s.runModelRegistrationTasks(registrationCtx, tasks) + if needsPluginSync { + s.syncPluginRuntime(registrationCtx) + } +} + +func coalesceAuthUpdates(updates []watcher.AuthUpdate) []watcher.AuthUpdate { + if len(updates) <= 1 { + return updates + } + order := make([]string, 0, len(updates)) + byID := make(map[string]watcher.AuthUpdate, len(updates)) + unkeyed := make([]watcher.AuthUpdate, 0) + for _, update := range updates { + id := authUpdateID(update) + if id == "" { + unkeyed = append(unkeyed, update) + continue + } + if _, exists := byID[id]; !exists { + order = append(order, id) + } + byID[id] = update + } + if len(byID) == 0 { + return unkeyed + } + out := make([]watcher.AuthUpdate, 0, len(byID)+len(unkeyed)) + for _, id := range order { + out = append(out, byID[id]) + } + out = append(out, unkeyed...) + return out +} + +func authUpdateID(update watcher.AuthUpdate) string { + if strings.TrimSpace(update.ID) != "" { + return strings.TrimSpace(update.ID) + } + if update.Auth != nil { + return strings.TrimSpace(update.Auth.ID) + } + return "" +} + +func (s *Service) ensureWebsocketGateway() { + if s == nil { + return + } + if s.wsGateway != nil { + return + } + opts := wsrelay.Options{ + Path: "/v1/ws", + OnConnected: s.wsOnConnected, + OnDisconnected: s.wsOnDisconnected, + LogDebugf: log.Debugf, + LogInfof: log.Infof, + LogWarnf: log.Warnf, + } + s.wsGateway = wsrelay.NewManager(opts) +} + +func (s *Service) wsOnConnected(channelID string) { + if s == nil || channelID == "" { + return + } + if !strings.HasPrefix(strings.ToLower(channelID), "aistudio-") { + return + } + if s.coreManager != nil { + if existing, ok := s.coreManager.GetByID(channelID); ok && existing != nil { + if !existing.Disabled && existing.Status == coreauth.StatusActive { + return + } + } + } + now := time.Now().UTC() + auth := &coreauth.Auth{ + ID: channelID, // keep channel identifier as ID + Provider: "aistudio", // logical provider for switch routing + Label: channelID, // display original channel id + Status: coreauth.StatusActive, + CreatedAt: now, + UpdatedAt: now, + Attributes: map[string]string{"runtime_only": "true"}, + Metadata: map[string]any{"email": channelID}, // metadata drives logging and usage tracking + } + log.Infof("websocket provider connected: %s", channelID) + s.emitAuthUpdate(context.Background(), watcher.AuthUpdate{ + Action: watcher.AuthUpdateActionAdd, + ID: auth.ID, + Auth: auth, + }) +} + +func (s *Service) wsOnDisconnected(channelID string, reason error) { + if s == nil || channelID == "" { + return + } + if reason != nil { + if strings.Contains(reason.Error(), "replaced by new connection") { + log.Infof("websocket provider replaced: %s", channelID) + return + } + log.Warnf("websocket provider disconnected: %s (%v)", channelID, reason) + } else { + log.Infof("websocket provider disconnected: %s", channelID) + } + ctx := context.Background() + s.emitAuthUpdate(ctx, watcher.AuthUpdate{ + Action: watcher.AuthUpdateActionDelete, + ID: channelID, + }) +} + +func (s *Service) applyCoreAuthAddOrUpdate(ctx context.Context, auth *coreauth.Auth) { + auth = s.prepareCoreAuthForModelRegistration(ctx, auth) + if auth == nil { + return + } + s.completeModelRegistrationForAuth(ctx, auth) + s.syncPluginRuntime(ctx) +} + +func (s *Service) prepareCoreAuthForModelRegistration(ctx context.Context, auth *coreauth.Auth) *coreauth.Auth { + if s == nil || s.coreManager == nil || auth == nil || auth.ID == "" { + return nil + } + auth = auth.Clone() + s.ensureExecutorsForAuthWithContext(ctx, auth, false) + + // IMPORTANT: Update coreManager FIRST, before model registration. + // This ensures that configuration changes (proxy_url, prefix, etc.) take effect + // immediately for API calls, rather than waiting for model registration to complete. + op := "register" + var err error + if existing, ok := s.coreManager.GetByID(auth.ID); ok { + auth.CreatedAt = existing.CreatedAt + if !existing.Disabled && existing.Status != coreauth.StatusDisabled && !auth.Disabled && auth.Status != coreauth.StatusDisabled { + auth.LastRefreshedAt = existing.LastRefreshedAt + auth.NextRefreshAfter = existing.NextRefreshAfter + if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { + auth.ModelStates = existing.ModelStates + } + } + op = "update" + _, err = s.coreManager.Update(ctx, auth) + } else { + _, err = s.coreManager.Register(ctx, auth) + } + if err != nil { + log.Errorf("failed to %s auth %s: %v", op, auth.ID, err) + current, ok := s.coreManager.GetByID(auth.ID) + if !ok || current.Disabled { + GlobalModelRegistry().UnregisterClient(auth.ID) + return nil + } + auth = current + } + return auth +} + +func (s *Service) completeModelRegistrationForAuth(ctx context.Context, auth *coreauth.Auth) { + s.completeModelRegistrationForAuthWithCache(ctx, auth, nil) +} + +func (s *Service) completeModelRegistrationForAuthWithCache(ctx context.Context, auth *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) { + if s == nil || s.coreManager == nil || auth == nil || auth.ID == "" { + return + } + if ctx != nil && ctx.Err() != nil { + return + } + s.registerModelsForAuthWithCache(ctx, auth, compatCache) + if ctx != nil && ctx.Err() != nil { + return + } + s.coreManager.ReconcileRegistryModelStates(ctx, auth.ID) + + // Refresh the scheduler entry so that the auth's supportedModelSet is rebuilt + // from the now-populated global model registry. Without this, newly added auths + // have an empty supportedModelSet (because Register/Update upserts into the + // scheduler before registerModelsForAuth runs) and are invisible to the scheduler. + s.coreManager.RefreshSchedulerEntry(auth.ID) +} + +func (s *Service) applyCoreAuthRemoval(ctx context.Context, id string) { + if s == nil || id == "" { + return + } + if s.coreManager == nil { + return + } + id = strings.TrimSpace(id) + var provider string + if existing, ok := s.coreManager.GetByID(id); ok && existing != nil { + provider = strings.TrimSpace(existing.Provider) + } + GlobalModelRegistry().UnregisterClient(id) + s.coreManager.Remove(ctx, id) + if strings.EqualFold(provider, "codex") { + executor.CloseCodexWebsocketSessionsForAuthID(id, "auth_removed") + } + if strings.EqualFold(provider, "xai") { + executor.CloseXAIWebsocketSessionsForAuthID(id, "auth_removed") + } + s.syncPluginRuntime(ctx) +} + +func (s *Service) applyRetryConfig(cfg *config.Config) { + if s == nil || s.coreManager == nil || cfg == nil { + return + } + maxInterval := time.Duration(cfg.MaxRetryInterval) * time.Second + s.coreManager.SetRetryConfig(cfg.RequestRetry, maxInterval, cfg.MaxRetryCredentials) + coreauth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) +} + +func (s *Service) configureCooldownStateStore(cfg *config.Config) { + _ = s.configureCooldownStateStoreContext(context.Background(), cfg, false) +} + +func (s *Service) configureCooldownStateStoreContext(ctx context.Context, cfg *config.Config, persistOld bool) bool { + if s == nil || s.coreManager == nil { + return true + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + return s.coreManager.SwapCooldownStateStore(ctx, s.resolveCooldownStateStore(cfg), persistOld) +} + +func (s *Service) resolveCooldownStateStore(cfg *config.Config) coreauth.CooldownStateStore { + if cfg == nil || !cfg.SaveCooldownStatus || cfg.Home.Enabled { + return nil + } + authDir, errResolve := resolveCooldownStateAuthDir(cfg) + if errResolve != nil { + log.Warnf("failed to resolve cooldown state directory: %v", errResolve) + return nil + } + if authDir == "" { + return nil + } + return coreauth.NewFileCooldownStateStoreWithAuthDir(authDir, authDir) +} + +func resolveCooldownStateAuthDir(cfg *config.Config) (string, error) { + if cfg == nil { + return "", nil + } + authDir, errAuthDir := util.ResolveAuthDir(cfg.AuthDir) + if errAuthDir != nil { + return "", errAuthDir + } + return authDir, nil +} + +func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName string, ok bool) { + if a == nil { + return "", "", false + } + if len(a.Attributes) > 0 { + providerKey = strings.TrimSpace(a.Attributes["provider_key"]) + compatName = strings.TrimSpace(a.Attributes["compat_name"]) + if compatName != "" { + if providerKey == "" { + providerKey = compatName + } + return util.OpenAICompatibleProviderKey(providerKey), compatName, true + } + } + if strings.EqualFold(strings.TrimSpace(a.Provider), "openai-compatibility") { + compatName = strings.TrimSpace(a.Label) + providerKey = compatName + if providerKey == "" { + providerKey = "openai-compatibility" + } + return util.OpenAICompatibleProviderKey(providerKey), compatName, true + } + return "", "", false +} diff --git a/sdk/cliproxy/service_config.go b/sdk/cliproxy/service_config.go new file mode 100644 index 000000000..c0e74eab6 --- /dev/null +++ b/sdk/cliproxy/service_config.go @@ -0,0 +1,287 @@ +package cliproxy + +import ( + "context" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/synthesizer" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + log "github.com/sirupsen/logrus" +) + +func (s *Service) applyConfigUpdate(newCfg *config.Config) { + s.applyConfigUpdateWithAuthSynthesis(context.Background(), newCfg, true) +} + +func (s *Service) applyWatcherConfigUpdate(newCfg *config.Config) { + s.applyConfigUpdateWithAuthSynthesis(context.Background(), newCfg, false) +} + +type configCommit struct { + cfg *config.Config + sequence uint64 +} + +type routingRuntimeState struct { + strategy string + sessionAffinity bool + sessionAffinityTTL time.Duration +} + +func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { + state := routingRuntimeState{ + strategy: "round-robin", + sessionAffinityTTL: time.Hour, + } + if cfg == nil { + return state + } + + switch strings.ToLower(strings.TrimSpace(cfg.Routing.Strategy)) { + case "fill-first", "fillfirst", "ff": + state.strategy = "fill-first" + } + state.sessionAffinity = cfg.Routing.SessionAffinity + if ttl := strings.TrimSpace(cfg.Routing.SessionAffinityTTL); ttl != "" { + if parsed, errParse := time.ParseDuration(ttl); errParse == nil && parsed > 0 { + state.sessionAffinityTTL = parsed + } + } + return state +} + +func newRoutingSelector(state routingRuntimeState) coreauth.Selector { + var selector coreauth.Selector + if state.strategy == "fill-first" { + selector = &coreauth.FillFirstSelector{} + } else { + selector = &coreauth.RoundRobinSelector{} + } + if state.sessionAffinity { + selector = coreauth.NewSessionAffinitySelectorWithConfig(coreauth.SessionAffinityConfig{ + Fallback: selector, + TTL: state.sessionAffinityTTL, + }) + } + return selector +} + +func (s *Service) applyConfigUpdateWithAuthSynthesis(ctx context.Context, newCfg *config.Config, synthesizeConfigAuths bool) bool { + commit := s.commitConfigUpdate(newCfg) + if commit.cfg == nil { + return false + } + return s.applyConfigRuntime(ctx, commit, synthesizeConfigAuths) +} + +// commitConfigUpdate applies only in-memory configuration state. Runtime work that +// may block on plugins, models, storage, or networking is deliberately deferred. +func (s *Service) commitConfigUpdate(newCfg *config.Config) configCommit { + if s == nil { + return configCommit{} + } + + s.configUpdateMu.Lock() + defer s.configUpdateMu.Unlock() + + if newCfg == nil { + s.cfgMu.RLock() + newCfg = s.cfg + s.cfgMu.RUnlock() + } + if newCfg == nil { + return configCommit{} + } + + s.cfgMu.Lock() + s.cfg = newCfg + s.cfgMu.Unlock() + s.configSequence++ + return configCommit{cfg: newCfg, sequence: s.configSequence} +} + +func (s *Service) configCommitCurrent(commit configCommit) bool { + if s == nil || commit.sequence == 0 { + return false + } + s.configUpdateMu.Lock() + current := s.configSequence == commit.sequence + s.configUpdateMu.Unlock() + return current +} + +func (s *Service) applyConfigRuntime(ctx context.Context, commit configCommit, synthesizeConfigAuths bool) bool { + cfg := commit.cfg + if s == nil || cfg == nil { + return false + } + s.configRuntimeMu.Lock() + defer s.configRuntimeMu.Unlock() + if !s.configCommitCurrent(commit) { + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + + if !s.applyManagerConfig(ctx, commit) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + if !s.applyPprofConfigContext(ctx, cfg) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + if !s.updateServerClientsContext(ctx, cfg) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + + registrationCtx := coreauth.WithSkipPersist(ctx) + s.syncPluginRuntimeConfigForConfig(registrationCtx, cfg) + if errContext := ctx.Err(); errContext != nil { + return false + } + var auths []*coreauth.Auth + if s.coreManager != nil { + auths = s.coreManager.List() + } + s.registerAvailableExecutors(registrationCtx, executorRegistrationOptions{ + includeBaseline: cfg.Home.Enabled, + forceReplaceAuths: true, + auths: auths, + }) + if errContext := ctx.Err(); errContext != nil { + return false + } + if synthesizeConfigAuths { + s.registerConfigAPIKeyAuths(registrationCtx, cfg) + } + if errContext := ctx.Err(); errContext != nil { + return false + } + if s.coreManager != nil && !cfg.Home.Enabled && cfg.SaveCooldownStatus { + if errRestoreCooldown := s.coreManager.RestoreCooldownStates(registrationCtx); errRestoreCooldown != nil && ctx.Err() == nil { + log.Warnf("failed to restore cooldown state after config update: %v", errRestoreCooldown) + } + } + if errContext := ctx.Err(); errContext != nil { + return false + } + s.syncPluginModelRuntime(registrationCtx) + return ctx.Err() == nil +} + +func (s *Service) applyManagerConfig(ctx context.Context, commit configCommit) bool { + if s == nil || s.coreManager == nil || commit.cfg == nil { + return s != nil && commit.cfg != nil + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + routingState := normalizedRoutingRuntimeState(commit.cfg) + if s.appliedRoutingState == nil || *s.appliedRoutingState != routingState { + s.coreManager.SetSelector(newRoutingSelector(routingState)) + s.appliedRoutingState = &routingState + } + s.applyRetryConfig(commit.cfg) + store := s.resolveCooldownStateStore(commit.cfg) + if !s.coreManager.ApplyConfigWithCooldownStateStore(ctx, commit.cfg, store) { + return false + } + s.coreManager.SetOAuthModelAlias(commit.cfg.OAuthModelAlias) + return true +} + +func (s *Service) updateServerClientsContext(ctx context.Context, cfg *config.Config) bool { + if s == nil || cfg == nil || (ctx != nil && ctx.Err() != nil) { + return false + } + if s.updateServerClientsContextFn != nil { + return s.updateServerClientsContextFn(ctx, cfg) + } + if s.server == nil { + return true + } + return s.server.UpdateClientsContext(ctx, cfg) +} + +func (s *Service) reloadConfigFromWatcher() bool { + if s == nil || s.watcher == nil { + return false + } + return s.watcher.ReloadConfigIfChanged() +} + +func (s *Service) registerConfigAPIKeyAuths(ctx context.Context, cfg *config.Config) { + if s == nil || s.coreManager == nil || cfg == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + configSynth := synthesizer.NewConfigSynthesizer() + auths, errSynthesize := configSynth.Synthesize(&synthesizer.SynthesisContext{ + Config: cfg, + Now: time.Now(), + IDGenerator: synthesizer.NewStableIDGenerator(), + }) + if errSynthesize != nil { + log.Warnf("failed to synthesize config API key auths: %v", errSynthesize) + return + } + + registrationCtx := coreauth.WithDeferredAPIKeyModelAliasRebuild(ctx) + tasks := make([]modelRegistrationTask, 0, len(auths)) + needsAliasRebuild := false + for _, auth := range auths { + if !coreauth.IsConfigAPIKeyAuth(auth) { + continue + } + prepared := s.prepareCoreAuthForModelRegistration(registrationCtx, auth) + if prepared == nil { + continue + } + needsAliasRebuild = true + authForRegistration := prepared + tasks = append(tasks, modelRegistrationTask{ + phase: modelRegistrationPhaseConfigAPIKey, + category: modelRegistrationCategory(authForRegistration), + run: func(compatCache *openAICompatibilityRegistrationCache) { + s.completeModelRegistrationForAuthWithCache(registrationCtx, authForRegistration, compatCache) + }, + }) + } + if needsAliasRebuild { + s.coreManager.RefreshAPIKeyModelAlias() + } + s.runModelRegistrationTasks(registrationCtx, tasks) +} + +func forceHomeRuntimeConfig(cfg *config.Config) { + if cfg == nil { + return + } + cfg.APIKeys = nil + cfg.UsageStatisticsEnabled = true + cfg.DisableCooling = true + cfg.SaveCooldownStatus = false + cfg.WebsocketAuth = false + cfg.RemoteManagement.AllowRemote = false + cfg.RemoteManagement.DisableControlPanel = true + cfg.Plugins.StoreAuth = nil +} diff --git a/sdk/cliproxy/service_executors.go b/sdk/cliproxy/service_executors.go new file mode 100644 index 000000000..6dd93ca1c --- /dev/null +++ b/sdk/cliproxy/service_executors.go @@ -0,0 +1,458 @@ +package cliproxy + +import ( + "context" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/constant" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" +) + +type openAICompatibilityRegistrationCache struct { + byName map[string]*openAICompatibilityRegistrationEntry +} + +type openAICompatibilityRegistrationEntry struct { + providerKey string + models []*ModelInfo +} + +func (s *Service) newOpenAICompatibilityRegistrationCache() *openAICompatibilityRegistrationCache { + if s == nil { + return nil + } + s.cfgMu.RLock() + cfg := s.cfg + s.cfgMu.RUnlock() + if cfg == nil || len(cfg.OpenAICompatibility) == 0 { + return nil + } + + cache := &openAICompatibilityRegistrationCache{ + byName: make(map[string]*openAICompatibilityRegistrationEntry, len(cfg.OpenAICompatibility)), + } + for i := range cfg.OpenAICompatibility { + compat := &cfg.OpenAICompatibility[i] + if compat.Disabled { + continue + } + compatName := strings.TrimSpace(compat.Name) + key := strings.ToLower(compatName) + if _, exists := cache.byName[key]; exists { + continue + } + providerName := strings.ToLower(compatName) + if providerName == "" { + providerName = "openai-compatibility" + } + cache.byName[key] = &openAICompatibilityRegistrationEntry{ + providerKey: util.OpenAICompatibleProviderKey(providerName), + models: buildOpenAICompatibilityConfigModels(compat), + } + } + if len(cache.byName) == 0 { + return nil + } + return cache +} + +func (c *openAICompatibilityRegistrationCache) lookup(compatName string) (*openAICompatibilityRegistrationEntry, bool) { + if c == nil || len(c.byName) == 0 { + return nil, false + } + entry, ok := c.byName[strings.ToLower(strings.TrimSpace(compatName))] + return entry, ok +} + +func (s *Service) hasNativeOpenAICompatExecutorConfig(a *coreauth.Auth, providerKey string, cfg *config.Config) bool { + if a == nil { + return false + } + providerKey = strings.ToLower(strings.TrimSpace(providerKey)) + if a.Attributes != nil { + if strings.TrimSpace(a.Attributes["base_url"]) != "" { + return true + } + if strings.TrimSpace(a.Attributes["compat_name"]) != "" { + return true + } + } + if strings.EqualFold(strings.TrimSpace(a.Provider), "openai-compatibility") { + return true + } + if s == nil || cfg == nil { + return false + } + + candidates := make([]string, 0, 3) + if providerKey != "" { + candidates = append(candidates, providerKey) + } + if a.Attributes != nil { + if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" { + candidates = append(candidates, strings.ToLower(v)) + } + } + if provider := strings.TrimSpace(a.Provider); provider != "" { + candidates = append(candidates, strings.ToLower(provider)) + } + + for i := range cfg.OpenAICompatibility { + compat := &cfg.OpenAICompatibility[i] + if compat.Disabled { + continue + } + name := strings.ToLower(strings.TrimSpace(compat.Name)) + if name == "" { + continue + } + for _, candidate := range candidates { + if candidate != "" && candidate == name { + return true + } + } + } + return false +} + +func (s *Service) unregisterOpenAICompatExecutor(providerKey string) { + if s == nil || s.coreManager == nil { + return + } + providerKey = strings.ToLower(strings.TrimSpace(providerKey)) + if providerKey == "" { + return + } + existing, okExecutor := s.coreManager.Executor(providerKey) + if !okExecutor || existing == nil { + return + } + if _, okOpenAICompat := existing.(*executor.OpenAICompatExecutor); !okOpenAICompat { + return + } + s.coreManager.UnregisterExecutor(providerKey) +} + +func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) { + s.ensureExecutorsForAuthWithContext(context.Background(), a, false) +} + +func (s *Service) ensureExecutorsForAuthWithMode(a *coreauth.Auth, forceReplace bool) { + s.ensureExecutorsForAuthWithContext(context.Background(), a, forceReplace) +} + +func (s *Service) ensureExecutorsForAuthWithContext(ctx context.Context, a *coreauth.Auth, forceReplace bool) { + if a == nil || (ctx != nil && ctx.Err() != nil) { + return + } + s.registerAvailableExecutors(ctx, executorRegistrationOptions{ + auths: []*coreauth.Auth{a}, + forceReplaceAuths: forceReplace, + }) +} + +func (s *Service) registerAvailableExecutors(ctx context.Context, opts executorRegistrationOptions) { + if s == nil || s.coreManager == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + s.executorRegistrationMu.Lock() + defer s.executorRegistrationMu.Unlock() + if ctx.Err() != nil { + return + } + // Keep all Service-owned executor registration paths here so native, Home, + // auth-derived, and plugin executors stay in the same binding order. + if opts.includeBaseline { + s.registerExecutorsForAuths(baselineExecutorAuths(), opts.forceReplaceAuths) + } + if len(opts.auths) > 0 { + s.registerExecutorsForAuths(opts.auths, opts.forceReplaceAuths) + } + if opts.includePlugins && s.pluginHost != nil { + registerPluginExecutors(s.pluginHost, s.coreManager) + } +} + +func baselineExecutorAuths() []*coreauth.Auth { + providers := []string{ + "codex", + "claude", + constant.Gemini, + constant.GeminiInteractions, + "vertex", + "aistudio", + "antigravity", + "kimi", + "xai", + "openai-compatibility", + } + auths := make([]*coreauth.Auth, 0, len(providers)) + for _, provider := range providers { + auth := &coreauth.Auth{ + ID: provider, + Provider: provider, + } + if provider == "openai-compatibility" { + auth.Attributes = map[string]string{"compat_name": "openai-compatibility"} + } + auths = append(auths, auth) + } + return auths +} + +func (s *Service) registerExecutorsForAuths(auths []*coreauth.Auth, forceReplace bool) { + reboundCodex := false + for _, auth := range auths { + if auth != nil && strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + if reboundCodex && forceReplace { + continue + } + reboundCodex = true + } + s.registerExecutorForAuth(auth, forceReplace) + } +} + +func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { + if s == nil || s.coreManager == nil || a == nil { + return + } + s.cfgMu.RLock() + cfg := s.cfg + s.cfgMu.RUnlock() + if strings.EqualFold(strings.TrimSpace(a.Provider), "codex") { + if !forceReplace { + existingExecutor, hasExecutor := s.coreManager.Executor("codex") + if hasExecutor { + _, isCodexAutoExecutor := existingExecutor.(*executor.CodexAutoExecutor) + if isCodexAutoExecutor { + return + } + } + } + s.coreManager.RegisterExecutor(executor.NewCodexAutoExecutor(cfg)) + return + } + // Skip disabled auth entries when (re)binding executors. + // Disabled auths can linger during config reloads (e.g., removed OpenAI-compat entries) + // and must not override active provider executors. + if a.Disabled { + return + } + if compatProviderKey, _, isCompat := openAICompatInfoFromAuth(a); isCompat { + if compatProviderKey == "" { + compatProviderKey = strings.ToLower(strings.TrimSpace(a.Provider)) + } + if compatProviderKey == "" { + compatProviderKey = "openai-compatibility" + } + if !forceReplace { + if existingExecutor, hasExecutor := s.coreManager.Executor(compatProviderKey); hasExecutor { + if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { + return + } + } + } + s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(compatProviderKey, cfg)) + return + } + switch strings.ToLower(a.Provider) { + case constant.Gemini: + s.coreManager.RegisterExecutor(executor.NewGeminiExecutor(cfg)) + case constant.GeminiInteractions: + s.coreManager.RegisterExecutor(executor.NewGeminiInteractionsExecutor(cfg)) + case "vertex": + s.coreManager.RegisterExecutor(executor.NewGeminiVertexExecutor(cfg)) + case "aistudio": + if s.wsGateway != nil { + s.coreManager.RegisterExecutor(executor.NewAIStudioExecutor(cfg, a.ID, s.wsGateway)) + } + return + case "antigravity": + s.coreManager.RegisterExecutor(executor.NewAntigravityExecutor(cfg)) + case "claude": + s.coreManager.RegisterExecutor(executor.NewClaudeExecutor(cfg)) + case "kimi": + s.coreManager.RegisterExecutor(executor.NewKimiExecutor(cfg)) + case "xai": + if !forceReplace { + existingExecutor, hasExecutor := s.coreManager.Executor("xai") + if hasExecutor { + existingXAIAutoExecutor, isXAIAutoExecutor := existingExecutor.(*executor.XAIAutoExecutor) + if isXAIAutoExecutor && existingXAIAutoExecutor.UsesConfig(cfg) { + return + } + } + } + s.coreManager.RegisterExecutor(executor.NewXAIAutoExecutor(cfg)) + default: + providerKey := strings.ToLower(strings.TrimSpace(a.Provider)) + if providerKey == "" { + providerKey = "openai-compatibility" + } + if s.pluginHost != nil && + s.pluginHost.HasExecutorCandidateProvider(providerKey) && + !s.hasNativeOpenAICompatExecutorConfig(a, providerKey, cfg) { + s.unregisterOpenAICompatExecutor(providerKey) + return + } + if !forceReplace { + if existingExecutor, hasExecutor := s.coreManager.Executor(providerKey); hasExecutor { + if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { + return + } + } + } + s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(providerKey, cfg)) + } +} + +func (s *Service) registerResolvedModelsForAuth(a *coreauth.Auth, providerKey string, models []*ModelInfo) { + if a == nil || a.ID == "" { + return + } + providerKey = strings.ToLower(strings.TrimSpace(providerKey)) + if providerKey == "" { + GlobalModelRegistry().UnregisterClient(a.ID) + return + } + normalizedModels := make([]*ModelInfo, 0, len(models)) + for _, model := range models { + if model == nil { + continue + } + modelID := strings.TrimSpace(model.ID) + if modelID == "" { + continue + } + clone := *model + clone.ID = modelID + normalizedModels = append(normalizedModels, &clone) + } + if len(normalizedModels) == 0 { + GlobalModelRegistry().UnregisterClient(a.ID) + return + } + GlobalModelRegistry().RegisterClient(a.ID, providerKey, normalizedModels) +} + +func (s *Service) pluginModelsForProvider(providerKey string) []*ModelInfo { + if s == nil || s.pluginHost == nil { + return nil + } + return s.pluginHost.ModelsForProvider(providerKey) +} + +func (s *Service) appendPluginModels(providerKey string, models []*ModelInfo) []*ModelInfo { + pluginModels := s.pluginModelsForProvider(providerKey) + if len(pluginModels) == 0 { + return models + } + out := make([]*ModelInfo, 0, len(models)+len(pluginModels)) + seen := make(map[string]struct{}, len(models)+len(pluginModels)) + for _, model := range models { + if model == nil { + continue + } + modelID := strings.TrimSpace(model.ID) + if modelID != "" { + seen[modelID] = struct{}{} + } + out = append(out, model) + } + for _, model := range pluginModels { + if model == nil { + continue + } + modelID := strings.TrimSpace(model.ID) + if modelID == "" { + continue + } + if _, exists := seen[modelID]; exists { + continue + } + seen[modelID] = struct{}{} + out = append(out, model) + } + return out +} + +func (s *Service) tryRegisterPluginModelsForAuth(ctx context.Context, a *coreauth.Auth, provider, authKind string, excluded []string) bool { + if s == nil || s.pluginHost == nil || a == nil { + return false + } + if ctx != nil && ctx.Err() != nil { + return true + } + result := s.pluginHost.ModelsForAuth(ctx, a) + if ctx != nil && ctx.Err() != nil { + return true + } + if !result.Handled { + return false + } + if result.Err != nil { + return true + } + activeAuth := a + providerKey := strings.ToLower(strings.TrimSpace(result.Provider)) + if providerKey == "" { + providerKey = strings.ToLower(strings.TrimSpace(provider)) + } + if result.Auth != nil && s.coreManager != nil { + result.Auth.ID = a.ID + if result.Auth.Provider == "" { + result.Auth.Provider = a.Provider + } + if result.Auth.FileName == "" { + result.Auth.FileName = a.FileName + } + if result.Auth.Attributes == nil { + result.Auth.Attributes = make(map[string]string) + } + for key, value := range a.Attributes { + if _, exists := result.Auth.Attributes[key]; !exists { + result.Auth.Attributes[key] = value + } + } + if updated, errUpdate := s.coreManager.Update(ctx, result.Auth); errUpdate == nil && updated != nil { + activeAuth = updated.Clone() + } + } + if activeAuth == nil { + activeAuth = a + } + if activeProvider := strings.ToLower(strings.TrimSpace(activeAuth.Provider)); activeProvider != "" { + providerKey = activeProvider + } + if providerKey == "" { + providerKey = strings.ToLower(strings.TrimSpace(provider)) + } + activeAuthKind := activeAuth.AuthKind() + activeExcluded := s.oauthExcludedModels(providerKey, activeAuthKind) + if a == activeAuth && len(activeExcluded) == 0 { + activeExcluded = excluded + } + if activeAuth.Attributes != nil { + if val, ok := activeAuth.Attributes["excluded_models"]; ok && strings.TrimSpace(val) != "" { + activeExcluded = strings.Split(val, ",") + } + } + if ctx != nil && ctx.Err() != nil { + return true + } + models := applyExcludedModels(result.Models, activeExcluded) + models = applyOAuthModelAliasForAuth(s.cfg, providerKey, activeAuthKind, activeAuth.Attributes, models) + if len(models) > 0 { + s.registerResolvedModelsForAuth(activeAuth, providerKey, applyModelPrefixes(models, activeAuth.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix)) + return true + } + GlobalModelRegistry().UnregisterClient(activeAuth.ID) + return true +} diff --git a/sdk/cliproxy/service_home.go b/sdk/cliproxy/service_home.go new file mode 100644 index 000000000..f13be7494 --- /dev/null +++ b/sdk/cliproxy/service_home.go @@ -0,0 +1,743 @@ +package cliproxy + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/homeplugins" + "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + "github.com/router-for-me/CLIProxyAPI/v7/internal/watcher/diff" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + log "github.com/sirupsen/logrus" +) + +type homeSubscriberSupervisor struct { + cancel context.CancelFunc + done chan struct{} + + publisherMu sync.Mutex + publisherDone <-chan struct{} +} + +func (s *homeSubscriberSupervisor) setPublisherCompletion(done <-chan struct{}) { + if s == nil { + return + } + s.publisherMu.Lock() + s.publisherDone = done + s.publisherMu.Unlock() +} + +func (s *homeSubscriberSupervisor) publisherCompletion() <-chan struct{} { + if s == nil { + return nil + } + s.publisherMu.Lock() + defer s.publisherMu.Unlock() + return s.publisherDone +} + +type homeConfigWorkQueue struct { + mu sync.Mutex + items [][]byte + wake chan struct{} +} + +func newHomeConfigWorkQueue() *homeConfigWorkQueue { + return &homeConfigWorkQueue{wake: make(chan struct{}, 1)} +} + +func (q *homeConfigWorkQueue) enqueue(raw []byte) { + if q == nil { + return + } + item := append([]byte(nil), raw...) + q.mu.Lock() + q.items = append(q.items, item) + q.mu.Unlock() + select { + case q.wake <- struct{}{}: + default: + } +} + +func (q *homeConfigWorkQueue) dequeue(ctx context.Context) ([]byte, bool) { + if q == nil || ctx == nil { + return nil, false + } + for { + if ctx.Err() != nil { + return nil, false + } + q.mu.Lock() + if ctx.Err() != nil { + q.mu.Unlock() + return nil, false + } + if len(q.items) > 0 { + item := q.items[0] + q.items[0] = nil + q.items = q.items[1:] + q.mu.Unlock() + return item, true + } + q.mu.Unlock() + select { + case <-ctx.Done(): + return nil, false + case <-q.wake: + } + } +} + +type homeLogForwarder interface { + Bind(*home.Client) + Deactivate(*home.Client) + Stop() +} + +var startHomeLogForwarder = func(queueSize int) homeLogForwarder { + return logging.StartHomeAppLogForwarder(queueSize) +} + +func (s *Service) applyHomeOverlay(remoteCfg *config.Config) { + if errApply := s.applyHomeOverlayContext(context.Background(), remoteCfg); errApply != nil { + log.Warnf("failed to apply home config payload: %v", errApply) + } +} + +func (s *Service) applyHomeOverlayContext(ctx context.Context, remoteCfg *config.Config) error { + return s.applyHomeOverlayWithClient(ctx, remoteCfg, nil) +} + +func (s *Service) applyHomeOverlayWithClient(ctx context.Context, remoteCfg *config.Config, client *home.Client) error { + work, errStage := s.stageHomeOverlayWithClient(ctx, remoteCfg, client) + if errStage != nil { + return errStage + } + if ctx != nil { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + } + if work.config != nil { + if !s.applyConfigUpdateWithAuthSynthesis(ctx, work.config, true) { + return context.Canceled + } + work.committed = true + } + if errFinalize := s.finalizeHomePluginWork(ctx, client, work); errFinalize != nil { + return errFinalize + } + return nil +} + +func (s *Service) stageHomeOverlayWithClient(ctx context.Context, remoteCfg *config.Config, client *home.Client) (*homePluginFinalization, error) { + work := &homePluginFinalization{} + if s == nil || remoteCfg == nil { + return work, nil + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return nil, errContext + } + + s.cfgMu.RLock() + baseCfg := s.cfg + s.cfgMu.RUnlock() + if baseCfg == nil { + return work, nil + } + + merged := *remoteCfg + merged.Host = baseCfg.Host + merged.Port = baseCfg.Port + merged.TLS = baseCfg.TLS + merged.Home = baseCfg.Home + storeAuth := merged.Plugins.StoreAuth + forceHomeRuntimeConfig(&merged) + syncCfg := merged + syncCfg.Plugins.StoreAuth = storeAuth + + logHomeConfigChanges(baseCfg, &merged) + report, syncKey, didSync, errSync := s.syncHomePluginsWithClient(ctx, &syncCfg, client) + if errSync != nil { + return nil, fmt.Errorf("sync home plugins: %w", errSync) + } + if errContext := ctx.Err(); errContext != nil { + return nil, errContext + } + if didSync { + if errLoad := homeplugins.MarkLoadResults(&report, s.pluginHost); errLoad != nil { + return nil, fmt.Errorf("load home plugins: %w", errLoad) + } + } + if strings.TrimSpace(report.Task) != "" { + work.syncKey = syncKey + work.markSynced = true + if strings.TrimSpace(merged.Home.NodeID) != "" { + work.statusWork = append(work.statusWork, homePluginStatusWork{cfg: &merged, report: report}) + } + } + taskWork, errTasks := s.stageHomePluginTasksWithClient(ctx, &merged, client) + if errTasks != nil { + return nil, fmt.Errorf("stage home plugin tasks: %w", errTasks) + } + work.taskWork = append(work.taskWork, taskWork...) + if errContext := ctx.Err(); errContext != nil { + return nil, errContext + } + work.config = &merged + return work, nil +} + +func (s *Service) commitHomeConfig(lifetimeCtx, homeCtx context.Context, generation uint64, work *homePluginFinalization) bool { + if s == nil || work == nil || work.config == nil { + return false + } + + s.homeConfigCommitMu.Lock() + defer s.homeConfigCommitMu.Unlock() + if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) { + return false + } + if s.homeConfigCommitHook != nil { + s.homeConfigCommitHook() + } + if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) { + return false + } + commit := s.commitConfigUpdate(work.config) + if commit.cfg == nil { + return false + } + work.config = commit.cfg + work.configCommit = commit + work.committed = true + return true +} + +func (s *Service) homeLifetimeActive(homeCtx, lifetimeCtx context.Context, generation uint64) bool { + if s == nil || homeCtx.Err() != nil || lifetimeCtx.Err() != nil { + return false + } + s.homeMu.Lock() + active := s.homeGeneration == generation + s.homeMu.Unlock() + return active +} + +func (s *Service) finalizeHomePluginWorkUntilDone(ctx, homeCtx context.Context, generation uint64, client *home.Client, work *homePluginFinalization, publish func() bool) error { + stopClose := closeHomeClientOnCancellation(ctx, client) + defer stopClose() + for { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + + s.homeOwnershipMu.Lock() + if !s.homeLifetimeActive(homeCtx, ctx, generation) { + s.homeOwnershipMu.Unlock() + return context.Canceled + } + errFinalize := s.finalizeHomePluginWork(ctx, client, work) + if errFinalize == nil && (publish == nil || publish()) { + s.homeOwnershipMu.Unlock() + return nil + } + s.homeOwnershipMu.Unlock() + if errFinalize == nil { + return context.Canceled + } + + log.WithError(errFinalize).Warn("failed to finalize home plugins; retrying") + timer := time.NewTimer(homeSubscriberPreAckRetryBackoff) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case <-timer.C: + } + } +} + +func closeHomeClientOnCancellation(ctx context.Context, client *home.Client) func() { + if ctx == nil || client == nil { + return func() {} + } + stop := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + client.Close() + case <-stop: + } + }() + return func() { close(stop) } +} + +func logHomeConfigChanges(oldCfg, newCfg *config.Config) { + if oldCfg == nil || newCfg == nil || !newCfg.Home.Enabled || (!oldCfg.Debug && !newCfg.Debug) { + return + } + + details := diff.BuildConfigChangeDetails(oldCfg, newCfg) + if len(details) == 0 { + return + } + + if newCfg.Debug && !log.IsLevelEnabled(log.DebugLevel) { + util.SetLogLevel(newCfg) + } + + log.Debugf("home config changes detected:") + for _, detail := range details { + log.Debugf(" %s", detail) + } +} + +func (s *Service) startHomeUsageForwarder(ctx context.Context, client *home.Client) { + if s == nil || client == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + + sleep := func(d time.Duration) bool { + if d <= 0 { + return true + } + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } + } + + go func() { + for { + select { + case <-ctx.Done(): + return + default: + } + + if !client.HeartbeatOK() { + if !sleep(time.Second) { + return + } + continue + } + + items := redisqueue.PopOldest(64) + if len(items) == 0 { + if !sleep(500 * time.Millisecond) { + return + } + continue + } + + for i := range items { + if errPush := client.LPushUsage(ctx, items[i]); errPush != nil { + for j := i; j < len(items); j++ { + redisqueue.Enqueue(items[j]) + } + if !sleep(time.Second) { + return + } + break + } + } + } + }() +} + +func applyHomeObservationBarrier(registry *executionregistry.Registry, revision int64) { + if registry != nil { + registry.ObserveBarrier(revision) + } +} + +func applyHomeInFlightPublisherConfig(manager *coreauth.Manager, cfg internalconfig.CredentialInFlightConfig) error { + publisherCfg, errConfig := coreauth.HomeInFlightPublisherConfigFromConfig(cfg) + if errConfig != nil { + return errConfig + } + if manager != nil { + manager.ApplyHomeInFlightPublisherConfig(publisherCfg) + } + return nil +} + +func (s *Service) startHomeSubscriber(ctx context.Context) { + if s == nil { + return + } + s.cfgMu.RLock() + cfg := s.cfg + s.cfgMu.RUnlock() + if cfg == nil || !cfg.Home.Enabled { + return + } + + parentCtx := ctx + if parentCtx == nil { + parentCtx = context.Background() + } + + s.homeLifecycleMu.Lock() + defer s.homeLifecycleMu.Unlock() + + if previousSupervisor := s.homeSupervisor; previousSupervisor != nil { + s.homeConfigCommitMu.Lock() + previousSupervisor.cancel() + s.homeConfigCommitMu.Unlock() + <-previousSupervisor.done + } + if !s.drainDetachedHomeLifetime(parentCtx) { + return + } + if parentCtx.Err() != nil { + return + } + + homeCtx, cancel := context.WithCancel(parentCtx) + done := make(chan struct{}) + s.homeMu.Lock() + s.homeGeneration++ + generation := s.homeGeneration + s.homeCancel = cancel + s.homeMu.Unlock() + supervisor := &homeSubscriberSupervisor{cancel: cancel, done: done} + s.homeSupervisor = supervisor + go s.runHomeSubscriber(homeCtx, parentCtx, cfg.Home, generation, supervisor) +} + +func (s *Service) drainDetachedHomeLifetime(parentCtx context.Context) bool { + s.homeMu.Lock() + previousCancel := s.homeCancel + previousClient := s.homeClient + previousRegistry := s.homeRegistry + previousBundle := s.homeDispatchBundle + previousDrainBound := s.homeDrainBound + previousForwarder := s.homeLogForwarder + previousForwarderClient := s.homeLogForwarderClient + s.homeCancel = nil + s.homeClient = nil + s.homeRegistry = nil + s.homeDispatchBundle = nil + s.homeDrainBound = 0 + s.homeLogForwarderClient = nil + s.homeMu.Unlock() + + if s.coreManager != nil { + s.coreManager.ClearHomeDispatchBundle(previousBundle) + } + home.ClearCurrentIf(previousClient) + if previousCancel != nil { + previousCancel() + } + if previousForwarder != nil && previousForwarderClient == previousClient { + previousForwarder.Deactivate(previousClient) + } + if previousRegistry != nil { + if previousDrainBound <= 0 { + previousDrainBound = internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound + } + drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), previousDrainBound) + errDrain := previousRegistry.Drain(drainCtx) + cancelDrain() + if errDrain != nil { + if previousClient != nil { + previousClient.Close() + } + if parentCtx.Err() == nil { + log.WithError(errDrain).Error("failed to drain replaced Home execution registry") + s.cancelServiceRun() + } + return false + } + } + if previousClient != nil { + previousClient.Close() + } + return true +} + +func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.Context, homeCfg internalconfig.HomeConfig, generation uint64, supervisor *homeSubscriberSupervisor) { + defer func() { + s.homeMu.Lock() + if s.homeGeneration == generation { + s.homeCancel = nil + } + s.homeMu.Unlock() + close(supervisor.done) + }() + + for homeCtx.Err() == nil { + supervisor.setPublisherCompletion(nil) + client := home.New(homeCfg) + client.SetManagedLifetime(true) + registry := executionregistry.New() + releaseCtx, releaseCancel := context.WithCancel(context.WithoutCancel(homeCtx)) + releaseFlusher := home.NewReleaseFlusher(client.LimiterConfig, client.PushConcurrencyRelease) + registry.SetReleaseSink(releaseFlusher.MarkDirty) + releaseDone := make(chan struct{}) + go func() { + defer close(releaseDone) + releaseFlusher.Run(releaseCtx) + }() + lifetimeCtx, lifetimeCancel := context.WithCancel(homeCtx) + cancelBound := atomic.Int64{} + cancelBound.Store(int64(internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound)) + queue := newHomeConfigWorkQueue() + ready := make(chan struct{}) + var readyOnce sync.Once + var published atomic.Bool + workerDone := make(chan struct{}) + + go func() { + defer close(workerDone) + s.runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx, generation, client, registry, queue, ready, &published, &cancelBound, supervisor) + }() + + errRun := client.RunConfigSubscriberLifetime(lifetimeCtx, func(raw []byte) error { + parsed, errParse := config.ParseConfigBytes(raw) + if errParse != nil { + log.Warnf("failed to parse home config payload: %v", errParse) + return errParse + } + if errSetLifecycle := client.SetLifecycleConfig(parsed.CredentialConcurrency); errSetLifecycle != nil { + log.Warnf("failed to apply Home lifecycle config: %v", errSetLifecycle) + return errSetLifecycle + } + if errPublisherConfig := applyHomeInFlightPublisherConfig(s.coreManager, parsed.CredentialInFlight); errPublisherConfig != nil { + log.Warnf("failed to apply Home in-flight publisher config: %v", errPublisherConfig) + return errPublisherConfig + } + applyHomeObservationBarrier(registry, parsed.CredentialConcurrency.ObservationBarrierRevision) + cancelBound.Store(int64(parsed.CredentialConcurrency.WithDefaults().CPACancelBound)) + queue.enqueue(raw) + return nil + }, func() { + readyOnce.Do(func() { close(ready) }) + }) + lifetimeCancel() + <-workerDone + if publisherDone := supervisor.publisherCompletion(); publisherDone != nil { + <-publisherDone + } + + s.detachHomeSubscriberLifetime(client, registry) + drainBound := time.Duration(cancelBound.Load()) + drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), drainBound) + errDrain := registry.Drain(drainCtx) + var errFlush error + if errDrain == nil { + errFlush = releaseFlusher.Flush(drainCtx) + } + cancelDrain() + releaseCancel() + <-releaseDone + client.Close() + if errDrain != nil { + if parentCtx.Err() == nil { + log.WithError(errDrain).Error("failed to drain Home execution registry") + s.cancelServiceRun() + } + return + } + if errFlush != nil { + if parentCtx.Err() == nil { + log.WithError(errFlush).Error("failed to flush Home concurrency releases") + s.cancelServiceRun() + } + return + } + if errRun != nil && homeCtx.Err() == nil { + log.WithError(errRun).Warn("home config subscription lifetime ended") + } + if !published.Load() && errRun != nil && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) { + return + } + } +} + +func (s *Service) runHomeConfigWorker(lifetimeCtx, homeCtx context.Context, generation uint64, client *home.Client, registry *executionregistry.Registry, queue *homeConfigWorkQueue, ready <-chan struct{}, published *atomic.Bool, cancelBound *atomic.Int64) { + s.runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx, generation, client, registry, queue, ready, published, cancelBound, nil) +} + +func (s *Service) runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx context.Context, generation uint64, client *home.Client, registry *executionregistry.Registry, queue *homeConfigWorkQueue, ready <-chan struct{}, published *atomic.Bool, cancelBound *atomic.Int64, supervisor *homeSubscriberSupervisor) { + select { + case <-lifetimeCtx.Done(): + return + case <-ready: + } + + for { + if lifetimeCtx.Err() != nil { + return + } + raw, ok := queue.dequeue(lifetimeCtx) + if !ok { + return + } + if lifetimeCtx.Err() != nil { + return + } + + var work *homePluginFinalization + for { + if lifetimeCtx.Err() != nil { + return + } + parsed, errParse := config.ParseConfigBytes(raw) + if errParse == nil { + work, errParse = s.stageHomeOverlayWithClient(lifetimeCtx, parsed, client) + } + if errParse == nil { + break + } + if lifetimeCtx.Err() != nil { + return + } + log.WithError(errParse).Warn("failed to stage home config; retrying") + if !waitForHomeSubscriberRetry(lifetimeCtx, homeSubscriberPreAckRetryBackoff) { + return + } + } + + var publish func() bool + if !published.Load() { + publish = func() bool { + s.homeMu.Lock() + defer s.homeMu.Unlock() + if homeCtx.Err() != nil || lifetimeCtx.Err() != nil || s.homeGeneration != generation { + return false + } + s.homeClient = client + s.homeRegistry = registry + s.homeDrainBound = time.Duration(cancelBound.Load()) + if s.coreManager != nil { + s.homeDispatchBundle = s.coreManager.PublishHomeDispatch(client, registry, generation) + } + home.SetCurrent(client) + if s.homeLogForwarder == nil { + s.homeLogForwarder = startHomeLogForwarder(0) + } + s.homeLogForwarder.Bind(client) + s.homeLogForwarderClient = client + published.Store(true) + return true + } + } + if s.homeConfigStageHook != nil { + s.homeConfigStageHook() + } + if !s.commitHomeConfig(lifetimeCtx, homeCtx, generation, work) { + return + } + if s.homeConfigRuntimeHook != nil { + s.homeConfigRuntimeHook() + } + if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) || !s.applyConfigRuntime(lifetimeCtx, work.configCommit, true) { + return + } + if errFinalize := s.finalizeHomePluginWorkUntilDone(lifetimeCtx, homeCtx, generation, client, work, publish); errFinalize != nil { + if !errors.Is(errFinalize, context.Canceled) { + log.WithError(errFinalize).Warn("home plugin finalization ended") + } + return + } + if publish != nil { + s.startHomeInFlightPublisher(lifetimeCtx, client, registry, supervisor) + s.startHomeUsageForwarder(lifetimeCtx, client) + } + } +} + +func (s *Service) startHomeInFlightPublisher(ctx context.Context, client *home.Client, registry *executionregistry.Registry, supervisor *homeSubscriberSupervisor) { + if s == nil || s.coreManager == nil { + return + } + done := make(chan struct{}) + if supervisor != nil { + supervisor.setPublisherCompletion(done) + } + go func() { + defer close(done) + s.coreManager.StartHomeInFlightPublisher(ctx, client, registry) + }() +} + +func waitForHomeSubscriberRetry(ctx context.Context, delay time.Duration) bool { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func (s *Service) detachHomeSubscriberLifetime(client *home.Client, registry *executionregistry.Registry) { + if s == nil { + return + } + s.homeMu.Lock() + var bundle *coreauth.HomeDispatchBundle + if s.homeClient == client && s.homeRegistry == registry { + bundle = s.homeDispatchBundle + s.homeClient = nil + s.homeRegistry = nil + s.homeDispatchBundle = nil + s.homeDrainBound = 0 + } + forwarder := s.homeLogForwarder + if s.homeLogForwarderClient == client { + s.homeLogForwarderClient = nil + } else { + forwarder = nil + } + s.homeMu.Unlock() + if s.coreManager != nil { + s.coreManager.ClearHomeDispatchBundle(bundle) + } + home.ClearCurrentIf(client) + if forwarder != nil { + forwarder.Deactivate(client) + } +} + +func (s *Service) cancelServiceRun() { + if s == nil { + return + } + s.homeMu.Lock() + cancel := s.runCancel + if cancel == nil { + cancel = s.homeCancel + } + s.homeMu.Unlock() + if cancel != nil { + cancel() + } +} diff --git a/sdk/cliproxy/service_lifecycle.go b/sdk/cliproxy/service_lifecycle.go new file mode 100644 index 000000000..e16b1433d --- /dev/null +++ b/sdk/cliproxy/service_lifecycle.go @@ -0,0 +1,369 @@ +package cliproxy + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/api" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" + sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" +) + +// Run starts the service and blocks until the context is cancelled or the server stops. +// It initializes all components including authentication, file watching, HTTP server, +// and starts processing requests. The method blocks until the context is cancelled. +// +// Parameters: +// - ctx: The context for controlling the service lifecycle +// +// Returns: +// - error: An error if the service fails to start or run +func (s *Service) Run(ctx context.Context) error { + if s == nil { + return fmt.Errorf("cliproxy: service is nil") + } + if ctx == nil { + ctx = context.Background() + } + ctx, runCancel := context.WithCancel(ctx) + s.homeMu.Lock() + s.runCancel = runCancel + s.homeMu.Unlock() + defer func() { + runCancel() + s.homeMu.Lock() + if s.runCancel != nil { + s.runCancel = nil + } + s.homeMu.Unlock() + }() + + usage.StartDefault(ctx) + homeEnabled := s.cfg != nil && s.cfg.Home.Enabled + if homeEnabled { + forceHomeRuntimeConfig(s.cfg) + redisqueue.SetUsageStatisticsEnabled(true) + } + + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) + defer shutdownCancel() + defer func() { + if err := s.Shutdown(shutdownCtx); err != nil { + log.Errorf("service shutdown returned error: %v", err) + } + }() + + if !homeEnabled { + if errEnsureAuthDir := s.ensureAuthDir(); errEnsureAuthDir != nil { + return errEnsureAuthDir + } + } + + s.applyRetryConfig(s.cfg) + s.configureCooldownStateStore(s.cfg) + + s.registerPluginAuthParser() + if s.coreManager != nil && !homeEnabled { + if errLoad := s.coreManager.Load(ctx); errLoad != nil { + log.Warnf("failed to load auth store: %v", errLoad) + } + s.registerConfigAPIKeyAuths(coreauth.WithSkipPersist(ctx), s.cfg) + if s.cfg.SaveCooldownStatus { + if errRestoreCooldown := s.coreManager.RestoreCooldownStates(ctx); errRestoreCooldown != nil { + log.Warnf("failed to restore cooldown state: %v", errRestoreCooldown) + } + } + } + + if !homeEnabled { + tokenResult, err := s.tokenProvider.Load(ctx, s.cfg) + if err != nil && !errors.Is(err, context.Canceled) { + return err + } + if tokenResult == nil { + tokenResult = &TokenClientResult{} + } + + apiKeyResult, err := s.apiKeyProvider.Load(ctx, s.cfg) + if err != nil && !errors.Is(err, context.Canceled) { + return err + } + if apiKeyResult == nil { + apiKeyResult = &APIKeyClientResult{} + } + } + + // legacy clients removed; no caches to refresh + + s.ensureWebsocketGateway() + if homeEnabled { + s.registerAvailableExecutors(ctx, executorRegistrationOptions{ + includeBaseline: true, + }) + // Home mode does not expose in-process Redis RESP usage output; usage is forwarded to home instead. + redisqueue.SetEnabled(true) + } + + // handlers no longer depend on legacy clients; pass nil slice initially + s.server = api.NewServer(s.cfg, s.coreManager, s.accessManager, s.configPath, s.serverOptions...) + s.syncPluginRuntimeConfig(ctx) + if homeEnabled { + s.syncPluginModelRuntime(ctx) + } + + if s.authManager == nil { + s.authManager = newDefaultAuthManager() + } + + if homeEnabled { + s.startHomeSubscriber(ctx) + } + + if s.server != nil && s.wsGateway != nil { + s.server.AttachWebsocketRoute(s.wsGateway.Path(), s.wsGateway.Handler()) + s.server.SetWebsocketAuthChangeHandler(func(oldEnabled, newEnabled bool) { + if oldEnabled == newEnabled { + return + } + if !oldEnabled && newEnabled { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if errStop := s.wsGateway.Stop(ctx); errStop != nil { + log.Warnf("failed to reset websocket connections after ws-auth change %t -> %t: %v", oldEnabled, newEnabled, errStop) + return + } + log.Debugf("ws-auth enabled; existing websocket sessions terminated to enforce authentication") + return + } + log.Debugf("ws-auth disabled; existing websocket sessions remain connected") + }) + } + + if s.hooks.OnBeforeStart != nil { + s.hooks.OnBeforeStart(s.cfg) + } + + s.serverErr = make(chan error, 1) + go func() { + if errStart := s.server.Start(); errStart != nil { + s.serverErr <- errStart + } else { + s.serverErr <- nil + } + }() + + time.Sleep(100 * time.Millisecond) + fmt.Printf("API server started successfully on: %s:%d\n", s.cfg.Host, s.cfg.Port) + + s.applyPprofConfig(s.cfg) + + if s.hooks.OnAfterStart != nil { + s.hooks.OnAfterStart(s) + } + + if !homeEnabled { + var watcherWrapper *WatcherWrapper + reloadCallback := func(newCfg *config.Config) { s.applyWatcherConfigUpdate(newCfg) } + + watcherWrapper, errCreate := s.watcherFactory(s.configPath, s.cfg.AuthDir, reloadCallback) + if errCreate != nil { + return fmt.Errorf("cliproxy: failed to create watcher: %w", errCreate) + } + s.watcher = watcherWrapper + s.ensureAuthUpdateQueue(ctx) + if s.authUpdates != nil { + watcherWrapper.SetAuthUpdateQueue(s.authUpdates) + } + watcherWrapper.SetConfig(s.cfg) + s.registerPluginAuthParser() + + watcherCtx, watcherCancel := context.WithCancel(context.Background()) + s.watcherCancel = watcherCancel + if errStart := watcherWrapper.Start(watcherCtx); errStart != nil { + return fmt.Errorf("cliproxy: failed to start watcher: %w", errStart) + } + log.Info("file watcher started for config and auth directory changes") + s.syncPluginModelRuntime(ctx) + } + + s.registerModelRefreshCallback() + + // Prefer core auth manager auto refresh if available. + if s.coreManager != nil && !homeEnabled { + interval := 15 * time.Minute + s.coreManager.StartAutoRefresh(context.Background(), interval) + log.Infof("core auth auto-refresh started (interval=%s)", interval) + } + + select { + case <-ctx.Done(): + log.Debug("service context cancelled, shutting down...") + return ctx.Err() + case errServer := <-s.serverErr: + return errServer + } +} + +// Shutdown gracefully stops background workers and the HTTP server. +// It ensures all resources are properly cleaned up and connections are closed. +// The shutdown is idempotent and can be called multiple times safely. +// +// Parameters: +// - ctx: The context for controlling the shutdown timeout +// +// Returns: +// - error: An error if shutdown fails +func (s *Service) Shutdown(ctx context.Context) error { + if s == nil { + return nil + } + var shutdownErr error + s.shutdownOnce.Do(func() { + if ctx == nil { + ctx = context.Background() + } + + s.homeLifecycleMu.Lock() + if supervisor := s.homeSupervisor; supervisor != nil { + s.homeConfigCommitMu.Lock() + supervisor.cancel() + s.homeConfigCommitMu.Unlock() + <-supervisor.done + } + s.homeMu.Lock() + homeCancel := s.homeCancel + homeClient := s.homeClient + homeRegistry := s.homeRegistry + homeDispatchBundle := s.homeDispatchBundle + homeForwarder := s.homeLogForwarder + homeForwarderClient := s.homeLogForwarderClient + s.homeGeneration++ + s.homeCancel = nil + s.homeClient = nil + s.homeRegistry = nil + s.homeDispatchBundle = nil + s.homeDrainBound = 0 + s.homeLogForwarder = nil + s.homeLogForwarderClient = nil + s.homeMu.Unlock() + if s.coreManager != nil { + s.coreManager.ClearHomeDispatchBundle(homeDispatchBundle) + } + home.ClearCurrentIf(homeClient) + if homeCancel != nil { + homeCancel() + } + if homeRegistry != nil { + if errClose := homeRegistry.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close Home execution registry during shutdown") + } + } + if homeClient != nil { + homeClient.Close() + } + if homeForwarder != nil { + if homeForwarderClient == homeClient { + homeForwarder.Deactivate(homeClient) + } + homeForwarder.Stop() + } + s.homeLifecycleMu.Unlock() + + // legacy refresh loop removed; only stopping core auth manager below + + if s.watcherCancel != nil { + s.watcherCancel() + } + if s.coreManager != nil { + s.coreManager.StopAutoRefresh() + } + if s.watcher != nil { + if err := s.watcher.Stop(); err != nil { + log.Errorf("failed to stop file watcher: %v", err) + shutdownErr = err + } + } + if s.wsGateway != nil { + if err := s.wsGateway.Stop(ctx); err != nil { + log.Errorf("failed to stop websocket gateway: %v", err) + if shutdownErr == nil { + shutdownErr = err + } + } + } + if s.authQueueStop != nil { + s.authQueueStop() + s.authQueueStop = nil + } + + if errShutdownPprof := s.shutdownPprof(ctx); errShutdownPprof != nil { + log.Errorf("failed to stop pprof server: %v", errShutdownPprof) + if shutdownErr == nil { + shutdownErr = errShutdownPprof + } + } + + // no legacy clients to persist + + if s.server != nil { + shutdownCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + if err := s.server.Stop(shutdownCtx); err != nil { + log.Errorf("error stopping API server: %v", err) + if shutdownErr == nil { + shutdownErr = err + } + } + } + + if s.pluginHost != nil { + sdktranslator.SetPluginHooks(nil) + sdkAuth.RegisterPluginAuthParser(nil) + if s.watcher != nil { + s.watcher.SetPluginAuthParser(nil) + } + s.pluginHost.ApplyConfig(ctx, &config.Config{}) + s.pluginHost.RegisterModels(ctx, registry.GetGlobalRegistry()) + s.registerAvailableExecutors(ctx, executorRegistrationOptions{ + includePlugins: true, + }) + s.pluginHost.RegisterFrontendAuthProviders() + s.pluginHost.ShutdownAllContext(ctx) + if s.accessManager != nil { + s.accessManager.SetProviders(sdkaccess.RegisteredProviders()) + } + } + + usage.StopDefault() + }) + return shutdownErr +} + +func (s *Service) ensureAuthDir() error { + info, err := os.Stat(s.cfg.AuthDir) + if err != nil { + if os.IsNotExist(err) { + if mkErr := os.MkdirAll(s.cfg.AuthDir, 0o755); mkErr != nil { + return fmt.Errorf("cliproxy: failed to create auth directory %s: %w", s.cfg.AuthDir, mkErr) + } + log.Infof("created missing auth directory: %s", s.cfg.AuthDir) + return nil + } + return fmt.Errorf("cliproxy: error checking auth directory %s: %w", s.cfg.AuthDir, err) + } + if !info.IsDir() { + return fmt.Errorf("cliproxy: auth path exists but is not a directory: %s", s.cfg.AuthDir) + } + return nil +} diff --git a/sdk/cliproxy/service_models.go b/sdk/cliproxy/service_models.go new file mode 100644 index 000000000..f54aa069b --- /dev/null +++ b/sdk/cliproxy/service_models.go @@ -0,0 +1,987 @@ +package cliproxy + +import ( + "context" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/constant" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" +) + +// registerModelsForAuth (re)binds provider models in the global registry using the core auth ID as client identifier. +func (s *Service) registerModelsForAuth(ctx context.Context, a *coreauth.Auth) { + s.registerModelsForAuthWithCache(ctx, a, nil) +} + +func (s *Service) registerModelsForAuthWithCache(ctx context.Context, a *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) { + if a == nil || a.ID == "" { + return + } + if ctx == nil { + ctx = context.Background() + } + if ctx.Err() != nil { + return + } + if a.Disabled { + GlobalModelRegistry().UnregisterClient(a.ID) + return + } + authKind := a.AuthKind() + // Unregister legacy client ID (if present) to avoid double counting + if a.Runtime != nil { + if idGetter, ok := a.Runtime.(interface{ GetClientID() string }); ok { + if rid := idGetter.GetClientID(); rid != "" && rid != a.ID { + GlobalModelRegistry().UnregisterClient(rid) + } + } + } + provider := strings.ToLower(strings.TrimSpace(a.Provider)) + compatProviderKey, compatDisplayName, compatDetected := openAICompatInfoFromAuth(a) + if compatDetected { + provider = "openai-compatibility" + } + excluded := s.oauthExcludedModels(provider, authKind) + // The synthesizer pre-merges per-account and global exclusions into the "excluded_models" attribute. + // If this attribute is present, it represents the complete list of exclusions and overrides the global config. + if a.Attributes != nil { + if val, ok := a.Attributes["excluded_models"]; ok && strings.TrimSpace(val) != "" { + excluded = strings.Split(val, ",") + } + } + if s.tryRegisterPluginModelsForAuth(ctx, a, provider, authKind, excluded) { + return + } + if ctx.Err() != nil { + return + } + var models []*ModelInfo + switch provider { + case constant.Gemini: + models = registry.GetGeminiModels() + if entry := s.resolveConfigGeminiKey(a); entry != nil { + if len(entry.Models) > 0 { + models = buildGeminiConfigModels(entry) + } + if authKind == "apikey" { + excluded = entry.ExcludedModels + } + } + models = applyExcludedModels(models, excluded) + case constant.GeminiInteractions: + models = registry.GetGeminiModels() + if entry := s.resolveConfigInteractionsKey(a); entry != nil { + if len(entry.Models) > 0 { + models = buildGeminiConfigModels(entry) + } + if authKind == "apikey" { + excluded = entry.ExcludedModels + } + } + models = applyExcludedModels(models, excluded) + case "vertex": + // Vertex AI Gemini supports the same model identifiers as Gemini. + models = registry.GetGeminiVertexModels() + if entry := s.resolveConfigVertexCompatKey(a); entry != nil { + if len(entry.Models) > 0 { + models = buildVertexCompatConfigModels(entry) + } + if authKind == "apikey" { + excluded = entry.ExcludedModels + } + } + models = applyExcludedModels(models, excluded) + case "aistudio": + models = registry.GetAIStudioModels() + models = applyExcludedModels(models, excluded) + case "antigravity": + models = registry.GetAntigravityModels() + models = applyAntigravityFetchedModelCapabilities(models, s.fetchAntigravityModelCapabilityHintsForAuth(ctx, a)) + models = applyExcludedModels(models, excluded) + case "claude": + models = registry.GetClaudeModels() + if entry := s.resolveConfigClaudeKey(a); entry != nil { + if len(entry.Models) > 0 { + models = buildClaudeConfigModels(entry) + } + if authKind == "apikey" { + excluded = entry.ExcludedModels + } + } + models = applyExcludedModels(models, excluded) + case "codex": + codexPlanType := "" + if a.Attributes != nil { + codexPlanType = strings.TrimSpace(a.Attributes["plan_type"]) + } + switch strings.ToLower(codexPlanType) { + case "pro": + models = registry.GetCodexProModels() + case "plus": + models = registry.GetCodexPlusModels() + case "team", "business", "go": + models = registry.GetCodexTeamModels() + case "free": + models = registry.GetCodexFreeModels() + default: + models = registry.GetCodexProModels() + } + if entry := s.resolveConfigCodexKey(a); entry != nil { + if len(entry.Models) > 0 { + models = buildCodexConfigModels(entry) + } + if authKind == "apikey" { + excluded = entry.ExcludedModels + } + } + models = applyExcludedModels(models, excluded) + case "kimi": + models = registry.GetKimiModels() + models = applyExcludedModels(models, excluded) + case "xai": + models = registry.GetXAIModels() + if entry := s.resolveConfigXAIKey(a); entry != nil { + if len(entry.Models) > 0 { + models = buildXAIConfigModels(entry) + } + if authKind == "apikey" { + excluded = entry.ExcludedModels + } + } + models = applyExcludedModels(models, excluded) + default: + // Handle OpenAI-compatibility providers by name using config + if s.cfg != nil { + providerKey := provider + compatName := strings.TrimSpace(a.Provider) + isCompatAuth := false + if compatDetected { + if compatProviderKey != "" { + providerKey = compatProviderKey + } + if compatDisplayName != "" { + compatName = compatDisplayName + } + isCompatAuth = true + } + if strings.EqualFold(providerKey, "openai-compatibility") { + isCompatAuth = true + if a.Attributes != nil { + if v := strings.TrimSpace(a.Attributes["compat_name"]); v != "" { + compatName = v + } + if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" { + providerKey = strings.ToLower(v) + isCompatAuth = true + } + } + if providerKey == "openai-compatibility" && compatName != "" { + providerKey = strings.ToLower(compatName) + } + } else if a.Attributes != nil { + if v := strings.TrimSpace(a.Attributes["compat_name"]); v != "" { + compatName = v + isCompatAuth = true + } + if v := strings.TrimSpace(a.Attributes["provider_key"]); v != "" { + providerKey = strings.ToLower(v) + isCompatAuth = true + } + } + if cached, ok := compatCache.lookup(compatName); ok { + isCompatAuth = true + if providerKey == "" { + providerKey = cached.providerKey + } + if providerKey == "" { + providerKey = "openai-compatibility" + } + ms := cached.models + if len(ms) > 0 { + ms = s.appendPluginModels(providerKey, ms) + s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) + } else { + ms = s.appendPluginModels(providerKey, nil) + if len(ms) > 0 { + s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) + } else { + GlobalModelRegistry().UnregisterClient(a.ID) + } + } + return + } + for i := range s.cfg.OpenAICompatibility { + compat := &s.cfg.OpenAICompatibility[i] + if compat.Disabled { + continue + } + if strings.EqualFold(compat.Name, compatName) { + isCompatAuth = true + ms := buildOpenAICompatibilityConfigModels(compat) + // Register and return + if len(ms) > 0 { + if providerKey == "" { + providerKey = "openai-compatibility" + } + ms = s.appendPluginModels(providerKey, ms) + s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) + } else { + // Ensure stale registrations are cleared when model list becomes empty. + ms = s.appendPluginModels(providerKey, nil) + if len(ms) > 0 { + s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(ms, a.Prefix, s.cfg.ForceModelPrefix)) + } else { + GlobalModelRegistry().UnregisterClient(a.ID) + } + } + return + } + } + if isCompatAuth { + models = s.appendPluginModels(providerKey, nil) + if len(models) > 0 { + s.registerResolvedModelsForAuth(a, providerKey, applyModelPrefixes(models, a.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix)) + } else { + // No matching provider found or models removed entirely; drop any prior registration. + GlobalModelRegistry().UnregisterClient(a.ID) + } + return + } + } + } + if ctx.Err() != nil { + return + } + models = applyOAuthModelAliasForAuth(s.cfg, provider, authKind, a.Attributes, models) + if ctx.Err() != nil { + return + } + key := provider + if key == "" { + key = strings.ToLower(strings.TrimSpace(a.Provider)) + } + models = s.appendPluginModels(key, models) + if len(models) > 0 { + s.registerResolvedModelsForAuth(a, key, applyModelPrefixes(models, a.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix)) + return + } + + GlobalModelRegistry().UnregisterClient(a.ID) +} + +// refreshModelRegistrationForAuth re-applies the latest model registration for +// one auth and reconciles any concurrent auth changes that race with the +// refresh. Callers are expected to pre-filter provider membership. +// +// Re-registration is deliberate: registry cooldown/suspension state is treated +// as part of the previous registration snapshot and is cleared when the auth is +// rebound to the refreshed model catalog. +func (s *Service) refreshModelRegistrationForAuth(current *coreauth.Auth) bool { + return s.refreshModelRegistrationForAuthWithContext(context.Background(), current, nil) +} + +func (s *Service) refreshModelRegistrationForAuthWithCache(current *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) bool { + return s.refreshModelRegistrationForAuthWithContext(context.Background(), current, compatCache) +} + +func (s *Service) refreshModelRegistrationForAuthWithContext(ctx context.Context, current *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) bool { + if s == nil || s.coreManager == nil || current == nil || current.ID == "" { + return false + } + if ctx == nil { + ctx = context.Background() + } + if ctx.Err() != nil { + return false + } + if !current.Disabled { + s.ensureExecutorsForAuthWithContext(ctx, current, false) + } + s.registerModelsForAuthWithCache(ctx, current, compatCache) + s.coreManager.ReconcileRegistryModelStates(ctx, current.ID) + if ctx.Err() != nil { + return false + } + + latest, ok := s.latestAuthForModelRegistration(current.ID) + if !ok || latest.Disabled { + GlobalModelRegistry().UnregisterClient(current.ID) + s.coreManager.RefreshSchedulerEntry(current.ID) + return false + } + + // Re-apply the latest auth snapshot so concurrent auth updates cannot leave + // stale model registrations behind. This may duplicate registration work when + // no auth fields changed, but keeps the refresh path simple and correct. + s.ensureExecutorsForAuthWithContext(ctx, latest, false) + s.registerModelsForAuthWithCache(ctx, latest, compatCache) + if ctx.Err() != nil { + return false + } + s.coreManager.ReconcileRegistryModelStates(ctx, latest.ID) + s.coreManager.RefreshSchedulerEntry(current.ID) + return true +} + +// latestAuthForModelRegistration returns the latest auth snapshot regardless of +// provider membership. Callers use this after a registration attempt to restore +// whichever state currently owns the client ID in the global registry. +func (s *Service) latestAuthForModelRegistration(authID string) (*coreauth.Auth, bool) { + if s == nil || s.coreManager == nil || authID == "" { + return nil, false + } + auth, ok := s.coreManager.GetByID(authID) + if !ok || auth == nil || auth.ID == "" { + return nil, false + } + return auth, true +} + +func (s *Service) resolveConfigClaudeKey(auth *coreauth.Auth) *config.ClaudeKey { + if auth == nil || s.cfg == nil { + return nil + } + var attrKey, attrBase string + if auth.Attributes != nil { + attrKey = strings.TrimSpace(auth.Attributes["api_key"]) + attrBase = strings.TrimSpace(auth.Attributes["base_url"]) + } + for i := range s.cfg.ClaudeKey { + entry := &s.cfg.ClaudeKey[i] + cfgKey := strings.TrimSpace(entry.APIKey) + cfgBase := strings.TrimSpace(entry.BaseURL) + if attrKey != "" && attrBase != "" { + if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) { + return entry + } + continue + } + if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { + if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey != "" { + for i := range s.cfg.ClaudeKey { + entry := &s.cfg.ClaudeKey[i] + if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { + return entry + } + } + } + return nil +} + +func (s *Service) resolveConfigGeminiKey(auth *coreauth.Auth) *config.GeminiKey { + if s == nil || s.cfg == nil { + return nil + } + return s.resolveConfigGeminiKeyEntry(auth, s.cfg.GeminiKey) +} + +func (s *Service) resolveConfigInteractionsKey(auth *coreauth.Auth) *config.GeminiKey { + if s == nil || s.cfg == nil { + return nil + } + return s.resolveConfigGeminiKeyEntry(auth, s.cfg.InteractionsKey) +} + +func (s *Service) resolveConfigGeminiKeyEntry(auth *coreauth.Auth, entries []config.GeminiKey) *config.GeminiKey { + if auth == nil || s.cfg == nil { + return nil + } + var attrKey, attrBase string + if auth.Attributes != nil { + attrKey = strings.TrimSpace(auth.Attributes["api_key"]) + attrBase = strings.TrimSpace(auth.Attributes["base_url"]) + } + for i := range entries { + entry := &entries[i] + cfgKey := strings.TrimSpace(entry.APIKey) + cfgBase := strings.TrimSpace(entry.BaseURL) + if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { + if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { + return entry + } + continue + } + if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + return nil +} + +func (s *Service) resolveConfigVertexCompatKey(auth *coreauth.Auth) *config.VertexCompatKey { + if auth == nil || s.cfg == nil { + return nil + } + var attrKey, attrBase string + if auth.Attributes != nil { + attrKey = strings.TrimSpace(auth.Attributes["api_key"]) + attrBase = strings.TrimSpace(auth.Attributes["base_url"]) + } + for i := range s.cfg.VertexCompatAPIKey { + entry := &s.cfg.VertexCompatAPIKey[i] + cfgKey := strings.TrimSpace(entry.APIKey) + cfgBase := strings.TrimSpace(entry.BaseURL) + if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { + if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { + return entry + } + continue + } + if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + if attrKey != "" { + for i := range s.cfg.VertexCompatAPIKey { + entry := &s.cfg.VertexCompatAPIKey[i] + if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) { + return entry + } + } + } + return nil +} + +func (s *Service) resolveConfigCodexKey(auth *coreauth.Auth) *config.CodexKey { + if s == nil || s.cfg == nil { + return nil + } + return resolveConfigCodexStyleKey(auth, s.cfg.CodexKey) +} + +func (s *Service) resolveConfigXAIKey(auth *coreauth.Auth) *config.XAIKey { + if s == nil || s.cfg == nil { + return nil + } + return resolveConfigCodexStyleKey(auth, s.cfg.XAIKey) +} + +func resolveConfigCodexStyleKey(auth *coreauth.Auth, entries []config.CodexKey) *config.CodexKey { + if auth == nil { + return nil + } + var attrKey, attrBase string + if auth.Attributes != nil { + attrKey = strings.TrimSpace(auth.Attributes["api_key"]) + attrBase = strings.TrimSpace(auth.Attributes["base_url"]) + } + for i := range entries { + entry := &entries[i] + cfgKey := strings.TrimSpace(entry.APIKey) + cfgBase := strings.TrimSpace(entry.BaseURL) + if attrKey != "" && strings.EqualFold(cfgKey, attrKey) { + if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) { + return entry + } + continue + } + if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) { + return entry + } + } + return nil +} + +func (s *Service) oauthExcludedModels(provider, authKind string) []string { + cfg := s.cfg + if cfg == nil { + return nil + } + authKindKey := strings.ToLower(strings.TrimSpace(authKind)) + providerKey := strings.ToLower(strings.TrimSpace(provider)) + if authKindKey == "apikey" { + return nil + } + return cfg.OAuthExcludedModels[providerKey] +} + +func applyExcludedModels(models []*ModelInfo, excluded []string) []*ModelInfo { + if len(models) == 0 || len(excluded) == 0 { + return models + } + + patterns := make([]string, 0, len(excluded)) + for _, item := range excluded { + if trimmed := strings.TrimSpace(item); trimmed != "" { + patterns = append(patterns, strings.ToLower(trimmed)) + } + } + if len(patterns) == 0 { + return models + } + + filtered := make([]*ModelInfo, 0, len(models)) + for _, model := range models { + if model == nil { + continue + } + modelID := strings.ToLower(strings.TrimSpace(model.ID)) + blocked := false + for _, pattern := range patterns { + if matchWildcard(pattern, modelID) { + blocked = true + break + } + } + if !blocked { + filtered = append(filtered, model) + } + } + return filtered +} + +func applyModelPrefixes(models []*ModelInfo, prefix string, forceModelPrefix bool) []*ModelInfo { + trimmedPrefix := strings.TrimSpace(prefix) + if trimmedPrefix == "" || len(models) == 0 { + return models + } + + out := make([]*ModelInfo, 0, len(models)*2) + seen := make(map[string]struct{}, len(models)*2) + + addModel := func(model *ModelInfo) { + if model == nil { + return + } + id := strings.TrimSpace(model.ID) + if id == "" { + return + } + if _, exists := seen[id]; exists { + return + } + seen[id] = struct{}{} + out = append(out, model) + } + + for _, model := range models { + if model == nil { + continue + } + baseID := strings.TrimSpace(model.ID) + if baseID == "" { + continue + } + if !forceModelPrefix || trimmedPrefix == baseID { + addModel(model) + } + clone := *model + clone.ID = trimmedPrefix + "/" + baseID + addModel(&clone) + } + return out +} + +// matchWildcard performs case-insensitive wildcard matching where '*' matches any substring. +func matchWildcard(pattern, value string) bool { + if pattern == "" { + return false + } + + // Fast path for exact match (no wildcard present). + if !strings.Contains(pattern, "*") { + return pattern == value + } + + parts := strings.Split(pattern, "*") + // Handle prefix. + if prefix := parts[0]; prefix != "" { + if !strings.HasPrefix(value, prefix) { + return false + } + value = value[len(prefix):] + } + + // Handle suffix. + if suffix := parts[len(parts)-1]; suffix != "" { + if !strings.HasSuffix(value, suffix) { + return false + } + value = value[:len(value)-len(suffix)] + } + + // Handle middle segments in order. + for i := 1; i < len(parts)-1; i++ { + segment := parts[i] + if segment == "" { + continue + } + idx := strings.Index(value, segment) + if idx < 0 { + return false + } + value = value[idx+len(segment):] + } + + return true +} + +type modelEntry interface { + GetName() string + GetAlias() string + GetDisplayName() string +} + +func buildConfiguredModelInfo(model modelEntry, ownedBy, modelType string, created int64, fallbackDisplayName string, userDefined bool) *ModelInfo { + name := strings.TrimSpace(model.GetName()) + alias := strings.TrimSpace(model.GetAlias()) + if alias == "" { + alias = name + } + if alias == "" { + return nil + } + displayName := strings.TrimSpace(model.GetDisplayName()) + if displayName == "" { + displayName = fallbackDisplayName + } + if displayName == "" { + displayName = alias + } + return &ModelInfo{ + ID: alias, + Object: "model", + Created: created, + OwnedBy: ownedBy, + Type: modelType, + DisplayName: displayName, + UserDefined: userDefined, + } +} + +func buildOpenAICompatibilityConfigModels(compat *config.OpenAICompatibility) []*ModelInfo { + if compat == nil || len(compat.Models) == 0 { + return nil + } + now := time.Now().Unix() + models := make([]*ModelInfo, 0, len(compat.Models)) + for i := range compat.Models { + model := compat.Models[i] + modelType := "openai-compatibility" + if model.Image { + modelType = registry.OpenAIImageModelType + } + info := buildConfiguredModelInfo(model, compat.Name, modelType, now, strings.TrimSpace(model.Alias), false) + if info == nil { + continue + } + thinking := model.Thinking + if thinking == nil && !model.Image { + thinking = ®istry.ThinkingSupport{Levels: []string{"low", "medium", "high"}} + } + info.Thinking = thinking + info.SupportedInputModalities = normalizeCompatConfigModalities(model.InputModalities) + info.SupportedOutputModalities = normalizeCompatConfigModalities(model.OutputModalities) + models = append(models, info) + } + return models +} + +func normalizeCompatConfigModalities(raw []string) []string { + if len(raw) == 0 { + return nil + } + out := make([]string, 0, len(raw)) + seen := make(map[string]struct{}, len(raw)) + for _, item := range raw { + modality := strings.ToLower(strings.TrimSpace(item)) + if modality == "" { + continue + } + if _, exists := seen[modality]; exists { + continue + } + seen[modality] = struct{}{} + out = append(out, modality) + } + if len(out) == 0 { + return nil + } + return out +} + +func buildConfigModels[T modelEntry](models []T, ownedBy, modelType string) []*ModelInfo { + if len(models) == 0 { + return nil + } + now := time.Now().Unix() + out := make([]*ModelInfo, 0, len(models)) + seen := make(map[string]struct{}, len(models)) + for i := range models { + model := models[i] + name := strings.TrimSpace(model.GetName()) + info := buildConfiguredModelInfo(model, ownedBy, modelType, now, name, true) + if info == nil { + continue + } + alias := info.ID + key := strings.ToLower(alias) + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + if name != "" { + if upstream := registry.LookupStaticModelInfo(name); upstream != nil && upstream.Thinking != nil { + info.Thinking = upstream.Thinking + } + } + out = append(out, info) + } + return out +} + +func buildVertexCompatConfigModels(entry *config.VertexCompatKey) []*ModelInfo { + if entry == nil { + return nil + } + return buildConfigModels(entry.Models, "google", "vertex") +} + +func buildGeminiConfigModels(entry *config.GeminiKey) []*ModelInfo { + if entry == nil { + return nil + } + return buildConfigModels(entry.Models, "google", "gemini") +} + +func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo { + if entry == nil { + return nil + } + return buildConfigModels(entry.Models, "anthropic", "claude") +} + +func buildXAIConfigModels(entry *config.XAIKey) []*ModelInfo { + if entry == nil { + return nil + } + return buildConfigModels(entry.Models, "xai", "xai") +} + +func buildCodexConfigModels(entry *config.CodexKey) []*ModelInfo { + if entry == nil { + return nil + } + + models := registry.WithCodexBuiltins(buildConfigModels(entry.Models, "openai", "openai")) + configuredDisplayNames := make(map[string]string, len(entry.Models)) + seenConfiguredModels := make(map[string]struct{}, len(entry.Models)) + for i := range entry.Models { + model := entry.Models[i] + alias := strings.TrimSpace(model.Alias) + if alias == "" { + alias = strings.TrimSpace(model.Name) + } + if alias == "" { + continue + } + key := strings.ToLower(alias) + if _, exists := seenConfiguredModels[key]; exists { + continue + } + seenConfiguredModels[key] = struct{}{} + + displayName := strings.TrimSpace(model.DisplayName) + if displayName != "" { + configuredDisplayNames[key] = displayName + } + } + for _, model := range models { + if model == nil { + continue + } + if displayName, ok := configuredDisplayNames[strings.ToLower(model.ID)]; ok { + model.DisplayName = displayName + } + } + return models +} + +func rewriteModelInfoName(name, oldID, newID string) string { + trimmed := strings.TrimSpace(name) + if trimmed == "" { + return name + } + oldID = strings.TrimSpace(oldID) + newID = strings.TrimSpace(newID) + if oldID == "" || newID == "" { + return name + } + if strings.EqualFold(oldID, newID) { + return name + } + if strings.EqualFold(trimmed, oldID) { + return newID + } + if strings.HasSuffix(trimmed, "/"+oldID) { + prefix := strings.TrimSuffix(trimmed, oldID) + return prefix + newID + } + if trimmed == "models/"+oldID { + return "models/" + newID + } + return name +} + +func applyOAuthModelAlias(cfg *config.Config, provider, authKind string, models []*ModelInfo) []*ModelInfo { + return applyOAuthModelAliasForAuth(cfg, provider, authKind, nil, models) +} + +func applyOAuthModelAliasForAuth(cfg *config.Config, provider, authKind string, attributes map[string]string, models []*ModelInfo) []*ModelInfo { + if len(models) == 0 { + return models + } + channel := coreauth.OAuthModelAliasChannel(provider, authKind) + if channel == "" { + return models + } + aliases := oauthModelAliasesForAuth(cfg, channel, attributes) + if len(aliases) == 0 { + return models + } + return applyOAuthModelAliasEntries(aliases, models) +} + +func oauthModelAliasesForAuth(cfg *config.Config, channel string, attributes map[string]string) []config.OAuthModelAlias { + perAuthAliases := coreauth.OAuthModelAliasesFromAttributes(attributes) + if cfg == nil || len(cfg.OAuthModelAlias) == 0 { + return perAuthAliases + } + globalAliases := cfg.OAuthModelAlias[channel] + if len(perAuthAliases) == 0 { + return globalAliases + } + if len(globalAliases) == 0 { + return perAuthAliases + } + out := make([]config.OAuthModelAlias, 0, len(perAuthAliases)+len(globalAliases)) + seenAlias := make(map[string]struct{}, len(perAuthAliases)+len(globalAliases)) + add := func(aliases []config.OAuthModelAlias) { + for _, entry := range aliases { + alias := strings.TrimSpace(entry.Alias) + if alias == "" { + continue + } + key := strings.ToLower(alias) + if _, exists := seenAlias[key]; exists { + continue + } + seenAlias[key] = struct{}{} + out = append(out, entry) + } + } + add(perAuthAliases) + add(globalAliases) + return out +} + +func applyOAuthModelAliasEntries(aliases []config.OAuthModelAlias, models []*ModelInfo) []*ModelInfo { + type aliasEntry struct { + alias string + displayName string + fork bool + } + + forward := make(map[string][]aliasEntry, len(aliases)) + for i := range aliases { + name := strings.TrimSpace(aliases[i].Name) + alias := strings.TrimSpace(aliases[i].Alias) + if name == "" || alias == "" { + continue + } + if strings.EqualFold(name, alias) { + continue + } + key := strings.ToLower(name) + forward[key] = append(forward[key], aliasEntry{ + alias: alias, + displayName: strings.TrimSpace(aliases[i].DisplayName), + fork: aliases[i].Fork, + }) + } + if len(forward) == 0 { + return models + } + + out := make([]*ModelInfo, 0, len(models)) + seen := make(map[string]struct{}, len(models)) + for _, model := range models { + if model == nil { + continue + } + id := strings.TrimSpace(model.ID) + if id == "" { + continue + } + key := strings.ToLower(id) + entries := forward[key] + if len(entries) == 0 { + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + out = append(out, model) + continue + } + + keepOriginal := false + for _, entry := range entries { + if entry.fork { + keepOriginal = true + break + } + } + if keepOriginal { + if _, exists := seen[key]; !exists { + seen[key] = struct{}{} + out = append(out, model) + } + } + + addedAlias := false + for _, entry := range entries { + mappedID := strings.TrimSpace(entry.alias) + if mappedID == "" { + continue + } + if strings.EqualFold(mappedID, id) { + continue + } + aliasKey := strings.ToLower(mappedID) + if _, exists := seen[aliasKey]; exists { + continue + } + seen[aliasKey] = struct{}{} + clone := *model + clone.ID = mappedID + if entry.displayName != "" { + clone.DisplayName = entry.displayName + } + if clone.Name != "" { + clone.Name = rewriteModelInfoName(clone.Name, id, mappedID) + } + out = append(out, &clone) + addedAlias = true + } + + if !keepOriginal && !addedAlias { + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + out = append(out, model) + } + } + return out +} diff --git a/sdk/cliproxy/service_plugins.go b/sdk/cliproxy/service_plugins.go new file mode 100644 index 000000000..d1e5490a4 --- /dev/null +++ b/sdk/cliproxy/service_plugins.go @@ -0,0 +1,357 @@ +package cliproxy + +import ( + "context" + "strings" + "sync" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" + sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" + log "github.com/sirupsen/logrus" +) + +const ( + modelRegistrationMaxWorkersPerCategory = 5 + modelRegistrationMaxWorkersOpenAICompatibility = 20 + homeSubscriberPreAckRetryBackoff = 100 * time.Millisecond +) + +const ( + modelRegistrationPhaseConfigAPIKey = iota + modelRegistrationPhaseOther +) + +type modelRegistrationTask struct { + phase int + category string + run func(*openAICompatibilityRegistrationCache) +} + +type executorRegistrationOptions struct { + includeBaseline bool + includePlugins bool + forceReplaceAuths bool + auths []*coreauth.Auth +} + +var registerPluginExecutors = func(host *pluginhost.Host, manager *coreauth.Manager) { + if host == nil || manager == nil { + return + } + host.RegisterExecutors(manager, registry.GetGlobalRegistry()) +} + +// RegisterUsagePlugin registers a usage plugin on the global usage manager. +// This allows external code to monitor API usage and token consumption. +// +// Parameters: +// - plugin: The usage plugin to register +func (s *Service) RegisterUsagePlugin(plugin usage.Plugin) { + usage.RegisterPlugin(plugin) +} + +func (s *Service) registerPluginAuthParser() { + var parser PluginAuthParser + if s != nil && s.pluginHost != nil { + parser = s.pluginHost + } + sdkAuth.RegisterPluginAuthParser(parser) + if s != nil && s.watcher != nil { + s.watcher.SetPluginAuthParser(parser) + } +} + +func (s *Service) syncPluginRuntime(ctx context.Context) { + if !s.syncPluginRuntimeConfig(ctx) { + return + } + s.syncPluginModelRuntime(ctx) +} + +func (s *Service) syncPluginRuntimeConfig(ctx context.Context) bool { + if s == nil { + sdkAuth.RegisterPluginAuthParser(nil) + return false + } + s.cfgMu.RLock() + cfg := s.cfg + s.cfgMu.RUnlock() + return s.syncPluginRuntimeConfigForConfig(ctx, cfg) +} + +func (s *Service) syncPluginRuntimeConfigForConfig(ctx context.Context, cfg *config.Config) bool { + if s == nil { + sdkAuth.RegisterPluginAuthParser(nil) + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + + if s.pluginHost != nil { + s.pluginHost.ApplyConfig(ctx, cfg) + } + if errContext := ctx.Err(); errContext != nil { + return false + } + if s.coreManager != nil { + s.coreManager.SetPluginScheduler(s.pluginHost) + } + s.registerPluginAuthParser() + if s.pluginHost == nil { + return false + } + s.pluginHost.RegisterFrontendAuthProviders() + if errContext := ctx.Err(); errContext != nil { + return false + } + if s.accessManager != nil { + s.accessManager.SetProviders(sdkaccess.RegisteredProviders()) + } + s.pluginHost.RegisterUsagePlugins() + sdktranslator.SetPluginHooks(s.pluginHost) + if s.server != nil { + s.server.RefreshPluginManagementRoutes() + } + return ctx.Err() == nil +} + +func (s *Service) syncPluginModelRuntime(ctx context.Context) { + if s == nil || s.pluginHost == nil || s.coreManager == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + s.pluginHost.RegisterModels(ctx, registry.GetGlobalRegistry()) + if ctx.Err() != nil { + return + } + s.cfgMu.RLock() + homeEnabled := s.cfg != nil && s.cfg.Home.Enabled + s.cfgMu.RUnlock() + s.registerAvailableExecutors(ctx, executorRegistrationOptions{ + includeBaseline: homeEnabled, + includePlugins: true, + forceReplaceAuths: false, + auths: s.coreManager.List(), + }) + s.refreshPluginModelRegistrations(ctx) + if ctx.Err() != nil { + return + } + s.coreManager.RefreshSchedulerAll() +} + +func (s *Service) refreshPluginModelRegistrations(ctx context.Context) { + if s == nil || s.pluginHost == nil || s.coreManager == nil { + return + } + s.registerModelsForAuthBatch(ctx, s.coreManager.List()) +} + +func (s *Service) registerModelsForAuthBatch(ctx context.Context, auths []*coreauth.Auth) { + if s == nil || s.coreManager == nil || len(auths) == 0 { + return + } + tasks := make([]modelRegistrationTask, 0, len(auths)) + for _, auth := range auths { + if auth == nil { + continue + } + authForRegistration := auth.Clone() + tasks = append(tasks, modelRegistrationTask{ + phase: modelRegistrationPhase(authForRegistration), + category: modelRegistrationCategory(authForRegistration), + run: func(compatCache *openAICompatibilityRegistrationCache) { + s.completeModelRegistrationForAuthWithCache(ctx, authForRegistration, compatCache) + }, + }) + } + s.runModelRegistrationTasks(ctx, tasks) +} + +func (s *Service) runModelRegistrationTasks(ctx context.Context, tasks []modelRegistrationTask) { + if len(tasks) == 0 { + return + } + if ctx == nil { + ctx = context.Background() + } + + configAPIKeyTasks := make([]modelRegistrationTask, 0) + otherTasks := make([]modelRegistrationTask, 0) + for _, task := range tasks { + if task.phase == modelRegistrationPhaseConfigAPIKey { + configAPIKeyTasks = append(configAPIKeyTasks, task) + continue + } + otherTasks = append(otherTasks, task) + } + + compatCache := s.newOpenAICompatibilityRegistrationCache() + s.runModelRegistrationTaskPhase(ctx, configAPIKeyTasks, compatCache) + s.runModelRegistrationTaskPhase(ctx, otherTasks, compatCache) +} + +func (s *Service) runModelRegistrationTaskPhase(ctx context.Context, tasks []modelRegistrationTask, compatCache *openAICompatibilityRegistrationCache) { + if len(tasks) == 0 { + return + } + + grouped := make(map[string][]modelRegistrationTask) + order := make([]string, 0) + for _, task := range tasks { + if task.run == nil { + continue + } + category := strings.ToLower(strings.TrimSpace(task.category)) + if category == "" { + category = "unknown" + } + if _, exists := grouped[category]; !exists { + order = append(order, category) + } + grouped[category] = append(grouped[category], task) + } + + var wg sync.WaitGroup + for _, category := range order { + group := grouped[category] + workers := len(group) + maxWorkers := modelRegistrationMaxWorkersForCategory(category) + if workers > maxWorkers { + workers = maxWorkers + } + if workers <= 0 { + continue + } + + taskCh := make(chan modelRegistrationTask) + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for task := range taskCh { + select { + case <-ctx.Done(): + return + default: + } + task.run(compatCache) + } + }() + } + go func(group []modelRegistrationTask) { + defer close(taskCh) + for _, task := range group { + select { + case <-ctx.Done(): + return + case taskCh <- task: + } + } + }(group) + } + wg.Wait() +} + +func modelRegistrationPhase(auth *coreauth.Auth) int { + if coreauth.IsConfigAPIKeyAuth(auth) { + return modelRegistrationPhaseConfigAPIKey + } + return modelRegistrationPhaseOther +} + +func modelRegistrationCategory(auth *coreauth.Auth) string { + if auth == nil { + return "unknown" + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if compatProviderKey, _, compatDetected := openAICompatInfoFromAuth(auth); compatDetected { + if compatProviderKey != "" { + provider = compatProviderKey + } else { + provider = "openai-compatibility" + } + } + if provider == "" { + provider = "unknown" + } + + authKind := auth.AuthKind() + if authKind == "" { + return provider + } + return provider + ":" + authKind +} + +func modelRegistrationMaxWorkersForCategory(category string) int { + category = strings.ToLower(strings.TrimSpace(category)) + if strings.HasPrefix(category, "openai-compatible-") || strings.HasPrefix(category, "openai-compatibility") { + return modelRegistrationMaxWorkersOpenAICompatibility + } + return modelRegistrationMaxWorkersPerCategory +} + +func (s *Service) registerModelRefreshCallback() { + // Register callback for startup and periodic model catalog refresh. + // When remote model definitions change, re-register models for affected providers. + // This intentionally rebuilds per-auth model availability from the latest catalog + // snapshot instead of preserving prior registry suppression state. + registry.SetModelRefreshCallback(func(changedProviders []string) { + if s == nil || s.coreManager == nil || len(changedProviders) == 0 { + return + } + + providerSet := make(map[string]bool, len(changedProviders)) + for _, p := range changedProviders { + providerSet[strings.ToLower(strings.TrimSpace(p))] = true + } + + auths := s.coreManager.List() + refreshed := 0 + var refreshedMu sync.Mutex + tasks := make([]modelRegistrationTask, 0, len(auths)) + for _, item := range auths { + if item == nil || item.ID == "" { + continue + } + auth, ok := s.coreManager.GetByID(item.ID) + if !ok || auth == nil || auth.Disabled { + continue + } + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if !providerSet[provider] { + continue + } + authForRefresh := auth + tasks = append(tasks, modelRegistrationTask{ + phase: modelRegistrationPhase(authForRefresh), + category: modelRegistrationCategory(authForRefresh), + run: func(compatCache *openAICompatibilityRegistrationCache) { + if s.refreshModelRegistrationForAuthWithCache(authForRefresh, compatCache) { + refreshedMu.Lock() + refreshed++ + refreshedMu.Unlock() + } + }, + }) + } + s.runModelRegistrationTasks(context.Background(), tasks) + + if refreshed > 0 { + log.Infof("re-registered models for %d auth(s) due to model catalog changes: %v", refreshed, changedProviders) + } + }) +}