diff --git a/frontend/src/hooks/useSettingsData.ts b/frontend/src/hooks/useSettingsData.ts index 8b4342f..ff6490c 100644 --- a/frontend/src/hooks/useSettingsData.ts +++ b/frontend/src/hooks/useSettingsData.ts @@ -41,6 +41,7 @@ interface SystemSettings { smartRouting?: SmartRoutingConfig; mcpRouter?: MCPRouterConfig; nameSeparator?: string; + enableSessionRebuild?: boolean; }; } @@ -86,6 +87,7 @@ export const useSettingsData = () => { }); const [nameSeparator, setNameSeparator] = useState('-'); + const [enableSessionRebuild, setEnableSessionRebuild] = useState(false); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); @@ -141,6 +143,9 @@ export const useSettingsData = () => { if (data.success && data.data?.systemConfig?.nameSeparator !== undefined) { setNameSeparator(data.data.systemConfig.nameSeparator); } + if (data.success && data.data?.systemConfig?.enableSessionRebuild !== undefined) { + setEnableSessionRebuild(data.data.systemConfig.enableSessionRebuild); + } } catch (error) { console.error('Failed to fetch settings:', error); setError(error instanceof Error ? error.message : 'Failed to fetch settings'); @@ -420,6 +425,36 @@ export const useSettingsData = () => { } }; + // Update session rebuild setting + const updateSessionRebuild = async (value: boolean) => { + setLoading(true); + setError(null); + + try { + const data = await apiPut('/system-config', { + enableSessionRebuild: value, + }); + + if (data.success) { + setEnableSessionRebuild(value); + showToast(t('settings.restartRequired'), 'info'); + return true; + } else { + showToast(data.message || t('errors.failedToUpdateSystemConfig')); + return false; + } + } catch (error) { + console.error('Failed to update session rebuild setting:', error); + const errorMessage = + error instanceof Error ? error.message : 'Failed to update session rebuild setting'; + setError(errorMessage); + showToast(errorMessage); + return false; + } finally { + setLoading(false); + } + }; + const exportMCPSettings = async (serverName?: string) => { setLoading(true); setError(null); @@ -456,6 +491,7 @@ export const useSettingsData = () => { smartRoutingConfig, mcpRouterConfig, nameSeparator, + enableSessionRebuild, loading, error, setError, @@ -469,6 +505,7 @@ export const useSettingsData = () => { updateMCPRouterConfig, updateMCPRouterConfigBatch, updateNameSeparator, + updateSessionRebuild, exportMCPSettings, }; }; diff --git a/frontend/src/pages/SettingsPage.tsx b/frontend/src/pages/SettingsPage.tsx index 0e1449a..ccfc557 100644 --- a/frontend/src/pages/SettingsPage.tsx +++ b/frontend/src/pages/SettingsPage.tsx @@ -59,6 +59,7 @@ const SettingsPage: React.FC = () => { smartRoutingConfig, mcpRouterConfig, nameSeparator, + enableSessionRebuild, loading, updateRoutingConfig, updateRoutingConfigBatch, @@ -67,6 +68,7 @@ const SettingsPage: React.FC = () => { updateSmartRoutingConfigBatch, updateMCPRouterConfig, updateNameSeparator, + updateSessionRebuild, exportMCPSettings, } = useSettingsData() @@ -599,6 +601,18 @@ const SettingsPage: React.FC = () => { + +
+
+

{t('settings.enableSessionRebuild')}

+

{t('settings.enableSessionRebuildDescription')}

