From c3a19895b2c7a9a11dd8bbdcf85747d1e2c8de07 Mon Sep 17 00:00:00 2001 From: Danil Pismenny Date: Mon, 28 Sep 2026 22:55:21 +0300 Subject: [PATCH] fix: drain concurrent MCP requests before shutdown Verify independent MCP sessions can overlap Telegram and archive calls while sharing one server. Wait for active tool responses before closing transports and shared services; preserve the sequential sync worker. --- README.md | 8 ++ mcp-server.js | 55 +++++++-- tests/mcp-concurrency.test.js | 161 +++++++++++++++++++++++++ tests/message-sync-concurrency.test.js | 35 ++++++ 4 files changed, 246 insertions(+), 13 deletions(-) create mode 100644 tests/mcp-concurrency.test.js create mode 100644 tests/message-sync-concurrency.test.js 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); + }); +});