Skip to content

Commit 314421b

Browse files
mattzcareymattzcarey
committed
session store
1 parent 0fb4c44 commit 314421b

File tree

3 files changed

+397
-12
lines changed

3 files changed

+397
-12
lines changed

‎src/experimental/fetch-streamable-http/fetchStreamableHttpServerTransport.test.ts‎

Lines changed: 263 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@ import {
1212
EventStore,
1313
EventId,
1414
StreamId,
15-
AuthenticatedRequest
15+
AuthenticatedRequest,
16+
SessionStore,
17+
SessionState
1618
} from './fetchStreamableHttpServerTransport.js';
1719
import { McpServer } from '../../server/mcp.js';
1820
import { CallToolResult, JSONRPCMessage } from '../../types.js';
@@ -130,6 +132,7 @@ interface TestServerConfig {
130132
sessionIdGenerator: (() => string) | undefined;
131133
enableJsonResponse?: boolean;
132134
eventStore?: EventStore;
135+
sessionStore?: SessionStore;
133136
onsessioninitialized?: (sessionId: string) => void | Promise<void>;
134137
onsessionclosed?: (sessionId: string) => void | Promise<void>;
135138
retryInterval?: number;
@@ -250,6 +253,7 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
250253
sessionIdGenerator: config.sessionIdGenerator,
251254
enableJsonResponse: config.enableJsonResponse ?? false,
252255
eventStore: config.eventStore,
256+
sessionStore: config.sessionStore,
253257
onsessioninitialized: config.onsessioninitialized,
254258
onsessionclosed: config.onsessionclosed,
255259
retryInterval: config.retryInterval,
@@ -2231,4 +2235,262 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
22312235
});
22322236
});
22332237
});
2238+
2239+
/**
2240+
* Tests for SessionStore functionality (distributed/serverless mode)
2241+
*/
2242+
describe('FetchStreamableHTTPServerTransport with SessionStore', () => {
2243+
let server: Server;
2244+
let transport: FetchStreamableHTTPServerTransport;
2245+
let baseUrl: URL;
2246+
2247+
afterEach(async () => {
2248+
if (server && transport) {
2249+
await stopTestServer({ server, transport });
2250+
}
2251+
});
2252+
2253+
/**
2254+
* Creates an in-memory session store for testing
2255+
*/
2256+
function createInMemorySessionStore(): SessionStore & { sessions: Map<string, SessionState> } {
2257+
const sessions = new Map<string, SessionState>();
2258+
return {
2259+
sessions,
2260+
get: async (sessionId: string) => sessions.get(sessionId),
2261+
save: async (sessionId: string, state: SessionState) => {
2262+
sessions.set(sessionId, state);
2263+
},
2264+
delete: async (sessionId: string) => {
2265+
sessions.delete(sessionId);
2266+
}
2267+
};
2268+
}
2269+
2270+
it('should save session state to store on initialization', async () => {
2271+
const sessionStore = createInMemorySessionStore();
2272+
const result = await createTestServer({
2273+
sessionIdGenerator: () => 'test-session-123',
2274+
sessionStore
2275+
});
2276+
server = result.server;
2277+
transport = result.transport;
2278+
baseUrl = result.baseUrl;
2279+
2280+
// Initialize the session
2281+
const response = await fetch(baseUrl, {
2282+
method: 'POST',
2283+
headers: {
2284+
'Content-Type': 'application/json',
2285+
Accept: 'application/json, text/event-stream'
2286+
},
2287+
body: JSON.stringify(TEST_MESSAGES.initialize)
2288+
});
2289+
2290+
expect(response.status).toBe(200);
2291+
2292+
// Verify session was saved to store
2293+
const savedSession = await sessionStore.get('test-session-123');
2294+
expect(savedSession).toBeDefined();
2295+
expect(savedSession?.initialized).toBe(true);
2296+
expect(savedSession?.protocolVersion).toBeDefined();
2297+
expect(savedSession?.createdAt).toBeGreaterThan(0);
2298+
});
2299+
2300+
it('should validate session from store for subsequent requests', async () => {
2301+
const sessionStore = createInMemorySessionStore();
2302+
const result = await createTestServer({
2303+
sessionIdGenerator: () => 'test-session-456',
2304+
sessionStore
2305+
});
2306+
server = result.server;
2307+
transport = result.transport;
2308+
baseUrl = result.baseUrl;
2309+
2310+
// Initialize the session
2311+
const initResponse = await fetch(baseUrl, {
2312+
method: 'POST',
2313+
headers: {
2314+
'Content-Type': 'application/json',
2315+
Accept: 'application/json, text/event-stream'
2316+
},
2317+
body: JSON.stringify(TEST_MESSAGES.initialize)
2318+
});
2319+
expect(initResponse.status).toBe(200);
2320+
const sessionId = initResponse.headers.get('mcp-session-id');
2321+
2322+
// Make a subsequent request with valid session ID
2323+
const listResponse = await fetch(baseUrl, {
2324+
method: 'POST',
2325+
headers: {
2326+
'Content-Type': 'application/json',
2327+
Accept: 'application/json, text/event-stream',
2328+
'mcp-session-id': sessionId!
2329+
},
2330+
body: JSON.stringify(TEST_MESSAGES.toolsList)
2331+
});
2332+
expect(listResponse.status).toBe(200);
2333+
});
2334+
2335+
it('should reject requests with invalid session ID when using session store', async () => {
2336+
const sessionStore = createInMemorySessionStore();
2337+
const result = await createTestServer({
2338+
sessionIdGenerator: () => 'test-session-789',
2339+
sessionStore
2340+
});
2341+
server = result.server;
2342+
transport = result.transport;
2343+
baseUrl = result.baseUrl;
2344+
2345+
// Initialize the session first
2346+
const initResponse = await fetch(baseUrl, {
2347+
method: 'POST',
2348+
headers: {
2349+
'Content-Type': 'application/json',
2350+
Accept: 'application/json, text/event-stream'
2351+
},
2352+
body: JSON.stringify(TEST_MESSAGES.initialize)
2353+
});
2354+
expect(initResponse.status).toBe(200);
2355+
2356+
// Try to make a request with invalid session ID
2357+
const response = await fetch(baseUrl, {
2358+
method: 'POST',
2359+
headers: {
2360+
'Content-Type': 'application/json',
2361+
Accept: 'application/json, text/event-stream',
2362+
'mcp-session-id': 'invalid-session-id'
2363+
},
2364+
body: JSON.stringify(TEST_MESSAGES.toolsList)
2365+
});
2366+
2367+
expect(response.status).toBe(404);
2368+
const body = await response.json();
2369+
expect(body.error.message).toBe('Session not found');
2370+
});
2371+
2372+
it('should delete session from store on DELETE request', async () => {
2373+
const sessionStore = createInMemorySessionStore();
2374+
const result = await createTestServer({
2375+
sessionIdGenerator: () => 'test-session-delete',
2376+
sessionStore
2377+
});
2378+
server = result.server;
2379+
transport = result.transport;
2380+
baseUrl = result.baseUrl;
2381+
2382+
// Initialize the session
2383+
const initResponse = await fetch(baseUrl, {
2384+
method: 'POST',
2385+
headers: {
2386+
'Content-Type': 'application/json',
2387+
Accept: 'application/json, text/event-stream'
2388+
},
2389+
body: JSON.stringify(TEST_MESSAGES.initialize)
2390+
});
2391+
expect(initResponse.status).toBe(200);
2392+
const sessionId = initResponse.headers.get('mcp-session-id');
2393+
2394+
// Verify session exists in store
2395+
expect(await sessionStore.get(sessionId!)).toBeDefined();
2396+
2397+
// Delete the session
2398+
const deleteResponse = await fetch(baseUrl, {
2399+
method: 'DELETE',
2400+
headers: {
2401+
'mcp-session-id': sessionId!
2402+
}
2403+
});
2404+
expect(deleteResponse.status).toBe(200);
2405+
2406+
// Verify session was deleted from store
2407+
expect(await sessionStore.get(sessionId!)).toBeUndefined();
2408+
});
2409+
2410+
it('should allow new transport instances to validate existing sessions (serverless mode)', async () => {
2411+
// This test simulates serverless behavior where each request
2412+
// is handled by a fresh transport instance
2413+
const sessionStore = createInMemorySessionStore();
2414+
2415+
// First, initialize using one transport instance
2416+
const result1 = await createTestServer({
2417+
sessionIdGenerator: () => 'serverless-session-123',
2418+
sessionStore
2419+
});
2420+
server = result1.server;
2421+
transport = result1.transport;
2422+
baseUrl = result1.baseUrl;
2423+
2424+
const initResponse = await fetch(baseUrl, {
2425+
method: 'POST',
2426+
headers: {
2427+
'Content-Type': 'application/json',
2428+
Accept: 'application/json, text/event-stream'
2429+
},
2430+
body: JSON.stringify(TEST_MESSAGES.initialize)
2431+
});
2432+
expect(initResponse.status).toBe(200);
2433+
const sessionId = initResponse.headers.get('mcp-session-id');
2434+
2435+
// Stop the first server
2436+
await stopTestServer({ server, transport });
2437+
2438+
// Create a NEW transport instance with same sessionStore (simulates new serverless invocation)
2439+
const result2 = await createTestServer({
2440+
sessionIdGenerator: () => crypto.randomUUID(), // Different generator, doesn't matter
2441+
sessionStore // Same session store
2442+
});
2443+
server = result2.server;
2444+
transport = result2.transport;
2445+
baseUrl = result2.baseUrl;
2446+
2447+
// The new transport should be able to validate the existing session from the store
2448+
const listResponse = await fetch(baseUrl, {
2449+
method: 'POST',
2450+
headers: {
2451+
'Content-Type': 'application/json',
2452+
Accept: 'application/json, text/event-stream',
2453+
'mcp-session-id': sessionId!
2454+
},
2455+
body: JSON.stringify(TEST_MESSAGES.toolsList)
2456+
});
2457+
2458+
expect(listResponse.status).toBe(200);
2459+
});
2460+
2461+
it('should work with GET SSE stream when session is hydrated from store', async () => {
2462+
const sessionStore = createInMemorySessionStore();
2463+
const result = await createTestServer({
2464+
sessionIdGenerator: () => 'sse-session-123',
2465+
sessionStore
2466+
});
2467+
server = result.server;
2468+
transport = result.transport;
2469+
baseUrl = result.baseUrl;
2470+
2471+
// Initialize session
2472+
const initResponse = await fetch(baseUrl, {
2473+
method: 'POST',
2474+
headers: {
2475+
'Content-Type': 'application/json',
2476+
Accept: 'application/json, text/event-stream'
2477+
},
2478+
body: JSON.stringify(TEST_MESSAGES.initialize)
2479+
});
2480+
expect(initResponse.status).toBe(200);
2481+
const sessionId = initResponse.headers.get('mcp-session-id');
2482+
2483+
// Open SSE stream with session ID
2484+
const sseResponse = await fetch(baseUrl, {
2485+
method: 'GET',
2486+
headers: {
2487+
Accept: 'text/event-stream',
2488+
'mcp-session-id': sessionId!
2489+
}
2490+
});
2491+
2492+
expect(sseResponse.status).toBe(200);
2493+
expect(sseResponse.headers.get('content-type')).toBe('text/event-stream');
2494+
});
2495+
});
22342496
});

0 commit comments

Comments
 (0)