fix(mcp): persist immediate oauth connections (#22376)
This commit is contained in:
parent
9a5178e4ac
commit
4626458175
2 changed files with 108 additions and 14 deletions
|
|
@ -204,6 +204,12 @@ export namespace MCP {
|
||||||
defs?: MCPToolDef[]
|
defs?: MCPToolDef[]
|
||||||
}
|
}
|
||||||
|
|
||||||
|
interface AuthResult {
|
||||||
|
authorizationUrl: string
|
||||||
|
oauthState: string
|
||||||
|
client?: MCPClient
|
||||||
|
}
|
||||||
|
|
||||||
// --- Effect Service ---
|
// --- Effect Service ---
|
||||||
|
|
||||||
interface State {
|
interface State {
|
||||||
|
|
@ -552,6 +558,21 @@ export namespace MCP {
|
||||||
return Effect.tryPromise(() => client.close()).pipe(Effect.ignore)
|
return Effect.tryPromise(() => client.close()).pipe(Effect.ignore)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const storeClient = Effect.fnUntraced(function* (
|
||||||
|
s: State,
|
||||||
|
name: string,
|
||||||
|
client: MCPClient,
|
||||||
|
listed: MCPToolDef[],
|
||||||
|
timeout?: number,
|
||||||
|
) {
|
||||||
|
yield* closeClient(s, name)
|
||||||
|
s.status[name] = { status: "connected" }
|
||||||
|
s.clients[name] = client
|
||||||
|
s.defs[name] = listed
|
||||||
|
watch(s, name, client, timeout)
|
||||||
|
return s.status[name]
|
||||||
|
})
|
||||||
|
|
||||||
const status = Effect.fn("MCP.status")(function* () {
|
const status = Effect.fn("MCP.status")(function* () {
|
||||||
const s = yield* InstanceState.get(state)
|
const s = yield* InstanceState.get(state)
|
||||||
|
|
||||||
|
|
@ -583,11 +604,7 @@ export namespace MCP {
|
||||||
return result.status
|
return result.status
|
||||||
}
|
}
|
||||||
|
|
||||||
yield* closeClient(s, name)
|
return yield* storeClient(s, name, result.mcpClient, result.defs!, mcp.timeout)
|
||||||
s.clients[name] = result.mcpClient
|
|
||||||
s.defs[name] = result.defs!
|
|
||||||
watch(s, name, result.mcpClient, mcp.timeout)
|
|
||||||
return result.status
|
|
||||||
})
|
})
|
||||||
|
|
||||||
const add = Effect.fn("MCP.add")(function* (name: string, mcp: Config.Mcp) {
|
const add = Effect.fn("MCP.add")(function* (name: string, mcp: Config.Mcp) {
|
||||||
|
|
@ -753,14 +770,16 @@ export namespace MCP {
|
||||||
return yield* Effect.tryPromise({
|
return yield* Effect.tryPromise({
|
||||||
try: () => {
|
try: () => {
|
||||||
const client = new Client({ name: "opencode", version: Installation.VERSION })
|
const client = new Client({ name: "opencode", version: Installation.VERSION })
|
||||||
return client.connect(transport).then(() => ({ authorizationUrl: "", oauthState }))
|
return client
|
||||||
|
.connect(transport)
|
||||||
|
.then(() => ({ authorizationUrl: "", oauthState, client }) satisfies AuthResult)
|
||||||
},
|
},
|
||||||
catch: (error) => error,
|
catch: (error) => error,
|
||||||
}).pipe(
|
}).pipe(
|
||||||
Effect.catch((error) => {
|
Effect.catch((error) => {
|
||||||
if (error instanceof UnauthorizedError && capturedUrl) {
|
if (error instanceof UnauthorizedError && capturedUrl) {
|
||||||
pendingOAuthTransports.set(mcpName, transport)
|
pendingOAuthTransports.set(mcpName, transport)
|
||||||
return Effect.succeed({ authorizationUrl: capturedUrl.toString(), oauthState })
|
return Effect.succeed({ authorizationUrl: capturedUrl.toString(), oauthState } satisfies AuthResult)
|
||||||
}
|
}
|
||||||
return Effect.die(error)
|
return Effect.die(error)
|
||||||
}),
|
}),
|
||||||
|
|
@ -768,14 +787,31 @@ export namespace MCP {
|
||||||
})
|
})
|
||||||
|
|
||||||
const authenticate = Effect.fn("MCP.authenticate")(function* (mcpName: string) {
|
const authenticate = Effect.fn("MCP.authenticate")(function* (mcpName: string) {
|
||||||
const { authorizationUrl, oauthState } = yield* startAuth(mcpName)
|
const result = yield* startAuth(mcpName)
|
||||||
if (!authorizationUrl) return { status: "connected" } as Status
|
if (!result.authorizationUrl) {
|
||||||
|
const client = "client" in result ? result.client : undefined
|
||||||
|
const mcpConfig = yield* getMcpConfig(mcpName)
|
||||||
|
if (!mcpConfig) {
|
||||||
|
yield* Effect.tryPromise(() => client?.close() ?? Promise.resolve()).pipe(Effect.ignore)
|
||||||
|
return { status: "failed", error: "MCP config not found after auth" } as Status
|
||||||
|
}
|
||||||
|
|
||||||
log.info("opening browser for oauth", { mcpName, url: authorizationUrl, state: oauthState })
|
const listed = client ? yield* defs(mcpName, client, mcpConfig.timeout) : undefined
|
||||||
|
if (!client || !listed) {
|
||||||
|
yield* Effect.tryPromise(() => client?.close() ?? Promise.resolve()).pipe(Effect.ignore)
|
||||||
|
return { status: "failed", error: "Failed to get tools" } as Status
|
||||||
|
}
|
||||||
|
|
||||||
const callbackPromise = McpOAuthCallback.waitForCallback(oauthState, mcpName)
|
const s = yield* InstanceState.get(state)
|
||||||
|
yield* auth.clearOAuthState(mcpName)
|
||||||
|
return yield* storeClient(s, mcpName, client, listed, mcpConfig.timeout)
|
||||||
|
}
|
||||||
|
|
||||||
yield* Effect.tryPromise(() => open(authorizationUrl)).pipe(
|
log.info("opening browser for oauth", { mcpName, url: result.authorizationUrl, state: result.oauthState })
|
||||||
|
|
||||||
|
const callbackPromise = McpOAuthCallback.waitForCallback(result.oauthState, mcpName)
|
||||||
|
|
||||||
|
yield* Effect.tryPromise(() => open(result.authorizationUrl)).pipe(
|
||||||
Effect.flatMap((subprocess) =>
|
Effect.flatMap((subprocess) =>
|
||||||
Effect.callback<void, Error>((resume) => {
|
Effect.callback<void, Error>((resume) => {
|
||||||
const timer = setTimeout(() => resume(Effect.void), 500)
|
const timer = setTimeout(() => resume(Effect.void), 500)
|
||||||
|
|
@ -793,14 +829,14 @@ export namespace MCP {
|
||||||
),
|
),
|
||||||
Effect.catch(() => {
|
Effect.catch(() => {
|
||||||
log.warn("failed to open browser, user must open URL manually", { mcpName })
|
log.warn("failed to open browser, user must open URL manually", { mcpName })
|
||||||
return bus.publish(BrowserOpenFailed, { mcpName, url: authorizationUrl }).pipe(Effect.ignore)
|
return bus.publish(BrowserOpenFailed, { mcpName, url: result.authorizationUrl }).pipe(Effect.ignore)
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
|
|
||||||
const code = yield* Effect.promise(() => callbackPromise)
|
const code = yield* Effect.promise(() => callbackPromise)
|
||||||
|
|
||||||
const storedState = yield* auth.getOAuthState(mcpName)
|
const storedState = yield* auth.getOAuthState(mcpName)
|
||||||
if (storedState !== oauthState) {
|
if (storedState !== result.oauthState) {
|
||||||
yield* auth.clearOAuthState(mcpName)
|
yield* auth.clearOAuthState(mcpName)
|
||||||
throw new Error("OAuth state mismatch - potential CSRF attack")
|
throw new Error("OAuth state mismatch - potential CSRF attack")
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -20,6 +20,7 @@ const transportCalls: Array<{
|
||||||
// Controls whether the mock transport simulates a 401 that triggers the SDK
|
// Controls whether the mock transport simulates a 401 that triggers the SDK
|
||||||
// auth flow (which calls provider.state()) or a simple UnauthorizedError.
|
// auth flow (which calls provider.state()) or a simple UnauthorizedError.
|
||||||
let simulateAuthFlow = true
|
let simulateAuthFlow = true
|
||||||
|
let connectSucceedsImmediately = false
|
||||||
|
|
||||||
// Mock the transport constructors to simulate OAuth auto-auth on 401
|
// Mock the transport constructors to simulate OAuth auto-auth on 401
|
||||||
mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
|
mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
|
||||||
|
|
@ -40,6 +41,8 @@ mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
async start() {
|
async start() {
|
||||||
|
if (connectSucceedsImmediately) return
|
||||||
|
|
||||||
// Simulate what the real SDK transport does on 401:
|
// Simulate what the real SDK transport does on 401:
|
||||||
// It calls auth() which eventually calls provider.state(), then
|
// It calls auth() which eventually calls provider.state(), then
|
||||||
// provider.redirectToAuthorization(), then throws UnauthorizedError.
|
// provider.redirectToAuthorization(), then throws UnauthorizedError.
|
||||||
|
|
@ -85,6 +88,14 @@ mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({
|
||||||
async connect(transport: { start: () => Promise<void> }) {
|
async connect(transport: { start: () => Promise<void> }) {
|
||||||
await transport.start()
|
await transport.start()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
setNotificationHandler() {}
|
||||||
|
|
||||||
|
async listTools() {
|
||||||
|
return { tools: [{ name: "test_tool", inputSchema: { type: "object", properties: {} } }] }
|
||||||
|
}
|
||||||
|
|
||||||
|
async close() {}
|
||||||
},
|
},
|
||||||
}))
|
}))
|
||||||
|
|
||||||
|
|
@ -96,6 +107,7 @@ mock.module("@modelcontextprotocol/sdk/client/auth.js", () => ({
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
transportCalls.length = 0
|
transportCalls.length = 0
|
||||||
simulateAuthFlow = true
|
simulateAuthFlow = true
|
||||||
|
connectSucceedsImmediately = false
|
||||||
})
|
})
|
||||||
|
|
||||||
// Import modules after mocking
|
// Import modules after mocking
|
||||||
|
|
@ -222,3 +234,49 @@ test("state() returns existing state when one is saved", async () => {
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
|
test("authenticate() stores a connected client when auth completes without redirect", async () => {
|
||||||
|
await using tmp = await tmpdir({
|
||||||
|
init: async (dir) => {
|
||||||
|
await Bun.write(
|
||||||
|
`${dir}/opencode.json`,
|
||||||
|
JSON.stringify({
|
||||||
|
$schema: "https://opencode.ai/config.json",
|
||||||
|
mcp: {
|
||||||
|
"test-oauth-connect": {
|
||||||
|
type: "remote",
|
||||||
|
url: "https://example.com/mcp",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
await Instance.provide({
|
||||||
|
directory: tmp.path,
|
||||||
|
fn: async () => {
|
||||||
|
await Effect.runPromise(
|
||||||
|
MCP.Service.use((mcp) =>
|
||||||
|
Effect.gen(function* () {
|
||||||
|
const added = yield* mcp.add("test-oauth-connect", {
|
||||||
|
type: "remote",
|
||||||
|
url: "https://example.com/mcp",
|
||||||
|
})
|
||||||
|
const before = added.status as Record<string, { status: string; error?: string }>
|
||||||
|
expect(before["test-oauth-connect"]?.status).toBe("needs_auth")
|
||||||
|
|
||||||
|
simulateAuthFlow = false
|
||||||
|
connectSucceedsImmediately = true
|
||||||
|
|
||||||
|
const result = yield* mcp.authenticate("test-oauth-connect")
|
||||||
|
expect(result.status).toBe("connected")
|
||||||
|
|
||||||
|
const after = yield* mcp.status()
|
||||||
|
expect(after["test-oauth-connect"]?.status).toBe("connected")
|
||||||
|
}),
|
||||||
|
).pipe(Effect.provide(MCP.defaultLayer)),
|
||||||
|
)
|
||||||
|
},
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue