diff --git a/config.example.yaml b/config.example.yaml --- a/config.example.yaml +++ b/config.example.yaml @@ -49,6 +49,34 @@ enable: false addr: "127.0.0.1:8316" +# Credential concurrency is configured by Home in Home mode. The synthesized Home config is +# authoritative and local values, including the values below, are ignored. Do not use local +# configuration to override a Home concurrency policy. +# credential-concurrency: +# lifecycle-config-revision: 1 +# observation-barrier-revision: 0 +# cpa-heartbeat-timeout: "3s" +# cpa-cancel-bound: "5s" +# reclaim-grace: "5s" +# cleanup-interval: "5s" +# release-flush-interval: 250ms +# release-max-backoff: 2s +# busy-retry-min: 250ms +# busy-retry-max: 1s +# max-limit: 1000000 + +# Credential in-flight observation snapshot contract. +# credential-in-flight: +# snapshot-interval: 2s +# stale-after: 10s +# max-part-bytes: 262144 +# max-part-count: 64 +# max-revision-bytes: 16777216 +# max-aggregate-groups: 100000 +# max-details: 10000 +# max-string-bytes: 256 +# staging-retention: 1m + # Standard dynamic library plugins are trusted in-process code. They are disabled by default. # Build Go examples with go build -buildmode=c-shared for the target GOOS/GOARCH. # Other languages can implement the same C ABI and JSON method protocol. diff --git a/testdata/credential-concurrency-lifecycle.json b/testdata/credential-concurrency-lifecycle.json new file mode 100644 --- /dev/null +++ b/testdata/credential-concurrency-lifecycle.json @@ -0,0 +1,45 @@ +{ + "defaults": { + "lifecycle-config-revision": 1, + "observation-barrier-revision": 0, + "cpa-heartbeat-timeout": 3000000000, + "cpa-cancel-bound": 5000000000, + "reclaim-grace": 5000000000, + "cleanup-interval": 5000000000, + "release-flush-interval": "250ms", + "release-max-backoff": "2s", + "busy-retry-min": "250ms", + "busy-retry-max": "1s", + "max-limit": 1000000 + }, + "invalid": [ + { + "node_heartbeat_timeout": 3000000000, + "config": { + "cpa-heartbeat-timeout": 3000000000, + "cpa-cancel-bound": 5000000000, + "reclaim-grace": 5000000000, + "cleanup-interval": 5000000000, + "release-flush-interval": "250ms", + "release-max-backoff": "2s", + "busy-retry-min": "250ms", + "busy-retry-max": "1s", + "max-limit": 1000000 + } + }, + { + "node_heartbeat_timeout": 20000000000, + "config": { + "cpa-heartbeat-timeout": 0, + "cpa-cancel-bound": 5000000000, + "reclaim-grace": 5000000000, + "cleanup-interval": 5000000000, + "release-flush-interval": "250ms", + "release-max-backoff": "2s", + "busy-retry-min": "250ms", + "busy-retry-max": "1s", + "max-limit": 1000000 + } + } + ] +} diff --git a/internal/api/server.go b/internal/api/server.go --- a/internal/api/server.go +++ b/internal/api/server.go @@ -652,6 +652,13 @@ return sanitizedBody } +func homeSelectionAttemptContext(ctx context.Context, selection *auth.HomeDispatchSelection) (context.Context, func(), error) { + if selection == nil { + return nil, func() {}, errors.New("Home dispatch selection is nil") + } + return selection.AttemptContext(ctx) +} + // codexAlphaSearch forwards the standalone search endpoint used by current // Codex clients. Unlike /responses, this payload is already in Codex search // format and must not pass through a protocol translator. @@ -679,17 +686,46 @@ selectionHeaders.Set("X-Session-ID", sessionID) } ctx := context.WithValue(c.Request.Context(), "gin", c) - selected, err := s.handlers.AuthManager.SelectAuthByKind(ctx, "codex", strings.TrimSpace(routing.Model), auth.AuthKindOAuth, coreexecutor.Options{ - Headers: selectionHeaders, - OriginalRequest: body, - }) + selectionOpts := coreexecutor.Options{Headers: selectionHeaders, OriginalRequest: body} + var selection *auth.HomeDispatchSelection + var selected *auth.Auth + if s.handlers.AuthManager.HomeEnabled() { + selection, err = s.handlers.AuthManager.SelectHomeAuthByKind(ctx, "codex", strings.TrimSpace(routing.Model), auth.AuthKindOAuth, selectionOpts) + if selection != nil { + selected = selection.CloneAuth() + } + } else { + selected, err = s.handlers.AuthManager.SelectAuthByKind(ctx, "codex", strings.TrimSpace(routing.Model), auth.AuthKindOAuth, selectionOpts) + } if err != nil { status := http.StatusServiceUnavailable if statusError, ok := err.(interface{ StatusCode() int }); ok && statusError.StatusCode() > 0 { status = statusError.StatusCode() } + for _, value := range auth.SafeResponseHeaders(err).Values("Retry-After") { + c.Writer.Header().Add("Retry-After", value) + } c.JSON(status, gin.H{"error": err.Error()}) return + } + if selected == nil { + if selection != nil { + selection.End("missing_auth") + } + c.JSON(http.StatusServiceUnavailable, gin.H{"error": "Codex auth unavailable"}) + return + } + var releaseAttempt func() + if selection != nil { + attemptCtx, release, errBind := homeSelectionAttemptContext(ctx, selection) + if errBind != nil { + selection.End("attempt_bind_failed") + c.JSON(http.StatusServiceUnavailable, gin.H{"error": errBind.Error()}) + return + } + ctx = attemptCtx + releaseAttempt = release + defer releaseAttempt() } logging.SetGinCPATraceID(c, selected.EnsureIndex()) @@ -711,6 +747,9 @@ ctx, selected, http.MethodPost, upstreamURL, upstreamRequestBody, headers, ) if err != nil { + if selection != nil { + selection.End("request_build_failed") + } c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) return } @@ -734,17 +773,39 @@ AuthValue: authValue, }) + if errCtx := ctx.Err(); errCtx != nil { + if selection != nil { + selection.End("attempt_canceled") + } + c.JSON(http.StatusRequestTimeout, gin.H{"error": errCtx.Error()}) + return + } resp, err := s.handlers.AuthManager.HttpRequest(ctx, selected, req) if err != nil { + if selection != nil { + selection.End("request_failed") + } helps.RecordAPIResponseError(ctx, s.cfg, err) c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) return } - defer func() { - if errClose := resp.Body.Close(); errClose != nil { + closeResponseBody := func() error { + errClose := resp.Body.Close() + if errClose != nil { log.Errorf("codex alpha search: close response body error: %v", errClose) } - }() + return errClose + } + if selection != nil { + if errBind := selection.Bind(closeResponseBody); errBind != nil { + selection.End("response_bind_failed") + c.JSON(http.StatusServiceUnavailable, gin.H{"error": errBind.Error()}) + return + } + defer selection.End("response_closed") + } else { + defer func() { _ = closeResponseBody() }() + } helps.RecordAPIResponseMetadata(ctx, s.cfg, resp.StatusCode, resp.Header.Clone()) upstreamBody, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20)) if err != nil { @@ -1812,6 +1873,21 @@ // - clients: The new slice of AI service clients // - cfg: The new application configuration func (s *Server) UpdateClients(cfg *config.Config) { + s.UpdateClientsContext(context.Background(), cfg) +} + +// UpdateClientsContext updates runtime clients while honoring cancellation between +// short configuration and filesystem operations. +func (s *Server) UpdateClientsContext(ctx context.Context, cfg *config.Config) bool { + if s == nil || cfg == nil { + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } // Reconstruct old config from YAML snapshot to avoid reference sharing issues var oldCfg *config.Config if len(s.oldConfigYaml) > 0 { @@ -1840,6 +1916,9 @@ if oldCfg == nil || oldCfg.LoggingToFile != cfg.LoggingToFile || oldCfg.LogsMaxTotalSizeMB != cfg.LogsMaxTotalSizeMB { if err := logging.ConfigureLogOutput(cfg); err != nil { log.Errorf("failed to reconfigure log output: %v", err) + } + if errContext := ctx.Err(); errContext != nil { + return false } } @@ -1926,6 +2005,9 @@ s.wsAuthChanged(oldCfg.WebsocketAuth, cfg.WebsocketAuth) } managementasset.SetCurrentConfig(cfg) + if errContext := ctx.Err(); errContext != nil { + return false + } // Save YAML snapshot for next comparison s.oldConfigYaml, _ = yaml.Marshal(cfg) @@ -1950,7 +2032,10 @@ if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok { dirSetter.SetBaseDir(cfg.AuthDir) } - authEntries = util.CountAuthFiles(context.Background(), tokenStore) + authEntries = util.CountAuthFiles(ctx, tokenStore) + if errContext := ctx.Err(); errContext != nil { + return false + } } geminiAPIKeyCount := len(cfg.GeminiKey) interactionsAPIKeyCount := len(cfg.InteractionsKey) @@ -1979,6 +2064,7 @@ vertexAICompatCount, openAICompatCount, ) + return ctx.Err() == nil } func (s *Server) SetWebsocketAuthChangeHandler(fn func(bool, bool)) { diff --git a/internal/api/server_test.go b/internal/api/server_test.go --- a/internal/api/server_test.go +++ b/internal/api/server_test.go @@ -3,12 +3,15 @@ import ( "context" "encoding/json" + "errors" "io" "net/http" "net/http/httptest" "os" "path/filepath" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -21,14 +24,18 @@ "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" ) type codexSearchCaptureExecutor struct { - request *http.Request - body []byte - authIDs []string + request *http.Request + body []byte + authIDs []string + prepareErr error + httpErr error + responseBody io.ReadCloser } func (e *codexSearchCaptureExecutor) Identifier() string { return "codex" } @@ -50,6 +57,9 @@ } func (e *codexSearchCaptureExecutor) PrepareRequest(req *http.Request, a *auth.Auth) error { + if e.prepareErr != nil { + return e.prepareErr + } token, _ := a.Metadata["access_token"].(string) req.Header.Set("Authorization", "Bearer "+token) return nil @@ -82,6 +92,9 @@ } func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *auth.Auth, req *http.Request) (*http.Response, error) { + if e.httpErr != nil { + return nil, e.httpErr + } e.request = req.Clone(req.Context()) e.authIDs = append(e.authIDs, selected.ID) body, err := io.ReadAll(req.Body) @@ -89,11 +102,248 @@ return nil, err } e.body = body + responseBody := e.responseBody + if responseBody == nil { + responseBody = io.NopCloser(strings.NewReader(`{"results":[{"url":"https://example.com"}]}`)) + } return &http.Response{ StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: io.NopCloser(strings.NewReader(`{"results":[{"url":"https://example.com"}]}`)), + Body: responseBody, }, nil +} + +type codexSearchHomeDispatcher struct { + calls atomic.Int32 +} + +func (*codexSearchHomeDispatcher) HeartbeatOK() bool { return true } + +func (d *codexSearchHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + d.calls.Add(1) + return json.Marshal(map[string]any{ + "model": model, + "auth_index": "home-codex-search", + "auth": map[string]any{ + "id": "home-codex-search", + "provider": "codex", + "status": "active", + "metadata": map[string]any{"access_token": "home-search-token"}, + }, + "concurrency": map[string]any{ + "accounted": true, + "credential_id": "home-codex-search", + "model": model, + }, + }) +} + +func (*codexSearchHomeDispatcher) AbortAmbiguousDispatch() {} + +type codexSearchBusyHomeDispatcher struct{} + +func (*codexSearchBusyHomeDispatcher) HeartbeatOK() bool { return true } +func (*codexSearchBusyHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + return []byte(`{"error":{"type":"credential_concurrency_exceeded","message":"busy","retry_after_ms":750}}`), nil +} +func (*codexSearchBusyHomeDispatcher) AbortAmbiguousDispatch() {} + +type trackedSearchResponseBody struct { + io.Reader + closed atomic.Bool +} + +func (b *trackedSearchResponseBody) Close() error { + b.closed.Store(true) + return nil +} + +type drainAwareSearchResponseBody struct { + started chan struct{} + closed chan struct{} + startOnce sync.Once + closeOnce sync.Once +} + +func newDrainAwareSearchResponseBody() *drainAwareSearchResponseBody { + return &drainAwareSearchResponseBody{started: make(chan struct{}), closed: make(chan struct{})} +} + +func (b *drainAwareSearchResponseBody) Read([]byte) (int, error) { + b.startOnce.Do(func() { close(b.started) }) + <-b.closed + return 0, io.EOF +} + +func (b *drainAwareSearchResponseBody) Close() error { + b.closeOnce.Do(func() { close(b.closed) }) + return nil +} + +func TestAuditHomeBusyNormalAndStream429Headers(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "normal", true: "stream"}[stream], func(t *testing.T) { + server := newTestServer(t) + server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}}) + server.handlers.AuthManager.PublishHomeDispatch(&codexSearchBusyHomeDispatcher{}, executionregistry.New(), 1) + + body := `{"model":"gpt-5-codex","input":[]}` + if stream { + body = `{"model":"gpt-5-codex","input":[],"stream":true}` + } + rr := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer test-key") + server.engine.ServeHTTP(rr, req) + + if rr.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusTooManyRequests, rr.Body.String()) + } + if got := rr.Header().Get("Retry-After"); got != "1" { + t.Fatalf("Retry-After = %q, want 1", got) + } + }) + } +} + +func TestAuditHomeCodexSearchBusyReturnsTrustedRetryAfter(t *testing.T) { + server := newTestServer(t) + server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}}) + server.handlers.AuthManager.PublishHomeDispatch(&codexSearchBusyHomeDispatcher{}, executionregistry.New(), 1) + + rr := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5-codex","query":"test"}`)) + req.Header.Set("Authorization", "Bearer test-key") + server.engine.ServeHTTP(rr, req) + + if rr.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusTooManyRequests, rr.Body.String()) + } + if got := rr.Header().Get("Retry-After"); got != "1" { + t.Fatalf("Retry-After = %q, want 1", got) + } + if !strings.Contains(rr.Body.String(), "busy") { + t.Fatalf("body = %q, want busy error", rr.Body.String()) + } +} + +func TestAuditHomeCodexSearchBodyCloseBeforeRelease(t *testing.T) { + server := newTestServer(t) + dispatcher := &codexSearchHomeDispatcher{} + registry := executionregistry.New() + body := newDrainAwareSearchResponseBody() + var releaseAfterBodyClose atomic.Bool + var releaseCount atomic.Int32 + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { + if group != (executionregistry.ReleaseGroup{CredentialID: "home-codex-search", Model: "gpt-5-codex"}) { + t.Errorf("release group = %#v", group) + } + select { + case <-body.closed: + releaseAfterBodyClose.Store(true) + default: + } + releaseCount.Add(1) + }) + server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}}) + server.handlers.AuthManager.PublishHomeDispatch(dispatcher, registry, 1) + executor := &codexSearchCaptureExecutor{responseBody: body} + server.handlers.AuthManager.RegisterExecutor(executor) + + rr := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"id":"home-search-drain","model":"gpt-5-codex","query":"test"}`)) + req.Header.Set("Authorization", "Bearer test-key") + handlerDone := make(chan struct{}) + go func() { + server.engine.ServeHTTP(rr, req) + close(handlerDone) + }() + + select { + case <-body.started: + case <-time.After(time.Second): + t.Fatal("search handler did not start reading the response body") + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } + if got := releaseCount.Load(); got != 1 { + t.Fatalf("accounted releases = %d, want 1", got) + } + if !releaseAfterBodyClose.Load() { + t.Fatal("accounted Home selection released before the search response body closed") + } + select { + case <-handlerDone: + case <-time.After(time.Second): + t.Fatal("search handler remained blocked after Home drain") + } + if rr.Code != http.StatusOK { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String()) + } +} + +func TestHomeCodexAlphaSearchEndsSelectionAcrossDirectHTTPPaths(t *testing.T) { + tests := []struct { + name string + configure func(*codexSearchCaptureExecutor, *trackedSearchResponseBody) + wantStatus int + wantClosed bool + }{ + { + name: "request build failure", + configure: func(executor *codexSearchCaptureExecutor, _ *trackedSearchResponseBody) { + executor.prepareErr = errors.New("request preparation failed") + }, + wantStatus: http.StatusBadGateway, + }, + { + name: "HTTP error", + configure: func(executor *codexSearchCaptureExecutor, _ *trackedSearchResponseBody) { + executor.httpErr = errors.New("upstream unavailable") + }, + wantStatus: http.StatusBadGateway, + }, + { + name: "response body close", + configure: func(executor *codexSearchCaptureExecutor, body *trackedSearchResponseBody) { + executor.responseBody = body + }, + wantStatus: http.StatusOK, + wantClosed: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := newTestServer(t) + dispatcher := &codexSearchHomeDispatcher{} + registry := executionregistry.New() + server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}}) + server.handlers.AuthManager.PublishHomeDispatch(dispatcher, registry, 1) + body := &trackedSearchResponseBody{Reader: strings.NewReader(`{"results":[]}`)} + executor := &codexSearchCaptureExecutor{} + test.configure(executor, body) + server.handlers.AuthManager.RegisterExecutor(executor) + + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"id":"home-search-session","model":"gpt-5-codex","query":"test"}`)) + req.Header.Set("Authorization", "Bearer test-key") + rr := httptest.NewRecorder() + server.engine.ServeHTTP(rr, req) + if rr.Code != test.wantStatus { + t.Fatalf("status = %d, want %d; body=%s", rr.Code, test.wantStatus, rr.Body.String()) + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1", got) + } + if got := body.closed.Load(); got != test.wantClosed { + t.Fatalf("response body closed = %t, want %t", got, test.wantClosed) + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } + }) + } } func newTestServer(t *testing.T) *Server { diff --git a/internal/config/config.go b/internal/config/config.go --- a/internal/config/config.go +++ b/internal/config/config.go @@ -41,6 +41,12 @@ // Home config is runtime-only and is populated from -home-jwt. Home HomeConfig `yaml:"-" json:"-"` + // CredentialConcurrency contains Home-authoritative credential lifecycle settings. + CredentialConcurrency CredentialConcurrencyConfig `yaml:"credential-concurrency" json:"credential-concurrency"` + + // CredentialInFlight configures credential observation snapshots. + CredentialInFlight CredentialInFlightConfig `yaml:"credential-in-flight" json:"credential-in-flight"` + // RemoteManagement nests management-related options under 'remote-management'. RemoteManagement RemoteManagement `yaml:"remote-management" json:"-"` @@ -730,7 +736,7 @@ if optional { if os.IsNotExist(err) || errors.Is(err, syscall.EISDIR) { // Missing and optional: return empty config (cloud deploy standby). - cfg := &Config{} + cfg := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} cfg.NormalizePluginsConfig() return cfg, nil } @@ -739,8 +745,8 @@ } // In cloud deploy mode (optional=true), if file is empty or contains only whitespace, return empty config. - if optional && len(data) == 0 { - cfg := &Config{} + if optional && len(bytes.TrimSpace(data)) == 0 { + cfg := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} cfg.NormalizePluginsConfig() return cfg, nil } @@ -762,14 +768,20 @@ cfg.Pprof.Enable = false cfg.Pprof.Addr = DefaultPprofAddr cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository + cfg.CredentialInFlight = DefaultCredentialInFlightConfig() if err = yaml.Unmarshal(data, &cfg); err != nil { if optional { // In cloud deploy mode, if YAML parsing fails, return empty config instead of error. - cfgOptional := &Config{} + cfgOptional := &Config{CredentialInFlight: DefaultCredentialInFlightConfig()} cfgOptional.NormalizePluginsConfig() return cfgOptional, nil } return nil, fmt.Errorf("failed to parse config file: %w", err) + } + + cfg.CredentialConcurrency = cfg.CredentialConcurrency.WithDefaults() + if errValidate := cfg.CredentialInFlight.Validate(); errValidate != nil { + return nil, errValidate } // Hash remote management key if plaintext is detected (nested) diff --git a/internal/config/credential_concurrency.go b/internal/config/credential_concurrency.go new file mode 100644 --- /dev/null +++ b/internal/config/credential_concurrency.go @@ -0,0 +1,194 @@ +package config + +import ( + "fmt" + "time" + + "gopkg.in/yaml.v3" +) + +const ( + defaultCPAHeartbeatTimeout = 3 * time.Second + defaultCPACancelBound = 5 * time.Second + defaultReclaimGrace = 5 * time.Second + defaultCleanupInterval = 5 * time.Second + defaultReleaseFlushInterval = 250 * time.Millisecond + defaultReleaseMaxBackoff = 2 * time.Second + defaultBusyRetryMin = 250 * time.Millisecond + defaultBusyRetryMax = time.Second + maxCredentialConcurrencyLimit int64 = 1_000_000 +) + +// CredentialConcurrencyConfig controls the credential concurrency lifecycle managed by Home. +type CredentialConcurrencyConfig struct { + LifecycleConfigRevision int64 `yaml:"lifecycle-config-revision" json:"lifecycle-config-revision"` + ObservationBarrierRevision int64 `yaml:"observation-barrier-revision" json:"observation-barrier-revision"` + CPAHeartbeatTimeout time.Duration `yaml:"cpa-heartbeat-timeout" json:"cpa-heartbeat-timeout"` + CPACancelBound time.Duration `yaml:"cpa-cancel-bound" json:"cpa-cancel-bound"` + ReclaimGrace time.Duration `yaml:"reclaim-grace" json:"reclaim-grace"` + CleanupInterval time.Duration `yaml:"cleanup-interval" json:"cleanup-interval"` + ReleaseFlushInterval time.Duration `yaml:"release-flush-interval" json:"release-flush-interval"` + ReleaseMaxBackoff time.Duration `yaml:"release-max-backoff" json:"release-max-backoff"` + BusyRetryMin time.Duration `yaml:"busy-retry-min" json:"busy-retry-min"` + BusyRetryMax time.Duration `yaml:"busy-retry-max" json:"busy-retry-max"` + MaxLimit int64 `yaml:"max-limit" json:"max-limit"` + + lifecycleConfigRevisionPresent bool + observationBarrierRevisionPresent bool + cpaHeartbeatTimeoutPresent bool + cpaCancelBoundPresent bool + reclaimGracePresent bool + cleanupIntervalPresent bool + releaseFlushIntervalPresent bool + releaseMaxBackoffPresent bool + busyRetryMinPresent bool + busyRetryMaxPresent bool + maxLimitPresent bool +} + +// UnmarshalYAML preserves field presence so only absent lifecycle values receive legacy defaults. +func (c *CredentialConcurrencyConfig) UnmarshalYAML(value *yaml.Node) error { + type rawCredentialConcurrencyConfig struct { + LifecycleConfigRevision int64 `yaml:"lifecycle-config-revision"` + ObservationBarrierRevision int64 `yaml:"observation-barrier-revision"` + CPAHeartbeatTimeout time.Duration `yaml:"cpa-heartbeat-timeout"` + CPACancelBound time.Duration `yaml:"cpa-cancel-bound"` + ReclaimGrace time.Duration `yaml:"reclaim-grace"` + CleanupInterval time.Duration `yaml:"cleanup-interval"` + ReleaseFlushInterval time.Duration `yaml:"release-flush-interval"` + ReleaseMaxBackoff time.Duration `yaml:"release-max-backoff"` + BusyRetryMin time.Duration `yaml:"busy-retry-min"` + BusyRetryMax time.Duration `yaml:"busy-retry-max"` + MaxLimit int64 `yaml:"max-limit"` + } + + var raw rawCredentialConcurrencyConfig + if errDecode := value.Decode(&raw); errDecode != nil { + return errDecode + } + + *c = CredentialConcurrencyConfig{ + LifecycleConfigRevision: raw.LifecycleConfigRevision, + ObservationBarrierRevision: raw.ObservationBarrierRevision, + CPAHeartbeatTimeout: raw.CPAHeartbeatTimeout, + CPACancelBound: raw.CPACancelBound, + ReclaimGrace: raw.ReclaimGrace, + CleanupInterval: raw.CleanupInterval, + ReleaseFlushInterval: raw.ReleaseFlushInterval, + ReleaseMaxBackoff: raw.ReleaseMaxBackoff, + BusyRetryMin: raw.BusyRetryMin, + BusyRetryMax: raw.BusyRetryMax, + MaxLimit: raw.MaxLimit, + lifecycleConfigRevisionPresent: credentialConcurrencyFieldPresent(value, "lifecycle-config-revision"), + observationBarrierRevisionPresent: credentialConcurrencyFieldPresent(value, "observation-barrier-revision"), + cpaHeartbeatTimeoutPresent: credentialConcurrencyFieldPresent(value, "cpa-heartbeat-timeout"), + cpaCancelBoundPresent: credentialConcurrencyFieldPresent(value, "cpa-cancel-bound"), + reclaimGracePresent: credentialConcurrencyFieldPresent(value, "reclaim-grace"), + cleanupIntervalPresent: credentialConcurrencyFieldPresent(value, "cleanup-interval"), + releaseFlushIntervalPresent: credentialConcurrencyFieldPresent(value, "release-flush-interval"), + releaseMaxBackoffPresent: credentialConcurrencyFieldPresent(value, "release-max-backoff"), + busyRetryMinPresent: credentialConcurrencyFieldPresent(value, "busy-retry-min"), + busyRetryMaxPresent: credentialConcurrencyFieldPresent(value, "busy-retry-max"), + maxLimitPresent: credentialConcurrencyFieldPresent(value, "max-limit"), + } + return nil +} + +func credentialConcurrencyFieldPresent(value *yaml.Node, field string) bool { + if value == nil || value.Kind != yaml.MappingNode { + return false + } + for index := 0; index+1 < len(value.Content); index += 2 { + if value.Content[index].Value == field { + return true + } + } + return false +} + +// WithDefaults applies the lifecycle defaults required for compatibility with older Home versions. +func (c CredentialConcurrencyConfig) WithDefaults() CredentialConcurrencyConfig { + if !c.cpaHeartbeatTimeoutPresent && c.CPAHeartbeatTimeout == 0 { + c.CPAHeartbeatTimeout = defaultCPAHeartbeatTimeout + } + if !c.cpaCancelBoundPresent && c.CPACancelBound == 0 { + c.CPACancelBound = defaultCPACancelBound + } + if !c.reclaimGracePresent && c.ReclaimGrace == 0 { + c.ReclaimGrace = defaultReclaimGrace + } + if !c.cleanupIntervalPresent && c.CleanupInterval == 0 { + c.CleanupInterval = defaultCleanupInterval + } + if !c.releaseFlushIntervalPresent && c.ReleaseFlushInterval == 0 { + c.ReleaseFlushInterval = defaultReleaseFlushInterval + } + if !c.releaseMaxBackoffPresent && c.ReleaseMaxBackoff == 0 { + c.ReleaseMaxBackoff = defaultReleaseMaxBackoff + } + if !c.busyRetryMinPresent && c.BusyRetryMin == 0 { + c.BusyRetryMin = defaultBusyRetryMin + } + if !c.busyRetryMaxPresent && c.BusyRetryMax == 0 { + c.BusyRetryMax = defaultBusyRetryMax + } + if !c.maxLimitPresent && c.MaxLimit == 0 { + c.MaxLimit = maxCredentialConcurrencyLimit + } + return c +} + +// ValidateCredentialConcurrency validates values intrinsic to a credential concurrency configuration. +func ValidateCredentialConcurrency(cfg CredentialConcurrencyConfig) error { + if cfg.LifecycleConfigRevision < 0 || (cfg.lifecycleConfigRevisionPresent && cfg.LifecycleConfigRevision == 0) { + return fmt.Errorf("lifecycle configuration revision must be positive when present") + } + if cfg.ObservationBarrierRevision < 0 { + return fmt.Errorf("observation barrier revision must not be negative") + } + if cfg.CPAHeartbeatTimeout <= 0 || cfg.CPACancelBound <= 0 || cfg.ReclaimGrace <= 0 || cfg.CleanupInterval <= 0 { + return fmt.Errorf("credential concurrency lifecycle durations must be positive") + } + if cfg.ReleaseFlushInterval <= 0 || cfg.ReleaseMaxBackoff <= 0 || cfg.BusyRetryMin <= 0 || cfg.BusyRetryMax <= 0 { + return fmt.Errorf("credential concurrency limiter durations must be positive") + } + if cfg.ReleaseMaxBackoff < cfg.ReleaseFlushInterval { + return fmt.Errorf("credential concurrency release max backoff must not be less than release flush interval") + } + if cfg.BusyRetryMin%time.Millisecond != 0 || cfg.BusyRetryMax%time.Millisecond != 0 { + return fmt.Errorf("credential concurrency busy retry durations must be whole milliseconds") + } + if cfg.BusyRetryMax < cfg.BusyRetryMin { + return fmt.Errorf("credential concurrency busy retry max must not be less than busy retry min") + } + if cfg.MaxLimit < 1 || cfg.MaxLimit > maxCredentialConcurrencyLimit { + return fmt.Errorf("credential concurrency max limit must be between 1 and %d", maxCredentialConcurrencyLimit) + } + return nil +} + +// ValidateCredentialConcurrencyLifecycle verifies the Home lifecycle timing safety invariant. +func ValidateCredentialConcurrencyLifecycle(nodeHeartbeatTimeout time.Duration, cfg CredentialConcurrencyConfig) error { + if nodeHeartbeatTimeout <= 0 { + return fmt.Errorf("credential concurrency lifecycle durations must be positive") + } + if errValidate := ValidateCredentialConcurrency(cfg); errValidate != nil { + return errValidate + } + left, leftOverflow := addCredentialConcurrencyDuration(nodeHeartbeatTimeout, cfg.ReclaimGrace) + right, rightOverflow := addCredentialConcurrencyDuration(cfg.CPAHeartbeatTimeout, cfg.CPACancelBound) + if leftOverflow || rightOverflow { + return fmt.Errorf("credential concurrency lifecycle timing safety invariant overflows") + } + if left <= right { + return fmt.Errorf("node heartbeat timeout plus reclaim grace must exceed CPA heartbeat timeout plus cancel bound") + } + return nil +} + +func addCredentialConcurrencyDuration(left time.Duration, right time.Duration) (time.Duration, bool) { + if right > 0 && left > time.Duration(1<<63-1)-right { + return 0, true + } + return left + right, false +} diff --git a/internal/config/credential_concurrency_fixture_test.go b/internal/config/credential_concurrency_fixture_test.go new file mode 100644 --- /dev/null +++ b/internal/config/credential_concurrency_fixture_test.go @@ -0,0 +1,156 @@ +package config + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "testing" + "time" + + "gopkg.in/yaml.v3" +) + +type credentialConcurrencyFixtureWireConfig struct { + LifecycleConfigRevision int64 `json:"lifecycle-config-revision"` + ObservationBarrierRevision int64 `json:"observation-barrier-revision"` + CPAHeartbeatTimeout time.Duration `json:"cpa-heartbeat-timeout"` + CPACancelBound time.Duration `json:"cpa-cancel-bound"` + ReclaimGrace time.Duration `json:"reclaim-grace"` + CleanupInterval time.Duration `json:"cleanup-interval"` + ReleaseFlushInterval string `json:"release-flush-interval" yaml:"release-flush-interval"` + ReleaseMaxBackoff string `json:"release-max-backoff" yaml:"release-max-backoff"` + BusyRetryMin string `json:"busy-retry-min" yaml:"busy-retry-min"` + BusyRetryMax string `json:"busy-retry-max" yaml:"busy-retry-max"` + MaxLimit int64 `json:"max-limit"` +} + +type credentialConcurrencyFixtureHotDurations struct { + ReleaseFlushInterval time.Duration `yaml:"release-flush-interval"` + ReleaseMaxBackoff time.Duration `yaml:"release-max-backoff"` + BusyRetryMin time.Duration `yaml:"busy-retry-min"` + BusyRetryMax time.Duration `yaml:"busy-retry-max"` +} + +func (c credentialConcurrencyFixtureWireConfig) config() (CredentialConcurrencyConfig, error) { + raw, errMarshal := yaml.Marshal(c) + if errMarshal != nil { + return CredentialConcurrencyConfig{}, fmt.Errorf("marshal fixture hot durations as YAML: %w", errMarshal) + } + var hot credentialConcurrencyFixtureHotDurations + if errUnmarshal := yaml.Unmarshal(raw, &hot); errUnmarshal != nil { + return CredentialConcurrencyConfig{}, fmt.Errorf("parse fixture hot durations as YAML: %w", errUnmarshal) + } + return CredentialConcurrencyConfig{ + LifecycleConfigRevision: c.LifecycleConfigRevision, + ObservationBarrierRevision: c.ObservationBarrierRevision, + CPAHeartbeatTimeout: c.CPAHeartbeatTimeout, + CPACancelBound: c.CPACancelBound, + ReclaimGrace: c.ReclaimGrace, + CleanupInterval: c.CleanupInterval, + ReleaseFlushInterval: hot.ReleaseFlushInterval, + ReleaseMaxBackoff: hot.ReleaseMaxBackoff, + BusyRetryMin: hot.BusyRetryMin, + BusyRetryMax: hot.BusyRetryMax, + MaxLimit: c.MaxLimit, + }, nil +} + +func TestCredentialConcurrencyLifecycleFixture(t *testing.T) { + raw, errRead := os.ReadFile(filepath.Join("..", "..", "testdata", "credential-concurrency-lifecycle.json")) + if errRead != nil { + t.Fatal(errRead) + } + var fixture struct { + Defaults credentialConcurrencyFixtureWireConfig `json:"defaults"` + Invalid []struct { + NodeHeartbeatTimeout time.Duration `json:"node_heartbeat_timeout"` + Config credentialConcurrencyFixtureWireConfig `json:"config"` + } `json:"invalid"` + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if errDecode := decoder.Decode(&fixture); errDecode != nil { + t.Fatal(errDecode) + } + if errTrailing := decoder.Decode(&struct{}{}); errTrailing != io.EOF { + t.Fatalf("fixture contains trailing JSON: %v", errTrailing) + } + + defaults, errConfig := fixture.Defaults.config() + if errConfig != nil { + t.Fatal(errConfig) + } + + expectedDefaults := CredentialConcurrencyConfig{ + LifecycleConfigRevision: 1, + ObservationBarrierRevision: 0, + CPAHeartbeatTimeout: 3 * time.Second, + CPACancelBound: 5 * time.Second, + ReclaimGrace: 5 * time.Second, + CleanupInterval: 5 * time.Second, + ReleaseFlushInterval: 250 * time.Millisecond, + ReleaseMaxBackoff: 2 * time.Second, + BusyRetryMin: 250 * time.Millisecond, + BusyRetryMax: time.Second, + MaxLimit: 1_000_000, + } + if defaults != expectedDefaults { + t.Fatalf("defaults = %#v, want %#v", defaults, expectedDefaults) + } + if errValidate := ValidateCredentialConcurrency(defaults); errValidate != nil { + t.Fatalf("ValidateCredentialConcurrency(defaults) error = %v", errValidate) + } + + expectedInvalid := []struct { + nodeHeartbeatTimeout time.Duration + config CredentialConcurrencyConfig + }{ + { + nodeHeartbeatTimeout: 3 * time.Second, + config: CredentialConcurrencyConfig{ + CPAHeartbeatTimeout: 3 * time.Second, + CPACancelBound: 5 * time.Second, + ReclaimGrace: 5 * time.Second, + CleanupInterval: 5 * time.Second, + ReleaseFlushInterval: 250 * time.Millisecond, + ReleaseMaxBackoff: 2 * time.Second, + BusyRetryMin: 250 * time.Millisecond, + BusyRetryMax: time.Second, + MaxLimit: 1_000_000, + }, + }, + { + nodeHeartbeatTimeout: 20 * time.Second, + config: CredentialConcurrencyConfig{ + CPAHeartbeatTimeout: 0, + CPACancelBound: 5 * time.Second, + ReclaimGrace: 5 * time.Second, + CleanupInterval: 5 * time.Second, + ReleaseFlushInterval: 250 * time.Millisecond, + ReleaseMaxBackoff: 2 * time.Second, + BusyRetryMin: 250 * time.Millisecond, + BusyRetryMax: time.Second, + MaxLimit: 1_000_000, + }, + }, + } + if len(fixture.Invalid) != len(expectedInvalid) { + t.Fatalf("invalid fixture count = %d, want %d", len(fixture.Invalid), len(expectedInvalid)) + } + for index, expected := range expectedInvalid { + item := fixture.Invalid[index] + itemConfig, errConfig := item.Config.config() + if errConfig != nil { + t.Fatalf("invalid fixture %d config() error = %v", index, errConfig) + } + if item.NodeHeartbeatTimeout != expected.nodeHeartbeatTimeout || itemConfig != expected.config { + t.Fatalf("invalid fixture %d = %#v, want node heartbeat timeout %s and config %#v", index, itemConfig, expected.nodeHeartbeatTimeout, expected.config) + } + if errValidate := ValidateCredentialConcurrencyLifecycle(item.NodeHeartbeatTimeout, itemConfig); errValidate == nil { + t.Fatalf("invalid fixture %d passed", index) + } + } +} diff --git a/internal/config/credential_concurrency_test.go b/internal/config/credential_concurrency_test.go new file mode 100644 --- /dev/null +++ b/internal/config/credential_concurrency_test.go @@ -0,0 +1,124 @@ +package config + +import ( + "testing" + "time" +) + +func TestCredentialConcurrencyLimiterConfig(t *testing.T) { + got := (CredentialConcurrencyConfig{}).WithDefaults() + if got.LifecycleConfigRevision != 0 || got.ObservationBarrierRevision != 0 { + t.Fatalf("default revisions = %d, %d, want 0, 0", got.LifecycleConfigRevision, got.ObservationBarrierRevision) + } + if got.CPAHeartbeatTimeout != 3*time.Second || got.CPACancelBound != 5*time.Second || got.ReclaimGrace != 5*time.Second || got.CleanupInterval != 5*time.Second { + t.Fatalf("default lifecycle config = %#v", got) + } + if got.ReleaseFlushInterval != 250*time.Millisecond || got.ReleaseMaxBackoff != 2*time.Second || got.BusyRetryMin != 250*time.Millisecond || got.BusyRetryMax != time.Second || got.MaxLimit != 1_000_000 { + t.Fatalf("default limiter config = %#v", got) + } + if errValidate := ValidateCredentialConcurrencyLifecycle(20*time.Second, got); errValidate != nil { + t.Fatalf("ValidateCredentialConcurrencyLifecycle() error = %v", errValidate) + } + if errValidate := ValidateCredentialConcurrencyLifecycle(2*time.Second, got); errValidate == nil { + t.Fatal("ValidateCredentialConcurrencyLifecycle() error = nil, want timing invariant failure") + } +} + +func TestValidateCredentialConcurrencyAcceptsHomeAuthoritativeHeartbeat(t *testing.T) { + cfg := (CredentialConcurrencyConfig{}).WithDefaults() + cfg.CPAHeartbeatTimeout = 20 * time.Second + + if errValidate := ValidateCredentialConcurrency(cfg); errValidate != nil { + t.Fatalf("ValidateCredentialConcurrency() error = %v", errValidate) + } + if errValidate := ValidateCredentialConcurrencyLifecycle(20*time.Second, cfg); errValidate == nil { + t.Fatal("ValidateCredentialConcurrencyLifecycle() error = nil, want Home timing invariant failure") + } +} + +func TestCredentialConcurrencyConfigDefaultsOnlyMissingFields(t *testing.T) { + tests := []struct { + name string + payload string + }{ + { + name: "explicit zero revision", + payload: "credential-concurrency:\n" + + " lifecycle-config-revision: 0\n" + + " cpa-heartbeat-timeout: 3s\n" + + " cpa-cancel-bound: 5s\n" + + " reclaim-grace: 5s\n" + + " cleanup-interval: 5s\n", + }, + { + name: "explicit zero duration", + payload: "credential-concurrency:\n" + + " lifecycle-config-revision: 1\n" + + " cpa-heartbeat-timeout: 0s\n" + + " cpa-cancel-bound: 5s\n" + + " reclaim-grace: 5s\n" + + " cleanup-interval: 5s\n", + }, + { + name: "explicit null duration", + payload: "credential-concurrency:\n" + + " lifecycle-config-revision: 1\n" + + " cpa-heartbeat-timeout: null\n" + + " cpa-cancel-bound: 5s\n" + + " reclaim-grace: 5s\n" + + " cleanup-interval: 5s\n", + }, + { + name: "negative observation barrier", + payload: "credential-concurrency:\n" + + " lifecycle-config-revision: 1\n" + + " observation-barrier-revision: -1\n" + + " cpa-heartbeat-timeout: 3s\n" + + " cpa-cancel-bound: 5s\n" + + " reclaim-grace: 5s\n" + + " cleanup-interval: 5s\n", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + parsed, errParse := ParseConfigBytes([]byte(test.payload)) + if errParse != nil { + t.Fatalf("ParseConfigBytes() error = %v", errParse) + } + if errValidate := ValidateCredentialConcurrencyLifecycle(20*time.Second, parsed.CredentialConcurrency); errValidate == nil { + t.Fatal("ValidateCredentialConcurrencyLifecycle() error = nil, want explicit invalid lifecycle value rejection") + } + }) + } +} + +func TestCredentialConcurrencyConfigRejectsInvalidLimiter(t *testing.T) { + tests := []CredentialConcurrencyConfig{ + {ReleaseFlushInterval: time.Second, ReleaseMaxBackoff: 500 * time.Millisecond, BusyRetryMin: time.Millisecond, BusyRetryMax: time.Millisecond, MaxLimit: 1}, + {ReleaseFlushInterval: time.Millisecond, ReleaseMaxBackoff: time.Millisecond, BusyRetryMin: 1500 * time.Microsecond, BusyRetryMax: 2 * time.Millisecond, MaxLimit: 1}, + {ReleaseFlushInterval: time.Millisecond, ReleaseMaxBackoff: time.Millisecond, BusyRetryMin: time.Millisecond, BusyRetryMax: time.Millisecond, MaxLimit: 1_000_001}, + } + for _, cfg := range tests { + cfg.CPAHeartbeatTimeout = 3 * time.Second + cfg.CPACancelBound = 5 * time.Second + cfg.ReclaimGrace = 5 * time.Second + cfg.CleanupInterval = 5 * time.Second + if errValidate := ValidateCredentialConcurrencyLifecycle(20*time.Second, cfg); errValidate == nil { + t.Fatalf("ValidateCredentialConcurrencyLifecycle(%#v) error = nil", cfg) + } + } +} + +func TestValidateCredentialConcurrencyLifecycleRejectsSafetyOverflow(t *testing.T) { + cfg := CredentialConcurrencyConfig{ + LifecycleConfigRevision: 1, + CPAHeartbeatTimeout: time.Duration(1<<63 - 1), + CPACancelBound: time.Nanosecond, + ReclaimGrace: time.Second, + CleanupInterval: time.Second, + } + if errValidate := ValidateCredentialConcurrencyLifecycle(time.Second, cfg); errValidate == nil { + t.Fatal("ValidateCredentialConcurrencyLifecycle() error = nil, want overflow rejection") + } +} diff --git a/internal/config/credential_in_flight.go b/internal/config/credential_in_flight.go new file mode 100644 --- /dev/null +++ b/internal/config/credential_in_flight.go @@ -0,0 +1,87 @@ +package config + +import ( + "fmt" + "time" +) + +const ( + DefaultInFlightMaxPartBytes = 256 * 1024 + DefaultInFlightMaxPartCount = 64 + DefaultInFlightMaxRevisionBytes = 16 * 1024 * 1024 + DefaultInFlightMaxAggregateGroups = 100000 + DefaultInFlightMaxDetails = 10000 + DefaultInFlightMaxStringBytes = 256 +) + +// CredentialInFlightConfig controls in-flight credential observation snapshots. +type CredentialInFlightConfig struct { + SnapshotInterval string `yaml:"snapshot-interval" json:"snapshot-interval"` + StaleAfter string `yaml:"stale-after" json:"stale-after"` + MaxPartBytes int `yaml:"max-part-bytes" json:"max-part-bytes"` + MaxPartCount int `yaml:"max-part-count" json:"max-part-count"` + MaxRevisionBytes int `yaml:"max-revision-bytes" json:"max-revision-bytes"` + MaxAggregateGroups int `yaml:"max-aggregate-groups" json:"max-aggregate-groups"` + MaxDetails int `yaml:"max-details" json:"max-details"` + MaxStringBytes int `yaml:"max-string-bytes" json:"max-string-bytes"` + StagingRetention string `yaml:"staging-retention" json:"staging-retention"` +} + +// DefaultCredentialInFlightConfig returns the in-flight observation defaults. +func DefaultCredentialInFlightConfig() CredentialInFlightConfig { + return CredentialInFlightConfig{ + SnapshotInterval: "2s", + StaleAfter: "10s", + MaxPartBytes: DefaultInFlightMaxPartBytes, + MaxPartCount: DefaultInFlightMaxPartCount, + MaxRevisionBytes: DefaultInFlightMaxRevisionBytes, + MaxAggregateGroups: DefaultInFlightMaxAggregateGroups, + MaxDetails: DefaultInFlightMaxDetails, + MaxStringBytes: DefaultInFlightMaxStringBytes, + StagingRetention: "1m", + } +} + +// Durations parses and validates the in-flight observation durations. +func (c CredentialInFlightConfig) Durations() (time.Duration, time.Duration, time.Duration, error) { + snapshotInterval, errSnapshot := time.ParseDuration(c.SnapshotInterval) + if errSnapshot != nil || snapshotInterval <= 0 { + return 0, 0, 0, fmt.Errorf("credential-in-flight.snapshot-interval must be positive") + } + staleAfter, errStale := time.ParseDuration(c.StaleAfter) + if errStale != nil || staleAfter <= 0 || snapshotInterval > staleAfter/3 { + return 0, 0, 0, fmt.Errorf("credential-in-flight.stale-after must be at least three snapshot intervals") + } + stagingRetention, errRetention := time.ParseDuration(c.StagingRetention) + if errRetention != nil || stagingRetention <= 0 { + return 0, 0, 0, fmt.Errorf("credential-in-flight.staging-retention must be positive") + } + return snapshotInterval, staleAfter, stagingRetention, nil +} + +// Validate verifies the in-flight observation bounds. +func (c CredentialInFlightConfig) Validate() error { + if _, _, _, errDurations := c.Durations(); errDurations != nil { + return errDurations + } + if c.MaxPartBytes < 1024 || c.MaxPartCount <= 0 || c.MaxPartCount > DefaultInFlightMaxPartCount { + return fmt.Errorf("credential-in-flight part bounds are invalid") + } + if c.MaxRevisionBytes < c.MaxPartBytes || c.MaxRevisionBytes > DefaultInFlightMaxRevisionBytes { + return fmt.Errorf("credential-in-flight.max-revision-bytes is outside hard bounds") + } + requiredParts := (c.MaxRevisionBytes + c.MaxPartBytes - 1) / c.MaxPartBytes + if requiredParts > c.MaxPartCount { + return fmt.Errorf("credential-in-flight.max-revision-bytes exceeds part capacity") + } + if c.MaxAggregateGroups <= 0 || c.MaxAggregateGroups > DefaultInFlightMaxAggregateGroups { + return fmt.Errorf("credential-in-flight.max-aggregate-groups is invalid") + } + if c.MaxDetails < 0 || c.MaxDetails > DefaultInFlightMaxDetails { + return fmt.Errorf("credential-in-flight.max-details is invalid") + } + if c.MaxStringBytes <= 0 || c.MaxStringBytes > DefaultInFlightMaxStringBytes { + return fmt.Errorf("credential-in-flight.max-string-bytes is invalid") + } + return nil +} diff --git a/internal/config/credential_in_flight_test.go b/internal/config/credential_in_flight_test.go new file mode 100644 --- /dev/null +++ b/internal/config/credential_in_flight_test.go @@ -0,0 +1,234 @@ +package config + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "math" + "os" + "path/filepath" + "reflect" + "testing" + "time" +) + +func TestLoadConfigOptionalMissingFallbackAppliesCredentialInFlightDefaults(t *testing.T) { + cfg, errLoad := LoadConfigOptional(filepath.Join(t.TempDir(), "missing.yaml"), true) + if errLoad != nil { + t.Fatalf("LoadConfigOptional() error = %v", errLoad) + } + assertOptionalConfigFallback(t, cfg) +} + +func TestLoadConfigOptionalEmptyFallbackAppliesCredentialInFlightDefaults(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.yaml") + if errWrite := os.WriteFile(configPath, nil, 0o600); errWrite != nil { + t.Fatal(errWrite) + } + cfg, errLoad := LoadConfigOptional(configPath, true) + if errLoad != nil { + t.Fatalf("LoadConfigOptional() error = %v", errLoad) + } + assertOptionalConfigFallback(t, cfg) +} + +func TestLoadConfigOptionalWhitespaceFallbackAppliesCredentialInFlightDefaults(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.yaml") + if errWrite := os.WriteFile(configPath, []byte(" \t\n\r "), 0o600); errWrite != nil { + t.Fatal(errWrite) + } + cfg, errLoad := LoadConfigOptional(configPath, true) + if errLoad != nil { + t.Fatalf("LoadConfigOptional() error = %v", errLoad) + } + assertOptionalConfigFallback(t, cfg) +} + +func TestLoadConfigOptionalInvalidFallbackAppliesCredentialInFlightDefaults(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.yaml") + if errWrite := os.WriteFile(configPath, []byte(":"), 0o600); errWrite != nil { + t.Fatal(errWrite) + } + cfg, errLoad := LoadConfigOptional(configPath, true) + if errLoad != nil { + t.Fatalf("LoadConfigOptional() error = %v", errLoad) + } + assertOptionalConfigFallback(t, cfg) +} + +func assertOptionalConfigFallback(t *testing.T, cfg *Config) { + t.Helper() + if cfg.CredentialInFlight != DefaultCredentialInFlightConfig() { + t.Fatalf("CredentialInFlight = %#v, want %#v", cfg.CredentialInFlight, DefaultCredentialInFlightConfig()) + } + if errValidate := cfg.CredentialInFlight.Validate(); errValidate != nil { + t.Fatalf("CredentialInFlight.Validate() error = %v", errValidate) + } + if cfg.ErrorLogsMaxFiles != 0 || cfg.WebsocketAuth || cfg.CredentialConcurrency != (CredentialConcurrencyConfig{}) { + t.Fatalf("fallback config changed existing empty-config defaults: %#v", cfg) + } +} + +func TestCredentialInFlightConfigContractFixture(t *testing.T) { + raw, errRead := os.ReadFile(filepath.Join("..", "home", "testdata", "credential_in_flight_contract.json")) + if errRead != nil { + t.Fatal(errRead) + } + fixture, errDecode := decodeCredentialInFlightConfigFixture(raw) + if errDecode != nil { + t.Fatal(errDecode) + } + if fixture.Config != DefaultCredentialInFlightConfig() { + t.Fatalf("default config = %#v, want %#v", DefaultCredentialInFlightConfig(), fixture.Config) + } + if errValidate := fixture.Config.Validate(); errValidate != nil { + t.Fatalf("Validate() error = %v", errValidate) + } + assertCredentialInFlightConfigFields(t) + assertRequiredJSONKeys(t, raw, []string{"config", "part", "overflow"}) + assertRequiredJSONKeys(t, fixture.ConfigJSON, []string{"snapshot-interval", "stale-after", "max-part-bytes", "max-part-count", "max-revision-bytes", "max-aggregate-groups", "max-details", "max-string-bytes", "staging-retention"}) +} + +func TestCredentialInFlightConfigFixtureRejectsInvalidJSON(t *testing.T) { + raw, errRead := os.ReadFile(filepath.Join("..", "home", "testdata", "credential_in_flight_contract.json")) + if errRead != nil { + t.Fatal(errRead) + } + for _, test := range []struct { + name string + raw []byte + }{ + {name: "unknown config field", raw: bytes.Replace(raw, []byte(`"snapshot-interval": "2s"`), []byte(`"snapshot-interval": "2s", "secret": "secret"`), 1)}, + {name: "trailing JSON", raw: append(append([]byte{}, raw...), []byte(` {"config": {}}`)...)}, + } { + t.Run(test.name, func(t *testing.T) { + if _, errDecode := decodeCredentialInFlightConfigFixture(test.raw); errDecode == nil { + t.Fatal("decodeCredentialInFlightConfigFixture() error = nil") + } + }) + } +} + +func TestCredentialInFlightConfigDurationBounds(t *testing.T) { + for _, test := range []struct { + name string + stale string + every string + valid bool + }{ + {name: "exact three intervals", every: "1s", stale: "3s", valid: true}, + {name: "below three intervals", every: "1s", stale: "2999999999ns", valid: false}, + {name: "near duration maximum", every: time.Duration(math.MaxInt64 / 2).String(), stale: time.Duration(math.MaxInt64).String(), valid: false}, + } { + t.Run(test.name, func(t *testing.T) { + cfg := DefaultCredentialInFlightConfig() + cfg.SnapshotInterval = test.every + cfg.StaleAfter = test.stale + errValidate := cfg.Validate() + if (errValidate == nil) != test.valid { + t.Fatalf("Validate() error = %v, want valid = %t", errValidate, test.valid) + } + }) + } +} + +func TestCredentialInFlightConfigRejectsUnsafeBounds(t *testing.T) { + cfg := DefaultCredentialInFlightConfig() + cfg.StaleAfter = "5s" + if errValidate := cfg.Validate(); errValidate == nil { + t.Fatal("Validate() error = nil, want stale-after error") + } + cfg = DefaultCredentialInFlightConfig() + cfg.MaxRevisionBytes = 16*1024*1024 + 1 + if errValidate := cfg.Validate(); errValidate == nil { + t.Fatal("Validate() error = nil, want hard revision bound error") + } + cfg = DefaultCredentialInFlightConfig() + cfg.MaxPartBytes = math.MaxInt + if errValidate := cfg.Validate(); errValidate == nil { + t.Fatal("Validate() error = nil, want overflow-safe part bound error") + } +} + +type credentialInFlightConfigFixture struct { + Config CredentialInFlightConfig `json:"config"` + ConfigJSON json.RawMessage `json:"-"` +} + +func decodeCredentialInFlightConfigFixture(raw []byte) (credentialInFlightConfigFixture, error) { + var fixture credentialInFlightConfigFixture + var document struct { + Config json.RawMessage `json:"config"` + Part json.RawMessage `json:"part"` + Overflow json.RawMessage `json:"overflow"` + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if errDecode := decoder.Decode(&document); errDecode != nil { + return fixture, errDecode + } + if errDecode := decoder.Decode(&struct{}{}); errDecode == nil { + return fixture, errors.New("unexpected trailing JSON") + } else if errDecode != io.EOF { + return fixture, errDecode + } + decoder = json.NewDecoder(bytes.NewReader(document.Config)) + decoder.DisallowUnknownFields() + if errDecode := decoder.Decode(&fixture.Config); errDecode != nil { + return fixture, errDecode + } + if errDecode := decoder.Decode(&struct{}{}); errDecode == nil { + return fixture, errors.New("unexpected trailing config JSON") + } else if errDecode != io.EOF { + return fixture, errDecode + } + fixture.ConfigJSON = document.Config + return fixture, nil +} + +func assertCredentialInFlightConfigFields(t *testing.T) { + t.Helper() + assertOrderedJSONFields(t, reflect.TypeOf(CredentialInFlightConfig{}), []jsonField{ + {name: "SnapshotInterval", tag: "snapshot-interval"}, + {name: "StaleAfter", tag: "stale-after"}, + {name: "MaxPartBytes", tag: "max-part-bytes"}, + {name: "MaxPartCount", tag: "max-part-count"}, + {name: "MaxRevisionBytes", tag: "max-revision-bytes"}, + {name: "MaxAggregateGroups", tag: "max-aggregate-groups"}, + {name: "MaxDetails", tag: "max-details"}, + {name: "MaxStringBytes", tag: "max-string-bytes"}, + {name: "StagingRetention", tag: "staging-retention"}, + }) +} + +type jsonField struct { + name string + tag string +} + +func assertOrderedJSONFields(t *testing.T, structType reflect.Type, want []jsonField) { + t.Helper() + if structType.NumField() != len(want) { + t.Fatalf("%s field count = %d, want %d", structType.Name(), structType.NumField(), len(want)) + } + for index, expected := range want { + field := structType.Field(index) + if field.Name != expected.name || field.Tag.Get("json") != expected.tag { + t.Fatalf("%s field %d = (%q, %q), want (%q, %q)", structType.Name(), index, field.Name, field.Tag.Get("json"), expected.name, expected.tag) + } + } +} + +func assertRequiredJSONKeys(t *testing.T, raw json.RawMessage, required []string) { + t.Helper() + var fields map[string]json.RawMessage + if errDecode := json.Unmarshal(raw, &fields); errDecode != nil { + t.Fatalf("json.Unmarshal() error = %v", errDecode) + } + for _, key := range required { + if _, ok := fields[key]; !ok { + t.Fatalf("required JSON key %q is missing", key) + } + } +} diff --git a/internal/config/parse.go b/internal/config/parse.go --- a/internal/config/parse.go +++ b/internal/config/parse.go @@ -32,9 +32,15 @@ cfg.Pprof.Enable = false cfg.Pprof.Addr = DefaultPprofAddr cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository + cfg.CredentialInFlight = DefaultCredentialInFlightConfig() if err := yaml.Unmarshal(data, &cfg); err != nil { return nil, fmt.Errorf("parse config payload: %w", err) + } + + cfg.CredentialConcurrency = cfg.CredentialConcurrency.WithDefaults() + if errValidate := cfg.CredentialInFlight.Validate(); errValidate != nil { + return nil, errValidate } // Hash remote management key if plaintext is detected (nested), but do NOT persist. diff --git a/internal/home/client.go b/internal/home/client.go --- a/internal/home/client.go +++ b/internal/home/client.go @@ -26,24 +26,61 @@ ) const ( - redisKeyConfig = "config" - redisChannelConfig = "config" - redisKeyUsage = "usage" - redisKeyRequestLog = "request-log" - redisKeyAppLog = "app-log" - redisKeyPluginStatus = "plugin-status" - redisKeyPluginTasks = "plugin-tasks" - redisKeyPluginSync = "plugin-sync" + redisKeyConfig = "config" + redisChannelConfig = "config" + redisKeyUsage = "usage" + redisKeyInFlightSnapshot = "in-flight-snapshot" + redisKeyConcurrencyRelease = "concurrency-release" + redisKeyRequestLog = "request-log" + redisKeyAppLog = "app-log" + redisKeyPluginStatus = "plugin-status" + redisKeyPluginTasks = "plugin-tasks" + redisKeyPluginSync = "plugin-sync" - homeReconnectInterval = time.Second - homeReconnectFailoverThreshold = 3 - homeRedisOperationTimeout = 3 * time.Second - homePluginSyncOperationTimeout = 2 * time.Minute - homeSubscriptionReceiveTimeout = 3 * time.Second - redisChannelCluster = "cluster" + homeReconnectInterval = time.Second + homeReconnectFailoverThreshold = 3 + homeRedisOperationTimeout = 3 * time.Second + homePluginSyncOperationTimeout = 2 * time.Minute + homeSubscriptionReceiveTimeout = 3 * time.Second + credentialConcurrencyNodeHeartbeatTimeout = 20 * time.Second + redisChannelCluster = "cluster" ) const pluginSyncUnsupportedErrorType = "plugin_sync_unsupported" + +// DispatchError classifies whether Home may have processed an auth dispatch request. +type DispatchError struct { + Err error + Ambiguous bool +} + +func (e *DispatchError) Error() string { + if e == nil || e.Err == nil { + return "home auth dispatch failed" + } + return e.Err.Error() +} + +func (e *DispatchError) Unwrap() error { + if e == nil { + return nil + } + return e.Err +} + +// NewAmbiguousDispatchError marks a post-send transport failure as requiring a client abort. +func NewAmbiguousDispatchError(err error) error { + if err == nil { + return nil + } + return &DispatchError{Err: err, Ambiguous: true} +} + +// IsAmbiguousDispatchError reports whether Home may have processed the dispatch request. +func IsAmbiguousDispatchError(err error) bool { + var dispatchErr *DispatchError + return errors.As(err, &dispatchErr) && dispatchErr.Ambiguous +} var ( ErrDisabled = errors.New("home client disabled") @@ -53,6 +90,7 @@ ErrConfigNotFound = errors.New("home config not found") ErrModelsNotFound = errors.New("home models not found") ErrPluginSyncUnsupported = errors.New("home plugin sync is unsupported") + ErrDispatchFenced = errors.New("home auth dispatch is fenced") ) type clusterNode struct { @@ -85,6 +123,10 @@ XX bool } +type subscriptionCloser interface { + Close() error +} + type Client struct { mu sync.Mutex @@ -92,11 +134,17 @@ seedHost string seedPort int - cmd *redis.Client - cmdOptions *redis.Options - sub *redis.Client + cmd *redis.Client + cmdOptions *redis.Options + sub *redis.Client + release *redis.Client + connections map[*homeDispatchConn]struct{} + lifecycle config.CredentialConcurrencyConfig + limiter atomic.Pointer[config.CredentialConcurrencyConfig] + managed bool heartbeatOK atomic.Bool + dispatchFenced atomic.Bool clusterNodes []clusterNode reconnectFailures int } @@ -128,26 +176,123 @@ return c.heartbeatOK.Load() } +// Close permanently ends this client's dispatch lifetime. func (c *Client) Close() { + if c == nil { + return + } + c.dispatchFenced.Store(true) + c.heartbeatOK.Store(false) + c.mu.Lock() + commandClient, subscriptionClient, connections := c.detachClientsLocked() + releaseClient := c.release + c.release = nil + c.mu.Unlock() + closeDetachedClients(commandClient, subscriptionClient, connections) + if releaseClient != nil { + _ = releaseClient.Close() + } +} + +// closeBootstrapPools replaces private bootstrap pools without ending the client lifetime. +func (c *Client) closeBootstrapPools() { if c == nil { return } c.heartbeatOK.Store(false) c.mu.Lock() - defer c.mu.Unlock() - c.closeClientsLocked() + commandClient, subscriptionClient, connections := c.detachClientsLocked() + c.mu.Unlock() + closeDetachedClients(commandClient, subscriptionClient, connections) } -func (c *Client) closeClientsLocked() { - if c.cmd != nil { - _ = c.cmd.Close() +// AbortAmbiguousDispatch fences this client after an auth dispatch response is ambiguous. +func (c *Client) AbortAmbiguousDispatch() { + if c == nil { + return } - if c.sub != nil { - _ = c.sub.Close() + c.dispatchFenced.Store(true) + c.heartbeatOK.Store(false) + c.mu.Lock() + commandClient, subscriptionClient, connections := c.detachClientsLocked() + releaseClient := c.release + c.release = nil + c.mu.Unlock() + for _, conn := range connections { + _ = conn.Close() } + if commandClient != nil { + go func() { + _ = commandClient.Close() + }() + } + if subscriptionClient != nil { + go func() { + _ = subscriptionClient.Close() + }() + } + if releaseClient != nil { + go func() { + _ = releaseClient.Close() + }() + } +} + +func (c *Client) detachClientsLocked() (*redis.Client, *redis.Client, []*homeDispatchConn) { + connections := make([]*homeDispatchConn, 0, len(c.connections)) + for conn := range c.connections { + connections = append(connections, conn) + } + commandClient := c.cmd + subscriptionClient := c.sub c.cmd = nil c.cmdOptions = nil c.sub = nil + c.connections = nil + return commandClient, subscriptionClient, connections +} + +func closeDetachedClients(commandClient *redis.Client, subscriptionClient *redis.Client, connections []*homeDispatchConn) { + for _, conn := range connections { + _ = conn.Close() + } + if commandClient != nil { + _ = commandClient.Close() + } + if subscriptionClient != nil { + _ = subscriptionClient.Close() + } +} + +func (c *Client) closeClientsLocked() { + commandClient, subscriptionClient, connections := c.detachClientsLocked() + releaseClient := c.release + c.release = nil + go func() { + closeDetachedClients(commandClient, subscriptionClient, connections) + if releaseClient != nil { + _ = releaseClient.Close() + } + }() +} + +// SetManagedLifetime defers client shutdown to the Service lifetime owner. +func (c *Client) SetManagedLifetime(managed bool) { + if c == nil { + return + } + c.mu.Lock() + c.managed = managed + c.mu.Unlock() +} + +func (c *Client) managedLifetime() bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + return c.managed } func (c *Client) addr() (string, bool) { @@ -174,11 +319,17 @@ if c == nil { return ErrDisabled } + if c.dispatchFenced.Load() { + return ErrDispatchFenced + } if !c.Enabled() { return ErrDisabled } c.mu.Lock() defer c.mu.Unlock() + if c.dispatchFenced.Load() { + return ErrDispatchFenced + } addr, ok := c.addrLocked() if !ok { @@ -208,7 +359,7 @@ if errTLS != nil { return nil, errTLS } - return &redis.Options{ + options := &redis.Options{ Addr: addr, TLSConfig: tlsConfig, DialTimeout: homeRedisOperationTimeout, @@ -217,7 +368,54 @@ MaxRetries: -1, DialerRetries: 1, ContextTimeoutEnabled: true, - }, nil + } + options.Dialer = c.trackedRedisDialer(redis.NewDialer(options)) + return options, nil +} + +type homeDispatchConn struct { + net.Conn + client *Client + once sync.Once +} + +func (c *Client) trackedRedisDialer(dialer func(context.Context, string, string) (net.Conn, error)) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, network string, address string) (net.Conn, error) { + conn, errDial := dialer(ctx, network, address) + if errDial != nil { + return nil, errDial + } + wrapped := &homeDispatchConn{Conn: conn, client: c} + if c == nil { + return wrapped, nil + } + c.mu.Lock() + if c.dispatchFenced.Load() { + c.mu.Unlock() + _ = wrapped.Close() + return nil, ErrDispatchFenced + } + if c.connections == nil { + c.connections = make(map[*homeDispatchConn]struct{}) + } + c.connections[wrapped] = struct{}{} + c.mu.Unlock() + return wrapped, nil + } +} + +func (c *homeDispatchConn) Close() error { + if c == nil || c.Conn == nil { + return net.ErrClosed + } + c.once.Do(func() { + if c.client != nil { + c.client.mu.Lock() + delete(c.client.connections, c) + c.client.mu.Unlock() + } + }) + return c.Conn.Close() } func cloneRedisOptions(options *redis.Options) *redis.Options { @@ -310,16 +508,21 @@ } func (c *Client) commandClient() (*redis.Client, error) { + if c == nil || c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } if errEnsure := c.ensureClients(); errEnsure != nil { return nil, errEnsure } c.mu.Lock() - cmd := c.cmd - c.mu.Unlock() - if cmd == nil { + defer c.mu.Unlock() + if c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } + if c.cmd == nil { return nil, ErrNotConnected } - return cmd, nil + return c.cmd, nil } func (c *Client) pluginSyncCommandOptions() (*redis.Options, error) { @@ -869,40 +1072,78 @@ count = 1 } return authDispatchRequest{ - Type: "auth", - Model: requestedModel, - Count: count, - SessionID: strings.TrimSpace(sessionID), - Headers: headersToLowerMap(headers), + Type: "auth", + Model: requestedModel, + Count: count, + ConcurrencyProtocol: 1, + SessionID: strings.TrimSpace(sessionID), + Headers: headersToLowerMap(headers), } } func (c *Client) RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error) { - cmd, errClient := c.commandClient() - if errClient != nil { - return nil, errClient + if c == nil || c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return nil, errContext } requestedModel = strings.TrimSpace(requestedModel) if requestedModel == "" { return nil, fmt.Errorf("home: requested model is empty") } req := newAuthDispatchRequest(requestedModel, sessionID, headers, count) - keyBytes, err := json.Marshal(&req) - if err != nil { - return nil, err + keyBytes, errMarshal := json.Marshal(&req) + if errMarshal != nil { + return nil, errMarshal } - - raw, err := cmd.RPop(ctx, string(keyBytes)).Bytes() - if errors.Is(err, redis.Nil) { + if c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } + cmd, errClient := c.commandClient() + if errClient != nil { + return nil, errClient + } + if c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } + conn := cmd.Conn() + defer func() { + if errClose := conn.Close(); errClose != nil { + log.WithError(errClose).Debug("Home auth dispatch connection close failed") + } + }() + if errProbe := conn.Ping(ctx).Err(); errProbe != nil { + return nil, errProbe + } + if c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } + raw, errRPop := conn.RPop(ctx, string(keyBytes)).Bytes() + if errors.Is(errRPop, redis.Nil) { return nil, ErrAuthNotFound } - if err != nil { - return nil, err + if errRPop != nil { + if isAmbiguousIssuedRPopAuthError(errRPop) { + return nil, NewAmbiguousDispatchError(errRPop) + } + return nil, errRPop } if len(raw) == 0 { return nil, ErrEmptyResponse } return raw, nil +} + +func isAmbiguousIssuedRPopAuthError(err error) bool { + if err == nil || errors.Is(err, redis.Nil) { + return false + } + var redisErr redis.Error + return !errors.As(err, &redisErr) } func (c *Client) GetRefreshAuth(ctx context.Context, authIndex string) ([]byte, error) { @@ -945,6 +1186,60 @@ return nil } return cmd.LPush(ctx, redisKeyUsage, payload).Err() +} + +// LPushInFlightSnapshot publishes a bounded in-flight observation frame. +func (c *Client) LPushInFlightSnapshot(ctx context.Context, payload []byte) error { + cmd, errClient := c.commandClient() + if errClient != nil { + return errClient + } + return cmd.LPush(ctx, redisKeyInFlightSnapshot, payload).Err() +} + +// PushConcurrencyRelease sends one cumulative concurrency release frame through an independent client. +func (c *Client) PushConcurrencyRelease(ctx context.Context, frame ConcurrencyReleaseFrame) error { + if frame.CredentialID == "" || frame.Model == "" || frame.ReleaseSeq <= 0 { + return fmt.Errorf("invalid concurrency release frame") + } + cmd, errClient := c.concurrencyReleaseClient() + if errClient != nil { + return errClient + } + payload, errMarshal := json.Marshal(frame) + if errMarshal != nil { + return fmt.Errorf("marshal concurrency release frame: %w", errMarshal) + } + return cmd.Do(ctx, "LPUSH", redisKeyConcurrencyRelease, payload).Err() +} + +func (c *Client) concurrencyReleaseClient() (*redis.Client, error) { + if c == nil || c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } + if !c.Enabled() { + return nil, ErrDisabled + } + + c.mu.Lock() + defer c.mu.Unlock() + if c.dispatchFenced.Load() { + return nil, ErrDispatchFenced + } + if c.release != nil { + return c.release, nil + } + addr, ok := c.addrLocked() + if !ok { + return nil, fmt.Errorf("home: invalid address (host=%q port=%d)", c.homeCfg.Host, c.homeCfg.Port) + } + options, errOptions := c.redisOptionsLocked(addr) + if errOptions != nil { + return nil, errOptions + } + options.Dialer = redis.NewDialer(options) + c.release = redis.NewClient(options) + return c.release, nil } func (c *Client) RPushRequestLog(ctx context.Context, payload []byte) error { @@ -1177,6 +1472,68 @@ } } +func (c *Client) SetLifecycleConfig(cfg config.CredentialConcurrencyConfig) error { + if c == nil { + return ErrDisabled + } + cfg = cfg.WithDefaults() + if errValidate := config.ValidateCredentialConcurrency(cfg); errValidate != nil { + return fmt.Errorf("validate credential concurrency lifecycle config: %w", errValidate) + } + c.mu.Lock() + c.lifecycle = cfg + c.mu.Unlock() + c.limiter.Store(&cfg) + return nil +} + +// LimiterConfig returns the latest immutable, validated Home limiter configuration. +func (c *Client) LimiterConfig() config.CredentialConcurrencyConfig { + if c == nil { + return config.CredentialConcurrencyConfig{}.WithDefaults() + } + if cfg := c.limiter.Load(); cfg != nil { + return *cfg + } + return config.CredentialConcurrencyConfig{}.WithDefaults() +} + +func (c *Client) subscriptionParameters() ([]string, time.Duration) { + if c == nil { + return []string{redisChannelConfig}, config.CredentialConcurrencyConfig{}.WithDefaults().CPAHeartbeatTimeout + } + c.mu.Lock() + cfg := c.lifecycle.WithDefaults() + c.mu.Unlock() + + args := []string{redisChannelConfig} + if cfg.LifecycleConfigRevision > 0 { + args = append(args, strconv.FormatInt(cfg.LifecycleConfigRevision, 10)) + } + return args, cfg.CPAHeartbeatTimeout +} + +func (c *Client) rebuildCommandPoolAndProbe(ctx context.Context) error { + c.promoteSubscription() + return c.Ping(ctx) +} + +func (c *Client) promoteSubscription() { + if c == nil { + return + } + c.mu.Lock() + commandClient := c.cmd + c.cmd = nil + c.cmdOptions = nil + c.mu.Unlock() + if commandClient != nil { + if errClose := commandClient.Close(); errClose != nil { + log.WithError(errClose).Warn("Home bootstrap command client close failed") + } + } +} + func (c *Client) handleSubscriptionPayload(ctx context.Context, channel string, payload string, onConfig func([]byte) error) error { payload = strings.TrimSpace(payload) if payload == "" { @@ -1196,121 +1553,124 @@ } } -// StartConfigSubscriber connects to home, fetches config once via GET config, then subscribes to -// the "config" channel to receive runtime config updates. -// -// The subscription connection is treated as the home heartbeat. HeartbeatOK is set to true only -// after the initial GET config succeeds and the SUBSCRIBE connection is established. When the -// subscription ends unexpectedly, HeartbeatOK becomes false and the loop reconnects. -func (c *Client) StartConfigSubscriber(ctx context.Context, onConfig func([]byte) error) { - if c == nil { - return - } - if !c.Enabled() { - return +// RunConfigSubscriberLifetime runs one GET, SUBSCRIBE, and receive lifetime. +// Reconnection is owned by the service so each replacement can install a new client lifetime. +func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func([]byte) error, onReady func()) error { + if c == nil || !c.Enabled() { + return ErrDisabled } if onConfig == nil { - return + return fmt.Errorf("home config subscriber callback is nil") + } + if ctx == nil { + ctx = context.Background() + } + + c.closeBootstrapPools() + if errEnsure := c.ensureClients(); errEnsure != nil { + return c.endConfigSubscriberLifetime(errEnsure) + } + + raw, errGet := c.GetConfig(ctx) + if errGet != nil { + return c.endConfigSubscriberLifetime(errGet) + } + if errApply := onConfig(raw); errApply != nil { + return c.endConfigSubscriberLifetime(errApply) + } + + sub, errSubClient := c.subscriptionClient() + if errSubClient != nil { + return c.endConfigSubscriberLifetime(errSubClient) + } + args, receiveTimeout := c.subscriptionParameters() + pubsub := sub.Subscribe(ctx, args...) + if pubsub == nil { + return c.endConfigSubscriberLifetime(ErrNotConnected) + } + + if errACK := receiveSubscriptionACKs(ctx, pubsub, receiveTimeout, args[:1]); errACK != nil { + return c.endConfigSubscriberLifetimeWithSubscription(errACK, pubsub, "failed ACK") + } + + if errProbe := c.rebuildCommandPoolAndProbe(ctx); errProbe != nil { + return c.endConfigSubscriberLifetimeWithSubscription(errProbe, pubsub, "fresh command probe failure") + } + c.heartbeatOK.Store(true) + if onReady != nil { + onReady() } for { - if ctx != nil { - select { - case <-ctx.Done(): - c.heartbeatOK.Store(false) - return - default: - } + _, receiveTimeout = c.subscriptionParameters() + event, errReceive := pubsub.ReceiveTimeout(ctx, receiveTimeout) + if errReceive != nil { + return c.endConfigSubscriberLifetimeWithSubscription(errReceive, pubsub, "heartbeat loss") } - - c.heartbeatOK.Store(false) - c.Close() - - if errEnsure := c.ensureClients(); errEnsure != nil { - log.Warn("unable to connect to home control center, retrying in 1 second") - c.markReconnectFailure("connect") - sleepWithContext(ctx, homeReconnectInterval) - continue - } - - if errPing := c.Ping(ctx); errPing != nil { - log.Warn("unable to connect to home control center, retrying in 1 second") - c.markReconnectFailure("ping") - sleepWithContext(ctx, homeReconnectInterval) - continue - } - - raw, errGet := c.GetConfig(ctx) - if errGet != nil { - log.Warn("unable to fetch config from home control center, retrying in 1 second") - c.markReconnectFailure("config fetch") - sleepWithContext(ctx, homeReconnectInterval) - continue - } - if errApply := onConfig(raw); errApply != nil { - log.Warn("unable to apply config from home control center, retrying in 1 second") - sleepWithContext(ctx, homeReconnectInterval) - continue - } - - sub, errSubClient := c.subscriptionClient() - if errSubClient != nil { - c.markReconnectFailure("subscribe client") - sleepWithContext(ctx, homeReconnectInterval) - continue - } - - pubsub := sub.Subscribe(ctx, redisChannelConfig) - if pubsub == nil { - c.markReconnectFailure("subscribe") - sleepWithContext(ctx, homeReconnectInterval) - continue - } - - // Ensure the subscription is established before marking heartbeat OK. - if _, errReceive := pubsub.ReceiveTimeout(ctx, homeSubscriptionReceiveTimeout); errReceive != nil { - _ = pubsub.Close() - c.markReconnectFailure("subscribe") - sleepWithContext(ctx, homeReconnectInterval) - continue - } - - c.resetReconnectFailures() - c.heartbeatOK.Store(true) - - for { - event, errMsg := pubsub.ReceiveTimeout(ctx, homeSubscriptionReceiveTimeout) - if errMsg != nil { - _ = pubsub.Close() - c.heartbeatOK.Store(false) - if isTimeoutError(errMsg) { - c.markSubscriptionTimeout() - } else { - c.markReconnectFailure("subscription") - } - sleepWithContext(ctx, homeReconnectInterval) - break - } - 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") - } - } - case *redis.Pong: - c.resetReconnectFailures() - case *redis.Subscription: + switch msg := event.(type) { + case *redis.Message: + if msg == nil { continue - default: - log.Debugf("home subscription returned unsupported message type %T", event) } + 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") + } + } + case *redis.Pong: + c.resetReconnectFailures() + case *redis.Subscription: + continue + default: + log.Debugf("home subscription returned unsupported message type %T", event) } + } +} + +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") + } + for index, channel := range channels { + event, errReceive := pubsub.ReceiveTimeout(ctx, receiveTimeout) + if errReceive != nil { + return errReceive + } + ack, ok := event.(*redis.Subscription) + if !ok || ack == nil || ack.Kind != "subscribe" || ack.Channel != channel || ack.Count != index+1 { + return fmt.Errorf("invalid Home subscription ACK") + } + } + return nil +} + +func (c *Client) endConfigSubscriberLifetime(err error) error { + c.heartbeatOK.Store(false) + if !c.managedLifetime() { + c.Close() + } + return err +} + +func (c *Client) endConfigSubscriberLifetimeWithSubscription(err error, subscription subscriptionCloser, reason string) error { + c.heartbeatOK.Store(false) + if subscription != nil { + if errClose := subscription.Close(); errClose != nil { + log.WithError(errClose).Debugf("Home subscription close after %s", reason) + } + } + if !c.managedLifetime() { + c.Close() + } + return err +} + +// 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") } } diff --git a/internal/home/client_test.go b/internal/home/client_test.go --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -4,6 +4,7 @@ "bufio" "context" "crypto/tls" + "crypto/x509" "encoding/json" "errors" "fmt" @@ -37,6 +38,9 @@ } if got := int(payload["count"].(float64)); got != 2 { t.Fatalf("count = %d, want 2", got) + } + if got := int(payload["concurrency_protocol"].(float64)); got != 1 { + t.Fatalf("concurrency_protocol = %d, want 1", got) } } @@ -194,6 +198,66 @@ } if _, errConflict := buildKVSetArgs("key", []byte("value"), KVSetOptions{NX: true, XX: true}); errConflict == nil { t.Fatalf("buildKVSetArgs(NX XX) error = nil, want error") + } +} + +func TestClientLPushInFlightSnapshotUsesDedicatedKeyWithoutChangingHeartbeat(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "LPUSH") { + return ":1\r\n" + } + return "-ERR unexpected command\r\n" + }) + client.heartbeatOK.Store(true) + + if errPush := client.LPushInFlightSnapshot(context.Background(), []byte(`{"revision":1}`)); errPush != nil { + t.Fatalf("LPushInFlightSnapshot() error = %v", errPush) + } + if !client.HeartbeatOK() { + t.Fatal("LPushInFlightSnapshot() changed heartbeat state") + } + last := commands.Last() + if len(last) != 3 || !strings.EqualFold(last[0], "LPUSH") || last[1] != redisKeyInFlightSnapshot || last[2] != `{"revision":1}` { + t.Fatalf("LPushInFlightSnapshot() command = %#v", last) + } +} + +func TestClientPushConcurrencyReleaseUsesIndependentClient(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "LPUSH") { + return ":1\r\n" + } + return "-ERR unexpected command\r\n" + }) + commandClient := client.cmd + + frame := concurrencyReleaseFrameFromFixture(t) + if errPush := client.PushConcurrencyRelease(context.Background(), frame); errPush != nil { + t.Fatalf("PushConcurrencyRelease() error = %v", errPush) + } + if client.release == nil || client.release == commandClient { + t.Fatal("PushConcurrencyRelease() did not create an independent client") + } + last := commands.Last() + if want := []string{"LPUSH", redisKeyConcurrencyRelease, `{"credential_id":"cred-1","model":"gpt","release_seq":1}`}; !reflect.DeepEqual(last, want) { + t.Fatalf("PushConcurrencyRelease() command = %#v, want %#v", last, want) + } +} + +func TestClientLPushInFlightSnapshotErrorKeepsHeartbeat(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + if len(args) > 0 && strings.EqualFold(args[0], "LPUSH") { + return "-ERR unavailable\r\n" + } + return "-ERR unexpected command\r\n" + }) + client.heartbeatOK.Store(true) + + if errPush := client.LPushInFlightSnapshot(context.Background(), []byte(`{"revision":1}`)); errPush == nil { + t.Fatal("LPushInFlightSnapshot() error = nil") + } + if !client.HeartbeatOK() { + t.Fatal("LPushInFlightSnapshot() changed heartbeat state after an error") } } @@ -666,12 +730,34 @@ return append([]string(nil), l.commands[len(l.commands)-1]...) } +func (l *redisCommandLog) All() [][]string { + l.mu.Lock() + defer l.mu.Unlock() + out := make([][]string, len(l.commands)) + for index := range l.commands { + out[index] = append([]string(nil), l.commands[index]...) + } + return out +} + func (l *redisCommandLog) CountKey(key string) int { l.mu.Lock() defer l.mu.Unlock() count := 0 for _, command := range l.commands { if len(command) >= 2 && command[1] == key { + count++ + } + } + return count +} + +func (l *redisCommandLog) CountCommandKey(commandName string, key string) int { + l.mu.Lock() + defer l.mu.Unlock() + count := 0 + for _, command := range l.commands { + if len(command) >= 2 && strings.EqualFold(command[0], commandName) && command[1] == key { count++ } } @@ -734,6 +820,85 @@ client.Close() }) return client, log +} + +func newBlockingRPopTestClient(t *testing.T) (*Client, <-chan struct{}, chan struct{}) { + t.Helper() + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + requestRead := make(chan struct{}) + release := make(chan struct{}) + serverDone := make(chan struct{}) + var handlers sync.WaitGroup + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + handlers.Add(1) + go func(conn net.Conn) { + defer handlers.Done() + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRedisCommand(reader) + if errRead != nil { + return + } + if len(args) > 0 && strings.EqualFold(args[0], "HELLO") { + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + continue + } + if len(args) > 0 && strings.EqualFold(args[0], "RPOP") { + select { + case <-requestRead: + default: + close(requestRead) + } + <-release + return + } + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + }(conn) + } + }() + + options := &redis.Options{ + Addr: listener.Addr().String(), + Protocol: 2, + DisableIdentity: true, + DialTimeout: time.Second, + ReadTimeout: time.Second, + WriteTimeout: time.Second, + MaxRetries: -1, + ContextTimeoutEnabled: true, + } + client := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 1, DisableClusterDiscovery: true}) + options.Dialer = client.trackedRedisDialer(redis.NewDialer(options)) + client.cmdOptions = cloneRedisOptions(options) + client.cmd = redis.NewClient(options) + client.sub = redis.NewClient(cloneRedisOptions(options)) + t.Cleanup(func() { + select { + case <-release: + default: + close(release) + } + client.Close() + _ = listener.Close() + <-serverDone + handlers.Wait() + }) + return client, requestRead, release } func serveRedisCommandTestConn(conn net.Conn, log *redisCommandLog, handler func([]string) string) { @@ -863,4 +1028,757 @@ if nilMap := queryToLowerMap(nil); nilMap != nil { t.Fatalf("queryToLowerMap(nil) = %v, want nil", nilMap) } +} + +func TestClientSetLifecycleConfigAcceptsHomeAuthoritativeHeartbeat(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 6379}) + cfg := (config.CredentialConcurrencyConfig{}).WithDefaults() + cfg.CPAHeartbeatTimeout = 20 * time.Second + + if errSet := client.SetLifecycleConfig(cfg); errSet != nil { + t.Fatalf("SetLifecycleConfig() error = %v", errSet) + } + if got := client.LimiterConfig().CPAHeartbeatTimeout; got != cfg.CPAHeartbeatTimeout { + t.Fatalf("LimiterConfig().CPAHeartbeatTimeout = %s, want %s", got, cfg.CPAHeartbeatTimeout) + } +} + +func TestConfigSubscriberUsesAppliedLifecycleRevisionAndRebuildsCommands(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 6379}) + client.mu.Lock() + client.cmd = redis.NewClient(&redis.Options{Addr: "127.0.0.1:6379"}) + client.mu.Unlock() + if errSet := client.SetLifecycleConfig(config.CredentialConcurrencyConfig{ + LifecycleConfigRevision: 9, + CPAHeartbeatTimeout: 4 * time.Second, + CPACancelBound: 5 * time.Second, + }); errSet != nil { + t.Fatalf("SetLifecycleConfig() error = %v", errSet) + } + args, timeout := client.subscriptionParameters() + if !reflect.DeepEqual(args, []string{"config", "9"}) { + t.Fatalf("subscribe args = %#v", args) + } + if timeout != 4*time.Second { + t.Fatalf("receive timeout = %s", timeout) + } + client.promoteSubscription() + client.mu.Lock() + commandClient := client.cmd + client.mu.Unlock() + if commandClient != nil { + t.Fatal("bootstrap command client was retained after subscription") + } +} + +func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) { + configPayload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 20ms\n" + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + commands := &redisCommandLog{} + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go func() { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRedisCommand(reader) + if errRead != nil { + return + } + commands.Append(args) + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(configPayload), configPayload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } + }() + } + }() + t.Cleanup(func() { + _ = listener.Close() + <-serverDone + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse listener port: %v", errPort) + } + client := New(config.HomeConfig{Enabled: true, Host: host, Port: port, DisableClusterDiscovery: true}) + + ready := make(chan struct{}, 1) + errRun := client.RunConfigSubscriberLifetime(context.Background(), func(raw []byte) error { + parsed, errParse := config.ParseConfigBytes(raw) + if errParse != nil { + return errParse + } + if errSet := client.SetLifecycleConfig(parsed.CredentialConcurrency); errSet != nil { + return errSet + } + return nil + }, func() { ready <- struct{}{} }) + if errRun == nil { + t.Fatal("RunConfigSubscriberLifetime() error = nil after heartbeat loss") + } + select { + case <-ready: + default: + t.Fatalf("RunConfigSubscriberLifetime() did not invoke onReady after subscription ACK: %v; commands=%#v", errRun, commands.All()) + } + if client.HeartbeatOK() { + t.Fatal("HeartbeatOK() = true after heartbeat loss") + } + client.mu.Lock() + commandClient, subscriptionClient := client.cmd, client.sub + client.mu.Unlock() + if commandClient != nil || subscriptionClient != nil { + t.Fatalf("clients retained after heartbeat loss: command=%v subscription=%v", commandClient != nil, subscriptionClient != nil) + } + if count := commands.CountCommandKey("GET", redisKeyConfig); count != 1 { + t.Fatalf("GET config count = %d, want 1", count) + } + if count := commands.CountCommandKey("SUBSCRIBE", redisChannelConfig); count != 1 { + t.Fatalf("SUBSCRIBE config count = %d, want 1", count) + } + if got := findRedisCommand(commands.All(), "SUBSCRIBE"); !reflect.DeepEqual(got, []string{"subscribe", "config", "1"}) { + t.Fatalf("SUBSCRIBE wire command = %#v, want []string{\"subscribe\", \"config\", \"1\"}", got) + } +} + +func TestRunConfigSubscriberLifetimeRejectsInvalidSubscriptionACK(t *testing.T) { + for name, ack := range map[string]string{ + "message": "*3\r\n$7\r\nmessage\r\n$6\r\nconfig\r\n$2\r\n{}\r\n", + "wrong-channel": "*3\r\n$9\r\nsubscribe\r\n$5\r\nother\r\n:1\r\n", + "wrong-count": "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:2\r\n", + } { + t.Run(name, func(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: + return "$16\r\nhost: 127.0.0.1\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: + return ack + default: + return "+OK\r\n" + } + }) + errRun := client.RunConfigSubscriberLifetime(context.Background(), func([]byte) error { return nil }, nil) + if errRun == nil { + t.Fatal("RunConfigSubscriberLifetime() error = nil, want invalid ACK rejection") + } + if command := findRedisCommand(commands.All(), "PING"); command != nil { + t.Fatalf("PING command = %#v, want no command pool exposure before valid ACK", command) + } + }) + } +} + +func TestReceiveSubscriptionACKsForMultipleChannels(t *testing.T) { + firstACK := "*3\r\n$9\r\nsubscribe\r\n$5\r\nfirst\r\n:1\r\n" + secondACK := "*3\r\n$9\r\nsubscribe\r\n$6\r\nsecond\r\n:2\r\n" + tests := []struct { + name string + response string + wantErr bool + }{ + {name: "ordered final count", response: firstACK + secondACK}, + {name: "missing final ACK", response: firstACK, wantErr: true}, + {name: "wrong second channel", response: firstACK + "*3\r\n$9\r\nsubscribe\r\n$5\r\nother\r\n:2\r\n", wantErr: true}, + {name: "wrong second kind", response: firstACK + "*3\r\n$11\r\nunsubscribe\r\n$6\r\nsecond\r\n:2\r\n", wantErr: true}, + {name: "wrong second count", response: firstACK + "*3\r\n$9\r\nsubscribe\r\n$6\r\nsecond\r\n:1\r\n", wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) == 3 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "first" && args[2] == "second": + return tt.response + default: + return "-ERR unexpected command\r\n" + } + }) + pubsub := client.cmd.Subscribe(context.Background(), "first", "second") + t.Cleanup(func() { + if errClose := pubsub.Close(); errClose != nil { + t.Errorf("close PubSub: %v", errClose) + } + }) + + errACK := receiveSubscriptionACKs(context.Background(), pubsub, homeRedisTestOperationTimeout, []string{"first", "second"}) + if (errACK != nil) != tt.wantErr { + t.Fatalf("receiveSubscriptionACKs() error = %v, wantErr %t", errACK, tt.wantErr) + } + }) + } +} + +func TestRunConfigSubscriberLifetimeRejectsNonPositiveLifecycleDuration(t *testing.T) { + configPayload := "credential-concurrency:\n" + + " lifecycle-config-revision: 1\n" + + " cpa-heartbeat-timeout: 0s\n" + + " cpa-cancel-bound: 5s\n" + + " reclaim-grace: 5s\n" + + " cleanup-interval: 5s\n" + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: + return fmt.Sprintf("$%d\r\n%s\r\n", len(configPayload), configPayload) + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + default: + return "+OK\r\n" + } + }) + + errRun := client.RunConfigSubscriberLifetime(context.Background(), func(raw []byte) error { + parsed, errParse := config.ParseConfigBytes(raw) + if errParse != nil { + return errParse + } + return client.SetLifecycleConfig(parsed.CredentialConcurrency) + }, nil) + if errRun == nil { + t.Fatal("RunConfigSubscriberLifetime() error = nil, want invalid lifecycle duration rejection") + } + if got := findRedisCommand(commands.All(), "SUBSCRIBE"); got != nil { + t.Fatalf("SUBSCRIBE wire command = %#v, want no subscription after invalid GET config", got) + } +} + +func TestRunConfigSubscriberLifetimeRejectsExplicitInvalidLifecycleConfig(t *testing.T) { + configPayload := "credential-concurrency:\n" + + " lifecycle-config-revision: 0\n" + + " cpa-heartbeat-timeout: 20ms\n" + + " cpa-cancel-bound: 5s\n" + + " reclaim-grace: 5s\n" + + " cleanup-interval: 5s\n" + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: + return fmt.Sprintf("$%d\r\n%s\r\n", len(configPayload), configPayload) + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + default: + return "+OK\r\n" + } + }) + + errRun := client.RunConfigSubscriberLifetime(context.Background(), func(raw []byte) error { + parsed, errParse := config.ParseConfigBytes(raw) + if errParse != nil { + return errParse + } + return client.SetLifecycleConfig(parsed.CredentialConcurrency) + }, nil) + if errRun == nil { + t.Fatal("RunConfigSubscriberLifetime() error = nil, want invalid lifecycle config rejection") + } + if got := findRedisCommand(commands.All(), "SUBSCRIBE"); got != nil { + t.Fatalf("SUBSCRIBE wire command = %#v, want no subscription after invalid GET config", got) + } +} + +type blockingSubscriptionCloser struct { + started chan struct{} + release chan struct{} +} + +func (c *blockingSubscriptionCloser) Close() error { + close(c.started) + <-c.release + return nil +} + +func TestEndConfigSubscriberLifetimeClearsHeartbeatBeforeCloseBlocks(t *testing.T) { + client := New(config.HomeConfig{Enabled: true}) + client.heartbeatOK.Store(true) + closer := &blockingSubscriptionCloser{started: make(chan struct{}), release: make(chan struct{})} + + done := make(chan error, 1) + go func() { + done <- client.endConfigSubscriberLifetimeWithSubscription(errors.New("heartbeat lost"), closer, "heartbeat loss") + }() + + select { + case <-closer.started: + case <-time.After(time.Second): + t.Fatal("subscription close did not start") + } + if client.heartbeatOK.Load() { + close(closer.release) + t.Fatal("HeartbeatOK() remained true while subscription close was blocked") + } + select { + case errEnd := <-done: + close(closer.release) + t.Fatalf("endConfigSubscriberLifetimeWithSubscription() returned before subscription close unblocked: %v", errEnd) + default: + } + close(closer.release) + if errEnd := <-done; errEnd == nil { + t.Fatal("endConfigSubscriberLifetimeWithSubscription() error = nil, want heartbeat loss") + } +} + +func TestRunConfigSubscriberLifetimeUsesLegacySubscribeWithoutLifecycleConfig(t *testing.T) { + configPayload := "host: 127.0.0.1\n" + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: + return fmt.Sprintf("$%d\r\n%s\r\n", len(configPayload), configPayload) + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + default: + return "+OK\r\n" + } + }) + + errRun := client.RunConfigSubscriberLifetime(context.Background(), func(raw []byte) error { + parsed, errParse := config.ParseConfigBytes(raw) + if errParse != nil { + return errParse + } + if errSet := client.SetLifecycleConfig(parsed.CredentialConcurrency); errSet != nil { + return errSet + } + return nil + }, nil) + if errRun == nil { + t.Fatal("RunConfigSubscriberLifetime() error = nil after heartbeat loss") + } + if got := findRedisCommand(commands.All(), "SUBSCRIBE"); !reflect.DeepEqual(got, []string{"subscribe", "config"}) { + t.Fatalf("SUBSCRIBE wire command = %#v, want []string{\"subscribe\", \"config\"}", got) + } +} + +func TestRPopAuthLeavesCompleteServerErrorDeterministic(t *testing.T) { + client, _ := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 1 && strings.EqualFold(args[0], "RPOP"): + return "-ERR dispatch denied\r\n" + default: + return "+OK\r\n" + } + }) + client.heartbeatOK.Store(true) + + _, errRPop := client.RPopAuth(context.Background(), "gpt-5.4", "", nil, 1) + if errRPop == nil { + t.Fatal("RPopAuth() error = nil, want server failure") + } + if IsAmbiguousDispatchError(errRPop) { + t.Fatalf("RPopAuth() error = %v, want deterministic server error", errRPop) + } + if client.dispatchFenced.Load() || !client.heartbeatOK.Load() { + t.Fatalf("client fence/heartbeat = %v/%v, want false/true", client.dispatchFenced.Load(), client.heartbeatOK.Load()) + } +} + +type testRedisServerError string + +func (e testRedisServerError) Error() string { return string(e) } +func (testRedisServerError) RedisError() {} + +func TestIssuedRPopAuthErrorClassification(t *testing.T) { + tests := []struct { + name string + err error + ambiguous bool + }{ + {name: "redis server error", err: testRedisServerError("ERR denied"), ambiguous: false}, + {name: "redis nil", err: redis.Nil, ambiguous: false}, + {name: "closed connection", err: redis.ErrClosed, ambiguous: true}, + {name: "pool timeout", err: redis.ErrPoolTimeout, ambiguous: true}, + {name: "dial interruption", err: &net.OpError{Op: "dial", Err: errors.New("connection refused")}, ambiguous: true}, + {name: "tls interruption", err: x509.UnknownAuthorityError{}, ambiguous: true}, + {name: "write interruption", err: &net.OpError{Op: "write", Err: io.ErrClosedPipe}, ambiguous: true}, + {name: "partial response", err: io.ErrUnexpectedEOF, ambiguous: true}, + {name: "unknown transport", err: errors.New("unknown transport state"), ambiguous: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isAmbiguousIssuedRPopAuthError(tt.err); got != tt.ambiguous { + t.Fatalf("isAmbiguousIssuedRPopAuthError(%v) = %v, want %v", tt.err, got, tt.ambiguous) + } + }) + } +} + +func TestRPopAuthRejectsPreCanceledContextBeforeRequest(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func([]string) string { return "+OK\r\n" }) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, errRPop := client.RPopAuth(ctx, "gpt-5.4", "", nil, 1) + if !errors.Is(errRPop, context.Canceled) { + t.Fatalf("RPopAuth() error = %v, want context.Canceled", errRPop) + } + if IsAmbiguousDispatchError(errRPop) { + t.Fatalf("RPopAuth() error = %v, want deterministic pre-send cancellation", errRPop) + } + if commands.CountCommandKey("RPOP", "") != 0 { + t.Fatalf("commands = %#v, want no RPOP", commands.All()) + } +} + +func TestRPopAuthMarksRequestReadThenCloseAmbiguous(t *testing.T) { + client, requestRead, release := newBlockingRPopTestClient(t) + result := make(chan error, 1) + go func() { + _, errRPop := client.RPopAuth(context.Background(), "gpt-5.4", "", nil, 1) + result <- errRPop + }() + select { + case <-requestRead: + case <-time.After(time.Second): + t.Fatal("server did not read RPOP request") + } + close(release) + if errRPop := <-result; !IsAmbiguousDispatchError(errRPop) { + t.Fatalf("RPopAuth() error = %v, want ambiguous response interruption", errRPop) + } +} + +func TestRPopAuthLeavesHELLOSetupInterruptionDeterministic(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + commands := &redisCommandLog{} + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go func() { + defer func() { _ = conn.Close() }() + args, errRead := readRedisCommand(bufio.NewReader(conn)) + if errRead == nil { + commands.Append(args) + } + }() + } + }() + t.Cleanup(func() { + _ = listener.Close() + <-serverDone + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse listener port: %v", errPort) + } + client := New(config.HomeConfig{Enabled: true, Host: host, Port: port, DisableClusterDiscovery: true}) + t.Cleanup(client.Close) + + _, errRPop := client.RPopAuth(context.Background(), "gpt-5.4", "", nil, 1) + if errRPop == nil { + t.Fatal("RPopAuth() error = nil, want setup interruption") + } + if IsAmbiguousDispatchError(errRPop) { + t.Fatalf("RPopAuth() error = %v, want deterministic setup interruption", errRPop) + } + if client.dispatchFenced.Load() { + t.Fatal("RPopAuth() fenced the client after setup interruption") + } + allCommands := commands.All() + if len(allCommands) == 0 || len(allCommands[0]) == 0 || !strings.EqualFold(allCommands[0][0], "HELLO") { + t.Fatalf("commands = %#v, want HELLO setup before interruption", allCommands) + } + for _, command := range allCommands { + if len(command) > 0 && strings.EqualFold(command[0], "RPOP") { + t.Fatalf("commands = %#v, want no RPOP after setup interruption", allCommands) + } + } +} + +func TestTrackedRedisConnectionCloseRemovesContendedEntries(t *testing.T) { + client := New(config.HomeConfig{Enabled: true}) + const connectionCount = 32 + connections := make([]*homeDispatchConn, 0, connectionCount) + peers := make([]net.Conn, 0, connectionCount) + for range connectionCount { + local, peer := net.Pipe() + connections = append(connections, &homeDispatchConn{Conn: local, client: client}) + peers = append(peers, peer) + } + t.Cleanup(func() { + for _, peer := range peers { + _ = peer.Close() + } + }) + + client.mu.Lock() + client.connections = make(map[*homeDispatchConn]struct{}, len(connections)) + for _, conn := range connections { + client.connections[conn] = struct{}{} + } + started := make(chan struct{}, len(connections)) + closed := make(chan error, len(connections)) + for _, conn := range connections { + go func(conn *homeDispatchConn) { + started <- struct{}{} + closed <- conn.Close() + }(conn) + } + for range connections { + <-started + } + time.Sleep(20 * time.Millisecond) + client.mu.Unlock() + for range connections { + if errClose := <-closed; errClose != nil && !errors.Is(errClose, net.ErrClosed) { + t.Fatalf("tracked connection close: %v", errClose) + } + } + client.mu.Lock() + remaining := len(client.connections) + client.mu.Unlock() + if remaining != 0 { + t.Fatalf("tracked connection count = %d, want 0 after contended close churn", remaining) + } +} + +func TestAbortAmbiguousDispatchClosesBlockedRPopWithoutWaitingForResponse(t *testing.T) { + client, requestRead, release := newBlockingRPopTestClient(t) + client.heartbeatOK.Store(true) + result := make(chan error, 1) + go func() { + _, errRPop := client.RPopAuth(context.Background(), "gpt-5.4", "", nil, 1) + result <- errRPop + }() + select { + case <-requestRead: + case <-time.After(time.Second): + t.Fatal("server did not read RPOP request") + } + + aborted := make(chan struct{}) + go func() { + client.AbortAmbiguousDispatch() + close(aborted) + }() + select { + case <-aborted: + case <-time.After(time.Second): + close(release) + t.Fatal("AbortAmbiguousDispatch() waited for blocked RPOP response") + } + if client.heartbeatOK.Load() { + close(release) + t.Fatal("HeartbeatOK() remained true after abort") + } + client.mu.Lock() + commandClient, subscriptionClient := client.cmd, client.sub + client.mu.Unlock() + if commandClient != nil || subscriptionClient != nil { + close(release) + t.Fatalf("clients retained after abort: command=%v subscription=%v", commandClient != nil, subscriptionClient != nil) + } + select { + case errRPop := <-result: + if errRPop == nil { + close(release) + t.Fatal("RPopAuth() error = nil after client abort") + } + case <-time.After(time.Second): + close(release) + t.Fatal("RPopAuth() remained blocked after abort closed its client") + } + close(release) +} + +func TestRPopAuthLeavesPreSendFailureDeterministic(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 6379}) + + _, errRPop := client.RPopAuth(context.Background(), "", "", nil, 1) + if errRPop == nil { + t.Fatal("RPopAuth() error = nil, want requested model validation failure") + } + if IsAmbiguousDispatchError(errRPop) { + t.Fatalf("RPopAuth() error = %v, want deterministic pre-send failure", errRPop) + } +} + +func TestClientClosePermanentlyFencesDispatch(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 6379}) + client.mu.Lock() + client.cmd = redis.NewClient(&redis.Options{Addr: "127.0.0.1:6379"}) + client.mu.Unlock() + + client.Close() + if _, errClient := client.commandClient(); !errors.Is(errClient, ErrDispatchFenced) { + t.Fatalf("commandClient() error = %v, want ErrDispatchFenced", errClient) + } + client.mu.Lock() + commandClient := client.cmd + client.mu.Unlock() + if commandClient != nil { + t.Fatal("commandClient() recreated a command pool after Close") + } +} + +func TestAbortAmbiguousDispatchFencesConcurrentRPop(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "127.0.0.1", Port: 6379}) + client.AbortAmbiguousDispatch() + + const attempts = 32 + errs := make(chan error, attempts) + var workers sync.WaitGroup + for range attempts { + workers.Add(1) + go func() { + defer workers.Done() + _, errRPop := client.RPopAuth(context.Background(), "gpt-5.4", "", nil, 1) + errs <- errRPop + }() + } + workers.Wait() + close(errs) + + for errRPop := range errs { + if !errors.Is(errRPop, ErrDispatchFenced) { + t.Fatalf("RPopAuth() error = %v, want ErrDispatchFenced", errRPop) + } + } + client.mu.Lock() + commandClient := client.cmd + client.mu.Unlock() + if commandClient != nil { + t.Fatal("RPopAuth() recreated a command pool after AbortAmbiguousDispatch") + } +} + +func TestRunConfigSubscriberLifetimeRebuildsFreshCommandPoolBeforeReady(t *testing.T) { + configPayload := "host: 127.0.0.1\n" + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: + return fmt.Sprintf("$%d\r\n%s\r\n", len(configPayload), configPayload) + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + case len(args) >= 1 && strings.EqualFold(args[0], "PING"): + return "+PONG\r\n" + default: + return "+OK\r\n" + } + }) + var bootstrap *redis.Client + var freshCommandClient *redis.Client + ready := make(chan struct{}, 1) + errRun := client.RunConfigSubscriberLifetime(context.Background(), func([]byte) error { + client.mu.Lock() + bootstrap = client.cmd + client.mu.Unlock() + return nil + }, func() { + client.mu.Lock() + freshCommandClient = client.cmd + client.mu.Unlock() + ready <- struct{}{} + }) + if errRun == nil { + t.Fatal("RunConfigSubscriberLifetime() error = nil after heartbeat loss") + } + select { + case <-ready: + default: + t.Fatalf("RunConfigSubscriberLifetime() did not invoke onReady: %v", errRun) + } + if bootstrap == nil || freshCommandClient == nil || freshCommandClient == bootstrap { + t.Fatalf("command pools bootstrap=%p fresh=%p, want distinct non-nil pools", bootstrap, freshCommandClient) + } + if got := findRedisCommand(commands.All(), "PING"); got == nil { + t.Fatalf("commands = %#v, want fresh command PING before onReady", commands.All()) + } +} + +func TestRunConfigSubscriberLifetimeDoesNotReadyWhenFreshCommandProbeFails(t *testing.T) { + configPayload := "host: 127.0.0.1\n" + client, _ := newRedisCommandTestClient(t, func(args []string) string { + switch { + case len(args) >= 1 && strings.EqualFold(args[0], "HELLO"): + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == redisKeyConfig: + return fmt.Sprintf("$%d\r\n%s\r\n", len(configPayload), configPayload) + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == redisChannelConfig: + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + case len(args) >= 1 && strings.EqualFold(args[0], "PING"): + return "-ERR fresh command probe failed\r\n" + default: + return "+OK\r\n" + } + }) + ready := make(chan struct{}, 1) + errRun := client.RunConfigSubscriberLifetime(context.Background(), func([]byte) error { return nil }, func() { ready <- struct{}{} }) + if errRun == nil { + t.Fatal("RunConfigSubscriberLifetime() error = nil, want fresh command probe failure") + } + select { + case <-ready: + t.Fatalf("RunConfigSubscriberLifetime() invoked onReady after fresh command probe failure: %v", errRun) + default: + } + client.mu.Lock() + commandClient, subscriptionClient := client.cmd, client.sub + client.mu.Unlock() + if commandClient != nil || subscriptionClient != nil { + t.Fatalf("clients retained after fresh command probe failure: command=%v subscription=%v", commandClient != nil, subscriptionClient != nil) + } +} + +func findRedisCommand(commands [][]string, commandName string) []string { + for _, command := range commands { + if len(command) > 0 && strings.EqualFold(command[0], commandName) { + return command + } + } + return nil } diff --git a/internal/home/concurrency_release.go b/internal/home/concurrency_release.go new file mode 100644 --- /dev/null +++ b/internal/home/concurrency_release.go @@ -0,0 +1,272 @@ +package home + +import ( + "context" + "sync" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" +) + +// ConcurrencyReleaseFrame is the cumulative release accepted by Home for one credential and model. +type ConcurrencyReleaseFrame struct { + CredentialID string `json:"credential_id"` + Model string `json:"model"` + ReleaseSeq int64 `json:"release_seq"` +} + +type releaseState struct { + Latest int64 + Acked int64 + waiters map[int64][]chan struct{} +} + +type releaseFlusher struct { + mu sync.Mutex + groups map[executionregistry.ReleaseGroup]releaseState + flushInterval time.Duration + maxBackoff time.Duration + configProvider func() internalconfig.CredentialConcurrencyConfig + send func(context.Context, ConcurrencyReleaseFrame) error + wake chan struct{} + force chan context.Context +} + +func newReleaseFlusher(flushInterval, maxBackoff time.Duration, send func(context.Context, ConcurrencyReleaseFrame) error) *releaseFlusher { + return &releaseFlusher{ + groups: make(map[executionregistry.ReleaseGroup]releaseState), + flushInterval: flushInterval, + maxBackoff: maxBackoff, + send: send, + wake: make(chan struct{}, 1), + force: make(chan context.Context, 1), + } +} + +// NewReleaseFlusher creates a flusher that reads timing updates from the current limiter configuration. +func NewReleaseFlusher(configProvider func() internalconfig.CredentialConcurrencyConfig, send func(context.Context, ConcurrencyReleaseFrame) error) *releaseFlusher { + flusher := newReleaseFlusher(0, 0, send) + flusher.SetConfigProvider(configProvider) + return flusher +} + +func (f *releaseFlusher) SetConfigProvider(provider func() internalconfig.CredentialConcurrencyConfig) { + if f == nil { + return + } + f.mu.Lock() + f.configProvider = provider + f.mu.Unlock() + f.signal() +} + +// MarkDirty records the latest cumulative sequence for one release group and +// returns a ticket completed when Home acknowledges that sequence. +func (f *releaseFlusher) MarkDirty(group executionregistry.ReleaseGroup, sequence int64) *executionregistry.ReleaseTicket { + if f == nil || sequence <= 0 || group.CredentialID == "" || group.Model == "" { + return nil + } + + done := make(chan struct{}) + f.mu.Lock() + state := f.groups[group] + if sequence <= state.Acked { + close(done) + } else { + if state.waiters == nil { + state.waiters = make(map[int64][]chan struct{}) + } + state.waiters[sequence] = append(state.waiters[sequence], done) + if sequence > state.Latest { + state.Latest = sequence + } + f.groups[group] = state + } + f.mu.Unlock() + f.signal() + return executionregistry.NewReleaseTicket(group, sequence, done) +} + +// Run sends dirty groups until its lifetime is cancelled. +func (f *releaseFlusher) Run(ctx context.Context) { + if f == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + + timer := time.NewTimer(0) + defer timer.Stop() + delay := f.timings().flushInterval + backingOff := false + for { + select { + case <-ctx.Done(): + return + case <-f.wake: + if !backingOff { + resetReleaseTimer(timer, 0) + } + case forceCtx := <-f.force: + resetReleaseTimer(timer, 0) + failed := f.flush(forceCtx) + delay, backingOff = f.nextDelay(delay, failed) + resetReleaseTimer(timer, delay) + case <-timer.C: + failed := f.flush(ctx) + delay, backingOff = f.nextDelay(delay, failed) + timer.Reset(delay) + } + } +} + +func (f *releaseFlusher) nextDelay(delay time.Duration, failed bool) (time.Duration, bool) { + timings := f.timings() + if !failed { + return timings.flushInterval, false + } + delay *= 2 + if delay < timings.flushInterval { + delay = timings.flushInterval + } + if delay > timings.maxBackoff { + delay = timings.maxBackoff + } + return delay, true +} + +func resetReleaseTimer(timer *time.Timer, delay time.Duration) { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timer.Reset(delay) +} + +type releaseFlusherTimings struct { + flushInterval time.Duration + maxBackoff time.Duration +} + +func (f *releaseFlusher) timings() releaseFlusherTimings { + defaults := internalconfig.CredentialConcurrencyConfig{}.WithDefaults() + timings := releaseFlusherTimings{flushInterval: f.flushInterval, maxBackoff: f.maxBackoff} + + f.mu.Lock() + provider := f.configProvider + f.mu.Unlock() + if provider != nil { + cfg := provider().WithDefaults() + timings.flushInterval = cfg.ReleaseFlushInterval + timings.maxBackoff = cfg.ReleaseMaxBackoff + } + if timings.flushInterval <= 0 { + timings.flushInterval = defaults.ReleaseFlushInterval + } + if timings.maxBackoff < timings.flushInterval { + timings.maxBackoff = timings.flushInterval + } + return timings +} + +func (f *releaseFlusher) flush(ctx context.Context) bool { + if f == nil || f.send == nil { + return false + } + + f.mu.Lock() + pending := make(map[executionregistry.ReleaseGroup]int64, len(f.groups)) + for group, state := range f.groups { + if state.Latest > state.Acked { + pending[group] = state.Latest + } + } + f.mu.Unlock() + + failed := false + for group, sequence := range pending { + errSend := f.send(ctx, ConcurrencyReleaseFrame{ + CredentialID: group.CredentialID, + Model: group.Model, + ReleaseSeq: sequence, + }) + if errSend != nil { + failed = true + continue + } + f.mu.Lock() + state := f.groups[group] + if sequence > state.Acked { + state.Acked = sequence + for waiterSequence, waiters := range state.waiters { + if waiterSequence <= state.Acked { + for _, done := range waiters { + close(done) + } + delete(state.waiters, waiterSequence) + } + } + } + f.groups[group] = state + f.mu.Unlock() + } + return failed +} + +// Flush waits for all currently dirty groups to be acknowledged within ctx. +func (f *releaseFlusher) Flush(ctx context.Context) error { + if f == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + f.forceFlush(ctx) + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + for { + if f.idle() { + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + } + } +} + +func (f *releaseFlusher) idle() bool { + f.mu.Lock() + defer f.mu.Unlock() + for _, state := range f.groups { + if state.Latest > state.Acked { + return false + } + } + return true +} + +func (f *releaseFlusher) signal() { + if f == nil { + return + } + select { + case f.wake <- struct{}{}: + default: + } +} + +func (f *releaseFlusher) forceFlush(ctx context.Context) { + if f == nil { + return + } + select { + case f.force <- ctx: + default: + } +} diff --git a/internal/home/concurrency_release_test.go b/internal/home/concurrency_release_test.go new file mode 100644 --- /dev/null +++ b/internal/home/concurrency_release_test.go @@ -0,0 +1,476 @@ +package home + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + "sync" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" +) + +func concurrencyReleaseFrameFromFixture(t *testing.T) ConcurrencyReleaseFrame { + t.Helper() + raw, errRead := os.ReadFile(filepath.Join("testdata", "concurrency_release.json")) + if errRead != nil { + t.Fatal(errRead) + } + + var frame ConcurrencyReleaseFrame + if errUnmarshal := json.Unmarshal(raw, &frame); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + return frame +} + +func TestConcurrencyReleaseFrameFixture(t *testing.T) { + raw, errRead := os.ReadFile(filepath.Join("testdata", "concurrency_release.json")) + if errRead != nil { + t.Fatal(errRead) + } + frame := concurrencyReleaseFrameFromFixture(t) + if frame != (ConcurrencyReleaseFrame{CredentialID: "cred-1", Model: "gpt", ReleaseSeq: 1}) { + t.Fatalf("fixture frame = %#v", frame) + } + marshaled, errMarshal := json.Marshal(frame) + if errMarshal != nil { + t.Fatal(errMarshal) + } + if !bytes.Equal(marshaled, bytes.TrimSpace(raw)) { + t.Fatalf("marshaled frame = %q, want fixture %q", marshaled, bytes.TrimSpace(raw)) + } +} + +type recordingReleaseSender struct { + mu sync.Mutex + failures int + frames []ConcurrencyReleaseFrame + acked []ConcurrencyReleaseFrame + sent chan struct{} +} + +func (s *recordingReleaseSender) Send(_ context.Context, frame ConcurrencyReleaseFrame) error { + s.mu.Lock() + s.frames = append(s.frames, frame) + failed := s.failures > 0 + if failed { + s.failures-- + } else { + s.acked = append(s.acked, frame) + } + s.mu.Unlock() + select { + case s.sent <- struct{}{}: + default: + } + if failed { + return errors.New("temporary Home failure") + } + return nil +} + +func (s *recordingReleaseSender) LastSequence() int64 { + s.mu.Lock() + defer s.mu.Unlock() + if len(s.acked) == 0 { + return 0 + } + return s.acked[len(s.acked)-1].ReleaseSeq +} + +func (s *recordingReleaseSender) WaitForSequence(sequence int64, timeout time.Duration) bool { + timer := time.NewTimer(timeout) + defer timer.Stop() + for { + if s.LastSequence() == sequence { + return true + } + select { + case <-timer.C: + return false + case <-s.sent: + } + } +} + +func TestReleaseFlusherRetriesLatestCumulativeSequence(t *testing.T) { + sender := &recordingReleaseSender{failures: 1, sent: make(chan struct{}, 8)} + flusher := newReleaseFlusher(10*time.Millisecond, 40*time.Millisecond, sender.Send) + group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"} + flusher.MarkDirty(group, 1) + flusher.MarkDirty(group, 3) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + go flusher.Run(ctx) + + if !sender.WaitForSequence(3, 500*time.Millisecond) { + t.Fatalf("last sequence = %d, want 3", sender.LastSequence()) + } + if sender.LastSequence() != 3 { + t.Fatalf("last sequence = %d, want 3", sender.LastSequence()) + } +} + +type blockingReleaseSender struct { + started chan struct{} + release chan struct{} + frames chan ConcurrencyReleaseFrame + once sync.Once +} + +func (s *blockingReleaseSender) Send(_ context.Context, frame ConcurrencyReleaseFrame) error { + s.once.Do(func() { close(s.started) }) + select { + case s.frames <- frame: + default: + } + <-s.release + return nil +} + +func TestReleaseFlusherDoesNotLoseASequenceMarkedDuringSend(t *testing.T) { + sender := &blockingReleaseSender{ + started: make(chan struct{}), + release: make(chan struct{}), + frames: make(chan ConcurrencyReleaseFrame, 4), + } + flusher := newReleaseFlusher(time.Millisecond, 10*time.Millisecond, sender.Send) + group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"} + flusher.MarkDirty(group, 1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go flusher.Run(ctx) + + select { + case <-sender.started: + case <-time.After(time.Second): + t.Fatal("release flusher did not begin sending") + } + flusher.MarkDirty(group, 2) + close(sender.release) + + deadline := time.NewTimer(time.Second) + defer deadline.Stop() + for { + select { + case frame := <-sender.frames: + if frame.ReleaseSeq == 2 { + return + } + case <-deadline.C: + t.Fatal("release flusher did not send the latest sequence") + } + } +} + +func TestReleaseFlusherUsesCurrentLimiterConfig(t *testing.T) { + flusher := newReleaseFlusher(time.Hour, 2*time.Hour, func(context.Context, ConcurrencyReleaseFrame) error { return nil }) + flusher.SetConfigProvider(func() internalconfig.CredentialConcurrencyConfig { + return internalconfig.CredentialConcurrencyConfig{ + ReleaseFlushInterval: 5 * time.Millisecond, + ReleaseMaxBackoff: 25 * time.Millisecond, + } + }) + if got := flusher.timings(); got.flushInterval != 5*time.Millisecond || got.maxBackoff != 25*time.Millisecond { + t.Fatalf("timings = %#v", got) + } +} + +func TestReleaseFlusherStopsWithLifetime(t *testing.T) { + sender := &recordingReleaseSender{sent: make(chan struct{}, 1)} + flusher := newReleaseFlusher(time.Hour, time.Hour, sender.Send) + done := make(chan struct{}) + ctx, cancel := context.WithCancel(context.Background()) + go func() { + defer close(done) + flusher.Run(ctx) + }() + cancel() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("release flusher did not stop with its lifetime") + } +} + +type timedReleaseAttempt struct { + at time.Time + frame ConcurrencyReleaseFrame + failed bool +} + +type outageReleaseSender struct { + mu sync.Mutex + outage bool + attempts []timedReleaseAttempt + sent chan struct{} +} + +func (s *outageReleaseSender) Send(_ context.Context, frame ConcurrencyReleaseFrame) error { + s.mu.Lock() + failed := s.outage + s.attempts = append(s.attempts, timedReleaseAttempt{at: time.Now(), frame: frame, failed: failed}) + s.mu.Unlock() + select { + case s.sent <- struct{}{}: + default: + } + if failed { + return errors.New("temporary Home outage") + } + return nil +} + +func (s *outageReleaseSender) SetOutage(outage bool) { + s.mu.Lock() + s.outage = outage + s.mu.Unlock() +} + +func (s *outageReleaseSender) WaitForAttempts(count int, timeout time.Duration) []timedReleaseAttempt { + timer := time.NewTimer(timeout) + defer timer.Stop() + for { + s.mu.Lock() + attempts := append([]timedReleaseAttempt(nil), s.attempts...) + s.mu.Unlock() + if len(attempts) >= count { + return attempts + } + select { + case <-timer.C: + return attempts + case <-s.sent: + } + } +} + +func TestReleaseFlusherCoalescesDirtyWakesDuringFailureBackoff(t *testing.T) { + const ( + flushInterval = 20 * time.Millisecond + maxBackoff = 80 * time.Millisecond + tolerance = 10 * time.Millisecond + ) + + sender := &outageReleaseSender{outage: true, sent: make(chan struct{}, 32)} + flusher := newReleaseFlusher(flushInterval, maxBackoff, sender.Send) + group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"} + flusher.MarkDirty(group, 1) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + flusher.Run(ctx) + }() + defer func() { + cancel() + <-done + }() + + stopReleases := make(chan struct{}) + producerDone := make(chan struct{}) + latest := int64(1) + go func() { + defer close(producerDone) + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + for { + select { + case <-stopReleases: + return + case <-ticker.C: + latest++ + flusher.MarkDirty(group, latest) + } + } + }() + + attempts := sender.WaitForAttempts(3, time.Second) + close(stopReleases) + <-producerDone + if len(attempts) < 3 { + t.Fatalf("attempt count = %d, want at least 3", len(attempts)) + } + for _, attempt := range attempts[:3] { + if !attempt.failed { + t.Fatal("release unexpectedly succeeded during outage") + } + } + if got := attempts[1].at.Sub(attempts[0].at); got < 2*flushInterval-tolerance { + t.Fatalf("first retry delay = %s, want at least %s", got, 2*flushInterval-tolerance) + } + if got := attempts[2].at.Sub(attempts[1].at); got < maxBackoff-tolerance { + t.Fatalf("second retry delay = %s, want at least %s", got, maxBackoff-tolerance) + } + + latest++ + recoverySequence := latest + recoveryStart := attempts[2].at + sender.SetOutage(false) + flusher.MarkDirty(group, recoverySequence) + + attempts = sender.WaitForAttempts(4, time.Second) + if len(attempts) < 4 { + t.Fatalf("attempt count after recovery = %d, want at least 4", len(attempts)) + } + recovered := attempts[3] + if recovered.failed || recovered.frame.ReleaseSeq != recoverySequence { + t.Fatalf("recovery attempt = %#v, want successful sequence %d", recovered, recoverySequence) + } + if got := recovered.at.Sub(recoveryStart); got < maxBackoff-tolerance { + t.Fatalf("recovery retry delay = %s, want at least %s", got, maxBackoff-tolerance) + } +} + +type boundedForceReleaseSender struct { + attempts chan context.Context + calls int +} + +func (s *boundedForceReleaseSender) Send(ctx context.Context, _ ConcurrencyReleaseFrame) error { + s.calls++ + select { + case s.attempts <- ctx: + default: + } + if s.calls == 1 { + return errors.New("temporary Home failure") + } + <-ctx.Done() + return ctx.Err() +} + +func TestReleaseFlusherFlushForceUsesBoundedContext(t *testing.T) { + sender := &boundedForceReleaseSender{attempts: make(chan context.Context, 2)} + flusher := newReleaseFlusher(time.Second, time.Second, sender.Send) + group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"} + flusher.MarkDirty(group, 1) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + flusher.Run(ctx) + }() + defer func() { + cancel() + <-done + }() + + select { + case <-sender.attempts: + case <-time.After(time.Second): + t.Fatal("release flusher did not make the initial failed attempt") + } + + flushCtx, cancelFlush := context.WithTimeout(context.Background(), 40*time.Millisecond) + defer cancelFlush() + if errFlush := flusher.Flush(flushCtx); !errors.Is(errFlush, context.DeadlineExceeded) { + t.Fatalf("Flush() error = %v, want deadline exceeded", errFlush) + } + + select { + case forceCtx := <-sender.attempts: + if _, ok := forceCtx.Deadline(); !ok { + t.Fatal("forced release attempt did not receive the bounded Flush context") + } + case <-time.After(time.Second): + t.Fatal("Flush() did not bypass the normal retry interval") + } +} + +func TestScopeEndBlocksDrainUntilReleaseSinkFlushesFinalSequence(t *testing.T) { + sender := &recordingReleaseSender{sent: make(chan struct{}, 2)} + flusher := newReleaseFlusher(time.Hour, time.Hour, sender.Send) + releaseCtx, cancelRelease := context.WithCancel(context.Background()) + releaseDone := make(chan struct{}) + go func() { + defer close(releaseDone) + flusher.Run(releaseCtx) + }() + defer func() { + cancelRelease() + <-releaseDone + }() + + registry := executionregistry.New() + sinkStarted := make(chan struct{}) + unblockSink := make(chan struct{}) + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, sequence int64) { + close(sinkStarted) + <-unblockSink + flusher.MarkDirty(group, sequence) + }) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{CredentialID: "cred-1", Model: "gpt", Accounted: true}) + if errInstall != nil { + t.Fatal(errInstall) + } + + endDone := make(chan struct{}) + go func() { + defer close(endDone) + scope.End("complete") + }() + select { + case <-sinkStarted: + case <-time.After(time.Second): + t.Fatal("Scope.End() did not call the release sink") + } + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + drainDone := make(chan error, 1) + go func() { drainDone <- registry.Drain(drainCtx) }() + + select { + case errDrain := <-drainDone: + t.Fatalf("Drain() returned before the release sink completed: %v", errDrain) + case <-time.After(20 * time.Millisecond): + } + + mutexAvailable := make(chan struct{}) + go func() { + registry.SetReleaseSink(nil) + close(mutexAvailable) + }() + select { + case <-mutexAvailable: + case <-time.After(time.Second): + t.Fatal("release sink blocked the registry mutex") + } + if _, errBegin := registry.BeginDispatch(); !errors.Is(errBegin, executionregistry.ErrRegistryNotAccepting) { + t.Fatalf("BeginDispatch() error = %v, want ErrRegistryNotAccepting", errBegin) + } + + close(unblockSink) + select { + case <-endDone: + case <-time.After(time.Second): + t.Fatal("Scope.End() did not complete after the release sink unblocked") + } + if errDrain := <-drainDone; errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } + + flushCtx, cancelFlush := context.WithTimeout(context.Background(), time.Second) + defer cancelFlush() + if errFlush := flusher.Flush(flushCtx); errFlush != nil { + t.Fatalf("Flush() error = %v", errFlush) + } + if got := sender.LastSequence(); got != 1 { + t.Fatalf("final flushed sequence = %d, want 1", got) + } +} diff --git a/internal/home/global.go b/internal/home/global.go --- a/internal/home/global.go +++ b/internal/home/global.go @@ -2,7 +2,7 @@ import "sync/atomic" -var currentClient atomic.Value // *Client +var currentClient atomic.Pointer[Client] // SetCurrent sets the active home client used by runtime integrations. func SetCurrent(client *Client) { @@ -11,15 +11,17 @@ // Current returns the active home client instance, if any. func Current() *Client { - if v := currentClient.Load(); v != nil { - if client, ok := v.(*Client); ok { - return client - } - } - return nil + return currentClient.Load() } // ClearCurrent removes the active home client. func ClearCurrent() { - currentClient.Store((*Client)(nil)) + currentClient.Store(nil) +} + +// ClearCurrentIf removes the active client only when it is client. +func ClearCurrentIf(client *Client) { + if client != nil { + currentClient.CompareAndSwap(client, nil) + } } diff --git a/internal/home/in_flight_contract_test.go b/internal/home/in_flight_contract_test.go new file mode 100644 --- /dev/null +++ b/internal/home/in_flight_contract_test.go @@ -0,0 +1,182 @@ +package home + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "os" + "path/filepath" + "reflect" + "testing" +) + +func TestCredentialInFlightWireContractFixture(t *testing.T) { + raw, errRead := os.ReadFile(filepath.Join("testdata", "credential_in_flight_contract.json")) + if errRead != nil { + t.Fatalf("ReadFile() error = %v", errRead) + } + fixture, errDecode := decodeInFlightContractFixture(raw) + if errDecode != nil { + t.Fatalf("decodeInFlightContractFixture() error = %v", errDecode) + } + if fixture.Part.Kind != InFlightFramePart || fixture.Part.PartIndex == nil || *fixture.Part.PartIndex != 0 || fixture.Part.PartCount == nil || *fixture.Part.PartCount != 1 { + t.Fatalf("part = %#v", fixture.Part) + } + if fixture.Part.Aggregates[0].Status != InFlightAccounted || fixture.Part.Aggregates[1].Status != InFlightUnaccounted { + t.Fatalf("statuses = %#v", fixture.Part.Aggregates) + } + if fixture.Overflow.Kind != InFlightFrameOverflow || fixture.Overflow.AggregateGroupCount != 100001 { + t.Fatalf("overflow = %#v", fixture.Overflow) + } + assertInFlightContractFields(t) + assertRequiredInFlightJSONKeys(t, raw, []string{"config", "part", "overflow"}) + assertInFlightFixtureKeys(t, fixture) +} + +func TestCredentialInFlightWireContractRejectsInvalidJSON(t *testing.T) { + raw, errRead := os.ReadFile(filepath.Join("testdata", "credential_in_flight_contract.json")) + if errRead != nil { + t.Fatalf("ReadFile() error = %v", errRead) + } + for _, test := range []struct { + name string + raw []byte + }{ + {name: "unknown frame owner field", raw: bytes.Replace(raw, []byte(`"kind": "part"`), []byte(`"kind": "part", "node_id": "node-a"`), 1)}, + {name: "unknown aggregate owner field", raw: bytes.Replace(raw, []byte(`"credential_id": "cred-a"`), []byte(`"credential_id": "cred-a", "fingerprint": "owner"`), 1)}, + {name: "unknown detail secret field", raw: bytes.Replace(raw, []byte(`"request_id": "req-1"`), []byte(`"request_id": "req-1", "secret": "secret"`), 1)}, + {name: "unknown overflow secret field", raw: bytes.Replace(raw, []byte(`"aggregate_group_count": 100001`), []byte(`"aggregate_group_count": 100001, "api_key": "secret"`), 1)}, + {name: "trailing JSON", raw: append(append([]byte{}, raw...), []byte(` {"part": {}}`)...)}, + } { + t.Run(test.name, func(t *testing.T) { + if _, errDecode := decodeInFlightContractFixture(test.raw); errDecode == nil { + t.Fatal("decodeInFlightContractFixture() error = nil") + } + }) + } +} + +type inFlightContractFixture struct { + Part InFlightSnapshotFrame + Overflow InFlightSnapshotFrame + PartJSON json.RawMessage + OverflowJSON json.RawMessage +} + +func decodeInFlightContractFixture(raw []byte) (inFlightContractFixture, error) { + var fixture inFlightContractFixture + var document struct { + Config json.RawMessage `json:"config"` + Part InFlightSnapshotFrame `json:"part"` + Overflow InFlightSnapshotFrame `json:"overflow"` + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if errDecode := decoder.Decode(&document); errDecode != nil { + return fixture, errDecode + } + if errDecode := decoder.Decode(&struct{}{}); errDecode == nil { + return fixture, errors.New("unexpected trailing JSON") + } else if errDecode != io.EOF { + return fixture, errDecode + } + documentRaw := struct { + Part json.RawMessage `json:"part"` + Overflow json.RawMessage `json:"overflow"` + }{} + if errDecode := json.Unmarshal(raw, &documentRaw); errDecode != nil { + return fixture, errDecode + } + fixture.Part = document.Part + fixture.Overflow = document.Overflow + fixture.PartJSON = documentRaw.Part + fixture.OverflowJSON = documentRaw.Overflow + return fixture, nil +} + +func assertInFlightContractFields(t *testing.T) { + t.Helper() + assertOrderedInFlightJSONFields(t, reflect.TypeOf(InFlightSnapshotFrame{}), []inFlightJSONField{ + {name: "Kind", tag: "kind"}, + {name: "Revision", tag: "revision"}, + {name: "ObservedAt", tag: "observed_at"}, + {name: "BarrierRevision", tag: "barrier_revision"}, + {name: "PartIndex", tag: "part_index,omitempty"}, + {name: "PartCount", tag: "part_count,omitempty"}, + {name: "DetailsTruncated", tag: "details_truncated,omitempty"}, + {name: "Aggregates", tag: "aggregates,omitempty"}, + {name: "Details", tag: "details,omitempty"}, + {name: "AggregateGroupCount", tag: "aggregate_group_count,omitempty"}, + }) + assertOrderedInFlightJSONFields(t, reflect.TypeOf(InFlightAggregate{}), []inFlightJSONField{ + {name: "CredentialID", tag: "credential_id"}, + {name: "Model", tag: "model"}, + {name: "Status", tag: "status"}, + {name: "Count", tag: "count"}, + }) + assertOrderedInFlightJSONFields(t, reflect.TypeOf(InFlightRequestDetail{}), []inFlightJSONField{ + {name: "RequestID", tag: "request_id"}, + {name: "CredentialID", tag: "credential_id"}, + {name: "Model", tag: "model"}, + {name: "RequestKind", tag: "request_kind"}, + {name: "StartedAt", tag: "started_at"}, + }) +} + +func assertInFlightFixtureKeys(t *testing.T, fixture inFlightContractFixture) { + t.Helper() + assertRequiredInFlightJSONKeys(t, fixture.PartJSON, []string{"kind", "revision", "observed_at", "barrier_revision", "part_index", "part_count", "details_truncated", "aggregates", "details"}) + assertRequiredInFlightJSONKeys(t, fixture.OverflowJSON, []string{"kind", "revision", "observed_at", "barrier_revision", "aggregate_group_count"}) + + var part struct { + Aggregates []json.RawMessage `json:"aggregates"` + Details []json.RawMessage `json:"details"` + } + if errDecode := json.Unmarshal(fixture.PartJSON, &part); errDecode != nil { + t.Fatalf("json.Unmarshal() error = %v", errDecode) + } + for index, aggregate := range part.Aggregates { + assertRequiredInFlightJSONKeys(t, aggregate, []string{"credential_id", "model", "status", "count"}) + if len(aggregate) == 0 { + t.Fatalf("aggregate %d is empty", index) + } + } + for index, detail := range part.Details { + assertRequiredInFlightJSONKeys(t, detail, []string{"request_id", "credential_id", "model", "request_kind", "started_at"}) + if len(detail) == 0 { + t.Fatalf("detail %d is empty", index) + } + } +} + +type inFlightJSONField struct { + name string + tag string +} + +func assertOrderedInFlightJSONFields(t *testing.T, structType reflect.Type, want []inFlightJSONField) { + t.Helper() + if structType.NumField() != len(want) { + t.Fatalf("%s field count = %d, want %d", structType.Name(), structType.NumField(), len(want)) + } + for index, expected := range want { + field := structType.Field(index) + if field.Name != expected.name || field.Tag.Get("json") != expected.tag { + t.Fatalf("%s field %d = (%q, %q), want (%q, %q)", structType.Name(), index, field.Name, field.Tag.Get("json"), expected.name, expected.tag) + } + } +} + +func assertRequiredInFlightJSONKeys(t *testing.T, raw json.RawMessage, required []string) { + t.Helper() + var fields map[string]json.RawMessage + if errDecode := json.Unmarshal(raw, &fields); errDecode != nil { + t.Fatalf("json.Unmarshal() error = %v", errDecode) + } + for _, key := range required { + if _, ok := fields[key]; !ok { + t.Fatalf("required JSON key %q is missing", key) + } + } +} diff --git a/internal/home/requests.go b/internal/home/requests.go --- a/internal/home/requests.go +++ b/internal/home/requests.go @@ -1,11 +1,14 @@ package home +import "time" + type authDispatchRequest struct { - Type string `json:"type"` - Model string `json:"model"` - Count int `json:"count"` - SessionID string `json:"session_id,omitempty"` - Headers map[string]string `json:"headers,omitempty"` + Type string `json:"type"` + Model string `json:"model"` + Count int `json:"count"` + ConcurrencyProtocol int `json:"concurrency_protocol,omitempty"` + SessionID string `json:"session_id,omitempty"` + Headers map[string]string `json:"headers,omitempty"` } type modelsRequest struct { @@ -17,4 +20,42 @@ type refreshRequest struct { Type string `json:"type"` AuthIndex string `json:"auth_index"` +} + +type InFlightFrameKind string +type InFlightAccountedStatus string + +const ( + InFlightFramePart InFlightFrameKind = "part" + InFlightFrameOverflow InFlightFrameKind = "overflow" + InFlightAccounted InFlightAccountedStatus = "accounted" + InFlightUnaccounted InFlightAccountedStatus = "unaccounted" +) + +type InFlightAggregate struct { + CredentialID string `json:"credential_id"` + Model string `json:"model"` + Status InFlightAccountedStatus `json:"status"` + Count int64 `json:"count"` +} + +type InFlightRequestDetail struct { + RequestID string `json:"request_id"` + CredentialID string `json:"credential_id"` + Model string `json:"model"` + RequestKind string `json:"request_kind"` + StartedAt time.Time `json:"started_at"` +} + +type InFlightSnapshotFrame struct { + Kind InFlightFrameKind `json:"kind"` + Revision int64 `json:"revision"` + ObservedAt time.Time `json:"observed_at"` + BarrierRevision int64 `json:"barrier_revision"` + PartIndex *int `json:"part_index,omitempty"` + PartCount *int `json:"part_count,omitempty"` + DetailsTruncated bool `json:"details_truncated,omitempty"` + Aggregates []InFlightAggregate `json:"aggregates,omitempty"` + Details []InFlightRequestDetail `json:"details,omitempty"` + AggregateGroupCount int `json:"aggregate_group_count,omitempty"` } diff --git a/internal/homeplugins/sync.go b/internal/homeplugins/sync.go --- a/internal/homeplugins/sync.go +++ b/internal/homeplugins/sync.go @@ -33,6 +33,10 @@ PluginRegistered(id string) bool } +type contextualPluginUnloader interface { + UnloadPluginContext(ctx context.Context, id string) bool +} + type SyncReport struct { SchemaVersion int `json:"schema_version"` TaskID uint `json:"task_id,omitempty"` @@ -364,7 +368,9 @@ } func DeleteWithReport(ctx context.Context, cfg *config.Config, pluginRuntime PluginRuntime, taskID uint, pluginID string) SyncReport { - _ = ctx + if ctx == nil { + ctx = context.Background() + } platform := CurrentPlatform() report := newSyncReport(platform) report.TaskID = taskID @@ -372,6 +378,13 @@ report.Phase = pluginTaskPhaseDelete pluginID = strings.TrimSpace(pluginID) status := PluginInstallStatus{ID: pluginID} + if errContext := ctx.Err(); errContext != nil { + status.InstallStatus = pluginInstallStatusFailed + status.Error = errContext.Error() + report.Plugins = append(report.Plugins, status) + finishReport(&report, errContext) + return report + } if cfg == nil { status.InstallStatus = pluginInstallStatusFailed status.Error = "home plugins: config is nil" @@ -388,7 +401,14 @@ finishReport(&report, errPluginsDir) return report } - path, deleted, errDelete := deletePluginArtifact(root, pluginID, pluginRuntime) + if errContext := ctx.Err(); errContext != nil { + status.InstallStatus = pluginInstallStatusFailed + status.Error = errContext.Error() + report.Plugins = append(report.Plugins, status) + finishReport(&report, errContext) + return report + } + path, deleted, errDelete := deletePluginArtifact(ctx, root, pluginID, pluginRuntime) status.Path = strings.TrimSpace(path) switch { case errDelete != nil: @@ -404,7 +424,13 @@ return report } -func deletePluginArtifact(root string, id string, pluginRuntime PluginRuntime) (string, bool, error) { +func deletePluginArtifact(ctx context.Context, root string, id string, pluginRuntime PluginRuntime) (string, bool, error) { + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return "", false, errContext + } id = strings.TrimSpace(id) if !validPluginFileID(id) { return "", false, fmt.Errorf("invalid plugin id %q", id) @@ -413,16 +439,31 @@ if errPaths != nil { return "", false, errPaths } + if errContext := ctx.Err(); errContext != nil { + return "", false, errContext + } if len(paths) == 0 { return "", false, nil } if pluginRuntime != nil && pluginRuntime.PluginBusy(id) { - if !pluginRuntime.UnloadPlugin(id) && pluginRuntime.PluginBusy(id) { + if errContext := ctx.Err(); errContext != nil { + return paths[0], false, errContext + } + unloaded := false + if contextual, ok := pluginRuntime.(contextualPluginUnloader); ok { + unloaded = contextual.UnloadPluginContext(ctx, id) + } else { + unloaded = pluginRuntime.UnloadPlugin(id) + } + if !unloaded && pluginRuntime.PluginBusy(id) { return paths[0], false, sdkpluginstore.ErrLoadedPluginLocked } } deleted := false for _, path := range paths { + if errContext := ctx.Err(); errContext != nil { + return paths[0], deleted, errContext + } if errRemove := os.Remove(path); errRemove != nil { if errors.Is(errRemove, os.ErrNotExist) { continue @@ -430,6 +471,9 @@ return paths[0], deleted, errRemove } deleted = true + if errContext := ctx.Err(); errContext != nil { + return paths[0], deleted, errContext + } } return paths[0], deleted, nil } diff --git a/internal/homeplugins/sync_test.go b/internal/homeplugins/sync_test.go --- a/internal/homeplugins/sync_test.go +++ b/internal/homeplugins/sync_test.go @@ -43,6 +43,16 @@ return i[id] } +type contextPluginRuntime struct { + fakePluginRuntime + unloadContext context.Context +} + +func (r *contextPluginRuntime) UnloadPluginContext(ctx context.Context, id string) bool { + r.unloadContext = ctx + return r.UnloadPlugin(id) +} + func TestSyncPlatformInstallsManifestArtifact(t *testing.T) { root := t.TempDir() archiveData := makeZip(t, map[string]string{"sample.dll": "library-data"}) @@ -650,6 +660,54 @@ } if _, errStat := os.Stat(otherTarget); errStat != nil { t.Fatalf("other plugin stat error = %v, want retained", errStat) + } +} + +func TestDeleteWithReportStopsBeforeUnloadWhenContextCanceled(t *testing.T) { + root := t.TempDir() + path := pluginTestPath(root, runtime.GOOS, runtime.GOARCH, "sample", "1.0.0") + if errMkdir := os.MkdirAll(filepath.Dir(path), 0o755); errMkdir != nil { + t.Fatal(errMkdir) + } + if errWrite := os.WriteFile(path, []byte("plugin"), 0o644); errWrite != nil { + t.Fatal(errWrite) + } + runtimeHost := &contextPluginRuntime{fakePluginRuntime: fakePluginRuntime{busy: true}} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + report := DeleteWithReport(ctx, syncTestConfig(t, root), runtimeHost, 44, "sample") + + if report.OK || !strings.Contains(report.Error, context.Canceled.Error()) { + t.Fatalf("canceled delete report = %+v, want context cancellation", report) + } + if runtimeHost.unloadContext != nil || len(runtimeHost.unloaded) != 0 { + t.Fatalf("canceled delete unloaded plugin: context=%v unloads=%v", runtimeHost.unloadContext, runtimeHost.unloaded) + } + if _, errStat := os.Stat(path); errStat != nil { + t.Fatalf("canceled delete removed plugin artifact: %v", errStat) + } +} + +func TestDeleteWithReportUsesContextualUnload(t *testing.T) { + root := t.TempDir() + path := pluginTestPath(root, runtime.GOOS, runtime.GOARCH, "sample", "1.0.0") + if errMkdir := os.MkdirAll(filepath.Dir(path), 0o755); errMkdir != nil { + t.Fatal(errMkdir) + } + if errWrite := os.WriteFile(path, []byte("plugin"), 0o644); errWrite != nil { + t.Fatal(errWrite) + } + runtimeHost := &contextPluginRuntime{fakePluginRuntime: fakePluginRuntime{busy: true}} + ctx := context.WithValue(context.Background(), struct{}{}, "contextual") + + report := DeleteWithReport(ctx, syncTestConfig(t, root), runtimeHost, 45, "sample") + + if !report.OK { + t.Fatalf("contextual delete report = %+v", report) + } + if runtimeHost.unloadContext != ctx || len(runtimeHost.unloaded) != 1 || runtimeHost.unloaded[0] != "sample" { + t.Fatalf("contextual unload = context=%v unloads=%v", runtimeHost.unloadContext, runtimeHost.unloaded) } } diff --git a/internal/logging/home_app_log_forwarder.go b/internal/logging/home_app_log_forwarder.go --- a/internal/logging/home_app_log_forwarder.go +++ b/internal/logging/home_app_log_forwarder.go @@ -25,10 +25,7 @@ Level string `json:"level,omitempty"` Timestamp string `json:"timestamp,omitempty"` RequestID string `json:"request_id,omitempty"` -} - -var currentHomeAppLogClient = func() homeAppLogClient { - return home.Current() + client homeAppLogClient } // HomeAppLogForwarder forwards application logs to Home after the control connection is healthy. @@ -39,9 +36,69 @@ stopOnce sync.Once wg sync.WaitGroup enabled atomic.Bool + stopped atomic.Bool + ownerMu sync.Mutex + owner homeAppLogClient } -// StartHomeAppLogForwarder installs a logrus hook that forwards future application logs to Home. +type homeAppLogMux struct { + mu sync.Mutex + targets map[*HomeAppLogForwarder]struct{} +} + +func (h *homeAppLogMux) Levels() []log.Level { + return log.AllLevels +} + +func (h *homeAppLogMux) Fire(entry *log.Entry) error { + h.mu.Lock() + targets := make([]*HomeAppLogForwarder, 0, len(h.targets)) + for target := range h.targets { + targets = append(targets, target) + } + h.mu.Unlock() + for _, target := range targets { + if errFire := target.Fire(entry); errFire != nil { + return errFire + } + } + return nil +} + +func (h *homeAppLogMux) register(target *HomeAppLogForwarder) { + if target == nil { + return + } + h.mu.Lock() + defer h.mu.Unlock() + if h.targets == nil { + h.targets = make(map[*HomeAppLogForwarder]struct{}) + } + h.targets[target] = struct{}{} +} + +func (h *homeAppLogMux) unregister(target *HomeAppLogForwarder) { + if target == nil { + return + } + h.mu.Lock() + delete(h.targets, target) + h.mu.Unlock() +} + +var ( + homeAppLogMuxHook = &homeAppLogMux{} + homeAppLogMuxInstallOnce sync.Once +) + +func registerHomeAppLogForwarder(forwarder *HomeAppLogForwarder) { + homeAppLogMuxInstallOnce.Do(func() { + log.AddHook(homeAppLogMuxHook) + }) + homeAppLogMuxHook.register(forwarder) +} + +// StartHomeAppLogForwarder registers a Home log forwarding target with the process-wide logrus hook. func StartHomeAppLogForwarder(queueSize int) *HomeAppLogForwarder { if queueSize <= 0 { queueSize = defaultHomeAppLogQueueSize @@ -54,7 +111,7 @@ forwarder.enabled.Store(true) forwarder.wg.Add(1) go forwarder.run() - log.AddHook(forwarder) + registerHomeAppLogForwarder(forwarder) return forwarder } @@ -64,10 +121,55 @@ return } f.stopOnce.Do(func() { + f.stopped.Store(true) + f.ownerMu.Lock() + f.owner = nil + f.ownerMu.Unlock() f.enabled.Store(false) + homeAppLogMuxHook.unregister(f) close(f.stop) f.wg.Wait() }) +} + +// Bind activates forwarding to client. +func (f *HomeAppLogForwarder) Bind(client *home.Client) { + f.bind(client) +} + +func (f *HomeAppLogForwarder) bind(client homeAppLogClient) { + if f == nil || client == nil || f.stopped.Load() { + return + } + f.ownerMu.Lock() + defer f.ownerMu.Unlock() + if f.stopped.Load() { + return + } + f.owner = client + f.enabled.Store(true) +} + +// Deactivate stops forwarding only when client owns the forwarder. +func (f *HomeAppLogForwarder) Deactivate(client *home.Client) { + f.deactivate(client) +} + +func (f *HomeAppLogForwarder) deactivate(client homeAppLogClient) { + if f == nil || client == nil { + return + } + f.ownerMu.Lock() + if f.owner == client { + f.owner = nil + } + f.ownerMu.Unlock() +} + +func (f *HomeAppLogForwarder) client() homeAppLogClient { + f.ownerMu.Lock() + defer f.ownerMu.Unlock() + return f.owner } // Levels implements logrus.Hook. @@ -80,7 +182,7 @@ if f == nil || entry == nil || !f.enabled.Load() { return nil } - client := currentHomeAppLogClient() + client := f.client() if client == nil || !client.HeartbeatOK() { return nil } @@ -94,6 +196,7 @@ Level: entry.Level.String(), Timestamp: entry.Time.Format(time.RFC3339Nano), RequestID: appLogRequestID(entry), + client: client, } select { case f.queue <- payload: @@ -139,11 +242,14 @@ } func (f *HomeAppLogForwarder) forward(payload homeAppLogPayload) { - if !f.enabled.Load() { + client := payload.client + if client == nil { + client = f.client() + } + if !f.enabled.Load() || client == nil || f.client() != client { return } - client := currentHomeAppLogClient() - if client == nil || !client.HeartbeatOK() { + if !client.HeartbeatOK() { return } raw, errMarshal := json.Marshal(&payload) @@ -151,8 +257,17 @@ return } if errPush := client.RPushAppLog(context.Background(), raw); errPush != nil && isHomeAppLogUnsupported(errPush) { - f.enabled.Store(false) + f.disableIfCurrentOwner(client) } +} + +func (f *HomeAppLogForwarder) disableIfCurrentOwner(client homeAppLogClient) { + f.ownerMu.Lock() + defer f.ownerMu.Unlock() + if f.owner != client { + return + } + f.enabled.Store(false) } func isHomeAppLogUnsupported(err error) bool { diff --git a/internal/logging/home_app_log_forwarder_test.go b/internal/logging/home_app_log_forwarder_test.go --- a/internal/logging/home_app_log_forwarder_test.go +++ b/internal/logging/home_app_log_forwarder_test.go @@ -10,6 +10,8 @@ "testing" "time" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" log "github.com/sirupsen/logrus" ) @@ -47,23 +49,15 @@ return bytes.Clone(c.pushed[index]) } -func TestHomeAppLogForwarder_ForwardsFormattedLogWhenHomeHealthy(t *testing.T) { - original := currentHomeAppLogClient - defer func() { - currentHomeAppLogClient = original - }() - +func TestHomeAppLogForwarder_ForwardsFormattedLogWhenBoundOwnerIsHealthy(t *testing.T) { stub := &stubHomeAppLogClient{heartbeatOK: true} - currentHomeAppLogClient = func() homeAppLogClient { - return stub - } - forwarder := &HomeAppLogForwarder{ formatter: &LogFormatter{}, queue: make(chan homeAppLogPayload, 4), stop: make(chan struct{}), } forwarder.enabled.Store(true) + forwarder.bind(stub) forwarder.wg.Add(1) go forwarder.run() defer forwarder.Stop() @@ -107,6 +101,237 @@ } } +func TestHomeAppLogForwarder_StopUnregistersMuxTarget(t *testing.T) { + beforeHooks := homeAppLogForwarderHookCount() + beforeTargets := homeAppLogForwarderTargetCount() + forwarder := StartHomeAppLogForwarder(1) + if got := homeAppLogForwarderHookCount(); got != beforeHooks { + forwarder.Stop() + t.Fatalf("direct Home log forwarder hooks = %d, want %d", got, beforeHooks) + } + if got := homeAppLogForwarderTargetCount(); got != beforeTargets+1 { + forwarder.Stop() + t.Fatalf("Home log forwarder targets = %d, want %d", got, beforeTargets+1) + } + forwarder.Stop() + if got := homeAppLogForwarderTargetCount(); got != beforeTargets { + t.Fatalf("Home log forwarder targets after Stop = %d, want %d", got, beforeTargets) + } +} + +func TestHomeAppLogForwardersUseOneProcessWideMuxHook(t *testing.T) { + first := StartHomeAppLogForwarder(1) + second := StartHomeAppLogForwarder(1) + t.Cleanup(first.Stop) + t.Cleanup(second.Stop) + + if got := homeAppLogForwarderHookCount(); got != 0 { + t.Fatalf("direct Home log forwarder hooks = %d, want 0", got) + } + if got := homeAppLogMuxHookCount(); got != 1 { + t.Fatalf("Home log mux hooks = %d, want 1", got) + } +} + +func homeAppLogForwarderHookCount() int { + count := 0 + for _, hooks := range log.StandardLogger().Hooks { + for _, hook := range hooks { + if _, ok := hook.(*HomeAppLogForwarder); ok { + count++ + } + } + } + return count / len(log.AllLevels) +} + +func homeAppLogMuxHookCount() int { + count := 0 + for _, hooks := range log.StandardLogger().Hooks { + for _, hook := range hooks { + if _, ok := hook.(*homeAppLogMux); ok { + count++ + } + } + } + return count / len(log.AllLevels) +} + +func homeAppLogForwarderTargetCount() int { + homeAppLogMuxHook.mu.Lock() + defer homeAppLogMuxHook.mu.Unlock() + return len(homeAppLogMuxHook.targets) +} + +func TestHomeAppLogForwarder_RebindsOnlyToCurrentOwner(t *testing.T) { + first := &stubHomeAppLogClient{heartbeatOK: true} + second := &stubHomeAppLogClient{heartbeatOK: true} + forwarder := &HomeAppLogForwarder{ + formatter: &LogFormatter{}, + queue: make(chan homeAppLogPayload, 4), + stop: make(chan struct{}), + } + forwarder.enabled.Store(true) + forwarder.wg.Add(1) + go forwarder.run() + t.Cleanup(forwarder.Stop) + + forwarder.bind(first) + if errFire := forwarder.Fire(log.NewEntry(log.StandardLogger())); errFire != nil { + t.Fatalf("Fire() error = %v", errFire) + } + waitForHomeAppLogPush(t, first, 1) + + forwarder.bind(second) + forwarder.deactivate(first) + if errFire := forwarder.Fire(log.NewEntry(log.StandardLogger())); errFire != nil { + t.Fatalf("Fire() error = %v", errFire) + } + waitForHomeAppLogPush(t, second, 1) + if first.pushedCount() != 1 { + t.Fatalf("stale owner received %d records, want 1", first.pushedCount()) + } + + forwarder.deactivate(first) + if errFire := forwarder.Fire(log.NewEntry(log.StandardLogger())); errFire != nil { + t.Fatalf("Fire() error = %v", errFire) + } + waitForHomeAppLogPush(t, second, 2) + + forwarder.deactivate(second) + if errFire := forwarder.Fire(log.NewEntry(log.StandardLogger())); errFire != nil { + t.Fatalf("Fire() error = %v", errFire) + } + time.Sleep(20 * time.Millisecond) + if second.pushedCount() != 2 { + t.Fatalf("detached owner received %d records, want 2", second.pushedCount()) + } +} + +func waitForHomeAppLogPush(t *testing.T, client *stubHomeAppLogClient, want int) { + t.Helper() + deadline := time.Now().Add(time.Second) + for client.pushedCount() < want && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := client.pushedCount(); got != want { + t.Fatalf("pushed records = %d, want %d", got, want) + } +} + +type delayedUnsupportedHomeAppLogClient struct { + started chan struct{} + startedOnce sync.Once + release <-chan struct{} +} + +func (c *delayedUnsupportedHomeAppLogClient) HeartbeatOK() bool { return true } + +func (c *delayedUnsupportedHomeAppLogClient) RPushAppLog(_ context.Context, _ []byte) error { + c.startedOnce.Do(func() { close(c.started) }) + <-c.release + return errors.New("ERR unsupported key") +} + +func TestHomeAppLogForwarder_DelayedOldOwnerUnsupportedDoesNotDisableNewOwner(t *testing.T) { + release := make(chan struct{}) + oldOwner := &delayedUnsupportedHomeAppLogClient{started: make(chan struct{}), release: release} + newOwner := &stubHomeAppLogClient{heartbeatOK: true} + forwarder := &HomeAppLogForwarder{ + formatter: &LogFormatter{}, + queue: make(chan homeAppLogPayload, 1), + stop: make(chan struct{}), + } + forwarder.enabled.Store(true) + forwarder.wg.Add(1) + go forwarder.run() + t.Cleanup(forwarder.Stop) + + forwarder.bind(oldOwner) + forwardDone := make(chan struct{}) + go func() { + forwarder.forward(homeAppLogPayload{Line: "old owner", client: oldOwner}) + close(forwardDone) + }() + + select { + case <-oldOwner.started: + case <-time.After(time.Second): + t.Fatal("old owner did not start forwarding") + } + + forwarder.bind(newOwner) + close(release) + select { + case <-forwardDone: + case <-time.After(time.Second): + t.Fatal("old owner forwarding did not finish") + } + if !forwarder.enabled.Load() { + t.Fatal("old owner unsupported response disabled the new owner") + } + + forwarder.forward(homeAppLogPayload{Line: "new owner", client: newOwner}) + waitForHomeAppLogPush(t, newOwner, 1) +} + +func TestHomeAppLogForwarder_UnboundNeverUsesGlobalFallbackClient(t *testing.T) { + fallback := home.New(internalconfig.HomeConfig{Enabled: true}) + home.SetCurrent(fallback) + t.Cleanup(home.ClearCurrent) + + forwarder := &HomeAppLogForwarder{ + formatter: &LogFormatter{}, + queue: make(chan homeAppLogPayload, 1), + stop: make(chan struct{}), + } + forwarder.enabled.Store(true) + + if client := forwarder.client(); client != nil { + t.Fatalf("unbound client = %v, want nil", client) + } + if errFire := forwarder.Fire(log.NewEntry(log.StandardLogger())); errFire != nil { + t.Fatalf("Fire() error = %v", errFire) + } + if queued := len(forwarder.queue); queued != 0 { + t.Fatalf("unbound queued records = %d, want 0", queued) + } +} + +func TestHomeAppLogForwarder_DropsPreACKAndReconnectGapLogs(t *testing.T) { + oldClient := home.New(internalconfig.HomeConfig{Enabled: true}) + newClient := home.New(internalconfig.HomeConfig{Enabled: true}) + home.SetCurrent(oldClient) + t.Cleanup(home.ClearCurrent) + + preACKForwarder := &HomeAppLogForwarder{ + formatter: &LogFormatter{}, + queue: make(chan homeAppLogPayload, 1), + stop: make(chan struct{}), + } + preACKForwarder.enabled.Store(true) + if client := preACKForwarder.client(); client != nil { + t.Fatalf("pre-ACK client = %v, want nil", client) + } + if errFire := preACKForwarder.Fire(log.NewEntry(log.StandardLogger())); errFire != nil { + t.Fatalf("pre-ACK Fire() error = %v", errFire) + } + + preACKForwarder.bind(oldClient) + preACKForwarder.deactivate(oldClient) + home.SetCurrent(newClient) + if errFire := preACKForwarder.Fire(log.NewEntry(log.StandardLogger())); errFire != nil { + t.Fatalf("reconnect-gap Fire() error = %v", errFire) + } + + if got := len(preACKForwarder.queue); got != 0 { + t.Fatalf("pre-ACK/reconnect-gap queued records = %d, want 0", got) + } + if client := preACKForwarder.client(); client != nil { + t.Fatalf("reconnect-gap client = %v, want nil", client) + } +} + func TestHomeAppLogForwarder_OmitsPlaceholderRequestID(t *testing.T) { entry := log.NewEntry(log.StandardLogger()) entry.Data["request_id"] = "--------" @@ -116,23 +341,15 @@ } } -func TestHomeAppLogForwarder_SkipsWhenHomeHeartbeatIsDown(t *testing.T) { - original := currentHomeAppLogClient - defer func() { - currentHomeAppLogClient = original - }() - +func TestHomeAppLogForwarder_SkipsWhenBoundOwnerHeartbeatIsDown(t *testing.T) { stub := &stubHomeAppLogClient{heartbeatOK: false} - currentHomeAppLogClient = func() homeAppLogClient { - return stub - } - forwarder := &HomeAppLogForwarder{ formatter: &LogFormatter{}, queue: make(chan homeAppLogPayload, 4), stop: make(chan struct{}), } forwarder.enabled.Store(true) + forwarder.bind(stub) entry := log.NewEntry(log.StandardLogger()) entry.Time = time.Now() @@ -147,26 +364,18 @@ } } -func TestHomeAppLogForwarder_DisablesForwardingWhenHomeDoesNotSupportAppLog(t *testing.T) { - original := currentHomeAppLogClient - defer func() { - currentHomeAppLogClient = original - }() - +func TestHomeAppLogForwarder_DisablesForwardingWhenBoundOwnerDoesNotSupportAppLog(t *testing.T) { stub := &stubHomeAppLogClient{ heartbeatOK: true, err: errors.New("ERR unsupported key"), } - currentHomeAppLogClient = func() homeAppLogClient { - return stub - } - forwarder := &HomeAppLogForwarder{ formatter: &LogFormatter{}, queue: make(chan homeAppLogPayload, 4), stop: make(chan struct{}), } forwarder.enabled.Store(true) + forwarder.bind(stub) forwarder.forward(homeAppLogPayload{Line: "legacy home cannot receive app logs"}) if forwarder.enabled.Load() { diff --git a/internal/pluginhost/client_guard.go b/internal/pluginhost/client_guard.go --- a/internal/pluginhost/client_guard.go +++ b/internal/pluginhost/client_guard.go @@ -7,15 +7,16 @@ ) type guardedPluginClient struct { - mu sync.Mutex - cond *sync.Cond - inner pluginClient - calls int - closed bool + mu sync.Mutex + cond *sync.Cond + inner pluginClient + calls int + closed bool + shutdownDone chan struct{} } -func newGuardedPluginClient(inner pluginClient) pluginClient { - client := &guardedPluginClient{inner: inner} +func newGuardedPluginClient(inner pluginClient) *guardedPluginClient { + client := &guardedPluginClient{inner: inner, shutdownDone: make(chan struct{})} client.cond = sync.NewCond(&client.mu) return client } @@ -25,8 +26,35 @@ if errAcquire != nil { return nil, errAcquire } - defer c.release() - return inner.Call(ctx, method, request) + if ctx == nil { + ctx = context.Background() + } + result := make(chan guardedPluginCallResult, 1) + go func() { + defer c.release() + defer func() { + if recovered := recover(); recovered != nil { + result <- guardedPluginCallResult{recovered: recovered} + } + }() + response, errCall := inner.Call(ctx, method, request) + result <- guardedPluginCallResult{response: response, err: errCall} + }() + select { + case callResult := <-result: + if callResult.recovered != nil { + panic(callResult.recovered) + } + return callResult.response, callResult.err + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +type guardedPluginCallResult struct { + response []byte + err error + recovered any } func (c *guardedPluginClient) acquire() (pluginClient, error) { @@ -52,28 +80,49 @@ } func (c *guardedPluginClient) Shutdown() { + c.ShutdownContext(context.Background()) +} + +// ShutdownContext detaches the client immediately and waits for active calls only +// until ctx is canceled. Detached cleanup continues asynchronously when needed. +func (c *guardedPluginClient) ShutdownContext(ctx context.Context) { if c == nil { return } + if ctx == nil { + ctx = context.Background() + } - var inner pluginClient c.mu.Lock() if c.closed { + done := c.shutdownDone + c.mu.Unlock() + select { + case <-done: + case <-ctx.Done(): + } + return + } + c.closed = true + inner := c.inner + c.inner = nil + done := c.shutdownDone + c.mu.Unlock() + + go func() { + c.mu.Lock() for c.calls > 0 { c.cond.Wait() } c.mu.Unlock() - return - } - c.closed = true - for c.calls > 0 { - c.cond.Wait() - } - inner = c.inner - c.inner = nil - c.mu.Unlock() + if inner != nil { + inner.Shutdown() + } + close(done) + }() - if inner != nil { - inner.Shutdown() + select { + case <-done: + case <-ctx.Done(): } } diff --git a/internal/pluginhost/client_guard_test.go b/internal/pluginhost/client_guard_test.go new file mode 100644 --- /dev/null +++ b/internal/pluginhost/client_guard_test.go @@ -0,0 +1,70 @@ +package pluginhost + +import ( + "context" + "sync/atomic" + "testing" + "time" +) + +type blockingGuardPluginClient struct { + started chan struct{} + release chan struct{} + shutdown atomic.Int32 +} + +func (c *blockingGuardPluginClient) Call(context.Context, string, []byte) ([]byte, error) { + close(c.started) + <-c.release + return nil, nil +} + +func (c *blockingGuardPluginClient) Shutdown() { + c.shutdown.Add(1) +} + +func TestGuardedPluginClientShutdownContextDetachesBlockedCall(t *testing.T) { + inner := &blockingGuardPluginClient{started: make(chan struct{}), release: make(chan struct{})} + guarded := newGuardedPluginClient(inner) + + callDone := make(chan struct{}) + go func() { + _, _ = guarded.Call(context.Background(), "blocked", nil) + close(callDone) + }() + select { + case <-inner.started: + case <-time.After(time.Second): + t.Fatal("guarded call did not start") + } + + shutdownCtx, cancelShutdown := context.WithCancel(context.Background()) + cancelShutdown() + shutdownDone := make(chan struct{}) + go func() { + guarded.ShutdownContext(shutdownCtx) + close(shutdownDone) + }() + select { + case <-shutdownDone: + case <-time.After(time.Second): + t.Fatal("context-canceled guarded shutdown waited for the active call") + } + if got := inner.shutdown.Load(); got != 0 { + t.Fatalf("shutdown calls before active call exits = %d, want 0", got) + } + + close(inner.release) + select { + case <-callDone: + case <-time.After(time.Second): + t.Fatal("guarded call did not exit") + } + deadline := time.Now().Add(time.Second) + for inner.shutdown.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := inner.shutdown.Load(); got != 1 { + t.Fatalf("shutdown calls after active call exits = %d, want 1", got) + } +} diff --git a/internal/pluginhost/host.go b/internal/pluginhost/host.go --- a/internal/pluginhost/host.go +++ b/internal/pluginhost/host.go @@ -40,13 +40,25 @@ client pluginClient } +type pluginLoadRequest struct { + result chan pluginLoadResult + cleanupStarted bool +} + +type pluginLoadResult struct { + loaded *loadedPlugin + plugin pluginapi.Plugin + initialized bool + err error +} + type Host struct { - applyMu sync.Mutex + applyMu chan struct{} mu sync.Mutex loader pluginLoader loaded map[string]*loadedPlugin retired map[string][]*loadedPlugin - loading map[string]struct{} + loading map[string]*pluginLoadRequest fused map[string]string pluginFileVersions map[string]string activePluginVersions map[string]string @@ -75,10 +87,11 @@ func New() *Host { h := &Host{ + applyMu: make(chan struct{}, 1), loader: defaultPluginLoader(), loaded: make(map[string]*loadedPlugin), retired: make(map[string][]*loadedPlugin), - loading: make(map[string]struct{}), + loading: make(map[string]*pluginLoadRequest), fused: make(map[string]string), pluginFileVersions: make(map[string]string), activePluginVersions: make(map[string]string), @@ -180,11 +193,16 @@ } func (h *Host) ApplyConfig(ctx context.Context, cfg *config.Config) { - if h == nil { + if h == nil || !h.lockApply(ctx) { return } - h.applyMu.Lock() - defer h.applyMu.Unlock() + defer h.unlockApply() + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return + } rc, errRuntimeConfig := runtimeConfigFromConfig(cfg) if errRuntimeConfig != nil { @@ -247,22 +265,42 @@ loadedNow := false var hotReloadFields log.Fields + var plugin pluginapi.Plugin + registeredNow := false if lp == nil { + request := &pluginLoadRequest{result: make(chan pluginLoadResult, 1)} h.mu.Lock() - h.loading[file.ID] = struct{}{} - h.mu.Unlock() - - loaded, errLoad := h.load(file) - h.mu.Lock() - delete(h.loading, file.ID) - if errLoad != nil { + if _, loading := h.loading[file.ID]; loading { h.mu.Unlock() - log.Warnf("pluginhost: failed to load plugin %s from %s: %v", file.ID, file.Path, errLoad) continue } - // ApplyConfig, UnloadPlugin, and ShutdownAll are serialized by applyMu, - // so a nil read cannot race into a duplicate load. - lp = loaded + h.loading[file.ID] = request + h.mu.Unlock() + h.startPluginLoad(ctx, file, item, request) + + loadResult, completed := h.waitForPluginLoad(ctx, file.ID, request) + if !completed { + return + } + if loadResult.err != nil { + h.cleanupPluginLoad(file.ID, request, loadResult.loaded) + log.Warnf("pluginhost: failed to load plugin %s from %s: %v", file.ID, file.Path, loadResult.err) + continue + } + + h.mu.Lock() + if h.loading[file.ID] != request { + h.mu.Unlock() + h.discardLoadedPlugin(loadResult.loaded) + return + } + if errContext := ctx.Err(); errContext != nil { + h.mu.Unlock() + h.cleanupPluginLoad(file.ID, request, loadResult.loaded) + return + } + delete(h.loading, file.ID) + lp = loadResult.loaded if replaced != nil { hotReloadFields = pluginHotReloadLogFields(file.ID, file.Version, file.Path, replaced.version, replaced.path) h.retireLoadedPluginLocked(replaced) @@ -271,13 +309,21 @@ } h.loaded[file.ID] = lp loadedNow = true + plugin = loadResult.plugin + registeredNow = loadResult.initialized h.mu.Unlock() log.WithFields(pluginLogFields(file.ID, "", file.Version, file.Path)).Info("pluginhost: plugin loaded") } - plugin, okCall := h.callRegister(ctx, lp, item) - if !okCall { - continue + if !registeredNow { + if loadedNow { + continue + } + var okCall bool + plugin, okCall = h.callRegister(ctx, lp, item) + if !okCall { + continue + } } plugin.Metadata = clonePluginMetadata(plugin.Metadata) h.mu.Lock() @@ -325,18 +371,108 @@ } } -func (h *Host) load(file pluginFile) (*loadedPlugin, error) { - client, errOpen := h.loader.Open(file, h) - if errOpen != nil { - return nil, errOpen +func (h *Host) startPluginLoad(ctx context.Context, file pluginFile, item runtimeItemConfig, request *pluginLoadRequest) { + if h == nil || request == nil || request.result == nil { + return } + if ctx == nil { + ctx = context.Background() + } + go func() { + client, errOpen := h.loader.Open(file, h) + if errOpen != nil { + request.result <- pluginLoadResult{err: errOpen} + return + } + if client == nil { + request.result <- pluginLoadResult{err: fmt.Errorf("plugin loader returned nil client")} + return + } + loaded := &loadedPlugin{ + id: file.ID, + path: file.Path, + version: file.Version, + client: newGuardedPluginClient(client), + } + plugin, okCall := h.callRegister(ctx, loaded, item) + request.result <- pluginLoadResult{loaded: loaded, plugin: plugin, initialized: okCall} + }() +} - return &loadedPlugin{ - id: file.ID, - path: file.Path, - version: file.Version, - client: newGuardedPluginClient(client), - }, nil +func (h *Host) waitForPluginLoad(ctx context.Context, id string, request *pluginLoadRequest) (pluginLoadResult, bool) { + if h == nil || request == nil || request.result == nil { + return pluginLoadResult{}, false + } + if ctx == nil { + ctx = context.Background() + } + select { + case result := <-request.result: + return result, true + case <-ctx.Done(): + h.cleanupCanceledPluginLoad(id, request) + return pluginLoadResult{}, false + } +} + +func (h *Host) cleanupCanceledPluginLoad(id string, request *pluginLoadRequest) { + if h == nil || request == nil || request.result == nil { + return + } + h.mu.Lock() + if h.loading[id] != request || request.cleanupStarted { + h.mu.Unlock() + return + } + request.cleanupStarted = true + h.mu.Unlock() + + go func() { + result := <-request.result + h.finishPluginLoadCleanup(id, request, result.loaded) + }() +} + +// cleanupPluginLoad retains the matching load token until the client has physically +// shut down, preventing a replacement ApplyConfig from opening a second client. +func (h *Host) cleanupPluginLoad(id string, request *pluginLoadRequest, loaded *loadedPlugin) { + if h == nil || request == nil { + return + } + h.mu.Lock() + if h.loading[id] != request || request.cleanupStarted { + h.mu.Unlock() + return + } + request.cleanupStarted = true + h.mu.Unlock() + + h.finishPluginLoadCleanup(id, request, loaded) +} + +func (h *Host) finishPluginLoadCleanup(id string, request *pluginLoadRequest, loaded *loadedPlugin) { + go func() { + h.discardLoadedPlugin(loaded) + h.clearLoadingRequest(id, request) + }() +} + +func (h *Host) clearLoadingRequest(id string, request *pluginLoadRequest) { + if h == nil || request == nil { + return + } + h.mu.Lock() + if h.loading[id] == request { + delete(h.loading, id) + } + h.mu.Unlock() +} + +func (h *Host) discardLoadedPlugin(loaded *loadedPlugin) { + if loaded == nil || loaded.client == nil { + return + } + shutdownPluginClient(context.Background(), loaded.client) } func (h *Host) withLoadedPluginFallbacks(files []pluginFile, items map[string]runtimeItemConfig, desired map[string]string) []pluginFile { @@ -381,16 +517,20 @@ // UnloadPlugin removes one plugin from the active runtime and closes its dynamic library. func (h *Host) UnloadPlugin(id string) bool { + return h.UnloadPluginContext(context.Background(), id) +} + +// UnloadPluginContext detaches a plugin from the runtime before waiting for its +// active calls. Physical client cleanup continues after cancellation if needed. +func (h *Host) UnloadPluginContext(ctx context.Context, id string) bool { if h == nil { return false } id = strings.TrimSpace(id) - if id == "" { + if id == "" || !h.lockApply(ctx) { return false } - - h.applyMu.Lock() - defer h.applyMu.Unlock() + defer h.unlockApply() targets := make([]pluginUnloadTarget, 0) h.mu.Lock() @@ -425,7 +565,7 @@ h.RegisterFrontendAuthProviders() for _, target := range targets { if target.client != nil { - target.client.Shutdown() + shutdownPluginClient(ctx, target.client) } log.WithFields(pluginLogFields(target.id, target.name, target.version, target.path)).Info("pluginhost: plugin unloaded") } @@ -434,15 +574,24 @@ // ShutdownAll removes active plugin capabilities and closes all loaded dynamic libraries. func (h *Host) ShutdownAll() { - if h == nil { + h.ShutdownAllContext(context.Background()) +} + +// ShutdownAllContext detaches all plugin runtime state without waiting beyond ctx +// for active plugin calls to complete. +func (h *Host) ShutdownAllContext(ctx context.Context) { + if h == nil || !h.lockApply(ctx) { return } - - h.applyMu.Lock() - defer h.applyMu.Unlock() + defer h.unlockApply() targets := make([]pluginUnloadTarget, 0) + var loading map[string]*pluginLoadRequest h.mu.Lock() + loading = make(map[string]*pluginLoadRequest, len(h.loading)) + for id, request := range h.loading { + loading[id] = request + } for _, lp := range h.loaded { if lp == nil || lp.client == nil { continue @@ -471,7 +620,6 @@ } h.loaded = make(map[string]*loadedPlugin) h.retired = make(map[string][]*loadedPlugin) - h.loading = make(map[string]struct{}) h.modelClientIDs = make(map[string]struct{}) h.executorModelClientIDs = make(map[string]struct{}) h.modelProviders = make(map[string]string) @@ -490,10 +638,48 @@ h.refreshThinkingProviders(nil) h.RegisterFrontendAuthProviders() + for id, request := range loading { + h.cleanupCanceledPluginLoad(id, request) + } for _, target := range targets { - target.client.Shutdown() + shutdownPluginClient(ctx, target.client) log.WithFields(pluginLogFields(target.id, target.name, target.version, target.path)).Info("pluginhost: plugin unloaded") } +} + +func (h *Host) lockApply(ctx context.Context) bool { + if h == nil { + return false + } + if ctx == nil { + ctx = context.Background() + } + select { + case h.applyMu <- struct{}{}: + return true + default: + } + select { + case h.applyMu <- struct{}{}: + return true + case <-ctx.Done(): + return false + } +} + +func (h *Host) unlockApply() { + <-h.applyMu +} + +func shutdownPluginClient(ctx context.Context, client pluginClient) { + if client == nil { + return + } + if guarded, ok := client.(*guardedPluginClient); ok { + guarded.ShutdownContext(ctx) + return + } + client.Shutdown() } func cleanPluginPath(path string) string { diff --git a/internal/pluginhost/host_test.go b/internal/pluginhost/host_test.go --- a/internal/pluginhost/host_test.go +++ b/internal/pluginhost/host_test.go @@ -4,6 +4,7 @@ "bytes" "context" "encoding/json" + "fmt" "net/http" "path/filepath" "strings" @@ -1091,6 +1092,233 @@ } } +func TestHostCanceledInitializationDiscardsBlockedClient(t *testing.T) { + client := &blockingInitializationClient{ + started: make(chan struct{}), + release: make(chan struct{}), + registration: validTestPlugin("alpha"), + } + h := NewForTest(&blockingHostCallLoader{client: client}) + cfg := &config.Config{Plugins: config.PluginsConfig{ + Enabled: true, + Dir: makePluginDir(t, "alpha"), + Configs: enabledPluginConfigs("alpha"), + }} + ctx, cancel := context.WithCancel(context.Background()) + applyDone := make(chan struct{}) + go func() { + h.ApplyConfig(ctx, cfg) + close(applyDone) + }() + waitForHostTestSignal(t, client.started, "plugin initialization") + cancel() + waitForHostTestSignal(t, applyDone, "canceled plugin initialization") + if !h.PluginBusy("alpha") || h.PluginLoaded("alpha") { + t.Fatal("canceled initialization did not retain only its in-flight load token") + } + + close(client.release) + deadline := time.Now().Add(time.Second) + for client.shutdown.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := client.shutdown.Load(); got != 1 { + t.Fatalf("blocked initialization client shutdown calls = %d, want 1", got) + } + if h.PluginBusy("alpha") || h.PluginLoaded("alpha") { + t.Fatal("canceled initialization remained in the host after late cleanup") + } +} + +func TestHostCancellationUnderMutationLockDoesNotInsertLoadedPlugin(t *testing.T) { + client := &blockingInitializationClient{ + started: make(chan struct{}), + release: make(chan struct{}), + completed: make(chan struct{}), + registration: validTestPlugin("alpha"), + } + h := NewForTest(&blockingHostCallLoader{client: client}) + cfg := &config.Config{Plugins: config.PluginsConfig{ + Enabled: true, + Dir: makePluginDir(t, "alpha"), + Configs: enabledPluginConfigs("alpha"), + }} + ctx, cancel := context.WithCancel(context.Background()) + applyDone := make(chan struct{}) + go func() { + h.ApplyConfig(ctx, cfg) + close(applyDone) + }() + waitForHostTestSignal(t, client.started, "plugin initialization") + + h.mu.Lock() + close(client.release) + waitForHostTestSignal(t, client.completed, "plugin initialization completion") + cancel() + h.mu.Unlock() + waitForHostTestSignal(t, applyDone, "canceled plugin apply") + + deadline := time.Now().Add(time.Second) + for client.shutdown.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := client.shutdown.Load(); got != 1 { + t.Fatalf("late client shutdown calls = %d, want 1", got) + } + if h.PluginLoaded("alpha") || h.PluginBusy("alpha") { + t.Fatal("canceled load inserted or retained a completed plugin") + } +} + +func TestHostCanceledLoadDiscardsLateClientWithoutReplacingCurrentPlugin(t *testing.T) { + first := &lateLoadClient{registration: validTestPlugin("alpha")} + second := &lateLoadClient{registration: validTestPlugin("alpha")} + loader := &lateLoadPluginLoader{ + first: first, + second: second, + firstStarted: make(chan struct{}), + firstRelease: make(chan struct{}), + secondStarted: make(chan struct{}), + } + h := NewForTest(loader) + cfg := &config.Config{Plugins: config.PluginsConfig{ + Enabled: true, + Dir: makePluginDir(t, "alpha"), + Configs: enabledPluginConfigs("alpha"), + }} + ctx, cancel := context.WithCancel(context.Background()) + firstDone := make(chan struct{}) + go func() { + h.ApplyConfig(ctx, cfg) + close(firstDone) + }() + waitForHostTestSignal(t, loader.firstStarted, "first plugin load") + cancel() + waitForHostTestSignal(t, firstDone, "canceled plugin load") + if !h.PluginBusy("alpha") { + t.Fatal("PluginBusy(alpha) = false after canceled load, want retained load token") + } + + secondDone := make(chan struct{}) + go func() { + h.ApplyConfig(context.Background(), cfg) + close(secondDone) + }() + waitForHostTestSignal(t, secondDone, "replacement apply completion") + if got := loader.calls.Load(); got != 1 { + t.Fatalf("Open calls = %d, want 1 while canceled load is still blocked", got) + } + select { + case <-loader.secondStarted: + t.Fatal("replacement started a second load before the canceled load completed") + default: + } + + close(loader.firstRelease) + deadline := time.Now().Add(time.Second) + for first.shutdown.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := first.shutdown.Load(); got != 1 { + t.Fatalf("late client shutdown calls = %d, want 1", got) + } + if h.PluginBusy("alpha") || h.PluginLoaded("alpha") { + t.Fatal("late canceled client remained in the host") + } + h.ShutdownAll() +} + +func TestHostCanceledBlockedLoadKeepsOneLoaderAndCleanupPerPlugin(t *testing.T) { + first := &lateLoadClient{registration: validTestPlugin("alpha")} + loader := &lateLoadPluginLoader{ + first: first, + second: &lateLoadClient{registration: validTestPlugin("alpha")}, + firstStarted: make(chan struct{}), + firstRelease: make(chan struct{}), + secondStarted: make(chan struct{}), + } + h := NewForTest(loader) + cfg := &config.Config{Plugins: config.PluginsConfig{ + Enabled: true, + Dir: makePluginDir(t, "alpha"), + Configs: enabledPluginConfigs("alpha"), + }} + + ctx, cancel := context.WithCancel(context.Background()) + firstDone := make(chan struct{}) + go func() { + h.ApplyConfig(ctx, cfg) + close(firstDone) + }() + waitForHostTestSignal(t, loader.firstStarted, "first plugin load") + cancel() + waitForHostTestSignal(t, firstDone, "canceled plugin load") + + for range 8 { + h.ApplyConfig(context.Background(), cfg) + } + if got := loader.calls.Load(); got != 1 { + t.Fatalf("Open calls = %d, want one blocked loader", got) + } + + close(loader.firstRelease) + deadline := time.Now().Add(time.Second) + for first.shutdown.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := first.shutdown.Load(); got != 1 { + t.Fatalf("late client shutdown calls = %d, want one cleanup", got) + } + if h.PluginBusy("alpha") { + t.Fatal("PluginBusy(alpha) = true after blocked load cleanup") + } +} + +func TestHostUnloadPluginContextDetachesBlockedCall(t *testing.T) { + plugin := validTestPlugin("alpha") + client := &blockingHostCallClient{started: make(chan struct{}), release: make(chan struct{}), registration: plugin} + loader := &blockingHostCallLoader{client: client} + h := NewForTest(loader) + cfg := &config.Config{Plugins: config.PluginsConfig{ + Enabled: true, + Dir: makePluginDir(t, "alpha"), + Configs: enabledPluginConfigs("alpha"), + }} + h.ApplyConfig(context.Background(), cfg) + + h.mu.Lock() + loaded := h.loaded["alpha"] + h.mu.Unlock() + if loaded == nil { + t.Fatal("plugin did not load") + } + go func() { _, _ = loaded.client.Call(context.Background(), pluginabi.MethodUsageHandle, nil) }() + waitForHostTestSignal(t, client.started, "blocked plugin call") + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + unloadDone := make(chan bool, 1) + go func() { unloadDone <- h.UnloadPluginContext(ctx, "alpha") }() + if ok := waitForHostTestBool(t, unloadDone, "contextual unload"); !ok { + t.Fatal("UnloadPluginContext() = false, want true after detaching runtime") + } + if h.PluginBusy("alpha") { + t.Fatal("PluginBusy(alpha) = true after contextual unload detached runtime") + } + if got := client.shutdown.Load(); got != 0 { + t.Fatalf("shutdown calls before blocked plugin call exits = %d, want 0", got) + } + + close(client.release) + deadline := time.Now().Add(time.Second) + for client.shutdown.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := client.shutdown.Load(); got != 1 { + t.Fatalf("shutdown calls after blocked plugin call exits = %d, want 1", got) + } +} + func TestHostUnloadWaitsForBlockingLoad(t *testing.T) { h, cfg, openStarted, releaseOpen := newBlockingOpenHost(t) applyDone := make(chan struct{}) @@ -1214,6 +1442,117 @@ func (c *capturePluginClient) Shutdown() {} +type blockingInitializationClient struct { + started chan struct{} + release chan struct{} + completed chan struct{} + registration pluginapi.Plugin + shutdown atomic.Int32 + shutdownStarted chan struct{} + shutdownRelease chan struct{} +} + +func (c *blockingInitializationClient) Call(_ context.Context, method string, _ []byte) ([]byte, error) { + if method != pluginabi.MethodPluginRegister { + return nil, fmt.Errorf("unexpected plugin method %s", method) + } + close(c.started) + <-c.release + if c.completed != nil { + close(c.completed) + } + return marshalRPCResult(rpcRegistration{ + SchemaVersion: pluginabi.SchemaVersion, + Metadata: c.registration.Metadata, + Capabilities: rpcCapabilitiesFromPlugin(c.registration), + }) +} + +func (c *blockingInitializationClient) Shutdown() { + c.shutdown.Add(1) + if c.shutdownStarted != nil { + close(c.shutdownStarted) + } + if c.shutdownRelease != nil { + <-c.shutdownRelease + } +} + +type lateLoadPluginLoader struct { + first pluginClient + second pluginClient + firstStarted chan struct{} + firstRelease chan struct{} + secondStarted chan struct{} + calls atomic.Int32 +} + +func (l *lateLoadPluginLoader) Open(pluginFile, *Host) (pluginClient, error) { + if l.calls.Add(1) == 1 { + close(l.firstStarted) + <-l.firstRelease + return l.first, nil + } + close(l.secondStarted) + return l.second, nil +} + +type lateLoadClient struct { + registration pluginapi.Plugin + shutdown atomic.Int32 +} + +func (c *lateLoadClient) Call(_ context.Context, method string, _ []byte) ([]byte, error) { + if method != pluginabi.MethodPluginRegister { + return nil, fmt.Errorf("unexpected plugin method %s", method) + } + return marshalRPCResult(rpcRegistration{ + SchemaVersion: pluginabi.SchemaVersion, + Metadata: c.registration.Metadata, + Capabilities: rpcCapabilitiesFromPlugin(c.registration), + }) +} + +func (c *lateLoadClient) Shutdown() { + c.shutdown.Add(1) +} + +type blockingHostCallLoader struct { + client pluginClient +} + +func (l *blockingHostCallLoader) Open(pluginFile, *Host) (pluginClient, error) { + return l.client, nil +} + +type blockingHostCallClient struct { + started chan struct{} + release chan struct{} + registration pluginapi.Plugin + shutdown atomic.Int32 +} + +func (c *blockingHostCallClient) Call(_ context.Context, method string, _ []byte) ([]byte, error) { + switch method { + case pluginabi.MethodPluginRegister: + return marshalRPCResult(rpcRegistration{ + SchemaVersion: pluginabi.SchemaVersion, + Metadata: c.registration.Metadata, + Capabilities: rpcCapabilitiesFromPlugin(c.registration), + }) + case pluginabi.MethodUsageHandle: + close(c.started) + <-c.release + return marshalRPCResult(rpcEmptyResponse{}) + default: + return nil, fmt.Errorf("unexpected plugin method %s", method) + } +} + +func (c *blockingHostCallClient) Shutdown() { + c.shutdown.Add(1) +} + type blockingOpenLoader struct { inner *testSymbolLoader started chan struct{} @@ -1308,5 +1647,120 @@ case <-time.After(time.Second): t.Fatalf("timed out waiting for %s", name) return false + } +} + +type countingPluginLoader struct { + client pluginClient + replacement pluginClient + calls atomic.Int32 +} + +func (l *countingPluginLoader) Open(pluginFile, *Host) (pluginClient, error) { + if l.calls.Add(1) == 1 { + return l.client, nil + } + return l.replacement, nil +} + +func TestHostShutdownAllRetainsBlockedLoadTokenUntilCleanup(t *testing.T) { + client := &blockingInitializationClient{ + started: make(chan struct{}), + release: make(chan struct{}), + registration: validTestPlugin("alpha"), + shutdownStarted: make(chan struct{}), + shutdownRelease: make(chan struct{}), + } + loader := &countingPluginLoader{client: client, replacement: &lateLoadClient{registration: validTestPlugin("alpha")}} + h := NewForTest(loader) + cfg := &config.Config{Plugins: config.PluginsConfig{ + Enabled: true, + Dir: makePluginDir(t, "alpha"), + Configs: enabledPluginConfigs("alpha"), + }} + + ctx, cancel := context.WithCancel(context.Background()) + firstDone := make(chan struct{}) + go func() { + h.ApplyConfig(ctx, cfg) + close(firstDone) + }() + waitForHostTestSignal(t, client.started, "plugin registration") + cancel() + waitForHostTestSignal(t, firstDone, "canceled plugin apply") + close(client.release) + waitForHostTestSignal(t, client.shutdownStarted, "plugin shutdown") + + h.ShutdownAllContext(context.Background()) + var applies sync.WaitGroup + for range 8 { + applies.Add(1) + go func() { + defer applies.Done() + h.ApplyConfig(context.Background(), cfg) + }() + } + applies.Wait() + if got := loader.calls.Load(); got != 1 { + t.Fatalf("Open calls while ShutdownAll cleanup is blocked = %d, want 1", got) + } + if !h.PluginBusy("alpha") { + t.Fatal("PluginBusy(alpha) = false before physical shutdown returns") + } + + close(client.shutdownRelease) + deadline := time.Now().Add(time.Second) + for h.PluginBusy("alpha") && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if h.PluginBusy("alpha") { + t.Fatal("PluginBusy(alpha) = true after physical shutdown returned") + } +} + +func TestHostCanceledRegisterRetainsLoadTokenUntilShutdownReturns(t *testing.T) { + client := &blockingInitializationClient{ + started: make(chan struct{}), + release: make(chan struct{}), + registration: validTestPlugin("alpha"), + shutdownStarted: make(chan struct{}), + shutdownRelease: make(chan struct{}), + } + loader := &countingPluginLoader{client: client, replacement: &lateLoadClient{registration: validTestPlugin("alpha")}} + h := NewForTest(loader) + cfg := &config.Config{Plugins: config.PluginsConfig{ + Enabled: true, + Dir: makePluginDir(t, "alpha"), + Configs: enabledPluginConfigs("alpha"), + }} + ctx, cancel := context.WithCancel(context.Background()) + applyDone := make(chan struct{}) + go func() { + h.ApplyConfig(ctx, cfg) + close(applyDone) + }() + waitForHostTestSignal(t, client.started, "plugin registration") + cancel() + waitForHostTestSignal(t, applyDone, "canceled plugin apply") + close(client.release) + waitForHostTestSignal(t, client.shutdownStarted, "plugin shutdown") + + for range 8 { + h.ApplyConfig(context.Background(), cfg) + } + if got := loader.calls.Load(); got != 1 { + t.Fatalf("Open calls while shutdown is blocked = %d, want 1", got) + } + if !h.PluginBusy("alpha") { + t.Fatal("PluginBusy(alpha) = false before physical shutdown returns") + } + + close(client.shutdownRelease) + deadline := time.Now().Add(time.Second) + for h.PluginBusy("alpha") && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if h.PluginBusy("alpha") { + t.Fatal("PluginBusy(alpha) = true after physical shutdown returned") } } diff --git a/internal/watcher/watcher_test.go b/internal/watcher/watcher_test.go --- a/internal/watcher/watcher_test.go +++ b/internal/watcher/watcher_test.go @@ -181,8 +181,9 @@ configPath := filepath.Join(tmpDir, "config.yaml") writeConfig := func(port int, allowRemote bool) { cfg := &config.Config{ - Port: port, - AuthDir: authDir, + Port: port, + AuthDir: authDir, + CredentialInFlight: config.DefaultCredentialInFlightConfig(), RemoteManagement: config.RemoteManagement{ AllowRemote: allowRemote, }, @@ -1356,13 +1357,15 @@ } oldCfg := &config.Config{ - AuthDir: authDir, + AuthDir: authDir, + CredentialInFlight: config.DefaultCredentialInFlightConfig(), OAuthExcludedModels: map[string][]string{ "provider-a": {"m1"}, }, } newCfg := &config.Config{ - AuthDir: authDir, + AuthDir: authDir, + CredentialInFlight: config.DefaultCredentialInFlightConfig(), OAuthExcludedModels: map[string][]string{ "provider-a": {"m2"}, }, @@ -1418,12 +1421,14 @@ oldCfg := &config.Config{ AuthDir: authDir, + CredentialInFlight: config.DefaultCredentialInFlightConfig(), MaxRetryCredentials: 0, RequestRetry: 1, MaxRetryInterval: 5, } newCfg := &config.Config{ AuthDir: authDir, + CredentialInFlight: config.DefaultCredentialInFlightConfig(), MaxRetryCredentials: 2, RequestRetry: 1, MaxRetryInterval: 5, diff --git a/sdk/cliproxy/builder.go b/sdk/cliproxy/builder.go --- a/sdk/cliproxy/builder.go +++ b/sdk/cliproxy/builder.go @@ -6,8 +6,6 @@ import ( "context" "fmt" - "strings" - "time" configaccess "github.com/router-for-me/CLIProxyAPI/v7/internal/access/config_access" "github.com/router-for-me/CLIProxyAPI/v7/internal/api" @@ -229,42 +227,16 @@ accessManager.SetProviders(sdkaccess.RegisteredProviders()) coreManager := b.coreManager + var appliedRoutingState *routingRuntimeState if coreManager == nil { tokenStore := sdkAuth.GetTokenStore() if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok && b.cfg != nil { dirSetter.SetBaseDir(b.cfg.AuthDir) } - strategy := "" - sessionAffinity := false - sessionAffinityTTL := time.Hour - if b.cfg != nil { - strategy = strings.ToLower(strings.TrimSpace(b.cfg.Routing.Strategy)) - // Support both legacy ClaudeCodeSessionAffinity and new universal SessionAffinity - sessionAffinity = b.cfg.Routing.SessionAffinity - if ttlStr := strings.TrimSpace(b.cfg.Routing.SessionAffinityTTL); ttlStr != "" { - if parsed, err := time.ParseDuration(ttlStr); err == nil && parsed > 0 { - sessionAffinityTTL = parsed - } - } - } - var selector coreauth.Selector - switch strategy { - case "fill-first", "fillfirst", "ff": - selector = &coreauth.FillFirstSelector{} - default: - selector = &coreauth.RoundRobinSelector{} - } - - // Wrap with session affinity if enabled (failover is always on) - if sessionAffinity { - selector = coreauth.NewSessionAffinitySelectorWithConfig(coreauth.SessionAffinityConfig{ - Fallback: selector, - TTL: sessionAffinityTTL, - }) - } - - coreManager = coreauth.NewManager(tokenStore, selector, nil) + routingState := normalizedRoutingRuntimeState(b.cfg) + coreManager = coreauth.NewManager(tokenStore, newRoutingSelector(routingState), nil) + appliedRoutingState = &routingState } // Attach a default RoundTripper provider so providers can opt-in per-auth transports. coreManager.SetRoundTripperProvider(newDefaultRoundTripperProvider()) @@ -275,17 +247,18 @@ } service := &Service{ - cfg: b.cfg, - configPath: b.configPath, - tokenProvider: tokenProvider, - apiKeyProvider: apiKeyProvider, - watcherFactory: watcherFactory, - hooks: b.hooks, - authManager: authManager, - accessManager: accessManager, - coreManager: coreManager, - pluginHost: pluginHost, - serverOptions: append([]api.ServerOption(nil), b.serverOptions...), + cfg: b.cfg, + configPath: b.configPath, + tokenProvider: tokenProvider, + apiKeyProvider: apiKeyProvider, + watcherFactory: watcherFactory, + hooks: b.hooks, + authManager: authManager, + accessManager: accessManager, + coreManager: coreManager, + pluginHost: pluginHost, + appliedRoutingState: appliedRoutingState, + serverOptions: append([]api.ServerOption(nil), b.serverOptions...), } if b.postAuthHook != nil { service.serverOptions = append(service.serverOptions, api.WithPostAuthHook(b.postAuthHook)) diff --git a/sdk/cliproxy/home_plugins.go b/sdk/cliproxy/home_plugins.go --- a/sdk/cliproxy/home_plugins.go +++ b/sdk/cliproxy/home_plugins.go @@ -21,7 +21,34 @@ const homePluginStatusReportTimeout = 10 * time.Second +type homePluginStatusWork struct { + cfg *config.Config + report homeplugins.SyncReport +} + +type homePluginTaskWork struct { + cfg *config.Config + task home.PluginTask + report *homeplugins.SyncReport +} + +type homePluginFinalization struct { + config *config.Config + configCommit configCommit + committed bool + statusWork []homePluginStatusWork + nextStatus int + taskWork []homePluginTaskWork + nextTask int + syncKey string + markSynced bool +} + func (s *Service) syncHomePlugins(ctx context.Context, cfg *config.Config) (homeplugins.SyncReport, string, bool, error) { + return s.syncHomePluginsWithClient(ctx, cfg, nil) +} + +func (s *Service) syncHomePluginsWithClient(ctx context.Context, cfg *config.Config, client *home.Client) (homeplugins.SyncReport, string, bool, error) { if s == nil || cfg == nil || !cfg.Home.Enabled { return homeplugins.SyncReport{}, "", false, nil } @@ -49,7 +76,7 @@ InstalledVersions: installedVersions, } defer request.Clear() - response, errFetch := s.fetchHomePluginSync(ctx, request) + response, errFetch := s.fetchHomePluginSyncWithClient(ctx, client, request) if errors.Is(errFetch, home.ErrPluginSyncUnsupported) { response.Clear() report, errSync := homeplugins.SyncWithReport(ctx, cfg, s.pluginHost) @@ -63,14 +90,19 @@ return report, syncKey, true, errSync } -func (s *Service) fetchHomePluginSync(ctx context.Context, request sdkpluginstore.PluginSyncRequest) (sdkpluginstore.PluginSyncResponse, error) { +func (s *Service) fetchHomePluginSyncWithClient(ctx context.Context, client *home.Client, request sdkpluginstore.PluginSyncRequest) (sdkpluginstore.PluginSyncResponse, error) { if s.homePluginSyncFetch != nil { return s.homePluginSyncFetch(ctx, request) } - if s.homeClient == nil { + if client == nil { + s.homeMu.Lock() + client = s.homeClient + s.homeMu.Unlock() + } + if client == nil { return sdkpluginstore.PluginSyncResponse{}, fmt.Errorf("home client is unavailable") } - return s.homeClient.GetPluginSync(ctx, request) + return client.GetPluginSync(ctx, request) } func (s *Service) markHomePluginsSynced(syncKey string) { @@ -83,60 +115,139 @@ } func (s *Service) reportHomePluginStatus(ctx context.Context, cfg *config.Config, report homeplugins.SyncReport) { - if s == nil || cfg == nil { - return + s.reportHomePluginStatusWithClient(ctx, cfg, report, nil) +} + +func (s *Service) reportHomePluginStatusWithClient(ctx context.Context, cfg *config.Config, report homeplugins.SyncReport, client *home.Client) { + if errReport := s.pushHomePluginStatusWithClient(ctx, cfg, report, client); errReport != nil { + log.Warnf("failed to report home plugin status: %v", errReport) } - if s.homeClient == nil { - log.Warn("failed to report home plugin status: home client is unavailable") - return +} + +func (s *Service) pushHomePluginStatusWithClient(ctx context.Context, cfg *config.Config, report homeplugins.SyncReport, client *home.Client) error { + if s == nil || cfg == nil { + return nil + } + if client == nil { + s.homeMu.Lock() + client = s.homeClient + s.homeMu.Unlock() + } + if client == nil { + return fmt.Errorf("home client is unavailable") } nodeID := strings.TrimSpace(cfg.Home.NodeID) if nodeID == "" { - log.Warn("failed to report home plugin status: node id is empty") - return + return fmt.Errorf("home node id is empty") } report.NodeID = nodeID report.UpdatedAt = time.Now().UTC() raw, errMarshal := json.Marshal(report) if errMarshal != nil { - log.Warnf("failed to marshal home plugin status: %v", errMarshal) - return + return fmt.Errorf("marshal home plugin status: %w", errMarshal) } if ctx == nil { ctx = context.Background() } reportCtx, cancel := context.WithTimeout(ctx, homePluginStatusReportTimeout) defer cancel() - if errReport := s.homeClient.RPushPluginStatus(reportCtx, raw); errReport != nil { - log.Warnf("failed to report home plugin status: %v", errReport) + if errReport := client.RPushPluginStatus(reportCtx, raw); errReport != nil { + return fmt.Errorf("push home plugin status: %w", errReport) } + return nil } func (s *Service) processHomePluginTasks(ctx context.Context, cfg *config.Config) { - if s == nil || cfg == nil || !cfg.Home.Enabled || s.homeClient == nil { + s.processHomePluginTasksWithClient(ctx, cfg, nil) +} + +func (s *Service) processHomePluginTasksWithClient(ctx context.Context, cfg *config.Config, client *home.Client) { + tasks, errStage := s.stageHomePluginTasksWithClient(ctx, cfg, client) + if errStage != nil { + log.Warnf("failed to fetch home plugin tasks: %v", errStage) return + } + work := &homePluginFinalization{taskWork: tasks} + if errFinalize := s.finalizeHomePluginWork(ctx, client, work); errFinalize != nil { + log.Warnf("failed to finalize home plugin tasks: %v", errFinalize) + } +} + +func (s *Service) stageHomePluginTasksWithClient(ctx context.Context, cfg *config.Config, client *home.Client) ([]homePluginTaskWork, error) { + if s == nil || cfg == nil || !cfg.Home.Enabled { + return nil, nil + } + if client == nil { + s.homeMu.Lock() + client = s.homeClient + s.homeMu.Unlock() + } + if client == nil { + return nil, fmt.Errorf("home client is unavailable") } if ctx == nil { ctx = context.Background() } - tasks, errTasks := s.homeClient.GetPluginTasks(ctx) + tasks, errTasks := client.GetPluginTasks(ctx) if errTasks != nil { - log.Warnf("failed to fetch home plugin tasks: %v", errTasks) - return + return nil, errTasks } + staged := make([]homePluginTaskWork, 0, len(tasks)) for _, task := range tasks { if !strings.EqualFold(strings.TrimSpace(task.Operation), "delete") { continue } - report := s.processHomePluginDeleteTask(ctx, cfg, task) - if !report.OK && strings.TrimSpace(report.Error) != "" { - log.Warnf("failed to process home plugin delete task %d for %s: %v", task.ID, task.PluginID, report.Error) - } - s.reportHomePluginStatus(ctx, cfg, report) + staged = append(staged, homePluginTaskWork{cfg: cfg, task: task}) } + return staged, nil +} + +func (s *Service) finalizeHomePluginWork(ctx context.Context, client *home.Client, work *homePluginFinalization) error { + if work == nil { + return nil + } + if ctx != nil { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + } + for work.nextStatus < len(work.statusWork) { + status := work.statusWork[work.nextStatus] + if errReport := s.pushHomePluginStatusWithClient(ctx, status.cfg, status.report, client); errReport != nil { + return errReport + } + work.nextStatus++ + } + for work.nextTask < len(work.taskWork) { + taskWork := &work.taskWork[work.nextTask] + if taskWork.report == nil { + report := s.processHomePluginDeleteTask(ctx, taskWork.cfg, taskWork.task) + taskWork.report = &report + if !report.OK && strings.TrimSpace(report.Error) != "" { + log.Warnf("failed to process home plugin delete task %d for %s: %v", taskWork.task.ID, taskWork.task.PluginID, report.Error) + } + } + if errReport := s.pushHomePluginStatusWithClient(ctx, taskWork.cfg, *taskWork.report, client); errReport != nil { + return errReport + } + work.nextTask++ + } + if work.markSynced { + if ctx != nil { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + } + s.markHomePluginsSynced(work.syncKey) + work.markSynced = false + } + return nil } func (s *Service) processHomePluginDeleteTask(ctx context.Context, cfg *config.Config, task home.PluginTask) homeplugins.SyncReport { + if s != nil && s.homePluginDeleteTask != nil { + return s.homePluginDeleteTask(ctx, cfg, task) + } return homeplugins.DeleteWithReport(ctx, cfg, s.pluginHost, task.ID, task.PluginID) } diff --git a/sdk/cliproxy/home_plugins_test.go b/sdk/cliproxy/home_plugins_test.go --- a/sdk/cliproxy/home_plugins_test.go +++ b/sdk/cliproxy/home_plugins_test.go @@ -1,13 +1,21 @@ package cliproxy import ( + "bufio" "context" + "encoding/json" "errors" + "io" + "net" + "strconv" + "strings" + "sync/atomic" "testing" "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/homeplugins" sdkpluginstore "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginstore" "gopkg.in/yaml.v3" ) @@ -141,7 +149,7 @@ } } -func TestApplyHomeOverlayWarnsOnRuntimePluginSyncFailure(t *testing.T) { +func TestApplyHomeOverlayReturnsRuntimePluginSyncFailureWithoutApplyingConfig(t *testing.T) { base := &config.Config{} base.Home.Enabled = true base.Plugins.Enabled = true @@ -171,11 +179,11 @@ }, } - if errApply := service.applyHomeOverlayContext(context.Background(), remote); errApply != nil { - t.Fatalf("applyHomeOverlayContext() error = %v, want warning-only plugin sync failure", errApply) + if errApply := service.applyHomeOverlayContext(context.Background(), remote); errApply == nil { + t.Fatal("applyHomeOverlayContext() error = nil, want plugin sync failure") } - if service.cfg == nil || !service.cfg.Home.Enabled || !service.cfg.Plugins.Enabled { - t.Fatalf("service cfg = %+v, want applied home config despite plugin sync failure", service.cfg) + if service.cfg == nil || !service.cfg.Home.Enabled || len(service.cfg.Plugins.Configs) != 0 { + t.Fatalf("service cfg = %+v, want unchanged config after plugin sync failure", service.cfg) } if service.homePluginSyncKey != "" { t.Fatalf("homePluginSyncKey = %q, want empty after plugin sync failure", service.homePluginSyncKey) @@ -207,6 +215,454 @@ if service.homePluginSyncKey != "" { t.Fatalf("homePluginSyncKey = %q, want empty before a successful plugin sync", service.homePluginSyncKey) } +} + +func TestFinalizeHomePluginWorkRetriesFailedStatusWithoutMarkingSynced(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + var writes atomic.Int32 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go func(conn net.Conn) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "RPUSH") && args[1] == "plugin-status": + if writes.Add(1) == 1 { + if _, errWrite := io.WriteString(conn, "-ERR blocked\r\n"); errWrite != nil { + return + } + continue + } + if _, errWrite := io.WriteString(conn, ":1\r\n"); errWrite != nil { + return + } + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } + }(conn) + } + }() + t.Cleanup(func() { + _ = listener.Close() + <-serverDone + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + client := home.New(config.HomeConfig{Enabled: true, Host: host, Port: port}) + t.Cleanup(client.Close) + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.NodeID = "node-1" + service := &Service{} + work := &homePluginFinalization{ + statusWork: []homePluginStatusWork{{cfg: cfg, report: homeplugins.CompletedSyncReport(homeplugins.CurrentPlatform(), nil)}}, + syncKey: "sync-key", + markSynced: true, + } + if errFinalize := service.finalizeHomePluginWork(context.Background(), client, work); errFinalize == nil { + t.Fatal("first plugin status finalization succeeded, want Home rejection") + } + if service.homePluginSyncKey != "" || work.nextStatus != 0 || !work.markSynced { + t.Fatalf("failed finalization marked or advanced work: key=%q next=%d marked=%v", service.homePluginSyncKey, work.nextStatus, work.markSynced) + } + if errFinalize := service.finalizeHomePluginWork(context.Background(), client, work); errFinalize != nil { + t.Fatalf("retry finalization error = %v", errFinalize) + } + if service.homePluginSyncKey != "sync-key" || work.nextStatus != 1 || work.markSynced { + t.Fatalf("successful finalization state: key=%q next=%d marked=%v", service.homePluginSyncKey, work.nextStatus, work.markSynced) + } + if errFinalize := service.finalizeHomePluginWork(context.Background(), client, work); errFinalize != nil { + t.Fatalf("duplicate finalization error = %v", errFinalize) + } + if got := writes.Load(); got != 2 { + t.Fatalf("plugin status writes = %d, want one failed write and one successful retry", got) + } +} + +func TestStageHomePluginTasksDefersDeleteUntilFinalization(t *testing.T) { + client, _ := newHomePluginTaskTestClient(t, []home.PluginTask{{ID: 7, Operation: "delete", PluginID: "plugin-a"}}, 0) + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.NodeID = "node-1" + var deletes atomic.Int32 + service := &Service{homePluginDeleteTask: func(_ context.Context, _ *config.Config, task home.PluginTask) homeplugins.SyncReport { + deletes.Add(1) + return homeplugins.DeleteWithReport(context.Background(), nil, nil, task.ID, task.PluginID) + }} + + taskWork, errStage := service.stageHomePluginTasksWithClient(context.Background(), cfg, client) + if errStage != nil { + t.Fatalf("stageHomePluginTasksWithClient() error = %v", errStage) + } + if got := deletes.Load(); got != 0 { + t.Fatalf("staged plugin deletes = %d, want 0 before controlled finalization", got) + } + if len(taskWork) != 1 || taskWork[0].task.ID != 7 { + t.Fatalf("staged task work = %#v, want delete task 7", taskWork) + } + + if errFinalize := service.finalizeHomePluginWork(context.Background(), client, &homePluginFinalization{taskWork: taskWork}); errFinalize != nil { + t.Fatalf("finalizeHomePluginWork() error = %v", errFinalize) + } + if got := deletes.Load(); got != 1 { + t.Fatalf("finalized plugin deletes = %d, want 1", got) + } +} + +func TestFinalizeHomePluginTaskStatusRetryDoesNotRepeatDelete(t *testing.T) { + client, writes := newHomePluginTaskTestClient(t, nil, 1) + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.NodeID = "node-1" + var deletes atomic.Int32 + service := &Service{homePluginDeleteTask: func(_ context.Context, _ *config.Config, task home.PluginTask) homeplugins.SyncReport { + deletes.Add(1) + return homeplugins.DeleteWithReport(context.Background(), nil, nil, task.ID, task.PluginID) + }} + work := &homePluginFinalization{taskWork: []homePluginTaskWork{{cfg: cfg, task: home.PluginTask{ID: 8, Operation: "delete", PluginID: "plugin-b"}}}} + + if errFinalize := service.finalizeHomePluginWork(context.Background(), client, work); errFinalize == nil { + t.Fatal("first task report finalization succeeded, want Home rejection") + } + if got := deletes.Load(); got != 1 { + t.Fatalf("first finalization deletes = %d, want 1", got) + } + if work.nextTask != 0 || work.taskWork[0].report == nil { + t.Fatalf("failed task status did not retain action result: next=%d report=%#v", work.nextTask, work.taskWork[0].report) + } + if errFinalize := service.finalizeHomePluginWork(context.Background(), client, work); errFinalize != nil { + t.Fatalf("retry finalization error = %v", errFinalize) + } + if got := deletes.Load(); got != 1 { + t.Fatalf("retried finalization deletes = %d, want 1", got) + } + if work.nextTask != 1 { + t.Fatalf("task finalization next = %d, want 1", work.nextTask) + } + if gotWrites := writes.Load(); gotWrites != 2 { + t.Fatalf("task status writes = %d, want 2", gotWrites) + } +} + +func newHomePluginTaskTestClient(t *testing.T, tasks []home.PluginTask, failStatuses int32) (*home.Client, *atomic.Int32) { + t.Helper() + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + rawTasks, errMarshal := json.Marshal(tasks) + if errMarshal != nil { + t.Fatalf("marshal tasks: %v", errMarshal) + } + var writes atomic.Int32 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go func(conn net.Conn) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + _, _ = io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n") + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + _, _ = io.WriteString(conn, "$"+strconv.Itoa(len(rawTasks))+"\r\n") + _, _ = conn.Write(rawTasks) + _, _ = io.WriteString(conn, "\r\n") + case len(args) >= 2 && strings.EqualFold(args[0], "RPUSH") && args[1] == "plugin-status": + if writes.Add(1) <= failStatuses { + _, _ = io.WriteString(conn, "-ERR blocked\r\n") + continue + } + _, _ = io.WriteString(conn, ":1\r\n") + default: + _, _ = io.WriteString(conn, "+OK\r\n") + } + } + }(conn) + } + }() + t.Cleanup(func() { + _ = listener.Close() + <-serverDone + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + client := home.New(config.HomeConfig{Enabled: true, Host: host, Port: port}) + t.Cleanup(client.Close) + return client, &writes +} + +func TestStageHomeOverlayDoesNotApplyConfigAfterStageFailure(t *testing.T) { + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + baseCfg.Routing.Strategy = "round-robin" + remoteCfg := &config.Config{} + remoteCfg.Home.Enabled = true + remoteCfg.Routing.Strategy = "fill-first" + remoteCfg.Plugins.Enabled = true + service := &Service{ + cfg: baseCfg, + homePluginSyncFetch: func(context.Context, sdkpluginstore.PluginSyncRequest) (sdkpluginstore.PluginSyncResponse, error) { + return sdkpluginstore.PluginSyncResponse{}, errors.New("plugin sync unavailable") + }, + } + + if _, errStage := service.stageHomeOverlayWithClient(context.Background(), remoteCfg, nil); errStage == nil { + t.Fatal("stageHomeOverlayWithClient() error = nil, want plugin sync failure") + } + + service.cfgMu.RLock() + strategy := service.cfg.Routing.Strategy + service.cfgMu.RUnlock() + if strategy != "round-robin" { + t.Fatalf("failed stage applied routing strategy %q", strategy) + } +} + +func TestReadyHomePluginFinalizationRetriesUntilStatusSucceeds(t *testing.T) { + client, writes := newHomePluginTaskTestClient(t, nil, 1) + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.NodeID = "node-1" + service := &Service{homeGeneration: 1} + work := &homePluginFinalization{ + statusWork: []homePluginStatusWork{{cfg: cfg, report: homeplugins.CompletedSyncReport(homeplugins.CurrentPlatform(), nil)}}, + syncKey: "sync-key", + markSynced: true, + } + + if errFinalize := service.finalizeHomePluginWorkUntilDone(context.Background(), context.Background(), 1, client, work, nil); errFinalize != nil { + t.Fatalf("finalizeHomePluginWorkUntilDone() error = %v", errFinalize) + } + if gotWrites := writes.Load(); gotWrites != 2 { + t.Fatalf("plugin status writes = %d, want 2 after retry", gotWrites) + } + if service.homePluginSyncKey != "sync-key" || work.nextStatus != 1 || work.markSynced { + t.Fatalf("retried ready finalization state: key=%q next=%d marked=%v", service.homePluginSyncKey, work.nextStatus, work.markSynced) + } +} + +func TestReplacementWaitsForHomePluginFinalizationOwnership(t *testing.T) { + client, statusStarted, releaseStatus := newBlockingHomePluginStatusClient(t) + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.NodeID = "node-1" + parentCtx, cancelParent := context.WithCancel(context.Background()) + t.Cleanup(cancelParent) + homeCtx, cancelHome := context.WithCancel(parentCtx) + t.Cleanup(cancelHome) + lifetimeCtx, cancelLifetime := context.WithCancel(homeCtx) + t.Cleanup(cancelLifetime) + previousDone := make(chan struct{}) + cancelled := make(chan struct{}) + service := &Service{ + cfg: cfg, + homeGeneration: 1, + homeSupervisor: &homeSubscriberSupervisor{cancel: func() { + cancelLifetime() + close(cancelled) + close(previousDone) + }, done: previousDone}, + } + work := &homePluginFinalization{statusWork: []homePluginStatusWork{{cfg: cfg, report: homeplugins.CompletedSyncReport(homeplugins.CurrentPlatform(), nil)}}} + finalized := make(chan error, 1) + go func() { + finalized <- service.finalizeHomePluginWorkUntilDone(lifetimeCtx, homeCtx, 1, client, work, func() bool { return true }) + }() + select { + case <-statusStarted: + case <-time.After(time.Second): + t.Fatal("plugin status finalization did not start") + } + + replacementReturned := make(chan struct{}) + go func() { + service.startHomeSubscriber(parentCtx) + close(replacementReturned) + }() + select { + case <-cancelled: + case <-time.After(time.Second): + t.Fatal("replacement did not cancel the blocked controlled finalization") + } + select { + case errFinalize := <-finalized: + if !errors.Is(errFinalize, context.Canceled) { + t.Fatalf("finalization error = %v, want context cancellation", errFinalize) + } + case <-time.After(time.Second): + t.Fatal("blocked finalization did not exit after replacement cancellation") + } + + close(releaseStatus) + cancelParent() + select { + case <-replacementReturned: + case <-time.After(time.Second): + t.Fatal("replacement did not return after cancellation") + } +} + +func TestShutdownCancelsBlockedHomePluginFinalization(t *testing.T) { + client, statusStarted, releaseStatus := newBlockingHomePluginStatusClient(t) + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.NodeID = "node-1" + parentCtx, cancelParent := context.WithCancel(context.Background()) + t.Cleanup(cancelParent) + homeCtx, cancelHome := context.WithCancel(parentCtx) + t.Cleanup(cancelHome) + lifetimeCtx, cancelLifetime := context.WithCancel(homeCtx) + t.Cleanup(cancelLifetime) + previousDone := make(chan struct{}) + cancelled := make(chan struct{}) + service := &Service{ + cfg: cfg, + homeGeneration: 1, + homeSupervisor: &homeSubscriberSupervisor{cancel: func() { + cancelLifetime() + close(cancelled) + close(previousDone) + }, done: previousDone}, + } + work := &homePluginFinalization{statusWork: []homePluginStatusWork{{cfg: cfg, report: homeplugins.CompletedSyncReport(homeplugins.CurrentPlatform(), nil)}}} + finalized := make(chan error, 1) + go func() { + finalized <- service.finalizeHomePluginWorkUntilDone(lifetimeCtx, homeCtx, 1, client, work, nil) + }() + select { + case <-statusStarted: + case <-time.After(time.Second): + t.Fatal("plugin status finalization did not start") + } + + shutdownDone := make(chan error, 1) + go func() { + shutdownDone <- service.Shutdown(context.Background()) + }() + select { + case <-cancelled: + case <-time.After(time.Second): + t.Fatal("shutdown did not cancel the blocked controlled finalization") + } + select { + case errFinalize := <-finalized: + if !errors.Is(errFinalize, context.Canceled) { + t.Fatalf("finalization error = %v, want context cancellation", errFinalize) + } + case <-time.After(time.Second): + t.Fatal("blocked finalization did not exit after shutdown cancellation") + } + + close(releaseStatus) + select { + case errShutdown := <-shutdownDone: + if errShutdown != nil { + t.Fatalf("Shutdown() error = %v", errShutdown) + } + case <-time.After(time.Second): + t.Fatal("shutdown did not return after finalization cancellation") + } +} + +func newBlockingHomePluginStatusClient(t *testing.T) (*home.Client, <-chan struct{}, chan<- struct{}) { + t.Helper() + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + statusStarted := make(chan struct{}) + releaseStatus := make(chan struct{}) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go func(conn net.Conn) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + _, _ = io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n") + case len(args) >= 2 && strings.EqualFold(args[0], "RPUSH") && args[1] == "plugin-status": + close(statusStarted) + <-releaseStatus + _, _ = io.WriteString(conn, ":1\r\n") + default: + _, _ = io.WriteString(conn, "+OK\r\n") + } + } + }(conn) + } + }() + t.Cleanup(func() { + _ = listener.Close() + <-serverDone + }) + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + client := home.New(config.HomeConfig{Enabled: true, Host: host, Port: port}) + t.Cleanup(client.Close) + return client, statusStarted, releaseStatus } func TestHomePluginSyncKeyIncludesCredentialRevision(t *testing.T) { diff --git a/sdk/cliproxy/pprof_server.go b/sdk/cliproxy/pprof_server.go --- a/sdk/cliproxy/pprof_server.go +++ b/sdk/cliproxy/pprof_server.go @@ -18,6 +18,7 @@ server *http.Server addr string enabled bool + owner uint64 } func newPprofServer() *pprofServer { @@ -25,13 +26,20 @@ } func (s *Service) applyPprofConfig(cfg *config.Config) { - if s == nil || cfg == nil { - return + s.applyPprofConfigContext(context.Background(), cfg) +} + +func (s *Service) applyPprofConfigContext(ctx context.Context, cfg *config.Config) bool { + if s == nil || cfg == nil || (ctx != nil && ctx.Err() != nil) { + return false + } + if s.applyPprofConfigContextFn != nil { + return s.applyPprofConfigContextFn(ctx, cfg) } if s.pprofServer == nil { s.pprofServer = newPprofServer() } - s.pprofServer.Apply(cfg) + return s.pprofServer.ApplyContext(ctx, cfg) } func (s *Service) shutdownPprof(ctx context.Context) error { @@ -42,8 +50,18 @@ } func (p *pprofServer) Apply(cfg *config.Config) { + p.ApplyContext(context.Background(), cfg) +} + +func (p *pprofServer) ApplyContext(ctx context.Context, cfg *config.Config) bool { if p == nil || cfg == nil { - return + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false } addr := strings.TrimSpace(cfg.Pprof.Addr) if addr == "" { @@ -52,6 +70,8 @@ enabled := cfg.Pprof.Enable p.mu.Lock() + p.owner++ + owner := p.owner currentServer := p.server currentAddr := p.addr p.addr = addr @@ -60,22 +80,38 @@ p.server = nil p.mu.Unlock() if currentServer != nil { - p.stopServer(currentServer, currentAddr, "disabled") + if errStop := p.stopServerWithContext(ctx, currentServer, currentAddr, "disabled"); errStop != nil { + return false + } } - return + return ctx.Err() == nil } if currentServer != nil && currentAddr == addr { p.mu.Unlock() - return + return ctx.Err() == nil } p.server = nil p.mu.Unlock() if currentServer != nil { - p.stopServer(currentServer, currentAddr, "restarted") + if errStop := p.stopServerWithContext(ctx, currentServer, currentAddr, "restarted"); errStop != nil { + return false + } + } + if errContext := ctx.Err(); errContext != nil { + return false } - p.startServer(addr) + startedServer := p.startServer(addr, owner) + if errContext := ctx.Err(); errContext != nil { + if startedServer != nil { + go func() { + _ = p.stopOwnedServerWithContext(context.Background(), startedServer, addr, "canceled", owner) + }() + } + return false + } + return true } func (p *pprofServer) Shutdown(ctx context.Context) error { @@ -85,6 +121,7 @@ p.mu.Lock() currentServer := p.server currentAddr := p.addr + p.owner++ p.server = nil p.enabled = false p.mu.Unlock() @@ -95,7 +132,7 @@ return p.stopServerWithContext(ctx, currentServer, currentAddr, "shutdown") } -func (p *pprofServer) startServer(addr string) { +func (p *pprofServer) startServer(addr string, owner uint64) *http.Server { mux := newPprofMux() server := &http.Server{ Addr: addr, @@ -104,9 +141,9 @@ } p.mu.Lock() - if !p.enabled || p.addr != addr || p.server != nil { + if !p.enabled || p.addr != addr || p.owner != owner || p.server != nil { p.mu.Unlock() - return + return nil } p.server = server p.mu.Unlock() @@ -115,17 +152,41 @@ go func() { if errServe := server.ListenAndServe(); errServe != nil && !errors.Is(errServe, http.ErrServerClosed) { log.Errorf("pprof server failed on %s: %v", addr, errServe) - p.mu.Lock() - if p.server == server { - p.server = nil - } - p.mu.Unlock() + p.clearFailedServer(server) } }() + return server +} + +// clearFailedServer removes a failed physical server even if a same-address +// ApplyContext transferred lifecycle ownership while ListenAndServe was starting. +func (p *pprofServer) clearFailedServer(server *http.Server) { + if p == nil || server == nil { + return + } + p.mu.Lock() + if p.server == server { + p.server = nil + } + p.mu.Unlock() } func (p *pprofServer) stopServer(server *http.Server, addr string, reason string) { _ = p.stopServerWithContext(context.Background(), server, addr, reason) +} + +func (p *pprofServer) stopOwnedServerWithContext(ctx context.Context, server *http.Server, addr string, reason string, owner uint64) error { + if p == nil || server == nil { + return nil + } + p.mu.Lock() + if p.server != server || p.owner != owner { + p.mu.Unlock() + return nil + } + p.server = nil + p.mu.Unlock() + return p.stopServerWithContext(ctx, server, addr, reason) } func (p *pprofServer) stopServerWithContext(ctx context.Context, server *http.Server, addr string, reason string) error { diff --git a/sdk/cliproxy/pprof_server_test.go b/sdk/cliproxy/pprof_server_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/pprof_server_test.go @@ -0,0 +1,74 @@ +package cliproxy + +import ( + "context" + "net/http" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +func TestPprofServerStopOwnedServerKeepsReplacement(t *testing.T) { + pprof := newPprofServer() + oldServer := &http.Server{} + replacement := &http.Server{} + pprof.server = replacement + + if errStop := pprof.stopOwnedServerWithContext(context.Background(), oldServer, "old", "canceled", 1); errStop != nil { + t.Fatalf("stopOwnedServerWithContext() error = %v", errStop) + } + pprof.mu.Lock() + current := pprof.server + pprof.mu.Unlock() + if current != replacement { + t.Fatal("stopping a stale pprof server removed the replacement server") + } +} + +func TestPprofServerSamePointerOwnerTransferKeepsCurrentServer(t *testing.T) { + pprof := newPprofServer() + server := &http.Server{} + pprof.server = server + pprof.addr = "127.0.0.1:6060" + pprof.enabled = true + pprof.owner = 1 + + cfg := &config.Config{} + cfg.Pprof.Enable = true + cfg.Pprof.Addr = "127.0.0.1:6060" + if !pprof.ApplyContext(context.Background(), cfg) { + t.Fatal("ApplyContext() = false, want same-pointer owner transfer") + } + + pprof.mu.Lock() + owner := pprof.owner + pprof.mu.Unlock() + if owner == 1 { + t.Fatal("ApplyContext() did not transfer same-server ownership") + } + if errStop := pprof.stopOwnedServerWithContext(context.Background(), server, cfg.Pprof.Addr, "canceled", 1); errStop != nil { + t.Fatalf("stopOwnedServerWithContext() error = %v", errStop) + } + pprof.mu.Lock() + current := pprof.server + pprof.mu.Unlock() + if current != server { + t.Fatal("stale owner stopped the current same-pointer server") + } +} + +func TestPprofServerServeFailureClearsTransferredOwner(t *testing.T) { + pprof := newPprofServer() + server := &http.Server{} + pprof.server = server + pprof.owner = 2 + + pprof.clearFailedServer(server) + + pprof.mu.Lock() + current := pprof.server + pprof.mu.Unlock() + if current != nil { + t.Fatal("serve failure retained a server after ownership transferred") + } +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -10,9 +10,11 @@ "os" "strings" "sync" + "sync/atomic" "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/api" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/constant" "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/homeplugins" @@ -29,6 +31,7 @@ sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access" sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" sdkpluginstore "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginstore" @@ -48,6 +51,11 @@ // configUpdateMu serializes config updates across watcher + home. configUpdateMu sync.Mutex + + // configRuntimeMu orders side-effecting runtime application after config commits. + configRuntimeMu sync.Mutex + configSequence uint64 + appliedRoutingState *routingRuntimeState // configPath is the path to the configuration file. configPath string @@ -106,17 +114,124 @@ // wsGateway manages websocket Gemini providers. wsGateway *wsrelay.Manager - homeClient *home.Client - homeCancel context.CancelFunc - homeLogForwarder *logging.HomeAppLogForwarder - homePluginSyncMu sync.Mutex - homePluginSyncKey string - homePluginSyncFetch func(context.Context, sdkpluginstore.PluginSyncRequest) (sdkpluginstore.PluginSyncResponse, error) + homeLifecycleMu sync.Mutex + homeOwnershipMu sync.Mutex + homeConfigCommitMu sync.Mutex + homeConfigStageHook func() + homeConfigCommitHook func() + homeConfigRuntimeHook func() + applyPprofConfigContextFn func(context.Context, *config.Config) bool + updateServerClientsContextFn func(context.Context, *config.Config) bool + homeSupervisor *homeSubscriberSupervisor + homeMu sync.Mutex + homeGeneration uint64 + homeClient *home.Client + homeRegistry *executionregistry.Registry + homeDispatchBundle *coreauth.HomeDispatchBundle + homeDrainBound time.Duration + homeCancel context.CancelFunc + runCancel context.CancelFunc + homeLogForwarder homeLogForwarder + homeLogForwarderClient *home.Client + homePluginSyncMu sync.Mutex + homePluginSyncKey string + homePluginSyncFetch func(context.Context, sdkpluginstore.PluginSyncRequest) (sdkpluginstore.PluginSyncResponse, error) + homePluginDeleteTask func(context.Context, *config.Config, home.PluginTask) homeplugins.SyncReport +} + +type homeSubscriberSupervisor struct { + cancel context.CancelFunc + done chan struct{} + + publisherMu sync.Mutex + publisherDone <-chan struct{} +} + +func (s *homeSubscriberSupervisor) setPublisherCompletion(done <-chan struct{}) { + if s == nil { + return + } + s.publisherMu.Lock() + s.publisherDone = done + s.publisherMu.Unlock() +} + +func (s *homeSubscriberSupervisor) publisherCompletion() <-chan struct{} { + if s == nil { + return nil + } + s.publisherMu.Lock() + defer s.publisherMu.Unlock() + return s.publisherDone +} + +type homeConfigWorkQueue struct { + mu sync.Mutex + items [][]byte + wake chan struct{} +} + +func newHomeConfigWorkQueue() *homeConfigWorkQueue { + return &homeConfigWorkQueue{wake: make(chan struct{}, 1)} +} + +func (q *homeConfigWorkQueue) enqueue(raw []byte) { + if q == nil { + return + } + item := append([]byte(nil), raw...) + q.mu.Lock() + q.items = append(q.items, item) + q.mu.Unlock() + select { + case q.wake <- struct{}{}: + default: + } +} + +func (q *homeConfigWorkQueue) dequeue(ctx context.Context) ([]byte, bool) { + if q == nil || ctx == nil { + return nil, false + } + for { + if ctx.Err() != nil { + return nil, false + } + q.mu.Lock() + if ctx.Err() != nil { + q.mu.Unlock() + return nil, false + } + if len(q.items) > 0 { + item := q.items[0] + q.items[0] = nil + q.items = q.items[1:] + q.mu.Unlock() + return item, true + } + q.mu.Unlock() + select { + case <-ctx.Done(): + return nil, false + case <-q.wake: + } + } +} + +type homeLogForwarder interface { + Bind(*home.Client) + Deactivate(*home.Client) + Stop() +} + +var startHomeLogForwarder = func(queueSize int) homeLogForwarder { + return logging.StartHomeAppLogForwarder(queueSize) } const ( modelRegistrationMaxWorkersPerCategory = 5 modelRegistrationMaxWorkersOpenAICompatibility = 20 + homeSubscriberPreAckRetryBackoff = 100 * time.Millisecond ) const ( @@ -176,16 +291,29 @@ sdkAuth.RegisterPluginAuthParser(nil) return false } - if ctx == nil { - ctx = context.Background() - } - s.cfgMu.RLock() cfg := s.cfg s.cfgMu.RUnlock() + return s.syncPluginRuntimeConfigForConfig(ctx, cfg) +} + +func (s *Service) syncPluginRuntimeConfigForConfig(ctx context.Context, cfg *config.Config) bool { + if s == nil { + sdkAuth.RegisterPluginAuthParser(nil) + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } if s.pluginHost != nil { s.pluginHost.ApplyConfig(ctx, cfg) + } + if errContext := ctx.Err(); errContext != nil { + return false } if s.coreManager != nil { s.coreManager.SetPluginScheduler(s.pluginHost) @@ -195,6 +323,9 @@ return false } s.pluginHost.RegisterFrontendAuthProviders() + if errContext := ctx.Err(); errContext != nil { + return false + } if s.accessManager != nil { s.accessManager.SetProviders(sdkaccess.RegisteredProviders()) } @@ -203,7 +334,7 @@ if s.server != nil { s.server.RefreshPluginManagementRoutes() } - return true + return ctx.Err() == nil } func (s *Service) syncPluginModelRuntime(ctx context.Context) { @@ -214,6 +345,9 @@ ctx = context.Background() } s.pluginHost.RegisterModels(ctx, registry.GetGlobalRegistry()) + if ctx.Err() != nil { + return + } s.registerAvailableExecutors(ctx, executorRegistrationOptions{ includeBaseline: s.cfg != nil && s.cfg.Home.Enabled, includePlugins: true, @@ -221,6 +355,9 @@ auths: s.coreManager.List(), }) s.refreshPluginModelRegistrations(ctx) + if ctx.Err() != nil { + return + } s.coreManager.RefreshSchedulerAll() } @@ -686,7 +823,7 @@ return nil } auth = auth.Clone() - s.ensureExecutorsForAuth(auth) + s.ensureExecutorsForAuthWithContext(ctx, auth, false) // IMPORTANT: Update coreManager FIRST, before model registration. // This ensures that configuration changes (proxy_url, prefix, etc.) take effect @@ -727,7 +864,13 @@ if s == nil || s.coreManager == nil || auth == nil || auth.ID == "" { return } + if ctx != nil && ctx.Err() != nil { + return + } s.registerModelsForAuthWithCache(ctx, auth, compatCache) + if ctx != nil && ctx.Err() != nil { + return + } s.coreManager.ReconcileRegistryModelStates(ctx, auth.ID) // Refresh the scheduler entry so that the auth's supportedModelSet is rebuilt @@ -770,24 +913,35 @@ } func (s *Service) configureCooldownStateStore(cfg *config.Config) { + _ = s.configureCooldownStateStoreContext(context.Background(), cfg, false) +} + +func (s *Service) configureCooldownStateStoreContext(ctx context.Context, cfg *config.Config, persistOld bool) bool { if s == nil || s.coreManager == nil { - return + return true } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + return s.coreManager.SwapCooldownStateStore(ctx, s.resolveCooldownStateStore(cfg), persistOld) +} + +func (s *Service) resolveCooldownStateStore(cfg *config.Config) coreauth.CooldownStateStore { if cfg == nil || !cfg.SaveCooldownStatus || cfg.Home.Enabled { - s.coreManager.SetCooldownStateStore(nil) - return + return nil } authDir, errResolve := resolveCooldownStateAuthDir(cfg) if errResolve != nil { log.Warnf("failed to resolve cooldown state directory: %v", errResolve) - s.coreManager.SetCooldownStateStore(nil) - return + return nil } if authDir == "" { - s.coreManager.SetCooldownStateStore(nil) - return + return nil } - s.coreManager.SetCooldownStateStore(coreauth.NewFileCooldownStateStoreWithAuthDir(authDir, authDir)) + return coreauth.NewFileCooldownStateStoreWithAuthDir(authDir, authDir) } func resolveCooldownStateAuthDir(cfg *config.Config) (string, error) { @@ -952,14 +1106,18 @@ } func (s *Service) ensureExecutorsForAuth(a *coreauth.Auth) { - s.ensureExecutorsForAuthWithMode(a, false) + s.ensureExecutorsForAuthWithContext(context.Background(), a, false) } func (s *Service) ensureExecutorsForAuthWithMode(a *coreauth.Auth, forceReplace bool) { - if a == nil { + s.ensureExecutorsForAuthWithContext(context.Background(), a, forceReplace) +} + +func (s *Service) ensureExecutorsForAuthWithContext(ctx context.Context, a *coreauth.Auth, forceReplace bool) { + if a == nil || (ctx != nil && ctx.Err() != nil) { return } - s.registerAvailableExecutors(context.Background(), executorRegistrationOptions{ + s.registerAvailableExecutors(ctx, executorRegistrationOptions{ auths: []*coreauth.Auth{a}, forceReplaceAuths: forceReplace, }) @@ -1181,7 +1339,13 @@ if s == nil || s.pluginHost == nil || a == nil { return false } + if ctx != nil && ctx.Err() != nil { + return true + } result := s.pluginHost.ModelsForAuth(ctx, a) + if ctx != nil && ctx.Err() != nil { + return true + } if !result.Handled { return false } @@ -1209,7 +1373,7 @@ result.Auth.Attributes[key] = value } } - if updated, errUpdate := s.coreManager.Update(context.Background(), result.Auth); errUpdate == nil && updated != nil { + if updated, errUpdate := s.coreManager.Update(ctx, result.Auth); errUpdate == nil && updated != nil { activeAuth = updated.Clone() } } @@ -1232,6 +1396,9 @@ activeExcluded = strings.Split(val, ",") } } + if ctx != nil && ctx.Err() != nil { + return true + } models := applyExcludedModels(result.Models, activeExcluded) models = applyOAuthModelAliasForAuth(s.cfg, providerKey, activeAuthKind, activeAuth.Attributes, models) if len(models) > 0 { @@ -1243,31 +1410,79 @@ } func (s *Service) applyConfigUpdate(newCfg *config.Config) { - s.applyConfigUpdateWithAuthSynthesis(newCfg, true) + s.applyConfigUpdateWithAuthSynthesis(context.Background(), newCfg, true) } func (s *Service) applyWatcherConfigUpdate(newCfg *config.Config) { - s.applyConfigUpdateWithAuthSynthesis(newCfg, false) + s.applyConfigUpdateWithAuthSynthesis(context.Background(), newCfg, false) } -func (s *Service) applyConfigUpdateWithAuthSynthesis(newCfg *config.Config, synthesizeConfigAuths bool) { +type configCommit struct { + cfg *config.Config + sequence uint64 +} + +type routingRuntimeState struct { + strategy string + sessionAffinity bool + sessionAffinityTTL time.Duration +} + +func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { + state := routingRuntimeState{ + strategy: "round-robin", + sessionAffinityTTL: time.Hour, + } + if cfg == nil { + return state + } + + switch strings.ToLower(strings.TrimSpace(cfg.Routing.Strategy)) { + case "fill-first", "fillfirst", "ff": + state.strategy = "fill-first" + } + state.sessionAffinity = cfg.Routing.SessionAffinity + if ttl := strings.TrimSpace(cfg.Routing.SessionAffinityTTL); ttl != "" { + if parsed, errParse := time.ParseDuration(ttl); errParse == nil && parsed > 0 { + state.sessionAffinityTTL = parsed + } + } + return state +} + +func newRoutingSelector(state routingRuntimeState) coreauth.Selector { + var selector coreauth.Selector + if state.strategy == "fill-first" { + selector = &coreauth.FillFirstSelector{} + } else { + selector = &coreauth.RoundRobinSelector{} + } + if state.sessionAffinity { + selector = coreauth.NewSessionAffinitySelectorWithConfig(coreauth.SessionAffinityConfig{ + Fallback: selector, + TTL: state.sessionAffinityTTL, + }) + } + return selector +} + +func (s *Service) applyConfigUpdateWithAuthSynthesis(ctx context.Context, newCfg *config.Config, synthesizeConfigAuths bool) bool { + commit := s.commitConfigUpdate(newCfg) + if commit.cfg == nil { + return false + } + return s.applyConfigRuntime(ctx, commit, synthesizeConfigAuths) +} + +// commitConfigUpdate applies only in-memory configuration state. Runtime work that +// may block on plugins, models, storage, or networking is deliberately deferred. +func (s *Service) commitConfigUpdate(newCfg *config.Config) configCommit { if s == nil { - return + return configCommit{} } s.configUpdateMu.Lock() defer s.configUpdateMu.Unlock() - - previousStrategy := "" - var previousSessionAffinity bool - var previousSessionAffinityTTL string - s.cfgMu.RLock() - if s.cfg != nil { - previousStrategy = strings.ToLower(strings.TrimSpace(s.cfg.Routing.Strategy)) - previousSessionAffinity = s.cfg.Routing.SessionAffinity - previousSessionAffinityTTL = s.cfg.Routing.SessionAffinityTTL - } - s.cfgMu.RUnlock() if newCfg == nil { s.cfgMu.RLock() @@ -1275,86 +1490,132 @@ s.cfgMu.RUnlock() } if newCfg == nil { - return + return configCommit{} } - nextStrategy := strings.ToLower(strings.TrimSpace(newCfg.Routing.Strategy)) - normalizeStrategy := func(strategy string) string { - switch strategy { - case "fill-first", "fillfirst", "ff": - return "fill-first" - default: - return "round-robin" - } - } - previousStrategy = normalizeStrategy(previousStrategy) - nextStrategy = normalizeStrategy(nextStrategy) - - nextSessionAffinity := newCfg.Routing.SessionAffinity - nextSessionAffinityTTL := newCfg.Routing.SessionAffinityTTL - - selectorChanged := previousStrategy != nextStrategy || - previousSessionAffinity != nextSessionAffinity || - previousSessionAffinityTTL != nextSessionAffinityTTL - - if s.coreManager != nil && selectorChanged { - var selector coreauth.Selector - switch nextStrategy { - case "fill-first": - selector = &coreauth.FillFirstSelector{} - default: - selector = &coreauth.RoundRobinSelector{} - } - - if nextSessionAffinity { - ttl := time.Hour - if ttlStr := strings.TrimSpace(nextSessionAffinityTTL); ttlStr != "" { - if parsed, err := time.ParseDuration(ttlStr); err == nil && parsed > 0 { - ttl = parsed - } - } - selector = coreauth.NewSessionAffinitySelectorWithConfig(coreauth.SessionAffinityConfig{ - Fallback: selector, - TTL: ttl, - }) - } - - s.coreManager.SetSelector(selector) - } - - s.applyRetryConfig(newCfg) - s.configureCooldownStateStore(newCfg) - s.applyPprofConfig(newCfg) - if s.server != nil { - s.server.UpdateClients(newCfg) - } s.cfgMu.Lock() s.cfg = newCfg s.cfgMu.Unlock() - if s.coreManager != nil { - s.coreManager.SetConfig(newCfg) - s.coreManager.SetOAuthModelAlias(newCfg.OAuthModelAlias) + s.configSequence++ + return configCommit{cfg: newCfg, sequence: s.configSequence} +} + +func (s *Service) configCommitCurrent(commit configCommit) bool { + if s == nil || commit.sequence == 0 { + return false } - ctx := coreauth.WithSkipPersist(context.Background()) - s.syncPluginRuntimeConfig(ctx) + s.configUpdateMu.Lock() + current := s.configSequence == commit.sequence + s.configUpdateMu.Unlock() + return current +} + +func (s *Service) applyConfigRuntime(ctx context.Context, commit configCommit, synthesizeConfigAuths bool) bool { + cfg := commit.cfg + if s == nil || cfg == nil { + return false + } + s.configRuntimeMu.Lock() + defer s.configRuntimeMu.Unlock() + if !s.configCommitCurrent(commit) { + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + + if !s.applyManagerConfig(ctx, commit) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + if !s.applyPprofConfigContext(ctx, cfg) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + if !s.updateServerClientsContext(ctx, cfg) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + + registrationCtx := coreauth.WithSkipPersist(ctx) + s.syncPluginRuntimeConfigForConfig(registrationCtx, cfg) + if errContext := ctx.Err(); errContext != nil { + return false + } var auths []*coreauth.Auth if s.coreManager != nil { auths = s.coreManager.List() } - s.registerAvailableExecutors(context.Background(), executorRegistrationOptions{ - includeBaseline: newCfg.Home.Enabled, + s.registerAvailableExecutors(registrationCtx, executorRegistrationOptions{ + includeBaseline: cfg.Home.Enabled, forceReplaceAuths: true, auths: auths, }) - if synthesizeConfigAuths { - s.registerConfigAPIKeyAuths(ctx, newCfg) + if errContext := ctx.Err(); errContext != nil { + return false } - if s.coreManager != nil && !newCfg.Home.Enabled && newCfg.SaveCooldownStatus { - if errRestoreCooldown := s.coreManager.RestoreCooldownStates(context.Background()); errRestoreCooldown != nil { + if synthesizeConfigAuths { + s.registerConfigAPIKeyAuths(registrationCtx, cfg) + } + if errContext := ctx.Err(); errContext != nil { + return false + } + if s.coreManager != nil && !cfg.Home.Enabled && cfg.SaveCooldownStatus { + if errRestoreCooldown := s.coreManager.RestoreCooldownStates(registrationCtx); errRestoreCooldown != nil && ctx.Err() == nil { log.Warnf("failed to restore cooldown state after config update: %v", errRestoreCooldown) } } - s.syncPluginModelRuntime(ctx) + if errContext := ctx.Err(); errContext != nil { + return false + } + s.syncPluginModelRuntime(registrationCtx) + return ctx.Err() == nil +} + +func (s *Service) applyManagerConfig(ctx context.Context, commit configCommit) bool { + if s == nil || s.coreManager == nil || commit.cfg == nil { + return s != nil && commit.cfg != nil + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + routingState := normalizedRoutingRuntimeState(commit.cfg) + if s.appliedRoutingState == nil || *s.appliedRoutingState != routingState { + s.coreManager.SetSelector(newRoutingSelector(routingState)) + s.appliedRoutingState = &routingState + } + s.applyRetryConfig(commit.cfg) + store := s.resolveCooldownStateStore(commit.cfg) + if !s.coreManager.ApplyConfigWithCooldownStateStore(ctx, commit.cfg, store) { + return false + } + s.coreManager.SetOAuthModelAlias(commit.cfg.OAuthModelAlias) + return true +} + +func (s *Service) updateServerClientsContext(ctx context.Context, cfg *config.Config) bool { + if s == nil || cfg == nil || (ctx != nil && ctx.Err() != nil) { + return false + } + if s.updateServerClientsContextFn != nil { + return s.updateServerClientsContextFn(ctx, cfg) + } + if s.server == nil { + return true + } + return s.server.UpdateClientsContext(ctx, cfg) } func (s *Service) reloadConfigFromWatcher() bool { @@ -1430,18 +1691,48 @@ } func (s *Service) applyHomeOverlayContext(ctx context.Context, remoteCfg *config.Config) error { + return s.applyHomeOverlayWithClient(ctx, remoteCfg, nil) +} + +func (s *Service) applyHomeOverlayWithClient(ctx context.Context, remoteCfg *config.Config, client *home.Client) error { + work, errStage := s.stageHomeOverlayWithClient(ctx, remoteCfg, client) + if errStage != nil { + return errStage + } + if ctx != nil { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + } + if work.config != nil { + if !s.applyConfigUpdateWithAuthSynthesis(ctx, work.config, true) { + return context.Canceled + } + work.committed = true + } + if errFinalize := s.finalizeHomePluginWork(ctx, client, work); errFinalize != nil { + return errFinalize + } + return nil +} + +func (s *Service) stageHomeOverlayWithClient(ctx context.Context, remoteCfg *config.Config, client *home.Client) (*homePluginFinalization, error) { + work := &homePluginFinalization{} if s == nil || remoteCfg == nil { - return nil + return work, nil } if ctx == nil { ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return nil, errContext } s.cfgMu.RLock() baseCfg := s.cfg s.cfgMu.RUnlock() if baseCfg == nil { - return nil + return work, nil } merged := *remoteCfg @@ -1455,26 +1746,120 @@ syncCfg.Plugins.StoreAuth = storeAuth logHomeConfigChanges(baseCfg, &merged) - report, syncKey, didSync, errSync := s.syncHomePlugins(ctx, &syncCfg) + report, syncKey, didSync, errSync := s.syncHomePluginsWithClient(ctx, &syncCfg, client) if errSync != nil { - log.Warnf("failed to sync home plugins: %v", errSync) + return nil, fmt.Errorf("sync home plugins: %w", errSync) } - s.applyConfigUpdate(&merged) - var errLoad error + if errContext := ctx.Err(); errContext != nil { + return nil, errContext + } if didSync { - errLoad = homeplugins.MarkLoadResults(&report, s.pluginHost) - if errLoad != nil { - log.Warnf("failed to load home plugins after config update: %v", errLoad) + if errLoad := homeplugins.MarkLoadResults(&report, s.pluginHost); errLoad != nil { + return nil, fmt.Errorf("load home plugins: %w", errLoad) } } if strings.TrimSpace(report.Task) != "" { - s.reportHomePluginStatus(ctx, &merged, report) - if errSync == nil && errLoad == nil { - s.markHomePluginsSynced(syncKey) + work.syncKey = syncKey + work.markSynced = true + if strings.TrimSpace(merged.Home.NodeID) != "" { + work.statusWork = append(work.statusWork, homePluginStatusWork{cfg: &merged, report: report}) } } - s.processHomePluginTasks(ctx, &merged) - return nil + taskWork, errTasks := s.stageHomePluginTasksWithClient(ctx, &merged, client) + if errTasks != nil { + return nil, fmt.Errorf("stage home plugin tasks: %w", errTasks) + } + work.taskWork = append(work.taskWork, taskWork...) + if errContext := ctx.Err(); errContext != nil { + return nil, errContext + } + work.config = &merged + return work, nil +} + +func (s *Service) commitHomeConfig(lifetimeCtx, homeCtx context.Context, generation uint64, work *homePluginFinalization) bool { + if s == nil || work == nil || work.config == nil { + return false + } + + s.homeConfigCommitMu.Lock() + defer s.homeConfigCommitMu.Unlock() + if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) { + return false + } + if s.homeConfigCommitHook != nil { + s.homeConfigCommitHook() + } + if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) { + return false + } + commit := s.commitConfigUpdate(work.config) + if commit.cfg == nil { + return false + } + work.config = commit.cfg + work.configCommit = commit + work.committed = true + return true +} + +func (s *Service) homeLifetimeActive(homeCtx, lifetimeCtx context.Context, generation uint64) bool { + if s == nil || homeCtx.Err() != nil || lifetimeCtx.Err() != nil { + return false + } + s.homeMu.Lock() + active := s.homeGeneration == generation + s.homeMu.Unlock() + return active +} + +func (s *Service) finalizeHomePluginWorkUntilDone(ctx, homeCtx context.Context, generation uint64, client *home.Client, work *homePluginFinalization, publish func() bool) error { + stopClose := closeHomeClientOnCancellation(ctx, client) + defer stopClose() + for { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + + s.homeOwnershipMu.Lock() + if !s.homeLifetimeActive(homeCtx, ctx, generation) { + s.homeOwnershipMu.Unlock() + return context.Canceled + } + errFinalize := s.finalizeHomePluginWork(ctx, client, work) + if errFinalize == nil && (publish == nil || publish()) { + s.homeOwnershipMu.Unlock() + return nil + } + s.homeOwnershipMu.Unlock() + if errFinalize == nil { + return context.Canceled + } + + log.WithError(errFinalize).Warn("failed to finalize home plugins; retrying") + timer := time.NewTimer(homeSubscriberPreAckRetryBackoff) + select { + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + case <-timer.C: + } + } +} + +func closeHomeClientOnCancellation(ctx context.Context, client *home.Client) func() { + if ctx == nil || client == nil { + return func() {} + } + stop := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + client.Close() + case <-stop: + } + }() + return func() { close(stop) } } func logHomeConfigChanges(oldCfg, newCfg *config.Config) { @@ -1557,6 +1942,23 @@ }() } +func applyHomeObservationBarrier(registry *executionregistry.Registry, revision int64) { + if registry != nil { + registry.ObserveBarrier(revision) + } +} + +func applyHomeInFlightPublisherConfig(manager *coreauth.Manager, cfg internalconfig.CredentialInFlightConfig) error { + publisherCfg, errConfig := coreauth.HomeInFlightPublisherConfigFromConfig(cfg) + if errConfig != nil { + return errConfig + } + if manager != nil { + manager.ApplyHomeInFlightPublisherConfig(publisherCfg) + } + return nil +} + func (s *Service) startHomeSubscriber(ctx context.Context) { if s == nil { return @@ -1568,40 +1970,349 @@ return } - if s.homeCancel != nil { - s.homeCancel() - s.homeCancel = nil - } - if s.homeClient != nil { - s.homeClient.Close() - s.homeClient = nil - } - if s.homeLogForwarder != nil { - s.homeLogForwarder.Stop() - s.homeLogForwarder = nil + parentCtx := ctx + if parentCtx == nil { + parentCtx = context.Background() } - homeCtx := ctx - if homeCtx == nil { - homeCtx = context.Background() + s.homeLifecycleMu.Lock() + defer s.homeLifecycleMu.Unlock() + + if previousSupervisor := s.homeSupervisor; previousSupervisor != nil { + s.homeConfigCommitMu.Lock() + previousSupervisor.cancel() + s.homeConfigCommitMu.Unlock() + <-previousSupervisor.done } - homeCtx, cancel := context.WithCancel(homeCtx) + if !s.drainDetachedHomeLifetime(parentCtx) { + return + } + if parentCtx.Err() != nil { + return + } + + homeCtx, cancel := context.WithCancel(parentCtx) + done := make(chan struct{}) + s.homeMu.Lock() + s.homeGeneration++ + generation := s.homeGeneration s.homeCancel = cancel + s.homeMu.Unlock() + supervisor := &homeSubscriberSupervisor{cancel: cancel, done: done} + s.homeSupervisor = supervisor + go s.runHomeSubscriber(homeCtx, parentCtx, cfg.Home, generation, supervisor) +} - client := home.New(cfg.Home) - s.homeClient = client - home.SetCurrent(client) +func (s *Service) drainDetachedHomeLifetime(parentCtx context.Context) bool { + s.homeMu.Lock() + previousCancel := s.homeCancel + previousClient := s.homeClient + previousRegistry := s.homeRegistry + previousBundle := s.homeDispatchBundle + previousDrainBound := s.homeDrainBound + previousForwarder := s.homeLogForwarder + previousForwarderClient := s.homeLogForwarderClient + s.homeCancel = nil + s.homeClient = nil + s.homeRegistry = nil + s.homeDispatchBundle = nil + s.homeDrainBound = 0 + s.homeLogForwarderClient = nil + s.homeMu.Unlock() - go client.StartConfigSubscriber(homeCtx, func(raw []byte) error { - parsed, err := config.ParseConfigBytes(raw) - if err != nil { - log.Warnf("failed to parse home config payload: %v", err) - return err + if s.coreManager != nil { + s.coreManager.ClearHomeDispatchBundle(previousBundle) + } + home.ClearCurrentIf(previousClient) + if previousCancel != nil { + previousCancel() + } + if previousForwarder != nil && previousForwarderClient == previousClient { + previousForwarder.Deactivate(previousClient) + } + if previousRegistry != nil { + if previousDrainBound <= 0 { + previousDrainBound = internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound } - return s.applyHomeOverlayContext(homeCtx, parsed) - }) - s.startHomeUsageForwarder(homeCtx, client) - s.homeLogForwarder = logging.StartHomeAppLogForwarder(0) + drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), previousDrainBound) + errDrain := previousRegistry.Drain(drainCtx) + cancelDrain() + if errDrain != nil { + if previousClient != nil { + previousClient.Close() + } + if parentCtx.Err() == nil { + log.WithError(errDrain).Error("failed to drain replaced Home execution registry") + s.cancelServiceRun() + } + return false + } + } + if previousClient != nil { + previousClient.Close() + } + return true +} + +func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.Context, homeCfg internalconfig.HomeConfig, generation uint64, supervisor *homeSubscriberSupervisor) { + defer func() { + s.homeMu.Lock() + if s.homeGeneration == generation { + s.homeCancel = nil + } + s.homeMu.Unlock() + close(supervisor.done) + }() + + for homeCtx.Err() == nil { + supervisor.setPublisherCompletion(nil) + client := home.New(homeCfg) + client.SetManagedLifetime(true) + registry := executionregistry.New() + releaseCtx, releaseCancel := context.WithCancel(context.WithoutCancel(homeCtx)) + releaseFlusher := home.NewReleaseFlusher(client.LimiterConfig, client.PushConcurrencyRelease) + registry.SetReleaseSink(releaseFlusher.MarkDirty) + releaseDone := make(chan struct{}) + go func() { + defer close(releaseDone) + releaseFlusher.Run(releaseCtx) + }() + lifetimeCtx, lifetimeCancel := context.WithCancel(homeCtx) + cancelBound := atomic.Int64{} + cancelBound.Store(int64(internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound)) + queue := newHomeConfigWorkQueue() + ready := make(chan struct{}) + var readyOnce sync.Once + var published atomic.Bool + workerDone := make(chan struct{}) + + go func() { + defer close(workerDone) + s.runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx, generation, client, registry, queue, ready, &published, &cancelBound, supervisor) + }() + + errRun := client.RunConfigSubscriberLifetime(lifetimeCtx, func(raw []byte) error { + parsed, errParse := config.ParseConfigBytes(raw) + if errParse != nil { + log.Warnf("failed to parse home config payload: %v", errParse) + return errParse + } + if errSetLifecycle := client.SetLifecycleConfig(parsed.CredentialConcurrency); errSetLifecycle != nil { + log.Warnf("failed to apply Home lifecycle config: %v", errSetLifecycle) + return errSetLifecycle + } + if errPublisherConfig := applyHomeInFlightPublisherConfig(s.coreManager, parsed.CredentialInFlight); errPublisherConfig != nil { + log.Warnf("failed to apply Home in-flight publisher config: %v", errPublisherConfig) + return errPublisherConfig + } + applyHomeObservationBarrier(registry, parsed.CredentialConcurrency.ObservationBarrierRevision) + cancelBound.Store(int64(parsed.CredentialConcurrency.WithDefaults().CPACancelBound)) + queue.enqueue(raw) + return nil + }, func() { + readyOnce.Do(func() { close(ready) }) + }) + lifetimeCancel() + <-workerDone + if publisherDone := supervisor.publisherCompletion(); publisherDone != nil { + <-publisherDone + } + + s.detachHomeSubscriberLifetime(client, registry) + drainBound := time.Duration(cancelBound.Load()) + drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), drainBound) + errDrain := registry.Drain(drainCtx) + var errFlush error + if errDrain == nil { + errFlush = releaseFlusher.Flush(drainCtx) + } + cancelDrain() + releaseCancel() + <-releaseDone + client.Close() + if errDrain != nil { + if parentCtx.Err() == nil { + log.WithError(errDrain).Error("failed to drain Home execution registry") + s.cancelServiceRun() + } + return + } + if errFlush != nil { + if parentCtx.Err() == nil { + log.WithError(errFlush).Error("failed to flush Home concurrency releases") + s.cancelServiceRun() + } + return + } + if errRun != nil && homeCtx.Err() == nil { + log.WithError(errRun).Warn("home config subscription lifetime ended") + } + if !published.Load() && errRun != nil && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) { + return + } + } +} + +func (s *Service) runHomeConfigWorker(lifetimeCtx, homeCtx context.Context, generation uint64, client *home.Client, registry *executionregistry.Registry, queue *homeConfigWorkQueue, ready <-chan struct{}, published *atomic.Bool, cancelBound *atomic.Int64) { + s.runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx, generation, client, registry, queue, ready, published, cancelBound, nil) +} + +func (s *Service) runHomeConfigWorkerWithSupervisor(lifetimeCtx, homeCtx context.Context, generation uint64, client *home.Client, registry *executionregistry.Registry, queue *homeConfigWorkQueue, ready <-chan struct{}, published *atomic.Bool, cancelBound *atomic.Int64, supervisor *homeSubscriberSupervisor) { + select { + case <-lifetimeCtx.Done(): + return + case <-ready: + } + + for { + if lifetimeCtx.Err() != nil { + return + } + raw, ok := queue.dequeue(lifetimeCtx) + if !ok { + return + } + if lifetimeCtx.Err() != nil { + return + } + + var work *homePluginFinalization + for { + if lifetimeCtx.Err() != nil { + return + } + parsed, errParse := config.ParseConfigBytes(raw) + if errParse == nil { + work, errParse = s.stageHomeOverlayWithClient(lifetimeCtx, parsed, client) + } + if errParse == nil { + break + } + if lifetimeCtx.Err() != nil { + return + } + log.WithError(errParse).Warn("failed to stage home config; retrying") + if !waitForHomeSubscriberRetry(lifetimeCtx, homeSubscriberPreAckRetryBackoff) { + return + } + } + + var publish func() bool + if !published.Load() { + publish = func() bool { + s.homeMu.Lock() + defer s.homeMu.Unlock() + if homeCtx.Err() != nil || lifetimeCtx.Err() != nil || s.homeGeneration != generation { + return false + } + s.homeClient = client + s.homeRegistry = registry + s.homeDrainBound = time.Duration(cancelBound.Load()) + if s.coreManager != nil { + s.homeDispatchBundle = s.coreManager.PublishHomeDispatch(client, registry, generation) + } + home.SetCurrent(client) + if s.homeLogForwarder == nil { + s.homeLogForwarder = startHomeLogForwarder(0) + } + s.homeLogForwarder.Bind(client) + s.homeLogForwarderClient = client + published.Store(true) + return true + } + } + if s.homeConfigStageHook != nil { + s.homeConfigStageHook() + } + if !s.commitHomeConfig(lifetimeCtx, homeCtx, generation, work) { + return + } + if s.homeConfigRuntimeHook != nil { + s.homeConfigRuntimeHook() + } + if !s.homeLifetimeActive(homeCtx, lifetimeCtx, generation) || !s.applyConfigRuntime(lifetimeCtx, work.configCommit, true) { + return + } + if errFinalize := s.finalizeHomePluginWorkUntilDone(lifetimeCtx, homeCtx, generation, client, work, publish); errFinalize != nil { + if !errors.Is(errFinalize, context.Canceled) { + log.WithError(errFinalize).Warn("home plugin finalization ended") + } + return + } + if publish != nil { + s.startHomeInFlightPublisher(lifetimeCtx, client, registry, supervisor) + s.startHomeUsageForwarder(lifetimeCtx, client) + } + } +} + +func (s *Service) startHomeInFlightPublisher(ctx context.Context, client *home.Client, registry *executionregistry.Registry, supervisor *homeSubscriberSupervisor) { + if s == nil || s.coreManager == nil { + return + } + done := make(chan struct{}) + if supervisor != nil { + supervisor.setPublisherCompletion(done) + } + go func() { + defer close(done) + s.coreManager.StartHomeInFlightPublisher(ctx, client, registry) + }() +} + +func waitForHomeSubscriberRetry(ctx context.Context, delay time.Duration) bool { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func (s *Service) detachHomeSubscriberLifetime(client *home.Client, registry *executionregistry.Registry) { + if s == nil { + return + } + s.homeMu.Lock() + var bundle *coreauth.HomeDispatchBundle + if s.homeClient == client && s.homeRegistry == registry { + bundle = s.homeDispatchBundle + s.homeClient = nil + s.homeRegistry = nil + s.homeDispatchBundle = nil + s.homeDrainBound = 0 + } + forwarder := s.homeLogForwarder + if s.homeLogForwarderClient == client { + s.homeLogForwarderClient = nil + } else { + forwarder = nil + } + s.homeMu.Unlock() + if s.coreManager != nil { + s.coreManager.ClearHomeDispatchBundle(bundle) + } + home.ClearCurrentIf(client) + if forwarder != nil { + forwarder.Deactivate(client) + } +} + +func (s *Service) cancelServiceRun() { + if s == nil { + return + } + s.homeMu.Lock() + cancel := s.runCancel + if cancel == nil { + cancel = s.homeCancel + } + s.homeMu.Unlock() + if cancel != nil { + cancel() + } } // Run starts the service and blocks until the context is cancelled or the server stops. @@ -1620,6 +2331,18 @@ if ctx == nil { ctx = context.Background() } + ctx, runCancel := context.WithCancel(ctx) + s.homeMu.Lock() + s.runCancel = runCancel + s.homeMu.Unlock() + defer func() { + runCancel() + s.homeMu.Lock() + if s.runCancel != nil { + s.runCancel = nil + } + s.homeMu.Unlock() + }() usage.StartDefault(ctx) homeEnabled := s.cfg != nil && s.cfg.Home.Enabled @@ -1806,19 +2529,51 @@ ctx = context.Background() } - if s.homeCancel != nil { - s.homeCancel() - s.homeCancel = nil + s.homeLifecycleMu.Lock() + if supervisor := s.homeSupervisor; supervisor != nil { + s.homeConfigCommitMu.Lock() + supervisor.cancel() + s.homeConfigCommitMu.Unlock() + <-supervisor.done } - if s.homeClient != nil { - s.homeClient.Close() - s.homeClient = nil + s.homeMu.Lock() + homeCancel := s.homeCancel + homeClient := s.homeClient + homeRegistry := s.homeRegistry + homeDispatchBundle := s.homeDispatchBundle + homeForwarder := s.homeLogForwarder + homeForwarderClient := s.homeLogForwarderClient + s.homeGeneration++ + s.homeCancel = nil + s.homeClient = nil + s.homeRegistry = nil + s.homeDispatchBundle = nil + s.homeDrainBound = 0 + s.homeLogForwarder = nil + s.homeLogForwarderClient = nil + s.homeMu.Unlock() + if s.coreManager != nil { + s.coreManager.ClearHomeDispatchBundle(homeDispatchBundle) } - if s.homeLogForwarder != nil { - s.homeLogForwarder.Stop() - s.homeLogForwarder = nil + home.ClearCurrentIf(homeClient) + if homeCancel != nil { + homeCancel() } - home.ClearCurrent() + if homeRegistry != nil { + if errClose := homeRegistry.Close(); errClose != nil { + log.WithError(errClose).Warn("failed to close Home execution registry during shutdown") + } + } + if homeClient != nil { + homeClient.Close() + } + if homeForwarder != nil { + if homeForwarderClient == homeClient { + homeForwarder.Deactivate(homeClient) + } + homeForwarder.Stop() + } + s.homeLifecycleMu.Unlock() // legacy refresh loop removed; only stopping core auth manager below @@ -1879,7 +2634,7 @@ includePlugins: true, }) s.pluginHost.RegisterFrontendAuthProviders() - s.pluginHost.ShutdownAll() + s.pluginHost.ShutdownAllContext(ctx) if s.accessManager != nil { s.accessManager.SetProviders(sdkaccess.RegisteredProviders()) } @@ -1920,6 +2675,9 @@ if ctx == nil { ctx = context.Background() } + if ctx.Err() != nil { + return + } if a.Disabled { GlobalModelRegistry().UnregisterClient(a.ID) return @@ -1947,6 +2705,9 @@ } } if s.tryRegisterPluginModelsForAuth(ctx, a, provider, authKind, excluded) { + return + } + if ctx.Err() != nil { return } var models []*ModelInfo @@ -2143,7 +2904,13 @@ } } } + if ctx.Err() != nil { + return + } models = applyOAuthModelAliasForAuth(s.cfg, provider, authKind, a.Attributes, models) + if ctx.Err() != nil { + return + } key := provider if key == "" { key = strings.ToLower(strings.TrimSpace(a.Provider)) @@ -2165,20 +2932,31 @@ // as part of the previous registration snapshot and is cleared when the auth is // rebound to the refreshed model catalog. func (s *Service) refreshModelRegistrationForAuth(current *coreauth.Auth) bool { - return s.refreshModelRegistrationForAuthWithCache(current, nil) + return s.refreshModelRegistrationForAuthWithContext(context.Background(), current, nil) } func (s *Service) refreshModelRegistrationForAuthWithCache(current *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) bool { + return s.refreshModelRegistrationForAuthWithContext(context.Background(), current, compatCache) +} + +func (s *Service) refreshModelRegistrationForAuthWithContext(ctx context.Context, current *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) bool { if s == nil || s.coreManager == nil || current == nil || current.ID == "" { return false } - - ctx := context.Background() + if ctx == nil { + ctx = context.Background() + } + if ctx.Err() != nil { + return false + } if !current.Disabled { - s.ensureExecutorsForAuth(current) + s.ensureExecutorsForAuthWithContext(ctx, current, false) } s.registerModelsForAuthWithCache(ctx, current, compatCache) s.coreManager.ReconcileRegistryModelStates(ctx, current.ID) + if ctx.Err() != nil { + return false + } latest, ok := s.latestAuthForModelRegistration(current.ID) if !ok || latest.Disabled { @@ -2190,8 +2968,11 @@ // Re-apply the latest auth snapshot so concurrent auth updates cannot leave // stale model registrations behind. This may duplicate registration work when // no auth fields changed, but keeps the refresh path simple and correct. - s.ensureExecutorsForAuth(latest) + s.ensureExecutorsForAuthWithContext(ctx, latest, false) s.registerModelsForAuthWithCache(ctx, latest, compatCache) + if ctx.Err() != nil { + return false + } s.coreManager.ReconcileRegistryModelStates(ctx, latest.ID) s.coreManager.RefreshSchedulerEntry(current.ID) return true diff --git a/sdk/cliproxy/service_executionregistry_test.go b/sdk/cliproxy/service_executionregistry_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/service_executionregistry_test.go @@ -0,0 +1,2776 @@ +package cliproxy + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/homeplugins" + "github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + sdkpluginstore "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginstore" +) + +type blockingServiceCooldownStore struct { + started chan struct{} +} + +func (s *blockingServiceCooldownStore) Load(context.Context) ([]coreauth.CooldownStateRecord, error) { + return nil, nil +} + +func (s *blockingServiceCooldownStore) Save(ctx context.Context, _ []coreauth.CooldownStateRecord) error { + close(s.started) + <-ctx.Done() + return ctx.Err() +} + +func TestConfigCommitDoesNotHoldCommitMutexDuringCooldownPersistence(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ID: "auth-1", Provider: "xai", Status: coreauth.StatusActive} + if _, errRegister := manager.Register(coreauth.WithSkipPersist(context.Background()), auth); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + manager.MarkResult(context.Background(), coreauth.Result{ + AuthID: auth.ID, Provider: auth.Provider, Model: "grok-4", Success: false, + Error: &coreauth.Error{Message: "rate limited", HTTPStatus: http.StatusTooManyRequests}, + }) + store := &blockingServiceCooldownStore{started: make(chan struct{})} + manager.SetCooldownStateStore(store) + service := &Service{cfg: &config.Config{}, coreManager: manager} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + applyDone := make(chan bool, 1) + go func() { + applyDone <- service.applyConfigUpdateWithAuthSynthesis(ctx, &config.Config{DisableCooling: true}, false) + }() + select { + case <-store.started: + case <-time.After(time.Second): + t.Fatal("old cooldown store persistence did not start") + } + + commitDone := make(chan struct{}) + go func() { + service.commitConfigUpdate(&config.Config{}) + close(commitDone) + }() + select { + case <-commitDone: + case <-time.After(time.Second): + t.Fatal("config commit mutex remained locked during cooldown persistence") + } + + cancel() + select { + case applied := <-applyDone: + if applied { + t.Fatal("config runtime apply succeeded after cooldown persistence cancellation") + } + case <-time.After(time.Second): + t.Fatal("config runtime apply did not honor cooldown persistence cancellation") + } +} + +func TestServiceShutdownPreservesReplacementHomeClient(t *testing.T) { + staleClient := home.New(internalconfig.HomeConfig{Enabled: true}) + replacementClient := home.New(internalconfig.HomeConfig{Enabled: true}) + home.SetCurrent(replacementClient) + t.Cleanup(home.ClearCurrent) + + service := &Service{homeClient: staleClient} + if errShutdown := service.Shutdown(context.Background()); errShutdown != nil { + t.Fatalf("Shutdown() error = %v", errShutdown) + } + if current := home.Current(); current != replacementClient { + t.Fatal("Shutdown() cleared the replacement Home client") + } +} + +func TestServiceConcurrentReplacementWaitsForInFlightDrain(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + _, oldCancel := context.WithCancel(context.Background()) + t.Cleanup(oldCancel) + cfg := &config.Config{} + cfg.Home.Enabled = true + service := &Service{ + cfg: cfg, + homeCancel: oldCancel, + homeClient: home.New(internalconfig.HomeConfig{Enabled: true}), + homeRegistry: registry, + homeDrainBound: time.Second, + } + + firstReturned := make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(firstReturned) + }() + deadline := time.Now().Add(time.Second) + for { + if _, errLate := registry.BeginDispatch(); errLate != nil { + break + } + if time.Now().After(deadline) { + t.Fatal("first replacement did not begin draining") + } + time.Sleep(time.Millisecond) + } + + secondReturned := make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(secondReturned) + }() + select { + case <-secondReturned: + t.Fatal("concurrent replacement returned before the first drain completed") + case <-time.After(50 * time.Millisecond): + } + + pending.End() + select { + case <-firstReturned: + case <-time.After(time.Second): + t.Fatal("first replacement did not complete after its drain") + } + select { + case <-secondReturned: + case <-time.After(time.Second): + t.Fatal("second replacement did not complete after the first drain") + } +} + +func TestServiceReplacementWaitsForPreACKSupervisorExit(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstSubscribed := make(chan struct{}) + secondStarted := make(chan struct{}) + secondStartedBeforeFirstDone := make(chan struct{}) + stop := make(chan struct{}) + firstDoneForServer := make(chan (<-chan struct{}), 1) + var configRequests atomic.Int32 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go servePreACKReplacementConnection(conn, &configRequests, firstSubscribed, secondStarted, secondStartedBeforeFirstDone, firstDoneForServer, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + select { + case <-firstSubscribed: + case <-time.After(time.Second): + t.Fatal("first subscriber did not reach pre-ACK state") + } + + service.homeLifecycleMu.Lock() + firstDone := service.homeSupervisor.done + service.homeLifecycleMu.Unlock() + if firstDone == nil { + t.Fatal("first subscriber has no supervisor completion signal") + } + firstDoneForServer <- firstDone + + replaced := make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(replaced) + }() + + select { + case <-secondStartedBeforeFirstDone: + t.Fatal("replacement subscriber started before the pre-ACK supervisor exited") + case <-secondStarted: + case <-time.After(time.Second): + t.Fatal("replacement subscriber did not start") + } + select { + case <-firstDone: + case <-time.After(time.Second): + t.Fatal("pre-ACK supervisor did not exit") + } + select { + case <-replaced: + case <-time.After(time.Second): + t.Fatal("replacement start did not return") + } +} + +func TestServiceReplacementWaitsForPublisherExitAndPinsACKedLifetimeDependencies(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + frames := make(chan home.InFlightSnapshotFrame, 64) + var configRequests atomic.Int32 + firstPublisherDoneForServer := make(chan (<-chan struct{}), 1) + secondConfigResult := make(chan error, 1) + allowSecondConfig := make(chan struct{}) + stop := make(chan struct{}) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go servePublisherReplacementConnection(conn, &configRequests, frames, firstPublisherDoneForServer, secondConfigResult, allowSecondConfig, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + service.coreManager = coreauth.NewManager(nil, nil, nil) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + + firstFrame := waitForPublisherReplacementFrame(t, frames, 11) + firstClient := waitForServiceHomeClient(t, service, time.Second) + firstRegistry := waitForServiceRegistry(t, service, time.Second) + service.homeLifecycleMu.Lock() + firstPublisherDone := service.homeSupervisor.publisherCompletion() + service.homeLifecycleMu.Unlock() + if firstPublisherDone == nil { + t.Fatal("first subscriber did not record publisher completion") + } + firstPublisherDoneForServer <- firstPublisherDone + if firstFrame.BarrierRevision != 11 { + t.Fatalf("first publisher frame = %#v", firstFrame) + } + + replaced := make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(replaced) + }() + + deadline := time.NewTimer(time.Second) + defer deadline.Stop() + select { + case errSecondConfig := <-secondConfigResult: + if errSecondConfig != nil { + t.Fatal(errSecondConfig) + } + case <-deadline.C: + t.Fatal("replacement did not begin its config lifetime") + } + close(allowSecondConfig) + + secondFrame := waitForPublisherReplacementFrame(t, frames, 22) + secondClient := waitForServiceHomeClient(t, service, time.Second) + secondRegistry := waitForServiceRegistry(t, service, time.Second) + if secondFrame.BarrierRevision != 22 { + t.Fatalf("replacement publisher frame = %#v", secondFrame) + } + if secondClient == firstClient || secondRegistry == firstRegistry { + t.Fatal("replacement publisher reused the previous lifetime dependencies") + } + select { + case <-replaced: + case <-time.After(time.Second): + t.Fatal("replacement subscriber did not finish setup") + } +} + +func TestHomeConfigWorkerDoesNotApplyCanceledQueuedConfig(t *testing.T) { + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + baseCfg.Routing.Strategy = "round-robin" + service := &Service{cfg: baseCfg} + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("routing:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + lifetimeCtx, cancelLifetime := context.WithCancel(context.Background()) + cancelLifetime() + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + + service.runHomeConfigWorker(lifetimeCtx, context.Background(), 1, nil, executionregistry.New(), queue, ready, &atomic.Bool{}, &cancelBound) + + service.cfgMu.RLock() + strategy := service.cfg.Routing.Strategy + service.cfgMu.RUnlock() + if strategy != "round-robin" { + t.Fatalf("canceled queued config changed routing strategy to %q", strategy) + } +} + +func TestHomeConfigWorkerSkipsStagedConfigWhenReplacementCancels(t *testing.T) { + client, _ := newHomePluginTaskTestClient(t, nil, 0) + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + baseCfg.Routing.Strategy = "round-robin" + parentCtx, cancelParent := context.WithCancel(context.Background()) + t.Cleanup(cancelParent) + homeCtx, cancelHome := context.WithCancel(parentCtx) + t.Cleanup(cancelHome) + lifetimeCtx, cancelLifetime := context.WithCancel(homeCtx) + t.Cleanup(cancelLifetime) + stagePaused := make(chan struct{}) + releaseStage := make(chan struct{}) + var releaseStageOnce sync.Once + t.Cleanup(func() { releaseStageOnce.Do(func() { close(releaseStage) }) }) + cancelled := make(chan struct{}) + workerDone := make(chan struct{}) + service := &Service{ + cfg: baseCfg, + homeGeneration: 1, + homeConfigStageHook: func() { + close(stagePaused) + <-releaseStage + }, + homeSupervisor: &homeSubscriberSupervisor{cancel: func() { + cancelLifetime() + close(cancelled) + }, done: workerDone}, + } + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("routing:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + go func() { + defer close(workerDone) + service.runHomeConfigWorker(lifetimeCtx, homeCtx, 1, client, executionregistry.New(), queue, ready, &atomic.Bool{}, &cancelBound) + }() + select { + case <-stagePaused: + case <-time.After(time.Second): + t.Fatal("config worker did not pause after staging") + } + + replacementDone := make(chan struct{}) + go func() { + service.startHomeSubscriber(parentCtx) + close(replacementDone) + }() + select { + case <-cancelled: + case <-time.After(time.Second): + t.Fatal("replacement did not cancel the staged Home config") + } + releaseStageOnce.Do(func() { close(releaseStage) }) + select { + case <-workerDone: + case <-time.After(time.Second): + t.Fatal("canceled config worker did not exit") + } + + service.cfgMu.RLock() + strategy := service.cfg.Routing.Strategy + service.cfgMu.RUnlock() + if strategy != "round-robin" { + t.Fatalf("canceled staged config changed routing strategy to %q", strategy) + } + select { + case <-replacementDone: + case <-time.After(time.Second): + t.Fatal("replacement deadlocked after canceling staged config") + } +} + +func TestHomeConfigWorkerCommitCompletesBeforeReplacementCancellation(t *testing.T) { + client, _ := newHomePluginTaskTestClient(t, nil, 0) + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + baseCfg.Routing.Strategy = "round-robin" + parentCtx, cancelParent := context.WithCancel(context.Background()) + t.Cleanup(cancelParent) + homeCtx, cancelHome := context.WithCancel(parentCtx) + t.Cleanup(cancelHome) + lifetimeCtx, cancelLifetime := context.WithCancel(homeCtx) + t.Cleanup(cancelLifetime) + commitPaused := make(chan struct{}) + releaseCommit := make(chan struct{}) + var releaseCommitOnce sync.Once + t.Cleanup(func() { releaseCommitOnce.Do(func() { close(releaseCommit) }) }) + cancelled := make(chan struct{}) + workerDone := make(chan struct{}) + service := &Service{ + cfg: baseCfg, + homeGeneration: 1, + homeConfigCommitHook: func() { + close(commitPaused) + <-releaseCommit + }, + homeSupervisor: &homeSubscriberSupervisor{cancel: func() { + cancelLifetime() + close(cancelled) + }, done: workerDone}, + } + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("routing:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + go func() { + defer close(workerDone) + service.runHomeConfigWorker(lifetimeCtx, homeCtx, 1, client, executionregistry.New(), queue, ready, &atomic.Bool{}, &cancelBound) + }() + select { + case <-commitPaused: + case <-time.After(time.Second): + t.Fatal("config worker did not pause inside commit") + } + + replacementDone := make(chan struct{}) + go func() { + service.startHomeSubscriber(parentCtx) + close(replacementDone) + }() + select { + case <-cancelled: + t.Fatal("replacement canceled while config commit owned the commit mutex") + case <-time.After(50 * time.Millisecond): + } + releaseCommitOnce.Do(func() { close(releaseCommit) }) + select { + case <-cancelled: + case <-time.After(time.Second): + t.Fatal("replacement did not cancel after config commit completed") + } + select { + case <-workerDone: + case <-time.After(time.Second): + t.Fatal("config worker deadlocked after committed config was canceled") + } + + service.cfgMu.RLock() + strategy := service.cfg.Routing.Strategy + service.cfgMu.RUnlock() + if strategy != "fill-first" { + t.Fatalf("committed config routing strategy = %q, want fill-first", strategy) + } + select { + case <-replacementDone: + case <-time.After(time.Second): + t.Fatal("replacement deadlocked after committed config") + } +} + +func TestHomeConfigWorkerCancellationAtPostCommitBoundarySkipsRuntimePublish(t *testing.T) { + for _, testCase := range []struct { + name string + cancel func(context.CancelFunc, context.CancelFunc) + }{ + {name: "parent", cancel: func(cancelParent, _ context.CancelFunc) { cancelParent() }}, + {name: "transport", cancel: func(_, cancelLifetime context.CancelFunc) { cancelLifetime() }}, + } { + t.Run(testCase.name, func(t *testing.T) { + client, _ := newHomePluginTaskTestClient(t, nil, 0) + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + baseCfg.Routing.Strategy = "round-robin" + parentCtx, cancelParent := context.WithCancel(context.Background()) + t.Cleanup(cancelParent) + homeCtx, cancelHome := context.WithCancel(parentCtx) + t.Cleanup(cancelHome) + lifetimeCtx, cancelLifetime := context.WithCancel(homeCtx) + t.Cleanup(cancelLifetime) + runtimePaused := make(chan struct{}) + releaseRuntime := make(chan struct{}) + var releaseRuntimeOnce sync.Once + t.Cleanup(func() { releaseRuntimeOnce.Do(func() { close(releaseRuntime) }) }) + service := &Service{ + cfg: baseCfg, + homeGeneration: 1, + homeConfigRuntimeHook: func() { + close(runtimePaused) + <-releaseRuntime + }, + } + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("routing:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + published := atomic.Bool{} + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + workerDone := make(chan struct{}) + go func() { + defer close(workerDone) + service.runHomeConfigWorker(lifetimeCtx, homeCtx, 1, client, executionregistry.New(), queue, ready, &published, &cancelBound) + }() + select { + case <-runtimePaused: + case <-time.After(time.Second): + t.Fatal("Home config worker did not reach post-commit boundary") + } + + testCase.cancel(cancelParent, cancelLifetime) + releaseRuntimeOnce.Do(func() { close(releaseRuntime) }) + select { + case <-workerDone: + case <-time.After(time.Second): + t.Fatal("canceled Home config worker did not exit") + } + service.cfgMu.RLock() + strategy := service.cfg.Routing.Strategy + service.cfgMu.RUnlock() + if strategy != "fill-first" { + t.Fatalf("post-commit cancellation changed committed routing strategy to %q", strategy) + } + if published.Load() { + t.Fatal("canceled post-commit work published Home runtime") + } + }) + } +} + +func TestHomeConfigWorkerShutdownCancelsBlockedRuntimeUpdatesBeforePublish(t *testing.T) { + for _, testCase := range []struct { + name string + apply func(*Service, func(context.Context, *config.Config) bool) + }{ + { + name: "pprof", + apply: func(service *Service, blocked func(context.Context, *config.Config) bool) { + service.applyPprofConfigContextFn = blocked + }, + }, + { + name: "server", + apply: func(service *Service, blocked func(context.Context, *config.Config) bool) { + service.updateServerClientsContextFn = blocked + }, + }, + } { + t.Run(testCase.name, func(t *testing.T) { + client, _ := newHomePluginTaskTestClient(t, nil, 0) + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + baseCfg.Home.NodeID = "node-1" + parentCtx, cancelParent := context.WithCancel(context.Background()) + t.Cleanup(cancelParent) + homeCtx, cancelHome := context.WithCancel(parentCtx) + t.Cleanup(cancelHome) + lifetimeCtx, cancelLifetime := context.WithCancel(homeCtx) + t.Cleanup(cancelLifetime) + started := make(chan struct{}) + workerDone := make(chan struct{}) + service := &Service{ + cfg: baseCfg, + homeGeneration: 1, + homeSupervisor: &homeSubscriberSupervisor{cancel: cancelLifetime, done: workerDone}, + } + testCase.apply(service, func(ctx context.Context, _ *config.Config) bool { + close(started) + <-ctx.Done() + return false + }) + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("routing:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + published := atomic.Bool{} + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + go func() { + defer close(workerDone) + service.runHomeConfigWorker(lifetimeCtx, homeCtx, 1, client, executionregistry.New(), queue, ready, &published, &cancelBound) + }() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("Home config worker did not start blocked runtime update") + } + + shutdownDone := make(chan error, 1) + go func() { shutdownDone <- service.Shutdown(context.Background()) }() + select { + case <-workerDone: + case <-time.After(time.Second): + t.Fatal("shutdown did not cancel blocked runtime update") + } + select { + case errShutdown := <-shutdownDone: + if errShutdown != nil { + t.Fatalf("Shutdown() error = %v", errShutdown) + } + case <-time.After(time.Second): + t.Fatal("shutdown waited for blocked runtime update") + } + if published.Load() { + t.Fatal("canceled runtime update published Home state") + } + }) + } +} + +func TestHomeConfigWorkerCancelsBlockedAntigravityModelRefreshBeforePublish(t *testing.T) { + modelRefreshStarted := make(chan struct{}) + releaseModelRefresh := make(chan struct{}) + var releaseModelRefreshOnce sync.Once + modelServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + close(modelRefreshStarted) + select { + case <-r.Context().Done(): + case <-releaseModelRefresh: + } + })) + t.Cleanup(modelServer.Close) + t.Cleanup(func() { releaseModelRefreshOnce.Do(func() { close(releaseModelRefresh) }) }) + + client, _ := newHomePluginTaskTestClient(t, nil, 0) + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ + ID: "blocked-antigravity-refresh", + Provider: "antigravity", + Metadata: map[string]any{"access_token": "test-token"}, + Attributes: map[string]string{ + "base_url": modelServer.URL, + }, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatal(errRegister) + } + t.Cleanup(func() { GlobalModelRegistry().UnregisterClient(auth.ID) }) + + parentCtx, cancelParent := context.WithCancel(context.Background()) + t.Cleanup(cancelParent) + homeCtx, cancelHome := context.WithCancel(parentCtx) + t.Cleanup(cancelHome) + lifetimeCtx, cancelLifetime := context.WithCancel(homeCtx) + t.Cleanup(cancelLifetime) + service := &Service{ + cfg: baseCfg, + coreManager: manager, + pluginHost: pluginhost.New(), + homeGeneration: 1, + } + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("routing:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + published := atomic.Bool{} + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + workerDone := make(chan struct{}) + go func() { + defer close(workerDone) + service.runHomeConfigWorker(lifetimeCtx, homeCtx, 1, client, executionregistry.New(), queue, ready, &published, &cancelBound) + }() + + select { + case <-modelRefreshStarted: + case <-time.After(time.Second): + t.Fatal("Home config worker did not start Antigravity model refresh") + } + cancelLifetime() + select { + case <-workerDone: + case <-time.After(time.Second): + t.Fatal("Home config worker did not stop after model refresh cancellation") + } + if published.Load() { + t.Fatal("canceled model refresh published Home runtime") + } + service.homeMu.Lock() + publishedClient := service.homeClient + publishedRegistry := service.homeRegistry + service.homeMu.Unlock() + if publishedClient != nil || publishedRegistry != nil { + t.Fatal("canceled model refresh exposed Home runtime state") + } +} + +func TestHomeConfigWorkerRetriesStageFailureForSameQueuedConfig(t *testing.T) { + client, _ := newHomePluginTaskTestClient(t, nil, 0) + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + baseCfg.Routing.Strategy = "round-robin" + var attempts atomic.Int32 + service := &Service{ + cfg: baseCfg, + homeGeneration: 1, + homePluginSyncFetch: func(context.Context, sdkpluginstore.PluginSyncRequest) (sdkpluginstore.PluginSyncResponse, error) { + if attempts.Add(1) == 1 { + return sdkpluginstore.PluginSyncResponse{}, fmt.Errorf("plugin sync unavailable") + } + return sdkpluginstore.PluginSyncResponse{ + SchemaVersion: sdkpluginstore.PluginSyncSchemaVersion, + ExpiresAt: time.Now().Add(time.Minute), + }, nil + }, + } + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("plugins:\n enabled: true\nrouting:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + lifetimeCtx, cancelLifetime := context.WithCancel(context.Background()) + t.Cleanup(cancelLifetime) + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + published := atomic.Bool{} + published.Store(true) + workerDone := make(chan struct{}) + go func() { + defer close(workerDone) + service.runHomeConfigWorker(lifetimeCtx, context.Background(), 1, client, executionregistry.New(), queue, ready, &published, &cancelBound) + }() + + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + service.cfgMu.RLock() + strategy := service.cfg.Routing.Strategy + service.cfgMu.RUnlock() + if attempts.Load() >= 2 && strategy == "fill-first" { + cancelLifetime() + select { + case <-workerDone: + case <-time.After(time.Second): + t.Fatal("config worker did not stop after cancellation") + } + return + } + time.Sleep(time.Millisecond) + } + cancelLifetime() + <-workerDone + t.Fatalf("stage attempts = %d and config was not applied after retry", attempts.Load()) +} + +func TestServiceInitialOverlayStagesPluginWritesUntilReady(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + pluginSync := make(chan struct{}) + pluginStatus := make(chan struct{}, 2) + pluginTasks := make(chan struct{}) + freshCommandProbe := make(chan struct{}) + allowAck := make(chan struct{}) + stop := make(chan struct{}) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveInitialOverlayPluginConnection(conn, pluginSync, pluginStatus, pluginTasks, freshCommandProbe, allowAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.Host = host + cfg.Home.Port = port + cfg.Home.NodeID = "node-1" + cfg.Home.DisableClusterDiscovery = true + cfg.Plugins.Enabled = true + cfg.Plugins.Dir = t.TempDir() + var deletes atomic.Int32 + service := &Service{cfg: cfg, homePluginDeleteTask: func(_ context.Context, _ *config.Config, task home.PluginTask) homeplugins.SyncReport { + deletes.Add(1) + return homeplugins.DeleteWithReport(context.Background(), nil, nil, task.ID, task.PluginID) + }} + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + + for name, observed := range map[string]<-chan struct{}{ + "plugin sync": pluginSync, + "plugin tasks": pluginTasks, + "plugin status": pluginStatus, + } { + select { + case <-observed: + t.Fatalf("initial overlay staged %s before subscription ACK and fresh command probe", name) + case <-time.After(50 * time.Millisecond): + } + } + if gotDeletes := deletes.Load(); gotDeletes != 0 { + t.Fatalf("initial overlay executed %d plugin deletes before subscription ACK and fresh command probe", gotDeletes) + } + service.homeMu.Lock() + client := service.homeClient + registry := service.homeRegistry + service.homeMu.Unlock() + if client != nil || registry != nil || home.Current() != nil { + t.Fatal("initial overlay exposed its Home client or registry before subscription ACK") + } + + close(allowAck) + select { + case <-freshCommandProbe: + case <-time.After(time.Second): + t.Fatal("subscription ACK did not rebuild and probe a fresh command connection") + } + for name, observed := range map[string]<-chan struct{}{ + "plugin sync": pluginSync, + "plugin tasks": pluginTasks, + } { + select { + case <-observed: + case <-time.After(time.Second): + t.Fatalf("ready Home lifetime did not stage %s after subscription ACK and fresh command probe", name) + } + } + for range 2 { + select { + case <-pluginStatus: + case <-time.After(time.Second): + t.Fatal("ready Home lifetime did not flush staged plugin reports") + } + } + if gotDeletes := deletes.Load(); gotDeletes != 1 { + t.Fatalf("ready Home lifetime executed %d plugin deletes, want 1", gotDeletes) + } + if waitForServiceRegistry(t, service, time.Second) == nil || home.Current() == nil { + t.Fatal("subscription ACK did not expose the Home client and registry") + } +} + +func TestServiceDiscardsStalePreACKPluginWork(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstSubscribed := make(chan struct{}) + secondSubscribed := make(chan struct{}) + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + serverDone := make(chan struct{}) + var subscriptions atomic.Int32 + var pluginWrites atomic.Int32 + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveStalePreACKPluginConnection(conn, &subscriptions, &pluginWrites, firstSubscribed, secondSubscribed, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.Host = host + cfg.Home.Port = port + cfg.Home.NodeID = "node-1" + cfg.Home.DisableClusterDiscovery = true + cfg.Plugins.Enabled = true + cfg.Plugins.Dir = t.TempDir() + service := &Service{cfg: cfg} + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + select { + case <-firstSubscribed: + case <-time.After(time.Second): + t.Fatal("first subscriber did not stage plugin work before ACK") + } + + replaced := make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(replaced) + }() + select { + case <-secondSubscribed: + case <-time.After(time.Second): + t.Fatal("replacement subscriber did not reach subscription ACK") + } + if got := pluginWrites.Load(); got != 0 { + t.Fatalf("stale pre-ACK lifetime flushed %d plugin reports", got) + } + close(allowSecondAck) + deadline := time.Now().Add(time.Second) + for pluginWrites.Load() != 1 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := pluginWrites.Load(); got != 1 { + t.Fatalf("replacement lifetime plugin reports = %d, want 1", got) + } + if waitForServiceRegistry(t, service, time.Second) == nil { + t.Fatal("replacement subscription did not expose a ready registry") + } + select { + case <-replaced: + case <-time.After(time.Second): + t.Fatal("replacement subscriber did not finish setup") + } +} + +func TestServiceExplicitReplacementDrainsPendingAndScopeBeforeStartingNewLifetime(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + secondSubscribe := make(chan struct{}) + var secondSubscribeOnce sync.Once + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + var subscriptionMu sync.Mutex + subscriptions := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + registry := waitForServiceRegistry(t, service, time.Second) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scopePending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(scopePending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + resourceClosed := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(resourceClosed) + go scope.End("canceled") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + replaced := make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(replaced) + }() + select { + case <-resourceClosed: + case <-time.After(time.Second): + t.Fatal("explicit replacement did not start draining the active scope") + } + select { + case <-secondSubscribe: + t.Fatal("new subscriber started before the old pending dispatch drained") + case <-time.After(50 * time.Millisecond): + } + pending.End() + select { + case <-replaced: + case <-time.After(time.Second): + t.Fatal("explicit replacement did not finish after pending dispatch ended") + } + select { + case <-secondSubscribe: + case <-time.After(time.Second): + t.Fatal("new subscriber did not start after successful drain") + } +} + +func TestServiceReplacementWaitsForBlockedDrainSupervisorExit(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + secondSubscribe := make(chan struct{}) + var secondSubscribeOnce sync.Once + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + var subscriptionMu sync.Mutex + subscriptions := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + service.homeLifecycleMu.Lock() + firstDone := service.homeSupervisor.done + service.homeLifecycleMu.Unlock() + registry := waitForServiceRegistry(t, service, time.Second) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scopePending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(scopePending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + resourceClosed := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(resourceClosed) + go scope.End("canceled") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + replaced := make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(replaced) + }() + select { + case <-resourceClosed: + case <-time.After(time.Second): + t.Fatal("replacement did not begin draining the active scope") + } + select { + case <-firstDone: + t.Fatal("supervisor exited before the pending dispatch drained") + case <-secondSubscribe: + t.Fatal("replacement subscriber started before the old supervisor exited") + case <-time.After(50 * time.Millisecond): + } + + pending.End() + select { + case <-firstDone: + case <-time.After(time.Second): + t.Fatal("old supervisor did not exit after drain completed") + } + select { + case <-secondSubscribe: + case <-time.After(time.Second): + t.Fatal("replacement subscriber did not start after old supervisor exit") + } + select { + case <-replaced: + case <-time.After(time.Second): + t.Fatal("replacement start did not return") + } +} + +func TestServiceExplicitReplacementCancelsRunWhenDrainTimesOut(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + secondSubscribe := make(chan struct{}) + var secondSubscribeOnce sync.Once + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + var subscriptionMu sync.Mutex + subscriptions := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + serviceCtx, cancelService := context.WithCancel(context.Background()) + t.Cleanup(cancelService) + service.homeMu.Lock() + service.runCancel = cancelService + service.homeMu.Unlock() + service.startHomeSubscriber(serviceCtx) + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + registry := waitForServiceRegistry(t, service, time.Second) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + resourceClosed := make(chan struct{}) + release := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(resourceClosed) + <-release + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + go service.startHomeSubscriber(serviceCtx) + select { + case <-resourceClosed: + case <-time.After(time.Second): + t.Fatal("explicit replacement did not start draining the blocking scope") + } + select { + case <-serviceCtx.Done(): + case <-time.After(time.Second): + t.Fatal("explicit replacement did not cancel the Service run after drain timeout") + } + select { + case <-secondSubscribe: + t.Fatal("new subscriber started after explicit replacement drain timeout") + case <-time.After(50 * time.Millisecond): + } + + close(release) + scope.End("test cleanup") +} + +func TestServiceReplacesRegistryOnlyAfterNewSubscriptionAck(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + secondSubscribe := make(chan struct{}) + var secondSubscribeOnce sync.Once + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + var subscriptionMu sync.Mutex + subscriptions := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.Host = host + cfg.Home.Port = port + cfg.Home.DisableClusterDiscovery = true + service := &Service{cfg: cfg} + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + firstRegistry := waitForServiceRegistry(t, service, time.Second) + if home.Current() == nil { + t.Fatal("first client was not exposed after subscription ACK") + } + + close(loseFirst) + select { + case <-secondSubscribe: + case <-time.After(time.Second): + t.Fatal("second subscription did not start after heartbeat loss") + } + service.homeMu.Lock() + exposedRegistry := service.homeRegistry + exposedClient := service.homeClient + service.homeMu.Unlock() + if exposedRegistry != nil || exposedClient != nil || home.Current() != nil { + t.Fatal("old subscriber lifetime remained exposed before the replacement ACK") + } + + close(allowSecondAck) + secondRegistry := waitForServiceRegistry(t, service, time.Second) + if secondRegistry == firstRegistry { + t.Fatal("replacement subscription reused the old registry") + } + if home.Current() == nil { + t.Fatal("replacement client was not exposed after the replacement ACK") + } +} + +func TestServiceDrainsBeforePreAckRetriesAndExposesOnlyAfterNewAck(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + resourceClosed := make(chan struct{}) + preAckAttempts := make(chan time.Time, 2) + finalSubscribe := make(chan struct{}) + allowFinalAck := make(chan struct{}) + stop := make(chan struct{}) + var configMu sync.Mutex + configRequests := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveSuccessChainHomeConnection(conn, &configMu, &configRequests, firstAck, loseFirst, preAckAttempts, finalSubscribe, allowFinalAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + firstRegistry := waitForServiceRegistry(t, service, time.Second) + pending, errBegin := firstRegistry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := firstRegistry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + if errBind := scope.Bind(func() error { + close(resourceClosed) + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + close(loseFirst) + select { + case <-resourceClosed: + case <-time.After(time.Second): + t.Fatal("heartbeat loss did not close the active scope resource") + } + select { + case <-preAckAttempts: + t.Fatal("pre-ACK retry started before the active scope owner ended") + case <-time.After(50 * time.Millisecond): + } + scope.End("canceled") + + firstPreAck := <-preAckAttempts + secondPreAck := <-preAckAttempts + if retryDelay := secondPreAck.Sub(firstPreAck); retryDelay < 75*time.Millisecond { + t.Fatalf("pre-ACK retry delay = %v, want at least 75ms", retryDelay) + } + select { + case <-finalSubscribe: + case <-time.After(time.Second): + t.Fatal("subscriber did not retry after pre-ACK rejections") + } + service.homeMu.Lock() + exposedRegistry := service.homeRegistry + exposedClient := service.homeClient + service.homeMu.Unlock() + if exposedRegistry != nil || exposedClient != nil || home.Current() != nil { + t.Fatal("new Home lifetime was exposed before its subscription ACK") + } + + close(allowFinalAck) + secondRegistry := waitForServiceRegistry(t, service, time.Second) + if secondRegistry == firstRegistry || home.Current() == nil { + t.Fatal("new Home lifetime was not exposed only after its subscription ACK") + } +} + +func TestServiceCancelsRunWhenBlockingScopeExceedsDrainBound(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + secondSubscribe := make(chan struct{}) + var secondSubscribeOnce sync.Once + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + var subscriptionMu sync.Mutex + subscriptions := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.Host = host + cfg.Home.Port = port + cfg.Home.DisableClusterDiscovery = true + service := &Service{cfg: cfg} + serviceCtx, cancelService := context.WithCancel(context.Background()) + t.Cleanup(cancelService) + service.homeMu.Lock() + service.runCancel = cancelService + service.homeMu.Unlock() + service.startHomeSubscriber(serviceCtx) + + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + registry := waitForServiceRegistry(t, service, time.Second) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + started := make(chan struct{}) + release := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(started) + <-release + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + close(loseFirst) + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("drain did not start closing the blocking scope") + } + select { + case <-secondSubscribe: + t.Fatal("new subscription started before the old registry drained") + case <-time.After(50 * time.Millisecond): + } + service.homeMu.Lock() + exposedRegistry := service.homeRegistry + service.homeMu.Unlock() + if exposedRegistry != nil { + t.Fatal("new registry was exposed before the old registry drained") + } + select { + case <-serviceCtx.Done(): + case <-time.After(time.Second): + t.Fatal("service run was not canceled after drain timeout") + } + + close(release) + scope.End("test cleanup") +} + +func TestServiceBacksOffAfterRepeatedPreAckFailures(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + attempts := make(chan time.Time, 8) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go servePreAckFailureConnection(conn, attempts) + } + }() + t.Cleanup(func() { + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.Host = host + cfg.Home.Port = port + cfg.Home.DisableClusterDiscovery = true + service := &Service{cfg: cfg} + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + + firstAttempt := <-attempts + secondAttempt := <-attempts + if retryDelay := secondAttempt.Sub(firstAttempt); retryDelay < 75*time.Millisecond { + t.Fatalf("pre-ACK retry delay = %v, want at least 75ms", retryDelay) + } + cancel() + select { + case thirdAttempt := <-attempts: + t.Fatalf("pre-ACK retry continued after cancellation at %v", thirdAttempt) + case <-time.After(150 * time.Millisecond): + } +} + +func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationBeforeDrain(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + update := make(chan struct{}) + statusStarted := make(chan struct{}) + statusRelease := make(chan struct{}) + secondConfig := make(chan struct{}) + var configRequests atomic.Int32 + var statusWrites atomic.Int32 + stop := make(chan struct{}) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveBlockedFinalizationConnection(conn, &configRequests, &statusWrites, update, statusStarted, statusRelease, secondConfig, stop) + } + }() + t.Cleanup(func() { + close(stop) + close(statusRelease) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + service.cfg.Home.NodeID = "node-1" + service.homePluginSyncKey = homePluginSyncKey(service.cfg) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + registry := waitForServiceRegistry(t, service, time.Second) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + resourceClosed := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(resourceClosed) + go scope.End("canceled") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + close(update) + select { + case <-statusStarted: + case <-time.After(time.Second): + t.Fatal("updated config did not enter blocked finalization") + } + select { + case <-resourceClosed: + case <-time.After(500 * time.Millisecond): + t.Fatal("heartbeat loss did not cancel the worker and drain the active execution") + } + select { + case <-secondConfig: + case <-time.After(time.Second): + t.Fatal("subscriber did not retry after heartbeat loss") + } + service.homeMu.Lock() + currentRegistry := service.homeRegistry + currentClient := service.homeClient + service.homeMu.Unlock() + if currentRegistry != nil || currentClient != nil || home.Current() != nil { + t.Fatal("heartbeat-lost lifetime left a published Home client or registry") + } +} + +func TestServiceConfigWorkerFinalizesRapidUpdatesInOrder(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + updates := make(chan struct{}) + statuses := make(chan homeplugins.SyncReport, 4) + var taskRequests atomic.Int32 + stop := make(chan struct{}) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveOrderedConfigUpdatesConnection(conn, &taskRequests, updates, statuses, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + service.cfg.Home.NodeID = "node-1" + service.homePluginSyncKey = homePluginSyncKey(service.cfg) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + waitForServiceRegistry(t, service, time.Second) + close(updates) + + gotTaskIDs := make([]uint, 0, 2) + for len(gotTaskIDs) < 2 { + select { + case report := <-statuses: + if report.TaskID != 0 { + gotTaskIDs = append(gotTaskIDs, report.TaskID) + } + case <-time.After(time.Second): + t.Fatal("rapid config updates did not finalize all ordered task work") + } + } + wantTaskIDs := []uint{1, 2} + for index := range wantTaskIDs { + if gotTaskIDs[index] != wantTaskIDs[index] { + t.Fatalf("plugin task status IDs = %v, want %v", gotTaskIDs, wantTaskIDs) + } + } +} + +func serveBlockedFinalizationConnection(conn net.Conn, configRequests, statusWrites *atomic.Int32, update <-chan struct{}, statusStarted chan<- struct{}, statusRelease <-chan struct{}, secondConfig chan<- struct{}, stop <-chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + if configRequests.Add(1) > 1 { + select { + case secondConfig <- struct{}{}: + case <-stop: + } + _, _ = io.WriteString(conn, "-ERR unavailable\r\n") + return + } + writeRegistryTestConfig(conn, "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\n") + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + _, _ = io.WriteString(conn, "$-1\r\n") + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-sync": + payload := fmt.Sprintf(`{"schema_version":1,"expires_at":%q,"items":[]}`, time.Now().Add(time.Minute).UTC().Format(time.RFC3339Nano)) + writeRegistryTestConfig(conn, payload) + case len(args) >= 2 && strings.EqualFold(args[0], "RPUSH") && args[1] == "plugin-status": + if statusWrites.Add(1) == 1 { + if _, errWrite := io.WriteString(conn, ":1\r\n"); errWrite != nil { + return + } + continue + } + select { + case statusStarted <- struct{}{}: + case <-stop: + return + } + select { + case <-statusRelease: + return + case <-stop: + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + select { + case <-update: + writeRegistryTestMessage(conn, "credential-concurrency:\n lifecycle-config-revision: 2\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\nplugins:\n enabled: true\n") + case <-stop: + return + } + <-stop + return + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func serveOrderedConfigUpdatesConnection(conn net.Conn, taskRequests *atomic.Int32, updates <-chan struct{}, statuses chan<- homeplugins.SyncReport, stop <-chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + _, _ = io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n") + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + writeRegistryTestConfig(conn, "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 1s\n cpa-cancel-bound: 100ms\n") + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-sync": + payload := fmt.Sprintf(`{"schema_version":1,"expires_at":%q,"items":[]}`, time.Now().Add(time.Minute).UTC().Format(time.RFC3339Nano)) + writeRegistryTestConfig(conn, payload) + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + request := taskRequests.Add(1) + if request == 1 { + _, _ = io.WriteString(conn, "$-1\r\n") + continue + } + payload := fmt.Sprintf(`[{"id":%d,"operation":"delete","plugin_id":"plugin-%d"}]`, request-1, request-1) + writeRegistryTestConfig(conn, payload) + case len(args) >= 3 && strings.EqualFold(args[0], "RPUSH") && args[1] == "plugin-status": + var report homeplugins.SyncReport + if errUnmarshal := json.Unmarshal([]byte(args[2]), &report); errUnmarshal != nil { + return + } + select { + case statuses <- report: + case <-stop: + return + } + _, _ = io.WriteString(conn, ":1\r\n") + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + select { + case <-updates: + writeRegistryTestMessage(conn, "credential-concurrency:\n lifecycle-config-revision: 2\n cpa-heartbeat-timeout: 1s\n cpa-cancel-bound: 100ms\nplugins:\n enabled: true\n") + writeRegistryTestMessage(conn, "credential-concurrency:\n lifecycle-config-revision: 3\n cpa-heartbeat-timeout: 1s\n cpa-cancel-bound: 100ms\n") + case <-stop: + return + } + <-stop + return + default: + _, _ = io.WriteString(conn, "+OK\r\n") + } + } +} + +func writeRegistryTestConfig(conn net.Conn, payload string) { + _, _ = io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)) +} + +func writeRegistryTestMessage(conn net.Conn, payload string) { + _, _ = io.WriteString(conn, fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$6\r\nconfig\r\n$%d\r\n%s\r\n", len(payload), payload)) +} + +func newRegistryTestService(t *testing.T, listener net.Listener) *Service { + t.Helper() + host, portText, errSplit := net.SplitHostPort(listener.Addr().String()) + if errSplit != nil { + t.Fatalf("split listener address: %v", errSplit) + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + t.Fatalf("parse port: %v", errPort) + } + cfg := &config.Config{} + cfg.Home.Enabled = true + cfg.Home.Host = host + cfg.Home.Port = port + cfg.Home.DisableClusterDiscovery = true + return &Service{cfg: cfg} +} + +func waitForServiceRegistry(t *testing.T, service *Service, timeout time.Duration) *executionregistry.Registry { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + service.homeMu.Lock() + registry := service.homeRegistry + service.homeMu.Unlock() + if registry != nil { + return registry + } + time.Sleep(time.Millisecond) + } + t.Fatal("service did not expose a ready execution registry") + return nil +} + +type testHomeLogForwarder struct { + mu sync.Mutex + owner *home.Client + binds int + deactivations int + stops atomic.Int32 +} + +func (f *testHomeLogForwarder) Bind(client *home.Client) { + f.mu.Lock() + defer f.mu.Unlock() + f.owner = client + f.binds++ +} + +func (f *testHomeLogForwarder) Deactivate(client *home.Client) { + f.mu.Lock() + defer f.mu.Unlock() + if f.owner == client { + f.owner = nil + } + f.deactivations++ +} + +func (f *testHomeLogForwarder) currentOwner() *home.Client { + f.mu.Lock() + defer f.mu.Unlock() + return f.owner +} + +func (f *testHomeLogForwarder) bindCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.binds +} + +func (f *testHomeLogForwarder) Stop() { + f.stops.Add(1) +} + +func TestServiceReusesHomeLogForwarderAcrossReconnects(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + acks := make(chan struct{}, 3) + stop := make(chan struct{}) + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveHomeLogForwarderReconnectConnection(conn, acks, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + forwarder := &testHomeLogForwarder{} + originalStart := startHomeLogForwarder + var starts atomic.Int32 + startHomeLogForwarder = func(int) homeLogForwarder { + starts.Add(1) + return forwarder + } + t.Cleanup(func() { startHomeLogForwarder = originalStart }) + + service := newRegistryTestService(t, listener) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + service.startHomeSubscriber(ctx) + waitForHomeLogForwarderACK(t, acks) + first := waitForServiceHomeClient(t, service, time.Second) + + service.startHomeSubscriber(ctx) + waitForHomeLogForwarderACK(t, acks) + second := waitForServiceHomeClient(t, service, time.Second) + if second == first { + t.Fatal("first reconnect reused the previous Home client") + } + + service.startHomeSubscriber(ctx) + waitForHomeLogForwarderACK(t, acks) + third := waitForServiceHomeClient(t, service, time.Second) + if third == second { + t.Fatal("second reconnect reused the previous Home client") + } + if got := starts.Load(); got != 1 { + t.Fatalf("Home log forwarder starts = %d, want 1", got) + } + if got := forwarder.bindCount(); got != 3 { + t.Fatalf("Home log forwarder binds = %d, want 3", got) + } + if owner := forwarder.currentOwner(); owner != third { + t.Fatal("Home log forwarder does not target the current Home client") + } + if current := home.Current(); current != third { + t.Fatal("current Home client does not match log forwarder owner") + } + + if errShutdown := service.Shutdown(context.Background()); errShutdown != nil { + t.Fatalf("Shutdown() error = %v", errShutdown) + } + if got := forwarder.stops.Load(); got != 1 { + t.Fatalf("Home log forwarder stops = %d, want 1", got) + } +} + +func waitForHomeLogForwarderACK(t *testing.T, acks <-chan struct{}) { + t.Helper() + select { + case <-acks: + case <-time.After(time.Second): + t.Fatal("Home subscription was not acknowledged") + } +} + +func waitForServiceHomeClient(t *testing.T, service *Service, timeout time.Duration) *home.Client { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + service.homeMu.Lock() + client := service.homeClient + service.homeMu.Unlock() + if client != nil { + return client + } + time.Sleep(time.Millisecond) + } + t.Fatal("service did not expose a Home client") + return nil +} + +func TestDetachHomeSubscriberLifetimeKeepsNewForwarderForStaleClient(t *testing.T) { + staleClient := home.New(internalconfig.HomeConfig{Enabled: true}) + currentClient := home.New(internalconfig.HomeConfig{Enabled: true}) + staleRegistry := executionregistry.New() + currentRegistry := executionregistry.New() + staleForwarder := &testHomeLogForwarder{} + currentForwarder := &testHomeLogForwarder{} + service := &Service{ + homeClient: currentClient, + homeRegistry: currentRegistry, + homeLogForwarder: currentForwarder, + homeLogForwarderClient: currentClient, + } + + staleForwarder.Stop() + service.detachHomeSubscriberLifetime(staleClient, staleRegistry) + + service.homeMu.Lock() + forwarder := service.homeLogForwarder + forwarderClient := service.homeLogForwarderClient + client := service.homeClient + registry := service.homeRegistry + service.homeMu.Unlock() + if forwarder != currentForwarder || forwarderClient != currentClient || client != currentClient || registry != currentRegistry { + t.Fatal("stale detach cleared the replacement Home lifetime") + } + if currentForwarder.stops.Load() != 0 { + t.Fatal("stale detach stopped the replacement log forwarder") + } + if staleForwarder.stops.Load() != 1 { + t.Fatal("stale forwarder ownership changed during stale detach") + } +} + +func serveHomeLogForwarderReconnectConnection(conn net.Conn, acks chan<- struct{}, stop <-chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + payload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\n" + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + if _, errWrite := io.WriteString(conn, "$2\r\n[]\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + acks <- struct{}{} + <-stop + return + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func waitForPublisherReplacementFrame(t *testing.T, frames <-chan home.InFlightSnapshotFrame, barrierRevision int64) home.InFlightSnapshotFrame { + t.Helper() + timer := time.NewTimer(time.Second) + defer timer.Stop() + for { + select { + case frame := <-frames: + if frame.BarrierRevision == barrierRevision { + return frame + } + case <-timer.C: + t.Fatalf("publisher did not send barrier revision %d", barrierRevision) + return home.InFlightSnapshotFrame{} + } + } +} + +func servePublisherReplacementConnection(conn net.Conn, configRequests *atomic.Int32, frames chan<- home.InFlightSnapshotFrame, firstPublisherDoneForServer <-chan (<-chan struct{}), secondConfigResult chan<- error, allowSecondConfig <-chan struct{}, stop <-chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + request := int(configRequests.Add(1)) + if request == 2 { + var firstPublisherDone <-chan struct{} + select { + case firstPublisherDone = <-firstPublisherDoneForServer: + case <-stop: + return + } + select { + case <-firstPublisherDone: + secondConfigResult <- nil + default: + secondConfigResult <- errors.New("replacement began its config lifetime before the previous publisher exited") + } + select { + case <-allowSecondConfig: + case <-stop: + return + } + } + barrierRevision := 11 + if request == 2 { + barrierRevision = 22 + } + payload := fmt.Sprintf("credential-concurrency:\n lifecycle-config-revision: %d\n observation-barrier-revision: %d\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\ncredential-in-flight:\n snapshot-interval: 10ms\n", request, barrierRevision) + writeRegistryTestConfig(conn, payload) + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + if _, errWrite := io.WriteString(conn, "$2\r\n[]\r\n"); errWrite != nil { + return + } + case len(args) > 0 && strings.EqualFold(args[0], "PING"): + if _, errWrite := io.WriteString(conn, "+PONG\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + select { + case <-stop: + return + case <-time.After(time.Second): + return + } + case len(args) >= 3 && strings.EqualFold(args[0], "LPUSH") && args[1] == "in-flight-snapshot": + var frame home.InFlightSnapshotFrame + if errUnmarshal := json.Unmarshal([]byte(args[2]), &frame); errUnmarshal != nil { + return + } + select { + case frames <- frame: + case <-stop: + return + } + if _, errWrite := io.WriteString(conn, ":1\r\n"); errWrite != nil { + return + } + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func servePreACKReplacementConnection(conn net.Conn, configRequests *atomic.Int32, firstSubscribed chan struct{}, secondStarted chan struct{}, secondStartedBeforeFirstDone chan struct{}, firstDone <-chan (<-chan struct{}), stop chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + if configRequests.Add(1) > 1 { + supervisorDone := <-firstDone + select { + case <-supervisorDone: + default: + close(secondStartedBeforeFirstDone) + } + close(secondStarted) + } + payload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\n" + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + if configRequests.Load() == 1 { + close(firstSubscribed) + } + select { + case <-stop: + return + case <-time.After(time.Second): + return + } + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func serveSuccessChainHomeConnection(conn net.Conn, configMu *sync.Mutex, configRequests *int, firstAck chan struct{}, loseFirst chan struct{}, preAckAttempts chan time.Time, finalSubscribe chan struct{}, allowFinalAck chan struct{}, stop chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + configMu.Lock() + *configRequests++ + request := *configRequests + configMu.Unlock() + if request == 2 || request == 3 { + preAckAttempts <- time.Now() + _, _ = io.WriteString(conn, "-ERR unavailable\r\n") + return + } + payload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\n" + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + if _, errWrite := io.WriteString(conn, "$2\r\n[]\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + configMu.Lock() + request := *configRequests + configMu.Unlock() + if request == 1 { + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + close(firstAck) + select { + case <-loseFirst: + <-stop + case <-stop: + } + return + } + close(finalSubscribe) + select { + case <-allowFinalAck: + case <-stop: + return + } + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + <-stop + return + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func serveStalePreACKPluginConnection(conn net.Conn, subscriptions *atomic.Int32, pluginWrites *atomic.Int32, firstSubscribed chan struct{}, secondSubscribed chan struct{}, allowSecondAck chan struct{}, stop chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + payload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\nplugins:\n enabled: true\n" + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-sync": + payload := fmt.Sprintf(`{"schema_version":1,"expires_at":%q,"items":[]}`, time.Now().Add(time.Minute).UTC().Format(time.RFC3339Nano)) + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + if _, errWrite := io.WriteString(conn, "$-1\r\n"); errWrite != nil { + return + } + case len(args) > 0 && strings.EqualFold(args[0], "PING"): + if _, errWrite := io.WriteString(conn, "+PONG\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "RPUSH") && args[1] == "plugin-status": + pluginWrites.Add(1) + if _, errWrite := io.WriteString(conn, ":1\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + subscription := subscriptions.Add(1) + switch subscription { + case 1: + close(firstSubscribed) + <-stop + return + case 2: + close(secondSubscribed) + select { + case <-allowSecondAck: + case <-stop: + return + } + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + <-stop + return + } + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func serveInitialOverlayPluginConnection(conn net.Conn, pluginSync chan struct{}, pluginStatus chan struct{}, pluginTasks chan struct{}, freshCommandProbe chan struct{}, allowAck chan struct{}, stop chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + payload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\nplugins:\n enabled: true\n" + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-sync": + if home.Current() != nil { + return + } + close(pluginSync) + payload := fmt.Sprintf(`{"schema_version":1,"expires_at":%q,"items":[]}`, time.Now().Add(time.Minute).UTC().Format(time.RFC3339Nano)) + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "RPUSH") && args[1] == "plugin-status": + select { + case <-freshCommandProbe: + default: + return + } + pluginStatus <- struct{}{} + if _, errWrite := io.WriteString(conn, ":1\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + close(pluginTasks) + payload := `[{"id":1,"operation":"delete","plugin_id":"plugin-a"}]` + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) > 0 && strings.EqualFold(args[0], "PING"): + close(freshCommandProbe) + if _, errWrite := io.WriteString(conn, "+PONG\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + select { + case <-allowAck: + case <-stop: + return + } + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + <-stop + return + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func serveRegistryTestHomeConnection(conn net.Conn, subscriptionMu *sync.Mutex, subscriptions *int, firstAck chan struct{}, loseFirst chan struct{}, secondSubscribe chan struct{}, secondSubscribeOnce *sync.Once, allowSecondAck chan struct{}, stop chan struct{}) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + payload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 100ms\n cpa-cancel-bound: 100ms\n" + if _, errWrite := io.WriteString(conn, fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "plugin-tasks": + if _, errWrite := io.WriteString(conn, "$2\r\n[]\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "SUBSCRIBE") && args[1] == "config": + subscriptionMu.Lock() + *subscriptions++ + subscription := *subscriptions + subscriptionMu.Unlock() + if subscription == 1 { + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + close(firstAck) + select { + case <-loseFirst: + <-stop + return + case <-stop: + return + } + } + secondSubscribeOnce.Do(func() { close(secondSubscribe) }) + select { + case <-allowSecondAck: + case <-stop: + return + } + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + <-stop + return + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func servePreAckFailureConnection(conn net.Conn, attempts chan<- time.Time) { + defer func() { _ = conn.Close() }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + switch { + case len(args) > 0 && strings.EqualFold(args[0], "HELLO"): + if _, errWrite := io.WriteString(conn, "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n"); errWrite != nil { + return + } + case len(args) >= 2 && strings.EqualFold(args[0], "GET") && args[1] == "config": + attempts <- time.Now() + _, _ = io.WriteString(conn, "-ERR unavailable\r\n") + return + default: + if _, errWrite := io.WriteString(conn, "+OK\r\n"); errWrite != nil { + return + } + } + } +} + +func readRegistryTestRedisCommand(reader *bufio.Reader) ([]string, error) { + line, errRead := reader.ReadString('\n') + if errRead != nil { + return nil, errRead + } + if !strings.HasPrefix(line, "*") { + return nil, fmt.Errorf("unexpected RESP command header %q", line) + } + count, errCount := strconv.Atoi(strings.TrimSpace(strings.TrimPrefix(line, "*"))) + if errCount != nil { + return nil, errCount + } + args := make([]string, 0, count) + for range count { + lengthLine, errLength := reader.ReadString('\n') + if errLength != nil { + return nil, errLength + } + length, errParseLength := strconv.Atoi(strings.TrimSpace(strings.TrimPrefix(lengthLine, "$"))) + if errParseLength != nil { + return nil, errParseLength + } + raw := make([]byte, length+2) + if _, errReadRaw := io.ReadFull(reader, raw); errReadRaw != nil { + return nil, errReadRaw + } + args = append(args, string(raw[:length])) + } + return args, nil +} + +func TestServiceSkipsStaleLocalConfigRuntimeApply(t *testing.T) { + service := &Service{cfg: &config.Config{}} + var applied []string + service.applyPprofConfigContextFn = func(_ context.Context, cfg *config.Config) bool { + applied = append(applied, cfg.Routing.Strategy) + return true + } + first := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{Strategy: "fill-first"}}) + second := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{Strategy: "round-robin"}}) + if !service.applyConfigRuntime(context.Background(), second, false) { + t.Fatal("newest config runtime apply failed") + } + if service.applyConfigRuntime(context.Background(), first, false) { + t.Fatal("stale config runtime apply succeeded") + } + if got, want := strings.Join(applied, ","), "round-robin"; got != want { + t.Fatalf("runtime apply order = %q, want %q", got, want) + } +} + +func TestServiceAppliesSameValueNewestSelectorCommit(t *testing.T) { + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + manager.RegisterExecutor(serviceTestPluginExecutor{}) + for _, id := range []string{"auth-b", "auth-a"} { + if _, errRegister := manager.Register(context.Background(), &coreauth.Auth{ID: id, Provider: "plugin-provider", Status: coreauth.StatusActive}); errRegister != nil { + t.Fatalf("Register(%s) error = %v", id, errRegister) + } + } + + service := &Service{cfg: &config.Config{}, coreManager: manager} + older := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{Strategy: "fill-first"}}) + newer := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{Strategy: "fill-first"}}) + if !service.applyConfigRuntime(context.Background(), newer, false) { + t.Fatal("newest same-value config runtime apply failed") + } + if service.applyConfigRuntime(context.Background(), older, false) { + t.Fatal("stale same-value config runtime apply succeeded") + } + + for range 2 { + selected, errSelect := manager.SelectAuth(context.Background(), "plugin-provider", "", cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectAuth() error = %v", errSelect) + } + if selected == nil || selected.ID != "auth-a" { + t.Fatalf("selector picked = %+v, want auth-a from fill-first", selected) + } + } +} + +func TestBuilderPreservesInitialSelectorForSameRouting(t *testing.T) { + cfg := &config.Config{ + AuthDir: t.TempDir(), + Routing: internalconfig.RoutingConfig{ + Strategy: "fill-first", + SessionAffinity: true, + SessionAffinityTTL: "1h", + }, + } + service, errBuild := NewBuilder(). + WithConfig(cfg). + WithConfigPath(t.TempDir() + "/config.yaml"). + Build() + if errBuild != nil { + t.Fatalf("Build() error = %v", errBuild) + } + + initialSelector := service.coreManager.Selector() + initialAffinity, ok := initialSelector.(*coreauth.SessionAffinitySelector) + if !ok { + t.Fatalf("initial selector = %T, want *SessionAffinitySelector", initialSelector) + } + defer initialAffinity.Stop() + commit := service.commitConfigUpdate(cfg) + if !service.applyConfigRuntime(context.Background(), commit, false) { + t.Fatal("same-routing config runtime apply failed") + } + if got := service.coreManager.Selector(); got != initialSelector { + t.Fatalf("same-routing selector = %p, want initial selector %p", got, initialSelector) + } +} + +func TestServiceApplyConfigRuntimePreservesSelectorForUnchangedRouting(t *testing.T) { + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + service := &Service{cfg: &config.Config{}, coreManager: manager} + + initial := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "fill-first", + SessionAffinity: true, + SessionAffinityTTL: "1h", + }}) + if !service.applyConfigRuntime(context.Background(), initial, false) { + t.Fatal("initial config runtime apply failed") + } + initialSelector := manager.Selector() + initialAffinity, ok := initialSelector.(*coreauth.SessionAffinitySelector) + if !ok { + t.Fatalf("initial selector = %T, want *SessionAffinitySelector", initialSelector) + } + defer initialAffinity.Stop() + + older := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: " FILLFIRST ", + SessionAffinity: true, + SessionAffinityTTL: "60m", + }}) + newer := service.commitConfigUpdate(&config.Config{ + Routing: internalconfig.RoutingConfig{ + Strategy: "fill-first", + SessionAffinity: true, + SessionAffinityTTL: "1h", + }, + UsageStatisticsEnabled: true, + }) + if !service.applyConfigRuntime(context.Background(), newer, false) { + t.Fatal("newest same-routing config runtime apply failed") + } + if got := manager.Selector(); got != initialSelector { + t.Fatalf("same-routing selector = %p, want original %p", got, initialSelector) + } + if service.applyConfigRuntime(context.Background(), older, false) { + t.Fatal("stale same-routing config runtime apply succeeded") + } + if got := manager.Selector(); got != initialSelector { + t.Fatalf("stale same-routing selector = %p, want original %p", got, initialSelector) + } + + changed := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: true, + SessionAffinityTTL: "1h", + }}) + if !service.applyConfigRuntime(context.Background(), changed, false) { + t.Fatal("changed-routing config runtime apply failed") + } + changedSelector := manager.Selector() + if changedSelector == initialSelector { + t.Fatal("changed-routing selector retained original identity") + } + changedAffinity, ok := changedSelector.(*coreauth.SessionAffinitySelector) + if !ok { + t.Fatalf("changed selector = %T, want *SessionAffinitySelector", changedSelector) + } + defer changedAffinity.Stop() + + unrelated := service.commitConfigUpdate(&config.Config{ + Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: true, + SessionAffinityTTL: "1h", + }, + UsageStatisticsEnabled: false, + }) + if !service.applyConfigRuntime(context.Background(), unrelated, false) { + t.Fatal("unrelated config runtime apply failed") + } + if got := manager.Selector(); got != changedSelector { + t.Fatalf("unrelated-update selector = %p, want changed selector %p", got, changedSelector) + } +} + +func TestServiceSerializesHomeAndWatcherConfigRuntimeApply(t *testing.T) { + baseCfg := &config.Config{} + baseCfg.Home.Enabled = true + service := &Service{cfg: baseCfg, homeGeneration: 1} + firstStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + var appliedMu sync.Mutex + var applied []string + service.applyPprofConfigContextFn = func(_ context.Context, cfg *config.Config) bool { + if cfg.Routing.Strategy == "fill-first" { + close(firstStarted) + <-releaseFirst + } + appliedMu.Lock() + applied = append(applied, cfg.Routing.Strategy) + appliedMu.Unlock() + return true + } + client, _ := newHomePluginTaskTestClient(t, nil, 0) + queue := newHomeConfigWorkQueue() + queue.enqueue([]byte("routing:\n strategy: fill-first\n")) + ready := make(chan struct{}) + close(ready) + lifetimeCtx, cancelLifetime := context.WithCancel(context.Background()) + defer cancelLifetime() + cancelBound := atomic.Int64{} + cancelBound.Store(int64(time.Second)) + workerDone := make(chan struct{}) + go func() { + defer close(workerDone) + service.runHomeConfigWorker(lifetimeCtx, context.Background(), 1, client, executionregistry.New(), queue, ready, &atomic.Bool{}, &cancelBound) + }() + select { + case <-firstStarted: + case <-time.After(time.Second): + t.Fatal("Home config runtime apply did not start") + } + + watcherDone := make(chan struct{}) + go func() { + service.applyWatcherConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{Strategy: "round-robin"}}) + close(watcherDone) + }() + select { + case <-watcherDone: + t.Fatal("watcher runtime apply completed before the older Home apply") + case <-time.After(100 * time.Millisecond): + } + close(releaseFirst) + select { + case <-watcherDone: + case <-time.After(time.Second): + t.Fatal("watcher runtime apply did not finish") + } + appliedMu.Lock() + got := strings.Join(applied, ",") + appliedMu.Unlock() + if want := "fill-first,round-robin"; got != want { + t.Fatalf("runtime completion order = %q, want %q", got, want) + } + cancelLifetime() + select { + case <-workerDone: + case <-time.After(time.Second): + t.Fatal("Home config worker did not stop") + } +} diff --git a/sdk/cliproxy/service_stale_state_test.go b/sdk/cliproxy/service_stale_state_test.go --- a/sdk/cliproxy/service_stale_state_test.go +++ b/sdk/cliproxy/service_stale_state_test.go @@ -5,8 +5,10 @@ "testing" "time" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" ) @@ -87,8 +89,27 @@ } } -func TestApplyHomeOverlayForcesUsageStatisticsEnabled(t *testing.T) { - baseCfg := &config.Config{} +func TestLifetimeRegistryObservesBarrierFromAppliedHomeConfig(t *testing.T) { + registry := executionregistry.New() + manager := coreauth.NewManager(nil, nil, nil) + cfg := internalconfig.DefaultCredentialInFlightConfig() + cfg.SnapshotInterval = "30ms" + + if errApply := applyHomeInFlightPublisherConfig(manager, cfg); errApply != nil { + t.Fatal(errApply) + } + applyHomeObservationBarrier(registry, 14) + + if freeze := registry.FreezeInFlight(time.Now().UTC()); freeze.BarrierRevision != 14 { + t.Fatalf("barrier revision = %d, want 14", freeze.BarrierRevision) + } + if got := manager.HomeInFlightPublisherConfig(); got.SnapshotInterval != 30*time.Millisecond { + t.Fatalf("publisher interval = %v, want 30ms", got.SnapshotInterval) + } +} + +func TestApplyHomeOverlayDoesNotApplyWithoutReadyClient(t *testing.T) { + baseCfg := &config.Config{UsageStatisticsEnabled: false, SaveCooldownStatus: true} baseCfg.Home.Enabled = true service := &Service{cfg: baseCfg} @@ -97,13 +118,13 @@ SaveCooldownStatus: true, }) - if service.cfg == nil || !service.cfg.UsageStatisticsEnabled { - t.Fatal("expected home overlay to force usage statistics enabled") + if service.cfg == nil || service.cfg.UsageStatisticsEnabled { + t.Fatal("unready home overlay changed usage statistics") } if !service.cfg.Home.Enabled { - t.Fatal("expected home overlay to preserve local home settings") + t.Fatal("unready home overlay changed local home settings") } - if service.cfg.SaveCooldownStatus { - t.Fatal("expected home overlay to force cooldown status persistence disabled") + if !service.cfg.SaveCooldownStatus { + t.Fatal("unready home overlay changed cooldown status persistence") } } diff --git a/sdk/pluginhost/host.go b/sdk/pluginhost/host.go --- a/sdk/pluginhost/host.go +++ b/sdk/pluginhost/host.go @@ -79,10 +79,15 @@ // ShutdownAll unloads every active plugin. func (h *Host) ShutdownAll() { + h.ShutdownAllContext(context.Background()) +} + +// ShutdownAllContext detaches every active plugin and bounds waiting for active calls by ctx. +func (h *Host) ShutdownAllContext(ctx context.Context) { if h == nil || h.inner == nil { return } - h.inner.ShutdownAll() + h.inner.ShutdownAllContext(ctx) } // PluginBusy reports whether a plugin dynamic library is loaded or being loaded. @@ -92,10 +97,15 @@ // UnloadPlugin removes one plugin from the active runtime and closes its dynamic library. func (h *Host) UnloadPlugin(id string) bool { + return h.UnloadPluginContext(context.Background(), id) +} + +// UnloadPluginContext detaches one plugin and bounds waiting for active calls by ctx. +func (h *Host) UnloadPluginContext(ctx context.Context, id string) bool { if h == nil || h.inner == nil { return false } - return h.inner.UnloadPlugin(id) + return h.inner.UnloadPluginContext(ctx, id) } // ParseAuth lets plugin auth providers parse a credential payload. diff --git a/internal/home/testdata/concurrency_dispatch_accounted.json b/internal/home/testdata/concurrency_dispatch_accounted.json new file mode 100644 --- /dev/null +++ b/internal/home/testdata/concurrency_dispatch_accounted.json @@ -0,0 +1,27 @@ +{ + "model": "gpt", + "provider": "codex", + "auth_index": "cred-1", + "user_api_key": "user-key", + "auth": { + "id": "cred-1", + "provider": "codex", + "status": "active", + "disabled": false, + "unavailable": false, + "quota": { + "exceeded": false, + "next_recover_at": "0001-01-01T00:00:00Z" + }, + "created_at": "0001-01-01T00:00:00Z", + "updated_at": "0001-01-01T00:00:00Z", + "last_refreshed_at": "0001-01-01T00:00:00Z", + "next_refresh_after": "0001-01-01T00:00:00Z", + "next_retry_after": "0001-01-01T00:00:00Z" + }, + "concurrency": { + "accounted": true, + "credential_id": "cred-1", + "model": "gpt" + } +} diff --git a/internal/home/testdata/concurrency_dispatch_busy.json b/internal/home/testdata/concurrency_dispatch_busy.json new file mode 100644 --- /dev/null +++ b/internal/home/testdata/concurrency_dispatch_busy.json @@ -0,0 +1,8 @@ +{ + "error": { + "type": "credential_concurrency_exceeded", + "message": "credential concurrency limit reached", + "retryable": true, + "retry_after_ms": 750 + } +} diff --git a/internal/home/testdata/concurrency_release.json b/internal/home/testdata/concurrency_release.json new file mode 100644 --- /dev/null +++ b/internal/home/testdata/concurrency_release.json @@ -0,0 +1,1 @@ +{"credential_id":"cred-1","model":"gpt","release_seq":1} diff --git a/internal/home/testdata/credential_in_flight_contract.json b/internal/home/testdata/credential_in_flight_contract.json new file mode 100644 --- /dev/null +++ b/internal/home/testdata/credential_in_flight_contract.json @@ -0,0 +1,52 @@ +{ + "config": { + "snapshot-interval": "2s", + "stale-after": "10s", + "max-part-bytes": 262144, + "max-part-count": 64, + "max-revision-bytes": 16777216, + "max-aggregate-groups": 100000, + "max-details": 10000, + "max-string-bytes": 256, + "staging-retention": "1m" + }, + "part": { + "kind": "part", + "revision": 7, + "observed_at": "2026-07-21T12:00:00Z", + "barrier_revision": 11, + "part_index": 0, + "part_count": 1, + "details_truncated": false, + "aggregates": [ + { + "credential_id": "cred-a", + "model": "gpt-5", + "status": "accounted", + "count": 2 + }, + { + "credential_id": "cred-a", + "model": "gpt-5", + "status": "unaccounted", + "count": 1 + } + ], + "details": [ + { + "request_id": "req-1", + "credential_id": "cred-a", + "model": "gpt-5", + "request_kind": "sse", + "started_at": "2026-07-21T11:59:58Z" + } + ] + }, + "overflow": { + "kind": "overflow", + "revision": 8, + "observed_at": "2026-07-21T12:00:02Z", + "barrier_revision": 12, + "aggregate_group_count": 100001 + } +} diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -4,6 +4,7 @@ "context" "encoding/json" "errors" + "fmt" "io" "net/http" "net/http/httptest" @@ -25,6 +26,17 @@ antigravityShortCooldownByAuth = sync.Map{} antigravityCreditsBalanceByAuth = sync.Map{} antigravityCreditsHintRefreshByID = sync.Map{} +} + +type closeSignalReadCloser struct { + io.ReadCloser + closed chan<- struct{} +} + +func (c *closeSignalReadCloser) Close() error { + errClose := c.ReadCloser.Close() + close(c.closed) + return errClose } type fakeAntigravityKVClient struct { @@ -338,7 +350,7 @@ QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true}, }) auth := &cliproxyauth.Auth{ - ID: "auth-credits-conductor", + ID: fmt.Sprintf("auth-credits-conductor-%d", time.Now().UnixNano()), Attributes: map[string]string{ "base_url": server.URL, }, @@ -361,6 +373,16 @@ if err != nil { t.Fatalf("Execute() error = %v", err) } + stateValue, ok := antigravityCreditsHintRefreshByID.Load(auth.ID) + if !ok { + t.Fatal("expected credits refresh state") + } + state, ok := stateValue.(*antigravityCreditsHintRefreshState) + if !ok || state == nil { + t.Fatal("credits refresh state has unexpected type") + } + state.mu.Lock() + state.mu.Unlock() if len(resp.Payload) == 0 { t.Fatal("Execute() returned empty payload") } @@ -624,12 +646,13 @@ QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true}, }) auth := &cliproxyauth.Auth{ - ID: "auth-warm-token-credits", + ID: fmt.Sprintf("auth-warm-token-credits-%d", time.Now().UnixNano()), Metadata: map[string]any{ "access_token": "token", "expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339), }, } + refreshDone := make(chan struct{}) ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) { if req.URL.String() != "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" { t.Fatalf("unexpected request url %s", req.URL.String()) @@ -637,7 +660,10 @@ return &http.Response{ StatusCode: http.StatusOK, Header: make(http.Header), - Body: io.NopCloser(strings.NewReader(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`)), + Body: &closeSignalReadCloser{ + ReadCloser: io.NopCloser(strings.NewReader(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`)), + closed: refreshDone, + }, }, nil })) @@ -651,9 +677,10 @@ if updatedAuth != nil { t.Fatalf("ensureAccessToken() updatedAuth = %v, want nil", updatedAuth) } - deadline := time.Now().Add(2 * time.Second) - for time.Now().Before(deadline) && !cliproxyauth.HasKnownAntigravityCreditsHint(auth.ID) { - time.Sleep(10 * time.Millisecond) + select { + case <-refreshDone: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for background credits refresh") } if !cliproxyauth.HasKnownAntigravityCreditsHint(auth.ID) { t.Fatal("expected credits hint to be populated for warm token auth") diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -59,15 +59,42 @@ sessions: make(map[string]*codexWebsocketSession), } +type websocketConnectionCloser struct { + conn *websocket.Conn + once sync.Once + err error +} + +func newWebsocketConnectionCloser(conn *websocket.Conn) *websocketConnectionCloser { + if conn == nil { + return nil + } + return &websocketConnectionCloser{conn: conn} +} + +func (c *websocketConnectionCloser) Close() error { + if c == nil || c.conn == nil { + return nil + } + c.once.Do(func() { + c.err = c.conn.Close() + }) + return c.err +} + type codexWebsocketSession struct { sessionID string reqMu sync.Mutex - connMu sync.Mutex - conn *websocket.Conn - wsURL string - authID string + connMu sync.Mutex + conn *websocket.Conn + connCloser *websocketConnectionCloser + wsURL string + authID string + lifecycleBindMu sync.Mutex + lifecycle cliproxyexecutor.ExecutionLifecycle + lifecycleModel string writeMu sync.Mutex @@ -141,6 +168,13 @@ return s.activeCh, s.activeDone } +func clearRetryActiveState(sess *codexWebsocketSession, conn *websocket.Conn, ch chan codexWebsocketRead) bool { + if sess == nil { + return false + } + return sess.clearActive(conn, ch) +} + func (s *codexWebsocketSession) clearActive(conn *websocket.Conn, ch chan codexWebsocketRead) bool { if s == nil { return false @@ -211,6 +245,100 @@ }) } +func (s *codexWebsocketSession) bindExecutionLifecycle(opts cliproxyexecutor.Options, conn *websocket.Conn, closer *websocketConnectionCloser, model string) error { + if closer == nil { + return fmt.Errorf("codex websockets executor: websocket connection closer is nil") + } + if s == nil { + return cliproxyexecutor.BindExecutionResource(opts, closer) + } + lifecycle := opts.ExecutionLifecycle + if lifecycle == nil || conn == nil { + return nil + } + + s.lifecycleBindMu.Lock() + defer s.lifecycleBindMu.Unlock() + + s.connMu.Lock() + if s.conn == conn && s.connCloser == nil { + s.connCloser = closer + } + alreadyBound := s.conn == conn && s.connCloser == closer && s.lifecycle == lifecycle + s.connMu.Unlock() + if alreadyBound { + return nil + } + + if errBind := lifecycle.Bind(func() error { + return s.closeBoundConnection(conn, closer, lifecycle) + }); errBind != nil { + return errBind + } + if retained, ok := lifecycle.(interface{ Retain() }); ok { + retained.Retain() + } + + s.connMu.Lock() + if s.conn != conn || s.connCloser != closer { + s.connMu.Unlock() + return fmt.Errorf("codex websockets executor: websocket connection closed during lifecycle bind") + } + previous := s.lifecycle + s.lifecycle = lifecycle + s.lifecycleModel = strings.TrimSpace(model) + s.connMu.Unlock() + if previous != nil && previous != lifecycle { + previous.End("target_replaced") + } + return nil +} + +func (s *codexWebsocketSession) closeBoundConnection(conn *websocket.Conn, closer *websocketConnectionCloser, lifecycle cliproxyexecutor.ExecutionLifecycle) error { + if s == nil || conn == nil { + return nil + } + s.detachConnection(conn, lifecycle) + errClose := closer.Close() + go lifecycle.End("connection_closed") + return errClose +} + +func (s *codexWebsocketSession) detachConnection(conn *websocket.Conn, lifecycle cliproxyexecutor.ExecutionLifecycle) *websocketConnectionCloser { + if s == nil || conn == nil { + return nil + } + s.connMu.Lock() + var closer *websocketConnectionCloser + matched := s.conn == conn + if matched { + closer = s.connCloser + s.conn = nil + s.connCloser = nil + if s.readerConn == conn { + s.readerConn = nil + } + } + if (lifecycle == nil && matched) || (lifecycle != nil && s.lifecycle == lifecycle) { + s.lifecycle = nil + s.lifecycleModel = "" + } + s.connMu.Unlock() + return closer +} + +func closeWebsocketAfterBindFailure(sess *codexWebsocketSession, conn *websocket.Conn, closer *websocketConnectionCloser) { + if conn == nil || closer == nil { + return + } + if sess != nil { + sess.detachConnection(conn, nil) + } + if errClose := closer.Close(); errClose != nil { + log.Errorf("websockets executor: close lifecycle bind failure connection error: %v", errClose) + } +} + func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, wsURL string) bool { if sess == nil { return false @@ -224,25 +352,30 @@ return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) || strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL) } -func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, string, string) { +func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, *websocketConnectionCloser, string, string, cliproxyexecutor.ExecutionLifecycle) { if sess == nil { - return nil, "", "" + return nil, nil, "", "", nil } sess.connMu.Lock() defer sess.connMu.Unlock() conn := sess.conn if conn == nil || (strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL)) { - return nil, "", "" + return nil, nil, "", "", nil } previousAuthID := sess.authID previousWSURL := sess.wsURL + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" sess.conn = nil + sess.connCloser = nil if sess.readerConn == conn { sess.readerConn = nil } - return conn, previousAuthID, previousWSURL + return conn, closer, previousAuthID, previousWSURL, lifecycle } func (s *codexWebsocketSession) resetUpstreamDisconnectError(conn *websocket.Conn) { @@ -371,10 +504,18 @@ executionSessionID := executionSessionIDFromOptions(opts) var sess *codexWebsocketSession + sessionLocked := false + unlockSession := func() { + if sess != nil && sessionLocked { + sess.reqMu.Unlock() + sessionLocked = false + } + } if executionSessionID != "" { sess = e.getOrCreateSession(executionSessionID) sess.reqMu.Lock() - defer sess.reqMu.Unlock() + sessionLocked = true + defer unlockSession() } wsReqBody := buildCodexWebsocketRequestBody(upstreamBody) @@ -391,13 +532,16 @@ } helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) - conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDial != nil { bodyErr := websocketHandshakeBody(respHS) if respHS != nil { helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) } if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { + if opts.ExecutionLifecycle != nil { + return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } return e.CodexExecutor.Execute(ctx, auth, req, opts) } if respHS != nil && respHS.StatusCode > 0 { @@ -405,6 +549,11 @@ } helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) return resp, errDial + } + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + unlockSession() + closeWebsocketAfterBindFailure(sess, conn, closer) + return resp, errBind } recordAPIWebsocketHandshake(ctx, e.cfg, respHS) reporter.StartResponseTTFT() @@ -416,7 +565,7 @@ reason = "error" } logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, reason, err) - if errClose := conn.Close(); errClose != nil { + if errClose := closer.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } }() @@ -442,9 +591,17 @@ // Retry once with a fresh websocket connection. This is mainly to handle // upstream closing the socket between sequential requests within the same // execution session. - connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry == nil && connRetry != nil { + previousConn, previousReadCh := conn, readCh conn = connRetry + closer = closerRetry + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + clearRetryActiveState(sess, previousConn, previousReadCh) + unlockSession() + closeWebsocketAfterBindFailure(sess, conn, closer) + return resp, errBind + } readCh = sess.activate(conn) wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ @@ -522,6 +679,10 @@ return resp, wsErr } if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { + if sess != nil { + unlockSession() + e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) + } if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { return resp, errClearReplay } @@ -627,6 +788,13 @@ sess.reqMu.Lock() } } + streamSessionLocked := sess != nil + unlockStreamSession := func() { + if sess != nil && streamSessionLocked { + sess.reqMu.Unlock() + streamSessionLocked = false + } + } wsReqBody := buildCodexWebsocketRequestBody(upstreamBody) wsReqLog := helps.UpstreamRequestLog{ @@ -642,7 +810,7 @@ } helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) - conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) var upstreamHeaders http.Header if respHS != nil { upstreamHeaders = respHS.Header.Clone() @@ -653,9 +821,18 @@ helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) } if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { + if sess != nil { + sess.reqMu.Unlock() + } + if opts.ExecutionLifecycle != nil { + return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} + } return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts) } if respHS != nil && respHS.StatusCode > 0 { + if sess != nil { + sess.reqMu.Unlock() + } return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} } helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) @@ -663,6 +840,13 @@ sess.reqMu.Unlock() } return nil, errDial + } + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + if sess != nil { + sess.reqMu.Unlock() + } + closeWebsocketAfterBindFailure(sess, conn, closer) + return nil, errBind } recordAPIWebsocketHandshake(ctx, e.cfg, respHS) reporter.StartResponseTTFT() @@ -688,7 +872,7 @@ } // Retry once with a new websocket connection for the same execution session. - connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry != nil || connRetry == nil { closeHTTPResponseBody(respHSRetry, "codex websockets executor: close handshake response body error") helps.RecordAPIWebsocketError(ctx, e.cfg, "dial_retry", errDialRetry) @@ -696,7 +880,15 @@ sess.reqMu.Unlock() return nil, errDialRetry } + previousConn, previousReadCh := conn, readCh conn = connRetry + closer = closerRetry + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + clearRetryActiveState(sess, previousConn, previousReadCh) + sess.reqMu.Unlock() + closeWebsocketAfterBindFailure(sess, conn, closer) + return nil, errBind + } readCh = sess.activate(conn) wsReqBodyRetry := buildCodexWebsocketRequestBody(upstreamBody) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ @@ -723,7 +915,7 @@ wsReqBody = wsReqBodyRetry } else { logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) - if errClose := conn.Close(); errClose != nil { + if errClose := closer.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } return nil, errSend @@ -739,11 +931,11 @@ defer func() { if sess != nil { sess.clearActive(conn, readCh) - sess.reqMu.Unlock() + unlockStreamSession() return } logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, terminateReason, terminateErr) - if errClose := conn.Close(); errClose != nil { + if errClose := closer.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } }() @@ -832,6 +1024,10 @@ if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { terminateReason = "upstream_error" terminateErr = streamErr + if sess != nil { + unlockStreamSession() + e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) + } if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { terminateErr = errClearReplay helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) @@ -897,7 +1093,7 @@ return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil } -func (e *CodexWebsocketsExecutor) dialCodexWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *http.Response, error) { +func (e *CodexWebsocketsExecutor) dialCodexWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { dialer := newProxyAwareWebsocketDialer(e.cfg, auth) dialer.HandshakeTimeout = codexResponsesWebsocketHandshakeTO dialer.EnableCompression = true @@ -905,12 +1101,13 @@ ctx = context.Background() } conn, resp, err := dialer.DialContext(ctx, wsURL, headers) + closer := newWebsocketConnectionCloser(conn) if conn != nil { // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. // Negotiating permessage-deflate is fine; we just don't compress outbound messages. conn.EnableWriteCompression(false) } - return conn, resp, err + return conn, closer, resp, err } func writeCodexWebsocketMessage(sess *codexWebsocketSession, conn *websocket.Conn, payload []byte) error { @@ -1646,20 +1843,26 @@ return sess.upstreamDisconnectCh } -func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *http.Response, error) { +func (e *CodexWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { if sess == nil { return e.dialCodexWebsocket(ctx, auth, wsURL, headers) } - if staleConn, staleAuthID, staleWSURL := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { + if staleConn, staleCloser, staleAuthID, staleWSURL, staleLifecycle := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { logCodexWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) - if errClose := staleConn.Close(); errClose != nil { - log.Errorf("codex websockets executor: close stale websocket error: %v", errClose) + if staleCloser != nil { + if errClose := staleCloser.Close(); errClose != nil { + log.Errorf("codex websockets executor: close stale websocket error: %v", errClose) + } + } + if staleLifecycle != nil { + staleLifecycle.End("target_changed") } } sess.connMu.Lock() conn := sess.conn + closer := sess.connCloser readerConn := sess.readerConn sess.connMu.Unlock() if conn != nil { @@ -1670,24 +1873,26 @@ sess.configureConn(conn) go e.readUpstreamLoop(sess, conn) } - return conn, nil, nil + return conn, closer, nil, nil } - conn, resp, errDial := e.dialCodexWebsocket(ctx, auth, wsURL, headers) + conn, closer, resp, errDial := e.dialCodexWebsocket(ctx, auth, wsURL, headers) if errDial != nil { - return nil, resp, errDial + return nil, closer, resp, errDial } sess.connMu.Lock() if sess.conn != nil { previous := sess.conn + previousCloser := sess.connCloser sess.connMu.Unlock() - if errClose := conn.Close(); errClose != nil { + if errClose := closer.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } - return previous, nil, nil + return previous, previousCloser, nil, nil } sess.conn = conn + sess.connCloser = closer sess.wsURL = wsURL sess.authID = authID sess.readerConn = conn @@ -1696,7 +1901,7 @@ sess.configureConn(conn) go e.readUpstreamLoop(sess, conn) logCodexWebsocketConnected(sess.sessionID, authID, wsURL) - return conn, resp, nil + return conn, closer, resp, nil } func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, conn *websocket.Conn) { @@ -1771,7 +1976,12 @@ sess.connMu.Unlock() return } + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" sess.conn = nil + sess.connCloser = nil if sess.readerConn == conn { sess.readerConn = nil } @@ -1779,8 +1989,13 @@ logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err) sess.notifyUpstreamDisconnect(err) - if errClose := conn.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) + if closer != nil { + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + } + if lifecycle != nil { + lifecycle.End(reason) } } @@ -1793,9 +2008,7 @@ return } if sessionID == cliproxyauth.CloseAllExecutionSessionsID { - // Executor replacement can happen during hot reload (config/credential changes). - // Do not force-close upstream websocket sessions here, otherwise in-flight - // downstream websocket requests get interrupted. + e.closeAllExecutionSessions("executor_shutdown") return } @@ -1852,19 +2065,28 @@ conn := sess.conn authID := sess.authID wsURL := sess.wsURL + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" sess.conn = nil + sess.connCloser = nil if sess.readerConn == conn { sess.readerConn = nil } sessionID := sess.sessionID sess.connMu.Unlock() - if conn == nil { - return + if conn != nil { + logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) + if closer != nil { + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + } } - logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) - if errClose := conn.Close(); errClose != nil { - log.Errorf("codex websockets executor: close websocket error: %v", errClose) + if lifecycle != nil { + lifecycle.End(reason) } } diff --git a/internal/runtime/executor/codex_websockets_executor_store_test.go b/internal/runtime/executor/codex_websockets_executor_store_test.go --- a/internal/runtime/executor/codex_websockets_executor_store_test.go +++ b/internal/runtime/executor/codex_websockets_executor_store_test.go @@ -6,7 +6,7 @@ cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" ) -func TestCodexWebsocketsExecutor_SessionStoreSurvivesExecutorReplacement(t *testing.T) { +func TestCodexWebsocketsExecutor_CloseAllReleasesSessions(t *testing.T) { sessionID := "test-session-store-survives-replace" globalCodexWebsocketSessionStore.mu.Lock() @@ -33,16 +33,9 @@ globalCodexWebsocketSessionStore.mu.Lock() _, stillPresent := globalCodexWebsocketSessionStore.sessions[sessionID] globalCodexWebsocketSessionStore.mu.Unlock() - if !stillPresent { - t.Fatalf("expected session to remain after executor replacement close marker") + if stillPresent { + t.Fatalf("expected session to be removed after executor shutdown") } exec2.CloseExecutionSession(sessionID) - - globalCodexWebsocketSessionStore.mu.Lock() - _, presentAfterClose := globalCodexWebsocketSessionStore.sessions[sessionID] - globalCodexWebsocketSessionStore.mu.Unlock() - if presentAfterClose { - t.Fatalf("expected session to be removed after explicit close") - } } diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -7,6 +7,7 @@ "net/http" "net/http/httptest" "strings" + "sync/atomic" "testing" "time" @@ -1450,5 +1451,291 @@ if dialer.Proxy != nil { t.Fatal("expected websocket proxy function to be nil for direct mode") + } +} + +func TestCodexWebsocketUpgradeRequiredDoesNotFallbackToHTTPWithLifecycle(t *testing.T) { + var httpFallbackCalls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + httpFallbackCalls.Add(1) + http.Error(w, "unexpected HTTP fallback", http.StatusInternalServerError) + return + } + http.Error(w, "websocket upgrade required", http.StatusUpgradeRequired) + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + ExecutionLifecycle: newTerminalFailureLifecycle(), + } + + if _, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts); errExecute == nil { + t.Fatal("ExecuteStream() error = nil, want failed Home lifecycle attempt") + } + if got := httpFallbackCalls.Load(); got != 0 { + t.Fatalf("HTTP fallback calls = %d, want 0 with an execution lifecycle", got) + } +} + +func TestCodexWebsocketHandshakeFailureReleasesSessionRequestLock(t *testing.T) { + for _, statusCode := range []int{http.StatusUpgradeRequired, http.StatusBadGateway} { + t.Run(http.StatusText(statusCode), func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "upstream rejected websocket", statusCode) + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "failed-handshake", + }, + } + + _, _ = exec.ExecuteStream(context.Background(), auth, req, opts) + sess := exec.getOrCreateSession("failed-handshake") + acquired := make(chan struct{}) + go func() { + sess.reqMu.Lock() + close(acquired) + sess.reqMu.Unlock() + }() + select { + case <-acquired: + case <-time.After(time.Second): + t.Fatal("websocket handshake failure left the session request lock held") + } + }) + } +} + +type terminalFailureLifecycle struct { + active atomic.Bool + ends atomic.Int32 +} + +func newTerminalFailureLifecycle() *terminalFailureLifecycle { + lifecycle := &terminalFailureLifecycle{} + lifecycle.active.Store(true) + return lifecycle +} + +func (*terminalFailureLifecycle) Bind(func() error) error { return nil } +func (l *terminalFailureLifecycle) End(string) { + l.ends.Add(1) + l.active.Store(false) +} +func (*terminalFailureLifecycle) Retain() {} + +func TestCodexWebsocketTerminalFailureInvalidatesRetainedLifecycle(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + var connections atomic.Int32 + firstRelease := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + defer func() { _ = conn.Close() }() + connection := connections.Add(1) + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + terminal := []byte(`{"type":"response.failed","response":{"error":{"type":"authentication_error","code":"invalid_api_key","message":"Invalid token."}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, terminal); errWrite != nil { + t.Errorf("write terminal response: %v", errWrite) + } + if connection == 1 { + <-firstRelease + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + ExecutionLifecycle: newTerminalFailureLifecycle(), + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "terminal-failure", + }, + } + + result, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts) + if errExecute != nil { + t.Fatalf("first ExecuteStream() error = %v", errExecute) + } + for chunk := range result.Chunks { + if chunk.Err == nil { + continue + } + } + lifecycle := opts.ExecutionLifecycle.(*terminalFailureLifecycle) + if lifecycle.active.Load() { + t.Fatal("terminal failure left the retained lifecycle active") + } + if got := lifecycle.ends.Load(); got != 1 { + t.Fatalf("retained lifecycle End calls = %d, want 1", got) + } + sess := exec.getOrCreateSession("terminal-failure") + sess.connMu.Lock() + connected := sess.conn != nil + sess.connMu.Unlock() + if connected { + t.Fatal("terminal failure left the upstream session connection cached") + } + close(firstRelease) + + opts.ExecutionLifecycle = newTerminalFailureLifecycle() + result, errExecute = exec.ExecuteStream(context.Background(), auth, req, opts) + if errExecute != nil { + t.Fatalf("second ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + if got := connections.Load(); got != 2 { + t.Fatalf("websocket connections = %d, want 2 after terminal invalidation", got) + } +} + +type rejectingExecutionLifecycle struct{} + +func (rejectingExecutionLifecycle) Bind(func() error) error { + return errors.New("lifecycle bind rejected") +} +func (rejectingExecutionLifecycle) End(string) {} + +func TestCodexWebsocketNonstreamLifecycleBindFailureDetachesConnection(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + var connections atomic.Int32 + closed := make(chan struct{}, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + connection := connections.Add(1) + defer func() { + _ = conn.Close() + if connection == 1 { + closed <- struct{}{} + } + }() + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + completed := []byte(`{"type":"response.completed","response":{"id":"resp-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Errorf("write completed response: %v", errWrite) + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + ExecutionLifecycle: rejectingExecutionLifecycle{}, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "nonstream-bind-failed", + }, + } + if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute == nil { + t.Fatal("Execute() error = nil, want lifecycle bind failure") + } + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("nonstream lifecycle bind failure did not close the upstream websocket") + } + sess := exec.getOrCreateSession("nonstream-bind-failed") + sess.connMu.Lock() + connected := sess.conn != nil + sess.connMu.Unlock() + if connected { + t.Fatal("nonstream lifecycle bind failure left the closed connection attached to the session") + } + + opts.ExecutionLifecycle = nil + if _, errExecute := exec.Execute(context.Background(), auth, req, opts); errExecute != nil { + t.Fatalf("second Execute() error = %v", errExecute) + } + if got := connections.Load(); got != 2 { + t.Fatalf("websocket connections = %d, want 2 after bind failure", got) + } +} + +func TestCodexWebsocketLifecycleBindFailureReleasesSessionRequestLock(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + closed := make(chan struct{}, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + defer func() { + _ = conn.Close() + closed <- struct{}{} + }() + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + } + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ID: "auth-a", Provider: "codex", Attributes: map[string]string{"api_key": "sk-test", "base_url": server.URL}} + req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + ExecutionLifecycle: rejectingExecutionLifecycle{}, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "bind-failed", + }, + } + if _, errExecute := exec.ExecuteStream(context.Background(), auth, req, opts); errExecute == nil { + t.Fatal("ExecuteStream() error = nil, want lifecycle bind failure") + } + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("lifecycle bind failure did not close the upstream websocket") + } + + sess := exec.getOrCreateSession("bind-failed") + acquired := make(chan struct{}) + go func() { + sess.reqMu.Lock() + close(acquired) + sess.reqMu.Unlock() + }() + select { + case <-acquired: + case <-time.After(time.Second): + t.Fatal("lifecycle bind failure left the session request lock held") } } diff --git a/internal/runtime/executor/home_codex_terminal_test.go b/internal/runtime/executor/home_codex_terminal_test.go new file mode 100644 --- /dev/null +++ b/internal/runtime/executor/home_codex_terminal_test.go @@ -0,0 +1,104 @@ +package executor + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +type terminalCodexHomeDispatcher struct { + auth cliproxyauth.Auth + calls atomic.Int32 +} + +func (*terminalCodexHomeDispatcher) HeartbeatOK() bool { return true } +func (d *terminalCodexHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + d.calls.Add(1) + return json.Marshal(d.auth) +} +func (*terminalCodexHomeDispatcher) AbortAmbiguousDispatch() {} + +func TestHomeCodexTerminalStreamFailureUsesFreshDispatchOnNextRequest(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + var connections atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + defer func() { _ = conn.Close() }() + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + if connections.Add(1) == 1 { + _ = conn.WriteJSON(map[string]any{"type": "response.created", "response": map[string]any{"id": "response-1"}}) + _ = conn.WriteJSON(map[string]any{"type": "error", "status": http.StatusBadGateway, "error": map[string]any{"message": "terminal failure"}}) + } else { + _ = conn.WriteJSON(map[string]any{"type": "response.completed", "response": map[string]any{"id": "response-2", "output": []any{}}}) + } + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + } + })) + defer server.Close() + + dispatcher := &terminalCodexHomeDispatcher{auth: cliproxyauth.Auth{ + ID: "home-codex", + Provider: "codex", + Status: cliproxyauth.StatusActive, + Attributes: map[string]string{ + "api_key": "test-key", + "base_url": server.URL, + }, + }} + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(NewCodexWebsocketsExecutor(&config.Config{})) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "terminal-home-session", + }, + } + request := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[]}`)} + + first, errFirst := manager.ExecuteStream(ctx, []string{"codex"}, request, opts) + if errFirst != nil { + t.Fatalf("first ExecuteStream() error = %v", errFirst) + } + for range first.Chunks { + } + + second, errSecond := manager.ExecuteStream(ctx, []string{"codex"}, request, opts) + if errSecond != nil { + t.Fatalf("second ExecuteStream() error = %v", errSecond) + } + for range second.Chunks { + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 after terminal failure", got) + } + if got := connections.Load(); got != 2 { + t.Fatalf("websocket connections = %d, want 2", got) + } + + manager.CloseExecutionSession("terminal-home-session") +} diff --git a/internal/runtime/executor/websocket_lifecycle_bind_test.go b/internal/runtime/executor/websocket_lifecycle_bind_test.go new file mode 100644 --- /dev/null +++ b/internal/runtime/executor/websocket_lifecycle_bind_test.go @@ -0,0 +1,38 @@ +package executor + +import ( + "sync/atomic" + "testing" + + "github.com/gorilla/websocket" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type countingWebsocketLifecycle struct { + binds atomic.Int32 +} + +func (l *countingWebsocketLifecycle) Bind(func() error) error { + l.binds.Add(1) + return nil +} + +func (*countingWebsocketLifecycle) End(string) {} + +func TestCodexWebsocketSessionBindsSameLifecycleAndConnectionOnce(t *testing.T) { + conn := &websocket.Conn{} + closer := newWebsocketConnectionCloser(conn) + sess := &codexWebsocketSession{conn: conn, connCloser: closer} + lifecycle := &countingWebsocketLifecycle{} + opts := cliproxyexecutor.Options{ExecutionLifecycle: lifecycle} + + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, "gpt-5-codex"); errBind != nil { + t.Fatalf("first bindExecutionLifecycle() error = %v", errBind) + } + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, "gpt-5-codex"); errBind != nil { + t.Fatalf("second bindExecutionLifecycle() error = %v", errBind) + } + if got := lifecycle.binds.Load(); got != 1 { + t.Fatalf("lifecycle Bind calls = %d, want 1 for the same lifecycle and connection", got) + } +} diff --git a/internal/runtime/executor/websocket_session_target_test.go b/internal/runtime/executor/websocket_session_target_test.go --- a/internal/runtime/executor/websocket_session_target_test.go +++ b/internal/runtime/executor/websocket_session_target_test.go @@ -2,15 +2,41 @@ import ( "context" + "encoding/json" + "fmt" + "net" "net/http" "net/http/httptest" + "net/url" + "reflect" + "strconv" "strings" + "sync" + "sync/atomic" "testing" + "time" "github.com/gorilla/websocket" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + internalhome "github.com/router-for-me/CLIProxyAPI/v7/internal/home" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" ) + +type rejectSecondBindLifecycle struct { + binds atomic.Int32 +} + +func (l *rejectSecondBindLifecycle) Bind(func() error) error { + if l.binds.Add(1) > 1 { + return fmt.Errorf("retry lifecycle bind rejected") + } + return nil +} + +func (*rejectSecondBindLifecycle) End(string) {} func TestCodexWebsocketSessionActiveChannelBelongsToConnection(t *testing.T) { sess := &codexWebsocketSession{} @@ -54,6 +80,425 @@ } } +type trackedWebsocketLifecycle struct { + mu sync.Mutex + close func() error + once sync.Once + ends atomic.Int32 +} + +type drainDuringBindWebsocketLifecycle struct{} + +func (drainDuringBindWebsocketLifecycle) Bind(closeFn func() error) error { + if errClose := closeFn(); errClose != nil { + return errClose + } + return fmt.Errorf("execution lifecycle drained during Bind") +} + +func (drainDuringBindWebsocketLifecycle) End(string) {} + +func (l *trackedWebsocketLifecycle) Bind(closeFn func() error) error { + l.mu.Lock() + l.close = closeFn + l.mu.Unlock() + return nil +} + +func (l *trackedWebsocketLifecycle) End(string) { + l.once.Do(func() { + l.ends.Add(1) + l.mu.Lock() + closeFn := l.close + l.mu.Unlock() + if closeFn != nil { + _ = closeFn() + } + }) +} + +func TestClearRetryActiveStateClearsOriginalConnection(t *testing.T) { + sess := &codexWebsocketSession{} + originalConn := &websocket.Conn{} + originalCh := sess.activate(originalConn) + if !clearRetryActiveState(sess, originalConn, originalCh) { + t.Fatal("clearRetryActiveState() = false, want true") + } + if ch, done := sess.activeForConn(originalConn); ch != nil || done != nil { + t.Fatalf("original active state = %v/%v, want nil", ch, done) + } +} + +func TestWebsocketRetryBindFailureClearsActiveSessionState(t *testing.T) { + tests := []struct { + name string + run func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) + }{ + { + name: "Codex nonstream", + run: func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) { + executor := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ID: "retry-bind-codex", Provider: "codex", Attributes: map[string]string{"api_key": "test-key", "base_url": baseURL}} + req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} + primed := false + return func(runOpts cliproxyexecutor.Options) error { + if !primed { + wsURL := "ws" + strings.TrimPrefix(baseURL, "http") + "/responses" + conn, _, _, errEnsure := executor.ensureUpstreamConn(context.Background(), auth, executor.getOrCreateSession("retry-bind"), auth.ID, wsURL, http.Header{}) + if errEnsure != nil { + return errEnsure + } + if errDeadline := conn.SetWriteDeadline(time.Now().Add(-time.Second)); errDeadline != nil { + return errDeadline + } + primed = true + } + _, errExecute := executor.Execute(context.Background(), auth, req, runOpts) + return errExecute + }, executor.getOrCreateSession("retry-bind") + }, + }, + { + name: "Codex stream", + run: func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) { + executor := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ID: "retry-bind-codex", Provider: "codex", Attributes: map[string]string{"api_key": "test-key", "base_url": baseURL}} + req := cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)} + primed := false + return func(runOpts cliproxyexecutor.Options) error { + if !primed { + wsURL := "ws" + strings.TrimPrefix(baseURL, "http") + "/responses" + conn, _, _, errEnsure := executor.ensureUpstreamConn(context.Background(), auth, executor.getOrCreateSession("retry-bind"), auth.ID, wsURL, http.Header{}) + if errEnsure != nil { + return errEnsure + } + if errDeadline := conn.SetWriteDeadline(time.Now().Add(-time.Second)); errDeadline != nil { + return errDeadline + } + primed = true + } + result, errExecute := executor.ExecuteStream(context.Background(), auth, req, runOpts) + if errExecute != nil { + return errExecute + } + for chunk := range result.Chunks { + if chunk.Err != nil { + return chunk.Err + } + } + return nil + }, executor.getOrCreateSession("retry-bind") + }, + }, + { + name: "xAI stream", + run: func(t *testing.T, baseURL string) (func(cliproxyexecutor.Options) error, *codexWebsocketSession) { + executor := NewXAIWebsocketsExecutor(&config.Config{}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + auth := &cliproxyauth.Auth{ID: "retry-bind-xai", Provider: "xai", Attributes: map[string]string{"base_url": baseURL, "websockets": "true"}, Metadata: map[string]any{"access_token": "test-token"}} + req := cliproxyexecutor.Request{Model: "grok-4", Payload: []byte(`{"model":"grok-4","input":[{"type":"message","role":"user","content":"hello"}]}`)} + primed := false + return func(runOpts cliproxyexecutor.Options) error { + if !primed { + wsURL := "ws" + strings.TrimPrefix(baseURL, "http") + "/responses" + conn, _, _, errEnsure := executor.ensureUpstreamConn(context.Background(), auth, executor.getOrCreateSession("retry-bind"), auth.ID, wsURL, http.Header{}) + if errEnsure != nil { + return errEnsure + } + if errDeadline := conn.SetWriteDeadline(time.Now().Add(-time.Second)); errDeadline != nil { + return errDeadline + } + primed = true + } + result, errExecute := executor.ExecuteStream(context.Background(), auth, req, runOpts) + if errExecute != nil { + return errExecute + } + for chunk := range result.Chunks { + if chunk.Err != nil { + return chunk.Err + } + } + return nil + }, executor.getOrCreateSession("retry-bind") + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + var connections atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + connection := connections.Add(1) + defer func() { _ = conn.Close() }() + if connection == 1 { + _, _, _ = conn.ReadMessage() + return + } + if connection == 2 { + return + } + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + completed := []byte(`{"type":"response.completed","response":{"id":"response-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + t.Errorf("write websocket completion: %v", errWrite) + } + })) + defer server.Close() + + lifecycle := &rejectSecondBindLifecycle{} + opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, ExecutionLifecycle: lifecycle, Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "retry-bind"}} + run, sess := test.run(t, server.URL) + if errRun := run(opts); errRun == nil { + t.Fatal("first request error = nil, want retry lifecycle bind rejection") + } + if got := lifecycle.binds.Load(); got != 2 { + t.Fatalf("lifecycle binds = %d, want 2", got) + } + sess.activeMu.Lock() + active := sess.activeConn != nil || sess.activeCh != nil || sess.activeDone != nil || sess.activeCancel != nil + sess.activeMu.Unlock() + if active { + t.Fatal("retry bind failure left the old active websocket state") + } + + opts.ExecutionLifecycle = nil + if errRun := run(opts); errRun != nil { + t.Fatalf("second request error = %v", errRun) + } + if got := connections.Load(); got != 3 { + t.Fatalf("websocket connections = %d, want 3 after retry bind failure", got) + } + }) + } +} + +func TestWebsocketSessionCloseEndsRetainedLifecycleOnce(t *testing.T) { + exec := NewCodexWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + server, closed := newWebsocketTargetServer(t) + defer server.Close() + + sess := exec.getOrCreateSession("retained-lifecycle") + auth := &cliproxyauth.Auth{ID: "auth-a"} + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn := ensureWebsocketTargetConn(t, exec.ensureUpstreamConn, auth, sess, auth.ID, wsURL) + lifecycle := &trackedWebsocketLifecycle{} + if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, sess.connCloser, "model-a"); errBind != nil { + t.Fatalf("bind execution lifecycle: %v", errBind) + } + + exec.CloseExecutionSession("retained-lifecycle") + lifecycle.End("duplicate_close") + if got := lifecycle.ends.Load(); got != 1 { + t.Fatalf("lifecycle End calls = %d, want 1", got) + } + if got := <-closed; got != auth.ID { + t.Fatalf("closed server auth = %q, want %q", got, auth.ID) + } +} + +type closeCountingNetConn struct { + net.Conn + closes atomic.Int32 +} + +func (c *closeCountingNetConn) Close() error { + c.closes.Add(1) + return c.Conn.Close() +} + +func newCloseCountingWebsocketConn(t *testing.T, rawURL string) (*websocket.Conn, *closeCountingNetConn) { + t.Helper() + parsed, errParse := url.Parse(rawURL) + if errParse != nil { + t.Fatalf("parse websocket URL: %v", errParse) + } + conn, errDial := net.Dial("tcp", parsed.Host) + if errDial != nil { + t.Fatalf("dial websocket: %v", errDial) + } + counting := &closeCountingNetConn{Conn: conn} + wsConn, _, errClient := websocket.NewClient(counting, parsed, nil, 1024, 1024) + if errClient != nil { + _ = counting.Close() + t.Fatalf("create websocket client: %v", errClient) + } + return wsConn, counting +} + +func TestSessionlessWebsocketSelectionEndAndDirectCloseRaceClosesOnce(t *testing.T) { + server, _ := newWebsocketTargetServer(t) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, physical := newCloseCountingWebsocketConn(t, wsURL) + closer := newWebsocketConnectionCloser(conn) + lifecycle := &trackedWebsocketLifecycle{} + if errBind := (*codexWebsocketSession)(nil).bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, closer, "model-a"); errBind != nil { + t.Fatalf("bind sessionless lifecycle: %v", errBind) + } + + var wait sync.WaitGroup + wait.Add(2) + go func() { + defer wait.Done() + lifecycle.End("selection_ended") + }() + go func() { + defer wait.Done() + if errClose := closer.Close(); errClose != nil { + t.Errorf("direct close: %v", errClose) + } + }() + wait.Wait() + + if got := physical.closes.Load(); got != 1 { + t.Fatalf("physical websocket closes = %d, want 1", got) + } +} + +func TestWebsocketDrainDuringBindClosesOwnedConnectionOnce(t *testing.T) { + server, _ := newWebsocketTargetServer(t) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, physical := newCloseCountingWebsocketConn(t, wsURL) + closer := newWebsocketConnectionCloser(conn) + sess := &codexWebsocketSession{conn: conn, connCloser: closer, wsURL: wsURL, authID: "auth-a", readerConn: conn} + + errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: drainDuringBindWebsocketLifecycle{}}, conn, closer, "model-a") + if errBind == nil { + t.Fatal("bind execution lifecycle error = nil, want drain error") + } + closeWebsocketAfterBindFailure(sess, conn, closer) + + if got := physical.closes.Load(); got != 1 { + t.Fatalf("physical websocket closes = %d, want 1", got) + } + sess.connMu.Lock() + defer sess.connMu.Unlock() + if sess.conn != nil || sess.connCloser != nil || sess.lifecycle != nil { + t.Fatalf("drained session state = conn:%v closer:%v lifecycle:%v, want detached", sess.conn, sess.connCloser, sess.lifecycle) + } +} + +func TestWebsocketTargetReplacementPhysicallyClosesOwnedConnectionOnce(t *testing.T) { + tests := []struct { + name string + }{ + {name: "Codex"}, + {name: "xAI"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + serverA, _ := newWebsocketTargetServer(t) + defer serverA.Close() + serverB, _ := newWebsocketTargetServer(t) + defer serverB.Close() + + var ensure func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) + var closeSession func(string) + var sess *codexWebsocketSession + switch test.name { + case "Codex": + exec := NewCodexWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + ensure = exec.ensureUpstreamConn + closeSession = exec.CloseExecutionSession + sess = exec.getOrCreateSession("counted-target-change") + case "xAI": + exec := NewXAIWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + ensure = exec.ensureUpstreamConn + closeSession = exec.CloseExecutionSession + sess = exec.getOrCreateSession("counted-target-change") + } + defer closeSession("counted-target-change") + + wsURLA := "ws" + strings.TrimPrefix(serverA.URL, "http") + wsURLB := "ws" + strings.TrimPrefix(serverB.URL, "http") + connA, physical := newCloseCountingWebsocketConn(t, wsURLA) + sess.connMu.Lock() + sess.conn = connA + sess.connCloser = newWebsocketConnectionCloser(connA) + sess.wsURL = wsURLA + sess.authID = "auth-a" + sess.readerConn = connA + sess.connMu.Unlock() + lifecycle := &trackedWebsocketLifecycle{} + if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, connA, sess.connCloser, "model-a"); errBind != nil { + t.Fatalf("bind execution lifecycle: %v", errBind) + } + + if _, _, _, errEnsure := ensure(context.Background(), &cliproxyauth.Auth{ID: "auth-b"}, sess, "auth-b", wsURLB, nil); errEnsure != nil { + t.Fatalf("replace websocket target: %v", errEnsure) + } + if got := physical.closes.Load(); got != 1 { + t.Fatalf("physical websocket closes = %d, want 1", got) + } + }) + } +} + +func TestWebsocketLifecycleEndThenInvalidateAndCloseAllPhysicallyClosesOnce(t *testing.T) { + tests := []struct { + name string + run func(*codexWebsocketSession, *websocket.Conn, *trackedWebsocketLifecycle) + }{ + { + name: "Codex", + run: func(sess *codexWebsocketSession, conn *websocket.Conn, lifecycle *trackedWebsocketLifecycle) { + exec := NewCodexWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: map[string]*codexWebsocketSession{sess.sessionID: sess}} + lifecycle.End("lifecycle_ended") + exec.invalidateUpstreamConn(sess, conn, "invalidated", nil) + exec.CloseExecutionSession(cliproxyauth.CloseAllExecutionSessionsID) + }, + }, + { + name: "xAI", + run: func(sess *codexWebsocketSession, conn *websocket.Conn, lifecycle *trackedWebsocketLifecycle) { + exec := NewXAIWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: map[string]*codexWebsocketSession{sess.sessionID: sess}} + lifecycle.End("lifecycle_ended") + exec.invalidateUpstreamConn(sess, conn, "invalidated", nil) + exec.CloseExecutionSession(cliproxyauth.CloseAllExecutionSessionsID) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server, _ := newWebsocketTargetServer(t) + defer server.Close() + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, physical := newCloseCountingWebsocketConn(t, wsURL) + sess := &codexWebsocketSession{sessionID: "counted-lifecycle", conn: conn, connCloser: newWebsocketConnectionCloser(conn), wsURL: wsURL, authID: "auth-a", readerConn: conn} + lifecycle := &trackedWebsocketLifecycle{} + if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, sess.connCloser, "model-a"); errBind != nil { + t.Fatalf("bind execution lifecycle: %v", errBind) + } + + test.run(sess, conn, lifecycle) + if got := physical.closes.Load(); got != 1 { + t.Fatalf("physical websocket closes = %d, want 1", got) + } + }) + } +} + func TestWebsocketExecutorsReconnectWhenSessionTargetChanges(t *testing.T) { t.Run("Codex", func(t *testing.T) { exec := NewCodexWebsocketsExecutor(&config.Config{}) @@ -84,7 +529,7 @@ t *testing.T, disconnectChan func(string) <-chan error, getSession func(string) *codexWebsocketSession, - ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *http.Response, error), + ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error), closeSession func(string), ) { t.Helper() @@ -146,7 +591,7 @@ func ensureWebsocketTargetConn( t *testing.T, - ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *http.Response, error), + ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error), auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, @@ -154,7 +599,7 @@ ) *websocket.Conn { t.Helper() headers := http.Header{"X-Test-Auth": []string{authID}} - conn, resp, errEnsure := ensureConn(context.Background(), auth, sess, authID, wsURL, headers) + conn, _, resp, errEnsure := ensureConn(context.Background(), auth, sess, authID, wsURL, headers) if resp != nil && resp.Body != nil { defer func() { if errClose := resp.Body.Close(); errClose != nil { @@ -195,4 +640,477 @@ } })) return server, closed +} + +type registryDrainWebsocketLifecycle struct { + scope *executionregistry.Scope + ends atomic.Int32 +} + +func (l *registryDrainWebsocketLifecycle) Bind(closeFn func() error) error { + return l.scope.Bind(closeFn) +} + +func (l *registryDrainWebsocketLifecycle) End(string) { + l.ends.Add(1) + l.scope.End("websocket_closed") +} + +func (l *registryDrainWebsocketLifecycle) Retain() {} + +type websocketHomeDispatcher struct { + provider string +} + +func (d websocketHomeDispatcher) HeartbeatOK() bool { return true } + +func (d websocketHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + return json.Marshal(map[string]any{"auth": map[string]any{ + "id": "home-websocket-auth", + "provider": d.provider, + "status": "active", + "attributes": map[string]string{ + "api_key": "home-key", + }, + }}) +} + +func (websocketHomeDispatcher) AbortAmbiguousDispatch() {} + +type accountedWebsocketHomeDispatcher struct { + provider string + baseURL string + calls atomic.Int32 + releases atomic.Int32 + before atomic.Bool +} + +func (*accountedWebsocketHomeDispatcher) HeartbeatOK() bool { return true } + +func (d *accountedWebsocketHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + call := d.calls.Add(1) + if call > 1 && d.releases.Load() != call-1 { + d.before.Store(false) + } else if call > 1 { + d.before.Store(true) + } + upstreamModel := "model-a" + if strings.Contains(strings.ToLower(model), "(custom)") { + upstreamModel = "model-a(custom)" + } + return json.Marshal(map[string]any{ + "model": upstreamModel, + "auth_index": "accounted-websocket-auth", + "auth": map[string]any{ + "id": "accounted-websocket-auth", + "provider": d.provider, + "status": "active", + "attributes": map[string]string{ + "api_key": "test-key", + "base_url": d.baseURL, + "websockets": "true", + }, + }, + "concurrency": map[string]any{ + "accounted": true, + "credential_id": "accounted-websocket-auth", + "model": upstreamModel, + }, + }) +} + +func (*accountedWebsocketHomeDispatcher) AbortAmbiguousDispatch() {} + +func TestAuditAccountedCodexXAIReconnectReuseAndTargetChange(t *testing.T) { + tests := []struct { + name string + provider string + newExecutor func() cliproxyauth.ProviderExecutor + }{ + { + name: "Codex", + provider: "codex", + newExecutor: func() cliproxyauth.ProviderExecutor { + executor := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + return executor + }, + }, + { + name: "xAI", + provider: "xai", + newExecutor: func() cliproxyauth.ProviderExecutor { + executor := NewXAIWebsocketsExecutor(&config.Config{}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + executor.idStore = &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + return executor + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + var connections atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, errUpgrade := upgrader.Upgrade(w, r, nil) + if errUpgrade != nil { + t.Errorf("upgrade websocket: %v", errUpgrade) + return + } + connections.Add(1) + defer func() { _ = conn.Close() }() + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + completed := []byte(`{"type":"response.completed","response":{"id":"response-1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, completed); errWrite != nil { + return + } + } + })) + defer server.Close() + + registry := executionregistry.New() + dispatcher := &accountedWebsocketHomeDispatcher{provider: test.provider, baseURL: server.URL} + var releaseGroups []executionregistry.ReleaseGroup + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { + dispatcher.releases.Add(1) + releaseGroups = append(releaseGroups, group) + }) + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(test.newExecutor()) + t.Cleanup(func() { manager.CloseExecutionSession("accounted-websocket-session") }) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{ + Stream: true, + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "accounted-websocket-session", + cliproxyexecutor.PinnedAuthMetadataKey: "accounted-websocket-auth", + }, + } + execute := func(model string) { + t.Helper() + result, errExecute := manager.ExecuteStream(ctx, []string{test.provider}, cliproxyexecutor.Request{Model: model, Payload: []byte(`{"model":"model-a","input":[]}`)}, opts) + if errExecute != nil { + t.Fatalf("ExecuteStream(%q) error = %v", model, errExecute) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("ExecuteStream(%q) chunk error = %v", model, chunk.Err) + } + } + } + + execute(" MODEL-A(HIGH) ") + execute("model-a") + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1 for canonical retained reuse", got) + } + manager.CloseExecutionSession("accounted-websocket-session") + execute("model-a") + execute("model-a(custom)") + if got := dispatcher.calls.Load(); got != 3 { + t.Fatalf("Home RPOP calls = %d, want 3 after reconnect and target change", got) + } + if !dispatcher.before.Load() { + t.Fatal("previous accounted selection was not released before redispatch") + } + manager.CloseExecutionSession("accounted-websocket-session") + wantGroups := []executionregistry.ReleaseGroup{ + {CredentialID: "accounted-websocket-auth", Model: "model-a"}, + {CredentialID: "accounted-websocket-auth", Model: "model-a"}, + {CredentialID: "accounted-websocket-auth", Model: "model-a(custom)"}, + } + if !reflect.DeepEqual(releaseGroups, wantGroups) { + t.Fatalf("release groups = %#v, want %#v", releaseGroups, wantGroups) + } + if got := connections.Load(); got != 3 { + t.Fatalf("upstream websocket connections = %d, want 3", got) + } + }) + } +} + +func TestHomeSelectionRegistryDrainClosesRealWebsocketSessions(t *testing.T) { + tests := []struct { + name string + provider string + newExecutor func() (cliproxyauth.ProviderExecutor, func(string) *codexWebsocketSession, func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error)) + }{ + { + name: "Codex", + provider: "codex", + newExecutor: func() (cliproxyauth.ProviderExecutor, func(string) *codexWebsocketSession, func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error)) { + executor := NewCodexWebsocketsExecutor(&config.Config{}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + return executor, executor.getOrCreateSession, executor.ensureUpstreamConn + }, + }, + { + name: "xAI", + provider: "xai", + newExecutor: func() (cliproxyauth.ProviderExecutor, func(string) *codexWebsocketSession, func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error)) { + executor := NewXAIWebsocketsExecutor(&config.Config{}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + return executor, executor.getOrCreateSession, executor.ensureUpstreamConn + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server, closed := newWebsocketTargetServer(t) + defer server.Close() + + executor, getSession, ensureConn := test.newExecutor() + registry := executionregistry.New() + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(websocketHomeDispatcher{provider: test.provider}, registry, 1) + manager.RegisterExecutor(executor) + selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), test.provider, "model-a", cliproxyauth.AuthKindAPIKey, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect) + } + auth := selection.CloneAuth() + sess := getSession("real-home-drain") + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn := ensureWebsocketTargetConn(t, ensureConn, auth, sess, auth.ID, wsURL) + if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: selection}, conn, sess.connCloser, "model-a"); errBind != nil { + t.Fatalf("bind execution lifecycle: %v", errBind) + } + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := registry.Drain(drainCtx); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } + if selection.Active() { + t.Fatal("registry drain did not end the Home dispatch selection") + } + if got := <-closed; got != auth.ID { + t.Fatalf("closed server auth = %q, want %q", got, auth.ID) + } + }) + } +} + +type codex426RetryDispatcher struct { + calls atomic.Int32 + baseURLs []string + websockets []bool + releases atomic.Int32 + releasedBeforeSecondRPop atomic.Bool +} + +func (d *codex426RetryDispatcher) HeartbeatOK() bool { return true } + +func (d *codex426RetryDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + call := int(d.calls.Add(1)) + if call > len(d.baseURLs) { + return nil, fmt.Errorf("unexpected Home dispatch %d", call) + } + if call == 2 { + d.releasedBeforeSecondRPop.Store(d.releases.Load() == 1) + } + credentialID := "codex-home-" + strconv.Itoa(call) + attributes := map[string]string{ + "api_key": "home-key", + "base_url": d.baseURLs[call-1], + } + if call <= len(d.websockets) && d.websockets[call-1] { + attributes["websockets"] = "true" + } + return json.Marshal(map[string]any{ + "model": model, + "auth_index": credentialID, + "auth": map[string]any{ + "id": credentialID, + "provider": "codex", + "status": "active", + "attributes": attributes, + }, + "concurrency": map[string]any{ + "accounted": true, + "credential_id": credentialID, + "model": model, + }, + }) +} + +func (*codex426RetryDispatcher) AbortAmbiguousDispatch() {} + +func TestAuditHomeCodex426WebsocketToHTTPFreshSelection(t *testing.T) { + upgradeRequired := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "websocket upgrade required", http.StatusUpgradeRequired) + })) + defer upgradeRequired.Close() + + var httpFallbackCalls atomic.Int32 + httpFallback := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost || r.URL.Path != "/responses" { + http.Error(w, "unexpected fallback request", http.StatusBadRequest) + return + } + httpFallbackCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"type\":\"response.completed\",\"response\":{\"id\":\"response-1\",\"output\":[],\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n")) + })) + defer httpFallback.Close() + + executor := NewCodexAutoExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + executor.wsExec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + dispatcher := &codex426RetryDispatcher{ + baseURLs: []string{upgradeRequired.URL, httpFallback.URL}, + websockets: []bool{true, false}, + } + registry := executionregistry.New() + var releaseGroups []executionregistry.ReleaseGroup + var releaseGroupsMu sync.Mutex + releaseFlusher := internalhome.NewReleaseFlusher(func() config.CredentialConcurrencyConfig { + return config.CredentialConcurrencyConfig{ + ReleaseFlushInterval: time.Millisecond, + ReleaseMaxBackoff: 10 * time.Millisecond, + } + }, func(_ context.Context, frame internalhome.ConcurrencyReleaseFrame) error { + dispatcher.releases.Add(1) + releaseGroupsMu.Lock() + releaseGroups = append(releaseGroups, executionregistry.ReleaseGroup{CredentialID: frame.CredentialID, Model: frame.Model}) + releaseGroupsMu.Unlock() + return nil + }) + registry.SetReleaseSink(releaseFlusher.MarkDirty) + releaseCtx, cancelRelease := context.WithCancel(context.Background()) + releaseDone := make(chan struct{}) + go func() { + defer close(releaseDone) + releaseFlusher.Run(releaseCtx) + }() + defer func() { + cancelRelease() + <-releaseDone + }() + manager := cliproxyauth.NewManager(nil, nil, nil) + manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(executor) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + result, errExecute := manager.ExecuteStream(ctx, []string{"codex"}, cliproxyexecutor.Request{Model: "gpt-5-codex", Payload: []byte(`{"model":"gpt-5-codex","input":[{"type":"message","role":"user","content":"hello"}]}`)}, cliproxyexecutor.Options{Stream: true, SourceFormat: sdktranslator.FormatOpenAIResponse, ResponseFormat: sdktranslator.FormatOpenAIResponse, Metadata: map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: "home-426"}}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + if !dispatcher.releasedBeforeSecondRPop.Load() { + t.Fatal("first accounted selection was not released before the 426 retry RPOP") + } + if got := dispatcher.releases.Load(); got != 1 { + t.Fatalf("accounted releases before response completion = %d, want 1", got) + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk error = %v", chunk.Err) + } + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 after 426", got) + } + if got := httpFallbackCalls.Load(); got != 1 { + t.Fatalf("HTTP fallback calls = %d, want 1 on the fresh Home selection", got) + } + deadline := time.NewTimer(time.Second) + defer deadline.Stop() + for dispatcher.releases.Load() != 2 { + select { + case <-deadline.C: + t.Fatalf("accounted releases after response completion = %d, want 2", dispatcher.releases.Load()) + case <-time.After(time.Millisecond): + } + } + releaseGroupsMu.Lock() + gotReleaseGroups := append([]executionregistry.ReleaseGroup(nil), releaseGroups...) + releaseGroupsMu.Unlock() + wantReleaseGroups := []executionregistry.ReleaseGroup{ + {CredentialID: "codex-home-1", Model: "gpt-5-codex"}, + {CredentialID: "codex-home-2", Model: "gpt-5-codex"}, + } + if !reflect.DeepEqual(gotReleaseGroups, wantReleaseGroups) { + t.Fatalf("accounted release groups = %#v, want %#v", gotReleaseGroups, wantReleaseGroups) + } +} + +func TestWebsocketRegistryDrainClosesAndEndsRetainedSession(t *testing.T) { + tests := []struct { + name string + getSession func(string) *codexWebsocketSession + ensureConn func(context.Context, *cliproxyauth.Auth, *codexWebsocketSession, string, string, http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) + }{ + { + name: "Codex", + getSession: func(sessionID string) *codexWebsocketSession { + executor := NewCodexWebsocketsExecutor(&config.Config{}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + return executor.getOrCreateSession(sessionID) + }, + ensureConn: func(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { + executor := NewCodexWebsocketsExecutor(&config.Config{}) + return executor.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, headers) + }, + }, + { + name: "xAI shared session", + getSession: func(sessionID string) *codexWebsocketSession { + executor := NewXAIWebsocketsExecutor(&config.Config{}) + executor.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + return executor.getOrCreateSession(sessionID) + }, + ensureConn: func(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { + executor := NewXAIWebsocketsExecutor(&config.Config{}) + return executor.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, headers) + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server, closed := newWebsocketTargetServer(t) + defer server.Close() + + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatalf("BeginDispatch() error = %v", errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{Kind: "websocket"}) + if errInstall != nil { + t.Fatalf("Install() error = %v", errInstall) + } + lifecycle := ®istryDrainWebsocketLifecycle{scope: scope} + auth := &cliproxyauth.Auth{ID: "auth-a"} + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + sess := test.getSession("drain-retained-session") + conn := ensureWebsocketTargetConn(t, test.ensureConn, auth, sess, auth.ID, wsURL) + if errBind := sess.bindExecutionLifecycle(cliproxyexecutor.Options{ExecutionLifecycle: lifecycle}, conn, sess.connCloser, "model-a"); errBind != nil { + t.Fatalf("bind execution lifecycle: %v", errBind) + } + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := registry.Drain(drainCtx); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } + if got := lifecycle.ends.Load(); got != 1 { + t.Fatalf("lifecycle End calls = %d, want 1", got) + } + if got := <-closed; got != auth.ID { + t.Fatalf("closed server auth = %q, want %q", got, auth.ID) + } + }) + } } diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -464,7 +464,7 @@ helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBody) - conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + conn, closer, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) var upstreamHeaders http.Header if respHS != nil { upstreamHeaders = respHS.Header.Clone() @@ -485,6 +485,13 @@ sess.reqMu.Unlock() } return nil, errDial + } + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + if sess != nil { + sess.reqMu.Unlock() + } + closeWebsocketAfterBindFailure(sess, conn, closer) + return nil, errBind } recordAPIWebsocketHandshake(ctx, e.cfg, respHS) reporter.StartResponseTTFT() @@ -508,7 +515,7 @@ sess.reqMu.Unlock() return nil, errSend } - connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + connRetry, closerRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry != nil || connRetry == nil { bodyErrRetry := websocketHandshakeBody(respHSRetry) closeHTTPResponseBody(respHSRetry, "xai websockets executor: close handshake response body error") @@ -520,7 +527,15 @@ } return nil, errDialRetry } + previousConn, previousReadCh := conn, readCh conn = connRetry + closer = closerRetry + if errBind := sess.bindExecutionLifecycle(opts, conn, closer, req.Model); errBind != nil { + clearRetryActiveState(sess, previousConn, previousReadCh) + sess.reqMu.Unlock() + closeWebsocketAfterBindFailure(sess, conn, closer) + return nil, errBind + } readCh = sess.activate(conn) wsReqBodyRetry := buildXAIWebsocketRequestBody(prepared.body) helps.RecordAPIWebsocketRequest(ctx, e.cfg, helps.UpstreamRequestLog{ @@ -548,7 +563,7 @@ wsReqBody = wsReqBodyRetry } else { logXAIWebsocketDisconnected(executionSessionID, authID, wsURL, "send_error", errSend) - if errClose := conn.Close(); errClose != nil { + if errClose := closer.Close(); errClose != nil { log.Errorf("xai websockets executor: close websocket error: %v", errClose) } return nil, errSend @@ -568,7 +583,7 @@ return } logXAIWebsocketDisconnected(executionSessionID, authID, wsURL, terminateReason, terminateErr) - if errClose := conn.Close(); errClose != nil { + if errClose := closer.Close(); errClose != nil { log.Errorf("xai websockets executor: close websocket error: %v", errClose) } }() @@ -897,7 +912,7 @@ return prepared, nil } -func (e *XAIWebsocketsExecutor) dialXAIWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *http.Response, error) { +func (e *XAIWebsocketsExecutor) dialXAIWebsocket(ctx context.Context, auth *cliproxyauth.Auth, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { dialer := newProxyAwareWebsocketDialer(e.cfg, auth) dialer.HandshakeTimeout = codexResponsesWebsocketHandshakeTO dialer.EnableCompression = true @@ -905,11 +920,12 @@ ctx = context.Background() } conn, resp, err := dialer.DialContext(ctx, wsURL, headers) + closer := newWebsocketConnectionCloser(conn) if conn != nil { // Avoid gorilla/websocket flate tail validation issues on some upstreams/Go versions. conn.EnableWriteCompression(false) } - return conn, resp, err + return conn, closer, resp, err } func (e *XAIWebsocketsExecutor) getOrCreateSession(sessionID string) *codexWebsocketSession { @@ -945,20 +961,26 @@ return sess.upstreamDisconnectCh } -func (e *XAIWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *http.Response, error) { +func (e *XAIWebsocketsExecutor) ensureUpstreamConn(ctx context.Context, auth *cliproxyauth.Auth, sess *codexWebsocketSession, authID string, wsURL string, headers http.Header) (*websocket.Conn, *websocketConnectionCloser, *http.Response, error) { if sess == nil { return e.dialXAIWebsocket(ctx, auth, wsURL, headers) } - if staleConn, staleAuthID, staleWSURL := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { + if staleConn, staleCloser, staleAuthID, staleWSURL, staleLifecycle := detachMismatchedWebsocketSessionConn(sess, authID, wsURL); staleConn != nil { logXAIWebsocketDisconnected(sess.sessionID, staleAuthID, staleWSURL, "target_changed", nil) - if errClose := staleConn.Close(); errClose != nil { - log.Errorf("xai websockets executor: close stale websocket error: %v", errClose) + if staleCloser != nil { + if errClose := staleCloser.Close(); errClose != nil { + log.Errorf("xai websockets executor: close stale websocket error: %v", errClose) + } + } + if staleLifecycle != nil { + staleLifecycle.End("target_changed") } } sess.connMu.Lock() conn := sess.conn + closer := sess.connCloser readerConn := sess.readerConn sess.connMu.Unlock() if conn != nil { @@ -969,24 +991,26 @@ configureXAIWebsocketConn(sess, conn) go e.readUpstreamLoop(sess, conn) } - return conn, nil, nil + return conn, closer, nil, nil } - conn, resp, errDial := e.dialXAIWebsocket(ctx, auth, wsURL, headers) + conn, closer, resp, errDial := e.dialXAIWebsocket(ctx, auth, wsURL, headers) if errDial != nil { - return nil, resp, errDial + return nil, closer, resp, errDial } sess.connMu.Lock() if sess.conn != nil { previous := sess.conn + previousCloser := sess.connCloser sess.connMu.Unlock() - if errClose := conn.Close(); errClose != nil { + if errClose := closer.Close(); errClose != nil { log.Errorf("xai websockets executor: close websocket error: %v", errClose) } - return previous, nil, nil + return previous, previousCloser, nil, nil } sess.conn = conn + sess.connCloser = closer sess.wsURL = wsURL sess.authID = authID sess.readerConn = conn @@ -995,7 +1019,7 @@ configureXAIWebsocketConn(sess, conn) go e.readUpstreamLoop(sess, conn) logXAIWebsocketConnected(sess.sessionID, authID, wsURL) - return conn, resp, nil + return conn, closer, resp, nil } func configureXAIWebsocketConn(sess *codexWebsocketSession, conn *websocket.Conn) { @@ -1142,7 +1166,12 @@ sess.connMu.Unlock() return } + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" sess.conn = nil + sess.connCloser = nil if sess.readerConn == conn { sess.readerConn = nil } @@ -1152,8 +1181,13 @@ if notify { sess.notifyUpstreamDisconnect(err) } - if errClose := conn.Close(); errClose != nil { - log.Errorf("xai websockets executor: close websocket error: %v", errClose) + if closer != nil { + if errClose := closer.Close(); errClose != nil { + log.Errorf("xai websockets executor: close websocket error: %v", errClose) + } + } + if lifecycle != nil { + lifecycle.End(reason) } } @@ -1163,6 +1197,7 @@ return } if sessionID == cliproxyauth.CloseAllExecutionSessionsID { + e.closeAllExecutionSessions("executor_shutdown") return } @@ -1183,6 +1218,28 @@ closeXAIWebsocketSession(sess, reason) } +func (e *XAIWebsocketsExecutor) closeAllExecutionSessions(reason string) { + if e == nil { + return + } + store := e.store + if store == nil { + store = globalXAIWebsocketSessionStore + } + store.mu.Lock() + sessions := make([]*codexWebsocketSession, 0, len(store.sessions)) + for sessionID, sess := range store.sessions { + delete(store.sessions, sessionID) + if sess != nil { + sessions = append(sessions, sess) + } + } + store.mu.Unlock() + for _, sess := range sessions { + closeXAIWebsocketSession(sess, reason) + } +} + func closeXAIWebsocketSession(sess *codexWebsocketSession, reason string) { if sess == nil { return @@ -1196,19 +1253,28 @@ conn := sess.conn authID := sess.authID wsURL := sess.wsURL + lifecycle := sess.lifecycle + closer := sess.connCloser + sess.lifecycle = nil + sess.lifecycleModel = "" sess.conn = nil + sess.connCloser = nil if sess.readerConn == conn { sess.readerConn = nil } sessionID := sess.sessionID sess.connMu.Unlock() - if conn == nil { - return + if conn != nil { + logXAIWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) + if closer != nil { + if errClose := closer.Close(); errClose != nil { + log.Errorf("xai websockets executor: close websocket error: %v", errClose) + } + } } - logXAIWebsocketDisconnected(sessionID, authID, wsURL, reason, nil) - if errClose := conn.Close(); errClose != nil { - log.Errorf("xai websockets executor: close websocket error: %v", errClose) + if lifecycle != nil { + lifecycle.End(reason) } } diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -874,6 +874,9 @@ } func (h *BaseAPIHandler) executeWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { + if h.AuthManager != nil && h.AuthManager.HomeEnabled() { + return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} + } host := h.pluginExecutorHost() if host == nil { return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} @@ -892,6 +895,9 @@ } func (h *BaseAPIHandler) countWithPluginExecutor(ctx context.Context, handlerType, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) ([]byte, http.Header, *interfaces.ErrorMessage) { + if h.AuthManager != nil && h.AuthManager.HomeEnabled() { + return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} + } host := h.pluginExecutorHost() if host == nil { return nil, nil, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("plugin executor host is unavailable")} @@ -992,6 +998,12 @@ } func (h *BaseAPIHandler) streamWithPluginExecutor(ctx context.Context, entryProtocol, responseProtocol, modelName, originalRequestedModel string, rawJSON []byte, alt, executorPluginID string, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { + if h.AuthManager != nil && h.AuthManager.HomeEnabled() { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable, Error: fmt.Errorf("plugin executor routing is unavailable while Home is enabled")} + close(errChan) + return nil, nil, errChan + } host := h.pluginExecutorHost() if host == nil { errChan := make(chan *interfaces.ErrorMessage, 1) @@ -1021,18 +1033,8 @@ streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost) rawStreamHeaders := cloneHeader(streamResult.Headers) baseStreamHeaders := cloneHeader(streamResult.Headers) - upstreamHeaders := downstreamHeadersFromExecutor(rawStreamHeaders, passthroughHeadersEnabled) - if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { - upstreamHeaders = make(http.Header) - } - streamHeadersCommitted := false applyStreamHeaders := func(headers http.Header) { rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers) - if streamHeadersCommitted || upstreamHeaders == nil { - return - } - nextHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) - replaceHeader(upstreamHeaders, nextHeaders) } if streamInterceptorsActive { intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ @@ -1047,6 +1049,10 @@ Metadata: opts.Metadata, }, execOptions.SkipInterceptorPluginID) applyStreamHeaders(intercepted.Headers) + } + upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) + if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { + upstreamHeaders = make(http.Header) } dataChan := make(chan []byte) @@ -1119,7 +1125,6 @@ return } } - streamHeadersCommitted = true select { case dataChan <- payload: if streamInterceptorsActive { @@ -1206,33 +1211,34 @@ close(errChan) return nil, nil, errChan } + if streamResult == nil { + errChan := make(chan *interfaces.ErrorMessage, 1) + errChan <- &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: fmt.Errorf("auth manager returned nil stream")} + close(errChan) + return nil, nil, errChan + } executedRequest := func() (coreexecutor.Request, coreexecutor.Options) { return afterAuthCapture.apply(req, opts) } passthroughHeadersEnabled := PassthroughHeadersEnabled(h.Cfg) interceptorHost := h.interceptorHost() streamInterceptorsActive := streamInterceptorsEnabled(interceptorHost) - // Capture upstream headers from the initial connection synchronously before the goroutine starts. - // Keep a mutable map so bootstrap retries can replace it before first payload is sent. + // Resolve bootstrap retries and header initialization before returning so the + // returned header snapshot is never modified by the stream goroutine. rawStreamHeaders := cloneHeader(streamResult.Headers) baseStreamHeaders := cloneHeader(streamResult.Headers) - upstreamHeaders := downstreamHeadersFromExecutor(rawStreamHeaders, passthroughHeadersEnabled) - if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { - upstreamHeaders = make(http.Header) - } chunks := streamResult.Chunks - dataChan := make(chan []byte) - errChan := make(chan *interfaces.ErrorMessage, 1) + if chunks == nil { + closed := make(chan coreexecutor.StreamChunk) + close(closed) + chunks = closed + } + streamClosedBeforeRead := false + streamCanceledBeforeRead := false streamHeaderInitialized := false - streamHeadersCommitted := false applyStreamHeaders := func(headers http.Header) { rawStreamHeaders = finalInterceptorHeaders(rawStreamHeaders, headers) - if streamHeadersCommitted { - return - } - nextHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) - replaceHeader(upstreamHeaders, nextHeaders) } applyStreamHeaderInit := func() { @@ -1255,9 +1261,48 @@ streamHeaderInitialized = true } - pendingChunks := make([]coreexecutor.StreamChunk, 0, 1) - streamClosedBeforeRead := false - streamCanceledBeforeRead := false + transformStreamPayload := func(payload []byte, chunkIndex *int, historyChunks [][]byte) ([]byte, bool, *interfaces.ErrorMessage) { + applyStreamHeaderInit() + payload = cloneBytes(payload) + if streamInterceptorsActive { + executedReq, executedOpts := executedRequest() + intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ + SourceFormat: responseProtocol, + Model: normalizedModel, + RequestedModel: originalRequestedModel, + RequestHeaders: cloneHeader(executedOpts.Headers), + ResponseHeaders: cloneHeader(rawStreamHeaders), + OriginalRequest: cloneBytes(executedOpts.OriginalRequest), + RequestBody: cloneBytes(executedReq.Payload), + Body: payload, + HistoryChunks: cloneByteSlices(historyChunks), + ChunkIndex: *chunkIndex, + Metadata: executedOpts.Metadata, + }, execOptions.SkipInterceptorPluginID) + applyStreamHeaders(intercepted.Headers) + if len(intercepted.Body) > 0 { + payload = cloneBytes(intercepted.Body) + } + (*chunkIndex)++ + if intercepted.DropChunk { + return nil, false, nil + } + } else { + (*chunkIndex)++ + } + if responseProtocol == "openai-response" { + if errValidate := validateSSEDataJSON(payload); errValidate != nil { + return nil, false, &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate} + } + } + return payload, true, nil + } + + var bootstrapPayload []byte + bootstrapChunkIndex := 0 + var bootstrapHistoryChunks [][]byte + var bootstrapStreamErr error + var bootstrapErr *interfaces.ErrorMessage readInitialStreamChunks := func() { for { var chunk coreexecutor.StreamChunk @@ -1277,17 +1322,85 @@ applyStreamHeaderInit() return } - pendingChunks = append(pendingChunks, chunk) if chunk.Err != nil { + bootstrapStreamErr = chunk.Err return } - if len(chunk.Payload) > 0 { - applyStreamHeaderInit() + if len(chunk.Payload) == 0 { + continue + } + payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &bootstrapChunkIndex, bootstrapHistoryChunks) + if errMsg != nil { + bootstrapErr = errMsg return } + if !deliverable { + continue + } + bootstrapPayload = payload + return } } - readInitialStreamChunks() + + bootstrapEligible := func(err error) bool { + status := statusFromError(err) + if status == 0 { + return true + } + switch status { + case http.StatusUnauthorized, http.StatusForbidden, http.StatusPaymentRequired, + http.StatusRequestTimeout, http.StatusTooManyRequests: + return true + default: + return status >= http.StatusInternalServerError + } + } + + maxBootstrapRetries := StreamingBootstrapRetries(h.Cfg) + if h.AuthManager.HomeEnabled() { + maxBootstrapRetries = 0 + } + for bootstrapRetries := 0; !streamCanceledBeforeRead; { + readInitialStreamChunks() + if streamCanceledBeforeRead || bootstrapErr != nil || bootstrapStreamErr == nil { + break + } + if bootstrapRetries >= maxBootstrapRetries || !bootstrapEligible(bootstrapStreamErr) { + bootstrapErr = executionErrorMessage(bootstrapStreamErr) + break + } + bootstrapRetries++ + retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) + if retryErr != nil { + bootstrapErr = executionErrorMessage(enrichAuthSelectionError(retryErr, providers, normalizedModel)) + break + } + if retryResult == nil { + bootstrapErr = executionErrorMessage(fmt.Errorf("auth manager returned nil stream")) + break + } + rawStreamHeaders = cloneHeader(retryResult.Headers) + baseStreamHeaders = cloneHeader(retryResult.Headers) + streamHeaderInitialized = false + streamClosedBeforeRead = false + bootstrapStreamErr = nil + bootstrapPayload = nil + bootstrapChunkIndex = 0 + bootstrapHistoryChunks = nil + chunks = retryResult.Chunks + if chunks == nil { + closed := make(chan coreexecutor.StreamChunk) + close(closed) + chunks = closed + } + } + + upstreamHeaders := downstreamHeadersAfterInterceptors(baseStreamHeaders, rawStreamHeaders, passthroughHeadersEnabled) + if upstreamHeaders == nil && (passthroughHeadersEnabled || streamInterceptorsActive) { + upstreamHeaders = make(http.Header) + } + dataChan := make(chan []byte) + errChan := make(chan *interfaces.ErrorMessage, 1) go func() { defer close(dataChan) @@ -1295,11 +1408,6 @@ if streamCanceledBeforeRead { return } - sentPayload := false - bootstrapRetries := 0 - chunkIndex := 0 - var historyChunks [][]byte - maxBootstrapRetries := StreamingBootstrapRetries(h.Cfg) sendErr := func(msg *interfaces.ErrorMessage) bool { if ctx == nil { @@ -1327,116 +1435,47 @@ } } - bootstrapEligible := func(err error) bool { - status := statusFromError(err) - if status == 0 { - return true - } - switch status { - case http.StatusUnauthorized, http.StatusForbidden, http.StatusPaymentRequired, - http.StatusRequestTimeout, http.StatusTooManyRequests: - return true - default: - return status >= http.StatusInternalServerError - } + if bootstrapErr != nil { + _ = sendErr(bootstrapErr) + return } - outer: - for { - for { - chunk, ok, canceled := nextStreamChunk(ctx, &pendingChunks, &streamClosedBeforeRead, chunks) - if canceled { - return - } - if !ok { - applyStreamHeaderInit() - return - } - if chunk.Err != nil { - streamErr := chunk.Err - // Safe bootstrap recovery: if the upstream fails before any payload bytes are sent, - // retry a few times (to allow auth rotation / transient recovery) and then attempt model fallback. - if !sentPayload { - if bootstrapRetries < maxBootstrapRetries && bootstrapEligible(streamErr) { - bootstrapRetries++ - retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) - if retryErr == nil { - rawStreamHeaders = cloneHeader(retryResult.Headers) - baseStreamHeaders = cloneHeader(retryResult.Headers) - replaceHeader(upstreamHeaders, downstreamHeadersFromExecutor(rawStreamHeaders, passthroughHeadersEnabled)) - streamHeaderInitialized = false - streamHeadersCommitted = false - pendingChunks = nil - streamClosedBeforeRead = false - chunks = retryResult.Chunks - continue outer - } - streamErr = enrichAuthSelectionError(retryErr, providers, normalizedModel) - } - } - - status := http.StatusInternalServerError - if se, ok := streamErr.(interface{ StatusCode() int }); ok && se != nil { - if code := se.StatusCode(); code > 0 { - status = code - } - } - var addon http.Header - if he, ok := streamErr.(interface{ Headers() http.Header }); ok && he != nil { - if hdr := he.Headers(); hdr != nil { - addon = hdr.Clone() - } - } - _ = sendErr(&interfaces.ErrorMessage{StatusCode: status, Error: streamErr, Addon: addon}) - return - } - if len(chunk.Payload) > 0 { - applyStreamHeaderInit() - payload := cloneBytes(chunk.Payload) - if streamInterceptorsActive { - executedReq, executedOpts := executedRequest() - intercepted := interceptStreamChunk(ctx, interceptorHost, pluginapi.StreamChunkInterceptRequest{ - SourceFormat: responseProtocol, - Model: normalizedModel, - RequestedModel: originalRequestedModel, - RequestHeaders: cloneHeader(executedOpts.Headers), - ResponseHeaders: cloneHeader(rawStreamHeaders), - OriginalRequest: cloneBytes(executedOpts.OriginalRequest), - RequestBody: cloneBytes(executedReq.Payload), - Body: payload, - HistoryChunks: cloneByteSlices(historyChunks), - ChunkIndex: chunkIndex, - Metadata: executedOpts.Metadata, - }, execOptions.SkipInterceptorPluginID) - applyStreamHeaders(intercepted.Headers) - if len(intercepted.Body) > 0 { - payload = cloneBytes(intercepted.Body) - } - chunkIndex++ - if intercepted.DropChunk { - continue - } - } else { - chunkIndex++ - } - if responseProtocol == "openai-response" { - if errValidate := validateSSEDataJSON(payload); errValidate != nil { - _ = sendErr(&interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: errValidate}) - return - } - } - sentPayload = true - streamHeadersCommitted = true - if okSendData := sendData(payload); !okSendData { - return - } - if streamInterceptorsActive { - historyChunks = appendStreamInterceptorHistory(historyChunks, payload) - } - } + chunkIndex := bootstrapChunkIndex + historyChunks := bootstrapHistoryChunks + if bootstrapPayload != nil { + if okSendData := sendData(bootstrapPayload); !okSendData { + return } - applyStreamHeaderInit() - return + if streamInterceptorsActive { + historyChunks = appendStreamInterceptorHistory(historyChunks, bootstrapPayload) + } + } + for { + chunk, ok, canceled := nextStreamChunk(ctx, nil, &streamClosedBeforeRead, chunks) + if canceled || !ok { + return + } + if chunk.Err != nil { + _ = sendErr(executionErrorMessage(chunk.Err)) + return + } + if len(chunk.Payload) == 0 { + continue + } + payload, deliverable, errMsg := transformStreamPayload(chunk.Payload, &chunkIndex, historyChunks) + if errMsg != nil { + _ = sendErr(errMsg) + return + } + if !deliverable { + continue + } + if okSendData := sendData(payload); !okSendData { + return + } + if streamInterceptorsActive { + historyChunks = appendStreamInterceptorHistory(historyChunks, payload) + } } }() return dataChan, upstreamHeaders, errChan @@ -1766,15 +1805,6 @@ total += len(item) } return total -} - -func replaceHeader(dst http.Header, src http.Header) { - for key := range dst { - delete(dst, key) - } - for key, values := range src { - dst[key] = append([]string(nil), values...) - } } func finalInterceptorHeaders(current, intercepted http.Header) http.Header { @@ -2205,6 +2235,11 @@ status := http.StatusInternalServerError if msg != nil && msg.StatusCode > 0 { status = msg.StatusCode + } + if msg != nil && msg.Error != nil { + for _, value := range coreauth.SafeResponseHeaders(msg.Error).Values("Retry-After") { + c.Writer.Header().Add("Retry-After", value) + } } if msg != nil && msg.Addon != nil && PassthroughHeadersEnabled(h.Cfg) { for key, values := range msg.Addon { diff --git a/sdk/api/handlers/handlers_error_response_test.go b/sdk/api/handlers/handlers_error_response_test.go --- a/sdk/api/handlers/handlers_error_response_test.go +++ b/sdk/api/handlers/handlers_error_response_test.go @@ -7,6 +7,7 @@ "reflect" "strings" "testing" + "time" "github.com/gin-gonic/gin" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" @@ -38,6 +39,52 @@ } if got := recorder.Header().Get("X-Request-Id"); got != "" { t.Fatalf("X-Request-Id should be empty when passthrough is disabled, got %q", got) + } +} + +func TestInternalConcurrencyBusyWritesRetryAfterWithoutPassthrough(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + + handler := NewBaseAPIHandlers(nil, nil) + handler.WriteErrorResponse(c, &interfaces.ErrorMessage{ + StatusCode: http.StatusTooManyRequests, + Error: coreauth.NewHomeConcurrencyBusyError("busy", 750*time.Millisecond), + }) + + if recorder.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusTooManyRequests) + } + if got := recorder.Header().Get("Retry-After"); got != "1" { + t.Fatalf("Retry-After = %q, want 1", got) + } +} + +func TestWriteErrorResponseHomeBusyNormalAndStreamHeaders(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "normal", true: "stream"}[stream], func(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + if stream { + c.Request.Header.Set("Accept", "text/event-stream") + } + + handler := NewBaseAPIHandlers(nil, nil) + handler.WriteErrorResponse(c, &interfaces.ErrorMessage{ + StatusCode: http.StatusTooManyRequests, + Error: coreauth.NewHomeConcurrencyBusyError("busy", 750*time.Millisecond), + }) + if recorder.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusTooManyRequests) + } + if got := recorder.Header().Get("Retry-After"); got != "1" { + t.Fatalf("Retry-After = %q, want 1", got) + } + }) } } diff --git a/sdk/api/handlers/handlers_interceptors_test.go b/sdk/api/handlers/handlers_interceptors_test.go --- a/sdk/api/handlers/handlers_interceptors_test.go +++ b/sdk/api/handlers/handlers_interceptors_test.go @@ -840,7 +840,7 @@ t.Fatalf("first chunk = %q, want first", firstChunk) } if upstreamHeaders.Get("X-Chunk") != "first" || upstreamHeaders.Get("X-Stage") != "init" { - t.Fatalf("upstream headers after first chunk = %#v, want first chunk headers", upstreamHeaders) + t.Fatalf("upstream headers after first chunk = %#v, want first transformed chunk headers", upstreamHeaders) } close(releaseSecond) @@ -857,7 +857,80 @@ t.Fatalf("stream payload = %q, want firstsecond", got) } if upstreamHeaders.Get("X-Chunk") != "first" { - t.Fatalf("upstream headers changed after first payload: %#v", upstreamHeaders) + t.Fatalf("upstream headers changed after return: %#v", upstreamHeaders) + } +} + +func TestHandlerStreamInterceptorReturnedHeadersImmutableAfterReturn(t *testing.T) { + model := "handler-interceptor-stream-immutable-headers-model" + releaseSecond := make(chan struct{}) + bodyStarted := make(chan struct{}) + releaseBody := make(chan struct{}) + executor := &interceptorCaptureExecutor{ + stream: func(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk) + go func() { + defer close(chunks) + chunks <- coreexecutor.StreamChunk{Payload: []byte("first")} + <-releaseSecond + chunks <- coreexecutor.StreamChunk{Payload: []byte("second")} + }() + return &coreexecutor.StreamResult{ + Headers: http.Header{"X-Upstream": []string{"stream"}}, + Chunks: chunks, + }, nil + }, + } + handler := newInterceptorHandler(t, model, executor, &sdkconfig.SDKConfig{PassthroughHeaders: true}) + handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptStreamChunk: func(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + headers := cloneHeader(req.ResponseHeaders) + switch req.ChunkIndex { + case pluginapi.StreamChunkHeaderInitIndex: + headers.Set("X-Init", "plugin") + case 1: + close(bodyStarted) + <-releaseBody + headers.Set("X-Body", "plugin") + } + return pluginapi.StreamChunkInterceptResponse{Headers: headers, Body: cloneBytes(req.Body)} + }, + }) + + dataChan, upstreamHeaders, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", model, []byte(fmt.Sprintf(`{"model":%q}`, model)), "") + dataDone := make(chan struct{}) + go func() { + defer close(dataDone) + for range dataChan { + } + }() + stopReading := make(chan struct{}) + readerDone := make(chan struct{}) + go func() { + defer close(readerDone) + for { + select { + case <-stopReading: + return + default: + _ = upstreamHeaders.Get("X-Init") + } + } + }() + + close(releaseSecond) + <-bodyStarted + close(releaseBody) + <-dataDone + for msg := range errChan { + if msg != nil { + t.Fatalf("unexpected stream error: %+v", msg) + } + } + close(stopReading) + <-readerDone + if upstreamHeaders.Get("X-Init") != "plugin" || upstreamHeaders.Get("X-Body") != "" { + t.Fatalf("returned headers mutated after return: %#v", upstreamHeaders) } } diff --git a/sdk/api/handlers/handlers_model_router_test.go b/sdk/api/handlers/handlers_model_router_test.go --- a/sdk/api/handlers/handlers_model_router_test.go +++ b/sdk/api/handlers/handlers_model_router_test.go @@ -10,6 +10,7 @@ "time" "github.com/gin-gonic/gin" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" @@ -89,6 +90,7 @@ lastPluginID string lastRequest coreexecutor.Request lastOptions coreexecutor.Options + stream func(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) } func (h *handlerDirectExecutorRouteHost) ExecutePluginExecutor(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { @@ -102,6 +104,9 @@ h.lastPluginID = pluginID h.lastRequest = req h.lastOptions = opts + if h.stream != nil { + return h.stream(ctx, pluginID, req, opts) + } chunks := make(chan coreexecutor.StreamChunk, 1) chunks <- coreexecutor.StreamChunk{Payload: []byte("direct-stream")} close(chunks) @@ -226,6 +231,39 @@ } if string(host.lastOptions.OriginalRequest) != `{"after":true}` { t.Fatalf("original request = %q, want after-auth body", host.lastOptions.OriginalRequest) + } +} + +func TestHandlerModelRouterPluginExecutorFailsClosedWhenHomeEnabled(t *testing.T) { + originalModel := "home-plugin-route" + targetPluginID := "plugin-executor" + host := &handlerDirectExecutorRouteHost{} + host.hasRouters = true + host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { + return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true + } + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + handler.SetModelRouterHost(host) + + body, _, errMsg := handler.ExecuteWithAuthManager(context.Background(), "openai", originalModel, []byte(`{"model":"home-plugin-route"}`), "") + if body != nil || errMsg == nil || errMsg.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("ExecuteWithAuthManager() = %q, %#v; want 503", body, errMsg) + } + body, _, errMsg = handler.ExecuteCountWithAuthManager(context.Background(), "openai", originalModel, []byte(`{"model":"home-plugin-route"}`), "") + if body != nil || errMsg == nil || errMsg.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("ExecuteCountWithAuthManager() = %q, %#v; want 503", body, errMsg) + } + data, _, errors := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(`{"model":"home-plugin-route","stream":true}`), "") + if data != nil { + t.Fatalf("ExecuteStreamWithAuthManager() data = %v, want nil", data) + } + if errMsg = <-errors; errMsg == nil || errMsg.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("ExecuteStreamWithAuthManager() error = %#v, want 503", errMsg) + } + if host.lastPluginID != "" { + t.Fatalf("plugin executor was invoked with %q while Home was enabled", host.lastPluginID) } } @@ -619,6 +657,84 @@ case <-deadline: t.Fatal("plugin executor stream goroutine did not exit after context cancel") } + } +} + +func TestStreamWithPluginExecutorReturnedHeadersImmutableAfterReturn(t *testing.T) { + originalModel := "handler-router-plugin-immutable-headers-model" + targetPluginID := "immutable-headers-plugin" + releaseSecond := make(chan struct{}) + bodyStarted := make(chan struct{}) + releaseBody := make(chan struct{}) + host := &handlerDirectExecutorRouteHost{} + host.stream = func(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk) + go func() { + defer close(chunks) + chunks <- coreexecutor.StreamChunk{Payload: []byte("first")} + <-releaseSecond + chunks <- coreexecutor.StreamChunk{Payload: []byte("second")} + }() + return &coreexecutor.StreamResult{Chunks: chunks}, nil + } + host.hasRouters = true + host.route = func(ctx context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { + return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true + } + handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{PassthroughHeaders: true}, nil) + handler.SetModelRouterHost(host) + handler.SetPluginHost(&handlerInterceptorTestHost{ + interceptStreamChunk: func(ctx context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + headers := cloneHeader(req.ResponseHeaders) + if headers == nil { + headers = make(http.Header) + } + switch req.ChunkIndex { + case pluginapi.StreamChunkHeaderInitIndex: + headers.Set("X-Init", "plugin") + case 1: + close(bodyStarted) + <-releaseBody + headers.Set("X-Body", "plugin") + } + return pluginapi.StreamChunkInterceptResponse{Headers: headers, Body: cloneBytes(req.Body)} + }, + }) + + dataChan, upstreamHeaders, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", originalModel, []byte(fmt.Sprintf(`{"model":%q,"stream":true}`, originalModel)), "") + dataDone := make(chan struct{}) + go func() { + defer close(dataDone) + for range dataChan { + } + }() + stopReading := make(chan struct{}) + readerDone := make(chan struct{}) + go func() { + defer close(readerDone) + for { + select { + case <-stopReading: + return + default: + _ = upstreamHeaders.Get("X-Init") + } + } + }() + + close(releaseSecond) + <-bodyStarted + close(releaseBody) + <-dataDone + for msg := range errChan { + if msg != nil { + t.Fatalf("unexpected stream error: %+v", msg) + } + } + close(stopReading) + <-readerDone + if upstreamHeaders.Get("X-Init") != "plugin" || upstreamHeaders.Get("X-Body") != "" { + t.Fatalf("returned headers mutated after return: %#v", upstreamHeaders) } } diff --git a/sdk/api/handlers/handlers_stream_bootstrap_test.go b/sdk/api/handlers/handlers_stream_bootstrap_test.go --- a/sdk/api/handlers/handlers_stream_bootstrap_test.go +++ b/sdk/api/handlers/handlers_stream_bootstrap_test.go @@ -2,21 +2,26 @@ import ( "context" + "encoding/json" "errors" "net/http" "net/http/httptest" "strings" "sync" + "sync/atomic" "testing" "time" "github.com/gin-gonic/gin" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) type failOnceStreamExecutor struct { @@ -81,6 +86,56 @@ e.mu.Lock() defer e.mu.Unlock() return e.calls +} + +type blockingRetryStreamExecutor struct { + mu sync.Mutex + calls int + retryStarted chan struct{} + allowRetry chan struct{} +} + +func (e *blockingRetryStreamExecutor) Identifier() string { return "codex" } + +func (e *blockingRetryStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} +} + +func (e *blockingRetryStreamExecutor) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { + e.mu.Lock() + e.calls++ + call := e.calls + e.mu.Unlock() + + if call == 1 { + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{Code: "unauthorized", Message: "unauthorized", HTTPStatus: http.StatusUnauthorized}} + close(chunks) + return &coreexecutor.StreamResult{Headers: http.Header{"X-Upstream-Attempt": {"1"}}, Chunks: chunks}, nil + } + + close(e.retryStarted) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-e.allowRetry: + } + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Payload: []byte("ok")} + close(chunks) + return &coreexecutor.StreamResult{Headers: http.Header{"X-Upstream-Attempt": {"2"}}, Chunks: chunks}, nil +} + +func (e *blockingRetryStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (e *blockingRetryStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} +} + +func (e *blockingRetryStreamExecutor) HttpRequest(ctx context.Context, auth *coreauth.Auth, req *http.Request) (*http.Response, error) { + return nil, &coreauth.Error{Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented} } type payloadThenErrorStreamExecutor struct { @@ -337,6 +392,328 @@ upstreamAttemptHeader := upstreamHeaders.Get("X-Upstream-Attempt") if upstreamAttemptHeader != "2" { t.Fatalf("expected upstream header from retry attempt, got %q", upstreamAttemptHeader) + } +} + +func TestExecuteStreamWithAuthManager_ResolvesBootstrapRetryHeadersBeforeReturn(t *testing.T) { + executor := &blockingRetryStreamExecutor{ + retryStarted: make(chan struct{}), + allowRetry: make(chan struct{}), + } + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth1 := &coreauth.Auth{ID: "auth1", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test1@example.com"}} + if _, err := manager.Register(context.Background(), auth1); err != nil { + t.Fatalf("manager.Register(auth1): %v", err) + } + auth2 := &coreauth.Auth{ID: "auth2", Provider: "codex", Status: coreauth.StatusActive, Metadata: map[string]any{"email": "test2@example.com"}} + if _, err := manager.Register(context.Background(), auth2); err != nil { + t.Fatalf("manager.Register(auth2): %v", err) + } + registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}}) + registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth1.ID) + registry.GetGlobalRegistry().UnregisterClient(auth2.ID) + }) + + handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{PassthroughHeaders: true, Streaming: sdkconfig.StreamingConfig{BootstrapRetries: 1}}, manager) + type streamResult struct { + dataChan <-chan []byte + upstreamHeaders http.Header + errChan <-chan *interfaces.ErrorMessage + } + resultChan := make(chan streamResult, 1) + go func() { + dataChan, upstreamHeaders, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "") + resultChan <- streamResult{dataChan: dataChan, upstreamHeaders: upstreamHeaders, errChan: errChan} + }() + + select { + case result := <-resultChan: + t.Fatalf("ExecuteStreamWithAuthManager returned before bootstrap retry completed: %#v", result.upstreamHeaders) + case <-executor.retryStarted: + } + select { + case result := <-resultChan: + t.Fatalf("ExecuteStreamWithAuthManager returned while bootstrap retry was blocked: %#v", result.upstreamHeaders) + default: + } + close(executor.allowRetry) + + result := <-resultChan + if result.upstreamHeaders.Get("X-Upstream-Attempt") != "2" { + t.Fatalf("upstream headers = %#v, want retry attempt headers", result.upstreamHeaders) + } + for range result.dataChan { + } + for msg := range result.errChan { + if msg != nil { + t.Fatalf("unexpected stream error: %+v", msg) + } + } +} + +type bootstrapStreamExecutor struct { + mu sync.Mutex + calls int + stream func(context.Context, int) (*coreexecutor.StreamResult, error) +} + +func (*bootstrapStreamExecutor) Identifier() string { return "bootstrap-test" } + +func (e *bootstrapStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"} +} + +func (e *bootstrapStreamExecutor) ExecuteStream(ctx context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { + e.mu.Lock() + e.calls++ + call := e.calls + e.mu.Unlock() + return e.stream(ctx, call) +} + +func (e *bootstrapStreamExecutor) Refresh(context.Context, *coreauth.Auth) (*coreauth.Auth, error) { + return nil, nil +} + +func (e *bootstrapStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"} +} + +func (e *bootstrapStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, &coreauth.Error{Code: "not_implemented", Message: "HttpRequest not implemented", HTTPStatus: http.StatusNotImplemented} +} + +func (e *bootstrapStreamExecutor) Calls() int { + e.mu.Lock() + defer e.mu.Unlock() + return e.calls +} + +func registerBootstrapExecutor(t *testing.T, executor *bootstrapStreamExecutor) (*BaseAPIHandler, *coreauth.Manager) { + t.Helper() + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &coreauth.Auth{ID: "bootstrap-auth", Provider: executor.Identifier(), Status: coreauth.StatusActive, Metadata: map[string]any{"email": "bootstrap@example.com"}} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("manager.Register(): %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: "bootstrap-model"}}) + authRetry := &coreauth.Auth{ID: "bootstrap-auth-retry", Provider: executor.Identifier(), Status: coreauth.StatusActive, Metadata: map[string]any{"email": "bootstrap-retry@example.com"}} + if _, errRegister := manager.Register(context.Background(), authRetry); errRegister != nil { + t.Fatalf("manager.Register(retry): %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(authRetry.ID, authRetry.Provider, []*registry.ModelInfo{{ID: "bootstrap-model"}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + registry.GetGlobalRegistry().UnregisterClient(authRetry.ID) + }) + return NewBaseAPIHandlers(&sdkconfig.SDKConfig{Streaming: sdkconfig.StreamingConfig{BootstrapRetries: 1}}, manager), manager +} + +func TestExecuteStreamWithAuthManager_RetriesAfterDroppedBootstrapPayload(t *testing.T) { + executor := &bootstrapStreamExecutor{stream: func(_ context.Context, call int) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 2) + if call == 1 { + chunks <- coreexecutor.StreamChunk{Payload: []byte("drop")} + chunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}} + } else { + chunks <- coreexecutor.StreamChunk{Payload: []byte("ok")} + } + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }} + handler, _ := registerBootstrapExecutor(t, executor) + var intercepted []string + handler.SetPluginHost(&handlerInterceptorTestHost{interceptStreamChunk: func(_ context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + if req.ChunkIndex >= 0 { + intercepted = append(intercepted, string(req.Body)) + } + return pluginapi.StreamChunkInterceptResponse{Body: cloneBytes(req.Body), DropChunk: string(req.Body) == "drop"} + }}) + + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "bootstrap-model", []byte(`{"model":"bootstrap-model"}`), "") + var got []byte + for chunk := range dataChan { + got = append(got, chunk...) + } + for msg := range errChan { + if msg != nil { + t.Fatalf("unexpected stream error: %+v", msg) + } + } + if string(got) != "ok" { + t.Fatalf("stream payload = %q, want ok", got) + } + if executor.Calls() != 2 { + t.Fatalf("stream attempts = %d, want 2", executor.Calls()) + } + if strings.Join(intercepted, ",") != "drop,ok" { + t.Fatalf("intercepted payloads = %v, want [drop ok] without double interception", intercepted) + } +} + +func TestExecuteStreamWithAuthManager_CancelDuringSynchronousBootstrap(t *testing.T) { + started := make(chan struct{}) + executor := &bootstrapStreamExecutor{stream: func(_ context.Context, _ int) (*coreexecutor.StreamResult, error) { + close(started) + return &coreexecutor.StreamResult{Chunks: make(chan coreexecutor.StreamChunk)}, nil + }} + handler, _ := registerBootstrapExecutor(t, executor) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + type result struct { + data <-chan []byte + errs <-chan *interfaces.ErrorMessage + } + results := make(chan result, 1) + go func() { + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", "bootstrap-model", []byte(`{"model":"bootstrap-model"}`), "") + results <- result{data: dataChan, errs: errChan} + }() + <-started + cancel() + select { + case got := <-results: + if got.data != nil { + if _, ok := <-got.data; ok { + t.Fatal("data channel remains open after bootstrap cancellation") + } + } + if got.errs != nil { + for range got.errs { + } + } + case <-time.After(time.Second): + t.Fatal("bootstrap cancellation did not return") + } +} + +func TestExecuteStreamWithAuthManager_EmptyClosedStream(t *testing.T) { + executor := &bootstrapStreamExecutor{stream: func(_ context.Context, _ int) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk) + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }} + handler, _ := registerBootstrapExecutor(t, executor) + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "bootstrap-model", []byte(`{"model":"bootstrap-model"}`), "") + if _, ok := <-dataChan; ok { + t.Fatal("empty stream produced data") + } + var streamErr *interfaces.ErrorMessage + for msg := range errChan { + if msg != nil { + streamErr = msg + } + } + if streamErr == nil || streamErr.StatusCode != http.StatusInternalServerError { + t.Fatalf("empty stream error = %+v, want terminal internal-server error", streamErr) + } +} + +type handlerReleaseNotification struct { + group executionregistry.ReleaseGroup + sequence int64 +} + +type handlerReleaseSink struct { + mu sync.Mutex + notifications []handlerReleaseNotification + notified chan struct{} +} + +func newHandlerReleaseSink() *handlerReleaseSink { + return &handlerReleaseSink{notified: make(chan struct{}, 1)} +} + +func (s *handlerReleaseSink) MarkDirty(group executionregistry.ReleaseGroup, sequence int64) { + s.mu.Lock() + s.notifications = append(s.notifications, handlerReleaseNotification{group: group, sequence: sequence}) + s.mu.Unlock() + select { + case s.notified <- struct{}{}: + default: + } +} + +func (s *handlerReleaseSink) Notifications() []handlerReleaseNotification { + s.mu.Lock() + defer s.mu.Unlock() + return append([]handlerReleaseNotification(nil), s.notifications...) +} + +type handlerAccountedHomeDispatcher struct { + calls atomic.Int32 +} + +func (*handlerAccountedHomeDispatcher) HeartbeatOK() bool { return true } +func (d *handlerAccountedHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + d.calls.Add(1) + return json.Marshal(map[string]any{ + "concurrency": map[string]any{"accounted": true, "credential_id": "handler-cred", "model": model}, + "model": model, + "auth_index": "handler-cred", + "auth": map[string]any{"id": "handler-cred", "provider": "bootstrap-test", "status": coreauth.StatusActive}, + }) +} +func (*handlerAccountedHomeDispatcher) AbortAmbiguousDispatch() {} + +func TestExecuteStreamWithAuthManager_HomeBootstrapFailureDoesNotRedispatch(t *testing.T) { + executor := &bootstrapStreamExecutor{stream: func(_ context.Context, _ int) (*coreexecutor.StreamResult, error) { + chunks := make(chan coreexecutor.StreamChunk, 2) + chunks <- coreexecutor.StreamChunk{Payload: []byte("drop")} + chunks <- coreexecutor.StreamChunk{Err: &coreauth.Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + }} + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.RegisterExecutor(executor) + registry := executionregistry.New() + releaseSink := newHandlerReleaseSink() + registry.SetReleaseSink(releaseSink.MarkDirty) + dispatcher := &handlerAccountedHomeDispatcher{} + manager.PublishHomeDispatch(dispatcher, registry, 1) + handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{Streaming: sdkconfig.StreamingConfig{BootstrapRetries: 1}}, manager) + handler.SetPluginHost(&handlerInterceptorTestHost{interceptStreamChunk: func(_ context.Context, req pluginapi.StreamChunkInterceptRequest) pluginapi.StreamChunkInterceptResponse { + return pluginapi.StreamChunkInterceptResponse{Body: cloneBytes(req.Body), DropChunk: string(req.Body) == "drop"} + }}) + + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "home-model", []byte(`{"model":"home-model"}`), "") + for range dataChan { + t.Fatal("Home bootstrap failure produced data") + } + var streamErr *interfaces.ErrorMessage + for msg := range errChan { + if msg != nil { + streamErr = msg + } + } + if streamErr == nil || streamErr.StatusCode != http.StatusUnauthorized { + t.Fatalf("stream error = %+v, want unauthorized terminal error", streamErr) + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1", got) + } + select { + case <-releaseSink.notified: + case <-time.After(time.Second): + t.Fatal("accounted Home selection was not released") + } + wantRelease := handlerReleaseNotification{ + group: executionregistry.ReleaseGroup{CredentialID: "handler-cred", Model: "home-model"}, + sequence: 1, + } + if got := releaseSink.Notifications(); len(got) != 1 || got[0] != wantRelease { + t.Fatalf("release notifications = %#v, want [%#v]", got, wantRelease) + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("registry.Drain(): %v", errDrain) + } + if got := releaseSink.Notifications(); len(got) != 1 || got[0] != wantRelease { + t.Fatalf("release notifications after drain = %#v, want [%#v]", got, wantRelease) } } diff --git a/sdk/cliproxy/auth/antigravity_credits_test.go b/sdk/cliproxy/auth/antigravity_credits_test.go --- a/sdk/cliproxy/auth/antigravity_credits_test.go +++ b/sdk/cliproxy/auth/antigravity_credits_test.go @@ -128,7 +128,7 @@ } } -func TestManagerExecuteStream_AntigravityCreditsHomeKVUnavailableFailsRequest(t *testing.T) { +func TestManagerExecuteStream_AntigravityCreditsHomeModeFailsClosedWithoutDispatch(t *testing.T) { const model = "claude-opus-4-6-thinking" executor := &antigravityCreditsFallbackExecutor{} manager := NewManager(nil, nil, nil) @@ -152,8 +152,8 @@ if status := statusCodeFromError(errExecute); status != http.StatusServiceUnavailable { t.Fatalf("ExecuteStream() status = %d, want %d; err=%v", status, http.StatusServiceUnavailable, errExecute) } - if !strings.Contains(errExecute.Error(), "home kv store unavailable") { - t.Fatalf("ExecuteStream() error = %v, want home kv store unavailable", errExecute) + if !strings.Contains(errExecute.Error(), "home dispatch bundle unavailable") { + t.Fatalf("ExecuteStream() error = %v, want home dispatch bundle unavailable", errExecute) } } diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -24,6 +24,7 @@ "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/internal/thinking" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" @@ -217,21 +218,29 @@ // Manager orchestrates auth lifecycle, selection, execution, and persistence. type Manager struct { - store Store - cooldownStore CooldownStateStore - executors map[string]ProviderExecutor - selector Selector - hook Hook - mu sync.RWMutex - auths map[string]*Auth - scheduler *authScheduler + store Store + cooldownStore CooldownStateStore + pendingCooldownStateStore CooldownStateStore + executors map[string]ProviderExecutor + selector Selector + hook Hook + mu sync.RWMutex + configCooldownMu sync.Mutex + auths map[string]*Auth + scheduler *authScheduler // pluginScheduler runs outside m.mu before falling back to native selection. pluginScheduler PluginScheduler - // homeRuntimeAuths caches auths returned by Home so websocket sessions can - // reuse an established upstream credential without dispatching every turn. + // homeRuntimeAuths retains legacy session auth lookups for non-execution callers. homeRuntimeAuths map[string]map[string]*Auth + // homeRuntimeAuthOwners prevents a stale selection from clearing a replacement auth. + homeRuntimeAuthOwners map[string]map[string]*HomeDispatchSelection + // homeSessionSelections owns retained Home selections for websocket sessions. + homeSessionSelections map[string]map[homeSessionSelectionKey]*HomeDispatchSelection + homeSessionLocks sync.Map // providerOffsets tracks per-model provider rotation state for multi-provider routing. - providerOffsets map[string]int + providerOffsets map[string]int + homeDispatchBundle atomic.Pointer[HomeDispatchBundle] + homeInFlightPublisherConfig atomic.Pointer[HomeInFlightPublisherConfig] // Retry controls request retry behavior. requestRetry atomic.Int32 @@ -274,20 +283,85 @@ hook = NoopHook{} } manager := &Manager{ - store: store, - executors: make(map[string]ProviderExecutor), - selector: selector, - hook: hook, - auths: make(map[string]*Auth), - homeRuntimeAuths: make(map[string]map[string]*Auth), - providerOffsets: make(map[string]int), - modelPoolOffsets: make(map[string]int), + store: store, + executors: make(map[string]ProviderExecutor), + selector: selector, + hook: hook, + auths: make(map[string]*Auth), + homeRuntimeAuths: make(map[string]map[string]*Auth), + homeRuntimeAuthOwners: make(map[string]map[string]*HomeDispatchSelection), + homeSessionSelections: make(map[string]map[homeSessionSelectionKey]*HomeDispatchSelection), + providerOffsets: make(map[string]int), + modelPoolOffsets: make(map[string]int), } // atomic.Value requires non-nil initial value. manager.runtimeConfig.Store(&internalconfig.Config{}) manager.apiKeyModelAlias.Store(apiKeyModelAliasTable(nil)) + defaultInFlightConfig, errInFlightConfig := HomeInFlightPublisherConfigFromConfig(internalconfig.DefaultCredentialInFlightConfig()) + if errInFlightConfig == nil { + manager.ApplyHomeInFlightPublisherConfig(defaultInFlightConfig) + } manager.scheduler = newAuthScheduler(selector) return manager +} + +// HomeDispatchBundle is the immutable client and registry pair for one Home lifetime. +type HomeDispatchBundle struct { + client homeAuthDispatcher + registry *executionregistry.Registry + generation uint64 +} + +// PublishHomeDispatch publishes the selectable Home lifetime as one atomic bundle. +func (m *Manager) PublishHomeDispatch(client homeAuthDispatcher, registry *executionregistry.Registry, generation uint64) *HomeDispatchBundle { + if m == nil || client == nil || registry == nil { + return nil + } + bundle := &HomeDispatchBundle{client: client, registry: registry, generation: generation} + m.homeDispatchBundle.Store(bundle) + return bundle +} + +// ClearHomeDispatchBundle removes bundle only when it still belongs to the active lifetime. +func (m *Manager) ClearHomeDispatchBundle(bundle *HomeDispatchBundle) bool { + if m == nil || bundle == nil { + return false + } + return m.homeDispatchBundle.CompareAndSwap(bundle, nil) +} + +// HomeDispatchBundle returns the active Home lifetime bundle. +func (m *Manager) HomeDispatchBundle() *HomeDispatchBundle { + if m == nil { + return nil + } + return m.homeDispatchBundle.Load() +} + +// SetHomeExecutionRegistry preserves the legacy registry API for callers that also install the current dispatcher. +func (m *Manager) SetHomeExecutionRegistry(registry *executionregistry.Registry) { + if m == nil { + return + } + m.PublishHomeDispatch(currentHomeDispatcher(), registry, 0) +} + +// ClearHomeExecutionRegistry removes a matching legacy registry bundle. +func (m *Manager) ClearHomeExecutionRegistry(registry *executionregistry.Registry) bool { + bundle := m.HomeDispatchBundle() + if bundle == nil || bundle.registry != registry { + return false + } + return m.ClearHomeDispatchBundle(bundle) +} + +// HomeExecutionRegistry returns the registry from the active Home lifetime bundle. +func (m *Manager) HomeExecutionRegistry() *executionregistry.Registry { + bundle := m.HomeDispatchBundle() + if bundle == nil { + return nil + } + return bundle.registry } func (m *Manager) SetPluginScheduler(scheduler PluginScheduler) { @@ -478,6 +552,16 @@ } } +// Selector returns the current credential selector. +func (m *Manager) Selector() Selector { + if m == nil { + return nil + } + m.mu.RLock() + defer m.mu.RUnlock() + return m.selector +} + // SetStore swaps the underlying persistence store. func (m *Manager) SetStore(store Store) { m.mu.Lock() @@ -490,6 +574,8 @@ if m == nil { return } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() m.mu.Lock() defer m.mu.Unlock() m.cooldownStore = store @@ -508,18 +594,128 @@ if m == nil { return } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + if m.setConfigSnapshotLocked(cfg) { + m.persistCooldownStatesLocked(context.Background()) + } +} + +// SetConfigSnapshot updates only in-memory configuration state. It reports whether +// a caller must persist cleared cooldown state after its commit critical section. +func (m *Manager) SetConfigSnapshot(cfg *internalconfig.Config) bool { + if m == nil { + return false + } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + return m.setConfigSnapshotLocked(cfg) +} + +func (m *Manager) setConfigSnapshotLocked(cfg *internalconfig.Config) bool { if cfg == nil { cfg = &internalconfig.Config{} } + m.mu.RLock() + oldCooldownStore := m.cooldownStore + m.mu.RUnlock() m.runtimeConfig.Store(cfg) clearedCooldowns := m.clearDisabledCooldownStates(cfg) + if clearedCooldowns && oldCooldownStore != nil { + m.mu.Lock() + if m.cooldownStore == oldCooldownStore { + m.pendingCooldownStateStore = oldCooldownStore + } + m.mu.Unlock() + } if !cfg.Home.Enabled { m.clearHomeRuntimeAuths() } m.rebuildAPIKeyModelAliasFromRuntimeConfig() - if clearedCooldowns { - m.persistCooldownStates(context.Background()) + return clearedCooldowns +} + +// ApplyConfigWithCooldownStateStore serializes a config update with its cooldown +// store transition. It persists the resulting state to the captured old store before +// exposing the resolved replacement store. +func (m *Manager) ApplyConfigWithCooldownStateStore(ctx context.Context, cfg *internalconfig.Config, store CooldownStateStore) bool { + if m == nil { + return false } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + m.mu.RLock() + oldStore := m.cooldownStore + m.mu.RUnlock() + m.setConfigSnapshotLocked(cfg) + if oldStore != nil && !m.persistCooldownStatesToLocked(ctx, oldStore) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if m.cooldownStore != oldStore { + return false + } + if m.pendingCooldownStateStore == oldStore { + m.pendingCooldownStateStore = nil + } + m.cooldownStore = store + return true +} + +// PersistCooldownStates writes the current cooldown snapshot using ctx. +func (m *Manager) PersistCooldownStates(ctx context.Context) { + m.persistCooldownStates(ctx) +} + +// SwapCooldownStateStore persists cleared state to the old store before replacing it. +// Persistence is deliberately performed without holding the manager lock. +func (m *Manager) SwapCooldownStateStore(ctx context.Context, store CooldownStateStore, persistOld bool) bool { + if m == nil { + return false + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + m.mu.RLock() + oldStore := m.cooldownStore + pendingStore := m.pendingCooldownStateStore + m.mu.RUnlock() + storeToPersist := pendingStore + if storeToPersist == nil && persistOld { + storeToPersist = oldStore + } + if storeToPersist != nil && !m.persistCooldownStatesToLocked(ctx, storeToPersist) { + return false + } + if errContext := ctx.Err(); errContext != nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if m.cooldownStore != oldStore { + return false + } + if m.pendingCooldownStateStore == storeToPersist { + m.pendingCooldownStateStore = nil + } + m.cooldownStore = store + return true } func (m *Manager) cooldownDisabledForAuth(auth *Auth) bool { @@ -808,28 +1004,47 @@ if m == nil { return } - if ctx == nil { - ctx = context.Background() - } - records, store := m.cooldownStateSnapshot() - if store == nil { - return - } - if errSave := store.Save(ctx, records); errSave != nil { - logEntryWithRequestID(ctx).Warnf("failed to persist cooldown state: %v", errSave) + m.configCooldownMu.Lock() + defer m.configCooldownMu.Unlock() + m.persistCooldownStatesLocked(ctx) +} + +func (m *Manager) persistCooldownStatesLocked(ctx context.Context) { + m.mu.RLock() + store := m.cooldownStore + m.mu.RUnlock() + if m.persistCooldownStatesToLocked(ctx, store) { + m.mu.Lock() + if m.pendingCooldownStateStore == store { + m.pendingCooldownStateStore = nil + } + m.mu.Unlock() } } -func (m *Manager) cooldownStateSnapshot() ([]CooldownStateRecord, CooldownStateStore) { +func (m *Manager) persistCooldownStatesToLocked(ctx context.Context, store CooldownStateStore) bool { + if m == nil || store == nil { + return true + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return false + } + records := m.cooldownStateRecordsSnapshot() + if errSave := store.Save(ctx, records); errSave != nil { + logEntryWithRequestID(ctx).Warnf("failed to persist cooldown state: %v", errSave) + return false + } + return ctx.Err() == nil +} + +func (m *Manager) cooldownStateRecordsSnapshot() []CooldownStateRecord { now := time.Now() records := make([]CooldownStateRecord, 0) m.mu.RLock() - store := m.cooldownStore - if store == nil { - m.mu.RUnlock() - return nil, nil - } for _, auth := range m.auths { records = append(records, m.cooldownStateRecordsForAuthLocked(auth, now)...) } @@ -844,7 +1059,7 @@ } return records[i].Model < records[j].Model }) - return records, store + return records } func (m *Manager) cooldownStateRecordsForAuthLocked(auth *Auth, now time.Time) []CooldownStateRecord { @@ -971,6 +1186,23 @@ } cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) return cfg != nil && cfg.Home.Enabled +} + +func (m *Manager) localExecutionAllowed() bool { + return m != nil && !m.HomeEnabled() +} + +func (m *Manager) localFallbackAuth(authID string) *Auth { + if !m.localExecutionAllowed() { + return nil + } + m.mu.RLock() + auth := m.auths[strings.TrimSpace(authID)] + m.mu.RUnlock() + if auth == nil { + return nil + } + return auth.Clone() } func (m *Manager) lookupAPIKeyUpstreamModel(authID, requestedModel string) string { @@ -1214,6 +1446,9 @@ func (m *Manager) executionModelCandidatesWithAlias(auth *Auth, routeModel string) ([]string, bool, OAuthModelAliasResult) { requestedModel := rewriteModelForAuth(routeModel, auth) aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel) + if aliasResult.ForceMapping && auth != nil && auth.Attributes != nil && strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") { + aliasResult.OriginalAlias = strings.TrimSpace(routeModel) + } upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult) var candidates []string @@ -1262,7 +1497,9 @@ return OAuthModelAliasResult{} } originalAlias := strings.TrimSpace(auth.Attributes[homeOriginalAliasAttributeKey]) - if originalAlias == "" { + canonicalOriginalAlias := canonicalHomeConcurrencyModelKey(auth.Attributes[homeOriginalAliasAttributeKey]) + canonicalRequestedModel := canonicalHomeConcurrencyModelKey(requestedModel) + if canonicalOriginalAlias == "" || canonicalOriginalAlias != canonicalRequestedModel { return OAuthModelAliasResult{} } upstreamModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]) @@ -1772,7 +2009,7 @@ } } -func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult) *cliproxyexecutor.StreamResult { +func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool) *cliproxyexecutor.StreamResult { out := make(chan cliproxyexecutor.StreamChunk) go func() { defer close(out) @@ -1786,7 +2023,7 @@ if chunk.Err != nil && !failed { failed = true rerr := resultErrorFromError(chunk.Err) - m.MarkResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}) + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr}, auth, ephemeralResult) } if !forward { return false @@ -1843,13 +2080,13 @@ } } if !failed { - m.MarkResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}) + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}, auth, ephemeralResult) } }() return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} } -func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult) (*cliproxyexecutor.StreamResult, error) { +func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor ProviderExecutor, auth *Auth, provider string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, routeModel, executionModel string, execModels []string, pooled bool, aliasResult OAuthModelAliasResult, allowRetry bool, ephemeralResult bool) (*cliproxyexecutor.StreamResult, error) { if executor == nil { return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} } @@ -1865,27 +2102,35 @@ } execOpts := opts execReq, execOpts = applyRequestAfterAuthInterceptor(ctx, executor, provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { return nil, errCtx } - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, errStream, didRefreshOnUnauthorized); okRefresh { - auth = refreshed - didRefreshOnUnauthorized = true - streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) - if errStream != nil { - if errCtx := ctx.Err(); errCtx != nil { - return nil, errCtx + if allowRetry { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, errStream, didRefreshOnUnauthorized); okRefresh { + auth = refreshed + didRefreshOnUnauthorized = true + streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) + if errStream != nil { + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } } } } + } + if errStream == nil && (streamResult == nil || streamResult.Chunks == nil) { + errStream = &Error{Code: "empty_stream", Message: "upstream stream has no source", Retryable: true} } if errStream != nil { rerr := resultErrorFromError(errStream) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} result.RetryAfter = retryAfterFromError(errStream) - m.MarkResult(ctx, result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) if isRequestInvalidError(errStream) { return nil, errStream } @@ -1899,20 +2144,22 @@ discardStreamChunks(streamResult.Chunks) return nil, errCtx } - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, bootstrapErr, didRefreshOnUnauthorized); okRefresh { - discardStreamChunks(streamResult.Chunks) - auth = refreshed - didRefreshOnUnauthorized = true - retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) - if retryErr != nil { - if errCtx := ctx.Err(); errCtx != nil { - return nil, errCtx + if allowRetry { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(ctx, auth, bootstrapErr, didRefreshOnUnauthorized); okRefresh { + discardStreamChunks(streamResult.Chunks) + auth = refreshed + didRefreshOnUnauthorized = true + retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) + if retryErr != nil { + if errCtx := ctx.Err(); errCtx != nil { + return nil, errCtx + } + bootstrapErr = retryErr + streamResult = &cliproxyexecutor.StreamResult{} + } else { + streamResult = retryStream + buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks) } - bootstrapErr = retryErr - streamResult = &cliproxyexecutor.StreamResult{} - } else { - streamResult = retryStream - buffered, closed, bootstrapErr = readStreamBootstrap(ctx, streamResult.Chunks) } } } @@ -1921,7 +2168,7 @@ rerr := resultErrorFromError(bootstrapErr) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} result.RetryAfter = retryAfterFromError(bootstrapErr) - m.MarkResult(ctx, result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) discardStreamChunks(streamResult.Chunks) return nil, bootstrapErr } @@ -1929,7 +2176,7 @@ rerr := resultErrorFromError(bootstrapErr) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} result.RetryAfter = retryAfterFromError(bootstrapErr) - m.MarkResult(ctx, result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) discardStreamChunks(streamResult.Chunks) lastErr = bootstrapErr continue @@ -1937,7 +2184,7 @@ rerr := resultErrorFromError(bootstrapErr) result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: rerr} result.RetryAfter = retryAfterFromError(bootstrapErr) - m.MarkResult(ctx, result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) discardStreamChunks(streamResult.Chunks) return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) } @@ -1945,7 +2192,7 @@ if closed && len(buffered) == 0 { emptyErr := &Error{Code: "empty_stream", Message: "upstream stream closed before first payload", Retryable: true} result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: false, Error: emptyErr} - m.MarkResult(ctx, result) + m.recordExecutionResult(ctx, result, auth, ephemeralResult) if idx < len(execModels)-1 { lastErr = emptyErr continue @@ -1959,7 +2206,7 @@ close(closedCh) remaining = closedCh } - return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, aliasResult), nil + return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, streamResult.Headers, buffered, remaining, aliasResult, ephemeralResult), nil } if lastErr == nil { lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"} @@ -2334,6 +2581,9 @@ if len(normalized) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } + if m.HomeEnabled() { + return m.executeHome(ctx, normalized, req, opts, false) + } _, maxRetryCredentials, maxWait := m.retrySettings() @@ -2372,6 +2622,9 @@ if len(normalized) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } + if m.HomeEnabled() { + return m.executeHome(ctx, normalized, req, opts, true) + } _, maxRetryCredentials, maxWait := m.retrySettings() @@ -2400,6 +2653,11 @@ // ExecuteStream performs a streaming execution using the configured selector and executor. // It supports multiple providers for the same model and round-robins the starting provider per model. func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if m.HomeEnabled() { + if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { + defer unlockSession() + } + } normalized := m.normalizeProviders(providers) if len(normalized) == 0 { return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} @@ -2438,6 +2696,129 @@ return nil, lastErr } return nil, &Error{Code: "auth_not_found", Message: "no auth available"} +} + +func (m *Manager) executeHome(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool) (cliproxyexecutor.Response, error) { + if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { + defer unlockSession() + } + routeModel := authSelectionModelFromOptions(opts, req.Model) + responseAlias := requestedModelAliasFromOptions(opts, routeModel) + executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) + opts = ensureRequestedModelMetadata(opts, routeModel) + tried := make(map[string]struct{}) + var lastErr error + for homeAuthCount := 1; ; homeAuthCount++ { + selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, withHomeAuthCount(opts, homeAuthCount)) + if errSelection != nil { + if lastErr != nil && isHomeRequestRetryExceededError(errSelection) { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, errSelection + } + auth := selection.CloneAuthForRoute(routeModel) + if auth == nil || selection.Executor == nil { + selection.End("missing_execution_target") + return cliproxyexecutor.Response{}, &Error{Code: "executor_not_found", Message: "executor not registered"} + } + if _, seen := tried[auth.ID]; seen { + selection.End("repeated_auth") + if lastErr != nil { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, repeatedHomeAuthError() + } + if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil { + selection.End("runtime_auth_bind_failed") + return cliproxyexecutor.Response{}, errRuntimeAuth + } + publishSelectedAuthMetadata(opts.Metadata, auth) + tried[auth.ID] = struct{}{} + execCtx, releaseAttempt, errBind := homeExecutionAttemptContext(ctx, selection) + if errBind != nil { + selection.End("attempt_bind_failed") + return cliproxyexecutor.Response{}, errBind + } + if rt := m.roundTripperFor(auth); rt != nil { + execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) + execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) + } + models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) + if aliasResult.ForceMapping && responseAlias != "" { + aliasResult.OriginalAlias = responseAlias + } + if len(models) > 1 { + models = models[:1] + pooled = false + } + if len(models) == 0 { + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil { + return cliproxyexecutor.Response{}, errEnd + } + lastErr = &Error{Code: "auth_not_found", Message: "no execution models available"} + continue + } + preparedAuth, errPrepare := m.prepareHomeRequestAuth(execCtx, selection.Executor, selection) + if errPrepare != nil { + m.reportHomeResult(execCtx, Result{AuthID: auth.ID, Provider: selection.Provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}, auth) + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { + return cliproxyexecutor.Response{}, errEnd + } + lastErr = errPrepare + continue + } + for _, upstreamModel := range models { + resultModel := m.stateModelForExecution(preparedAuth, routeModel, upstreamModel, pooled) + execReq := req + execReq.Model = upstreamModel + if restoreExecutionModel { + execReq.Model = executionModel + } + execOpts := opts + execOpts.ExecutionLifecycle = selection + execReq, execOpts = applyRequestAfterAuthInterceptor(execCtx, selection.Executor, selection.Provider, execReq, execOpts, requestedModelAliasFromOptions(execOpts, routeModel)) + if errCtx := execCtx.Err(); errCtx != nil { + releaseAttempt() + selection.End("attempt_canceled") + return cliproxyexecutor.Response{}, errCtx + } + var response cliproxyexecutor.Response + var errExecute error + if countTokens { + response, errExecute = selection.Executor.CountTokens(execCtx, preparedAuth, execReq, execOpts) + } else { + response, errExecute = selection.Executor.Execute(execCtx, preparedAuth, execReq, execOpts) + } + result := Result{AuthID: preparedAuth.ID, Provider: selection.Provider, Model: resultModel, Success: errExecute == nil} + if errExecute == nil { + m.reportHomeResult(execCtx, result, preparedAuth) + releaseAttempt() + rewriteForceMappedResponse(&response, aliasResult) + if !m.retainHomeWebsocketSelection(ctx, opts, routeModel, selection) { + selection.End("completed") + } + return response, nil + } + result.Error = resultErrorFromError(errExecute) + result.RetryAfter = retryAfterFromError(errExecute) + m.reportHomeResult(execCtx, result, preparedAuth) + lastErr = errExecute + if isRequestInvalidError(errExecute) { + releaseAttempt() + selection.End("request_invalid") + return cliproxyexecutor.Response{}, errExecute + } + } + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "execution_failed"); errEnd != nil { + return cliproxyexecutor.Response{}, errEnd + } + if errCtx := execCtx.Err(); errCtx != nil && ctx != nil && ctx.Err() != nil { + return cliproxyexecutor.Response{}, errCtx + } + } } type requestToFormatResolver interface { @@ -2770,6 +3151,7 @@ return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } routeModel := authSelectionModelFromOptions(opts, req.Model) + responseAlias := requestedModelAliasFromOptions(opts, routeModel) executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) opts = ensureRequestedModelMetadata(opts, routeModel) homeMode := m.HomeEnabled() @@ -2788,35 +3170,94 @@ if homeMode { pickOpts = withHomeAuthCount(opts, homeAuthCount) } - auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + + var selection *HomeDispatchSelection + var auth *Auth + var executor ProviderExecutor + var provider string + var errPick error + if homeMode { + selection, errPick = m.pickHomeDispatchSelection(ctx, routeModel, pickOpts) + if selection != nil { + auth = selection.CloneAuthForRoute(routeModel) + executor = selection.Executor + provider = selection.Provider + } + } else { + auth, executor, provider, errPick = m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) + } if errPick != nil { if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { return nil, lastErr } return nil, errPick } + if auth == nil || executor == nil { + if selection != nil { + selection.End("missing_execution_target") + } + return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} + } entry := logEntryWithRequestID(ctx) debugLogAuthSelection(entry, auth, provider, routeModel) + if selection != nil { + if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil { + selection.End("runtime_auth_bind_failed") + return nil, errRuntimeAuth + } + } publishSelectedAuthMetadata(opts.Metadata, auth) tried[auth.ID] = struct{}{} execCtx := ctx + releaseAttempt := func() {} + if selection != nil { + var errBind error + execCtx, releaseAttempt, errBind = homeExecutionAttemptContext(ctx, selection) + if errBind != nil { + selection.End("attempt_bind_failed") + return nil, errBind + } + } if rt := m.roundTripperFor(auth); rt != nil { execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } models, pooled, aliasResult := m.preparedExecutionModelsWithAlias(auth, routeModel) + if selection != nil && aliasResult.ForceMapping && responseAlias != "" { + aliasResult.OriginalAlias = responseAlias + } if len(models) == 0 { + if selection != nil { + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil { + return nil, errEnd + } + } continue } attempted[auth.ID] = struct{}{} var errPrepare error - auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) + if selection != nil { + auth, errPrepare = m.prepareHomeRequestAuth(execCtx, executor, selection) + } else { + auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) + } if errPrepare != nil { result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} - m.MarkResult(execCtx, result) + if selection != nil { + m.reportHomeResult(execCtx, result, auth) + releaseAttempt() + } else { + m.MarkResult(execCtx, result) + } lastErr = errPrepare + if selection != nil { + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { + return nil, errEnd + } + } continue } execReq := sanitizeDownstreamWebsocketFallbackRequest(execCtx, auth, req) @@ -2824,9 +3265,23 @@ if restoreExecutionModel { streamExecutionModel = executionModel } - streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, opts, routeModel, streamExecutionModel, models, pooled, aliasResult) + execOpts := opts + if selection != nil { + execOpts.ExecutionLifecycle = selection + } + if homeMode && len(models) > 1 { + models = models[:1] + pooled = false + } + streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, !homeMode, selection != nil) if errStream != nil { - if errCtx := execCtx.Err(); errCtx != nil { + if selection != nil { + releaseAttempt() + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "stream_start_failed"); errEnd != nil { + return nil, errEnd + } + } + if errCtx := execCtx.Err(); errCtx != nil && ctx != nil && ctx.Err() != nil { return nil, errCtx } if isRequestInvalidError(errStream) { @@ -2838,8 +3293,63 @@ } continue } + if selection != nil { + if m.retainHomeWebsocketSelection(ctx, opts, routeModel, selection) { + return wrapHomeStream(ctx, streamResult, nil, releaseAttempt), nil + } + return wrapHomeStream(ctx, streamResult, selection, releaseAttempt), nil + } return streamResult, nil } +} + +func homeExecutionAttemptContext(ctx context.Context, selection *HomeDispatchSelection) (context.Context, func(), error) { + if selection == nil { + return nil, func() {}, fmt.Errorf("Home dispatch selection is nil") + } + return selection.AttemptContext(ctx) +} + +func wrapHomeStream(ctx context.Context, result *cliproxyexecutor.StreamResult, selection *HomeDispatchSelection, releaseAttempt func()) *cliproxyexecutor.StreamResult { + if result == nil || result.Chunks == nil { + if releaseAttempt != nil { + releaseAttempt() + } + return result + } + out := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(out) + if releaseAttempt != nil { + defer releaseAttempt() + } + if selection != nil { + defer selection.End("stream_closed") + } + forward := true + for { + select { + case <-ctx.Done(): + return + case chunk, ok := <-result.Chunks: + if !ok { + return + } + if !forward { + continue + } + select { + case <-ctx.Done(): + return + case out <- chunk: + } + if chunk.Err != nil && selection != nil { + forward = false + } + } + } + }() + return &cliproxyexecutor.StreamResult{Headers: result.Headers, Chunks: out} } func sanitizeDownstreamWebsocketFallbackRequest(ctx context.Context, auth *Auth, req cliproxyexecutor.Request) cliproxyexecutor.Request { @@ -2963,6 +3473,49 @@ type requestAuthPrepareLock struct { mu sync.Mutex +} + +// prepareHomeRequestAuth prepares a dispatch auth without reading or updating local auth state. +func (m *Manager) prepareHomeRequestAuth(ctx context.Context, executor ProviderExecutor, selection *HomeDispatchSelection) (*Auth, error) { + if m == nil || executor == nil || selection == nil { + return nil, nil + } + auth := selection.CloneAuth() + if auth == nil { + return nil, nil + } + preparer, ok := executor.(RequestAuthPreparer) + if !ok || preparer == nil || !preparer.ShouldPrepareRequestAuth(auth) { + return auth, nil + } + + prepare := func() (*Auth, error) { + target := auth.Clone() + if !preparer.ShouldPrepareRequestAuth(target) { + return target, nil + } + updated, errPrepare := preparer.PrepareRequestAuth(ctx, target) + if errPrepare != nil { + return auth, errPrepare + } + if updated == nil { + return target, nil + } + return updated, nil + } + + id := strings.TrimSpace(auth.ID) + if id == "" { + return prepare() + } + lockValue, _ := m.requestPrepareLocks.LoadOrStore(id, &requestAuthPrepareLock{}) + lock, ok := lockValue.(*requestAuthPrepareLock) + if !ok || lock == nil { + return prepare() + } + lock.mu.Lock() + defer lock.mu.Unlock() + return prepare() } func (m *Manager) prepareRequestAuth(ctx context.Context, executor ProviderExecutor, auth *Auth) (*Auth, error) { @@ -3639,6 +4192,10 @@ if err == nil { return 0, false } + var homeBusy *HomeConcurrencyBusyError + if errors.As(err, &homeBusy) && homeBusy != nil { + return 0, false + } if maxWait <= 0 { return 0, false } @@ -3897,6 +4454,27 @@ m.hook.OnResult(ctx, result) m.publishErrorEvent(result, authSnapshot) +} + +func (m *Manager) recordExecutionResult(ctx context.Context, result Result, auth *Auth, ephemeral bool) { + if !ephemeral { + m.MarkResult(ctx, result) + return + } + m.reportHomeResult(ctx, result, auth) +} + +// reportHomeResult only observes a Home dispatch result and never updates local auth state. +func (m *Manager) reportHomeResult(ctx context.Context, result Result, auth *Auth) { + if m == nil || result.AuthID == "" { + return + } + var snapshot *Auth + if auth != nil { + snapshot = auth.Clone() + } + m.hook.OnResult(ctx, result) + m.publishErrorEvent(result, snapshot) } func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Result) { @@ -4185,8 +4763,8 @@ type retryAfterProvider interface { RetryAfter() *time.Duration } - rap, ok := err.(retryAfterProvider) - if !ok || rap == nil { + var rap retryAfterProvider + if !errors.As(err, &rap) || rap == nil { return nil } retryAfter := rap.RetryAfter() @@ -4770,10 +5348,15 @@ } m.mu.Lock() + var selections []*HomeDispatchSelection if sessionID == CloseAllExecutionSessionsID { m.clearHomeRuntimeAuthsLocked() + selections = m.takeAllHomeSessionSelectionsLocked() + m.clearHomeSessionLocks() } else { m.clearHomeRuntimeAuthsForSessionLocked(sessionID) + selections = m.takeHomeSessionSelectionsLocked(sessionID) + m.homeSessionLocks.Delete(sessionID) } executors := make([]ProviderExecutor, 0, len(m.executors)) for _, exec := range m.executors { @@ -4781,6 +5364,9 @@ } m.mu.Unlock() + for _, selection := range selections { + selection.End("session_closed") + } for i := range executors { if closer, ok := executors[i].(ExecutionSessionCloser); ok && closer != nil { closer.CloseExecutionSession(sessionID) @@ -4902,9 +5488,15 @@ // SelectAuth selects one credential through the configured scheduling strategy. // It does not execute or alter the selected credential's result state. func (m *Manager) SelectAuth(ctx context.Context, provider, model string, opts cliproxyexecutor.Options) (*Auth, error) { - selected, _, errPick := m.pickNext(ctx, provider, model, opts, nil) + if m != nil && m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } + selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, nil) if errPick != nil { return nil, errPick + } + if m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} } return selected, nil } @@ -4912,20 +5504,17 @@ // SelectAuthByKind selects one credential of the required kind through the // configured scheduling strategy. Credentials of other kinds are skipped. func (m *Manager) SelectAuthByKind(ctx context.Context, provider, model, requiredKind string, opts cliproxyexecutor.Options) (*Auth, error) { + if m != nil && m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } requiredKind = normalizeAuthKind(requiredKind) if requiredKind == "" { return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} } - homeMode := m.HomeEnabled() - homeAuthCount := homeAuthCountFromMetadata(opts.Metadata) tried := make(map[string]struct{}) for { - pickOpts := opts - if homeMode { - pickOpts = withHomeAuthCount(opts, homeAuthCount) - } - selected, _, errPick := m.pickNext(ctx, provider, model, pickOpts, tried) + selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, tried) if errPick != nil { return nil, errPick } @@ -4933,6 +5522,9 @@ return nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} } if selected.AuthKind() == requiredKind { + if m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} + } return selected, nil } authID := strings.TrimSpace(selected.ID) @@ -4943,9 +5535,52 @@ return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} } tried[authID] = struct{}{} - if homeMode { - homeAuthCount++ + } +} + +// SelectHomeAuthByKind selects a Home dispatch while retaining its execution scope. +func (m *Manager) SelectHomeAuthByKind(ctx context.Context, provider string, model string, requiredKind string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { + requiredKind = normalizeAuthKind(requiredKind) + if requiredKind == "" { + return nil, &Error{Code: "invalid_auth_kind", Message: "required auth kind is invalid", HTTPStatus: http.StatusBadRequest} + } + if m == nil || !m.HomeEnabled() { + return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + + homeAuthCount := homeAuthCountFromMetadata(opts.Metadata) + tried := make(map[string]struct{}) + for { + selectionOpts := withHomeAuthCount(opts, homeAuthCount) + selection, errSelection := m.pickHomeDispatchSelection(ctx, model, selectionOpts) + if errSelection != nil { + return nil, errSelection } + providerMatches := strings.TrimSpace(provider) == "" || strings.EqualFold(strings.TrimSpace(selection.Provider), strings.TrimSpace(provider)) + kindMatches := selection.Auth != nil && selection.Auth.AuthKind() == requiredKind + if providerMatches && kindMatches { + return selection, nil + } + + authID := "" + if selection.Auth != nil { + authID = strings.TrimSpace(selection.Auth.ID) + } + reason := "auth_kind_mismatch" + if !providerMatches { + reason = "provider_mismatch" + } + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, reason); errEnd != nil { + return nil, errEnd + } + if authID == "" { + return nil, &Error{Code: "auth_not_found", Message: "selected auth has no ID"} + } + if _, alreadyTried := tried[authID]; alreadyTried { + return nil, &Error{Code: "auth_not_found", Message: "selector repeatedly returned an ineligible auth"} + } + tried[authID] = struct{}{} + homeAuthCount++ } } @@ -5208,9 +5843,11 @@ } type homeErrorDetail struct { - Type string `json:"type"` - Message string `json:"message"` - Code string `json:"code,omitempty"` + Type string `json:"type"` + Message string `json:"message"` + Code string `json:"code,omitempty"` + Retryable bool `json:"retryable,omitempty"` + RetryAfterMS int64 `json:"retry_after_ms,omitempty"` } const ( @@ -5268,6 +5905,7 @@ type homeAuthDispatcher interface { HeartbeatOK() bool RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error) + AbortAmbiguousDispatch() } var currentHomeDispatcher = func() homeAuthDispatcher { @@ -5377,13 +6015,231 @@ } } +type homeSessionSelectionKey struct { + credentialID string + routeModel string +} + +func (m *Manager) lockHomeWebsocketSession(ctx context.Context, opts cliproxyexecutor.Options) func() { + if m == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { + return nil + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + if sessionID == "" { + return nil + } + lock, _ := m.homeSessionLocks.LoadOrStore(sessionID, &sync.Mutex{}) + mutex, ok := lock.(*sync.Mutex) + if !ok || mutex == nil { + return nil + } + mutex.Lock() + return mutex.Unlock +} + +func (m *Manager) retainedHomeSessionSelection(ctx context.Context, opts cliproxyexecutor.Options, model string) (*HomeDispatchSelection, bool, error) { + if m == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { + return nil, false, nil + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + credentialID := pinnedAuthIDFromMetadata(opts.Metadata) + if sessionID == "" { + return nil, false, nil + } + + routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) + var retained *HomeDispatchSelection + var ended []*HomeDispatchSelection + fallbackAttempt := homeAuthCountFromMetadata(opts.Metadata) > 1 + m.mu.Lock() + selections := m.homeSessionSelections[sessionID] + for key, selection := range selections { + if selection == nil { + delete(selections, key) + continue + } + matchesCredential := credentialID == "" || key.credentialID == credentialID + matchesRoute := validRouteModel && key.routeModel == routeModel + if !fallbackAttempt && matchesCredential && selection.Active() && matchesRoute && retained == nil { + retained = selection + continue + } + delete(selections, key) + ended = append(ended, selection) + } + if len(selections) == 0 { + delete(m.homeSessionSelections, sessionID) + } + m.mu.Unlock() + + for _, selection := range ended { + if errWait := m.endHomeSelectionBeforeRedispatch(ctx, selection, "target_changed"); errWait != nil { + return nil, false, errWait + } + } + return retained, retained != nil, nil +} + +func (m *Manager) predictedHomeConcurrencyModel(auth *Auth, routeModel string) (string, bool) { + requestedModel := rewriteModelForAuth(routeModel, auth) + aliasResult := m.resolveExecutionAliasResultForRequested(auth, requestedModel) + upstreamModel := executionAliasPoolModel(auth, requestedModel, aliasResult) + if pool := m.resolveOpenAICompatUpstreamModelPool(auth, upstreamModel); len(pool) != 0 { + if len(pool) != 1 { + return "", false + } + upstreamModel = pool[0] + } else { + upstreamModel = m.applyAPIKeyModelAlias(auth, upstreamModel) + } + return validCanonicalHomeConcurrencyModelKey(upstreamModel) +} + +func (m *Manager) endMismatchedHomeSessionSelections(ctx context.Context, sessionID, credentialID, model string, waitForAck bool) error { + if m == nil || sessionID == "" { + return nil + } + routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) + var ended []*HomeDispatchSelection + m.mu.Lock() + selections := m.homeSessionSelections[sessionID] + for key, selection := range selections { + if selection == nil { + delete(selections, key) + continue + } + matchesRoute := validRouteModel && key.routeModel == routeModel + if key.credentialID == credentialID && matchesRoute { + continue + } + delete(selections, key) + ended = append(ended, selection) + } + if len(selections) == 0 { + delete(m.homeSessionSelections, sessionID) + } + m.mu.Unlock() + for _, selection := range ended { + if !waitForAck { + selection.End("target_changed") + continue + } + if errWait := m.endHomeSelectionBeforeRedispatch(ctx, selection, "target_changed"); errWait != nil { + return errWait + } + } + return nil +} + +func (m *Manager) endHomeSelectionBeforeRedispatch(ctx context.Context, selection *HomeDispatchSelection, reason string) error { + if selection == nil { + return nil + } + ticket := selection.EndWithRelease(reason) + if ticket == nil { + return nil + } + + bound := internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound + if m != nil { + if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil { + bound = cfg.CredentialConcurrency.WithDefaults().CPACancelBound + } + } + waitCtx := ctx + if waitCtx == nil { + waitCtx = context.Background() + } + waitCtx, cancelWait := context.WithTimeout(waitCtx, bound) + defer cancelWait() + if errWait := ticket.Wait(waitCtx); errWait != nil { + return &Error{Code: "home_unavailable", Message: "Home did not acknowledge credential release: " + errWait.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} + } + return nil +} + +func (m *Manager) retainHomeWebsocketSelection(ctx context.Context, opts cliproxyexecutor.Options, model string, selection *HomeDispatchSelection) bool { + if m == nil || selection == nil || !selection.Retained() || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil { + return false + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + credentialID := strings.TrimSpace(selection.Auth.ID) + routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) + if selection.accountedModel == "" { + selection.accountedModel, _ = m.predictedHomeConcurrencyModel(selection.Auth, model) + } + if sessionID == "" || credentialID == "" || !validRouteModel || selection.accountedModel == "" { + return false + } + _ = m.endMismatchedHomeSessionSelections(ctx, sessionID, credentialID, routeModel, false) + key := homeSessionSelectionKey{credentialID: credentialID, routeModel: routeModel} + m.mu.Lock() + if m.homeSessionSelections == nil { + m.homeSessionSelections = make(map[string]map[homeSessionSelectionKey]*HomeDispatchSelection) + } + selections := m.homeSessionSelections[sessionID] + if selections == nil { + selections = make(map[homeSessionSelectionKey]*HomeDispatchSelection) + m.homeSessionSelections[sessionID] = selections + } + previous := selections[key] + selections[key] = selection + m.mu.Unlock() + m.rememberHomeRuntimeAuth(sessionID, selection.Auth) + if previous != nil && previous != selection { + previous.End("target_replaced") + } + return true +} + +func (m *Manager) clearHomeSessionLocks() { + if m == nil { + return + } + m.homeSessionLocks.Range(func(key, _ any) bool { + m.homeSessionLocks.Delete(key) + return true + }) +} + +func (m *Manager) takeHomeSessionSelectionsLocked(sessionID string) []*HomeDispatchSelection { + if m == nil { + return nil + } + selections := m.homeSessionSelections[sessionID] + delete(m.homeSessionSelections, sessionID) + result := make([]*HomeDispatchSelection, 0, len(selections)) + for _, selection := range selections { + result = append(result, selection) + } + return result +} + +func (m *Manager) takeAllHomeSessionSelectionsLocked() []*HomeDispatchSelection { + if m == nil { + return nil + } + result := make([]*HomeDispatchSelection, 0) + for sessionID, selections := range m.homeSessionSelections { + delete(m.homeSessionSelections, sessionID) + for _, selection := range selections { + result = append(result, selection) + } + } + return result +} + func (m *Manager) clearHomeRuntimeAuths() { if m == nil { return } m.mu.Lock() m.clearHomeRuntimeAuthsLocked() + selections := m.takeAllHomeSessionSelectionsLocked() m.mu.Unlock() + for _, selection := range selections { + selection.End("home_disabled") + } } func (m *Manager) clearHomeRuntimeAuthsLocked() { @@ -5391,6 +6247,7 @@ return } m.homeRuntimeAuths = make(map[string]map[string]*Auth) + m.homeRuntimeAuthOwners = make(map[string]map[string]*HomeDispatchSelection) } func (m *Manager) clearHomeRuntimeAuthsForSessionLocked(sessionID string) { @@ -5399,6 +6256,79 @@ return } delete(m.homeRuntimeAuths, sessionID) + delete(m.homeRuntimeAuthOwners, sessionID) +} + +func (m *Manager) bindHomeSelectionRuntimeAuth(ctx context.Context, opts cliproxyexecutor.Options, selection *HomeDispatchSelection) error { + if m == nil || selection == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) || selection.Auth == nil || !authWebsocketsEnabled(selection.Auth) { + return nil + } + sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) + authID := strings.TrimSpace(selection.Auth.ID) + if sessionID == "" || authID == "" || !selection.runtimeAuthBound.CompareAndSwap(false, true) { + return nil + } + m.rememberHomeSelectionRuntimeAuth(sessionID, selection) + if errBind := selection.Bind(func() error { + m.forgetHomeRuntimeAuth(sessionID, authID, selection) + return nil + }); errBind != nil { + selection.runtimeAuthBound.Store(false) + m.forgetHomeRuntimeAuth(sessionID, authID, selection) + return errBind + } + return nil +} + +func (m *Manager) rememberHomeSelectionRuntimeAuth(sessionID string, selection *HomeDispatchSelection) { + if m == nil || selection == nil || selection.Auth == nil { + return + } + sessionID = strings.TrimSpace(sessionID) + authID := strings.TrimSpace(selection.Auth.ID) + if sessionID == "" || authID == "" { + return + } + m.mu.Lock() + if m.homeRuntimeAuths == nil { + m.homeRuntimeAuths = make(map[string]map[string]*Auth) + } + if m.homeRuntimeAuthOwners == nil { + m.homeRuntimeAuthOwners = make(map[string]map[string]*HomeDispatchSelection) + } + if m.homeRuntimeAuths[sessionID] == nil { + m.homeRuntimeAuths[sessionID] = make(map[string]*Auth) + } + if m.homeRuntimeAuthOwners[sessionID] == nil { + m.homeRuntimeAuthOwners[sessionID] = make(map[string]*HomeDispatchSelection) + } + m.homeRuntimeAuths[sessionID][authID] = selection.Auth.Clone() + m.homeRuntimeAuthOwners[sessionID][authID] = selection + m.mu.Unlock() +} + +func (m *Manager) forgetHomeRuntimeAuth(sessionID string, authID string, owner *HomeDispatchSelection) { + sessionID = strings.TrimSpace(sessionID) + authID = strings.TrimSpace(authID) + if m == nil || sessionID == "" || authID == "" { + return + } + m.mu.Lock() + owners := m.homeRuntimeAuthOwners[sessionID] + if owner != nil && owners[authID] != owner { + m.mu.Unlock() + return + } + sessionAuths := m.homeRuntimeAuths[sessionID] + delete(sessionAuths, authID) + delete(owners, authID) + if len(sessionAuths) == 0 { + delete(m.homeRuntimeAuths, sessionID) + } + if len(owners) == 0 { + delete(m.homeRuntimeAuthOwners, sessionID) + } + m.mu.Unlock() } func (m *Manager) rememberHomeRuntimeAuth(sessionID string, auth *Auth) { @@ -5436,21 +6366,19 @@ if auth == nil || !authWebsocketsEnabled(auth) { return nil, nil, "", false } - providerKey := executorKeyFromAuth(auth) - if providerKey == "" { + logicalProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) + executorKey := executorKeyFromAuth(auth) + if logicalProvider == "" || executorKey == "" { return nil, nil, "", false } - executor, ok := m.Executor(providerKey) + executor, ok := m.Executor(executorKey) if !ok && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" { executor, ok = m.Executor("openai-compatibility") - if ok { - providerKey = "openai-compatibility" - } } if !ok { return nil, nil, "", false } - return auth.Clone(), executor, providerKey, true + return auth.Clone(), executor, logicalProvider, true } func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { @@ -5460,68 +6388,154 @@ if ctx == nil { ctx = context.Background() } - executionSessionID := homeExecutionSessionIDFromMetadata(opts.Metadata) - count := homeAuthCountFromMetadata(opts.Metadata) - if cliproxyexecutor.DownstreamWebsocket(ctx) && executionSessionID != "" && count <= 1 { - if pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata); pinnedAuthID != "" { - _, alreadyTried := tried[pinnedAuthID] - if !alreadyTried { - if auth, executor, providerKey, ok := m.homeRuntimeAuthByID(executionSessionID, pinnedAuthID); ok { - return auth, executor, providerKey, nil - } + selection, errSelection := m.pickHomeDispatchSelection(ctx, model, opts) + if errSelection != nil { + return nil, nil, "", errSelection + } + if selection.Auth == nil || homeAuthAlreadyTried(tried, selection.Auth.ID) { + selection.End("repeated_auth") + return nil, nil, "", repeatedHomeAuthError() + } + auth := selection.CloneAuthForRoute(model) + executor := selection.Executor + provider := selection.Provider + selection.End("legacy_selection_unbound") + return auth, executor, provider, nil +} + +func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, opts cliproxyexecutor.Options) (*HomeDispatchSelection, error) { + if m == nil { + return nil, &Error{Code: "auth_not_found", Message: "no auth available"} + } + if ctx == nil { + ctx = context.Background() + } + + requestedModel := strings.TrimSpace(model) + if requestedModel == "" { + requestedModel = requestedModelFromMetadata(opts.Metadata, model) + } + retained, retainedOK, errRetained := m.retainedHomeSessionSelection(ctx, opts, requestedModel) + if errRetained != nil { + return nil, errRetained + } + if retainedOK { + return retained, nil + } + if sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata); sessionID != "" { + if credentialID := pinnedAuthIDFromMetadata(opts.Metadata); credentialID != "" { + if errEnd := m.endMismatchedHomeSessionSelections(ctx, sessionID, credentialID, requestedModel, true); errEnd != nil { + return nil, errEnd } } } - client := currentHomeDispatcher() - if client == nil || !client.HeartbeatOK() { - return nil, nil, "", &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} + bundle := m.HomeDispatchBundle() + if bundle == nil || bundle.client == nil || bundle.registry == nil { + return nil, &Error{Code: "home_unavailable", Message: "home dispatch bundle unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + client := bundle.client + registry := bundle.registry + if !client.HeartbeatOK() { + return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} + } + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} } - requestedModel := requestedModelFromMetadata(opts.Metadata, model) sessionID := ExtractSessionID(opts.Headers, opts.OriginalRequest, opts.Metadata) dispatchHeaders := homeDispatchHeaders(ctx, opts.Headers) - - raw, err := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, count) - if err != nil { - if errors.Is(err, home.ErrAuthNotFound) { - return nil, nil, "", &Error{Code: "auth_not_found", Message: err.Error(), HTTPStatus: http.StatusServiceUnavailable} + raw, errRPop := client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) + if errRPop != nil { + if home.IsAmbiguousDispatchError(errRPop) { + client.AbortAmbiguousDispatch() } - return nil, nil, "", &Error{Code: "home_unavailable", Message: err.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} + pending.End() + if errors.Is(errRPop, home.ErrAuthNotFound) { + return nil, &Error{Code: "auth_not_found", Message: errRPop.Error(), HTTPStatus: http.StatusServiceUnavailable} + } + return nil, &Error{Code: "home_unavailable", Message: errRPop.Error(), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} } - var env homeErrorEnvelope - if errUnmarshal := json.Unmarshal(raw, &env); errUnmarshal == nil && env.Error != nil { - code := strings.TrimSpace(env.Error.Type) - if code == "" { - code = strings.TrimSpace(env.Error.Code) + envelope, errEnvelope := decodeHomeDispatchConcurrencyEnvelope(raw) + if errEnvelope != nil { + if envelope.Present { + client.AbortAmbiguousDispatch() } - msg := strings.TrimSpace(env.Error.Message) - if msg == "" { - msg = "home returned error" + pending.End() + if envelope.Present { + return nil, invalidHomeConcurrencyResponse("Home returned malformed concurrency tuple") } - status := http.StatusBadGateway - switch strings.ToLower(code) { - case "model_not_found": - status = http.StatusNotFound - case "authentication_error", "unauthorized", "no_credentials", "invalid_credential": - status = http.StatusUnauthorized + return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} + } + + kind := "http" + if cliproxyexecutor.DownstreamWebsocket(ctx) { + kind = "websocket" + } else if opts.Stream { + kind = "stream" + } + baseScope := executionregistry.ScopeSpec{ + RequestID: logging.GetRequestID(ctx), + Model: requestedModel, + Kind: kind, + StartedAt: time.Now(), + } + var scope *executionregistry.Scope + if envelope.Present { + var errInstall error + scope, errInstall = installHomeConcurrencyScope(registry, pending, envelope.Tuple, baseScope) + if errInstall != nil { + client.AbortAmbiguousDispatch() + pending.End() + return nil, homeConcurrencyInstallError(errInstall) } - return nil, nil, "", &Error{Code: code, Message: msg, HTTPStatus: status} + } + endScope := func() { + if scope != nil { + scope.End("local_validation_failed") + return + } + pending.End() + } + if errHome := decodeHomeDispatchError(raw); errHome != nil { + if envelope.Present { + client.AbortAmbiguousDispatch() + endScope() + return nil, invalidHomeConcurrencyResponse("Home returned both accounted concurrency and an error") + } + pending.End() + return nil, errHome } var dispatch homeAuthDispatchResponse if errUnmarshal := json.Unmarshal(raw, &dispatch); errUnmarshal != nil { - return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} } - setHomeUserAPIKeyOnGinContext(ctx, dispatch.UserAPIKey) auth := dispatch.Auth if strings.TrimSpace(auth.ID) == "" { - // Backward compatibility: older home instances returned the auth directly. + // Backward compatibility: older Home instances returned the auth directly. if errUnmarshal := json.Unmarshal(raw, &auth); errUnmarshal != nil { - return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned invalid auth payload", HTTPStatus: http.StatusBadGateway} } } + observedModel := canonicalHomeDispatchModel(dispatch.Model, requestedModel) + if envelope.Present { + observedConcurrencyModel, validModel := validCanonicalHomeConcurrencyModelKey(observedModel) + if !validModel || envelope.Tuple.Model != observedConcurrencyModel { + client.AbortAmbiguousDispatch() + endScope() + return nil, invalidHomeConcurrencyResponse("Home concurrency model does not match dispatched model") + } + } + if !envelope.Present { + baseScope.Model = observedModel + } + + setHomeUserAPIKeyOnGinContext(ctx, dispatch.UserAPIKey) if upstreamModel := strings.TrimSpace(dispatch.Model); upstreamModel != "" { if auth.Attributes == nil { auth.Attributes = make(map[string]string, 3) @@ -5536,14 +6550,18 @@ auth.Attributes[homeOriginalAliasAttributeKey] = originalAlias } if strings.TrimSpace(auth.ID) == "" { - return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned auth without id", HTTPStatus: http.StatusBadGateway} + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned auth without id", HTTPStatus: http.StatusBadGateway} } - if homeAuthAlreadyTried(tried, auth.ID) { - return nil, nil, "", repeatedHomeAuthError() + if errIdentity := verifyAccountedHomeConcurrencyIdentity(envelope.Tuple, &auth, dispatch.AuthIndex); errIdentity != nil { + endScope() + return nil, errIdentity } - providerKey := executorKeyFromAuth(&auth) - if providerKey == "" { - return nil, nil, "", &Error{Code: "invalid_auth", Message: "home returned auth without provider", HTTPStatus: http.StatusBadGateway} + logicalProvider := strings.ToLower(strings.TrimSpace(auth.Provider)) + executorKey := executorKeyFromAuth(&auth) + if logicalProvider == "" || executorKey == "" { + endScope() + return nil, &Error{Code: "invalid_auth", Message: "home returned auth without provider", HTTPStatus: http.StatusBadGateway} } homeAuthIndex := strings.TrimSpace(dispatch.AuthIndex) @@ -5554,22 +6572,45 @@ auth.EnsureIndex() } - executor, ok := m.Executor(providerKey) - if !ok && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" { - executor, ok = m.Executor("openai-compatibility") - if ok { - providerKey = "openai-compatibility" + executor, okExecutor := m.Executor(executorKey) + if !okExecutor && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["base_url"]) != "" { + executor, okExecutor = m.Executor("openai-compatibility") + } + if !okExecutor { + endScope() + return nil, &Error{Code: "executor_not_found", Message: "executor not registered", HTTPStatus: http.StatusBadGateway} + } + if scope == nil { + var errInstall error + scope, errInstall = installHomeConcurrencyScope(registry, pending, homeConcurrencyTuple{}, executionregistry.ScopeSpec{ + RequestID: baseScope.RequestID, + CredentialID: strings.TrimSpace(auth.ID), + Model: baseScope.Model, + Kind: baseScope.Kind, + StartedAt: baseScope.StartedAt, + }) + if errInstall != nil { + client.AbortAmbiguousDispatch() + pending.End() + return nil, homeConcurrencyInstallError(errInstall) } } - if !ok { - return nil, nil, "", &Error{Code: "executor_not_found", Message: "executor not registered", HTTPStatus: http.StatusBadGateway} - } - authCopy := auth.Clone() - if cliproxyexecutor.DownstreamWebsocket(ctx) && executionSessionID != "" && authWebsocketsEnabled(authCopy) { - m.rememberHomeRuntimeAuth(executionSessionID, authCopy) + selection, errSelection := newHomeDispatchSelection(auth.Clone(), executor, logicalProvider, scope) + if errSelection != nil { + endScope() + return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} } - return authCopy, executor, providerKey, nil + if envelope.Present { + selection.accountedModel = envelope.Tuple.Model + } + if executionSessionID := homeExecutionSessionIDFromMetadata(opts.Metadata); executionSessionID != "" && cliproxyexecutor.DownstreamWebsocket(ctx) { + if errEnd := m.endMismatchedHomeSessionSelections(ctx, executionSessionID, strings.TrimSpace(auth.ID), requestedModel, true); errEnd != nil { + selection.End("target_change_release_failed") + return nil, errEnd + } + } + return selection, nil } func requestedModelFromMetadata(metadata map[string]any, fallback string) string { @@ -5595,7 +6636,7 @@ } func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, routeModel string, opts cliproxyexecutor.Options) ([]creditsCandidateEntry, error) { - if m == nil { + if m == nil || !m.localExecutionAllowed() { return nil, nil } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) @@ -5674,7 +6715,7 @@ "status": status, "providers": providers, }).Debug("shouldAttemptAntigravityCreditsFallback") - if m == nil || lastErr == nil { + if m == nil || lastErr == nil || m.HomeEnabled() { return false } cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) @@ -5700,6 +6741,12 @@ } func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, bool, error) { + if m != nil && m.HomeEnabled() { + return cliproxyexecutor.Response{}, false, &Error{Code: "home_fallback_unsupported", Message: "Home does not support Antigravity credits fallback", HTTPStatus: http.StatusServiceUnavailable} + } + if !m.localExecutionAllowed() { + return cliproxyexecutor.Response{}, false, nil + } routeModel := req.Model candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) if errCandidates != nil { @@ -5749,6 +6796,12 @@ } func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, bool, error) { + if m != nil && m.HomeEnabled() { + return nil, false, &Error{Code: "home_fallback_unsupported", Message: "Home does not support Antigravity credits fallback", HTTPStatus: http.StatusServiceUnavailable} + } + if !m.localExecutionAllowed() { + return nil, false, nil + } routeModel := req.Model candidates, errCandidates := m.findAllAntigravityCreditsCandidateAuths(ctx, routeModel, opts) if errCandidates != nil { @@ -5774,7 +6827,7 @@ if len(models) == 0 { continue } - result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, "", models, pooled, aliasResult) + result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, "", models, pooled, aliasResult, true, false) if errStream != nil { continue } diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -67,6 +67,21 @@ } } +func TestManager_ShouldRetryAfterError_SkipsWrappedHomeConcurrencyBusy(t *testing.T) { + m := NewManager(nil, nil, nil) + m.SetRetryConfig(1, 30*time.Second, 0) + if _, errRegister := m.Register(context.Background(), &Auth{ID: "retry-auth", Provider: "codex"}); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + _, _, maxWait := m.retrySettings() + errBusy := fmt.Errorf("outer retry: %w", NewHomeConcurrencyBusyError("busy", 20*time.Second)) + wait, shouldRetry := m.shouldRetryAfterError(errBusy, 0, []string{"codex"}, "gpt", maxWait) + if shouldRetry || wait != 0 { + t.Fatalf("wrapped Home busy retry = (%v, %t), want (0, false)", wait, shouldRetry) + } +} + func TestManager_ShouldRetryAfterError_UsesOAuthModelAliasForCooldown(t *testing.T) { m := NewManager(nil, nil, nil) m.SetRetryConfig(3, 30*time.Second, 0) diff --git a/sdk/cliproxy/auth/cooldown_state_test.go b/sdk/cliproxy/auth/cooldown_state_test.go --- a/sdk/cliproxy/auth/cooldown_state_test.go +++ b/sdk/cliproxy/auth/cooldown_state_test.go @@ -9,6 +9,8 @@ "sync/atomic" "testing" "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" ) type recordingCooldownStateStore struct { @@ -253,6 +255,196 @@ } } +func TestManagerSetConfigSnapshotDefersCooldownPersistence(t *testing.T) { + store := &recordingCooldownStateStore{} + manager := NewManager(nil, nil, nil) + manager.SetCooldownStateStore(store) + auth := &Auth{ID: "auth-1", Provider: "xai", Status: StatusActive} + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil { + t.Fatalf("Register() returned error: %v", errRegister) + } + manager.MarkResult(context.Background(), Result{ + AuthID: auth.ID, + Provider: auth.Provider, + Model: "grok-4", + Success: false, + Error: &Error{Message: "rate limited", HTTPStatus: 429}, + }) + store.saveCount.Store(0) + + if changed := manager.SetConfigSnapshot(&internalconfig.Config{DisableCooling: true}); !changed { + t.Fatal("SetConfigSnapshot() = false, want cleared cooldown state") + } + if got := store.saveCount.Load(); got != 0 { + t.Fatalf("SetConfigSnapshot() persisted cooldown state %d times, want 0", got) + } + manager.PersistCooldownStates(context.Background()) + if got := store.saveCount.Load(); got != 1 { + t.Fatalf("PersistCooldownStates() saved cooldown state %d times, want 1", got) + } +} + +type blockingCooldownStateStore struct { + started chan struct{} + release chan struct{} +} + +func (s *blockingCooldownStateStore) Load(context.Context) ([]CooldownStateRecord, error) { + return nil, nil +} + +func (s *blockingCooldownStateStore) Save(ctx context.Context, _ []CooldownStateRecord) error { + select { + case <-s.started: + default: + close(s.started) + } + select { + case <-s.release: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func TestManagerSwapCooldownStateStorePersistsOldStoreBeforeSwap(t *testing.T) { + oldStore := &recordingCooldownStateStore{} + newStore := &recordingCooldownStateStore{} + manager := NewManager(nil, nil, nil) + manager.SetCooldownStateStore(oldStore) + auth := &Auth{ID: "auth-1", Provider: "xai", Status: StatusActive} + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil { + t.Fatalf("Register() returned error: %v", errRegister) + } + manager.MarkResult(context.Background(), Result{ + AuthID: auth.ID, Provider: auth.Provider, Model: "grok-4", Success: false, + Error: &Error{Message: "rate limited", HTTPStatus: 429}, + }) + oldStore.saveCount.Store(0) + if changed := manager.SetConfigSnapshot(&internalconfig.Config{DisableCooling: true}); !changed { + t.Fatal("SetConfigSnapshot() = false, want cleared cooldown state") + } + + if swapped := manager.SwapCooldownStateStore(context.Background(), newStore, true); !swapped { + t.Fatal("SwapCooldownStateStore() = false, want true") + } + if got := oldStore.saveCount.Load(); got != 1 { + t.Fatalf("old store save count = %d, want 1", got) + } + if len(oldStore.records) != 0 { + t.Fatalf("old store records = %+v, want cleared cooldown state", oldStore.records) + } + manager.mu.RLock() + currentStore := manager.cooldownStore + manager.mu.RUnlock() + if currentStore != newStore { + t.Fatal("cooldown store swapped before the old store was persisted") + } +} + +func TestManagerApplyConfigWithCooldownStoreSerializesTransitions(t *testing.T) { + oldStore := &blockingCooldownStateStore{started: make(chan struct{}), release: make(chan struct{})} + firstStore := &recordingCooldownStateStore{} + secondStore := &recordingCooldownStateStore{} + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-1", Provider: "xai", Status: StatusActive} + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil { + t.Fatalf("Register() returned error: %v", errRegister) + } + manager.MarkResult(context.Background(), Result{ + AuthID: auth.ID, Provider: auth.Provider, Model: "grok-4", Success: false, + Error: &Error{Message: "rate limited", HTTPStatus: 429}, + }) + manager.SetCooldownStateStore(oldStore) + + firstDone := make(chan bool, 1) + go func() { + firstDone <- manager.ApplyConfigWithCooldownStateStore(context.Background(), &internalconfig.Config{DisableCooling: true}, firstStore) + }() + select { + case <-oldStore.started: + case <-time.After(time.Second): + t.Fatal("first old-store persistence did not start") + } + + secondDone := make(chan bool, 1) + go func() { + secondDone <- manager.ApplyConfigWithCooldownStateStore(context.Background(), &internalconfig.Config{}, secondStore) + }() + select { + case <-secondDone: + t.Fatal("concurrent config transition completed while old-store persistence was blocked") + case <-time.After(100 * time.Millisecond): + } + + close(oldStore.release) + if applied := waitForCooldownTransition(t, firstDone, "first config transition"); !applied { + t.Fatal("first config transition returned false") + } + if applied := waitForCooldownTransition(t, secondDone, "second config transition"); !applied { + t.Fatal("second config transition returned false") + } + manager.mu.RLock() + currentStore := manager.cooldownStore + manager.mu.RUnlock() + if currentStore != secondStore { + t.Fatal("concurrent config transitions did not leave the final resolved store installed") + } +} + +func waitForCooldownTransition(t *testing.T, done <-chan bool, name string) bool { + t.Helper() + select { + case applied := <-done: + return applied + case <-time.After(time.Second): + t.Fatalf("timed out waiting for %s", name) + return false + } +} + +func TestManagerSwapCooldownStateStoreKeepsOldStoreWhenCanceled(t *testing.T) { + oldStore := &blockingCooldownStateStore{started: make(chan struct{}), release: make(chan struct{})} + newStore := &recordingCooldownStateStore{} + manager := NewManager(nil, nil, nil) + manager.SetCooldownStateStore(oldStore) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan bool, 1) + go func() { done <- manager.SwapCooldownStateStore(ctx, newStore, true) }() + select { + case <-oldStore.started: + case <-time.After(time.Second): + t.Fatal("old cooldown store persistence did not start") + } + manager.mu.RLock() + currentStore := manager.cooldownStore + manager.mu.RUnlock() + if currentStore != oldStore { + t.Fatal("cooldown store swapped while old store persistence was blocked") + } + cancel() + select { + case swapped := <-done: + if swapped { + t.Fatal("SwapCooldownStateStore() = true after cancellation") + } + case <-time.After(time.Second): + t.Fatal("SwapCooldownStateStore() did not honor cancellation") + } + + close(oldStore.release) + if swapped := manager.SwapCooldownStateStore(context.Background(), newStore, false); !swapped { + t.Fatal("SwapCooldownStateStore() = false, want retry to persist the old store before swapping") + } + manager.mu.RLock() + currentStore = manager.cooldownStore + manager.mu.RUnlock() + if currentStore != newStore { + t.Fatal("cooldown store was not swapped after pending persistence completed") + } +} + func TestManager_RestoreCooldownStates(t *testing.T) { nextRetry := time.Now().Add(time.Hour).UTC().Truncate(time.Second) store := &recordingCooldownStateStore{ @@ -300,5 +492,53 @@ } if got := store.saveCount.Load(); got != 1 { t.Fatalf("restore cleanup saved cooldown state %d times, want 1", got) + } +} + +func TestManagerResultSaveWaitsForCooldownStoreTransition(t *testing.T) { + oldStore := &blockingCooldownStateStore{started: make(chan struct{}), release: make(chan struct{})} + newStore := &recordingCooldownStateStore{} + manager := NewManager(nil, nil, nil) + auth := &Auth{ID: "auth-1", Provider: "xai", Status: StatusActive} + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil { + t.Fatalf("Register() returned error: %v", errRegister) + } + manager.SetCooldownStateStore(oldStore) + + transitionDone := make(chan bool, 1) + go func() { + transitionDone <- manager.SwapCooldownStateStore(context.Background(), newStore, true) + }() + select { + case <-oldStore.started: + case <-time.After(time.Second): + t.Fatal("old-store transition save did not start") + } + + resultDone := make(chan struct{}) + go func() { + manager.MarkResult(context.Background(), Result{ + AuthID: auth.ID, Provider: auth.Provider, Model: "grok-4", Success: false, + Error: &Error{Message: "rate limited", HTTPStatus: 429}, + }) + close(resultDone) + }() + select { + case <-resultDone: + t.Fatal("result save completed while the store transition was blocked") + case <-time.After(100 * time.Millisecond): + } + + close(oldStore.release) + if swapped := waitForCooldownTransition(t, transitionDone, "cooldown store transition"); !swapped { + t.Fatal("SwapCooldownStateStore() = false") + } + select { + case <-resultDone: + case <-time.After(time.Second): + t.Fatal("result save did not complete after store transition") + } + if got := newStore.saveCount.Load(); got != 1 { + t.Fatalf("new store save count = %d, want 1", got) } } diff --git a/sdk/cliproxy/auth/home_concurrency.go b/sdk/cliproxy/auth/home_concurrency.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_concurrency.go @@ -0,0 +1,293 @@ +package auth + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" + "unicode/utf8" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" +) + +const ( + maxHomeConcurrencyTupleFieldLength = 256 + asciiWhitespace = " \t\r\n\v\f" +) + +var ErrMalformedHomeConcurrencyTuple = errors.New("malformed Home concurrency tuple") + +// HomeConcurrencyBusyError is a trusted, Home-originated concurrency admission failure. +type HomeConcurrencyBusyError struct { + cause *Error + retryAfter time.Duration +} + +// NewHomeConcurrencyBusyError creates a typed Home concurrency busy error. +func NewHomeConcurrencyBusyError(message string, retryAfter time.Duration) error { + message = strings.TrimSpace(message) + if message == "" { + message = "credential concurrency limit exceeded" + } + return newHomeConcurrencyBusyError(&Error{ + Code: "credential_concurrency_exceeded", + Message: message, + Retryable: true, + HTTPStatus: http.StatusTooManyRequests, + }, retryAfter) +} + +func newHomeConcurrencyBusyError(cause *Error, retryAfter time.Duration) *HomeConcurrencyBusyError { + return &HomeConcurrencyBusyError{cause: cause, retryAfter: retryAfter} +} + +func (e *HomeConcurrencyBusyError) Error() string { + if e == nil || e.cause == nil { + return "" + } + return e.cause.Error() +} + +// Unwrap preserves the Home error's code, retryability, and status for errors.As callers. +func (e *HomeConcurrencyBusyError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func (e *HomeConcurrencyBusyError) StatusCode() int { + if e == nil || e.cause == nil { + return 0 + } + return e.cause.StatusCode() +} + +func (e *HomeConcurrencyBusyError) RetryAfter() *time.Duration { + if e == nil || e.retryAfter <= 0 { + return nil + } + value := e.retryAfter + return &value +} + +func (e *HomeConcurrencyBusyError) SafeResponseHeaders() http.Header { + if e == nil { + return nil + } + return safeRetryAfterHeader(e.retryAfter) +} + +type homeConcurrencyTuple struct { + Accounted bool `json:"accounted"` + CredentialID string `json:"credential_id"` + Model string `json:"model"` +} + +func validateAccountedHomeConcurrencyTuple(tuple homeConcurrencyTuple) error { + model, validModel := validCanonicalHomeConcurrencyModelKey(tuple.Model) + if !tuple.Accounted || !validHomeConcurrencyTupleField(tuple.CredentialID) || !validModel || tuple.Model != model { + return ErrMalformedHomeConcurrencyTuple + } + return nil +} + +// canonicalHomeConcurrencyModelKey removes recognized reasoning suffixes from a Home limiter model key. +func canonicalHomeConcurrencyModelKey(model string) string { + if !utf8.ValidString(model) { + return "" + } + trimmed := strings.ToLower(strings.Trim(model, asciiWhitespace)) + if !strings.HasSuffix(trimmed, ")") { + return trimmed + } + open := strings.LastIndexByte(trimmed, '(') + if open < 0 { + return trimmed + } + suffix := trimmed[open+1 : len(trimmed)-1] + if !recognizedHomeConcurrencySuffix(suffix) { + return trimmed + } + base := strings.Trim(trimmed[:open], asciiWhitespace) + if base == "" { + return trimmed + } + return base +} + +func validCanonicalHomeConcurrencyModelKey(model string) (string, bool) { + key := canonicalHomeConcurrencyModelKey(model) + return key, key != "" && utf8.ValidString(key) && len(key) <= maxHomeConcurrencyTupleFieldLength +} + +func recognizedHomeConcurrencySuffix(value string) bool { + if value == "-1" { + return true + } + switch strings.ToLower(value) { + case "none", "auto", "minimal", "low", "medium", "high", "xhigh", "max": + return true + } + if value == "" || len(value) > 10 { + return false + } + var parsed int64 + for index := 0; index < len(value); index++ { + if value[index] < '0' || value[index] > '9' { + return false + } + parsed = parsed*10 + int64(value[index]-'0') + if parsed > 2_147_483_647 { + return false + } + } + return true +} + +func validHomeConcurrencyTupleField(value string) bool { + return value != "" && utf8.ValidString(value) && strings.TrimSpace(value) == value && len(value) <= maxHomeConcurrencyTupleFieldLength +} + +func installHomeConcurrencyScope(registry *executionregistry.Registry, pending *executionregistry.PendingDispatch, tuple homeConcurrencyTuple, base executionregistry.ScopeSpec) (*executionregistry.Scope, error) { + if registry == nil || pending == nil { + return nil, executionregistry.ErrInvalidPendingDispatch + } + if !tuple.Accounted { + base.Accounted = false + return registry.Install(pending, base) + } + if errValidate := validateAccountedHomeConcurrencyTuple(tuple); errValidate != nil { + return nil, errValidate + } + + base.CredentialID = tuple.CredentialID + base.Model = tuple.Model + base.Accounted = true + return registry.Install(pending, base) +} + +type homeDispatchConcurrencyEnvelope struct { + Tuple homeConcurrencyTuple + Present bool +} + +func decodeHomeDispatchConcurrencyEnvelope(raw []byte) (homeDispatchConcurrencyEnvelope, error) { + if !utf8.Valid(raw) { + return homeDispatchConcurrencyEnvelope{}, errors.New("Home response is not valid UTF-8") + } + + var fields map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(raw, &fields); errUnmarshal != nil || fields == nil { + return homeDispatchConcurrencyEnvelope{}, errors.New("Home response is not a JSON object") + } + + envelope := homeDispatchConcurrencyEnvelope{} + rawTuple, present := fields["concurrency"] + if !present { + return envelope, nil + } + envelope.Present = true + if errUnmarshal := json.Unmarshal(rawTuple, &envelope.Tuple); errUnmarshal != nil { + return envelope, errUnmarshal + } + if errValidate := validateAccountedHomeConcurrencyTuple(envelope.Tuple); errValidate != nil { + return envelope, errValidate + } + return envelope, nil +} + +func canonicalHomeDispatchModel(responseModel, requestedModel string) string { + if model := strings.TrimSpace(responseModel); model != "" { + return model + } + return requestedModel +} + +func decodeHomeDispatchError(raw []byte) error { + var fields map[string]json.RawMessage + if errUnmarshal := json.Unmarshal(raw, &fields); errUnmarshal != nil || fields == nil { + return nil + } + rawError, present := fields["error"] + if !present { + return nil + } + + var detail *homeErrorDetail + if errUnmarshal := json.Unmarshal(rawError, &detail); errUnmarshal != nil || detail == nil { + return &Error{Code: "invalid_auth", Message: "home returned malformed error payload", HTTPStatus: http.StatusBadGateway} + } + code := strings.TrimSpace(detail.Type) + if code == "" { + code = strings.TrimSpace(detail.Code) + } + if code == "" { + return &Error{Code: "invalid_auth", Message: "home returned malformed error payload", HTTPStatus: http.StatusBadGateway} + } + message := strings.TrimSpace(detail.Message) + if message == "" { + message = "home returned error" + } + + result := &Error{Code: code, Message: message, Retryable: detail.Retryable, HTTPStatus: http.StatusBadGateway} + switch strings.ToLower(code) { + case "model_not_found": + result.HTTPStatus = http.StatusNotFound + case "authentication_error", "unauthorized", "no_credentials", "invalid_credential": + result.HTTPStatus = http.StatusUnauthorized + case "credential_concurrency_exceeded", "credential_model_concurrency_exceeded": + result.HTTPStatus = http.StatusTooManyRequests + return newHomeConcurrencyBusyError(result, time.Duration(detail.RetryAfterMS)*time.Millisecond) + case "concurrency_protocol_required", "concurrency_tracker_unavailable", "concurrency_node_unavailable": + result.HTTPStatus = http.StatusServiceUnavailable + } + return result +} + +func invalidHomeConcurrencyResponse(message string) error { + return &Error{Code: "invalid_home_concurrency", Message: message, HTTPStatus: http.StatusBadGateway} +} + +func verifyAccountedHomeConcurrencyIdentity(tuple homeConcurrencyTuple, auth *Auth, authIndex string) error { + if !tuple.Accounted { + return nil + } + if auth == nil || auth.ID != tuple.CredentialID || authIndex != tuple.CredentialID { + return invalidHomeConcurrencyResponse("Home concurrency identity does not match dispatched auth") + } + return nil +} + +// SafeResponseHeaders returns trusted response headers only for CPA's concrete Home busy error. +func SafeResponseHeaders(err error) http.Header { + var busy *HomeConcurrencyBusyError + if !errors.As(err, &busy) || busy == nil { + return nil + } + return busy.SafeResponseHeaders() +} + +func safeRetryAfterHeader(retryAfter time.Duration) http.Header { + if retryAfter <= 0 { + return nil + } + seconds := int64(retryAfter / time.Second) + if retryAfter%time.Second != 0 { + seconds++ + } + if seconds < 1 { + seconds = 1 + } + return http.Header{"Retry-After": []string{strconv.FormatInt(seconds, 10)}} +} + +func homeConcurrencyInstallError(err error) error { + if errors.Is(err, ErrMalformedHomeConcurrencyTuple) { + return invalidHomeConcurrencyResponse(err.Error()) + } + return &Error{Code: "home_unavailable", Message: fmt.Sprintf("home execution registry unavailable: %v", err), Retryable: true, HTTPStatus: http.StatusServiceUnavailable} +} diff --git a/sdk/cliproxy/auth/home_concurrency_test.go b/sdk/cliproxy/auth/home_concurrency_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_concurrency_test.go @@ -0,0 +1,544 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "strings" + "sync/atomic" + "testing" + "time" + "unicode/utf8" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type fixtureHomeDispatcher struct { + payload []byte + payloads [][]byte + calls int + closedForAmbiguity bool + onAbort func() +} + +func (d *fixtureHomeDispatcher) HeartbeatOK() bool { return true } + +func (d *fixtureHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + if len(d.payloads) == 0 { + return d.payload, nil + } + if d.calls >= len(d.payloads) { + return nil, errors.New("unexpected Home dispatch") + } + payload := d.payloads[d.calls] + d.calls++ + return payload, nil +} + +func (d *fixtureHomeDispatcher) AbortAmbiguousDispatch() { + d.closedForAmbiguity = true + if d.onAbort != nil { + d.onAbort() + } +} + +func newHomeSelectionTestManager(t *testing.T, dispatcher homeAuthDispatcher) *Manager { + t.Helper() + manager := NewManager(nil, nil, nil) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + return manager +} + +type busyHomeRetryDispatcher struct { + calls atomic.Int32 +} + +func (*busyHomeRetryDispatcher) HeartbeatOK() bool { return true } + +func (d *busyHomeRetryDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + d.calls.Add(1) + return []byte(`{"error":{"type":"credential_concurrency_exceeded","message":"busy","retryable":true,"retry_after_ms":20000}}`), nil +} + +func (*busyHomeRetryDispatcher) AbortAmbiguousDispatch() {} + +func TestHomeBusySkipsNormalAndStreamOuterRetries(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "normal", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &busyHomeRetryDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, 30*time.Second, 0) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + if _, errRegister := manager.Register(context.Background(), &Auth{ID: "retry-auth", Provider: "home-busy"}); errRegister != nil { + t.Fatalf("register retry auth: %v", errRegister) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + started := time.Now() + go func() { + if stream { + _, errExecute := manager.ExecuteStream(ctx, []string{"home-busy"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + result <- errExecute + return + } + _, errExecute := manager.Execute(ctx, []string{"home-busy"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + result <- errExecute + }() + + select { + case errExecute := <-result: + var busy *HomeConcurrencyBusyError + if !errors.As(errExecute, &busy) { + t.Fatalf("execution error = %v, want HomeConcurrencyBusyError", errExecute) + } + case <-time.After(250 * time.Millisecond): + t.Fatal("Home busy waited for its retry hint") + } + if elapsed := time.Since(started); elapsed >= 250*time.Millisecond { + t.Fatalf("Home busy returned after %v, want prompt return", elapsed) + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1", got) + } + }) + } +} + +func TestPickHomeDispatchSelectionReleasesAccountedScopeAfterAuthValidationFailure(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"auth":{"id":"","provider":"codex"}}`)} + manager := newHomeSelectionTestManager(t, dispatcher) + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil { + t.Fatalf("selection=%#v error=%v", selection, errPick) + } + if dispatcher.closedForAmbiguity { + t.Fatal("accounted local auth validation failure fenced Home") + } + if freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()); len(freeze.Executions) != 0 { + t.Fatalf("scope was not released: %#v", freeze) + } +} + +func TestPickHomeDispatchSelectionReleasesAccountedScopeAfterPayloadDecodeFailure(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"model":123,"auth":{"id":"cred-1","provider":"codex"}}`)} + manager := newHomeSelectionTestManager(t, dispatcher) + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil { + t.Fatalf("selection=%#v error=%v", selection, errPick) + } + if dispatcher.closedForAmbiguity { + t.Fatal("accounted payload decode failure fenced Home") + } + if freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()); len(freeze.Executions) != 0 { + t.Fatalf("scope was not released: %#v", freeze) + } +} + +func TestPickHomeDispatchSelectionRejectsMalformedErrorPresence(t *testing.T) { + tests := []struct { + name string + payload string + wantCode string + wantFence bool + }{ + {name: "string without tuple", payload: `{"error":"busy","auth":{"id":"cred-1","provider":"codex"}}`, wantCode: "invalid_auth"}, + {name: "empty object without tuple", payload: `{"error":{},"auth":{"id":"cred-1","provider":"codex"}}`, wantCode: "invalid_auth"}, + {name: "null without tuple", payload: `{"error":null,"auth":{"id":"cred-1","provider":"codex"}}`, wantCode: "invalid_auth"}, + {name: "empty type and code without tuple", payload: `{"error":{"type":" ","code":""},"auth":{"id":"cred-1","provider":"codex"}}`, wantCode: "invalid_auth"}, + {name: "string with tuple", payload: `{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"error":"busy","auth":{"id":"cred-1","provider":"codex"}}`, wantCode: "invalid_home_concurrency", wantFence: true}, + {name: "empty object with tuple", payload: `{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"error":{},"auth":{"id":"cred-1","provider":"codex"}}`, wantCode: "invalid_home_concurrency", wantFence: true}, + {name: "null with tuple", payload: `{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"error":null,"auth":{"id":"cred-1","provider":"codex"}}`, wantCode: "invalid_home_concurrency", wantFence: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(tt.payload)} + manager := newHomeSelectionTestManager(t, dispatcher) + manager.executors["codex"] = schedulerTestExecutor{provider: "codex"} + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil { + t.Fatalf("selection=%#v error=%v, want malformed error rejection", selection, errPick) + } + var authErr *Error + if !errors.As(errPick, &authErr) || authErr.Code != tt.wantCode { + t.Fatalf("error=%#v, want code %q", errPick, tt.wantCode) + } + if dispatcher.closedForAmbiguity != tt.wantFence { + t.Fatalf("fenced=%t, want %t", dispatcher.closedForAmbiguity, tt.wantFence) + } + if freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()); len(freeze.Executions) != 0 { + t.Fatalf("scope was not released: %#v", freeze) + } + }) + } +} + +func TestPickHomeDispatchSelectionValidAccountedLocalValidationReleasesAndKeepsHomeHealthy(t *testing.T) { + validPayload := []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"auth_index":"cred-1","auth":{"id":"cred-1","provider":"codex"}}`) + tests := map[string][]byte{ + "auth validation": []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"auth":{"id":"","provider":"codex"}}`), + "payload decode": []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"model":123,"auth":{"id":"cred-1","provider":"codex"}}`), + "auth decode": []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"auth":"invalid"}`), + "identity mismatch": []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"auth_index":"other","auth":{"id":"cred-1","provider":"codex"}}`), + } + for name, invalidPayload := range tests { + t.Run(name, func(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payloads: [][]byte{invalidPayload, validPayload}} + manager := newHomeSelectionTestManager(t, dispatcher) + manager.executors["codex"] = schedulerTestExecutor{provider: "codex"} + releases := make(map[executionregistry.ReleaseGroup]int64) + registry := manager.HomeDispatchBundle().registry + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, sequence int64) { + releases[group] = sequence + }) + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil { + t.Fatalf("first selection=%#v error=%v, want local validation failure", selection, errPick) + } + if dispatcher.closedForAmbiguity { + t.Fatal("valid accounted local validation failure fenced Home") + } + group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"} + if len(releases) != 1 || releases[group] != 1 { + t.Fatalf("first releases=%#v, want exactly %v:1", releases, group) + } + + selection, errPick = manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if errPick != nil || selection == nil { + t.Fatalf("second selection=%#v error=%v, want healthy dispatch", selection, errPick) + } + selection.End("test_complete") + if dispatcher.closedForAmbiguity { + t.Fatal("second dispatch fenced Home") + } + if len(releases) != 1 || releases[group] != 2 { + t.Fatalf("cumulative releases=%#v, want exactly %v:2", releases, group) + } + }) + } +} + +func TestPickHomeDispatchSelectionFencesAccountedBusyError(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"error":{"type":"credential_concurrency_exceeded","message":"busy","retry_after_ms":750}}`)} + manager := newHomeSelectionTestManager(t, dispatcher) + abortSawScope := make(chan bool, 1) + dispatcher.onAbort = func() { + freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()) + abortSawScope <- len(freeze.Executions) == 1 + } + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil { + t.Fatalf("selection=%#v error=%v", selection, errPick) + } + var busy *HomeConcurrencyBusyError + if errors.As(errPick, &busy) { + t.Fatalf("accounted error returned ordinary busy response: %v", errPick) + } + if !dispatcher.closedForAmbiguity { + t.Fatal("accounted busy error did not fence Home") + } + if sawScope := <-abortSawScope; !sawScope { + t.Fatal("accounted scope ended before Home dispatch was aborted") + } + if freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()); len(freeze.Executions) != 0 { + t.Fatalf("scope was not released: %#v", freeze) + } +} + +func TestMalformedAccountedTupleClosesHomeClient(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":""},"auth":{"id":"cred-1","provider":"codex"}}`)} + manager := newHomeSelectionTestManager(t, dispatcher) + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil || !dispatcher.closedForAmbiguity { + t.Fatalf("selection=%#v error=%v closed=%t", selection, errPick, dispatcher.closedForAmbiguity) + } +} + +func TestConcurrencyDispatchFixture(t *testing.T) { + t.Run("accounted", func(t *testing.T) { + raw, errRead := os.ReadFile("../../../internal/home/testdata/concurrency_dispatch_accounted.json") + if errRead != nil { + t.Fatalf("ReadFile(accounted fixture) error = %v", errRead) + } + + var fixture struct { + Model string `json:"model"` + Provider string `json:"provider"` + AuthIndex string `json:"auth_index"` + Auth struct { + ID string `json:"id"` + Provider string `json:"provider"` + } `json:"auth"` + Concurrency homeConcurrencyTuple `json:"concurrency"` + } + if errUnmarshal := json.Unmarshal(raw, &fixture); errUnmarshal != nil { + t.Fatalf("Unmarshal(accounted fixture) error = %v", errUnmarshal) + } + wantTuple := homeConcurrencyTuple{Accounted: true, CredentialID: "cred-1", Model: "gpt"} + if fixture.Concurrency != wantTuple { + t.Fatalf("accounted concurrency = %#v, want %#v", fixture.Concurrency, wantTuple) + } + if fixture.Model != "gpt" || fixture.Provider != "codex" || fixture.AuthIndex != "cred-1" || fixture.Auth.ID != "cred-1" || fixture.Auth.Provider != "codex" { + t.Fatalf("accounted identity model=%q provider=%q auth_index=%q auth=%#v", fixture.Model, fixture.Provider, fixture.AuthIndex, fixture.Auth) + } + + envelope, errEnvelope := decodeHomeDispatchConcurrencyEnvelope(raw) + if errEnvelope != nil { + t.Fatalf("decodeHomeDispatchConcurrencyEnvelope(accounted fixture) error = %v", errEnvelope) + } + if !envelope.Present || envelope.Tuple != wantTuple { + t.Fatalf("accounted envelope = %#v, want present tuple %#v", envelope, wantTuple) + } + + dispatcher := &fixtureHomeDispatcher{payload: raw} + manager := newHomeSelectionTestManager(t, dispatcher) + manager.RegisterExecutor(schedulerTestExecutor{provider: "codex"}) + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if errPick != nil || selection == nil { + t.Fatalf("pickHomeDispatchSelection(accounted fixture) selection=%#v error=%v", selection, errPick) + } + defer selection.End("fixture_complete") + if selection.Auth == nil || selection.Auth.ID != "cred-1" || selection.Auth.Index != "cred-1" || selection.Auth.Provider != "codex" { + t.Fatalf("selected auth = %#v", selection.Auth) + } + + bundle := manager.HomeDispatchBundle() + if bundle == nil || bundle.registry == nil { + t.Fatal("accounted fixture did not retain a Home dispatch registry") + } + freeze := bundle.registry.FreezeInFlight(time.Now()) + if len(freeze.Executions) != 1 { + t.Fatalf("accounted fixture executions = %#v", freeze.Executions) + } + gotScope := freeze.Executions[0] + if !gotScope.Accounted || gotScope.CredentialID != "cred-1" || gotScope.Model != "gpt" { + t.Fatalf("accounted fixture scope = %#v", gotScope) + } + }) + + t.Run("busy", func(t *testing.T) { + raw, errRead := os.ReadFile("../../../internal/home/testdata/concurrency_dispatch_busy.json") + if errRead != nil { + t.Fatalf("ReadFile(busy fixture) error = %v", errRead) + } + + var fixture struct { + Error *struct { + Type string `json:"type"` + Message string `json:"message"` + Retryable bool `json:"retryable"` + RetryAfterMS int64 `json:"retry_after_ms"` + } `json:"error"` + } + if errUnmarshal := json.Unmarshal(raw, &fixture); errUnmarshal != nil { + t.Fatalf("Unmarshal(busy fixture) error = %v", errUnmarshal) + } + if fixture.Error == nil { + t.Fatal("busy fixture has no error object") + } + if fixture.Error.Type != "credential_concurrency_exceeded" || fixture.Error.Message != "credential concurrency limit reached" || !fixture.Error.Retryable || fixture.Error.RetryAfterMS != 750 { + t.Fatalf("busy fixture error = %#v", fixture.Error) + } + + errBusy := decodeHomeDispatchError(raw) + var busy *HomeConcurrencyBusyError + if !errors.As(errBusy, &busy) || busy == nil { + t.Fatalf("decodeHomeDispatchError(busy fixture) error = %#v, want *HomeConcurrencyBusyError", errBusy) + } + if got := busy.StatusCode(); got != http.StatusTooManyRequests { + t.Fatalf("busy status = %d, want %d", got, http.StatusTooManyRequests) + } + retryAfter := busy.RetryAfter() + if retryAfter == nil || *retryAfter != 750*time.Millisecond { + t.Fatalf("busy retry after = %v, want 750ms", retryAfter) + } + var cause *Error + if !errors.As(errBusy, &cause) || cause == nil || cause.Code != fixture.Error.Type || cause.Message != fixture.Error.Message || !cause.Retryable || cause.HTTPStatus != http.StatusTooManyRequests { + t.Fatalf("busy typed cause = %#v", cause) + } + }) +} + +func TestHomeBusyErrorMaps429AndRetryAfter(t *testing.T) { + errBusy := decodeHomeDispatchError([]byte(`{"error":{"type":"credential_concurrency_exceeded","message":"busy","retryable":true,"retry_after_ms":750}}`)) + statusError, ok := errBusy.(interface{ StatusCode() int }) + if !ok || statusError.StatusCode() != http.StatusTooManyRequests { + t.Fatalf("error = %#v", errBusy) + } + retryError, ok := errBusy.(interface{ RetryAfter() *time.Duration }) + if !ok || retryError.RetryAfter() == nil || *retryError.RetryAfter() != 750*time.Millisecond { + t.Fatalf("retry after = %v", retryError.RetryAfter()) + } +} + +func TestHomeConcurrencyTupleAuthMismatchEndsScope(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"auth_index":"other","auth":{"id":"cred-1","provider":"codex"}}`)} + manager := newHomeSelectionTestManager(t, dispatcher) + manager.executors["codex"] = schedulerTestExecutor{provider: "codex"} + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil { + t.Fatalf("selection=%#v error=%v", selection, errPick) + } + if dispatcher.closedForAmbiguity { + t.Fatal("accounted auth identity mismatch fenced Home") + } + if freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()); len(freeze.Executions) != 0 { + t.Fatalf("scope was not released: %#v", freeze) + } +} + +func TestOldHomeDispatchIsUnaccounted(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(`{"auth":{"id":"cred-1","provider":"codex"}}`)} + manager := newHomeSelectionTestManager(t, dispatcher) + manager.executors["codex"] = schedulerTestExecutor{provider: "codex"} + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if errPick != nil || selection == nil { + t.Fatalf("selection=%#v error=%v", selection, errPick) + } + defer selection.End("test") + freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()) + if len(freeze.Executions) != 1 || freeze.Executions[0].Accounted { + t.Fatalf("old Home dispatch freeze = %#v", freeze) + } +} + +func TestHomeBusyErrorHeadersRoundUpMilliseconds(t *testing.T) { + errBusy := decodeHomeDispatchError([]byte(`{"error":{"type":"credential_concurrency_exceeded","message":"busy","retry_after_ms":750}}`)) + headers, ok := errBusy.(interface{ SafeResponseHeaders() http.Header }) + if !ok { + t.Fatalf("error has no safe headers: %#v", errBusy) + } + if got := headers.SafeResponseHeaders().Get("Retry-After"); got != "1" { + t.Fatalf("Retry-After = %q, want 1", got) + } +} + +func TestInstallHomeConcurrencyScopeRejectsNonCanonicalTuple(t *testing.T) { + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + defer pending.End() + + _, errInstall := installHomeConcurrencyScope(registry, pending, homeConcurrencyTuple{ + Accounted: true, CredentialID: " cred-1 ", Model: "gpt", + }, executionregistry.ScopeSpec{Kind: "http", StartedAt: time.Now()}) + if !errors.Is(errInstall, ErrMalformedHomeConcurrencyTuple) { + t.Fatalf("install error = %v, want malformed tuple", errInstall) + } +} + +func TestPickHomeDispatchSelectionFencesInvalidExplicitConcurrency(t *testing.T) { + tests := []string{ + `{"concurrency":{"accounted":false,"credential_id":"cred-1","model":"gpt"},"auth":{"id":"cred-1","provider":"codex"}}`, + `{"concurrency":{"accounted":true,"credential_id":" cred-1","model":"gpt"},"auth":{"id":"cred-1","provider":"codex"}}`, + `{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"other"},"model":"gpt","auth":{"id":"cred-1","provider":"codex"}}`, + } + for _, payload := range tests { + dispatcher := &fixtureHomeDispatcher{payload: []byte(payload)} + manager := newHomeSelectionTestManager(t, dispatcher) + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil || !dispatcher.closedForAmbiguity { + t.Fatalf("payload=%s selection=%#v error=%v closed=%t", payload, selection, errPick, dispatcher.closedForAmbiguity) + } + } +} + +func TestPickHomeDispatchSelectionReleasesAccountedScopeAfterAuthDecodeFailure(t *testing.T) { + dispatcher := &fixtureHomeDispatcher{payload: []byte(`{"concurrency":{"accounted":true,"credential_id":"cred-1","model":"gpt"},"auth":"invalid"}`)} + manager := newHomeSelectionTestManager(t, dispatcher) + + selection, errPick := manager.pickHomeDispatchSelection(context.Background(), "gpt", cliproxyexecutor.Options{}) + if selection != nil || errPick == nil || dispatcher.closedForAmbiguity { + t.Fatalf("selection=%#v error=%v closed=%t", selection, errPick, dispatcher.closedForAmbiguity) + } + if freeze := manager.HomeDispatchBundle().registry.FreezeInFlight(time.Now()); len(freeze.Executions) != 0 { + t.Fatalf("scope was not released: %#v", freeze) + } +} + +func TestHomeConcurrencyBusyErrorsRemainTypedWhenWrapped(t *testing.T) { + for _, code := range []string{"credential_concurrency_exceeded", "credential_model_concurrency_exceeded"} { + errBusy := decodeHomeDispatchError([]byte(fmt.Sprintf(`{"error":{"type":%q,"message":"busy","retryable":false}}`, code))) + var busy *HomeConcurrencyBusyError + if !errors.As(errBusy, &busy) { + t.Fatalf("code=%s error=%#v, want typed busy error", code, errBusy) + } + if busy.RetryAfter() != nil { + t.Fatalf("code=%s retry after = %v, want nil", code, busy.RetryAfter()) + } + var cause *Error + if !errors.As(fmt.Errorf("wrapped: %w", errBusy), &cause) || cause.Code != code || cause.Retryable { + t.Fatalf("code=%s cause=%#v", code, cause) + } + } +} + +func TestRetryAfterFromWrappedHomeBusyError(t *testing.T) { + errBusy := NewHomeConcurrencyBusyError("busy", 750*time.Millisecond) + if got := retryAfterFromError(fmt.Errorf("wrapped: %w", errBusy)); got == nil || *got != 750*time.Millisecond { + t.Fatalf("retry after = %v, want 750ms", got) + } +} + +func TestCanonicalHomeConcurrencyModelKeyMatchesHomeLimiter(t *testing.T) { + cases := map[string]string{ + " gpt(high) ": "gpt", + "gpt(8192)": "gpt", + "gpt(-1)": "gpt", + " GPT(AUTO) ": "gpt", + "model(custom)": "model(custom)", + "model(+1)": "model(+1)", + "model(2147483648)": "model(2147483648)", + "(high)": "(high)", + } + for input, want := range cases { + if got := canonicalHomeConcurrencyModelKey(input); got != want { + t.Fatalf("canonicalHomeConcurrencyModelKey(%q) = %q, want %q", input, got, want) + } + } + if got := canonicalHomeConcurrencyModelKey("gpt\xff(high)"); got != "" { + t.Fatalf("canonicalHomeConcurrencyModelKey() = %q, want empty for malformed UTF-8", got) + } +} + +func TestAccountedHomeConcurrencyTupleRequiresCanonicalLimiterModel(t *testing.T) { + for _, model := range []string{"GPT", "gpt(high)", "model(custom) "} { + errValidate := validateAccountedHomeConcurrencyTuple(homeConcurrencyTuple{Accounted: true, CredentialID: "cred-1", Model: model}) + if !errors.Is(errValidate, ErrMalformedHomeConcurrencyTuple) { + t.Fatalf("model=%q validation error = %v, want malformed tuple", model, errValidate) + } + } +} + +func TestHomeConcurrencyTupleStringsAreValidUTF8(t *testing.T) { + if utf8.ValidString(string([]byte{0xff})) { + t.Fatal("test setup expected invalid UTF-8") + } + if _, errDecode := decodeHomeDispatchConcurrencyEnvelope([]byte{'{', 0xff, '}'}); errDecode == nil { + t.Fatal("raw non-UTF-8 Home envelope was accepted") + } + if !errors.Is(validateAccountedHomeConcurrencyTuple(homeConcurrencyTuple{Accounted: true, CredentialID: string([]byte{0xff}), Model: "gpt"}), ErrMalformedHomeConcurrencyTuple) { + t.Fatal("invalid UTF-8 credential was accepted") + } + if !errors.Is(validateAccountedHomeConcurrencyTuple(homeConcurrencyTuple{Accounted: true, CredentialID: "cred-1", Model: strings.Repeat("g", 257)}), ErrMalformedHomeConcurrencyTuple) { + t.Fatal("oversized model was accepted") + } +} diff --git a/sdk/cliproxy/auth/home_execution_paths_test.go b/sdk/cliproxy/auth/home_execution_paths_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_execution_paths_test.go @@ -0,0 +1,1128 @@ +package auth + +import ( + "context" + "encoding/json" + "net/http" + "strconv" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type homeExecutionDispatcher struct{} + +func (homeExecutionDispatcher) HeartbeatOK() bool { return true } + +func (homeExecutionDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive}}) +} + +func (homeExecutionDispatcher) AbortAmbiguousDispatch() {} + +type homeExecutionStreamExecutor struct { + chunks <-chan cliproxyexecutor.StreamChunk +} + +type homeExecutionExecutor struct { + ctx context.Context +} + +func (*homeExecutionExecutor) Identifier() string { return "home-execution" } +func (e *homeExecutionExecutor) Execute(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.ctx = ctx + if errCtx := ctx.Err(); errCtx != nil { + return cliproxyexecutor.Response{}, errCtx + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} +func (*homeExecutionExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*homeExecutionExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*homeExecutionExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*homeExecutionExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (*homeExecutionStreamExecutor) Identifier() string { return "home-execution" } +func (*homeExecutionStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *homeExecutionStreamExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return &cliproxyexecutor.StreamResult{Chunks: e.chunks}, nil +} +func (*homeExecutionStreamExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*homeExecutionStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*homeExecutionStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeModeNeverAuthorizesLocalAuthFallback(t *testing.T) { + manager := NewManager(nil, nil, nil) + cfg := &internalconfig.Config{} + cfg.Home.Enabled = true + manager.runtimeConfig.Store(cfg) + manager.auths["local-antigravity"] = &Auth{ID: "local-antigravity", Provider: "antigravity", Status: StatusActive} + + if manager.localExecutionAllowed() { + t.Fatal("local execution allowed in Home mode") + } + if selected := manager.localFallbackAuth("local-antigravity"); selected != nil { + t.Fatalf("local fallback auth = %#v", selected) + } +} + +func TestHomeSelectionEndsAfterExecute(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(homeExecutionDispatcher{}, executionregistry.New(), 1) + executor := &homeExecutionExecutor{} + manager.RegisterExecutor(executor) + + if _, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + if executor.ctx == nil { + t.Fatal("executor did not receive an attempt context") + } + if errCtx := executor.ctx.Err(); errCtx == nil { + t.Fatal("attempt context was not canceled after execution") + } +} + +func TestHomeSelectionEndsOnMissingExecutor(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1) + + if _, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errExecute == nil { + t.Fatal("Execute() error = nil, want missing executor") + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestHomeSelectionClosesAttemptAndWebSocketResources(t *testing.T) { + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + selection, errSelection := newHomeDispatchSelection(&Auth{ID: "home-auth"}, nil, "home-execution", scope) + if errSelection != nil { + t.Fatal(errSelection) + } + attemptCtx, releaseAttempt, errBind := homeExecutionAttemptContext(context.Background(), selection) + if errBind != nil { + t.Fatal(errBind) + } + var closeCalls atomic.Int32 + if errBind = selection.Bind(func() error { + closeCalls.Add(1) + return nil + }); errBind != nil { + t.Fatal(errBind) + } + selection.End("completed") + releaseAttempt() + if errCtx := attemptCtx.Err(); errCtx == nil { + t.Fatal("attempt context was not canceled") + } + if got := closeCalls.Load(); got != 1 { + t.Fatalf("resource close calls = %d, want 1", got) + } +} + +func TestHomeStreamConsumerCancelEndsSelection(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1) + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")} + manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks}) + + result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + cancel() + for range result.Chunks { + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +type retainingHomeExecutionDispatcher struct { + calls atomic.Int32 +} + +func (d *retainingHomeExecutionDispatcher) HeartbeatOK() bool { return true } + +func (d *retainingHomeExecutionDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + d.calls.Add(1) + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-auth", + Provider: "home-execution", + Status: StatusActive, + Attributes: map[string]string{ + "websockets": "true", + }, + }}) +} + +func (*retainingHomeExecutionDispatcher) AbortAmbiguousDispatch() {} + +type retainingHomeExecutionExecutor struct { + calls atomic.Int32 +} + +func (*retainingHomeExecutionExecutor) Identifier() string { return "home-execution" } + +func (e *retainingHomeExecutionExecutor) Execute(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.calls.Add(1) + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (*retainingHomeExecutionExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*retainingHomeExecutionExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} +func (*retainingHomeExecutionExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*retainingHomeExecutionExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeWebsocketSessionReusesRetainedSelection(t *testing.T) { + dispatcher := &retainingHomeExecutionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &retainingHomeExecutionExecutor{} + manager.RegisterExecutor(executor) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "session-1", + cliproxyexecutor.PinnedAuthMetadataKey: "home-auth", + }} + for range 2 { + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1 for one retained session target", got) + } + if got := executor.calls.Load(); got != 2 { + t.Fatalf("executor calls = %d, want 2", got) + } +} + +type changingHomeTargetDispatcher struct { + calls atomic.Int32 + firstSelection *HomeDispatchSelection + oldEndedBeforeRPop atomic.Bool +} + +func (d *changingHomeTargetDispatcher) HeartbeatOK() bool { return true } +func (d *changingHomeTargetDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + if d.calls.Add(1) == 2 && d.firstSelection != nil { + d.oldEndedBeforeRPop.Store(!d.firstSelection.Active()) + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive, Attributes: map[string]string{"websockets": "true"}}}) +} +func (*changingHomeTargetDispatcher) AbortAmbiguousDispatch() {} + +type selectionRecordingExecutor struct { + first *HomeDispatchSelection +} + +func (*selectionRecordingExecutor) Identifier() string { return "home-execution" } +func (e *selectionRecordingExecutor) Execute(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + selection, _ := opts.ExecutionLifecycle.(*HomeDispatchSelection) + if e.first == nil { + e.first = selection + } + if selection != nil { + selection.Retain() + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} +func (*selectionRecordingExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*selectionRecordingExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*selectionRecordingExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*selectionRecordingExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeWebsocketTargetChangeEndsSelectionBeforeRedispatch(t *testing.T) { + dispatcher := &changingHomeTargetDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &selectionRecordingExecutor{} + manager.RegisterExecutor(executor) + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "session-1", + cliproxyexecutor.PinnedAuthMetadataKey: "home-auth", + }} + + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil { + t.Fatalf("first Execute() error = %v", errExecute) + } + dispatcher.firstSelection = executor.first + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-b"}, opts); errExecute != nil { + t.Fatalf("second Execute() error = %v", errExecute) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 after target change", got) + } + if !dispatcher.oldEndedBeforeRPop.Load() { + t.Fatal("previous selection remained active when target-change RPOP started") + } +} + +type unpinnedTargetChangeDispatcher struct { + calls atomic.Int32 + first *HomeDispatchSelection + oldClosedBeforeDispatch atomic.Bool + closeCalls *atomic.Int32 +} + +func (d *unpinnedTargetChangeDispatcher) HeartbeatOK() bool { return true } +func (d *unpinnedTargetChangeDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, _ int) ([]byte, error) { + call := d.calls.Add(1) + if call == 2 && d.first != nil { + d.oldClosedBeforeDispatch.Store(!d.first.Active() && d.closeCalls.Load() == 1) + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-auth-" + strconv.Itoa(int(call)), + Provider: "home-execution", + Status: StatusActive, + Attributes: map[string]string{ + "websockets": "true", + }, + }}) +} +func (*unpinnedTargetChangeDispatcher) AbortAmbiguousDispatch() {} + +type bindingSelectionRecordingExecutor struct { + first *HomeDispatchSelection + closeCalls *atomic.Int32 +} + +func (*bindingSelectionRecordingExecutor) Identifier() string { return "home-execution" } +func (e *bindingSelectionRecordingExecutor) Execute(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + selection, _ := opts.ExecutionLifecycle.(*HomeDispatchSelection) + if e.first == nil { + e.first = selection + } + if selection != nil { + if errBind := selection.Bind(func() error { + e.closeCalls.Add(1) + return nil + }); errBind != nil { + return cliproxyexecutor.Response{}, errBind + } + selection.Retain() + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} +func (*bindingSelectionRecordingExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*bindingSelectionRecordingExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} +func (*bindingSelectionRecordingExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*bindingSelectionRecordingExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeWebsocketUnpinnedModelChangeClosesSelectionBeforeRedispatch(t *testing.T) { + var closeCalls atomic.Int32 + dispatcher := &unpinnedTargetChangeDispatcher{closeCalls: &closeCalls} + registry := executionregistry.New() + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + executor := &bindingSelectionRecordingExecutor{closeCalls: &closeCalls} + manager.RegisterExecutor(executor) + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "session-1", + }} + + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil { + t.Fatalf("first Execute() error = %v", errExecute) + } + dispatcher.first = executor.first + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-b"}, opts); errExecute != nil { + t.Fatalf("second Execute() error = %v", errExecute) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2", got) + } + if !dispatcher.oldClosedBeforeDispatch.Load() { + t.Fatal("old unpinned selection was not ended and closed before the second RPOP") + } + manager.CloseExecutionSession("session-1") + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +type lifecycleRetryDispatcher struct { + calls atomic.Int32 + executor *lifecycleRetryExecutor + firstEndedBeforeRedispatch atomic.Bool +} + +func (d *lifecycleRetryDispatcher) HeartbeatOK() bool { return true } +func (d *lifecycleRetryDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, _ int) ([]byte, error) { + if d.calls.Add(1) == 2 && d.executor.first != nil { + d.firstEndedBeforeRedispatch.Store(!d.executor.first.Active() && d.executor.firstCtx.Err() != nil) + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive, Attributes: map[string]string{"websockets": "true"}}}) +} +func (*lifecycleRetryDispatcher) AbortAmbiguousDispatch() {} + +type lifecycleRetryExecutor struct { + calls atomic.Int32 + first *HomeDispatchSelection + firstCtx context.Context +} + +func (*lifecycleRetryExecutor) Identifier() string { return "home-execution" } +func (*lifecycleRetryExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *lifecycleRetryExecutor) ExecuteStream(ctx context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if e.calls.Add(1) == 1 { + e.first, _ = opts.ExecutionLifecycle.(*HomeDispatchSelection) + e.firstCtx = ctx + return nil, &Error{HTTPStatus: http.StatusUpgradeRequired, Message: "websocket upgrade required"} + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed"}`)} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} +func (*lifecycleRetryExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*lifecycleRetryExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*lifecycleRetryExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeStreamLifecycleFailureEndsBeforeFreshDispatch(t *testing.T) { + executor := &lifecycleRetryExecutor{} + dispatcher := &lifecycleRetryDispatcher{executor: executor} + registry := executionregistry.New() + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(executor) + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Stream: true, Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "session-426", + }} + + result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + if got := executor.calls.Load(); got != 2 { + t.Fatalf("executor invocations = %d, want 2", got) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2", got) + } + if !dispatcher.firstEndedBeforeRedispatch.Load() { + t.Fatal("failed stream attempt remained active when the fresh Home selection was dispatched") + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestHomeSelectionCancellationPreventsExecute(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(homeExecutionDispatcher{}, executionregistry.New(), 1) + executor := &homeExecutionExecutor{} + manager.RegisterExecutor(executor) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("Execute() error = nil, want canceled context") + } + if executor.ctx != nil { + t.Fatal("executor was invoked after attempt context cancellation") + } +} + +type retryingHomeStreamExecutor struct { + calls atomic.Int32 +} + +func (*retryingHomeStreamExecutor) Identifier() string { return "home-execution" } +func (*retryingHomeStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *retryingHomeStreamExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if e.calls.Add(1) == 1 { + return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired"} + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("data: {\"type\":\"response.completed\"}\n\n")} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} +func (*retryingHomeStreamExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (*retryingHomeStreamExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*retryingHomeStreamExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeStreamRetryUsesFreshSelection(t *testing.T) { + dispatcher := &retainingHomeExecutionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &retryingHomeStreamExecutor{} + manager.RegisterExecutor(executor) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 for retrying stream invocations", got) + } +} + +type cancellationBarrierExecutor struct { + executeCalls atomic.Int32 + countCalls atomic.Int32 + streamCalls atomic.Int32 +} + +func (*cancellationBarrierExecutor) Identifier() string { return "home-execution" } +func (e *cancellationBarrierExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.executeCalls.Add(1) + return cliproxyexecutor.Response{}, nil +} +func (e *cancellationBarrierExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.countCalls.Add(1) + return cliproxyexecutor.Response{}, nil +} +func (e *cancellationBarrierExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.streamCalls.Add(1) + return nil, nil +} +func (*cancellationBarrierExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*cancellationBarrierExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeCancellationBarrierPreventsEveryExecutorInvocation(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(homeExecutionDispatcher{}, executionregistry.New(), 1) + executor := &cancellationBarrierExecutor{} + manager.RegisterExecutor(executor) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errExecute == nil { + t.Fatal("Execute() error = nil, want canceled context") + } + if _, errCount := manager.ExecuteCount(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{}); errCount == nil { + t.Fatal("ExecuteCount() error = nil, want canceled context") + } + if _, errStream := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true}); errStream == nil { + t.Fatal("ExecuteStream() error = nil, want canceled context") + } + if got := executor.executeCalls.Load(); got != 0 { + t.Fatalf("Execute calls = %d, want 0", got) + } + if got := executor.countCalls.Load(); got != 0 { + t.Fatalf("CountTokens calls = %d, want 0", got) + } + if got := executor.streamCalls.Load(); got != 0 { + t.Fatalf("ExecuteStream calls = %d, want 0", got) + } +} + +func TestHomeStreamEndsOnTerminalChunk(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1) + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")} + manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks}) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + + close(chunks) + for range result.Chunks { + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestHomeWebsocketSessionReusesSelectionWithoutPinnedMetadataAndCachesRuntimeAuth(t *testing.T) { + dispatcher := &retainingHomeExecutionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &retainingHomeExecutionExecutor{} + manager.RegisterExecutor(executor) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "session-without-pin", + }} + for range 2 { + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts); errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1 for a retained session without a pin", got) + } + if auth, ok := manager.GetExecutionSessionAuthByID("session-without-pin", "home-auth"); !ok || auth == nil { + t.Fatal("retained selection did not populate the handler runtime auth cache") + } +} + +func TestCloseExecutionSessionReclaimsHomeSessionLock(t *testing.T) { + manager := NewManager(nil, nil, nil) + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "reclaim-lock", + }} + unlock := manager.lockHomeWebsocketSession(ctx, opts) + if unlock == nil { + t.Fatal("lockHomeWebsocketSession() = nil") + } + unlock() + if _, ok := manager.homeSessionLocks.Load("reclaim-lock"); !ok { + t.Fatal("session lock was not created") + } + + manager.CloseExecutionSession("reclaim-lock") + if _, ok := manager.homeSessionLocks.Load("reclaim-lock"); ok { + t.Fatal("closed session retained its mutex entry") + } +} + +type homePerSelectionDispatcher struct { + auths []Auth + calls atomic.Int32 + first *HomeDispatchSelection + firstEndedBefore2 atomic.Bool +} + +func (*homePerSelectionDispatcher) HeartbeatOK() bool { return true } +func (d *homePerSelectionDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + call := d.calls.Add(1) + if call == 2 && d.first != nil { + d.firstEndedBefore2.Store(!d.first.Active()) + } + if int(call) > len(d.auths) { + return nil, home.ErrAuthNotFound + } + return json.Marshal(homeAuthDispatchResponse{Auth: d.auths[call-1]}) +} +func (*homePerSelectionDispatcher) AbortAmbiguousDispatch() {} + +type homePerSelectionFailureExecutor struct { + dispatcher *homePerSelectionDispatcher + selections []*HomeDispatchSelection + invocations []string +} + +func (*homePerSelectionFailureExecutor) Identifier() string { return openAICompatPoolProviderKey } +func (e *homePerSelectionFailureExecutor) invoke(auth *Auth, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + selection, _ := opts.ExecutionLifecycle.(*HomeDispatchSelection) + if e.selections == nil { + e.selections = append(e.selections, selection) + } + if selection != nil && len(e.selections) == 1 { + e.selections[0] = selection + if e.dispatcher != nil { + e.dispatcher.first = selection + } + } + e.invocations = append(e.invocations, auth.ID) + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"} +} +func (e *homePerSelectionFailureExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return e.invoke(auth, opts) +} +func (*homePerSelectionFailureExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*homePerSelectionFailureExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} +func (e *homePerSelectionFailureExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return e.invoke(auth, opts) +} +func (*homePerSelectionFailureExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeNonstreamAndCountUseOneModelPerSelection(t *testing.T) { + for _, countTokens := range []bool{false, true} { + t.Run(map[bool]string{false: "Execute", true: "CountTokens"}[countTokens], func(t *testing.T) { + dispatcher := &homePerSelectionDispatcher{auths: []Auth{ + {ID: "home-auth-a", Provider: "home-pool", Status: StatusActive, Attributes: map[string]string{"api_key": "test-key", "compat_name": "pool", "provider_key": "pool"}}, + {ID: "home-auth-b", Provider: "home-pool", Status: StatusActive, Attributes: map[string]string{"api_key": "test-key", "compat_name": "pool", "provider_key": "pool"}}, + }} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + OpenAICompatibility: []internalconfig.OpenAICompatibility{{ + Name: "pool", + Models: []internalconfig.OpenAICompatibilityModel{{Name: "upstream-a", Alias: "requested"}, {Name: "upstream-b", Alias: "requested"}}, + }}, + }) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &homePerSelectionFailureExecutor{dispatcher: dispatcher} + manager.RegisterExecutor(executor) + + var errExecute error + if countTokens { + _, errExecute = manager.ExecuteCount(context.Background(), []string{openAICompatPoolProviderKey}, cliproxyexecutor.Request{Model: "requested"}, cliproxyexecutor.Options{}) + } else { + _, errExecute = manager.Execute(context.Background(), []string{openAICompatPoolProviderKey}, cliproxyexecutor.Request{Model: "requested"}, cliproxyexecutor.Options{}) + } + if errExecute == nil { + t.Fatal("execution error = nil, want upstream failure") + } + if len(executor.invocations) != 2 { + t.Fatalf("execution error = %v; upstream invocations = %v, want one per Home selection", errExecute, executor.invocations) + } + if !dispatcher.firstEndedBefore2.Load() { + t.Fatal("first Home selection was not ended before the next dispatch") + } + }) + } +} + +func TestHomeStreamEndsOnErrorChunk(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(homeExecutionDispatcher{}, registry, 1) + chunks := make(chan cliproxyexecutor.StreamChunk, 2) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")} + chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}} + close(chunks) + manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks}) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + sawError := false + for chunk := range result.Chunks { + if chunk.Err != nil { + sawError = true + } + } + if !sawError { + t.Fatal("stream did not preserve the upstream error chunk") + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +type missingHomeStreamSourceExecutor struct{} + +func (*missingHomeStreamSourceExecutor) Identifier() string { return "home-execution" } +func (*missingHomeStreamSourceExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*missingHomeStreamSourceExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*missingHomeStreamSourceExecutor) Refresh(context.Context, *Auth) (*Auth, error) { + return nil, nil +} +func (*missingHomeStreamSourceExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*missingHomeStreamSourceExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +type accountedHomeExecutionDispatcher struct { + calls atomic.Int32 + auths []Auth +} + +func (*accountedHomeExecutionDispatcher) HeartbeatOK() bool { return true } +func (d *accountedHomeExecutionDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + index := int(d.calls.Add(1)) - 1 + if index >= len(d.auths) { + return nil, home.ErrAuthNotFound + } + auth := d.auths[index] + return json.Marshal(struct { + Concurrency homeConcurrencyTuple `json:"concurrency"` + Model string `json:"model"` + AuthIndex string `json:"auth_index"` + Auth Auth `json:"auth"` + }{ + Concurrency: homeConcurrencyTuple{Accounted: true, CredentialID: auth.ID, Model: model}, + Model: model, + AuthIndex: auth.ID, + Auth: auth, + }) +} +func (*accountedHomeExecutionDispatcher) AbortAmbiguousDispatch() {} + +func TestAccountedHomeExecuteAndCountReleaseOnce(t *testing.T) { + for _, countTokens := range []bool{false, true} { + t.Run(map[bool]string{false: "Execute", true: "Count"}[countTokens], func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + releases := make(chan executionregistry.ReleaseGroup, 2) + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group }) + manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{ + ID: "cred-1", Provider: "home-execution", Status: StatusActive, + }}}, registry, 1) + manager.RegisterExecutor(&homeExecutionExecutor{}) + + var errExecute error + if countTokens { + _, errExecute = manager.ExecuteCount(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) + } else { + _, errExecute = manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}) + } + if errExecute != nil { + t.Fatalf("execution error = %v", errExecute) + } + select { + case group := <-releases: + if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) { + t.Fatalf("release group = %#v", group) + } + default: + t.Fatal("accounted selection did not release") + } + select { + case group := <-releases: + t.Fatalf("duplicate release = %#v", group) + default: + } + }) + } +} + +func TestAccountedHomeStreamEndsOnlyAfterSourceTerminates(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + releases := make(chan executionregistry.ReleaseGroup, 1) + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group }) + manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{ + ID: "cred-1", Provider: "home-execution", Status: StatusActive, + }}}, registry, 1) + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")} + manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks}) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + if _, ok := <-result.Chunks; !ok { + t.Fatal("stream closed before initial chunk") + } + select { + case group := <-releases: + t.Fatalf("stream released before source termination: %#v", group) + default: + } + + close(chunks) + for range result.Chunks { + } + select { + case group := <-releases: + if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) { + t.Fatalf("release group = %#v", group) + } + case <-time.After(time.Second): + t.Fatal("stream did not release after source termination") + } +} + +func TestAccountedHomeStreamErrorDrainsUntilSourceClosesBeforeRelease(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + releases := make(chan executionregistry.ReleaseGroup, 1) + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group }) + manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{ + ID: "cred-1", Provider: "home-execution", Status: StatusActive, + }}}, registry, 1) + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")} + manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks}) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + if chunk, ok := <-result.Chunks; !ok || string(chunk.Payload) != "initial" { + t.Fatalf("initial chunk = %#v, open = %v", chunk, ok) + } + chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}} + if chunk, ok := <-result.Chunks; !ok || chunk.Err == nil { + t.Fatalf("error chunk = %#v, open = %v", chunk, ok) + } + + sent := make(chan struct{}) + go func() { + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("after-error-1")} + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("after-error-2")} + close(sent) + }() + select { + case <-sent: + case <-time.After(time.Second): + t.Fatal("stream source was not drained after its error chunk") + } + select { + case group := <-releases: + t.Fatalf("stream released while source remained open: %#v", group) + default: + } + select { + case chunk, ok := <-result.Chunks: + t.Fatalf("chunk after error = %#v, open = %v", chunk, ok) + case <-time.After(50 * time.Millisecond): + } + + close(chunks) + for range result.Chunks { + } + select { + case group := <-releases: + if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) { + t.Fatalf("release group = %#v", group) + } + case <-time.After(time.Second): + t.Fatal("stream did not release after the source closed") + } +} + +func TestAccountedHomeStreamErrorCancellationReleasesSelection(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + releases := make(chan executionregistry.ReleaseGroup, 1) + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group }) + manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{ + ID: "cred-1", Provider: "home-execution", Status: StatusActive, + }}}, registry, 1) + chunks := make(chan cliproxyexecutor.StreamChunk, 2) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")} + chunks <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"}} + manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks}) + + result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + if _, ok := <-result.Chunks; !ok { + t.Fatal("stream closed before initial chunk") + } + if chunk, ok := <-result.Chunks; !ok || chunk.Err == nil { + t.Fatalf("error chunk = %#v, open = %v", chunk, ok) + } + select { + case group := <-releases: + t.Fatalf("stream released before cancellation: %#v", group) + default: + } + + cancel() + for range result.Chunks { + } + select { + case group := <-releases: + if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) { + t.Fatalf("release group = %#v", group) + } + case <-time.After(time.Second): + t.Fatal("stream did not release after cancellation") + } + close(chunks) +} + +func TestAccountedHomeStreamConsumerCancellationEndsSelection(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + releases := make(chan executionregistry.ReleaseGroup, 1) + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group }) + manager.PublishHomeDispatch(&accountedHomeExecutionDispatcher{auths: []Auth{{ + ID: "cred-1", Provider: "home-execution", Status: StatusActive, + }}}, registry, 1) + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("initial")} + manager.RegisterExecutor(&homeExecutionStreamExecutor{chunks: chunks}) + + result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + cancel() + for range result.Chunks { + } + select { + case group := <-releases: + if group != (executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "model-a"}) { + t.Fatalf("release group = %#v", group) + } + case <-time.After(time.Second): + t.Fatal("stream did not release after consumer cancellation") + } +} + +type retryingAccountedHomeExecutor struct{ calls atomic.Int32 } + +func (*retryingAccountedHomeExecutor) Identifier() string { return "home-execution" } +func (e *retryingAccountedHomeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if e.calls.Add(1) == 1 { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream failed"} + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} +func (*retryingAccountedHomeExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*retryingAccountedHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*retryingAccountedHomeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*retryingAccountedHomeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestAccountedHomeRetrySelectsAndReleasesEveryAttempt(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + releases := make(chan executionregistry.ReleaseGroup, 2) + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { releases <- group }) + dispatcher := &accountedHomeExecutionDispatcher{auths: []Auth{ + {ID: "cred-1", Provider: "home-execution", Status: StatusActive}, + {ID: "cred-2", Provider: "home-execution", Status: StatusActive}, + }} + manager.PublishHomeDispatch(dispatcher, registry, 1) + executor := &retryingAccountedHomeExecutor{} + manager.RegisterExecutor(executor) + + if _, errExecute := manager.Execute(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{}); errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home selections = %d, want 2", got) + } + if got := executor.calls.Load(); got != 2 { + t.Fatalf("executor attempts = %d, want 2", got) + } + groups := map[executionregistry.ReleaseGroup]bool{} + for range 2 { + groups[<-releases] = true + } + for _, credentialID := range []string{"cred-1", "cred-2"} { + if !groups[executionregistry.ReleaseGroup{CredentialID: credentialID, Model: "model-a"}] { + t.Fatalf("missing release for %s: %#v", credentialID, groups) + } + } +} + +func TestHomeStreamWithoutSourceEndsSelection(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(&homePerSelectionDispatcher{auths: []Auth{{ + ID: "home-auth", Provider: "home-execution", Status: StatusActive, + }}}, registry, 1) + manager.RegisterExecutor(&missingHomeStreamSourceExecutor{}) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-execution"}, cliproxyexecutor.Request{Model: "test"}, cliproxyexecutor.Options{Stream: true}) + if errExecute == nil { + t.Fatalf("ExecuteStream() result = %#v, want error", result) + } + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := registry.Drain(drainCtx); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} diff --git a/sdk/cliproxy/auth/home_fallback_audit_test.go b/sdk/cliproxy/auth/home_fallback_audit_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_fallback_audit_test.go @@ -0,0 +1,54 @@ +package auth + +import ( + "context" + "errors" + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestHomeWebsocketReusesCanonicalModelSelection(t *testing.T) { + dispatcher := &retainingHomeExecutionDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(&retainingHomeExecutionExecutor{}) + t.Cleanup(func() { manager.CloseExecutionSession("canonical-model-session") }) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "canonical-model-session", + cliproxyexecutor.PinnedAuthMetadataKey: "home-auth", + }} + for _, model := range []string{"model-a(high)", "model-a"} { + if _, errExecute := manager.Execute(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: model}, opts); errExecute != nil { + t.Fatalf("Execute(%q) error = %v", model, errExecute) + } + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1 for one credential and canonical model", got) + } +} + +func TestAuditHomeCreditsFailClosed(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.auths["local-credits"] = &Auth{ID: "local-credits", Provider: "antigravity", Status: StatusActive} + + _, _, errExecute := manager.tryAntigravityCreditsExecute(context.Background(), cliproxyexecutor.Request{Model: "claude-test"}, cliproxyexecutor.Options{}) + assertHomeCreditsFallbackUnsupported(t, errExecute) + + _, _, errStream := manager.tryAntigravityCreditsExecuteStream(context.Background(), cliproxyexecutor.Request{Model: "claude-test"}, cliproxyexecutor.Options{Stream: true}) + assertHomeCreditsFallbackUnsupported(t, errStream) +} + +func assertHomeCreditsFallbackUnsupported(t *testing.T, err error) { + t.Helper() + var authErr *Error + if !errors.As(err, &authErr) || authErr.Code != "home_fallback_unsupported" { + t.Fatalf("error = %v, want home_fallback_unsupported", err) + } +} diff --git a/sdk/cliproxy/auth/home_force_mapping_test.go b/sdk/cliproxy/auth/home_force_mapping_test.go --- a/sdk/cliproxy/auth/home_force_mapping_test.go +++ b/sdk/cliproxy/auth/home_force_mapping_test.go @@ -1,6 +1,20 @@ package auth -import "testing" +import ( + "context" + "encoding/json" + "net/http" + "reflect" + "sync" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + internalhome "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) func TestHomeForceMappingAliasResult(t *testing.T) { auth := &Auth{ @@ -15,6 +29,592 @@ result := homeForceMappingAliasResult(auth, "grok-latest") if result.UpstreamModel != "grok-4.5" || !result.ForceMapping || result.OriginalAlias != "grok-latest" { t.Fatalf("homeForceMappingAliasResult() = %+v", result) + } +} + +func TestHomeForceMappingAliasResultRequiresSameOriginalAlias(t *testing.T) { + auth := &Auth{ + Provider: "xai", + Attributes: map[string]string{ + homeUpstreamModelAttributeKey: "grok-4.5", + homeForceMappingAttributeKey: "true", + homeOriginalAliasAttributeKey: "grok-latest", + }, + } + + if result := homeForceMappingAliasResult(auth, " GROK-LATEST "); !result.ForceMapping { + t.Fatalf("homeForceMappingAliasResult() = %+v, want same alias force mapping", result) + } + if result := homeForceMappingAliasResult(auth, "grok-latest(high)"); !result.ForceMapping { + t.Fatalf("homeForceMappingAliasResult() = %+v, want reasoning suffix force mapping", result) + } + if result := homeForceMappingAliasResult(auth, "grok-latest(custom)"); result.ForceMapping || result.OriginalAlias != "" { + t.Fatalf("homeForceMappingAliasResult() = %+v, want no force mapping for a custom suffix", result) + } + if result := homeForceMappingAliasResult(auth, "grok-other"); result.ForceMapping || result.OriginalAlias != "" { + t.Fatalf("homeForceMappingAliasResult() = %+v, want no force mapping for a different alias", result) + } +} + +func TestHomeNonForceAliasSessionReuseAndTargetChangeReleasesAccountedModel(t *testing.T) { + registry := executionregistry.New() + dispatcher := &accountedAliasTargetDispatcher{} + var releases []executionregistry.ReleaseGroup + var releasesMu sync.Mutex + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) { + releasesMu.Lock() + releases = append(releases, group) + releasesMu.Unlock() + dispatcher.releases.Add(1) + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(forceMappingAliasChangeExecutor{}) + t.Cleanup(func() { manager.CloseExecutionSession("non-force-alias-session") }) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "non-force-alias-session", + cliproxyexecutor.PinnedAuthMetadataKey: "non-force-alias-auth", + }} + for _, model := range []string{"alias-a(high)", "alias-a", "alias-b"} { + if _, errExecute := manager.Execute(ctx, []string{"force-mapping"}, cliproxyexecutor.Request{Model: model}, opts); errExecute != nil { + t.Fatalf("Execute(%q) error = %v", model, errExecute) + } + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 for same-route reuse and target change", got) + } + if !dispatcher.releasedBeforeSecondRPop.Load() { + t.Fatal("previous accounted selection was not released before the different-alias redispatch") + } + + manager.CloseExecutionSession("non-force-alias-session") + releasesMu.Lock() + gotReleases := append([]executionregistry.ReleaseGroup(nil), releases...) + releasesMu.Unlock() + wantReleases := []executionregistry.ReleaseGroup{ + {CredentialID: "non-force-alias-auth", Model: "target-a"}, + {CredentialID: "non-force-alias-auth", Model: "target-b"}, + } + if !reflect.DeepEqual(gotReleases, wantReleases) { + t.Fatalf("accounted release groups = %#v, want %#v", gotReleases, wantReleases) + } +} + +type accountedAliasTargetDispatcher struct { + calls atomic.Int32 + releases atomic.Int32 + releasedBeforeSecondRPop atomic.Bool +} + +func (*accountedAliasTargetDispatcher) HeartbeatOK() bool { return true } + +func (d *accountedAliasTargetDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + call := d.calls.Add(1) + if call == 2 { + d.releasedBeforeSecondRPop.Store(d.releases.Load() == 1) + } + target := "target-a" + if canonicalHomeConcurrencyModelKey(model) == "alias-b" { + target = "target-b" + } + return json.Marshal(map[string]any{ + "model": target, + "auth_index": "non-force-alias-auth", + "auth": Auth{ + ID: "non-force-alias-auth", + Provider: "force-mapping", + Status: StatusActive, + Attributes: map[string]string{ + "websockets": "true", + }, + }, + "concurrency": homeConcurrencyTuple{ + Accounted: true, + CredentialID: "non-force-alias-auth", + Model: target, + }, + }) +} + +func (*accountedAliasTargetDispatcher) AbortAmbiguousDispatch() {} + +func TestHomeAuthSelectionRouteRetainsRequestedResponseAliasAcrossWebsocketReuse(t *testing.T) { + registry := executionregistry.New() + dispatcher := &authSelectionAliasDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(authSelectionAliasExecutor{}) + t.Cleanup(func() { manager.CloseExecutionSession("auth-selection-route") }) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.AuthSelectionModelMetadataKey: "route-model", + cliproxyexecutor.RequestedModelMetadataKey: "client-alias", + cliproxyexecutor.ExecutionSessionMetadataKey: "auth-selection-route", + cliproxyexecutor.PinnedAuthMetadataKey: "auth-selection-route-auth", + }} + for attempt := 0; attempt < 2; attempt++ { + response, errExecute := manager.Execute(ctx, []string{"force-mapping"}, cliproxyexecutor.Request{Model: "execution-model"}, opts) + if errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + if got := string(response.Payload); got != `{"model":"client-alias"}` { + t.Fatalf("response = %s, want requested response alias", got) + } + } + if got := dispatcher.Models(); !reflect.DeepEqual(got, []string{"route-model"}) { + t.Fatalf("Home RPOP models = %#v, want canonical auth-selection route", got) + } +} + +type authSelectionAliasDispatcher struct { + mu sync.Mutex + models []string +} + +func (*authSelectionAliasDispatcher) HeartbeatOK() bool { return true } + +func (d *authSelectionAliasDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + d.mu.Lock() + d.models = append(d.models, model) + d.mu.Unlock() + return json.Marshal(map[string]any{ + "model": "target-model", + "force_mapping": true, + "original_alias": "route-model", + "auth_index": "auth-selection-route-auth", + "auth": Auth{ + ID: "auth-selection-route-auth", + Provider: "force-mapping", + Status: StatusActive, + Attributes: map[string]string{ + "websockets": "true", + }, + }, + "concurrency": homeConcurrencyTuple{Accounted: true, CredentialID: "auth-selection-route-auth", Model: "target-model"}, + }) +} + +func (*authSelectionAliasDispatcher) AbortAmbiguousDispatch() {} + +func (d *authSelectionAliasDispatcher) Models() []string { + d.mu.Lock() + defer d.mu.Unlock() + return append([]string(nil), d.models...) +} + +type authSelectionAliasExecutor struct{} + +func (authSelectionAliasExecutor) Identifier() string { return "force-mapping" } +func (authSelectionAliasExecutor) Execute(_ context.Context, _ *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + return cliproxyexecutor.Response{Payload: []byte(`{"model":"` + req.Model + `"}`)}, nil +} +func (authSelectionAliasExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (authSelectionAliasExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (authSelectionAliasExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (authSelectionAliasExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeForceMappingAliasChangeEndsAndFlushesBeforeRedispatch(t *testing.T) { + registry := executionregistry.New() + dispatcher := &forceMappingAliasChangeDispatcher{} + registry.SetReleaseSink(func(executionregistry.ReleaseGroup, int64) { + dispatcher.releases.Add(1) + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(forceMappingAliasChangeExecutor{}) + t.Cleanup(func() { manager.CloseExecutionSession("force-mapping-alias-change") }) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "force-mapping-alias-change", + cliproxyexecutor.PinnedAuthMetadataKey: "force-mapping-auth", + }} + for _, model := range []string{"alias-a", "alias-b"} { + if _, errExecute := manager.Execute(ctx, []string{"force-mapping"}, cliproxyexecutor.Request{Model: model}, opts); errExecute != nil { + t.Fatalf("Execute(%q) error = %v", model, errExecute) + } + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 after original alias changes", got) + } + if !dispatcher.releasedBeforeSecondRPop.Load() { + t.Fatal("previous selection was not ended and released before the second Home RPOP") + } +} + +type forceMappingAliasChangeDispatcher struct { + calls atomic.Int32 + releases atomic.Int32 + releasedBeforeSecondRPop atomic.Bool +} + +func (*forceMappingAliasChangeDispatcher) HeartbeatOK() bool { return true } + +func (d *forceMappingAliasChangeDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, _ int) ([]byte, error) { + if d.calls.Add(1) == 2 { + d.releasedBeforeSecondRPop.Store(d.releases.Load() == 1) + } + return json.Marshal(map[string]any{ + "model": "upstream-a", + "auth_index": "force-mapping-auth", + "auth": Auth{ + ID: "force-mapping-auth", + Provider: "force-mapping", + Status: StatusActive, + Attributes: map[string]string{ + "websockets": "true", + homeForceMappingAttributeKey: "true", + homeOriginalAliasAttributeKey: "alias-a", + }, + }, + "concurrency": homeConcurrencyTuple{ + Accounted: true, + CredentialID: "force-mapping-auth", + Model: "upstream-a", + }, + }) +} + +func (*forceMappingAliasChangeDispatcher) AbortAmbiguousDispatch() {} + +type forceMappingAliasChangeExecutor struct{} + +func (forceMappingAliasChangeExecutor) Identifier() string { return "force-mapping" } +func (forceMappingAliasChangeExecutor) Execute(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} +func (forceMappingAliasChangeExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (forceMappingAliasChangeExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (forceMappingAliasChangeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (forceMappingAliasChangeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestHomeRetainedRouteRewritesReasoningSuffixAndWaitsForReleaseACK(t *testing.T) { + registry := executionregistry.New() + dispatcher := &ackOrderedRouteDispatcher{} + flusher := internalhome.NewReleaseFlusher(func() internalconfig.CredentialConcurrencyConfig { + return internalconfig.CredentialConcurrencyConfig{ + ReleaseFlushInterval: time.Millisecond, + ReleaseMaxBackoff: 10 * time.Millisecond, + } + }, func(_ context.Context, _ internalhome.ConcurrencyReleaseFrame) error { + dispatcher.acks.Add(1) + return nil + }) + registry.SetReleaseSink(flusher.MarkDirty) + releaseCtx, cancelRelease := context.WithCancel(context.Background()) + releaseDone := make(chan struct{}) + go func() { + defer close(releaseDone) + flusher.Run(releaseCtx) + }() + defer func() { + cancelRelease() + <-releaseDone + }() + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + executor := &retainedRouteModelExecutor{} + manager.RegisterExecutor(executor) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "retained-route-ack", + cliproxyexecutor.PinnedAuthMetadataKey: "retained-route-auth", + }} + for _, model := range []string{"alias-a", "alias-a(high)", "alias-a", "alias-a(custom)"} { + response, errExecute := manager.Execute(ctx, []string{"retained-route"}, cliproxyexecutor.Request{Model: model}, opts) + if errExecute != nil { + t.Fatalf("Execute(%q) error = %v", model, errExecute) + } + if got := string(response.Payload); got != `{"model":"`+model+`"}` { + t.Fatalf("Execute(%q) response = %s, want response alias", model, got) + } + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home RPOP calls = %d, want 2 because custom suffix must redispatch", got) + } + if !dispatcher.ackedBeforeSecondRPop.Load() { + t.Fatal("Home release PUSH was not acknowledged before the second RPOP") + } + if got := executor.Models(); !reflect.DeepEqual(got, []string{"target-a", "target-a(high)", "target-a", "target-custom"}) { + t.Fatalf("executor models = %#v", got) + } + + manager.CloseExecutionSession("retained-route-ack") + deadline := time.NewTimer(time.Second) + defer deadline.Stop() + for dispatcher.acks.Load() != 2 { + select { + case <-deadline.C: + t.Fatalf("final release acknowledgements = %d, want 2", dispatcher.acks.Load()) + case <-time.After(time.Millisecond): + } + } +} + +type ackOrderedRouteDispatcher struct { + calls atomic.Int32 + acks atomic.Int32 + ackedBeforeSecondRPop atomic.Bool +} + +func (*ackOrderedRouteDispatcher) HeartbeatOK() bool { return true } + +func (d *ackOrderedRouteDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + call := d.calls.Add(1) + if call == 2 { + d.ackedBeforeSecondRPop.Store(d.acks.Load() == 1) + } + target := "target-a" + if canonicalHomeConcurrencyModelKey(model) != "alias-a" { + target = "target-custom" + } + return json.Marshal(map[string]any{ + "model": target, + "force_mapping": true, + "original_alias": model, + "auth_index": "retained-route-auth", + "auth": Auth{ + ID: "retained-route-auth", + Provider: "retained-route", + Status: StatusActive, + Attributes: map[string]string{ + "websockets": "true", + }, + }, + "concurrency": homeConcurrencyTuple{Accounted: true, CredentialID: "retained-route-auth", Model: target}, + }) +} + +func (*ackOrderedRouteDispatcher) AbortAmbiguousDispatch() {} + +type retainedRouteModelExecutor struct { + mu sync.Mutex + models []string +} + +func (*retainedRouteModelExecutor) Identifier() string { return "retained-route" } +func (e *retainedRouteModelExecutor) Execute(_ context.Context, _ *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.models = append(e.models, req.Model) + e.mu.Unlock() + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + return cliproxyexecutor.Response{Payload: []byte(`{"model":"` + req.Model + `"}`)}, nil +} +func (*retainedRouteModelExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*retainedRouteModelExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (*retainedRouteModelExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*retainedRouteModelExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} +func (e *retainedRouteModelExecutor) Models() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.models...) +} + +func TestHomeRetainedPrefixedRouteRewritesSuffixAndResponse(t *testing.T) { + registry := executionregistry.New() + dispatcher := &prefixedRetainedRouteDispatcher{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + executor := &prefixedRetainedRouteExecutor{} + manager.RegisterExecutor(executor) + t.Cleanup(func() { manager.CloseExecutionSession("prefixed-retained-route") }) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "prefixed-retained-route", + cliproxyexecutor.PinnedAuthMetadataKey: "prefixed-retained-route-auth", + }} + for _, model := range []string{"team/alias-a", "team/alias-a(high)"} { + response, errExecute := manager.Execute(ctx, []string{"prefixed-retained-route"}, cliproxyexecutor.Request{Model: model}, opts) + if errExecute != nil { + t.Fatalf("Execute(%q) error = %v", model, errExecute) + } + if got := string(response.Payload); got != `{"model":"`+model+`"}` { + t.Fatalf("Execute(%q) response = %s, want external response alias", model, got) + } + } + if got := dispatcher.Models(); !reflect.DeepEqual(got, []string{"team/alias-a"}) { + t.Fatalf("Home RPOP models = %#v, want external canonical route only", got) + } + if got := executor.Models(); !reflect.DeepEqual(got, []string{"target-a", "target-a(high)"}) { + t.Fatalf("executor models = %#v, want upstream suffix rewrite", got) + } + manager.mu.RLock() + selection := manager.homeSessionSelections["prefixed-retained-route"][homeSessionSelectionKey{ + credentialID: "prefixed-retained-route-auth", + routeModel: "team/alias-a", + }] + manager.mu.RUnlock() + if selection == nil { + t.Fatal("retained selection missing external route key") + } + retainedAuth := selection.CloneAuthForRoute("team/alias-a(high)") + if got := retainedAuth.Attributes[homeOriginalAliasAttributeKey]; got != "alias-a(high)" { + t.Fatalf("retained original alias = %q, want prefix-stripped alias-a(high)", got) + } +} + +type prefixedRetainedRouteDispatcher struct { + mu sync.Mutex + models []string +} + +func (*prefixedRetainedRouteDispatcher) HeartbeatOK() bool { return true } + +func (d *prefixedRetainedRouteDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) { + d.mu.Lock() + d.models = append(d.models, model) + d.mu.Unlock() + return json.Marshal(map[string]any{ + "model": "target-a", + "force_mapping": true, + "original_alias": "alias-a", + "auth_index": "prefixed-retained-route-auth", + "auth": Auth{ + ID: "prefixed-retained-route-auth", + Provider: "prefixed-retained-route", + Prefix: "team", + Status: StatusActive, + Attributes: map[string]string{ + "websockets": "true", + }, + }, + "concurrency": homeConcurrencyTuple{Accounted: true, CredentialID: "prefixed-retained-route-auth", Model: "target-a"}, + }) +} + +func (*prefixedRetainedRouteDispatcher) AbortAmbiguousDispatch() {} + +func (d *prefixedRetainedRouteDispatcher) Models() []string { + d.mu.Lock() + defer d.mu.Unlock() + return append([]string(nil), d.models...) +} + +type prefixedRetainedRouteExecutor struct { + mu sync.Mutex + models []string +} + +func (*prefixedRetainedRouteExecutor) Identifier() string { return "prefixed-retained-route" } +func (e *prefixedRetainedRouteExecutor) Execute(_ context.Context, _ *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.models = append(e.models, req.Model) + e.mu.Unlock() + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + return cliproxyexecutor.Response{Payload: []byte(`{"model":"` + req.Model + `"}`)}, nil +} +func (*prefixedRetainedRouteExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} +func (*prefixedRetainedRouteExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} +func (*prefixedRetainedRouteExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*prefixedRetainedRouteExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} +func (e *prefixedRetainedRouteExecutor) Models() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.models...) +} +func TestHomeRedispatchStopsWhenReleaseAcknowledgementFails(t *testing.T) { + registry := executionregistry.New() + dispatcher := &ackOrderedRouteDispatcher{} + flusher := internalhome.NewReleaseFlusher(func() internalconfig.CredentialConcurrencyConfig { + return internalconfig.CredentialConcurrencyConfig{ + CPACancelBound: 20 * time.Millisecond, + ReleaseFlushInterval: time.Millisecond, + ReleaseMaxBackoff: time.Millisecond, + } + }, func(context.Context, internalhome.ConcurrencyReleaseFrame) error { + return context.DeadlineExceeded + }) + registry.SetReleaseSink(flusher.MarkDirty) + releaseCtx, cancelRelease := context.WithCancel(context.Background()) + releaseDone := make(chan struct{}) + go func() { + defer close(releaseDone) + flusher.Run(releaseCtx) + }() + defer func() { + cancelRelease() + <-releaseDone + }() + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + CredentialConcurrency: internalconfig.CredentialConcurrencyConfig{ + CPACancelBound: 20 * time.Millisecond, + ReleaseFlushInterval: time.Millisecond, + ReleaseMaxBackoff: time.Millisecond, + }, + }) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(&retainedRouteModelExecutor{}) + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "release-failure", + cliproxyexecutor.PinnedAuthMetadataKey: "retained-route-auth", + }} + if _, errExecute := manager.Execute(ctx, []string{"retained-route"}, cliproxyexecutor.Request{Model: "alias-a"}, opts); errExecute != nil { + t.Fatalf("first Execute() error = %v", errExecute) + } + if _, errExecute := manager.Execute(ctx, []string{"retained-route"}, cliproxyexecutor.Request{Model: "alias-a(custom)"}, opts); errExecute == nil { + t.Fatal("redispatch after unacknowledged release unexpectedly succeeded") + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want no second RPOP after release failure", got) } } diff --git a/sdk/cliproxy/auth/home_in_flight_publisher.go b/sdk/cliproxy/auth/home_in_flight_publisher.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_in_flight_publisher.go @@ -0,0 +1,399 @@ +package auth + +import ( + "context" + "encoding/json" + "sort" + "strings" + "time" + "unicode/utf8" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + log "github.com/sirupsen/logrus" +) + +// HomeInFlightTransport publishes in-flight observation frames for one Home lifetime. +type HomeInFlightTransport interface { + HeartbeatOK() bool + LPushInFlightSnapshot(context.Context, []byte) error +} + +// HomeInFlightPublisherConfig bounds in-flight observation frames. +type HomeInFlightPublisherConfig struct { + SnapshotInterval time.Duration + MaxPartBytes int + MaxPartCount int + MaxRevisionBytes int + MaxAggregateGroups int + MaxDetails int + MaxStringBytes int +} + +type homeInFlightAggregateKey struct { + CredentialID string + Model string + Accounted bool +} + +func homeInFlightStatus(accounted bool) home.InFlightAccountedStatus { + if accounted { + return home.InFlightAccounted + } + return home.InFlightUnaccounted +} + +// HomeInFlightPublisherConfigFromConfig converts validated runtime config into publisher bounds. +func HomeInFlightPublisherConfigFromConfig(cfg internalconfig.CredentialInFlightConfig) (HomeInFlightPublisherConfig, error) { + snapshotInterval, _, _, errDurations := cfg.Durations() + if errDurations != nil { + return HomeInFlightPublisherConfig{}, errDurations + } + if errValidate := cfg.Validate(); errValidate != nil { + return HomeInFlightPublisherConfig{}, errValidate + } + return HomeInFlightPublisherConfig{ + SnapshotInterval: snapshotInterval, + MaxPartBytes: cfg.MaxPartBytes, + MaxPartCount: cfg.MaxPartCount, + MaxRevisionBytes: cfg.MaxRevisionBytes, + MaxAggregateGroups: cfg.MaxAggregateGroups, + MaxDetails: cfg.MaxDetails, + MaxStringBytes: cfg.MaxStringBytes, + }, nil +} + +// ApplyHomeInFlightPublisherConfig stores an immutable validated publisher config snapshot. +func (m *Manager) ApplyHomeInFlightPublisherConfig(cfg HomeInFlightPublisherConfig) { + if m == nil || !validHomeInFlightPublisherConfig(cfg) { + return + } + snapshot := cfg + m.homeInFlightPublisherConfig.Store(&snapshot) +} + +// HomeInFlightPublisherConfig returns the current immutable publisher config snapshot. +func (m *Manager) HomeInFlightPublisherConfig() HomeInFlightPublisherConfig { + if m == nil { + return HomeInFlightPublisherConfig{} + } + cfg := m.homeInFlightPublisherConfig.Load() + if cfg == nil { + return HomeInFlightPublisherConfig{} + } + return *cfg +} + +func validHomeInFlightPublisherConfig(cfg HomeInFlightPublisherConfig) bool { + if cfg.SnapshotInterval <= 0 || cfg.MaxPartBytes < 1024 || cfg.MaxPartCount <= 0 || cfg.MaxPartCount > internalconfig.DefaultInFlightMaxPartCount || + cfg.MaxRevisionBytes < cfg.MaxPartBytes || cfg.MaxRevisionBytes > internalconfig.DefaultInFlightMaxRevisionBytes || + cfg.MaxAggregateGroups <= 0 || cfg.MaxAggregateGroups > internalconfig.DefaultInFlightMaxAggregateGroups || + cfg.MaxDetails < 0 || cfg.MaxDetails > internalconfig.DefaultInFlightMaxDetails || + cfg.MaxStringBytes <= 0 || cfg.MaxStringBytes > internalconfig.DefaultInFlightMaxStringBytes { + return false + } + return (cfg.MaxRevisionBytes+cfg.MaxPartBytes-1)/cfg.MaxPartBytes <= cfg.MaxPartCount +} + +func validHomeInFlightPublisherBounds(cfg HomeInFlightPublisherConfig) bool { + return cfg.MaxPartBytes > 0 && cfg.MaxPartCount > 0 && cfg.MaxRevisionBytes >= cfg.MaxPartBytes && + cfg.MaxAggregateGroups > 0 && cfg.MaxDetails >= 0 && cfg.MaxStringBytes > 0 +} + +// StartHomeInFlightPublisher publishes periodic snapshots for the supplied lifetime registry. +func (m *Manager) StartHomeInFlightPublisher(ctx context.Context, transport HomeInFlightTransport, registry *executionregistry.Registry) { + if m == nil || transport == nil || registry == nil { + return + } + if ctx == nil { + ctx = context.Background() + } + + timer := time.NewTimer(0) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case observedAt := <-timer.C: + cfg := m.HomeInFlightPublisherConfig() + interval := cfg.SnapshotInterval + if interval <= 0 { + interval = 2 * time.Second + } + timer.Reset(interval) + if !transport.HeartbeatOK() { + continue + } + freeze := registry.FreezeInFlight(observedAt.UTC()) + frames := encodeHomeInFlightFreeze(freeze, observedAt.UTC(), cfg) + for index := range frames { + raw, errMarshal := json.Marshal(frames[index]) + if errMarshal != nil { + log.Warn("failed to encode in-flight snapshot frame") + break + } + if errPush := transport.LPushInFlightSnapshot(ctx, raw); errPush != nil { + log.Warn("failed to publish in-flight snapshot frame") + break + } + } + } + } +} + +func encodeHomeInFlightFreeze(freeze executionregistry.Freeze, observedAt time.Time, cfg HomeInFlightPublisherConfig) []home.InFlightSnapshotFrame { + observedAt = observedAt.UTC() + aggregateCounts := make(map[homeInFlightAggregateKey]int64, len(freeze.Executions)) + aggregateKeysValid := true + for _, observation := range freeze.Executions { + key := homeInFlightAggregateKey{ + CredentialID: observation.CredentialID, + Model: homeInFlightObservationModel(observation), + Accounted: observation.Accounted, + } + if len(key.CredentialID) > cfg.MaxStringBytes || len(key.Model) > cfg.MaxStringBytes { + aggregateKeysValid = false + } + aggregateCounts[key]++ + } + aggregates := make([]home.InFlightAggregate, 0, len(aggregateCounts)) + for key, count := range aggregateCounts { + aggregates = append(aggregates, home.InFlightAggregate{ + CredentialID: key.CredentialID, + Model: key.Model, + Status: homeInFlightStatus(key.Accounted), + Count: count, + }) + } + sort.Slice(aggregates, func(left, right int) bool { + if aggregates[left].CredentialID != aggregates[right].CredentialID { + return aggregates[left].CredentialID < aggregates[right].CredentialID + } + if aggregates[left].Model != aggregates[right].Model { + return aggregates[left].Model < aggregates[right].Model + } + return aggregates[left].Status < aggregates[right].Status + }) + if !validHomeInFlightPublisherBounds(cfg) || !aggregateKeysValid || len(aggregates) > cfg.MaxAggregateGroups { + return homeInFlightOverflow(freeze, observedAt, len(aggregates)) + } + + details := make([]home.InFlightRequestDetail, 0, len(freeze.Executions)) + detailsTruncated := false + for _, observation := range freeze.Executions { + detail, bounded := homeInFlightBoundDetail(home.InFlightRequestDetail{ + RequestID: observation.RequestID, + CredentialID: observation.CredentialID, + Model: homeInFlightObservationModel(observation), + RequestKind: observation.RequestKind, + StartedAt: observation.StartedAt.UTC(), + }, cfg.MaxStringBytes) + if !validHomeInFlightDetail(detail, cfg.MaxStringBytes) { + detailsTruncated = true + continue + } + detailsTruncated = detailsTruncated || bounded + details = append(details, detail) + } + sort.Slice(details, func(left, right int) bool { + if !details[left].StartedAt.Equal(details[right].StartedAt) { + return details[left].StartedAt.Before(details[right].StartedAt) + } + if details[left].RequestID != details[right].RequestID { + return details[left].RequestID < details[right].RequestID + } + if details[left].CredentialID != details[right].CredentialID { + return details[left].CredentialID < details[right].CredentialID + } + if details[left].Model != details[right].Model { + return details[left].Model < details[right].Model + } + return details[left].RequestKind < details[right].RequestKind + }) + + if len(details) > cfg.MaxDetails { + details = details[:cfg.MaxDetails] + detailsTruncated = true + } + + for { + frames, aggregatesPacked, includedDetails := packHomeInFlightFrames(freeze, observedAt, cfg, aggregates, details, detailsTruncated) + if !aggregatesPacked { + return homeInFlightOverflow(freeze, observedAt, len(aggregates)) + } + if includedDetails < len(details) { + details = details[:includedDetails] + detailsTruncated = true + continue + } + if homeInFlightFramesWithinBounds(frames, cfg) { + return frames + } + if len(details) == 0 { + return homeInFlightOverflow(freeze, observedAt, len(aggregates)) + } + details = details[:len(details)-1] + detailsTruncated = true + } +} + +func homeInFlightObservationModel(observation executionregistry.Observation) string { + if observation.Accounted { + return observation.Model + } + if model, valid := validCanonicalHomeConcurrencyModelKey(observation.Model); valid { + return model + } + return "unknown" +} + +func validHomeInFlightDetail(detail home.InFlightRequestDetail, maxStringBytes int) bool { + validString := func(value string) bool { + return utf8.ValidString(value) && strings.TrimSpace(value) != "" && len(value) <= maxStringBytes + } + return validString(detail.RequestID) && validString(detail.CredentialID) && validString(detail.Model) && validString(detail.RequestKind) && !detail.StartedAt.IsZero() && detail.StartedAt.Location() == time.UTC +} + +func homeInFlightBoundDetail(detail home.InFlightRequestDetail, maxBytes int) (home.InFlightRequestDetail, bool) { + truncated := false + bound := func(value string) string { + bounded := homeInFlightTruncateString(value, maxBytes) + truncated = truncated || bounded != value + return bounded + } + detail.RequestID = bound(detail.RequestID) + detail.CredentialID = bound(detail.CredentialID) + detail.Model = bound(detail.Model) + detail.RequestKind = bound(detail.RequestKind) + return detail, truncated +} + +func homeInFlightTruncateString(value string, maxBytes int) string { + if maxBytes <= 0 || len(value) <= maxBytes { + return value + } + value = value[:maxBytes] + for len(value) > 0 && !utf8.ValidString(value) { + value = value[:len(value)-1] + } + return value +} + +func packHomeInFlightFrames(freeze executionregistry.Freeze, observedAt time.Time, cfg HomeInFlightPublisherConfig, aggregates []home.InFlightAggregate, details []home.InFlightRequestDetail, detailsTruncated bool) ([]home.InFlightSnapshotFrame, bool, int) { + frames := make([]home.InFlightSnapshotFrame, 0, cfg.MaxPartCount) + current := homeInFlightPartFrame(freeze, observedAt, cfg.MaxPartCount, detailsTruncated) + appendCurrent := func() bool { + if len(frames) >= cfg.MaxPartCount { + return false + } + frames = append(frames, current) + current = homeInFlightPartFrame(freeze, observedAt, cfg.MaxPartCount, detailsTruncated) + return true + } + for _, aggregate := range aggregates { + candidate := current + candidate.Aggregates = append(candidate.Aggregates, aggregate) + if homeInFlightFrameWithinPartLimit(candidate, cfg.MaxPartBytes) { + current = candidate + continue + } + if len(current.Aggregates) == 0 && len(current.Details) == 0 { + return nil, false, 0 + } + if !appendCurrent() { + return nil, false, 0 + } + candidate = current + candidate.Aggregates = append(candidate.Aggregates, aggregate) + if !homeInFlightFrameWithinPartLimit(candidate, cfg.MaxPartBytes) { + return nil, false, 0 + } + current = candidate + } + + includedDetails := 0 + for _, detail := range details { + candidate := current + candidate.Details = append(candidate.Details, detail) + if homeInFlightFrameWithinPartLimit(candidate, cfg.MaxPartBytes) { + current = candidate + includedDetails++ + continue + } + if len(current.Aggregates) == 0 && len(current.Details) == 0 { + return frames, true, includedDetails + } + if !appendCurrent() { + return frames, true, includedDetails - len(current.Details) + } + candidate = current + candidate.Details = append(candidate.Details, detail) + if !homeInFlightFrameWithinPartLimit(candidate, cfg.MaxPartBytes) { + return frames, true, includedDetails + } + current = candidate + includedDetails++ + } + if len(current.Aggregates) != 0 || len(current.Details) != 0 || len(frames) == 0 { + if !appendCurrent() { + if len(current.Aggregates) != 0 { + return nil, false, 0 + } + return frames, true, includedDetails - len(current.Details) + } + } + for index := range frames { + partIndex, partCount := index, len(frames) + frames[index].PartIndex = &partIndex + frames[index].PartCount = &partCount + } + return frames, true, includedDetails +} + +func homeInFlightPartFrame(freeze executionregistry.Freeze, observedAt time.Time, partCount int, detailsTruncated bool) home.InFlightSnapshotFrame { + partIndex := 0 + return home.InFlightSnapshotFrame{ + Kind: home.InFlightFramePart, + Revision: freeze.Revision, + ObservedAt: observedAt, + BarrierRevision: freeze.BarrierRevision, + PartIndex: &partIndex, + PartCount: &partCount, + DetailsTruncated: detailsTruncated, + } +} + +func homeInFlightFrameWithinPartLimit(frame home.InFlightSnapshotFrame, maxPartBytes int) bool { + raw, errMarshal := json.Marshal(frame) + return errMarshal == nil && len(raw) <= maxPartBytes +} + +func homeInFlightFramesWithinBounds(frames []home.InFlightSnapshotFrame, cfg HomeInFlightPublisherConfig) bool { + if len(frames) == 0 || len(frames) > cfg.MaxPartCount { + return false + } + totalBytes := 0 + for _, frame := range frames { + raw, errMarshal := json.Marshal(frame) + if errMarshal != nil || len(raw) > cfg.MaxPartBytes { + return false + } + totalBytes += len(raw) + if totalBytes > cfg.MaxRevisionBytes { + return false + } + } + return true +} + +func homeInFlightOverflow(freeze executionregistry.Freeze, observedAt time.Time, aggregateGroupCount int) []home.InFlightSnapshotFrame { + return []home.InFlightSnapshotFrame{{ + Kind: home.InFlightFrameOverflow, + Revision: freeze.Revision, + ObservedAt: observedAt, + BarrierRevision: freeze.BarrierRevision, + AggregateGroupCount: aggregateGroupCount, + }} +} diff --git a/sdk/cliproxy/auth/home_in_flight_publisher_test.go b/sdk/cliproxy/auth/home_in_flight_publisher_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_in_flight_publisher_test.go @@ -0,0 +1,476 @@ +package auth + +import ( + "context" + "encoding/json" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestEncodeHomeInFlightFreezePreservesPartitionsAndBarrier(t *testing.T) { + freeze := executionregistry.Freeze{ + Revision: 9, + BarrierRevision: 14, + Executions: []executionregistry.Observation{ + {RequestID: "req-a", CredentialID: "cred", Model: "gpt-5", RequestKind: "http", StartedAt: time.Unix(10, 0).UTC(), Accounted: true}, + {RequestID: "req-b", CredentialID: "cred", Model: "gpt-5", RequestKind: "sse", StartedAt: time.Unix(11, 0).UTC(), Accounted: false}, + }, + } + frames := encodeHomeInFlightFreeze(freeze, time.Unix(12, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 1024, MaxPartCount: 64, MaxRevisionBytes: 16384, + MaxAggregateGroups: 100000, MaxDetails: 1, MaxStringBytes: 256, + }) + if len(frames) != 1 || frames[0].Kind != home.InFlightFramePart { + t.Fatalf("frames = %#v", frames) + } + if frames[0].BarrierRevision != 14 || !frames[0].DetailsTruncated { + t.Fatalf("metadata = %#v", frames[0]) + } + if got := frames[0].Aggregates; len(got) != 2 || got[0].Count != 1 || got[1].Count != 1 { + t.Fatalf("aggregates = %#v", got) + } +} + +func TestEncodeHomeInFlightFreezeUsesOverflowWithoutPartialAggregates(t *testing.T) { + freeze := executionregistry.Freeze{Revision: 10, BarrierRevision: 15, Executions: []executionregistry.Observation{ + {CredentialID: "a", Model: "m1", RequestKind: "http", Accounted: false}, + {CredentialID: "b", Model: "m2", RequestKind: "http", Accounted: true}, + }} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 256, MaxPartCount: 1, MaxRevisionBytes: 256, + MaxAggregateGroups: 1, MaxDetails: 0, MaxStringBytes: 256, + }) + if len(frames) != 1 || frames[0].Kind != home.InFlightFrameOverflow || frames[0].AggregateGroupCount != 2 { + t.Fatalf("frames = %#v", frames) + } + if len(frames[0].Aggregates) != 0 || len(frames[0].Details) != 0 { + t.Fatalf("overflow leaked partial data: %#v", frames[0]) + } + if frames[0].PartIndex != nil || frames[0].PartCount != nil { + t.Fatalf("overflow contains part metadata: %#v", frames[0]) + } +} + +func TestEncodeHomeInFlightFreezeUsesDeterministicBoundedMultipartFrames(t *testing.T) { + freeze := executionregistry.Freeze{Revision: 4, Executions: []executionregistry.Observation{ + {RequestID: "req-c", CredentialID: "cred", Model: "model", RequestKind: "http", StartedAt: time.Unix(12, 0).UTC()}, + {RequestID: "req-a", CredentialID: "cred", Model: "model", RequestKind: "http", StartedAt: time.Unix(10, 0).UTC()}, + {RequestID: "req-b", CredentialID: "cred", Model: "model", RequestKind: "http", StartedAt: time.Unix(11, 0).UTC()}, + }} + cfg := HomeInFlightPublisherConfig{ + MaxPartBytes: 300, MaxPartCount: 8, MaxRevisionBytes: 2048, + MaxAggregateGroups: 8, MaxDetails: 3, MaxStringBytes: 256, + } + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), cfg) + if len(frames) < 2 { + t.Fatalf("frames = %#v, want multipart", frames) + } + for index, frame := range frames { + raw, errMarshal := json.Marshal(frame) + if errMarshal != nil { + t.Fatal(errMarshal) + } + if len(raw) > cfg.MaxPartBytes || frame.PartIndex == nil || frame.PartCount == nil || *frame.PartIndex != index || *frame.PartCount != len(frames) { + t.Fatalf("frame %d = %s", index, raw) + } + } + requestIDs := make([]string, 0, 3) + for _, frame := range frames { + for _, detail := range frame.Details { + requestIDs = append(requestIDs, detail.RequestID) + } + } + if strings.Join(requestIDs, ",") != "req-a,req-b,req-c" { + t.Fatalf("details are not sorted: %#v", frames) + } +} + +func TestEncodeHomeInFlightFreezeOverflowsWhenFinalAggregatePartExceedsPartCount(t *testing.T) { + freeze := executionregistry.Freeze{Revision: 13, Executions: []executionregistry.Observation{ + {CredentialID: strings.Repeat("a", 300), Model: strings.Repeat("a", 300), Accounted: true}, + {CredentialID: strings.Repeat("b", 300), Model: strings.Repeat("b", 300), Accounted: true}, + {CredentialID: strings.Repeat("c", 300), Model: strings.Repeat("c", 300), Accounted: true}, + }} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 1024, MaxPartCount: 2, MaxRevisionBytes: 2048, + MaxAggregateGroups: 3, MaxDetails: 0, MaxStringBytes: 512, + }) + if len(frames) != 1 || frames[0].Kind != home.InFlightFrameOverflow || frames[0].AggregateGroupCount != 3 { + t.Fatalf("frames = %#v", frames) + } + if len(frames[0].Aggregates) != 0 || len(frames[0].Details) != 0 { + t.Fatalf("overflow leaked aggregate prefix: %#v", frames[0]) + } +} + +func TestEncodeHomeInFlightFreezeTruncatesDetailsBeforeTotalOverflow(t *testing.T) { + freeze := executionregistry.Freeze{Revision: 12} + for index := 0; index < 5; index++ { + freeze.Executions = append(freeze.Executions, executionregistry.Observation{ + RequestID: strings.Repeat(string(rune('a'+index)), 60), CredentialID: "cred", Model: "model", RequestKind: "http", + StartedAt: time.Unix(int64(index), 0).UTC(), + }) + } + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 512, MaxPartCount: 8, MaxRevisionBytes: 1000, + MaxAggregateGroups: 8, MaxDetails: 5, MaxStringBytes: 128, + }) + if len(frames) == 1 && frames[0].Kind == home.InFlightFrameOverflow { + t.Fatalf("details overflowed complete aggregates: %#v", frames) + } + if !frames[0].DetailsTruncated || len(frames[0].Aggregates) != 1 { + t.Fatalf("frames = %#v", frames) + } +} + +func TestEncodeHomeInFlightFreezeBoundsStringsAndExcludesSensitiveFields(t *testing.T) { + freeze := executionregistry.Freeze{Revision: 3, Executions: []executionregistry.Observation{{ + RequestID: strings.Repeat("request", 20), CredentialID: strings.Repeat("credential", 20), + Model: strings.Repeat("model", 20), RequestKind: strings.Repeat("kind", 20), + }}} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 1024, MaxPartCount: 2, MaxRevisionBytes: 2048, + MaxAggregateGroups: 8, MaxDetails: 1, MaxStringBytes: 8, + }) + raw, errMarshal := json.Marshal(frames) + if errMarshal != nil { + t.Fatal(errMarshal) + } + if strings.Contains(string(raw), "credentialcredential") || strings.Contains(string(raw), "token") { + t.Fatalf("snapshot leaked unbounded or sensitive data: %s", raw) + } +} + +func TestHomeInFlightPublisherConfigFromConfigValidatesAndUpdates(t *testing.T) { + cfg := internalconfig.DefaultCredentialInFlightConfig() + cfg.SnapshotInterval = "25ms" + publisherCfg, errConfig := HomeInFlightPublisherConfigFromConfig(cfg) + if errConfig != nil || publisherCfg.SnapshotInterval != 25*time.Millisecond { + t.Fatalf("config = %#v, error = %v", publisherCfg, errConfig) + } + + manager := NewManager(nil, nil, nil) + manager.ApplyHomeInFlightPublisherConfig(publisherCfg) + if got := manager.HomeInFlightPublisherConfig(); got.SnapshotInterval != 25*time.Millisecond { + t.Fatalf("manager config = %#v", got) + } +} + +type homeInFlightTransportStub struct { + heartbeat bool + payloads chan []byte +} + +func (t *homeInFlightTransportStub) HeartbeatOK() bool { return t.heartbeat } +func (t *homeInFlightTransportStub) LPushInFlightSnapshot(_ context.Context, payload []byte) error { + t.payloads <- append([]byte(nil), payload...) + return nil +} + +func TestHomeInFlightPublisherPinsLifetimeRegistry(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.ApplyHomeInFlightPublisherConfig(HomeInFlightPublisherConfig{SnapshotInterval: time.Hour, MaxPartBytes: 1024, MaxPartCount: 1, MaxRevisionBytes: 1024, MaxAggregateGroups: 1, MaxDetails: 0, MaxStringBytes: 8}) + registry := executionregistry.New() + registry.ObserveBarrier(14) + transport := &homeInFlightTransportStub{heartbeat: true, payloads: make(chan []byte, 1)} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go manager.StartHomeInFlightPublisher(ctx, transport, registry) + + select { + case raw := <-transport.payloads: + var frame home.InFlightSnapshotFrame + if errUnmarshal := json.Unmarshal(raw, &frame); errUnmarshal != nil { + t.Fatal(errUnmarshal) + } + if frame.BarrierRevision != 14 { + t.Fatalf("frame = %#v", frame) + } + case <-time.After(time.Second): + t.Fatal("publisher did not send lifetime snapshot") + } +} + +type homeInFlightModelDispatcher struct{} + +func (homeInFlightModelDispatcher) HeartbeatOK() bool { return true } +func (homeInFlightModelDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + return json.Marshal(homeAuthDispatchResponse{ + Model: "final-upstream-model", + Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive}, + }) +} +func (homeInFlightModelDispatcher) AbortAmbiguousDispatch() {} + +func TestHomeInFlightObservationUsesFinalDispatchModel(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.PublishHomeDispatch(homeInFlightModelDispatcher{}, registry, 1) + manager.RegisterExecutor(&homeExecutionExecutor{}) + + selection, errSelection := manager.pickHomeDispatchSelection(context.Background(), "requested-model", cliproxyexecutor.Options{}) + if errSelection != nil { + t.Fatalf("pickHomeDispatchSelection() error = %v", errSelection) + } + defer selection.End("test_complete") + + freeze := registry.FreezeInFlight(time.Now()) + if len(freeze.Executions) != 1 || freeze.Executions[0].Model != "final-upstream-model" { + t.Fatalf("observation = %#v", freeze.Executions) + } +} + +func TestEncodeHomeInFlightFreezeOverflowsForRawAggregateKey(t *testing.T) { + freeze := executionregistry.Freeze{Executions: []executionregistry.Observation{{ + CredentialID: "credential-id-exceeds-limit", Model: "model", RequestKind: "http", + }}} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 1024, MaxPartCount: 2, MaxRevisionBytes: 2048, + MaxAggregateGroups: 2, MaxDetails: 1, MaxStringBytes: 8, + }) + if len(frames) != 1 || frames[0].Kind != home.InFlightFrameOverflow || frames[0].AggregateGroupCount != 1 { + t.Fatalf("frames = %#v", frames) + } +} + +func TestEncodeHomeInFlightFreezeKeepsRawAggregateGroupsDistinct(t *testing.T) { + freeze := executionregistry.Freeze{Executions: []executionregistry.Observation{ + {CredentialID: "credential-a", Model: "model", RequestKind: "http"}, + {CredentialID: "credential-b", Model: "model", RequestKind: "http"}, + }} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 1024, MaxPartCount: 2, MaxRevisionBytes: 2048, + MaxAggregateGroups: 1, MaxDetails: 0, MaxStringBytes: 8, + }) + if len(frames) != 1 || frames[0].Kind != home.InFlightFrameOverflow || frames[0].AggregateGroupCount != 2 { + t.Fatalf("frames = %#v", frames) + } +} + +func TestEncodeHomeInFlightFreezeDropsInvalidDetailsWithoutDiscardingAggregates(t *testing.T) { + freeze := executionregistry.Freeze{Revision: 21, Executions: []executionregistry.Observation{ + {RequestID: "", CredentialID: "cred-a", Model: "model-a", RequestKind: "http", StartedAt: time.Unix(1, 0).UTC()}, + {RequestID: "request-b", CredentialID: "cred-a", Model: "model-a", RequestKind: "http", StartedAt: time.Unix(2, 0).UTC()}, + }} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 1024, MaxPartCount: 2, MaxRevisionBytes: 2048, + MaxAggregateGroups: 2, MaxDetails: 2, MaxStringBytes: 64, + }) + if len(frames) != 1 || frames[0].Kind != home.InFlightFramePart { + t.Fatalf("frames = %#v, want one part", frames) + } + if !frames[0].DetailsTruncated || len(frames[0].Aggregates) != 1 || frames[0].Aggregates[0].Count != 2 { + t.Fatalf("frame = %#v, want preserved aggregate and truncated details", frames[0]) + } + if len(frames[0].Details) != 1 || frames[0].Details[0].RequestID != "request-b" { + t.Fatalf("details = %#v, want only valid request-b", frames[0].Details) + } +} + +func TestEncodeHomeInFlightFreezeCanonicalizesUnaccountedModelsWithFallback(t *testing.T) { + freeze := executionregistry.Freeze{Revision: 22, Executions: []executionregistry.Observation{ + {RequestID: "request-a", CredentialID: "cred-a", Model: "GPT-5(HIGH)", RequestKind: "http", StartedAt: time.Unix(1, 0).UTC()}, + {RequestID: "request-b", CredentialID: "cred-b", Model: " ", RequestKind: "http", StartedAt: time.Unix(2, 0).UTC()}, + }} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 1024, MaxPartCount: 2, MaxRevisionBytes: 2048, + MaxAggregateGroups: 3, MaxDetails: 2, MaxStringBytes: 64, + }) + if len(frames) != 1 || frames[0].Kind != home.InFlightFramePart { + t.Fatalf("frames = %#v, want one part", frames) + } + models := make([]string, 0, len(frames[0].Aggregates)) + for _, aggregate := range frames[0].Aggregates { + models = append(models, aggregate.Model) + } + if strings.Join(models, ",") != "gpt-5,unknown" { + t.Fatalf("aggregate models = %v, want canonical valid models", models) + } + if frames[0].Details[0].Model != "gpt-5" || frames[0].Details[1].Model != "unknown" { + t.Fatalf("detail models = %#v, want canonical valid models", frames[0].Details) + } +} + +func TestEncodeHomeInFlightFreezeSetsGlobalDetailTruncationMetadata(t *testing.T) { + freeze := executionregistry.Freeze{Executions: []executionregistry.Observation{ + {RequestID: strings.Repeat("r", 32), CredentialID: "cred-a", Model: "model-a", RequestKind: "http", StartedAt: time.Unix(1, 0)}, + {RequestID: "request-b", CredentialID: "cred-b", Model: "model-b", RequestKind: "http", StartedAt: time.Unix(2, 0)}, + {RequestID: "request-c", CredentialID: "cred-c", Model: "model-c", RequestKind: "http", StartedAt: time.Unix(3, 0)}, + }} + frames := encodeHomeInFlightFreeze(freeze, time.Unix(20, 0).UTC(), HomeInFlightPublisherConfig{ + MaxPartBytes: 300, MaxPartCount: 8, MaxRevisionBytes: 2048, + MaxAggregateGroups: 4, MaxDetails: 2, MaxStringBytes: 8, + }) + if len(frames) < 2 { + t.Fatalf("frames = %#v, want multipart", frames) + } + for index, frame := range frames { + if !frame.DetailsTruncated { + t.Fatalf("frame %d missing global truncation metadata: %#v", index, frame) + } + } +} + +type homeInFlightPublisherPayload struct { + observedAt time.Time + raw []byte +} + +type homeInFlightLifecycleTransport struct { + heartbeat atomic.Bool + payloads chan homeInFlightPublisherPayload +} + +func newHomeInFlightLifecycleTransport(heartbeat bool) *homeInFlightLifecycleTransport { + transport := &homeInFlightLifecycleTransport{payloads: make(chan homeInFlightPublisherPayload, 32)} + transport.heartbeat.Store(heartbeat) + return transport +} + +func (t *homeInFlightLifecycleTransport) HeartbeatOK() bool { return t.heartbeat.Load() } +func (t *homeInFlightLifecycleTransport) LPushInFlightSnapshot(_ context.Context, raw []byte) error { + t.payloads <- homeInFlightPublisherPayload{observedAt: time.Now(), raw: append([]byte(nil), raw...)} + return nil +} + +func homeInFlightPublisherTestConfig(interval time.Duration) HomeInFlightPublisherConfig { + return HomeInFlightPublisherConfig{ + SnapshotInterval: interval, MaxPartBytes: 1024, MaxPartCount: 2, MaxRevisionBytes: 2048, + MaxAggregateGroups: 2, MaxDetails: 1, MaxStringBytes: 32, + } +} + +func waitForHomeInFlightPublisherPayload(t *testing.T, payloads <-chan homeInFlightPublisherPayload) homeInFlightPublisherPayload { + t.Helper() + select { + case payload := <-payloads: + return payload + case <-time.After(time.Second): + t.Fatal("publisher did not send a payload") + return homeInFlightPublisherPayload{} + } +} + +func TestHomeInFlightPublisherSkipsFreezeAndPublishWithoutHeartbeat(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.ApplyHomeInFlightPublisherConfig(homeInFlightPublisherTestConfig(10 * time.Millisecond)) + registry := executionregistry.New() + transport := newHomeInFlightLifecycleTransport(false) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + manager.StartHomeInFlightPublisher(ctx, transport, registry) + close(done) + }() + time.Sleep(30 * time.Millisecond) + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("publisher did not exit after cancellation") + } + select { + case published := <-transport.payloads: + t.Fatalf("publisher sent payload without heartbeat at %v", published.observedAt) + default: + } + if freeze := registry.FreezeInFlight(time.Now()); freeze.Revision != 1 { + t.Fatalf("publisher froze registry without heartbeat: %#v", freeze) + } +} + +func TestHomeInFlightPublisherCancellationExits(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.ApplyHomeInFlightPublisherConfig(homeInFlightPublisherTestConfig(time.Hour)) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + manager.StartHomeInFlightPublisher(ctx, newHomeInFlightLifecycleTransport(false), executionregistry.New()) + close(done) + }() + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("publisher did not exit after cancellation") + } +} + +func TestHomeInFlightPublisherReplacementStopsOldLifetimeAndPinsDependencies(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.ApplyHomeInFlightPublisherConfig(homeInFlightPublisherTestConfig(10 * time.Millisecond)) + oldRegistry := executionregistry.New() + oldRegistry.ObserveBarrier(11) + oldTransport := newHomeInFlightLifecycleTransport(true) + oldCtx, cancelOld := context.WithCancel(context.Background()) + oldDone := make(chan struct{}) + go func() { + manager.StartHomeInFlightPublisher(oldCtx, oldTransport, oldRegistry) + close(oldDone) + }() + oldPayload := waitForHomeInFlightPublisherPayload(t, oldTransport.payloads) + var oldFrame home.InFlightSnapshotFrame + if errUnmarshal := json.Unmarshal(oldPayload.raw, &oldFrame); errUnmarshal != nil || oldFrame.BarrierRevision != 11 { + t.Fatalf("old publisher frame = %#v, error = %v", oldFrame, errUnmarshal) + } + cancelOld() + select { + case <-oldDone: + case <-time.After(time.Second): + t.Fatal("old publisher did not stop") + } + + newRegistry := executionregistry.New() + newRegistry.ObserveBarrier(22) + newTransport := newHomeInFlightLifecycleTransport(true) + newCtx, cancelNew := context.WithCancel(context.Background()) + defer cancelNew() + go manager.StartHomeInFlightPublisher(newCtx, newTransport, newRegistry) + newPayload := waitForHomeInFlightPublisherPayload(t, newTransport.payloads) + var newFrame home.InFlightSnapshotFrame + if errUnmarshal := json.Unmarshal(newPayload.raw, &newFrame); errUnmarshal != nil || newFrame.BarrierRevision != 22 { + t.Fatalf("new publisher frame = %#v, error = %v", newFrame, errUnmarshal) + } + time.Sleep(30 * time.Millisecond) + select { + case published := <-oldTransport.payloads: + t.Fatalf("replaced publisher sent payload at %v", published.observedAt) + default: + } + + freeze := newRegistry.FreezeInFlight(time.Now()) + if freeze.BarrierRevision != 22 { + t.Fatalf("new publisher did not use replacement registry: %#v", freeze) + } +} + +func TestHomeInFlightPublisherAppliesConfigUpdateAtNextTimerCycle(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.ApplyHomeInFlightPublisherConfig(homeInFlightPublisherTestConfig(60 * time.Millisecond)) + transport := newHomeInFlightLifecycleTransport(true) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go manager.StartHomeInFlightPublisher(ctx, transport, executionregistry.New()) + waitForHomeInFlightPublisherPayload(t, transport.payloads) + + manager.ApplyHomeInFlightPublisherConfig(homeInFlightPublisherTestConfig(10 * time.Millisecond)) + select { + case published := <-transport.payloads: + t.Fatalf("publisher applied hot interval before the next timer cycle at %v", published) + case <-time.After(30 * time.Millisecond): + } + second := waitForHomeInFlightPublisherPayload(t, transport.payloads) + third := waitForHomeInFlightPublisherPayload(t, transport.payloads) + if elapsed := third.observedAt.Sub(second.observedAt); elapsed > 35*time.Millisecond { + t.Fatalf("publisher interval after update = %v, want <= 35ms", elapsed) + } +} diff --git a/sdk/cliproxy/auth/home_retry_loop_test.go b/sdk/cliproxy/auth/home_retry_loop_test.go --- a/sdk/cliproxy/auth/home_retry_loop_test.go +++ b/sdk/cliproxy/auth/home_retry_loop_test.go @@ -9,6 +9,7 @@ "time" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) @@ -32,6 +33,8 @@ }) return raw, nil } + +func (*repeatedHomeAuthDispatcher) AbortAmbiguousDispatch() {} type unauthorizedHomeExecutor struct { calls atomic.Int32 @@ -75,6 +78,7 @@ executor := &unauthorizedHomeExecutor{} manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(executionregistry.New()) manager.RegisterExecutor(executor) ctx, cancel := context.WithTimeout(context.Background(), time.Second) diff --git a/sdk/cliproxy/auth/home_selected_auth_callback_test.go b/sdk/cliproxy/auth/home_selected_auth_callback_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_selected_auth_callback_test.go @@ -0,0 +1,97 @@ +package auth + +import ( + "context" + "encoding/json" + "net/http" + "sync/atomic" + "testing" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type selectedAuthCallbackDispatcher struct { + calls atomic.Int32 +} + +func (*selectedAuthCallbackDispatcher) HeartbeatOK() bool { return true } +func (d *selectedAuthCallbackDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + if d.calls.Add(1) > 2 { + return json.Marshal(homeErrorEnvelope{Error: &homeErrorDetail{Code: homeRequestRetryExceededErrorCode, Message: "no more auths"}}) + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ID: "home-auth", Provider: "home-execution", Status: StatusActive, Attributes: map[string]string{"websockets": "true"}}}) +} +func (*selectedAuthCallbackDispatcher) AbortAmbiguousDispatch() {} + +type callbackPinHomeExecutor struct { + manager *Manager + session string + calls atomic.Int32 +} + +func (*callbackPinHomeExecutor) Identifier() string { return "home-execution" } +func (e *callbackPinHomeExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (e *callbackPinHomeExecutor) ExecuteStream(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if e.calls.Add(1) == 2 { + return nil, errSelectedAuthCallbackFailure + } + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed"}`)} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} +func (*callbackPinHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } +func (*callbackPinHomeExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} +func (*callbackPinHomeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +var errSelectedAuthCallbackFailure = &Error{HTTPStatus: 502, Message: "selected auth failed"} + +func TestHomeSelectedAuthCallbackPinsFirstHandlerSelectionAndCleansFailure(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(&selectedAuthCallbackDispatcher{}, executionregistry.New(), 1) + executor := &callbackPinHomeExecutor{manager: manager, session: "callback-session"} + manager.RegisterExecutor(executor) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + callbackSawRuntimeAuth := false + opts := cliproxyexecutor.Options{Stream: true, Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: executor.session, + cliproxyexecutor.SelectedAuthCallbackMetadataKey: func(authID string) { + _, callbackSawRuntimeAuth = manager.GetExecutionSessionAuthByID(executor.session, authID) + }, + }} + result, errExecute := manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-a"}, opts) + if errExecute != nil { + t.Fatalf("first ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + if !callbackSawRuntimeAuth { + t.Fatal("first selected-auth callback could not resolve the Home runtime auth") + } + + manager.CloseExecutionSession(executor.session) + callbackSawRuntimeAuth = false + _, errExecute = manager.ExecuteStream(ctx, []string{"home-execution"}, cliproxyexecutor.Request{Model: "model-b"}, opts) + if errExecute == nil { + t.Fatal("failed ExecuteStream() error = nil") + } + if !callbackSawRuntimeAuth { + t.Fatal("failed selected-auth callback could not resolve the Home runtime auth") + } + if _, ok := manager.GetExecutionSessionAuthByID(executor.session, "home-auth"); ok { + t.Fatal("failed selection retained Home runtime auth") + } +} diff --git a/sdk/cliproxy/auth/home_selection.go b/sdk/cliproxy/auth/home_selection.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_selection.go @@ -0,0 +1,300 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "slices" + "strings" + "sync" + "sync/atomic" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" +) + +type executionResources struct { + mu sync.Mutex + closed bool + closers []func() error +} + +type attemptCancel struct { + cancel context.CancelFunc + once sync.Once +} + +func (a *attemptCancel) Cancel() { + if a == nil || a.cancel == nil { + return + } + a.once.Do(a.cancel) +} + +type attemptCancels struct { + mu sync.Mutex + closed bool + next uint64 + cancels map[uint64]*attemptCancel +} + +func (a *attemptCancels) Add(cancel context.CancelFunc) (func(), error) { + if a == nil || cancel == nil { + return func() {}, executionregistry.ErrInvalidExecutionResource + } + + a.mu.Lock() + if a.closed { + a.mu.Unlock() + cancel() + return func() {}, executionregistry.ErrRegistryNotAccepting + } + if a.cancels == nil { + a.cancels = make(map[uint64]*attemptCancel) + } + a.next++ + token := a.next + attempt := &attemptCancel{cancel: cancel} + a.cancels[token] = attempt + a.mu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + a.mu.Lock() + delete(a.cancels, token) + a.mu.Unlock() + attempt.Cancel() + }) + }, nil +} + +func (a *attemptCancels) Close() error { + if a == nil { + return nil + } + + a.mu.Lock() + if a.closed { + a.mu.Unlock() + return nil + } + a.closed = true + cancels := a.cancels + a.cancels = nil + a.mu.Unlock() + + for _, cancel := range cancels { + cancel.Cancel() + } + return nil +} + +func (a *attemptCancels) Len() int { + if a == nil { + return 0 + } + a.mu.Lock() + defer a.mu.Unlock() + return len(a.cancels) +} + +func (r *executionResources) Add(closeFn func() error) error { + if closeFn == nil { + return executionregistry.ErrInvalidExecutionResource + } + + r.mu.Lock() + if !r.closed { + r.closers = append(r.closers, closeFn) + r.mu.Unlock() + return nil + } + r.mu.Unlock() + + if errClose := closeFn(); errClose != nil { + return errors.Join(executionregistry.ErrRegistryNotAccepting, errClose) + } + return executionregistry.ErrRegistryNotAccepting +} + +func (r *executionResources) Close() error { + r.mu.Lock() + if r.closed { + r.mu.Unlock() + return nil + } + r.closed = true + closers := slices.Clone(r.closers) + r.closers = nil + r.mu.Unlock() + + var result error + for index := len(closers) - 1; index >= 0; index-- { + result = errors.Join(result, closers[index]()) + } + return result +} + +// HomeDispatchSelection keeps a Home execution scope separate from its auth. +type HomeDispatchSelection struct { + Auth *Auth + Executor ProviderExecutor + Provider string + + scope *executionregistry.Scope + accountedModel string + resources *executionResources + attemptCancels *attemptCancels + once sync.Once + retained atomic.Bool + runtimeAuthBound atomic.Bool + ended atomic.Bool +} + +func newHomeDispatchSelection(auth *Auth, executor ProviderExecutor, provider string, scope *executionregistry.Scope) (*HomeDispatchSelection, error) { + if scope == nil { + return nil, fmt.Errorf("Home dispatch selection has no execution scope") + } + + resources := &executionResources{} + attemptCancels := &attemptCancels{} + if errBind := resources.Add(attemptCancels.Close); errBind != nil { + _ = attemptCancels.Close() + scope.End("attempt_cancel_bind_failed") + return nil, errBind + } + if errBind := scope.Bind(resources.Close); errBind != nil { + _ = resources.Close() + scope.End("resource_controller_bind_failed") + return nil, errBind + } + + return &HomeDispatchSelection{ + Auth: auth, + Executor: executor, + Provider: strings.TrimSpace(provider), + scope: scope, + resources: resources, + attemptCancels: attemptCancels, + }, nil +} + +// Bind adds a resource to be closed when this selection ends or drains. +func (s *HomeDispatchSelection) Bind(closeFn func() error) error { + if s == nil || s.resources == nil { + if closeFn != nil { + _ = closeFn() + } + return fmt.Errorf("Home dispatch selection has no execution resources") + } + return s.resources.Add(closeFn) +} + +// AttemptContext creates a selection-owned context and returns its release function. +func (s *HomeDispatchSelection) AttemptContext(ctx context.Context) (context.Context, func(), error) { + if ctx == nil { + ctx = context.Background() + } + attemptCtx, cancelAttempt := context.WithCancel(ctx) + if s == nil || s.attemptCancels == nil { + cancelAttempt() + return nil, func() {}, fmt.Errorf("Home dispatch selection has no attempt cancels") + } + release, errAdd := s.attemptCancels.Add(cancelAttempt) + if errAdd != nil { + cancelAttempt() + return nil, func() {}, errAdd + } + return attemptCtx, release, nil +} + +// Retain transfers selection ownership from a request to an execution session. +func (s *HomeDispatchSelection) Retain() { + if s == nil || s.ended.Load() { + return + } + s.retained.Store(true) +} + +// Retained reports whether an executor transferred this selection to a session. +func (s *HomeDispatchSelection) Retained() bool { + return s != nil && s.retained.Load() && !s.ended.Load() +} + +// Active reports whether the selection has not ended. +func (s *HomeDispatchSelection) Active() bool { + return s != nil && !s.ended.Load() +} + +// End closes all bound resources and releases the Home execution scope once. +func (s *HomeDispatchSelection) End(reason string) { + _ = s.EndWithRelease(reason) +} + +// EndWithRelease closes all bound resources and returns the Home release ticket. +func (s *HomeDispatchSelection) EndWithRelease(reason string) *executionregistry.ReleaseTicket { + if s == nil { + return nil + } + var ticket *executionregistry.ReleaseTicket + s.once.Do(func() { + s.ended.Store(true) + if s.scope != nil { + ticket = s.scope.EndWithRelease(strings.TrimSpace(reason)) + } + }) + if ticket != nil || s.scope == nil { + return ticket + } + return s.scope.EndWithRelease("") +} + +// CloneAuth returns a standalone auth copy without the selection handle. +func (s *HomeDispatchSelection) CloneAuth() *Auth { + if s == nil || s.Auth == nil { + return nil + } + return s.Auth.Clone() +} + +// CloneAuthForRoute returns an auth copy adapted for a retained canonical route. +func (s *HomeDispatchSelection) CloneAuthForRoute(routeModel string) *Auth { + auth := s.CloneAuth() + if auth == nil || !s.Retained() { + return auth + } + return cloneRetainedHomeAuthForRoute(auth, routeModel) +} + +func cloneRetainedHomeAuthForRoute(auth *Auth, routeModel string) *Auth { + if auth == nil || auth.Attributes == nil { + return auth + } + upstreamModel := strings.TrimSpace(auth.Attributes[homeUpstreamModelAttributeKey]) + if upstreamModel == "" { + return auth + } + upstreamBase, _ := splitRecognizedHomeReasoningSuffix(upstreamModel) + _, routeSuffix := splitRecognizedHomeReasoningSuffix(routeModel) + auth.Attributes[homeUpstreamModelAttributeKey] = upstreamBase + routeSuffix + if strings.EqualFold(strings.TrimSpace(auth.Attributes[homeForceMappingAttributeKey]), "true") { + auth.Attributes[homeOriginalAliasAttributeKey] = strings.TrimSpace(rewriteModelForAuth(routeModel, auth)) + } + return auth +} + +func splitRecognizedHomeReasoningSuffix(model string) (string, string) { + model = strings.Trim(model, asciiWhitespace) + if !strings.HasSuffix(model, ")") { + return model, "" + } + open := strings.LastIndexByte(model, '(') + if open < 0 || !recognizedHomeConcurrencySuffix(model[open+1:len(model)-1]) { + return model, "" + } + base := strings.Trim(model[:open], asciiWhitespace) + if base == "" { + return model, "" + } + return base, model[open:] +} diff --git a/sdk/cliproxy/auth/home_selection_attempt_test.go b/sdk/cliproxy/auth/home_selection_attempt_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_selection_attempt_test.go @@ -0,0 +1,107 @@ +package auth + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" +) + +func TestHomeDispatchSelectionReleasesAttemptCancelTokensWithoutGrowingResources(t *testing.T) { + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + selection, errSelection := newHomeDispatchSelection(&Auth{ID: "home-auth"}, nil, "home", scope) + if errSelection != nil { + t.Fatal(errSelection) + } + + for range 100 { + _, release, errAttempt := selection.AttemptContext(context.Background()) + if errAttempt != nil { + t.Fatalf("AttemptContext() error = %v", errAttempt) + } + release() + } + + selection.resources.mu.Lock() + resourceCount := len(selection.resources.closers) + selection.resources.mu.Unlock() + if resourceCount != 1 { + t.Fatalf("bound resources = %d, want 1 attempt cancel registry", resourceCount) + } + if got := selection.attemptCancels.Len(); got != 0 { + t.Fatalf("active attempt cancel tokens = %d, want 0", got) + } + + selection.End("completed") + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := registry.Drain(drainCtx); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestAttemptCancelReleaseAfterCloseCancelsOnce(t *testing.T) { + cancels := &attemptCancels{} + var cancelCalls atomic.Int32 + release, errAdd := cancels.Add(func() { cancelCalls.Add(1) }) + if errAdd != nil { + t.Fatalf("Add() error = %v", errAdd) + } + if errClose := cancels.Close(); errClose != nil { + t.Fatalf("Close() error = %v", errClose) + } + release() + if got := cancelCalls.Load(); got != 1 { + t.Fatalf("cancel calls = %d, want 1", got) + } +} + +func TestHomeDispatchSelectionAttemptReleaseRacesDrainExactlyOnce(t *testing.T) { + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + selection, errSelection := newHomeDispatchSelection(&Auth{ID: "home-auth"}, nil, "home", scope) + if errSelection != nil { + t.Fatal(errSelection) + } + + _, release, errAttempt := selection.AttemptContext(context.Background()) + if errAttempt != nil { + t.Fatal(errAttempt) + } + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + release() + }() + go func() { + defer wg.Done() + selection.End("draining") + }() + wg.Wait() + + if got := selection.attemptCancels.Len(); got != 0 { + t.Fatalf("active attempt cancel tokens = %d, want 0", got) + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} diff --git a/sdk/cliproxy/auth/home_selection_test.go b/sdk/cliproxy/auth/home_selection_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/auth/home_selection_test.go @@ -0,0 +1,186 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestHomeDispatchSelectionOwnsScopeOutsideAuth(t *testing.T) { + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{RequestID: "req-1", CredentialID: "cred-1", Model: "gpt", Kind: "http", StartedAt: time.Now()}) + if errInstall != nil { + t.Fatal(errInstall) + } + selection, errSelection := newHomeDispatchSelection(&Auth{ID: "cred-1", Provider: "codex"}, nil, "codex", scope) + if errSelection != nil { + t.Fatal(errSelection) + } + clone := selection.CloneAuth() + if clone == nil || clone.ID != "cred-1" || clone.Runtime != nil { + t.Fatalf("clone = %#v", clone) + } + closed := atomic.Int32{} + if errBind := selection.Bind(func() error { closed.Add(1); return nil }); errBind != nil { + t.Fatal(errBind) + } + selection.End("completed") + selection.End("duplicate") + if closed.Load() != 1 { + t.Fatalf("close calls = %d", closed.Load()) + } +} + +func TestHomeDispatchSelectionDrainsResourcesAddedDuringEnd(t *testing.T) { + registry := executionregistry.New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + selection, errSelection := newHomeDispatchSelection(&Auth{ID: "cred-1"}, nil, "test", scope) + if errSelection != nil { + t.Fatal(errSelection) + } + + started := make(chan struct{}) + release := make(chan struct{}) + if errBind := selection.Bind(func() error { + close(started) + <-release + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + done := make(chan struct{}) + go func() { + selection.End("draining") + close(done) + }() + <-started + + closedLate := atomic.Int32{} + errLate := selection.Bind(func() error { + closedLate.Add(1) + return errors.New("late close") + }) + if !errors.Is(errLate, executionregistry.ErrRegistryNotAccepting) { + t.Fatalf("late Bind() error = %v, want ErrRegistryNotAccepting", errLate) + } + if closedLate.Load() != 1 { + t.Fatalf("late close calls = %d, want 1", closedLate.Load()) + } + + close(release) + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("End did not complete") + } + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := registry.Drain(drainCtx); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +type gatedHomeDispatcher struct { + loaded chan struct{} + release chan struct{} + rpop atomic.Int32 +} + +func (d *gatedHomeDispatcher) HeartbeatOK() bool { + select { + case <-d.loaded: + default: + close(d.loaded) + } + <-d.release + return true +} + +func (d *gatedHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + d.rpop.Add(1) + return nil, errors.New("old Home dispatcher was used") +} + +func (*gatedHomeDispatcher) AbortAmbiguousDispatch() {} + +func TestManagerHomeDispatchBundleCompareAndClearDoesNotRemoveReplacement(t *testing.T) { + manager := NewManager(nil, nil, nil) + first := manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 1) + second := manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 2) + + if manager.ClearHomeDispatchBundle(first) { + t.Fatal("ClearHomeDispatchBundle() cleared a replacement bundle") + } + if got := manager.HomeDispatchBundle(); got != second { + t.Fatalf("HomeDispatchBundle() = %p, want %p", got, second) + } + if !manager.ClearHomeDispatchBundle(second) { + t.Fatal("ClearHomeDispatchBundle() = false, want true") + } + if got := manager.HomeDispatchBundle(); got != nil { + t.Fatalf("HomeDispatchBundle() = %p, want nil", got) + } +} + +func TestPickHomeDispatchSelectionDoesNotMixDetachedBundleWithReplacement(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + oldDispatcher := &gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})} + oldRegistry := executionregistry.New() + oldBundle := manager.PublishHomeDispatch(oldDispatcher, oldRegistry, 1) + + result := make(chan error, 1) + go func() { + _, errSelect := manager.pickHomeDispatchSelection(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}) + result <- errSelect + }() + select { + case <-oldDispatcher.loaded: + case <-time.After(time.Second): + t.Fatal("selection did not load the old dispatch bundle") + } + + if !manager.ClearHomeDispatchBundle(oldBundle) { + t.Fatal("ClearHomeDispatchBundle() = false, want true") + } + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := oldRegistry.Drain(drainCtx); errDrain != nil { + t.Fatalf("old registry Drain() error = %v", errDrain) + } + manager.PublishHomeDispatch(&gatedHomeDispatcher{loaded: make(chan struct{}), release: make(chan struct{})}, executionregistry.New(), 2) + close(oldDispatcher.release) + + select { + case errSelect := <-result: + var authErr *Error + if !errors.As(errSelect, &authErr) || authErr.Code != "home_unavailable" { + t.Fatalf("pickHomeDispatchSelection() error = %v, want home_unavailable", errSelect) + } + case <-time.After(time.Second): + t.Fatal("selection did not resume after the old bundle was detached") + } + if got := oldDispatcher.rpop.Load(); got != 0 { + t.Fatalf("old dispatcher RPopAuth() calls = %d, want 0", got) + } +} diff --git a/sdk/cliproxy/auth/home_websocket_reuse_test.go b/sdk/cliproxy/auth/home_websocket_reuse_test.go --- a/sdk/cliproxy/auth/home_websocket_reuse_test.go +++ b/sdk/cliproxy/auth/home_websocket_reuse_test.go @@ -4,13 +4,17 @@ "context" "errors" "net/http" + "sync/atomic" "testing" + "time" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) -func TestPickNextViaHomeReusesPinnedWebsocketAuthWithoutHomeDispatch(t *testing.T) { +func TestPickNextViaHomeDoesNotReusePinnedWebsocketAuthWithoutSelection(t *testing.T) { manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) manager.RegisterExecutor(schedulerTestExecutor{}) @@ -42,21 +46,15 @@ } got, executor, provider, errPick := manager.pickNextViaHome(ctx, "gpt-5.4", opts, nil) - if errPick != nil { - t.Fatalf("pickNextViaHome() error = %v", errPick) + if errPick == nil { + t.Fatal("pickNextViaHome() unexpectedly reused an auth without a Home selection") } - if got == nil || got.ID != "home-auth-1" { - t.Fatalf("pickNextViaHome() auth = %#v, want home-auth-1", got) - } - if executor == nil { - t.Fatal("pickNextViaHome() executor is nil") - } - if provider != "test" { - t.Fatalf("pickNextViaHome() provider = %q, want test", provider) + if got != nil || executor != nil || provider != "" { + t.Fatalf("pickNextViaHome() returned unbound execution target: auth=%#v executor=%#v provider=%q", got, executor, provider) } } -func TestPickNextViaHomeKeepsSameAuthIDPayloadSessionScoped(t *testing.T) { +func TestPickNextViaHomeRejectsSessionScopedAuthCache(t *testing.T) { manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) manager.RegisterExecutor(schedulerTestExecutor{}) @@ -94,20 +92,11 @@ }, } - gotSession1, _, _, errSession1 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession1, nil) - if errSession1 != nil { - t.Fatalf("pickNextViaHome(session-1) error = %v", errSession1) + if _, _, _, errSession1 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession1, nil); errSession1 == nil { + t.Fatal("pickNextViaHome(session-1) unexpectedly reused a session auth cache") } - if got := gotSession1.Attributes[homeUpstreamModelAttributeKey]; got != "upstream-model-a" { - t.Fatalf("pickNextViaHome(session-1) upstream model = %q, want upstream-model-a", got) - } - - gotSession2, _, _, errSession2 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession2, nil) - if errSession2 != nil { - t.Fatalf("pickNextViaHome(session-2) error = %v", errSession2) - } - if got := gotSession2.Attributes[homeUpstreamModelAttributeKey]; got != "upstream-model-b" { - t.Fatalf("pickNextViaHome(session-2) upstream model = %q, want upstream-model-b", got) + if _, _, _, errSession2 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession2, nil); errSession2 == nil { + t.Fatal("pickNextViaHome(session-2) unexpectedly reused a session auth cache") } } @@ -222,19 +211,28 @@ } type homeAuthTransportErrorDispatcher struct { - err error + err error + aborts atomic.Int32 + onAbort func() } -func (d homeAuthTransportErrorDispatcher) HeartbeatOK() bool { +func (d *homeAuthTransportErrorDispatcher) HeartbeatOK() bool { return true } -func (d homeAuthTransportErrorDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { +func (d *homeAuthTransportErrorDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { return nil, d.err } +func (d *homeAuthTransportErrorDispatcher) AbortAmbiguousDispatch() { + d.aborts.Add(1) + if d.onAbort != nil { + d.onAbort() + } +} + func TestPickNextViaHomeClassifiesTransportErrorsAsHomeUnavailable(t *testing.T) { - dispatcher := homeAuthTransportErrorDispatcher{err: errors.New("read tcp 127.0.0.1:46704->127.0.0.1:8327: i/o timeout")} + dispatcher := &homeAuthTransportErrorDispatcher{err: errors.New("read tcp 127.0.0.1:46704->127.0.0.1:8327: i/o timeout")} oldCurrentHomeDispatcher := currentHomeDispatcher currentHomeDispatcher = func() homeAuthDispatcher { return dispatcher @@ -245,6 +243,7 @@ manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(executionregistry.New()) _, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil) if errPick == nil { @@ -262,6 +261,91 @@ } if !authErr.Retryable { t.Fatal("pickNextViaHome() retryable = false, want true") + } +} + +func TestPickNextViaHomeAbortsBeforeEndingPendingDispatch(t *testing.T) { + registry := executionregistry.New() + abortSawPending := make(chan bool, 1) + dispatcher := &homeAuthTransportErrorDispatcher{ + err: home.NewAmbiguousDispatchError(errors.New("response connection closed")), + onAbort: func() { + cancelledCtx, cancel := context.WithCancel(context.Background()) + cancel() + abortSawPending <- errors.Is(registry.Drain(cancelledCtx), context.Canceled) + }, + } + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(registry) + + _, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil) + if errPick == nil { + t.Fatal("pickNextViaHome() error = nil, want home unavailable") + } + if sawPending := <-abortSawPending; !sawPending { + t.Fatal("AbortAmbiguousDispatch() observed an already-ended pending dispatch") + } +} + +func TestPickNextViaHomeDoesNotAbortDeterministicDispatchFailure(t *testing.T) { + dispatcher := &homeAuthTransportErrorDispatcher{err: home.ErrNotConnected} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(executionregistry.New()) + + _, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil) + if errPick == nil { + t.Fatal("pickNextViaHome() error = nil, want home unavailable") + } + if got := dispatcher.aborts.Load(); got != 0 { + t.Fatalf("AbortAmbiguousDispatch() calls = %d, want 0 for deterministic failure", got) + } +} + +func TestPickNextViaHomeAbortsAmbiguousTransport(t *testing.T) { + dispatcher := &homeAuthTransportErrorDispatcher{err: home.NewAmbiguousDispatchError(errors.New("response connection closed"))} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.SetHomeExecutionRegistry(registry) + + _, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil) + if errPick == nil { + t.Fatal("pickNextViaHome() error = nil, want home unavailable") + } + if got := dispatcher.aborts.Load(); got != 1 { + t.Fatalf("AbortAmbiguousDispatch() calls = %d, want 1", got) + } + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := registry.Drain(drainCtx); errDrain != nil { + t.Fatalf("Drain() error = %v, ambiguous pending dispatch was not ended", errDrain) } } diff --git a/sdk/cliproxy/auth/request_auth_prepare_test.go b/sdk/cliproxy/auth/request_auth_prepare_test.go --- a/sdk/cliproxy/auth/request_auth_prepare_test.go +++ b/sdk/cliproxy/auth/request_auth_prepare_test.go @@ -2,13 +2,19 @@ import ( "context" + "encoding/json" + "errors" "net/http" + "reflect" "strings" "sync" "sync/atomic" "testing" + "time" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) @@ -39,6 +45,10 @@ type requestPrepareExecutor struct { prepareCalls atomic.Int32 executeCalls atomic.Int32 + prepareErr error + executeErr error + mu sync.Mutex + observed []*Auth } func (e *requestPrepareExecutor) Identifier() string { return "antigravity" } @@ -49,6 +59,9 @@ func (e *requestPrepareExecutor) PrepareRequestAuth(_ context.Context, auth *Auth) (*Auth, error) { e.prepareCalls.Add(1) + if e.prepareErr != nil { + return nil, e.prepareErr + } updated := auth.Clone() if updated.Metadata == nil { updated.Metadata = make(map[string]any) @@ -57,28 +70,285 @@ return updated, nil } -func (e *requestPrepareExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { +func (e *requestPrepareExecutor) recordPreparedAuth(auth *Auth) error { e.executeCalls.Add(1) if got := testStringValue(auth.Metadata["project_id"]); got != "prepared-project" { - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusBadRequest, Message: "missing prepared project"} + return &Error{HTTPStatus: http.StatusBadRequest, Message: "missing prepared project"} + } + e.mu.Lock() + e.observed = append(e.observed, auth.Clone()) + e.mu.Unlock() + return nil +} + +func (e *requestPrepareExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if errPrepared := e.recordPreparedAuth(auth); errPrepared != nil { + return cliproxyexecutor.Response{}, errPrepared + } + if e.executeErr != nil { + return cliproxyexecutor.Response{}, e.executeErr } return cliproxyexecutor.Response{Payload: []byte("ok")}, nil } -func (e *requestPrepareExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { - return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "stream not implemented"} +func (e *requestPrepareExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if errPrepared := e.recordPreparedAuth(auth); errPrepared != nil { + return nil, errPrepared + } + if e.executeErr != nil { + return nil, e.executeErr + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed"}`)} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil } func (e *requestPrepareExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { return auth, nil } -func (e *requestPrepareExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { - return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "count not implemented"} +func (e *requestPrepareExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if errPrepared := e.recordPreparedAuth(auth); errPrepared != nil { + return cliproxyexecutor.Response{}, errPrepared + } + if e.executeErr != nil { + return cliproxyexecutor.Response{}, e.executeErr + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (e *requestPrepareExecutor) lastObservedAuth() *Auth { + e.mu.Lock() + defer e.mu.Unlock() + if len(e.observed) == 0 { + return nil + } + return e.observed[len(e.observed)-1].Clone() } func (e *requestPrepareExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "http not implemented"} +} + +type homeRequestPrepareDispatcher struct { + calls atomic.Int32 +} + +func (*homeRequestPrepareDispatcher) HeartbeatOK() bool { return true } + +func (d *homeRequestPrepareDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + if d.calls.Add(1) > 1 { + return json.Marshal(homeErrorEnvelope{Error: &homeErrorDetail{Code: homeRequestRetryExceededErrorCode, Message: "no more Home auths"}}) + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "same-id", + Provider: "antigravity", + Status: StatusActive, + Metadata: map[string]any{"access_token": "home-token", "source": "home"}, + }}) +} + +func (*homeRequestPrepareDispatcher) AbortAmbiguousDispatch() {} + +func TestHomePrepareUsesEphemeralDispatchAuthAcrossExecutionPaths(t *testing.T) { + for _, path := range []struct { + name string + run func(*Manager, context.Context) error + }{ + { + name: "Execute", + run: func(manager *Manager, ctx context.Context) error { + _, errExecute := manager.Execute(ctx, []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "Count", + run: func(manager *Manager, ctx context.Context) error { + _, errCount := manager.ExecuteCount(ctx, []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "Stream", + run: func(manager *Manager, ctx context.Context) error { + result, errStream := manager.ExecuteStream(ctx, []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{Stream: true}) + if errStream != nil { + return errStream + } + for range result.Chunks { + } + return nil + }, + }, + } { + t.Run(path.name, func(t *testing.T) { + store := &requestPrepareStore{} + executor := &requestPrepareExecutor{} + manager := NewManager(store, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(&homeRequestPrepareDispatcher{}, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + localAuth := &Auth{ID: "same-id", Provider: "antigravity", Status: StatusActive, Metadata: map[string]any{"access_token": "local-token", "source": "local"}} + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), localAuth); errRegister != nil { + t.Fatalf("register local auth: %v", errRegister) + } + if errRun := path.run(manager, context.Background()); errRun != nil { + t.Fatalf("%s error: %v", path.name, errRun) + } + observed := executor.lastObservedAuth() + if observed == nil { + t.Fatal("executor did not receive prepared auth") + } + if got := testStringValue(observed.Metadata["access_token"]); got != "home-token" { + t.Fatalf("executor access token = %q, want Home token", got) + } + if got := testStringValue(observed.Metadata["source"]); got != "home" { + t.Fatalf("executor source = %q, want Home metadata", got) + } + current, ok := manager.GetByID("same-id") + if !ok { + t.Fatal("local auth disappeared") + } + if got := testStringValue(current.Metadata["access_token"]); got != "local-token" { + t.Fatalf("local access token = %q, want unchanged local token", got) + } + if got := testStringValue(current.Metadata["source"]); got != "local" { + t.Fatalf("local source = %q, want unchanged local metadata", got) + } + }) + } +} + +func TestHomeExecutionResultsDoNotMutateSameIDLocalAuth(t *testing.T) { + paths := []struct { + name string + run func(*Manager, context.Context) error + }{ + { + name: "Execute", + run: func(manager *Manager, ctx context.Context) error { + _, errExecute := manager.Execute(ctx, []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "Count", + run: func(manager *Manager, ctx context.Context) error { + _, errCount := manager.ExecuteCount(ctx, []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "Stream", + run: func(manager *Manager, ctx context.Context) error { + result, errStream := manager.ExecuteStream(ctx, []string{"antigravity"}, cliproxyexecutor.Request{Model: "test-model"}, cliproxyexecutor.Options{Stream: true}) + if errStream != nil { + return errStream + } + for range result.Chunks { + } + return nil + }, + }, + } + outcomes := []struct { + name string + prepareErr error + executeErr error + }{ + {name: "success"}, + {name: "execution failure", executeErr: errors.New("upstream failed")}, + {name: "prepare failure", prepareErr: errors.New("prepare failed")}, + } + + for _, path := range paths { + for _, outcome := range outcomes { + t.Run(path.name+"/"+outcome.name, func(t *testing.T) { + store := &requestPrepareStore{} + hook := &resultCaptureHook{} + executor := &requestPrepareExecutor{prepareErr: outcome.prepareErr, executeErr: outcome.executeErr} + manager := NewManager(store, nil, hook) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(&homeRequestPrepareDispatcher{}, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + localAuth := &Auth{ + ID: "same-id", + Provider: "antigravity", + Status: StatusActive, + Success: 7, + Failed: 4, + UpdatedAt: time.Unix(123, 0), + Metadata: map[string]any{"access_token": "local-token", "source": "local"}, + ModelStates: map[string]*ModelState{ + "test-model": {Status: StatusError, Unavailable: true, StatusMessage: "local failure", UpdatedAt: time.Unix(122, 0)}, + }, + } + if _, errRegister := manager.Register(WithSkipPersist(context.Background()), localAuth); errRegister != nil { + t.Fatalf("register local auth: %v", errRegister) + } + registry.GetGlobalRegistry().RegisterClient(localAuth.ID, localAuth.Provider, []*registry.ModelInfo{{ID: "test-model"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(localAuth.ID) }) + + beforeLocal, ok := manager.GetByID(localAuth.ID) + if !ok { + t.Fatal("local auth is missing before Home execution") + } + beforeScheduler := homeExecutionSchedulerAuthSnapshot(t, manager, localAuth.ID) + beforeModels := registry.GetGlobalRegistry().GetModelsForClient(localAuth.ID) + failed := outcome.prepareErr != nil || outcome.executeErr != nil + if errRun := path.run(manager, context.Background()); failed != (errRun != nil) { + t.Fatalf("%s error = %v, want failure=%t", path.name, errRun, failed) + } + if outcome.prepareErr == nil { + observed := executor.lastObservedAuth() + if observed == nil { + t.Fatal("executor did not receive prepared auth") + } + if got := testStringValue(observed.Metadata["access_token"]); got != "home-token" { + t.Fatalf("executor access token = %q, want Home token", got) + } + } + assertHomeExecutionResultStateUnchanged(t, manager, store, hook, beforeLocal, beforeScheduler, beforeModels) + }) + } + } +} + +func homeExecutionSchedulerAuthSnapshot(t *testing.T, manager *Manager, authID string) *Auth { + t.Helper() + manager.scheduler.mu.Lock() + defer manager.scheduler.mu.Unlock() + provider := manager.scheduler.authProviders[authID] + entry := manager.scheduler.providers[provider] + if entry == nil || entry.auths[authID] == nil || entry.auths[authID].auth == nil { + t.Fatalf("scheduler auth %q is missing", authID) + } + return entry.auths[authID].auth.Clone() +} + +func assertHomeExecutionResultStateUnchanged(t *testing.T, manager *Manager, store *requestPrepareStore, hook *resultCaptureHook, beforeLocal, beforeScheduler *Auth, beforeModels []*registry.ModelInfo) { + t.Helper() + current, ok := manager.GetByID(beforeLocal.ID) + if !ok { + t.Fatal("local auth disappeared") + } + if !reflect.DeepEqual(current, beforeLocal) { + t.Fatalf("Home execution mutated local auth:\n got %#v\nwant %#v", current, beforeLocal) + } + if currentScheduler := homeExecutionSchedulerAuthSnapshot(t, manager, beforeLocal.ID); !reflect.DeepEqual(currentScheduler, beforeScheduler) { + t.Fatalf("Home execution mutated scheduler auth:\n got %#v\nwant %#v", currentScheduler, beforeScheduler) + } + if afterModels := registry.GetGlobalRegistry().GetModelsForClient(beforeLocal.ID); !reflect.DeepEqual(afterModels, beforeModels) { + t.Fatalf("Home execution mutated global model state:\n got %#v\nwant %#v", afterModels, beforeModels) + } + if got := store.saveCount.Load(); got != 0 { + t.Fatalf("Home execution save count = %d, want 0", got) + } + if results := hook.Results(); len(results) != 1 { + t.Fatalf("Home execution hook results = %#v, want exactly one ephemeral result", results) + } } func TestManagerExecute_PreparesAndPersistsMissingRequestAuthMetadata(t *testing.T) { diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -11,13 +11,21 @@ internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/home" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) -type schedulerTestExecutor struct{} +type schedulerTestExecutor struct { + provider string +} -func (schedulerTestExecutor) Identifier() string { return "test" } +func (e schedulerTestExecutor) Identifier() string { + if e.provider != "" { + return e.provider + } + return "test" +} func (schedulerTestExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { return cliproxyexecutor.Response{}, nil @@ -77,6 +85,8 @@ } return json.Marshal(homeAuthDispatchResponse{Auth: d.auths[count-1]}) } + +func (*authKindHomeDispatcher) AbortAmbiguousDispatch() {} func (s *inactivePluginScheduler) HasScheduler() bool { return false @@ -509,66 +519,173 @@ } } -func TestManagerSelectAuthByKindAdvancesHomeAuthCount(t *testing.T) { - tests := []struct { - name string - auths []Auth - wantAuthID string - wantError string - }{ - { - name: "skips API key for OAuth", - auths: []Auth{ - {ID: "home-api-key", Provider: "test", Attributes: map[string]string{AttributeAPIKey: "test-key"}}, - {ID: "home-oauth", Provider: "test", Metadata: map[string]any{"access_token": "test-token"}}, - }, - wantAuthID: "home-oauth", +func TestManagerLegacySelectAuthFailsClosedWhenHomeEnabled(t *testing.T) { + dispatcher := &authKindHomeDispatcher{auths: []Auth{{ + ID: "home-oauth", + Provider: "test", + Metadata: map[string]any{"access_token": "test-token"}, + }}} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { return dispatcher } + t.Cleanup(func() { currentHomeDispatcher = oldCurrentHomeDispatcher }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(executionregistry.New()) + manager.RegisterExecutor(schedulerTestExecutor{}) + + for name, selectAuth := range map[string]func() (*Auth, error){ + "SelectAuth": func() (*Auth, error) { + return manager.SelectAuth(context.Background(), "test", "model", cliproxyexecutor.Options{}) }, - { - name: "returns not found without OAuth", - auths: []Auth{ - {ID: "home-api-key", Provider: "test", Attributes: map[string]string{AttributeAPIKey: "test-key"}}, - }, - wantError: "auth_not_found", + "SelectAuthByKind": func() (*Auth, error) { + return manager.SelectAuthByKind(context.Background(), "test", "model", AuthKindOAuth, cliproxyexecutor.Options{}) }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - dispatcher := &authKindHomeDispatcher{auths: tt.auths} - oldCurrentHomeDispatcher := currentHomeDispatcher - currentHomeDispatcher = func() homeAuthDispatcher { - return dispatcher + } { + t.Run(name, func(t *testing.T) { + selected, errSelect := selectAuth() + if selected != nil { + t.Fatalf("%s() auth = %#v, want nil", name, selected) } - t.Cleanup(func() { - currentHomeDispatcher = oldCurrentHomeDispatcher - }) - - manager := NewManager(nil, nil, nil) - manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) - manager.RegisterExecutor(schedulerTestExecutor{}) - - selected, errSelect := manager.SelectAuthByKind(context.Background(), "codex", "", AuthKindOAuth, cliproxyexecutor.Options{}) - if tt.wantError != "" { - if selected != nil { - t.Fatalf("SelectAuthByKind() auth = %#v, want nil", selected) - } - var authErr *Error - if !errors.As(errSelect, &authErr) || authErr.Code != tt.wantError { - t.Fatalf("SelectAuthByKind() error = %#v, want %s", errSelect, tt.wantError) - } - } else { - if errSelect != nil { - t.Fatalf("SelectAuthByKind() error = %v", errSelect) - } - if selected == nil || selected.ID != tt.wantAuthID { - t.Fatalf("SelectAuthByKind() auth = %#v, want %s", selected, tt.wantAuthID) - } - } - if len(dispatcher.counts) != 2 || dispatcher.counts[0] != 1 || dispatcher.counts[1] != 2 { - t.Fatalf("home auth counts = %v, want [1 2]", dispatcher.counts) + var authErr *Error + if !errors.As(errSelect, &authErr) || authErr.Code != "home_unavailable" || authErr.HTTPStatus != http.StatusServiceUnavailable { + t.Fatalf("%s() error = %#v, want home_unavailable", name, errSelect) } }) + } + if len(dispatcher.counts) != 0 { + t.Fatalf("legacy selection issued Home RPOP calls: %v", dispatcher.counts) + } +} + +func TestSelectHomeAuthByKindReturnsHomeSelection(t *testing.T) { + dispatcher := &authKindHomeDispatcher{auths: []Auth{{ + ID: "home-oauth", + Provider: "test", + Metadata: map[string]any{"access_token": "test-token"}, + }}} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(executionregistry.New()) + manager.RegisterExecutor(schedulerTestExecutor{}) + + selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), "test", "gpt-5.4", AuthKindOAuth, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect) + } + if selection == nil || selection.Auth == nil || selection.Auth.ID != "home-oauth" { + t.Fatalf("SelectHomeAuthByKind() = %#v, want home-oauth", selection) + } + if selection.Executor == nil || selection.Provider != "test" { + t.Fatalf("selection executor/provider = %#v/%q, want test", selection.Executor, selection.Provider) + } + selection.End("test_complete") +} + +func TestSelectHomeAuthByKindSkipsProviderMismatch(t *testing.T) { + dispatcher := &authKindHomeDispatcher{auths: []Auth{ + {ID: "wrong-provider", Provider: "other", Metadata: map[string]any{"access_token": "test-token"}}, + {ID: "matching-provider", Provider: "test", Metadata: map[string]any{"access_token": "test-token"}}, + }} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(executionregistry.New()) + manager.RegisterExecutor(schedulerTestExecutor{}) + manager.RegisterExecutor(schedulerTestExecutor{provider: "other"}) + + selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), "test", "gpt-5.4", AuthKindOAuth, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect) + } + if selection == nil || selection.Auth == nil || selection.Auth.ID != "matching-provider" { + t.Fatalf("SelectHomeAuthByKind() = %#v, want matching provider auth", selection) + } + if got := dispatcher.counts; len(got) != 2 || got[0] != 1 || got[1] != 2 { + t.Fatalf("home auth counts = %v, want [1 2]", got) + } + selection.End("test_complete") +} + +func TestSelectHomeAuthByKindKeepsLogicalProviderWhenUsingCompatibilityExecutor(t *testing.T) { + dispatcher := &authKindHomeDispatcher{auths: []Auth{{ + ID: "compat-auth", + Provider: "base-url-provider", + Attributes: map[string]string{ + "base_url": "https://compat.example.com", + AttributeAPIKey: "test-key", + }, + }}} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetHomeExecutionRegistry(executionregistry.New()) + manager.RegisterExecutor(schedulerTestExecutor{provider: "openai-compatibility"}) + + selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), "base-url-provider", "gpt-5.4", AuthKindAPIKey, cliproxyexecutor.Options{}) + if errSelect != nil { + t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect) + } + if selection == nil || selection.Auth == nil || selection.Auth.ID != "compat-auth" { + t.Fatalf("SelectHomeAuthByKind() = %#v, want compat-auth", selection) + } + if selection.Provider != "base-url-provider" { + t.Fatalf("selection.Provider = %q, want logical provider base-url-provider", selection.Provider) + } + if selection.Executor == nil || selection.Executor.Identifier() != "openai-compatibility" { + t.Fatalf("selection.Executor = %#v, want openai-compatibility", selection.Executor) + } + selection.End("test_complete") +} + +func TestPickNextViaHomeEndsPendingOnInvalidAuth(t *testing.T) { + dispatcher := &authKindHomeDispatcher{auths: []Auth{{Provider: "test"}}} + oldCurrentHomeDispatcher := currentHomeDispatcher + currentHomeDispatcher = func() homeAuthDispatcher { + return dispatcher + } + t.Cleanup(func() { + currentHomeDispatcher = oldCurrentHomeDispatcher + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + registry := executionregistry.New() + manager.SetHomeExecutionRegistry(registry) + manager.RegisterExecutor(schedulerTestExecutor{}) + + _, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil) + var authErr *Error + if !errors.As(errPick, &authErr) || authErr.Code != "invalid_auth" { + t.Fatalf("pickNextViaHome() error = %v, want invalid_auth", errPick) + } + + drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second) + defer cancelDrain() + if errDrain := registry.Drain(drainCtx); errDrain != nil { + t.Fatalf("Drain() error = %v, pending dispatch was not ended", errDrain) } } diff --git a/sdk/cliproxy/executionregistry/concurrency_release_test.go b/sdk/cliproxy/executionregistry/concurrency_release_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/executionregistry/concurrency_release_test.go @@ -0,0 +1,85 @@ +package executionregistry + +import ( + "sync" + "testing" +) + +type recordingReleaseSink struct { + mu sync.Mutex + sequences map[ReleaseGroup]int64 +} + +func (s *recordingReleaseSink) MarkDirty(group ReleaseGroup, sequence int64) { + s.mu.Lock() + defer s.mu.Unlock() + if s.sequences == nil { + s.sequences = make(map[ReleaseGroup]int64) + } + if sequence > s.sequences[group] { + s.sequences[group] = sequence + } +} + +func (s *recordingReleaseSink) Sequence(credentialID, model string) int64 { + s.mu.Lock() + defer s.mu.Unlock() + return s.sequences[ReleaseGroup{CredentialID: credentialID, Model: model}] +} + +func installAccountedScope(t *testing.T, registry *Registry, credentialID, model string) *Scope { + t.Helper() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{CredentialID: credentialID, Model: model, Accounted: true}) + if errInstall != nil { + t.Fatal(errInstall) + } + return scope +} + +func TestRegistryEndMarksOneDirtyGroup(t *testing.T) { + sink := &recordingReleaseSink{} + registry := New() + registry.SetReleaseSink(sink.MarkDirty) + + scope := installAccountedScope(t, registry, "cred-1", "gpt") + scope.End("complete") + scope.End("duplicate") + + if got := sink.Sequence("cred-1", "gpt"); got != 1 { + t.Fatalf("release sequence = %d, want 1", got) + } +} + +func TestUnaccountedScopeDoesNotRelease(t *testing.T) { + sink := &recordingReleaseSink{} + registry := New() + registry.SetReleaseSink(sink.MarkDirty) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{CredentialID: "cred-1", Model: "gpt", Accounted: false}) + if errInstall != nil { + t.Fatal(errInstall) + } + scope.End("observation_complete") + + if got := sink.Sequence("cred-1", "gpt"); got != 0 { + t.Fatalf("release sequence = %d, want 0", got) + } +} + +func TestSetReleaseSinkReplaysExistingSequences(t *testing.T) { + registry := New() + installAccountedScope(t, registry, "cred-1", "gpt").End("complete") + + sink := &recordingReleaseSink{} + registry.SetReleaseSink(sink.MarkDirty) + if got := sink.Sequence("cred-1", "gpt"); got != 1 { + t.Fatalf("replayed release sequence = %d, want 1", got) + } +} diff --git a/sdk/cliproxy/executionregistry/observation.go b/sdk/cliproxy/executionregistry/observation.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/executionregistry/observation.go @@ -0,0 +1,74 @@ +package executionregistry + +import "time" + +// Observation is an immutable in-flight execution snapshot entry. +type Observation struct { + RequestID string + CredentialID string + Model string + RequestKind string + StartedAt time.Time + Accounted bool +} + +// Freeze is an immutable in-flight execution snapshot. +type Freeze struct { + Revision int64 + BarrierRevision int64 + Executions []Observation +} + +// ObserveBarrier records the latest Home observation barrier. +func (r *Registry) ObserveBarrier(revision int64) { + if r == nil || revision <= 0 { + return + } + + r.mu.Lock() + defer r.mu.Unlock() + if revision > r.observedBarrier { + r.observedBarrier = revision + r.pendingBarrierSequence = r.next + } +} + +// FreezeInFlight copies all active executions into an immutable snapshot. +func (r *Registry) FreezeInFlight(_ time.Time) Freeze { + if r == nil { + return Freeze{} + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.observedBarrier > r.publishedBarrier { + blocked := false + for sequence := range r.pending { + if sequence <= r.pendingBarrierSequence { + blocked = true + break + } + } + if !blocked { + r.publishedBarrier = r.observedBarrier + } + } + + r.snapshotRevision++ + freeze := Freeze{ + Revision: r.snapshotRevision, + BarrierRevision: r.publishedBarrier, + Executions: make([]Observation, 0, len(r.scopes)), + } + for _, scope := range r.scopes { + freeze.Executions = append(freeze.Executions, Observation{ + RequestID: scope.spec.RequestID, + CredentialID: scope.spec.CredentialID, + Model: scope.spec.Model, + RequestKind: scope.spec.Kind, + StartedAt: scope.spec.StartedAt, + Accounted: scope.spec.Accounted, + }) + } + return freeze +} diff --git a/sdk/cliproxy/executionregistry/observation_test.go b/sdk/cliproxy/executionregistry/observation_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/executionregistry/observation_test.go @@ -0,0 +1,45 @@ +package executionregistry + +import ( + "testing" + "time" +) + +func TestFreezeInFlightWaitsForPendingBarrierAndCopiesScopes(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + registry.ObserveBarrier(14) + + before := registry.FreezeInFlight(time.Unix(12, 0).UTC()) + if before.BarrierRevision != 0 { + t.Fatalf("barrier before install = %d", before.BarrierRevision) + } + + scope, errInstall := registry.Install(pending, ScopeSpec{ + RequestID: "req-a", CredentialID: "cred", Model: "gpt-5", + Kind: "http", StartedAt: time.Unix(10, 0).UTC(), Accounted: true, + }) + if errInstall != nil { + t.Fatal(errInstall) + } + + after := registry.FreezeInFlight(time.Unix(13, 0).UTC()) + if after.BarrierRevision != 14 || len(after.Executions) != 1 || !after.Executions[0].Accounted { + t.Fatalf("freeze after install = %#v", after) + } + after.Executions[0].RequestID = "mutated" + + copied := registry.FreezeInFlight(time.Unix(13, 0).UTC()) + if len(copied.Executions) != 1 || copied.Executions[0].RequestID != "req-a" { + t.Fatalf("freeze did not copy scope = %#v", copied) + } + + scope.End("completed") + ended := registry.FreezeInFlight(time.Unix(14, 0).UTC()) + if len(ended.Executions) != 0 || ended.Revision <= after.Revision { + t.Fatalf("freeze after end = %#v", ended) + } +} diff --git a/sdk/cliproxy/executionregistry/registry.go b/sdk/cliproxy/executionregistry/registry.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/executionregistry/registry.go @@ -0,0 +1,446 @@ +// Package executionregistry tracks Home-dispatched executions for one subscriber lifetime. +package executionregistry + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "time" + + log "github.com/sirupsen/logrus" +) + +var ( + ErrRegistryNotAccepting = errors.New("execution registry is not accepting dispatches") + ErrRegistryClosed = errors.New("execution registry is closed") + ErrInvalidPendingDispatch = errors.New("invalid pending dispatch") + ErrInvalidExecutionResource = errors.New("invalid execution resource") + ErrExecutionResourceAlreadyBound = errors.New("execution resource is already bound") +) + +// State is the lifecycle state of a Registry. +type State uint32 + +const ( + StateAccepting State = iota + StateDraining + StateClosed +) + +// Registry owns all dispatches accepted during one Home subscriber lifetime. +type Registry struct { + state atomic.Uint32 + + mu sync.Mutex + next uint64 + snapshotRevision int64 + observedBarrier int64 + pendingBarrierSequence uint64 + publishedBarrier int64 + pending map[uint64]*PendingDispatch + scopes map[uint64]*Scope + releaseSequences map[ReleaseGroup]int64 + releaseSink ReleaseSink + changed chan struct{} + + closeMu sync.Mutex + closeStarted bool + closeDone chan struct{} + closeErr error +} + +// PendingDispatch reserves an execution slot until it is installed or ended. +type PendingDispatch struct { + id uint64 + registry *Registry + mu sync.Mutex + once sync.Once +} + +// ScopeSpec describes a Home-dispatched execution. +type ScopeSpec struct { + RequestID string + CredentialID string + Model string + Kind string + StartedAt time.Time + Accounted bool +} + +// ReleaseGroup identifies the cumulative release sequence for one accounted credential and model. +type ReleaseGroup struct { + CredentialID string + Model string +} + +// ReleaseTicket completes after Home acknowledges a cumulative release sequence. +type ReleaseTicket struct { + Group ReleaseGroup + Sequence int64 + done <-chan struct{} +} + +// NewReleaseTicket creates a ticket backed by done. A nil done channel represents +// a release sink that does not support acknowledgements. +func NewReleaseTicket(group ReleaseGroup, sequence int64, done <-chan struct{}) *ReleaseTicket { + if sequence <= 0 || done == nil { + return nil + } + return &ReleaseTicket{Group: group, Sequence: sequence, done: done} +} + +// Wait blocks until Home acknowledges the release or ctx expires. +func (t *ReleaseTicket) Wait(ctx context.Context) error { + if t == nil || t.done == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + select { + case <-t.done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// ReleaseSink receives the latest cumulative sequence for a release group and +// optionally returns an acknowledgement ticket. +type ReleaseSink func(ReleaseGroup, int64) *ReleaseTicket + +// Scope owns the resource for one installed execution. +type Scope struct { + id uint64 + registry *Registry + spec ScopeSpec + + mu sync.Mutex + closeFn func() error + closeDone chan struct{} + releaseTicket *ReleaseTicket + active bool + ended sync.Once +} + +// New creates an accepting registry. +func New() *Registry { + registry := &Registry{ + pending: make(map[uint64]*PendingDispatch), + scopes: make(map[uint64]*Scope), + releaseSequences: make(map[ReleaseGroup]int64), + changed: make(chan struct{}), + } + registry.state.Store(uint32(StateAccepting)) + return registry +} + +// BeginDispatch reserves a dispatch token while the registry accepts traffic. +func (r *Registry) BeginDispatch() (*PendingDispatch, error) { + if r == nil || State(r.state.Load()) != StateAccepting { + return nil, ErrRegistryNotAccepting + } + + r.mu.Lock() + defer r.mu.Unlock() + if State(r.state.Load()) != StateAccepting { + return nil, ErrRegistryNotAccepting + } + + r.next++ + pending := &PendingDispatch{id: r.next, registry: r} + r.pending[pending.id] = pending + return pending, nil +} + +// End releases a dispatch token that was not installed. +func (p *PendingDispatch) End() { + if p == nil || p.registry == nil { + return + } + + p.mu.Lock() + defer p.mu.Unlock() + p.once.Do(func() { + p.registry.mu.Lock() + delete(p.registry.pending, p.id) + p.registry.signalLocked() + p.registry.mu.Unlock() + }) +} + +// Install atomically turns a pending dispatch token into an active execution scope. +func (r *Registry) Install(pending *PendingDispatch, spec ScopeSpec) (*Scope, error) { + if r == nil || pending == nil || pending.registry != r { + return nil, ErrInvalidPendingDispatch + } + + pending.mu.Lock() + defer pending.mu.Unlock() + r.mu.Lock() + defer r.mu.Unlock() + + if State(r.state.Load()) != StateAccepting { + pending.once.Do(func() {}) + delete(r.pending, pending.id) + r.signalLocked() + return nil, ErrRegistryNotAccepting + } + if _, exists := r.pending[pending.id]; !exists { + return nil, ErrInvalidPendingDispatch + } + + pending.once.Do(func() {}) + delete(r.pending, pending.id) + scope := &Scope{id: pending.id, registry: r, spec: spec, active: true} + r.scopes[scope.id] = scope + r.signalLocked() + return scope, nil +} + +// SetReleaseSink replaces the cumulative release sink and replays every known group. +// Legacy callbacks remain supported but cannot provide acknowledgement tickets. +func (r *Registry) SetReleaseSink(rawSink any) { + if r == nil { + return + } + + var sink ReleaseSink + switch typed := rawSink.(type) { + case nil: + case ReleaseSink: + sink = typed + case func(ReleaseGroup, int64) *ReleaseTicket: + sink = ReleaseSink(typed) + case func(ReleaseGroup, int64): + sink = func(group ReleaseGroup, sequence int64) *ReleaseTicket { + typed(group, sequence) + return nil + } + default: + return + } + + r.mu.Lock() + r.releaseSink = sink + sequences := make(map[ReleaseGroup]int64, len(r.releaseSequences)) + for group, sequence := range r.releaseSequences { + sequences[group] = sequence + } + r.mu.Unlock() + + if sink == nil { + return + } + for group, sequence := range sequences { + if sequence > 0 { + sink(group, sequence) + } + } +} + +// Bind attaches the execution resource. A scope accepts exactly one resource. +func (s *Scope) Bind(closeFn func() error) error { + if s == nil || s.registry == nil || closeFn == nil { + return ErrInvalidExecutionResource + } + + s.registry.mu.Lock() + defer s.registry.mu.Unlock() + if State(s.registry.state.Load()) != StateAccepting || !s.active { + return ErrRegistryNotAccepting + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.closeFn != nil || s.closeDone != nil { + return ErrExecutionResourceAlreadyBound + } + s.closeFn = closeFn + return nil +} + +// End closes the bound resource and releases this execution scope exactly once. +func (s *Scope) End(reason string) { + _ = s.EndWithRelease(reason) +} + +// EndWithRelease closes the scope and returns the release acknowledgement ticket. +// The release sink is invoked without the registry mutex held. +func (s *Scope) EndWithRelease(_ string) *ReleaseTicket { + if s == nil || s.registry == nil { + return nil + } + + var ticket *ReleaseTicket + s.ended.Do(func() { + s.registry.mu.Lock() + s.mu.Lock() + s.active = false + s.mu.Unlock() + s.registry.mu.Unlock() + + s.waitForBoundResourceClose() + + s.registry.mu.Lock() + releaseSink, releaseGroup, releaseSequence := s.registry.markReleasedLocked(s) + s.registry.mu.Unlock() + + if releaseSink != nil && releaseSequence > 0 { + ticket = releaseSink(releaseGroup, releaseSequence) + } + + s.mu.Lock() + s.releaseTicket = ticket + s.mu.Unlock() + + s.registry.mu.Lock() + delete(s.registry.scopes, s.id) + s.registry.signalLocked() + s.registry.mu.Unlock() + }) + + s.mu.Lock() + ticket = s.releaseTicket + s.mu.Unlock() + return ticket +} + +func (r *Registry) markReleasedLocked(scope *Scope) (ReleaseSink, ReleaseGroup, int64) { + if scope == nil || !scope.spec.Accounted { + return nil, ReleaseGroup{}, 0 + } + group := ReleaseGroup{CredentialID: scope.spec.CredentialID, Model: scope.spec.Model} + r.releaseSequences[group]++ + return r.releaseSink, group, r.releaseSequences[group] +} + +func (s *Scope) startBoundResourceClose() <-chan struct{} { + s.mu.Lock() + defer s.mu.Unlock() + if s.closeDone != nil { + return s.closeDone + } + closeFn := s.closeFn + if closeFn == nil { + return nil + } + closeDone := make(chan struct{}) + s.closeFn = nil + s.closeDone = closeDone + go func() { + s.closeResource(closeFn) + close(closeDone) + }() + return closeDone +} + +func (s *Scope) waitForBoundResourceClose() { + if closeDone := s.startBoundResourceClose(); closeDone != nil { + <-closeDone + } +} + +func (s *Scope) closeResource(closeFn func() error) { + if closeFn == nil { + return + } + if errClose := closeFn(); errClose != nil { + log.WithError(errClose).Warn("Home execution resource close failed") + } +} + +// Drain rejects new work, cancels active resources, and waits for all owners to end. +func (r *Registry) Drain(ctx context.Context) error { + if r == nil { + return ErrRegistryClosed + } + if ctx == nil { + ctx = context.Background() + } + + if !r.state.CompareAndSwap(uint32(StateAccepting), uint32(StateDraining)) && State(r.state.Load()) != StateDraining { + return ErrRegistryClosed + } + + r.mu.Lock() + scopes := make([]*Scope, 0, len(r.scopes)) + for _, scope := range r.scopes { + scopes = append(scopes, scope) + } + r.mu.Unlock() + + for _, scope := range scopes { + scope.startBoundResourceClose() + } + + r.mu.Lock() + for len(r.pending) != 0 || len(r.scopes) != 0 { + changed := r.changed + r.mu.Unlock() + select { + case <-ctx.Done(): + return ctx.Err() + case <-changed: + } + r.mu.Lock() + } + r.state.Store(uint32(StateClosed)) + r.mu.Unlock() + return nil +} + +// Close permanently rejects new work and closes every currently bound resource. +func (r *Registry) Close() error { + if r == nil { + return ErrRegistryClosed + } + + r.closeMu.Lock() + if r.closeStarted { + closeDone := r.closeDone + r.closeMu.Unlock() + <-closeDone + r.closeMu.Lock() + errClose := r.closeErr + r.closeMu.Unlock() + return errClose + } + if State(r.state.Load()) == StateClosed { + r.closeMu.Unlock() + return nil + } + r.closeStarted = true + r.closeDone = make(chan struct{}) + closeDone := r.closeDone + r.closeMu.Unlock() + + for { + state := State(r.state.Load()) + if state == StateClosed || r.state.CompareAndSwap(uint32(state), uint32(StateClosed)) { + break + } + } + + r.mu.Lock() + scopes := make([]*Scope, 0, len(r.scopes)) + for _, scope := range r.scopes { + scopes = append(scopes, scope) + } + r.mu.Unlock() + for _, scope := range scopes { + scope.waitForBoundResourceClose() + } + + r.closeMu.Lock() + errClose := r.closeErr + close(closeDone) + r.closeMu.Unlock() + return errClose +} + +func (r *Registry) signalLocked() { + close(r.changed) + r.changed = make(chan struct{}) +} diff --git a/sdk/cliproxy/executionregistry/registry_test.go b/sdk/cliproxy/executionregistry/registry_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/executionregistry/registry_test.go @@ -0,0 +1,349 @@ +package executionregistry + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" +) + +func TestDrainRejectsLateInstallAndCancelsBoundScopes(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{RequestID: "req-1", CredentialID: "cred-1", Model: "gpt", Kind: "http", StartedAt: time.Now()}) + if errInstall != nil { + t.Fatal(errInstall) + } + closed := atomic.Int32{} + if errBind := scope.Bind(func() error { + closed.Add(1) + go scope.End("canceled") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if errDrain := registry.Drain(ctx); errDrain != nil { + t.Fatal(errDrain) + } + if closed.Load() != 1 { + t.Fatalf("close calls = %d", closed.Load()) + } + if _, errLate := registry.BeginDispatch(); !errors.Is(errLate, ErrRegistryNotAccepting) { + t.Fatalf("late dispatch error = %v", errLate) + } +} + +func TestScopeEndIsExactlyOnce(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + closed := atomic.Int32{} + if errBind := scope.Bind(func() error { + closed.Add(1) + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + done := make(chan struct{}) + go func() { + scope.End("complete") + close(done) + }() + scope.End("duplicate") + <-done + if closed.Load() != 1 { + t.Fatalf("close calls = %d, want 1", closed.Load()) + } +} + +func TestDrainWaitsForPendingDispatch(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + done := make(chan error, 1) + go func() { done <- registry.Drain(ctx) }() + + select { + case errDrain := <-done: + t.Fatalf("Drain() returned before pending dispatch ended: %v", errDrain) + case <-time.After(20 * time.Millisecond): + } + pending.End() + if errDrain := <-done; errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestDrainReturnsWhenBlockingResourceCloseExceedsContext(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + started := make(chan struct{}) + release := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(started) + <-release + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + errDrain := registry.Drain(ctx) + if !errors.Is(errDrain, context.DeadlineExceeded) { + t.Fatalf("Drain() error = %v, want context deadline exceeded", errDrain) + } + select { + case <-started: + default: + t.Fatal("Drain() did not start closing the bound resource") + } + if state := State(registry.state.Load()); state != StateDraining { + t.Fatalf("registry state = %v, want draining", state) + } + + ended := make(chan struct{}) + go func() { + scope.End("canceled") + close(ended) + }() + close(release) + select { + case <-ended: + case <-time.After(time.Second): + t.Fatal("Scope.End() did not wait for resource close completion") + } + if errDrain = registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() after resource close = %v", errDrain) + } +} + +func TestDrainWaitsForBlockingResourceClose(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + started := make(chan struct{}) + release := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(started) + <-release + return nil + }); errBind != nil { + t.Fatal(errBind) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + done := make(chan error, 1) + go func() { done <- registry.Drain(ctx) }() + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("Drain() did not close the bound resource") + } + go scope.End("canceled") + select { + case errDrain := <-done: + t.Fatalf("Drain() returned before the resource close completed: %v", errDrain) + case <-time.After(20 * time.Millisecond): + } + close(release) + if errDrain := <-done; errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestConcurrentDrainWaitsForBlockingResourceClose(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + started := make(chan struct{}) + release := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(started) + <-release + return nil + }); errBind != nil { + t.Fatal(errBind) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + firstDrain := make(chan error, 1) + go func() { firstDrain <- registry.Drain(ctx) }() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("first Drain() did not close the bound resource") + } + ended := make(chan struct{}) + go func() { + scope.End("canceled") + close(ended) + }() + select { + case <-ended: + t.Fatal("Scope.End() returned before the resource close completed") + case <-time.After(20 * time.Millisecond): + } + secondDrain := make(chan error, 1) + go func() { secondDrain <- registry.Drain(ctx) }() + select { + case errDrain := <-secondDrain: + t.Fatalf("second Drain() returned before resource close completed: %v", errDrain) + case <-time.After(20 * time.Millisecond): + } + close(release) + select { + case <-ended: + case <-time.After(time.Second): + t.Fatal("Scope.End() did not complete after the resource close") + } + if errDrain := <-firstDrain; errDrain != nil { + t.Fatalf("first Drain() error = %v", errDrain) + } + if errDrain := <-secondDrain; errDrain != nil { + t.Fatalf("second Drain() error = %v", errDrain) + } +} + +func TestConcurrentCloseWaitsForBlockingResourceClose(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + started := make(chan struct{}) + release := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(started) + <-release + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + firstClose := make(chan error, 1) + go func() { firstClose <- registry.Close() }() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("first Close() did not close the bound resource") + } + + secondClose := make(chan error, 1) + go func() { secondClose <- registry.Close() }() + select { + case errClose := <-secondClose: + t.Fatalf("second Close() returned before resource close completed: %v", errClose) + case <-time.After(20 * time.Millisecond): + } + + close(release) + if errClose := <-firstClose; errClose != nil { + t.Fatalf("first Close() error = %v", errClose) + } + if errClose := <-secondClose; errClose != nil { + t.Fatalf("second Close() error = %v", errClose) + } +} + +func TestDrainRejectsLateBind(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + done := make(chan error, 1) + go func() { done <- registry.Drain(ctx) }() + + deadline := time.After(time.Second) + for State(registry.state.Load()) == StateAccepting { + select { + case <-deadline: + t.Fatal("registry did not begin draining") + default: + time.Sleep(time.Millisecond) + } + } + if errBind := scope.Bind(func() error { return nil }); !errors.Is(errBind, ErrRegistryNotAccepting) { + t.Fatalf("Bind() error = %v, want ErrRegistryNotAccepting", errBind) + } + scope.End("canceled") + if errDrain := <-done; errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestDrainRejectsLateInstall(t *testing.T) { + registry := New() + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + done := make(chan error, 1) + go func() { done <- registry.Drain(ctx) }() + + deadline := time.After(time.Second) + for State(registry.state.Load()) == StateAccepting { + select { + case <-deadline: + t.Fatal("registry did not begin draining") + default: + time.Sleep(time.Millisecond) + } + } + if _, errInstall := registry.Install(pending, ScopeSpec{}); !errors.Is(errInstall, ErrRegistryNotAccepting) { + t.Fatalf("Install() error = %v, want ErrRegistryNotAccepting", errInstall) + } + if errDrain := <-done; errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} diff --git a/sdk/cliproxy/executor/lifecycle.go b/sdk/cliproxy/executor/lifecycle.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/executor/lifecycle.go @@ -0,0 +1,33 @@ +package executor + +import ( + "errors" + "io" + "sync" +) + +// ExecutionLifecycle owns resources associated with an execution attempt. +type ExecutionLifecycle interface { + Bind(func() error) error + End(string) +} + +// BindExecutionResource binds a closer to the execution lifecycle. +func BindExecutionResource(opts Options, closer io.Closer) error { + if opts.ExecutionLifecycle == nil || closer == nil { + return nil + } + + var closeOnce sync.Once + var closeErr error + closeResource := func() error { + closeOnce.Do(func() { + closeErr = closer.Close() + }) + return closeErr + } + if errBind := opts.ExecutionLifecycle.Bind(closeResource); errBind != nil { + return errors.Join(errBind, closeResource()) + } + return nil +} diff --git a/sdk/cliproxy/executor/lifecycle_test.go b/sdk/cliproxy/executor/lifecycle_test.go new file mode 100644 --- /dev/null +++ b/sdk/cliproxy/executor/lifecycle_test.go @@ -0,0 +1,69 @@ +package executor + +import ( + "errors" + "sync/atomic" + "testing" +) + +type lifecycleRecorder struct { + closeFn func() error +} + +func (r *lifecycleRecorder) Bind(closeFn func() error) error { + r.closeFn = closeFn + return nil +} + +func (*lifecycleRecorder) End(string) {} + +type lifecycleCloser struct { + calls atomic.Int32 +} + +func (c *lifecycleCloser) Close() error { + c.calls.Add(1) + return nil +} + +func TestBindExecutionResourceClosesResourceOnce(t *testing.T) { + lifecycle := &lifecycleRecorder{} + closer := &lifecycleCloser{} + + if errBind := BindExecutionResource(Options{ExecutionLifecycle: lifecycle}, closer); errBind != nil { + t.Fatalf("BindExecutionResource() error = %v", errBind) + } + if lifecycle.closeFn == nil { + t.Fatal("BindExecutionResource() did not bind a closer") + } + if errClose := lifecycle.closeFn(); errClose != nil { + t.Fatalf("first close error = %v", errClose) + } + if errClose := lifecycle.closeFn(); errClose != nil { + t.Fatalf("second close error = %v", errClose) + } + if got := closer.calls.Load(); got != 1 { + t.Fatalf("closer calls = %d, want 1", got) + } +} + +func TestBindExecutionResourceClosesWhenBindFails(t *testing.T) { + want := errors.New("selection ended") + lifecycle := &failingLifecycle{err: want} + closer := &lifecycleCloser{} + + errBind := BindExecutionResource(Options{ExecutionLifecycle: lifecycle}, closer) + if !errors.Is(errBind, want) { + t.Fatalf("BindExecutionResource() error = %v, want %v", errBind, want) + } + if got := closer.calls.Load(); got != 1 { + t.Fatalf("closer calls = %d, want 1", got) + } +} + +type failingLifecycle struct { + err error +} + +func (l *failingLifecycle) Bind(func() error) error { return l.err } +func (*failingLifecycle) End(string) {} diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -112,6 +112,8 @@ Metadata map[string]any // RequestAfterAuthInterceptor runs after credential selection and before executor translation. RequestAfterAuthInterceptor RequestAfterAuthInterceptor + // ExecutionLifecycle owns Home-dispatched execution resources. Executors must not add it to request metadata. + ExecutionLifecycle ExecutionLifecycle } // ResponseFormatOrSource returns the response target format for an execution. diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -3,27 +3,162 @@ import ( "bytes" "context" + "encoding/json" "errors" "fmt" + "maps" "net/http" "net/http/httptest" "strings" "sync" + "sync/atomic" "testing" "time" "unicode/utf8" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" requestlogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" "github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" "github.com/tidwall/gjson" ) + +type homeResponsesWebsocketDispatcher struct { + calls atomic.Int32 +} + +func (*homeResponsesWebsocketDispatcher) HeartbeatOK() bool { return true } + +func (d *homeResponsesWebsocketDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + d.calls.Add(1) + return json.Marshal(coreauth.Auth{ + ID: "home-responses-websocket-auth", + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{ + "websockets": "true", + }, + }) +} + +func (*homeResponsesWebsocketDispatcher) AbortAmbiguousDispatch() {} + +type homeResponsesWebsocketExecutor struct { + calls atomic.Int32 + metadata []map[string]any + mu sync.Mutex +} + +func (*homeResponsesWebsocketExecutor) Identifier() string { return "codex" } + +func (*homeResponsesWebsocketExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (e *homeResponsesWebsocketExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, _ coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { + e.calls.Add(1) + e.mu.Lock() + e.metadata = append(e.metadata, maps.Clone(opts.Metadata)) + e.mu.Unlock() + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + chunks := make(chan coreexecutor.StreamChunk, 1) + chunks <- coreexecutor.StreamChunk{Payload: []byte(`{"type":"response.completed","response":{"id":"home-response","output":[]}}`)} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil +} + +func (*homeResponsesWebsocketExecutor) Refresh(context.Context, *coreauth.Auth) (*coreauth.Auth, error) { + return nil, errors.New("not implemented") +} + +func (*homeResponsesWebsocketExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (*homeResponsesWebsocketExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +func TestResponsesWebsocketHomeSelectedAuthCallbackPinsAndReusesFirstSelection(t *testing.T) { + gin.SetMode(gin.TestMode) + + dispatcher := &homeResponsesWebsocketDispatcher{} + executor := &homeResponsesWebsocketExecutor{} + manager := coreauth.NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + registry.GetGlobalRegistry().RegisterClient("home-responses-websocket-auth", "codex", []*registry.ModelInfo{{ID: "gpt-5.4"}}) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIResponsesAPIHandler(base) + router := gin.New() + router.GET("/v1/responses/ws", h.ResponsesWebsocket) + server := httptest.NewServer(router) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws" + conn, _, errDial := websocket.DefaultDialer.Dial(wsURL, nil) + if errDial != nil { + t.Fatalf("dial websocket: %v", errDial) + } + defer func() { + if errClose := conn.Close(); errClose != nil { + t.Errorf("close websocket: %v", errClose) + } + }() + + requests := []string{ + `{"type":"response.create","model":"gpt-5.4","input":[]}`, + `{"type":"response.create","model":"gpt-5.4","input":[]}`, + } + for index, request := range requests { + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil { + t.Fatalf("write websocket request %d: %v", index+1, errWrite) + } + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read websocket response %d: %v", index+1, errRead) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeCompleted { + t.Fatalf("response %d type = %q, want %q: %s", index+1, got, wsEventTypeCompleted, payload) + } + if index == 0 { + executor.mu.Lock() + firstMetadata := maps.Clone(executor.metadata[0]) + executor.mu.Unlock() + sessionID, _ := firstMetadata[coreexecutor.ExecutionSessionMetadataKey].(string) + if _, ok := manager.GetExecutionSessionAuthByID(sessionID, "home-responses-websocket-auth"); !ok { + t.Fatal("first selected-auth callback did not stage the session runtime auth") + } + } + } + + executor.mu.Lock() + metadata := append([]map[string]any(nil), executor.metadata...) + executor.mu.Unlock() + if len(metadata) != 2 { + t.Fatalf("executor metadata calls = %d, want 2", len(metadata)) + } + if got := metadata[1][coreexecutor.PinnedAuthMetadataKey]; got != "home-responses-websocket-auth" { + t.Fatalf("second turn pinned auth metadata = %#v, want home selected auth (first metadata: %#v, second metadata: %#v)", got, metadata[0], metadata[1]) + } + if got := dispatcher.calls.Load(); got != 1 { + t.Fatalf("Home RPOP calls = %d, want 1 after selected-auth callback pin", got) + } + if got := executor.calls.Load(); got != 2 { + t.Fatalf("executor calls = %d, want 2", got) + } +} func TestWriteWebsocketCloseForUpstreamErrorMirrorsMessageTooBig(t *testing.T) { tests := []struct {