diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index 06204da81..aeb059d3e 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -3,6 +3,7 @@ package executor import ( "bytes" "context" + "errors" "fmt" "net/http" "strings" @@ -26,6 +27,42 @@ type ClaudeExecutor struct { oauthProfileFetcher claudeOAuthProfileFetcher } +type claudeOAuthCancellationError struct { + cause error +} + +func (e *claudeOAuthCancellationError) Error() string { + if e == nil || e.cause == nil { + return "" + } + return e.cause.Error() +} + +func (e *claudeOAuthCancellationError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func (e *claudeOAuthCancellationError) IsRequestScoped() bool { + return e != nil +} + +func newClaudeOAuthCancellationError(ctx context.Context, oauth bool, err error) error { + if !oauth { + return nil + } + cause := err + if ctx != nil && ctx.Err() != nil { + cause = ctx.Err() + } + if !errors.Is(cause, context.Canceled) { + return nil + } + return &claudeOAuthCancellationError{cause: cause} +} + func shouldSanitizeClaudeMessagesForUpstream(baseModel string) bool { return sigcompat.SignatureProviderFromModelName(baseModel) == sigcompat.SignatureProviderClaude } diff --git a/internal/runtime/executor/claude_executor_stream.go b/internal/runtime/executor/claude_executor_stream.go index 4b060ea54..6f69674ee 100644 --- a/internal/runtime/executor/claude_executor_stream.go +++ b/internal/runtime/executor/claude_executor_stream.go @@ -31,6 +31,11 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL) oauthToken := isClaudeOAuthToken(apiKey) + defer func() { + if cancelErr := newClaudeOAuthCancellationError(ctx, oauthToken, err); cancelErr != nil { + err = cancelErr + } + }() cchSigning := claudeCCHSigningEnabled(apiKey, claudeCCHUpstreamAnthropic, url) reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth) @@ -210,7 +215,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A } return nil, err } - out := make(chan cliproxyexecutor.StreamChunk) + out := make(chan cliproxyexecutor.StreamChunk, 1) go func() { defer close(out) defer func() { @@ -218,6 +223,19 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A log.Errorf("response body close error: %v", errClose) } }() + emitCancellation := func(cause error) bool { + cancelErr := newClaudeOAuthCancellationError(ctx, oauthToken, cause) + if cancelErr == nil { + return false + } + helps.RecordAPIResponseError(ctx, e.cfg, cancelErr) + reporter.PublishFailure(ctx, cancelErr) + select { + case out <- cliproxyexecutor.StreamChunk{Err: cancelErr}: + default: + } + return true + } // If the response target is Claude, directly forward complete SSE events without translation. if responseFormat == to { @@ -251,10 +269,15 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A event.Write(line) event.WriteByte('\n') if len(bytes.TrimSpace(line)) == 0 && !flushEvent() { + emitCancellation(ctx.Err()) return } } if !flushEvent() { + emitCancellation(ctx.Err()) + return + } + if emitCancellation(scanner.Err()) { return } if errScan := scanner.Err(); errScan != nil { @@ -301,10 +324,14 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A select { case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}: case <-ctx.Done(): + emitCancellation(ctx.Err()) return } } } + if emitCancellation(scanner.Err()) { + return + } if errScan := scanner.Err(); errScan != nil { helps.RecordAPIResponseError(ctx, e.cfg, errScan) reporter.PublishFailure(ctx, errScan) diff --git a/internal/runtime/executor/claude_executor_test.go b/internal/runtime/executor/claude_executor_test.go index b20238a04..62763af9d 100644 --- a/internal/runtime/executor/claude_executor_test.go +++ b/internal/runtime/executor/claude_executor_test.go @@ -6,6 +6,7 @@ import ( "context" "encoding/base64" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -1719,6 +1720,124 @@ func TestClaudeExecutor_ExecuteStreamStripsOpenAIEncryptedThinkingBeforeUpstream } } +func claudeOAuthCancellationTestMetadata() map[string]any { + return map[string]any{ + "account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", + claudeauth.ClaudeDeviceIDsMetadataKey: []string{ + "0000000000000000000000000000000000000000000000000000000000000000", + }, + } +} + +func TestClaudeExecutor_ExecuteStreamOAuthStartupCancellationIsRequestScoped(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + close(started) + <-release + })) + defer server.Close() + defer close(release) + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "oauth-stream-startup-cancellation", + Attributes: map[string]string{ + "api_key": "sk-ant-oat-stream-startup-cancellation", + "base_url": server.URL, + }, + Metadata: claudeOAuthCancellationTestMetadata(), + } + ctx, cancel := context.WithCancel(context.Background()) + errCh := make(chan error, 1) + go func() { + _, errStream := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "claude-opus-5", + Payload: []byte(`{"model":"claude-opus-5","messages":[{"role":"user","content":"hello"}],"stream":true}`), + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) + errCh <- errStream + }() + <-started + cancel() + + select { + case errStream := <-errCh: + if !errors.Is(errStream, context.Canceled) { + t.Fatalf("ExecuteStream() error = %v, want context.Canceled", errStream) + } + var requestErr cliproxyexecutor.RequestScopedError + if !errors.As(errStream, &requestErr) || requestErr == nil || !requestErr.IsRequestScoped() { + t.Fatalf("ExecuteStream() error = %T %v, want request-scoped cancellation", errStream, errStream) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for startup cancellation") + } +} + +func TestClaudeExecutor_ExecuteStreamOAuthCancellationIsRequestScoped(t *testing.T) { + started := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("data")) + if flusher, ok := w.(http.Flusher); ok { + flusher.Flush() + } + close(started) + <-r.Context().Done() + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "oauth-stream-cancellation", + Attributes: map[string]string{ + "api_key": "sk-ant-oat-stream-cancellation", + "base_url": server.URL, + }, + Metadata: claudeOAuthCancellationTestMetadata(), + } + payload := []byte(`{"model":"claude-opus-5","system":"system prompt","messages":[{"role":"user","content":"hello"}],"stream":true}`) + ctx, cancel := context.WithCancel(context.Background()) + result, errStream := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "claude-opus-5", + Payload: payload, + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) + if errStream != nil { + cancel() + t.Fatalf("ExecuteStream() error = %v", errStream) + } + <-started + cancel() + + var cancellationErr error + deadline := time.After(2 * time.Second) + for cancellationErr == nil { + select { + case chunk, ok := <-result.Chunks: + if !ok { + t.Fatal("stream closed without a cancellation result") + } + cancellationErr = chunk.Err + case <-deadline: + t.Fatal("timed out waiting for cancellation result") + } + } + if !errors.Is(cancellationErr, context.Canceled) { + t.Fatalf("stream error = %v, want context.Canceled", cancellationErr) + } + var requestErr cliproxyexecutor.RequestScopedError + if !errors.As(cancellationErr, &requestErr) || requestErr == nil || !requestErr.IsRequestScoped() { + t.Fatalf("stream error = %T %v, want request-scoped cancellation", cancellationErr, cancellationErr) + } + var statusErr interface{ StatusCode() int } + if errors.As(cancellationErr, &statusErr) { + t.Fatalf("stream cancellation unexpectedly exposes HTTP status %d", statusErr.StatusCode()) + } + for range result.Chunks { + } +} + func TestClaudeExecutor_ExecuteStreamDirectPassthroughEmitsCompleteSSEEvents(t *testing.T) { firstData := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}` secondData := `{"type":"message_stop"}` diff --git a/sdk/cliproxy/auth/conductor_claude_cancellation_test.go b/sdk/cliproxy/auth/conductor_claude_cancellation_test.go new file mode 100644 index 000000000..1bcf6dd05 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_claude_cancellation_test.go @@ -0,0 +1,317 @@ +package auth + +import ( + "context" + "errors" + "net/http" + "sync/atomic" + "testing" + + "github.com/google/uuid" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type claudeCancellationTestExecutor struct { + prepareFn func(context.Context, *Auth) (*Auth, error) + executeFn func(context.Context, *Auth) (cliproxyexecutor.Response, error) + countFn func(context.Context, *Auth) (cliproxyexecutor.Response, error) + streamFn func(context.Context, *Auth) (*cliproxyexecutor.StreamResult, error) + refreshFn func(context.Context, *Auth) (*Auth, error) + + prepareCalls atomic.Int32 + executeCalls atomic.Int32 + countCalls atomic.Int32 + streamCalls atomic.Int32 + refreshCalls atomic.Int32 +} + +func (*claudeCancellationTestExecutor) Identifier() string { return "claude" } + +func (e *claudeCancellationTestExecutor) ShouldPrepareRequestAuth(*Auth) bool { + return e.prepareFn != nil +} + +func (e *claudeCancellationTestExecutor) PrepareRequestAuth(ctx context.Context, auth *Auth) (*Auth, error) { + e.prepareCalls.Add(1) + return e.prepareFn(ctx, auth) +} + +func (e *claudeCancellationTestExecutor) Execute(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.executeCalls.Add(1) + if e.executeFn != nil { + return e.executeFn(ctx, auth) + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (e *claudeCancellationTestExecutor) CountTokens(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.countCalls.Add(1) + if e.countFn != nil { + return e.countFn(ctx, auth) + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (e *claudeCancellationTestExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.streamCalls.Add(1) + if e.streamFn != nil { + return e.streamFn(ctx, auth) + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("ok")} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *claudeCancellationTestExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { + e.refreshCalls.Add(1) + if e.refreshFn != nil { + return e.refreshFn(ctx, auth) + } + return auth, nil +} + +func (*claudeCancellationTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +type claudeRequestScopedCancellation struct{} + +func (claudeRequestScopedCancellation) Error() string { return context.Canceled.Error() } +func (claudeRequestScopedCancellation) Unwrap() error { return context.Canceled } +func (claudeRequestScopedCancellation) IsRequestScoped() bool { return true } + +func newClaudeCancellationTestManager(t *testing.T, executor *claudeCancellationTestExecutor, hook Hook) (*Manager, *Auth, string) { + t.Helper() + if hook == nil { + hook = NoopHook{} + } + model := "claude-cancel-model-" + uuid.NewString() + auth := &Auth{ + ID: "claude-cancel-auth-" + uuid.NewString(), + Provider: "claude", + Attributes: map[string]string{"auth_kind": "oauth"}, + Metadata: map[string]any{ + "access_token": "access-token", + "refresh_token": "refresh-token", + "request_retry": float64(0), + }, + } + manager := NewManager(nil, nil, hook) + manager.SetRetryConfig(0, 0, 0) + manager.RegisterExecutor(executor) + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + return manager, auth, model +} + +func requireClaudeCancellationNeutral(t *testing.T, manager *Manager, authID, model string) { + t.Helper() + auth, ok := manager.GetByID(authID) + if !ok || auth == nil { + t.Fatalf("GetByID(%q) did not return auth", authID) + } + if auth.Unavailable || !auth.NextRetryAfter.IsZero() { + t.Fatalf("auth was cooled: unavailable=%t next=%v", auth.Unavailable, auth.NextRetryAfter) + } + if state := auth.ModelStates[model]; state != nil && (state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded) { + t.Fatalf("model was cooled: %#v", state) + } +} + +func TestManagerClaudePrepareCancellationStopsWithoutCooldown(t *testing.T) { + tests := []struct { + name string + run func(context.Context, *Manager, string) error + }{ + { + name: "execute", + run: func(ctx context.Context, manager *Manager, model string) error { + _, errExecute := manager.Execute(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count tokens", + run: func(ctx context.Context, manager *Manager, model string) error { + _, errCount := manager.ExecuteCount(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "stream", + run: func(ctx context.Context, manager *Manager, model string) error { + _, errStream := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + return errStream + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + executor := &claudeCancellationTestExecutor{} + executor.prepareFn = func(ctx context.Context, auth *Auth) (*Auth, error) { + cancel() + return auth, ctx.Err() + } + manager, auth, model := newClaudeCancellationTestManager(t, executor, nil) + + errExecute := tt.run(ctx, manager, model) + if !errors.Is(errExecute, context.Canceled) { + t.Fatalf("error = %v, want context.Canceled", errExecute) + } + if got := executor.prepareCalls.Load(); got != 1 { + t.Fatalf("PrepareRequestAuth calls = %d, want 1", got) + } + if executor.executeCalls.Load()+executor.countCalls.Load()+executor.streamCalls.Load() != 0 { + t.Fatal("executor ran after request preparation was canceled") + } + requireClaudeCancellationNeutral(t, manager, auth.ID, model) + }) + } +} + +func TestManagerClaudeRefreshCancellationStopsWithoutCooldown(t *testing.T) { + unauthorized := &Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"} + tests := []struct { + name string + configure func(*claudeCancellationTestExecutor) + run func(context.Context, *Manager, string) error + }{ + { + name: "execute", + configure: func(executor *claudeCancellationTestExecutor) { + executor.executeFn = func(context.Context, *Auth) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, unauthorized + } + }, + run: func(ctx context.Context, manager *Manager, model string) error { + _, errExecute := manager.Execute(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count tokens", + configure: func(executor *claudeCancellationTestExecutor) { + executor.countFn = func(context.Context, *Auth) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, unauthorized + } + }, + run: func(ctx context.Context, manager *Manager, model string) error { + _, errCount := manager.ExecuteCount(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + return errCount + }, + }, + { + name: "stream", + configure: func(executor *claudeCancellationTestExecutor) { + executor.streamFn = func(context.Context, *Auth) (*cliproxyexecutor.StreamResult, error) { + return nil, unauthorized + } + }, + run: func(ctx context.Context, manager *Manager, model string) error { + _, errStream := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + return errStream + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + executor := &claudeCancellationTestExecutor{} + tt.configure(executor) + executor.refreshFn = func(ctx context.Context, _ *Auth) (*Auth, error) { + cancel() + return nil, ctx.Err() + } + manager, auth, model := newClaudeCancellationTestManager(t, executor, nil) + + errExecute := tt.run(ctx, manager, model) + if !errors.Is(errExecute, context.Canceled) { + t.Fatalf("error = %v, want context.Canceled", errExecute) + } + if got := executor.refreshCalls.Load(); got != 1 { + t.Fatalf("Refresh calls = %d, want 1", got) + } + if upstreamCalls := executor.executeCalls.Load() + executor.countCalls.Load() + executor.streamCalls.Load(); upstreamCalls != 1 { + t.Fatalf("upstream calls = %d, want 1", upstreamCalls) + } + requireClaudeCancellationNeutral(t, manager, auth.ID, model) + }) + } +} + +func TestManagerClaudeStreamTailCancellationIsAvailabilityNeutral(t *testing.T) { + source := make(chan cliproxyexecutor.StreamChunk, 1) + source <- cliproxyexecutor.StreamChunk{Payload: []byte("first")} + executor := &claudeCancellationTestExecutor{ + streamFn: func(context.Context, *Auth) (*cliproxyexecutor.StreamResult, error) { + return &cliproxyexecutor.StreamResult{Chunks: source}, nil + }, + } + hook := &resultCaptureHook{} + manager, auth, model := newClaudeCancellationTestManager(t, executor, hook) + ctx, cancel := context.WithCancel(context.Background()) + + stream, errStream := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + if errStream != nil { + t.Fatalf("ExecuteStream() error = %v", errStream) + } + if chunk := <-stream.Chunks; chunk.Err != nil || string(chunk.Payload) != "first" { + t.Fatalf("first chunk = %#v", chunk) + } + cancel() + source <- cliproxyexecutor.StreamChunk{Err: claudeRequestScopedCancellation{}} + close(source) + for range stream.Chunks { + } + + results := hook.Results() + if len(results) != 1 || results[0].Success || results[0].Error == nil { + t.Fatalf("results = %#v, want one failed cancellation result", results) + } + if results[0].Error.Code != requestScopedErrorCode || results[0].Error.StatusCode() != 0 { + t.Fatalf("cancellation result = %#v, want request-scoped status 0", results[0].Error) + } + requireClaudeCancellationNeutral(t, manager, auth.ID, model) +} + +func TestManagerClaudeUpstreamFailureStillCoolsCredential(t *testing.T) { + executor := &claudeCancellationTestExecutor{ + executeFn: func(context.Context, *Auth) (cliproxyexecutor.Response, error) { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusInternalServerError, Message: "upstream failure"} + }, + } + manager, auth, model := newClaudeCancellationTestManager(t, executor, nil) + + _, errExecute := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if statusCodeFromError(errExecute) != http.StatusInternalServerError { + t.Fatalf("Execute() error = %v, want HTTP 500", errExecute) + } + got, ok := manager.GetByID(auth.ID) + if !ok || got == nil { + t.Fatalf("GetByID(%q) did not return auth", auth.ID) + } + state := got.ModelStates[model] + if state == nil || !state.Unavailable || state.NextRetryAfter.IsZero() { + t.Fatalf("upstream failure did not cool model: %#v", state) + } +} + +func TestClaudeRequestCancellationDoesNotChangeOtherProviders(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + tests := []*Auth{ + {Provider: "codex", Attributes: map[string]string{"auth_kind": "oauth"}}, + {Provider: "claude", Attributes: map[string]string{"auth_kind": "api_key"}}, + } + for _, auth := range tests { + if errCancel := claudeOAuthRequestCancellation(ctx, auth, context.Canceled); errCancel != nil { + t.Fatalf("auth %#v was classified as Claude OAuth cancellation: %v", auth.Attributes, errCancel) + } + } +} diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index a9ca51654..442e5d2f8 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -20,6 +20,19 @@ import ( log "github.com/sirupsen/logrus" ) +func claudeOAuthRequestCancellation(ctx context.Context, auth *Auth, err error) error { + if auth == nil || !strings.EqualFold(strings.TrimSpace(auth.Provider), "claude") || !strings.EqualFold(strings.TrimSpace(auth.Attributes["auth_kind"]), "oauth") { + return nil + } + if ctx != nil && errors.Is(ctx.Err(), context.Canceled) { + return ctx.Err() + } + if errors.Is(err, context.Canceled) { + return err + } + return nil +} + // Execute performs a non-streaming execution using the configured selector and executor. // It supports multiple providers for the same model and round-robins the starting provider per model. func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { @@ -310,6 +323,9 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req var errPrepare error auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) if errPrepare != nil { + if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil { + return cliproxyexecutor.Response{}, errCancel + } result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} m.MarkResult(execCtx, result) lastErr = errPrepare @@ -349,6 +365,9 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } } } + if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { + return cliproxyexecutor.Response{}, errCancel + } result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} if errExec != nil { result.Error = resultErrorFromError(errExec) @@ -431,6 +450,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, var errPrepare error auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) if errPrepare != nil { + if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil { + return cliproxyexecutor.Response{}, errCancel + } result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} m.MarkResult(execCtx, result) lastErr = errPrepare @@ -470,6 +492,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } } } + if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { + return cliproxyexecutor.Response{}, errCancel + } result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil} if errExec != nil { result.Error = resultErrorFromError(errExec) @@ -608,6 +633,11 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) } if errPrepare != nil { + if selection == nil { + if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil { + return nil, errCancel + } + } result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)} if selection != nil { m.reportHomeResult(execCtx, result, auth) diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index be6784af3..f32099638 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -173,7 +173,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re return } } - if !failed { + if !failed && (ephemeralResult || claudeOAuthRequestCancellation(ctx, auth, nil) == nil) { m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}, auth, ephemeralResult) } }() @@ -224,6 +224,11 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } } } + if !ephemeralResult { + if errCancel := claudeOAuthRequestCancellation(ctx, auth, errStream); errCancel != nil { + return nil, errCancel + } + } if errStream == nil && (streamResult == nil || streamResult.Chunks == nil) { errStream = &Error{Code: "empty_stream", Message: "upstream stream has no source", Retryable: true} } @@ -264,6 +269,12 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } } } + if !ephemeralResult { + if errCancel := claudeOAuthRequestCancellation(ctx, auth, bootstrapErr); errCancel != nil { + discardStreamChunks(streamResult.Chunks) + return nil, errCancel + } + } if bootstrapErr != nil { if isRequestInvalidError(bootstrapErr) { rerr := resultErrorFromError(bootstrapErr)