Unify auth and MCP servers into single modular codebase
- Merge auth-server and mcp-server into unified architecture under src/modules/ - Add configurable auth modes: internal (single process) and external (separate processes) - Update example scripts to support OAuth discovery and both auth modes - Consolidate OAuth documentation into comprehensive implementation guide - Add configuration options for development vs production deployments - Support both Redis-backed and in-memory session storage
This commit is contained in:
Basil Hosmer
parent
22c7f7476b
commit
bc3619c73c
4 changed files with 3946 additions and 17945 deletions
694
src/modules/mcp/handlers/shttp.integration.test.ts
Normal file
694
src/modules/mcp/handlers/shttp.integration.test.ts
Normal file
|
|
@ -0,0 +1,694 @@
|
|||
import { jest } from '@jest/globals';
|
||||
import { Request, Response } from 'express';
|
||||
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { MockRedisClient, setRedisClient } from '../../shared/redis.js';
|
||||
import { handleStreamableHTTP } from './shttp.js';
|
||||
import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types.js';
|
||||
// import { randomUUID } from 'crypto'; // Currently unused but may be needed for future tests
|
||||
import { shutdownSession } from '../services/redisTransport.js';
|
||||
|
||||
// Type for MCP initialization response
|
||||
interface MCPInitResponse {
|
||||
jsonrpc: string;
|
||||
id: string | number;
|
||||
result?: {
|
||||
_meta?: {
|
||||
sessionId?: string;
|
||||
};
|
||||
[key: string]: unknown;
|
||||
};
|
||||
}
|
||||
|
||||
describe('Streamable HTTP Handler Integration Tests', () => {
|
||||
let mockRedis: MockRedisClient;
|
||||
let mockReq: Partial<Request>;
|
||||
let mockRes: Partial<Response>;
|
||||
|
||||
beforeEach(() => {
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
jest.resetAllMocks();
|
||||
|
||||
// Create mock response with chainable methods
|
||||
mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn().mockReturnThis(),
|
||||
on: jest.fn().mockReturnThis(),
|
||||
once: jest.fn().mockReturnThis(),
|
||||
emit: jest.fn().mockReturnThis(),
|
||||
headersSent: false,
|
||||
setHeader: jest.fn().mockReturnThis(),
|
||||
writeHead: jest.fn().mockReturnThis(),
|
||||
write: jest.fn().mockReturnThis(),
|
||||
end: jest.fn().mockReturnThis(),
|
||||
getHeader: jest.fn(),
|
||||
removeHeader: jest.fn().mockReturnThis(),
|
||||
socket: {
|
||||
setTimeout: jest.fn(),
|
||||
},
|
||||
} as unknown as Partial<Response>;
|
||||
|
||||
// Create mock request
|
||||
mockReq = {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'content-type': 'application/json',
|
||||
'accept': 'application/json, text/event-stream',
|
||||
'mcp-protocol-version': '2024-11-05',
|
||||
},
|
||||
body: {},
|
||||
};
|
||||
});
|
||||
|
||||
// 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 = (): string | undefined => {
|
||||
// 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] as string;
|
||||
}
|
||||
|
||||
// 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', () => {
|
||||
it('should clean up Redis subscriptions after shttp response completes', async () => {
|
||||
// Set up initialization request (no session ID for new initialization)
|
||||
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 = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
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));
|
||||
|
||||
// get the sessionId from the response
|
||||
const sessionId = getSessionIdFromTest();
|
||||
expect(sessionId).toBeDefined();
|
||||
|
||||
// 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());
|
||||
const totalSubscriptions = allChannels.reduce((sum, channel) => sum + (mockRedis.subscribers.get(channel)?.length || 0), 0);
|
||||
|
||||
// Should have created at least one subscription (server channel)
|
||||
expect(totalSubscriptions).toBeGreaterThan(0);
|
||||
expect(allChannels.some(channel => channel.includes('mcp:shttp:toserver:'))).toBe(true);
|
||||
|
||||
// Find the finish handler that was registered
|
||||
const finishHandler = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
expect(finishHandler).toBeDefined();
|
||||
|
||||
// Simulate response completion to trigger cleanup
|
||||
if (finishHandler) {
|
||||
await finishHandler();
|
||||
}
|
||||
|
||||
// Verify cleanup handler was registered
|
||||
expect(mockRes.on).toHaveBeenCalledWith('finish', expect.any(Function));
|
||||
|
||||
if (sessionId) {
|
||||
await shutdownSession(sessionId)
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle cleanup errors gracefully', async () => {
|
||||
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 = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'test-user-123' }
|
||||
} as AuthInfo;
|
||||
|
||||
// Call the handler
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Use Redis error simulation to test error handling
|
||||
const finishHandler = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
// Simulate error during cleanup
|
||||
|
||||
// Cleanup should not throw error even if Redis operations fail
|
||||
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';
|
||||
if (mockReq.headers) {
|
||||
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 trigger onsessionclosed callback which sends shutdown control message', async () => {
|
||||
// First, create a session with an initialization request
|
||||
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 = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'test-user-123' }
|
||||
} as AuthInfo;
|
||||
|
||||
// Initialize session
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// 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] as MCPInitResponse;
|
||||
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) as MCPInitResponse;
|
||||
if (parsed?.result?._meta?.sessionId) {
|
||||
sessionId = parsed.result._meta.sessionId;
|
||||
}
|
||||
} catch {
|
||||
// Not valid JSON, continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to getting from Redis channels
|
||||
if (!sessionId) {
|
||||
sessionId = getSessionIdFromTest();
|
||||
}
|
||||
|
||||
expect(sessionId).toBeDefined();
|
||||
|
||||
// Reset mocks but keep the session
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Now test DELETE request
|
||||
mockReq.method = 'DELETE';
|
||||
mockReq.headers = {
|
||||
...mockReq.headers,
|
||||
'mcp-session-id': sessionId
|
||||
};
|
||||
mockReq.body = {};
|
||||
|
||||
// Track control messages sent to Redis
|
||||
const publishSpy = jest.spyOn(mockRedis, 'publish');
|
||||
|
||||
// Call DELETE handler - StreamableHTTPServerTransport should handle it
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait for async processing and onsessionclosed callback
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
|
||||
// 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}`
|
||||
);
|
||||
if (controlCall) {
|
||||
const message = JSON.parse(controlCall[1]);
|
||||
expect(message.type).toBe('control');
|
||||
expect(message.action).toBe('SHUTDOWN');
|
||||
}
|
||||
});
|
||||
|
||||
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 = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
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] as MCPInitResponse;
|
||||
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) as MCPInitResponse;
|
||||
if (parsed?.result?._meta?.sessionId) {
|
||||
sessionId = parsed.result._meta.sessionId;
|
||||
}
|
||||
} catch {
|
||||
// Ignore JSON parse errors
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 401 for unauthorized access to another user's session
|
||||
expect(mockRes.status).toHaveBeenCalledWith(401);
|
||||
|
||||
// shutdown the session
|
||||
if (sessionId) {
|
||||
await shutdownSession(sessionId)
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('User Session Isolation', () => {
|
||||
it('should prevent users from accessing sessions created by other users', async () => {
|
||||
// Create session for user 1
|
||||
const user1Auth: AuthInfo = {
|
||||
clientId: 'user1-client',
|
||||
token: 'user1-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'user1' }
|
||||
};
|
||||
|
||||
const user2Auth: AuthInfo = {
|
||||
clientId: 'user2-client',
|
||||
token: 'user2-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'user2' }
|
||||
};
|
||||
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'user1-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
// User 1 creates session
|
||||
mockReq.body = initRequest;
|
||||
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] as MCPInitResponse;
|
||||
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) as MCPInitResponse;
|
||||
if (parsed?.result?._meta?.sessionId) {
|
||||
actualSessionId = parsed.result._meta.sessionId;
|
||||
}
|
||||
} catch {
|
||||
// Ignore JSON parse errors
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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': actualSessionId
|
||||
};
|
||||
mockReq.body = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'user2-request',
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
mockReq.auth = user2Auth;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Should return 401 for unauthorized access to another user's session
|
||||
expect(mockRes.status).toHaveBeenCalledWith(401);
|
||||
|
||||
// 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 () => {
|
||||
// This test shows that different users should be able to use sessions
|
||||
// without interfering with each other, even if session IDs might collide
|
||||
|
||||
const user1Auth: AuthInfo = {
|
||||
clientId: 'user1-client',
|
||||
token: 'user1-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
const user2Auth: AuthInfo = {
|
||||
clientId: 'user2-client',
|
||||
token: 'user2-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
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' }
|
||||
}
|
||||
};
|
||||
|
||||
// User 1 creates session
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = user1Auth;
|
||||
|
||||
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 = {
|
||||
...initRequest,
|
||||
id: 'init-2',
|
||||
params: {
|
||||
...initRequest.params,
|
||||
clientInfo: { name: 'user2-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
mockReq.auth = user2Auth;
|
||||
delete mockReq.headers!['mcp-session-id']; // New initialization
|
||||
|
||||
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);
|
||||
expect(mockRes.status).not.toHaveBeenCalledWith(403);
|
||||
});
|
||||
|
||||
it('should clean up only the requesting user\'s session on DELETE', async () => {
|
||||
// Create sessions for both users
|
||||
const user1Auth: AuthInfo = {
|
||||
clientId: 'user1-client',
|
||||
token: 'user1-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
const user2Auth: AuthInfo = {
|
||||
clientId: 'user2-client',
|
||||
token: 'user2-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
// Create session for user 1
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'user1-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = user1Auth;
|
||||
|
||||
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
|
||||
|
||||
// Create session for user 2
|
||||
jest.clearAllMocks();
|
||||
mockReq.body = {
|
||||
...initRequest,
|
||||
id: 'init-2',
|
||||
params: {
|
||||
...initRequest.params,
|
||||
clientInfo: { name: 'user2-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
mockReq.auth = user2Auth;
|
||||
delete mockReq.headers!['mcp-session-id'];
|
||||
|
||||
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 (placeholder for actual implementation)
|
||||
|
||||
// User 1 deletes their session
|
||||
jest.clearAllMocks();
|
||||
mockReq.method = 'DELETE';
|
||||
mockReq.headers = {
|
||||
...mockReq.headers,
|
||||
'mcp-session-id': session1Id
|
||||
};
|
||||
mockReq.body = {};
|
||||
mockReq.auth = user1Auth;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Only user 1's session should be cleaned up
|
||||
// User 2's session should remain active
|
||||
// This test documents expected behavior for proper user isolation
|
||||
});
|
||||
});
|
||||
});
|
||||
229
src/modules/mcp/handlers/shttp.test.ts
Normal file
229
src/modules/mcp/handlers/shttp.test.ts
Normal file
|
|
@ -0,0 +1,229 @@
|
|||
import { jest } from '@jest/globals';
|
||||
import { Request, Response } from 'express';
|
||||
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { MockRedisClient, setRedisClient } from '../../shared/redis.js';
|
||||
|
||||
describe('Streamable HTTP Handler', () => {
|
||||
let mockRedis: MockRedisClient;
|
||||
|
||||
beforeEach(() => {
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
jest.resetAllMocks();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
mockRedis.clear();
|
||||
});
|
||||
|
||||
describe('Helper function tests', () => {
|
||||
it('should verify Redis mock is working', async () => {
|
||||
await mockRedis.set('test-key', 'test-value');
|
||||
const value = await mockRedis.get('test-key');
|
||||
expect(value).toBe('test-value');
|
||||
});
|
||||
|
||||
it('should handle Redis pub/sub', async () => {
|
||||
const messageHandler = jest.fn();
|
||||
const cleanup = await mockRedis.createSubscription(
|
||||
'test-channel',
|
||||
messageHandler,
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
await mockRedis.publish('test-channel', 'test-message');
|
||||
|
||||
expect(messageHandler).toHaveBeenCalledWith('test-message');
|
||||
|
||||
await cleanup();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Request validation', () => {
|
||||
it('should identify initialize requests correctly', async () => {
|
||||
const { isInitializeRequest } = await import('@modelcontextprotocol/sdk/types.js');
|
||||
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 1,
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test', version: '1.0' }
|
||||
}
|
||||
};
|
||||
|
||||
const nonInitRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 2,
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
expect(isInitializeRequest(initRequest)).toBe(true);
|
||||
expect(isInitializeRequest(nonInitRequest)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('HTTP response mock behavior', () => {
|
||||
it('should create proper response mock with chainable methods', () => {
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn().mockReturnThis(),
|
||||
on: jest.fn().mockReturnThis(),
|
||||
headersSent: false,
|
||||
} as Partial<Response>;
|
||||
|
||||
// Test chaining
|
||||
const result = mockRes.status!(400).json!({
|
||||
jsonrpc: '2.0',
|
||||
error: { code: -32000, message: 'Bad Request' },
|
||||
id: null
|
||||
});
|
||||
|
||||
expect(mockRes.status).toHaveBeenCalledWith(400);
|
||||
expect(mockRes.json).toHaveBeenCalledWith({
|
||||
jsonrpc: '2.0',
|
||||
error: { code: -32000, message: 'Bad Request' },
|
||||
id: null
|
||||
});
|
||||
expect(result).toBe(mockRes);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Session ID generation', () => {
|
||||
it('should generate valid UUIDs', async () => {
|
||||
const { randomUUID } = await import('crypto');
|
||||
|
||||
const sessionId = randomUUID();
|
||||
|
||||
// UUID v4 format: xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx
|
||||
expect(sessionId).toMatch(/^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Redis channel naming', () => {
|
||||
it('should create correct channel names for server communication', () => {
|
||||
const sessionId = 'test-session-123';
|
||||
const requestId = 'req-456';
|
||||
|
||||
const toServerChannel = `mcp:shttp:toserver:${sessionId}`;
|
||||
const toClientChannel = `mcp:shttp:toclient:${sessionId}:${requestId}`;
|
||||
const notificationChannel = `mcp:shttp:toclient:${sessionId}:__GET_stream`;
|
||||
|
||||
expect(toServerChannel).toBe('mcp:shttp:toserver:test-session-123');
|
||||
expect(toClientChannel).toBe('mcp:shttp:toclient:test-session-123:req-456');
|
||||
expect(notificationChannel).toBe('mcp:shttp:toclient:test-session-123:__GET_stream');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error response formatting', () => {
|
||||
it('should format JSON-RPC error responses correctly', () => {
|
||||
const errorResponse = {
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
code: -32000,
|
||||
message: 'Bad Request: No valid session ID provided',
|
||||
},
|
||||
id: null,
|
||||
};
|
||||
|
||||
expect(errorResponse.jsonrpc).toBe('2.0');
|
||||
expect(errorResponse.error.code).toBe(-32000);
|
||||
expect(errorResponse.error.message).toContain('Bad Request');
|
||||
expect(errorResponse.id).toBe(null);
|
||||
});
|
||||
|
||||
it('should format internal error responses correctly', () => {
|
||||
const internalErrorResponse = {
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
code: -32603,
|
||||
message: 'Internal error',
|
||||
},
|
||||
id: null,
|
||||
};
|
||||
|
||||
expect(internalErrorResponse.jsonrpc).toBe('2.0');
|
||||
expect(internalErrorResponse.error.code).toBe(-32603);
|
||||
expect(internalErrorResponse.error.message).toBe('Internal error');
|
||||
expect(internalErrorResponse.id).toBe(null);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Request/Response patterns', () => {
|
||||
it('should handle typical MCP message structures', () => {
|
||||
const initializeRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 1,
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
const toolsListRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 2,
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
const toolsListResponse: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 2,
|
||||
result: {
|
||||
tools: [
|
||||
{
|
||||
name: 'echo',
|
||||
description: 'Echo the input',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
text: { type: 'string' }
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
};
|
||||
|
||||
expect(initializeRequest.method).toBe('initialize');
|
||||
expect(toolsListRequest.method).toBe('tools/list');
|
||||
expect(toolsListResponse.result).toBeDefined();
|
||||
expect(Array.isArray(toolsListResponse.result?.tools)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('HTTP header handling', () => {
|
||||
it('should extract session ID from headers', () => {
|
||||
const mockReq = {
|
||||
headers: {
|
||||
'mcp-session-id': 'test-session-123',
|
||||
'content-type': 'application/json'
|
||||
},
|
||||
body: {}
|
||||
} as Partial<Request>;
|
||||
|
||||
const sessionId = mockReq.headers!['mcp-session-id'] as string;
|
||||
|
||||
expect(sessionId).toBe('test-session-123');
|
||||
});
|
||||
|
||||
it('should handle missing session ID in headers', () => {
|
||||
const mockReq = {
|
||||
headers: {
|
||||
'content-type': 'application/json'
|
||||
},
|
||||
body: {}
|
||||
} as Partial<Request>;
|
||||
|
||||
const sessionId = mockReq.headers!['mcp-session-id'] as string | undefined;
|
||||
|
||||
expect(sessionId).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
177
src/modules/mcp/handlers/shttp.ts
Normal file
177
src/modules/mcp/handlers/shttp.ts
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
||||
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
|
||||
import { Request, Response } from "express";
|
||||
import { getShttpTransport, isSessionOwnedBy, redisRelayToMcpServer, ServerRedisTransport, setSessionOwner, shutdownSession } from "../services/redisTransport.js";
|
||||
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
|
||||
import { randomUUID } from "crypto";
|
||||
import { createMcpServer } from "../services/mcp.js";
|
||||
import { logger } from "../../shared/logger.js";
|
||||
|
||||
|
||||
declare module "express-serve-static-core" {
|
||||
interface Request {
|
||||
/**
|
||||
* Information about the validated access token, if the `requireBearerAuth` middleware was used.
|
||||
*/
|
||||
auth?: AuthInfo;
|
||||
}
|
||||
}
|
||||
|
||||
function getUserIdFromAuth(auth?: AuthInfo): string | null {
|
||||
return auth?.extra?.userId as string || null;
|
||||
}
|
||||
|
||||
// TODO: Document Streamable HTTP implementation choices:
|
||||
// 1. STATEFUL: Requires clients to initialize sessions and track session IDs
|
||||
// - First request must be 'initialize' without Mcp-Session-Id header
|
||||
// - Server returns session ID, client must include it in subsequent requests
|
||||
// - Alternative: Could implement STATELESS mode (each request independent)
|
||||
// 2. SSE RESPONSES: Returns results via Server-Sent Events stream, not JSON responses
|
||||
// - Requires Accept: application/json, text/event-stream header
|
||||
// - Responses formatted as: event: message\ndata: {...}
|
||||
// - Alternative: Could use JSON response mode (check StreamableHTTPServerTransport options)
|
||||
|
||||
export async function handleStreamableHTTP(req: Request, res: Response) {
|
||||
let shttpTransport: StreamableHTTPServerTransport | undefined = undefined;
|
||||
|
||||
res.on('finish', async () => {
|
||||
await shttpTransport?.close();
|
||||
});
|
||||
|
||||
const onsessionclosed = async (sessionId: string) => {
|
||||
logger.info('Session closed callback triggered', {
|
||||
sessionId,
|
||||
userId: getUserIdFromAuth(req.auth)
|
||||
});
|
||||
await shutdownSession(sessionId);
|
||||
}
|
||||
|
||||
try {
|
||||
// Check for existing session ID
|
||||
const sessionId = req.headers['mcp-session-id'] as string | undefined;
|
||||
const userId = getUserIdFromAuth(req.auth);
|
||||
|
||||
logger.debug('SHTTP request received', {
|
||||
method: req.method,
|
||||
sessionId,
|
||||
userId,
|
||||
hasAuth: !!req.auth,
|
||||
authExtra: req.auth?.extra
|
||||
});
|
||||
|
||||
// if no userid, return 401, we shouldn't get here ideally
|
||||
if (!userId) {
|
||||
logger.warning('Request without user ID', {
|
||||
sessionId,
|
||||
hasAuth: !!req.auth
|
||||
});
|
||||
res.status(401).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32002,
|
||||
"message": "User ID required"
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const isGetRequest = req.method === 'GET';
|
||||
|
||||
// incorrect session for the authed user, return 401
|
||||
if (sessionId) {
|
||||
if (!(await isSessionOwnedBy(sessionId, userId))) {
|
||||
logger.warning('Session ownership mismatch', {
|
||||
sessionId,
|
||||
userId,
|
||||
requestMethod: req.method
|
||||
});
|
||||
res.status(401).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32001,
|
||||
"message": "Session not found or access denied"
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
// Reuse existing transport for owned session
|
||||
logger.info('Reusing existing session', {
|
||||
sessionId,
|
||||
userId,
|
||||
isGetRequest
|
||||
});
|
||||
shttpTransport = await getShttpTransport(sessionId, onsessionclosed, isGetRequest);
|
||||
} else if (isInitializeRequest(req.body)) {
|
||||
// New initialization request - use JSON response mode
|
||||
logger.debug('Processing initialize request', {
|
||||
body: req.body,
|
||||
userId,
|
||||
headerSessionId: sessionId, // This is the sessionId from header (should be undefined for init)
|
||||
isInitializeRequest: true
|
||||
});
|
||||
|
||||
const onsessioninitialized = async (sessionId: string) => {
|
||||
logger.info('Initializing new session', {
|
||||
sessionId,
|
||||
userId
|
||||
});
|
||||
|
||||
const { server, cleanup: mcpCleanup } = createMcpServer();
|
||||
|
||||
const serverRedisTransport = new ServerRedisTransport(sessionId);
|
||||
serverRedisTransport.onclose = mcpCleanup;
|
||||
await server.connect(serverRedisTransport)
|
||||
|
||||
// Set session ownership
|
||||
await setSessionOwner(sessionId, userId);
|
||||
|
||||
logger.info('Session initialized successfully', {
|
||||
sessionId,
|
||||
userId
|
||||
});
|
||||
}
|
||||
|
||||
const newSessionId = randomUUID();
|
||||
shttpTransport = new StreamableHTTPServerTransport({
|
||||
sessionIdGenerator: () => newSessionId,
|
||||
onsessionclosed,
|
||||
onsessioninitialized,
|
||||
});
|
||||
shttpTransport.onclose = await redisRelayToMcpServer(newSessionId, shttpTransport);
|
||||
} else {
|
||||
// Invalid request - no session ID and not initialization request
|
||||
logger.warning('Invalid request: no session ID and not initialization', {
|
||||
hasSessionId: !!sessionId,
|
||||
isInitRequest: false,
|
||||
userId,
|
||||
method: req.method
|
||||
});
|
||||
res.status(400).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32600,
|
||||
"message": "Invalid request method for existing session"
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
// Handle the request with existing transport - no need to reconnect
|
||||
await shttpTransport.handleRequest(req, res, req.body);
|
||||
} catch (error) {
|
||||
logger.error('Error handling MCP request', error as Error, {
|
||||
sessionId: req.headers['mcp-session-id'] as string | undefined,
|
||||
method: req.method,
|
||||
userId: getUserIdFromAuth(req.auth)
|
||||
});
|
||||
|
||||
if (!res.headersSent) {
|
||||
res.status(500).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32603,
|
||||
"message": "Internal error during request processing"
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
100
src/modules/mcp/handlers/sse.ts
Normal file
100
src/modules/mcp/handlers/sse.ts
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
||||
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
|
||||
import contentType from "content-type";
|
||||
import { Request, Response } from "express";
|
||||
import { redisClient } from "../../shared/redis.js";
|
||||
import { createMcpServer } from "../services/mcp.js";
|
||||
import { logger, logMcpMessage } from "../../shared/logger.js";
|
||||
|
||||
declare module "express-serve-static-core" {
|
||||
interface Request {
|
||||
/**
|
||||
* Information about the validated access token, if the `requireBearerAuth` middleware was used.
|
||||
*/
|
||||
auth?: AuthInfo;
|
||||
}
|
||||
}
|
||||
|
||||
function redisChannelForSession(sessionId: string): string {
|
||||
return `mcp:${sessionId}`;
|
||||
}
|
||||
|
||||
export async function handleSSEConnection(req: Request, res: Response) {
|
||||
const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer();
|
||||
const transport = new SSEServerTransport("/message", res);
|
||||
logger.info('Received MCP SSE connection', {
|
||||
sessionId: transport.sessionId
|
||||
});
|
||||
|
||||
const redisCleanup = await redisClient.createSubscription(
|
||||
redisChannelForSession(transport.sessionId),
|
||||
(json) => {
|
||||
// TODO handle DELETE messages
|
||||
// TODO set timeout to kill the session
|
||||
|
||||
const message = JSON.parse(json);
|
||||
logMcpMessage(message, transport.sessionId);
|
||||
transport.handleMessage(message).catch((error) => {
|
||||
logger.error('Error handling message', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
});
|
||||
});
|
||||
},
|
||||
(error) => {
|
||||
logger.error('Disconnecting due to error in Redis subscriber', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
});
|
||||
transport
|
||||
.close()
|
||||
.catch((error) =>
|
||||
logger.error('Error closing transport', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
}),
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
const cleanup = () => {
|
||||
void mcpCleanup();
|
||||
redisCleanup().catch((error) =>
|
||||
logger.error('Error disconnecting Redis subscriber', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Clean up Redis subscription when the connection closes
|
||||
mcpServer.onclose = cleanup
|
||||
|
||||
logger.info('Listening on Redis channel', {
|
||||
sessionId: transport.sessionId,
|
||||
channel: redisChannelForSession(transport.sessionId)
|
||||
});
|
||||
await mcpServer.connect(transport);
|
||||
}
|
||||
|
||||
export async function handleMessage(req: Request, res: Response) {
|
||||
const sessionId = req.query.sessionId;
|
||||
let body: string;
|
||||
try {
|
||||
if (typeof sessionId !== "string") {
|
||||
throw new Error("Only one sessionId allowed");
|
||||
}
|
||||
|
||||
const ct = contentType.parse(req.headers["content-type"] ?? "");
|
||||
if (ct.type !== "application/json") {
|
||||
throw new Error(`Unsupported content-type: ${ct}`);
|
||||
}
|
||||
|
||||
body = JSON.stringify(req.body);
|
||||
} catch (error) {
|
||||
res.status(400).json({ error: "Bad request" });
|
||||
logger.error('Bad POST request', error as Error, {
|
||||
sessionId,
|
||||
contentType: req.headers['content-type']
|
||||
});
|
||||
return;
|
||||
}
|
||||
await redisClient.publish(redisChannelForSession(sessionId), body);
|
||||
res.status(202).end();
|
||||
}
|
||||
Loading…
Reference in a new issue