History

fix(claude): keep client cancellation availability-neutral

sususu committed 3 months ago

Commit ce7fcd91 parent 22288476 files · +543 −2Browse files

Unified · Split

6 files
  1. internal/runtime/executor/claude_executor.go +37 −0
  2. internal/runtime/executor/claude_executor_stream.go +28 −1
  3. internal/runtime/executor/claude_executor_test.go +119 −0
  4. sdk/cliproxy/auth/conductor_claude_cancellation_test.go +317 −0
  5. sdk/cliproxy/auth/conductor_execution.go +30 −0
  6. sdk/cliproxy/auth/conductor_stream.go +12 −1
internal/runtime/executor/claude_executor.go+37−0
@@ -3,6 +3,7 @@ package executor · 2 unmodified lines
3import (3import (
4 "bytes"4 "bytes"
5 "context"5 "context"
6 "errors"
6 "fmt"7 "fmt"
7 "net/http"8 "net/http"
8 "strings"9 "strings"
@@ -26,6 +27,42 @@ type ClaudeExecutor struct { · 17 unmodified lines
26 oauthProfileFetcher claudeOAuthProfileFetcher27 oauthProfileFetcher claudeOAuthProfileFetcher
27}28}
28 29
30type claudeOAuthCancellationError struct {
31 cause error
32}
33
34func (e *claudeOAuthCancellationError) Error() string {
35 if e == nil || e.cause == nil {
36 return ""
37 }
38 return e.cause.Error()
39}
40
41func (e *claudeOAuthCancellationError) Unwrap() error {
42 if e == nil {
43 return nil
44 }
45 return e.cause
46}
47
48func (e *claudeOAuthCancellationError) IsRequestScoped() bool {
49 return e != nil
50}
51
52func newClaudeOAuthCancellationError(ctx context.Context, oauth bool, err error) error {
53 if !oauth {
54 return nil
55 }
56 cause := err
57 if ctx != nil && ctx.Err() != nil {
58 cause = ctx.Err()
59 }
60 if !errors.Is(cause, context.Canceled) {
61 return nil
62 }
63 return &claudeOAuthCancellationError{cause: cause}
64}
65
29func shouldSanitizeClaudeMessagesForUpstream(baseModel string) bool {66func shouldSanitizeClaudeMessagesForUpstream(baseModel string) bool {
30 return sigcompat.SignatureProviderFromModelName(baseModel) == sigcompat.SignatureProviderClaude67 return sigcompat.SignatureProviderFromModelName(baseModel) == sigcompat.SignatureProviderClaude
31}68}
internal/runtime/executor/claude_executor_stream.go+28−1
@@ -31,6 +31,11 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 30 unmodified lines
31 }31 }
32 url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL)32 url := fmt.Sprintf("%s/v1/messages?beta=true", baseURL)
33 oauthToken := isClaudeOAuthToken(apiKey)33 oauthToken := isClaudeOAuthToken(apiKey)
34 defer func() {
35 if cancelErr := newClaudeOAuthCancellationError(ctx, oauthToken, err); cancelErr != nil {
36 err = cancelErr
37 }
38 }()
34 cchSigning := claudeCCHSigningEnabled(apiKey, claudeCCHUpstreamAnthropic, url)39 cchSigning := claudeCCHSigningEnabled(apiKey, claudeCCHUpstreamAnthropic, url)
35 40
36 reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)41 reporter := helps.NewExecutorUsageReporter(ctx, e, baseModel, auth)
@@ -210,7 +215,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 173 unmodified lines
210 }215 }
211 return nil, err216 return nil, err
212 }217 }
213 out := make(chan cliproxyexecutor.StreamChunk)218 out := make(chan cliproxyexecutor.StreamChunk, 1)
214 go func() {219 go func() {
215 defer close(out)220 defer close(out)
216 defer func() {221 defer func() {
@@ -218,6 +223,19 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 1 unmodified line
218 log.Errorf("response body close error: %v", errClose)223 log.Errorf("response body close error: %v", errClose)
219 }224 }
220 }()225 }()
226 emitCancellation := func(cause error) bool {
227 cancelErr := newClaudeOAuthCancellationError(ctx, oauthToken, cause)
228 if cancelErr == nil {
229 return false
230 }
231 helps.RecordAPIResponseError(ctx, e.cfg, cancelErr)
232 reporter.PublishFailure(ctx, cancelErr)
233 select {
234 case out <- cliproxyexecutor.StreamChunk{Err: cancelErr}:
235 default:
236 }
237 return true
238 }
221 239
222 // If the response target is Claude, directly forward complete SSE events without translation.240 // If the response target is Claude, directly forward complete SSE events without translation.
223 if responseFormat == to {241 if responseFormat == to {
@@ -251,10 +269,15 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 27 unmodified lines
251 event.Write(line)269 event.Write(line)
252 event.WriteByte('\n')270 event.WriteByte('\n')
253 if len(bytes.TrimSpace(line)) == 0 && !flushEvent() {271 if len(bytes.TrimSpace(line)) == 0 && !flushEvent() {
272 emitCancellation(ctx.Err())
254 return273 return
255 }274 }
256 }275 }
257 if !flushEvent() {276 if !flushEvent() {
277 emitCancellation(ctx.Err())
278 return
279 }
280 if emitCancellation(scanner.Err()) {
258 return281 return
259 }282 }
260 if errScan := scanner.Err(); errScan != nil {283 if errScan := scanner.Err(); errScan != nil {
@@ -301,10 +324,14 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A · 40 unmodified lines
301 select {324 select {
302 case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}:325 case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}:
303 case <-ctx.Done():326 case <-ctx.Done():
327 emitCancellation(ctx.Err())
304 return328 return
305 }329 }
306 }330 }
307 }331 }
332 if emitCancellation(scanner.Err()) {
333 return
334 }
308 if errScan := scanner.Err(); errScan != nil {335 if errScan := scanner.Err(); errScan != nil {
309 helps.RecordAPIResponseError(ctx, e.cfg, errScan)336 helps.RecordAPIResponseError(ctx, e.cfg, errScan)
310 reporter.PublishFailure(ctx, errScan)337 reporter.PublishFailure(ctx, errScan)
internal/runtime/executor/claude_executor_test.go+119−0
@@ -6,6 +6,7 @@ import ( · 5 unmodified lines
6 "context"6 "context"
7 "encoding/base64"7 "encoding/base64"
8 "encoding/json"8 "encoding/json"
9 "errors"
9 "fmt"10 "fmt"
10 "io"11 "io"
11 "net/http"12 "net/http"
@@ -1719,6 +1720,124 @@ func TestClaudeExecutor_ExecuteStreamStripsOpenAIEncryptedThinkingBeforeUpstream · 1707 unmodified lines
1719 }1720 }
1720}1721}
1721 1722
1723func claudeOAuthCancellationTestMetadata() map[string]any {
1724 return map[string]any{
1725 "account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa",
1726 claudeauth.ClaudeDeviceIDsMetadataKey: []string{
1727 "0000000000000000000000000000000000000000000000000000000000000000",
1728 },
1729 }
1730}
1731
1732func TestClaudeExecutor_ExecuteStreamOAuthStartupCancellationIsRequestScoped(t *testing.T) {
1733 started := make(chan struct{})
1734 release := make(chan struct{})
1735 server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
1736 close(started)
1737 <-release
1738 }))
1739 defer server.Close()
1740 defer close(release)
1741
1742 executor := NewClaudeExecutor(&config.Config{})
1743 auth := &cliproxyauth.Auth{
1744 ID: "oauth-stream-startup-cancellation",
1745 Attributes: map[string]string{
1746 "api_key": "sk-ant-oat-stream-startup-cancellation",
1747 "base_url": server.URL,
1748 },
1749 Metadata: claudeOAuthCancellationTestMetadata(),
1750 }
1751 ctx, cancel := context.WithCancel(context.Background())
1752 errCh := make(chan error, 1)
1753 go func() {
1754 _, errStream := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
1755 Model: "claude-opus-5",
1756 Payload: []byte(`{"model":"claude-opus-5","messages":[{"role":"user","content":"hello"}],"stream":true}`),
1757 }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})
1758 errCh <- errStream
1759 }()
1760 <-started
1761 cancel()
1762
1763 select {
1764 case errStream := <-errCh:
1765 if !errors.Is(errStream, context.Canceled) {
1766 t.Fatalf("ExecuteStream() error = %v, want context.Canceled", errStream)
1767 }
1768 var requestErr cliproxyexecutor.RequestScopedError
1769 if !errors.As(errStream, &requestErr) || requestErr == nil || !requestErr.IsRequestScoped() {
1770 t.Fatalf("ExecuteStream() error = %T %v, want request-scoped cancellation", errStream, errStream)
1771 }
1772 case <-time.After(2 * time.Second):
1773 t.Fatal("timed out waiting for startup cancellation")
1774 }
1775}
1776
1777func TestClaudeExecutor_ExecuteStreamOAuthCancellationIsRequestScoped(t *testing.T) {
1778 started := make(chan struct{})
1779 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
1780 w.Header().Set("Content-Type", "text/event-stream")
1781 w.WriteHeader(http.StatusOK)
1782 _, _ = w.Write([]byte("data"))
1783 if flusher, ok := w.(http.Flusher); ok {
1784 flusher.Flush()
1785 }
1786 close(started)
1787 <-r.Context().Done()
1788 }))
1789 defer server.Close()
1790
1791 executor := NewClaudeExecutor(&config.Config{})
1792 auth := &cliproxyauth.Auth{
1793 ID: "oauth-stream-cancellation",
1794 Attributes: map[string]string{
1795 "api_key": "sk-ant-oat-stream-cancellation",
1796 "base_url": server.URL,
1797 },
1798 Metadata: claudeOAuthCancellationTestMetadata(),
1799 }
1800 payload := []byte(`{"model":"claude-opus-5","system":"system prompt","messages":[{"role":"user","content":"hello"}],"stream":true}`)
1801 ctx, cancel := context.WithCancel(context.Background())
1802 result, errStream := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
1803 Model: "claude-opus-5",
1804 Payload: payload,
1805 }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude})
1806 if errStream != nil {
1807 cancel()
1808 t.Fatalf("ExecuteStream() error = %v", errStream)
1809 }
1810 <-started
1811 cancel()
1812
1813 var cancellationErr error
1814 deadline := time.After(2 * time.Second)
1815 for cancellationErr == nil {
1816 select {
1817 case chunk, ok := <-result.Chunks:
1818 if !ok {
1819 t.Fatal("stream closed without a cancellation result")
1820 }
1821 cancellationErr = chunk.Err
1822 case <-deadline:
1823 t.Fatal("timed out waiting for cancellation result")
1824 }
1825 }
1826 if !errors.Is(cancellationErr, context.Canceled) {
1827 t.Fatalf("stream error = %v, want context.Canceled", cancellationErr)
1828 }
1829 var requestErr cliproxyexecutor.RequestScopedError
1830 if !errors.As(cancellationErr, &requestErr) || requestErr == nil || !requestErr.IsRequestScoped() {
1831 t.Fatalf("stream error = %T %v, want request-scoped cancellation", cancellationErr, cancellationErr)
1832 }
1833 var statusErr interface{ StatusCode() int }
1834 if errors.As(cancellationErr, &statusErr) {
1835 t.Fatalf("stream cancellation unexpectedly exposes HTTP status %d", statusErr.StatusCode())
1836 }
1837 for range result.Chunks {
1838 }
1839}
1840
1722func TestClaudeExecutor_ExecuteStreamDirectPassthroughEmitsCompleteSSEEvents(t *testing.T) {1841func TestClaudeExecutor_ExecuteStreamDirectPassthroughEmitsCompleteSSEEvents(t *testing.T) {
1723 firstData := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`1842 firstData := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`
1724 secondData := `{"type":"message_stop"}`1843 secondData := `{"type":"message_stop"}`
sdk/cliproxy/auth/conductor_claude_cancellation_test.go+317−0
@@ -0,0 +1,317 @@
1package auth
2
3import (
4 "context"
5 "errors"
6 "net/http"
7 "sync/atomic"
8 "testing"
9
10 "github.com/google/uuid"
11 "github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
12 cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
13)
14
15type claudeCancellationTestExecutor struct {
16 prepareFn func(context.Context, *Auth) (*Auth, error)
17 executeFn func(context.Context, *Auth) (cliproxyexecutor.Response, error)
18 countFn func(context.Context, *Auth) (cliproxyexecutor.Response, error)
19 streamFn func(context.Context, *Auth) (*cliproxyexecutor.StreamResult, error)
20 refreshFn func(context.Context, *Auth) (*Auth, error)
21
22 prepareCalls atomic.Int32
23 executeCalls atomic.Int32
24 countCalls atomic.Int32
25 streamCalls atomic.Int32
26 refreshCalls atomic.Int32
27}
28
29func (*claudeCancellationTestExecutor) Identifier() string { return "claude" }
30
31func (e *claudeCancellationTestExecutor) ShouldPrepareRequestAuth(*Auth) bool {
32 return e.prepareFn != nil
33}
34
35func (e *claudeCancellationTestExecutor) PrepareRequestAuth(ctx context.Context, auth *Auth) (*Auth, error) {
36 e.prepareCalls.Add(1)
37 return e.prepareFn(ctx, auth)
38}
39
40func (e *claudeCancellationTestExecutor) Execute(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
41 e.executeCalls.Add(1)
42 if e.executeFn != nil {
43 return e.executeFn(ctx, auth)
44 }
45 return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
46}
47
48func (e *claudeCancellationTestExecutor) CountTokens(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
49 e.countCalls.Add(1)
50 if e.countFn != nil {
51 return e.countFn(ctx, auth)
52 }
53 return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
54}
55
56func (e *claudeCancellationTestExecutor) ExecuteStream(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
57 e.streamCalls.Add(1)
58 if e.streamFn != nil {
59 return e.streamFn(ctx, auth)
60 }
61 chunks := make(chan cliproxyexecutor.StreamChunk, 1)
62 chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("ok")}
63 close(chunks)
64 return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil
65}
66
67func (e *claudeCancellationTestExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) {
68 e.refreshCalls.Add(1)
69 if e.refreshFn != nil {
70 return e.refreshFn(ctx, auth)
71 }
72 return auth, nil
73}
74
75func (*claudeCancellationTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
76 return nil, errors.New("not implemented")
77}
78
79type claudeRequestScopedCancellation struct{}
80
81func (claudeRequestScopedCancellation) Error() string { return context.Canceled.Error() }
82func (claudeRequestScopedCancellation) Unwrap() error { return context.Canceled }
83func (claudeRequestScopedCancellation) IsRequestScoped() bool { return true }
84
85func newClaudeCancellationTestManager(t *testing.T, executor *claudeCancellationTestExecutor, hook Hook) (*Manager, *Auth, string) {
86 t.Helper()
87 if hook == nil {
88 hook = NoopHook{}
89 }
90 model := "claude-cancel-model-" + uuid.NewString()
91 auth := &Auth{
92 ID: "claude-cancel-auth-" + uuid.NewString(),
93 Provider: "claude",
94 Attributes: map[string]string{"auth_kind": "oauth"},
95 Metadata: map[string]any{
96 "access_token": "access-token",
97 "refresh_token": "refresh-token",
98 "request_retry": float64(0),
99 },
100 }
101 manager := NewManager(nil, nil, hook)
102 manager.SetRetryConfig(0, 0, 0)
103 manager.RegisterExecutor(executor)
104 registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}})
105 t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) })
106 if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
107 t.Fatalf("Register() error = %v", errRegister)
108 }
109 return manager, auth, model
110}
111
112func requireClaudeCancellationNeutral(t *testing.T, manager *Manager, authID, model string) {
113 t.Helper()
114 auth, ok := manager.GetByID(authID)
115 if !ok || auth == nil {
116 t.Fatalf("GetByID(%q) did not return auth", authID)
117 }
118 if auth.Unavailable || !auth.NextRetryAfter.IsZero() {
119 t.Fatalf("auth was cooled: unavailable=%t next=%v", auth.Unavailable, auth.NextRetryAfter)
120 }
121 if state := auth.ModelStates[model]; state != nil && (state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded) {
122 t.Fatalf("model was cooled: %#v", state)
123 }
124}
125
126func TestManagerClaudePrepareCancellationStopsWithoutCooldown(t *testing.T) {
127 tests := []struct {
128 name string
129 run func(context.Context, *Manager, string) error
130 }{
131 {
132 name: "execute",
133 run: func(ctx context.Context, manager *Manager, model string) error {
134 _, errExecute := manager.Execute(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{})
135 return errExecute
136 },
137 },
138 {
139 name: "count tokens",
140 run: func(ctx context.Context, manager *Manager, model string) error {
141 _, errCount := manager.ExecuteCount(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{})
142 return errCount
143 },
144 },
145 {
146 name: "stream",
147 run: func(ctx context.Context, manager *Manager, model string) error {
148 _, errStream := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true})
149 return errStream
150 },
151 },
152 }
153 for _, tt := range tests {
154 t.Run(tt.name, func(t *testing.T) {
155 ctx, cancel := context.WithCancel(context.Background())
156 executor := &claudeCancellationTestExecutor{}
157 executor.prepareFn = func(ctx context.Context, auth *Auth) (*Auth, error) {
158 cancel()
159 return auth, ctx.Err()
160 }
161 manager, auth, model := newClaudeCancellationTestManager(t, executor, nil)
162
163 errExecute := tt.run(ctx, manager, model)
164 if !errors.Is(errExecute, context.Canceled) {
165 t.Fatalf("error = %v, want context.Canceled", errExecute)
166 }
167 if got := executor.prepareCalls.Load(); got != 1 {
168 t.Fatalf("PrepareRequestAuth calls = %d, want 1", got)
169 }
170 if executor.executeCalls.Load()+executor.countCalls.Load()+executor.streamCalls.Load() != 0 {
171 t.Fatal("executor ran after request preparation was canceled")
172 }
173 requireClaudeCancellationNeutral(t, manager, auth.ID, model)
174 })
175 }
176}
177
178func TestManagerClaudeRefreshCancellationStopsWithoutCooldown(t *testing.T) {
179 unauthorized := &Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}
180 tests := []struct {
181 name string
182 configure func(*claudeCancellationTestExecutor)
183 run func(context.Context, *Manager, string) error
184 }{
185 {
186 name: "execute",
187 configure: func(executor *claudeCancellationTestExecutor) {
188 executor.executeFn = func(context.Context, *Auth) (cliproxyexecutor.Response, error) {
189 return cliproxyexecutor.Response{}, unauthorized
190 }
191 },
192 run: func(ctx context.Context, manager *Manager, model string) error {
193 _, errExecute := manager.Execute(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{})
194 return errExecute
195 },
196 },
197 {
198 name: "count tokens",
199 configure: func(executor *claudeCancellationTestExecutor) {
200 executor.countFn = func(context.Context, *Auth) (cliproxyexecutor.Response, error) {
201 return cliproxyexecutor.Response{}, unauthorized
202 }
203 },
204 run: func(ctx context.Context, manager *Manager, model string) error {
205 _, errCount := manager.ExecuteCount(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{})
206 return errCount
207 },
208 },
209 {
210 name: "stream",
211 configure: func(executor *claudeCancellationTestExecutor) {
212 executor.streamFn = func(context.Context, *Auth) (*cliproxyexecutor.StreamResult, error) {
213 return nil, unauthorized
214 }
215 },
216 run: func(ctx context.Context, manager *Manager, model string) error {
217 _, errStream := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true})
218 return errStream
219 },
220 },
221 }
222 for _, tt := range tests {
223 t.Run(tt.name, func(t *testing.T) {
224 ctx, cancel := context.WithCancel(context.Background())
225 executor := &claudeCancellationTestExecutor{}
226 tt.configure(executor)
227 executor.refreshFn = func(ctx context.Context, _ *Auth) (*Auth, error) {
228 cancel()
229 return nil, ctx.Err()
230 }
231 manager, auth, model := newClaudeCancellationTestManager(t, executor, nil)
232
233 errExecute := tt.run(ctx, manager, model)
234 if !errors.Is(errExecute, context.Canceled) {
235 t.Fatalf("error = %v, want context.Canceled", errExecute)
236 }
237 if got := executor.refreshCalls.Load(); got != 1 {
238 t.Fatalf("Refresh calls = %d, want 1", got)
239 }
240 if upstreamCalls := executor.executeCalls.Load() + executor.countCalls.Load() + executor.streamCalls.Load(); upstreamCalls != 1 {
241 t.Fatalf("upstream calls = %d, want 1", upstreamCalls)
242 }
243 requireClaudeCancellationNeutral(t, manager, auth.ID, model)
244 })
245 }
246}
247
248func TestManagerClaudeStreamTailCancellationIsAvailabilityNeutral(t *testing.T) {
249 source := make(chan cliproxyexecutor.StreamChunk, 1)
250 source <- cliproxyexecutor.StreamChunk{Payload: []byte("first")}
251 executor := &claudeCancellationTestExecutor{
252 streamFn: func(context.Context, *Auth) (*cliproxyexecutor.StreamResult, error) {
253 return &cliproxyexecutor.StreamResult{Chunks: source}, nil
254 },
255 }
256 hook := &resultCaptureHook{}
257 manager, auth, model := newClaudeCancellationTestManager(t, executor, hook)
258 ctx, cancel := context.WithCancel(context.Background())
259
260 stream, errStream := manager.ExecuteStream(ctx, []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true})
261 if errStream != nil {
262 t.Fatalf("ExecuteStream() error = %v", errStream)
263 }
264 if chunk := <-stream.Chunks; chunk.Err != nil || string(chunk.Payload) != "first" {
265 t.Fatalf("first chunk = %#v", chunk)
266 }
267 cancel()
268 source <- cliproxyexecutor.StreamChunk{Err: claudeRequestScopedCancellation{}}
269 close(source)
270 for range stream.Chunks {
271 }
272
273 results := hook.Results()
274 if len(results) != 1 || results[0].Success || results[0].Error == nil {
275 t.Fatalf("results = %#v, want one failed cancellation result", results)
276 }
277 if results[0].Error.Code != requestScopedErrorCode || results[0].Error.StatusCode() != 0 {
278 t.Fatalf("cancellation result = %#v, want request-scoped status 0", results[0].Error)
279 }
280 requireClaudeCancellationNeutral(t, manager, auth.ID, model)
281}
282
283func TestManagerClaudeUpstreamFailureStillCoolsCredential(t *testing.T) {
284 executor := &claudeCancellationTestExecutor{
285 executeFn: func(context.Context, *Auth) (cliproxyexecutor.Response, error) {
286 return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusInternalServerError, Message: "upstream failure"}
287 },
288 }
289 manager, auth, model := newClaudeCancellationTestManager(t, executor, nil)
290
291 _, errExecute := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{})
292 if statusCodeFromError(errExecute) != http.StatusInternalServerError {
293 t.Fatalf("Execute() error = %v, want HTTP 500", errExecute)
294 }
295 got, ok := manager.GetByID(auth.ID)
296 if !ok || got == nil {
297 t.Fatalf("GetByID(%q) did not return auth", auth.ID)
298 }
299 state := got.ModelStates[model]
300 if state == nil || !state.Unavailable || state.NextRetryAfter.IsZero() {
301 t.Fatalf("upstream failure did not cool model: %#v", state)
302 }
303}
304
305func TestClaudeRequestCancellationDoesNotChangeOtherProviders(t *testing.T) {
306 ctx, cancel := context.WithCancel(context.Background())
307 cancel()
308 tests := []*Auth{
309 {Provider: "codex", Attributes: map[string]string{"auth_kind": "oauth"}},
310 {Provider: "claude", Attributes: map[string]string{"auth_kind": "api_key"}},
311 }
312 for _, auth := range tests {
313 if errCancel := claudeOAuthRequestCancellation(ctx, auth, context.Canceled); errCancel != nil {
314 t.Fatalf("auth %#v was classified as Claude OAuth cancellation: %v", auth.Attributes, errCancel)
315 }
316 }
317}
sdk/cliproxy/auth/conductor_execution.go+30−0
@@ -20,6 +20,19 @@ import ( · 19 unmodified lines
20 log "github.com/sirupsen/logrus"20 log "github.com/sirupsen/logrus"
21)21)
22 22
23func claudeOAuthRequestCancellation(ctx context.Context, auth *Auth, err error) error {
24 if auth == nil || !strings.EqualFold(strings.TrimSpace(auth.Provider), "claude") || !strings.EqualFold(strings.TrimSpace(auth.Attributes["auth_kind"]), "oauth") {
25 return nil
26 }
27 if ctx != nil && errors.Is(ctx.Err(), context.Canceled) {
28 return ctx.Err()
29 }
30 if errors.Is(err, context.Canceled) {
31 return err
32 }
33 return nil
34}
35
23// Execute performs a non-streaming execution using the configured selector and executor.36// Execute performs a non-streaming execution using the configured selector and executor.
24// It supports multiple providers for the same model and round-robins the starting provider per model.37// It supports multiple providers for the same model and round-robins the starting provider per model.
25func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {38func (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 · 284 unmodified lines
310 var errPrepare error323 var errPrepare error
311 auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)324 auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)
312 if errPrepare != nil {325 if errPrepare != nil {
326 if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil {
327 return cliproxyexecutor.Response{}, errCancel
328 }
313 result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}329 result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}
314 m.MarkResult(execCtx, result)330 m.MarkResult(execCtx, result)
315 lastErr = errPrepare331 lastErr = errPrepare
@@ -349,6 +365,9 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req · 33 unmodified lines
349 }365 }
350 }366 }
351 }367 }
368 if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil {
369 return cliproxyexecutor.Response{}, errCancel
370 }
352 result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil}371 result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil}
353 if errExec != nil {372 if errExec != nil {
354 result.Error = resultErrorFromError(errExec)373 result.Error = resultErrorFromError(errExec)
@@ -431,6 +450,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, · 76 unmodified lines
431 var errPrepare error450 var errPrepare error
432 auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)451 auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)
433 if errPrepare != nil {452 if errPrepare != nil {
453 if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil {
454 return cliproxyexecutor.Response{}, errCancel
455 }
434 result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}456 result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}
435 m.MarkResult(execCtx, result)457 m.MarkResult(execCtx, result)
436 lastErr = errPrepare458 lastErr = errPrepare
@@ -470,6 +492,9 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, · 33 unmodified lines
470 }492 }
471 }493 }
472 }494 }
495 if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil {
496 return cliproxyexecutor.Response{}, errCancel
497 }
473 result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil}498 result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: errExec == nil}
474 if errExec != nil {499 if errExec != nil {
475 result.Error = resultErrorFromError(errExec)500 result.Error = resultErrorFromError(errExec)
@@ -608,6 +633,11 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string · 132 unmodified lines
608 auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)633 auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth)
609 }634 }
610 if errPrepare != nil {635 if errPrepare != nil {
636 if selection == nil {
637 if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil {
638 return nil, errCancel
639 }
640 }
611 result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}641 result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: false, Error: resultErrorFromError(errPrepare)}
612 if selection != nil {642 if selection != nil {
613 m.reportHomeResult(execCtx, result, auth)643 m.reportHomeResult(execCtx, result, auth)
sdk/cliproxy/auth/conductor_stream.go+12−1
@@ -173,7 +173,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re · 172 unmodified lines
173 return173 return
174 }174 }
175 }175 }
176 if !failed {176 if !failed && (ephemeralResult || claudeOAuthRequestCancellation(ctx, auth, nil) == nil) {
177 m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}, auth, ephemeralResult)177 m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, Success: true}, auth, ephemeralResult)
178 }178 }
179 }()179 }()
@@ -224,6 +224,11 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi · 44 unmodified lines
224 }224 }
225 }225 }
226 }226 }
227 if !ephemeralResult {
228 if errCancel := claudeOAuthRequestCancellation(ctx, auth, errStream); errCancel != nil {
229 return nil, errCancel
230 }
231 }
227 if errStream == nil && (streamResult == nil || streamResult.Chunks == nil) {232 if errStream == nil && (streamResult == nil || streamResult.Chunks == nil) {
228 errStream = &Error{Code: "empty_stream", Message: "upstream stream has no source", Retryable: true}233 errStream = &Error{Code: "empty_stream", Message: "upstream stream has no source", Retryable: true}
229 }234 }
@@ -264,6 +269,12 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi · 34 unmodified lines
264 }269 }
265 }270 }
266 }271 }
272 if !ephemeralResult {
273 if errCancel := claudeOAuthRequestCancellation(ctx, auth, bootstrapErr); errCancel != nil {
274 discardStreamChunks(streamResult.Chunks)
275 return nil, errCancel
276 }
277 }
267 if bootstrapErr != nil {278 if bootstrapErr != nil {
268 if isRequestInvalidError(bootstrapErr) {279 if isRequestInvalidError(bootstrapErr) {
269 rerr := resultErrorFromError(bootstrapErr)280 rerr := resultErrorFromError(bootstrapErr)