Compare commits

...

1 Commits

Author SHA1 Message Date
Aiden Cline 63a5b1c6dc fix(mcp): recover expired legacy sessions 2026-07-28 00:09:24 -05:00
4 changed files with 156 additions and 15 deletions
+20 -11
View File
@@ -1,4 +1,4 @@
import { Client, type CallToolResult, type Tool as MCPToolDef } from "@modelcontextprotocol/client"
import { Client, SdkHttpError, type CallToolResult, type Tool as MCPToolDef } from "@modelcontextprotocol/client"
import { dynamicTool, jsonSchema, type JSONSchema7, type Tool } from "ai"
import { Effect } from "effect"
@@ -8,6 +8,7 @@ export interface McpTool {
readonly def: MCPToolDef
readonly client: Client
readonly timeout?: number
readonly recover?: () => Promise<Client | undefined>
}
export async function callTool(
@@ -15,16 +16,24 @@ export async function callTool(
args: Record<string, unknown>,
signal?: AbortSignal,
): Promise<CallToolResult> {
const result = await tool.client.callTool(
{ name: tool.def.name, arguments: args },
{
resetTimeoutOnProgress: true,
signal,
timeout: tool.timeout,
// The MCP SDK only sends a progress token when this hook is present, enabling timeout resets.
onprogress: () => {},
},
)
const invoke = (client: Client) =>
client.callTool(
{ name: tool.def.name, arguments: args },
{
resetTimeoutOnProgress: true,
signal,
timeout: tool.timeout,
// The MCP SDK only sends a progress token when this hook is present, enabling timeout resets.
onprogress: () => {},
},
)
const result = await invoke(tool.client).catch(async (error) => {
if (!tool.recover || !(error instanceof SdkHttpError) || error.status !== 404) throw error
const client = await tool.recover()
if (!client) throw error
signal?.throwIfAborted()
return invoke(client)
})
if (result.isError)
throw new Error(
result.content
+49 -1
View File
@@ -593,8 +593,13 @@ const layer = Layer.effect(
listed: MCPToolDef[],
instructions: string | undefined,
timeout?: number,
expected?: MCPClient,
) {
const bridge = yield* EffectBridge.make()
if (expected && s.clients[name] !== expected) {
yield* Effect.tryPromise(() => client.close()).pipe(Effect.ignore)
return s.status[name] ?? { status: "disabled" }
}
const previous = s.clients[name]
s.status[name] = { status: "connected" }
s.clients[name] = client
@@ -656,6 +661,37 @@ const layer = Layer.effect(
return yield* storeClient(s, name, result.mcpClient, result.defs!, result.instructions, mcp.timeout)
})
const reconnectStale = Effect.fn("MCP.reconnectStale")(function* (name: string, stale: MCPClient) {
const s = yield* InstanceState.get(state)
if (s.clients[name] !== stale) return undefined
const mcp = yield* getMcpConfig(name)
if (!mcp) return undefined
yield* Effect.logWarning("reconnecting MCP server after stale session", { server: name })
const result = yield* create(name, { ...mcp, enabled: true })
if (s.clients[name] !== stale) {
if (result.mcpClient) yield* Effect.tryPromise(() => result.mcpClient!.close()).pipe(Effect.ignore)
return undefined
}
if (!result.mcpClient) {
s.status[name] = result.status
yield* closeClient(s, name)
return undefined
}
yield* storeClient(
s,
name,
result.mcpClient,
result.defs!,
result.instructions,
mcp.timeout,
stale,
)
return s.clients[name]
})
const recoveries = new WeakMap<MCPClient, Promise<MCPClient | undefined>>()
const add = Effect.fn("MCP.add")(function* (name: string, mcp: ConfigMCPV1.Info) {
const s = yield* InstanceState.get(state)
s.config[name] = mcp
@@ -684,6 +720,7 @@ const layer = Layer.effect(
const tools = Effect.fn("MCP.tools")(function* () {
const result: Record<string, McpTool> = {}
const s = yield* InstanceState.get(state)
const bridge = yield* EffectBridge.make()
const cfg = yield* cfgSvc.get()
const config = cfg.mcp ?? {}
@@ -698,8 +735,19 @@ const layer = Layer.effect(
continue
}
const timeout = requestTimeout(s, clientName, mcpConfig, defaultTimeout)
const transport = client.transport
const recover =
transport instanceof StreamableHTTPClientTransport && transport.sessionId && client.getProtocolEra() === "legacy"
? () => {
const current = recoveries.get(client)
if (current) return current
const created = bridge.promise(reconnectStale(clientName, client))
recoveries.set(client, created)
return created
}
: undefined
for (const def of listed) {
result[McpCatalog.toolName(clientName, def.name)] = { def, client, timeout }
result[McpCatalog.toolName(clientName, def.name)] = { def, client, timeout, recover }
}
}
return result
+46 -1
View File
@@ -1,5 +1,5 @@
import { describe, expect, test } from "bun:test"
import { Client, InMemoryTransport } from "@modelcontextprotocol/client"
import { Client, InMemoryTransport, SdkErrorCode, SdkHttpError } from "@modelcontextprotocol/client"
import { Server } from "@modelcontextprotocol/server"
import { McpCatalog } from "@/mcp/catalog"
import { Effect } from "effect"
@@ -86,6 +86,51 @@ describe("McpCatalog.callTool", () => {
await expect(McpCatalog.callTool({ def: mcpTool(), client }, {})).rejects.toThrow("first\n\nsecond")
})
test("retries a stale session only once", async () => {
const error = new SdkHttpError(SdkErrorCode.SendFailed, "expired", { status: 404 })
const client = { callTool: () => Promise.reject(error) } as unknown as Client
let recoveries = 0
await expect(
McpCatalog.callTool(
{
def: mcpTool(),
client,
recover: async () => {
recoveries++
return client
},
},
{},
),
).rejects.toBe(error)
expect(recoveries).toBe(1)
})
test("does not replay after cancellation during recovery", async () => {
const error = new SdkHttpError(SdkErrorCode.SendFailed, "expired", { status: 404 })
const controller = new AbortController()
let calls = 0
const stale = { callTool: () => Promise.reject(error) } as unknown as Client
const fresh = { callTool: async () => ({ content: [], calls: ++calls }) } as unknown as Client
await expect(
McpCatalog.callTool(
{
def: mcpTool(),
client: stale,
recover: async () => {
controller.abort()
return fresh
},
},
{},
controller.signal,
),
).rejects.toBeInstanceOf(DOMException)
expect(calls).toBe(0)
})
})
test("preserves output schema validation across paginated tool discovery", async () => {
+41 -2
View File
@@ -11,6 +11,7 @@ import { LayerNode } from "@opencode-ai/core/effect/layer-node"
import { Cause, Effect, Exit } from "effect"
import type { MCP as MCPNS } from "../../src/mcp/index"
import { MCP } from "../../src/mcp/index"
import { McpCatalog } from "../../src/mcp/catalog"
import { McpOAuthCallback } from "../../src/mcp/oauth-callback"
import { TestInstance } from "../fixture/fixture"
import { pollWithTimeout, testEffect } from "../lib/effect"
@@ -36,7 +37,12 @@ interface LifecycleServerState {
aborted: number
}
function lifecycleServer(input?: { capabilities?: ServerCapabilities; instructions?: string; requestRoots?: boolean }) {
function lifecycleServer(input?: {
capabilities?: ServerCapabilities
instructions?: string
requestRoots?: boolean
legacy?: boolean
}) {
const capabilities = input?.capabilities ?? { tools: {}, prompts: {}, resources: {} }
return Effect.acquireRelease(
Effect.promise(async () => {
@@ -52,7 +58,11 @@ function lifecycleServer(input?: { capabilities?: ServerCapabilities; instructio
const makeProtocol = async () => {
const protocol = new Server(
{ name: "mcp-lifecycle", version: "1.0.0" },
{ capabilities, instructions: input?.instructions },
{
capabilities,
instructions: input?.instructions,
supportedProtocolVersions: input?.legacy ? ["2025-11-25"] : undefined,
},
)
const transport = new WebStandardStreamableHTTPServerTransport({
sessionIdGenerator: () => crypto.randomUUID(),
@@ -65,6 +75,9 @@ function lifecycleServer(input?: { capabilities?: ServerCapabilities; instructio
const page = state.toolPages?.[request.params?.cursor ?? "initial"]
return Promise.resolve({ tools: page?.items ?? state.tools, nextCursor: page?.nextCursor })
})
protocol.setRequestHandler("tools/call", (request) =>
Promise.resolve({ content: [{ type: "text", text: request.params.name }] }),
)
}
if (capabilities.prompts) {
protocol.setRequestHandler("prompts/list", (request) => {
@@ -114,6 +127,10 @@ function lifecycleServer(input?: { capabilities?: ServerCapabilities; instructio
fetch(request) {
state.requests.push(request.method)
request.signal.addEventListener("abort", () => state.aborted++)
const session = request.headers.get("mcp-session-id")
if (request.method === "POST" && session && session !== current.transport.sessionId) {
return new Response("Session not found", { status: 404 })
}
return current.transport.handleRequest(request)
},
})
@@ -322,6 +339,28 @@ it.instance("disconnect removes protocol data and reconnect establishes a new se
}),
)
it.instance("recovers a legacy HTTP session after the server discards it", () =>
Effect.gen(function* () {
const server = yield* lifecycleServer({ legacy: true })
const mcp = yield* MCP.Service
yield* mcp.add("session-server", remote(server.url))
const tool = (yield* mcp.tools())["session-server_test_tool"]
const client = tool?.client
expect(tool).toBeDefined()
yield* Effect.promise(server.restart)
const results = yield* Effect.promise(() =>
Promise.all([McpCatalog.callTool(tool!, {}), McpCatalog.callTool(tool!, {})]),
)
expect(results.map((result) => result.content)).toEqual([
[{ type: "text", text: "test_tool" }],
[{ type: "text", text: "test_tool" }],
])
expect((yield* mcp.clients())["session-server"]).not.toBe(client)
}),
)
it.instance("add() closes the old protocol session when replacing a server", () =>
Effect.gen(function* () {
const first = yield* lifecycleServer()