From bb41d1f513581bb4e94c12f5cd987ae667444d1a Mon Sep 17 00:00:00 2001 From: Matteo Collina Date: Thu, 15 Oct 2026 16:49:26 +0000 Subject: [PATCH] feat: add configurable limits for sessions and resource subscriptions Add two opt-in limits so retained state cannot be grown without bound: - maxSessionsPerPrincipal caps the non-expired sessions a caller may keep active. Keyed by the session's (issuer, subject) pair when authorization identifies a caller, else a shared anonymous bucket. At the cap a new session is refused with SESSION_LIMIT_REACHED, and existing sessions keep working and expire on their own. The Redis store enforces it with an atomic Lua script (pruning expired ids on each create) so concurrent creates across instances cannot overshoot the cap. - maxResourceSubscriptionsPerSession caps the distinct URIs one session may subscribe to. Duplicates stay idempotent; only a new URI at the cap fails with SUBSCRIPTION_LIMIT_REACHED. Enforced atomically on Redis. The memory and Redis session stores both accept the options (default Infinity, no change for existing deployments). The subscribe handler and the POST/GET /mcp session-creation paths surface the errors as stable protocol errors. The GET SSE path now resolves the session before hijacking the reply, so a refused session answers with a 429 instead of a hung stream. Adds memory + Redis store tests and HTTP integration tests covering duplicate subscriptions, per-principal isolation, expiry cleanup, and concurrent Redis writes. Documents both options in the README and types. --- README.md | 29 +- src/handlers.ts | 10 +- src/index.ts | 38 ++- src/routes/mcp.ts | 43 ++- src/stores/memory-session-store.ts | 56 +++- src/stores/redis-session-store.ts | 147 +++++++++- src/stores/session-store.ts | 73 +++++ src/types.ts | 19 ++ test/session-subscription-limits.test.ts | 342 +++++++++++++++++++++++ 9 files changed, 726 insertions(+), 31 deletions(-) create mode 100644 test/session-subscription-limits.test.ts diff --git a/README.md b/README.md index bfd15b3..e21795e 100644 --- a/README.md +++ b/README.md @@ -263,6 +263,31 @@ await app.register(mcpPlugin, { }) ``` +### Session and subscription limits + +By default there is no bound on how many sessions one caller may keep open or on how many +resources one session may subscribe to. Set `maxSessionsPerPrincipal` and +`maxResourceSubscriptionsPerSession` to bound retained state; each defaults to `Infinity`. + +```typescript +await app.register(mcpPlugin, { + maxSessionsPerPrincipal: 10, // active sessions per caller + maxResourceSubscriptionsPerSession: 50 +}) +``` + +- **`maxSessionsPerPrincipal`** caps the non-expired sessions a caller may keep active. When +authorization identifies a caller the cap is keyed by the token's (issuer, subject) pair, so +each principal gets its own quota; without authorization, or for requests that identify no +caller, every session shares one anonymous bucket. At the cap a new session is refused with +`SESSION_LIMIT_REACHED` (`-32035`) — an existing session is never torn down, and sessions +still expire on their own via `sessionTtlMs`. The Redis store enforces the cap atomically, so +concurrent session creation across instances cannot overshoot it. +- **`maxResourceSubscriptionsPerSession`** caps the distinct resource URIs one session may +subscribe to via `resources/subscribe`. Subscribing to an already-subscribed URI stays +idempotent; only at the cap does a new URI fail with `SUBSCRIPTION_LIMIT_REACHED` (`-32036`). +The Redis store enforces the cap atomically, so concurrent subscribes cannot overshoot it. + ## Request Body Errors Body-parser failures and bodies that are not JSON-RPC messages on `POST /mcp` are answered @@ -916,9 +941,11 @@ const store: SessionStore = new MemorySessionStore() ``` - **`SessionStore`**: the interface the plugin uses to persist SSE session metadata, message history, token-to-session mappings, and per-session resource subscriptions (`addResourceSubscription` / `removeResourceSubscription`, exposed as `resourceSubscriptions` on the metadata). The optional `touch(sessionId)` records activity so the session is not expired while in use. -- **`MemorySessionStore`**: in-process implementation used automatically when no `redis` option is passed. Takes an optional `maxMessages` history cap (default 100) and `{ sessionTtlMs }` idle expiry (default 1 hour); call `cleanup()` periodically to free expired sessions when using it outside the plugin. +- **`MemorySessionStore`**: in-process implementation used automatically when no `redis` option is passed. Takes an optional `maxMessages` history cap (default 100) and `{ sessionTtlMs }` idle expiry (default 1 hour); call `cleanup()` periodically to free expired sessions when using it outside the plugin. Both stores also accept `maxSessionsPerPrincipal` and `maxResourceSubscriptionsPerSession` to bound active sessions and per-session resource subscriptions. - **`RedisSessionStore`**: distributed implementation used automatically when the `redis` option is passed; sessions and message history live in Redis so any instance can serve any client. +`SessionLimitError` and `SubscriptionLimitError` (with `SESSION_LIMIT_REACHED` and `SUBSCRIPTION_LIMIT_REACHED`) are also exported, so you can catch them when calling a store directly or distinguish those failures from an internal error. + These are the same store classes the plugin instantiates internally based on the `redis` option. Exporting them lets you construct or inspect one directly in tests and scripts. ### Multi-Instance Deployment diff --git a/src/handlers.ts b/src/handlers.ts index 985240c..244dec0 100644 --- a/src/handlers.ts +++ b/src/handlers.ts @@ -37,6 +37,7 @@ import type { RequestId } from './schema.ts' import type { MCPTool, MCPResource, MCPPrompt, MCPPluginOptions, ResourceHandlers, McpCallToolOutcome, ToolAccessOperation, MCPToolCallCompleteEvent, TracerLike } from './types.ts' import type { SessionStore } from './stores/session-store.ts' +import { SubscriptionLimitError, SUBSCRIPTION_LIMIT_REACHED } from './stores/session-store.ts' import type { TaskStore, TaskRecord, TaskOwner, TaskWaiters, TaskLeases } from './stores/task-store.ts' import { InvalidTaskCursorError, TaskLimitError, TASK_LIMIT_REACHED, TASK_OWNER_LOST, decodeTaskCursor, isTaskOwner, isTerminal, ownerLostOutcome, toWireTask } from './stores/task-store.ts' import type { AuthorizationContext } from './types/auth-types.ts' @@ -1474,7 +1475,14 @@ async function handleResourcesSubscribe ( // Remember the subscription so resources/updated broadcasts only reach // the sessions that asked for this URI if (sessionId && dependencies.sessionStore) { - await dependencies.sessionStore.addResourceSubscription(sessionId, params.uri) + try { + await dependencies.sessionStore.addResourceSubscription(sessionId, params.uri) + } catch (error: any) { + if (error instanceof SubscriptionLimitError) { + return createError(request.id, SUBSCRIPTION_LIMIT_REACHED, `Resource subscription limit reached (${error.limit})`) + } + throw error + } } return createResponse(request.id, result) } catch (error: any) { diff --git a/src/index.ts b/src/index.ts index a57f9b3..62938cc 100644 --- a/src/index.ts +++ b/src/index.ts @@ -2,6 +2,7 @@ import type { FastifyInstance, preHandlerHookHandler } from 'fastify' import fp from 'fastify-plugin' import { Redis } from 'ioredis' import type { SessionStore } from './stores/session-store.ts' +import { validateStoreLimit } from './stores/session-store.ts' import type { MessageBroker } from './brokers/message-broker.ts' import { MemorySessionStore } from './stores/memory-session-store.ts' import { MemoryMessageBroker } from './brokers/memory-message-broker.ts' @@ -98,6 +99,17 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption if (opts.sessionTtlMs !== undefined && !(Number.isFinite(opts.sessionTtlMs) && opts.sessionTtlMs > 0)) { throw new Error('sessionTtlMs must be a positive number of milliseconds') } + const validateSessionLimit = (name: string, value: unknown): void => { + // Validation happens here (before a Redis connection is opened) and inside + // the stores, so a bad value cannot leak a connection. + validateStoreLimit(name, value) + } + if (opts.maxSessionsPerPrincipal !== undefined) { + validateSessionLimit('maxSessionsPerPrincipal', opts.maxSessionsPerPrincipal) + } + if (opts.maxResourceSubscriptionsPerSession !== undefined) { + validateSessionLimit('maxResourceSubscriptionsPerSession', opts.maxResourceSubscriptionsPerSession) + } if (opts.taskMaxRetainedPerPrincipal !== undefined) { // Checked before any Redis connection is opened, so a bad value cannot leak one validateTaskLimit('taskMaxRetainedPerPrincipal', opts.taskMaxRetainedPerPrincipal) @@ -121,7 +133,13 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption redis.on('error', (err) => { app.log.error({ err }, 'Redis connection error') }) - sessionStore = new RedisSessionStore({ redis, maxMessages: 100, sessionTtlMs: opts.sessionTtlMs }) + sessionStore = new RedisSessionStore({ + redis, + maxMessages: 100, + sessionTtlMs: opts.sessionTtlMs, + maxSessionsPerPrincipal: opts.maxSessionsPerPrincipal, + maxResourceSubscriptionsPerSession: opts.maxResourceSubscriptionsPerSession + }) messageBroker = new RedisMessageBroker(redis, { onCloseTimeout: (closeTimeoutMs) => { app.log.warn({ closeTimeoutMs }, 'Redis message broker close timed out; forcing disconnect') @@ -139,7 +157,11 @@ const mcpPlugin = fp(async function (app: FastifyInstance, opts: MCPPluginOption } } else { // Memory implementations for single-instance deployment - sessionStore = new MemorySessionStore(100, { sessionTtlMs: opts.sessionTtlMs }) + sessionStore = new MemorySessionStore(100, { + sessionTtlMs: opts.sessionTtlMs, + maxSessionsPerPrincipal: opts.maxSessionsPerPrincipal, + maxResourceSubscriptionsPerSession: opts.maxResourceSubscriptionsPerSession + }) messageBroker = new MemoryMessageBroker() if (enableTasks) { taskStore = new MemoryTaskStore({ maxTasksPerPrincipal: opts.taskMaxRetainedPerPrincipal }) @@ -462,8 +484,18 @@ export { RedisTaskStore } from './stores/redis-task-store.ts' // Session storage, for callers that want to supply or inspect a backend export type { SessionStore, SessionMetadata } from './stores/session-store.ts' -export { MemorySessionStore } from './stores/memory-session-store.ts' +export { + MemorySessionStore +} from './stores/memory-session-store.ts' export { RedisSessionStore } from './stores/redis-session-store.ts' +export { + SessionLimitError, + SubscriptionLimitError, + SESSION_LIMIT_REACHED, + SUBSCRIPTION_LIMIT_REACHED, + DEFAULT_MAX_SESSIONS_PER_PRINCIPAL, + DEFAULT_MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION +} from './stores/session-store.ts' // In-process MCP client, for exercising a registered server via app.inject() export { createMcpClient } from './client.ts' diff --git a/src/routes/mcp.ts b/src/routes/mcp.ts index efdd63b..f8a2de5 100644 --- a/src/routes/mcp.ts +++ b/src/routes/mcp.ts @@ -14,12 +14,13 @@ import type { MCPRouteSchemaTransformer } from '../types.ts' import type { SessionStore, SessionMetadata } from '../stores/session-store.ts' +import { SessionLimitError, SESSION_LIMIT_REACHED } from '../stores/session-store.ts' import { DEFAULT_SESSION_TTL_MS } 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 { processMessage, checkRequestId, createError } from '../handlers.ts' import { isStdioRequest } from '../stdio-marker.ts' const MAX_TIMER_DELAY_MS = 2 ** 31 - 1 @@ -493,6 +494,15 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async reply.code(202) } } catch (error) { + if (error instanceof SessionLimitError) { + // A stable protocol error rather than an opaque 500, so a client that + // hit the active-session cap can react (e.g. free a session) instead + // of treating it as a transient server failure. + const id = Array.isArray(request.body) + ? null + : (request.body as { id?: string | number | null })?.id ?? null + return reply.type('application/json').code(429).send(createError(id, SESSION_LIMIT_REACHED, `Session limit reached (${error.limit})`)) + } app.log.error({ err: error }, 'Error processing MCP message') reply.type('application/json').code(500).send({ jsonrpc: JSONRPC_VERSION, @@ -542,16 +552,11 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async request.log.info({ sessionId }, 'Handling SSE request') - // We are opting out of Fastify proper - reply.hijack() - - const raw = reply.raw - - // Set up SSE stream - raw.setHeader('Content-type', 'text/event-stream') - raw.setHeader('Cache-Control', 'no-cache') - + // Resolve (or create) the session before hijacking the reply, so a + // failed session creation (e.g. the active-session cap) can still be + // answered with a normal HTTP response rather than a hung stream. let session: SessionMetadata + let isNewSession = false if (sessionId) { const existingSession = await sessionStore.get(sessionId) if (existingSession) { @@ -559,10 +564,22 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async await sessionStore.touch?.(session.id) } else { session = await createSession(request) - raw.setHeader('Mcp-Session-Id', session.id) + isNewSession = true } } else { session = await createSession(request) + isNewSession = true + } + + // We are opting out of Fastify proper + reply.hijack() + + const raw = reply.raw + + // Set up SSE stream + raw.setHeader('Content-type', 'text/event-stream') + raw.setHeader('Cache-Control', 'no-cache') + if (isNewSession) { raw.setHeader('Mcp-Session-Id', session.id) } @@ -685,6 +702,10 @@ const mcpPubSubRoutesPlugin: FastifyPluginAsync = async clearInterval(heartbeatInterval) }) } catch (error) { + if (error instanceof SessionLimitError) { + reply.type('application/json').code(429).send({ error: `Session limit reached (${error.limit})` }) + return + } app.log.error({ err: error }, 'Error setting up SSE stream') reply.type('application/json').code(500).send({ error: 'Internal server error' }) } diff --git a/src/stores/memory-session-store.ts b/src/stores/memory-session-store.ts index 658b14f..af40f69 100644 --- a/src/stores/memory-session-store.ts +++ b/src/stores/memory-session-store.ts @@ -1,5 +1,6 @@ import type { JSONRPCMessage } from '../schema.ts' import type { SessionStore, SessionMetadata } from './session-store.ts' +import { SessionLimitError, SubscriptionLimitError, DEFAULT_MAX_SESSIONS_PER_PRINCIPAL, DEFAULT_MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION, validateStoreLimit } from './session-store.ts' import type { AuthorizationContext, TokenRefreshInfo } from '../types/auth-types.ts' interface MessageHistoryEntry { @@ -16,6 +17,17 @@ export interface MemorySessionStoreOptions { * (default 1 hour, like the Redis store's TTL). */ sessionTtlMs?: number + /** + * Most non-expired sessions one principal may keep active (default + * Infinity). Keyed by the session's (ownerIss, ownerSub) pair; sessions with + * no owner share one anonymous bucket. + */ + maxSessionsPerPrincipal?: number + /** + * Most distinct resource URIs one session may subscribe to (default + * Infinity). Duplicate subscriptions are idempotent. + */ + maxResourceSubscriptionsPerSession?: number } export class MemorySessionStore implements SessionStore { @@ -24,10 +36,40 @@ export class MemorySessionStore implements SessionStore { private tokenToSession = new Map() // tokenHash -> sessionId private maxMessages: number readonly sessionTtlMs: number + private readonly maxSessionsPerPrincipal: number + private readonly maxResourceSubscriptionsPerSession: number constructor (maxMessages: number = 100, options: MemorySessionStoreOptions = {}) { this.maxMessages = maxMessages this.sessionTtlMs = options.sessionTtlMs ?? DEFAULT_SESSION_TTL_MS + this.maxSessionsPerPrincipal = options.maxSessionsPerPrincipal ?? DEFAULT_MAX_SESSIONS_PER_PRINCIPAL + this.maxResourceSubscriptionsPerSession = options.maxResourceSubscriptionsPerSession ?? DEFAULT_MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION + validateStoreLimit('maxSessionsPerPrincipal', this.maxSessionsPerPrincipal) + validateStoreLimit('maxResourceSubscriptionsPerSession', this.maxResourceSubscriptionsPerSession) + } + + /** + * The bucket a session counts against for the per-principal active-session cap. + * A session is owned when both the issuer and subject are known; otherwise all + * ownerless sessions share one anonymous bucket. + */ + private static ownerKey (metadata: SessionMetadata): string { + if (metadata.ownerSub === undefined && metadata.ownerIss === undefined) { + return '__anonymous__' + } + return `${metadata.ownerIss ?? ''}:${metadata.ownerSub ?? ''}` + } + + /** Count non-expired sessions in the same owner bucket. */ + private activeSessionCountFor (metadata: SessionMetadata): number { + const key = MemorySessionStore.ownerKey(metadata) + let count = 0 + for (const [otherId, other] of this.sessions) { + if (otherId === metadata.id) continue + if (this.isExpired(other)) continue + if (MemorySessionStore.ownerKey(other) === key) count++ + } + return count } private isExpired (session: SessionMetadata): boolean { @@ -50,6 +92,10 @@ export class MemorySessionStore implements SessionStore { } async create (metadata: SessionMetadata): Promise { + if (this.maxSessionsPerPrincipal !== Infinity && + this.activeSessionCountFor(metadata) >= this.maxSessionsPerPrincipal) { + throw new SessionLimitError(this.maxSessionsPerPrincipal) + } this.sessions.set(metadata.id, { ...metadata }) this.messageHistory.set(metadata.id, []) } @@ -209,9 +255,15 @@ export class MemorySessionStore implements SessionStore { const session = this.live(sessionId) if (!session) return const subscriptions = session.resourceSubscriptions ??= [] - if (!subscriptions.includes(uri)) { - subscriptions.push(uri) + // Duplicate subscriptions stay idempotent: the URI is already present. + if (subscriptions.includes(uri)) { + return + } + if (this.maxResourceSubscriptionsPerSession !== Infinity && + subscriptions.length >= this.maxResourceSubscriptionsPerSession) { + throw new SubscriptionLimitError(this.maxResourceSubscriptionsPerSession) } + subscriptions.push(uri) } async removeResourceSubscription (sessionId: string, uri: string): Promise { diff --git a/src/stores/redis-session-store.ts b/src/stores/redis-session-store.ts index efbd79f..fd2bd13 100644 --- a/src/stores/redis-session-store.ts +++ b/src/stores/redis-session-store.ts @@ -1,6 +1,7 @@ import type { Redis } from 'ioredis' import type { JSONRPCMessage } from '../schema.ts' import type { SessionStore, SessionMetadata } from './session-store.ts' +import { SessionLimitError, SubscriptionLimitError, DEFAULT_MAX_SESSIONS_PER_PRINCIPAL, DEFAULT_MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION, validateStoreLimit } from './session-store.ts' import type { AuthorizationContext, TokenRefreshInfo } from '../types/auth-types.ts' // Each subscribed URI is its own hash field, so concurrent subscribe and @@ -11,19 +12,101 @@ const RESOURCE_SUBSCRIPTION_PREFIX = 'resourceSubscription:' /** Idle time after which a session key expires (1 hour) */ const DEFAULT_SESSION_TTL_MS = 60 * 60 * 1000 +// Per-owner set of session ids, so `create` can enforce the active-session cap +// atomically. Session keys expire independently, so stale ids (whose key is +// gone) are pruned on every create for that owner, just like the task store. +const SESSION_OWNER_INDEX_PREFIX = 'sessionOwner:' +const SESSION_OWNER_INDEX_ANONYMOUS = 'sessionOwner:__anonymous__' + +// Create a session while enforcing its owner's active-session cap, atomically. +// +// KEYS[1] = session key, KEYS[2] = owner index (set of session ids) +// ARGV[1] = session id, ARGV[2] = sessionTtlMs, ARGV[3] = cap (-1 = unlimited), +// ARGV[4] = session key prefix, ARGV[5...] = hash field/value pairs +// +// Ids whose session key has expired are pruned first; then, at the cap, the +// create is refused. Returns 0 (storing nothing) when at the cap, else 1. +const CREATE_WITH_LIMIT_SCRIPT = ` +local prefix = ARGV[4] +local cap = tonumber(ARGV[3]) +local live = {} +for _, id in ipairs(redis.call('SMEMBERS', KEYS[2])) do + if redis.call('EXISTS', prefix .. id) == 1 then + table.insert(live, id) + else + redis.call('SREM', KEYS[2], id) + end +end + +if cap >= 0 and #live >= cap then + return 0 +end + +redis.call('DEL', KEYS[1]) +redis.call('HSET', KEYS[1], unpack(ARGV, 5)) +redis.call('PEXPIRE', KEYS[1], tonumber(ARGV[2])) +redis.call('SADD', KEYS[2], ARGV[1]) +return 1 +` + +// Add a resource subscription while enforcing the per-session cap, atomically, +// so concurrent subscribes cannot overshoot it. +// +// KEYS[1] = session key +// ARGV[1] = subscription field, ARGV[2] = subscription prefix, ARGV[3] = cap +// +// Returns: 0 = session gone (no-op), 1 = added, 2 = already subscribed +// (idempotent), -1 = at the subscription cap. +const ADD_SUBSCRIPTION_SCRIPT = ` +if redis.call('EXISTS', KEYS[1]) == 0 then + return 0 +end +if redis.call('HEXISTS', KEYS[1], ARGV[1]) == 1 then + return 2 +end +local prefix = ARGV[2] +local cap = tonumber(ARGV[3]) +if cap >= 0 then + local count = 0 + for _, field in ipairs(redis.call('HKEYS', KEYS[1])) do + if string.sub(field, 1, #prefix) == prefix then + count = count + 1 + end + end + if count >= cap then + return -1 + end +end +redis.call('HSET', KEYS[1], ARGV[1], '1') +return 1 +` + export class RedisSessionStore implements SessionStore { private redis: Redis private maxMessages: number readonly sessionTtlMs: number + private readonly maxSessionsPerPrincipal: number + private readonly maxResourceSubscriptionsPerSession: number - constructor (options: { redis: Redis, maxMessages?: number, sessionTtlMs?: number }) { + constructor (options: { redis: Redis, maxMessages?: number, sessionTtlMs?: number, maxSessionsPerPrincipal?: number, maxResourceSubscriptionsPerSession?: number }) { this.redis = options.redis this.maxMessages = options.maxMessages || 100 this.sessionTtlMs = options.sessionTtlMs ?? DEFAULT_SESSION_TTL_MS + this.maxSessionsPerPrincipal = options.maxSessionsPerPrincipal ?? DEFAULT_MAX_SESSIONS_PER_PRINCIPAL + this.maxResourceSubscriptionsPerSession = options.maxResourceSubscriptionsPerSession ?? DEFAULT_MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION + validateStoreLimit('maxSessionsPerPrincipal', this.maxSessionsPerPrincipal) + validateStoreLimit('maxResourceSubscriptionsPerSession', this.maxResourceSubscriptionsPerSession) } - async create (metadata: SessionMetadata): Promise { - const sessionKey = `session:${metadata.id}` + /** The set key holding the session ids of one owner bucket. */ + private ownerIndexKey (metadata: SessionMetadata): string { + if (metadata.ownerSub === undefined && metadata.ownerIss === undefined) { + return SESSION_OWNER_INDEX_ANONYMOUS + } + return `${SESSION_OWNER_INDEX_PREFIX}${encodeURIComponent(metadata.ownerIss ?? '')}:${encodeURIComponent(metadata.ownerSub ?? '')}` + } + + private sessionData (metadata: SessionMetadata): Record { const sessionData: Record = { id: metadata.id, eventId: metadata.eventId.toString(), @@ -57,10 +140,38 @@ export class RedisSessionStore implements SessionStore { sessionData.authSession = JSON.stringify(metadata.authSession) } - await this.redis.hset(sessionKey, sessionData) + return sessionData + } - // Expire the session after sessionTtlMs without activity - await this.redis.pexpire(sessionKey, this.sessionTtlMs) + async create (metadata: SessionMetadata): Promise { + const sessionKey = `session:${metadata.id}` + const sessionData = this.sessionData(metadata) + const args = Object.entries(sessionData).flat() + + if (this.maxSessionsPerPrincipal === Infinity) { + await this.redis.hset(sessionKey, sessionData) + + // Expire the session after sessionTtlMs without activity + await this.redis.pexpire(sessionKey, this.sessionTtlMs) + } else { + // Enforce the owner's active-session cap atomically: the check and the + // write happen in one script, so concurrent creates on different + // instances cannot overshoot the cap. + const stored = await this.redis.eval( + CREATE_WITH_LIMIT_SCRIPT, + 2, + sessionKey, + this.ownerIndexKey(metadata), + metadata.id, + String(this.sessionTtlMs), + String(this.maxSessionsPerPrincipal), + 'session:', + ...args + ) + if (stored !== 1) { + throw new SessionLimitError(this.maxSessionsPerPrincipal) + } + } // Add token mapping if present if (metadata.authorization?.tokenHash) { @@ -185,11 +296,16 @@ export class RedisSessionStore implements SessionStore { const sessionKey = `session:${sessionId}` const historyKey = `session:${sessionId}:history` - // Get session to clean up token mappings + // Get session to clean up token mappings and pull its id out of the owner + // index (a no-op when the session is already gone; the index prunes ids + // whose key has expired the next time that owner creates a session). const session = await this.get(sessionId) if (session?.authorization?.tokenHash) { await this.removeTokenMapping(session.authorization.tokenHash) } + if (session) { + await this.redis.srem(this.ownerIndexKey(session), sessionId) + } await this.redis.del(sessionKey, historyKey) } @@ -312,15 +428,20 @@ export class RedisSessionStore implements SessionStore { async addResourceSubscription (sessionId: string, uri: string): Promise { // Skip the write when the session is gone so we never recreate an expired - // session key without a TTL. - await this.redis.eval( - `if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end - redis.call('HSET', KEYS[1], ARGV[1], '1') - return 1`, + // session key without a TTL. Duplicate subscriptions stay idempotent, and + // the per-session cap (when configured) is enforced atomically so + // concurrent subscribes cannot overshoot it. + const result = await this.redis.eval( + ADD_SUBSCRIPTION_SCRIPT, 1, `session:${sessionId}`, - RESOURCE_SUBSCRIPTION_PREFIX + uri + RESOURCE_SUBSCRIPTION_PREFIX + uri, + RESOURCE_SUBSCRIPTION_PREFIX, + String(this.maxResourceSubscriptionsPerSession === Infinity ? -1 : this.maxResourceSubscriptionsPerSession) ) + if (result === -1) { + throw new SubscriptionLimitError(this.maxResourceSubscriptionsPerSession) + } } async removeResourceSubscription (sessionId: string, uri: string): Promise { diff --git a/src/stores/session-store.ts b/src/stores/session-store.ts index b6b9d6b..885e139 100644 --- a/src/stores/session-store.ts +++ b/src/stores/session-store.ts @@ -32,7 +32,74 @@ export interface SessionMetadata { resourceSubscriptions?: string[] } +/** + * How many sessions one principal may keep active, by default: unlimited. + * A session's principal is the (issuer, subject) pair that created it, or the + * shared anonymous bucket when no caller can be identified. + */ +export const DEFAULT_MAX_SESSIONS_PER_PRINCIPAL = Infinity + +/** + * How many distinct resource URIs one session may subscribe to, by default: + * unlimited. + */ +export const DEFAULT_MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION = Infinity + +/** + * JSON-RPC error code returned when a session cannot be created because its + * owner is at the active-session cap. Implementation-defined, in the + * -32000..-32099 server error range (matching TASK_LIMIT_REACHED). + */ +export const SESSION_LIMIT_REACHED = -32035 + +/** + * JSON-RPC error code returned when a resource subscription cannot be added + * because the session is at its subscription cap. + */ +export const SUBSCRIPTION_LIMIT_REACHED = -32036 + +/** + * Thrown by `SessionStore.create` when the owner is at the active-session cap. + */ +export class SessionLimitError extends Error { + readonly limit: number + + constructor (limit: number) { + super(`Session limit reached (${limit})`) + this.name = 'SessionLimitError' + this.limit = limit + } +} + +/** + * Thrown by `SessionStore.addResourceSubscription` when the session is at its + * subscription cap and the URI is not already subscribed. + */ +export class SubscriptionLimitError extends Error { + readonly limit: number + + constructor (limit: number) { + super(`Resource subscription limit reached (${limit})`) + this.name = 'SubscriptionLimitError' + this.limit = limit + } +} + +/** Validate a session/subscription limit option, accepting Infinity or a positive integer. */ +export function validateStoreLimit (name: string, value: unknown): void { + if (value === Infinity) return + if (typeof value !== 'number' || !Number.isInteger(value) || value < 1) { + throw new TypeError(`${name} must be a positive integer or Infinity`) + } +} + export interface SessionStore { + /** + * Store a new session, enforcing the owner's active-session cap when one is + * configured. The cap is keyed by the session's (ownerIss, ownerSub) pair, or + * by a shared anonymous bucket when the session has no owner. + * @throws SessionLimitError when the owner is at the active-session cap + */ create(metadata: SessionMetadata): Promise /** * Persist changes to an existing session without touching its message history. @@ -60,6 +127,12 @@ export interface SessionStore { updateAuthorization(sessionId: string, authorization: AuthorizationContext, tokenRefresh?: TokenRefreshInfo): Promise // Resource subscription operations (no-ops when the session is gone) + /** + * Record a resource subscription. Adding an already-subscribed URI is a + * no-op. When a per-session cap is configured and the session is at that cap, + * adding a new URI fails. + * @throws SubscriptionLimitError when the session is at the subscription cap + */ addResourceSubscription(sessionId: string, uri: string): Promise removeResourceSubscription(sessionId: string, uri: string): Promise } diff --git a/src/types.ts b/src/types.ts index 35d21db..2753938 100644 --- a/src/types.ts +++ b/src/types.ts @@ -303,6 +303,25 @@ export interface MCPPluginOptions { * the Redis session store. */ sessionTtlMs?: number + /** + * Most non-expired sessions one principal may keep active (default + * Infinity, no limit). A session's principal is the (issuer, subject) pair + * that created it; when authorization is off (or a caller cannot be + * identified) every session shares one anonymous bucket. At the cap a new + * session cannot be created and the request fails with a stable + * `SESSION_LIMIT_REACHED` error; existing sessions keep working and expire + * on their own. Uses a per-principal index so it can be enforced atomically + * across instances with the Redis store. + */ + maxSessionsPerPrincipal?: number + /** + * Most distinct resource URIs one session may subscribe to via + * `resources/subscribe` (default Infinity, no limit). Subscribing to an + * already-subscribed URI stays idempotent; only at the cap does a new URI + * fail with a stable `SUBSCRIPTION_LIMIT_REACHED` error. Enforced atomically + * on the Redis store so concurrent subscribes cannot overshoot it. + */ + maxResourceSubscriptionsPerSession?: number /** * Enable task-augmented execution (2025-11-25, experimental). Tools opt in * individually via `execution.taskSupport`. diff --git a/test/session-subscription-limits.test.ts b/test/session-subscription-limits.test.ts new file mode 100644 index 0000000..e0624ed --- /dev/null +++ b/test/session-subscription-limits.test.ts @@ -0,0 +1,342 @@ +import { describe, test } from 'node:test' +import type { TestContext } from 'node:test' +import fastify from 'fastify' +import type { FastifyInstance } from 'fastify' +import mcpPlugin from '../src/index.ts' +import { MemorySessionStore } from '../src/stores/memory-session-store.ts' +import { RedisSessionStore } from '../src/stores/redis-session-store.ts' +import { + SessionLimitError, + SubscriptionLimitError, + SESSION_LIMIT_REACHED, + SUBSCRIPTION_LIMIT_REACHED +} from '../src/stores/session-store.ts' +import type { SessionMetadata } from '../src/stores/session-store.ts' +import { JSONRPC_VERSION } from '../src/schema.ts' +import { testWithRedis } from './redis-test-utils.ts' + +function newSession (id: string, owner?: { ownerSub?: string, ownerIss?: string }): SessionMetadata { + return { + id, + eventId: 0, + createdAt: new Date(), + lastActivity: new Date(), + ...owner + } +} + +describe('MemorySessionStore session limit', () => { + test('rejects a new session when the owner is at the cap', async (t: TestContext) => { + const store = new MemorySessionStore(100, { maxSessionsPerPrincipal: 2 }) + await store.create(newSession('a', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + await store.create(newSession('b', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + + await t.assert.rejects( + store.create(newSession('c', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })), + (err: unknown) => err instanceof SessionLimitError && err.limit === 2 + ) + }) + + test('is isolated between principals', async (t: TestContext) => { + const store = new MemorySessionStore(100, { maxSessionsPerPrincipal: 1 }) + await store.create(newSession('a', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + + // A different subject is a different bucket and is not blocked + await store.create(newSession('b', { ownerSub: 'bob', ownerIss: 'https://auth.example.com' })) + t.assert.strictEqual((await store.get('b'))?.id, 'b') + + // A different issuer for the same subject is also a different bucket + await store.create(newSession('c', { ownerSub: 'alice', ownerIss: 'https://other.example.com' })) + t.assert.strictEqual((await store.get('c'))?.id, 'c') + + // Same (issuer, subject) again is blocked + await t.assert.rejects( + store.create(newSession('d', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })), + SessionLimitError + ) + }) + + test('anonymous sessions share one bucket', async (t: TestContext) => { + const store = new MemorySessionStore(100, { maxSessionsPerPrincipal: 1 }) + await store.create(newSession('a')) + + await t.assert.rejects( + store.create(newSession('b')), + SessionLimitError + ) + }) + + test('expired sessions do not count against the cap', async (t: TestContext) => { + const store = new MemorySessionStore(100, { sessionTtlMs: 1000, maxSessionsPerPrincipal: 1 }) + await store.create(newSession('a', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + + // Age the session past its TTL, then let the store observe it. + const aged = (await store.get('a'))! + aged.lastActivity = new Date(Date.now() - 5000) + await store.update(aged) + + // The expired session no longer counts, so a new one fits + await store.create(newSession('b', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + t.assert.strictEqual((await store.get('b'))?.id, 'b') + }) + + test('rejects invalid limit values', (t: TestContext) => { + t.assert.throws(() => new MemorySessionStore(100, { maxSessionsPerPrincipal: 0 }), TypeError) + t.assert.throws(() => new MemorySessionStore(100, { maxResourceSubscriptionsPerSession: 1.5 }), TypeError) + }) +}) + +describe('MemorySessionStore subscription limit', () => { + test('rejects a new URI when the session is at the cap, keeping duplicates idempotent', async (t: TestContext) => { + const store = new MemorySessionStore(100, { maxResourceSubscriptionsPerSession: 2 }) + await store.create(newSession('s1')) + + await store.addResourceSubscription('s1', 'file:///a.txt') + await store.addResourceSubscription('s1', 'file:///b.txt') + + // Duplicate is a no-op, not an error + await store.addResourceSubscription('s1', 'file:///a.txt') + + // A third distinct URI is rejected + await t.assert.rejects( + store.addResourceSubscription('s1', 'file:///c.txt'), + (err: unknown) => err instanceof SubscriptionLimitError && err.limit === 2 + ) + + // The stored subscriptions still only hold the two accepted URIs + t.assert.deepStrictEqual((await store.get('s1'))?.resourceSubscriptions, ['file:///a.txt', 'file:///b.txt']) + }) + + test('a removed URI frees a slot', async (t: TestContext) => { + const store = new MemorySessionStore(100, { maxResourceSubscriptionsPerSession: 1 }) + await store.create(newSession('s1')) + await store.addResourceSubscription('s1', 'file:///a.txt') + await store.removeResourceSubscription('s1', 'file:///a.txt') + await store.addResourceSubscription('s1', 'file:///b.txt') + t.assert.deepStrictEqual((await store.get('s1'))?.resourceSubscriptions, ['file:///b.txt']) + }) +}) + +describe('RedisSessionStore session limit', () => { + testWithRedis('rejects a new session when the owner is at the cap', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxSessionsPerPrincipal: 2 }) + await store.create(newSession('a', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + await store.create(newSession('b', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + + await t.assert.rejects( + store.create(newSession('c', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })), + (err: unknown) => err instanceof SessionLimitError && err.limit === 2 + ) + }) + + testWithRedis('is isolated between principals', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxSessionsPerPrincipal: 1 }) + await store.create(newSession('a', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + + await store.create(newSession('b', { ownerSub: 'bob', ownerIss: 'https://auth.example.com' })) + await store.create(newSession('c', { ownerSub: 'alice', ownerIss: 'https://other.example.com' })) + + await t.assert.rejects( + store.create(newSession('d', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })), + SessionLimitError + ) + }) + + testWithRedis('concurrent creates cannot exceed the cap', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxSessionsPerPrincipal: 3 }) + const attempts = Array.from({ length: 8 }, (_, i) => + store.create(newSession(`c${i}`, { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + ) + const results = await Promise.allSettled(attempts) + + const succeeded = results.filter(r => r.status === 'fulfilled') + const failed = results.filter(r => r.status === 'rejected') + // The Lua script is atomic, so exactly `cap` sessions survive + t.assert.strictEqual(succeeded.length, 3) + t.assert.strictEqual(failed.length, 5) + for (const r of failed) { + t.assert.ok(r.reason instanceof SessionLimitError) + } + }) + + testWithRedis('deleting a session frees a slot for its owner', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxSessionsPerPrincipal: 1 }) + await store.create(newSession('a', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + await store.delete('a') + await store.create(newSession('b', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + t.assert.strictEqual((await store.get('b'))?.id, 'b') + }) + + testWithRedis('an expired session is pruned from the owner index', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxSessionsPerPrincipal: 1, sessionTtlMs: 5000 }) + await store.create(newSession('a', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + + // Simulate expiry by removing the session key but leaving the owner index + await redis.del('session:a') + + // The stale id must be pruned before the cap check, so a new session fits + await store.create(newSession('b', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + t.assert.strictEqual((await store.get('b'))?.id, 'b') + }) + + testWithRedis('anonymous sessions share one bucket, separate from owned sessions', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxSessionsPerPrincipal: 1 }) + await store.create(newSession('anon')) + + // An owned session is a different bucket and is not blocked + await store.create(newSession('owned', { ownerSub: 'alice', ownerIss: 'https://auth.example.com' })) + + // A second anonymous session is blocked + await t.assert.rejects(store.create(newSession('anon2')), SessionLimitError) + }) +}) + +describe('RedisSessionStore subscription limit', () => { + testWithRedis('rejects a new URI at the cap, keeping duplicates idempotent', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxResourceSubscriptionsPerSession: 2 }) + await store.create(newSession('s1')) + + await store.addResourceSubscription('s1', 'file:///a.txt') + await store.addResourceSubscription('s1', 'file:///b.txt') + await store.addResourceSubscription('s1', 'file:///a.txt') // duplicate, idempotent + + await t.assert.rejects( + store.addResourceSubscription('s1', 'file:///c.txt'), + (err: unknown) => err instanceof SubscriptionLimitError && err.limit === 2 + ) + + t.assert.deepStrictEqual( + (await store.get('s1'))?.resourceSubscriptions?.sort(), + ['file:///a.txt', 'file:///b.txt'] + ) + }) + + testWithRedis('concurrent subscribes cannot exceed the cap', async (redis, t: TestContext) => { + const store = new RedisSessionStore({ redis, maxResourceSubscriptionsPerSession: 3 }) + await store.create(newSession('s1')) + + const attempts = Array.from({ length: 8 }, (_, i) => + store.addResourceSubscription('s1', `file:///c${i}.txt`) + ) + const results = await Promise.allSettled(attempts) + const succeeded = results.filter(r => r.status === 'fulfilled') + const failed = results.filter(r => r.status === 'rejected') + t.assert.strictEqual(succeeded.length, 3) + t.assert.strictEqual(failed.length, 5) + for (const r of failed) { + t.assert.ok(r.reason instanceof SubscriptionLimitError) + } + }) +}) + +// Integration: the limits surface as stable protocol errors over HTTP +async function createApp (t: TestContext, pluginOptions: Record = {}): Promise { + const app = fastify() + t.after(() => app.close()) + await app.register(mcpPlugin, { + serverInfo: { name: 'test', version: '1.0.0' }, + capabilities: { resources: { subscribe: true } }, + enableSSE: true, + ...pluginOptions + }) + app.mcpSetResourceSubscribeHandler(async () => ({})) + app.mcpSetResourceUnsubscribeHandler(async () => ({})) + return app +} + +async function initialize (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: '2025-06-18', + capabilities: {}, + clientInfo: { name: 'test-client', version: '1.0.0' } + } + } + }) + if (response.statusCode !== 200) { + throw new Error(`initialize failed: ${response.statusCode} ${response.body}`) + } + return response.headers['mcp-session-id'] as string +} + +describe('session and subscription limits over HTTP', () => { + test('creating a session beyond the per-principal cap fails with a protocol error', async (t: TestContext) => { + const app = await createApp(t, { maxSessionsPerPrincipal: 1 }) + await initialize(app) + + const response = await app.inject({ + method: 'POST', + url: '/mcp', + headers: { 'content-type': 'application/json', accept: 'application/json' }, + payload: { + jsonrpc: JSONRPC_VERSION, + id: 7, + method: 'initialize', + params: { + protocolVersion: '2025-06-18', + capabilities: {}, + clientInfo: { name: 'test-client', version: '1.0.0' } + } + } + }) + + t.assert.strictEqual(response.statusCode, 429) + const body = response.json() + t.assert.strictEqual(body.error.code, SESSION_LIMIT_REACHED) + t.assert.strictEqual(body.id, 7) + }) + + test('opening an SSE stream beyond the per-principal cap fails without hanging', async (t: TestContext) => { + const app = await createApp(t, { maxSessionsPerPrincipal: 1 }) + await initialize(app) + + const response = await app.inject({ + method: 'GET', + url: '/mcp', + headers: { accept: 'text/event-stream' }, + payloadAsStream: true + }) + + t.assert.strictEqual(response.statusCode, 429) + t.assert.match(response.headers['content-type'] ?? '', /application\/json/) + response.stream().destroy() + }) + + test('subscribing beyond the per-session cap fails with a protocol error', async (t: TestContext) => { + const app = await createApp(t, { maxResourceSubscriptionsPerSession: 1 }) + const sessionId = await initialize(app) + + async function subscribe (uri: string) { + return app.inject({ + method: 'POST', + url: '/mcp', + headers: { + 'content-type': 'application/json', + accept: 'application/json', + 'mcp-session-id': sessionId + }, + payload: { jsonrpc: JSONRPC_VERSION, id: 2, method: 'resources/subscribe', params: { uri } } + }) + } + + const first = await subscribe('file:///a.txt') + t.assert.strictEqual(first.statusCode, 200) + t.assert.strictEqual(first.json().error, undefined) + + const second = await subscribe('file:///b.txt') + t.assert.strictEqual(second.statusCode, 200) + t.assert.strictEqual(second.json().error.code, SUBSCRIPTION_LIMIT_REACHED) + + // A duplicate of the subscribed URI stays idempotent and succeeds + const dup = await subscribe('file:///a.txt') + t.assert.strictEqual(dup.statusCode, 200) + t.assert.strictEqual(dup.json().error, undefined) + }) +})