feat(McpHub): handle OAuth refresh mid tool call

This commit is contained in:
Elliott de Launay 2026-03-11 19:16:31 -04:00
parent e580b77084
commit 947ae61a63
No known key found for this signature in database
GPG key ID: BB899BED766D1806
5 changed files with 263 additions and 28 deletions

View file

@ -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
}
}
/**

View file

@ -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
}
}
}

View file

@ -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> {

View file

@ -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)

View file

@ -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")
})
})
})