diff --git a/packages/cloud/src/index.ts b/packages/cloud/src/index.ts index 65b92ebc13c..2ed696a03c2 100644 --- a/packages/cloud/src/index.ts +++ b/packages/cloud/src/index.ts @@ -6,3 +6,6 @@ export { BridgeOrchestrator } from "./bridge/index.js" export { RetryQueue } from "./retry-queue/index.js" export type { QueuedRequest, QueueStats, RetryQueueConfig, RetryQueueEvents } from "./retry-queue/index.js" + +export { RefreshTimer } from "./RefreshTimer.js" +export type { RefreshTimerOptions } from "./RefreshTimer.js" diff --git a/src/extension.ts b/src/extension.ts index dcb941fa581..752906ed8cb 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -96,7 +96,7 @@ export async function activate(context: vscode.ExtensionContext) { TerminalRegistry.initialize() // Initialize Claude Code OAuth manager for direct API access. - claudeCodeOAuthManager.initialize(context) + await claudeCodeOAuthManager.initialize(context) // Get default commands from configuration. const defaultCommands = vscode.workspace.getConfiguration(Package.name).get("allowedCommands") || [] @@ -424,6 +424,9 @@ export async function deactivate() { await bridge.disconnect() } + // Cleanup Claude Code OAuth manager + claudeCodeOAuthManager.dispose() + await McpServerManager.cleanup(extensionContext) TelemetryService.instance.shutdown() TerminalRegistry.cleanup() diff --git a/src/integrations/claude-code/__tests__/oauth.spec.ts b/src/integrations/claude-code/__tests__/oauth.spec.ts index 526ef2f6f7d..3d9cd01b21a 100644 --- a/src/integrations/claude-code/__tests__/oauth.spec.ts +++ b/src/integrations/claude-code/__tests__/oauth.spec.ts @@ -5,10 +5,22 @@ import { generateUserId, buildAuthorizationUrl, isTokenExpired, + refreshAccessToken, CLAUDE_CODE_OAUTH_CONFIG, + ClaudeCodeOAuthManager, type ClaudeCodeCredentials, } from "../oauth" +// Mock vscode module +vi.mock("vscode", () => ({ + window: { + showWarningMessage: vi.fn(), + }, +})) + +// Mock fetch for token refresh tests +global.fetch = vi.fn() + describe("Claude Code OAuth", () => { describe("generateCodeVerifier", () => { test("should generate a base64url encoded verifier", () => { @@ -195,4 +207,239 @@ describe("Claude Code OAuth", () => { expect(CLAUDE_CODE_OAUTH_CONFIG.callbackPort).toBe(54545) }) }) + + describe("refreshAccessToken", () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + test("should successfully refresh token", async () => { + const mockResponse = { + access_token: "new-access-token", + refresh_token: "new-refresh-token", + expires_in: 3600, + email: "test@example.com", + } + + ;(global.fetch as any).mockResolvedValue({ + ok: true, + json: async () => mockResponse, + }) + + const credentials = await refreshAccessToken("old-refresh-token") + + expect(credentials.access_token).toBe("new-access-token") + expect(credentials.refresh_token).toBe("new-refresh-token") + expect(credentials.email).toBe("test@example.com") + expect(credentials.type).toBe("claude") + expect(new Date(credentials.expired).getTime()).toBeGreaterThan(Date.now()) + }) + + test("should throw error on failed refresh", async () => { + ;(global.fetch as any).mockResolvedValue({ + ok: false, + status: 401, + statusText: "Unauthorized", + text: async () => "Invalid refresh token", + }) + + await expect(refreshAccessToken("invalid-refresh-token")).rejects.toThrow( + /Token refresh failed.*401.*Unauthorized/, + ) + }) + }) + + describe("ClaudeCodeOAuthManager", () => { + let manager: ClaudeCodeOAuthManager + let mockContext: any + + beforeEach(() => { + vi.clearAllMocks() + manager = new ClaudeCodeOAuthManager() + mockContext = { + secrets: { + get: vi.fn(), + store: vi.fn(), + delete: vi.fn(), + }, + } + }) + + afterEach(() => { + manager.dispose() + }) + + describe("initialize", () => { + test("should initialize without credentials", async () => { + mockContext.secrets.get.mockResolvedValue(null) + + await manager.initialize(mockContext) + + expect(mockContext.secrets.get).toHaveBeenCalled() + }) + + test("should load and start refresh timer with existing credentials", async () => { + const futureDate = new Date(Date.now() + 60 * 60 * 1000) + const credentials: ClaudeCodeCredentials = { + type: "claude", + access_token: "test-token", + refresh_token: "test-refresh", + expired: futureDate.toISOString(), + } + + mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials)) + + await manager.initialize(mockContext) + + // Credentials should be loaded + expect(manager.getCredentials()).toMatchObject(credentials) + }) + }) + + describe("getAccessToken", () => { + test("should return null when not authenticated", async () => { + mockContext.secrets.get.mockResolvedValue(null) + await manager.initialize(mockContext) + + const token = await manager.getAccessToken() + + expect(token).toBeNull() + }) + + test("should return existing valid token", async () => { + const futureDate = new Date(Date.now() + 60 * 60 * 1000) + const credentials: ClaudeCodeCredentials = { + type: "claude", + access_token: "valid-token", + refresh_token: "test-refresh", + expired: futureDate.toISOString(), + } + + mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials)) + await manager.initialize(mockContext) + + const token = await manager.getAccessToken() + + expect(token).toBe("valid-token") + }) + + test("should refresh expired token automatically", async () => { + const pastDate = new Date(Date.now() - 60 * 60 * 1000) + const credentials: ClaudeCodeCredentials = { + type: "claude", + access_token: "expired-token", + refresh_token: "test-refresh", + expired: pastDate.toISOString(), + } + + mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials)) + + const mockResponse = { + access_token: "new-token", + refresh_token: "new-refresh", + expires_in: 3600, + } + + ;(global.fetch as any).mockResolvedValue({ + ok: true, + json: async () => mockResponse, + }) + + await manager.initialize(mockContext) + + const token = await manager.getAccessToken() + + expect(token).toBe("new-token") + expect(mockContext.secrets.store).toHaveBeenCalled() + }) + }) + + describe("saveCredentials", () => { + test("should save credentials and start refresh timer", async () => { + await manager.initialize(mockContext) + + const credentials: ClaudeCodeCredentials = { + type: "claude", + access_token: "new-token", + refresh_token: "new-refresh", + expired: new Date(Date.now() + 3600 * 1000).toISOString(), + } + + await manager.saveCredentials(credentials) + + expect(mockContext.secrets.store).toHaveBeenCalledWith( + "claude-code-oauth-credentials", + JSON.stringify(credentials), + ) + expect(manager.getCredentials()).toMatchObject(credentials) + }) + }) + + describe("clearCredentials", () => { + test("should clear credentials and stop refresh timer", async () => { + const credentials: ClaudeCodeCredentials = { + type: "claude", + access_token: "test-token", + refresh_token: "test-refresh", + expired: new Date(Date.now() + 3600 * 1000).toISOString(), + } + + mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials)) + await manager.initialize(mockContext) + + await manager.clearCredentials() + + expect(mockContext.secrets.delete).toHaveBeenCalled() + expect(manager.getCredentials()).toBeNull() + }) + }) + + describe("isAuthenticated", () => { + test("should return false when no credentials", async () => { + mockContext.secrets.get.mockResolvedValue(null) + await manager.initialize(mockContext) + + const isAuth = await manager.isAuthenticated() + + expect(isAuth).toBe(false) + }) + + test("should return true with valid credentials", async () => { + const futureDate = new Date(Date.now() + 60 * 60 * 1000) + const credentials: ClaudeCodeCredentials = { + type: "claude", + access_token: "valid-token", + refresh_token: "test-refresh", + expired: futureDate.toISOString(), + } + + mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials)) + await manager.initialize(mockContext) + + const isAuth = await manager.isAuthenticated() + + expect(isAuth).toBe(true) + }) + }) + + describe("dispose", () => { + test("should stop refresh timer and cancel auth flow", async () => { + const credentials: ClaudeCodeCredentials = { + type: "claude", + access_token: "test-token", + refresh_token: "test-refresh", + expired: new Date(Date.now() + 3600 * 1000).toISOString(), + } + + mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials)) + await manager.initialize(mockContext) + + // Should not throw + manager.dispose() + + // Calling dispose multiple times should be safe + manager.dispose() + }) + }) + }) }) diff --git a/src/integrations/claude-code/oauth.ts b/src/integrations/claude-code/oauth.ts index 036ad70b631..46ad8750591 100644 --- a/src/integrations/claude-code/oauth.ts +++ b/src/integrations/claude-code/oauth.ts @@ -3,6 +3,7 @@ import * as http from "http" import { URL } from "url" import type { ExtensionContext } from "vscode" import { z } from "zod" +import { RefreshTimer } from "@roo-code/cloud" // OAuth Configuration export const CLAUDE_CODE_OAUTH_CONFIG = { @@ -205,12 +206,21 @@ export class ClaudeCodeOAuthManager { state: string server?: http.Server } | null = null + private refreshTimer: RefreshTimer | null = null + private refreshAttempts: number = 0 + private maxRefreshAttempts: number = 3 /** * Initialize the OAuth manager with VS Code extension context */ - initialize(context: ExtensionContext): void { + async initialize(context: ExtensionContext): Promise { this.context = context + + // Load existing credentials and start refresh timer if authenticated + await this.loadCredentials() + if (this.credentials) { + this.startRefreshTimer() + } } /** @@ -246,6 +256,14 @@ export class ClaudeCodeOAuthManager { await this.context.secrets.store(CLAUDE_CODE_CREDENTIALS_KEY, JSON.stringify(credentials)) this.credentials = credentials + + // Reset refresh attempts on successful save + this.refreshAttempts = 0 + + // Start refresh timer if not already running + if (!this.refreshTimer) { + this.startRefreshTimer() + } } /** @@ -258,6 +276,8 @@ export class ClaudeCodeOAuthManager { await this.context.secrets.delete(CLAUDE_CODE_CREDENTIALS_KEY) this.credentials = null + this.stopRefreshTimer() + this.refreshAttempts = 0 } /** @@ -278,11 +298,25 @@ export class ClaudeCodeOAuthManager { try { const newCredentials = await refreshAccessToken(this.credentials.refresh_token) await this.saveCredentials(newCredentials) + console.log("[claude-code-oauth] Token refreshed successfully") } catch (error) { console.error("[claude-code-oauth] Failed to refresh token:", error) - // Clear invalid credentials - await this.clearCredentials() - return null + + // Increment refresh attempts + this.refreshAttempts++ + + // Only clear credentials after max attempts to allow retry + if (this.refreshAttempts >= this.maxRefreshAttempts) { + console.error( + `[claude-code-oauth] Max refresh attempts (${this.maxRefreshAttempts}) reached, clearing credentials`, + ) + await this.clearCredentials() + await this.notifyReauthRequired() + return null + } + + // Don't return null immediately - let the old token be tried + // It might still work for some grace period } } @@ -473,6 +507,120 @@ export class ClaudeCodeOAuthManager { getCredentials(): ClaudeCodeCredentials | null { return this.credentials } + + /** + * Start the background token refresh timer + * Refreshes tokens proactively before they expire + */ + private startRefreshTimer(): void { + // Don't start if already running + if (this.refreshTimer) { + return + } + + this.refreshTimer = new RefreshTimer({ + callback: async () => { + try { + await this.performBackgroundRefresh() + return true // Success + } catch (error) { + console.error("[claude-code-oauth] Background refresh failed:", error) + return false // Failure - will trigger backoff + } + }, + // Refresh every 50 minutes (tokens typically last 1 hour, we refresh with 10min buffer) + successInterval: 50 * 60 * 1000, + initialBackoffMs: 1000, + maxBackoffMs: 5 * 60 * 1000, // Max 5 minutes backoff + }) + + this.refreshTimer.start() + console.log("[claude-code-oauth] Started background token refresh timer") + } + + /** + * Stop the background token refresh timer + */ + private stopRefreshTimer(): void { + if (this.refreshTimer) { + this.refreshTimer.stop() + this.refreshTimer = null + console.log("[claude-code-oauth] Stopped background token refresh timer") + } + } + + /** + * Perform a background token refresh + */ + private async performBackgroundRefresh(): Promise { + if (!this.credentials) { + throw new Error("No credentials available for refresh") + } + + // Only refresh if token is close to expiring (within 10 minutes) + const expiryTime = new Date(this.credentials.expired).getTime() + const timeUntilExpiry = expiryTime - Date.now() + const tenMinutes = 10 * 60 * 1000 + + if (timeUntilExpiry > tenMinutes) { + // Token is still fresh, no need to refresh yet + console.log( + `[claude-code-oauth] Token still valid for ${Math.round(timeUntilExpiry / 60000)} minutes, skipping refresh`, + ) + return + } + + console.log("[claude-code-oauth] Performing background token refresh...") + + try { + const newCredentials = await refreshAccessToken(this.credentials.refresh_token) + await this.saveCredentials(newCredentials) + console.log("[claude-code-oauth] Background token refresh successful") + } catch (error) { + this.refreshAttempts++ + + if (this.refreshAttempts >= this.maxRefreshAttempts) { + console.error( + `[claude-code-oauth] Max refresh attempts (${this.maxRefreshAttempts}) reached in background refresh`, + ) + await this.clearCredentials() + await this.notifyReauthRequired() + throw new Error("Failed to refresh token after multiple attempts") + } + + throw error + } + } + + /** + * Notify the user that re-authentication is required + */ + private async notifyReauthRequired(): Promise { + try { + const vscode = await import("vscode") + const action = await vscode.window.showWarningMessage( + "Your Claude.ai authentication has expired. Please sign in again to continue using Claude models.", + "Sign In", + "Later", + ) + + if (action === "Sign In") { + // Trigger the sign-in flow through the webview message handler + // This will be handled by the UI layer + console.log("[claude-code-oauth] User requested re-authentication") + } + } catch (error) { + console.error("[claude-code-oauth] Failed to show re-auth notification:", error) + } + } + + /** + * Dispose of resources (call this when extension deactivates) + */ + dispose(): void { + this.stopRefreshTimer() + this.cancelAuthorizationFlow() + } } // Singleton instance