fix(mcp): cancel timed out oauth refresh
This commit is contained in:
parent
ee5ee61db5
commit
ecf550c88c
3 changed files with 27 additions and 8 deletions
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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) =>
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue