From 4a6304440b855702b05b5924e6bad048fa4f2310 Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 11:31:24 +0800 Subject: [PATCH 01/11] feat(sdk): add headless embedded runtime Expose a lifecycle for host-owned auth managers that reuses native executor and model registration without constructing the HTTP server or file watcher. Support targeted auth synchronization, active-model reconciliation, and cleanup of runtime-owned projections.\n\nVerified with focused SDK tests, race tests, the server build, and go test ./.... --- sdk/cliproxy/embedded_runtime.go | 291 ++++++++++++++++++++++++++ sdk/cliproxy/embedded_runtime_test.go | 149 +++++++++++++ sdk/cliproxy/service.go | 28 +++ 3 files changed, 468 insertions(+) create mode 100644 sdk/cliproxy/embedded_runtime.go create mode 100644 sdk/cliproxy/embedded_runtime_test.go diff --git a/sdk/cliproxy/embedded_runtime.go b/sdk/cliproxy/embedded_runtime.go new file mode 100644 index 00000000000..ad823f6e553 --- /dev/null +++ b/sdk/cliproxy/embedded_runtime.go @@ -0,0 +1,291 @@ +package cliproxy + +import ( + "context" + "fmt" + "reflect" + "strings" + "sync" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" +) + +// EmbeddedRuntime provides native provider execution and model registration +// without starting the CLIProxyAPI HTTP server or file watcher. +type EmbeddedRuntime struct { + manager *coreauth.Manager + service *Service + + mu sync.Mutex + started bool + closed bool + modelsReconciled bool + activeModels map[string]struct{} + registeredAuthIDs map[string]struct{} + installedExecutors map[string]coreauth.ProviderExecutor +} + +// NewEmbeddedRuntime creates a headless runtime backed by a host-owned auth manager. +func NewEmbeddedRuntime(cfg *config.Config, manager *coreauth.Manager) (*EmbeddedRuntime, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy: embedded runtime configuration is required") + } + if manager == nil { + return nil, fmt.Errorf("cliproxy: embedded runtime auth manager is required") + } + + manager.SetRoundTripperProvider(newDefaultRoundTripperProvider()) + manager.SetConfig(cfg) + manager.SetOAuthModelAlias(cfg.OAuthModelAlias) + + return &EmbeddedRuntime{ + manager: manager, + service: &Service{cfg: cfg, coreManager: manager}, + registeredAuthIDs: make(map[string]struct{}), + installedExecutors: make(map[string]coreauth.ProviderExecutor), + }, nil +} + +// Start registers native executors and model projections for existing auths. +func (r *EmbeddedRuntime) Start(ctx context.Context) error { + if r == nil { + return fmt.Errorf("cliproxy: embedded runtime is nil") + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return errContext + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.closed { + return fmt.Errorf("cliproxy: embedded runtime is closed") + } + if r.started { + return nil + } + + auths := r.manager.List() + r.service.registerAvailableExecutors(ctx, executorRegistrationOptions{ + auths: auths, + forceReplaceAuths: true, + }) + for _, auth := range auths { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + r.trackExecutorForAuthLocked(auth) + r.registerAuthModelsLocked(ctx, auth) + } + r.started = true + return nil +} + +// SyncAuth adds or updates one auth in the native runtime without persisting it. +func (r *EmbeddedRuntime) SyncAuth(ctx context.Context, auth *coreauth.Auth) error { + if r == nil { + return fmt.Errorf("cliproxy: embedded runtime is nil") + } + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return fmt.Errorf("cliproxy: embedded runtime auth ID is required") + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return errContext + } + + r.mu.Lock() + defer r.mu.Unlock() + if errState := r.requireStartedLocked(); errState != nil { + return errState + } + + runtimeCtx := coreauth.WithSkipPersist(ctx) + prepared := r.service.prepareCoreAuthForModelRegistration(runtimeCtx, auth) + if prepared == nil { + return fmt.Errorf("cliproxy: failed to synchronize auth %q", auth.ID) + } + r.trackExecutorForAuthLocked(prepared) + r.registerAuthModelsLocked(runtimeCtx, prepared) + return nil +} + +// RemoveAuth removes one auth and its model projection from the native runtime. +func (r *EmbeddedRuntime) RemoveAuth(ctx context.Context, authID string) error { + if r == nil { + return fmt.Errorf("cliproxy: embedded runtime is nil") + } + authID = strings.TrimSpace(authID) + if authID == "" { + return fmt.Errorf("cliproxy: embedded runtime auth ID is required") + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return errContext + } + + r.mu.Lock() + defer r.mu.Unlock() + if errState := r.requireStartedLocked(); errState != nil { + return errState + } + + r.service.applyCoreAuthRemoval(coreauth.WithSkipPersist(ctx), authID) + delete(r.registeredAuthIDs, authID) + return nil +} + +// ReconcileModels limits each auth's native model projection to active model IDs. +func (r *EmbeddedRuntime) ReconcileModels(ctx context.Context, activeModels []string) error { + if r == nil { + return fmt.Errorf("cliproxy: embedded runtime is nil") + } + if ctx == nil { + ctx = context.Background() + } + if errContext := ctx.Err(); errContext != nil { + return errContext + } + + r.mu.Lock() + defer r.mu.Unlock() + if errState := r.requireStartedLocked(); errState != nil { + return errState + } + + r.activeModels = make(map[string]struct{}, len(activeModels)) + for _, modelID := range activeModels { + modelID = strings.TrimSpace(modelID) + if modelID != "" { + r.activeModels[modelID] = struct{}{} + } + } + r.modelsReconciled = true + + for _, auth := range r.manager.List() { + if errContext := ctx.Err(); errContext != nil { + return errContext + } + r.registerAuthModelsLocked(ctx, auth) + } + return nil +} + +// Close removes model projections owned by the runtime. Host-owned auths remain registered. +func (r *EmbeddedRuntime) Close(ctx context.Context) error { + if r == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + + r.mu.Lock() + defer r.mu.Unlock() + if r.closed { + return nil + } + for authID := range r.registeredAuthIDs { + registry.GetGlobalRegistry().UnregisterClient(authID) + r.manager.RefreshSchedulerEntry(authID) + } + for provider, installed := range r.installedExecutors { + current, okExecutor := r.manager.Executor(provider) + if !okExecutor || !sameProviderExecutor(current, installed) { + continue + } + if closer, okCloser := current.(coreauth.ExecutionSessionCloser); okCloser { + closer.CloseExecutionSession(coreauth.CloseAllExecutionSessionsID) + } + r.manager.UnregisterExecutor(provider) + } + r.registeredAuthIDs = make(map[string]struct{}) + r.installedExecutors = make(map[string]coreauth.ProviderExecutor) + r.closed = true + return nil +} + +// ServerStarted reports whether an HTTP server was constructed by this runtime. +func (r *EmbeddedRuntime) ServerStarted() bool { + if r == nil { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + return r.service != nil && r.service.server != nil +} + +// WatcherStarted reports whether a file watcher was constructed by this runtime. +func (r *EmbeddedRuntime) WatcherStarted() bool { + if r == nil { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + return r.service != nil && r.service.watcher != nil +} + +func (r *EmbeddedRuntime) requireStartedLocked() error { + if r.closed { + return fmt.Errorf("cliproxy: embedded runtime is closed") + } + if !r.started { + return fmt.Errorf("cliproxy: embedded runtime is not started") + } + return nil +} + +func (r *EmbeddedRuntime) registerAuthModelsLocked(ctx context.Context, auth *coreauth.Auth) { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + return + } + var activeModels map[string]struct{} + if r.modelsReconciled { + activeModels = r.activeModels + } + r.service.reconcileModelsForAuth(ctx, auth, activeModels) + r.registeredAuthIDs[auth.ID] = struct{}{} +} + +func (r *EmbeddedRuntime) trackExecutorForAuthLocked(auth *coreauth.Auth) { + provider := embeddedRuntimeExecutorProvider(auth) + if provider == "" { + return + } + executor, okExecutor := r.manager.Executor(provider) + if okExecutor && executor != nil { + r.installedExecutors[provider] = executor + } +} + +func embeddedRuntimeExecutorProvider(auth *coreauth.Auth) string { + if auth == nil { + return "" + } + if providerKey, _, isCompat := openAICompatInfoFromAuth(auth); isCompat { + if providerKey != "" { + return providerKey + } + return "openai-compatibility" + } + return strings.ToLower(strings.TrimSpace(auth.Provider)) +} + +func sameProviderExecutor(first, second coreauth.ProviderExecutor) bool { + if first == nil || second == nil { + return first == nil && second == nil + } + firstType := reflect.TypeOf(first) + if firstType != reflect.TypeOf(second) || !firstType.Comparable() { + return false + } + return first == second +} diff --git a/sdk/cliproxy/embedded_runtime_test.go b/sdk/cliproxy/embedded_runtime_test.go new file mode 100644 index 00000000000..1c29827fa05 --- /dev/null +++ b/sdk/cliproxy/embedded_runtime_test.go @@ -0,0 +1,149 @@ +package cliproxy_test + +import ( + "context" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" +) + +func TestEmbeddedRuntimeRegistersNativeExecutorsWithoutServer(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ID: "codex-1", Provider: "codex"}); errRegister != nil { + t.Fatalf("register codex auth: %v", errRegister) + } + if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ID: "claude-1", Provider: "claude"}); errRegister != nil { + t.Fatalf("register claude auth: %v", errRegister) + } + + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + t.Cleanup(func() { + if errClose := runtime.Close(context.Background()); errClose != nil { + t.Errorf("close embedded runtime: %v", errClose) + } + }) + + codexExecutor, okCodex := manager.Executor("codex") + if !okCodex { + t.Fatal("codex executor was not registered") + } + if _, okNative := codexExecutor.(*runtimeexecutor.CodexAutoExecutor); !okNative { + t.Fatalf("codex executor type = %T, want *executor.CodexAutoExecutor", codexExecutor) + } + + claudeExecutor, okClaude := manager.Executor("claude") + if !okClaude { + t.Fatal("claude executor was not registered") + } + if _, okNative := claudeExecutor.(*runtimeexecutor.ClaudeExecutor); !okNative { + t.Fatalf("claude executor type = %T, want *executor.ClaudeExecutor", claudeExecutor) + } + + if runtime.ServerStarted() { + t.Fatal("embedded runtime started an HTTP server") + } + if runtime.WatcherStarted() { + t.Fatal("embedded runtime started a file watcher") + } +} + +func TestEmbeddedRuntimeSyncRemoveAndReconcileModels(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + t.Cleanup(func() { + if errClose := runtime.Close(context.Background()); errClose != nil { + t.Errorf("close embedded runtime: %v", errClose) + } + }) + + auth := &coreauth.Auth{ID: "codex-sync-1", Provider: "codex"} + if errSync := runtime.SyncAuth(t.Context(), auth); errSync != nil { + t.Fatalf("sync auth: %v", errSync) + } + if _, okAuth := manager.GetByID(auth.ID); !okAuth { + t.Fatal("synced auth is missing from manager") + } + + models := registry.GetCodexProModels() + if len(models) < 2 { + t.Fatalf("codex model fixture count = %d, want at least 2", len(models)) + } + keptModel := models[0].ID + removedModel := models[1].ID + if !cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, keptModel) { + t.Fatalf("synced auth does not support native model %q", keptModel) + } + + if errReconcile := runtime.ReconcileModels(t.Context(), []string{keptModel}); errReconcile != nil { + t.Fatalf("reconcile models: %v", errReconcile) + } + if !cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, keptModel) { + t.Fatalf("reconciled auth does not support active model %q", keptModel) + } + if cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, removedModel) { + t.Fatalf("reconciled auth still supports inactive model %q", removedModel) + } + + if errRemove := runtime.RemoveAuth(t.Context(), auth.ID); errRemove != nil { + t.Fatalf("remove auth: %v", errRemove) + } + if _, okAuth := manager.GetByID(auth.ID); okAuth { + t.Fatal("removed auth is still present in manager") + } + if cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, keptModel) { + t.Fatalf("removed auth still supports model %q", keptModel) + } +} + +func TestEmbeddedRuntimeCloseUnregistersModelsWithoutRemovingAuth(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ID: "claude-close-1", Provider: "claude"} + if _, errRegister := manager.Register(t.Context(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + + models := registry.GetClaudeModels() + if len(models) == 0 { + t.Fatal("claude model fixture is empty") + } + model := models[0].ID + if !cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, model) { + t.Fatalf("started runtime does not register model %q", model) + } + + if errClose := runtime.Close(t.Context()); errClose != nil { + t.Fatalf("close embedded runtime: %v", errClose) + } + if _, okAuth := manager.GetByID(auth.ID); !okAuth { + t.Fatal("closing runtime removed host-owned auth") + } + if cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, model) { + t.Fatalf("closed runtime still registers model %q", model) + } + if _, okExecutor := manager.Executor("claude"); okExecutor { + t.Fatal("closed runtime still exposes its claude executor") + } +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 61c6dad812d..c82190d8d46 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -1906,6 +1906,34 @@ func (s *Service) registerModelsForAuth(ctx context.Context, a *coreauth.Auth) { s.registerModelsForAuthWithCache(ctx, a, nil) } +// reconcileModelsForAuth rebuilds one auth's native model projection and, +// when activeModels is non-nil, prunes models the embedding host has not activated. +func (s *Service) reconcileModelsForAuth(ctx context.Context, a *coreauth.Auth, activeModels map[string]struct{}) { + if s == nil || s.coreManager == nil || a == nil || strings.TrimSpace(a.ID) == "" { + return + } + s.registerModelsForAuth(ctx, a) + if activeModels != nil { + models := registry.GetGlobalRegistry().GetModelsForClient(a.ID) + filtered := make([]*ModelInfo, 0, len(models)) + for _, model := range models { + if model == nil { + continue + } + if _, active := activeModels[strings.TrimSpace(model.ID)]; active { + filtered = append(filtered, model) + } + } + provider := strings.ToLower(strings.TrimSpace(a.Provider)) + if providerKey, _, isCompat := openAICompatInfoFromAuth(a); isCompat && providerKey != "" { + provider = providerKey + } + s.registerResolvedModelsForAuth(a, provider, filtered) + } + s.coreManager.ReconcileRegistryModelStates(ctx, a.ID) + s.coreManager.RefreshSchedulerEntry(a.ID) +} + func (s *Service) registerModelsForAuthWithCache(ctx context.Context, a *coreauth.Auth, compatCache *openAICompatibilityRegistrationCache) { if a == nil || a.ID == "" { return From bd56e0815c33db3b4620c8c42cab7d3acf19f8fc Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 16:05:41 +0800 Subject: [PATCH 02/11] fix(sdk): harden embedded runtime ownership Preserve host-owned executors and transports, make executor ownership atomic, roll back canceled starts, avoid persistence and hooks during runtime-only reconciliation, and prevent model projection without an executable provider. --- sdk/cliproxy/auth/conductor.go | 102 +++++++++++++ .../auth/conductor_executor_replace_test.go | 29 ++++ sdk/cliproxy/embedded_runtime.go | 139 +++++++++++------ sdk/cliproxy/embedded_runtime_test.go | 140 ++++++++++++++++++ sdk/cliproxy/service.go | 65 ++++---- 5 files changed, 389 insertions(+), 86 deletions(-) diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index fedf4951862..9ccdf22a3bd 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -10,6 +10,7 @@ import ( "math/rand/v2" "net/http" "path/filepath" + "reflect" "sort" "strconv" "strings" @@ -502,6 +503,21 @@ func (m *Manager) SetRoundTripperProvider(p RoundTripperProvider) { m.mu.Unlock() } +// SetRoundTripperProviderIfAbsent installs p only when the host has not +// already configured a RoundTripper provider. It reports whether p was stored. +func (m *Manager) SetRoundTripperProviderIfAbsent(p RoundTripperProvider) bool { + if m == nil || p == nil { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if m.rtProvider != nil { + return false + } + m.rtProvider = p + return true +} + // SetConfig updates the runtime config snapshot used by request-time helpers. // Callers should provide the latest config on reload so per-credential alias mapping stays in sync. func (m *Manager) SetConfig(cfg *internalconfig.Config) { @@ -2144,6 +2160,57 @@ func (m *Manager) RegisterExecutor(executor ProviderExecutor) { } } +// RegisterExecutorIfAbsent registers executor only when its provider key has +// no current executor. It reports whether the executor was installed. +func (m *Manager) RegisterExecutorIfAbsent(executor ProviderExecutor) bool { + if m == nil || executor == nil { + return false + } + provider := strings.ToLower(strings.TrimSpace(executor.Identifier())) + if provider == "" { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + if current := m.executors[provider]; current != nil { + return false + } + m.executors[provider] = executor + return true +} + +// UnregisterExecutorIfMatch removes provider only when expected is still the +// registered executor. This lets embedding hosts avoid deleting replacements +// installed concurrently by another owner. +func (m *Manager) UnregisterExecutorIfMatch(provider string, expected ProviderExecutor) bool { + if m == nil || expected == nil { + return false + } + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" { + return false + } + m.mu.Lock() + defer m.mu.Unlock() + current := m.executors[provider] + if !sameProviderExecutor(current, expected) { + return false + } + delete(m.executors, provider) + return true +} + +func sameProviderExecutor(first, second ProviderExecutor) bool { + if first == nil || second == nil { + return first == nil && second == nil + } + firstType := reflect.TypeOf(first) + if firstType != reflect.TypeOf(second) || !firstType.Comparable() { + return false + } + return first == second +} + // UnregisterExecutor removes the executor associated with the provider key. func (m *Manager) UnregisterExecutor(provider string) { provider = strings.ToLower(strings.TrimSpace(provider)) @@ -2235,6 +2302,41 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { return auth.Clone(), nil } +// SyncRuntimeAuth adds or replaces auth in runtime memory without persistence +// or host hooks. It is intended for embedded runtimes whose host owns durable +// storage and mutation notifications. +func (m *Manager) SyncRuntimeAuth(auth *Auth) *Auth { + if m == nil || auth == nil || strings.TrimSpace(auth.ID) == "" { + return nil + } + auth = auth.Clone() + m.mu.Lock() + if existing := m.auths[auth.ID]; existing != nil { + auth.CreatedAt = existing.CreatedAt + auth.Success = existing.Success + auth.Failed = existing.Failed + auth.recentRequests = existing.recentRequests + if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { + auth.LastRefreshedAt = existing.LastRefreshedAt + auth.NextRefreshAfter = existing.NextRefreshAfter + if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { + auth.ModelStates = existing.ModelStates + } + } + } + auth.EnsureIndex() + stored := auth.Clone() + m.auths[auth.ID] = stored + m.mu.Unlock() + + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + if m.scheduler != nil { + m.scheduler.upsertAuth(stored) + } + m.queueRefreshReschedule(auth.ID) + return stored.Clone() +} + // Remove deletes an auth from runtime state without persisting. // Disk and token-store deletion must be handled by the caller. func (m *Manager) Remove(ctx context.Context, id string) { diff --git a/sdk/cliproxy/auth/conductor_executor_replace_test.go b/sdk/cliproxy/auth/conductor_executor_replace_test.go index 99ecf466a6e..138f7f4eab3 100644 --- a/sdk/cliproxy/auth/conductor_executor_replace_test.go +++ b/sdk/cliproxy/auth/conductor_executor_replace_test.go @@ -102,3 +102,32 @@ func TestManagerExecutorReturnsRegisteredExecutor(t *testing.T) { t.Fatal("expected unknown provider lookup to fail") } } + +func TestManagerRegisterAndUnregisterExecutorOwnershipIsAtomic(t *testing.T) { + t.Parallel() + + manager := NewManager(nil, nil, nil) + owned := &replaceAwareExecutor{id: "codex"} + host := &replaceAwareExecutor{id: "codex"} + + if !manager.RegisterExecutorIfAbsent(owned) { + t.Fatal("expected first executor registration to succeed") + } + if manager.RegisterExecutorIfAbsent(host) { + t.Fatal("expected second executor registration to preserve current owner") + } + manager.RegisterExecutor(host) + if manager.UnregisterExecutorIfMatch("codex", owned) { + t.Fatal("stale owner removed host replacement") + } + resolved, okResolved := manager.Executor("codex") + if !okResolved || resolved != host { + t.Fatalf("executor after stale unregister = %T, want host executor", resolved) + } + if !manager.UnregisterExecutorIfMatch("codex", host) { + t.Fatal("current owner could not unregister its executor") + } + if _, okResolved := manager.Executor("codex"); okResolved { + t.Fatal("executor remains registered after matching unregister") + } +} diff --git a/sdk/cliproxy/embedded_runtime.go b/sdk/cliproxy/embedded_runtime.go index ad823f6e553..0faaa7e07e5 100644 --- a/sdk/cliproxy/embedded_runtime.go +++ b/sdk/cliproxy/embedded_runtime.go @@ -3,7 +3,6 @@ package cliproxy import ( "context" "fmt" - "reflect" "strings" "sync" @@ -18,6 +17,7 @@ type EmbeddedRuntime struct { manager *coreauth.Manager service *Service + opMu sync.Mutex mu sync.Mutex started bool closed bool @@ -28,6 +28,9 @@ type EmbeddedRuntime struct { } // NewEmbeddedRuntime creates a headless runtime backed by a host-owned auth manager. +// The supplied configuration becomes the manager's runtime configuration and OAuth +// alias source. A default RoundTripper provider is installed only when the host has +// not already configured one. func NewEmbeddedRuntime(cfg *config.Config, manager *coreauth.Manager) (*EmbeddedRuntime, error) { if cfg == nil { return nil, fmt.Errorf("cliproxy: embedded runtime configuration is required") @@ -36,7 +39,7 @@ func NewEmbeddedRuntime(cfg *config.Config, manager *coreauth.Manager) (*Embedde return nil, fmt.Errorf("cliproxy: embedded runtime auth manager is required") } - manager.SetRoundTripperProvider(newDefaultRoundTripperProvider()) + manager.SetRoundTripperProviderIfAbsent(newDefaultRoundTripperProvider()) manager.SetConfig(cfg) manager.SetOAuthModelAlias(cfg.OAuthModelAlias) @@ -60,28 +63,42 @@ func (r *EmbeddedRuntime) Start(ctx context.Context) error { return errContext } + r.opMu.Lock() + defer r.opMu.Unlock() + if errContext := ctx.Err(); errContext != nil { + return errContext + } r.mu.Lock() - defer r.mu.Unlock() if r.closed { + r.mu.Unlock() return fmt.Errorf("cliproxy: embedded runtime is closed") } if r.started { + r.mu.Unlock() return nil } + r.mu.Unlock() + runtimeCtx := coreauth.WithSkipPersist(ctx) auths := r.manager.List() - r.service.registerAvailableExecutors(ctx, executorRegistrationOptions{ - auths: auths, - forceReplaceAuths: true, - }) for _, auth := range auths { + installed, owned := r.service.registerExecutorForAuth(auth, false) + if owned { + r.installedExecutors[embeddedRuntimeExecutorProvider(auth)] = installed + } if errContext := ctx.Err(); errContext != nil { + r.rollbackStartLocked() + return errContext + } + r.registerAuthModelsLocked(runtimeCtx, auth) + if errContext := ctx.Err(); errContext != nil { + r.rollbackStartLocked() return errContext } - r.trackExecutorForAuthLocked(auth) - r.registerAuthModelsLocked(ctx, auth) } + r.mu.Lock() r.started = true + r.mu.Unlock() return nil } @@ -100,19 +117,28 @@ func (r *EmbeddedRuntime) SyncAuth(ctx context.Context, auth *coreauth.Auth) err return errContext } - r.mu.Lock() - defer r.mu.Unlock() - if errState := r.requireStartedLocked(); errState != nil { + r.opMu.Lock() + defer r.opMu.Unlock() + if errContext := ctx.Err(); errContext != nil { + return errContext + } + if errState := r.requireStarted(); errState != nil { return errState } + installed, owned := r.service.registerExecutorForAuth(auth, false) + if owned { + r.installedExecutors[embeddedRuntimeExecutorProvider(auth)] = installed + } runtimeCtx := coreauth.WithSkipPersist(ctx) - prepared := r.service.prepareCoreAuthForModelRegistration(runtimeCtx, auth) + prepared := r.manager.SyncRuntimeAuth(auth) if prepared == nil { return fmt.Errorf("cliproxy: failed to synchronize auth %q", auth.ID) } - r.trackExecutorForAuthLocked(prepared) r.registerAuthModelsLocked(runtimeCtx, prepared) + if errContext := ctx.Err(); errContext != nil { + return errContext + } return nil } @@ -132,9 +158,12 @@ func (r *EmbeddedRuntime) RemoveAuth(ctx context.Context, authID string) error { return errContext } - r.mu.Lock() - defer r.mu.Unlock() - if errState := r.requireStartedLocked(); errState != nil { + r.opMu.Lock() + defer r.opMu.Unlock() + if errContext := ctx.Err(); errContext != nil { + return errContext + } + if errState := r.requireStarted(); errState != nil { return errState } @@ -155,26 +184,33 @@ func (r *EmbeddedRuntime) ReconcileModels(ctx context.Context, activeModels []st return errContext } - r.mu.Lock() - defer r.mu.Unlock() - if errState := r.requireStartedLocked(); errState != nil { + r.opMu.Lock() + defer r.opMu.Unlock() + if errContext := ctx.Err(); errContext != nil { + return errContext + } + if errState := r.requireStarted(); errState != nil { return errState } r.activeModels = make(map[string]struct{}, len(activeModels)) for _, modelID := range activeModels { - modelID = strings.TrimSpace(modelID) + modelID = strings.ToLower(strings.TrimSpace(modelID)) if modelID != "" { r.activeModels[modelID] = struct{}{} } } r.modelsReconciled = true + runtimeCtx := coreauth.WithSkipPersist(ctx) for _, auth := range r.manager.List() { if errContext := ctx.Err(); errContext != nil { return errContext } - r.registerAuthModelsLocked(ctx, auth) + r.registerAuthModelsLocked(runtimeCtx, auth) + if errContext := ctx.Err(); errContext != nil { + return errContext + } } return nil } @@ -188,28 +224,30 @@ func (r *EmbeddedRuntime) Close(ctx context.Context) error { ctx = context.Background() } + r.opMu.Lock() + defer r.opMu.Unlock() r.mu.Lock() - defer r.mu.Unlock() if r.closed { + r.mu.Unlock() return nil } + r.mu.Unlock() for authID := range r.registeredAuthIDs { registry.GetGlobalRegistry().UnregisterClient(authID) r.manager.RefreshSchedulerEntry(authID) } for provider, installed := range r.installedExecutors { - current, okExecutor := r.manager.Executor(provider) - if !okExecutor || !sameProviderExecutor(current, installed) { - continue - } - if closer, okCloser := current.(coreauth.ExecutionSessionCloser); okCloser { + if closer, okCloser := installed.(coreauth.ExecutionSessionCloser); okCloser { closer.CloseExecutionSession(coreauth.CloseAllExecutionSessionsID) } - r.manager.UnregisterExecutor(provider) + r.manager.UnregisterExecutorIfMatch(provider, installed) } r.registeredAuthIDs = make(map[string]struct{}) r.installedExecutors = make(map[string]coreauth.ProviderExecutor) + r.mu.Lock() r.closed = true + r.started = false + r.mu.Unlock() return nil } @@ -233,7 +271,9 @@ func (r *EmbeddedRuntime) WatcherStarted() bool { return r.service != nil && r.service.watcher != nil } -func (r *EmbeddedRuntime) requireStartedLocked() error { +func (r *EmbeddedRuntime) requireStarted() error { + r.mu.Lock() + defer r.mu.Unlock() if r.closed { return fmt.Errorf("cliproxy: embedded runtime is closed") } @@ -247,6 +287,16 @@ func (r *EmbeddedRuntime) registerAuthModelsLocked(ctx context.Context, auth *co if auth == nil || strings.TrimSpace(auth.ID) == "" { return } + provider := embeddedRuntimeExecutorProvider(auth) + if provider == "" { + return + } + if _, okExecutor := r.manager.Executor(provider); !okExecutor { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + r.manager.RefreshSchedulerEntry(auth.ID) + delete(r.registeredAuthIDs, auth.ID) + return + } var activeModels map[string]struct{} if r.modelsReconciled { activeModels = r.activeModels @@ -255,15 +305,19 @@ func (r *EmbeddedRuntime) registerAuthModelsLocked(ctx context.Context, auth *co r.registeredAuthIDs[auth.ID] = struct{}{} } -func (r *EmbeddedRuntime) trackExecutorForAuthLocked(auth *coreauth.Auth) { - provider := embeddedRuntimeExecutorProvider(auth) - if provider == "" { - return +func (r *EmbeddedRuntime) rollbackStartLocked() { + for authID := range r.registeredAuthIDs { + registry.GetGlobalRegistry().UnregisterClient(authID) + r.manager.RefreshSchedulerEntry(authID) } - executor, okExecutor := r.manager.Executor(provider) - if okExecutor && executor != nil { - r.installedExecutors[provider] = executor + for provider, installed := range r.installedExecutors { + if closer, okCloser := installed.(coreauth.ExecutionSessionCloser); okCloser { + closer.CloseExecutionSession(coreauth.CloseAllExecutionSessionsID) + } + r.manager.UnregisterExecutorIfMatch(provider, installed) } + r.registeredAuthIDs = make(map[string]struct{}) + r.installedExecutors = make(map[string]coreauth.ProviderExecutor) } func embeddedRuntimeExecutorProvider(auth *coreauth.Auth) string { @@ -278,14 +332,3 @@ func embeddedRuntimeExecutorProvider(auth *coreauth.Auth) string { } return strings.ToLower(strings.TrimSpace(auth.Provider)) } - -func sameProviderExecutor(first, second coreauth.ProviderExecutor) bool { - if first == nil || second == nil { - return first == nil && second == nil - } - firstType := reflect.TypeOf(first) - if firstType != reflect.TypeOf(second) || !firstType.Comparable() { - return false - } - return first == second -} diff --git a/sdk/cliproxy/embedded_runtime_test.go b/sdk/cliproxy/embedded_runtime_test.go index 1c29827fa05..5f7bb36797a 100644 --- a/sdk/cliproxy/embedded_runtime_test.go +++ b/sdk/cliproxy/embedded_runtime_test.go @@ -2,15 +2,46 @@ package cliproxy_test import ( "context" + "net/http" + "strings" "testing" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" ) +type embeddedRuntimeTestExecutor struct { + provider string +} + +func (e *embeddedRuntimeTestExecutor) Identifier() string { return e.provider } + +func (e *embeddedRuntimeTestExecutor) Execute(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *embeddedRuntimeTestExecutor) ExecuteStream(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + chunks := make(chan cliproxyexecutor.StreamChunk) + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *embeddedRuntimeTestExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (e *embeddedRuntimeTestExecutor) CountTokens(context.Context, *coreauth.Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *embeddedRuntimeTestExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + func TestEmbeddedRuntimeRegistersNativeExecutorsWithoutServer(t *testing.T) { manager := coreauth.NewManager(nil, nil, nil) if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ID: "codex-1", Provider: "codex"}); errRegister != nil { @@ -147,3 +178,112 @@ func TestEmbeddedRuntimeCloseUnregistersModelsWithoutRemovingAuth(t *testing.T) t.Fatal("closed runtime still exposes its claude executor") } } + +func TestEmbeddedRuntimePreservesHostExecutor(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + hostExecutor := &embeddedRuntimeTestExecutor{provider: "codex"} + manager.RegisterExecutor(hostExecutor) + if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ID: "codex-host-1", Provider: "codex"}); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + + resolved, okResolved := manager.Executor("codex") + if !okResolved || resolved != hostExecutor { + t.Fatalf("executor after start = %T, want host executor", resolved) + } + if errClose := runtime.Close(t.Context()); errClose != nil { + t.Fatalf("close embedded runtime: %v", errClose) + } + resolved, okResolved = manager.Executor("codex") + if !okResolved || resolved != hostExecutor { + t.Fatalf("executor after close = %T, want host executor", resolved) + } +} + +func TestEmbeddedRuntimeDoesNotProjectModelsWithoutExecutor(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ID: "aistudio-headless-1", Provider: "aistudio"} + if _, errRegister := manager.Register(t.Context(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + t.Cleanup(func() { _ = runtime.Close(context.Background()) }) + + if _, okExecutor := manager.Executor("aistudio"); okExecutor { + t.Fatal("headless runtime unexpectedly installed an AI Studio executor") + } + for _, model := range registry.GetAIStudioModels() { + if model != nil && cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, model.ID) { + t.Fatalf("auth without executor supports model %q", model.ID) + } + } +} + +func TestEmbeddedRuntimeReconcileModelsIsCaseInsensitive(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + auth := &coreauth.Auth{ID: "codex-case-1", Provider: "codex"} + if _, errRegister := manager.Register(t.Context(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + t.Cleanup(func() { _ = runtime.Close(context.Background()) }) + + models := registry.GetCodexProModels() + if len(models) == 0 { + t.Fatal("codex model fixture is empty") + } + keptModel := models[0].ID + if errReconcile := runtime.ReconcileModels(t.Context(), []string{strings.ToUpper(keptModel)}); errReconcile != nil { + t.Fatalf("reconcile models: %v", errReconcile) + } + if !cliproxy.GlobalModelRegistry().ClientSupportsModel(auth.ID, keptModel) { + t.Fatalf("case-insensitive reconciliation removed model %q", keptModel) + } +} + +func TestEmbeddedRuntimeCanceledStartRollsBackOwnedExecutors(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ID: "codex-cancel-1", Provider: "codex"}); errRegister != nil { + t.Fatalf("register codex auth: %v", errRegister) + } + if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ID: "claude-cancel-1", Provider: "claude"}); errRegister != nil { + t.Fatalf("register claude auth: %v", errRegister) + } + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + if errStart := runtime.Start(ctx); errStart == nil { + t.Fatal("canceled start returned nil error") + } + if _, okExecutor := manager.Executor("codex"); okExecutor { + t.Fatal("canceled start leaked codex executor") + } + if _, okExecutor := manager.Executor("claude"); okExecutor { + t.Fatal("canceled start leaked claude executor") + } +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index c82190d8d46..ad7d573cc60 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -1023,22 +1023,12 @@ func (s *Service) registerExecutorsForAuths(auths []*coreauth.Auth, forceReplace } } -func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { +func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) (installed coreauth.ProviderExecutor, owned bool) { if s == nil || s.coreManager == nil || a == nil { return } if strings.EqualFold(strings.TrimSpace(a.Provider), "codex") { - if !forceReplace { - existingExecutor, hasExecutor := s.coreManager.Executor("codex") - if hasExecutor { - _, isCodexAutoExecutor := existingExecutor.(*executor.CodexAutoExecutor) - if isCodexAutoExecutor { - return - } - } - } - s.coreManager.RegisterExecutor(executor.NewCodexAutoExecutor(s.cfg)) - return + return s.installExecutor(executor.NewCodexAutoExecutor(s.cfg), forceReplace) } // Skip disabled auth entries when (re)binding executors. // Disabled auths can linger during config reloads (e.g., removed OpenAI-compat entries) @@ -1053,36 +1043,28 @@ func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { if compatProviderKey == "" { compatProviderKey = "openai-compatibility" } - if !forceReplace { - if existingExecutor, hasExecutor := s.coreManager.Executor(compatProviderKey); hasExecutor { - if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { - return - } - } - } - s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(compatProviderKey, s.cfg)) - return + return s.installExecutor(executor.NewOpenAICompatExecutor(compatProviderKey, s.cfg), forceReplace) } switch strings.ToLower(a.Provider) { case constant.Gemini: - s.coreManager.RegisterExecutor(executor.NewGeminiExecutor(s.cfg)) + return s.installExecutor(executor.NewGeminiExecutor(s.cfg), forceReplace) case constant.GeminiInteractions: - s.coreManager.RegisterExecutor(executor.NewGeminiInteractionsExecutor(s.cfg)) + return s.installExecutor(executor.NewGeminiInteractionsExecutor(s.cfg), forceReplace) case "vertex": - s.coreManager.RegisterExecutor(executor.NewGeminiVertexExecutor(s.cfg)) + return s.installExecutor(executor.NewGeminiVertexExecutor(s.cfg), forceReplace) case "aistudio": if s.wsGateway != nil { - s.coreManager.RegisterExecutor(executor.NewAIStudioExecutor(s.cfg, a.ID, s.wsGateway)) + return s.installExecutor(executor.NewAIStudioExecutor(s.cfg, a.ID, s.wsGateway), forceReplace) } return case "antigravity": - s.coreManager.RegisterExecutor(executor.NewAntigravityExecutor(s.cfg)) + return s.installExecutor(executor.NewAntigravityExecutor(s.cfg), forceReplace) case "claude": - s.coreManager.RegisterExecutor(executor.NewClaudeExecutor(s.cfg)) + return s.installExecutor(executor.NewClaudeExecutor(s.cfg), forceReplace) case "kimi": - s.coreManager.RegisterExecutor(executor.NewKimiExecutor(s.cfg)) + return s.installExecutor(executor.NewKimiExecutor(s.cfg), forceReplace) case "xai": - s.coreManager.RegisterExecutor(executor.NewXAIAutoExecutor(s.cfg)) + return s.installExecutor(executor.NewXAIAutoExecutor(s.cfg), forceReplace) default: providerKey := strings.ToLower(strings.TrimSpace(a.Provider)) if providerKey == "" { @@ -1094,15 +1076,22 @@ func (s *Service) registerExecutorForAuth(a *coreauth.Auth, forceReplace bool) { s.unregisterOpenAICompatExecutor(providerKey) return } - if !forceReplace { - if existingExecutor, hasExecutor := s.coreManager.Executor(providerKey); hasExecutor { - if _, isOpenAICompatExecutor := existingExecutor.(*executor.OpenAICompatExecutor); isOpenAICompatExecutor { - return - } - } - } - s.coreManager.RegisterExecutor(executor.NewOpenAICompatExecutor(providerKey, s.cfg)) + return s.installExecutor(executor.NewOpenAICompatExecutor(providerKey, s.cfg), forceReplace) + } +} + +func (s *Service) installExecutor(candidate coreauth.ProviderExecutor, forceReplace bool) (coreauth.ProviderExecutor, bool) { + if candidate == nil { + return nil, false + } + if forceReplace { + s.coreManager.RegisterExecutor(candidate) + return candidate, true + } + if !s.coreManager.RegisterExecutorIfAbsent(candidate) { + return nil, false } + return candidate, true } func (s *Service) registerResolvedModelsForAuth(a *coreauth.Auth, providerKey string, models []*ModelInfo) { @@ -1920,7 +1909,7 @@ func (s *Service) reconcileModelsForAuth(ctx context.Context, a *coreauth.Auth, if model == nil { continue } - if _, active := activeModels[strings.TrimSpace(model.ID)]; active { + if _, active := activeModels[strings.ToLower(strings.TrimSpace(model.ID))]; active { filtered = append(filtered, model) } } From 056552f09d235606d0264099d5696127750e2340 Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 21:06:52 +0800 Subject: [PATCH 03/11] feat(sdk): own embedded auto-refresh lifecycle Add scoped auto-refresh leases so EmbeddedRuntime can start and drain refresh workers without stopping host-owned loops or selectors. --- sdk/cliproxy/auth/auto_refresh_lease_test.go | 171 +++++++++++++++++++ sdk/cliproxy/auth/auto_refresh_loop.go | 8 +- sdk/cliproxy/auth/conductor.go | 114 ++++++++++++- sdk/cliproxy/embedded_runtime.go | 23 ++- sdk/cliproxy/embedded_runtime_test.go | 91 ++++++++++ 5 files changed, 404 insertions(+), 3 deletions(-) create mode 100644 sdk/cliproxy/auth/auto_refresh_lease_test.go diff --git a/sdk/cliproxy/auth/auto_refresh_lease_test.go b/sdk/cliproxy/auth/auto_refresh_lease_test.go new file mode 100644 index 00000000000..66042eeb6c8 --- /dev/null +++ b/sdk/cliproxy/auth/auto_refresh_lease_test.go @@ -0,0 +1,171 @@ +package auth + +import ( + "context" + "net/http" + "sync/atomic" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type leaseTestSelector struct { + stopped atomic.Bool +} + +type leaseRefreshEvaluator struct{} + +func (leaseRefreshEvaluator) ShouldRefresh(time.Time, *Auth) bool { return true } + +func (s *leaseTestSelector) Pick(context.Context, string, string, cliproxyexecutor.Options, []*Auth) (*Auth, error) { + return nil, nil +} + +func (s *leaseTestSelector) Stop() { + s.stopped.Store(true) +} + +type blockingRefreshExecutor struct { + started chan struct{} + release chan struct{} +} + +func (e *blockingRefreshExecutor) Identifier() string { return "codex" } + +func (e *blockingRefreshExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *blockingRefreshExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + return nil, nil +} + +func (e *blockingRefreshExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + select { + case <-e.started: + default: + close(e.started) + } + <-e.release + return auth, nil +} + +func (e *blockingRefreshExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, nil +} + +func (e *blockingRefreshExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func TestAutoRefreshLeaseWaitsWithoutStoppingSelector(t *testing.T) { + selector := &leaseTestSelector{} + manager := NewManager(nil, selector, nil) + executor := &blockingRefreshExecutor{started: make(chan struct{}), release: make(chan struct{})} + manager.RegisterExecutor(executor) + _, errRegister := manager.Register(t.Context(), &Auth{ + ID: "codex-refresh-lease", + Provider: "codex", + Attributes: map[string]string{ + AttributeAuthKind: AuthKindOAuth, + }, + Metadata: map[string]any{ + "access_token": "expired-access", + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(-time.Minute).Format(time.RFC3339), + }, + Runtime: leaseRefreshEvaluator{}, + }) + if errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + lease := manager.EnsureAutoRefresh(t.Context(), time.Millisecond) + if lease == nil || !lease.owned { + t.Fatal("expected an owned auto-refresh lease") + } + select { + case <-executor.started: + case <-time.After(2 * time.Second): + t.Fatal("refresh worker did not start") + } + + closed := make(chan error, 1) + go func() { closed <- lease.Close(context.Background()) }() + select { + case errClose := <-closed: + t.Fatalf("lease closed before refresh worker exited: %v", errClose) + case <-time.After(50 * time.Millisecond): + } + close(executor.release) + select { + case errClose := <-closed: + if errClose != nil { + t.Fatalf("close lease: %v", errClose) + } + case <-time.After(2 * time.Second): + t.Fatal("lease did not wait for refresh worker shutdown") + } + if selector.stopped.Load() { + t.Fatal("closing refresh lease stopped host selector") + } + manager.mu.RLock() + defer manager.mu.RUnlock() + if manager.refreshLoop != nil || manager.refreshCancel != nil || manager.refreshDone != nil { + t.Fatal("owned refresh loop remains installed after lease close") + } +} + +func TestAutoRefreshLeasePreservesExistingHostLoop(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.StartAutoRefresh(t.Context(), time.Hour) + manager.mu.RLock() + hostLoop := manager.refreshLoop + manager.mu.RUnlock() + if hostLoop == nil { + t.Fatal("host refresh loop was not installed") + } + + lease := manager.EnsureAutoRefresh(t.Context(), time.Minute) + if lease == nil || lease.owned { + t.Fatal("expected a non-owning lease for the existing host loop") + } + if errClose := lease.Close(t.Context()); errClose != nil { + t.Fatalf("close non-owning lease: %v", errClose) + } + manager.mu.RLock() + current := manager.refreshLoop + manager.mu.RUnlock() + if current != hostLoop { + t.Fatal("closing non-owning lease replaced or removed host loop") + } + manager.StopAutoRefresh() +} + +func TestAutoRefreshLeaseReplacesExitedHostLoop(t *testing.T) { + manager := NewManager(nil, nil, nil) + hostCtx, cancelHost := context.WithCancel(t.Context()) + manager.StartAutoRefresh(hostCtx, time.Hour) + manager.mu.RLock() + hostLoop := manager.refreshLoop + hostDone := manager.refreshDone + manager.mu.RUnlock() + cancelHost() + select { + case <-hostDone: + case <-time.After(2 * time.Second): + t.Fatal("host refresh loop did not exit") + } + + lease := manager.EnsureAutoRefresh(t.Context(), time.Hour) + if lease == nil || !lease.owned { + t.Fatal("expected an owned lease after host loop exit") + } + if lease.loop == hostLoop { + t.Fatal("exited host loop was reused") + } + if errClose := lease.Close(t.Context()); errClose != nil { + t.Fatalf("close replacement lease: %v", errClose) + } +} diff --git a/sdk/cliproxy/auth/auto_refresh_loop.go b/sdk/cliproxy/auth/auto_refresh_loop.go index b4217b31636..e70ebaf3581 100644 --- a/sdk/cliproxy/auth/auto_refresh_loop.go +++ b/sdk/cliproxy/auth/auto_refresh_loop.go @@ -68,11 +68,17 @@ func (l *authAutoRefreshLoop) run(ctx context.Context) { if workers <= 0 { workers = refreshMaxConcurrency } + var workerWG sync.WaitGroup + workerWG.Add(workers) for i := 0; i < workers; i++ { - go l.worker(ctx) + go func() { + defer workerWG.Done() + l.worker(ctx) + }() } l.loop(ctx) + workerWG.Wait() } func (l *authAutoRefreshLoop) worker(ctx context.Context) { diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 9ccdf22a3bd..cf85a2ae7a4 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -259,6 +259,7 @@ type Manager struct { // Auto refresh state refreshCancel context.CancelFunc refreshLoop *authAutoRefreshLoop + refreshDone <-chan struct{} requestPrepareLocks sync.Map // refreshLocks serializes credential refresh per auth ID so concurrent @@ -5695,6 +5696,7 @@ func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duratio cancelPrev := m.refreshCancel m.refreshCancel = nil m.refreshLoop = nil + m.refreshDone = nil m.mu.Unlock() if cancelPrev != nil { cancelPrev() @@ -5706,14 +5708,123 @@ func (m *Manager) StartAutoRefresh(parent context.Context, interval time.Duratio workers = cfg.AuthAutoRefreshWorkers } loop := newAuthAutoRefreshLoop(m, interval, workers) + done := make(chan struct{}) m.mu.Lock() m.refreshCancel = cancelCtx m.refreshLoop = loop + m.refreshDone = done m.mu.Unlock() loop.rebuild(time.Now()) - go loop.run(ctx) + go m.runAutoRefreshLoop(ctx, loop, done) +} + +func (m *Manager) runAutoRefreshLoop(ctx context.Context, loop *authAutoRefreshLoop, done chan struct{}) { + loop.run(ctx) + close(done) + m.mu.Lock() + if m.refreshLoop == loop { + m.refreshCancel = nil + m.refreshLoop = nil + m.refreshDone = nil + } + m.mu.Unlock() +} + +// AutoRefreshLease represents an auto-refresh loop acquired for an embedding +// runtime. A lease only stops the exact loop it created; an existing host-owned +// loop is reused and remains untouched when the lease closes. +type AutoRefreshLease struct { + manager *Manager + loop *authAutoRefreshLoop + cancel context.CancelFunc + done <-chan struct{} + owned bool + once sync.Once +} + +// EnsureAutoRefresh reuses an existing manager loop or starts one owned by the +// returned lease. Closing a non-owning lease is a no-op. +func (m *Manager) EnsureAutoRefresh(parent context.Context, interval time.Duration) *AutoRefreshLease { + if m == nil { + return &AutoRefreshLease{} + } + if parent == nil { + parent = context.Background() + } + if interval <= 0 { + interval = refreshCheckInterval + } + + m.mu.Lock() + if m.refreshLoop != nil { + loopDone := m.refreshDone + loopRunning := true + if loopDone != nil { + select { + case <-loopDone: + loopRunning = false + default: + } + } + if loopRunning { + lease := &AutoRefreshLease{manager: m, loop: m.refreshLoop, done: loopDone} + m.mu.Unlock() + return lease + } + m.refreshCancel = nil + m.refreshLoop = nil + m.refreshDone = nil + } + ctx, cancelCtx := context.WithCancel(parent) + workers := refreshMaxConcurrency + if cfg, ok := m.runtimeConfig.Load().(*internalconfig.Config); ok && cfg != nil && cfg.AuthAutoRefreshWorkers > 0 { + workers = cfg.AuthAutoRefreshWorkers + } + loop := newAuthAutoRefreshLoop(m, interval, workers) + done := make(chan struct{}) + m.refreshCancel = cancelCtx + m.refreshLoop = loop + m.refreshDone = done + m.mu.Unlock() + + loop.rebuild(time.Now()) + go m.runAutoRefreshLoop(ctx, loop, done) + return &AutoRefreshLease{manager: m, loop: loop, cancel: cancelCtx, done: done, owned: true} +} + +// Close cancels an owned loop and waits for its scheduler and workers to exit. +func (l *AutoRefreshLease) Close(ctx context.Context) error { + if l == nil || !l.owned { + return nil + } + if ctx == nil { + ctx = context.Background() + } + l.once.Do(func() { + if l.manager != nil { + l.manager.mu.Lock() + if l.manager.refreshLoop == l.loop { + l.manager.refreshCancel = nil + l.manager.refreshLoop = nil + l.manager.refreshDone = nil + } + l.manager.mu.Unlock() + } + if l.cancel != nil { + l.cancel() + } + }) + if l.done == nil { + return nil + } + select { + case <-l.done: + return nil + case <-ctx.Done(): + return ctx.Err() + } } // StopAutoRefresh cancels the background refresh loop, if running. @@ -5723,6 +5834,7 @@ func (m *Manager) StopAutoRefresh() { cancel := m.refreshCancel m.refreshCancel = nil m.refreshLoop = nil + m.refreshDone = nil m.mu.Unlock() if cancel != nil { cancel() diff --git a/sdk/cliproxy/embedded_runtime.go b/sdk/cliproxy/embedded_runtime.go index 0faaa7e07e5..04c34b28c7b 100644 --- a/sdk/cliproxy/embedded_runtime.go +++ b/sdk/cliproxy/embedded_runtime.go @@ -5,6 +5,7 @@ import ( "fmt" "strings" "sync" + "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -25,8 +26,11 @@ type EmbeddedRuntime struct { activeModels map[string]struct{} registeredAuthIDs map[string]struct{} installedExecutors map[string]coreauth.ProviderExecutor + refreshLease *coreauth.AutoRefreshLease } +const embeddedRuntimeRefreshInterval = 15 * time.Minute + // NewEmbeddedRuntime creates a headless runtime backed by a host-owned auth manager. // The supplied configuration becomes the manager's runtime configuration and OAuth // alias source. A default RoundTripper provider is installed only when the host has @@ -96,6 +100,11 @@ func (r *EmbeddedRuntime) Start(ctx context.Context) error { return errContext } } + r.refreshLease = r.manager.EnsureAutoRefresh(context.Background(), embeddedRuntimeRefreshInterval) + if errContext := ctx.Err(); errContext != nil { + r.rollbackStartLocked() + return errContext + } r.mu.Lock() r.started = true r.mu.Unlock() @@ -232,6 +241,14 @@ func (r *EmbeddedRuntime) Close(ctx context.Context) error { return nil } r.mu.Unlock() + var refreshErr error + if r.refreshLease != nil { + refreshErr = r.refreshLease.Close(ctx) + if refreshErr != nil { + return refreshErr + } + r.refreshLease = nil + } for authID := range r.registeredAuthIDs { registry.GetGlobalRegistry().UnregisterClient(authID) r.manager.RefreshSchedulerEntry(authID) @@ -248,7 +265,7 @@ func (r *EmbeddedRuntime) Close(ctx context.Context) error { r.closed = true r.started = false r.mu.Unlock() - return nil + return refreshErr } // ServerStarted reports whether an HTTP server was constructed by this runtime. @@ -306,6 +323,10 @@ func (r *EmbeddedRuntime) registerAuthModelsLocked(ctx context.Context, auth *co } func (r *EmbeddedRuntime) rollbackStartLocked() { + if r.refreshLease != nil { + _ = r.refreshLease.Close(context.Background()) + r.refreshLease = nil + } for authID := range r.registeredAuthIDs { registry.GetGlobalRegistry().UnregisterClient(authID) r.manager.RefreshSchedulerEntry(authID) diff --git a/sdk/cliproxy/embedded_runtime_test.go b/sdk/cliproxy/embedded_runtime_test.go index 5f7bb36797a..989d44cd4cf 100644 --- a/sdk/cliproxy/embedded_runtime_test.go +++ b/sdk/cliproxy/embedded_runtime_test.go @@ -2,9 +2,12 @@ package cliproxy_test import ( "context" + "errors" "net/http" "strings" + "sync" "testing" + "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" @@ -16,6 +19,10 @@ import ( type embeddedRuntimeTestExecutor struct { provider string + + refreshOnce sync.Once + refreshed chan struct{} + release chan struct{} } func (e *embeddedRuntimeTestExecutor) Identifier() string { return e.provider } @@ -31,6 +38,12 @@ func (e *embeddedRuntimeTestExecutor) ExecuteStream(context.Context, *coreauth.A } func (e *embeddedRuntimeTestExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + if e.refreshed != nil { + e.refreshOnce.Do(func() { close(e.refreshed) }) + } + if e.release != nil { + <-e.release + } return auth, nil } @@ -287,3 +300,81 @@ func TestEmbeddedRuntimeCanceledStartRollsBackOwnedExecutors(t *testing.T) { t.Fatal("canceled start leaked claude executor") } } + +func TestEmbeddedRuntimeOwnsManagerAutoRefreshLifecycle(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + refreshed := make(chan struct{}) + manager.RegisterExecutor(&embeddedRuntimeTestExecutor{provider: "codex", refreshed: refreshed}) + if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ + ID: "codex-refresh-1", + Provider: "codex", + Attributes: map[string]string{ + coreauth.AttributeAuthKind: coreauth.AuthKindOAuth, + }, + Metadata: map[string]any{ + "access_token": "expired-access", + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(-time.Minute).Format(time.RFC3339), + }, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + t.Cleanup(func() { _ = runtime.Close(context.Background()) }) + + select { + case <-refreshed: + case <-time.After(2 * time.Second): + t.Fatal("embedded runtime did not start manager auto-refresh") + } +} + +func TestEmbeddedRuntimeCloseCanRetryRefreshDrain(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + refreshed := make(chan struct{}) + release := make(chan struct{}) + manager.RegisterExecutor(&embeddedRuntimeTestExecutor{provider: "codex", refreshed: refreshed, release: release}) + if _, errRegister := manager.Register(t.Context(), &coreauth.Auth{ + ID: "codex-refresh-close", + Provider: "codex", + Attributes: map[string]string{ + coreauth.AttributeAuthKind: coreauth.AuthKindOAuth, + }, + Metadata: map[string]any{ + "access_token": "expired-access", + "refresh_token": "refresh-token", + "expires_at": time.Now().Add(-time.Minute).Format(time.RFC3339), + }, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + runtime, errRuntime := cliproxy.NewEmbeddedRuntime(&config.Config{}, manager) + if errRuntime != nil { + t.Fatalf("new embedded runtime: %v", errRuntime) + } + if errStart := runtime.Start(t.Context()); errStart != nil { + t.Fatalf("start embedded runtime: %v", errStart) + } + select { + case <-refreshed: + case <-time.After(2 * time.Second): + t.Fatal("refresh worker did not start") + } + + closeCtx, cancelClose := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancelClose() + if errClose := runtime.Close(closeCtx); !errors.Is(errClose, context.DeadlineExceeded) { + t.Fatalf("first close error = %v, want context deadline exceeded", errClose) + } + close(release) + if errClose := runtime.Close(context.Background()); errClose != nil { + t.Fatalf("retry close: %v", errClose) + } +} From 9238d59a07a5a4c7ce3a77e7d3f23fc331d6de2f Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 21:52:28 +0800 Subject: [PATCH 04/11] feat(sdk): execute host-resolved provider candidates Add a headless resolved-execution facade that accepts host-authorized credential candidates and delegates transport, translation, retry, cooldown, and streaming behavior to CLIProxyAPI native executors and Manager semantics. Treat provider 400/422 responses as request-scoped so client errors do not rotate credentials or poison account health.\n\nValidated SDK/auth tests, race tests, and go vet. --- sdk/cliproxy/embedded_runtime.go | 8 + sdk/cliproxy/resolved_execution.go | 261 ++++++++++++++++++++++++ sdk/cliproxy/resolved_execution_test.go | 244 ++++++++++++++++++++++ 3 files changed, 513 insertions(+) create mode 100644 sdk/cliproxy/resolved_execution.go create mode 100644 sdk/cliproxy/resolved_execution_test.go diff --git a/sdk/cliproxy/embedded_runtime.go b/sdk/cliproxy/embedded_runtime.go index 04c34b28c7b..d9dd405aed1 100644 --- a/sdk/cliproxy/embedded_runtime.go +++ b/sdk/cliproxy/embedded_runtime.go @@ -27,6 +27,11 @@ type EmbeddedRuntime struct { registeredAuthIDs map[string]struct{} installedExecutors map[string]coreauth.ProviderExecutor refreshLease *coreauth.AutoRefreshLease + + resolvedMu sync.Mutex + resolvedClosing bool + resolvedInflight int + resolvedDrained chan struct{} } const embeddedRuntimeRefreshInterval = 15 * time.Minute @@ -241,6 +246,9 @@ func (r *EmbeddedRuntime) Close(ctx context.Context) error { return nil } r.mu.Unlock() + if err := r.closeResolvedExecutions(ctx); err != nil { + return err + } var refreshErr error if r.refreshLease != nil { refreshErr = r.refreshLease.Close(ctx) diff --git a/sdk/cliproxy/resolved_execution.go b/sdk/cliproxy/resolved_execution.go new file mode 100644 index 00000000000..2ffa324a0db --- /dev/null +++ b/sdk/cliproxy/resolved_execution.go @@ -0,0 +1,261 @@ +package cliproxy + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + + nativeexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type NativeExecutorKind string + +const ( + NativeExecutorOpenAICompatible NativeExecutorKind = "openai-compatible" + NativeExecutorAnthropic NativeExecutorKind = "anthropic" +) + +type ResolvedExecutionCandidate struct { + Auth *coreauth.Auth + RuntimeModel string + APIBase string +} + +type ResolvedExecutionRequest struct { + Executor NativeExecutorKind + Provider string + Candidates []ResolvedExecutionCandidate + Request cliproxyexecutor.Request + Options cliproxyexecutor.Options +} + +func (r *EmbeddedRuntime) ExecuteResolved(ctx context.Context, execution ResolvedExecutionRequest) (cliproxyexecutor.Response, error) { + ctx, finish, err := r.beginResolvedExecution(ctx) + if err != nil { + return cliproxyexecutor.Response{}, err + } + defer finish() + executor, provider, candidates, err := r.prepareResolvedCandidates(execution) + if err != nil { + return cliproxyexecutor.Response{}, err + } + var lastErr error + for _, candidate := range candidates { + req := execution.Request + req.Model = candidate.RuntimeModel + resp, errExec := executor.Execute(ctx, candidate.Auth, req, execution.Options) + if errExec == nil { + return resp, nil + } + if err := ctx.Err(); err != nil { + return cliproxyexecutor.Response{}, err + } + if resolvedRequestInvalid(errExec) { + return cliproxyexecutor.Response{}, errExec + } + lastErr = errExec + } + if lastErr != nil { + return cliproxyexecutor.Response{}, lastErr + } + return cliproxyexecutor.Response{}, &coreauth.Error{Code: "auth_not_found", Message: fmt.Sprintf("no %s auth available", provider)} +} + +func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution ResolvedExecutionRequest) (*cliproxyexecutor.StreamResult, error) { + ctx, finish, err := r.beginResolvedExecution(ctx) + if err != nil { + return nil, err + } + executor, provider, candidates, err := r.prepareResolvedCandidates(execution) + if err != nil { + finish() + return nil, err + } + var lastErr error + for _, candidate := range candidates { + req := execution.Request + req.Model = candidate.RuntimeModel + result, errExec := executor.ExecuteStream(ctx, candidate.Auth, req, execution.Options) + if errExec == nil { + first, ok := <-result.Chunks + if ok && first.Err != nil { + if resolvedRequestInvalid(first.Err) { + finish() + return nil, first.Err + } + lastErr = first.Err + continue + } + forwarded := make(chan cliproxyexecutor.StreamChunk) + go func() { + defer close(forwarded) + defer finish() + if ok { + select { + case forwarded <- first: + case <-ctx.Done(): + return + } + } + for chunk := range result.Chunks { + select { + case forwarded <- chunk: + case <-ctx.Done(): + return + } + } + }() + return &cliproxyexecutor.StreamResult{Headers: result.Headers, Chunks: forwarded}, nil + } + if err := ctx.Err(); err != nil { + finish() + return nil, err + } + if resolvedRequestInvalid(errExec) { + finish() + return nil, errExec + } + lastErr = errExec + } + finish() + if lastErr != nil { + return nil, lastErr + } + return nil, &coreauth.Error{Code: "auth_not_found", Message: fmt.Sprintf("no %s auth available", provider)} +} + +type preparedResolvedCandidate struct { + Auth *coreauth.Auth + RuntimeModel string +} + +func (r *EmbeddedRuntime) prepareResolvedCandidates(execution ResolvedExecutionRequest) (coreauth.ProviderExecutor, string, []preparedResolvedCandidate, error) { + if len(execution.Candidates) == 0 { + return nil, "", nil, &coreauth.Error{Code: "auth_not_found", Message: "no resolved auth candidates"} + } + provider := strings.ToLower(strings.TrimSpace(execution.Provider)) + if provider == "" { + provider = "resolved" + } + var executor coreauth.ProviderExecutor + switch execution.Executor { + case NativeExecutorOpenAICompatible: + executor = nativeexecutor.NewOpenAICompatExecutor(provider, r.service.cfg) + case NativeExecutorAnthropic: + provider = "claude" + executor = nativeexecutor.NewClaudeExecutor(r.service.cfg) + default: + return nil, "", nil, fmt.Errorf("cliproxy: unsupported resolved executor %q", execution.Executor) + } + candidates := make([]preparedResolvedCandidate, 0, len(execution.Candidates)) + for index, candidate := range execution.Candidates { + if candidate.Auth == nil { + return nil, "", nil, fmt.Errorf("cliproxy: resolved auth candidate %d is required", index) + } + model := strings.TrimSpace(candidate.RuntimeModel) + if model == "" { + return nil, "", nil, fmt.Errorf("cliproxy: resolved runtime model is required") + } + apiBase := resolvedExecutorBase(execution.Executor, candidate.APIBase) + if apiBase == "" { + return nil, "", nil, fmt.Errorf("cliproxy: resolved API base is required") + } + apiKey := strings.TrimSpace(candidate.Auth.Attributes["api_key"]) + if apiKey == "" { + return nil, "", nil, fmt.Errorf("cliproxy: resolved API key is required") + } + auth := &coreauth.Auth{ + ID: fmt.Sprintf("resolved-%d", index), + Provider: provider, + Label: candidate.Auth.Label, + Status: coreauth.StatusActive, + ProxyURL: candidate.Auth.ProxyURL, + Attributes: map[string]string{ + "api_key": apiKey, + "base_url": apiBase, + }, + } + candidates = append(candidates, preparedResolvedCandidate{Auth: auth, RuntimeModel: model}) + } + return executor, provider, candidates, nil +} + +func (r *EmbeddedRuntime) beginResolvedExecution(ctx context.Context) (context.Context, func(), error) { + if r == nil || r.service == nil || r.service.cfg == nil { + return nil, nil, fmt.Errorf("cliproxy: embedded runtime is unavailable") + } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return nil, nil, err + } + r.mu.Lock() + started, closed := r.started, r.closed + r.mu.Unlock() + if !started || closed { + return nil, nil, fmt.Errorf("cliproxy: embedded runtime is not started") + } + r.resolvedMu.Lock() + if r.resolvedClosing { + r.resolvedMu.Unlock() + return nil, nil, fmt.Errorf("cliproxy: embedded runtime is closing") + } + r.resolvedInflight++ + r.resolvedMu.Unlock() + return ctx, r.finishResolvedExecution, nil +} + +func (r *EmbeddedRuntime) finishResolvedExecution() { + r.resolvedMu.Lock() + r.resolvedInflight-- + if r.resolvedClosing && r.resolvedInflight == 0 && r.resolvedDrained != nil { + close(r.resolvedDrained) + r.resolvedDrained = nil + } + r.resolvedMu.Unlock() +} + +func (r *EmbeddedRuntime) closeResolvedExecutions(ctx context.Context) error { + r.resolvedMu.Lock() + r.resolvedClosing = true + if r.resolvedInflight == 0 { + r.resolvedMu.Unlock() + return nil + } + if r.resolvedDrained == nil { + r.resolvedDrained = make(chan struct{}) + } + drained := r.resolvedDrained + r.resolvedMu.Unlock() + select { + case <-drained: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func resolvedExecutorBase(kind NativeExecutorKind, apiBase string) string { + apiBase = strings.TrimRight(strings.TrimSpace(apiBase), "/") + if kind == NativeExecutorAnthropic { + apiBase = strings.TrimSuffix(apiBase, "/v1") + } + return apiBase +} + +func resolvedRequestInvalid(err error) bool { + var scoped interface{ IsRequestScoped() bool } + if errors.As(err, &scoped) && scoped.IsRequestScoped() { + return true + } + var status interface{ StatusCode() int } + if errors.As(err, &status) { + return status.StatusCode() == http.StatusBadRequest || status.StatusCode() == http.StatusUnprocessableEntity + } + return false +} diff --git a/sdk/cliproxy/resolved_execution_test.go b/sdk/cliproxy/resolved_execution_test.go new file mode 100644 index 00000000000..95d7b28d61e --- /dev/null +++ b/sdk/cliproxy/resolved_execution_test.go @@ -0,0 +1,244 @@ +package cliproxy + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + cliproxytranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +func TestEmbeddedRuntimeExecuteResolvedUsesNativeRetryAndExactModel(t *testing.T) { + var firstCalls, secondCalls atomic.Int32 + first := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + firstCalls.Add(1) + http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized) + })) + defer first.Close() + second := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + secondCalls.Add(1) + body, _ := io.ReadAll(r.Body) + if !strings.Contains(string(body), `"model":"deepseek-chat"`) { + t.Fatalf("body=%s, want exact runtime model", body) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"chatcmpl-ok","choices":[]}`)) + })) + defer second.Close() + + runtime := resolvedTestRuntime(t) + resp, err := runtime.ExecuteResolved(context.Background(), ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{ + {Auth: &coreauth.Auth{ID: "first", Attributes: map[string]string{"api_key": "one"}}, RuntimeModel: "deepseek-chat", APIBase: first.URL}, + {Auth: &coreauth.Auth{ID: "second", Attributes: map[string]string{"api_key": "two"}}, RuntimeModel: "deepseek-chat", APIBase: second.URL}, + }, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[]}`), Format: cliproxytranslator.FormatOpenAI}, + Options: cliproxyexecutor.Options{SourceFormat: cliproxytranslator.FormatOpenAI, ResponseFormat: cliproxytranslator.FormatOpenAI}, + }) + if err != nil { + t.Fatalf("ExecuteResolved: %v", err) + } + if !strings.Contains(string(resp.Payload), "chatcmpl-ok") || firstCalls.Load() != 1 || secondCalls.Load() != 1 { + t.Fatalf("payload=%s calls=%d/%d", resp.Payload, firstCalls.Load(), secondCalls.Load()) + } +} + +func TestEmbeddedRuntimeExecuteResolvedStopsOnClientRequestError(t *testing.T) { + var secondCalls atomic.Int32 + first := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, `{"error":"bad request"}`, http.StatusBadRequest) + })) + defer first.Close() + second := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + secondCalls.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer second.Close() + + runtime := resolvedTestRuntime(t) + _, err := runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{ + {Auth: &coreauth.Auth{ID: "first", Attributes: map[string]string{"api_key": "one"}}, RuntimeModel: "deepseek-chat", APIBase: first.URL}, + {Auth: &coreauth.Auth{ID: "second", Attributes: map[string]string{"api_key": "two"}}, RuntimeModel: "deepseek-chat", APIBase: second.URL}, + }, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[]}`), Format: cliproxytranslator.FormatOpenAI}, + Options: cliproxyexecutor.Options{SourceFormat: cliproxytranslator.FormatOpenAI, ResponseFormat: cliproxytranslator.FormatOpenAI}, + }) + if err == nil || secondCalls.Load() != 0 { + t.Fatalf("error=%v second_calls=%d, want terminal client error", err, secondCalls.Load()) + } +} + +func TestEmbeddedRuntimeExecuteResolvedAnthropicUsesExactVersionedPrefix(t *testing.T) { + var path string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + path = r.URL.Path + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg-ok","type":"message","content":[]}`)) + })) + defer upstream.Close() + + runtime := resolvedTestRuntime(t) + resp, err := runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorAnthropic, + Provider: "zhipu", + Candidates: []ResolvedExecutionCandidate{{ + Auth: &coreauth.Auth{ID: "zhipu", Attributes: map[string]string{"api_key": "key"}}, + RuntimeModel: "glm-5.2", + APIBase: upstream.URL + "/api/anthropic/v1", + }}, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[],"max_tokens":16}`), Format: cliproxytranslator.FormatClaude}, + Options: cliproxyexecutor.Options{SourceFormat: cliproxytranslator.FormatClaude, ResponseFormat: cliproxytranslator.FormatClaude}, + }) + if err != nil { + t.Fatalf("ExecuteResolved: %v", err) + } + if path != "/api/anthropic/v1/messages" || !strings.Contains(string(resp.Payload), "msg-ok") { + t.Fatalf("path=%q payload=%s", path, resp.Payload) + } +} + +func TestEmbeddedRuntimeExecuteResolvedStreamRetriesBeforePayloadWithoutServer(t *testing.T) { + var secondCalls atomic.Int32 + first := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized) + })) + defer first.Close() + second := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + secondCalls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-stream\",\"choices\":[]}\n\n")) + _, _ = w.Write([]byte("data: [DONE]\n\n")) + })) + defer second.Close() + + runtime := resolvedTestRuntime(t) + result, err := runtime.ExecuteResolvedStream(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{ + {Auth: &coreauth.Auth{ID: "first", Attributes: map[string]string{"api_key": "one"}}, RuntimeModel: "deepseek-chat", APIBase: first.URL}, + {Auth: &coreauth.Auth{ID: "second", Attributes: map[string]string{"api_key": "two"}}, RuntimeModel: "deepseek-chat", APIBase: second.URL}, + }, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[],"stream":true}`), Format: cliproxytranslator.FormatOpenAI}, + Options: cliproxyexecutor.Options{Stream: true, SourceFormat: cliproxytranslator.FormatOpenAI, ResponseFormat: cliproxytranslator.FormatOpenAI}, + }) + if err != nil { + t.Fatalf("ExecuteResolvedStream: %v", err) + } + var payload strings.Builder + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("stream chunk: %v", chunk.Err) + } + payload.Write(chunk.Payload) + } + if secondCalls.Load() != 1 || !strings.Contains(payload.String(), "chatcmpl-stream") { + t.Fatalf("calls=%d payload=%q", secondCalls.Load(), payload.String()) + } + if runtime.ServerStarted() || runtime.WatcherStarted() { + t.Fatal("resolved execution started CLIProxyAPI server or watcher") + } +} + +func TestEmbeddedRuntimeExecuteResolvedRejectsEmptyAnthropicBase(t *testing.T) { + runtime := resolvedTestRuntime(t) + _, err := runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorAnthropic, + Provider: "zhipu", + Candidates: []ResolvedExecutionCandidate{{Auth: &coreauth.Auth{Attributes: map[string]string{"api_key": "secret"}}, RuntimeModel: "glm-5.2"}}, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[]}`), Format: cliproxytranslator.FormatClaude}, + Options: cliproxyexecutor.Options{SourceFormat: cliproxytranslator.FormatClaude, ResponseFormat: cliproxytranslator.FormatClaude}, + }) + if err == nil || !strings.Contains(err.Error(), "API base") { + t.Fatalf("error=%v, want API base validation", err) + } +} + +func TestEmbeddedRuntimeExecuteResolvedDoesNotMutateGlobalModelRegistry(t *testing.T) { + before := len(registry.GetGlobalRegistry().GetAvailableModels("openai")) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"ok","choices":[]}`)) + })) + defer upstream.Close() + runtime := resolvedTestRuntime(t) + _, err := runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{{Auth: &coreauth.Auth{Attributes: map[string]string{"api_key": "secret"}}, RuntimeModel: "deepseek-chat", APIBase: upstream.URL}}, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[]}`), Format: cliproxytranslator.FormatOpenAI}, + Options: cliproxyexecutor.Options{SourceFormat: cliproxytranslator.FormatOpenAI, ResponseFormat: cliproxytranslator.FormatOpenAI}, + }) + if err != nil { + t.Fatalf("ExecuteResolved: %v", err) + } + after := len(registry.GetGlobalRegistry().GetAvailableModels("openai")) + if after != before { + t.Fatalf("global model count=%d, want %d", after, before) + } +} + +func TestEmbeddedRuntimeCloseWaitsForResolvedStreamAndRejectsNewCalls(t *testing.T) { + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"id\":\"started\",\"choices\":[]}\n\n")) + w.(http.Flusher).Flush() + <-release + _, _ = w.Write([]byte("data: [DONE]\n\n")) + })) + defer upstream.Close() + runtime := resolvedTestRuntime(t) + result, err := runtime.ExecuteResolvedStream(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{{Auth: &coreauth.Auth{Attributes: map[string]string{"api_key": "secret"}}, RuntimeModel: "deepseek-chat", APIBase: upstream.URL}}, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[],"stream":true}`), Format: cliproxytranslator.FormatOpenAI}, + Options: cliproxyexecutor.Options{Stream: true, SourceFormat: cliproxytranslator.FormatOpenAI, ResponseFormat: cliproxytranslator.FormatOpenAI}, + }) + if err != nil { + t.Fatalf("ExecuteResolvedStream: %v", err) + } + closeCtx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if err := runtime.Close(closeCtx); err == nil { + t.Fatal("Close returned before resolved stream drained") + } + _, err = runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{}) + if err == nil || !strings.Contains(err.Error(), "closing") { + t.Fatalf("post-close-start error=%v", err) + } + close(release) + for range result.Chunks { + } + if err := runtime.Close(context.Background()); err != nil { + t.Fatalf("retry Close: %v", err) + } +} + +func resolvedTestRuntime(t *testing.T) *EmbeddedRuntime { + t.Helper() + manager := coreauth.NewManager(nil, nil, nil) + runtime, err := NewEmbeddedRuntime(&config.Config{}, manager) + if err != nil { + t.Fatalf("NewEmbeddedRuntime: %v", err) + } + if err := runtime.Start(t.Context()); err != nil { + t.Fatalf("Start: %v", err) + } + return runtime +} From ab70ce333108b5548aedbfaca5955d531ea7796f Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 22:28:58 +0800 Subject: [PATCH 05/11] feat(sdk): harden host-resolved execution Preserve host auth identity and transport policy while executing request-local provider candidates. Carry interceptor metadata safely, suppress duplicate usage publication, fence resolved executions during shutdown, and cover retry, streaming, and mutation isolation behavior. --- sdk/cliproxy/auth/conductor.go | 15 +++ sdk/cliproxy/embedded_runtime.go | 2 + sdk/cliproxy/resolved_execution.go | 111 +++++++++++++++--- sdk/cliproxy/resolved_execution_test.go | 146 ++++++++++++++++++++++++ sdk/cliproxy/usage/context.go | 20 ++++ sdk/cliproxy/usage/context_test.go | 18 +++ sdk/cliproxy/usage/manager.go | 7 +- 7 files changed, 303 insertions(+), 16 deletions(-) create mode 100644 sdk/cliproxy/usage/context.go create mode 100644 sdk/cliproxy/usage/context_test.go diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index cf85a2ae7a4..81f8ebb8667 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -4630,6 +4630,21 @@ func (m *Manager) Executor(provider string) (ProviderExecutor, bool) { return executor, true } +// RoundTripperFor returns the host-configured transport for an auth. +// Embedded SDK hosts use WithRoundTripper to preserve egress policy without +// depending on the auth package's private context key. +func (m *Manager) RoundTripperFor(auth *Auth) http.RoundTripper { + return m.roundTripperFor(auth) +} + +func WithRoundTripper(ctx context.Context, rt http.RoundTripper) context.Context { + if ctx == nil { + ctx = context.Background() + } + ctx = context.WithValue(ctx, roundTripperContextKey{}, rt) + return context.WithValue(ctx, "cliproxy.roundtripper", rt) +} + // CloseExecutionSession asks all registered executors to release the supplied execution session. func (m *Manager) CloseExecutionSession(sessionID string) { sessionID = strings.TrimSpace(sessionID) diff --git a/sdk/cliproxy/embedded_runtime.go b/sdk/cliproxy/embedded_runtime.go index d9dd405aed1..0077226104b 100644 --- a/sdk/cliproxy/embedded_runtime.go +++ b/sdk/cliproxy/embedded_runtime.go @@ -5,6 +5,7 @@ import ( "fmt" "strings" "sync" + "sync/atomic" "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" @@ -32,6 +33,7 @@ type EmbeddedRuntime struct { resolvedClosing bool resolvedInflight int resolvedDrained chan struct{} + resolvedSequence atomic.Uint64 } const embeddedRuntimeRefreshInterval = 15 * time.Minute diff --git a/sdk/cliproxy/resolved_execution.go b/sdk/cliproxy/resolved_execution.go index 2ffa324a0db..c94ba5984d7 100644 --- a/sdk/cliproxy/resolved_execution.go +++ b/sdk/cliproxy/resolved_execution.go @@ -10,6 +10,8 @@ import ( nativeexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" ) type NativeExecutorKind string @@ -47,7 +49,8 @@ func (r *EmbeddedRuntime) ExecuteResolved(ctx context.Context, execution Resolve for _, candidate := range candidates { req := execution.Request req.Model = candidate.RuntimeModel - resp, errExec := executor.Execute(ctx, candidate.Auth, req, execution.Options) + execCtx, execReq, execOpts := r.prepareResolvedRequest(ctx, candidate, req, execution.Options, execution.Executor, execution.Request.Model) + resp, errExec := executor.Execute(execCtx, candidate.Auth, execReq, execOpts) if errExec == nil { return resp, nil } @@ -79,7 +82,8 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R for _, candidate := range candidates { req := execution.Request req.Model = candidate.RuntimeModel - result, errExec := executor.ExecuteStream(ctx, candidate.Auth, req, execution.Options) + execCtx, execReq, execOpts := r.prepareResolvedRequest(ctx, candidate, req, execution.Options, execution.Executor, execution.Request.Model) + result, errExec := executor.ExecuteStream(execCtx, candidate.Auth, execReq, execOpts) if errExec == nil { first, ok := <-result.Chunks if ok && first.Err != nil { @@ -129,8 +133,9 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R } type preparedResolvedCandidate struct { - Auth *coreauth.Auth - RuntimeModel string + Auth *coreauth.Auth + TransportAuth *coreauth.Auth + RuntimeModel string } func (r *EmbeddedRuntime) prepareResolvedCandidates(execution ResolvedExecutionRequest) (coreauth.ProviderExecutor, string, []preparedResolvedCandidate, error) { @@ -151,6 +156,7 @@ func (r *EmbeddedRuntime) prepareResolvedCandidates(execution ResolvedExecutionR default: return nil, "", nil, fmt.Errorf("cliproxy: unsupported resolved executor %q", execution.Executor) } + sequence := r.resolvedSequence.Add(1) candidates := make([]preparedResolvedCandidate, 0, len(execution.Candidates)) for index, candidate := range execution.Candidates { if candidate.Auth == nil { @@ -168,22 +174,97 @@ func (r *EmbeddedRuntime) prepareResolvedCandidates(execution ResolvedExecutionR if apiKey == "" { return nil, "", nil, fmt.Errorf("cliproxy: resolved API key is required") } - auth := &coreauth.Auth{ - ID: fmt.Sprintf("resolved-%d", index), - Provider: provider, - Label: candidate.Auth.Label, - Status: coreauth.StatusActive, - ProxyURL: candidate.Auth.ProxyURL, - Attributes: map[string]string{ - "api_key": apiKey, - "base_url": apiBase, - }, + transportAuth := candidate.Auth.Clone() + auth := candidate.Auth.Clone() + if strings.TrimSpace(auth.ID) == "" { + auth.ID = fmt.Sprintf("resolved-%p-%d-%d", r, sequence, index) + transportAuth.ID = auth.ID } - candidates = append(candidates, preparedResolvedCandidate{Auth: auth, RuntimeModel: model}) + auth.Provider = provider + auth.Status = coreauth.StatusActive + if auth.Attributes == nil { + auth.Attributes = make(map[string]string) + } + auth.Attributes["api_key"] = apiKey + auth.Attributes["base_url"] = apiBase + candidates = append(candidates, preparedResolvedCandidate{Auth: auth, TransportAuth: transportAuth, RuntimeModel: model}) } return executor, provider, candidates, nil } +func (r *EmbeddedRuntime) prepareResolvedRequest(ctx context.Context, candidate preparedResolvedCandidate, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, kind NativeExecutorKind, requestedModel string) (context.Context, cliproxyexecutor.Request, cliproxyexecutor.Options) { + ctx = coreusage.WithPublishingSuppressed(ctx) + if r.manager != nil { + if roundTripper := r.manager.RoundTripperFor(candidate.TransportAuth); roundTripper != nil { + ctx = coreauth.WithRoundTripper(ctx, roundTripper) + } + } + if opts.RequestAfterAuthInterceptor == nil { + return ctx, req, opts + } + opts.Headers = opts.Headers.Clone() + if opts.Headers == nil { + opts.Headers = make(http.Header) + } + toFormat := sdktranslator.FormatOpenAI + if kind == NativeExecutorAnthropic { + toFormat = sdktranslator.FormatClaude + } + intercepted := opts.RequestAfterAuthInterceptor(ctx, cliproxyexecutor.RequestAfterAuthInterceptRequest{ + SourceFormat: opts.SourceFormat, + ToFormat: toFormat, + Model: req.Model, + RequestedModel: resolvedRequestedModel(opts, requestedModel), + Stream: opts.Stream, + Headers: opts.Headers.Clone(), + Body: append([]byte(nil), req.Payload...), + Metadata: cloneResolvedMetadata(opts.Metadata), + }) + for _, key := range intercepted.ClearHeaders { + opts.Headers.Del(key) + } + for key, values := range intercepted.Headers { + opts.Headers.Del(key) + for _, value := range values { + opts.Headers.Add(key, value) + } + } + if len(intercepted.Body) > 0 { + req.Payload = append([]byte(nil), intercepted.Body...) + opts.OriginalRequest = append([]byte(nil), intercepted.Body...) + } + return ctx, req, opts +} + +func resolvedRequestedModel(opts cliproxyexecutor.Options, fallback string) string { + fallback = strings.TrimSpace(fallback) + if len(opts.Metadata) == 0 { + return fallback + } + switch value := opts.Metadata[cliproxyexecutor.RequestedModelMetadataKey].(type) { + case string: + if requested := strings.TrimSpace(value); requested != "" { + return requested + } + case []byte: + if requested := strings.TrimSpace(string(value)); requested != "" { + return requested + } + } + return fallback +} + +func cloneResolvedMetadata(metadata map[string]any) map[string]any { + if metadata == nil { + return nil + } + cloned := make(map[string]any, len(metadata)) + for key, value := range metadata { + cloned[key] = value + } + return cloned +} + func (r *EmbeddedRuntime) beginResolvedExecution(ctx context.Context) (context.Context, func(), error) { if r == nil || r.service == nil || r.service.cfg == nil { return nil, nil, fmt.Errorf("cliproxy: embedded runtime is unavailable") diff --git a/sdk/cliproxy/resolved_execution_test.go b/sdk/cliproxy/resolved_execution_test.go index 95d7b28d61e..68b66bbb60f 100644 --- a/sdk/cliproxy/resolved_execution_test.go +++ b/sdk/cliproxy/resolved_execution_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "net/url" "strings" "sync/atomic" "testing" @@ -230,6 +231,151 @@ func TestEmbeddedRuntimeCloseWaitsForResolvedStreamAndRejectsNewCalls(t *testing } } +func TestEmbeddedRuntimeExecuteResolvedAppliesInterceptorAndHostRoundTripper(t *testing.T) { + var gotBody string + var gotTransportAuth *coreauth.Auth + var gotIntercept cliproxyexecutor.RequestAfterAuthInterceptRequest + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + gotBody = string(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"ok","choices":[]}`)) + })) + defer upstream.Close() + upstreamURL, _ := url.Parse(upstream.URL) + manager := coreauth.NewManager(nil, nil, nil) + manager.SetRoundTripperProvider(roundTripperProviderFunc(func(auth *coreauth.Auth) http.RoundTripper { + gotTransportAuth = auth.Clone() + return roundTripperFunc(func(req *http.Request) (*http.Response, error) { + clone := req.Clone(req.Context()) + clone.URL.Scheme = upstreamURL.Scheme + clone.URL.Host = upstreamURL.Host + return http.DefaultTransport.RoundTrip(clone) + }) + })) + runtime, err := NewEmbeddedRuntime(&config.Config{}, manager) + if err != nil { + t.Fatalf("NewEmbeddedRuntime: %v", err) + } + if err := runtime.Start(t.Context()); err != nil { + t.Fatalf("Start: %v", err) + } + callerHeaders := http.Header{"X-Caller": []string{"preserved"}} + callerMetadata := map[string]any{ + cliproxyexecutor.RequestedModelMetadataKey: "canonical-model", + "project_id": "project-1", + } + candidateAuth := &coreauth.Auth{ + ID: "credential-17", + Provider: "deepseek", + Metadata: map[string]any{"tenant": "tenant-1"}, + Attributes: map[string]string{ + "api_key": "secret", + "egress_zone": "cn-east", + }, + } + _, err = runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{{Auth: candidateAuth, RuntimeModel: "deepseek-chat", APIBase: "https://blocked.invalid"}}, + Request: cliproxyexecutor.Request{Model: "request-model", Payload: []byte(`{"messages":[]}`), Format: cliproxytranslator.FormatOpenAI}, + Options: cliproxyexecutor.Options{ + SourceFormat: cliproxytranslator.FormatOpenAI, + ResponseFormat: cliproxytranslator.FormatOpenAI, + Headers: callerHeaders, + Metadata: callerMetadata, + RequestAfterAuthInterceptor: func(_ context.Context, input cliproxyexecutor.RequestAfterAuthInterceptRequest) cliproxyexecutor.RequestAfterAuthInterceptResponse { + gotIntercept = input + input.Metadata["project_id"] = "mutated" + return cliproxyexecutor.RequestAfterAuthInterceptResponse{Headers: http.Header{"X-Resolved-Policy": []string{"applied"}}, Body: []byte(`{"messages":[{"role":"user","content":"intercepted"}]}`)} + }, + }, + }) + if err != nil { + t.Fatalf("ExecuteResolved: %v", err) + } + if !strings.Contains(gotBody, "intercepted") { + t.Fatalf("body=%s", gotBody) + } + if callerHeaders.Get("X-Resolved-Policy") != "" || callerHeaders.Get("X-Caller") != "preserved" { + t.Fatalf("caller headers mutated: %#v", callerHeaders) + } + if gotTransportAuth == nil || gotTransportAuth.ID != "credential-17" || gotTransportAuth.Metadata["tenant"] != "tenant-1" || gotTransportAuth.Attributes["egress_zone"] != "cn-east" { + t.Fatalf("transport auth=%#v, want original host identity and metadata", gotTransportAuth) + } + if gotIntercept.Model != "deepseek-chat" || gotIntercept.RequestedModel != "canonical-model" || gotIntercept.Metadata["project_id"] != "mutated" { + t.Fatalf("intercept=%#v", gotIntercept) + } + if callerMetadata["project_id"] != "project-1" { + t.Fatalf("caller metadata mutated: %#v", callerMetadata) + } + if candidateAuth.Attributes["base_url"] != "" || candidateAuth.Attributes["egress_zone"] != "cn-east" { + t.Fatalf("candidate auth mutated: %#v", candidateAuth.Attributes) + } +} + +func TestEmbeddedRuntimeExecuteResolvedInterceptorHandlesNilHeaders(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"ok","choices":[]}`)) + })) + defer upstream.Close() + runtime := resolvedTestRuntime(t) + _, err := runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{{Auth: &coreauth.Auth{Attributes: map[string]string{"api_key": "secret"}}, RuntimeModel: "deepseek-chat", APIBase: upstream.URL}}, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[]}`), Format: cliproxytranslator.FormatOpenAI}, + Options: cliproxyexecutor.Options{ + SourceFormat: cliproxytranslator.FormatOpenAI, + ResponseFormat: cliproxytranslator.FormatOpenAI, + RequestAfterAuthInterceptor: func(_ context.Context, input cliproxyexecutor.RequestAfterAuthInterceptRequest) cliproxyexecutor.RequestAfterAuthInterceptResponse { + return cliproxyexecutor.RequestAfterAuthInterceptResponse{Headers: http.Header{"X-Test": []string{"value"}}} + }, + }, + }) + if err != nil { + t.Fatalf("ExecuteResolved: %v", err) + } +} + +func TestEmbeddedRuntimeResolvedCandidatesUseUniqueFallbackAuthIDs(t *testing.T) { + runtime := resolvedTestRuntime(t) + request := ResolvedExecutionRequest{ + Executor: NativeExecutorOpenAICompatible, + Provider: "deepseek", + Candidates: []ResolvedExecutionCandidate{{ + Auth: &coreauth.Auth{Attributes: map[string]string{"api_key": "secret"}}, + RuntimeModel: "deepseek-chat", + APIBase: "https://example.invalid", + }}, + } + _, _, first, err := runtime.prepareResolvedCandidates(request) + if err != nil { + t.Fatalf("first prepare: %v", err) + } + _, _, second, err := runtime.prepareResolvedCandidates(request) + if err != nil { + t.Fatalf("second prepare: %v", err) + } + if first[0].Auth.ID == "" || first[0].Auth.ID == second[0].Auth.ID { + t.Fatalf("fallback auth IDs=%q/%q, want unique non-empty IDs", first[0].Auth.ID, second[0].Auth.ID) + } + if first[0].TransportAuth.ID != first[0].Auth.ID { + t.Fatalf("transport auth ID=%q, want execution auth ID %q", first[0].TransportAuth.ID, first[0].Auth.ID) + } +} + +type roundTripperProviderFunc func(*coreauth.Auth) http.RoundTripper + +func (f roundTripperProviderFunc) RoundTripperFor(auth *coreauth.Auth) http.RoundTripper { + return f(auth) +} + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + func resolvedTestRuntime(t *testing.T) *EmbeddedRuntime { t.Helper() manager := coreauth.NewManager(nil, nil, nil) diff --git a/sdk/cliproxy/usage/context.go b/sdk/cliproxy/usage/context.go new file mode 100644 index 00000000000..846be5e4668 --- /dev/null +++ b/sdk/cliproxy/usage/context.go @@ -0,0 +1,20 @@ +package usage + +import "context" + +type suppressPublishingKey struct{} + +func WithPublishingSuppressed(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, suppressPublishingKey{}, true) +} + +func PublishingSuppressed(ctx context.Context) bool { + if ctx == nil { + return false + } + value, _ := ctx.Value(suppressPublishingKey{}).(bool) + return value +} diff --git a/sdk/cliproxy/usage/context_test.go b/sdk/cliproxy/usage/context_test.go new file mode 100644 index 00000000000..cad6ad287b2 --- /dev/null +++ b/sdk/cliproxy/usage/context_test.go @@ -0,0 +1,18 @@ +package usage + +import ( + "context" + "testing" +) + +func TestPublishingSuppressedContext(t *testing.T) { + if PublishingSuppressed(nil) { + t.Fatal("nil context unexpectedly suppresses usage") + } + if PublishingSuppressed(context.Background()) { + t.Fatal("background context unexpectedly suppresses usage") + } + if !PublishingSuppressed(WithPublishingSuppressed(context.Background())) { + t.Fatal("suppressed context was not recognized") + } +} diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index 0c9529f197b..d1903d103b2 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -376,7 +376,12 @@ func RegisterPlugin(plugin Plugin) { DefaultManager().Register(plugin) } func RegisterNamedPlugin(name string, plugin Plugin) { DefaultManager().RegisterNamed(name, plugin) } // PublishRecord publishes a record using the default manager. -func PublishRecord(ctx context.Context, record Record) { DefaultManager().Publish(ctx, record) } +func PublishRecord(ctx context.Context, record Record) { + if PublishingSuppressed(ctx) { + return + } + DefaultManager().Publish(ctx, record) +} // StartDefault starts the default manager's dispatcher. func StartDefault(ctx context.Context) { DefaultManager().Start(ctx) } From bf7ceb947a582d1163bc018e0be00ee3aee39dd6 Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 22:48:01 +0800 Subject: [PATCH 06/11] fix(runtime): cancel abandoned executions and normalize response models Make embedded resolved streams cancelable during runtime shutdown, isolate nested interceptor metadata, and apply usage suppression at the manager boundary. Normalize requested model names in native translator responses so host-resolved physical models do not leak through OpenAI, Claude, or Codex protocols. --- internal/runtime/executor/claude_executor.go | 4 +- internal/runtime/executor/codex_executor.go | 6 +- .../executor/openai_compat_executor.go | 6 +- .../openai_compat_executor_compact_test.go | 25 ++++++++ sdk/cliproxy/embedded_runtime.go | 5 ++ sdk/cliproxy/resolved_execution.go | 60 +++++++++++++++-- sdk/cliproxy/resolved_execution_test.go | 33 +++++----- sdk/cliproxy/usage/context_test.go | 10 +++ sdk/cliproxy/usage/manager.go | 6 +- sdk/translator/registry.go | 64 ++++++++++++++++++- sdk/translator/registry_test.go | 22 +++++++ 11 files changed, 209 insertions(+), 32 deletions(-) diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index b3e348abb5d..fbbabe51ae2 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -380,7 +380,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r ctx, to, responseFormat, - req.Model, + requestedModel, opts.OriginalRequest, bodyForTranslation, data, @@ -607,7 +607,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A ctx, to, responseFormat, - req.Model, + requestedModel, opts.OriginalRequest, bodyForTranslation, bytes.Clone(line), diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index c4922bec3dc..8167785de0d 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -986,7 +986,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re var param any clientCompletedData := applyCodexIdentityExposeResponsePayload(completedData, identityState) - out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, clientCompletedData, ¶m) + out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, requestedModel, originalPayload, body, clientCompletedData, ¶m) resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} return resp, nil } @@ -1096,7 +1096,7 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A reporter.EnsurePublished(ctx) var param any clientData := applyCodexIdentityExposeResponsePayload(upstreamData, identityState) - out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, originalPayload, body, clientData, ¶m) + out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, requestedModel, originalPayload, body, clientData, ¶m) resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} return resp, nil } @@ -1264,7 +1264,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au } translatedLine = applyCodexIdentityExposeResponsePayload(translatedLine, identityState) - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m) + chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, requestedModel, originalPayload, body, translatedLine, ¶m) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go index 7588161430c..abc233786bb 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -194,7 +194,7 @@ func (e *OpenAICompatExecutor) Execute(ctx context.Context, auth *cliproxyauth.A reporter.EnsurePublished(ctx) // Translate response back to source format when needed var param any - out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, body, ¶m) + out := sdktranslator.TranslateNonStream(ctx, to, responseFormat, requestedModel, opts.OriginalRequest, translated, body, ¶m) resp = cliproxyexecutor.Response{Payload: out, Headers: httpResp.Header.Clone()} return resp, nil } @@ -423,7 +423,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy } // OpenAI-compatible streams must use SSE data lines. - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, bytes.Clone(trimmedLine), ¶m) + chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, requestedModel, opts.OriginalRequest, translated, bytes.Clone(trimmedLine), ¶m) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: @@ -443,7 +443,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy // In case the upstream close the stream without a terminal [DONE] marker. // Feed a synthetic done marker through the translator so pending // response.completed events are still emitted exactly once. - chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, req.Model, opts.OriginalRequest, translated, []byte("data: [DONE]"), ¶m) + chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, requestedModel, opts.OriginalRequest, translated, []byte("data: [DONE]"), ¶m) for i := range chunks { select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: diff --git a/internal/runtime/executor/openai_compat_executor_compact_test.go b/internal/runtime/executor/openai_compat_executor_compact_test.go index cf5fe636b26..d3fb07f2752 100644 --- a/internal/runtime/executor/openai_compat_executor_compact_test.go +++ b/internal/runtime/executor/openai_compat_executor_compact_test.go @@ -106,6 +106,31 @@ func TestOpenAICompatExecutorPayloadOverrideWinsOverThinkingSuffix(t *testing.T) } } +func TestOpenAICompatExecutorNormalizesRequestedModelInResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"chatcmpl_1","object":"chat.completion","model":"physical-model","choices":[]}`)) + })) + defer server.Close() + + executor := NewOpenAICompatExecutor("deepseek", &config.Config{}) + auth := &cliproxyauth.Auth{Attributes: map[string]string{"base_url": server.URL, "api_key": "test"}} + resp, err := executor.Execute(t.Context(), auth, cliproxyexecutor.Request{ + Model: "physical-model", + Payload: []byte(`{"model":"canonical-model","messages":[]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + ResponseFormat: sdktranslator.FormatOpenAI, + Metadata: map[string]any{cliproxyexecutor.RequestedModelMetadataKey: "canonical-model"}, + }) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + if model := gjson.GetBytes(resp.Payload, "model").String(); model != "canonical-model" { + t.Fatalf("response model=%q payload=%s", model, resp.Payload) + } +} + func TestOpenAICompatExecutorImagesGenerationsPassthrough(t *testing.T) { var gotPath string var gotBody []byte diff --git a/sdk/cliproxy/embedded_runtime.go b/sdk/cliproxy/embedded_runtime.go index 0077226104b..e717b0a9046 100644 --- a/sdk/cliproxy/embedded_runtime.go +++ b/sdk/cliproxy/embedded_runtime.go @@ -34,6 +34,8 @@ type EmbeddedRuntime struct { resolvedInflight int resolvedDrained chan struct{} resolvedSequence atomic.Uint64 + resolvedContext context.Context + resolvedCancel context.CancelFunc } const embeddedRuntimeRefreshInterval = 15 * time.Minute @@ -54,11 +56,14 @@ func NewEmbeddedRuntime(cfg *config.Config, manager *coreauth.Manager) (*Embedde manager.SetConfig(cfg) manager.SetOAuthModelAlias(cfg.OAuthModelAlias) + resolvedContext, resolvedCancel := context.WithCancel(context.Background()) return &EmbeddedRuntime{ manager: manager, service: &Service{cfg: cfg, coreManager: manager}, registeredAuthIDs: make(map[string]struct{}), installedExecutors: make(map[string]coreauth.ProviderExecutor), + resolvedContext: resolvedContext, + resolvedCancel: resolvedCancel, }, nil } diff --git a/sdk/cliproxy/resolved_execution.go b/sdk/cliproxy/resolved_execution.go index c94ba5984d7..8fdcef721a2 100644 --- a/sdk/cliproxy/resolved_execution.go +++ b/sdk/cliproxy/resolved_execution.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "strings" + "sync" nativeexecutor "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -85,7 +86,14 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R execCtx, execReq, execOpts := r.prepareResolvedRequest(ctx, candidate, req, execution.Options, execution.Executor, execution.Request.Model) result, errExec := executor.ExecuteStream(execCtx, candidate.Auth, execReq, execOpts) if errExec == nil { - first, ok := <-result.Chunks + var first cliproxyexecutor.StreamChunk + var ok bool + select { + case first, ok = <-result.Chunks: + case <-ctx.Done(): + finish() + return nil, ctx.Err() + } if ok && first.Err != nil { if resolvedRequestInvalid(first.Err) { finish() @@ -105,7 +113,17 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R return } } - for chunk := range result.Chunks { + for { + var chunk cliproxyexecutor.StreamChunk + var more bool + select { + case chunk, more = <-result.Chunks: + if !more { + return + } + case <-ctx.Done(): + return + } select { case forwarded <- chunk: case <-ctx.Done(): @@ -260,11 +278,30 @@ func cloneResolvedMetadata(metadata map[string]any) map[string]any { } cloned := make(map[string]any, len(metadata)) for key, value := range metadata { - cloned[key] = value + cloned[key] = cloneResolvedMetadataValue(value) } return cloned } +func cloneResolvedMetadataValue(value any) any { + switch typed := value.(type) { + case map[string]any: + return cloneResolvedMetadata(typed) + case []any: + cloned := make([]any, len(typed)) + for index := range typed { + cloned[index] = cloneResolvedMetadataValue(typed[index]) + } + return cloned + case []string: + return append([]string(nil), typed...) + case []byte: + return append([]byte(nil), typed...) + default: + return value + } +} + func (r *EmbeddedRuntime) beginResolvedExecution(ctx context.Context) (context.Context, func(), error) { if r == nil || r.service == nil || r.service.cfg == nil { return nil, nil, fmt.Errorf("cliproxy: embedded runtime is unavailable") @@ -287,8 +324,19 @@ func (r *EmbeddedRuntime) beginResolvedExecution(ctx context.Context) (context.C return nil, nil, fmt.Errorf("cliproxy: embedded runtime is closing") } r.resolvedInflight++ + runtimeContext := r.resolvedContext r.resolvedMu.Unlock() - return ctx, r.finishResolvedExecution, nil + execContext, cancel := context.WithCancel(ctx) + stopRuntimeCancel := context.AfterFunc(runtimeContext, cancel) + var finishOnce sync.Once + finish := func() { + finishOnce.Do(func() { + stopRuntimeCancel() + cancel() + r.finishResolvedExecution() + }) + } + return execContext, finish, nil } func (r *EmbeddedRuntime) finishResolvedExecution() { @@ -304,6 +352,10 @@ func (r *EmbeddedRuntime) finishResolvedExecution() { func (r *EmbeddedRuntime) closeResolvedExecutions(ctx context.Context) error { r.resolvedMu.Lock() r.resolvedClosing = true + resolvedCancel := r.resolvedCancel + if resolvedCancel != nil { + resolvedCancel() + } if r.resolvedInflight == 0 { r.resolvedMu.Unlock() return nil diff --git a/sdk/cliproxy/resolved_execution_test.go b/sdk/cliproxy/resolved_execution_test.go index 68b66bbb60f..a49e72ac3c7 100644 --- a/sdk/cliproxy/resolved_execution_test.go +++ b/sdk/cliproxy/resolved_execution_test.go @@ -193,14 +193,14 @@ func TestEmbeddedRuntimeExecuteResolvedDoesNotMutateGlobalModelRegistry(t *testi } } -func TestEmbeddedRuntimeCloseWaitsForResolvedStreamAndRejectsNewCalls(t *testing.T) { - release := make(chan struct{}) - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { +func TestEmbeddedRuntimeCloseCancelsAbandonedResolvedStreamAndRejectsNewCalls(t *testing.T) { + upstreamCanceled := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte("data: {\"id\":\"started\",\"choices\":[]}\n\n")) w.(http.Flusher).Flush() - <-release - _, _ = w.Write([]byte("data: [DONE]\n\n")) + <-r.Context().Done() + close(upstreamCanceled) })) defer upstream.Close() runtime := resolvedTestRuntime(t) @@ -214,20 +214,21 @@ func TestEmbeddedRuntimeCloseWaitsForResolvedStreamAndRejectsNewCalls(t *testing if err != nil { t.Fatalf("ExecuteResolvedStream: %v", err) } - closeCtx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + closeCtx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() - if err := runtime.Close(closeCtx); err == nil { - t.Fatal("Close returned before resolved stream drained") + if err := runtime.Close(closeCtx); err != nil { + t.Fatalf("Close: %v", err) } _, err = runtime.ExecuteResolved(t.Context(), ResolvedExecutionRequest{}) - if err == nil || !strings.Contains(err.Error(), "closing") { + if err == nil || !strings.Contains(err.Error(), "not started") { t.Fatalf("post-close-start error=%v", err) } - close(release) for range result.Chunks { } - if err := runtime.Close(context.Background()); err != nil { - t.Fatalf("retry Close: %v", err) + select { + case <-upstreamCanceled: + case <-time.After(time.Second): + t.Fatal("upstream request was not canceled") } } @@ -263,7 +264,7 @@ func TestEmbeddedRuntimeExecuteResolvedAppliesInterceptorAndHostRoundTripper(t * callerHeaders := http.Header{"X-Caller": []string{"preserved"}} callerMetadata := map[string]any{ cliproxyexecutor.RequestedModelMetadataKey: "canonical-model", - "project_id": "project-1", + "project": map[string]any{"id": "project-1"}, } candidateAuth := &coreauth.Auth{ ID: "credential-17", @@ -286,7 +287,7 @@ func TestEmbeddedRuntimeExecuteResolvedAppliesInterceptorAndHostRoundTripper(t * Metadata: callerMetadata, RequestAfterAuthInterceptor: func(_ context.Context, input cliproxyexecutor.RequestAfterAuthInterceptRequest) cliproxyexecutor.RequestAfterAuthInterceptResponse { gotIntercept = input - input.Metadata["project_id"] = "mutated" + input.Metadata["project"].(map[string]any)["id"] = "mutated" return cliproxyexecutor.RequestAfterAuthInterceptResponse{Headers: http.Header{"X-Resolved-Policy": []string{"applied"}}, Body: []byte(`{"messages":[{"role":"user","content":"intercepted"}]}`)} }, }, @@ -303,10 +304,10 @@ func TestEmbeddedRuntimeExecuteResolvedAppliesInterceptorAndHostRoundTripper(t * if gotTransportAuth == nil || gotTransportAuth.ID != "credential-17" || gotTransportAuth.Metadata["tenant"] != "tenant-1" || gotTransportAuth.Attributes["egress_zone"] != "cn-east" { t.Fatalf("transport auth=%#v, want original host identity and metadata", gotTransportAuth) } - if gotIntercept.Model != "deepseek-chat" || gotIntercept.RequestedModel != "canonical-model" || gotIntercept.Metadata["project_id"] != "mutated" { + if gotIntercept.Model != "deepseek-chat" || gotIntercept.RequestedModel != "canonical-model" || gotIntercept.Metadata["project"].(map[string]any)["id"] != "mutated" { t.Fatalf("intercept=%#v", gotIntercept) } - if callerMetadata["project_id"] != "project-1" { + if callerMetadata["project"].(map[string]any)["id"] != "project-1" { t.Fatalf("caller metadata mutated: %#v", callerMetadata) } if candidateAuth.Attributes["base_url"] != "" || candidateAuth.Attributes["egress_zone"] != "cn-east" { diff --git a/sdk/cliproxy/usage/context_test.go b/sdk/cliproxy/usage/context_test.go index cad6ad287b2..0de68fa3cc1 100644 --- a/sdk/cliproxy/usage/context_test.go +++ b/sdk/cliproxy/usage/context_test.go @@ -16,3 +16,13 @@ func TestPublishingSuppressedContext(t *testing.T) { t.Fatal("suppressed context was not recognized") } } + +func TestManagerPublishHonorsSuppressedContext(t *testing.T) { + manager := NewManager(1) + manager.Publish(WithPublishingSuppressed(context.Background()), Record{}) + manager.mu.Lock() + defer manager.mu.Unlock() + if len(manager.queue) != 0 { + t.Fatalf("suppressed queue length=%d", len(manager.queue)) + } +} diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index d1903d103b2..d7743713ce6 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -310,6 +310,9 @@ func (m *Manager) Publish(ctx context.Context, record Record) { if m == nil { return } + if PublishingSuppressed(ctx) { + return + } // ensure worker is running even if Start was not called explicitly m.Start(context.Background()) m.mu.Lock() @@ -377,9 +380,6 @@ func RegisterNamedPlugin(name string, plugin Plugin) { DefaultManager().Register // PublishRecord publishes a record using the default manager. func PublishRecord(ctx context.Context, record Record) { - if PublishingSuppressed(ctx) { - return - } DefaultManager().Publish(ctx, record) } diff --git a/sdk/translator/registry.go b/sdk/translator/registry.go index ad4d351dbe5..f5215afb5aa 100644 --- a/sdk/translator/registry.go +++ b/sdk/translator/registry.go @@ -1,7 +1,9 @@ package translator import ( + "bytes" "context" + "strings" "sync" log "github.com/sirupsen/logrus" @@ -173,6 +175,9 @@ func (r *Registry) TranslateStream(ctx context.Context, from, to Format, model s outputs[i] = hooks.NormalizeResponseAfter(ctx, from, to, model, originalRequestRawJSON, requestRawJSON, output, true) } } + for i := range outputs { + outputs[i] = normalizeResponseModel(outputs[i], model, true) + } return outputs } @@ -200,7 +205,64 @@ func (r *Registry) TranslateNonStream(ctx context.Context, from, to Format, mode if hooks != nil { body = hooks.NormalizeResponseAfter(ctx, from, to, model, originalRequestRawJSON, requestRawJSON, body, false) } - return body + return normalizeResponseModel(body, model, false) +} + +func normalizeResponseModel(body []byte, model string, stream bool) []byte { + model = strings.TrimSpace(model) + if model == "" || len(body) == 0 { + return body + } + if !stream { + return normalizeJSONResponseModel(body, model) + } + trimmed := bytes.TrimSpace(body) + if !bytes.HasPrefix(trimmed, []byte("data:")) { + return normalizeJSONResponseModel(body, model) + } + payloadOffset := bytes.Index(body, []byte("data:")) + len("data:") + for payloadOffset < len(body) && (body[payloadOffset] == ' ' || body[payloadOffset] == '\t') { + payloadOffset++ + } + payloadEnd := len(body) + for payloadEnd > payloadOffset { + switch body[payloadEnd-1] { + case '\r', '\n', ' ', '\t': + payloadEnd-- + default: + goto payloadReady + } + } +payloadReady: + payload := body[payloadOffset:payloadEnd] + if bytes.Equal(payload, []byte("[DONE]")) { + return body + } + updated := normalizeJSONResponseModel(payload, model) + if bytes.Equal(updated, payload) { + return body + } + out := make([]byte, 0, len(body)-len(payload)+len(updated)) + out = append(out, body[:payloadOffset]...) + out = append(out, updated...) + out = append(out, body[payloadEnd:]...) + return out +} + +func normalizeJSONResponseModel(body []byte, model string) []byte { + if !gjson.ValidBytes(body) { + return body + } + updated := body + for _, path := range []string{"model", "response.model", "message.model"} { + if !gjson.GetBytes(updated, path).Exists() { + continue + } + if next, err := sjson.SetBytes(updated, path, model); err == nil { + updated = next + } + } + return updated } // TranslateTokenCount applies the registered token count response translator. diff --git a/sdk/translator/registry_test.go b/sdk/translator/registry_test.go index f154cb397ab..5b7838c0759 100644 --- a/sdk/translator/registry_test.go +++ b/sdk/translator/registry_test.go @@ -1,6 +1,7 @@ package translator import ( + "bytes" "context" "testing" @@ -402,3 +403,24 @@ func TestPluginNormalizersChainAfterNative(t *testing.T) { t.Fatalf("plugin translators should not run when native transformers exist, calls=%v", hooks.calls) } } + +func TestTranslateNonStreamNormalizesRequestedModelWithoutNativeTransformer(t *testing.T) { + r := NewRegistry() + got := r.TranslateNonStream(context.Background(), FormatOpenAI, FormatOpenAI, "canonical-model", nil, nil, []byte(`{"id":"chatcmpl-1","model":"physical-model","choices":[]}`), nil) + if model := gjson.GetBytes(got, "model").String(); model != "canonical-model" { + t.Fatalf("model=%q body=%s", model, got) + } +} + +func TestTranslateStreamNormalizesNestedRequestedModelAndPreservesFraming(t *testing.T) { + r := NewRegistry() + input := []byte("data: {\"type\":\"response.completed\",\"response\":{\"model\":\"physical-model\"}}\n\n") + got := r.TranslateStream(context.Background(), FormatOpenAIResponse, FormatOpenAIResponse, "canonical-model", nil, nil, input, nil) + if len(got) != 1 || !bytes.HasSuffix(got[0], []byte("\n\n")) { + t.Fatalf("stream framing=%q", got) + } + payload := bytes.TrimSpace(bytes.TrimPrefix(bytes.TrimSpace(got[0]), []byte("data:"))) + if model := gjson.GetBytes(payload, "response.model").String(); model != "canonical-model" { + t.Fatalf("model=%q payload=%s", model, payload) + } +} From 5704e35833a4e2794c9f5410022b55e08c49ee7d Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 23:03:17 +0800 Subject: [PATCH 07/11] fix(runtime): normalize complete SSE frames safely Route event-prefixed and Claude same-format SSE frames through canonical model normalization, deepen interceptor metadata isolation across typed containers, and make Antigravity test cleanup concurrency-safe so the full executor race suite is reliable. --- .../antigravity_executor_credits_test.go | 15 ++-- internal/runtime/executor/claude_executor.go | 14 ++-- .../runtime/executor/claude_executor_test.go | 12 +++- sdk/cliproxy/resolved_execution.go | 71 ++++++++++++++++--- sdk/cliproxy/resolved_execution_test.go | 11 ++- sdk/translator/registry.go | 37 +++++++--- sdk/translator/registry_test.go | 5 +- 7 files changed, 129 insertions(+), 36 deletions(-) diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go index e516483d999..186e0638719 100644 --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -21,10 +21,17 @@ import ( ) func resetAntigravityCreditsRetryState() { - antigravityCreditsFailureByAuth = sync.Map{} - antigravityShortCooldownByAuth = sync.Map{} - antigravityCreditsBalanceByAuth = sync.Map{} - antigravityCreditsHintRefreshByID = sync.Map{} + clearAntigravityTestSyncMap(&antigravityCreditsFailureByAuth) + clearAntigravityTestSyncMap(&antigravityShortCooldownByAuth) + clearAntigravityTestSyncMap(&antigravityCreditsBalanceByAuth) + clearAntigravityTestSyncMap(&antigravityCreditsHintRefreshByID) +} + +func clearAntigravityTestSyncMap(values *sync.Map) { + values.Range(func(key, _ any) bool { + values.Delete(key) + return true + }) } type fakeAntigravityKVClient struct { diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index fbbabe51ae2..0f56366c3fe 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -552,18 +552,22 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A scanner := bufio.NewScanner(decodedBody) scanner.Buffer(nil, 52_428_800) // 50MB var event bytes.Buffer + var param any flushEvent := func() bool { if event.Len() == 0 { return true } cloned := bytes.Clone(event.Bytes()) event.Reset() - select { - case out <- cliproxyexecutor.StreamChunk{Payload: cloned}: - return true - case <-ctx.Done(): - return false + chunks := sdktranslator.TranslateStream(ctx, to, responseFormat, requestedModel, opts.OriginalRequest, bodyForTranslation, cloned, ¶m) + for i := range chunks { + select { + case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: + case <-ctx.Done(): + return false + } } + return true } for scanner.Scan() { line := scanner.Bytes() diff --git a/internal/runtime/executor/claude_executor_test.go b/internal/runtime/executor/claude_executor_test.go index 1325edcae4e..51e25746826 100644 --- a/internal/runtime/executor/claude_executor_test.go +++ b/internal/runtime/executor/claude_executor_test.go @@ -1237,9 +1237,13 @@ func TestClaudeExecutor_ExecuteStreamStripsOpenAIEncryptedThinkingBeforeUpstream } func TestClaudeExecutor_ExecuteStreamDirectPassthroughEmitsCompleteSSEEvents(t *testing.T) { + startData := `{"type":"message_start","message":{"id":"msg_1","model":"physical-model"}}` firstData := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}` secondData := `{"type":"message_stop"}` - upstreamStream := "event: content_block_delta\n" + + upstreamStream := "event: message_start\n" + + "data: " + startData + "\n" + + "\n" + + "event: content_block_delta\n" + "data: " + firstData + "\n" + "\n" + "event: message_stop\n" + @@ -1262,7 +1266,10 @@ func TestClaudeExecutor_ExecuteStreamDirectPassthroughEmitsCompleteSSEEvents(t * result, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, - }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}) + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("claude"), + Metadata: map[string]any{cliproxyexecutor.RequestedModelMetadataKey: "canonical-model"}, + }) if err != nil { t.Fatalf("ExecuteStream() error = %v", err) } @@ -1276,6 +1283,7 @@ func TestClaudeExecutor_ExecuteStreamDirectPassthroughEmitsCompleteSSEEvents(t * } want := []string{ + "event: message_start\n" + "data: " + strings.Replace(startData, "physical-model", "canonical-model", 1) + "\n\n", "event: content_block_delta\n" + "data: " + firstData + "\n\n", "event: message_stop\n" + "data: " + secondData + "\n\n", } diff --git a/sdk/cliproxy/resolved_execution.go b/sdk/cliproxy/resolved_execution.go index 8fdcef721a2..d5741ebc828 100644 --- a/sdk/cliproxy/resolved_execution.go +++ b/sdk/cliproxy/resolved_execution.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "net/http" + "reflect" "strings" "sync" @@ -284,19 +285,67 @@ func cloneResolvedMetadata(metadata map[string]any) map[string]any { } func cloneResolvedMetadataValue(value any) any { - switch typed := value.(type) { - case map[string]any: - return cloneResolvedMetadata(typed) - case []any: - cloned := make([]any, len(typed)) - for index := range typed { - cloned[index] = cloneResolvedMetadataValue(typed[index]) + cloned := cloneResolvedReflectValue(reflect.ValueOf(value)) + if !cloned.IsValid() { + return nil + } + return cloned.Interface() +} + +func cloneResolvedReflectValue(value reflect.Value) reflect.Value { + if !value.IsValid() { + return reflect.Value{} + } + switch value.Kind() { + case reflect.Interface: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + cloned := cloneResolvedReflectValue(value.Elem()) + wrapped := reflect.New(value.Type()).Elem() + wrapped.Set(cloned) + return wrapped + case reflect.Pointer: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + cloned := reflect.New(value.Type().Elem()) + cloned.Elem().Set(cloneResolvedReflectValue(value.Elem())) + return cloned + case reflect.Map: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + cloned := reflect.MakeMapWithSize(value.Type(), value.Len()) + iterator := value.MapRange() + for iterator.Next() { + cloned.SetMapIndex(cloneResolvedReflectValue(iterator.Key()), cloneResolvedReflectValue(iterator.Value())) + } + return cloned + case reflect.Slice: + if value.IsNil() { + return reflect.Zero(value.Type()) + } + cloned := reflect.MakeSlice(value.Type(), value.Len(), value.Len()) + for index := 0; index < value.Len(); index++ { + cloned.Index(index).Set(cloneResolvedReflectValue(value.Index(index))) + } + return cloned + case reflect.Array: + cloned := reflect.New(value.Type()).Elem() + for index := 0; index < value.Len(); index++ { + cloned.Index(index).Set(cloneResolvedReflectValue(value.Index(index))) + } + return cloned + case reflect.Struct: + cloned := reflect.New(value.Type()).Elem() + cloned.Set(value) + for index := 0; index < value.NumField(); index++ { + if cloned.Field(index).CanSet() && value.Field(index).CanInterface() { + cloned.Field(index).Set(cloneResolvedReflectValue(value.Field(index))) + } } return cloned - case []string: - return append([]string(nil), typed...) - case []byte: - return append([]byte(nil), typed...) default: return value } diff --git a/sdk/cliproxy/resolved_execution_test.go b/sdk/cliproxy/resolved_execution_test.go index a49e72ac3c7..86f3edb1039 100644 --- a/sdk/cliproxy/resolved_execution_test.go +++ b/sdk/cliproxy/resolved_execution_test.go @@ -262,9 +262,13 @@ func TestEmbeddedRuntimeExecuteResolvedAppliesInterceptorAndHostRoundTripper(t * t.Fatalf("Start: %v", err) } callerHeaders := http.Header{"X-Caller": []string{"preserved"}} + nestedMap := map[string]string{"id": "project-1"} + nestedPointer := &struct{ Value string }{Value: "original"} callerMetadata := map[string]any{ cliproxyexecutor.RequestedModelMetadataKey: "canonical-model", - "project": map[string]any{"id": "project-1"}, + "project": map[string]any{"id": "project-1"}, + "typed_map": nestedMap, + "pointer": nestedPointer, } candidateAuth := &coreauth.Auth{ ID: "credential-17", @@ -288,6 +292,8 @@ func TestEmbeddedRuntimeExecuteResolvedAppliesInterceptorAndHostRoundTripper(t * RequestAfterAuthInterceptor: func(_ context.Context, input cliproxyexecutor.RequestAfterAuthInterceptRequest) cliproxyexecutor.RequestAfterAuthInterceptResponse { gotIntercept = input input.Metadata["project"].(map[string]any)["id"] = "mutated" + input.Metadata["typed_map"].(map[string]string)["id"] = "mutated" + input.Metadata["pointer"].(*struct{ Value string }).Value = "mutated" return cliproxyexecutor.RequestAfterAuthInterceptResponse{Headers: http.Header{"X-Resolved-Policy": []string{"applied"}}, Body: []byte(`{"messages":[{"role":"user","content":"intercepted"}]}`)} }, }, @@ -310,6 +316,9 @@ func TestEmbeddedRuntimeExecuteResolvedAppliesInterceptorAndHostRoundTripper(t * if callerMetadata["project"].(map[string]any)["id"] != "project-1" { t.Fatalf("caller metadata mutated: %#v", callerMetadata) } + if nestedMap["id"] != "project-1" || nestedPointer.Value != "original" { + t.Fatalf("typed caller metadata mutated: map=%#v pointer=%#v", nestedMap, nestedPointer) + } if candidateAuth.Attributes["base_url"] != "" || candidateAuth.Attributes["egress_zone"] != "cn-east" { t.Fatalf("candidate auth mutated: %#v", candidateAuth.Attributes) } diff --git a/sdk/translator/registry.go b/sdk/translator/registry.go index f5215afb5aa..354b1a0a032 100644 --- a/sdk/translator/registry.go +++ b/sdk/translator/registry.go @@ -216,24 +216,24 @@ func normalizeResponseModel(body []byte, model string, stream bool) []byte { if !stream { return normalizeJSONResponseModel(body, model) } - trimmed := bytes.TrimSpace(body) - if !bytes.HasPrefix(trimmed, []byte("data:")) { + dataOffset := sseDataOffset(body) + if dataOffset < 0 { return normalizeJSONResponseModel(body, model) } - payloadOffset := bytes.Index(body, []byte("data:")) + len("data:") + payloadOffset := dataOffset + len("data:") for payloadOffset < len(body) && (body[payloadOffset] == ' ' || body[payloadOffset] == '\t') { payloadOffset++ } payloadEnd := len(body) - for payloadEnd > payloadOffset { - switch body[payloadEnd-1] { - case '\r', '\n', ' ', '\t': - payloadEnd-- - default: - goto payloadReady - } + if lineEnd := bytes.IndexByte(body[payloadOffset:], '\n'); lineEnd >= 0 { + payloadEnd = payloadOffset + lineEnd + } + if payloadEnd > payloadOffset && body[payloadEnd-1] == '\r' { + payloadEnd-- + } + for payloadEnd > payloadOffset && (body[payloadEnd-1] == ' ' || body[payloadEnd-1] == '\t') { + payloadEnd-- } -payloadReady: payload := body[payloadOffset:payloadEnd] if bytes.Equal(payload, []byte("[DONE]")) { return body @@ -249,6 +249,21 @@ payloadReady: return out } +func sseDataOffset(body []byte) int { + for offset := 0; offset < len(body); { + lineEnd := bytes.IndexByte(body[offset:], '\n') + if lineEnd < 0 { + lineEnd = len(body) - offset + } + line := bytes.TrimSpace(body[offset : offset+lineEnd]) + if bytes.HasPrefix(line, []byte("data:")) { + return offset + bytes.Index(body[offset:offset+lineEnd], []byte("data:")) + } + offset += lineEnd + 1 + } + return -1 +} + func normalizeJSONResponseModel(body []byte, model string) []byte { if !gjson.ValidBytes(body) { return body diff --git a/sdk/translator/registry_test.go b/sdk/translator/registry_test.go index 5b7838c0759..bcd88a03b30 100644 --- a/sdk/translator/registry_test.go +++ b/sdk/translator/registry_test.go @@ -414,12 +414,13 @@ func TestTranslateNonStreamNormalizesRequestedModelWithoutNativeTransformer(t *t func TestTranslateStreamNormalizesNestedRequestedModelAndPreservesFraming(t *testing.T) { r := NewRegistry() - input := []byte("data: {\"type\":\"response.completed\",\"response\":{\"model\":\"physical-model\"}}\n\n") + input := []byte("event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"model\":\"physical-model\"}}\n\n") got := r.TranslateStream(context.Background(), FormatOpenAIResponse, FormatOpenAIResponse, "canonical-model", nil, nil, input, nil) if len(got) != 1 || !bytes.HasSuffix(got[0], []byte("\n\n")) { t.Fatalf("stream framing=%q", got) } - payload := bytes.TrimSpace(bytes.TrimPrefix(bytes.TrimSpace(got[0]), []byte("data:"))) + dataOffset := bytes.Index(got[0], []byte("data:")) + payload := bytes.TrimSpace(got[0][dataOffset+len("data:"):]) if model := gjson.GetBytes(payload, "response.model").String(); model != "canonical-model" { t.Fatalf("model=%q payload=%s", model, payload) } From eeb892f3523032db63ba2486d7f72d4f4622cd68 Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Thu, 16 Jul 2026 23:36:32 +0800 Subject: [PATCH 08/11] fix(runtime): normalize Zhipu Claude thinking streams --- sdk/cliproxy/resolved_execution.go | 178 ++++++++++++++++++++++++ sdk/cliproxy/resolved_execution_test.go | 116 ++++++++++++++- 2 files changed, 292 insertions(+), 2 deletions(-) diff --git a/sdk/cliproxy/resolved_execution.go b/sdk/cliproxy/resolved_execution.go index d5741ebc828..6f21eb1e90c 100644 --- a/sdk/cliproxy/resolved_execution.go +++ b/sdk/cliproxy/resolved_execution.go @@ -1,7 +1,9 @@ package cliproxy import ( + "bytes" "context" + "encoding/json" "errors" "fmt" "net/http" @@ -51,9 +53,11 @@ func (r *EmbeddedRuntime) ExecuteResolved(ctx context.Context, execution Resolve for _, candidate := range candidates { req := execution.Request req.Model = candidate.RuntimeModel + normalizer := newResolvedProviderResponseNormalizer(execution.Provider, execution.Executor, cliproxyexecutor.ResponseFormatOrSource(execution.Options)) execCtx, execReq, execOpts := r.prepareResolvedRequest(ctx, candidate, req, execution.Options, execution.Executor, execution.Request.Model) resp, errExec := executor.Execute(execCtx, candidate.Auth, execReq, execOpts) if errExec == nil { + resp.Payload = normalizer.Normalize(resp.Payload) return resp, nil } if err := ctx.Err(); err != nil { @@ -84,6 +88,7 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R for _, candidate := range candidates { req := execution.Request req.Model = candidate.RuntimeModel + normalizer := newResolvedProviderResponseNormalizer(execution.Provider, execution.Executor, cliproxyexecutor.ResponseFormatOrSource(execution.Options)) execCtx, execReq, execOpts := r.prepareResolvedRequest(ctx, candidate, req, execution.Options, execution.Executor, execution.Request.Model) result, errExec := executor.ExecuteStream(execCtx, candidate.Auth, execReq, execOpts) if errExec == nil { @@ -103,6 +108,22 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R lastErr = first.Err continue } + for ok { + first.Payload = normalizer.Normalize(first.Payload) + if len(first.Payload) > 0 { + break + } + select { + case first, ok = <-result.Chunks: + if ok && first.Err != nil { + finish() + return nil, first.Err + } + case <-ctx.Done(): + finish() + return nil, ctx.Err() + } + } forwarded := make(chan cliproxyexecutor.StreamChunk) go func() { defer close(forwarded) @@ -125,6 +146,17 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R case <-ctx.Done(): return } + if chunk.Err != nil { + select { + case forwarded <- chunk: + case <-ctx.Done(): + } + return + } + chunk.Payload = normalizer.Normalize(chunk.Payload) + if len(chunk.Payload) == 0 { + continue + } select { case forwarded <- chunk: case <-ctx.Done(): @@ -151,6 +183,152 @@ func (r *EmbeddedRuntime) ExecuteResolvedStream(ctx context.Context, execution R return nil, &coreauth.Error{Code: "auth_not_found", Message: fmt.Sprintf("no %s auth available", provider)} } +type resolvedProviderResponseNormalizer struct { + stripClaudeThinking bool + thinkingIndexes map[int64]struct{} + suppressedIndexes map[int64]struct{} +} + +func newResolvedProviderResponseNormalizer(provider string, executor NativeExecutorKind, responseFormat sdktranslator.Format) *resolvedProviderResponseNormalizer { + strip := strings.EqualFold(strings.TrimSpace(provider), "zhipu") && executor == NativeExecutorAnthropic && responseFormat == sdktranslator.FormatClaude + return &resolvedProviderResponseNormalizer{stripClaudeThinking: strip, thinkingIndexes: make(map[int64]struct{}), suppressedIndexes: make(map[int64]struct{})} +} + +func (n *resolvedProviderResponseNormalizer) Normalize(payload []byte) []byte { + if n == nil || !n.stripClaudeThinking { + return payload + } + return n.stripClaudeThinkingPayload(payload) +} + +func (n *resolvedProviderResponseNormalizer) stripClaudeThinkingPayload(payload []byte) []byte { + if len(payload) == 0 { + return payload + } + jsonStart := 0 + jsonEnd := len(payload) + if dataOffset := sseResolvedDataOffset(payload); dataOffset >= 0 { + jsonStart = bytes.IndexByte(payload[dataOffset:], '{') + if jsonStart < 0 { + return payload + } + jsonStart += dataOffset + lineEnd := bytes.IndexByte(payload[jsonStart:], '\n') + if lineEnd >= 0 { + jsonEnd = jsonStart + lineEnd + } + } + jsonPayload := bytes.TrimSpace(payload[jsonStart:jsonEnd]) + var message map[string]any + if err := json.Unmarshal(jsonPayload, &message); err != nil { + return payload + } + changed := false + eventType, _ := message["type"].(string) + index := resolvedJSONInt64(message["index"]) + switch eventType { + case "content_block_start": + block, _ := message["content_block"].(map[string]any) + blockType, _ := block["type"].(string) + if blockType == "thinking" || blockType == "redacted_thinking" { + n.thinkingIndexes[index] = struct{}{} + n.suppressedIndexes[index] = struct{}{} + return nil + } + case "content_block_delta": + delta, _ := message["delta"].(map[string]any) + deltaType, _ := delta["type"].(string) + _, suppressed := n.thinkingIndexes[index] + if suppressed || deltaType == "thinking_delta" || deltaType == "signature_delta" { + n.thinkingIndexes[index] = struct{}{} + return nil + } + case "content_block_stop": + if _, suppressed := n.thinkingIndexes[index]; suppressed { + delete(n.thinkingIndexes, index) + return nil + } + } + if _, hasIndex := message["index"]; hasIndex { + message["index"] = index - n.suppressedBefore(index) + changed = true + } + if content, ok := message["content"].([]any); ok { + filtered := make([]any, 0, len(content)) + for _, item := range content { + block, isBlock := item.(map[string]any) + blockType, _ := block["type"].(string) + if isBlock && (blockType == "thinking" || blockType == "redacted_thinking") { + changed = true + continue + } + filtered = append(filtered, item) + } + if changed { + message["content"] = filtered + } + } + if stopReason, ok := message["stop_reason"].(string); ok && strings.Contains(strings.ToLower(stopReason), "thinking") { + message["stop_reason"] = "end_turn" + changed = true + } + if delta, ok := message["delta"].(map[string]any); ok { + if stopReason, ok := delta["stop_reason"].(string); ok && strings.Contains(strings.ToLower(stopReason), "thinking") { + delta["stop_reason"] = "end_turn" + changed = true + } + } + if !changed { + return payload + } + cleaned, err := json.Marshal(message) + if err != nil { + return payload + } + out := make([]byte, 0, len(payload)-len(jsonPayload)+len(cleaned)) + out = append(out, payload[:jsonStart]...) + out = append(out, cleaned...) + out = append(out, payload[jsonEnd:]...) + return out +} + +func (n *resolvedProviderResponseNormalizer) suppressedBefore(index int64) int64 { + var count int64 + for suppressed := range n.suppressedIndexes { + if suppressed < index { + count++ + } + } + return count +} + +func resolvedJSONInt64(value any) int64 { + switch typed := value.(type) { + case float64: + return int64(typed) + case json.Number: + parsed, _ := typed.Int64() + return parsed + default: + return 0 + } +} + +func sseResolvedDataOffset(payload []byte) int { + for offset := 0; offset < len(payload); { + lineEnd := bytes.IndexByte(payload[offset:], '\n') + if lineEnd < 0 { + lineEnd = len(payload) - offset + } + line := bytes.TrimSpace(payload[offset : offset+lineEnd]) + if bytes.HasPrefix(line, []byte("data:")) { + return offset + } + offset += lineEnd + 1 + } + return -1 +} + type preparedResolvedCandidate struct { Auth *coreauth.Auth TransportAuth *coreauth.Auth diff --git a/sdk/cliproxy/resolved_execution_test.go b/sdk/cliproxy/resolved_execution_test.go index 86f3edb1039..eb0a17e291c 100644 --- a/sdk/cliproxy/resolved_execution_test.go +++ b/sdk/cliproxy/resolved_execution_test.go @@ -2,6 +2,7 @@ package cliproxy import ( "context" + "errors" "io" "net/http" "net/http/httptest" @@ -88,7 +89,7 @@ func TestEmbeddedRuntimeExecuteResolvedAnthropicUsesExactVersionedPrefix(t *test upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { path = r.URL.Path w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"id":"msg-ok","type":"message","content":[]}`)) + _, _ = w.Write([]byte(`{"id":"msg-ok","type":"message","content":[{"type":"thinking","thinking":"hidden"},{"type":"text","text":"ok"}],"stop_reason":"thinking_end"}`)) })) defer upstream.Close() @@ -107,11 +108,108 @@ func TestEmbeddedRuntimeExecuteResolvedAnthropicUsesExactVersionedPrefix(t *test if err != nil { t.Fatalf("ExecuteResolved: %v", err) } - if path != "/api/anthropic/v1/messages" || !strings.Contains(string(resp.Payload), "msg-ok") { + if path != "/api/anthropic/v1/messages" || !strings.Contains(string(resp.Payload), "msg-ok") || strings.Contains(string(resp.Payload), `"type":"thinking"`) || !strings.Contains(string(resp.Payload), `"stop_reason":"end_turn"`) { t.Fatalf("path=%q payload=%s", path, resp.Payload) } } +func TestEmbeddedRuntimeExecuteResolvedStreamStripsZhipuClaudeThinkingAndPreservesSSE(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg-1","model":"glm-5.2"}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"hidden"}}`, + ``, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"ok"}}`, + ``, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":1}`, + ``, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"thinking_end"}}`, + ``, + `event: message_stop`, + `data: {"type":"message_stop"}`, + ``, + }, "\n") + "\n")) + })) + defer upstream.Close() + runtime := resolvedTestRuntime(t) + result, err := runtime.ExecuteResolvedStream(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorAnthropic, + Provider: "zhipu", + Candidates: []ResolvedExecutionCandidate{{Auth: &coreauth.Auth{Attributes: map[string]string{"api_key": "key"}}, RuntimeModel: "glm-5.2", APIBase: upstream.URL}}, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[],"max_tokens":16,"stream":true}`), Format: cliproxytranslator.FormatClaude}, + Options: cliproxyexecutor.Options{Stream: true, SourceFormat: cliproxytranslator.FormatClaude}, + }) + if err != nil { + t.Fatalf("ExecuteResolvedStream: %v", err) + } + var payload strings.Builder + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("chunk: %v", chunk.Err) + } + payload.Write(chunk.Payload) + } + if strings.Contains(payload.String(), "thinking") || !strings.Contains(payload.String(), `"type":"text"`) || !strings.Contains(payload.String(), `"text":"ok"`) || !strings.Contains(payload.String(), `"stop_reason":"end_turn"`) || strings.Contains(payload.String(), `"index":1`) || !strings.Contains(payload.String(), `"index":0`) || !strings.HasSuffix(payload.String(), "\n\n") { + t.Fatalf("payload=%q", payload.String()) + } +} + +func TestEmbeddedRuntimeDoesNotRetryAfterHiddenZhipuStreamChunk(t *testing.T) { + var secondCalls atomic.Int32 + manager := coreauth.NewManager(nil, nil, nil) + manager.SetRoundTripperProvider(roundTripperProviderFunc(func(auth *coreauth.Auth) http.RoundTripper { + return roundTripperFunc(func(req *http.Request) (*http.Response, error) { + if auth.ID == "second" { + secondCalls.Add(1) + } + body := io.ReadCloser(&errorAfterReader{ + reader: strings.NewReader("event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\n"), + err: errors.New("upstream stream failed"), + }) + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: body, Request: req}, nil + }) + })) + runtime, err := NewEmbeddedRuntime(&config.Config{}, manager) + if err != nil { + t.Fatalf("NewEmbeddedRuntime: %v", err) + } + if err := runtime.Start(t.Context()); err != nil { + t.Fatalf("Start: %v", err) + } + _, err = runtime.ExecuteResolvedStream(t.Context(), ResolvedExecutionRequest{ + Executor: NativeExecutorAnthropic, + Provider: "zhipu", + Candidates: []ResolvedExecutionCandidate{ + {Auth: &coreauth.Auth{ID: "first", Attributes: map[string]string{"api_key": "one"}}, RuntimeModel: "glm-5.2", APIBase: "https://first.invalid"}, + {Auth: &coreauth.Auth{ID: "second", Attributes: map[string]string{"api_key": "two"}}, RuntimeModel: "glm-5.2", APIBase: "https://second.invalid"}, + }, + Request: cliproxyexecutor.Request{Payload: []byte(`{"messages":[],"max_tokens":16,"stream":true}`), Format: cliproxytranslator.FormatClaude}, + Options: cliproxyexecutor.Options{Stream: true, SourceFormat: cliproxytranslator.FormatClaude}, + }) + if err == nil || !strings.Contains(err.Error(), "upstream stream failed") { + t.Fatalf("error=%v", err) + } + if secondCalls.Load() != 0 { + t.Fatalf("second credential calls=%d", secondCalls.Load()) + } +} + func TestEmbeddedRuntimeExecuteResolvedStreamRetriesBeforePayloadWithoutServer(t *testing.T) { var secondCalls atomic.Int32 first := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { @@ -386,6 +484,20 @@ type roundTripperFunc func(*http.Request) (*http.Response, error) func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } +type errorAfterReader struct { + reader *strings.Reader + err error +} + +func (r *errorAfterReader) Read(payload []byte) (int, error) { + if r.reader.Len() > 0 { + return r.reader.Read(payload) + } + return 0, r.err +} + +func (*errorAfterReader) Close() error { return nil } + func resolvedTestRuntime(t *testing.T) *EmbeddedRuntime { t.Helper() manager := coreauth.NewManager(nil, nil, nil) From ef74f33a2f9e865e51bec00e531b8b6ce5f0c156 Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Fri, 17 Jul 2026 00:09:17 +0800 Subject: [PATCH 09/11] ci(release): preserve immutable fork artifacts --- .github/workflows/release.yaml | 39 +++++----------------------------- 1 file changed, 5 insertions(+), 34 deletions(-) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 92941ebac6f..97affa7d6a1 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -86,10 +86,10 @@ jobs: } > "$updated_notes_file" if gh release view "$GITHUB_REF_NAME" >/dev/null 2>&1; then - gh release edit "$GITHUB_REF_NAME" --title "$GITHUB_REF_NAME" --notes-file "$updated_notes_file" - else - gh release create "$GITHUB_REF_NAME" --title "$GITHUB_REF_NAME" --notes-file "$updated_notes_file" + echo "release already exists for immutable tag $GITHUB_REF_NAME" >&2 + exit 1 fi + gh release create "$GITHUB_REF_NAME" --title "$GITHUB_REF_NAME" --notes-file "$updated_notes_file" build-hosted: name: build ${{ matrix.target }} @@ -201,36 +201,7 @@ jobs: printf '%s\n' "${assets[@]}" >&2 exit 1 fi - gh release upload "$GITHUB_REF_NAME" "${assets[@]}" --clobber - - name: Refresh release checksums - continue-on-error: true - shell: bash - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - set -euo pipefail - tmp_dir="$(mktemp -d)" - trap 'rm -rf "$tmp_dir"' EXIT - gh release download "$GITHUB_REF_NAME" \ - --pattern 'CLIProxyAPI_*.tar.gz' \ - --pattern 'CLIProxyAPI_*.zip' \ - --dir "$tmp_dir" \ - --clobber - ( - cd "$tmp_dir" - shopt -s nullglob - archives=(CLIProxyAPI_*.tar.gz CLIProxyAPI_*.zip) - if [[ ${#archives[@]} -eq 0 ]]; then - echo "No release archives found" - exit 0 - fi - if command -v sha256sum >/dev/null 2>&1; then - sha256sum "${archives[@]}" | sort -k2 > checksums.txt - else - shasum -a 256 "${archives[@]}" | sort -k2 > checksums.txt - fi - ) - gh release upload "$GITHUB_REF_NAME" "$tmp_dir/checksums.txt" --clobber + gh release upload "$GITHUB_REF_NAME" "${assets[@]}" build-linux-glibc: name: build linux-${{ matrix.goarch }} glibc @@ -667,4 +638,4 @@ jobs: exit 0 fi sha256sum "${archives[@]}" | sort -k2 > checksums.txt - gh release upload "$GITHUB_REF_NAME" checksums.txt --clobber + gh release upload "$GITHUB_REF_NAME" checksums.txt From fd505b85afc2241649f330d2444e5075533dbfa8 Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Fri, 17 Jul 2026 00:33:16 +0800 Subject: [PATCH 10/11] ci: enforce compatibility fork verification --- .github/workflows/pr-test-build.yml | 6 ++++++ .gitlab-ci.yml | 31 +++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+) create mode 100644 .gitlab-ci.yml diff --git a/.github/workflows/pr-test-build.yml b/.github/workflows/pr-test-build.yml index f1f0e2879c7..2c2e1f9c851 100644 --- a/.github/workflows/pr-test-build.yml +++ b/.github/workflows/pr-test-build.yml @@ -23,3 +23,9 @@ jobs: run: | go build -o test-output ./cmd/server rm -f test-output + - name: Test + run: go test ./... -count=1 + - name: Focused race tests + run: go test -race ./sdk/cliproxy ./sdk/translator ./internal/runtime/executor -count=1 + - name: Vet + run: go vet ./... diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml new file mode 100644 index 00000000000..1cc13151eb4 --- /dev/null +++ b/.gitlab-ci.yml @@ -0,0 +1,31 @@ +stages: + - verify + +default: + image: docker.m.daocloud.io/library/golang:1.26-bookworm + variables: + GOPROXY: https://goproxy.cn,direct + GOFLAGS: -buildvcs=false + cache: + key: go-mod-${CI_COMMIT_REF_SLUG} + paths: + - .cache/go-build + - .cache/go-mod + before_script: + - export GOCACHE="$CI_PROJECT_DIR/.cache/go-build" + - export GOMODCACHE="$CI_PROJECT_DIR/.cache/go-mod" + +test: + stage: verify + script: + - go test ./... -count=1 + +race-focused: + stage: verify + script: + - go test -race ./sdk/cliproxy ./sdk/translator ./internal/runtime/executor -count=1 + +vet: + stage: verify + script: + - go vet ./... From a877371789bf5f608df506077f6ceb5a1de0bb23 Mon Sep 17 00:00:00 2001 From: Zhexuan Yang Date: Fri, 17 Jul 2026 00:39:16 +0800 Subject: [PATCH 11/11] ci: support canonical GitLab runner schema --- .gitlab-ci.yml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 1cc13151eb4..c276d2b7e73 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -1,11 +1,12 @@ stages: - verify +variables: + GOPROXY: https://goproxy.cn,direct + GOFLAGS: -buildvcs=false + default: image: docker.m.daocloud.io/library/golang:1.26-bookworm - variables: - GOPROXY: https://goproxy.cn,direct - GOFLAGS: -buildvcs=false cache: key: go-mod-${CI_COMMIT_REF_SLUG} paths: