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 { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js"
|
||||||
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"
|
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"
|
||||||
import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"
|
import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js"
|
||||||
|
import { createFetchWithInit } from "@modelcontextprotocol/sdk/shared/transport.js"
|
||||||
import {
|
import {
|
||||||
type LoggingMessageNotification,
|
type LoggingMessageNotification,
|
||||||
LoggingMessageNotificationSchema,
|
LoggingMessageNotificationSchema,
|
||||||
|
|
@ -209,10 +210,11 @@ export const layer = Layer.effect(
|
||||||
status: { status: "failed" as const, error: `Invalid MCP URL for "${key}"` },
|
status: { status: "failed" as const, error: `Invalid MCP URL for "${key}"` },
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
const connectTimeout = mcp.timeout ?? DEFAULT_TIMEOUT
|
||||||
let authProvider: McpOAuthProvider | undefined
|
let authProvider: McpOAuthProvider | undefined
|
||||||
|
|
||||||
if (!oauthDisabled) {
|
if (!oauthDisabled) {
|
||||||
authProvider = new McpOAuthProvider(
|
const provider = new McpOAuthProvider(
|
||||||
key,
|
key,
|
||||||
mcp.url,
|
mcp.url,
|
||||||
{
|
{
|
||||||
|
|
@ -227,6 +229,17 @@ export const layer = Layer.effect(
|
||||||
},
|
},
|
||||||
auth,
|
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 }> = [
|
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
|
let lastStatus: Status | undefined
|
||||||
|
|
||||||
for (const { name, transport } of transports) {
|
for (const { name, transport } of transports) {
|
||||||
|
|
|
||||||
|
|
@ -1,15 +1,18 @@
|
||||||
import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js"
|
import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js"
|
||||||
|
import { discoverOAuthServerInfo, refreshAuthorization, selectResourceURL } from "@modelcontextprotocol/sdk/client/auth.js"
|
||||||
import type {
|
import type {
|
||||||
OAuthClientMetadata,
|
OAuthClientMetadata,
|
||||||
OAuthTokens,
|
OAuthTokens,
|
||||||
OAuthClientInformation,
|
OAuthClientInformation,
|
||||||
OAuthClientInformationFull,
|
OAuthClientInformationFull,
|
||||||
} from "@modelcontextprotocol/sdk/shared/auth.js"
|
} from "@modelcontextprotocol/sdk/shared/auth.js"
|
||||||
|
import type { FetchLike } from "@modelcontextprotocol/sdk/shared/transport.js"
|
||||||
import { Effect } from "effect"
|
import { Effect } from "effect"
|
||||||
import { McpAuth } from "./auth"
|
import { McpAuth } from "./auth"
|
||||||
|
|
||||||
const OAUTH_CALLBACK_PORT = 19876
|
const OAUTH_CALLBACK_PORT = 19876
|
||||||
const OAUTH_CALLBACK_PATH = "/mcp/oauth/callback"
|
const OAUTH_CALLBACK_PATH = "/mcp/oauth/callback"
|
||||||
|
const TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||||
|
|
||||||
export interface McpOAuthConfig {
|
export interface McpOAuthConfig {
|
||||||
clientId?: string
|
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> {
|
async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
|
||||||
await this.callbacks.onRedirect(authorizationUrl)
|
await this.callbacks.onRedirect(authorizationUrl)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ import { expect, mock, beforeEach } from "bun:test"
|
||||||
import { ToolListChangedNotificationSchema } from "@modelcontextprotocol/sdk/types.js"
|
import { ToolListChangedNotificationSchema } from "@modelcontextprotocol/sdk/types.js"
|
||||||
import { Cause, Effect, Exit } from "effect"
|
import { Cause, Effect, Exit } from "effect"
|
||||||
import type { MCP as MCPNS } from "../../src/mcp/index"
|
import type { MCP as MCPNS } from "../../src/mcp/index"
|
||||||
|
import type { McpAuth } from "../../src/mcp/auth"
|
||||||
import { testEffect } from "../lib/effect"
|
import { testEffect } from "../lib/effect"
|
||||||
import { TestInstance } from "../fixture/fixture"
|
import { TestInstance } from "../fixture/fixture"
|
||||||
|
|
||||||
|
|
@ -52,6 +53,8 @@ let clientCreateCount = 0
|
||||||
let transportCloseCount = 0
|
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 refreshAborted = false
|
||||||
|
|
||||||
function getOrCreateClientState(name?: string): MockClientState {
|
function getOrCreateClientState(name?: string): MockClientState {
|
||||||
const key = name ?? "default"
|
const key = name ?? "default"
|
||||||
|
|
@ -141,6 +144,18 @@ void mock.module("@modelcontextprotocol/sdk/client/auth.js", () => ({
|
||||||
super("Unauthorized")
|
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
|
// Mock Client that delegates to per-name MockClientState
|
||||||
|
|
@ -238,11 +253,14 @@ beforeEach(() => {
|
||||||
connectError = "Mock transport cannot connect"
|
connectError = "Mock transport cannot connect"
|
||||||
clientCreateCount = 0
|
clientCreateCount = 0
|
||||||
transportCloseCount = 0
|
transportCloseCount = 0
|
||||||
|
refreshAuthorizationCalls = 0
|
||||||
|
refreshAborted = false
|
||||||
})
|
})
|
||||||
|
|
||||||
// Import after mocks
|
// Import after mocks
|
||||||
const { MCP } = await import("../../src/mcp/index")
|
const { MCP } = await import("../../src/mcp/index")
|
||||||
const { McpOAuthCallback } = await import("../../src/mcp/oauth-callback")
|
const { McpOAuthCallback } = await import("../../src/mcp/oauth-callback")
|
||||||
|
const { McpOAuthProvider } = await import("../../src/mcp/oauth-provider")
|
||||||
|
|
||||||
const it = testEffect(MCP.defaultLayer)
|
const it = testEffect(MCP.defaultLayer)
|
||||||
|
|
||||||
|
|
@ -251,6 +269,90 @@ function statusName(status: Record<string, MCPNS.Status> | MCPNS.Status, server:
|
||||||
return status[server]?.status
|
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(
|
it.instance(
|
||||||
"local mcp cwd resolves relative paths against instance directory",
|
"local mcp cwd resolves relative paths against instance directory",
|
||||||
() =>
|
() =>
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue