From c929bf90970b80d4c5f6e3f96f0cb507f9837229 Mon Sep 17 00:00:00 2001 From: Matteo Collina Date: Mon, 19 Oct 2026 13:14:29 +0200 Subject: [PATCH 1/2] fix(sse): subscribe legacy session topics wherever a stream attaches The per-session broker topic was subscribed only by the instance that created the session, and unsubscribed when its last stream closed. A GET attaching later, or on another instance, never subscribed again, so messages were neither delivered nor recorded in history. With Redis the subscription also outlived expired sessions, and RedisMessageBroker unsubscribe never sent UNSUBSCRIBE, so channels leaked in Redis too. - mcpSendToSession now records the message in history and then publishes it with its event id. History no longer depends on any subscription and stays exactly-once with several instances. - An instance subscribes to a session topic exactly while it holds one of the session's streams. Changes per session are serialized, so a quick close-and-reconnect can't drop the new subscription. - A GET subscribes before replaying Last-Event-ID; live events that arrive during the replay are queued and deduplicated against it. - DELETE publishes a termination on the session topic, so the instance holding the stream ends it. - The stale-session sweep runs for every store and closes streams whose session is gone. Memory store cleanup() still runs as before. - RedisMessageBroker waits for SUBSCRIBE and releases the Redis channel on unsubscribe. Fixes #327 Fixes #328 Fixes #329 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LmhsFjk8W96X2PvmShmDm9 --- README.md | 5 + src/brokers/redis-message-broker.ts | 35 +- src/decorators/pubsub.ts | 15 +- src/index.ts | 56 +--- src/routes/mcp.ts | 305 ++++++++++------- src/session-topic.ts | 151 +++++++++ test/broadcast-error-handling.test.ts | 6 +- test/legacy-session-subscriptions.test.ts | 377 ++++++++++++++++++++++ 8 files changed, 769 insertions(+), 181 deletions(-) create mode 100644 src/session-topic.ts create mode 100644 test/legacy-session-subscriptions.test.ts diff --git a/README.md b/README.md index bfd15b35..388b8b50 100644 --- a/README.md +++ b/README.md @@ -987,6 +987,11 @@ app.mcpSendToSession('session-xyz', { }) ``` +**How session messages are delivered:** +- `mcpSendToSession` (and everything built on it: elicitation, task status updates) first records the message in the session's history, then publishes it on `mcp/session//message`. It returns `false` if the session is gone or the message could not be recorded. +- An instance subscribes to that topic only while it holds one of the session's GET streams, wherever the session was created. A message sent while no stream is open is not lost: the client gets it on its next GET with `Last-Event-ID`. +- `DELETE /mcp` ends the session's streams on every instance. Streams and subscriptions held for a session that expired, or was deleted elsewhere, are released by a periodic sweep that runs with both backends. + ### Implementation Design Decisions This plugin implements certain behaviors that differ from the strict MCP specification to provide enhanced scalability and resilience in production environments: diff --git a/src/brokers/redis-message-broker.ts b/src/brokers/redis-message-broker.ts index 0da89ea5..859bf656 100644 --- a/src/brokers/redis-message-broker.ts +++ b/src/brokers/redis-message-broker.ts @@ -24,7 +24,6 @@ function suppressUnhandledQuitRejection (conn?: RedisConnLike): void { interface RedisMQEmitter extends ReturnType { state?: { on (event: 'error', listener: (err: Error) => void): unknown } - removeAllListeners? (topic: string, done?: (err?: Error) => void): void subConn?: RedisConnLike pubConn?: RedisConnLike _close?: (cb: () => void) => void @@ -44,6 +43,7 @@ export class RedisMessageBroker implements MessageBroker { private readonly closeTimeoutMs: number private readonly onCloseTimeout?: (closeTimeoutMs: number) => void private closePromise: Promise | null = null + private readonly listeners = new Map void) => void>>() /** * @param redis Redis connection duplicated (all ioredis options: host, port, auth, db, tls, …) for the MQEmitter pub/sub connections. @@ -88,20 +88,37 @@ export class RedisMessageBroker implements MessageBroker { } async subscribe (topic: string, handler: (message: JSONRPCMessage) => void): Promise { - return new Promise((resolve) => { - this.emitter.on(topic, (packet, cb) => { + const listener = (packet: any, cb: () => void) => { + try { handler(packet.message) + } finally { cb() - }) - resolve() + } + } + let listeners = this.listeners.get(topic) + if (!listeners) { + listeners = new Set() + this.listeners.set(topic, listeners) + } + listeners.add(listener) + + // Resolve once Redis has acknowledged the SUBSCRIBE, so a message published + // right after this returns is not missed + return new Promise((resolve) => { + this.emitter.on(topic, listener, () => resolve()) }) } async unsubscribe (topic: string): Promise { - return new Promise((resolve) => { - this.emitter.removeAllListeners?.(topic) - resolve() - }) + const listeners = this.listeners.get(topic) + if (!listeners) return + this.listeners.delete(topic) + + // removeListener, unlike removeAllListeners, also sends UNSUBSCRIBE once a + // channel has no listener left; otherwise the Redis subscription leaks + await Promise.all([...listeners].map(listener => new Promise((resolve) => { + this.emitter.removeListener(topic, listener, () => resolve()) + }))) } async close (): Promise { diff --git a/src/decorators/pubsub.ts b/src/decorators/pubsub.ts index 717b3641..81054baf 100644 --- a/src/decorators/pubsub.ts +++ b/src/decorators/pubsub.ts @@ -14,6 +14,7 @@ import { randomUUID } from 'node:crypto' import { supportsUrlElicitation } from '../protocol-version.ts' import type { SessionStore } from '../stores/session-store.ts' import type { MessageBroker } from '../brokers/message-broker.ts' +import { sendToSession } from '../session-topic.ts' /** * Whether the client declared support for an elicitation mode. Elicitation is a @@ -58,17 +59,11 @@ const mcpPubSubDecoratorsPlugin: FastifyPluginAsync return false } - // Check if session exists in store - const session = await sessionStore.get(sessionId) - if (!session) { - return false - } - - // Always publish to messageBroker to support cross-instance messaging in Redis deployments - // This ensures the message reaches the correct instance where the SSE connection exists + // Recorded in the session's history first, so it is replayed on reconnect + // even when no stream is open right now, then published so whichever + // instance holds the session's stream delivers it try { - await messageBroker.publish(`mcp/session/${sessionId}/message`, message) - return true + return await sendToSession(sessionStore, messageBroker, sessionId, message) } catch (error) { app.log.error({ err: error }, 'Failed to send message to session') return false diff --git a/src/index.ts b/src/index.ts index a57f9b31..d66c7917 100644 --- a/src/index.ts +++ b/src/index.ts @@ -15,6 +15,7 @@ import type { MCPPluginOptions, MCPTool, MCPResource, MCPPrompt, ResourceHandler import pubsubDecorators from './decorators/pubsub.ts' import metaDecorators from './decorators/meta.ts' import routes from './routes/mcp.ts' +import type { SessionLifecycle } from './routes/mcp.ts' import wellKnownRoutes from './routes/well-known.ts' import { TokenValidator } from './auth/token-validator.ts' import { createAuthPreHandler } from './auth/prehandler.ts' @@ -182,32 +183,8 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption // Local stream management per server instance const localStreams = new Map>() - // Redis expires idle sessions itself through key TTLs. The memory store has no - // such mechanism, so sweep it periodically: otherwise every session a client - // abandons without DELETE would stay in memory forever. Reads already treat an - // expired session as gone; this sweep is what actually frees the memory. - let sessionCleanupTimer: NodeJS.Timeout | undefined - if (sessionStore instanceof MemorySessionStore) { - const memoryStore = sessionStore - const sweep = async () => { - await memoryStore.cleanup() - // An SSE session created by POST that never opened a stream keeps an empty - // stream set and a broker subscription; drop both once the session is gone - for (const [sessionId, streams] of localStreams) { - if (streams.size === 0 && !(await memoryStore.get(sessionId))) { - localStreams.delete(sessionId) - await messageBroker.unsubscribe(`mcp/session/${sessionId}/message`) - } - } - } - sessionCleanupTimer = setInterval(() => { - sweep().catch((err: unknown) => { - app.log.error({ err }, 'Failed to clean up expired sessions') - }) - }, Math.min(memoryStore.sessionTtlMs, 60_000)) - // Never keep the process alive just to sweep sessions - sessionCleanupTimer.unref() - } + // Filled in by the routes plugin, which owns the SSE streams + const sessionLifecycle: SessionLifecycle = {} // Initialize authorization components if enabled let tokenValidator: TokenValidator | null = null @@ -282,37 +259,18 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption shutdownSignal: shutdownController.signal, taskLeases, jsonSchemaValidator, - authPreHandler: mcpAuthPreHandler + authPreHandler: mcpAuthPreHandler, + sessionLifecycle }) // Add close hook to clean up Redis connections and authorization components app.addHook('onClose', async () => { - clearInterval(sessionCleanupTimer) - // Tasks still running here are abandoned: stop vouching for them so their // leases lapse and other instances report them failed. taskLeases?.close() - // Clean up all SSE streams and sessions - const unsubscribePromises: Promise[] = [] - for (const [sessionId, streams] of localStreams.entries()) { - for (const stream of streams) { - try { - if (stream.raw && !stream.raw.destroyed) { - stream.raw.destroy() - } - } catch (error) { - app.log.debug({ error, sessionId }, 'Error destroying SSE stream') - } - } - streams.clear() - // Collect unsubscribe promises for parallel execution - unsubscribePromises.push(messageBroker.unsubscribe(`mcp/session/${sessionId}/message`)) - } - localStreams.clear() - - // Execute all unsubscribes in parallel - await Promise.all(unsubscribePromises) + // Close every SSE stream and drop the session subscriptions + await sessionLifecycle.close?.() try { await messageBroker.close() diff --git a/src/routes/mcp.ts b/src/routes/mcp.ts index efdd63ba..ce4c8b94 100644 --- a/src/routes/mcp.ts +++ b/src/routes/mcp.ts @@ -14,13 +14,15 @@ import type { MCPRouteSchemaTransformer } from '../types.ts' import type { SessionStore, SessionMetadata } from '../stores/session-store.ts' -import { DEFAULT_SESSION_TTL_MS } from '../stores/memory-session-store.ts' +import { DEFAULT_SESSION_TTL_MS, MemorySessionStore } from '../stores/memory-session-store.ts' import type { TaskStore, TaskWaiters, TaskLeases } from '../stores/task-store.ts' import type { MessageBroker } from '../brokers/message-broker.ts' import type { AuthorizationContext } from '../types/auth-types.ts' import type { JsonSchemaValidator } from '../validation/json-schema-validator.ts' import { processMessage, checkRequestId } from '../handlers.ts' import { isStdioRequest } from '../stdio-marker.ts' +import { appendToHistory, publishSessionTerminated, SessionTopicSubscriptions } from '../session-topic.ts' +import type { SessionTopicPayload } from '../session-topic.ts' const MAX_TIMER_DELAY_MS = 2 ** 31 - 1 @@ -42,6 +44,23 @@ interface MCPPubSubRoutesOptions { taskLeases?: TaskLeases jsonSchemaValidator?: JsonSchemaValidator authPreHandler?: preHandlerHookHandler + /** Filled in by this plugin; the parent calls `close` before closing the broker */ + sessionLifecycle: SessionLifecycle +} + +export interface SessionLifecycle { + close?: () => Promise +} + +/** + * Per GET stream. While a Last-Event-ID replay is running, live events are + * queued so they cannot overtake the replay; events already replayed are not + * written a second time when they also arrive live. + */ +interface StreamState { + replaying: boolean + queued: Array<{ eventId: string, frame: string }> + replayed: Set } function resolveRouteUrl (prefix: string | undefined, url: string): string { @@ -122,7 +141,7 @@ function buildMcpPostErrorHandler (app: FastifyInstance) { } const mcpPubSubRoutesPlugin: FastifyPluginAsync = async (app, options) => { - const { enableSSE, opts, capabilities, serverInfo, tools, resources, prompts, resourceHandlers, sessionStore, messageBroker, localStreams, taskStore, taskWaiters, shutdownSignal, taskLeases, jsonSchemaValidator, authPreHandler } = options + const { enableSSE, opts, capabilities, serverInfo, tools, resources, prompts, resourceHandlers, sessionStore, messageBroker, localStreams, taskStore, taskWaiters, shutdownSignal, taskLeases, jsonSchemaValidator, authPreHandler, sessionLifecycle } = options const mcpUrl = resolveRouteUrl(app.prefix, '/mcp') const allowedOrigins = opts.allowedOrigins @@ -297,38 +316,9 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async await sessionStore.create(session) - // Without SSE there is no stream to deliver to, so skip the subscription - if (!enableSSE) { - return session - } - - localStreams.set(sessionId, new Set()) - - // Subscribe to messages for this session - await messageBroker.subscribe(`mcp/session/${sessionId}/message`, async (message: JSONRPCMessage) => { - // Brokers do not await this handler, so any rejection must be caught - // here or it becomes an unhandled rejection that crashes the process - try { - const streams = localStreams.get(sessionId) - if (streams && streams.size > 0) { - app.log.debug({ sessionId, message }, 'Received message for session via broker, sending to streams') - await sendSSEToStreams(sessionId, message, streams) - } else { - app.log.debug({ sessionId }, 'Received message for session via broker, storing in history without active streams') - // Store message in history even without active streams for session persistence - const session = await sessionStore.get(sessionId) - if (session) { - const eventId = (++session.eventId).toString() - session.lastEventId = eventId - session.lastActivity = new Date() - await sessionStore.addMessage(sessionId, eventId, message) - } - } - } catch (error) { - app.log.error({ err: error, sessionId }, 'Failed to deliver message to session') - } - }) - + // No broker subscription here: an instance subscribes to the session's + // topic only while it holds one of its streams (see attachStream), and + // messages sent meanwhile are kept in history by the sender return session } @@ -343,50 +333,113 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async return streams ? streams.size > 0 : false } - async function sendSSEToStreams (sessionId: string, message: JSONRPCMessage, streams: Set): Promise { - const session = await sessionStore.get(sessionId) - if (!session) return + const streamStates = new WeakMap() - const eventId = (++session.eventId).toString() - const sseEvent = `id: ${eventId}\ndata: ${JSON.stringify(message)}\n\n` - session.lastEventId = eventId - session.lastActivity = new Date() + // This instance is subscribed to a session's topic exactly while it holds one + // of the session's streams; the subscription is what lets a stream on any + // instance receive what is sent to the session from any other. + const sessionTopics = new SessionTopicSubscriptions( + messageBroker, + (sessionId) => (localStreams.get(sessionId)?.size ?? 0) > 0, + (sessionId, payload) => { + // Brokers do not await this handler, so any rejection must be caught + // here or it becomes an unhandled rejection that crashes the process + onSessionTopicMessage(sessionId, payload).catch((error) => { + app.log.error({ err: error, sessionId }, 'Failed to deliver message to session') + }) + } + ) - // Store message in history - await sessionStore.addMessage(sessionId, eventId, message) + async function onSessionTopicMessage (sessionId: string, payload: SessionTopicPayload): Promise { + if (payload.kind === 'terminated') { + endLocalStreams(sessionId) + return + } - // Send to all connected streams in this session - const deadStreams = new Set() - for (const stream of streams) { - try { - stream.raw.write(sseEvent) - } catch (error) { - app.log.error({ err: error }, 'Failed to write SSE event') - deadStreams.add(stream) - } + const streams = localStreams.get(sessionId) + if (!streams || streams.size === 0) return + + if (payload.kind === 'event') { + app.log.debug({ sessionId, eventId: payload.eventId }, 'Received message for session via broker, sending to streams') + writeEventToStreams(sessionId, payload.eventId, payload.message, streams) + } else { + // Published to the topic directly instead of through mcpSendToSession, + // so it still needs to be recorded + await sendSSEToStreams(sessionId, payload.message, streams) } + } - // Clean up dead streams - for (const deadStream of deadStreams) { - streams.delete(deadStream) + async function attachStream (sessionId: string, reply: FastifyReply): Promise { + let streams = localStreams.get(sessionId) + if (!streams) { + streams = new Set() + localStreams.set(sessionId, streams) } + streams.add(reply) + await sessionTopics.sync(sessionId) + } - // Clean up session if no streams left + function detachStream (sessionId: string, reply: FastifyReply): void { + const streams = localStreams.get(sessionId) + if (!streams || !streams.delete(reply)) return if (streams.size === 0) { - app.log.info({ - sessionId - }, 'Session has no active streams, cleaning up') localStreams.delete(sessionId) - await messageBroker.unsubscribe(`mcp/session/${sessionId}/message`) + } + sessionTopics.sync(sessionId).catch((error) => { + app.log.error({ err: error, sessionId }, 'Failed to unsubscribe session from message broker') + }) + } + + /** End this instance's streams for a session that was terminated or is gone */ + function endLocalStreams (sessionId: string): void { + const streams = localStreams.get(sessionId) + if (!streams) return + for (const stream of [...streams]) { + try { + stream.raw.end() + } catch { + // stream may already be closed + } + detachStream(sessionId, stream) } } + function writeEvent (sessionId: string, stream: FastifyReply, eventId: string, frame: string): void { + const state = streamStates.get(stream) + if (state?.replayed.has(eventId)) return + if (state?.replaying) { + state.queued.push({ eventId, frame }) + return + } + try { + stream.raw.write(frame) + } catch (error) { + app.log.error({ err: error }, 'Failed to write SSE event') + detachStream(sessionId, stream) + } + } + + function writeEventToStreams (sessionId: string, eventId: string, message: JSONRPCMessage, streams: Set): void { + const frame = `id: ${eventId}\ndata: ${JSON.stringify(message)}\n\n` + for (const stream of [...streams]) { + writeEvent(sessionId, stream, eventId, frame) + } + } + + async function sendSSEToStreams (sessionId: string, message: JSONRPCMessage, streams: Set): Promise { + const eventId = await appendToHistory(sessionStore, sessionId, message) + if (eventId === null) return + writeEventToStreams(sessionId, eventId, message, streams) + } + async function replayMessagesFromEventId (sessionId: string, lastEventId: string, stream: FastifyReply): Promise { + const state = streamStates.get(stream) try { const messagesToReplay = await sessionStore.getMessagesFrom(sessionId, lastEventId) for (const entry of messagesToReplay) { const sseEvent = `id: ${entry.eventId}\ndata: ${JSON.stringify(entry.message)}\n\n` + state?.replayed.add(entry.eventId) try { stream.raw.write(sseEvent) } catch (error) { @@ -400,6 +453,14 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async } } catch (error) { app.log.warn({ err: error, lastEventId }, 'Failed to replay messages from event ID') + } finally { + // Deliver what arrived live during the replay, minus what it already sent + if (state) { + state.replaying = false + for (const { eventId, frame } of state.queued.splice(0)) { + writeEvent(sessionId, stream, eventId, frame) + } + } } } @@ -566,50 +627,34 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async raw.setHeader('Mcp-Session-Id', session.id) } - raw.writeHead(200) + const lastEventId = request.headers['last-event-id'] as string | undefined + streamStates.set(reply, { replaying: Boolean(lastEventId), queued: [], replayed: new Set() }) - let streams = localStreams.get(session.id) - if (!streams) { - streams = new Set() - localStreams.set(session.id, streams) - } - streams.add(reply) + // Registered before anything else is awaited, so a client that goes + // away during setup is still detached + reply.raw.on('close', () => { + detachStream(session.id, reply) + app.log.info({ sessionId: session.id }, 'SSE connection closed') + }) + + // Subscribe before reading history for the replay, so nothing sent in + // between can be missed; anything seen by both is written once + await attachStream(session.id, reply) + + raw.writeHead(200) app.log.info({ sessionId: session.id, - totalStreams: streams.size, + totalStreams: localStreams.get(session.id)?.size ?? 0, method: 'GET' }, 'Added new stream to session') // Handle resumability with Last-Event-ID - const lastEventId = request.headers['last-event-id'] as string if (lastEventId) { app.log.info(`Resuming SSE stream from event ID: ${lastEventId}`) await replayMessagesFromEventId(session.id, lastEventId, reply) } - // Handle connection close - reply.raw.on('close', () => { - const streams = localStreams.get(session.id) - if (streams) { - streams.delete(reply) - app.log.info({ - sessionId: session.id, - remainingStreams: streams.size - }, 'SSE connection closed') - - if (streams.size === 0) { - app.log.info({ - sessionId: session.id - }, 'Last SSE stream closed, cleaning up session') - localStreams.delete(session.id) - messageBroker.unsubscribe(`mcp/session/${session.id}/message`).catch((error) => { - app.log.error({ err: error, sessionId: session.id }, 'Failed to unsubscribe session from message broker') - }) - } - } - }) - // SEP-1699: servers may end an SSE stream whenever they like, turning the // stream into a polling channel. The client reconnects with Last-Event-ID // on GET and we replay whatever it missed, so closing here loses nothing. @@ -670,10 +715,7 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async }) } catch (error) { clearInterval(heartbeatInterval) - const streams = localStreams.get(session.id) - if (streams) { - streams.delete(reply) - } + detachStream(session.id, reply) } }, heartbeatIntervalMs) heartbeatInterval.unref() @@ -715,25 +757,19 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async return } - // Force-close any active SSE streams for this session - const streams = localStreams.get(sessionId) - if (streams) { - for (const stream of streams) { - try { - stream.raw.end() - } catch { - // stream may already be closed - } + // Delete first, so a stream reconnecting elsewhere no longer finds it + await sessionStore.delete(sessionId) + + // End this instance's streams, and tell the others to end theirs + endLocalStreams(sessionId) + if (enableSSE) { + try { + await publishSessionTerminated(messageBroker, sessionId) + } catch (error) { + app.log.error({ err: error, sessionId }, 'Failed to notify other instances of session termination') } - localStreams.delete(sessionId) } - // Unsubscribe from message broker - await messageBroker.unsubscribe(`mcp/session/${sessionId}/message`) - - // Delete session from store - await sessionStore.delete(sessionId) - app.log.info({ sessionId }, 'Session terminated via DELETE') reply.code(204).send() }) @@ -768,6 +804,53 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async } }) } + + // Sessions also end without a DELETE reaching this instance: they expire, are + // deleted through another instance, or (Redis) their key is evicted. Sweep + // periodically, whatever the store, so what this instance holds for such a + // session is released. The memory store has no TTL mechanism of its own, so + // the sweep is also what frees its expired sessions. + const sweepIntervalMs = Math.min(opts.sessionTtlMs ?? DEFAULT_SESSION_TTL_MS, 60_000) + let sweeping = false + async function sweepSessions (): Promise { + if (sessionStore instanceof MemorySessionStore) { + await sessionStore.cleanup() + } + for (const sessionId of [...localStreams.keys()]) { + if (!(await sessionStore.get(sessionId))) { + app.log.info({ sessionId }, 'Session is gone, closing its SSE streams') + endLocalStreams(sessionId) + } + } + } + const sweepTimer = setInterval(() => { + if (sweeping) return + sweeping = true + sweepSessions().catch((err: unknown) => { + app.log.error({ err }, 'Failed to clean up expired sessions') + }).finally(() => { + sweeping = false + }) + }, sweepIntervalMs) + // Never keep the process alive just to sweep sessions + sweepTimer.unref() + + sessionLifecycle.close = async () => { + clearInterval(sweepTimer) + for (const [sessionId, streams] of localStreams) { + for (const stream of streams) { + try { + if (stream.raw && !stream.raw.destroyed) { + stream.raw.destroy() + } + } catch (error) { + app.log.debug({ error, sessionId }, 'Error destroying SSE stream') + } + } + } + localStreams.clear() + await sessionTopics.close() + } } export default fp(mcpPubSubRoutesPlugin, { diff --git a/src/session-topic.ts b/src/session-topic.ts new file mode 100644 index 00000000..fd2d1ffc --- /dev/null +++ b/src/session-topic.ts @@ -0,0 +1,151 @@ +import type { JSONRPCMessage } from './schema.ts' +import type { MessageBroker } from './brokers/message-broker.ts' +import type { SessionStore } from './stores/session-store.ts' + +/** + * Per-session broker topic of the legacy (Streamable HTTP + SSE) transport. + * + * A message is written to the session's history by whoever sends it, and only + * then published here with the event id it was stored under. Instances that + * hold one of the session's GET streams subscribe and write it to their + * streams; nobody else needs to. Storing on the sending side is what keeps + * history exactly-once with several instances: if subscribers stored it, every + * subscribed instance would add its own copy. + */ +export function sessionTopic (sessionId: string): string { + return `mcp/session/${sessionId}/message` +} + +interface SessionEventPacket { + mcpSessionEvent: { eventId: string, message: JSONRPCMessage } +} + +interface SessionTerminatedPacket { + mcpSessionTerminated: true +} + +export type SessionTopicPayload = + | { kind: 'event', eventId: string, message: JSONRPCMessage } + | { kind: 'terminated' } + // Published straight to the topic by code that bypasses mcpSendToSession, + // so it is not in history yet + | { kind: 'raw', message: JSONRPCMessage } + +export function decodeSessionTopicPayload (payload: unknown): SessionTopicPayload { + if (typeof payload === 'object' && payload !== null) { + if ('mcpSessionEvent' in payload) { + const { eventId, message } = (payload as SessionEventPacket).mcpSessionEvent + return { kind: 'event', eventId, message } + } + if ('mcpSessionTerminated' in payload) { + return { kind: 'terminated' } + } + } + return { kind: 'raw', message: payload as JSONRPCMessage } +} + +// The broker interface is typed for JSON-RPC messages, but every broker passes +// payloads through untouched (or JSON-serialized), so packets travel as is. +function asBrokerPayload (packet: SessionEventPacket | SessionTerminatedPacket): JSONRPCMessage { + return packet as unknown as JSONRPCMessage +} + +/** + * Record a message in a session's history and return the event id it was + * stored under, or null when the session no longer exists. + */ +export async function appendToHistory (sessionStore: SessionStore, sessionId: string, message: JSONRPCMessage): Promise { + const session = await sessionStore.get(sessionId) + if (!session) return null + + const eventId = (session.eventId + 1).toString() + await sessionStore.addMessage(sessionId, eventId, message) + return eventId +} + +/** + * Store a message in the session's history, then hand it to whichever + * instances hold the session's streams. Returns false when the session is gone. + */ +export async function sendToSession ( + sessionStore: SessionStore, + messageBroker: MessageBroker, + sessionId: string, + message: JSONRPCMessage +): Promise { + const eventId = await appendToHistory(sessionStore, sessionId, message) + if (eventId === null) return false + + await messageBroker.publish(sessionTopic(sessionId), asBrokerPayload({ mcpSessionEvent: { eventId, message } })) + return true +} + +/** Tell every instance that the session was terminated, so they end its streams */ +export async function publishSessionTerminated (messageBroker: MessageBroker, sessionId: string): Promise { + await messageBroker.publish(sessionTopic(sessionId), asBrokerPayload({ mcpSessionTerminated: true })) +} + +/** + * Keeps this instance subscribed to a session's topic exactly while `wanted` + * says so (while it holds at least one of the session's streams). + * + * Changes for one session are applied one at a time, so a stream closing and + * another attaching in quick succession can never leave the subscription + * dropped by an unsubscribe that finished after the new subscribe. + */ +export class SessionTopicSubscriptions { + private readonly broker: MessageBroker + private readonly wanted: (sessionId: string) => boolean + private readonly handler: (sessionId: string, payload: SessionTopicPayload) => void + private readonly subscribed = new Set() + private readonly pending = new Map>() + private closed = false + + constructor ( + broker: MessageBroker, + wanted: (sessionId: string) => boolean, + handler: (sessionId: string, payload: SessionTopicPayload) => void + ) { + this.broker = broker + this.wanted = wanted + this.handler = handler + } + + /** Subscribe or unsubscribe so the subscription matches `wanted` */ + sync (sessionId: string): Promise { + const previous = this.pending.get(sessionId) ?? Promise.resolve() + const apply = () => this.apply(sessionId) + const next = previous.then(apply, apply) + const tracked = next.catch(() => {}).finally(() => { + if (this.pending.get(sessionId) === tracked) { + this.pending.delete(sessionId) + } + }) + this.pending.set(sessionId, tracked) + return next + } + + isSubscribed (sessionId: string): boolean { + return this.subscribed.has(sessionId) + } + + /** Stop following `wanted` and drop every subscription */ + async close (): Promise { + this.closed = true + await Promise.all([...this.pending.values()]) + await Promise.allSettled([...this.subscribed].map(sessionId => this.sync(sessionId))) + } + + private async apply (sessionId: string): Promise { + const want = !this.closed && this.wanted(sessionId) + if (want && !this.subscribed.has(sessionId)) { + await this.broker.subscribe(sessionTopic(sessionId), (payload: unknown) => { + this.handler(sessionId, decodeSessionTopicPayload(payload)) + }) + this.subscribed.add(sessionId) + } else if (!want && this.subscribed.has(sessionId)) { + this.subscribed.delete(sessionId) + await this.broker.unsubscribe(sessionTopic(sessionId)) + } + } +} diff --git a/test/broadcast-error-handling.test.ts b/test/broadcast-error-handling.test.ts index a503c88c..eb7b2497 100644 --- a/test/broadcast-error-handling.test.ts +++ b/test/broadcast-error-handling.test.ts @@ -130,7 +130,8 @@ test('session message does not cause an unhandled rejection when the session sto method: 'notifications/message', params: { level: 'info', data: 'hello' } }) - t.assert.strictEqual(sent, true) + // The message is stored by the sender, so a failing store is reported to the caller + t.assert.strictEqual(sent, false) await sleep(100) t.assert.deepStrictEqual(unhandled, []) @@ -150,7 +151,8 @@ test('session message does not cause an unhandled rejection when the session sto method: 'notifications/message', params: { level: 'info', data: 'hello' } }) - t.assert.strictEqual(sent, true) + // The message is stored by the sender, so a failing store is reported to the caller + t.assert.strictEqual(sent, false) await sleep(100) t.assert.deepStrictEqual(unhandled, []) diff --git a/test/legacy-session-subscriptions.test.ts b/test/legacy-session-subscriptions.test.ts new file mode 100644 index 00000000..7d23c35c --- /dev/null +++ b/test/legacy-session-subscriptions.test.ts @@ -0,0 +1,377 @@ +import { test, describe } from 'node:test' +import type { TestContext } from 'node:test' +import type { Readable } from 'node:stream' +import { setTimeout as sleep } from 'node:timers/promises' +import Fastify from 'fastify' +import type { FastifyInstance } from 'fastify' +import { Agent, request as undiciRequest } from 'undici' +import type { Redis } from 'ioredis' +import mcpPlugin from '../src/index.ts' +import { RedisSessionStore } from '../src/stores/redis-session-store.ts' +import { RedisMessageBroker } from '../src/brokers/redis-message-broker.ts' +import { SessionTopicSubscriptions } from '../src/session-topic.ts' +import type { MCPPluginOptions } from '../src/types.ts' +import { JSONRPC_VERSION, LATEST_PROTOCOL_VERSION } from '../src/schema.ts' +import type { JSONRPCNotification } from '../src/schema.ts' +import { testWithRedis } from './redis-test-utils.ts' + +type RedisConfig = { host: string, port: number, db: number } + +function redisConfig (redis: Redis): RedisConfig { + return { host: redis.options.host!, port: redis.options.port!, db: redis.options.db! } +} + +interface TestServer { + app: FastifyInstance + baseUrl: string + /** Resolves once the server has seen `count` GET streams close in total */ + waitForStreamCloses (count: number): Promise +} + +// All client connections go through one agent, destroyed before any app is +// closed: a server cannot finish closing while a client still holds a socket +// open, and SSE streams keep theirs open indefinitely +let agent = new Agent() +function request (url: string, options: Parameters[1] = {}) { + return undiciRequest(url, { ...options, dispatcher: agent }) +} + +async function startServer (t: TestContext, opts: Partial = {}): Promise { + const app = Fastify({ logger: false }) + t.after(async () => { + const current = agent + agent = new Agent() + await current.destroy() + await app.close() + }) + + await app.register(mcpPlugin, { + serverInfo: { name: 'test-server', version: '1.0.0' }, + enableSSE: true, + ...opts + }) + + await app.listen({ port: 0, host: '127.0.0.1' }) + const address = app.server.address() + const port = typeof address === 'object' && address ? address.port : 0 + + // Our listener is added after Fastify's, so by the time it runs the plugin's + // own close handler for the stream has run too + let closes = 0 + const waiters: Array<() => void> = [] + app.server.on('request', (req, res) => { + if (req.method !== 'GET') return + res.on('close', () => { + closes++ + for (const wake of waiters.splice(0)) wake() + }) + }) + + async function waitForStreamCloses (count: number): Promise { + if (closes >= count) return + await new Promise(resolve => waiters.push(resolve)) + return waitForStreamCloses(count) + } + + return { app, baseUrl: `http://127.0.0.1:${port}`, waitForStreamCloses } +} + +async function initialize (baseUrl: string): Promise { + const response = await request(`${baseUrl}/mcp`, { + method: 'POST', + headers: { 'content-type': 'application/json', accept: 'application/json' }, + body: JSON.stringify({ + jsonrpc: JSONRPC_VERSION, + id: 1, + method: 'initialize', + params: { + protocolVersion: LATEST_PROTOCOL_VERSION, + capabilities: {}, + clientInfo: { name: 'test-client', version: '1.0.0' } + } + }) + }) + await response.body.text() + const sessionId = response.headers['mcp-session-id'] as string + if (!sessionId) throw new Error('initialize did not return a session id') + return sessionId +} + +async function openStream (t: TestContext, baseUrl: string, sessionId: string, lastEventId?: string): Promise { + const headers: Record = { + accept: 'text/event-stream', + 'mcp-session-id': sessionId, + 'mcp-protocol-version': LATEST_PROTOCOL_VERSION + } + if (lastEventId) headers['last-event-id'] = lastEventId + const response = await request(`${baseUrl}/mcp`, { method: 'GET', headers }) + t.assert.strictEqual(response.statusCode, 200) + return response.body +} + +// Reads SSE data until `marker` shows up, or fails after `timeoutMs` +function readUntil (stream: Readable, marker: string, timeoutMs = 3000): Promise { + return new Promise((resolve, reject) => { + let data = '' + const cleanup = () => { + clearTimeout(timer) + stream.off('data', onData) + stream.off('error', onError) + stream.pause() + } + const timer = setTimeout(() => { + cleanup() + reject(new Error(`Timed out waiting for ${marker}; received: ${JSON.stringify(data)}`)) + }, timeoutMs) + const onData = (chunk: Buffer) => { + data += chunk.toString() + if (data.includes(marker)) { + cleanup() + resolve(data) + } + } + const onError = (err: Error) => { + cleanup() + reject(err) + } + stream.on('data', onData) + stream.on('error', onError) + stream.resume() + }) +} + +// Resolves when the server ends the stream, or 'timeout' +function waitForEnd (stream: Readable, timeoutMs = 3000): Promise<'ended' | 'timeout'> { + return new Promise(resolve => { + const timer = setTimeout(() => resolve('timeout'), timeoutMs) + stream.on('end', () => { + clearTimeout(timer) + resolve('ended') + }) + stream.on('error', () => { + clearTimeout(timer) + resolve('ended') + }) + stream.resume() + }) +} + +function note (text: string): JSONRPCNotification { + return { jsonrpc: JSONRPC_VERSION, method: 'notifications/message', params: { level: 'info', data: text } } +} + +function eventIds (data: string): string[] { + return [...data.matchAll(/^id: (\S+)$/gm)].map(match => match[1]) +} + +async function subscriberCount (redis: Redis, sessionId: string): Promise { + const [, count] = await redis.pubsub('NUMSUB', `mcp/session/${sessionId}/message`) as [string, number] + return Number(count) +} + +async function waitUntil (predicate: () => Promise, timeoutMs = 3000): Promise { + const deadline = Date.now() + timeoutMs + while (Date.now() < deadline) { + if (await predicate()) return true + await sleep(25) + } + return predicate() +} + +/** + * Close the first stream, send while no stream is open, then reconnect with + * Last-Event-ID: the message sent in the gap must be replayed, and the new + * stream must receive live messages again. + */ +async function reconnectScenario (t: TestContext, redis?: RedisConfig) { + const server = await startServer(t, { redis }) + const { app, baseUrl } = server + const sessionId = await initialize(baseUrl) + + const first = await openStream(t, baseUrl, sessionId) + t.assert.strictEqual(await app.mcpSendToSession(sessionId, note('msg-1')), true) + const firstData = await readUntil(first, 'msg-1') + t.assert.deepStrictEqual(eventIds(firstData), ['1']) + + first.destroy() + await server.waitForStreamCloses(1) + + t.assert.strictEqual(await app.mcpSendToSession(sessionId, note('msg-2')), true) + + const second = await openStream(t, baseUrl, sessionId, '1') + const replayed = await readUntil(second, 'msg-2') + t.assert.deepStrictEqual(eventIds(replayed), ['2']) + + t.assert.strictEqual(await app.mcpSendToSession(sessionId, note('msg-3')), true) + const live = await readUntil(second, 'msg-3') + t.assert.deepStrictEqual(eventIds(live), ['3']) +} + +describe('legacy SSE session subscriptions (#327)', () => { + test('messages sent after the first stream closes are kept and delivered on reconnect', async (t: TestContext) => { + await reconnectScenario(t) + }) + + testWithRedis('Redis: messages sent after the first stream closes are kept and delivered on reconnect', async (redis, t: TestContext) => { + await reconnectScenario(t, redisConfig(redis)) + }) + + test('messages sent before any stream is opened are replayed', async (t: TestContext) => { + const { app, baseUrl } = await startServer(t) + const sessionId = await initialize(baseUrl) + + await app.mcpSendToSession(sessionId, note('early-1')) + await app.mcpSendToSession(sessionId, note('early-2')) + + // Event 1 is acknowledged, so only event 2 is replayed + const stream = await openStream(t, baseUrl, sessionId, '1') + const data = await readUntil(stream, 'early-2') + t.assert.ok(!data.includes('early-1')) + t.assert.deepStrictEqual(eventIds(data), ['2']) + }) +}) + +describe('legacy SSE sessions across Redis instances (#328)', () => { + testWithRedis('a stream on another instance receives per-session messages sent from either instance', async (redis, t: TestContext) => { + const a = await startServer(t, { redis: redisConfig(redis) }) + const b = await startServer(t, { redis: redisConfig(redis) }) + + const sessionId = await initialize(a.baseUrl) + const stream = await openStream(t, b.baseUrl, sessionId) + + t.assert.strictEqual(await a.app.mcpSendToSession(sessionId, note('from-a')), true) + t.assert.strictEqual(await b.app.mcpSendToSession(sessionId, note('from-b')), true) + + const data = await readUntil(stream, 'from-b') + t.assert.ok(data.includes('from-a'), 'message sent from the other instance must be delivered') + t.assert.deepStrictEqual(eventIds(data), ['1', '2']) + + // Each message is recorded in history exactly once + const store = new RedisSessionStore({ redis }) + const history = await store.getMessagesFrom(sessionId, '0') + t.assert.deepStrictEqual(history.map(entry => entry.eventId), ['1', '2']) + }) + + testWithRedis('DELETE on one instance closes the stream held by another', async (redis, t: TestContext) => { + const a = await startServer(t, { redis: redisConfig(redis) }) + const b = await startServer(t, { redis: redisConfig(redis) }) + + const sessionId = await initialize(a.baseUrl) + const stream = await openStream(t, b.baseUrl, sessionId) + const ended = waitForEnd(stream) + + const response = await request(`${a.baseUrl}/mcp`, { + method: 'DELETE', + headers: { 'mcp-session-id': sessionId } + }) + await response.body.dump() + t.assert.strictEqual(response.statusCode, 204) + + t.assert.strictEqual(await ended, 'ended') + t.assert.ok(await waitUntil(async () => await subscriberCount(redis, sessionId) === 0), + 'no instance may stay subscribed to a deleted session') + }) +}) + +describe('legacy SSE session subscriptions do not leak (#329)', () => { + testWithRedis('expired sessions leave no broker subscription behind', async (redis, t: TestContext) => { + const { baseUrl } = await startServer(t, { redis: redisConfig(redis), sessionTtlMs: 200 }) + + const sessionIds: string[] = [] + for (let i = 0; i < 10; i++) { + sessionIds.push(await initialize(baseUrl)) + } + + await sleep(600) + + for (const sessionId of sessionIds) { + t.assert.strictEqual(await redis.exists(`session:${sessionId}`), 0) + t.assert.strictEqual(await subscriberCount(redis, sessionId), 0) + } + }) + + testWithRedis('a closed stream leaves no broker subscription behind', async (redis, t: TestContext) => { + const server = await startServer(t, { redis: redisConfig(redis) }) + const sessionId = await initialize(server.baseUrl) + + const stream = await openStream(t, server.baseUrl, sessionId) + t.assert.strictEqual(await subscriberCount(redis, sessionId), 1) + + stream.destroy() + await server.waitForStreamCloses(1) + t.assert.ok(await waitUntil(async () => await subscriberCount(redis, sessionId) === 0)) + }) + + testWithRedis('a stream whose session disappears is closed and unsubscribed by the sweep', async (redis, t: TestContext) => { + const server = await startServer(t, { redis: redisConfig(redis), sessionTtlMs: 300 }) + const sessionId = await initialize(server.baseUrl) + + const stream = await openStream(t, server.baseUrl, sessionId) + const ended = waitForEnd(stream) + + // Simulate the key vanishing (eviction, flush) while the stream is open + await redis.del(`session:${sessionId}`) + + t.assert.strictEqual(await ended, 'ended') + t.assert.ok(await waitUntil(async () => await subscriberCount(redis, sessionId) === 0)) + }) +}) + +describe('SessionTopicSubscriptions', () => { + test('a stream attaching while the last one detaches keeps the subscription', async (t: TestContext) => { + const calls: string[] = [] + const active = new Set() + let releaseUnsubscribe!: () => void + const broker = { + async publish () {}, + async subscribe (topic: string) { + calls.push(`subscribe ${topic}`) + active.add(topic) + }, + async unsubscribe (topic: string) { + calls.push(`unsubscribe ${topic}`) + // A slow UNSUBSCRIBE must not win over the subscribe queued after it + await new Promise(resolve => { releaseUnsubscribe = resolve }) + active.delete(topic) + }, + async close () {} + } + let streams = 1 + const subscriptions = new SessionTopicSubscriptions(broker, () => streams > 0, () => {}) + + await subscriptions.sync('s1') + streams = 0 + const detached = subscriptions.sync('s1') + // Let the unsubscribe start before the next stream attaches + await sleep(5) + t.assert.strictEqual(typeof releaseUnsubscribe, 'function') + streams = 1 + const attached = subscriptions.sync('s1') + await sleep(10) + releaseUnsubscribe() + await Promise.all([detached, attached]) + + t.assert.deepStrictEqual(calls, [ + 'subscribe mcp/session/s1/message', + 'unsubscribe mcp/session/s1/message', + 'subscribe mcp/session/s1/message' + ]) + t.assert.ok(active.has('mcp/session/s1/message')) + t.assert.strictEqual(subscriptions.isSubscribed('s1'), true) + }) + + testWithRedis('RedisMessageBroker.unsubscribe releases the Redis channel', async (redis, t: TestContext) => { + const broker = new RedisMessageBroker(redis) + t.after(() => broker.close()) + + await broker.subscribe('mcp/session/numsub-check/message', () => {}) + t.assert.strictEqual(await subscriberCount(redis, 'numsub-check'), 1) + + await broker.unsubscribe('mcp/session/numsub-check/message') + t.assert.strictEqual(await subscriberCount(redis, 'numsub-check'), 0) + + // And it can subscribe again afterwards + await broker.subscribe('mcp/session/numsub-check/message', () => {}) + t.assert.strictEqual(await subscriberCount(redis, 'numsub-check'), 1) + }) +}) From 6bdf51ebae1d09c73296edf2bec7085e9665a518 Mon Sep 17 00:00:00 2001 From: Matteo Collina Date: Mon, 19 Oct 2026 16:17:18 +0200 Subject: [PATCH 2/2] fix(resources): call the unsubscribe handler when a legacy session ends When a legacy session was DELETEd or expired, its resource subscriptions went away with it but the unsubscribe handler never ran. A handler that starts a watcher on subscribe leaked it. Each instance now keeps a per-URI holder count of the sessions it ran the subscribe handler for. When a session stops holding a URI without an explicit resources/unsubscribe, the handler runs once no other session on that instance holds the URI. That covers DELETE, a session termination published by another instance, and, through the sweep, expiry (including Redis key TTL expiry) and unsubscribes or DELETEs handled by another instance. An explicit resources/unsubscribe still calls the handler directly, since its result is the response. On session end there is no request to answer, so the unsubscribe handler context now types request and reply as optional (ResourceUnsubscribeContext). Fixes #337 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LmhsFjk8W96X2PvmShmDm9 --- README.md | 5 + src/handlers.ts | 8 + src/index.ts | 1 + src/resource-subscription-holders.ts | 59 ++++++ src/routes/mcp.ts | 40 ++++ src/types.ts | 13 +- test/resource-subscription-lifecycle.test.ts | 201 +++++++++++++++++++ 7 files changed, 326 insertions(+), 1 deletion(-) create mode 100644 src/resource-subscription-holders.ts create mode 100644 test/resource-subscription-lifecycle.test.ts diff --git a/README.md b/README.md index 388b8b50..f057c7ce 100644 --- a/README.md +++ b/README.md @@ -2333,6 +2333,11 @@ app.mcpAddPrompt({ - `app.mcpBroadcastNotification(notification)`: Broadcast a notification to all connected SSE clients (works across Redis instances). `notifications/resources/updated` only reaches sessions subscribed to its `params.uri` - `app.mcpSendToSession(sessionId, message)`: Send a message/request to a specific SSE session (works across Redis instances) +### Resource Subscription Handlers + +- `app.mcpSetResourceSubscribeHandler(handler)`: called for every `resources/subscribe`; its return value is the response result. +- `app.mcpSetResourceUnsubscribeHandler(handler)`: called for every `resources/unsubscribe`. It is also called when a session ends without unsubscribing (DELETE, or idle expiry), for each URI it held, but only on the instance that ran the subscribe handler and only once no other session on that instance still holds the URI. In that case there is no request to answer, so `context.request` and `context.reply` are undefined (`ResourceUnsubscribeContext`). Expiry, and sessions deleted or unsubscribed through another instance, are picked up by a periodic sweep (at most every `min(sessionTtlMs, 60s)`). + ## Handler API Reference All MCP handlers follow a consistent pattern where the context parameter is always provided as the final parameter: diff --git a/src/handlers.ts b/src/handlers.ts index 985240c9..f9962aa9 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -1,3 +1,4 @@ +import type { ResourceSubscriptionHolders } from './resource-subscription-holders.ts' import { randomUUID } from 'node:crypto' import type { FastifyInstance, FastifyRequest, FastifyReply } from 'fastify' import createFastifyError from 'fastify-error' @@ -79,6 +80,8 @@ export type HandlerDependencies = { shutdownSignal?: AbortSignal taskLeases?: TaskLeases jsonSchemaValidator?: JsonSchemaValidator + /** Sessions holding resource subscriptions made through this instance */ + resourceHolders?: ResourceSubscriptionHolders sessionId?: string /** The revision this client negotiated; responses are shaped to match it */ protocolVersion?: string @@ -1475,6 +1478,8 @@ async function handleResourcesSubscribe ( // the sessions that asked for this URI if (sessionId && dependencies.sessionStore) { await dependencies.sessionStore.addResourceSubscription(sessionId, params.uri) + // Recorded after the store, so the sweep never sees a hold the store lacks + dependencies.resourceHolders?.add(sessionId, params.uri) } return createResponse(request.id, result) } catch (error: any) { @@ -1506,6 +1511,9 @@ async function handleResourcesUnsubscribe ( authContext: dependencies.authContext }) if (sessionId && dependencies.sessionStore) { + // The handler just ran, so this hold must not be released a second time + // by the sweep once the store no longer lists it + dependencies.resourceHolders?.release(sessionId, [params.uri]) await dependencies.sessionStore.removeResourceSubscription(sessionId, params.uri) } return createResponse(request.id, result) diff --git a/src/index.ts b/src/index.ts index d66c7917..186fb457 100644 --- a/src/index.ts +++ b/src/index.ts @@ -350,6 +350,7 @@ export type { ResourceHandlers, ResourceSubscribeHandler, ResourceUnsubscribeHandler, + ResourceUnsubscribeContext, TracerLike } from './types.ts' diff --git a/src/resource-subscription-holders.ts b/src/resource-subscription-holders.ts new file mode 100644 index 00000000..2d487325 --- /dev/null +++ b/src/resource-subscription-holders.ts @@ -0,0 +1,59 @@ +/** + * Which sessions hold a resource subscription on this instance, per URI. + * + * A session is recorded here when this instance ran the subscribe handler for + * it, since that is where any watcher the handler started lives. When sessions + * stop holding a URI without an explicit `resources/unsubscribe` (DELETE, + * expiry), the unsubscribe handler must run, but only once no other session + * on this instance still holds that URI. + */ +export class ResourceSubscriptionHolders { + private readonly byUri = new Map>() + private readonly bySession = new Map>() + + add (sessionId: string, uri: string): void { + let sessions = this.byUri.get(uri) + if (!sessions) { + sessions = new Set() + this.byUri.set(uri, sessions) + } + sessions.add(sessionId) + + let uris = this.bySession.get(sessionId) + if (!uris) { + uris = new Set() + this.bySession.set(sessionId, uris) + } + uris.add(uri) + } + + /** + * Drop a session's hold on some URIs (all of them when `uris` is omitted). + * Returns the URIs that no session holds any more, for which the unsubscribe + * handler should run. A URI the session did not hold is ignored. + */ + release (sessionId: string, uris?: Iterable): string[] { + const held = this.bySession.get(sessionId) + if (!held) return [] + + const released: string[] = [] + for (const uri of [...(uris ?? held)]) { + if (!held.delete(uri)) continue + const sessions = this.byUri.get(uri) + sessions?.delete(sessionId) + if (!sessions || sessions.size === 0) { + this.byUri.delete(uri) + released.push(uri) + } + } + if (held.size === 0) { + this.bySession.delete(sessionId) + } + return released + } + + /** Snapshot of every session holding a URI here, with the URIs it holds */ + entries (): Array<[string, string[]]> { + return [...this.bySession].map(([sessionId, uris]) => [sessionId, [...uris]]) + } +} diff --git a/src/routes/mcp.ts b/src/routes/mcp.ts index ce4c8b94..f33f3579 100644 --- a/src/routes/mcp.ts +++ b/src/routes/mcp.ts @@ -23,6 +23,7 @@ import { processMessage, checkRequestId } from '../handlers.ts' import { isStdioRequest } from '../stdio-marker.ts' import { appendToHistory, publishSessionTerminated, SessionTopicSubscriptions } from '../session-topic.ts' import type { SessionTopicPayload } from '../session-topic.ts' +import { ResourceSubscriptionHolders } from '../resource-subscription-holders.ts' const MAX_TIMER_DELAY_MS = 2 ** 31 - 1 @@ -334,6 +335,25 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async } const streamStates = new WeakMap() + const resourceHolders = new ResourceSubscriptionHolders() + + /** + * A session stopped holding some URIs without asking (DELETE, expiry, or an + * unsubscribe handled by another instance): run the unsubscribe handler for + * each URI no other session on this instance still holds. + */ + async function releaseResourceSubscriptions (sessionId: string, uris?: string[], authContext?: AuthorizationContext): Promise { + const released = resourceHolders.release(sessionId, uris) + const handler = resourceHandlers.unsubscribeHandler + if (!handler) return + for (const uri of released) { + try { + await handler({ uri }, { sessionId, authContext }) + } catch (error) { + app.log.error({ err: error, sessionId, uri }, 'Resource unsubscribe handler failed while releasing a session subscription') + } + } + } // This instance is subscribed to a session's topic exactly while it holds one // of the session's streams; the subscription is what lets a stream on any @@ -353,6 +373,7 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async async function onSessionTopicMessage (sessionId: string, payload: SessionTopicPayload): Promise { if (payload.kind === 'terminated') { endLocalStreams(sessionId) + await releaseResourceSubscriptions(sessionId) return } @@ -545,6 +566,7 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async shutdownSignal, taskLeases, jsonSchemaValidator, + resourceHolders, sessionId, protocolVersion: (request as any).mcpProtocolVersion ?? DEFAULT_NEGOTIATED_PROTOCOL_VERSION }) @@ -762,6 +784,7 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async // End this instance's streams, and tell the others to end theirs endLocalStreams(sessionId) + await releaseResourceSubscriptions(sessionId, undefined, session.authorization) if (enableSSE) { try { await publishSessionTerminated(messageBroker, sessionId) @@ -822,6 +845,23 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async endLocalStreams(sessionId) } } + + // Resource subscriptions this instance ran the subscribe handler for, whose + // session ended or which were unsubscribed through another instance. The + // snapshot is taken before reading the store: a subscribe is recorded in the + // store before it is recorded here, so it cannot look stale. + for (const [sessionId, uris] of resourceHolders.entries()) { + const session = await sessionStore.get(sessionId) + if (!session) { + await releaseResourceSubscriptions(sessionId, uris) + continue + } + const stored = new Set(session.resourceSubscriptions ?? []) + const stale = uris.filter(uri => !stored.has(uri)) + if (stale.length > 0) { + await releaseResourceSubscriptions(sessionId, stale, session.authorization) + } + } } const sweepTimer = setInterval(() => { if (sweeping) return diff --git a/src/types.ts b/src/types.ts index 35d21dbd..1ad4cb39 100644 --- a/src/types.ts +++ b/src/types.ts @@ -43,9 +43,20 @@ export type ResourceSubscribeHandler = ( context: HandlerContext ) => Promise> | Record +/** + * Context for the unsubscribe handler. Besides explicit `resources/unsubscribe` + * requests, the handler also runs when a session ends (DELETE or expiry) and + * was the last one on this instance holding the URI; then there is no request + * to answer, so `request` and `reply` are absent. + */ +export interface ResourceUnsubscribeContext extends Omit { + request?: FastifyRequest + reply?: FastifyReply +} + export type ResourceUnsubscribeHandler = ( params: { uri: string }, - context: HandlerContext + context: ResourceUnsubscribeContext ) => Promise> | Record // Resource handlers container diff --git a/test/resource-subscription-lifecycle.test.ts b/test/resource-subscription-lifecycle.test.ts new file mode 100644 index 00000000..2b62b278 --- /dev/null +++ b/test/resource-subscription-lifecycle.test.ts @@ -0,0 +1,201 @@ +import { test, describe } from 'node:test' +import type { TestContext } from 'node:test' +import { setTimeout as sleep } from 'node:timers/promises' +import Fastify from 'fastify' +import type { FastifyInstance } from 'fastify' +import type { Redis } from 'ioredis' +import mcpPlugin from '../src/index.ts' +import type { MCPPluginOptions } from '../src/types.ts' +import { JSONRPC_VERSION, LATEST_PROTOCOL_VERSION } from '../src/schema.ts' +import { testWithRedis } from './redis-test-utils.ts' + +interface HandlerCall { + uri: string + sessionId?: string +} + +interface TestApp { + app: FastifyInstance + subscribed: HandlerCall[] + unsubscribed: HandlerCall[] +} + +async function createApp (t: TestContext, opts: Partial = {}): Promise { + const app = Fastify({ logger: false }) + t.after(() => app.close()) + + await app.register(mcpPlugin, { + serverInfo: { name: 'test', version: '1.0.0' }, + capabilities: { resources: { subscribe: true } }, + enableSSE: true, + ...opts + }) + + const subscribed: HandlerCall[] = [] + const unsubscribed: HandlerCall[] = [] + app.mcpSetResourceSubscribeHandler(async ({ uri }, { sessionId }) => { + subscribed.push({ uri, sessionId }) + return {} + }) + app.mcpSetResourceUnsubscribeHandler(async ({ uri }, { sessionId }) => { + unsubscribed.push({ uri, sessionId }) + return {} + }) + + await app.ready() + return { app, subscribed, unsubscribed } +} + +function redisOpts (redis: Redis): Partial { + return { redis: { host: redis.options.host!, port: redis.options.port!, db: redis.options.db! } } +} + +async function initSession (app: FastifyInstance): Promise { + const response = await app.inject({ + method: 'POST', + url: '/mcp', + headers: { 'content-type': 'application/json', accept: 'application/json' }, + payload: { + jsonrpc: JSONRPC_VERSION, + id: 1, + method: 'initialize', + params: { + protocolVersion: LATEST_PROTOCOL_VERSION, + capabilities: {}, + clientInfo: { name: 'test-client', version: '1.0.0' } + } + } + }) + return response.headers['mcp-session-id'] as string +} + +async function rpc (app: FastifyInstance, sessionId: string, method: string, uri: string): Promise { + const response = await app.inject({ + method: 'POST', + url: '/mcp', + headers: { + 'content-type': 'application/json', + accept: 'application/json', + 'mcp-session-id': sessionId, + 'mcp-protocol-version': LATEST_PROTOCOL_VERSION + }, + payload: { jsonrpc: JSONRPC_VERSION, id: 2, method, params: { uri } } + }) + const body = response.json() + if (body.error) throw new Error(body.error.message) +} + +async function deleteSession (app: FastifyInstance, sessionId: string): Promise { + const response = await app.inject({ + method: 'DELETE', + url: '/mcp', + headers: { 'mcp-session-id': sessionId, 'mcp-protocol-version': LATEST_PROTOCOL_VERSION } + }) + if (response.statusCode !== 204) throw new Error(`DELETE answered ${response.statusCode}`) +} + +async function waitFor (predicate: () => boolean, timeoutMs = 3000): Promise { + const deadline = Date.now() + timeoutMs + while (!predicate() && Date.now() < deadline) { + await sleep(20) + } +} + +describe('resource unsubscribe handler runs when a legacy session ends (#337)', () => { + test('DELETE calls the unsubscribe handler for every URI the session held', async (t: TestContext) => { + const { app, unsubscribed } = await createApp(t) + const sessionId = await initSession(app) + await rpc(app, sessionId, 'resources/subscribe', 'file:///a.txt') + await rpc(app, sessionId, 'resources/subscribe', 'file:///b.txt') + + await deleteSession(app, sessionId) + + t.assert.deepStrictEqual( + unsubscribed.sort((x, y) => x.uri.localeCompare(y.uri)), + [{ uri: 'file:///a.txt', sessionId }, { uri: 'file:///b.txt', sessionId }] + ) + }) + + test('DELETE releases subscriptions when SSE is disabled too', async (t: TestContext) => { + const { app, unsubscribed } = await createApp(t, { enableSSE: false }) + const sessionId = await initSession(app) + await rpc(app, sessionId, 'resources/subscribe', 'file:///a.txt') + + await deleteSession(app, sessionId) + + t.assert.deepStrictEqual(unsubscribed, [{ uri: 'file:///a.txt', sessionId }]) + }) + + test('expiry calls the unsubscribe handler', async (t: TestContext) => { + const { app, unsubscribed } = await createApp(t, { sessionTtlMs: 150 }) + const sessionId = await initSession(app) + await rpc(app, sessionId, 'resources/subscribe', 'file:///a.txt') + + await waitFor(() => unsubscribed.length > 0) + + t.assert.deepStrictEqual(unsubscribed, [{ uri: 'file:///a.txt', sessionId }]) + }) + + test('the handler runs only when the last session holding a URI ends', async (t: TestContext) => { + const { app, unsubscribed } = await createApp(t) + const first = await initSession(app) + const second = await initSession(app) + await rpc(app, first, 'resources/subscribe', 'file:///shared.txt') + await rpc(app, second, 'resources/subscribe', 'file:///shared.txt') + + await deleteSession(app, first) + t.assert.deepStrictEqual(unsubscribed, [], 'another session still holds the URI') + + await deleteSession(app, second) + t.assert.deepStrictEqual(unsubscribed, [{ uri: 'file:///shared.txt', sessionId: second }]) + }) + + test('a URI already unsubscribed explicitly is not released again', async (t: TestContext) => { + const { app, unsubscribed } = await createApp(t) + const sessionId = await initSession(app) + await rpc(app, sessionId, 'resources/subscribe', 'file:///a.txt') + await rpc(app, sessionId, 'resources/unsubscribe', 'file:///a.txt') + t.assert.strictEqual(unsubscribed.length, 1) + + await deleteSession(app, sessionId) + t.assert.strictEqual(unsubscribed.length, 1) + }) + + testWithRedis('Redis: expiry through the key TTL calls the unsubscribe handler', async (redis, t: TestContext) => { + const { app, unsubscribed } = await createApp(t, { ...redisOpts(redis), sessionTtlMs: 200 }) + const sessionId = await initSession(app) + await rpc(app, sessionId, 'resources/subscribe', 'file:///a.txt') + + await waitFor(() => unsubscribed.length > 0) + + t.assert.strictEqual(await redis.exists(`session:${sessionId}`), 0) + t.assert.deepStrictEqual(unsubscribed, [{ uri: 'file:///a.txt', sessionId }]) + }) + + testWithRedis('Redis: the instance that ran the subscribe handler releases it when another deletes the session', async (redis, t: TestContext) => { + const a = await createApp(t, { ...redisOpts(redis), sessionTtlMs: 400 }) + const b = await createApp(t, { ...redisOpts(redis), sessionTtlMs: 400 }) + const sessionId = await initSession(a.app) + await rpc(b.app, sessionId, 'resources/subscribe', 'file:///a.txt') + t.assert.strictEqual(b.subscribed.length, 1) + + await deleteSession(a.app, sessionId) + + await waitFor(() => b.unsubscribed.length > 0) + t.assert.deepStrictEqual(b.unsubscribed, [{ uri: 'file:///a.txt', sessionId }]) + t.assert.deepStrictEqual(a.unsubscribed, [], 'instance A never ran the subscribe handler') + }) + + testWithRedis('Redis: an explicit unsubscribe on another instance is released where the subscribe ran', async (redis, t: TestContext) => { + const a = await createApp(t, { ...redisOpts(redis), sessionTtlMs: 400 }) + const b = await createApp(t, { ...redisOpts(redis), sessionTtlMs: 400 }) + const sessionId = await initSession(a.app) + await rpc(b.app, sessionId, 'resources/subscribe', 'file:///a.txt') + + await rpc(a.app, sessionId, 'resources/unsubscribe', 'file:///a.txt') + t.assert.deepStrictEqual(a.unsubscribed, [{ uri: 'file:///a.txt', sessionId }]) + + await waitFor(() => b.unsubscribed.length > 0) + t.assert.deepStrictEqual(b.unsubscribed, [{ uri: 'file:///a.txt', sessionId }]) + }) +})