Compare commits
3 commits
dev
...
mcp-oauth-
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ecf550c88c | ||
|
|
ee5ee61db5 | ||
|
|
c5fff9fb68 |
3 changed files with 144 additions and 2 deletions
|
|
@ -8,6 +8,7 @@ import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/
|
|||
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js"
|
||||
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"
|
||||
import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"
|
||||
import { createFetchWithInit } from "@modelcontextprotocol/sdk/shared/transport.js"
|
||||
import {
|
||||
type LoggingMessageNotification,
|
||||
LoggingMessageNotificationSchema,
|
||||
|
|
@ -209,10 +210,11 @@ export const layer = Layer.effect(
|
|||
status: { status: "failed" as const, error: `Invalid MCP URL for "${key}"` },
|
||||
}
|
||||
}
|
||||
const connectTimeout = mcp.timeout ?? DEFAULT_TIMEOUT
|
||||
let authProvider: McpOAuthProvider | undefined
|
||||
|
||||
if (!oauthDisabled) {
|
||||
authProvider = new McpOAuthProvider(
|
||||
const provider = new McpOAuthProvider(
|
||||
key,
|
||||
mcp.url,
|
||||
{
|
||||
|
|
@ -227,6 +229,17 @@ export const layer = Layer.effect(
|
|||
},
|
||||
auth,
|
||||
)
|
||||
authProvider = provider
|
||||
const controller = new AbortController()
|
||||
yield* Effect.tryPromise(() =>
|
||||
withTimeout(
|
||||
provider.refreshTokensIfExpired(
|
||||
mcp.headers ? createFetchWithInit(fetch, { headers: mcp.headers }) : undefined,
|
||||
controller.signal,
|
||||
),
|
||||
connectTimeout,
|
||||
).finally(() => controller.abort()),
|
||||
).pipe(Effect.ignore)
|
||||
}
|
||||
|
||||
const transports: Array<{ name: string; transport: TransportWithAuth }> = [
|
||||
|
|
@ -246,7 +259,6 @@ export const layer = Layer.effect(
|
|||
},
|
||||
]
|
||||
|
||||
const connectTimeout = mcp.timeout ?? DEFAULT_TIMEOUT
|
||||
let lastStatus: Status | undefined
|
||||
|
||||
for (const { name, transport } of transports) {
|
||||
|
|
|
|||
|
|
@ -1,15 +1,18 @@
|
|||
import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js"
|
||||
import { discoverOAuthServerInfo, refreshAuthorization, selectResourceURL } from "@modelcontextprotocol/sdk/client/auth.js"
|
||||
import type {
|
||||
OAuthClientMetadata,
|
||||
OAuthTokens,
|
||||
OAuthClientInformation,
|
||||
OAuthClientInformationFull,
|
||||
} from "@modelcontextprotocol/sdk/shared/auth.js"
|
||||
import type { FetchLike } from "@modelcontextprotocol/sdk/shared/transport.js"
|
||||
import { Effect } from "effect"
|
||||
import { McpAuth } from "./auth"
|
||||
|
||||
const OAUTH_CALLBACK_PORT = 19876
|
||||
const OAUTH_CALLBACK_PATH = "/mcp/oauth/callback"
|
||||
const TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
|
||||
export interface McpOAuthConfig {
|
||||
clientId?: string
|
||||
|
|
@ -125,6 +128,31 @@ export class McpOAuthProvider implements OAuthClientProvider {
|
|||
)
|
||||
}
|
||||
|
||||
async refreshTokensIfExpired(fetchFn?: FetchLike, signal?: AbortSignal): Promise<boolean> {
|
||||
const entry = await Effect.runPromise(this.auth.getForUrl(this.mcpName, this.serverUrl))
|
||||
if (!entry?.tokens?.refreshToken) return false
|
||||
if (!entry.tokens.expiresAt) return false
|
||||
if (entry.tokens.expiresAt > Date.now() / 1000 + TOKEN_REFRESH_SKEW_SECONDS) return false
|
||||
|
||||
const clientInformation = await this.clientInformation()
|
||||
if (!clientInformation) return false
|
||||
|
||||
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 refreshAuthorization(info.authorizationServerUrl, {
|
||||
metadata: info.authorizationServerMetadata,
|
||||
clientInformation,
|
||||
refreshToken: entry.tokens.refreshToken,
|
||||
resource: await selectResourceURL(this.serverUrl, this, info.resourceMetadata),
|
||||
fetchFn: request,
|
||||
}),
|
||||
)
|
||||
return true
|
||||
}
|
||||
|
||||
async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
|
||||
await this.callbacks.onRedirect(authorizationUrl)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import { expect, mock, beforeEach } from "bun:test"
|
|||
import { ToolListChangedNotificationSchema } from "@modelcontextprotocol/sdk/types.js"
|
||||
import { Cause, Effect, Exit } from "effect"
|
||||
import type { MCP as MCPNS } from "../../src/mcp/index"
|
||||
import type { McpAuth } from "../../src/mcp/auth"
|
||||
import { testEffect } from "../lib/effect"
|
||||
import { TestInstance } from "../fixture/fixture"
|
||||
|
||||
|
|
@ -52,6 +53,8 @@ let clientCreateCount = 0
|
|||
let transportCloseCount = 0
|
||||
// Captures the opts passed to each MockStdioTransport, keyed by lastCreatedClientName
|
||||
const stdioOptsByName = new Map<string, any>()
|
||||
let refreshAuthorizationCalls = 0
|
||||
let refreshAborted = false
|
||||
|
||||
function getOrCreateClientState(name?: string): MockClientState {
|
||||
const key = name ?? "default"
|
||||
|
|
@ -141,6 +144,18 @@ void mock.module("@modelcontextprotocol/sdk/client/auth.js", () => ({
|
|||
super("Unauthorized")
|
||||
}
|
||||
},
|
||||
discoverOAuthServerInfo: async () => ({ authorizationServerUrl: "https://auth.example.com" }),
|
||||
selectResourceURL: async () => undefined,
|
||||
refreshAuthorization: async () => {
|
||||
refreshAuthorizationCalls++
|
||||
return {
|
||||
access_token: "new-access-token",
|
||||
token_type: "Bearer",
|
||||
refresh_token: "new-refresh-token",
|
||||
expires_in: 3600,
|
||||
scope: "read",
|
||||
}
|
||||
},
|
||||
}))
|
||||
|
||||
// Mock Client that delegates to per-name MockClientState
|
||||
|
|
@ -238,11 +253,14 @@ beforeEach(() => {
|
|||
connectError = "Mock transport cannot connect"
|
||||
clientCreateCount = 0
|
||||
transportCloseCount = 0
|
||||
refreshAuthorizationCalls = 0
|
||||
refreshAborted = false
|
||||
})
|
||||
|
||||
// Import after mocks
|
||||
const { MCP } = await import("../../src/mcp/index")
|
||||
const { McpOAuthCallback } = await import("../../src/mcp/oauth-callback")
|
||||
const { McpOAuthProvider } = await import("../../src/mcp/oauth-provider")
|
||||
|
||||
const it = testEffect(MCP.defaultLayer)
|
||||
|
||||
|
|
@ -251,6 +269,90 @@ function statusName(status: Record<string, MCPNS.Status> | MCPNS.Status, server:
|
|||
return status[server]?.status
|
||||
}
|
||||
|
||||
it.live("McpOAuthProvider refreshes expired stored tokens", () =>
|
||||
Effect.gen(function* () {
|
||||
const entry: McpAuth.Entry = {
|
||||
serverUrl: "https://mcp.example.com/mcp",
|
||||
clientInfo: { clientId: "client-id" },
|
||||
tokens: {
|
||||
accessToken: "old-access-token",
|
||||
refreshToken: "old-refresh-token",
|
||||
expiresAt: Date.now() / 1000 - 1,
|
||||
},
|
||||
}
|
||||
const auth: Pick<McpAuth.Interface, "getForUrl" | "updateTokens"> = {
|
||||
getForUrl: () => Effect.succeed(entry),
|
||||
updateTokens: (_mcpName: string, tokens: McpAuth.Tokens, serverUrl?: string) =>
|
||||
Effect.sync(() => {
|
||||
entry.tokens = tokens
|
||||
entry.serverUrl = serverUrl ?? entry.serverUrl
|
||||
}),
|
||||
}
|
||||
|
||||
const refreshed = yield* Effect.promise(() =>
|
||||
new McpOAuthProvider(
|
||||
"remote-server",
|
||||
"https://mcp.example.com/mcp",
|
||||
{},
|
||||
{ onRedirect: async () => {} },
|
||||
auth as McpAuth.Interface,
|
||||
).refreshTokensIfExpired(),
|
||||
)
|
||||
|
||||
expect(refreshed).toBe(true)
|
||||
expect(refreshAuthorizationCalls).toBe(1)
|
||||
if (!entry.tokens) throw new Error("tokens were not saved")
|
||||
expect(entry.tokens.accessToken).toBe("new-access-token")
|
||||
expect(entry.tokens.refreshToken).toBe("new-refresh-token")
|
||||
expect(entry.tokens.expiresAt).toBeGreaterThan(Date.now() / 1000)
|
||||
}),
|
||||
)
|
||||
|
||||
it.instance(
|
||||
"remote connect cancels expired token refresh after mcp timeout",
|
||||
() =>
|
||||
Effect.acquireUseRelease(
|
||||
Effect.sync(() => {
|
||||
const original = McpOAuthProvider.prototype.refreshTokensIfExpired
|
||||
McpOAuthProvider.prototype.refreshTokensIfExpired = (_fetchFn, signal) => {
|
||||
refreshAuthorizationCalls++
|
||||
return new Promise((_, reject) => {
|
||||
signal?.addEventListener(
|
||||
"abort",
|
||||
() => {
|
||||
refreshAborted = true
|
||||
reject(signal.reason)
|
||||
},
|
||||
{ once: true },
|
||||
)
|
||||
})
|
||||
}
|
||||
return original
|
||||
}),
|
||||
() =>
|
||||
MCP.Service.use((mcp: MCPNS.Interface) =>
|
||||
Effect.gen(function* () {
|
||||
lastCreatedClientName = "remote-timeout"
|
||||
|
||||
const result = yield* mcp.add("remote-timeout", {
|
||||
type: "remote",
|
||||
url: "https://mcp.example.com/mcp",
|
||||
timeout: 20,
|
||||
})
|
||||
|
||||
expect(statusName(result.status, "remote-timeout")).toBe("connected")
|
||||
expect(refreshAuthorizationCalls).toBe(1)
|
||||
expect(refreshAborted).toBe(true)
|
||||
}),
|
||||
),
|
||||
(original) =>
|
||||
Effect.sync(() => {
|
||||
McpOAuthProvider.prototype.refreshTokensIfExpired = original
|
||||
}),
|
||||
),
|
||||
{ config: { mcp: {} } },
|
||||
)
|
||||
|
||||
it.instance(
|
||||
"local mcp cwd resolves relative paths against instance directory",
|
||||
() =>
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue