diff --git a/cmd/fetch_codex_models/main.go b/cmd/fetch_codex_models/main.go index 147f02f6..76ddf4d2 100644 --- a/cmd/fetch_codex_models/main.go +++ b/cmd/fetch_codex_models/main.go @@ -37,6 +37,8 @@ import ( coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v8/sdk/proxyutil" log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" ) const ( @@ -46,6 +48,8 @@ const ( defaultCodexUserAgent = "codex_cli_rs/0.159.0 (Mac OS 26.3.1; arm64) iTerm.app/3.6.9" defaultCodexOriginator = "codex_cli_rs" accessTokenRefreshLeeway = 30 * time.Second + baselineModelSlug = "gpt-5.5" + baselineModelVisibility = "list" ) func init() { @@ -289,13 +293,85 @@ func fetchModels(ctx context.Context, auth *coreauth.Auth, accessToken, clientVe return nil, 0, fmt.Errorf("models request failed with status %d: %s", httpResp.StatusCode, strings.TrimSpace(string(bodyBytes))) } - count, errCount := countModels(bodyBytes) - if errCount != nil { - return nil, 0, errCount + bodyBytes, count, errEnsure := ensureBaselineModel(bodyBytes) + if errEnsure != nil { + return nil, 0, errEnsure } return bodyBytes, count, nil } +func ensureBaselineModel(raw []byte) ([]byte, int, error) { + var payload struct { + Models []json.RawMessage `json:"models"` + } + if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil { + return nil, 0, fmt.Errorf("failed to parse response JSON: %w", errUnmarshal) + } + // Keep this check intentionally loose: fetch_codex_models dumps the upstream + // Codex API payload. Strict CPA catalog validation belongs in + // cmd/validate_codex_models and registry.ValidateCodexClientModelsJSON. + if payload.Models == nil { + return nil, 0, fmt.Errorf("response JSON does not contain models array") + } + if len(payload.Models) == 0 { + return raw, 0, nil + } + + foundBaseline := false + for i, modelRaw := range payload.Models { + slugResult := gjson.GetBytes(modelRaw, "slug") + if slugResult.Type != gjson.String || strings.TrimSpace(slugResult.Str) != baselineModelSlug { + continue + } + foundBaseline = true + visResult := gjson.GetBytes(modelRaw, "visibility") + if visResult.Type != gjson.String || visResult.Str != baselineModelVisibility { + updated, errSet := sjson.SetBytes(raw, fmt.Sprintf("models.%d.visibility", i), baselineModelVisibility) + if errSet != nil { + return nil, 0, fmt.Errorf("failed to update %s visibility: %w", baselineModelSlug, errSet) + } + raw = updated + } + } + + count := len(payload.Models) + if !foundBaseline { + cloned := append([]byte(nil), payload.Models[0]...) + var errSet error + cloned, errSet = sjson.SetBytes(cloned, "slug", baselineModelSlug) + if errSet != nil { + return nil, 0, fmt.Errorf("failed to set %s slug on cloned model: %w", baselineModelSlug, errSet) + } + visResult := gjson.GetBytes(cloned, "visibility") + if visResult.Type != gjson.String || visResult.Str != baselineModelVisibility { + cloned, errSet = sjson.SetBytes(cloned, "visibility", baselineModelVisibility) + if errSet != nil { + return nil, 0, fmt.Errorf("failed to set %s visibility on cloned model: %w", baselineModelSlug, errSet) + } + } + for _, reset := range []struct { + path string + value any + }{ + {path: "tool_mode", value: "direct"}, + {path: "use_responses_lite", value: false}, + {path: "supports_reasoning_effort_updates", value: false}, + } { + cloned, errSet = sjson.SetBytes(cloned, reset.path, reset.value) + if errSet != nil { + return nil, 0, fmt.Errorf("failed to set %s %s on cloned model: %w", baselineModelSlug, reset.path, errSet) + } + } + raw, errSet = sjson.SetRawBytes(raw, "models.-1", cloned) + if errSet != nil { + return nil, 0, fmt.Errorf("failed to append %s model: %w", baselineModelSlug, errSet) + } + count++ + } + + return raw, count, nil +} + func codexModelsURL(clientVersion string) (string, error) { u, err := url.Parse(codexModelsBaseURL + codexModelsPath) if err != nil { diff --git a/cmd/fetch_devin_models/main.go b/cmd/fetch_devin_models/main.go index b6d11401..d0f2793a 100644 --- a/cmd/fetch_devin_models/main.go +++ b/cmd/fetch_devin_models/main.go @@ -429,6 +429,15 @@ var devinCompoundSuffixes = []struct { {suffix: "-thinking", effort: ""}, {suffix: "-max-1m", effort: "max", readd: "-1m"}, {suffix: "-none-1m", effort: "none", readd: "-1m"}, + // Lightning tier keeps the effort before the -lightning suffix upstream + // (e.g. swe-2-high-lightning), so it aggregates as -lightning. + {suffix: "-minimal-lightning", effort: "minimal", readd: "-lightning"}, + {suffix: "-low-lightning", effort: "low", readd: "-lightning"}, + {suffix: "-medium-lightning", effort: "medium", readd: "-lightning"}, + {suffix: "-high-lightning", effort: "high", readd: "-lightning"}, + {suffix: "-xhigh-lightning", effort: "xhigh", readd: "-lightning"}, + {suffix: "-max-lightning", effort: "max", readd: "-lightning"}, + {suffix: "-none-lightning", effort: "none", readd: "-lightning"}, } var devinSimpleEffortSuffixes = []struct { @@ -571,12 +580,19 @@ func aggregateModels(raw []rawDevinModel) []devinModelJSON { } } + // Lightning display names strip both the effort word and "Lightning"; + // re-append it so swe-2-lightning shows as "SWE-2 Lightning". + displayName := entry.displayName + if strings.HasSuffix(id, "-lightning") && !strings.HasSuffix(strings.ToLower(displayName), "lightning") { + displayName = strings.TrimSpace(displayName + " Lightning") + } + res = append(res, devinModelJSON{ ID: entry.baseID, Object: "model", Type: "devin", OwnedBy: vendorName(entry.vendorID, entry.baseID), - DisplayName: entry.displayName, + DisplayName: displayName, ContextLength: entry.contextLength, MaxCompletionTokens: 64000, SupportedInputModalities: modalities, @@ -595,7 +611,7 @@ var devinDisplayNameSuffixes = []string{ " Low Thinking", " Medium Thinking", " High Thinking", " XHigh Thinking", " Max Thinking", " No Thinking", " Low", " Medium", " High", " XHigh", " Max", " None", " Minimal", - " Thinking", " Fast", + " Thinking", " Fast", " Lightning", } func cleanDisplayName(label string) string { diff --git a/cmd/server/main.go b/cmd/server/main.go index 2b990778..960d5e53 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -23,6 +23,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/buildinfo" "github.com/router-for-me/CLIProxyAPI/v8/internal/cmd" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/githubauth" "github.com/router-for-me/CLIProxyAPI/v8/internal/home" "github.com/router-for-me/CLIProxyAPI/v8/internal/homeplugins" "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" @@ -123,6 +124,7 @@ func main() { var discoverExclude []string var vertexImport string var vertexImportPrefix string + var siwcBind string var configPath string var password string var homeJWT string @@ -153,13 +155,14 @@ func main() { flag.StringVar(&configPath, "config", DefaultConfigPath, "Configure File Path") flag.StringVar(&vertexImport, "vertex-import", "", "Import Vertex service account key JSON file") flag.StringVar(&vertexImportPrefix, "vertex-import-prefix", "", "Prefix for Vertex model namespacing (use with -vertex-import)") + flag.StringVar(&siwcBind, "siwc-bind", "", "Bind Sign in with ChatGPT (SIWC) credentials to a Codex credential file") flag.StringVar(&password, "password", "", "") flag.StringVar(&homeJWT, "home-jwt", "", "Home control plane JWT for mTLS certificate bootstrap and connection") flag.BoolVar(&homeDisableClusterDiscovery, "home-disable-cluster-discovery", false, "Disable Home CLUSTER NODES discovery and keep using the configured -home-jwt address") flag.BoolVar(&tuiMode, "tui", false, "Start with terminal management UI") flag.BoolVar(&standalone, "standalone", false, "In TUI mode, start an embedded local server") flag.StringVar(&managementBaseURL, "management-base-url", "", "Base URL of remote management API for TUI client mode (e.g. https://proxy.example.com)") - flag.BoolVar(&localModel, "local-model", false, "Use embedded models.json and codex_client_models.json only, skip remote model catalog fetching") + flag.BoolVar(&localModel, "local-model", false, "Use embedded model catalogs unless models.catalog, models.codex-catalog, or models.devin-catalog explicitly overrides the source") flag.CommandLine.Usage = func() { out := flag.CommandLine.Output() @@ -567,7 +570,7 @@ func main() { log.Errorf("failed to bootstrap git-backed config: %v", errCopy) return } - if errCommit := gitStoreInst.PersistConfig(context.Background()); errCommit != nil { + if errCommit := persistSeededConfig(context.Background(), configFilePath, gitStoreInst.PersistConfig); errCommit != nil { log.Errorf("failed to commit initial git-backed config: %v", errCommit) return } @@ -646,6 +649,7 @@ func main() { } else { cfg.AuthDir = resolvedAuthDir } + githubauth.SetToken(cfg.GitHubToken) managementasset.SetCurrentConfig(cfg) // Create login options to be used in authentication flows. @@ -654,7 +658,7 @@ func main() { CallbackPort: oauthCallbackPort, } - commandMode := vertexImport != "" || antigravityLogin || codexLogin || codexDeviceLogin || claudeLogin || kimiLogin || kimiAILogin || xaiLogin || devinLogin || metaLogin + commandMode := vertexImport != "" || siwcBind != "" || antigravityLogin || codexLogin || codexDeviceLogin || claudeLogin || kimiLogin || kimiAILogin || xaiLogin || devinLogin || metaLogin cloudConfigMissing := isCloudDeploy && !configFileExists homeMode := configLoadedFromHome || (cfg != nil && cfg.Home.Enabled) exampleAPIKeySafeMode := shouldEnableExampleAPIKeySafeMode(cfg, commandMode, tuiMode, standalone, cloudConfigMissing, homeMode) @@ -712,6 +716,9 @@ func main() { if vertexImport != "" { // Handle Vertex service account import cmd.DoVertexImport(cfg, vertexImport, vertexImportPrefix) + } else if siwcBind != "" { + // Handle Sign in with ChatGPT (SIWC) credential binding + cmd.DoSIWCBind(cfg, siwcBind, options) } else if antigravityLogin { // Handle Antigravity login cmd.DoAntigravityLogin(cfg, options) @@ -742,14 +749,14 @@ func main() { return } if localModel && (!tuiMode || standalone) { - log.Info("Local model mode: using embedded model catalogs, remote model updates disabled") + log.Info("Local model mode: using embedded catalogs unless an explicit catalog source is configured") } if tuiMode { if standalone { // Standalone mode: start an embedded local server and connect TUI client to it. managementasset.StartAutoUpdater(context.Background(), configFilePath) misc.StartAntigravityVersionUpdater(context.Background()) - startModelCatalogUpdaters(localModel, cfg.Home.Enabled) + registry.SetLocalModelCatalogs(localModel) hook := tui.NewLogHook(2000) hook.SetFormatter(&logging.LogFormatter{}) log.AddHook(hook) @@ -824,7 +831,7 @@ func main() { // Start the main proxy service managementasset.StartAutoUpdater(context.Background(), configFilePath) misc.StartAntigravityVersionUpdater(context.Background()) - startModelCatalogUpdaters(localModel, cfg.Home.Enabled) + registry.SetLocalModelCatalogs(localModel) cmd.StartServiceWithPluginHost(cfg, configFilePath, password, pluginHost, serverOptions...) } } @@ -850,29 +857,18 @@ func resolveManagementBaseURL(flagURL string, cfg *config.Config) string { return fmt.Sprintf("http://127.0.0.1:%d", port) } -// modelCatalogUpdaterPlan decides which remote model catalogs should refresh. -// Codex client and Devin catalogs still refresh under Home mode because -// template metadata and Devin models stay edge-local. -func modelCatalogUpdaterPlan(localModel, homeEnabled bool) (startModels, startCodexClient, startDevin bool) { - if localModel { - return false, false, false - } - return !homeEnabled, true, true -} - -func startModelCatalogUpdaters(localModel, homeEnabled bool) { - startModels, startCodexClient, startDevin := modelCatalogUpdaterPlan(localModel, homeEnabled) - if startCodexClient { - registry.StartCodexClientModelsUpdater(context.Background()) - } - if startDevin { - registry.StartDevinModelsUpdater(context.Background()) - } - if startModels { - registry.StartModelsUpdater(context.Background()) - } else if homeEnabled { - log.Info("Home mode: remote models.json updates disabled; Codex client model list follows Home model IDs") +// persistSeededConfig persists a config file freshly copied from the template. The file +// is first rewritten as the startup load would, so a plaintext management secret never +// reaches the remote store. When that fails, the copy is removed and an error returned, +// so every start refuses the template instead of running from an unpersisted config. +func persistSeededConfig(ctx context.Context, configFilePath string, persist func(context.Context) error) error { + if errPrepare := config.PrepareConfigFileForUpload(configFilePath); errPrepare != nil { + if errRemove := os.Remove(configFilePath); errRemove != nil { + log.WithError(errRemove).Warn("failed to remove the config copied from the template") + } + return fmt.Errorf("prepare config for upload: %w", errPrepare) } + return persist(ctx) } func pluginBootstrapConfigPath(args []string, defaultPath string) string { diff --git a/internal/cmd/discover.go b/internal/cmd/discover.go index 929915e5..c871abd6 100644 --- a/internal/cmd/discover.go +++ b/internal/cmd/discover.go @@ -8,6 +8,7 @@ import ( "net" "os" "path/filepath" + "sort" "strconv" "strings" "time" @@ -220,14 +221,6 @@ func runDiscoverWithOptions(opts DiscoverOptions, stdout, stderr io.Writer, newB } fmt.Fprintf(stdout, " Protocols: %s\n", strings.Join(safeProtocols, ", ")) } - if len(gw.Features) > 0 { - var safeFeatures []string - for _, f := range gw.Features { - safeFeatures = append(safeFeatures, sanitizeTerminal(f)) - } - fmt.Fprintf(stdout, " Features: %s\n", strings.Join(safeFeatures, ", ")) - } - authStatus := "No" if gw.AuthRequired { authStatus = "Required" @@ -241,24 +234,33 @@ func runDiscoverWithOptions(opts DiscoverOptions, stdout, stderr io.Writer, newB } fmt.Fprintf(stdout, " Auth: %s\n", authStatus) - // Endpoints display scheme := "http" - if gw.RawTXT["tls"] == "1" { + if gw.TLS { scheme = "https" } - - fmt.Fprintf(stdout, " Base URLs:\n") - if openAIPath, ok := gw.Endpoints["openai"]; ok { - fmt.Fprintf(stdout, " - OpenAI: %s://%s%s\n", scheme, hostPort, sanitizeTerminal(openAIPath)) - } else { - fmt.Fprintf(stdout, " - OpenAI: %s://%s/v1\n", scheme, hostPort) + if gw.Path != "" { + fmt.Fprintf(stdout, " Path: %s\n", sanitizeTerminal(gw.Path)) } - - if anthropicPath, ok := gw.Endpoints["anthropic"]; ok { - fmt.Fprintf(stdout, " - Anthropic: %s://%s%s\n", scheme, hostPort, sanitizeTerminal(anthropicPath)) - } - if geminiPath, ok := gw.Endpoints["gemini"]; ok { - fmt.Fprintf(stdout, " - Gemini: %s://%s%s\n", scheme, hostPort, sanitizeTerminal(geminiPath)) + fmt.Fprintf(stdout, " TLS: %t\n", gw.TLS) + if len(gw.Endpoints) > 0 && primaryIP != "" { + fmt.Fprintf(stdout, " Endpoints:\n") + protocols := make([]string, 0, len(gw.Endpoints)) + for protocol := range gw.Endpoints { + protocols = append(protocols, protocol) + } + sort.Strings(protocols) + for _, protocol := range protocols { + endpointScheme := scheme + if protocol == "responses-ws" { + endpointScheme = "ws" + if gw.TLS { + endpointScheme = "wss" + } + } + // RFC 6874 URI syntax escapes the zone separator, unlike host:port display. + urlHost := strings.ReplaceAll(hostPort, "%", "%25") + fmt.Fprintf(stdout, " - %-16s %s://%s%s\n", sanitizeTerminal(protocol)+":", endpointScheme, urlHost, sanitizeTerminal(gw.Endpoints[protocol])) + } } fmt.Fprintln(stdout) @@ -269,7 +271,7 @@ func runDiscoverWithOptions(opts DiscoverOptions, stdout, stderr io.Writer, newB func preferredDisplayAddresses(gw discovery.DiscoveredService) (primary string, all []string) { for _, ip := range gw.IPv4 { - if ip == nil || ip.IsLoopback() || ip.IsUnspecified() { + if !usableDisplayIP(ip) { continue } all = append(all, ip.String()) @@ -281,11 +283,13 @@ func preferredDisplayAddresses(gw discovery.DiscoveredService) (primary string, var routable []string var linkLocal []string for _, ip := range gw.IPv6 { - if ip == nil || ip.IsLoopback() || ip.IsUnspecified() { + if !usableDisplayIP(ip) { continue } if ip.IsLinkLocalUnicast() { - linkLocal = append(linkLocal, ip.String()) + if zone := gw.IPv6Zones[ip.String()]; validDisplayZone(zone) { + linkLocal = append(linkLocal, ip.String()+"%"+zone) + } continue } routable = append(routable, ip.String()) @@ -294,14 +298,30 @@ func preferredDisplayAddresses(gw discovery.DiscoveredService) (primary string, return routable[0], append(routable, linkLocal...) } + if len(linkLocal) > 0 { + return linkLocal[0], linkLocal + } host := sanitizeDisplayHost(gw.Host) if host != "" { return host, append([]string{host}, linkLocal...) } - if len(linkLocal) > 0 { - return linkLocal[0], linkLocal + return "", nil +} + +func usableDisplayIP(ip net.IP) bool { + return ip != nil && ip.To16() != nil && !ip.IsLoopback() && !ip.IsUnspecified() && !ip.IsMulticast() && !ip.Equal(net.IPv4bcast) +} + +func validDisplayZone(zone string) bool { + if zone == "" { + return false + } + for _, c := range zone { + if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '-' || c == '_' || c == '.') { + return false + } } - return "127.0.0.1", nil + return true } func sanitizeDisplayHost(host string) string { diff --git a/internal/cmd/siwc_bind.go b/internal/cmd/siwc_bind.go new file mode 100644 index 00000000..f186fb0e --- /dev/null +++ b/internal/cmd/siwc_bind.go @@ -0,0 +1,326 @@ +package cmd + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "path/filepath" + "strings" + "time" + + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v8/internal/browser" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/misc" + "github.com/router-for-me/CLIProxyAPI/v8/internal/util" + log "github.com/sirupsen/logrus" +) + +// ResolveSIWCBindPath resolves a target Codex credential file path (either an absolute path +// or a bare filename looked up under cfg.AuthDir) and validates that it is an existing Codex JSON file. +func ResolveSIWCBindPath(cfg *config.Config, target string) (string, error) { + raw := strings.TrimSpace(target) + if raw == "" { + return "", fmt.Errorf("siwc-bind: missing codex credential file path") + } + + if cfg == nil || strings.TrimSpace(cfg.AuthDir) == "" { + return "", fmt.Errorf("siwc-bind: auth-dir is not configured") + } + authDir, errResolve := util.ResolveAuthDir(cfg.AuthDir) + if errResolve != nil { + return "", fmt.Errorf("siwc-bind: failed to resolve auth-dir: %w", errResolve) + } + authDir, errAbs := filepath.Abs(authDir) + if errAbs != nil { + return "", fmt.Errorf("siwc-bind: failed to resolve auth-dir path: %w", errAbs) + } + authDir, errRealDir := filepath.EvalSymlinks(authDir) + if errRealDir != nil { + return "", fmt.Errorf("siwc-bind: failed to resolve real auth-dir: %w", errRealDir) + } + + candidate := raw + if !filepath.IsAbs(raw) { + if strings.ContainsAny(raw, `/\`) || raw == "." || raw == ".." { + return "", fmt.Errorf("siwc-bind: expected an absolute path or bare filename") + } + candidate = filepath.Join(authDir, raw) + } + candidate, errRealFile := filepath.EvalSymlinks(candidate) + if errRealFile != nil { + return "", fmt.Errorf("siwc-bind: failed to resolve credential file: %w", errRealFile) + } + relative, errRel := filepath.Rel(authDir, candidate) + if errRel != nil || relative == "." || relative == ".." || strings.HasPrefix(relative, ".."+string(os.PathSeparator)) || filepath.IsAbs(relative) { + return "", fmt.Errorf("siwc-bind: credential file must belong to auth-dir") + } + info, errStat := os.Stat(candidate) + if errStat != nil || !info.Mode().IsRegular() { + return "", fmt.Errorf("siwc-bind: credential target must be an existing regular file") + } + data, errRead := os.ReadFile(candidate) + if errRead != nil { + return "", fmt.Errorf("siwc-bind: failed to read credential file %s: %w", candidate, errRead) + } + var parsed map[string]any + if errUnmarshal := json.Unmarshal(data, &parsed); errUnmarshal != nil { + return "", fmt.Errorf("siwc-bind: invalid credential JSON in %s", candidate) + } + if errValidate := validateSIWCBindCredential(parsed); errValidate != nil { + return "", errValidate + } + return filepath.Clean(candidate), nil +} + +func validateSIWCBindCredential(metadata map[string]any) error { + provider, _ := metadata["type"].(string) + if !strings.EqualFold(strings.TrimSpace(provider), "codex") { + return fmt.Errorf("siwc-bind: credential type must be codex") + } + for _, key := range []string{"access_token", "refresh_token"} { + value, _ := metadata[key].(string) + if strings.TrimSpace(value) == "" { + return fmt.Errorf("siwc-bind: codex credential requires a non-empty %s", key) + } + } + return nil +} + +func ensureSIWCDeviceIDFromFile(authFilePath string) (string, error) { + data, errRead := os.ReadFile(authFilePath) + if errRead != nil { + return "", fmt.Errorf("failed to read credential file: %w", errRead) + } + var parsed map[string]any + if errUnmarshal := json.Unmarshal(data, &parsed); errUnmarshal != nil { + return "", fmt.Errorf("invalid credential JSON: %w", errUnmarshal) + } + if siwcMap, ok := parsed["siwc"].(map[string]any); ok && len(siwcMap) > 0 { + if existingID, okID := siwcMap["device_id"].(string); okID && codex.IsValidUUID(existingID) { + return strings.ToLower(strings.TrimSpace(existingID)), nil + } + } + return strings.ToLower(uuid.NewString()), nil +} + +func saveSIWCToCodexCredentialFile(authFilePath string, siwcData *codex.CodexSIWCTokenData) error { + if siwcData == nil { + return fmt.Errorf("siwc token data is nil") + } + misc.LogSavingCredentials(authFilePath) + return codex.WithFileLock(authFilePath, func() error { + existingBytes, errRead := os.ReadFile(authFilePath) + if errRead != nil { + return fmt.Errorf("failed to read existing credential file: %w", errRead) + } + var diskMap map[string]any + if errUnmarshal := json.Unmarshal(existingBytes, &diskMap); errUnmarshal != nil { + return fmt.Errorf("failed to parse existing credential file: %w", errUnmarshal) + } + if errValidate := validateSIWCBindCredential(diskMap); errValidate != nil { + return errValidate + } + // Binding is an explicit SIWC replacement, not a stale snapshot merge. + diskMap["use_siwc"] = true + diskMap["siwc"] = siwcData.ToMap() + + return codex.WriteJSONFileAtomic(authFilePath, diskMap) + }) +} + +// DoSIWCBind runs the OpenAI Sign in with ChatGPT (SIWC) OAuth flow and writes the resulting +// SIWC credential section and use_siwc=true into the specified Codex credential file. +func DoSIWCBind(cfg *config.Config, target string, options *LoginOptions) { + if cfg == nil { + cfg = &config.Config{} + } + if options == nil { + options = &LoginOptions{} + } + promptFn := options.Prompt + if promptFn == nil { + promptFn = defaultProjectPrompt() + } + + authFilePath, errResolve := ResolveSIWCBindPath(cfg, target) + if errResolve != nil { + log.Errorf("%v", errResolve) + return + } + + deviceID, errDevice := ensureSIWCDeviceIDFromFile(authFilePath) + if errDevice != nil { + log.Errorf("siwc-bind: %v", errDevice) + return + } + + callbackPort := codex.SIWCCallbackPort + if options.CallbackPort > 0 { + callbackPort = options.CallbackPort + } + redirectURI := fmt.Sprintf("http://%s:%d%s", codex.SIWCCallbackHost, callbackPort, codex.SIWCCallbackPath) + + pkceCodes, errPKCE := codex.GeneratePKCECodes() + if errPKCE != nil { + log.Errorf("siwc-bind: PKCE generation failed: %v", errPKCE) + return + } + state, errState := codex.GenerateRandomBase64URL() + if errState != nil { + log.Errorf("siwc-bind: state generation failed: %v", errState) + return + } + nonce, errNonce := codex.GenerateRandomBase64URL() + if errNonce != nil { + log.Errorf("siwc-bind: nonce generation failed: %v", errNonce) + return + } + + oauthServer := codex.NewOAuthServer(callbackPort) + if errStart := oauthServer.Start(); errStart != nil { + if strings.Contains(errStart.Error(), "already in use") { + authErr := codex.NewAuthenticationError(codex.ErrPortInUse, errStart) + log.Error(codex.GetUserFriendlyMessage(authErr)) + os.Exit(codex.ErrPortInUse.Code) + } + log.Errorf("siwc-bind: failed to start OAuth callback server: %v", errStart) + return + } + defer func() { + stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if errStop := oauthServer.Stop(stopCtx); errStop != nil { + log.Warnf("siwc-bind: oauth server stop error: %v", errStop) + } + }() + + authSvc := codex.NewCodexAuth(cfg) + authURL, errURL := authSvc.GenerateSIWCAuthURL(state, nonce, deviceID, redirectURI, pkceCodes) + if errURL != nil { + log.Errorf("siwc-bind: failed to generate authorization URL: %v", errURL) + return + } + + if !options.NoBrowser { + fmt.Println("Opening browser for OpenAI Sign in with ChatGPT authentication") + if !browser.IsAvailable() { + log.Warn("No browser available; please open the URL manually") + util.PrintSSHTunnelInstructions(callbackPort) + fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL) + } else if errOpen := browser.OpenURL(authURL); errOpen != nil { + log.Warnf("Failed to open browser automatically: %v", errOpen) + util.PrintSSHTunnelInstructions(callbackPort) + fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL) + } + } else { + util.PrintSSHTunnelInstructions(callbackPort) + fmt.Printf("Visit the following URL to continue authentication:\n%s\n", authURL) + } + + fmt.Println("Waiting for Sign in with ChatGPT authentication callback...") + + callbackCh := make(chan *codex.OAuthResult, 1) + callbackErrCh := make(chan error, 1) + manualDescription := "" + + go func() { + result, errWait := oauthServer.WaitForCallback(5 * time.Minute) + if errWait != nil { + callbackErrCh <- errWait + return + } + callbackCh <- result + }() + + var result *codex.OAuthResult + var manualPromptTimer *time.Timer + var manualPromptC <-chan time.Time + if promptFn != nil { + manualPromptTimer = time.NewTimer(15 * time.Second) + manualPromptC = manualPromptTimer.C + defer manualPromptTimer.Stop() + } + + var manualInputCh <-chan string + var manualInputErrCh <-chan error + +waitForCallback: + for { + select { + case result = <-callbackCh: + break waitForCallback + case errWait := <-callbackErrCh: + fmt.Printf("Sign in with ChatGPT authentication failed: %v\n", errWait) + return + case <-manualPromptC: + manualPromptC = nil + if manualPromptTimer != nil { + manualPromptTimer.Stop() + } + select { + case result = <-callbackCh: + break waitForCallback + case errWait := <-callbackErrCh: + fmt.Printf("Sign in with ChatGPT authentication failed: %v\n", errWait) + return + default: + } + manualInputCh, manualInputErrCh = misc.AsyncPrompt(promptFn, "Paste the callback URL (or press Enter to keep waiting): ") + continue + case input := <-manualInputCh: + manualInputCh = nil + manualInputErrCh = nil + parsed, errParse := misc.ParseOAuthCallback(input) + if errParse != nil { + fmt.Printf("Sign in with ChatGPT authentication failed: %v\n", errParse) + return + } + if parsed == nil { + continue + } + manualDescription = parsed.ErrorDescription + result = &codex.OAuthResult{ + Code: parsed.Code, + State: parsed.State, + ClientID: parsed.ClientID, + Error: parsed.Error, + } + break waitForCallback + case errManual := <-manualInputErrCh: + fmt.Printf("Sign in with ChatGPT authentication failed: %v\n", errManual) + return + } + } + + if result.Error != "" { + errOAuth := codex.NewOAuthError(result.Error, manualDescription, http.StatusBadRequest) + fmt.Printf("Sign in with ChatGPT authentication failed: %v\n", errOAuth) + return + } + if result.State != state { + fmt.Println("Sign in with ChatGPT authentication failed: state mismatch") + return + } + if strings.TrimSpace(result.ClientID) == "" { + fmt.Println("Sign in with ChatGPT authentication failed: callback did not contain an issued client ID") + return + } + + siwcData, errExchange := authSvc.ExchangeSIWCCodeForTokens(context.Background(), result.Code, result.ClientID, redirectURI, deviceID, pkceCodes) + if errExchange != nil { + fmt.Printf("Sign in with ChatGPT token exchange failed: %v\n", errExchange) + return + } + + if errSave := saveSIWCToCodexCredentialFile(authFilePath, siwcData); errSave != nil { + fmt.Printf("Failed to save SIWC credentials: %v\n", errSave) + return + } + + fmt.Printf("SIWC credentials bound and saved to %s\n", authFilePath) + fmt.Println("Sign in with ChatGPT authentication successful!") +} diff --git a/internal/cmd/vertex_import.go b/internal/cmd/vertex_import.go index 31946237..6199beb7 100644 --- a/internal/cmd/vertex_import.go +++ b/internal/cmd/vertex_import.go @@ -60,7 +60,7 @@ func DoVertexImport(cfg *config.Config, keyPath string, prefix string) { log.Warn("vertex-import: client_email missing in service account json") } // Default location if not provided by user. Can be edited in the saved file later. - location := "us-central1" + location := "global" // Normalize and validate prefix: must be a single segment (no "/" allowed). prefix = strings.TrimSpace(prefix) diff --git a/internal/home/certificate.go b/internal/home/certificate.go index 9ccb5075..716ab9d9 100644 --- a/internal/home/certificate.go +++ b/internal/home/certificate.go @@ -13,6 +13,7 @@ import ( "encoding/hex" "encoding/json" "encoding/pem" + "errors" "fmt" "io" "net" @@ -25,7 +26,11 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/config" ) -const homeCertificateRequestTimeout = 30 * time.Second +const ( + homeCertificateRequestTimeout = 30 * time.Second + maxHomeCertificateRESPBulkLen = 1 << 20 // 1 MiB + maxHomeCertificateRESPLineLen = 64 << 10 // 64 KiB +) type homeJWTClaims struct { CertificateID string `json:"certificate_id"` @@ -354,7 +359,7 @@ func readRESPBulk(reader *bufio.Reader) ([]byte, error) { } switch prefix { case '$': - line, errLine := reader.ReadString('\n') + line, errLine := readHomeCertificateRESPLine(reader) if errLine != nil { return nil, errLine } @@ -362,16 +367,25 @@ func readRESPBulk(reader *bufio.Reader) ([]byte, error) { if errSize != nil { return nil, errSize } + if size < -1 { + return nil, fmt.Errorf("home certificate response has invalid bulk length") + } if size < 0 { return nil, fmt.Errorf("home certificate request returned nil") } + if size > maxHomeCertificateRESPBulkLen { + return nil, fmt.Errorf("home certificate response is too large") + } payload := make([]byte, size+2) if _, errFull := io.ReadFull(reader, payload); errFull != nil { return nil, errFull } + if payload[size] != '\r' || payload[size+1] != '\n' { + return nil, fmt.Errorf("home certificate response has invalid terminator") + } return payload[:size], nil case '-': - line, errLine := reader.ReadString('\n') + line, errLine := readHomeCertificateRESPLine(reader) if errLine != nil { return nil, errLine } @@ -381,6 +395,36 @@ func readRESPBulk(reader *bufio.Reader) ([]byte, error) { } } +func readHomeCertificateRESPLine(reader *bufio.Reader) (string, error) { + var buf []byte + for { + fragment, errRead := reader.ReadSlice('\n') + if len(buf)+len(fragment) > maxHomeCertificateRESPLineLen+2 { + return "", fmt.Errorf("home certificate response line is too long") + } + if errRead == nil { + var line string + if len(buf) == 0 { + line = string(fragment) + } else { + buf = append(buf, fragment...) + line = string(buf) + } + line = strings.TrimSuffix(line, "\n") + line = strings.TrimSuffix(line, "\r") + if len(line) > maxHomeCertificateRESPLineLen { + return "", fmt.Errorf("home certificate response line is too long") + } + return line, nil + } + if errors.Is(errRead, bufio.ErrBufferFull) { + buf = append(buf, fragment...) + continue + } + return "", errRead + } +} + func fileExists(path string) bool { info, errStat := os.Stat(path) return errStat == nil && !info.IsDir() diff --git a/internal/home/client.go b/internal/home/client.go index 5f7777dc..30b7c5e5 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -7,6 +7,7 @@ import ( "encoding/json" "errors" "fmt" + "math/rand/v2" "net" "net/http" "net/url" @@ -190,9 +191,12 @@ type Client struct { lifecycle config.CredentialConcurrencyConfig limiter atomic.Pointer[config.CredentialConcurrencyConfig] managed bool + // releaseReady belongs to this acknowledged membership and is never copied to a replacement. + releaseReady bool heartbeatOK atomic.Bool dispatchFenced atomic.Bool + releaseFenced atomic.Bool ambiguousDispatch atomic.Bool // casUnsupported latches when Home does not implement the CAS command. // It is deliberately NOT carried across NewLifetime: CAS support is a @@ -207,6 +211,8 @@ type Client struct { legacyMembership bool clusterNodes []clusterNode reconnectFailures int + balance clusterBalanceState + latency *clusterLatencyHistory } func New(homeCfg config.HomeConfig) *Client { @@ -234,8 +240,20 @@ func (c *Client) NewLifetime() *Client { testOperationTimeout: c.testOperationTimeout, instanceID: c.instanceID, legacyMembership: c.legacyMembership, + balance: c.balance, + latency: c.latency, } + next.balance.registered = false + next.balance.unsupported = false + next.balance.latencySupported = false next.recoveryState.Store(c.recoveryState.Load()) + if next.balance.moving { + next.homeCfg.Host = next.balance.target.IP + next.homeCfg.Port = next.balance.target.Port + if !next.recoveryState.CompareAndSwap(uint32(recoveryStateTakeoverEligible), uint32(recoveryStateSwitchingTakeover)) { + next.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateSwitching)) + } + } return next } @@ -289,12 +307,20 @@ func (c *Client) HeartbeatOK() bool { return c.heartbeatOK.Load() } -// Close permanently ends this client's dispatch lifetime. +// FenceDispatch stops new credential acquisitions while preserving transports and final concurrency releases. +func (c *Client) FenceDispatch() { + if c != nil { + c.dispatchFenced.Store(true) + } +} + +// Close permanently ends this client's dispatch and release lifetime. func (c *Client) Close() { if c == nil { return } c.dispatchFenced.Store(true) + c.releaseFenced.Store(true) c.heartbeatOK.Store(false) c.mu.Lock() commandClient, subscriptionClient, connections := c.detachClientsLocked() @@ -318,6 +344,7 @@ func (c *Client) closeBootstrapPools() { } c.heartbeatOK.Store(false) c.mu.Lock() + c.releaseReady = false commandClient, subscriptionClient, connections := c.detachClientsLocked() c.mu.Unlock() closeDetachedClients(commandClient, subscriptionClient, connections) @@ -330,6 +357,7 @@ func (c *Client) AbortAmbiguousDispatch() { } c.ambiguousDispatch.Store(true) c.dispatchFenced.Store(true) + c.releaseFenced.Store(true) c.heartbeatOK.Store(false) c.mu.Lock() commandClient, subscriptionClient, connections := c.detachClientsLocked() @@ -515,10 +543,31 @@ func (c *Client) ensureClients() error { } func (c *Client) redisOptionsLocked(addr string) (*redis.Options, error) { - tlsConfig, errTLS := c.homeTLSConfigLocked(addr) + options, tlsSettings := c.redisOptionsSnapshotLocked(addr) + tlsConfig, errTLS := newHomeTLSConfig(tlsSettings, "") if errTLS != nil { return nil, errTLS } + options.TLSConfig = tlsConfig + options.Dialer = c.trackedRedisDialer(redis.NewDialer(options)) + return options, nil +} + +// redisOptionsSnapshotLocked copies only values; certificate IO can run after releasing c.mu. +func (c *Client) redisOptionsSnapshotLocked(addr string) (*redis.Options, config.HomeTLSConfig) { + tlsSettings := c.homeCfg.TLS + serverName := strings.TrimSpace(tlsSettings.ServerName) + if serverName == "" { + if tlsSettings.UseTargetServerName { + serverName = hostFromAddress(addr) + } else { + serverName = strings.TrimSpace(c.seedHost) + } + } + if serverName == "" { + serverName = strings.TrimSpace(c.homeCfg.Host) + } + tlsSettings.ServerName = serverName dialTimeout := homeRedisOperationTimeout readTimeout := homeRedisOperationTimeout writeTimeout := homeRedisOperationTimeout @@ -539,7 +588,6 @@ func (c *Client) redisOptionsLocked(addr string) (*redis.Options, error) { } options := &redis.Options{ Addr: addr, - TLSConfig: tlsConfig, DialTimeout: dialTimeout, ReadTimeout: readTimeout, WriteTimeout: writeTimeout, @@ -547,8 +595,7 @@ func (c *Client) redisOptionsLocked(addr string) (*redis.Options, error) { DialerRetries: 1, ContextTimeoutEnabled: true, } - options.Dialer = c.trackedRedisDialer(redis.NewDialer(options)) - return options, nil + return options, tlsSettings } type homeDispatchConn struct { @@ -625,21 +672,6 @@ func cloneRedisOptions(options *redis.Options) *redis.Options { return &cloned } -func (c *Client) homeTLSConfigLocked(addr string) (*tls.Config, error) { - serverName := strings.TrimSpace(c.homeCfg.TLS.ServerName) - if serverName == "" { - if c.homeCfg.TLS.UseTargetServerName { - serverName = hostFromAddress(addr) - } else { - serverName = strings.TrimSpace(c.seedHost) - } - } - if serverName == "" { - serverName = strings.TrimSpace(c.homeCfg.Host) - } - return newHomeTLSConfig(c.homeCfg.TLS, serverName) -} - func hostFromAddress(addr string) string { host, _, errSplit := net.SplitHostPort(strings.TrimSpace(addr)) if errSplit == nil { @@ -785,6 +817,7 @@ func (c *Client) refreshClusterNodes(ctx context.Context) (bool, error) { if !c.clusterDiscoveryEnabled() { return false, nil } + if ctx == nil { ctx = context.Background() } @@ -805,7 +838,6 @@ func (c *Client) refreshClusterNodes(ctx context.Context) (bool, error) { if errText != nil { return false, errText } - nodes, errParse := parseClusterNodesPayload([]byte(raw)) if errParse != nil { return false, errParse @@ -813,12 +845,16 @@ func (c *Client) refreshClusterNodes(ctx context.Context) (bool, error) { if len(nodes) == 0 { return false, nil } - c.mu.Lock() defer c.mu.Unlock() c.clusterNodes = nodes - c.reconnectFailures = 0 - return c.switchToNodeLocked(nodes[0]), nil + // Discovery alone does not establish recovery or reset reconnect failures. + // Keep planned handoffs pinned here; repeated failures may still trigger failover. + if c.balance.moving || c.balance.returning || c.heartbeatOK.Load() { + return false, nil + } + node, ok := c.chooseClusterNodeLocked(nodes, false, rand.Float64()) + return ok && c.switchToNodeLocked(node), nil } func parseClusterNodesPayload(raw []byte) ([]clusterNode, error) { @@ -845,11 +881,17 @@ func (c *Client) updateClusterNodesFromPayload(raw []byte) error { func normalizeClusterNodes(nodes []clusterNode) []clusterNode { out := make([]clusterNode, 0, len(nodes)) + seen := make(map[string]bool, len(nodes)) for _, node := range nodes { node.IP = strings.TrimSpace(node.IP) - if node.IP == "" || node.Port <= 0 { + if node.IP == "" || node.Port <= 0 || node.Port > 65535 { + continue + } + address := clusterNodeAddress(node) + if seen[address] { continue } + seen[address] = true if node.ClientCount < 0 { node.ClientCount = 0 } @@ -869,6 +911,7 @@ func (c *Client) switchToNodeLocked(node clusterNode) bool { if strings.TrimSpace(c.homeCfg.Host) == host && c.homeCfg.Port == node.Port { return false } + c.releaseReady = false c.homeCfg.Host = host c.homeCfg.Port = node.Port if !c.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateSwitching)) { @@ -923,22 +966,13 @@ func (c *Client) failoverAfterSubscriptionTimeout() (bool, string) { func (c *Client) switchToNextNodeLocked() (bool, string) { currentHost := strings.TrimSpace(c.homeCfg.Host) currentPort := c.homeCfg.Port - candidates := append([]clusterNode(nil), c.clusterNodes...) - if strings.TrimSpace(c.seedHost) != "" && c.seedPort > 0 { - candidates = append(candidates, clusterNode{IP: c.seedHost, Port: c.seedPort}) + node, ok := c.chooseClusterNodeLocked(c.clusterNodes, true, rand.Float64()) + if !ok && strings.TrimSpace(c.seedHost) != "" && c.seedPort > 0 && (c.seedHost != currentHost || c.seedPort != currentPort) { + node, ok = clusterNode{IP: c.seedHost, Port: c.seedPort}, true } - for _, node := range candidates { - host := strings.TrimSpace(node.IP) - if host == "" || node.Port <= 0 { - continue - } - if host == currentHost && node.Port == currentPort { - continue - } - if c.switchToNodeLocked(clusterNode{IP: host, Port: node.Port}) { - addr, _ := c.addrLocked() - return true, addr - } + if ok && c.switchToNodeLocked(node) { + addr, _ := c.addrLocked() + return true, addr } return false, "" } @@ -1403,7 +1437,7 @@ func (c *Client) rPopAuth(ctx context.Context, requestedModel string, sessionID conn := cmd.Conn() defer func() { if errClose := conn.Close(); errClose != nil { - log.WithError(errClose).Debug("Home auth dispatch connection close failed") + log.Debugf("Home auth dispatch connection close failed: %v", errClose) } }() if errProbe := conn.Ping(ctx).Err(); errProbe != nil { @@ -1505,24 +1539,23 @@ func (c *Client) PushConcurrencyRelease(ctx context.Context, frame ConcurrencyRe } func (c *Client) concurrencyReleaseClient() (*redis.Client, error) { - if c == nil || c.dispatchFenced.Load() { + if c == nil || c.releaseFenced.Load() { return nil, ErrDispatchFenced } - state := recoveryState(c.recoveryState.Load()) - if state == recoveryStateTakeoverEligible || state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover { - return nil, ErrNotConnected - } if !c.Enabled() { return nil, ErrDisabled } c.mu.Lock() defer c.mu.Unlock() - if c.dispatchFenced.Load() { + if c.releaseFenced.Load() { return nil, ErrDispatchFenced } - state = recoveryState(c.recoveryState.Load()) - if state == recoveryStateTakeoverEligible || state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover { + state := recoveryState(c.recoveryState.Load()) + // Preparing takeover does not move membership. The acknowledged source owns final releases + // until Close; an unacknowledged replacement or a switched address does not inherit that right. + ownsSource := state == recoveryStateTakeoverEligible && c.releaseReady + if state != recoveryStateStable && !ownsSource { return nil, ErrNotConnected } if c.release != nil { @@ -1564,14 +1597,38 @@ func (c *Client) RPushAppLog(ctx context.Context, payload []byte) error { } func (c *Client) RPushPluginStatus(ctx context.Context, payload []byte) error { - cmd, errClient := c.commandClient() - if errClient != nil { - return errClient + options, errOptions := c.pluginSyncCommandOptions() + if errOptions != nil { + return errOptions } if len(payload) == 0 { return nil } - return cmd.RPush(ctx, redisKeyPluginStatus, payload).Err() + if ctx == nil { + ctx = context.Background() + } + // Plugin cancellation owns only this command's transport, never dispatch or release connections. + baseDialer := options.Dialer + if baseDialer == nil { + baseDialer = pluginSyncDialer(options) + } + options.Dialer = func(dialCtx context.Context, network, address string) (net.Conn, error) { + conn, errDial := baseDialer(dialCtx, network, address) + if errDial != nil { + return nil, errDial + } + return newPluginSyncCancelableConn(ctx, conn), nil + } + statusClient := redis.NewClient(options) + errPush := statusClient.RPush(ctx, redisKeyPluginStatus, payload).Err() + errClose := statusClient.Close() + if errContext := ctx.Err(); errContext != nil { + return errContext + } + if errClose != nil { + return errors.Join(errPush, fmt.Errorf("close plugin status client: %w", errClose)) + } + return errPush } func (c *Client) GetPluginTasks(ctx context.Context) ([]PluginTask, error) { @@ -1707,22 +1764,36 @@ func pluginSyncDialer(options *redis.Options) func(context.Context, string, stri type pluginSyncCancelableConn struct { net.Conn - done chan struct{} - once sync.Once + ctx context.Context + stopCancel func() bool + once sync.Once + closeUnderlyingOnce sync.Once } func newPluginSyncCancelableConn(ctx context.Context, conn net.Conn) net.Conn { - wrapped := &pluginSyncCancelableConn{Conn: conn, done: make(chan struct{})} - go func() { - select { - case <-ctx.Done(): - _ = closeUnderlyingTransport(conn) - case <-wrapped.done: - } - }() + if ctx == nil { + ctx = context.Background() + } + wrapped := &pluginSyncCancelableConn{ + Conn: conn, + ctx: ctx, + } + wrapped.stopCancel = context.AfterFunc(ctx, wrapped.closeUnderlying) + if ctx.Err() != nil { + wrapped.closeUnderlying() + } return wrapped } +func (c *pluginSyncCancelableConn) closeUnderlying() { + if c == nil || c.Conn == nil { + return + } + c.closeUnderlyingOnce.Do(func() { + _ = closeUnderlyingTransport(c.Conn) + }) +} + func closeUnderlyingTransport(conn net.Conn) error { if conn == nil { return net.ErrClosed @@ -1760,7 +1831,14 @@ func (c *pluginSyncCancelableConn) Close() error { if c == nil || c.Conn == nil { return net.ErrClosed } - c.once.Do(func() { close(c.done) }) + defer c.once.Do(func() { + if c.stopCancel != nil { + c.stopCancel() + } + }) + if c.ctx != nil && c.ctx.Err() != nil { + c.closeUnderlying() + } return c.Conn.Close() } @@ -1873,7 +1951,10 @@ func (c *Client) rebuildCommandPoolAndProbe(ctx context.Context) error { if errPing := c.Ping(ctx); errPing != nil { return errPing } + c.mu.Lock() + c.releaseReady = true c.recoveryState.Store(uint32(recoveryStateStable)) + c.mu.Unlock() return nil } @@ -1888,7 +1969,7 @@ func (c *Client) promoteSubscription() { c.mu.Unlock() if commandClient != nil { if errClose := commandClient.Close(); errClose != nil { - log.WithError(errClose).Warn("Home bootstrap command client close failed") + log.Warnf("Home bootstrap command client close failed: %v", errClose) } } } @@ -1906,7 +1987,12 @@ func (c *Client) handleSubscriptionPayload(ctx context.Context, channel string, } return onConfig([]byte(payload)) case redisChannelCluster: - return c.updateClusterNodesFromPayload([]byte(payload)) + if errUpdate := c.updateClusterNodesFromPayload([]byte(payload)); errUpdate != nil { + return errUpdate + } + return c.syncClusterBalance(ctx) + case redisChannelRebalance: + return c.executeClusterRebalance(ctx, []byte(payload)) default: return nil } @@ -1979,48 +2065,150 @@ func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func( } c.resetReconnectFailures() c.heartbeatOK.Store(true) - if onReady != nil { - onReady() + // Reception must keep monitoring silence while the consumer performs migration RPCs. + receiveCtx, cancelReceive := context.WithCancelCause(ctx) + events := make(chan any, 16) + receiveDone := make(chan struct{}) + c.mu.Lock() + if c.latency == nil { + c.latency = &clusterLatencyHistory{entries: make(map[string]clusterLatencyEntry)} } - - for { - _, receiveTimeout = c.subscriptionParameters() - event, errReceive := pubsub.ReceiveTimeout(ctx, receiveTimeout) - if errReceive != nil { + sampler := &clusterRTTSampler{history: c.latency, node: clusterNode{IP: c.homeCfg.Host, Port: c.homeCfg.Port}} + c.mu.Unlock() + failReceive := func(errReceive error) { + if !errors.Is(context.Cause(receiveCtx), ErrRebalanceRequested) { if ctx.Err() == nil { if c.heartbeatOK.Load() { c.markMembershipTakeoverEligible() } - if isTimeoutError(errReceive) { - c.markSubscriptionTimeout() - } else { - c.markReconnectFailure("subscription") + c.releaseFenced.Store(true) + } + // Graceful cancellation stops new work but keeps final releases available until Close. + c.dispatchFenced.Store(true) + } + c.heartbeatOK.Store(false) + cancelReceive(errReceive) + } + stopReceive := context.AfterFunc(receiveCtx, func() { _ = pubsub.Close() }) + latencyDone := make(chan struct{}) + go func() { + defer close(latencyDone) + if errLatency := c.runClusterLatency(receiveCtx, pubsub, sampler); errLatency != nil { + failReceive(errLatency) + } + }() + go func() { + defer close(receiveDone) + for { + _, timeout := c.subscriptionParameters() + event, errReceive := receiveSubscriptionEvent(receiveCtx, pubsub, sampler, timeout) + if errReceive == nil { + select { + case events <- event: + continue + case <-receiveCtx.Done(): + errReceive = context.Cause(receiveCtx) + default: + // Bounded backpressure fails the lifetime rather than dropping control messages. + errReceive = fmt.Errorf("Home subscription consumer backlog exceeded") } } - return c.endConfigSubscriberLifetimeWithSubscription(errReceive, pubsub, "heartbeat loss") + failReceive(errReceive) + return + } + }() + defer func() { + cancelReceive(context.Canceled) + stopReceive() + _ = pubsub.Close() + <-receiveDone + <-latencyDone + sampler.expire(time.Now(), true) + }() + endReceiveFailure := func(err error) error { + if cause := context.Cause(receiveCtx); cause != nil { + err = cause + } + if ctx.Err() == nil { + // Processing failures must preserve the committed membership before cleanup clears heartbeatOK. + if len(args) > 1 { + c.markMembershipTakeoverEligible() + } + if isTimeoutError(err) { + c.markSubscriptionTimeout() + } else { + c.markReconnectFailure("subscription") + } + } + return c.endConfigSubscriberLifetimeWithSubscription(err, pubsub, "subscription failure") + } + if errFinish := c.finishClusterRebalance(receiveCtx); errFinish != nil || receiveCtx.Err() != nil { + return endReceiveFailure(errFinish) + } + if onReady != nil { + onReady() + } + + for { + var event any + select { + case <-receiveCtx.Done(): + return endReceiveFailure(context.Cause(receiveCtx)) + case event = <-events: + } + if receiveCtx.Err() != nil { + return endReceiveFailure(context.Cause(receiveCtx)) } switch msg := event.(type) { case *redis.Message: if msg == nil { continue } - if errApply := c.handleSubscriptionPayload(ctx, msg.Channel, msg.Payload, onConfig); errApply != nil { - if strings.EqualFold(strings.TrimSpace(msg.Channel), redisChannelCluster) { - log.Warn("failed to apply cluster update from home control center, ignoring") - } else { - log.Warn("failed to apply config update from home control center, ignoring") + if errApply := c.handleSubscriptionPayload(receiveCtx, msg.Channel, msg.Payload, onConfig); errApply != nil { + if receiveCtx.Err() != nil { + return endReceiveFailure(errApply) + } + if errors.Is(errApply, ErrRebalanceRequested) { + cancelReceive(ErrRebalanceRequested) + return c.endConfigSubscriberLifetimeWithSubscription(errApply, pubsub, "planned rebalance") + } + if strings.EqualFold(strings.TrimSpace(msg.Channel), redisChannelCluster) || strings.EqualFold(strings.TrimSpace(msg.Channel), redisChannelRebalance) { + return endReceiveFailure(errApply) } + log.Warn("failed to apply config update from home control center, ignoring") } case *redis.Pong: c.resetReconnectFailures() - case *redis.Subscription: - continue default: log.Debugf("home subscription returned unsupported message type %T", event) } } } +// receiveSubscriptionEvent consumes telemetry without extending the heartbeat deadline. +func receiveSubscriptionEvent(ctx context.Context, pubsub *redis.PubSub, sampler *clusterRTTSampler, timeout time.Duration) (any, error) { + deadline := time.Now().Add(timeout) + for { + remaining := time.Until(deadline) + if remaining <= 0 { + return nil, os.ErrDeadlineExceeded + } + event, errReceive := pubsub.ReceiveTimeout(ctx, remaining) + if errReceive != nil { + return nil, errReceive + } + if _, ok := event.(*redis.Subscription); ok { + return nil, fmt.Errorf("unexpected Home subscription ACK after initialization") + } + if pong, ok := event.(*redis.Pong); ok && pong != nil && strings.HasPrefix(pong.Payload, "home-rtt/") { + // Timestamp before consumer callbacks, and also exclude stale or duplicate telemetry from liveness. + sampler.pong(pong.Payload, time.Now()) + continue + } + return event, nil + } +} + func receiveSubscriptionACKs(ctx context.Context, pubsub *redis.PubSub, receiveTimeout time.Duration, channels []string) error { if pubsub == nil || len(channels) == 0 { return fmt.Errorf("Home subscription ACK is missing") @@ -2040,6 +2228,7 @@ func receiveSubscriptionACKs(ctx context.Context, pubsub *redis.PubSub, receiveT func (c *Client) endConfigSubscriberLifetime(err error) error { c.heartbeatOK.Store(false) + c.failClusterRebalance() if !c.managedLifetime() { c.Close() } @@ -2048,9 +2237,12 @@ func (c *Client) endConfigSubscriberLifetime(err error) error { func (c *Client) endConfigSubscriberLifetimeWithSubscription(err error, subscription subscriptionCloser, reason string) error { c.heartbeatOK.Store(false) + if !errors.Is(err, ErrRebalanceRequested) { + c.failClusterRebalance() + } if subscription != nil { if errClose := subscription.Close(); errClose != nil { - log.WithError(errClose).Debugf("Home subscription close after %s", reason) + log.Debugf("Home subscription close after %s: %v", reason, errClose) } } if !c.managedLifetime() { @@ -2062,7 +2254,7 @@ func (c *Client) endConfigSubscriberLifetimeWithSubscription(err error, subscrip // StartConfigSubscriber is retained for callers that do not need the lifetime error. func (c *Client) StartConfigSubscriber(ctx context.Context, onConfig func([]byte) error) { if errRun := c.RunConfigSubscriberLifetime(ctx, onConfig, nil); errRun != nil && !errors.Is(errRun, context.Canceled) { - log.WithError(errRun).Warn("Home config subscription lifetime ended") + log.Warnf("Home config subscription lifetime ended: %v", errRun) } } diff --git a/internal/home/cluster_balance.go b/internal/home/cluster_balance.go new file mode 100644 index 00000000..cac3ab2d --- /dev/null +++ b/internal/home/cluster_balance.go @@ -0,0 +1,267 @@ +package home + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "strconv" + "strings" + + "github.com/redis/go-redis/v9" + log "github.com/sirupsen/logrus" +) + +const redisChannelRebalance = "rebalance" + +// ErrRebalanceRequested asks the supervisor to reconnect without draining safe active executions. +var ErrRebalanceRequested = errors.New("home cluster rebalance requested") + +type rebalanceInstruction struct { + ID string `json:"id"` + Term int64 `json:"term"` + Target clusterNode `json:"target"` +} + +type clusterBalanceState struct { + moving bool + returning bool + target clusterNode + origin clusterNode + request rebalanceInstruction + result string + registered bool + unsupported bool + latencySupported bool +} + +func (c *Client) canExecuteRebalanceLocked() bool { + return c.managed && c.clusterDiscoveryEnabledLocked() && !c.legacyMembership && + c.lifecycle.LifecycleConfigRevision > 0 && c.heartbeatOK.Load() && + !c.dispatchFenced.Load() && !c.ambiguousDispatch.Load() && + recoveryState(c.recoveryState.Load()) == recoveryStateStable && !c.balance.moving +} + +// rebalanceCommand owns a cancellable transport independent of issued dispatches. +func (c *Client) rebalanceCommand(ctx context.Context, address string, args ...any) (string, error) { + if errContext := ctx.Err(); errContext != nil { + return "", errContext + } + c.mu.Lock() + if address == "" { + address, _ = c.addrLocked() + } + options, errOptions := c.redisOptionsLocked(address) + c.mu.Unlock() + if errOptions != nil { + return "", errOptions + } + options.ReadTimeout = -1 + options.WriteTimeout = -1 + baseDialer := redis.NewDialer(options) + options.Dialer = func(dialCtx context.Context, network, dialAddress string) (net.Conn, error) { + conn, errDial := baseDialer(dialCtx, network, dialAddress) + if errDial != nil { + return nil, errDial + } + if ctx.Err() != nil { + _ = closeUnderlyingTransport(conn) + return nil, context.Cause(ctx) + } + return newPluginSyncCancelableConn(ctx, conn), nil + } + command := redis.NewClient(options) + stopCancel := context.AfterFunc(ctx, func() { _ = command.Close() }) + defer func() { + stopCancel() + if errClose := command.Close(); errClose != nil && ctx.Err() == nil { + log.Debugf("Home rebalance transport close failed: %v", errClose) + } + }() + raw, errCommand := command.Do(ctx, args...).Text() + if errCommand != nil { + return "", fmt.Errorf("Home rebalance %v: %w", args[0], errCommand) + } + return raw, nil +} + +func (c *Client) syncClusterBalance(ctx context.Context) error { + c.mu.Lock() + eligible := c.canExecuteRebalanceLocked() && !c.balance.unsupported + registered, request, result := c.balance.registered, c.balance.request, c.balance.result + c.mu.Unlock() + if !eligible { + return nil + } + if !registered { + raw, errEnable := c.rebalanceCommand(ctx, "", "CLUSTER", "BALANCE", "ENABLE") + if errEnable != nil { + message := strings.ToLower(errEnable.Error()) + if isHomeCommandUnsupported(errEnable) || strings.Contains(message, "unsupported subcommand") || strings.Contains(message, "rebalance unsupported") { + // Legacy Home has no capability code; only an explicit unsupported reply disables balancing. + c.mu.Lock() + c.balance.unsupported = true + c.mu.Unlock() + return nil + } + return fmt.Errorf("register Home rebalance: %w", errEnable) + } + var response struct { + OK bool `json:"ok"` + Capabilities []string `json:"capabilities"` + } + if errJSON := json.Unmarshal([]byte(raw), &response); errJSON != nil { + return fmt.Errorf("decode Home rebalance ENABLE response: %w", errJSON) + } + if !response.OK { + return nil + } + c.mu.Lock() + c.balance.registered = true + for _, capability := range response.Capabilities { + if capability == clusterLatencyCapability { + c.balance.latencySupported = true + } + } + c.mu.Unlock() + } + if result == "" { + return nil + } + + raw, errResult := c.rebalanceCommand(ctx, "", "CLUSTER", "BALANCE", "RESULT", request.ID, result) + if errResult != nil { + return fmt.Errorf("report Home rebalance result: %w", errResult) + } + var response struct { + OK bool `json:"ok"` + } + if errJSON := json.Unmarshal([]byte(raw), &response); errJSON != nil { + return fmt.Errorf("decode Home rebalance RESULT response: %w", errJSON) + } + if !response.OK { + return nil + } + c.mu.Lock() + if c.balance.request.ID == request.ID && c.balance.result == result { + c.balance.result = "" + c.balance.request = rebalanceInstruction{} + } + c.mu.Unlock() + return nil +} + +// executeClusterRebalance claims an exact persisted plan before touching any transport. +// This lifetime keeps its original address; only the supervisor's replacement uses the target. +func (c *Client) executeClusterRebalance(ctx context.Context, payload []byte) error { + var request rebalanceInstruction + if errJSON := json.Unmarshal(payload, &request); errJSON != nil { + return fmt.Errorf("decode Home rebalance instruction: %w", errJSON) + } + request.Target.IP = strings.TrimSpace(request.Target.IP) + if request.ID == "" || request.Term <= 0 || request.Target.IP == "" || request.Target.Port <= 0 || request.Target.Port > 65535 { + return nil + } + c.mu.Lock() + eligible := c.canExecuteRebalanceLocked() && c.balance.registered && c.balance.result == "" && + (request.Target.IP != strings.TrimSpace(c.homeCfg.Host) || request.Target.Port != c.homeCfg.Port) + c.mu.Unlock() + if !eligible || ctx.Err() != nil { + return nil + } + raw, errClaim := c.rebalanceCommand(ctx, "", "CLUSTER", "BALANCE", "CLAIM", request.ID, request.Term, request.Target.IP, request.Target.Port) + if errClaim != nil { + return errClaim + } + var response struct { + OK bool `json:"ok"` + } + if errJSON := json.Unmarshal([]byte(raw), &response); errJSON != nil { + return fmt.Errorf("decode Home rebalance CLAIM response: %w", errJSON) + } + if !response.OK { + return nil + } + c.mu.Lock() + c.balance.request = request + c.balance.result = "failed" + c.mu.Unlock() + if _, errPing := c.rebalanceCommand(ctx, clusterNodeAddress(request.Target), "PING"); errPing != nil { + if errReport := c.syncClusterBalance(ctx); errReport != nil { + return errors.Join(errPing, errReport) + } + c.mu.Lock() + recovered := ctx.Err() == nil && c.canExecuteRebalanceLocked() && c.balance.result == "" + c.mu.Unlock() + if !recovered { + return errPing + } + // No switch occurred. An acknowledged failure and a healthy source safely end this plan, not the subscription. + log.Warnf("Home rebalance preflight failed; failure reported, keeping source: plan=%q target=%q error=%v", request.ID, clusterNodeAddress(request.Target), errPing) + return nil + } + c.mu.Lock() + if ctx.Err() != nil || !c.canExecuteRebalanceLocked() { + c.mu.Unlock() + return c.syncClusterBalance(ctx) + } + c.balance.origin = clusterNode{IP: strings.TrimSpace(c.homeCfg.Host), Port: c.homeCfg.Port} + c.balance.target = request.Target + c.balance.result = "" + c.balance.moving = true + c.markMembershipTakeoverEligible() + origin := c.balance.origin + c.mu.Unlock() + log.Infof("executing Home rebalance plan: plan=%q from=%q to=%q", request.ID, clusterNodeAddress(origin), clusterNodeAddress(request.Target)) + return ErrRebalanceRequested +} + +func (c *Client) finishClusterRebalance(ctx context.Context) error { + c.mu.Lock() + if c.balance.moving { + c.balance.moving = false + c.balance.result = "success" + } + c.balance.returning = false + c.mu.Unlock() + return c.syncClusterBalance(ctx) +} + +func (c *Client) failClusterRebalance() { + c.mu.Lock() + defer c.mu.Unlock() + if !c.balance.moving { + return + } + c.balance.moving = false + c.balance.returning = true + c.balance.result = "failed" + c.switchToNodeLocked(c.balance.origin) +} + +// chooseClusterNodeLocked only chooses bootstrap/fault-recovery targets, never active migrations. +func (c *Client) chooseClusterNodeLocked(nodes []clusterNode, excludeCurrent bool, draw float64) (clusterNode, bool) { + var choices []clusterNode + for _, node := range nodes { + if node.IP == "" || node.Port <= 0 || node.Port > 65535 { + continue + } + if excludeCurrent && node.IP == strings.TrimSpace(c.homeCfg.Host) && node.Port == c.homeCfg.Port { + continue + } + if len(choices) == 0 || node.ClientCount < choices[0].ClientCount { + choices = []clusterNode{node} + } else if node.ClientCount == choices[0].ClientCount { + choices = append(choices, node) + } + } + if len(choices) == 0 { + return clusterNode{}, false + } + return choices[int(draw*float64(len(choices)))], true +} + +func clusterNodeAddress(node clusterNode) string { + return net.JoinHostPort(node.IP, strconv.Itoa(node.Port)) +} diff --git a/internal/home/cluster_latency.go b/internal/home/cluster_latency.go new file mode 100644 index 00000000..46c2b974 --- /dev/null +++ b/internal/home/cluster_latency.go @@ -0,0 +1,317 @@ +package home + +import ( + "context" + "encoding/json" + "fmt" + "math/rand/v2" + "net" + "sync" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + log "github.com/sirupsen/logrus" +) + +const ( + clusterLatencyCapability = "cpa_latency_swap_v1" + clusterLatencyTTL = 2 * time.Minute + clusterLatencyTimeout = 3 * time.Second +) + +type clusterLatencySample struct { + IP string `json:"ip"` + Port int `json:"port"` + Sequence int64 `json:"sequence"` + ObservedAt time.Time `json:"observed_at"` + AgeMS int64 `json:"age_ms"` + RTTUS int64 `json:"rtt_us"` + Valid bool `json:"valid"` +} + +type clusterLatencyEntry struct { + sample clusterLatencySample + measuredAt time.Time + nextProbeCycle int64 + nextProbeAt time.Time + probeBackoff time.Duration +} + +// Shared across replacement lifetimes; pending PINGs belong only to a subscription sampler. +type clusterLatencyHistory struct { + mu sync.Mutex + sequence int64 + entries map[string]clusterLatencyEntry + probeBaseline time.Time +} + +func (h *clusterLatencyHistory) observe(node clusterNode, rtt time.Duration, valid bool, now time.Time) { + h.mu.Lock() + defer h.mu.Unlock() + if h.entries == nil { + h.entries = make(map[string]clusterLatencyEntry) + } + key := clusterNodeAddress(node) + previous := h.entries[key] + if len(h.entries) >= 64 && previous.measuredAt.IsZero() { + oldestKey := "" + var oldest time.Time + for address, entry := range h.entries { + if oldestKey == "" || entry.measuredAt.Before(oldest) { + oldestKey, oldest = address, entry.measuredAt + } + } + delete(h.entries, oldestKey) + } + h.sequence++ + if valid && rtt > 0 && previous.sample.Valid && now.Sub(previous.measuredAt) <= clusterLatencyTTL { + rtt = (time.Duration(previous.sample.RTTUS)*time.Microsecond*3 + rtt) / 4 + } + previous.sample = clusterLatencySample{IP: node.IP, Port: node.Port, Sequence: h.sequence, ObservedAt: now.UTC(), RTTUS: max(1, rtt.Microseconds()), Valid: valid && rtt > 0} + previous.measuredAt = now + h.entries[key] = previous +} + +func (h *clusterLatencyHistory) probeDue(node clusterNode, reportAt time.Time) bool { + h.mu.Lock() + defer h.mu.Unlock() + entry := h.entries[clusterNodeAddress(node)] + if reportAt.Before(entry.nextProbeAt) { + return false + } + // Round against one fixed baseline so ticker jitter cannot skip the next report cycle. + return h.probeBaseline.IsZero() || int64(reportAt.Sub(h.probeBaseline).Round(time.Minute)/time.Minute) >= entry.nextProbeCycle +} + +func (h *clusterLatencyHistory) observeProbe(node clusterNode, rtt time.Duration, valid bool, reportAt, now time.Time) { + h.observe(node, rtt, valid, now) + h.mu.Lock() + defer h.mu.Unlock() + address := clusterNodeAddress(node) + entry := h.entries[address] + if h.probeBaseline.IsZero() { + h.probeBaseline = reportAt + } + entry.nextProbeCycle = int64(reportAt.Sub(h.probeBaseline).Round(time.Minute)/time.Minute) + 1 + entry.nextProbeAt = time.Time{} + if !valid { + entry.probeBackoff = min(max(time.Minute, entry.probeBackoff*2), 10*time.Minute) + entry.nextProbeAt = now.Add(entry.probeBackoff) + } else { + entry.probeBackoff = 0 + } + h.entries[address] = entry +} + +func (h *clusterLatencyHistory) report(now time.Time) []clusterLatencySample { + h.mu.Lock() + defer h.mu.Unlock() + var samples []clusterLatencySample + for _, entry := range h.entries { + age := now.Sub(entry.measuredAt) + if age < 0 || age > clusterLatencyTTL { + continue + } + sample := entry.sample + sample.AgeMS = age.Milliseconds() + samples = append(samples, sample) + } + return samples +} + +type clusterRTTSampler struct { + mu sync.Mutex + history *clusterLatencyHistory + node clusterNode + token string + sentAt time.Time +} + +func (s *clusterRTTSampler) begin(now time.Time) string { + s.mu.Lock() + defer s.mu.Unlock() + if s.token != "" { + s.history.observe(s.node, 0, false, now) + } + s.token, s.sentAt = "home-rtt/"+uuid.NewString(), now + return s.token +} + +func (s *clusterRTTSampler) pong(token string, now time.Time) bool { + s.mu.Lock() + defer s.mu.Unlock() + if token == "" || token != s.token { + return false + } + rtt := now.Sub(s.sentAt) + s.history.observe(s.node, max(time.Microsecond, rtt), rtt >= 0 && rtt <= clusterLatencyTimeout, now) + s.token = "" + return true +} + +func (s *clusterRTTSampler) expire(now time.Time, force bool) { + s.mu.Lock() + defer s.mu.Unlock() + if s.token != "" && (force || now.Sub(s.sentAt) > clusterLatencyTimeout) { + s.history.observe(s.node, 0, false, now) + s.token = "" + } +} + +// runClusterLatency returns subscription transport failures; report and probe failures remain telemetry-only. +func (c *Client) runClusterLatency(ctx context.Context, pubsub *redis.PubSub, sampler *clusterRTTSampler) error { + sampleTimer := time.NewTimer(clusterLatencySampleDelay()) + defer sampleTimer.Stop() + // Spread simultaneously started clients across a minute, then keep a fixed report cadence. + reportTicker := time.NewTicker(clusterLatencyReportDelay()) + defer reportTicker.Stop() + firstReport := true + for { + select { + case <-ctx.Done(): + return nil + case <-sampleTimer.C: + c.mu.Lock() + enabled := c.balance.latencySupported && c.balance.registered + c.mu.Unlock() + if enabled { + token := sampler.begin(time.Now()) + if errPing := pubsub.Ping(ctx, token); errPing != nil { + sampler.expire(time.Now(), true) + // go-redis may already have resubscribed; only the lifecycle manager may replace membership. + return fmt.Errorf("Home subscription telemetry PING: %w", errPing) + } + } + sampleTimer.Reset(clusterLatencySampleDelay()) + case reportAt := <-reportTicker.C: + if firstReport { + // Reset before the RPC so report duration never adds to the recurring interval. + reportTicker.Reset(time.Minute) + firstReport = false + } + sampler.expire(time.Now(), false) + if errReport := c.reportClusterLatency(ctx, sampler.history, reportAt); errReport != nil { + // Telemetry failure must not end a healthy subscription or reclaim its membership. + log.Warnf("Home latency telemetry failed: %v", errReport) + } + } + } +} + +func clusterLatencySampleDelay() time.Duration { + return 27*time.Second + time.Duration(rand.Int64N(int64(6*time.Second))) +} + +func clusterLatencyReportDelay() time.Duration { + return time.Minute + time.Duration(1+rand.Int64N(int64(time.Minute))) +} + +func (c *Client) reportClusterLatency(ctx context.Context, history *clusterLatencyHistory, reportAt time.Time) error { + c.mu.Lock() + eligible := c.canExecuteRebalanceLocked() && c.balance.registered && c.balance.latencySupported + processID := c.instanceID + c.mu.Unlock() + if !eligible { + return nil + } + payload, errJSON := json.Marshal(struct { + ProcessID string `json:"process_id"` + Samples []clusterLatencySample `json:"samples"` + }{ProcessID: processID, Samples: history.report(time.Now())}) + if errJSON != nil { + return fmt.Errorf("encode Home latency report: %w", errJSON) + } + command, errClient := c.commandClient() + if errClient != nil { + return fmt.Errorf("get Home latency command client: %w", errClient) + } + // Reuse the command pool; telemetry must never close dispatch transports on cancellation. + commandCtx, cancelCommand := context.WithTimeout(ctx, clusterLatencyTimeout) + raw, errReport := command.Do(commandCtx, "CLUSTER", "BALANCE", "LATENCY", string(payload)).Text() + cancelCommand() + if errReport != nil { + return fmt.Errorf("report Home latency: %w", errReport) + } + var response struct { + OK bool `json:"ok"` + Probe *clusterNode `json:"probe"` + } + if errDecode := json.Unmarshal([]byte(raw), &response); errDecode != nil { + return fmt.Errorf("decode Home latency response: %w", errDecode) + } + if !response.OK { + return fmt.Errorf("Home latency report rejected") + } + if response.Probe == nil { + return nil + } + target := *response.Probe + c.mu.Lock() + known := false + for _, node := range c.clusterNodes { + known = known || (node.IP == target.IP && node.Port == target.Port) + } + current := target.IP == c.homeCfg.Host && target.Port == c.homeCfg.Port + c.mu.Unlock() + if !known || current || target.IP == "" || target.Port <= 0 || target.Port > 65535 { + return nil + } + // One worker, at most one probe per report cycle, with per-target failure backoff. + if !history.probeDue(target, reportAt) { + return nil + } + address := clusterNodeAddress(target) + probeCtx, cancelProbe := context.WithTimeout(ctx, clusterLatencyTimeout) + rtt, errProbe := c.probeClusterLatency(probeCtx, address) + cancelProbe() + history.observeProbe(target, rtt, errProbe == nil, reportAt, time.Now()) + if errProbe != nil { + return fmt.Errorf("probe Home latency: %w", errProbe) + } + return nil +} + +// The first PING establishes mTLS; only the second round trip is timed. Neither subscribes. +func (c *Client) probeClusterLatency(ctx context.Context, address string) (time.Duration, error) { + c.mu.Lock() + options, tlsSettings := c.redisOptionsSnapshotLocked(address) + c.mu.Unlock() + // Loading and parsing certificate files must not block dispatch on the client's mutex. + tlsConfig, errTLS := newHomeTLSConfig(tlsSettings, "") + if errTLS != nil { + return 0, errTLS + } + options.TLSConfig = tlsConfig + options.ReadTimeout, options.WriteTimeout, options.MaxRetries = -1, -1, -1 + options.PoolSize = 1 + baseDialer := redis.NewDialer(options) + options.Dialer = func(dialCtx context.Context, network, dialAddress string) (net.Conn, error) { + conn, errDial := baseDialer(dialCtx, network, dialAddress) + if errDial != nil { + return nil, errDial + } + if ctx.Err() != nil { + _ = closeUnderlyingTransport(conn) + return nil, context.Cause(ctx) + } + return newPluginSyncCancelableConn(ctx, conn), nil + } + probe := redis.NewClient(options) + stopCancel := context.AfterFunc(ctx, func() { _ = probe.Close() }) + defer func() { + stopCancel() + _ = probe.Close() + }() + if errWarmup := probe.Ping(ctx).Err(); errWarmup != nil { + return 0, errWarmup + } + token := "home-probe/" + uuid.NewString() + started := time.Now() + response, errPing := probe.Do(ctx, "PING", token).Text() + rtt := max(time.Microsecond, time.Since(started)) + if errPing == nil && response != token { + return 0, ErrEmptyResponse + } + return rtt, errPing +} diff --git a/internal/pluginhost/adapters.go b/internal/pluginhost/adapters.go index 88846201..7410b88e 100644 --- a/internal/pluginhost/adapters.go +++ b/internal/pluginhost/adapters.go @@ -150,6 +150,7 @@ func pluginThinkingSupportToRegistryThinkingSupport(thinking *pluginapi.Thinking Max: thinking.Max, ZeroAllowed: thinking.ZeroAllowed, DynamicAllowed: thinking.DynamicAllowed, + DisableType: thinking.DisableType, Levels: cloneStringSlice(thinking.Levels), } } @@ -190,6 +191,7 @@ func registryThinkingSupportToPluginThinkingSupport(thinking *registry.ThinkingS Max: thinking.Max, ZeroAllowed: thinking.ZeroAllowed, DynamicAllowed: thinking.DynamicAllowed, + DisableType: thinking.DisableType, Levels: cloneStringSlice(thinking.Levels), } } @@ -225,6 +227,8 @@ func cloneRegistryModels(in []*registry.ModelInfo) []*registry.ModelInfo { return out } +// RegisterModels collects static catalogs for executor and per-auth registration. +// Catalog-only plugins must not advertise availability independently of credentials. func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) { if h == nil || modelRegistry == nil { return @@ -232,8 +236,6 @@ func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) snap := h.Snapshot() records := h.activeRecordsFromSnapshot(snap) - registrations := make([]modelClientRegistration, 0) - nextClients := make(map[string]struct{}) nextProviders := make(map[string]string) nextModelRegistrations := make(map[string]pluginModelRegistration) for _, record := range records { @@ -288,17 +290,48 @@ func (h *Host) RegisterModels(ctx context.Context, modelRegistry modelRegistry) hasExecutor: record.plugin.Capabilities.Executor != nil, } nextProviders[record.id] = provider - if record.plugin.Capabilities.Executor == nil { - clientID := "plugin:" + record.id + ":" + provider - registrations = append(registrations, modelClientRegistration{ - clientID: clientID, - provider: provider, - models: models, - }) - nextClients[clientID] = struct{}{} + } + h.mu.Lock() + defer h.mu.Unlock() + if h.Snapshot() != snap { + return + } + h.modelProviders = nextProviders + h.modelRegistrations = nextModelRegistrations +} + +// HasAuthModelProvider reports whether an active plugin declares per-auth model +// discovery for provider. It does not fetch models or refresh credentials. +func (h *Host) HasAuthModelProvider(provider string) bool { + if h == nil { + return false + } + provider = normalizeProviderID(provider) + if provider == "" { + return false + } + for _, record := range h.activeRecords() { + if h.authModelProviderMatches(record, provider) { + return true } } - h.commitModelClients(snap, modelRegistry, registrations, nextClients, nextProviders, nextModelRegistrations) + return false +} + +func (h *Host) authModelProviderMatches(record capabilityRecord, provider string) bool { + caps := record.plugin.Capabilities + if caps.ModelProvider == nil || h.isPluginFused(record.id) || !executorScopeAllowsOAuthModels(caps) { + return false + } + if caps.AuthProvider != nil { + identifier, ok := h.callAuthProviderIdentifier(record.id, caps.AuthProvider) + return ok && normalizeProviderID(identifier) == provider + } + candidate := normalizeProviderID(h.modelProvider(record.id)) + if candidate == "" && caps.Executor != nil { + candidate, _ = h.executorProvider(record, caps.Executor) + } + return candidate == provider } func (h *Host) ModelsForAuth(ctx context.Context, auth *coreauth.Auth) AuthModelResult { @@ -310,35 +343,10 @@ func (h *Host) ModelsForAuth(ctx context.Context, auth *coreauth.Auth) AuthModel return AuthModelResult{} } for _, record := range h.activeRecords() { - modelProvider := record.plugin.Capabilities.ModelProvider - if modelProvider == nil || h.isPluginFused(record.id) { - continue - } - if !executorScopeAllowsOAuthModels(record.plugin.Capabilities) { + if !h.authModelProviderMatches(record, providerKey) { continue } - authProvider := record.plugin.Capabilities.AuthProvider - if authProvider != nil { - identifier, okIdentifier := h.callAuthProviderIdentifier(record.id, authProvider) - if !okIdentifier || normalizeProviderID(identifier) != providerKey { - continue - } - } else { - recordProvider := normalizeProviderID(h.modelProvider(record.id)) - if recordProvider == "" { - executor := record.plugin.Capabilities.Executor - if executor != nil { - candidate, okCandidate := h.executorProvider(record, executor) - if okCandidate { - recordProvider = candidate - } - } - } - if recordProvider != providerKey { - continue - } - } - resp, errModels := h.callModelsForAuth(ctx, record, modelProvider, auth) + resp, errModels := h.callModelsForAuth(ctx, record, record.plugin.Capabilities.ModelProvider, auth) if errModels != nil { log.Warnf("pluginhost: models for auth %s failed: %v", auth.ID, errModels) return AuthModelResult{Handled: true, Err: errModels} diff --git a/internal/pluginhost/adapters_auth.go b/internal/pluginhost/adapters_auth.go index e9ba8d14..f55d2984 100644 --- a/internal/pluginhost/adapters_auth.go +++ b/internal/pluginhost/adapters_auth.go @@ -3,13 +3,92 @@ package pluginhost import ( "bytes" "context" + "errors" + "fmt" + "io" "net/http" "strings" + "github.com/klauspost/compress/zstd" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" sdkaccess "github.com/router-for-me/CLIProxyAPI/v8/sdk/access" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" ) +// Decoding runs before authentication, so what it allocates, produces and runs at once is bounded. +// 8 MiB of window covers every zstd level up to 19; a larger declared window is allocated up front. +const ( + maxAuthZstdWindowBytes = 8 << 20 + maxAuthDecodedBodyBytes = 64 << 20 + // maxAuthSpeechBodyBytes is the speech handlers' own decoded cap (maxXAISpeechBody). + maxAuthSpeechBodyBytes = 1 << 20 + // maxAuthContentCodings bounds the decoder layers run before authentication; a longer + // Content-Encoding list is left to the handler and the plugin gets no DecodedBody. + maxAuthContentCodings = 2 +) + +var ( + authDecodeSlots = make(chan struct{}, 2) + errAuthBodyTooLarge = errors.New("request body too large") +) + +// decodeAuthBody decodes a Content-Encoding list the way handlers.ReadRequestBody does (zstd, +// last to first), but with the window bound the handler's decoder lacks and a decoded size of at +// most limit bytes. +func decodeAuthBody(body []byte, encoding string, limit int64) ([]byte, error) { + parts := strings.Split(encoding, ",") + for i := len(parts) - 1; i >= 0; i-- { + switch enc := strings.ToLower(strings.TrimSpace(parts[i])); enc { + case "", "identity": + case "zstd": + decoder, errNewReader := zstd.NewReader(bytes.NewReader(body), zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxWindow(maxAuthZstdWindowBytes)) + if errNewReader != nil { + return nil, errNewReader + } + out, errReadAll := io.ReadAll(io.LimitReader(decoder, limit+1)) + decoder.Close() + if errors.Is(errReadAll, zstd.ErrWindowSizeExceeded) || errors.Is(errReadAll, zstd.ErrDecoderSizeExceeded) || int64(len(out)) > limit { + return nil, errAuthBodyTooLarge + } + if errReadAll != nil { + return nil, errReadAll + } + body = out + default: + return nil, fmt.Errorf("unsupported request content encoding: %s", enc) + } + } + return body, nil +} + +// handlerDecodesBody reports whether the route's handler reads the body through +// handlers.ReadRequestBody. Other handlers read the raw bytes, and a plugin judging a +// decoded body there could approve a model the handler never runs. Only POST is decoded: +// GET on the responses routes is a websocket upgrade whose body the handler never reads. +// NOTE: keep in step with the ReadRequestBody callers. +func handlerDecodesBody(r *http.Request) bool { + if r.Method != http.MethodPost { + return false + } + switch r.URL.Path { + case "/v1/chat/completions", "/v1/completions", "/v1/responses", "/v1/responses/compact", + "/backend-api/codex/responses", "/backend-api/codex/responses/compact", + "/v1/images/generations", "/v1/audio/speech", "/v1/tts", + "/v1/videos", "/v1/videos/generations", "/v1/videos/edits", "/v1/videos/extensions": + return true + case "/v1/images/edits": // only the JSON branch; multipart is parsed raw + return strings.HasPrefix(strings.ToLower(strings.TrimSpace(r.Header.Get("Content-Type"))), "application/json") + case "/openai/v1/videos": // forms are read raw + contentType := r.Header.Get("Content-Type") + if i := strings.IndexAny(contentType, " ;"); i >= 0 { + contentType = contentType[:i] + } + contentType = strings.ToLower(strings.TrimSpace(contentType)) + return contentType != "multipart/form-data" && contentType != "application/x-www-form-urlencoded" + } + return false +} + func (h *Host) RegisterFrontendAuthProviders() { if h == nil { return @@ -125,16 +204,62 @@ func (a *accessAdapter) Authenticate(ctx context.Context, r *http.Request) (resu body, errReadAll := readAndRestoreRequestBody(r) if errReadAll != nil { + var maxBytesErr *http.MaxBytesError + if errors.As(errReadAll, &maxBytesErr) { + return nil, sdkaccess.NewRequestEntityTooLargeError(errReadAll) + } return nil, sdkaccess.NewInternalAuthError("failed to read plugin auth request body", errReadAll) } + headers := cloneHeader(r.Header) + var decodedBody []byte + cfg := a.host.currentRuntimeConfig() + // With image generation disabled the image handlers answer 404 before reading the body. + imagesOff := cfg != nil && cfg.DisableImageGeneration == config.DisableImageGenerationAll && strings.HasPrefix(r.URL.Path, "/v1/images/") + if encoding := strings.TrimSpace(headers.Get("Content-Encoding")); encoding != "" && !strings.EqualFold(encoding, "identity") && strings.Count(encoding, ",") < maxAuthContentCodings && !imagesOff && handlerDecodesBody(r) { + // Decoding runs before authentication, so it is bounded; a slot is held until the plugin call finishes. + var releaseSlot func() + select { + case authDecodeSlots <- struct{}{}: + ctx, releaseSlot = withPluginCallFinish(ctx, func() { <-authDecodeSlots }) + defer releaseSlot() + case <-ctx.Done(): + return nil, sdkaccess.NewInternalAuthError("plugin auth request canceled", ctx.Err()) + } + // The server middleware limits only compressed bytes; ReadRequestBody applies the configured + // limit to decoded bytes after authentication, so apply it here too. + limit := int64(maxAuthDecodedBodyBytes) + if cfg != nil && cfg.MaxRequestBodyBytes > 0 && cfg.MaxRequestBodyBytes < limit { + limit = cfg.MaxRequestBodyBytes + } + if (r.URL.Path == "/v1/audio/speech" || r.URL.Path == "/v1/tts") && limit > maxAuthSpeechBodyBytes { + limit = maxAuthSpeechBodyBytes + } + decoded, errDecode := decodeAuthBody(body, encoding, limit) + if errors.Is(errDecode, errAuthBodyTooLarge) { + return nil, sdkaccess.NewRequestEntityTooLargeError(errDecode) + } + // Any other failure leaves DecodedBody unset; the handler falls back to the raw body or rejects it. + if errDecode == nil { + decodedBody = decoded + } else { + releaseSlot() + } + } 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), + Method: r.Method, + Path: r.URL.Path, + Headers: headers, + Query: cloneValues(r.URL.Query()), + Body: bytes.Clone(body), + DecodedBody: decodedBody, }) - if errAuthenticate != nil || !resp.Authenticated { + if errAuthenticate != nil { + return nil, sdkaccess.NewNotHandledError() + } + if !resp.Authenticated { + if resp.DenyStatusCode > 0 || strings.TrimSpace(resp.DenyMessage) != "" { + return nil, sdkaccess.NewDeniedError(resp.DenyMessage, resp.DenyStatusCode) + } return nil, sdkaccess.NewNotHandledError() } providerID := a.Identifier() diff --git a/internal/pluginhost/adapters_executors.go b/internal/pluginhost/adapters_executors.go index a83f638e..46b719d0 100644 --- a/internal/pluginhost/adapters_executors.go +++ b/internal/pluginhost/adapters_executors.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "sort" + "strconv" "strings" "sync" @@ -18,6 +19,7 @@ import ( coreexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator" + "github.com/tidwall/gjson" ) type executorManager interface { @@ -44,8 +46,21 @@ func (h *Host) RegisterExecutors(manager executorManager, modelRegistry modelPro claimedModels := make(map[string]struct{}) claimedProviders := make(map[string]string) for _, registration := range registrations { - if !registration.hasExecutor { - appendModelsForProvider(providerModels, registration.provider, registration.models) + if registration.hasExecutor { + continue + } + appendModelsForProvider(providerModels, registration.provider, registration.models) + // Reserve native catalog models before plugin arbitration, independently of + // credential publication or exclusions. This preserves native priority without + // publishing availability, which remains credential-bound. + if h.providerHasNativeExecutor(manager, registration.provider) { + for _, model := range registration.models { + if model != nil { + if modelID := strings.TrimSpace(model.ID); modelID != "" { + claimedModels[modelID] = struct{}{} + } + } + } } } for _, record := range records { @@ -165,6 +180,7 @@ func (h *Host) commitExecutorState(snap *Snapshot, manager executorManager, mode } } h.executorModelClientIDs = nextModelClients + h.mu.Unlock() for _, registration := range registrations { if registration.adapter == nil || registration.provider == "" { @@ -179,7 +195,6 @@ func (h *Host) commitExecutorState(snap *Snapshot, manager executorManager, mode } manager.UnregisterExecutor(provider) } - h.mu.Unlock() if modelRegistry == nil { return @@ -408,10 +423,77 @@ func (a *executorAdapter) prepareExecutorCall(req coreexecutor.Request, opts cor return preparedExecutorCall{}, errOutput } + baseModel := strings.TrimSpace(req.Model) + modelInfo := a.lookupPluginModelInfo(baseModel) + if parsed := thinking.ParseSuffix(req.Model); parsed.HasSuffix && modelInfo == nil { + baseModel = strings.TrimSpace(parsed.ModelName) + modelInfo = a.lookupPluginModelInfo(baseModel) + } + if baseModel == "" { + baseModel = req.Model + } + nativeReq := req nativeOpts := opts + var updatesChanged bool if inputRequested != "" && inputRequested != inputFormat { - nativeReq.Payload = sdktranslator.TranslateRequest(inputRequested, inputFormat, req.Model, req.Payload, opts.Stream) + translated := sdktranslator.TranslateRequestEnvelope(context.Background(), inputRequested, inputFormat, sdktranslator.RequestEnvelope{ + Format: inputRequested, + Model: baseModel, + Stream: opts.Stream, + Body: req.Payload, + }) + if translated.Err != nil { + return preparedExecutorCall{}, translated.Err + } + nativeReq.Payload = translated.Body + updatesChanged = translated.ConfigurationUpdatesChanged + } + hasPluginThinkingTarget := modelInfo != nil || + registry.LookupModelInfo(baseModel, a.provider) != nil || + thinking.GetProviderApplier(a.provider) != nil + if hasPluginThinkingTarget && baseModel != "" && gjson.GetBytes(nativeReq.Payload, "model").Exists() { + nativeReq.Payload = helps.SetStringIfDifferent(nativeReq.Payload, "model", baseModel) + } + originalTranslated := nativeReq.Payload + if len(nativeReq.Payload) > 0 { + body, errThinking := helps.ApplyRequestThinkingWithModelInfo( + nativeReq.Payload, + req, + opts, + inputRequested.String(), + inputFormat.String(), + a.provider, + modelInfo, + updatesChanged, + ) + if errThinking != nil { + return preparedExecutorCall{}, errThinking + } + nativeReq.Payload = body + } + // User payload rules are the last host-side mutation before the plugin sees + // the request, matching on the suffix-free model like the built-in executors + // and using the canonical client format as source protocol. + if a.host != nil { + if cfg := a.host.currentRuntimeConfig(); cfg != nil { + finalizeOpts := opts + if inputRequested != "" { + finalizeOpts.SourceFormat = inputRequested + } + if len(req.Metadata) > 0 { + finalizeOpts.Metadata = mergeExecutorMetadata(req.Metadata, opts.Metadata) + } + root := "" + if inputFormat == sdktranslator.FormatAntigravity { + root = "request" + } + finalize := helps.NewPayloadFinalizer(cfg, a.provider, baseModel, inputFormat.String(), root, originalTranslated, req, finalizeOpts) + nativeReq.Payload = finalize(nativeReq.Payload) + } + } + if hasPluginThinkingTarget { + nativeReq.Model = baseModel } nativeReq.Format = outputFormat nativeOpts.SourceFormat = inputFormat @@ -427,6 +509,24 @@ func (a *executorAdapter) prepareExecutorCall(req coreexecutor.Request, opts cor }, nil } +func (a *executorAdapter) lookupPluginModelInfo(baseModel string) *registry.ModelInfo { + baseModel = strings.TrimSpace(baseModel) + if a == nil || a.host == nil || baseModel == "" { + return nil + } + for _, model := range a.host.ModelsForProvider(a.provider) { + if model != nil && strings.EqualFold(strings.TrimSpace(model.ID), baseModel) { + return model + } + } + for _, model := range a.host.modelRegistration(a.pluginID).models { + if model != nil && strings.EqualFold(strings.TrimSpace(model.ID), baseModel) { + return model + } + } + return nil +} + func (a *executorAdapter) RequestToFormat(req coreexecutor.Request, opts coreexecutor.Options) sdktranslator.Format { if a == nil { return "" @@ -785,7 +885,6 @@ func (a *executorAdapter) observeAndTranslateExecutorStream(ctx context.Context, var streamErr error var publishOnce sync.Once var lineBuffer []byte - const maxLineBufferSize = 64 * 1024 publishResult := func() { publishOnce.Do(func() { @@ -857,10 +956,6 @@ func (a *executorAdapter) observeAndTranslateExecutorStream(ctx context.Context, lineBuffer = nil } } - if len(lineBuffer) > maxLineBufferSize { - helps.ObservePluginExecutorStreamUsage(prepared.outputFormat.String(), lineBuffer, &streamUsage) - lineBuffer = nil - } } select { @@ -984,13 +1079,52 @@ func (a *executorAdapter) CountTokens(ctx context.Context, auth *coreauth.Auth, if errCountTokens != nil { return coreexecutor.Response{}, errCountTokens } + var payload []byte + translatedTokenCount := false + if prepared.outputFormat != prepared.requestedFormat { + // A token-count reply is not a model response; translating it as one + // wraps the count in a full completion message. + if count, okCount := pluginExecutorTokenCount(prepared.outputFormat, pluginResp.Payload); okCount { + if sdktranslator.HasTokenCountTransformer(prepared.outputFormat, prepared.requestedFormat) { + translated := sdktranslator.TranslateTokenCount(ctx, prepared.outputFormat, prepared.requestedFormat, count, pluginResp.Payload) + payload = translated + translatedTokenCount = true + } + } + } + if !translatedTokenCount { + payload = a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil) + } return coreexecutor.Response{ - Payload: a.translateExecutorResponse(ctx, prepared, pluginResp.Payload, false, nil), + Payload: payload, Metadata: cloneAnyMap(pluginResp.Metadata), Headers: cloneHeader(pluginResp.Headers), }, nil } +func pluginExecutorTokenCount(format sdktranslator.Format, payload []byte) (int64, bool) { + for _, path := range []string{"choices", "content", "candidates", "output", "response.output"} { + if gjson.GetBytes(payload, path).Exists() { + return 0, false + } + } + // response.usage.input_tokens is the native Codex count shape. + paths := []string{"usage.prompt_tokens", "input_tokens", "response.usage.input_tokens", "totalTokens", "total_tokens"} + if format == sdktranslator.FormatGemini { + paths = []string{"totalTokens", "usageMetadata.promptTokenCount", "input_tokens", "total_tokens"} + } + for _, path := range paths { + value := gjson.GetBytes(payload, path) + if value.Exists() && value.Type == gjson.Number { + count, errParse := strconv.ParseInt(value.Raw, 10, 64) + if errParse == nil && count >= 0 { + return count, true + } + } + } + return 0, false +} + 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()) diff --git a/internal/pluginhost/adapters_interceptors.go b/internal/pluginhost/adapters_interceptors.go index 055d5254..657dd5fc 100644 --- a/internal/pluginhost/adapters_interceptors.go +++ b/internal/pluginhost/adapters_interceptors.go @@ -92,7 +92,7 @@ func (h *Host) InterceptRequestBeforeAuth(ctx context.Context, req pluginapi.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 h.interceptRequest(ctx, req, true, "RequestInterceptor.InterceptRequestBeforeAuth", func(interceptor pluginapi.RequestInterceptor, ctx context.Context, req pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error) { return interceptor.InterceptRequestBeforeAuth(ctx, req) }, skipPluginID) } @@ -102,12 +102,12 @@ func (h *Host) InterceptRequestAfterAuth(ctx context.Context, req pluginapi.Requ } 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 h.interceptRequest(ctx, req, false, "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 { +func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterceptRequest, beforeAuth bool, method string, invoke func(pluginapi.RequestInterceptor, context.Context, pluginapi.RequestInterceptRequest) (pluginapi.RequestInterceptResponse, error), skipPluginID string) pluginapi.RequestInterceptResponse { current := pluginapi.RequestInterceptResponse{ Headers: cloneHeader(req.Headers), } @@ -122,12 +122,24 @@ func (h *Host) interceptRequest(ctx context.Context, req pluginapi.RequestInterc } nextReq := req nextReq.Headers = cloneHeader(current.Headers) - if len(currentBase) > 0 { + if requestInterceptorOmitsBody(record, beforeAuth) { + nextReq.Body = nil + } else if len(currentBase) > 0 { nextReq.Body = bytes.Clone(currentBase) } else { nextReq.Body = nil } nextReq.Metadata = cloneInterceptorMetadata(req.Metadata) + if nextReq.CallerScope == "" { + if scope, ok := nextReq.Metadata[coreexecutor.CallerScopeMetadataKey].(string); ok { + nextReq.CallerScope = strings.TrimSpace(scope) + } + } else if _, ok := nextReq.Metadata[coreexecutor.CallerScopeMetadataKey]; !ok { + if nextReq.Metadata == nil { + nextReq.Metadata = make(map[string]any, 1) + } + nextReq.Metadata[coreexecutor.CallerScopeMetadataKey] = strings.TrimSpace(nextReq.CallerScope) + } if current.Path != "" { if nextReq.Metadata == nil { nextReq.Metadata = make(map[string]any, 1) @@ -373,6 +385,17 @@ func (h *Host) StreamChunkPayloadIncludesHistory() bool { return false } +func requestInterceptorOmitsBody(record capabilityRecord, beforeAuth bool) bool { + if record.plugin.SchemaVersion < pluginabi.SchemaVersionRequestInterceptorOmitBody { + return false + } + caps := record.plugin.Capabilities + if caps.RequestInterceptor == nil { + return false + } + return requestInterceptorWantsOmitBody(caps, beforeAuth) +} + func streamChunkOmitsHistory(schemaVersion uint32) bool { return schemaVersion >= pluginabi.SchemaVersionStreamChunkOmitHistory } @@ -392,42 +415,20 @@ func (h *Host) HasRequestInterceptors() bool { 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) + origBody := r.Body + body, errReadAll := io.ReadAll(origBody) if errReadAll != nil { - r.Body = io.NopCloser(bytes.NewReader(body)) + r.Body = struct { + io.Reader + io.Closer + }{ + Reader: io.MultiReader(bytes.NewReader(body), origBody), + Closer: origBody, + } return nil, errReadAll } r.Body = io.NopCloser(bytes.NewReader(body)) diff --git a/internal/pluginhost/adapters_usage_translation.go b/internal/pluginhost/adapters_usage_translation.go index 5458302d..77236eb2 100644 --- a/internal/pluginhost/adapters_usage_translation.go +++ b/internal/pluginhost/adapters_usage_translation.go @@ -167,31 +167,55 @@ func (a *usageAdapter) HandleUsage(ctx context.Context, record coreusage.Record) if traceID == "" { traceID = strings.TrimSpace(logging.GetRequestID(ctx)) } + selectionID := strings.TrimSpace(record.SelectionID) + if selectionID == "" { + selectionID = strings.TrimSpace(coreusage.SelectionIDFromContext(ctx)) + } + accessProvider := strings.TrimSpace(record.AccessProvider) + if accessProvider == "" { + accessProvider = strings.TrimSpace(coreusage.AccessProviderFromContext(ctx)) + } + isNativeKey := record.IsNativeKey + if !isNativeKey { + if nativeExplicit, ok := coreusage.IsNativeKeyFromContext(ctx); ok { + isNativeKey = nativeExplicit + } else if accessProvider != "" { + isNativeKey = coreusage.IsNativeAccessProvider(accessProvider) + } + } + apiKey := strings.TrimSpace(record.APIKey) + if apiKey == "" { + apiKey = strings.TrimSpace(coreusage.APIKeyFromContext(ctx)) + } plugin.HandleUsage(ctx, pluginapi.UsageRecord{ - RequestID: requestID, - TraceID: traceID, - Provider: record.Provider, - BaseURL: record.BaseURL, - ExecutorType: record.ExecutorType, - Model: record.Model, - Alias: record.Alias, - APIKey: record.APIKey, - SessionID: sessionID, - ParentSessionID: parentSessionID, - AuthID: record.AuthID, - AuthIndex: record.AuthIndex, - AuthType: record.AuthType, - Source: record.Source, - ReasoningEffort: record.ReasoningEffort, - ServiceTier: record.ServiceTier, - ResponseServiceTier: record.ResponseServiceTier, - ResponseModel: record.ResponseModel, - Generate: coreusage.GenerateEnabled(record.Generate), - Stream: record.Stream, - RequestedAt: record.RequestedAt, - Latency: record.Latency, - TTFT: record.TTFT, - Failed: record.Failed, + RequestID: requestID, + TraceID: traceID, + SelectionID: selectionID, + Provider: record.Provider, + BaseURL: record.BaseURL, + ExecutorType: record.ExecutorType, + Model: record.Model, + Alias: record.Alias, + APIKey: apiKey, + IsNativeKey: isNativeKey, + AccessProvider: accessProvider, + SessionID: sessionID, + ParentSessionID: parentSessionID, + AuthID: record.AuthID, + AuthIndex: record.AuthIndex, + AuthType: record.AuthType, + Source: record.Source, + ReasoningEffort: record.ReasoningEffort, + ServiceTier: record.ServiceTier, + EffectiveServiceTier: record.EffectiveServiceTier, + ResponseServiceTier: record.ResponseServiceTier, + ResponseModel: record.ResponseModel, + Generate: coreusage.GenerateEnabled(record.Generate), + Stream: record.Stream, + RequestedAt: record.RequestedAt, + Latency: record.Latency, + TTFT: record.TTFT, + Failed: record.Failed, Failure: pluginapi.UsageFailure{ StatusCode: record.Fail.StatusCode, Body: record.Fail.Body, diff --git a/internal/pluginhost/auth_callbacks.go b/internal/pluginhost/auth_callbacks.go index 558b67d9..2e13a504 100644 --- a/internal/pluginhost/auth_callbacks.go +++ b/internal/pluginhost/auth_callbacks.go @@ -1,6 +1,7 @@ package pluginhost import ( + "bytes" "context" "encoding/json" "fmt" @@ -12,6 +13,7 @@ import ( "strings" "time" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" ) @@ -132,6 +134,78 @@ func (h *Host) callHostAuthSave(ctx context.Context, request []byte) ([]byte, er }) } +func (h *Host) callHostAuthModels(ctx context.Context, request []byte) ([]byte, error) { + _ = ctx + var req pluginapi.HostAuthModelsRequest + if len(bytesTrimSpace(request)) > 0 { + if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { + return nil, fmt.Errorf("decode host auth models request: %w", errUnmarshal) + } + } + reg := registry.GetGlobalRegistry() + authIndex := strings.TrimSpace(req.AuthIndex) + if authIndex != "" { + auth, errGet := h.authByIndex(authIndex) + if errGet != nil { + return nil, errGet + } + return marshalRPCResult(pluginapi.HostAuthModelsResponse{ + Credentials: []pluginapi.HostAuthModels{buildHostAuthModels(auth, reg)}, + }) + } + manager := h.currentAuthManager() + if manager == nil { + return marshalRPCResult(pluginapi.HostAuthModelsResponse{ + Credentials: []pluginapi.HostAuthModels{}, + }) + } + auths := manager.List() + credentials := make([]pluginapi.HostAuthModels, 0, len(auths)) + for _, auth := range auths { + if auth == nil { + continue + } + credentials = append(credentials, buildHostAuthModels(auth, reg)) + } + sort.Slice(credentials, func(i, j int) bool { + leftID := strings.ToLower(credentials[i].ID) + rightID := strings.ToLower(credentials[j].ID) + if leftID != rightID { + return leftID < rightID + } + if credentials[i].ID != credentials[j].ID { + return credentials[i].ID < credentials[j].ID + } + return credentials[i].AuthIndex < credentials[j].AuthIndex + }) + return marshalRPCResult(pluginapi.HostAuthModelsResponse{ + Credentials: credentials, + }) +} + +func buildHostAuthModels(auth *coreauth.Auth, reg *registry.ModelRegistry) pluginapi.HostAuthModels { + auth.EnsureIndex() + var regModels []*registry.ModelInfo + if reg != nil && strings.TrimSpace(auth.ID) != "" { + regModels = reg.GetModelsForClient(auth.ID) + } + models := make([]string, 0, len(regModels)) + for _, m := range regModels { + if m == nil { + continue + } + if id := strings.TrimSpace(m.ID); id != "" { + models = append(models, id) + } + } + return pluginapi.HostAuthModels{ + ID: auth.ID, + AuthIndex: auth.Index, + Provider: strings.TrimSpace(auth.Provider), + Models: models, + } +} + func (h *Host) listAuthFiles() ([]pluginapi.HostAuthFileEntry, error) { manager := h.currentAuthManager() if manager != nil { @@ -298,20 +372,20 @@ func (h *Host) saveAuthFile(ctx context.Context, name string, data []byte) (stri dst = abs } } - auth, errBuild := h.buildAuthFromFileData(dst, data) + auths, errBuild := h.buildAuthsFromFileData(ctx, dst, data) if errBuild != nil { return "", errBuild } if errWrite := os.WriteFile(dst, data, 0o600); errWrite != nil { return "", fmt.Errorf("failed to write auth file: %w", errWrite) } - if errUpsert := h.upsertAuthRecord(ctx, auth); errUpsert != nil { + if errUpsert := h.upsertAuthRecordsForPath(ctx, dst, auths); errUpsert != nil { return "", errUpsert } return dst, nil } -func (h *Host) buildAuthFromFileData(path string, data []byte) (*coreauth.Auth, error) { +func (h *Host) buildAuthsFromFileData(ctx context.Context, path string, data []byte) ([]*coreauth.Auth, error) { if strings.TrimSpace(path) == "" { return nil, fmt.Errorf("auth path is empty") } @@ -327,66 +401,282 @@ func (h *Host) buildAuthFromFileData(path string, data []byte) (*coreauth.Auth, return nil, fmt.Errorf("invalid auth file: %w", errUnmarshal) } coreauth.NormalizeCredentialMetadata(metadata) + if errWeight := coreauth.ValidateAuthWeight(&coreauth.Auth{Metadata: metadata}); errWeight != nil { + return nil, fmt.Errorf("invalid auth weight: %w", errWeight) + } provider, _ := metadata["type"].(string) - if strings.TrimSpace(provider) == "" { + provider = strings.TrimSpace(provider) + + fileName := filepath.Base(path) + defaultAuthID := h.authIDForPath(path) + if defaultAuthID == "" { + defaultAuthID = path + } + disabled, okDisabled := parseBoolValue(metadata["disabled"]) + now := time.Now().UTC() + + parsedAuths, handled, errParse := h.ParseAuths(ctx, pluginapi.AuthParseRequest{ + Provider: provider, + Path: path, + FileName: defaultAuthID, + RawJSON: data, + }) + if errParse != nil { + return nil, fmt.Errorf("invalid auth file: %w", errParse) + } + if handled { + out := make([]*coreauth.Auth, 0, len(parsedAuths)) + for _, auth := range parsedAuths { + if auth == nil { + continue + } + if errWeight := coreauth.ValidateAuthWeight(auth); errWeight != nil { + return nil, fmt.Errorf("invalid auth weight: %w", errWeight) + } + out = append(out, auth) + } + for index, auth := range out { + coreauth.NormalizeCredentialMetadata(auth.Metadata) + if len(out) == 1 { + if auth.Metadata == nil { + auth.Metadata = make(map[string]any, len(metadata)) + } + for key, value := range metadata { + if _, exists := auth.Metadata[key]; !exists { + auth.Metadata[key] = value + } + } + if pts, ok := auth.Storage.(*pluginTokenStorage); ok && pts != nil && len(bytes.TrimSpace(pts.rawJSON)) == 0 { + pts.rawJSON = bytes.Clone(data) + } + } + if len(out) > 1 { + coreauth.MarkPluginVirtualAuth(auth, path, index) + } + if strings.TrimSpace(auth.FileName) == "" { + auth.FileName = fileName + } + if len(out) == 1 && (strings.TrimSpace(auth.ID) == "" || auth.ID == defaultAuthID) { + if existing := h.findExistingAuthForSave(auth.ID, path, true); existing != nil && strings.TrimSpace(existing.ID) != "" { + auth.ID = strings.TrimSpace(existing.ID) + } + } + if strings.TrimSpace(auth.ID) == "" { + auth.ID = defaultAuthID + } + auth.CreatedAt = now + auth.UpdatedAt = now + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + auth.Attributes[coreauth.AttributePath] = path + auth.Attributes[coreauth.AttributeSource] = path + auth.Attributes[coreauth.AttributeSourceBackend] = coreauth.AuthSourceFile + if okDisabled && disabled { + auth.Disabled = true + auth.Status = coreauth.StatusDisabled + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["disabled"] = true + } + if p, ok := metadata["proxy_url"].(string); ok && auth.ProxyURL == "" { + auth.ProxyURL = strings.TrimSpace(p) + } + if pref, ok := metadata["prefix"].(string); ok && auth.Prefix == "" { + auth.Prefix = strings.Trim(strings.TrimSpace(pref), "/") + } + if errWeight := coreauth.ApplyAuthWeightMetadata(auth, metadata); errWeight != nil { + return nil, fmt.Errorf("invalid auth weight: %w", errWeight) + } + coreauth.ApplyAuthPriorityMetadata(auth, metadata) + if _, inherited := auth.Attributes[coreauth.AttributeFilePriority]; inherited { + if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { + setter.SetMetadata(auth.Metadata) + } + } + coreauth.ApplyCustomHeadersFromMetadata(auth) + if existing := h.findExistingAuthForSave(auth.ID, path, len(out) == 1); existing != nil { + preserveExistingAuthState(auth, existing) + } + } + return out, nil + } + + if provider == "" { provider = "unknown" } label := provider if email, ok := metadata["email"].(string); ok && strings.TrimSpace(email) != "" { label = strings.TrimSpace(email) } - authID := h.authIDForPath(path) - if authID == "" { - authID = path + authID := defaultAuthID + existing := h.findExistingAuthForSave(authID, path, true) + if existing != nil && strings.TrimSpace(existing.ID) != "" { + authID = strings.TrimSpace(existing.ID) } status := coreauth.StatusActive - disabled := false - if d, okDisabled := parseBoolValue(metadata["disabled"]); okDisabled && d { - disabled = true + if okDisabled && disabled { status = coreauth.StatusDisabled + } else { + disabled = false } auth := &coreauth.Auth{ ID: authID, Provider: provider, - FileName: filepath.Base(path), + FileName: fileName, Label: label, Status: status, Disabled: disabled, Attributes: map[string]string{ - "path": path, - "source": path, + coreauth.AttributePath: path, + coreauth.AttributeSource: path, + coreauth.AttributeSourceBackend: coreauth.AuthSourceFile, }, Metadata: metadata, - CreatedAt: time.Now().UTC(), - UpdatedAt: time.Now().UTC(), + CreatedAt: now, + UpdatedAt: now, } - if manager := h.currentAuthManager(); manager != nil { - if existing, ok := manager.GetByID(authID); ok { - auth.CreatedAt = existing.CreatedAt - auth.LastRefreshedAt = existing.LastRefreshedAt - auth.NextRetryAfter = existing.NextRetryAfter - auth.Runtime = existing.Runtime - } + if existing != nil { + preserveExistingAuthState(auth, existing) } - if errWeight := coreauth.ValidateAuthWeight(auth); errWeight != nil { + if errWeight := coreauth.ApplyAuthWeightMetadata(auth, metadata); errWeight != nil { return nil, fmt.Errorf("invalid auth weight: %w", errWeight) } + coreauth.ApplyAuthPriorityMetadata(auth, metadata) coreauth.ApplyCustomHeadersFromMetadata(auth) - return auth, nil + return []*coreauth.Auth{auth}, nil } -func (h *Host) upsertAuthRecord(ctx context.Context, auth *coreauth.Auth) error { +func preserveExistingAuthState(dst, existing *coreauth.Auth) { + if dst == nil || existing == nil { + return + } + if !existing.CreatedAt.IsZero() { + dst.CreatedAt = existing.CreatedAt + } + dst.LastRefreshedAt = existing.LastRefreshedAt + dst.NextRetryAfter = existing.NextRetryAfter + dst.NextRefreshAfter = existing.NextRefreshAfter + dst.Runtime = existing.Runtime + if !existing.Disabled && existing.Status != coreauth.StatusDisabled && !dst.Disabled && dst.Status != coreauth.StatusDisabled { + if len(dst.ModelStates) == 0 && len(existing.ModelStates) > 0 { + dst.ModelStates = existing.ModelStates + } + } +} + +func (h *Host) findExistingAuthForSave(authID, path string, allowPathFallback bool) *coreauth.Auth { manager := h.currentAuthManager() - if manager == nil || auth == nil { + if manager == nil { return nil } - if existing, ok := manager.GetByID(auth.ID); ok { - auth.CreatedAt = existing.CreatedAt - _, errUpdate := manager.Update(ctx, auth) - return errUpdate + authID = strings.TrimSpace(authID) + defaultID := h.authIDForPath(path) + if authID != "" && (!allowPathFallback || authID != defaultID) { + if existing, ok := manager.GetByID(authID); ok && existing != nil { + return existing + } + } + if !allowPathFallback { + return nil + } + var matched *coreauth.Auth + for _, candidate := range manager.List() { + if candidate == nil || isRuntimeOnlyAuth(candidate) || coreauth.IsPluginVirtualAuth(candidate) { + continue + } + if !authMatchesPhysicalFile(candidate, path) { + continue + } + if matched == nil || (matched.ID == defaultID && candidate.ID != defaultID) { + matched = candidate + } } - _, errRegister := manager.Register(ctx, auth) - return errRegister + if matched != nil { + return matched + } + if authID != "" { + if existing, ok := manager.GetByID(authID); ok && existing != nil { + return existing + } + } + return nil +} + +func authMatchesPhysicalFile(auth *coreauth.Auth, path string) bool { + if auth == nil || strings.TrimSpace(path) == "" { + return false + } + if sameHostAuthFilePath(authAttribute(auth, coreauth.AttributePath), path) || + sameHostAuthFilePath(authAttribute(auth, coreauth.AttributeSource), path) || + sameHostAuthFilePath(authAttribute(auth, coreauth.AttributeVirtualSource), path) { + return true + } + if strings.TrimSpace(authAttribute(auth, coreauth.AttributePath)) == "" { + return strings.EqualFold(strings.TrimSpace(auth.FileName), filepath.Base(path)) + } + return false +} + +func sameHostAuthFilePath(left, right string) bool { + left = strings.TrimSpace(left) + right = strings.TrimSpace(right) + if left == "" || right == "" { + return false + } + if absLeft, errLeft := filepath.Abs(left); errLeft == nil { + left = absLeft + } + if absRight, errRight := filepath.Abs(right); errRight == nil { + right = absRight + } + left = filepath.Clean(left) + right = filepath.Clean(right) + if runtime.GOOS == "windows" { + return strings.EqualFold(left, right) + } + return left == right +} + +func (h *Host) upsertAuthRecordsForPath(ctx context.Context, path string, auths []*coreauth.Auth) error { + manager := h.currentAuthManager() + if manager == nil { + return nil + } + keepIDs := make(map[string]struct{}, len(auths)) + for _, auth := range auths { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + continue + } + keepIDs[auth.ID] = struct{}{} + if existing, ok := manager.GetByID(auth.ID); ok && existing != nil { + if !existing.CreatedAt.IsZero() { + auth.CreatedAt = existing.CreatedAt + } + if _, errUpdate := manager.Update(ctx, auth); errUpdate != nil { + return errUpdate + } + continue + } + if _, errRegister := manager.Register(ctx, auth); errRegister != nil { + return errRegister + } + } + skipPersistCtx := coreauth.WithSkipPersist(ctx) + for _, existing := range manager.List() { + if existing == nil || isRuntimeOnlyAuth(existing) { + continue + } + if _, keep := keepIDs[existing.ID]; keep { + continue + } + if authMatchesPhysicalFile(existing, path) { + manager.Remove(skipPersistCtx, existing.ID) + registry.GetGlobalRegistry().UnregisterClient(existing.ID) + } + } + return nil } func isUnsafeAuthFileName(name string) bool { @@ -526,26 +816,7 @@ func (h *Host) resolvedAuthDir() string { } func (h *Host) 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 authDir := h.resolvedAuthDir(); authDir != "" { - if rel, errRel := filepath.Rel(authDir, path); errRel == nil && rel != "" { - id = rel - } - } - if runtime.GOOS == "windows" { - id = strings.ToLower(id) - } - return id + return authIDForPath(path, h.resolvedAuthDir()) } func authEmail(auth *coreauth.Auth) string { diff --git a/internal/pluginhost/auth_provider.go b/internal/pluginhost/auth_provider.go index 34f96aaa..d316e05e 100644 --- a/internal/pluginhost/auth_provider.go +++ b/internal/pluginhost/auth_provider.go @@ -94,8 +94,24 @@ func authIDForPath(path, authDir string) string { } id := path if authDir = strings.TrimSpace(authDir); authDir != "" { - if rel, errRel := filepath.Rel(authDir, path); errRel == nil && rel != "" && !strings.HasPrefix(rel, "..") { + cleanDir := filepath.Clean(authDir) + cleanPath := filepath.Clean(path) + if rel, errRel := filepath.Rel(cleanDir, cleanPath); isRelativeWithinAuthDir(rel, errRel) { id = rel + } else { + if !filepath.IsAbs(cleanDir) { + if absDir, errAbs := filepath.Abs(cleanDir); errAbs == nil { + cleanDir = absDir + } + } + if !filepath.IsAbs(cleanPath) { + if absPath, errAbs := filepath.Abs(cleanPath); errAbs == nil { + cleanPath = absPath + } + } + if rel, errRel := filepath.Rel(cleanDir, cleanPath); isRelativeWithinAuthDir(rel, errRel) { + id = rel + } } } id = filepath.ToSlash(filepath.Clean(id)) @@ -105,6 +121,13 @@ func authIDForPath(path, authDir string) string { return id } +func isRelativeWithinAuthDir(rel string, errRel error) bool { + if errRel != nil || rel == "" || rel == "." || rel == ".." { + return false + } + return !strings.HasPrefix(rel, ".."+string(filepath.Separator)) +} + func (h *Host) AuthProviderIdentifiers() []string { if h == nil { return nil @@ -255,6 +278,28 @@ func pluginAuthParseResponseAuths(resp pluginapi.AuthParseResponse) []pluginapi. return []pluginapi.AuthData{resp.Auth} } +type oauthProxyURLContextKey struct{} + +// WithOAuthProxyURL attaches a session-scoped proxy URL to ctx for plugin OAuth login calls. +func WithOAuthProxyURL(ctx context.Context, proxyURL string) context.Context { + if ctx == nil { + ctx = context.Background() + } + proxyURL = strings.TrimSpace(proxyURL) + if proxyURL == "" { + return ctx + } + return context.WithValue(ctx, oauthProxyURLContextKey{}, proxyURL) +} + +func oauthProxyURLFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + value, _ := ctx.Value(oauthProxyURLContextKey{}).(string) + return strings.TrimSpace(value) +} + func (h *Host) StartLogin(ctx context.Context, provider string, baseURL string, metadata ...map[string]any) (pluginapi.AuthLoginStartResponse, bool, error) { record := h.authProviderRecord(provider) if record == nil { @@ -280,11 +325,16 @@ func (h *Host) callStartLogin(ctx context.Context, record capabilityRecord, prov err = fmt.Errorf("auth provider start login panic: %v", recovered) } }() + hostSummary := h.hostConfigSummary() + sessionProxyURL := oauthProxyURLFromContext(ctx) + if sessionProxyURL != "" { + hostSummary.ProxyURL = sessionProxyURL + } req := pluginapi.AuthLoginStartRequest{ Provider: normalizeProviderID(provider), BaseURL: strings.TrimSpace(baseURL), - Host: h.hostConfigSummary(), - HTTPClient: h.newHTTPClient(nil), + Host: hostSummary, + HTTPClient: h.newHTTPClientWithProxy(nil, sessionProxyURL, provider), Metadata: cloneAnyMap(metadata), } resp, errStart := authProvider.StartLogin(ctx, req) @@ -319,11 +369,16 @@ func (h *Host) callPollLogin(ctx context.Context, record capabilityRecord, provi err = fmt.Errorf("auth provider poll login panic: %v", recovered) } }() + hostSummary := h.hostConfigSummary() + sessionProxyURL := oauthProxyURLFromContext(ctx) + if sessionProxyURL != "" { + hostSummary.ProxyURL = sessionProxyURL + } req := pluginapi.AuthLoginPollRequest{ Provider: normalizeProviderID(provider), State: strings.TrimSpace(state), - Host: h.hostConfigSummary(), - HTTPClient: h.newHTTPClient(nil), + Host: hostSummary, + HTTPClient: h.newHTTPClientWithProxy(nil, sessionProxyURL, provider), Metadata: cloneAnyMap(metadata), } resp, errPoll := authProvider.PollLogin(ctx, req) diff --git a/internal/pluginhost/client_guard.go b/internal/pluginhost/client_guard.go index 3b6533d3..46dfa143 100644 --- a/internal/pluginhost/client_guard.go +++ b/internal/pluginhost/client_guard.go @@ -4,8 +4,42 @@ import ( "context" "fmt" "sync" + "sync/atomic" ) +type pluginCallFinishKey struct{} + +type pluginCallFinishHook struct { + once sync.Once + claimed atomic.Bool + fn func() +} + +func withPluginCallFinish(ctx context.Context, fn func()) (context.Context, func()) { + if ctx == nil { + ctx = context.Background() + } + hook := &pluginCallFinishHook{fn: fn} + return context.WithValue(ctx, pluginCallFinishKey{}, hook), func() { + if !hook.claimed.Load() { + hook.once.Do(hook.fn) + } + } +} + +func claimPluginCallFinish(ctx context.Context) func() { + if ctx == nil { + return nil + } + hook, ok := ctx.Value(pluginCallFinishKey{}).(*pluginCallFinishHook) + if !ok || hook == nil || !hook.claimed.CompareAndSwap(false, true) { + return nil + } + return func() { + hook.once.Do(hook.fn) + } +} + type guardedPluginClient struct { mu sync.Mutex cond *sync.Cond @@ -51,15 +85,22 @@ func (c *guardedPluginClient) Call(ctx context.Context, method string, request [ if ctx == nil { ctx = context.Background() } + onFinish := claimPluginCallFinish(ctx) result := make(chan guardedPluginCallResult, 1) go func() { defer c.release() defer func() { if recovered := recover(); recovered != nil { + if onFinish != nil { + onFinish() + } result <- guardedPluginCallResult{recovered: recovered} } }() response, errCall := inner.Call(ctx, method, request) + if onFinish != nil { + onFinish() + } result <- guardedPluginCallResult{response: response, err: errCall} }() select { diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go index b9cf9d0f..f23854cf 100644 --- a/internal/pluginhost/host.go +++ b/internal/pluginhost/host.go @@ -75,7 +75,7 @@ type Host struct { runtimeConfig *config.Config authManager *coreauth.Manager modelExecutor modelExecutor - modelClientIDs map[string]struct{} + modelClientIDs map[string]struct{} // Reserved for future use; catalog-only plugins do not populate it. executorModelClientIDs map[string]struct{} modelProviders map[string]string modelRegistrations map[string]pluginModelRegistration @@ -92,6 +92,8 @@ type Host struct { modelStreams *modelStreamBridge callbackContexts *callbackContextRegistry snapshot atomic.Value + // loadCleanupDoneHook, set by tests, runs after a late load cleanup dropped its load token. + loadCleanupDoneHook func() } func New() *Host { @@ -163,7 +165,7 @@ func (h *Host) Snapshot() *Snapshot { return emptySnapshot() } -// PluginLoaded reports whether a plugin dynamic library is still loaded by the host. +// PluginLoaded reports whether a plugin is still active in the host runtime. func (h *Host) PluginLoaded(id string) bool { if h == nil { return false @@ -181,7 +183,7 @@ func (h *Host) PluginLoaded(id string) bool { return len(h.retired[id]) > 0 } -// PluginBusy reports whether a plugin dynamic library is loaded or being loaded. +// PluginBusy reports whether a plugin is active in the host runtime or being loaded. func (h *Host) PluginBusy(id string) bool { if h == nil { return false @@ -224,6 +226,7 @@ func (h *Host) ApplyConfig(ctx context.Context, cfg *config.Config) { h.mu.Unlock() if !rc.Enabled { + h.detachUnselectedPlugins(ctx, nil) h.mu.Lock() h.managementRoutes = make(map[string]managementRouteRecord) h.resourceRoutes = make(map[string]resourceRouteRecord) @@ -248,6 +251,7 @@ func (h *Host) ApplyConfig(ctx context.Context, cfg *config.Config) { return } files = h.withLoadedPluginFallbacks(files, rc.Items, desiredVersions) + h.detachUnselectedPlugins(ctx, enabledPluginIDs(files, rc.Items)) records := make([]capabilityRecord, 0, len(files)) loadedFiles := make([]pluginFile, 0, len(files)) @@ -576,6 +580,9 @@ func (h *Host) finishPluginLoadCleanup(id string, request *pluginLoadRequest, lo go func() { h.discardLoadedPlugin(loaded) h.clearLoadingRequest(id, request) + if h.loadCleanupDoneHook != nil { + h.loadCleanupDoneHook() + } }() } @@ -638,7 +645,8 @@ func (h *Host) withLoadedPluginFallbacks(files []pluginFile, items map[string]ru return files } -// UnloadPlugin removes one plugin from the active runtime and closes its dynamic library. +// UnloadPlugin removes one plugin from the active runtime. +// The plugin image itself stays mapped for the lifetime of the process. func (h *Host) UnloadPlugin(id string) bool { return h.UnloadPluginContext(context.Background(), id) } @@ -654,6 +662,21 @@ func (h *Host) UnloadPluginContext(ctx context.Context, id string) bool { return false } defer h.unlockApply() + return h.detachPlugin(ctx, id) +} + +// detachPlugin shuts down one loaded plugin. Caller must hold the apply lock. +func (h *Host) detachPlugin(ctx context.Context, id string) bool { + if h == nil { + return false + } + id = strings.TrimSpace(id) + if id == "" { + return false + } + if ctx == nil { + ctx = context.Background() + } targets := make([]pluginUnloadTarget, 0) h.mu.Lock() @@ -710,7 +733,60 @@ func (h *Host) UnloadPluginContext(ctx context.Context, id string) bool { return true } -// ShutdownAll removes active plugin capabilities and closes all loaded dynamic libraries. +// detachUnselectedPlugins stops plugins that are no longer enabled. nil keeps none. +// Caller must hold the apply lock. +func (h *Host) detachUnselectedPlugins(ctx context.Context, selected map[string]struct{}) { + if h == nil { + return + } + h.mu.Lock() + ids := make([]string, 0, len(h.loaded)+len(h.retired)+len(h.loading)) + for id := range h.loaded { + ids = append(ids, id) + } + for id := range h.retired { + ids = append(ids, id) + } + for id := range h.loading { + ids = append(ids, id) + } + h.mu.Unlock() + sort.Strings(ids) + + seen := make(map[string]struct{}, len(ids)) + for _, id := range ids { + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + if _, ok := selected[id]; ok { + continue + } + h.detachPlugin(ctx, id) + } +} + +func enabledPluginIDs(files []pluginFile, items map[string]runtimeItemConfig) map[string]struct{} { + ids := make(map[string]struct{}, len(files)) + for _, file := range files { + id := strings.TrimSpace(file.ID) + if id == "" { + continue + } + item, ok := items[id] + if !ok { + item = defaultRuntimeItemConfig(id) + } + if !item.Enabled { + continue + } + ids[id] = struct{}{} + } + return ids +} + +// ShutdownAll removes active plugin capabilities and detaches all plugins. +// Plugin images stay mapped for the lifetime of the process. func (h *Host) ShutdownAll() { h.ShutdownAllContext(context.Background()) } diff --git a/internal/pluginhost/host_callbacks.go b/internal/pluginhost/host_callbacks.go index 4af872a9..8ab005f3 100644 --- a/internal/pluginhost/host_callbacks.go +++ b/internal/pluginhost/host_callbacks.go @@ -179,6 +179,8 @@ func (h *Host) callFromPlugin(ctx context.Context, method string, request []byte return h.callHostAuthGetRuntime(ctx, request) case pluginabi.MethodHostAuthSave: return h.callHostAuthSave(ctx, request) + case pluginabi.MethodHostAuthModels: + return h.callHostAuthModels(ctx, request) case pluginabi.MethodHostAffinityLookup: return h.callHostAffinityLookup(ctx, request) case pluginabi.MethodHostRoutingResetCooldown: @@ -205,7 +207,7 @@ func (h *Host) callHostHTTPDo(ctx context.Context, request []byte) ([]byte, erro return nil, errAcquire } defer operation.finish() - resp, errDo := h.newHTTPClient(nil).Do(operation.ctx, httpReq) + resp, errDo := h.newHTTPClientWithProxy(nil, oauthProxyURLFromContext(operation.ctx)).Do(operation.ctx, httpReq) if errDo != nil { return nil, errDo } @@ -238,7 +240,7 @@ func (h *Host) callHostHTTPDoStream(ctx context.Context, request []byte) ([]byte operation.finish() }() - resp, errDo := h.newHTTPClient(nil).DoStream(streamCtx, httpReq) + resp, errDo := h.newHTTPClientWithProxy(nil, oauthProxyURLFromContext(streamCtx)).DoStream(streamCtx, httpReq) if errDo != nil { return nil, errDo } @@ -386,9 +388,7 @@ func (h *Host) callHostStreamEmit(ctx context.Context, request []byte) ([]byte, return nil, fmt.Errorf("decode stream emit request: %w", errUnmarshal) } chunk := pluginapi.ExecutorStreamChunk{Payload: append([]byte(nil), req.Payload...)} - if req.Error != "" { - chunk.Err = fmt.Errorf("%s", req.Error) - } + chunk.Err = rpcStreamError(req.Error, req.ErrorDetail) if errEmit := h.streams.emit(ctx, req.StreamID, chunk); errEmit != nil { return nil, errEmit } @@ -400,7 +400,7 @@ func (h *Host) callHostStreamClose(request []byte) ([]byte, error) { if errUnmarshal := json.Unmarshal(request, &req); errUnmarshal != nil { return nil, fmt.Errorf("decode stream close request: %w", errUnmarshal) } - h.streams.close(req.StreamID, req.Error) + h.streams.closeWithError(req.StreamID, rpcStreamError(req.Error, req.ErrorDetail)) return marshalRPCResult(rpcEmptyResponse{}) } @@ -448,6 +448,7 @@ func modelExecutionRequestFromPlugin(req pluginapi.HostModelExecutionRequest, sk AuthID: req.AuthID, ProxyURL: strings.TrimSpace(req.ProxyURL), Path: strings.TrimSpace(req.Path), + CallerScope: strings.TrimSpace(req.CallerScope), } } diff --git a/internal/pluginhost/http_bridge.go b/internal/pluginhost/http_bridge.go index 40b15ef4..1a23fa23 100644 --- a/internal/pluginhost/http_bridge.go +++ b/internal/pluginhost/http_bridge.go @@ -221,6 +221,7 @@ func (c *hostHTTPClient) newHTTPClientForRequest(ctx context.Context, cfg *confi if client == nil { client = &http.Client{} } + installPluginRedirectPolicy(client) return client, nil, nil } @@ -514,16 +515,101 @@ func (c *hostHTTPClient) newHTTPClientForRequest(ctx context.Context, cfg *confi CheckRedirect: func(redirectReq *http.Request, via []*http.Request) error { reqHolder.set(redirectReq) baseTransport.CloseIdleConnections() - if len(via) >= 10 { - return errors.New("stopped after 10 redirects") - } - return nil + return checkPluginRedirect(redirectReq, via, nil) }, } return client, cleanup, nil } +func installPluginRedirectPolicy(client *http.Client) { + if client == nil { + return + } + previous := client.CheckRedirect + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { + return checkPluginRedirect(req, via, previous) + } +} + +func checkPluginRedirect(req *http.Request, via []*http.Request, previous func(*http.Request, []*http.Request) error) error { + // net/http rebuilds each redirect request from the initial request's headers, + // so strip credential headers whenever the target or any intermediate hop + // leaves the initial request's origin (scheme + host:port). + if crossOriginRedirect(req, via) { + for name := range req.Header { + if isPluginCredentialHeader(name) { + // Delete the exact map key because Header.Del canonicalizes and + // would miss non-canonical keys such as "x-api-key". + delete(req.Header, name) + } + } + } + if previous != nil { + return previous(req, via) + } + if len(via) >= 10 { + return errors.New("stopped after 10 redirects") + } + return nil +} + +func crossOriginRedirect(req *http.Request, via []*http.Request) bool { + if req == nil || len(via) == 0 || via[0] == nil { + return false + } + origin := via[0].URL + if !sameRedirectOrigin(origin, req.URL) { + return true + } + for _, hop := range via[1:] { + if hop == nil || !sameRedirectOrigin(origin, hop.URL) { + return true + } + } + return false +} + +func sameRedirectOrigin(a, b *url.URL) bool { + if a == nil || b == nil || a.Hostname() == "" || b.Hostname() == "" { + return false + } + return strings.EqualFold(a.Scheme, b.Scheme) && + strings.EqualFold(a.Hostname(), b.Hostname()) && + redirectOriginPort(a) == redirectOriginPort(b) +} + +func redirectOriginPort(u *url.URL) string { + if u == nil { + return "" + } + if port := u.Port(); port != "" { + return port + } + switch strings.ToLower(u.Scheme) { + case "http": + return "80" + case "https": + return "443" + default: + return "" + } +} + +func isPluginCredentialHeader(name string) bool { + lower := strings.ToLower(strings.TrimSpace(name)) + switch lower { + case "authorization", "proxy-authorization", "cookie", "cookie2", "www-authenticate": + return true + } + return strings.Contains(lower, "api-key") || + strings.Contains(lower, "apikey") || + strings.Contains(lower, "api_key") || + strings.Contains(lower, "token") || + strings.Contains(lower, "secret") || + strings.Contains(lower, "credential") +} + type requestHolder struct { mu sync.Mutex req *http.Request diff --git a/internal/pluginhost/loader_unix.go b/internal/pluginhost/loader_unix.go index 412aa467..636413a8 100644 --- a/internal/pluginhost/loader_unix.go +++ b/internal/pluginhost/loader_unix.go @@ -52,10 +52,6 @@ static const char* cliproxy_dlerror(void) { return dlerror(); } -static int cliproxy_dlclose(void* handle) { - return dlclose(handle); -} - static int cliproxy_call_init(void* fn, const cliproxy_host_api* host, cliproxy_plugin_api* plugin) { return ((cliproxy_plugin_init_fn)fn)(host, plugin); } @@ -99,6 +95,7 @@ type dynamicLibraryLoader struct{} type dynamicLibraryClient struct { handle unsafe.Pointer + tempPath string hostAPI *C.cliproxy_host_api hostCtx unsafe.Pointer api C.cliproxy_plugin_api @@ -112,31 +109,37 @@ func defaultPluginLoader() pluginLoader { } func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, error) { - cPath := C.CString(file.Path) + loadPath, errShadow := shadowCopyPlugin(file) + if errShadow != nil { + return nil, errShadow + } + cPath := C.CString(loadPath) defer C.free(unsafe.Pointer(cPath)) handle := C.cliproxy_dlopen(cPath) if handle == nil { + removeShadowPlugin(loadPath) return nil, fmt.Errorf("dlopen %s: %s", file.Path, dlerrorString()) } + // From here on the image must never be unloaded: dlopen may already + // have started the plugin runtime, and a Go c-shared runtime cannot be + // removed from the host process safely. Failure paths below only detach + // the client logically and deliberately leak the mapped image together + // with any host callback state the runtime may reference. cSymbol := C.CString("cliproxy_plugin_init") initSymbol := C.cliproxy_dlsym(handle, cSymbol) C.free(unsafe.Pointer(cSymbol)) if initSymbol == nil { - C.cliproxy_dlclose(handle) return nil, fmt.Errorf("missing cliproxy_plugin_init: %s", dlerrorString()) } hostAPI := (*C.cliproxy_host_api)(C.malloc(C.size_t(unsafe.Sizeof(C.cliproxy_host_api{})))) if hostAPI == nil { - C.cliproxy_dlclose(handle) return nil, fmt.Errorf("allocate host api") } hostCtx := C.malloc(C.size_t(unsafe.Sizeof(C.uintptr_t(0)))) if hostCtx == nil { - C.free(unsafe.Pointer(hostAPI)) - C.cliproxy_dlclose(handle) return nil, fmt.Errorf("allocate host context") } id := hostCallbackID.Add(1) @@ -148,6 +151,7 @@ func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, err client := &dynamicLibraryClient{ handle: handle, + tempPath: loadPath, hostAPI: hostAPI, hostCtx: hostCtx, host: host, @@ -156,15 +160,15 @@ func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, err } rc := C.cliproxy_call_init(initSymbol, hostAPI, &client.api) if rc != 0 { - client.Shutdown() + client.closeAfterOpenFailure() return nil, fmt.Errorf("cliproxy_plugin_init returned %d", int(rc)) } if uint32(client.api.abi_version) != pluginHostABIVersion { - client.Shutdown() + client.closeAfterOpenFailure() return nil, fmt.Errorf("plugin ABI version %d is not supported", uint32(client.api.abi_version)) } if client.api.call == nil || client.api.free_buffer == nil { - client.Shutdown() + client.closeAfterOpenFailure() return nil, fmt.Errorf("plugin function table is incomplete") } return client, nil @@ -215,6 +219,22 @@ func (c *dynamicLibraryClient) Call(ctx context.Context, method string, request } func (c *dynamicLibraryClient) Shutdown() { + c.close() +} + +func (c *dynamicLibraryClient) closeAfterOpenFailure() { + c.close() +} + +// close detaches the client logically without unloading the plugin image. +// Go plugin libraries are not safe to hot-unload from the host process: the +// plugin runtime, its threads and its signal handlers stay resident after +// dlclose, so a later load in the same process forwards signals into +// unmapped memory and spins. This also applies when initialization failed, +// because dlopen may already have started the runtime. The plugin was +// loaded from a shadow copy, so keeping the image mapped does not block +// deleting or replacing the source artifact. +func (c *dynamicLibraryClient) close() { if c == nil { return } @@ -226,19 +246,16 @@ func (c *dynamicLibraryClient) Shutdown() { C.cliproxy_shutdown_plugin(c.api.shutdown) c.api.shutdown = nil } + c.api.call = nil + c.api.free_buffer = nil if c.hostCtx != nil { + // Drop the callback registration so repeated reloads do not pin the + // host forever. The C host context and host API block themselves are + // deliberately leaked: the resident plugin runtime may still pass + // the context back into cliproxyHostCall, which reads the callback + // id from it before failing the lookup. id := uintptr(*(*C.uintptr_t)(c.hostCtx)) hostCallbackEntries.Delete(id) - C.free(c.hostCtx) - c.hostCtx = nil - } - if c.hostAPI != nil { - C.free(unsafe.Pointer(c.hostAPI)) - c.hostAPI = nil - } - if c.handle != nil { - C.cliproxy_dlclose(c.handle) - c.handle = nil } } diff --git a/internal/pluginhost/loader_windows.go b/internal/pluginhost/loader_windows.go index adb7c6a7..bfffeed7 100644 --- a/internal/pluginhost/loader_windows.go +++ b/internal/pluginhost/loader_windows.go @@ -4,15 +4,8 @@ package pluginhost import ( "context" - "crypto/sha256" - "encoding/hex" - "errors" "fmt" - "io" - "os" - "path/filepath" "runtime" - "strings" "sync" "sync/atomic" "syscall" @@ -27,6 +20,11 @@ type windowsBuffer struct { len uintptr } +type windowsResidentHostBlocks struct { + hostCtx *uintptr + hostAPI *windowsHostAPI +} + type windowsHostAPI struct { abiVersion uint32 hostCtx uintptr @@ -46,14 +44,13 @@ var ( windowsHostCallbackEntries sync.Map windowsHostCallCallback = syscall.NewCallback(windowsHostCall) windowsHostFreeCallback = syscall.NewCallback(windowsHostFree) - shadowPluginCleanupOnce sync.Once -) - -const ( - shadowPluginPrefix = "cliproxy-plugin-" - shadowPluginTempPrefix = ".cliproxy-plugin-" - shadowPluginProcessDirPrefix = "pid-" - shadowPluginDigestLength = 32 + // windowsResidentHostState pins the host context and host API blocks of + // every loaded plugin for the lifetime of the process. A resident + // plugin runtime may pass the context back into windowsHostCall at any + // time, long after its client was detached, and native pointers are + // not GC roots. Only these small blocks are pinned, never the client + // or the host. + windowsResidentHostState sync.Map ) type dynamicLibraryLoader struct{} @@ -85,8 +82,8 @@ func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, err } proc, errProc := dll.FindProc("cliproxy_plugin_init") if errProc != nil { - _ = dll.Release() - removeShadowPlugin(loadPath) + // Do not release the module: loading it may already have started + // the plugin runtime, which cannot be unloaded safely. return nil, errProc } id := windowsHostCallbackID.Add(1) @@ -94,6 +91,13 @@ func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, err host.registerHostCallbackInstance(file.ID, instance) hostCtx := new(uintptr) *hostCtx = id + hostAPI := &windowsHostAPI{ + abiVersion: pluginHostABIVersion, + hostCtx: uintptr(unsafe.Pointer(hostCtx)), + call: windowsHostCallCallback, + freeBuffer: windowsHostFreeCallback, + } + windowsResidentHostState.Store(id, windowsResidentHostBlocks{hostCtx: hostCtx, hostAPI: hostAPI}) windowsHostCallbackEntries.Store(id, dynamicHostCallbackEntry{host: host, pluginID: file.ID, instance: instance}) client := &dynamicLibraryClient{ dll: dll, @@ -102,12 +106,7 @@ func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, err host: host, pluginID: file.ID, instance: instance, - hostAPI: &windowsHostAPI{ - abiVersion: pluginHostABIVersion, - hostCtx: uintptr(unsafe.Pointer(hostCtx)), - call: windowsHostCallCallback, - freeBuffer: windowsHostFreeCallback, - }, + hostAPI: hostAPI, } rc, _, errCall := proc.Call(uintptr(unsafe.Pointer(client.hostAPI)), uintptr(unsafe.Pointer(&client.api))) if rc != 0 { @@ -125,141 +124,6 @@ func (dynamicLibraryLoader) Open(file pluginFile, host *Host) (pluginClient, err return client, nil } -func shadowCopyPlugin(file pluginFile) (string, error) { - dir, errDir := shadowPluginDir() - if errDir != nil { - return "", errDir - } - shadowPluginCleanupOnce.Do(func() { - removeStaleShadowPlugins(dir) - }) - return shadowCopyPluginToDir(file, dir) -} - -func shadowCopyPluginToDir(file pluginFile, dir string) (string, error) { - source := filepath.Clean(file.Path) - tmp, errTemp := os.CreateTemp(dir, shadowPluginTempPrefix+file.ID+"-*"+filepath.Ext(source)) - if errTemp != nil { - return "", errTemp - } - tmpName := tmp.Name() - removeTemp := true - defer func() { - if removeTemp { - removeShadowPlugin(tmpName) - } - }() - - in, errOpen := os.Open(source) - if errOpen != nil { - _ = tmp.Close() - return "", errOpen - } - defer func() { - _ = in.Close() - }() - hasher := sha256.New() - size, errCopy := io.Copy(io.MultiWriter(tmp, hasher), in) - if errCopy != nil { - _ = tmp.Close() - return "", errCopy - } - if errClose := tmp.Close(); errClose != nil { - return "", errClose - } - digest := hex.EncodeToString(hasher.Sum(nil)) - target := shadowPluginPath(dir, file.ID, digest, filepath.Ext(source)) - if shadowPluginMatches(target, size, digest) { - return target, nil - } - if errRemove := os.Remove(target); errRemove != nil && !errors.Is(errRemove, os.ErrNotExist) { - if shadowPluginMatches(target, size, digest) { - return target, nil - } - removeShadowPlugin(target) - return "", fmt.Errorf("remove stale shadow plugin: %w", errRemove) - } - if errRename := os.Rename(tmpName, target); errRename != nil { - if shadowPluginMatches(target, size, digest) { - return target, nil - } - return "", fmt.Errorf("move shadow plugin: %w", errRename) - } - removeTemp = false - return target, nil -} - -func shadowPluginDir() (string, error) { - dir := filepath.Join(os.TempDir(), "cliproxy-pluginhost", shadowPluginProcessDirName(os.Getpid())) - if errMkdir := os.MkdirAll(dir, 0o700); errMkdir != nil { - return "", errMkdir - } - return dir, nil -} - -func shadowPluginProcessDirName(pid int) string { - return fmt.Sprintf("%s%d", shadowPluginProcessDirPrefix, pid) -} - -func removeShadowPlugin(path string) { - if path == "" { - return - } - if errRemove := os.Remove(path); errRemove == nil { - return - } - pathPtr, errPath := windows.UTF16PtrFromString(path) - if errPath != nil { - return - } - _ = windows.MoveFileEx(pathPtr, nil, windows.MOVEFILE_DELAY_UNTIL_REBOOT) -} - -func removeStaleShadowPlugins(dir string) { - entries, errRead := os.ReadDir(dir) - if errRead != nil { - return - } - for _, entry := range entries { - if entry == nil || entry.IsDir() { - continue - } - name := entry.Name() - if strings.HasPrefix(name, shadowPluginPrefix) || strings.HasPrefix(name, shadowPluginTempPrefix) { - removeShadowPlugin(filepath.Join(dir, name)) - } - } -} - -func shadowPluginPath(dir string, id string, digest string, extension string) string { - if len(digest) > shadowPluginDigestLength { - digest = digest[:shadowPluginDigestLength] - } - return filepath.Join(dir, shadowPluginPrefix+id+"-"+digest+extension) -} - -func shadowPluginMatches(path string, size int64, digest string) bool { - info, errStat := os.Stat(path) - if errStat != nil { - return false - } - if !info.Mode().IsRegular() || info.Size() != size { - return false - } - file, errOpen := os.Open(path) - if errOpen != nil { - return false - } - defer func() { - _ = file.Close() - }() - hasher := sha256.New() - if _, errCopy := io.Copy(hasher, file); errCopy != nil { - return false - } - return hex.EncodeToString(hasher.Sum(nil)) == digest -} - func (c *dynamicLibraryClient) callbackInstance() *hostCallbackInstance { if c == nil { return nil @@ -327,17 +191,21 @@ func (c *dynamicLibraryClient) Call(ctx context.Context, method string, request } func (c *dynamicLibraryClient) Shutdown() { - // Windows Go DLLs are not safe to hot-unload from the host process. - // The plugin was loaded from a shadow copy, so keeping the module mapped - // does not block deleting or replacing the source artifact. - c.close(false) + c.close() } func (c *dynamicLibraryClient) closeAfterOpenFailure() { - c.close(true) + c.close() } -func (c *dynamicLibraryClient) close(releaseDLL bool) { +// close detaches the client logically without unloading the plugin module. +// Go plugin DLLs are not safe to hot-unload from the host process: the +// plugin runtime, its threads and its signal handlers stay resident after +// FreeLibrary. This also applies when initialization failed, because +// loading the module may already have started the runtime. The plugin was +// loaded from a shadow copy, so keeping the module mapped does not block +// deleting or replacing the source artifact. +func (c *dynamicLibraryClient) close() { if c == nil { return } @@ -349,18 +217,15 @@ func (c *dynamicLibraryClient) close(releaseDLL bool) { _, _, _ = syscall.SyscallN(c.api.shutdown) c.api.shutdown = 0 } + c.api.call = 0 + c.api.freeBuffer = 0 if c.hostCtx != nil { + // Drop the callback registration so repeated reloads do not pin + // the host forever. The context and host API blocks stay pinned in + // windowsResidentHostState because the resident plugin runtime may + // still dereference them on late callbacks. windowsHostCallbackEntries.Delete(*c.hostCtx) - c.hostCtx = nil - } - if c.dll != nil { - if releaseDLL { - _ = c.dll.Release() - } - c.dll = nil } - removeShadowPlugin(c.tempPath) - c.tempPath = "" } func windowsHostCall(hostCtx uintptr, methodPtr uintptr, requestPtr uintptr, requestLen uintptr, responsePtr uintptr) uintptr { diff --git a/internal/pluginhost/management.go b/internal/pluginhost/management.go index 72e9c12e..536e8b32 100644 --- a/internal/pluginhost/management.go +++ b/internal/pluginhost/management.go @@ -3,6 +3,7 @@ package pluginhost import ( "bytes" "context" + "errors" "fmt" "io" "net/http" @@ -246,7 +247,12 @@ func (h *Host) ServeManagementHTTP(w http.ResponseWriter, r *http.Request) bool var errRead error body, errRead = io.ReadAll(r.Body) if errRead != nil { - http.Error(w, "failed to read plugin management request body", http.StatusBadRequest) + status := http.StatusBadRequest + var maxBytesErr *http.MaxBytesError + if errors.As(errRead, &maxBytesErr) { + status = http.StatusRequestEntityTooLarge + } + http.Error(w, "failed to read plugin management request body", status) return true } if errClose := r.Body.Close(); errClose != nil { diff --git a/internal/pluginhost/quota_provider.go b/internal/pluginhost/quota_provider.go index 5acf04f9..f366560d 100644 --- a/internal/pluginhost/quota_provider.go +++ b/internal/pluginhost/quota_provider.go @@ -113,41 +113,67 @@ func (h *Host) HasQuotaProviderForPlugin(pluginID string) bool { return h.quotaProviderRecordByPlugin(pluginID) != nil } -// QuotaSupportedProvidersSet returns a point-in-time set of all provider keys supported by active quota providers. -// This allows callers (e.g. ListAuthFiles) to discover all quota-capable providers with zero per-credential RPCs. -func (h *Host) QuotaSupportedProvidersSet(ctx context.Context) map[string]struct{} { +// QuotaProviderPluginIDs returns a point-in-time map from normalized provider keys to the ID of the +// quota plugin that quotaProviderRecord would select for that provider, preserving its two-pass precedence. +// This allows callers (e.g. ListAuthFiles) to resolve quota plugins with zero per-credential RPCs. +func (h *Host) QuotaProviderPluginIDs(ctx context.Context) map[string]string { if h == nil { return nil } + if ctx == nil { + ctx = context.Background() + } snap := h.Snapshot() if snap == nil { return nil } records := h.activeRecordsFromSnapshot(snap) - out := make(map[string]struct{}) + // Keep candidates in selection order until probing finishes, because DescribeQuota may fuse a plugin. + candidates := make(map[string][]string) + addCandidate := func(provider, pluginID string) { + if provider != "" { + candidates[provider] = append(candidates[provider], pluginID) + } + } + quotaRecords := make([]capabilityRecord, 0, len(records)) + // First pass: QuotaProvider.Identifier(), plugin ID, or AuthProvider.Identifier(). for _, record := range records { if err := ctx.Err(); err != nil { - return out + break } quotaProvider := record.plugin.Capabilities.QuotaProvider if quotaProvider == nil || h.isPluginFused(record.id) { continue } - if identifier, okID := h.callQuotaIdentifier(record.id, quotaProvider); okID && identifier != "" { - out[identifier] = struct{}{} - } - if normID := normalizeProviderID(record.id); normID != "" { - out[normID] = struct{}{} + quotaRecords = append(quotaRecords, record) + if identifier, okID := h.callQuotaIdentifier(record.id, quotaProvider); okID { + addCandidate(identifier, record.id) } + addCandidate(normalizeProviderID(record.id), record.id) if record.plugin.Capabilities.AuthProvider != nil { - if authID, okAuth := h.callAuthProviderIdentifier(record.id, record.plugin.Capabilities.AuthProvider); okAuth && authID != "" { - out[authID] = struct{}{} + if authID, okAuth := h.callAuthProviderIdentifier(record.id, record.plugin.Capabilities.AuthProvider); okAuth { + addCandidate(authID, record.id) } } - supported := snap.cachedQuotaSupportedProviders(ctx, h, record, quotaProvider) - for _, p := range supported { - if clean := normalizeProviderID(p); clean != "" { - out[clean] = struct{}{} + } + // Second pass: SupportedProviders declared by DescribeQuota. + for _, record := range quotaRecords { + if err := ctx.Err(); err != nil { + break + } + if h.isPluginFused(record.id) { + continue + } + for _, p := range snap.cachedQuotaSupportedProviders(ctx, h, record, record.plugin.Capabilities.QuotaProvider) { + addCandidate(normalizeProviderID(p), record.id) + } + } + out := make(map[string]string, len(candidates)) + for provider, pluginIDs := range candidates { + for _, pluginID := range pluginIDs { + if !h.isPluginFused(pluginID) { + out[provider] = pluginID + break } } } diff --git a/internal/pluginhost/rpc_client.go b/internal/pluginhost/rpc_client.go index 96e3785d..4aaf3c66 100644 --- a/internal/pluginhost/rpc_client.go +++ b/internal/pluginhost/rpc_client.go @@ -7,7 +7,9 @@ import ( "fmt" "io" "net/http" + "strconv" "strings" + "time" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" @@ -44,9 +46,10 @@ type rpcQuotaProvider struct { } type rpcError struct { - Code string - message string - statusCode int + Code string + message string + statusCode int + retryAfterSeconds int } func (e rpcError) Error() string { @@ -57,6 +60,27 @@ func (e rpcError) StatusCode() int { return e.statusCode } +func (e rpcError) RetryAfter() *time.Duration { + if e.retryAfterSeconds <= 0 { + return nil + } + seconds := int64(e.retryAfterSeconds) + maxSeconds := int64(1<<63-1) / int64(time.Second) + if seconds > maxSeconds { + seconds = maxSeconds + } + retryAfter := time.Duration(seconds) * time.Second + return &retryAfter +} + +func (e rpcError) SafeResponseHeaders() http.Header { + retryAfter := e.RetryAfter() + if retryAfter == nil { + return nil + } + return http.Header{"Retry-After": []string{strconv.FormatInt(int64(*retryAfter/time.Second), 10)}} +} + type rpcResponseNormalizer struct { *rpcPluginAdapter method string @@ -132,6 +156,8 @@ func registerRPCPlugin(ctx context.Context, host *Host, id string, client plugin } if resp.Capabilities.RequestInterceptor { plugin.Capabilities.RequestInterceptor = adapter + plugin.Capabilities.RequestInterceptorOmitBodyBefore = resp.Capabilities.RequestInterceptorOmitBodyBefore + plugin.Capabilities.RequestInterceptorOmitBodyAfter = resp.Capabilities.RequestInterceptorOmitBodyAfter } if resp.Capabilities.RequestLifecyclePlugin { plugin.Capabilities.RequestLifecyclePlugin = adapter @@ -366,9 +392,10 @@ func decodeEnvelopeResult[T any](envelope pluginabi.Envelope) (T, error) { message = "plugin call failed" } return zero, rpcError{ - Code: strings.TrimSpace(envelope.Error.Code), - message: message, - statusCode: envelope.Error.HTTPStatus, + Code: strings.TrimSpace(envelope.Error.Code), + message: message, + statusCode: envelope.Error.HTTPStatus, + retryAfterSeconds: envelope.Error.RetryAfterSeconds, } } return zero, fmt.Errorf("plugin call failed") diff --git a/internal/pluginhost/rpc_schema.go b/internal/pluginhost/rpc_schema.go index 188ed246..f2887029 100644 --- a/internal/pluginhost/rpc_schema.go +++ b/internal/pluginhost/rpc_schema.go @@ -19,33 +19,35 @@ type rpcRegistration struct { } type rpcCapabilities struct { - ModelRegistrar bool `json:"model_registrar"` - ModelProvider bool `json:"model_provider"` - AuthProvider bool `json:"auth_provider"` - FrontendAuthProvider bool `json:"frontend_auth_provider"` - FrontendAuthProviderExclusive bool `json:"frontend_auth_provider_exclusive"` - Scheduler bool `json:"scheduler"` - SchedulerAcrossPriorities bool `json:"scheduler_across_priorities,omitempty"` - ModelRouter bool `json:"model_router"` - Executor bool `json:"executor"` - ExecutorModelScope pluginapi.ExecutorModelScope `json:"executor_model_scope"` - ExecutorInputFormats []string `json:"executor_input_formats,omitempty"` - ExecutorOutputFormats []string `json:"executor_output_formats,omitempty"` - RequestTranslator bool `json:"request_translator"` - RequestNormalizer bool `json:"request_normalizer"` - RequestInterceptor bool `json:"request_interceptor"` - RequestLifecyclePlugin bool `json:"request_lifecycle_plugin"` - ResponseTranslator bool `json:"response_translator"` - ResponseBeforeTranslator bool `json:"response_before_translator"` - ResponseAfterTranslator bool `json:"response_after_translator"` - ResponseInterceptor bool `json:"response_interceptor"` - StreamChunkInterceptor bool `json:"response_stream_interceptor"` - WebSocketResponseObserver bool `json:"websocket_response_observer"` - ThinkingApplier bool `json:"thinking_applier"` - UsagePlugin bool `json:"usage_plugin"` - CommandLinePlugin bool `json:"command_line_plugin"` - ManagementAPI bool `json:"management_api"` - QuotaProvider bool `json:"quota_provider"` + ModelRegistrar bool `json:"model_registrar"` + ModelProvider bool `json:"model_provider"` + AuthProvider bool `json:"auth_provider"` + FrontendAuthProvider bool `json:"frontend_auth_provider"` + FrontendAuthProviderExclusive bool `json:"frontend_auth_provider_exclusive"` + Scheduler bool `json:"scheduler"` + SchedulerAcrossPriorities bool `json:"scheduler_across_priorities,omitempty"` + ModelRouter bool `json:"model_router"` + Executor bool `json:"executor"` + ExecutorModelScope pluginapi.ExecutorModelScope `json:"executor_model_scope"` + ExecutorInputFormats []string `json:"executor_input_formats,omitempty"` + ExecutorOutputFormats []string `json:"executor_output_formats,omitempty"` + RequestTranslator bool `json:"request_translator"` + RequestNormalizer bool `json:"request_normalizer"` + RequestInterceptor bool `json:"request_interceptor"` + RequestInterceptorOmitBodyBefore bool `json:"request_interceptor_omit_body_before,omitempty"` + RequestInterceptorOmitBodyAfter bool `json:"request_interceptor_omit_body_after,omitempty"` + RequestLifecyclePlugin bool `json:"request_lifecycle_plugin"` + ResponseTranslator bool `json:"response_translator"` + ResponseBeforeTranslator bool `json:"response_before_translator"` + ResponseAfterTranslator bool `json:"response_after_translator"` + ResponseInterceptor bool `json:"response_interceptor"` + StreamChunkInterceptor bool `json:"response_stream_interceptor"` + WebSocketResponseObserver bool `json:"websocket_response_observer"` + ThinkingApplier bool `json:"thinking_applier"` + UsagePlugin bool `json:"usage_plugin"` + CommandLinePlugin bool `json:"command_line_plugin"` + ManagementAPI bool `json:"management_api"` + QuotaProvider bool `json:"quota_provider"` } type rpcIdentifierResponse struct { @@ -145,6 +147,22 @@ type rpcQuotaResetRequest struct { type rpcEmptyResponse struct{} +func requestInterceptorWantsOmitBody(caps pluginapi.Capabilities, before bool) bool { + if caps.RequestInterceptor == nil { + return false + } + if before && caps.RequestInterceptorOmitBodyBefore { + return true + } + if !before && caps.RequestInterceptorOmitBodyAfter { + return true + } + if opt, ok := caps.RequestInterceptor.(interface{ RequestInterceptorWantsOmitBody(bool) bool }); ok && opt != nil { + return opt.RequestInterceptorWantsOmitBody(before) + } + return false +} + func schedulerWantsAcrossPriorities(caps pluginapi.Capabilities) bool { if caps.Scheduler == nil { return false @@ -161,33 +179,35 @@ func schedulerWantsAcrossPriorities(caps pluginapi.Capabilities) bool { func rpcCapabilitiesFromPlugin(plugin pluginapi.Plugin) rpcCapabilities { caps := plugin.Capabilities return rpcCapabilities{ - ModelRegistrar: caps.ModelRegistrar != nil, - ModelProvider: caps.ModelProvider != nil, - AuthProvider: caps.AuthProvider != nil, - FrontendAuthProvider: caps.FrontendAuthProvider != nil, - FrontendAuthProviderExclusive: caps.FrontendAuthProvider != nil && caps.FrontendAuthProviderExclusive, - Scheduler: caps.Scheduler != nil, - SchedulerAcrossPriorities: schedulerWantsAcrossPriorities(caps), - ModelRouter: caps.ModelRouter != nil, - Executor: caps.Executor != nil, - ExecutorModelScope: normalizedExecutorModelScope(caps), - ExecutorInputFormats: append([]string(nil), caps.ExecutorInputFormats...), - ExecutorOutputFormats: append([]string(nil), caps.ExecutorOutputFormats...), - RequestTranslator: caps.RequestTranslator != nil, - RequestNormalizer: caps.RequestNormalizer != nil, - RequestInterceptor: caps.RequestInterceptor != nil, - RequestLifecyclePlugin: caps.RequestLifecyclePlugin != nil, - ResponseTranslator: caps.ResponseTranslator != nil, - ResponseBeforeTranslator: caps.ResponseBeforeTranslator != nil, - ResponseAfterTranslator: caps.ResponseAfterTranslator != nil, - ResponseInterceptor: caps.ResponseInterceptor != nil, - StreamChunkInterceptor: caps.StreamChunkInterceptor != nil, - WebSocketResponseObserver: caps.WebSocketResponseObserver != nil, - ThinkingApplier: caps.ThinkingApplier != nil, - UsagePlugin: caps.UsagePlugin != nil, - CommandLinePlugin: caps.CommandLinePlugin != nil, - ManagementAPI: caps.ManagementAPI != nil, - QuotaProvider: caps.QuotaProvider != nil, + ModelRegistrar: caps.ModelRegistrar != nil, + ModelProvider: caps.ModelProvider != nil, + AuthProvider: caps.AuthProvider != nil, + FrontendAuthProvider: caps.FrontendAuthProvider != nil, + FrontendAuthProviderExclusive: caps.FrontendAuthProvider != nil && caps.FrontendAuthProviderExclusive, + Scheduler: caps.Scheduler != nil, + SchedulerAcrossPriorities: schedulerWantsAcrossPriorities(caps), + ModelRouter: caps.ModelRouter != nil, + Executor: caps.Executor != nil, + ExecutorModelScope: normalizedExecutorModelScope(caps), + ExecutorInputFormats: append([]string(nil), caps.ExecutorInputFormats...), + ExecutorOutputFormats: append([]string(nil), caps.ExecutorOutputFormats...), + RequestTranslator: caps.RequestTranslator != nil, + RequestNormalizer: caps.RequestNormalizer != nil, + RequestInterceptor: caps.RequestInterceptor != nil, + RequestInterceptorOmitBodyBefore: requestInterceptorWantsOmitBody(caps, true), + RequestInterceptorOmitBodyAfter: requestInterceptorWantsOmitBody(caps, false), + RequestLifecyclePlugin: caps.RequestLifecyclePlugin != nil, + ResponseTranslator: caps.ResponseTranslator != nil, + ResponseBeforeTranslator: caps.ResponseBeforeTranslator != nil, + ResponseAfterTranslator: caps.ResponseAfterTranslator != nil, + ResponseInterceptor: caps.ResponseInterceptor != nil, + StreamChunkInterceptor: caps.StreamChunkInterceptor != nil, + WebSocketResponseObserver: caps.WebSocketResponseObserver != nil, + ThinkingApplier: caps.ThinkingApplier != nil, + UsagePlugin: caps.UsagePlugin != nil, + CommandLinePlugin: caps.CommandLinePlugin != nil, + ManagementAPI: caps.ManagementAPI != nil, + QuotaProvider: caps.QuotaProvider != nil, } } diff --git a/internal/pluginhost/shadow.go b/internal/pluginhost/shadow.go new file mode 100644 index 00000000..3884b11e --- /dev/null +++ b/internal/pluginhost/shadow.go @@ -0,0 +1,181 @@ +//go:build windows || (cgo && (linux || darwin || freebsd)) + +package pluginhost + +import ( + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strconv" + "strings" + "sync" +) + +// Plugin dynamic libraries are loaded from content-addressed shadow copies. +// A loaded Go plugin image must stay mapped for the lifetime of the host +// process (its runtime cannot be unloaded), so loading from a shadow copy +// keeps the original artifact replaceable and gives every distinct plugin +// build its own path, which dynamic loaders map as a fresh image. + +const ( + shadowPluginPrefix = "cliproxy-plugin-" + shadowPluginTempPrefix = ".cliproxy-plugin-" + shadowPluginProcessDirPrefix = "pid-" + shadowPluginDigestLength = 32 +) + +var shadowPluginCleanupOnce sync.Once + +func shadowCopyPlugin(file pluginFile) (string, error) { + dir, errDir := shadowPluginDir() + if errDir != nil { + return "", errDir + } + shadowPluginCleanupOnce.Do(func() { + removeStaleShadowPlugins(dir) + removeStaleShadowPluginDirs(filepath.Dir(dir), os.Getpid()) + }) + return shadowCopyPluginToDir(file, dir) +} + +// removeStaleShadowPluginDirs removes shadow directories left behind by +// host processes that already exited. Directories whose owner process is +// still alive are kept, because that process may reuse the resident plugin +// images inside them. +func removeStaleShadowPluginDirs(baseDir string, currentPID int) { + entries, errRead := os.ReadDir(baseDir) + if errRead != nil { + return + } + for _, entry := range entries { + if entry == nil || !entry.IsDir() { + continue + } + name := entry.Name() + if !strings.HasPrefix(name, shadowPluginProcessDirPrefix) { + continue + } + pid, errParse := strconv.Atoi(strings.TrimPrefix(name, shadowPluginProcessDirPrefix)) + if errParse != nil || pid <= 0 || pid == currentPID { + continue + } + if shadowPluginProcessAlive(pid) { + continue + } + removeShadowPluginDir(filepath.Join(baseDir, name)) + } +} + +func shadowCopyPluginToDir(file pluginFile, dir string) (string, error) { + source := filepath.Clean(file.Path) + tmp, errTemp := os.CreateTemp(dir, shadowPluginTempPrefix+file.ID+"-*"+filepath.Ext(source)) + if errTemp != nil { + return "", errTemp + } + tmpName := tmp.Name() + removeTemp := true + defer func() { + if removeTemp { + removeShadowPlugin(tmpName) + } + }() + + in, errOpen := os.Open(source) + if errOpen != nil { + _ = tmp.Close() + return "", errOpen + } + defer func() { + _ = in.Close() + }() + hasher := sha256.New() + size, errCopy := io.Copy(io.MultiWriter(tmp, hasher), in) + if errCopy != nil { + _ = tmp.Close() + return "", errCopy + } + if errClose := tmp.Close(); errClose != nil { + return "", errClose + } + digest := hex.EncodeToString(hasher.Sum(nil)) + target := shadowPluginPath(dir, file.ID, digest, filepath.Ext(source)) + if shadowPluginMatches(target, size, digest) { + return target, nil + } + if errRemove := os.Remove(target); errRemove != nil && !errors.Is(errRemove, os.ErrNotExist) { + if shadowPluginMatches(target, size, digest) { + return target, nil + } + removeShadowPlugin(target) + return "", fmt.Errorf("remove stale shadow plugin: %w", errRemove) + } + if errRename := os.Rename(tmpName, target); errRename != nil { + if shadowPluginMatches(target, size, digest) { + return target, nil + } + return "", fmt.Errorf("move shadow plugin: %w", errRename) + } + removeTemp = false + return target, nil +} + +func shadowPluginDir() (string, error) { + dir := filepath.Join(os.TempDir(), "cliproxy-pluginhost", shadowPluginProcessDirName(os.Getpid())) + if errMkdir := os.MkdirAll(dir, 0o700); errMkdir != nil { + return "", errMkdir + } + return dir, nil +} + +func shadowPluginProcessDirName(pid int) string { + return fmt.Sprintf("%s%d", shadowPluginProcessDirPrefix, pid) +} + +func removeStaleShadowPlugins(dir string) { + entries, errRead := os.ReadDir(dir) + if errRead != nil { + return + } + for _, entry := range entries { + if entry == nil || entry.IsDir() { + continue + } + name := entry.Name() + if strings.HasPrefix(name, shadowPluginPrefix) || strings.HasPrefix(name, shadowPluginTempPrefix) { + removeShadowPlugin(filepath.Join(dir, name)) + } + } +} + +func shadowPluginPath(dir string, id string, digest string, extension string) string { + if len(digest) > shadowPluginDigestLength { + digest = digest[:shadowPluginDigestLength] + } + return filepath.Join(dir, shadowPluginPrefix+id+"-"+digest+extension) +} + +func shadowPluginMatches(path string, size int64, digest string) bool { + info, errStat := os.Stat(path) + if errStat != nil { + return false + } + if !info.Mode().IsRegular() || info.Size() != size { + return false + } + file, errOpen := os.Open(path) + if errOpen != nil { + return false + } + defer func() { + _ = file.Close() + }() + hasher := sha256.New() + if _, errCopy := io.Copy(hasher, file); errCopy != nil { + return false + } + return hex.EncodeToString(hasher.Sum(nil)) == digest +} diff --git a/internal/pluginhost/shadow_unix.go b/internal/pluginhost/shadow_unix.go new file mode 100644 index 00000000..39e626e7 --- /dev/null +++ b/internal/pluginhost/shadow_unix.go @@ -0,0 +1,27 @@ +//go:build cgo && (linux || darwin || freebsd) + +package pluginhost + +import ( + "errors" + "os" + "syscall" +) + +func removeShadowPlugin(path string) { + if path == "" { + return + } + _ = os.Remove(path) +} + +func removeShadowPluginDir(dir string) { + _ = os.RemoveAll(dir) +} + +func shadowPluginProcessAlive(pid int) bool { + // Only ESRCH proves the process is gone; any other outcome keeps the + // directory so a live process never loses its resident plugin images. + errSignal := syscall.Kill(pid, 0) + return !errors.Is(errSignal, syscall.ESRCH) +} diff --git a/internal/pluginhost/shadow_windows.go b/internal/pluginhost/shadow_windows.go new file mode 100644 index 00000000..5efe995f --- /dev/null +++ b/internal/pluginhost/shadow_windows.go @@ -0,0 +1,62 @@ +//go:build windows + +package pluginhost + +import ( + "errors" + "os" + "path/filepath" + "strings" + + "golang.org/x/sys/windows" +) + +func removeShadowPlugin(path string) { + if path == "" { + return + } + if errRemove := os.Remove(path); errRemove == nil { + return + } + pathPtr, errPath := windows.UTF16PtrFromString(path) + if errPath != nil { + return + } + _ = windows.MoveFileEx(pathPtr, nil, windows.MOVEFILE_DELAY_UNTIL_REBOOT) +} + +func removeShadowPluginDir(dir string) { + if errRemove := os.RemoveAll(dir); errRemove == nil { + return + } + entries, errRead := os.ReadDir(dir) + if errRead != nil { + return + } + for _, entry := range entries { + if entry == nil || entry.IsDir() { + continue + } + name := entry.Name() + if strings.HasPrefix(name, shadowPluginPrefix) || strings.HasPrefix(name, shadowPluginTempPrefix) { + removeShadowPlugin(filepath.Join(dir, name)) + } + } + _ = os.Remove(dir) +} + +func shadowPluginProcessAlive(pid int) bool { + // Only an invalid parameter proves the process id no longer exists; + // any other error keeps the directory so a live process never loses + // its resident plugin images. + handle, errOpen := windows.OpenProcess(windows.SYNCHRONIZE, false, uint32(pid)) + if errOpen != nil { + return !errors.Is(errOpen, windows.ERROR_INVALID_PARAMETER) + } + event, errWait := windows.WaitForSingleObject(handle, 0) + _ = windows.CloseHandle(handle) + if errWait != nil { + return true + } + return event == uint32(windows.WAIT_TIMEOUT) +} diff --git a/internal/pluginhost/snapshot.go b/internal/pluginhost/snapshot.go index 601c9eb4..f5c7e9ad 100644 --- a/internal/pluginhost/snapshot.go +++ b/internal/pluginhost/snapshot.go @@ -1,6 +1,7 @@ package pluginhost import ( + "reflect" "sort" "strings" "sync" @@ -171,6 +172,15 @@ func cloneConfigFields(fields []pluginapi.ConfigField) []pluginapi.ConfigField { copy(out, fields) for index := range out { out[index].EnumValues = append([]string(nil), fields[index].EnumValues...) + out[index].Default = cloneConfigFieldValue(fields[index].Default) } return out } + +func cloneConfigFieldValue(value any) any { + if value == nil { + return nil + } + visited := make(map[metadataCloneVisit]reflect.Value) + return cloneInterceptorMetadataAny(reflect.ValueOf(value), visited) +} diff --git a/internal/pluginhost/stream_bridge.go b/internal/pluginhost/stream_bridge.go index 33a7a459..6cd5b163 100644 --- a/internal/pluginhost/stream_bridge.go +++ b/internal/pluginhost/stream_bridge.go @@ -5,9 +5,11 @@ import ( "errors" "fmt" "strconv" + "strings" "sync" "sync/atomic" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginabi" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" ) @@ -39,19 +41,50 @@ type streamBridgeEmit struct { } type streamBridgeClose struct { - errorMessage string - accepted chan struct{} + err error + accepted chan struct{} } type rpcStreamEmitRequest struct { StreamID string `json:"stream_id"` Payload []byte `json:"payload,omitempty"` Error string `json:"error,omitempty"` + // ErrorDetail optionally carries the structured plugin error (status, code, + // retry-after). It takes precedence over Error when present. + ErrorDetail *pluginabi.Error `json:"error_detail,omitempty"` } type rpcStreamCloseRequest struct { StreamID string `json:"stream_id"` Error string `json:"error,omitempty"` + // ErrorDetail optionally carries the structured plugin error (status, code, + // retry-after). It takes precedence over Error when present. + ErrorDetail *pluginabi.Error `json:"error_detail,omitempty"` +} + +// rpcStreamError builds the error reported to the client for a plugin stream +// failure. A structured detail keeps the HTTP status and retry metadata that a +// plain message cannot carry. +func rpcStreamError(message string, detail *pluginabi.Error) error { + if detail != nil { + text := strings.TrimSpace(detail.Message) + if text == "" { + text = strings.TrimSpace(message) + } + if text == "" { + text = "plugin stream failed" + } + return rpcError{ + Code: strings.TrimSpace(detail.Code), + message: text, + statusCode: detail.HTTPStatus, + retryAfterSeconds: detail.RetryAfterSeconds, + } + } + if message != "" { + return fmt.Errorf("%s", message) + } + return nil } func newStreamBridge() *streamBridge { @@ -97,8 +130,8 @@ func (s *streamBridgeStream) run() { case request := <-s.closes: s.markClosed() close(request.accepted) - if request.errorMessage != "" { - queue = append(queue, pluginapi.ExecutorStreamChunk{Err: fmt.Errorf("%s", request.errorMessage)}) + if request.err != nil { + queue = append(queue, pluginapi.ExecutorStreamChunk{Err: request.err}) } for len(queue) > 0 { select { @@ -161,13 +194,13 @@ func (s *streamBridgeStream) emit(ctx context.Context, chunk pluginapi.ExecutorS return <-request.done } -func (s *streamBridgeStream) close(errorMessage string) { +func (s *streamBridgeStream) closeWithError(err error) { if s == nil { return } request := streamBridgeClose{ - errorMessage: errorMessage, - accepted: make(chan struct{}), + err: err, + accepted: make(chan struct{}), } select { case <-s.finished: @@ -229,6 +262,10 @@ func (b *streamBridge) emit(ctx context.Context, id string, chunk pluginapi.Exec } func (b *streamBridge) close(id string, errorMessage string) { + b.closeWithError(id, rpcStreamError(errorMessage, nil)) +} + +func (b *streamBridge) closeWithError(id string, err error) { if b == nil || id == "" { return } @@ -239,5 +276,5 @@ func (b *streamBridge) close(id string, errorMessage string) { if stream == nil { return } - stream.close(errorMessage) + stream.closeWithError(err) } diff --git a/internal/pluginstore/auth.go b/internal/pluginstore/auth.go index 110c4eff..f01c6c6f 100644 --- a/internal/pluginstore/auth.go +++ b/internal/pluginstore/auth.go @@ -8,6 +8,8 @@ import ( "os" "strings" "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/githubauth" ) const ( @@ -228,6 +230,10 @@ func applyPluginStoreAuthForClient(headers http.Header, resolved []ResolvedAuthC } item, ok := matchingAuthConfig(auth, requestURL, kind) if !ok { + if token := githubauth.TokenForURL(requestURL); token != "" { + headers.Set("Authorization", "Bearer "+token) + return true, nil + } return false, nil } switch strings.ToLower(strings.TrimSpace(item.Type)) { diff --git a/internal/pluginstore/github.go b/internal/pluginstore/github.go index a8d612f4..f056d947 100644 --- a/internal/pluginstore/github.go +++ b/internal/pluginstore/github.go @@ -143,7 +143,11 @@ func (c Client) releaseAssetAPIAuthenticated(apiURL string) bool { if item, ok := matchingResolvedAuthConfig(c.ResolvedAuth, apiURL, RequestKindArtifact); ok { return resolvedAuthConfigured(item) } - return AuthConfigured(c.Auth, apiURL, RequestKindArtifact) + if _, ok := matchingAuthConfig(c.Auth, apiURL, RequestKindArtifact); ok { + return AuthConfigured(c.Auth, apiURL, RequestKindArtifact) + } + _, authenticated, errAuth := c.authHeaders(apiURL, RequestKindArtifact) + return errAuth == nil && authenticated } func (c Client) get(ctx context.Context, requestURL string, accept string, kind string, maxSize int64) ([]byte, error) { diff --git a/sdk/pluginabi/types.go b/sdk/pluginabi/types.go index e79fd8c1..5ec69479 100644 --- a/sdk/pluginabi/types.go +++ b/sdk/pluginabi/types.go @@ -17,7 +17,12 @@ const ( // Version 6 preserves raw JSON bodies for plugin management responses. // Plugins that still require HTML entity escaping on JSON response strings // should keep schema_version < 6. - SchemaVersion uint32 = 6 + // Version 7 lets request interceptors opt out of receiving the request body + // per phase via request_interceptor_omit_body_before/after capabilities; the + // host then omits RequestInterceptRequest.Body for that phase instead of + // cloning and Base64-encoding the full payload. Plugins that declare the + // flags on schema_version < 7 (or do not declare them) still receive it. + SchemaVersion uint32 = 7 // SchemaVersionStreamChunkOmitRequestBody is the first schema version that omits // request bodies on payload stream-chunk interceptor calls. SchemaVersionStreamChunkOmitRequestBody uint32 = 3 @@ -30,6 +35,9 @@ const ( // SchemaVersionRawManagementResponse is the first schema version where plugin // management JSON responses are preserved without HTML-escaping strings. SchemaVersionRawManagementResponse uint32 = 6 + // SchemaVersionRequestInterceptorOmitBody is the first schema version where + // request interceptors may opt out of receiving the request body per phase. + SchemaVersionRequestInterceptorOmitBody uint32 = 7 ) const ( @@ -109,6 +117,7 @@ const ( MethodHostAuthGet = "host.auth.get" MethodHostAuthGetRuntime = "host.auth.get_runtime" MethodHostAuthSave = "host.auth.save" + MethodHostAuthModels = "host.auth.models" MethodHostAffinityLookup = "host.affinity.lookup" MethodHostRoutingResetCooldown = "host.routing.reset_cooldown" @@ -127,6 +136,12 @@ type Error struct { // HTTPStatus is the HTTP status code (e.g. 401, 403, 429) to surface to the client. // When omitted or 0, CPA defaults to HTTP 500 (internal_server_error). HTTPStatus int `json:"http_status,omitempty"` + // RetryAfterSeconds carries upstream retry guidance for cooldown and response headers. + // It is an optional hint. Schema-version agreement does not negotiate support for + // this field, and older hosts silently ignore it, so plugins must not depend on it + // for correctness. The same pair is accepted as error_detail on host.stream.emit + // and host.stream.close. + RetryAfterSeconds int `json:"retry_after_seconds,omitempty"` } // Error implements the error interface for Error. diff --git a/sdk/pluginapi/types.go b/sdk/pluginapi/types.go index b1bd3686..714b6fac 100644 --- a/sdk/pluginapi/types.go +++ b/sdk/pluginapi/types.go @@ -62,6 +62,8 @@ type ConfigField struct { Name string // Type classifies the field value for management clients. Type ConfigFieldType + // Default is the value used when the configuration key is omitted. + Default any // EnumValues lists allowed values when Type is ConfigFieldTypeEnum. EnumValues []string // Description explains how the plugin uses the field. @@ -110,6 +112,14 @@ type Capabilities struct { ResponseAfterTranslator ResponseNormalizer // RequestInterceptor rewrites execution requests before and after credential selection. RequestInterceptor RequestInterceptor + // RequestInterceptorOmitBodyBefore opts out of receiving RequestInterceptRequest.Body + // in InterceptRequestBeforeAuth. It is honored only on schema version 7+; the plugin + // may still replace the body blindly via RequestInterceptResponse.Body. + RequestInterceptorOmitBodyBefore bool + // RequestInterceptorOmitBodyAfter opts out of receiving RequestInterceptRequest.Body + // in InterceptRequestAfterAuth. It is honored only on schema version 7+; the plugin + // may still replace the body blindly via RequestInterceptResponse.Body. + RequestInterceptorOmitBodyAfter bool // RequestLifecyclePlugin asynchronously receives one terminal event for each request that reached request interception. RequestLifecyclePlugin RequestLifecyclePlugin // ResponseInterceptor rewrites successful non-streaming HTTP execution responses before downstream delivery. @@ -194,6 +204,9 @@ type ThinkingSupport struct { ZeroAllowed bool // DynamicAllowed reports whether automatic reasoning budget selection is supported. DynamicAllowed bool + // DisableType is the Claude thinking.type that turns thinking off for this + // model. Empty means "disabled"; claude-sonnet-5-5 needs "between_tools". + DisableType string // Levels lists supported named reasoning levels. Levels []string } @@ -372,6 +385,8 @@ type AuthRefreshResponse struct { } // ModelRegistrar registers plugin-provided models with the host. +// Without an Executor, these models are published only through matching credentials +// after applying their model exclusions. type ModelRegistrar interface { RegisterModels(context.Context, ModelRegistrationRequest) (ModelRegistrationResponse, error) } @@ -391,6 +406,7 @@ type ModelRegistrationResponse struct { } // ModelProvider contributes provider-native static and per-auth model metadata. +// StaticModels alone does not make models available without an Executor or matching credentials. type ModelProvider interface { StaticModels(context.Context, StaticModelRequest) (ModelResponse, error) ModelsForAuth(context.Context, AuthModelRequest) (ModelResponse, error) @@ -450,8 +466,11 @@ type FrontendAuthRequest struct { Headers http.Header // Query contains inbound query parameters. Query url.Values - // Body contains the raw request body. + // Body contains the raw request body as received, before any Content-Encoding decoding. Body []byte + // DecodedBody contains the request body after Content-Encoding decoding. It is set only when + // the body was content-encoded and the route's handler decodes it, and is nil otherwise. + DecodedBody []byte } // FrontendAuthResponse reports the authentication decision and identity metadata. @@ -462,6 +481,10 @@ type FrontendAuthResponse struct { Principal string // Metadata carries plugin-defined identity attributes for downstream use. Metadata map[string]string + // DenyStatusCode optionally specifies an HTTP status code when denying the request. + DenyStatusCode int + // DenyMessage optionally specifies an error message when denying the request. + DenyMessage string } const ( @@ -469,6 +492,8 @@ const ( SchedulerBuiltinRoundRobin = "round-robin" // SchedulerBuiltinFillFirst delegates auth selection to the built-in fill-first scheduler. SchedulerBuiltinFillFirst = "fill-first" + // SelectionIDMetadataKey is the metadata key carrying the per-selection identifier in SchedulerOptions.Metadata. + SelectionIDMetadataKey = "selection_id" ) // Scheduler chooses an auth candidate before the built-in scheduler runs. @@ -695,6 +720,10 @@ type HostModelExecutionRequest struct { ProxyURL string `json:"proxy_url,omitempty"` // Path optionally specifies or overrides the request path (e.g. "/v1/images/generations" or "/v1/images/edits"). Path string `json:"path,omitempty"` + // CallerScope optionally carries the opaque caller scope handle from an intercepted + // request so the host execution resolves the same caller-scoped session and alias state + // without exposing the downstream client API key. + CallerScope string `json:"caller_scope,omitempty"` } // HostModelExecutionResponse describes a non-streaming host model execution response. @@ -852,6 +881,30 @@ type HostAuthSaveResponse struct { Path string `json:"path"` } +// HostAuthModelsRequest asks the host for the models registered for one or all credentials. +type HostAuthModelsRequest struct { + // AuthIndex identifies the credential index; empty requests all credentials. + AuthIndex string `json:"auth_index,omitempty"` +} + +// HostAuthModelsResponse returns the models registered for each matching credential. +type HostAuthModelsResponse struct { + // Credentials lists the credentials and their registered models. + Credentials []HostAuthModels `json:"credentials"` +} + +// HostAuthModels describes the models registered for one credential. +type HostAuthModels struct { + // ID identifies the credential record. + ID string `json:"id"` + // AuthIndex is the stable runtime credential index. + AuthIndex string `json:"auth_index"` + // Provider is the credential provider key. + Provider string `json:"provider"` + // Models lists the model identifiers registered for this credential in the global registry. + Models []string `json:"models"` +} + // Host affinity lookup status outcomes. const ( HostAffinityStatusBound = "bound" @@ -1095,6 +1148,9 @@ type WebSocketResponseObserver interface { // StreamChunkHeaderInitIndex marks the header-only stream initialization interceptor call. const StreamChunkHeaderInitIndex = -1 +// CallerScopeMetadataKey is the metadata key carrying the opaque downstream caller scope handle. +const CallerScopeMetadataKey = "caller_scope" + // RequestTransformRequest describes a request payload transformation. type RequestTransformRequest struct { // FromFormat is the source protocol format. @@ -1147,6 +1203,9 @@ type RequestInterceptRequest struct { Headers http.Header // Body contains the current request payload. Treat it as read-only; modifications must be returned in RequestInterceptResponse.Body. Body []byte + // CallerScope carries an opaque, irreversible handle identifying the downstream caller + // when available (mirrors Metadata[CallerScopeMetadataKey]). + CallerScope string // Metadata is a best-effort cloned context snapshot. Treat it as read-only and JSON-like. Metadata map[string]any } @@ -1488,6 +1547,8 @@ type UsageRecord struct { RequestID string // TraceID identifies the parent inbound HTTP request when available (8-character hex). TraceID string + // SelectionID identifies the scheduler selection decision that produced this usage record when available. + SelectionID string // Provider identifies the upstream provider. Provider string // BaseURL is the upstream base URL configured for the request/credential when available. @@ -1500,6 +1561,10 @@ type UsageRecord struct { Alias string // APIKey is the client API key identifier when available. APIKey string + // IsNativeKey reports whether the client API key was authenticated by CPA's native config provider. + IsNativeKey bool + // AccessProvider identifies the client request authentication provider (e.g. "config-inline" or plugin ID). + AccessProvider string // SessionID identifies the session when present. SessionID string // ParentSessionID identifies the parent session in a hierarchy or fork. @@ -1514,8 +1579,11 @@ type UsageRecord struct { Source string // ReasoningEffort records the requested reasoning effort. ReasoningEffort string - // ServiceTier records the requested or reported service tier. + // ServiceTier records the tier requested by the original client. ServiceTier string + // EffectiveServiceTier records the tier in the final outbound request after + // protocol translation and payload defaults or overrides. + EffectiveServiceTier string // ResponseServiceTier stores the final tier reported by the upstream response. ResponseServiceTier string // ResponseModel stores the model name reported by the upstream response, empty when unknown. diff --git a/sdk/pluginhost/host.go b/sdk/pluginhost/host.go index 0326b489..077ad81b 100644 --- a/sdk/pluginhost/host.go +++ b/sdk/pluginhost/host.go @@ -90,12 +90,13 @@ func (h *Host) ShutdownAllContext(ctx context.Context) { h.inner.ShutdownAllContext(ctx) } -// PluginBusy reports whether a plugin dynamic library is loaded or being loaded. +// PluginBusy reports whether a plugin is active in the host runtime or being loaded. func (h *Host) PluginBusy(id string) bool { return h != nil && h.inner != nil && h.inner.PluginBusy(id) } -// UnloadPlugin removes one plugin from the active runtime and closes its dynamic library. +// UnloadPlugin removes one plugin from the active runtime. +// The plugin image itself stays mapped for the lifetime of the process. func (h *Host) UnloadPlugin(id string) bool { return h.UnloadPluginContext(context.Background(), id) } @@ -321,6 +322,7 @@ func thinkingSupportToPlugin(thinking *internalregistry.ThinkingSupport) *Thinki Max: thinking.Max, ZeroAllowed: thinking.ZeroAllowed, DynamicAllowed: thinking.DynamicAllowed, + DisableType: thinking.DisableType, Levels: cloneStringSlice(thinking.Levels), } } diff --git a/sdk/pluginhost/quota.go b/sdk/pluginhost/quota.go index 8c98905b..37b306da 100644 --- a/sdk/pluginhost/quota.go +++ b/sdk/pluginhost/quota.go @@ -18,6 +18,15 @@ func (h *Host) QuotaProviders(ctx context.Context) []RegisteredQuotaProviderInfo return h.inner.QuotaProviders(ctx) } +// QuotaProviderPluginIDs maps normalized provider keys to the selected active quota plugin. +// Selection preserves host precedence and excludes plugins fused during capability discovery. +func (h *Host) QuotaProviderPluginIDs(ctx context.Context) map[string]string { + if h == nil || h.inner == nil { + return nil + } + return h.inner.QuotaProviderPluginIDs(ctx) +} + // HasQuotaProviderContext reports whether a provider handles credential quotas. func (h *Host) HasQuotaProviderContext(ctx context.Context, provider string) bool { return h != nil && h.inner != nil && h.inner.HasQuotaProviderContext(ctx, provider)