Browse Source

fix(claude-code): stop frequent sign-ins by hardening OAuth refresh (#10410)

* fix(claude-code): prevent sign-outs on oauth refresh

* test(claude-code): restore fetch after mocking

* refactor(claude-code): replace while(true) with bounded for loop for clarity

---------

Co-authored-by: Roo Code <[email protected]>
Hannes Rudolph 1 week ago
parent
commit
3074ccc278

+ 162 - 131
src/api/providers/claude-code.ts

@@ -122,153 +122,184 @@ export class ClaudeCodeHandler implements ApiHandler, SingleCompletionHandler {
 		// Reset per-request state that we persist into apiConversationHistory
 		this.lastThinkingSignature = undefined
 
-		// Get access token from OAuth manager
-		const accessToken = await claudeCodeOAuthManager.getAccessToken()
-
-		if (!accessToken) {
-			throw new Error(
+		const buildNotAuthenticatedError = () =>
+			new Error(
 				t("common:errors.claudeCode.notAuthenticated", {
 					defaultValue:
 						"Not authenticated with Claude Code. Please sign in using the Claude Code OAuth flow.",
 				}),
 			)
-		}
-
-		// Get user email for generating user_id metadata
-		const email = await claudeCodeOAuthManager.getEmail()
-
-		const model = this.getModel()
-
-		// Validate that the model ID is a valid ClaudeCodeModelId
-		const modelId = Object.hasOwn(claudeCodeModels, model.id)
-			? (model.id as ClaudeCodeModelId)
-			: claudeCodeDefaultModelId
 
-		// Generate user_id metadata in the format required by Claude Code API
-		const userId = generateUserId(email || undefined)
-
-		// Convert OpenAI tools to Anthropic format if provided and protocol is native
-		// Exclude tools when tool_choice is "none" since that means "don't use tools"
-		const shouldIncludeNativeTools =
-			metadata?.tools &&
-			metadata.tools.length > 0 &&
-			metadata?.toolProtocol !== "xml" &&
-			metadata?.tool_choice !== "none"
-
-		const anthropicTools = shouldIncludeNativeTools ? convertOpenAIToolsToAnthropic(metadata.tools!) : undefined
-
-		const anthropicToolChoice = shouldIncludeNativeTools
-			? convertOpenAIToolChoice(metadata.tool_choice, metadata.parallelToolCalls)
-			: undefined
-
-		// Determine reasoning effort and thinking configuration
-		const reasoningLevel = this.getReasoningEffort(model.info)
-
-		let thinking: ThinkingConfig
-		// With interleaved thinking (enabled via beta header), budget_tokens can exceed max_tokens
-		// as the token limit becomes the entire context window. We use the model's maxTokens.
-		// See: https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#interleaved-thinking
-		const maxTokens = model.info.maxTokens ?? 16384
-
-		if (reasoningLevel) {
-			// Use thinking mode with budget_tokens from config
-			const config = claudeCodeReasoningConfig[reasoningLevel]
-			thinking = {
-				type: "enabled",
-				budget_tokens: config.budgetTokens,
+		async function* streamOnce(this: ClaudeCodeHandler, accessToken: string): ApiStream {
+			// Get user email for generating user_id metadata
+			const email = await claudeCodeOAuthManager.getEmail()
+
+			const model = this.getModel()
+
+			// Validate that the model ID is a valid ClaudeCodeModelId
+			const modelId = Object.hasOwn(claudeCodeModels, model.id)
+				? (model.id as ClaudeCodeModelId)
+				: claudeCodeDefaultModelId
+
+			// Generate user_id metadata in the format required by Claude Code API
+			const userId = generateUserId(email || undefined)
+
+			// Convert OpenAI tools to Anthropic format if provided and protocol is native
+			// Exclude tools when tool_choice is "none" since that means "don't use tools"
+			const shouldIncludeNativeTools =
+				metadata?.tools &&
+				metadata.tools.length > 0 &&
+				metadata?.toolProtocol !== "xml" &&
+				metadata?.tool_choice !== "none"
+
+			const anthropicTools = shouldIncludeNativeTools ? convertOpenAIToolsToAnthropic(metadata.tools!) : undefined
+
+			const anthropicToolChoice = shouldIncludeNativeTools
+				? convertOpenAIToolChoice(metadata.tool_choice, metadata.parallelToolCalls)
+				: undefined
+
+			// Determine reasoning effort and thinking configuration
+			const reasoningLevel = this.getReasoningEffort(model.info)
+
+			let thinking: ThinkingConfig
+			// With interleaved thinking (enabled via beta header), budget_tokens can exceed max_tokens
+			// as the token limit becomes the entire context window. We use the model's maxTokens.
+			// See: https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#interleaved-thinking
+			const maxTokens = model.info.maxTokens ?? 16384
+
+			if (reasoningLevel) {
+				// Use thinking mode with budget_tokens from config
+				const config = claudeCodeReasoningConfig[reasoningLevel]
+				thinking = {
+					type: "enabled",
+					budget_tokens: config.budgetTokens,
+				}
+			} else {
+				// Explicitly disable thinking
+				thinking = { type: "disabled" }
 			}
-		} else {
-			// Explicitly disable thinking
-			thinking = { type: "disabled" }
-		}
 
-		// Create streaming request using OAuth
-		const stream = createStreamingMessage({
-			accessToken,
-			model: modelId,
-			systemPrompt,
-			messages,
-			maxTokens,
-			thinking,
-			tools: anthropicTools,
-			toolChoice: anthropicToolChoice,
-			metadata: {
-				user_id: userId,
-			},
-		})
-
-		// Track usage for cost calculation
-		let inputTokens = 0
-		let outputTokens = 0
-		let cacheReadTokens = 0
-		let cacheWriteTokens = 0
-
-		for await (const chunk of stream) {
-			switch (chunk.type) {
-				case "text":
-					yield {
-						type: "text",
-						text: chunk.text,
-					}
-					break
-
-				case "reasoning":
-					yield {
-						type: "reasoning",
-						text: chunk.text,
+			// Create streaming request using OAuth
+			const stream = createStreamingMessage({
+				accessToken,
+				model: modelId,
+				systemPrompt,
+				messages,
+				maxTokens,
+				thinking,
+				tools: anthropicTools,
+				toolChoice: anthropicToolChoice,
+				metadata: {
+					user_id: userId,
+				},
+			})
+
+			// Track usage for cost calculation
+			let inputTokens = 0
+			let outputTokens = 0
+			let cacheReadTokens = 0
+			let cacheWriteTokens = 0
+
+			for await (const chunk of stream) {
+				switch (chunk.type) {
+					case "text":
+						yield {
+							type: "text",
+							text: chunk.text,
+						}
+						break
+
+					case "reasoning":
+						yield {
+							type: "reasoning",
+							text: chunk.text,
+						}
+						break
+
+					case "thinking_complete":
+						// Capture the signature for persistence in api_conversation_history
+						// This enables tool use continuations where thinking blocks must be passed back
+						if (chunk.signature) {
+							this.lastThinkingSignature = chunk.signature
+						}
+						// Emit a complete thinking block with signature
+						// This is critical for interleaved thinking with tool use
+						// The signature must be included when passing thinking blocks back to the API
+						yield {
+							type: "reasoning",
+							text: chunk.thinking,
+							signature: chunk.signature,
+						}
+						break
+
+					case "tool_call_partial":
+						yield {
+							type: "tool_call_partial",
+							index: chunk.index,
+							id: chunk.id,
+							name: chunk.name,
+							arguments: chunk.arguments,
+						}
+						break
+
+					case "usage": {
+						inputTokens = chunk.inputTokens
+						outputTokens = chunk.outputTokens
+						cacheReadTokens = chunk.cacheReadTokens || 0
+						cacheWriteTokens = chunk.cacheWriteTokens || 0
+
+						// Claude Code is subscription-based, no per-token cost
+						const usageChunk: ApiStreamUsageChunk = {
+							type: "usage",
+							inputTokens,
+							outputTokens,
+							cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined,
+							cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined,
+							totalCost: 0,
+						}
+
+						yield usageChunk
+						break
 					}
-					break
 
-				case "thinking_complete":
-					// Capture the signature for persistence in api_conversation_history
-					// This enables tool use continuations where thinking blocks must be passed back
-					if (chunk.signature) {
-						this.lastThinkingSignature = chunk.signature
-					}
-					// Emit a complete thinking block with signature
-					// This is critical for interleaved thinking with tool use
-					// The signature must be included when passing thinking blocks back to the API
-					yield {
-						type: "reasoning",
-						text: chunk.thinking,
-						signature: chunk.signature,
-					}
-					break
-
-				case "tool_call_partial":
-					yield {
-						type: "tool_call_partial",
-						index: chunk.index,
-						id: chunk.id,
-						name: chunk.name,
-						arguments: chunk.arguments,
-					}
-					break
+					case "error":
+						throw new Error(chunk.error)
+				}
+			}
+		}
 
-				case "usage": {
-					inputTokens = chunk.inputTokens
-					outputTokens = chunk.outputTokens
-					cacheReadTokens = chunk.cacheReadTokens || 0
-					cacheWriteTokens = chunk.cacheWriteTokens || 0
-
-					// Claude Code is subscription-based, no per-token cost
-					const usageChunk: ApiStreamUsageChunk = {
-						type: "usage",
-						inputTokens,
-						outputTokens,
-						cacheReadTokens: cacheReadTokens > 0 ? cacheReadTokens : undefined,
-						cacheWriteTokens: cacheWriteTokens > 0 ? cacheWriteTokens : undefined,
-						totalCost: 0,
-					}
+		// Get access token from OAuth manager
+		let accessToken = await claudeCodeOAuthManager.getAccessToken()
+		if (!accessToken) {
+			throw buildNotAuthenticatedError()
+		}
 
-					yield usageChunk
-					break
+		// Try the request with at most one force-refresh retry on auth failure
+		for (let attempt = 0; attempt < 2; attempt++) {
+			try {
+				yield* streamOnce.call(this, accessToken)
+				return
+			} catch (error) {
+				const message = error instanceof Error ? error.message : String(error)
+				const isAuthFailure = /unauthorized|invalid token|not authenticated|authentication/i.test(message)
+
+				// Only retry on auth failure during first attempt
+				const canRetry = attempt === 0 && isAuthFailure
+				if (!canRetry) {
+					throw error
 				}
 
-				case "error":
-					throw new Error(chunk.error)
+				// Force refresh the token for retry
+				const refreshed = await claudeCodeOAuthManager.forceRefreshAccessToken()
+				if (!refreshed) {
+					throw buildNotAuthenticatedError()
+				}
+				accessToken = refreshed
 			}
 		}
+
+		// Unreachable: loop always returns on success or throws on failure
+		throw buildNotAuthenticatedError()
 	}
 
 	getModel(): { id: string; info: ModelInfo } {

+ 1 - 1
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)
+	claudeCodeOAuthManager.initialize(context, (message) => outputChannel.appendLine(message))
 
 	// Get default commands from configuration.
 	const defaultCommands = vscode.workspace.getConfiguration(Package.name).get<string[]>("allowedCommands") || []

+ 37 - 0
src/integrations/claude-code/__tests__/oauth.spec.ts

@@ -195,4 +195,41 @@ describe("Claude Code OAuth", () => {
 			expect(CLAUDE_CODE_OAUTH_CONFIG.callbackPort).toBe(54545)
 		})
 	})
+
+	describe("refresh token behavior", () => {
+		afterEach(() => {
+			vi.unstubAllGlobals()
+		})
+
+		test("refresh responses may omit refresh_token (should be tolerated)", async () => {
+			const { refreshAccessToken } = await import("../oauth")
+
+			// Mock fetch to return a refresh response with no refresh_token
+			const mockFetch = vi.fn().mockResolvedValue(
+				new Response(
+					JSON.stringify({
+						access_token: "new-access",
+						expires_in: 3600,
+						// refresh_token intentionally omitted
+					}),
+					{ status: 200, headers: { "Content-Type": "application/json" } },
+				),
+			)
+
+			vi.stubGlobal("fetch", mockFetch)
+
+			const creds: ClaudeCodeCredentials = {
+				type: "claude" as const,
+				access_token: "old-access",
+				refresh_token: "old-refresh",
+				expired: new Date(Date.now() - 1000).toISOString(),
+				email: "[email protected]",
+			}
+
+			const refreshed = await refreshAccessToken(creds)
+			expect(refreshed.access_token).toBe("new-access")
+			expect(refreshed.refresh_token).toBe("old-refresh")
+			expect(refreshed.email).toBe("[email protected]")
+		})
+	})
 })

+ 171 - 12
src/integrations/claude-code/oauth.ts

@@ -31,12 +31,74 @@ export type ClaudeCodeCredentials = z.infer<typeof claudeCodeCredentialsSchema>
 // Token response schema from Anthropic
 const tokenResponseSchema = z.object({
 	access_token: z.string(),
-	refresh_token: z.string(),
+	// Refresh responses may omit refresh_token (common OAuth behavior). When omitted,
+	// callers must preserve the existing refresh token.
+	refresh_token: z.string().min(1).optional(),
 	expires_in: z.number(),
 	email: z.string().optional(),
 	token_type: z.string().optional(),
 })
 
+class ClaudeCodeOAuthTokenError extends Error {
+	public readonly status?: number
+	public readonly errorCode?: string
+
+	constructor(message: string, opts?: { status?: number; errorCode?: string }) {
+		super(message)
+		this.name = "ClaudeCodeOAuthTokenError"
+		this.status = opts?.status
+		this.errorCode = opts?.errorCode
+	}
+
+	public isLikelyInvalidGrant(): boolean {
+		if (this.errorCode && /invalid_grant/i.test(this.errorCode)) {
+			return true
+		}
+		if (this.status === 400 || this.status === 401 || this.status === 403) {
+			return /invalid_grant|revoked|expired|invalid refresh/i.test(this.message)
+		}
+		return false
+	}
+}
+
+function parseOAuthErrorDetails(errorText: string): { errorCode?: string; errorMessage?: string } {
+	try {
+		const json: unknown = JSON.parse(errorText)
+		if (!json || typeof json !== "object") {
+			return {}
+		}
+
+		const obj = json as Record<string, unknown>
+		const errorField = obj.error
+
+		const errorCode: string | undefined =
+			typeof errorField === "string"
+				? errorField
+				: errorField &&
+					  typeof errorField === "object" &&
+					  typeof (errorField as Record<string, unknown>).type === "string"
+					? ((errorField as Record<string, unknown>).type as string)
+					: undefined
+
+		const errorDescription = obj.error_description
+		const errorMessageFromError =
+			errorField && typeof errorField === "object" ? (errorField as Record<string, unknown>).message : undefined
+
+		const errorMessage: string | undefined =
+			typeof errorDescription === "string"
+				? errorDescription
+				: typeof errorMessageFromError === "string"
+					? errorMessageFromError
+					: typeof obj.message === "string"
+						? obj.message
+						: undefined
+
+		return { errorCode, errorMessage }
+	} catch {
+		return {}
+	}
+}
+
 /**
  * Generates a cryptographically random PKCE code verifier
  * Must be 43-128 characters long using unreserved characters
@@ -134,6 +196,11 @@ export async function exchangeCodeForTokens(
 	const data = await response.json()
 	const tokenResponse = tokenResponseSchema.parse(data)
 
+	if (!tokenResponse.refresh_token) {
+		// The access token is unusable without a refresh token for persistence.
+		throw new Error("Token exchange did not return a refresh_token")
+	}
+
 	// Calculate expiry time
 	const expiresAt = new Date(Date.now() + tokenResponse.expires_in * 1000)
 
@@ -149,11 +216,11 @@ export async function exchangeCodeForTokens(
 /**
  * Refreshes the access token using the refresh token
  */
-export async function refreshAccessToken(refreshToken: string): Promise<ClaudeCodeCredentials> {
+export async function refreshAccessToken(credentials: ClaudeCodeCredentials): Promise<ClaudeCodeCredentials> {
 	const body = {
 		grant_type: "refresh_token",
 		client_id: CLAUDE_CODE_OAUTH_CONFIG.clientId,
-		refresh_token: refreshToken,
+		refresh_token: credentials.refresh_token,
 	}
 
 	const response = await fetch(CLAUDE_CODE_OAUTH_CONFIG.tokenEndpoint, {
@@ -167,7 +234,12 @@ export async function refreshAccessToken(refreshToken: string): Promise<ClaudeCo
 
 	if (!response.ok) {
 		const errorText = await response.text()
-		throw new Error(`Token refresh failed: ${response.status} ${response.statusText} - ${errorText}`)
+		const { errorCode, errorMessage } = parseOAuthErrorDetails(errorText)
+		const details = errorMessage ? errorMessage : errorText
+		throw new ClaudeCodeOAuthTokenError(
+			`Token refresh failed: ${response.status} ${response.statusText}${details ? ` - ${details}` : ""}`,
+			{ status: response.status, errorCode },
+		)
 	}
 
 	const data = await response.json()
@@ -179,9 +251,9 @@ export async function refreshAccessToken(refreshToken: string): Promise<ClaudeCo
 	return {
 		type: "claude",
 		access_token: tokenResponse.access_token,
-		refresh_token: tokenResponse.refresh_token,
+		refresh_token: tokenResponse.refresh_token ?? credentials.refresh_token,
 		expired: expiresAt.toISOString(),
-		email: tokenResponse.email,
+		email: tokenResponse.email ?? credentials.email,
 	}
 }
 
@@ -200,17 +272,80 @@ export function isTokenExpired(credentials: ClaudeCodeCredentials): boolean {
 export class ClaudeCodeOAuthManager {
 	private context: ExtensionContext | null = null
 	private credentials: ClaudeCodeCredentials | null = null
+	private logFn: ((message: string) => void) | null = null
+	private refreshPromise: Promise<ClaudeCodeCredentials> | null = null
 	private pendingAuth: {
 		codeVerifier: string
 		state: string
 		server?: http.Server
 	} | null = null
 
+	private log(message: string): void {
+		if (this.logFn) {
+			this.logFn(message)
+		} else {
+			console.log(message)
+		}
+	}
+
+	private logError(message: string, error?: unknown): void {
+		const details = error instanceof Error ? error.message : error !== undefined ? String(error) : undefined
+		const full = details ? `${message} ${details}` : message
+		this.log(full)
+		console.error(full)
+	}
+
 	/**
 	 * Initialize the OAuth manager with VS Code extension context
 	 */
-	initialize(context: ExtensionContext): void {
+	initialize(context: ExtensionContext, logFn?: (message: string) => void): void {
 		this.context = context
+		this.logFn = logFn ?? null
+	}
+
+	/**
+	 * Force a refresh using the stored refresh token even if the access token is not expired.
+	 * Useful when the server invalidates an access token early.
+	 */
+	async forceRefreshAccessToken(): Promise<string | null> {
+		if (!this.credentials) {
+			await this.loadCredentials()
+		}
+
+		if (!this.credentials) {
+			return null
+		}
+
+		try {
+			// De-dupe concurrent refreshes
+			if (!this.refreshPromise) {
+				const prevRefreshToken = this.credentials.refresh_token
+				this.log(`[claude-code-oauth] Forcing token refresh (expired=${this.credentials.expired})...`)
+				this.refreshPromise = refreshAccessToken(this.credentials).then((newCreds) => {
+					const rotated = newCreds.refresh_token !== prevRefreshToken
+					this.log(
+						`[claude-code-oauth] Forced refresh response received (expires_in≈${Math.round(
+							(new Date(newCreds.expired).getTime() - Date.now()) / 1000,
+						)}s, refresh_token_rotated=${rotated})`,
+					)
+					return newCreds
+				})
+			}
+
+			const newCredentials = await this.refreshPromise
+			this.refreshPromise = null
+			await this.saveCredentials(newCredentials)
+			this.log(`[claude-code-oauth] Forced token persisted (expired=${newCredentials.expired})`)
+			return newCredentials.access_token
+		} catch (error) {
+			this.refreshPromise = null
+			this.logError("[claude-code-oauth] Failed to force refresh token:", error)
+			if (error instanceof ClaudeCodeOAuthTokenError && error.isLikelyInvalidGrant()) {
+				this.log("[claude-code-oauth] Refresh token appears invalid; clearing stored credentials")
+				await this.clearCredentials()
+			}
+			return null
+		}
 	}
 
 	/**
@@ -231,7 +366,7 @@ export class ClaudeCodeOAuthManager {
 			this.credentials = claudeCodeCredentialsSchema.parse(parsed)
 			return this.credentials
 		} catch (error) {
-			console.error("[claude-code-oauth] Failed to load credentials:", error)
+			this.logError("[claude-code-oauth] Failed to load credentials:", error)
 			return null
 		}
 	}
@@ -276,12 +411,36 @@ export class ClaudeCodeOAuthManager {
 		// Check if token is expired and refresh if needed
 		if (isTokenExpired(this.credentials)) {
 			try {
-				const newCredentials = await refreshAccessToken(this.credentials.refresh_token)
+				// De-dupe concurrent refreshes
+				if (!this.refreshPromise) {
+					this.log(
+						`[claude-code-oauth] Access token expired (expired=${this.credentials.expired}). Refreshing...`,
+					)
+					const prevRefreshToken = this.credentials.refresh_token
+					this.refreshPromise = refreshAccessToken(this.credentials).then((newCreds) => {
+						const rotated = newCreds.refresh_token !== prevRefreshToken
+						this.log(
+							`[claude-code-oauth] Refresh response received (expires_in≈${Math.round(
+								(new Date(newCreds.expired).getTime() - Date.now()) / 1000,
+							)}s, refresh_token_rotated=${rotated})`,
+						)
+						return newCreds
+					})
+				}
+
+				const newCredentials = await this.refreshPromise
+				this.refreshPromise = null
 				await this.saveCredentials(newCredentials)
+				this.log(`[claude-code-oauth] Token persisted (expired=${newCredentials.expired})`)
 			} catch (error) {
-				console.error("[claude-code-oauth] Failed to refresh token:", error)
-				// Clear invalid credentials
-				await this.clearCredentials()
+				this.refreshPromise = null
+				this.logError("[claude-code-oauth] Failed to refresh token:", error)
+
+				// Only clear secrets when the refresh token is clearly invalid/revoked.
+				if (error instanceof ClaudeCodeOAuthTokenError && error.isLikelyInvalidGrant()) {
+					this.log("[claude-code-oauth] Refresh token appears invalid; clearing stored credentials")
+					await this.clearCredentials()
+				}
 				return null
 			}
 		}