mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-05 08:10:14 +00:00
feat(McpHub): handle OAuth refresh mid tool call
This commit is contained in:
parent
e580b77084
commit
947ae61a63
5 changed files with 263 additions and 28 deletions
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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> {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue