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
parent c3be6c4965
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

@ -1071,8 +1071,9 @@
],
"patchedDependencies": {
"@pierre/trees@1.0.0-beta.4": "patches/@pierre%2Ftrees@1.0.0-beta.4.patch",
"@tanstack/virtual-core@3.17.3": "patches/@tanstack%2Fvirtual-core@3.17.3.patch",
"@modelcontextprotocol/client@2.0.0-beta.5": "patches/@modelcontextprotocol%2Fclient@2.0.0-beta.5.patch",
"@ai-sdk/xai@3.0.102": "patches/@ai-sdk%2Fxai@3.0.102.patch",
"@tanstack/virtual-core@3.17.3": "patches/@tanstack%2Fvirtual-core@3.17.3.patch",
"gcp-metadata@8.1.2": "patches/gcp-metadata@8.1.2.patch",
"@standard-community/standard-openapi@0.2.9": "patches/@standard-community%2Fstandard-openapi@0.2.9.patch",
"effect@4.0.0-beta.83": "patches/effect@4.0.0-beta.83.patch",

View file

@ -155,6 +155,7 @@
"@ai-sdk/google@3.0.73": "patches/@ai-sdk%2Fgoogle@3.0.73.patch",
"@pierre/trees@1.0.0-beta.4": "patches/@pierre%2Ftrees@1.0.0-beta.4.patch",
"effect@4.0.0-beta.83": "patches/effect@4.0.0-beta.83.patch",
"@tanstack/virtual-core@3.17.3": "patches/@tanstack%2Fvirtual-core@3.17.3.patch"
"@tanstack/virtual-core@3.17.3": "patches/@tanstack%2Fvirtual-core@3.17.3.patch",
"@modelcontextprotocol/client@2.0.0-beta.5": "patches/@modelcontextprotocol%2Fclient@2.0.0-beta.5.patch"
}
}

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)
})
})

View file

