diff --git a/README.md b/README.md index ff84ca5..bce59cb 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,14 @@ tgcli config set mcp.port 8080 Then run `tgcli server` and point your client at the configured address. +One server can serve multiple MCP clients at the same time. Each client gets its +own MCP session, while all sessions for that account share one Telegram client +and one archive service. Sync jobs still run through a single sequential queue. +During shutdown, the server lets active MCP tool calls finish before closing +their sessions and the shared services. This concurrency model applies to MCP +clients connected to the same server process; separate `tgcli` processes still +open the account store directly. + ### Telegram proxy Telegram traffic can be routed through a SOCKS5, HTTP, or MTProto proxy. The diff --git a/mcp-server.js b/mcp-server.js index d9ea2e2..e7473f9 100644 --- a/mcp-server.js +++ b/mcp-server.js @@ -25,7 +25,7 @@ const mcpEnabled = Boolean(mcpConfig.enabled); const resolvedHost = mcpConfig.host ?? process.env.MCP_HOST ?? process.env.FASTMCP_HOST ?? "127.0.0.1"; const resolvedPort = Number(mcpConfig.port ?? process.env.MCP_PORT ?? process.env.FASTMCP_PORT ?? "8080"); const HOST = resolvedHost; -const PORT = Number.isFinite(resolvedPort) && resolvedPort > 0 ? resolvedPort : 8080; +const PORT = Number.isFinite(resolvedPort) && resolvedPort >= 0 ? resolvedPort : 8080; const { telegramClient, messageSyncService } = createServices({ storeDir, config }); let telegramReady = false; @@ -90,7 +90,19 @@ async function initializeTelegram() { * Represents an active MCP session – a transport plus its server instance. */ const sessions = new Map(); +const activeRequests = new Set(); let shuttingDown = false; +let shutdownPromise = null; + +async function trackRequest(task) { + const request = task(); + activeRequests.add(request); + try { + return await request; + } finally { + activeRequests.delete(request); + } +} function closeSessionRecord(record, context) { if (!record || record.closing) { @@ -2117,6 +2129,10 @@ let httpServer = null; if (mcpEnabled) { httpServer = http.createServer(async (req, res) => { try { + if (shuttingDown) { + res.writeHead(503).end("Server is shutting down"); + return; + } const url = new URL(req.url ?? "", `http://${req.headers.host ?? `${HOST}:${PORT}`}`); if (req.method === "OPTIONS") { @@ -2132,7 +2148,7 @@ if (mcpEnabled) { } if (req.method === "POST" && url.pathname === "/mcp") { - await handlePost(req, res); + await trackRequest(() => handlePost(req, res)); return; } @@ -2174,7 +2190,7 @@ if (mcpEnabled) { }); httpServer.listen(PORT, HOST, () => { - console.log(`[startup] MCP HTTP server listening on http://${HOST}:${PORT}/mcp`); + console.log(`[startup] MCP HTTP server listening on http://${HOST}:${httpServer.address().port}/mcp`); }); httpServer.on("error", (error) => { @@ -2184,12 +2200,19 @@ if (mcpEnabled) { console.log("[startup] MCP disabled; running sync-only service."); } -async function shutdown() { - if (shuttingDown) { - return; - } +async function performShutdown() { shuttingDown = true; console.log("[shutdown] received termination signal, closing resources..."); + const httpClosed = httpServer + ? new Promise((resolve) => httpServer.close(() => { + console.log("[shutdown] HTTP server closed"); + resolve(); + })) + : Promise.resolve(); + + // Let in-flight tool calls finish before closing their transports or the + // shared Telegram client and archive database. + await Promise.allSettled(Array.from(activeRequests)); const closeTasks = []; for (const record of sessions.values()) { const task = closeSessionRecord(record, "shutdown"); @@ -2200,12 +2223,11 @@ async function shutdown() { if (closeTasks.length) { await Promise.allSettled(closeTasks); } - if (httpServer) { - httpServer.closeAllConnections?.(); - httpServer.close(() => { - console.log("[shutdown] HTTP server closed"); - }); - } + // The SDK can leave an HTTP stream open after its transport has closed. + // Active tool responses have finished by this point, so remaining sockets + // can be closed without interrupting work. + httpServer?.closeAllConnections?.(); + await httpClosed; try { await messageSyncService.shutdown(); @@ -2225,9 +2247,16 @@ async function shutdown() { }); } +function shutdown() { + shutdownPromise ??= performShutdown(); + return shutdownPromise; +} + const handleShutdownSignal = () => { void shutdown().finally(() => process.exit(0)); }; process.prependListener("SIGINT", handleShutdownSignal); process.prependListener("SIGTERM", handleShutdownSignal); + +export { httpServer, shutdown }; diff --git a/tests/mcp-concurrency.test.js b/tests/mcp-concurrency.test.js new file mode 100644 index 0000000..a593ecb --- /dev/null +++ b/tests/mcp-concurrency.test.js @@ -0,0 +1,161 @@ +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'; + +const mocks = vi.hoisted(() => ({ + initializeDialogCache: vi.fn().mockResolvedValue(true), + ensureLogin: vi.fn().mockResolvedValue(true), + listDialogs: vi.fn().mockResolvedValue([]), + destroy: vi.fn().mockResolvedValue(undefined), + refreshChannelsFromDialogs: vi.fn().mockResolvedValue(0), + startRealtimeSync: vi.fn(), + resumePendingJobs: vi.fn(), + listArchivedMessages: vi.fn().mockReturnValue([]), + addJob: vi.fn().mockReturnValue({ id: 1, channel_id: '123', status: 'pending' }), + processQueue: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock('../core/config.js', () => ({ + loadConfig: () => ({ + config: { + apiId: '12345', + apiHash: 'test-hash', + phoneNumber: '+10000000000', + mcp: { enabled: true, host: '127.0.0.1', port: Number(process.env.TGCLI_TEST_PORT ?? 0) }, + }, + path: 'test-config.json', + }), + validateConfig: () => [], +})); + +vi.mock('../core/store.js', () => ({ + resolveStoreDir: () => process.env.TGCLI_TEST_STORE, +})); + +vi.mock('../core/services.js', () => ({ + createServices: () => ({ + telegramClient: { + initializeDialogCache: mocks.initializeDialogCache, + ensureLogin: mocks.ensureLogin, + listDialogs: mocks.listDialogs, + destroy: mocks.destroy, + }, + messageSyncService: { + refreshChannelsFromDialogs: mocks.refreshChannelsFromDialogs, + startRealtimeSync: mocks.startRealtimeSync, + resumePendingJobs: mocks.resumePendingJobs, + listArchivedMessages: mocks.listArchivedMessages, + addJob: mocks.addJob, + processQueue: mocks.processQueue, + shutdown: mocks.shutdown, + }, + }), +})); + +function deferred() { + let resolve; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} + +describe('concurrent MCP clients', () => { + let storeDir; + let server; + let baseUrl; + const clients = []; + + async function connectClient() { + const transport = new StreamableHTTPClientTransport(new URL(`${baseUrl}/mcp`)); + const client = new Client({ name: 'tgcli-concurrency-test', version: '1.0.0' }); + await client.connect(transport); + clients.push({ client, transport }); + return { client, transport }; + } + + beforeAll(async () => { + storeDir = fs.mkdtempSync(path.join(os.tmpdir(), 'tgcli-mcp-concurrency-')); + process.env.TGCLI_TEST_STORE = storeDir; + server = await import('../mcp-server.js'); + if (!server.httpServer.listening) { + await new Promise((resolve) => server.httpServer.once('listening', resolve)); + } + baseUrl = `http://127.0.0.1:${server.httpServer.address().port}`; + }); + + afterAll(async () => { + await server?.shutdown(); + await Promise.all(clients.map(({ client }) => client.close().catch(() => {}))); + fs.rmSync(storeDir, { recursive: true, force: true }); + delete process.env.TGCLI_TEST_STORE; + }); + + it('serves overlapping calls on independent MCP sessions', async () => { + const [first, second] = await Promise.all([connectClient(), connectClient()]); + expect(first.transport.sessionId).toBeTruthy(); + expect(second.transport.sessionId).toBeTruthy(); + expect(first.transport.sessionId).not.toBe(second.transport.sessionId); + + const gate = deferred(); + const bothEntered = deferred(); + let entered = 0; + mocks.listDialogs.mockImplementation(async () => { + entered += 1; + if (entered === 2) bothEntered.resolve(); + await gate.promise; + return [{ id: String(entered), title: 'Test dialog' }]; + }); + + const firstCall = first.client.callTool({ name: 'listChannels', arguments: { limit: 1 } }); + const archiveRead = await second.client.callTool({ + name: 'messagesList', + arguments: { channelId: '123', source: 'archive', limit: 1 }, + }); + expect(archiveRead.isError).not.toBe(true); + const syncJob = await second.client.callTool({ + name: 'scheduleMessageSync', + arguments: { channelId: '123', depth: 10 }, + }); + expect(syncJob.isError).not.toBe(true); + expect(mocks.processQueue).toHaveBeenCalledTimes(1); + const secondCall = second.client.callTool({ name: 'listChannels', arguments: { limit: 1 } }); + await bothEntered.promise; + gate.resolve(); + const [firstResult, secondResult] = await Promise.all([firstCall, secondCall]); + expect(firstResult.isError).not.toBe(true); + expect(secondResult.isError).not.toBe(true); + expect(mocks.listDialogs).toHaveBeenCalledTimes(2); + + await first.transport.terminateSession(); + await first.client.close(); + const remainingResult = await second.client.callTool({ name: 'listChannels', arguments: { limit: 1 } }); + expect(remainingResult.isError).not.toBe(true); + expect(mocks.destroy).not.toHaveBeenCalled(); + }); + + it('drains an active tool call before closing shared services', async () => { + const active = await connectClient(); + const gate = deferred(); + const entered = deferred(); + mocks.listDialogs.mockImplementation(async () => { + entered.resolve(); + await gate.promise; + return [{ id: 'last', title: 'Last dialog' }]; + }); + + const call = active.client.callTool({ name: 'listChannels', arguments: { limit: 1 } }); + await entered.promise; + const closing = server.shutdown(); + await new Promise((resolve) => setImmediate(resolve)); + expect(mocks.destroy).not.toHaveBeenCalled(); + gate.resolve(); + const result = await call; + expect(result.isError).not.toBe(true); + await closing; + expect(mocks.shutdown).toHaveBeenCalledTimes(1); + expect(mocks.destroy).toHaveBeenCalledTimes(1); + }); +}); diff --git a/tests/message-sync-concurrency.test.js b/tests/message-sync-concurrency.test.js new file mode 100644 index 0000000..e01f8de --- /dev/null +++ b/tests/message-sync-concurrency.test.js @@ -0,0 +1,35 @@ +import { describe, expect, it, vi } from 'vitest'; + +import MessageSyncService from '../message-sync-service.js'; + +function deferred() { + let resolve; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} + +describe('sync queue concurrency', () => { + it('processes a queued job once when workers are started concurrently', async () => { + const service = Object.create(MessageSyncService.prototype); + const entered = deferred(); + const finish = deferred(); + service.processing = false; + service.stopRequested = false; + service.interJobDelayMs = 0; + service._getNextJob = vi.fn().mockReturnValueOnce({ id: 1 }).mockReturnValue(null); + service._processJob = vi.fn(async () => { + entered.resolve(); + await finish.promise; + }); + + const firstWorker = service.processQueue(); + await entered.promise; + const secondWorker = service.processQueue(); + expect(service._processJob).toHaveBeenCalledTimes(1); + + finish.resolve(); + await Promise.all([firstWorker, secondWorker]); + expect(service._processJob).toHaveBeenCalledTimes(1); + expect(service.processing).toBe(false); + }); +});