mirror of
https://github.com/samanhappy/mcphub.git
synced 2025-12-23 18:29:21 -05:00
@@ -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<string>('-');
|
||||
const [enableSessionRebuild, setEnableSessionRebuild] = useState<boolean>(false);
|
||||
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [error, setError] = useState<string | null>(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,
|
||||
};
|
||||
};
|
||||
|
||||
@@ -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 = () => {
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center justify-between p-3 bg-gray-50 rounded-md">
|
||||
<div>
|
||||
<h3 className="font-medium text-gray-700">{t('settings.enableSessionRebuild')}</h3>
|
||||
<p className="text-sm text-gray-500">{t('settings.enableSessionRebuildDescription')}</p>
|
||||
</div>
|
||||
<Switch
|
||||
disabled={loading}
|
||||
checked={enableSessionRebuild}
|
||||
onCheckedChange={(checked) => updateSessionRebuild(checked)}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -576,6 +576,8 @@
|
||||
"systemSettings": "系统设置",
|
||||
"nameSeparatorLabel": "名称分隔符",
|
||||
"nameSeparatorDescription": "用于分隔服务器名称和工具/提示名称(默认:-)",
|
||||
"enableSessionRebuild": "启用服务端会话重建",
|
||||
"enableSessionRebuildDescription": "开启后会应用服务端会话重建的改进代码,提供更好的会话管理体验",
|
||||
"restartRequired": "配置已保存。为确保所有服务正确加载新设置,建议重启应用。",
|
||||
"exportMcpSettings": "导出配置",
|
||||
"mcpSettingsJson": "MCP 配置 JSON",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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> = {}): 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<typeof loadSettings>).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<typeof StreamableHTTPServerTransport>).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<typeof loadSettings>).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<typeof loadSettings>).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<any>).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<typeof loadSettings>).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<typeof loadSettings>).mockReturnValue({
|
||||
mcpServers: {},
|
||||
|
||||
@@ -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<StreamableHTTPServerTransport> } = {};
|
||||
|
||||
export const getGroup = (sessionId: string): string => {
|
||||
return transports[sessionId]?.group || '';
|
||||
@@ -144,6 +147,76 @@ export const handleSseMessage = async (req: Request, res: Response): Promise<voi
|
||||
}
|
||||
};
|
||||
|
||||
// Helper function to create a session with a specific sessionId
|
||||
async function createSessionWithId(sessionId: string, group: string, username?: string): Promise<StreamableHTTPServerTransport> {
|
||||
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<StreamableHTTPServerTransport> {
|
||||
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<void> => {
|
||||
// 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 };
|
||||
},
|
||||
});
|
||||
let transportInfo: typeof transports[string] | undefined;
|
||||
|
||||
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}`);
|
||||
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);
|
||||
};
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
});
|
||||
|
||||
|
||||
@@ -32,6 +32,7 @@ export const createMockSettings = (overrides: Partial<McpSettings> = {}): McpSet
|
||||
enableBearerAuth: true,
|
||||
bearerAuthKey: 'test-auth-token-123',
|
||||
},
|
||||
enableSessionRebuild: false,
|
||||
} as SystemConfig,
|
||||
users: [
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user