fix(mcp): cancel timed out oauth refresh

This commit is contained in:
Aiden Cline 2026-06-13 11:28:54 -05:00
commit ecf550c88c
3 changed files with 27 additions and 8 deletions

View file

@ -230,11 +230,15 @@ export const layer = Layer.effect(
auth, auth,
) )
authProvider = provider authProvider = provider
const controller = new AbortController()
yield* Effect.tryPromise(() => yield* Effect.tryPromise(() =>
withTimeout( withTimeout(
provider.refreshTokensIfExpired(mcp.headers ? createFetchWithInit(fetch, { headers: mcp.headers }) : undefined), provider.refreshTokensIfExpired(
mcp.headers ? createFetchWithInit(fetch, { headers: mcp.headers }) : undefined,
controller.signal,
),
connectTimeout, connectTimeout,
), ).finally(() => controller.abort()),
).pipe(Effect.ignore) ).pipe(Effect.ignore)
} }

View file

@ -128,7 +128,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
) )
} }
async refreshTokensIfExpired(fetchFn?: FetchLike): Promise<boolean> { async refreshTokensIfExpired(fetchFn?: FetchLike, signal?: AbortSignal): Promise<boolean> {
const entry = await Effect.runPromise(this.auth.getForUrl(this.mcpName, this.serverUrl)) const entry = await Effect.runPromise(this.auth.getForUrl(this.mcpName, this.serverUrl))
if (!entry?.tokens?.refreshToken) return false if (!entry?.tokens?.refreshToken) return false
if (!entry.tokens.expiresAt) return false if (!entry.tokens.expiresAt) return false
@ -137,14 +137,17 @@ export class McpOAuthProvider implements OAuthClientProvider {
const clientInformation = await this.clientInformation() const clientInformation = await this.clientInformation()
if (!clientInformation) return false if (!clientInformation) return false
const info = await discoverOAuthServerInfo(this.serverUrl, { fetchFn }) const request = signal
? (url: string | URL, init?: RequestInit) => (fetchFn ?? fetch)(url, { ...init, signal })
: fetchFn
const info = await discoverOAuthServerInfo(this.serverUrl, { fetchFn: request })
await this.saveTokens( await this.saveTokens(
await refreshAuthorization(info.authorizationServerUrl, { await refreshAuthorization(info.authorizationServerUrl, {
metadata: info.authorizationServerMetadata, metadata: info.authorizationServerMetadata,
clientInformation, clientInformation,
refreshToken: entry.tokens.refreshToken, refreshToken: entry.tokens.refreshToken,
resource: await selectResourceURL(this.serverUrl, this, info.resourceMetadata), resource: await selectResourceURL(this.serverUrl, this, info.resourceMetadata),
fetchFn, fetchFn: request,
}), }),
) )
return true return true

View file

@ -54,6 +54,7 @@ let transportCloseCount = 0
// Captures the opts passed to each MockStdioTransport, keyed by lastCreatedClientName // Captures the opts passed to each MockStdioTransport, keyed by lastCreatedClientName
const stdioOptsByName = new Map<string, any>() const stdioOptsByName = new Map<string, any>()
let refreshAuthorizationCalls = 0 let refreshAuthorizationCalls = 0
let refreshAborted = false
function getOrCreateClientState(name?: string): MockClientState { function getOrCreateClientState(name?: string): MockClientState {
const key = name ?? "default" const key = name ?? "default"
@ -253,6 +254,7 @@ beforeEach(() => {
clientCreateCount = 0 clientCreateCount = 0
transportCloseCount = 0 transportCloseCount = 0
refreshAuthorizationCalls = 0 refreshAuthorizationCalls = 0
refreshAborted = false
}) })
// Import after mocks // Import after mocks
@ -307,14 +309,23 @@ it.live("McpOAuthProvider refreshes expired stored tokens", () =>
) )
it.instance( it.instance(
"remote connect bounds expired token refresh by mcp timeout", "remote connect cancels expired token refresh after mcp timeout",
() => () =>
Effect.acquireUseRelease( Effect.acquireUseRelease(
Effect.sync(() => { Effect.sync(() => {
const original = McpOAuthProvider.prototype.refreshTokensIfExpired const original = McpOAuthProvider.prototype.refreshTokensIfExpired
McpOAuthProvider.prototype.refreshTokensIfExpired = () => { McpOAuthProvider.prototype.refreshTokensIfExpired = (_fetchFn, signal) => {
refreshAuthorizationCalls++ refreshAuthorizationCalls++
return new Promise(() => {}) return new Promise((_, reject) => {
signal?.addEventListener(
"abort",
() => {
refreshAborted = true
reject(signal.reason)
},
{ once: true },
)
})
} }
return original return original
}), }),
@ -331,6 +342,7 @@ it.instance(
expect(statusName(result.status, "remote-timeout")).toBe("connected") expect(statusName(result.status, "remote-timeout")).toBe("connected")
expect(refreshAuthorizationCalls).toBe(1) expect(refreshAuthorizationCalls).toBe(1)
expect(refreshAborted).toBe(true)
}), }),
), ),
(original) => (original) =>