@ -0,0 +1,214 @@
diff --git a/dist/index.cjs b/dist/index.cjs
index 1c43bac25a1037416fdf2ddfb4534ba0897a2e69..7e2108326f368ccbc48897195d57e370a4553855 100644
--- a/dist/index.cjs
+++ b/dist/index.cjs
@@ -3154,6 +3154,7 @@ var Client = class extends require_src.Protocol {
*/
async _connectPlainLegacy(transport, options) {
await super.connect(transport);
+ transport.onsessionexpired = () => this._legacyHandshake(transport, options);
if (transport.sessionId !== void 0) {
const negotiatedProtocolVersion = this._negotiatedProtocolVersion;
if (negotiatedProtocolVersion !== void 0) transport.setProtocolVersion?.(negotiatedProtocolVersion);
@@ -3170,6 +3171,7 @@ var Client = class extends require_src.Protocol {
* the handshake; its completion sets the negotiated (legacy) version.
*/
async _legacyHandshake(transport, options) {
+ transport.onsessionexpired = () => this._legacyHandshake(transport, options);
const legacyVersions = require_src.legacyProtocolVersions(this._supportedProtocolVersions);
try {
const offeredVersion = legacyVersions[0];
@@ -3208,6 +3210,7 @@ var Client = class extends require_src.Protocol {
await super.connect(transport);
const negotiatedProtocolVersion = this._negotiatedProtocolVersion;
if (negotiatedProtocolVersion !== void 0 && transport.setProtocolVersion) transport.setProtocolVersion(negotiatedProtocolVersion);
+ if (negotiatedProtocolVersion !== void 0 && !require_src.isModernProtocolVersion(negotiatedProtocolVersion)) transport.onsessionexpired = () => this._legacyHandshake(transport, options);
return;
}
this._resetConnectionState();
@@ -5211,10 +5214,32 @@ var StreamableHTTPClientTransport = class {
}
}
async send(message, options) {
- return this._send(message, options, false);
+ return this._send(message, options, false, 0, false);
+ }
+ async _recoverSession(expiredSessionId) {
+ if (this._sessionRecovery) return this._sessionRecovery;
+ if (!this.onsessionexpired) return false;
+ if (this._sessionId !== expiredSessionId) return true;
+ this._sessionId = void 0;
+ this._sessionRecovery = Promise.resolve().then(() => this.onsessionexpired()).then(() => true);
+ try {
+ return await this._sessionRecovery;
+ } catch (error) {
+ this._sessionId = void 0;
+ await this.close();
+ throw error;
+ } finally {
+ this._sessionRecovery = void 0;
+ }
}
- async _send(message, options, isAuthRetry, stepUpRetries = 0) {
+ async _send(message, options, isAuthRetry, stepUpRetries = 0, isSessionRetry = false) {
try {
+ const isHandshake = Array.isArray(message) ? message.some((m) => require_src.isInitializeRequest(m)) : require_src.isInitializeRequest(message);
+ const isInitialized = Array.isArray(message) ? message.some((m) => require_src.isInitializedNotification(m)) : require_src.isInitializedNotification(message);
+ if (this._sessionRecovery && !isHandshake && !isInitialized) {
+ await this._sessionRecovery;
+ options?.requestSignal?.throwIfAborted();
+ }
const { resumptionToken, onresumptiontoken } = options || {};
if (resumptionToken) {
this._startOrAuthSse({
@@ -5226,8 +5251,8 @@ var StreamableHTTPClientTransport = class {
}
const headers = await this._commonHeaders();
this._applyBodyDerivedHeaders(headers, message);
- const isHandshake = Array.isArray(message) ? message.some((m) => require_src.isInitializeRequest(m)) : require_src.isInitializeRequest(message);
if (isHandshake) headers.delete("mcp-session-id");
+ const requestSessionId = headers.get("mcp-session-id") || void 0;
if (options?.headers !== void 0) for (const [name, value] of Object.entries(options.headers)) {
if (RESERVED_REQUEST_HEADER_NAMES.has(name.toLowerCase())) continue;
headers.set(name, value);
@@ -5249,8 +5274,14 @@ var StreamableHTTPClientTransport = class {
signal
};
const response = await (this._fetch ?? fetch)(this._url, init);
- if (isHandshake && response.ok) this._sessionId = response.headers.get("mcp-session-id") || void 0;
+ if (isHandshake && response.ok && (requestSessionId === void 0 || this._sessionId === requestSessionId)) this._sessionId = response.headers.get("mcp-session-id") || void 0;
if (!response.ok) {
+ if (response.status === 404 && requestSessionId && !isSessionRetry && !isInitialized) {
+ if (await this._recoverSession(requestSessionId)) {
+ options?.requestSignal?.throwIfAborted();
+ return this._send(message, options, isAuthRetry, stepUpRetries, true);
+ }
+ }
if (response.status === 401 && this._authProvider) {
if (response.headers.has("www-authenticate")) {
const { resourceMetadataUrl, scope } = extractWWWAuthenticateParams(response);
@@ -5264,7 +5295,7 @@ var StreamableHTTPClientTransport = class {
fetchFn: this._fetchWithInit
});
await response.text?.().catch(() => {});
- return this._send(message, options, true, stepUpRetries);
+ return this._send(message, options, true, stepUpRetries, isSessionRetry);
}
await response.text?.().catch(() => {});
if (isAuthRetry) throw new require_src.SdkHttpError(require_src.SdkErrorCode.ClientHttpAuthentication, "Server returned 401 after re-authentication", {
@@ -5284,7 +5315,7 @@ var StreamableHTTPClientTransport = class {
statusText: response.statusText,
text
}, stepUpRetries) !== "AUTHORIZED") throw new UnauthorizedError();
- return this._send(message, options, isAuthRetry, stepUpRetries + 1);
+ return this._send(message, options, isAuthRetry, stepUpRetries + 1, isSessionRetry);
}
}
if (response.status === 400 && typeof text === "string" && this._isModernEnvelopedRequest(message)) try {
diff --git a/dist/index.mjs b/dist/index.mjs
index 77e2389913cb5c5c2b047f95d990ab2892bef923..4b5e4ff2869189d600ca644488a7749668c39747 100644
--- a/dist/index.mjs
+++ b/dist/index.mjs
@@ -3151,6 +3151,7 @@ var Client = class extends Protocol {
*/
async _connectPlainLegacy(transport, options) {
await super.connect(transport);
+ transport.onsessionexpired = () => this._legacyHandshake(transport, options);
if (transport.sessionId !== void 0) {
const negotiatedProtocolVersion = this._negotiatedProtocolVersion;
if (negotiatedProtocolVersion !== void 0) transport.setProtocolVersion?.(negotiatedProtocolVersion);
@@ -3167,6 +3168,7 @@ var Client = class extends Protocol {
* the handshake; its completion sets the negotiated (legacy) version.
*/
async _legacyHandshake(transport, options) {
+ transport.onsessionexpired = () => this._legacyHandshake(transport, options);
const legacyVersions = legacyProtocolVersions(this._supportedProtocolVersions);
try {
const offeredVersion = legacyVersions[0];
@@ -3205,6 +3207,7 @@ var Client = class extends Protocol {
await super.connect(transport);
const negotiatedProtocolVersion = this._negotiatedProtocolVersion;
if (negotiatedProtocolVersion !== void 0 && transport.setProtocolVersion) transport.setProtocolVersion(negotiatedProtocolVersion);
+ if (negotiatedProtocolVersion !== void 0 && !isModernProtocolVersion(negotiatedProtocolVersion)) transport.onsessionexpired = () => this._legacyHandshake(transport, options);
return;
}
this._resetConnectionState();
@@ -5208,10 +5211,32 @@ var StreamableHTTPClientTransport = class {
}
}
async send(message, options) {
- return this._send(message, options, false);
+ return this._send(message, options, false, 0, false);
+ }
+ async _recoverSession(expiredSessionId) {
+ if (this._sessionRecovery) return this._sessionRecovery;
+ if (!this.onsessionexpired) return false;
+ if (this._sessionId !== expiredSessionId) return true;
+ this._sessionId = void 0;
+ this._sessionRecovery = Promise.resolve().then(() => this.onsessionexpired()).then(() => true);
+ try {
+ return await this._sessionRecovery;
+ } catch (error) {
+ this._sessionId = void 0;
+ await this.close();
+ throw error;
+ } finally {
+ this._sessionRecovery = void 0;
+ }
}
- async _send(message, options, isAuthRetry, stepUpRetries = 0) {
+ async _send(message, options, isAuthRetry, stepUpRetries = 0, isSessionRetry = false) {
try {
+ const isHandshake = Array.isArray(message) ? message.some((m) => isInitializeRequest(m)) : isInitializeRequest(message);
+ const isInitialized = Array.isArray(message) ? message.some((m) => isInitializedNotification(m)) : isInitializedNotification(message);
+ if (this._sessionRecovery && !isHandshake && !isInitialized) {
+ await this._sessionRecovery;
+ options?.requestSignal?.throwIfAborted();
+ }
const { resumptionToken, onresumptiontoken } = options || {};
if (resumptionToken) {
this._startOrAuthSse({
@@ -5223,8 +5248,8 @@ var StreamableHTTPClientTransport = class {
}
const headers = await this._commonHeaders();
this._applyBodyDerivedHeaders(headers, message);
- const isHandshake = Array.isArray(message) ? message.some((m) => isInitializeRequest(m)) : isInitializeRequest(message);
if (isHandshake) headers.delete("mcp-session-id");
+ const requestSessionId = headers.get("mcp-session-id") || void 0;
if (options?.headers !== void 0) for (const [name, value] of Object.entries(options.headers)) {
if (RESERVED_REQUEST_HEADER_NAMES.has(name.toLowerCase())) continue;
headers.set(name, value);
@@ -5246,8 +5271,14 @@ var StreamableHTTPClientTransport = class {
signal
};
const response = await (this._fetch ?? fetch)(this._url, init);
- if (isHandshake && response.ok) this._sessionId = response.headers.get("mcp-session-id") || void 0;
+ if (isHandshake && response.ok && (requestSessionId === void 0 || this._sessionId === requestSessionId)) this._sessionId = response.headers.get("mcp-session-id") || void 0;
if (!response.ok) {
+ if (response.status === 404 && requestSessionId && !isSessionRetry && !isInitialized) {
+ if (await this._recoverSession(requestSessionId)) {
+ options?.requestSignal?.throwIfAborted();
+ return this._send(message, options, isAuthRetry, stepUpRetries, true);
+ }
+ }
if (response.status === 401 && this._authProvider) {
if (response.headers.has("www-authenticate")) {
const { resourceMetadataUrl, scope } = extractWWWAuthenticateParams(response);
@@ -5261,7 +5292,7 @@ var StreamableHTTPClientTransport = class {
fetchFn: this._fetchWithInit
});
await response.text?.().catch(() => {});
- return this._send(message, options, true, stepUpRetries);
+ return this._send(message, options, true, stepUpRetries, isSessionRetry);
}
await response.text?.().catch(() => {});
if (isAuthRetry) throw new SdkHttpError(SdkErrorCode.ClientHttpAuthentication, "Server returned 401 after re-authentication", {
@@ -5281,7 +5312,7 @@ var StreamableHTTPClientTransport = class {
statusText: response.statusText,
text
}, stepUpRetries) !== "AUTHORIZED") throw new UnauthorizedError();
- return this._send(message, options, isAuthRetry, stepUpRetries + 1);
+ return this._send(message, options, isAuthRetry, stepUpRetries + 1, isSessionRetry);
}
}
if (response.status === 400 && typeof text === "string" && this._isModernEnvelopedRequest(message)) try {