fix(mcp): recover expired SDK sessions (#39265)
This commit is contained in:
parent
c3be6c4965
commit
484f00ebf4
5 changed files with 247 additions and 5 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
})
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue