diff --git a/src/services/mcp/McpHub.ts b/src/services/mcp/McpHub.ts index aa1dc3abec..9c67c10edf 100644 --- a/src/services/mcp/McpHub.ts +++ b/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 = new Map() private initializationPromise: Promise private secretStorage?: SecretStorageService + private reauthPromises: Map> = 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 + } } /** diff --git a/src/services/mcp/McpOAuthClientProvider.ts b/src/services/mcp/McpOAuthClientProvider.ts index 0327214be9..48ba529ce5 100644 --- a/src/services/mcp/McpOAuthClientProvider.ts +++ b/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 | 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, + private _server: http.Server | null, + private _port: number, + private _authCodePromise: Promise | 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 { + 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 { 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 { @@ -209,6 +247,10 @@ export class McpOAuthClientProvider implements OAuthClientProvider { } async redirectToAuthorization(authorizationUrl: URL): Promise { + // 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 { - return this._authCodePromise + async waitForAuthCode(): Promise { + 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 { + 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 = { + 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 { - if (!this._closed) { + if (!this._closed && this._server) { this._closed = true await stopCallbackServer(this._server).catch(() => {}) + this._server = null + this._authCodePromise = null } } } diff --git a/src/services/mcp/SecretStorageService.ts b/src/services/mcp/SecretStorageService.ts index e320b2ba6d..6d49d4f7a1 100644 --- a/src/services/mcp/SecretStorageService.ts +++ b/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 { diff --git a/src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts b/src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts index b5863b3dd0..9486b75d14 100644 --- a/src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts +++ b/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) diff --git a/src/services/mcp/__tests__/SecretStorageService.spec.ts b/src/services/mcp/__tests__/SecretStorageService.spec.ts index 12709347d0..87d394c048 100644 --- a/src/services/mcp/__tests__/SecretStorageService.spec.ts +++ b/src/services/mcp/__tests__/SecretStorageService.spec.ts @@ -47,7 +47,7 @@ describe("SecretStorageService", () => { it("should return undefined for malformed JSON", async () => { // Manually store garbage via the underlying mock - context.secrets.store("mcp.oauth.example.com.data", "not-json") + context.secrets.store("mcp.oauth.example.com.mcp.data", "not-json") const result = await service.getOAuthData("https://example.com/mcp") expect(result).toBeUndefined() @@ -55,13 +55,23 @@ describe("SecretStorageService", () => { }) describe("saveOAuthData", () => { - it("should persist data under host-based key", async () => { + it("should persist data under host and path-based key", async () => { const data: StoredMcpOAuthData = { tokens: { access_token: "abc", token_type: "Bearer" }, expires_at: 12345, } await service.saveOAuthData("https://example.com/mcp", data) + expect(context.secrets.store).toHaveBeenCalledWith("mcp.oauth.example.com.mcp.data", JSON.stringify(data)) + }) + + it("should handle root path correctly", async () => { + const data: StoredMcpOAuthData = { + tokens: { access_token: "abc", token_type: "Bearer" }, + expires_at: 12345, + } + await service.saveOAuthData("https://example.com/", data) + expect(context.secrets.store).toHaveBeenCalledWith("mcp.oauth.example.com.data", JSON.stringify(data)) }) }) @@ -76,7 +86,7 @@ describe("SecretStorageService", () => { await service.deleteOAuthData("https://example.com/mcp") - expect(context.secrets.delete).toHaveBeenCalledWith("mcp.oauth.example.com.data") + expect(context.secrets.delete).toHaveBeenCalledWith("mcp.oauth.example.com.mcp.data") const result = await service.getOAuthData("https://example.com/mcp") expect(result).toBeUndefined() }) @@ -98,5 +108,21 @@ describe("SecretStorageService", () => { expect((await service.getOAuthData("https://host1.com/mcp"))?.tokens.access_token).toBe("a") expect((await service.getOAuthData("https://host2.com/mcp"))?.tokens.access_token).toBe("b") }) + + it("should isolate data by path on the same host", async () => { + const data1: StoredMcpOAuthData = { + tokens: { access_token: "path1", token_type: "Bearer" }, + expires_at: 1, + } + const data2: StoredMcpOAuthData = { + tokens: { access_token: "path2", token_type: "Bearer" }, + expires_at: 2, + } + await service.saveOAuthData("https://example.com/service1", data1) + await service.saveOAuthData("https://example.com/service2", data2) + + expect((await service.getOAuthData("https://example.com/service1"))?.tokens.access_token).toBe("path1") + expect((await service.getOAuthData("https://example.com/service2"))?.tokens.access_token).toBe("path2") + }) }) })