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/.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 diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml new file mode 100644 index 00000000000..c276d2b7e73 --- /dev/null +++ b/.gitlab-ci.yml @@ -0,0 +1,32 @@ +stages: + - verify + +variables: + GOPROXY: https://goproxy.cn,direct + GOFLAGS: -buildvcs=false + +default: + image: docker.m.daocloud.io/library/golang:1.26-bookworm + 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 ./... 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 b3e348abb5d..0f56366c3fe 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, @@ -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() @@ -607,7 +611,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/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/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/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 fedf4951862..81f8ebb8667 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" @@ -258,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 @@ -502,6 +504,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 +2161,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 +2303,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) { @@ -4527,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) @@ -5593,6 +5711,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() @@ -5604,14 +5723,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 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 loop.run(ctx) + 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. @@ -5621,6 +5849,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/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 new file mode 100644 index 00000000000..e717b0a9046 --- /dev/null +++ b/sdk/cliproxy/embedded_runtime.go @@ -0,0 +1,370 @@ +package cliproxy + +import ( + "context" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + "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 + + opMu sync.Mutex + mu sync.Mutex + started bool + closed bool + modelsReconciled bool + activeModels map[string]struct{} + registeredAuthIDs map[string]struct{} + installedExecutors map[string]coreauth.ProviderExecutor + refreshLease *coreauth.AutoRefreshLease + + resolvedMu sync.Mutex + resolvedClosing bool + resolvedInflight int + resolvedDrained chan struct{} + resolvedSequence atomic.Uint64 + resolvedContext context.Context + resolvedCancel context.CancelFunc +} + +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 +// 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") + } + if manager == nil { + return nil, fmt.Errorf("cliproxy: embedded runtime auth manager is required") + } + + manager.SetRoundTripperProviderIfAbsent(newDefaultRoundTripperProvider()) + 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 +} + +// 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.opMu.Lock() + defer r.opMu.Unlock() + if errContext := ctx.Err(); errContext != nil { + return errContext + } + r.mu.Lock() + 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() + 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.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() + 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.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.manager.SyncRuntimeAuth(auth) + if prepared == nil { + return fmt.Errorf("cliproxy: failed to synchronize auth %q", auth.ID) + } + r.registerAuthModelsLocked(runtimeCtx, prepared) + if errContext := ctx.Err(); errContext != nil { + return errContext + } + 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.opMu.Lock() + defer r.opMu.Unlock() + if errContext := ctx.Err(); errContext != nil { + return errContext + } + if errState := r.requireStarted(); 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.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.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(runtimeCtx, auth) + if errContext := ctx.Err(); errContext != nil { + return errContext + } + } + 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.opMu.Lock() + defer r.opMu.Unlock() + r.mu.Lock() + if r.closed { + r.mu.Unlock() + 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) + if refreshErr != nil { + return refreshErr + } + r.refreshLease = nil + } + for authID := range r.registeredAuthIDs { + registry.GetGlobalRegistry().UnregisterClient(authID) + r.manager.RefreshSchedulerEntry(authID) + } + 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) + r.mu.Lock() + r.closed = true + r.started = false + r.mu.Unlock() + return refreshErr +} + +// 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) requireStarted() error { + r.mu.Lock() + defer r.mu.Unlock() + 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 + } + 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 + } + r.service.reconcileModelsForAuth(ctx, auth, activeModels) + r.registeredAuthIDs[auth.ID] = struct{}{} +} + +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) + } + 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 { + if auth == nil { + return "" + } + if providerKey, _, isCompat := openAICompatInfoFromAuth(auth); isCompat { + if providerKey != "" { + return providerKey + } + return "openai-compatibility" + } + return strings.ToLower(strings.TrimSpace(auth.Provider)) +} diff --git a/sdk/cliproxy/embedded_runtime_test.go b/sdk/cliproxy/embedded_runtime_test.go new file mode 100644 index 00000000000..989d44cd4cf --- /dev/null +++ b/sdk/cliproxy/embedded_runtime_test.go @@ -0,0 +1,380 @@ +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" + "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 + + refreshOnce sync.Once + refreshed chan struct{} + release chan struct{} +} + +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) { + if e.refreshed != nil { + e.refreshOnce.Do(func() { close(e.refreshed) }) + } + if e.release != nil { + <-e.release + } + 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 { + 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") + } +} + +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") + } +} + +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) + } +} diff --git a/sdk/cliproxy/resolved_execution.go b/sdk/cliproxy/resolved_execution.go new file mode 100644 index 00000000000..6f21eb1e90c --- /dev/null +++ b/sdk/cliproxy/resolved_execution.go @@ -0,0 +1,621 @@ +package cliproxy + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "reflect" + "strings" + "sync" + + 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 + +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 + 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 { + 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 + 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 { + 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() + return nil, first.Err + } + 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) + defer finish() + if ok { + select { + case forwarded <- first: + case <-ctx.Done(): + return + } + } + for { + var chunk cliproxyexecutor.StreamChunk + var more bool + select { + case chunk, more = <-result.Chunks: + if !more { + return + } + 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(): + 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 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 + 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) + } + sequence := r.resolvedSequence.Add(1) + 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") + } + 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 + } + 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] = cloneResolvedMetadataValue(value) + } + return cloned +} + +func cloneResolvedMetadataValue(value any) any { + 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 + 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") + } + 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++ + runtimeContext := r.resolvedContext + r.resolvedMu.Unlock() + 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() { + 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 + resolvedCancel := r.resolvedCancel + if resolvedCancel != nil { + resolvedCancel() + } + 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..eb0a17e291c --- /dev/null +++ b/sdk/cliproxy/resolved_execution_test.go @@ -0,0 +1,512 @@ +package cliproxy + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "net/url" + "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":[{"type":"thinking","thinking":"hidden"},{"type":"text","text":"ok"}],"stop_reason":"thinking_end"}`)) + })) + 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") || 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) { + 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 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() + <-r.Context().Done() + close(upstreamCanceled) + })) + 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(), time.Second) + defer cancel() + 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(), "not started") { + t.Fatalf("post-close-start error=%v", err) + } + for range result.Chunks { + } + select { + case <-upstreamCanceled: + case <-time.After(time.Second): + t.Fatal("upstream request was not canceled") + } +} + +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"}} + 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"}, + "typed_map": nestedMap, + "pointer": nestedPointer, + } + 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"].(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"}]}`)} + }, + }, + }) + 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"].(map[string]any)["id"] != "mutated" { + t.Fatalf("intercept=%#v", gotIntercept) + } + 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) + } +} + +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) } + +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) + 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 +} diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index 61c6dad812d..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,17 +1076,24 @@ 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) { if a == nil || a.ID == "" { return @@ -1906,6 +1895,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.ToLower(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 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..0de68fa3cc1 --- /dev/null +++ b/sdk/cliproxy/usage/context_test.go @@ -0,0 +1,28 @@ +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") + } +} + +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 0c9529f197b..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() @@ -376,7 +379,9 @@ 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) { + DefaultManager().Publish(ctx, record) +} // StartDefault starts the default manager's dispatcher. func StartDefault(ctx context.Context) { DefaultManager().Start(ctx) } diff --git a/sdk/translator/registry.go b/sdk/translator/registry.go index ad4d351dbe5..354b1a0a032 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,79 @@ 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) + } + dataOffset := sseDataOffset(body) + if dataOffset < 0 { + return normalizeJSONResponseModel(body, model) + } + payloadOffset := dataOffset + len("data:") + for payloadOffset < len(body) && (body[payloadOffset] == ' ' || body[payloadOffset] == '\t') { + payloadOffset++ + } + payloadEnd := len(body) + 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-- + } + 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 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 + } + 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..bcd88a03b30 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,25 @@ 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("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) + } + 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) + } +}