fix(mcp): recover expired SDK sessions (#39265)

This commit is contained in:
Aiden Cline 2026-07-28 00:35:19 -05:00 committed by GitHub
commit 484f00ebf4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 247 additions and 5 deletions

View file

@ -1,8 +1,11 @@
import { Client, LATEST_PROTOCOL_VERSION, StreamableHTTPClientTransport } from "@modelcontextprotocol/client"
const posts: Array<{ method: string; session: string | null }> = []
const concurrent = process.env.MCP_RECOVERY_CONCURRENT === "1"
let initializeCount = 0
let pingCount = 0
let replacementStarted!: () => void
const replacement = new Promise<void>((resolve) => (replacementStarted = resolve))
const server = Bun.serve({
port: 0,
async fetch(request) {
@ -15,6 +18,7 @@ const server = Bun.serve({
if (message.method === "initialize") {
initializeCount++
if (initializeCount === 2) replacementStarted()
return Response.json(
{
jsonrpc: "2.0",
@ -32,7 +36,8 @@ const server = Bun.serve({
if (message.method === "notifications/initialized") return new Response(null, { status: 202 })
pingCount++
if (pingCount === 1) return new Response("Session not found", { status: 404 })
if (concurrent && pingCount === 2) await replacement
if (pingCount <= (concurrent ? 2 : 1)) return new Response("Session not found", { status: 404 })
return Response.json({ jsonrpc: "2.0", id: message.id, result: {} })
},
})
@ -40,7 +45,8 @@ const client = new Client({ name: "test", version: "1" })
try {
await client.connect(new StreamableHTTPClientTransport(server.url))
await client.ping()
if (concurrent) await Promise.all([client.ping(), client.ping()])
else await client.ping()
process.stdout.write(JSON.stringify(posts))
} finally {
await client.close()

View file

@ -2,7 +2,7 @@ import path from "node:path"
import { describe, expect, test } from "bun:test"
describe("mcp session recovery", () => {
test.skip("reinitializes and retries once after a session-bound POST returns 404", async () => {
test("reinitializes and retries once after a session-bound POST returns 404", async () => {
const child = Bun.spawn([process.execPath, path.join(import.meta.dir, "../fixture/mcp-session-recovery.ts")], {
cwd: path.join(import.meta.dir, "../.."),
stdout: "pipe",
@ -24,4 +24,24 @@ describe("mcp session recovery", () => {
{ method: "ping", session: "replacement" },
])
})
test("retries a concurrent stale response after recovery completes", async () => {
const child = Bun.spawn([process.execPath, path.join(import.meta.dir, "../fixture/mcp-session-recovery.ts")], {
cwd: path.join(import.meta.dir, "../.."),
env: { ...process.env, MCP_RECOVERY_CONCURRENT: "1" },
stdout: "pipe",
stderr: "pipe",
})
const [code, stdout, stderr] = await Promise.all([
child.exited,
Bun.readableStreamToText(child.stdout),
Bun.readableStreamToText(child.stderr),
])
expect(code, stderr).toBe(0)
const posts = JSON.parse(stdout) as Array<{ method: string; session: string | null }>
expect(posts.filter((post) => post.method === "initialize").map((post) => post.session)).toEqual([null, null])
expect(posts.filter((post) => post.method === "ping" && post.session === "expired")).toHaveLength(2)
expect(posts.filter((post) => post.method === "ping" && post.session === "replacement")).toHaveLength(2)
})
})