Fix test suite for refactored shttp and redisTransport

- Update tests to handle async session initialization
- Fix session ID retrieval from SSE/JSON responses
- Update DELETE request tests to match new behavior
- Remove references to deprecated startServerListeningToRedis
- Add proper waits for async operations

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Jerome 2025-07-09 14:50:23 +01:00
commit 88f683ef12
No known key found for this signature in database
5 changed files with 9953 additions and 94 deletions

View file

@ -24,7 +24,8 @@ function createMockResponse() {
redirect: jest.fn().mockReturnThis(),
status: jest.fn().mockReturnThis(),
json: jest.fn().mockReturnThis(),
send: jest.fn().mockReturnThis()
send: jest.fn().mockReturnThis(),
setHeader: jest.fn().mockReturnThis()
};
return res as unknown as jest.Mocked<Response>;
}
@ -136,7 +137,8 @@ describe("EverythingAuthProvider", () => {
// Verify HTML sent with redirect
expect(res.send).toHaveBeenCalled();
const sentHtml = (res.send as jest.Mock).mock.calls[0][0];
expect(sentHtml).toContain('MCP Auth Page');
expect(sentHtml).toContain('MCP Server Authorization');
expect(sentHtml).toContain('Authorization Required');
expect(sentHtml).toContain('fakeupstreamauth/authorize?redirect_uri=/fakeupstreamauth/callback&state=');
});
});
@ -328,6 +330,9 @@ describe("EverythingAuthProvider", () => {
clientId: mcpInstallation.clientId,
scopes: ['mcp'],
expiresAt: mcpInstallation.mcpTokens.expires_in! + mcpInstallation.issuedAt,
extra: {
userId: "test-user-id"
}
});
});

View file

@ -41,8 +41,41 @@ describe('Streamable HTTP Handler Integration Tests', () => {
};
});
afterEach(() => {
// Helper function to trigger cleanup after handleStreamableHTTP calls
const triggerResponseCleanup = async () => {
// Find all finish handlers registered during the test
const finishHandlers = (mockRes.on as jest.Mock).mock.calls
.filter(([event]) => event === 'finish')
.map(([, handler]) => handler);
// Trigger all finish handlers
for (const handler of finishHandlers) {
if (typeof handler === 'function') {
await handler();
}
}
};
// Helper to extract session ID from test context
const getSessionIdFromTest = () => {
// Try to get from response headers first
const setHeaderCalls = (mockRes.setHeader as jest.Mock).mock.calls;
const sessionIdHeader = setHeaderCalls.find(([name]) => name === 'mcp-session-id');
if (sessionIdHeader?.[1]) {
return sessionIdHeader[1];
}
// Fall back to extracting from Redis channels
const allChannels = Array.from(mockRedis.subscribers.keys());
const serverChannel = allChannels.find(channel => channel.includes('mcp:shttp:toserver:'));
return serverChannel?.split(':')[3];
};
afterEach(async () => {
// Always trigger cleanup for any MCP servers created during tests
await triggerResponseCleanup();
mockRedis.clear();
jest.clearAllMocks();
});
describe('Redis Subscription Cleanup', () => {
@ -63,12 +96,16 @@ describe('Streamable HTTP Handler Integration Tests', () => {
mockReq.auth = {
clientId: 'test-client-123',
token: 'test-token',
scopes: ['mcp']
scopes: ['mcp'],
extra: { userId: 'test-user-123' }
} as AuthInfo;
// Call the handler
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Wait longer for async initialization to complete
await new Promise(resolve => setTimeout(resolve, 200));
// Check if any subscriptions were created on any channels
// Since we don't know the exact sessionId generated, check all channels
const allChannels = Array.from(mockRedis.subscribers.keys());
@ -111,7 +148,8 @@ describe('Streamable HTTP Handler Integration Tests', () => {
mockReq.auth = {
clientId: 'test-client-123',
token: 'test-token',
scopes: ['mcp']
scopes: ['mcp'],
extra: { userId: 'test-user-123' }
} as AuthInfo;
// Call the handler
@ -129,11 +167,27 @@ describe('Streamable HTTP Handler Integration Tests', () => {
if (finishHandler) {
await expect(finishHandler()).resolves.not.toThrow();
}
// Clean up the MCP server by sending DELETE request
const cleanupSessionId = getSessionIdFromTest();
if (cleanupSessionId) {
// Send DELETE request to clean up MCP server
jest.clearAllMocks();
mockReq.method = 'DELETE';
mockReq.headers['mcp-session-id'] = cleanupSessionId;
mockReq.body = {};
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Wait a bit for cleanup to complete
await new Promise(resolve => setTimeout(resolve, 50));
}
});
});
describe('DELETE Request Session Cleanup', () => {
it('should shutdown MCP server and clean up Redis channels on DELETE request', async () => {
it('should trigger onsessionclosed callback which sends shutdown control message', async () => {
// First, create a session with an initialization request
const initRequest: JSONRPCMessage = {
jsonrpc: '2.0',
@ -150,20 +204,58 @@ describe('Streamable HTTP Handler Integration Tests', () => {
mockReq.auth = {
clientId: 'test-client-123',
token: 'test-token',
scopes: ['mcp']
scopes: ['mcp'],
extra: { userId: 'test-user-123' }
} as AuthInfo;
// Initialize session
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Get the actual session ID from created channels
const allChannels = Array.from(mockRedis.subscribers.keys());
const serverChannel = allChannels.find(channel => channel.includes('mcp:shttp:toserver:'));
const sessionId = serverChannel?.split(':').pop();
// Wait for async initialization
await new Promise(resolve => setTimeout(resolve, 100));
// For initialization requests with StreamableHTTPServerTransport,
// the handler might not immediately return a response if using SSE mode
// Let's check different possible locations for the session ID
// Check JSON responses
const jsonCalls = (mockRes.json as jest.Mock).mock.calls;
let sessionId: string | undefined;
if (jsonCalls.length > 0) {
const response = jsonCalls[0][0];
if (response?.result?._meta?.sessionId) {
sessionId = response.result._meta.sessionId;
}
}
// Check write calls (for SSE responses)
if (!sessionId) {
const writeCalls = (mockRes.write as jest.Mock).mock.calls;
for (const [data] of writeCalls) {
if (typeof data === 'string' && data.includes('sessionId')) {
try {
// SSE data format: "data: {...}\n\n"
const jsonStr = data.replace(/^data: /, '').trim();
const parsed = JSON.parse(jsonStr);
if (parsed?.result?._meta?.sessionId) {
sessionId = parsed.result._meta.sessionId;
}
} catch (e) {
// Not valid JSON, continue
}
}
}
}
// Fallback to getting from Redis channels
if (!sessionId) {
sessionId = getSessionIdFromTest();
}
expect(sessionId).toBeDefined();
// Reset mocks
// Reset mocks but keep the session
jest.clearAllMocks();
// Now test DELETE request
@ -177,23 +269,20 @@ describe('Streamable HTTP Handler Integration Tests', () => {
// Track control messages sent to Redis
const publishSpy = jest.spyOn(mockRedis, 'publish');
// Call DELETE handler
// Call DELETE handler - StreamableHTTPServerTransport should handle it
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Verify successful response
expect(mockRes.status).toHaveBeenCalledWith(200);
expect(mockRes.json).toHaveBeenCalledWith({
jsonrpc: '2.0',
result: { status: 'Session terminated successfully' },
id: null,
});
// Wait for async processing and onsessionclosed callback
await new Promise(resolve => setTimeout(resolve, 100));
// Verify shutdown control message was sent
expect(publishSpy).toHaveBeenCalledWith(
`mcp:control:${sessionId}`,
expect.stringContaining('"type":"control"')
// The StreamableHTTPServerTransport should handle the DELETE and trigger onsessionclosed
// which calls shutdownSession, sending the control message
const controlCalls = publishSpy.mock.calls.filter(call =>
call[0] === `mcp:control:${sessionId}`
);
expect(controlCalls.length).toBeGreaterThan(0);
// Verify the control message content
const controlCall = publishSpy.mock.calls.find(call =>
call[0] === `mcp:control:${sessionId}`
@ -205,33 +294,86 @@ describe('Streamable HTTP Handler Integration Tests', () => {
}
});
it('should return 404 for DELETE request with invalid session ID', async () => {
const invalidSessionId = 'invalid-session-id';
mockReq.method = 'DELETE';
mockReq.headers = {
...mockReq.headers,
'mcp-session-id': invalidSessionId
it('should return 401 for DELETE request with wrong user', async () => {
// First, create a session as user1
const initRequest: JSONRPCMessage = {
jsonrpc: '2.0',
id: 'init-1',
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: {},
clientInfo: { name: 'test-client', version: '1.0.0' }
}
};
mockReq.body = {};
mockReq.body = initRequest;
mockReq.auth = {
clientId: 'test-client-123',
token: 'test-token',
scopes: ['mcp']
scopes: ['mcp'],
extra: { userId: 'user1' }
} as AuthInfo;
// Initialize session as user1
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Wait for async initialization
await new Promise(resolve => setTimeout(resolve, 100));
// Get the session ID from response
let sessionId: string | undefined;
// Check JSON responses
const jsonCalls = (mockRes.json as jest.Mock).mock.calls;
if (jsonCalls.length > 0) {
const response = jsonCalls[0][0];
if (response?.result?._meta?.sessionId) {
sessionId = response.result._meta.sessionId;
}
}
// Check write calls (for SSE responses)
if (!sessionId) {
const writeCalls = (mockRes.write as jest.Mock).mock.calls;
for (const [data] of writeCalls) {
if (typeof data === 'string' && data.includes('sessionId')) {
try {
const jsonStr = data.replace(/^data: /, '').trim();
const parsed = JSON.parse(jsonStr);
if (parsed?.result?._meta?.sessionId) {
sessionId = parsed.result._meta.sessionId;
}
} catch (e) {}
}
}
}
if (!sessionId) {
sessionId = getSessionIdFromTest();
}
// Reset mocks
jest.clearAllMocks();
// Now test DELETE request as user2
mockReq.method = 'DELETE';
mockReq.headers = {
...mockReq.headers,
'mcp-session-id': sessionId
};
mockReq.body = {};
mockReq.auth = {
clientId: 'test-client-456',
token: 'test-token-2',
scopes: ['mcp'],
extra: { userId: 'user2' }
} as AuthInfo;
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Should return 404 for non-existent session
expect(mockRes.status).toHaveBeenCalledWith(404);
expect(mockRes.json).toHaveBeenCalledWith({
jsonrpc: '2.0',
error: {
code: -32001,
message: 'Session not found',
},
id: null,
});
// Should return 401 for unauthorized access to another user's session
expect(mockRes.status).toHaveBeenCalledWith(401);
});
});
@ -242,13 +384,15 @@ describe('Streamable HTTP Handler Integration Tests', () => {
const user1Auth: AuthInfo = {
clientId: 'user1-client',
token: 'user1-token',
scopes: ['mcp']
scopes: ['mcp'],
extra: { userId: 'user1' }
};
const user2Auth: AuthInfo = {
clientId: 'user2-client',
token: 'user2-token',
scopes: ['mcp']
scopes: ['mcp'],
extra: { userId: 'user2' }
};
const initRequest: JSONRPCMessage = {
@ -267,14 +411,61 @@ describe('Streamable HTTP Handler Integration Tests', () => {
mockReq.auth = user1Auth;
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Wait for async initialization to complete
await new Promise(resolve => setTimeout(resolve, 100));
// Get the actual session ID from response
let actualSessionId: string | undefined;
// Check JSON responses
const jsonCalls = (mockRes.json as jest.Mock).mock.calls;
if (jsonCalls.length > 0) {
const response = jsonCalls[0][0];
if (response?.result?._meta?.sessionId) {
actualSessionId = response.result._meta.sessionId;
}
}
// Check write calls (for SSE responses)
if (!actualSessionId) {
const writeCalls = (mockRes.write as jest.Mock).mock.calls;
for (const [data] of writeCalls) {
if (typeof data === 'string' && data.includes('sessionId')) {
try {
const jsonStr = data.replace(/^data: /, '').trim();
const parsed = JSON.parse(jsonStr);
if (parsed?.result?._meta?.sessionId) {
actualSessionId = parsed.result._meta.sessionId;
}
} catch (e) {}
}
}
}
if (!actualSessionId) {
actualSessionId = getSessionIdFromTest();
}
expect(actualSessionId).toBeDefined();
// Store finish handler before clearing mocks
const finishHandler1 = (mockRes.on as jest.Mock).mock.calls.find(
([event]) => event === 'finish'
)?.[1] as (() => Promise<void>) | undefined;
// Reset mocks
jest.clearAllMocks();
// Trigger cleanup for the MCP server created in this step
if (finishHandler1) {
await finishHandler1();
}
// User 2 tries to access user 1's session
mockReq.headers = {
...mockReq.headers,
'mcp-session-id': sessionId
'mcp-session-id': actualSessionId
};
mockReq.body = {
jsonrpc: '2.0',
@ -286,19 +477,22 @@ describe('Streamable HTTP Handler Integration Tests', () => {
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Should return 401 or 403 for unauthorized access
// Note: Current implementation might not enforce this
// This test documents what SHOULD happen for security
// Should return 401 for unauthorized access to another user's session
expect(mockRes.status).toHaveBeenCalledWith(401);
// expect(mockRes.status).toHaveBeenCalledWith(403);
// expect(mockRes.json).toHaveBeenCalledWith({
// jsonrpc: '2.0',
// error: {
// code: -32000,
// message: 'Unauthorized: Session belongs to different user'
// },
// id: null
// });
// Clean up the MCP server by sending DELETE request
if (actualSessionId) {
jest.clearAllMocks();
mockReq.method = 'DELETE';
mockReq.headers['mcp-session-id'] = actualSessionId;
mockReq.body = {};
mockReq.auth = user1Auth; // Use user1's auth to delete their session
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Wait a bit for cleanup to complete
await new Promise(resolve => setTimeout(resolve, 50));
}
});
it('should allow users to create separate sessions with same session ID pattern', async () => {
@ -334,8 +528,18 @@ describe('Streamable HTTP Handler Integration Tests', () => {
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Store finish handler before clearing mocks
const finishHandler1 = (mockRes.on as jest.Mock).mock.calls.find(
([event]) => event === 'finish'
)?.[1] as (() => Promise<void>) | undefined;
// Reset for user 2
jest.clearAllMocks();
// Trigger cleanup for User 1's MCP server
if (finishHandler1) {
await finishHandler1();
}
// User 2 creates their own session
mockReq.body = {
@ -351,6 +555,15 @@ describe('Streamable HTTP Handler Integration Tests', () => {
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Trigger cleanup for User 2's MCP server
const finishHandler2 = (mockRes.on as jest.Mock).mock.calls.find(
([event]) => event === 'finish'
)?.[1] as (() => Promise<void>) | undefined;
if (finishHandler2) {
await finishHandler2();
}
// Both users should be able to create sessions successfully
// Sessions should be isolated in Redis using user-scoped keys
expect(mockRes.status).not.toHaveBeenCalledWith(400);
@ -388,6 +601,15 @@ describe('Streamable HTTP Handler Integration Tests', () => {
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Trigger cleanup for User 1's MCP server
const finishHandler1 = (mockRes.on as jest.Mock).mock.calls.find(
([event]) => event === 'finish'
)?.[1] as (() => Promise<void>) | undefined;
if (finishHandler1) {
await finishHandler1();
}
// Track session 1 ID (would be returned in response headers)
const session1Id = 'user1-session-id'; // In real implementation, extract from response
@ -406,6 +628,15 @@ describe('Streamable HTTP Handler Integration Tests', () => {
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
// Trigger cleanup for User 2's MCP server
const finishHandler2 = (mockRes.on as jest.Mock).mock.calls.find(
([event]) => event === 'finish'
)?.[1] as (() => Promise<void>) | undefined;
if (finishHandler2) {
await finishHandler2();
}
// Track session 2 ID
const session2Id = 'user2-session-id';

View file

@ -59,7 +59,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
await setSessionOwner(sessionId, userId);
}
const onsessionclosed = async (sessionId: string) => {
const onsessionclosed = (sessionId: string) => {
console.log(`Session ${sessionId} closing`);
void shutdownSession(sessionId);
console.log(`Session ${sessionId} closed`);

View file

@ -2,9 +2,10 @@ import { jest } from '@jest/globals';
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
import { MockRedisClient, setRedisClient } from '../redis.js';
import {
startServerListeningToRedis,
ServerRedisTransport,
redisRelayToMcpServer,
shutdownSession
shutdownSession,
setSessionOwner
} from './redisTransport.js';
import { createMcpServer } from './mcp.js';
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
@ -27,8 +28,10 @@ describe('Redis Transport Integration', () => {
it('should relay initialization request from client to server through Redis', async () => {
// 1. Start the server listening to Redis
const serverWrapper = createMcpServer();
const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId);
const { server, cleanup: serverCleanup } = createMcpServer();
const serverTransport = new ServerRedisTransport(sessionId);
serverTransport.onclose = serverCleanup;
await server.connect(serverTransport);
// 2. Create a mock client transport (simulating the streamable HTTP client side)
const mockClientTransport: Transport = {
@ -121,13 +124,20 @@ describe('Redis Transport Integration', () => {
// Cleanup
await cleanup();
await shutdownSession(sessionId);
serverCleanup(); // Clean up MCP server intervals
// Ensure server transport is closed
await serverTransport.close();
await new Promise(resolve => setTimeout(resolve, 10));
});
it('should handle tools/list request through Redis relay', async () => {
// Set up server and mock client
const serverWrapper = createMcpServer();
const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId);
const { server, cleanup: serverCleanup } = createMcpServer();
const serverTransport = new ServerRedisTransport(sessionId);
serverTransport.onclose = serverCleanup;
await server.connect(serverTransport);
const mockClientTransport: Transport = {
onmessage: undefined,
@ -168,13 +178,20 @@ describe('Redis Transport Integration', () => {
// Cleanup
await cleanup();
await shutdownSession(sessionId);
serverCleanup(); // Clean up MCP server intervals
// Ensure server transport is closed
await serverTransport.close();
await new Promise(resolve => setTimeout(resolve, 10));
});
it('should handle notifications through Redis relay', async () => {
// Set up server and mock client
const serverWrapper = createMcpServer();
const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId);
const { server, cleanup: serverCleanup } = createMcpServer();
const serverTransport = new ServerRedisTransport(sessionId);
serverTransport.onclose = serverCleanup;
await server.connect(serverTransport);
const mockClientTransport: Transport = {
onmessage: undefined,
@ -187,6 +204,17 @@ describe('Redis Transport Integration', () => {
const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
// Set up notification subscription manually since notifications don't have an id
const notificationChannel = `mcp:shttp:toclient:${sessionId}:__GET_stream`;
const notificationCleanup = await mockRedis.createSubscription(notificationChannel, async (redisMessageJson) => {
const redisMessage = JSON.parse(redisMessageJson);
if (redisMessage.type === 'mcp') {
await mockClientTransport.send(redisMessage.message, redisMessage.options);
}
}, (error) => {
mockClientTransport.onerror?.(error);
});
// Send a notification from server (notifications don't have an id)
const notification: JSONRPCMessage = {
jsonrpc: '2.0',
@ -216,9 +244,17 @@ describe('Redis Transport Integration', () => {
undefined
);
// Cleanup notification subscription
await notificationCleanup();
// Cleanup
await cleanup();
await shutdownSession(sessionId);
serverCleanup(); // Clean up MCP server intervals
// Ensure server transport is closed
await serverTransport.close();
await new Promise(resolve => setTimeout(resolve, 10));
});

View file

@ -5,8 +5,11 @@ import {
ServerRedisTransport,
redisRelayToMcpServer,
isLive,
startServerListeningToRedis,
shutdownSession
shutdownSession,
setSessionOwner,
getSessionOwner,
validateSessionOwnership,
isSessionOwnedBy
} from './redisTransport.js';
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
import { createMcpServer } from './mcp.js';
@ -278,36 +281,45 @@ describe('Redis Transport', () => {
});
});
describe('startServerListeningToRedis', () => {
const sessionId = 'test-session-server';
describe('Session Ownership', () => {
const sessionId = 'test-session-ownership';
const userId = 'test-user-123';
it('should connect server with ServerRedisTransport', async () => {
const serverWrapper = createMcpServer();
const connectSpy = jest.spyOn(serverWrapper.server, 'connect');
await startServerListeningToRedis(serverWrapper, sessionId);
expect(connectSpy).toHaveBeenCalledWith(
expect.any(ServerRedisTransport)
);
// Clean up properly
await shutdownSession(sessionId);
await new Promise(resolve => setTimeout(resolve, 10));
it('should set and get session owner', async () => {
await setSessionOwner(sessionId, userId);
const owner = await getSessionOwner(sessionId);
expect(owner).toBe(userId);
});
it('should create transport that can send responses and shutdown properly', async () => {
const serverWrapper = createMcpServer();
it('should validate session ownership correctly', async () => {
await setSessionOwner(sessionId, userId);
const transport = await startServerListeningToRedis(serverWrapper, sessionId);
expect(await validateSessionOwnership(sessionId, userId)).toBe(true);
expect(await validateSessionOwnership(sessionId, 'different-user')).toBe(false);
});
// The server should now be connected and able to handle requests via Redis
expect(serverWrapper.server).toBeDefined();
expect(transport).toBeInstanceOf(ServerRedisTransport);
// Test proper shutdown
await shutdownSession(sessionId);
await new Promise(resolve => setTimeout(resolve, 10));
it('should check if session is owned by user including liveness', async () => {
// Session not live yet
expect(await isSessionOwnedBy(sessionId, userId)).toBe(false);
// Make session live
await mockRedis.createSubscription(
`mcp:shttp:toserver:${sessionId}`,
jest.fn(),
jest.fn()
);
// Still false because no owner set
expect(await isSessionOwnedBy(sessionId, userId)).toBe(false);
// Set owner
await setSessionOwner(sessionId, userId);
// Now should be true
expect(await isSessionOwnedBy(sessionId, userId)).toBe(true);
// False for different user
expect(await isSessionOwnedBy(sessionId, 'different-user')).toBe(false);
});
});