Skip to content

Commit f27e62c

Browse files
Muhammad-Bin-AliMuhammad Alimattzcarey
authored
fix: WorkerTransport doesn't persist clientCapabilities, breaking elicitation in serverless (#783)
* Modify TransportState and WorkerTransport to persist clientCapabilities * Add tests * Change type to use official SDK type * Fix linting error * changeset --------- Co-authored-by: Muhammad Ali <muhammadali@cloudflare.com> Co-authored-by: Matt <77928207+mattzcarey@users.noreply.github.com> Co-authored-by: Matt Carey <mcarey@cloudflare.com>
1 parent 70e5040 commit f27e62c

3 files changed

Lines changed: 323 additions & 2 deletions

File tree

‎.changeset/three-kings-lay.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"agents": patch
3+
---
4+
5+
fix saving initialize params for stateless MCP server (effects eliciations and other optional features)

‎packages/agents/src/mcp/worker-transport.ts‎

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@ import type {
1010
JSONRPCMessage,
1111
RequestId,
1212
RequestInfo,
13-
MessageExtraInfo
13+
MessageExtraInfo,
14+
InitializeRequestParams
1415
} from "@modelcontextprotocol/sdk/types.js";
1516
import {
1617
isInitializeRequest,
@@ -27,6 +28,7 @@ import type {
2728
} from "@modelcontextprotocol/sdk/server/streamableHttp.js";
2829

2930
const MCP_PROTOCOL_VERSION_HEADER = "MCP-Protocol-Version";
31+
const RESTORE_REQUEST_ID = "__restore__";
3032

3133
interface StreamMapping {
3234
writer?: WritableStreamDefaultWriter<Uint8Array>;
@@ -43,6 +45,7 @@ export interface MCPStorageApi {
4345
export interface TransportState {
4446
sessionId?: string;
4547
initialized: boolean;
48+
initializeParams?: InitializeRequestParams;
4649
}
4750

4851
export interface WorkerTransportOptions {
@@ -99,6 +102,7 @@ export class WorkerTransport implements Transport {
99102
private stateRestored = false;
100103
private eventStore?: EventStore;
101104
private retryInterval?: number;
105+
private initializeParams?: TransportState["initializeParams"];
102106

103107
sessionId?: string;
104108
onclose?: () => void;
@@ -130,6 +134,16 @@ export class WorkerTransport implements Transport {
130134
if (state) {
131135
this.sessionId = state.sessionId;
132136
this.initialized = state.initialized;
137+
138+
// Restore _clientCapabilities on the Server instance by replaying the original initialize request
139+
if (state.initializeParams && this.onmessage) {
140+
this.onmessage({
141+
jsonrpc: "2.0",
142+
id: RESTORE_REQUEST_ID,
143+
method: "initialize",
144+
params: state.initializeParams
145+
});
146+
}
133147
}
134148

135149
this.stateRestored = true;
@@ -145,7 +159,8 @@ export class WorkerTransport implements Transport {
145159

146160
const state: TransportState = {
147161
sessionId: this.sessionId,
148-
initialized: this.initialized
162+
initialized: this.initialized,
163+
initializeParams: this.initializeParams
149164
};
150165

151166
await Promise.resolve(this.storage.set(state));
@@ -538,6 +553,16 @@ export class WorkerTransport implements Transport {
538553

539554
this.sessionId = this.sessionIdGenerator?.();
540555
this.initialized = true;
556+
557+
const initMessage = messages.find(isInitializeRequest);
558+
if (initMessage && isInitializeRequest(initMessage)) {
559+
this.initializeParams = {
560+
capabilities: initMessage.params.capabilities,
561+
clientInfo: initMessage.params.clientInfo,
562+
protocolVersion: initMessage.params.protocolVersion
563+
};
564+
}
565+
541566
await this.saveState();
542567

543568
if (this.sessionId && this.onsessioninitialized) {
@@ -802,6 +827,10 @@ export class WorkerTransport implements Transport {
802827
requestId = message.id;
803828
}
804829

830+
if (requestId === RESTORE_REQUEST_ID) {
831+
return;
832+
}
833+
805834
if (requestId === undefined) {
806835
if (isJSONRPCResultResponse(message) || isJSONRPCErrorResponse(message)) {
807836
throw new Error(

‎packages/agents/src/tests/mcp/transports/worker-transport.test.ts‎

Lines changed: 287 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -858,6 +858,293 @@ describe("WorkerTransport", () => {
858858
});
859859
});
860860

861+
describe("Client Capabilities Persistence (Serverless Restart)", () => {
862+
it("should persist initializeParams when client sends capabilities", async () => {
863+
const server = createTestServer();
864+
let storedState: TransportState | undefined;
865+
866+
const mockStorage = {
867+
get: async () => storedState,
868+
set: async (state: TransportState) => {
869+
storedState = state;
870+
}
871+
};
872+
873+
const transport = await setupTransport(server, {
874+
sessionIdGenerator: () => "test-session",
875+
storage: mockStorage,
876+
enableJsonResponse: true
877+
});
878+
879+
const request = new Request("http://example.com/", {
880+
method: "POST",
881+
headers: {
882+
"Content-Type": "application/json",
883+
Accept: "application/json, text/event-stream"
884+
},
885+
body: JSON.stringify({
886+
jsonrpc: "2.0",
887+
id: "1",
888+
method: "initialize",
889+
params: {
890+
capabilities: {
891+
elicitation: { form: {} }
892+
},
893+
clientInfo: { name: "test-client", version: "1.0" },
894+
protocolVersion: "2025-06-18"
895+
}
896+
})
897+
});
898+
899+
const response = await transport.handleRequest(request);
900+
await response.json();
901+
902+
expect(response.status).toBe(200);
903+
expect(storedState).toBeDefined();
904+
expect(storedState?.initializeParams).toBeDefined();
905+
expect(
906+
storedState?.initializeParams?.capabilities?.elicitation?.form
907+
).toBeDefined();
908+
expect(storedState?.initializeParams?.clientInfo).toEqual({
909+
name: "test-client",
910+
version: "1.0"
911+
});
912+
expect(storedState?.initializeParams?.protocolVersion).toBe("2025-06-18");
913+
});
914+
915+
it("should restore client capabilities on Server instance after restart", async () => {
916+
// Phase 1: Initialize with capabilities
917+
let storedState: TransportState | undefined;
918+
const mockStorage = {
919+
get: async () => storedState,
920+
set: async (state: TransportState) => {
921+
storedState = state;
922+
}
923+
};
924+
925+
const server1 = createTestServer();
926+
const transport1 = await setupTransport(server1, {
927+
sessionIdGenerator: () => "test-session",
928+
storage: mockStorage,
929+
enableJsonResponse: true
930+
});
931+
932+
const initRequest = new Request("http://example.com/", {
933+
method: "POST",
934+
headers: {
935+
"Content-Type": "application/json",
936+
Accept: "application/json, text/event-stream"
937+
},
938+
body: JSON.stringify({
939+
jsonrpc: "2.0",
940+
id: "1",
941+
method: "initialize",
942+
params: {
943+
capabilities: {
944+
elicitation: { form: {} }
945+
},
946+
clientInfo: { name: "test-client", version: "1.0" },
947+
protocolVersion: "2025-06-18"
948+
}
949+
})
950+
});
951+
952+
await transport1.handleRequest(initRequest);
953+
954+
// Verify server1 has capabilities
955+
expect(
956+
server1.server.getClientCapabilities()?.elicitation?.form
957+
).toBeDefined();
958+
959+
// Phase 2: Simulate serverless restart with NEW instances
960+
const server2 = createTestServer();
961+
const transport2 = await setupTransport(server2, {
962+
sessionIdGenerator: () => "test-session",
963+
storage: mockStorage,
964+
enableJsonResponse: true
965+
});
966+
967+
// Trigger state restoration by making a request
968+
const listRequest = new Request("http://example.com/", {
969+
method: "POST",
970+
headers: {
971+
"Content-Type": "application/json",
972+
Accept: "application/json, text/event-stream",
973+
"mcp-session-id": "test-session"
974+
},
975+
body: JSON.stringify({
976+
jsonrpc: "2.0",
977+
id: "2",
978+
method: "tools/list",
979+
params: {}
980+
})
981+
});
982+
983+
await transport2.handleRequest(listRequest);
984+
985+
// Verify capabilities were restored on server2
986+
expect(transport2.sessionId).toBe("test-session");
987+
expect(server2.server.getClientCapabilities()).toBeDefined();
988+
expect(
989+
server2.server.getClientCapabilities()?.elicitation?.form
990+
).toBeDefined();
991+
});
992+
993+
it("should restore clientInfo on Server instance after restart", async () => {
994+
let storedState: TransportState | undefined;
995+
const mockStorage = {
996+
get: async () => storedState,
997+
set: async (state: TransportState) => {
998+
storedState = state;
999+
}
1000+
};
1001+
1002+
const server1 = createTestServer();
1003+
const transport1 = await setupTransport(server1, {
1004+
sessionIdGenerator: () => "test-session",
1005+
storage: mockStorage,
1006+
enableJsonResponse: true
1007+
});
1008+
1009+
const initRequest = new Request("http://example.com/", {
1010+
method: "POST",
1011+
headers: {
1012+
"Content-Type": "application/json",
1013+
Accept: "application/json, text/event-stream"
1014+
},
1015+
body: JSON.stringify({
1016+
jsonrpc: "2.0",
1017+
id: "1",
1018+
method: "initialize",
1019+
params: {
1020+
capabilities: {},
1021+
clientInfo: { name: "my-client", version: "2.0" },
1022+
protocolVersion: "2025-06-18"
1023+
}
1024+
})
1025+
});
1026+
1027+
await transport1.handleRequest(initRequest);
1028+
1029+
// Simulate restart
1030+
const server2 = createTestServer();
1031+
const transport2 = await setupTransport(server2, {
1032+
sessionIdGenerator: () => "test-session",
1033+
storage: mockStorage,
1034+
enableJsonResponse: true
1035+
});
1036+
1037+
const listRequest = new Request("http://example.com/", {
1038+
method: "POST",
1039+
headers: {
1040+
"Content-Type": "application/json",
1041+
Accept: "application/json, text/event-stream",
1042+
"mcp-session-id": "test-session"
1043+
},
1044+
body: JSON.stringify({
1045+
jsonrpc: "2.0",
1046+
id: "2",
1047+
method: "tools/list",
1048+
params: {}
1049+
})
1050+
});
1051+
1052+
await transport2.handleRequest(listRequest);
1053+
1054+
// Verify clientInfo was restored
1055+
expect(server2.server.getClientVersion()).toEqual({
1056+
name: "my-client",
1057+
version: "2.0"
1058+
});
1059+
});
1060+
1061+
it("should handle old storage format without initializeParams (backward compatibility)", async () => {
1062+
// Simulate old stored state without initializeParams field
1063+
const oldState: TransportState = {
1064+
sessionId: "old-session",
1065+
initialized: true
1066+
// No initializeParams - simulating old storage format
1067+
};
1068+
1069+
const mockStorage = {
1070+
get: async () => oldState,
1071+
set: async () => {}
1072+
};
1073+
1074+
const server = createTestServer();
1075+
const transport = await setupTransport(server, {
1076+
storage: mockStorage,
1077+
enableJsonResponse: true
1078+
});
1079+
1080+
const request = new Request("http://example.com/", {
1081+
method: "POST",
1082+
headers: {
1083+
"Content-Type": "application/json",
1084+
Accept: "application/json, text/event-stream",
1085+
"mcp-session-id": "old-session"
1086+
},
1087+
body: JSON.stringify({
1088+
jsonrpc: "2.0",
1089+
id: "1",
1090+
method: "tools/list",
1091+
params: {}
1092+
})
1093+
});
1094+
1095+
// Should not throw
1096+
const response = await transport.handleRequest(request);
1097+
expect(response.status).toBe(200);
1098+
1099+
// Session restored but capabilities not available (no initializeParams)
1100+
expect(transport.sessionId).toBe("old-session");
1101+
expect(server.server.getClientCapabilities()).toBeUndefined();
1102+
});
1103+
1104+
it("should persist initializeParams with empty capabilities", async () => {
1105+
const server = createTestServer();
1106+
let storedState: TransportState | undefined;
1107+
1108+
const mockStorage = {
1109+
get: async () => storedState,
1110+
set: async (state: TransportState) => {
1111+
storedState = state;
1112+
}
1113+
};
1114+
1115+
const transport = await setupTransport(server, {
1116+
sessionIdGenerator: () => "test-session",
1117+
storage: mockStorage,
1118+
enableJsonResponse: true
1119+
});
1120+
1121+
const request = new Request("http://example.com/", {
1122+
method: "POST",
1123+
headers: {
1124+
"Content-Type": "application/json",
1125+
Accept: "application/json, text/event-stream"
1126+
},
1127+
body: JSON.stringify({
1128+
jsonrpc: "2.0",
1129+
id: "1",
1130+
method: "initialize",
1131+
params: {
1132+
capabilities: {}, // Empty but present
1133+
clientInfo: { name: "test-client", version: "1.0" },
1134+
protocolVersion: "2025-06-18"
1135+
}
1136+
})
1137+
});
1138+
1139+
const response = await transport.handleRequest(request);
1140+
await response.json();
1141+
1142+
expect(response.status).toBe(200);
1143+
expect(storedState?.initializeParams).toBeDefined();
1144+
expect(storedState?.initializeParams?.capabilities).toEqual({});
1145+
});
1146+
});
1147+
8611148
describe("Session Management", () => {
8621149
it("should use custom sessionIdGenerator", async () => {
8631150
const server = createTestServer();

0 commit comments

Comments
 (0)