From 41c52b9df6e12efb9bb5ae9a6427e74d24a26446 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Thu, 18 Jun 2026 23:15:43 +0800 Subject: [PATCH] test(management): add concurrency test for Codex OAuth session handling - Introduced `TestRequestCodexTokenCompletionKeepsConcurrentSessionPending` to validate proper handling of concurrent OAuth sessions. - Refactored Codex OAuth logic to use `newCodexOAuthService` for improved testability. Closes: #3171 --- .../api/handlers/management/auth_files.go | 16 +-- .../oauth_codex_concurrency_test.go | 111 ++++++++++++++++++ 2 files changed, 120 insertions(+), 7 deletions(-) create mode 100644 internal/api/handlers/management/oauth_codex_concurrency_test.go diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index e2483690..162f1fa8 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -28,6 +28,7 @@ import ( geminiAuth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/gemini" "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/kimi" xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v7/internal/misc" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" @@ -58,11 +59,18 @@ type callbackForwarder struct { done chan struct{} } +type codexOAuthService interface { + GenerateAuthURL(state string, pkceCodes *codex.PKCECodes) (string, error) + ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *codex.PKCECodes) (*codex.CodexAuthBundle, error) + CreateTokenStorage(bundle *codex.CodexAuthBundle) *codex.CodexTokenStorage +} + var ( callbackForwardersMu sync.Mutex callbackForwarders = make(map[int]*callbackForwarder) errAuthFileMustBeJSON = errors.New("auth file must be .json") errAuthFileNotFound = errors.New("auth file not found") + newCodexOAuthService = func(cfg *config.Config) codexOAuthService { return codex.NewCodexAuth(cfg) } ) func extractLastRefreshTimestamp(meta map[string]any) (time.Time, bool) { @@ -1891,7 +1899,6 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) { } fmt.Println("You can now use Claude services through this CLI") CompleteOAuthSession(state) - CompleteOAuthSessionsByProvider("anthropic") }() c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) @@ -2149,7 +2156,6 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) { } CompleteOAuthSession(state) - CompleteOAuthSessionsByProvider("gemini") fmt.Printf("You can now use Gemini CLI services through this CLI; token saved to %s\n", savedPath) }() @@ -2179,7 +2185,7 @@ func (h *Handler) RequestCodexToken(c *gin.Context) { } // Initialize Codex auth service - openaiAuth := codex.NewCodexAuth(h.cfg) + openaiAuth := newCodexOAuthService(h.cfg) // Generate authorization URL authURL, err := openaiAuth.GenerateAuthURL(state, pkceCodes) @@ -2296,7 +2302,6 @@ func (h *Handler) RequestCodexToken(c *gin.Context) { } fmt.Println("You can now use Codex services through this CLI") CompleteOAuthSession(state) - CompleteOAuthSessionsByProvider("codex") }() c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) @@ -2456,7 +2461,6 @@ func (h *Handler) RequestAntigravityToken(c *gin.Context) { } CompleteOAuthSession(state) - CompleteOAuthSessionsByProvider("antigravity") fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) if projectID != "" { fmt.Printf("Using GCP project: %s\n", util.HideAPIKey(projectID)) @@ -2638,7 +2642,6 @@ func (h *Handler) RequestXAIToken(c *gin.Context) { } CompleteOAuthSession(state) - CompleteOAuthSessionsByProvider("xai") fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) fmt.Println("You can now use xAI services through this CLI") }() @@ -2717,7 +2720,6 @@ func (h *Handler) RequestKimiToken(c *gin.Context) { fmt.Printf("Authentication successful! Token saved to %s\n", savedPath) fmt.Println("You can now use Kimi services through this CLI") CompleteOAuthSession(state) - CompleteOAuthSessionsByProvider("kimi") }() c.JSON(200, gin.H{"status": "ok", "url": authURL, "state": state}) diff --git a/internal/api/handlers/management/oauth_codex_concurrency_test.go b/internal/api/handlers/management/oauth_codex_concurrency_test.go new file mode 100644 index 00000000..8d1e3a95 --- /dev/null +++ b/internal/api/handlers/management/oauth_codex_concurrency_test.go @@ -0,0 +1,111 @@ +package management + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +type fakeCodexOAuthService struct{} + +func (f *fakeCodexOAuthService) GenerateAuthURL(state string, pkceCodes *codex.PKCECodes) (string, error) { + return "https://auth.example.test/oauth?state=" + state, nil +} + +func (f *fakeCodexOAuthService) ExchangeCodeForTokens(ctx context.Context, code string, pkceCodes *codex.PKCECodes) (*codex.CodexAuthBundle, error) { + now := time.Now() + return &codex.CodexAuthBundle{ + TokenData: codex.CodexTokenData{ + IDToken: "invalid-test-id-token", + AccessToken: "access-" + code, + RefreshToken: "refresh-" + code, + Email: "codex-" + code + "@example.test", + Expire: now.Add(time.Hour).Format(time.RFC3339), + }, + LastRefresh: now.Format(time.RFC3339), + }, nil +} + +func (f *fakeCodexOAuthService) CreateTokenStorage(bundle *codex.CodexAuthBundle) *codex.CodexTokenStorage { + return &codex.CodexTokenStorage{ + IDToken: bundle.TokenData.IDToken, + AccessToken: bundle.TokenData.AccessToken, + RefreshToken: bundle.TokenData.RefreshToken, + AccountID: bundle.TokenData.AccountID, + LastRefresh: bundle.LastRefresh, + Email: bundle.TokenData.Email, + Expire: bundle.TokenData.Expire, + } +} + +func TestRequestCodexTokenCompletionKeepsConcurrentSessionPending(t *testing.T) { + originalNewCodexOAuthService := newCodexOAuthService + newCodexOAuthService = func(cfg *config.Config) codexOAuthService { + return &fakeCodexOAuthService{} + } + defer func() { + newCodexOAuthService = originalNewCodexOAuthService + }() + + authDir := filepath.Join(t.TempDir(), "auths") + handler := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: authDir}, nil) + router := gin.New() + router.GET("/codex-auth-url", handler.RequestCodexToken) + + firstState := requestCodexTokenState(t, router) + secondState := requestCodexTokenState(t, router) + defer CompleteOAuthSession(firstState) + defer CompleteOAuthSession(secondState) + + if _, errWrite := WriteOAuthCallbackFileForPendingSession(authDir, "codex", firstState, "first-code", ""); errWrite != nil { + t.Fatalf("write first callback file: %v", errWrite) + } + + waitForOAuthSessionDone(t, firstState) + if !IsOAuthSessionPending(secondState, "codex") { + t.Fatalf("expected concurrent codex session %s to remain pending after %s completed", secondState, firstState) + } +} + +func requestCodexTokenState(t *testing.T, router http.Handler) string { + t.Helper() + + req := httptest.NewRequest(http.MethodGet, "/codex-auth-url", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("expected status %d, got %d with body %s", http.StatusOK, w.Code, w.Body.String()) + } + + var payload struct { + State string `json:"state"` + } + if errDecode := json.Unmarshal(w.Body.Bytes(), &payload); errDecode != nil { + t.Fatalf("decode codex auth URL response: %v", errDecode) + } + if payload.State == "" { + t.Fatalf("expected codex auth URL response to include state") + } + return payload.State +} + +func waitForOAuthSessionDone(t *testing.T, state string) { + t.Helper() + + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if !IsOAuthSessionPending(state, "codex") { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatalf("timed out waiting for codex session %s to complete", state) +} -- 2.51.2