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) + }) +})