+
+ updateSessionRebuild(checked)} + /> +
)} diff --git a/locales/en.json b/locales/en.json index 20080c5..e633823 100644 --- a/locales/en.json +++ b/locales/en.json @@ -574,6 +574,8 @@ "systemSettings": "System Settings", "nameSeparatorLabel": "Name Separator", "nameSeparatorDescription": "Character used to separate server name and tool/prompt name (default: -)", + "enableSessionRebuild": "Enable Server Session Rebuild", + "enableSessionRebuildDescription": "When enabled, applies the improved server session rebuild code for better session management experience", "restartRequired": "Configuration saved. It is recommended to restart the application to ensure all services load the new settings correctly.", "exportMcpSettings": "Export Settings", "mcpSettingsJson": "MCP Settings JSON", diff --git a/locales/fr.json b/locales/fr.json index 793b90b..104233f 100644 --- a/locales/fr.json +++ b/locales/fr.json @@ -574,6 +574,8 @@ "systemSettings": "Paramètres système", "nameSeparatorLabel": "Séparateur de noms", "nameSeparatorDescription": "Caractère utilisé pour séparer le nom du serveur et le nom de l'outil/prompt (par défaut : -)", + "enableSessionRebuild": "Activer la reconstruction de session serveur", + "enableSessionRebuildDescription": "Lorsqu'il est activé, applique le code de reconstruction de session serveur amélioré pour une meilleure expérience de gestion de session", "restartRequired": "Configuration enregistrée. Il est recommandé de redémarrer l'application pour s'assurer que tous les services chargent correctement les nouveaux paramètres.", "exportMcpSettings": "Exporter les paramètres", "mcpSettingsJson": "JSON des paramètres MCP", diff --git a/locales/tr.json b/locales/tr.json index d67e9de..e9474e7 100644 --- a/locales/tr.json +++ b/locales/tr.json @@ -574,6 +574,8 @@ "systemSettings": "Sistem Ayarları", "nameSeparatorLabel": "İsim Ayırıcı", "nameSeparatorDescription": "Sunucu adı ile araç/istek adını ayırmak için kullanılan karakter (varsayılan: -)", + "enableSessionRebuild": "Sunucu Oturum Yeniden Oluşturmayı Etkinleştir", + "enableSessionRebuildDescription": "Etkinleştirildiğinde, daha iyi oturum yönetimi deneyimi için geliştirilmiş sunucu oturum yeniden oluşturma kodunu uygular", "restartRequired": "Yapılandırma kaydedildi. Tüm hizmetlerin yeni ayarları doğru şekilde yüklemesini sağlamak için uygulamayı yeniden başlatmanız önerilir.", "exportMcpSettings": "Ayarları Dışa Aktar", "mcpSettingsJson": "MCP Ayarları JSON", diff --git a/locales/zh.json b/locales/zh.json index 261b8e9..746725a 100644 --- a/locales/zh.json +++ b/locales/zh.json @@ -576,6 +576,8 @@ "systemSettings": "系统设置", "nameSeparatorLabel": "名称分隔符", "nameSeparatorDescription": "用于分隔服务器名称和工具/提示名称（默认：-）", + "enableSessionRebuild": "启用服务端会话重建", + "enableSessionRebuildDescription": "开启后会应用服务端会话重建的改进代码，提供更好的会话管理体验", "restartRequired": "配置已保存。为确保所有服务正确加载新设置，建议重启应用。", "exportMcpSettings": "导出配置", "mcpSettingsJson": "MCP 配置 JSON", diff --git a/src/controllers/serverController.ts b/src/controllers/serverController.ts index c36413c..0911b2e 100644 --- a/src/controllers/serverController.ts +++ b/src/controllers/serverController.ts @@ -508,7 +508,7 @@ export const updateToolDescription = async (req: Request, res: Response): Promis export const updateSystemConfig = (req: Request, res: Response): void => { try { - const { routing, install, smartRouting, mcpRouter, nameSeparator } = req.body; + const { routing, install, smartRouting, mcpRouter, nameSeparator, enableSessionRebuild } = req.body; const currentUser = (req as any).user; if ( @@ -533,7 +533,8 @@ export const updateSystemConfig = (req: Request, res: Response): void => { typeof mcpRouter.referer !== 'string' && typeof mcpRouter.title !== 'string' && typeof mcpRouter.baseUrl !== 'string')) && - typeof nameSeparator !== 'string' + typeof nameSeparator !== 'string' && + typeof enableSessionRebuild !== 'boolean' ) { res.status(400).json({ success: false, @@ -719,6 +720,10 @@ export const updateSystemConfig = (req: Request, res: Response): void => { settings.systemConfig.nameSeparator = nameSeparator; } + if (typeof enableSessionRebuild === 'boolean') { + settings.systemConfig.enableSessionRebuild = enableSessionRebuild; + } + if (saveSettings(settings, currentUser)) { res.json({ success: true, diff --git a/src/services/sseService.test.ts b/src/services/sseService.test.ts index 9a70793..d93a710 100644 --- a/src/services/sseService.test.ts +++ b/src/services/sseService.test.ts @@ -7,6 +7,7 @@ import { handleMcpOtherRequest, getGroup, getConnectionCount, + transports, } from './sseService.js'; // Mock dependencies @@ -33,6 +34,7 @@ jest.mock('../config/index.js', () => { enableBearerAuth: false, bearerAuthKey: 'test-key', }, + enableSessionRebuild: false, // Default to false for tests }, })), }; @@ -55,12 +57,7 @@ jest.mock('@modelcontextprotocol/sdk/server/sse.js', () => ({ })); jest.mock('@modelcontextprotocol/sdk/server/streamableHttp.js', () => ({ - StreamableHTTPServerTransport: jest.fn().mockImplementation(() => ({ - sessionId: 'test-session-id', - connect: jest.fn(), - handleRequest: jest.fn(), - onclose: null, - })), + StreamableHTTPServerTransport: jest.fn().mockImplementation(() => mockStreamableHTTPServerTransport), })); jest.mock('@modelcontextprotocol/sdk/types.js', () => ({ @@ -74,6 +71,14 @@ import { UserContextService } from './userContextService.js'; import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'; import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; +// Create mock instances for testing +const mockStreamableHTTPServerTransport = { + sessionId: 'test-session-id', + connect: jest.fn(), + handleRequest: jest.fn(), + onclose: null, +}; + // Mock Express Request and Response const createMockRequest = (overrides: Partial = {}): Request => ({ @@ -108,6 +113,7 @@ describe('sseService', () => { enableBearerAuth: false, bearerAuthKey: 'test-key', }, + enableSessionRebuild: false, // Default to false for tests }, }); }); @@ -383,7 +389,7 @@ describe('sseService', () => { expect(getMcpServer).toHaveBeenCalled(); }); - it('should return error for invalid session', async () => { + it('should return error when session rebuild is disabled and session is invalid', async () => { const req = createMockRequest({ params: { group: 'test-group' }, headers: { 'mcp-session-id': 'invalid-session' }, @@ -393,6 +399,7 @@ describe('sseService', () => { await handleMcpPostRequest(req, res); + // When session rebuild is disabled, invalid sessions should return an error expect(res.status).toHaveBeenCalledWith(400); expect(res.json).toHaveBeenCalledWith({ jsonrpc: '2.0', @@ -404,6 +411,36 @@ describe('sseService', () => { }); }); + it('should transparently rebuild invalid session when enabled', async () => { + // Enable session rebuild for this test + (loadSettings as jest.MockedFunction).mockReturnValue({ + mcpServers: {}, + systemConfig: { + routing: { + enableGlobalRoute: true, + enableGroupNameRoute: true, + enableBearerAuth: false, + bearerAuthKey: 'test-key', + }, + enableSessionRebuild: true, // Enable session rebuild + }, + }); + + const req = createMockRequest({ + params: { group: 'test-group' }, + headers: { 'mcp-session-id': 'invalid-session' }, + body: { method: 'someMethod' }, + }); + const res = createMockResponse(); + + await handleMcpPostRequest(req, res); + + // With session rebuild enabled, invalid sessions should be transparently rebuilt + expect(StreamableHTTPServerTransport).toHaveBeenCalled(); + const mockInstance = (StreamableHTTPServerTransport as jest.MockedClass).mock.results[0].value; + expect(mockInstance.handleRequest).toHaveBeenCalledWith(req, res, req.body); + }); + it('should return 401 when bearer auth fails', async () => { (loadSettings as jest.MockedFunction).mockReturnValue({ mcpServers: {}, @@ -442,19 +479,79 @@ describe('sseService', () => { expect(res.status).toHaveBeenCalledWith(400); expect(res.send).toHaveBeenCalledWith('Invalid or missing session ID'); }); + it('should return error when session rebuild is disabled in handleMcpOtherRequest', async () => { + // Clear transports before test + Object.keys(transports).forEach(key => delete transports[key]); + + // Enable bearer auth for this test + (loadSettings as jest.MockedFunction).mockReturnValue({ + mcpServers: {}, + systemConfig: { + routing: { + enableGlobalRoute: true, + enableGroupNameRoute: true, + enableBearerAuth: true, + bearerAuthKey: 'test-key', + }, + enableSessionRebuild: false, // Disable session rebuild + }, + }); + + // Mock user context to exist + const mockGetCurrentUser = jest.fn(() => ({ username: 'testuser' })); + (UserContextService.getInstance as jest.MockedFunction).mockReturnValue({ + getCurrentUser: mockGetCurrentUser, + }); - it('should return 400 for invalid session ID', async () => { const req = createMockRequest({ - headers: { 'mcp-session-id': 'invalid-session' }, + headers: { + 'mcp-session-id': 'invalid-session', + 'authorization': 'Bearer test-key' + }, + params: { group: 'test-group' }, }); const res = createMockResponse(); await handleMcpOtherRequest(req, res); + // Should return 400 error when session rebuild is disabled expect(res.status).toHaveBeenCalledWith(400); expect(res.send).toHaveBeenCalledWith('Invalid or missing session ID'); }); + it('should transparently rebuild invalid session in handleMcpOtherRequest when enabled', async () => { + // Enable bearer auth and session rebuild for this test + (loadSettings as jest.MockedFunction).mockReturnValue({ + mcpServers: {}, + systemConfig: { + routing: { + enableGlobalRoute: true, + enableGroupNameRoute: true, + enableBearerAuth: true, + bearerAuthKey: 'test-key', + }, + enableSessionRebuild: true, // Enable session rebuild + }, + }); + + const req = createMockRequest({ + headers: { + 'mcp-session-id': 'invalid-session', + 'authorization': 'Bearer test-key' + }, + }); + const res = createMockResponse(); + + await handleMcpOtherRequest(req, res); + + // Should not return 400 error, but instead transparently rebuild the session + expect(res.status).not.toHaveBeenCalledWith(400); + expect(res.send).not.toHaveBeenCalledWith('Invalid or missing session ID'); + + // Should attempt to handle the request (session was rebuilt) + expect(mockStreamableHTTPServerTransport.handleRequest).toHaveBeenCalled(); + }); + it('should return 401 when bearer auth fails', async () => { (loadSettings as jest.MockedFunction).mockReturnValue({ mcpServers: {}, diff --git a/src/services/sseService.ts b/src/services/sseService.ts index 05faa11..4fe22bc 100644 --- a/src/services/sseService.ts +++ b/src/services/sseService.ts @@ -10,7 +10,10 @@ import config from '../config/index.js'; import { UserContextService } from './userContextService.js'; import { RequestContextService } from './requestContextService.js'; -const transports: { [sessionId: string]: { transport: Transport; group: string } } = {}; +export const transports: { [sessionId: string]: { transport: Transport; group: string; needsInitialization?: boolean } } = {}; + +// Session creation locks to prevent concurrent session creation conflicts +const sessionCreationLocks: { [sessionId: string]: Promise } = {}; export const getGroup = (sessionId: string): string => { return transports[sessionId]?.group || ''; @@ -144,6 +147,76 @@ export const handleSseMessage = async (req: Request, res: Response): Promise { + console.log(`[SESSION REBUILD] Starting session rebuild for ID: ${sessionId}${username ? ` for user: ${username}` : ''}`); + + // Create a new server instance to ensure clean state + const server = getMcpServer(sessionId, group); + + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => sessionId, // Use the specified sessionId + onsessioninitialized: (initializedSessionId) => { + console.log(`[SESSION REBUILD] onsessioninitialized triggered for ID: ${initializedSessionId}`); // New log + if (initializedSessionId === sessionId) { + transports[sessionId] = { transport, group }; + console.log(`[SESSION REBUILD] Session ${sessionId} initialized successfully${username ? ` for user: ${username}` : ''}`); + } else { + console.warn(`[SESSION REBUILD] Session ID mismatch: expected ${sessionId}, got ${initializedSessionId}`); + } + }, + }); + + transport.onclose = () => { + console.log(`[SESSION REBUILD] Transport closed: ${sessionId}`); + delete transports[sessionId]; + deleteMcpServer(sessionId); + }; + + // Connect to MCP server + await server.connect(transport); + + // Wait for the server to fully initialize + await new Promise(resolve => setTimeout(resolve, 500)); + + // Ensure the transport is properly initialized + if (!transports[sessionId]) { + console.warn(`[SESSION REBUILD] Transport not found in transports after initialization, forcing registration`); + transports[sessionId] = { transport, group, needsInitialization: true }; + } else { + // Mark the session as needing initialization + transports[sessionId].needsInitialization = true; + } + + console.log(`[SESSION REBUILD] Session ${sessionId} created but not yet initialized. It will be initialized on first use.`); + + console.log(`[SESSION REBUILD] Successfully rebuilt session ${sessionId} in group: ${group}`); + return transport; +} +// Helper function to create a completely new session +async function createNewSession(group: string, username?: string): Promise { + const newSessionId = randomUUID(); + console.log(`[SESSION NEW] Creating new session with ID: ${newSessionId}${username ? ` for user: ${username}` : ''}`); + + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => newSessionId, + onsessioninitialized: (sessionId) => { + transports[sessionId] = { transport, group }; + console.log(`[SESSION NEW] New session ${sessionId} initialized successfully${username ? ` for user: ${username}` : ''}`); + }, + }); + + transport.onclose = () => { + console.log(`[SESSION NEW] Transport closed: ${newSessionId}`); + delete transports[newSessionId]; + deleteMcpServer(newSessionId); + }; + + await getMcpServer(newSessionId, group).connect(transport); + console.log(`[SESSION NEW] Successfully created new session ${newSessionId} in group: ${group}`); + return transport; +} + export const handleMcpPostRequest = async (req: Request, res: Response): Promise => { // User context is now set by sseUserContextMiddleware const userContextService = UserContextService.getInstance(); @@ -175,31 +248,63 @@ export const handleMcpPostRequest = async (req: Request, res: Response): Promise } let transport: StreamableHTTPServerTransport; - if (sessionId && transports[sessionId]) { - console.log(`Reusing existing transport for sessionId: ${sessionId}`); - transport = transports[sessionId].transport as StreamableHTTPServerTransport; - } else if (!sessionId && isInitializeRequest(req.body)) { - transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: () => randomUUID(), - onsessioninitialized: (sessionId) => { - transports[sessionId] = { transport, group }; - }, - }); - - transport.onclose = () => { - console.log(`Transport closed: ${transport.sessionId}`); - if (transport.sessionId) { - delete transports[transport.sessionId]; - deleteMcpServer(transport.sessionId); - console.log(`MCP connection closed: ${transport.sessionId}`); + let transportInfo: typeof transports[string] | undefined; + + if (sessionId) { + transportInfo = transports[sessionId]; + } + + if (sessionId && transportInfo) { + // Case 1: Session exists and is valid, reuse it + console.log(`[SESSION REUSE] Reusing existing session: ${sessionId}${username ? ` for user: ${username}` : ''}`); + transport = transportInfo.transport as StreamableHTTPServerTransport; + } else if (sessionId) { + // Case 2: SessionId exists but transport is missing (server restart), check if session rebuild is enabled + const settings = loadSettings(); + const enableSessionRebuild = settings.systemConfig?.enableSessionRebuild || false; + + if (enableSessionRebuild) { + console.log(`[SESSION AUTO-REBUILD] Session ${sessionId} not found, initiating transparent rebuild${username ? ` for user: ${username}` : ''}`); + // Prevent concurrent session creation + if (sessionCreationLocks[sessionId] !== undefined) { + console.log(`[SESSION AUTO-REBUILD] Session creation in progress for ${sessionId}, waiting...`); + transport = await sessionCreationLocks[sessionId]; + } else { + sessionCreationLocks[sessionId] = createSessionWithId(sessionId, group, username); + try { + transport = await sessionCreationLocks[sessionId]; + console.log(`[SESSION AUTO-REBUILD] Successfully transparently rebuilt session: ${sessionId}`); + } catch (error) { + console.error(`[SESSION AUTO-REBUILD] Failed to rebuild session ${sessionId}:`, error); + throw error; + } finally { + delete sessionCreationLocks[sessionId]; + } } - }; - - console.log( - `MCP connection established: ${transport.sessionId}${username ? ` for user: ${username}` : ''}`, - ); - await getMcpServer(transport.sessionId, group).connect(transport); + // Get the updated transport info after rebuild + if (sessionId) { + transportInfo = transports[sessionId]; + } + } else { + // Session rebuild is disabled, return error + console.warn(`[SESSION ERROR] Session ${sessionId} not found and session rebuild is disabled${username ? ` for user: ${username}` : ''}`); + res.status(400).json({ + jsonrpc: '2.0', + error: { + code: -32000, + message: 'Bad Request: No valid session ID provided', + }, + id: null, + }); + return; + } + } else if (isInitializeRequest(req.body)) { + // Case 3: No sessionId and this is an initialize request, create new session + console.log(`[SESSION CREATE] No session ID provided for initialize request, creating new session${username ? ` for user: ${username}` : ''}`); + transport = await createNewSession(group, username); } else { + // Case 4: No sessionId and not an initialize request, return error + console.warn(`[SESSION ERROR] No session ID provided for non-initialize request (method: ${req.body?.method})${username ? ` for user: ${username}` : ''}`); res.status(400).json({ jsonrpc: '2.0', error: { @@ -217,8 +322,118 @@ export const handleMcpPostRequest = async (req: Request, res: Response): Promise const requestContextService = RequestContextService.getInstance(); requestContextService.setRequestContext(req); + // Check if the session needs initialization (for rebuilt sessions) + if (transportInfo && transportInfo.needsInitialization) { + console.log(`[MCP] Session ${sessionId} needs initialization, performing proactive initialization`); + + try { + // Create a mock response object that doesn't actually send headers + const mockRes = { + writeHead: () => {}, + end: () => {}, + json: () => {}, + status: () => mockRes, + send: () => {}, + headersSent: false + } as any; + + // First, send the initialize request + const initializeRequest = { + method: 'initialize', + params: { + protocolVersion: '2025-03-26', + capabilities: {}, + clientInfo: { + name: 'MCPHub-Client', + version: '1.0.0' + } + }, + jsonrpc: '2.0', + id: `init-${sessionId}-${Date.now()}` + }; + + console.log(`[MCP] Sending initialize request for session ${sessionId}`); + // Use mock response to avoid sending actual HTTP response + await transport.handleRequest(req, mockRes, initializeRequest); + + // Then send the initialized notification + const initializedNotification = { + method: 'notifications/initialized', + jsonrpc: '2.0' + }; + + console.log(`[MCP] Sending initialized notification for session ${sessionId}`); + await transport.handleRequest(req, mockRes, initializedNotification); + + // Mark the session as initialized + transportInfo.needsInitialization = false; + console.log(`[MCP] Session ${sessionId} successfully initialized`); + } catch (initError) { + console.error(`[MCP] Failed to initialize session ${sessionId}:`, initError); + console.error(`[MCP] Initialization error details:`, initError); + // Don't return here, continue with the original request + } + } + try { await transport.handleRequest(req, res, req.body); + } catch (error: any) { + // Check if this is a "Server not initialized" error for a newly rebuilt session + if (sessionId && error.message && error.message.includes('Server not initialized')) { + console.log(`[SESSION AUTO-REBUILD] Server not initialized for ${sessionId}. Attempting to initialize with the current request.`); + + // Check if the current request is an 'initialize' request + if (isInitializeRequest(req.body)) { + // If it is, we can just retry it. The transport should now be in the transports map. + console.log(`[SESSION AUTO-REBUILD] Retrying initialize request for ${sessionId}.`); + await transport.handleRequest(req, res, req.body); + } else { + // If not, we need to send an initialize request first. + // We construct a mock initialize request, but use the REAL req/res objects. + const initializeRequest = { + method: 'initialize', + params: { + protocolVersion: '2025-03-26', + capabilities: {}, + clientInfo: { + name: 'MCPHub-Client', + version: '1.0.0' + } + }, + jsonrpc: '2.0', + id: `init-${sessionId}-${Date.now()}` + }; + + console.log(`[SESSION AUTO-REBUILD] Sending initialize request for ${sessionId} before handling the actual request.`); + try { + // Temporarily replace the body to send the initialize request + const originalBody = req.body; + req.body = initializeRequest; + await transport.handleRequest(req, res, req.body); + + // Now, send the notifications/initialized + const initializedNotification = { + method: 'notifications/initialized', + jsonrpc: '2.0' + }; + req.body = initializedNotification; + await transport.handleRequest(req, res, req.body); + + // Restore the original body and retry the original request + req.body = originalBody; + console.log(`[SESSION AUTO-REBUILD] Initialization complete for ${sessionId}. Retrying original request.`); + await transport.handleRequest(req, res, req.body); + + } catch (initError) { + console.error(`[SESSION AUTO-REBUILD] Failed to initialize session ${sessionId} on-the-fly:`, initError); + // Re-throw the original error if initialization fails + throw error; + } + } + } else { + // If it's a different error, just re-throw it + throw error; + } } finally { // Clean up request context after handling requestContextService.clearRequestContext(); @@ -240,12 +455,51 @@ export const handleMcpOtherRequest = async (req: Request, res: Response) => { } const sessionId = req.headers['mcp-session-id'] as string | undefined; - if (!sessionId || !transports[sessionId]) { + if (!sessionId) { res.status(400).send('Invalid or missing session ID'); return; } - const { transport } = transports[sessionId]; + let transportEntry = transports[sessionId]; + + // If session doesn't exist, attempt transparent rebuild if enabled + if (!transportEntry) { + const settings = loadSettings(); + const enableSessionRebuild = settings.systemConfig?.enableSessionRebuild || false; + + if (enableSessionRebuild) { + console.log(`[SESSION AUTO-REBUILD] Session ${sessionId} not found in handleMcpOtherRequest, initiating transparent rebuild`); + + try { + // Check if user context exists + if (!currentUser) { + res.status(401).send('User context not found'); + return; + } + + // Create session with same ID using existing function + const group = req.params.group; + const rebuiltSession = await createSessionWithId(sessionId, group, currentUser.username); + if (rebuiltSession) { + console.log(`[SESSION AUTO-REBUILD] Successfully transparently rebuilt session: ${sessionId}`); + transportEntry = transports[sessionId]; + } + } catch (error) { + console.error(`[SESSION AUTO-REBUILD] Failed to rebuild session ${sessionId}:`, error); + } + } else { + console.warn(`[SESSION ERROR] Session ${sessionId} not found and session rebuild is disabled in handleMcpOtherRequest`); + res.status(400).send('Invalid or missing session ID'); + return; + } + } + + if (!transportEntry) { + res.status(400).send('Invalid or missing session ID'); + return; + } + + const { transport } = transportEntry; await (transport as StreamableHTTPServerTransport).handleRequest(req, res); }; diff --git a/src/types/index.ts b/src/types/index.ts index 44baf23..93ff96f 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -171,6 +171,7 @@ export interface SystemConfig { }; nameSeparator?: string; // Separator used between server name and tool/prompt name (default: '-') oauth?: OAuthProviderConfig; // OAuth provider configuration for upstream MCP servers + enableSessionRebuild?: boolean; // Controls whether server session rebuild is enabled } export interface UserConfig { diff --git a/tests/services/oauthService.test.ts b/tests/services/oauthService.test.ts index 397ad7f..8201004 100644 --- a/tests/services/oauthService.test.ts +++ b/tests/services/oauthService.test.ts @@ -45,6 +45,7 @@ describe('OAuth Service', () => { tokenUrl: 'http://auth.example.com/token', }, }, + enableSessionRebuild: false, }, }); @@ -55,7 +56,9 @@ describe('OAuth Service', () => { it('should not initialize OAuth when not configured', () => { mockLoadSettings.mockReturnValue({ mcpServers: {}, - systemConfig: {}, + systemConfig: { + enableSessionRebuild: false, + }, }); initOAuthProvider(); @@ -80,6 +83,7 @@ describe('OAuth Service', () => { }, ], }, + enableSessionRebuild: false, }, }); diff --git a/tests/utils/mockSettings.ts b/tests/utils/mockSettings.ts index 87d09f1..187b6aa 100644 --- a/tests/utils/mockSettings.ts +++ b/tests/utils/mockSettings.ts @@ -32,6 +32,7 @@ export const createMockSettings = (overrides: Partial = {}): McpSet enableBearerAuth: true, bearerAuthKey: 'test-auth-token-123', }, + enableSessionRebuild: false, } as SystemConfig, users: [ {