Unified · Split
| @@ -7,6 +7,7 @@ package codex · 6 unmodified lines | |||
| 7 | import ( | 7 | import ( |
| 8 | "context" | 8 | "context" |
| 9 | "encoding/json" | 9 | "encoding/json" |
| 10 | "errors" | ||
| 10 | "fmt" | 11 | "fmt" |
| 11 | "io" | 12 | "io" |
| 12 | "net/http" | 13 | "net/http" |
| @@ -34,6 +35,33 @@ type CodexAuth struct { · 21 unmodified lines | |||
| 34 | httpClient *http.Client | 35 | httpClient *http.Client |
| 35 | } | 36 | } |
| 36 | 37 | ||
| 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 | |||
| 37 | // NewCodexAuth creates a new CodexAuth service instance. | 65 | // NewCodexAuth creates a new CodexAuth service instance. |
| 38 | // It initializes an HTTP client with proxy settings from the provided configuration. | 66 | // It initializes an HTTP client with proxy settings from the provided configuration. |
| 39 | func NewCodexAuth(cfg *config.Config) *CodexAuth { | 67 | func NewCodexAuth(cfg *config.Config) *CodexAuth { |
| @@ -202,7 +230,7 @@ func (o *CodexAuth) RefreshTokens(ctx context.Context, refreshToken string) (*Co · 162 unmodified lines | |||
| 202 | } | 230 | } |
| 203 | 231 | ||
| 204 | if resp.StatusCode != http.StatusOK { | 232 | 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)) |
| 206 | } | 234 | } |
| 207 | 235 | ||
| 208 | var tokenResp struct { | 236 | var tokenResp struct { |
| @@ -292,8 +320,40 @@ func isNonRetryableRefreshErr(err error) bool { · 83 unmodified lines | |||
| 292 | if err == nil { | 320 | if err == nil { |
| 293 | return false | 321 | return false |
| 294 | } | 322 | } |
| 323 | var terminal interface { | ||
| 324 | Terminal() bool | ||
| 325 | } | ||
| 326 | if errors.As(err, &terminal) && terminal != nil && terminal.Terminal() { | ||
| 327 | return true | ||
| 328 | } | ||
| 295 | raw := strings.ToLower(err.Error()) | 329 | 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 | } | ||
| 297 | } | 357 | } |
| 298 | 358 | ||
| 299 | // UpdateTokenStorage updates an existing CodexTokenStorage with new token data. | 359 | // UpdateTokenStorage updates an existing CodexTokenStorage with new token data. |
| @@ -15,6 +15,17 @@ func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { · 14 unmodified lines | |||
| 15 | return f(req) | 15 | return f(req) |
| 16 | } | 16 | } |
| 17 | 17 | ||
| 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 | |||
| 18 | func TestRefreshTokensWithRetry_NonRetryableOnlyAttemptsOnce(t *testing.T) { | 29 | func TestRefreshTokensWithRetry_NonRetryableOnlyAttemptsOnce(t *testing.T) { |
| 19 | var calls int32 | 30 | var calls int32 |
| 20 | auth := &CodexAuth{ | 31 | auth := &CodexAuth{ |
| @@ -38,6 +49,37 @@ func TestRefreshTokensWithRetry_NonRetryableOnlyAttemptsOnce(t *testing.T) { · 17 unmodified lines | |||
| 38 | if !strings.Contains(strings.ToLower(err.Error()), "refresh_token_reused") { | 49 | if !strings.Contains(strings.ToLower(err.Error()), "refresh_token_reused") { |
| 39 | t.Fatalf("expected refresh_token_reused in error, got: %v", err) | 50 | t.Fatalf("expected refresh_token_reused in error, got: %v", err) |
| 40 | } | 51 | } |
| 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 | } | ||
| 41 | if got := atomic.LoadInt32(&calls); got != 1 { | 83 | if got := atomic.LoadInt32(&calls); got != 1 { |
| 42 | t.Fatalf("expected 1 refresh attempt, got %d", got) | 84 | t.Fatalf("expected 1 refresh attempt, got %d", got) |
| 43 | } | 85 | } |
| @@ -1940,6 +1940,21 @@ func statusCodeFromError(err error) int { · 1939 unmodified lines | |||
| 1940 | return 0 | 1940 | return 0 |
| 1941 | } | 1941 | } |
| 1942 | 1942 | ||
| 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 | |||
| 1943 | func retryAfterFromError(err error) *time.Duration { | 1958 | func retryAfterFromError(err error) *time.Duration { |
| 1944 | if err == nil { | 1959 | if err == nil { |
| 1945 | return nil | 1960 | return nil |
| @@ -2755,10 +2770,15 @@ func (m *Manager) refreshAuth(ctx context.Context, id string) { · 809 unmodified lines | |||
| 2755 | log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err) | 2770 | log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err) |
| 2756 | now := time.Now() | 2771 | now := time.Now() |
| 2757 | if err != nil { | 2772 | if err != nil { |
| 2773 | refreshErr := &Error{Message: err.Error()} | ||
| 2774 | if status := terminalStatusCodeFromError(err); status > 0 { | ||
| 2775 | refreshErr.HTTPStatus = status | ||
| 2776 | } | ||
| 2758 | m.mu.Lock() | 2777 | m.mu.Lock() |
| 2759 | if current := m.auths[id]; current != nil { | 2778 | if current := m.auths[id]; current != nil { |
| 2760 | current.NextRefreshAfter = now.Add(refreshFailureBackoff) | 2779 | current.NextRefreshAfter = now.Add(refreshFailureBackoff) |
| 2761 | current.LastError = &Error{Message: err.Error()} | 2780 | current.LastError = refreshErr |
| 2781 | current.UpdatedAt = now | ||
| 2762 | m.auths[id] = current | 2782 | m.auths[id] = current |
| 2763 | if m.scheduler != nil { | 2783 | if m.scheduler != nil { |
| 2764 | m.scheduler.upsertAuth(current.Clone()) | 2784 | m.scheduler.upsertAuth(current.Clone()) |
| @@ -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 | } | ||