fix(sse): coherent legacy session subscription lifecycle (#327, #328, #329, #337)

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/<id>/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:
@@ -2328,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:

src/brokers/redis-message-broker.ts

@@ -24,7 +24,6 @@ function suppressUnhandledQuitRejection (conn?: RedisConnLike): void {
 
 interface RedisMQEmitter extends ReturnType<typeof MQEmitterRedis> {
   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<void> | null = null
+  private readonly listeners = new Map<string, Set<(packet: any, cb: () => 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<void> {
-    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<void> {
-    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<void>((resolve) => {
+      this.emitter.removeListener(topic, listener, () => resolve())
+    })))
   }
 
   async close (): Promise<void> {

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<MCPPubSubDecoratorsOptions>
       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

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)

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<string, Set<any>>()
 
-  // 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<void>[] = []
-    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()
@@ -392,6 +350,7 @@ export type {
   ResourceHandlers,
   ResourceSubscribeHandler,
   ResourceUnsubscribeHandler,
+  ResourceUnsubscribeContext,
   TracerLike
 } from './types.ts'
 

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<string, Set<string>>()
+  private readonly bySession = new Map<string, Set<string>>()
+
+  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>): 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]])
+  }
+}