feat(McpHub): handle OAuth refresh mid tool call

Author: edelauna

src/services/mcp/McpHub.ts

@@ -49,6 +49,7 @@ export type ConnectedMcpConnection = {
 	server: McpServer
 	client: Client
 	transport: StdioClientTransport | SSEClientTransport | StreamableHTTPClientTransport
+	authProvider?: McpOAuthClientProvider
 }
 
 export type DisconnectedMcpConnection = {
@@ -166,6 +167,7 @@ export class McpHub {
 	private sanitizedNameRegistry: Map<string, string> = new Map()
 	private initializationPromise: Promise<void>
 	private secretStorage?: SecretStorageService
+	private reauthPromises: Map<string, Promise<void>> = new Map()
 
 	constructor(provider: ClineProvider) {
 		this.providerRef = new WeakRef(provider)
@@ -822,6 +824,31 @@ export class McpHub {
 					console.error(`Transport error for "${name}" (streamable-http):`, error)
 					const connection = this.findConnection(name, source)
 					if (connection) {
+						if (error instanceof UnauthorizedError && authProvider) {
+							// Mid-session re-auth triggered by a tool call (401)
+							connection.server.status = "connecting"
+
+							const reauthKey = `${name}:${source}`
+							let reauthPromise = this.reauthPromises.get(reauthKey)
+							if (!reauthPromise) {
+								reauthPromise = this._completeOAuthFlow(
+									authProvider,
+									transport as StreamableHTTPClientTransport,
+									connection as ConnectedMcpConnection,
+									name,
+									source,
+								)
+									.catch((err) => {
+										console.error(`OAuth flow failed for "${name}":`, err)
+									})
+									.finally(() => {
+										this.reauthPromises.delete(reauthKey)
+									})
+								this.reauthPromises.set(reauthKey, reauthPromise)
+							}
+							return
+						}
+
 						connection.server.status = "disconnected"
 						this.appendErrorMessage(connection, error instanceof Error ? error.message : `${error}`)
 					}
@@ -905,6 +932,7 @@ export class McpHub {
 				},
 				client,
 				transport,
+				authProvider: streamableHttpAuthProvider,
 			}
 			this.connections.push(connection)
 
@@ -935,6 +963,7 @@ export class McpHub {
 			}
 
 			// Successful connection — close callback server if it was started.
+			// We keep the authProvider on the connection so it can handle mid-session 401s.
 			await streamableHttpAuthProvider?.close()
 
 			connection.server.status = "connected"
@@ -1203,9 +1232,10 @@ export class McpHub {
 				if (connection.type === "connected") {
 					await connection.transport.close()
 					await connection.client.close()
+					await connection.authProvider?.close()
 				}
 			} catch (error) {
-				console.error(`Failed to close transport for ${name}:`, error)
+				console.error(`Failed to close transport or auth provider for ${name}:`, error)
 			}
 		}
 
@@ -1876,19 +1906,66 @@ export class McpHub {
 			timeout = 60 * 1000
 		}
 
-		return await connection.client.request(
-			{
-				method: "tools/call",
-				params: {
-					name: toolName,
-					arguments: toolArguments,
+		try {
+			return await connection.client.request(
+				{
+					method: "tools/call",
+					params: {
+						name: toolName,
+						arguments: toolArguments,
+					},
 				},
-			},
-			CallToolResultSchema,
-			{
-				timeout,
-			},
-		)
+				CallToolResultSchema,
+				{
+					timeout,
+				},
+			)
+		} catch (error) {
+			if (error instanceof UnauthorizedError && connection.authProvider) {
+				// Mid-session re-auth triggered by a tool call (401)
+				connection.server.status = "connecting"
+
+				const reauthKey = `${serverName}:${source || connection.server.source || "global"}`
+				let reauthPromise = this.reauthPromises.get(reauthKey)
+
+				if (!reauthPromise) {
+					reauthPromise = this._completeOAuthFlow(
+						connection.authProvider,
+						connection.transport as StreamableHTTPClientTransport,
+						connection,
+						serverName,
+						source || connection.server.source || "global",
+					).finally(() => {
+						this.reauthPromises.delete(reauthKey)
+					})
+					this.reauthPromises.set(reauthKey, reauthPromise)
+				}
+
+				await reauthPromise
+
+				// After re-auth completes, the connection has been replaced.
+				// We need to find the new connection and retry the tool call.
+				const newConnection = this.findConnection(serverName, source)
+				if (!newConnection || newConnection.type !== "connected") {
+					throw new Error(`Failed to reconnect to server ${serverName} after OAuth`)
+				}
+
+				return await newConnection.client.request(
+					{
+						method: "tools/call",
+						params: {
+							name: toolName,
+							arguments: toolArguments,
+						},
+					},
+					CallToolResultSchema,
+					{
+						timeout,
+					},
+				)
+			}
+			throw error
+		}
 	}
 
 	/**

src/services/mcp/McpOAuthClientProvider.ts

@@ -36,13 +36,14 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
 	// when the redirect URI port changes between sessions.
 	private _clientInfo?: OAuthClientInformationFull
 	private _closed = false
+	private _refreshPromise: Promise<OAuthTokens> | null = null
 
 	private constructor(
 		private readonly _serverUrl: string,
 		private readonly _secretStorage: SecretStorageService,
-		private readonly _server: http.Server,
-		private readonly _port: number,
-		private readonly _authCodePromise: Promise<string>,
+		private _server: http.Server | null,
+		private _port: number,
+		private _authCodePromise: Promise<string> | null,
 		private readonly _tokenEndpointAuthMethod: string,
 		private readonly _grantTypes: string[],
 		private readonly _scopes: string[],
@@ -120,6 +121,20 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
 		return `http://localhost:${this._port}/callback`
 	}
 
+	private async _ensureCallbackServer(): Promise<void> {
+		if (this._server && !this._closed) return
+
+		this._closed = false
+		const { server, port, result } = await startCallbackServer(this._port, this._state)
+		this._server = server
+		this._port = port
+		this._authCodePromise = result.then((r) => {
+			if (r.error) throw new Error(`OAuth authorization failed: ${r.error}`)
+			if (!r.code) throw new Error("No authorization code received in callback")
+			return r.code
+		})
+	}
+
 	state(): string {
 		return this._state
 	}
@@ -192,10 +207,33 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
 	async tokens(): Promise<OAuthTokens | undefined> {
 		const data = await this._secretStorage.getOAuthData(this._serverUrl)
 		if (!data) return undefined
-		// Return undefined 5 minutes before expiry so the SDK triggers re-auth
-		// before the server actually rejects requests.
-		if (Date.now() >= data.expires_at - 5 * 60 * 1000) return undefined
-		return data.tokens
+
+		// If the access token is still valid (with 5m buffer), return it.
+		if (Date.now() < data.expires_at - 5 * 60 * 1000) {
+			return data.tokens
+		}
+
+		// Access token is expired or near expiry. Try to refresh if we have a refresh token.
+		if (data.tokens.refresh_token) {
+			if (this._refreshPromise) {
+				return this._refreshPromise
+			}
+
+			this._refreshPromise = this.refreshAccessToken(data.tokens.refresh_token).finally(() => {
+				this._refreshPromise = null
+			})
+
+			try {
+				return await this._refreshPromise
+			} catch (error) {
+				console.error(`Failed to refresh MCP OAuth token for ${this._serverUrl}:`, error)
+				// Clear stale tokens on refresh failure so we don't keep retrying a dead refresh token
+				await this._secretStorage.deleteOAuthData(this._serverUrl)
+				// Fall through to return undefined, which triggers full re-auth
+			}
+		}
+
+		return undefined
 	}
 
 	async saveTokens(tokens: OAuthTokens): Promise<void> {
@@ -209,6 +247,10 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
 	}
 
 	async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
+		// Ensure the callback server is running before opening the browser.
+		// This handles mid-session re-auth where the initial server was closed.
+		await this._ensureCallbackServer()
+
 		// Workaround for SDK metadata discovery bug (see utils/oauth.ts for issue links).
 		// The SDK's discoverOAuthMetadata() builds a wrong well-known URL for issuers
 		// with path components, causing it to fall back to a default "/authorize" path.
@@ -267,8 +309,11 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
 	 * browser flow and the local callback server receives the redirect.
 	 * Rejects on error or 5-minute timeout.
 	 */
-	waitForAuthCode(): Promise<string> {
-		return this._authCodePromise
+	async waitForAuthCode(): Promise<string> {
+		if (!this._authCodePromise) {
+			await this._ensureCallbackServer()
+		}
+		return this._authCodePromise!
 	}
 
 	/**
@@ -327,11 +372,54 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
 		await this.saveTokens(tokens)
 	}
 
+	/**
+	 * Refreshes the access token using a refresh token.
+	 * @param refreshToken The refresh token to use.
+	 * @returns The new tokens.
+	 */
+	async refreshAccessToken(refreshToken: string): Promise<OAuthTokens> {
+		if (!this._authServerMeta?.token_endpoint) {
+			throw new Error("No token_endpoint in auth server metadata — cannot refresh token")
+		}
+		if (!this._clientInfo) {
+			throw new Error("No client information — registerClientIfNeeded() must be called first")
+		}
+
+		const params: Record<string, string> = {
+			grant_type: "refresh_token",
+			refresh_token: refreshToken,
+			client_id: this._clientInfo.client_id,
+		}
+
+		if (this._tokenEndpointAuthMethod === "client_secret_post" && this._clientInfo.client_secret) {
+			params.client_secret = this._clientInfo.client_secret
+		}
+
+		const response = await fetch(this._authServerMeta.token_endpoint as string, {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/x-www-form-urlencoded",
+				Accept: "application/json",
+			},
+			body: new URLSearchParams(params).toString(),
+		})
+
+		if (!response.ok) {
+			throw new Error(`Token refresh failed: HTTP ${response.status}`)
+		}
+
+		const tokens = (await response.json()) as OAuthTokens
+		await this.saveTokens(tokens)
+		return tokens
+	}
+
 	/** Close the local callback server. Always call this when done. */
 	async close(): Promise<void> {
-		if (!this._closed) {
+		if (!this._closed && this._server) {
 			this._closed = true
 			await stopCallbackServer(this._server).catch(() => {})
+			this._server = null
+			this._authCodePromise = null
 		}
 	}
 }

src/services/mcp/SecretStorageService.ts

@@ -25,7 +25,12 @@ export class SecretStorageService {
 	}
 
 	private _key(serverUrl: string): string {
-		return `${this._namespace}${new URL(serverUrl).host}.data`
+		const url = new URL(serverUrl)
+		// Use host + pathname for stricter isolation between different MCP servers on the same host.
+		// We sanitize the pathname to ensure it's a valid key component.
+		const sanitizedPath = url.pathname.replace(/[^a-zA-Z0-9]/g, "_").replace(/^_+|_+$/g, "")
+		const pathSuffix = sanitizedPath ? `.${sanitizedPath}` : ""
+		return `${this._namespace}${url.host}${pathSuffix}.data`
 	}
 
 	async getOAuthData(serverUrl: string): Promise<StoredMcpOAuthData | undefined> {

src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts

@@ -178,7 +178,46 @@ describe("McpOAuthClientProvider", () => {
 			await provider.close()
 		})
 
-		it("should return undefined for expired tokens", async () => {
+		it("should refresh tokens when access token is expired but refresh token exists", async () => {
+			setupCallbackServerMock()
+			const secretStorage = createMockSecretStorage()
+			const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage)
+
+			const initialTokens = {
+				access_token: "expired-access",
+				refresh_token: "valid-refresh",
+				token_type: "Bearer",
+			}
+			const refreshedTokens = {
+				access_token: "new-access",
+				refresh_token: "new-refresh",
+				token_type: "Bearer",
+				expires_in: 3600,
+			}
+
+			await provider.saveClientInformation({ client_id: "id", redirect_uris: [] } as any)
+			await secretStorage.saveOAuthData("https://example.com/mcp", {
+				tokens: initialTokens,
+				expires_at: Date.now() - 1000,
+			})
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(refreshedTokens),
+			})
+
+			const result = await provider.tokens()
+			expect(result).toEqual(refreshedTokens)
+			expect(mockFetch).toHaveBeenCalledWith(
+				expect.stringContaining("/token"),
+				expect.objectContaining({
+					body: expect.stringContaining("grant_type=refresh_token"),
+				}),
+			)
+			await provider.close()
+		})
+
+		it("should return undefined for expired tokens without refresh token", async () => {
 			setupCallbackServerMock()
 			const secretStorage = createMockSecretStorage()
 			const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage)