History

Handle refresh 401 in auth maintenance

benjamin committed 7 months ago

Commit f438e591 parent d34cdd04 files · +276 −3Browse files

Unified · Split

4 files
  1. internal/auth/codex/openai_auth.go +62 −2
  2. internal/auth/codex/openai_auth_test.go +42 −0
  3. sdk/cliproxy/auth/conductor.go +21 −1
  4. sdk/cliproxy/auth/conductor_refresh_failure_test.go +151 −0
internal/auth/codex/openai_auth.go+62−2
@@ -7,6 +7,7 @@ package codex · 6 unmodified lines
77 import (
88 "context"
99 "encoding/json"
10+ "errors"
1011 "fmt"
1112 "io"
1213 "net/http"
@@ -34,6 +35,33 @@ type CodexAuth struct { · 21 unmodified lines
3435 httpClient *http.Client
3536 }
3637
38+type refreshFailureError struct {
39+ statusCode int
40+ terminal bool
41+ message string
42+}
43+
44+func (e *refreshFailureError) Error() string {
45+ if e == nil {
46+ return ""
47+ }
48+ return e.message
49+}
50+
51+func (e *refreshFailureError) StatusCode() int {
52+ if e == nil {
53+ return 0
54+ }
55+ return e.statusCode
56+}
57+
58+func (e *refreshFailureError) Terminal() bool {
59+ if e == nil {
60+ return false
61+ }
62+ return e.terminal
63+}
64+
3765 // NewCodexAuth creates a new CodexAuth service instance.
3866 // It initializes an HTTP client with proxy settings from the provided configuration.
3967 func NewCodexAuth(cfg *config.Config) *CodexAuth {
@@ -202,7 +230,7 @@ func (o *CodexAuth) RefreshTokens(ctx context.Context, refreshToken string) (*Co · 162 unmodified lines
202230 }
203231
204232 if resp.StatusCode != http.StatusOK {
205− return nil, fmt.Errorf("token refresh failed with status %d: %s", resp.StatusCode, string(body))
233+ return nil, newRefreshFailureError(resp.StatusCode, string(body))
206234 }
207235
208236 var tokenResp struct {
@@ -292,8 +320,40 @@ func isNonRetryableRefreshErr(err error) bool { · 83 unmodified lines
292320 if err == nil {
293321 return false
294322 }
323+ var terminal interface {
324+ Terminal() bool
325+ }
326+ if errors.As(err, &terminal) && terminal != nil && terminal.Terminal() {
327+ return true
328+ }
295329 raw := strings.ToLower(err.Error())
296− return strings.Contains(raw, "refresh_token_reused")
330+ return strings.Contains(raw, "refresh_token_reused") ||
331+ strings.Contains(raw, "token_invalidated") ||
332+ strings.Contains(raw, "token_revoked")
333+}
334+
335+func newRefreshFailureError(statusCode int, body string) error {
336+ message := fmt.Sprintf("token refresh failed with status %d: %s", statusCode, body)
337+ classifiedStatus := statusCode
338+ terminal := false
339+ raw := strings.ToLower(body)
340+
341+ switch {
342+ case statusCode == http.StatusUnauthorized:
343+ classifiedStatus = http.StatusUnauthorized
344+ terminal = true
345+ case strings.Contains(raw, "refresh_token_reused"),
346+ strings.Contains(raw, "token_invalidated"),
347+ strings.Contains(raw, "token_revoked"):
348+ classifiedStatus = http.StatusUnauthorized
349+ terminal = true
350+ }
351+
352+ return &refreshFailureError{
353+ statusCode: classifiedStatus,
354+ terminal: terminal,
355+ message: message,
356+ }
297357 }
298358
299359 // UpdateTokenStorage updates an existing CodexTokenStorage with new token data.
internal/auth/codex/openai_auth_test.go+42−0
@@ -15,6 +15,17 @@ func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { · 14 unmodified lines
1515 return f(req)
1616 }
1717
18+func testStatusCode(err error) int {
19+ if err == nil {
20+ return 0
21+ }
22+ statusErr, ok := err.(interface{ StatusCode() int })
23+ if !ok || statusErr == nil {
24+ return 0
25+ }
26+ return statusErr.StatusCode()
27+}
28+
1829 func TestRefreshTokensWithRetry_NonRetryableOnlyAttemptsOnce(t *testing.T) {
1930 var calls int32
2031 auth := &CodexAuth{
@@ -38,6 +49,37 @@ func TestRefreshTokensWithRetry_NonRetryableOnlyAttemptsOnce(t *testing.T) { · 17 unmodified lines
3849 if !strings.Contains(strings.ToLower(err.Error()), "refresh_token_reused") {
3950 t.Fatalf("expected refresh_token_reused in error, got: %v", err)
4051 }
52+ if got := testStatusCode(err); got != http.StatusUnauthorized {
53+ t.Fatalf("expected maintenance status code 401, got %d", got)
54+ }
55+ if got := atomic.LoadInt32(&calls); got != 1 {
56+ t.Fatalf("expected 1 refresh attempt, got %d", got)
57+ }
58+}
59+
60+func TestRefreshTokensWithRetry_UnauthorizedOnlyAttemptsOnce(t *testing.T) {
61+ var calls int32
62+ auth := &CodexAuth{
63+ httpClient: &http.Client{
64+ Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
65+ atomic.AddInt32(&calls, 1)
66+ return &http.Response{
67+ StatusCode: http.StatusUnauthorized,
68+ Body: io.NopCloser(strings.NewReader(`{"error":{"message":"Your authentication token has been invalidated. Please try signing in again.","type":"invalid_request_error","code":"token_invalidated"}}`)),
69+ Header: make(http.Header),
70+ Request: req,
71+ }, nil
72+ }),
73+ },
74+ }
75+
76+ _, err := auth.RefreshTokensWithRetry(context.Background(), "dummy_refresh_token", 3)
77+ if err == nil {
78+ t.Fatalf("expected error for unauthorized refresh failure")
79+ }
80+ if got := testStatusCode(err); got != http.StatusUnauthorized {
81+ t.Fatalf("expected maintenance status code 401, got %d", got)
82+ }
4183 if got := atomic.LoadInt32(&calls); got != 1 {
4284 t.Fatalf("expected 1 refresh attempt, got %d", got)
4385 }
sdk/cliproxy/auth/conductor.go+21−1
@@ -1940,6 +1940,21 @@ func statusCodeFromError(err error) int { · 1939 unmodified lines
19401940 return 0
19411941 }
19421942
1943+func terminalStatusCodeFromError(err error) int {
1944+ if err == nil {
1945+ return 0
1946+ }
1947+ type terminalStatusCoder interface {
1948+ StatusCode() int
1949+ Terminal() bool
1950+ }
1951+ var sc terminalStatusCoder
1952+ if errors.As(err, &sc) && sc != nil && sc.Terminal() {
1953+ return sc.StatusCode()
1954+ }
1955+ return 0
1956+}
1957+
19431958 func retryAfterFromError(err error) *time.Duration {
19441959 if err == nil {
19451960 return nil
@@ -2755,10 +2770,15 @@ func (m *Manager) refreshAuth(ctx context.Context, id string) { · 809 unmodified lines
27552770 log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err)
27562771 now := time.Now()
27572772 if err != nil {
2773+ refreshErr := &Error{Message: err.Error()}
2774+ if status := terminalStatusCodeFromError(err); status > 0 {
2775+ refreshErr.HTTPStatus = status
2776+ }
27582777 m.mu.Lock()
27592778 if current := m.auths[id]; current != nil {
27602779 current.NextRefreshAfter = now.Add(refreshFailureBackoff)
2761− current.LastError = &Error{Message: err.Error()}
2780+ current.LastError = refreshErr
2781+ current.UpdatedAt = now
27622782 m.auths[id] = current
27632783 if m.scheduler != nil {
27642784 m.scheduler.upsertAuth(current.Clone())
sdk/cliproxy/auth/conductor_refresh_failure_test.go+151−0
@@ -0,0 +1,151 @@
1+package auth
2+
3+import (
4+ "context"
5+ "net/http"
6+ "strings"
7+ "testing"
8+ "time"
9+
10+ cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
11+)
12+
13+type refreshFailureTestExecutor struct {
14+ provider string
15+ err error
16+}
17+
18+func (e refreshFailureTestExecutor) Identifier() string { return e.provider }
19+
20+func (e refreshFailureTestExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
21+ return cliproxyexecutor.Response{}, nil
22+}
23+
24+func (e refreshFailureTestExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
25+ return nil, nil
26+}
27+
28+func (e refreshFailureTestExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) {
29+ return auth, e.err
30+}
31+
32+func (e refreshFailureTestExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
33+ return cliproxyexecutor.Response{}, nil
34+}
35+
36+func (e refreshFailureTestExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
37+ return nil, nil
38+}
39+
40+type terminalRefreshTestError struct {
41+ status int
42+ msg string
43+}
44+
45+func (e terminalRefreshTestError) Error() string { return e.msg }
46+func (e terminalRefreshTestError) StatusCode() int { return e.status }
47+func (e terminalRefreshTestError) Terminal() bool { return true }
48+
49+type transientRefreshTestError struct {
50+ status int
51+ msg string
52+}
53+
54+func (e transientRefreshTestError) Error() string { return e.msg }
55+func (e transientRefreshTestError) StatusCode() int { return e.status }
56+func (e transientRefreshTestError) Terminal() bool { return false }
57+
58+func TestManagerRefreshAuth_PersistsTerminalRefresh401ForMaintenance(t *testing.T) {
59+ ctx := context.Background()
60+ manager := NewManager(nil, nil, nil)
61+ manager.RegisterExecutor(refreshFailureTestExecutor{
62+ provider: "codex",
63+ err: terminalRefreshTestError{
64+ status: http.StatusUnauthorized,
65+ msg: "token refresh failed with status 401: unauthorized",
66+ },
67+ })
68+
69+ auth := &Auth{
70+ ID: "refresh-401",
71+ Provider: "codex",
72+ Status: StatusActive,
73+ Metadata: map[string]any{
74+ "email": "user@example.com",
75+ "refresh_token": "refresh-token",
76+ },
77+ }
78+ if _, err := manager.Register(ctx, auth); err != nil {
79+ t.Fatalf("register auth: %v", err)
80+ }
81+
82+ started := time.Now()
83+ manager.refreshAuth(ctx, auth.ID)
84+
85+ updated, ok := manager.GetByID(auth.ID)
86+ if !ok || updated == nil {
87+ t.Fatalf("expected auth %q to remain registered", auth.ID)
88+ }
89+ if updated.LastError == nil {
90+ t.Fatal("expected refresh failure to persist LastError")
91+ }
92+ if updated.LastError.HTTPStatus != http.StatusUnauthorized {
93+ t.Fatalf("expected LastError.HTTPStatus = 401, got %d", updated.LastError.HTTPStatus)
94+ }
95+ if !strings.Contains(updated.LastError.Message, "status 401") {
96+ t.Fatalf("expected LastError.Message to preserve refresh failure details, got %q", updated.LastError.Message)
97+ }
98+ if updated.Status != StatusActive {
99+ t.Fatalf("expected auth status to remain active, got %q", updated.Status)
100+ }
101+ if updated.Unavailable {
102+ t.Fatal("expected auth to remain schedulable until maintenance handles deletion")
103+ }
104+ if updated.NextRefreshAfter.IsZero() || !updated.NextRefreshAfter.After(started) {
105+ t.Fatalf("expected NextRefreshAfter to be scheduled after refresh failure, got %v", updated.NextRefreshAfter)
106+ }
107+}
108+
109+func TestManagerRefreshAuth_DoesNotMarkTransientRefreshStatusForMaintenance(t *testing.T) {
110+ ctx := context.Background()
111+ manager := NewManager(nil, nil, nil)
112+ manager.RegisterExecutor(refreshFailureTestExecutor{
113+ provider: "codex",
114+ err: transientRefreshTestError{
115+ status: http.StatusTooManyRequests,
116+ msg: "token refresh failed with status 429: rate limited",
117+ },
118+ })
119+
120+ auth := &Auth{
121+ ID: "refresh-429",
122+ Provider: "codex",
123+ Status: StatusActive,
124+ Metadata: map[string]any{
125+ "email": "user@example.com",
126+ "refresh_token": "refresh-token",
127+ },
128+ }
129+ if _, err := manager.Register(ctx, auth); err != nil {
130+ t.Fatalf("register auth: %v", err)
131+ }
132+
133+ manager.refreshAuth(ctx, auth.ID)
134+
135+ updated, ok := manager.GetByID(auth.ID)
136+ if !ok || updated == nil {
137+ t.Fatalf("expected auth %q to remain registered", auth.ID)
138+ }
139+ if updated.LastError == nil {
140+ t.Fatal("expected refresh failure to persist LastError")
141+ }
142+ if updated.LastError.HTTPStatus != 0 {
143+ t.Fatalf("expected transient refresh error to avoid maintenance status code, got %d", updated.LastError.HTTPStatus)
144+ }
145+ if updated.Status != StatusActive {
146+ t.Fatalf("expected auth status to remain active, got %q", updated.Status)
147+ }
148+ if updated.Unavailable {
149+ t.Fatal("expected transient refresh failure to avoid blocking scheduler")
150+ }
151+}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t