@cyanheads/jinaai-mcp-server
Version:
A Model Context Protocol (MCP) server that provides intelligent web reading capabilities using the Jina AI Reader API. It extracts clean, LLM-ready content from any URL.
298 lines • 13.3 kB
JavaScript
/**
* @fileoverview Implements a stateful transport manager for the MCP SDK.
*
* This manager handles multiple, persistent MCP sessions. It creates and maintains
* a dedicated McpServer and StreamableHTTPServerTransport instance for each session,
* allowing for stateful, multi-turn interactions. It includes robust mechanisms for
* session lifecycle management, including garbage collection of stale sessions and
* concurrency controls to prevent race conditions.
*
* SCALABILITY NOTE: This manager maintains all session state in local process memory.
* For horizontal scaling across multiple server instances, a load balancer with
* sticky sessions (session affinity) is required to ensure that all requests for a
* given session are routed to the same process instance that holds that session's state.
*
* @module src/mcp-server/transports/core/statefulTransportManager
*/
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { randomUUID } from "node:crypto";
import { Readable } from "stream";
import { BaseErrorCode, McpError } from "../../../types-global/errors.js";
import { ErrorHandler, logger, requestContextService, } from "../../../utils/index.js";
import { BaseTransportManager } from "./baseTransportManager.js";
import { HonoStreamResponse } from "./honoNodeBridge.js";
import { convertNodeHeadersToWebHeaders } from "./headerUtils.js";
/**
* Manages persistent, stateful MCP sessions.
*/
export class StatefulTransportManager extends BaseTransportManager {
transports = new Map();
servers = new Map();
sessions = new Map();
garbageCollector;
options;
/**
* @param createServerInstanceFn - A factory function to create new McpServer instances.
* @param options - Configuration options for the manager.
*/
constructor(createServerInstanceFn, options) {
super(createServerInstanceFn);
this.options = options;
const context = requestContextService.createRequestContext({
operation: "StatefulTransportManager.constructor",
});
logger.info("Starting session garbage collector.", context);
this.garbageCollector = setInterval(() => this.cleanupStaleSessions(), this.options.staleSessionTimeoutMs);
}
/**
* Initializes a new stateful session and handles the first request.
*
* @param headers - The incoming request headers.
* @param body - The parsed body of the request.
* @param context - The request context.
* @returns A promise resolving to a streaming TransportResponse with a session ID.
*/
async initializeAndHandle(headers, body, context) {
const opContext = {
...context,
operation: "StatefulTransportManager.initializeAndHandle",
};
logger.debug("Initializing new stateful session.", opContext);
let server;
let transport;
try {
server = await this.createServerInstanceFn();
const mockRes = new HonoStreamResponse();
const currentServer = server;
transport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => randomUUID(),
onsessioninitialized: (sessionId) => {
const sessionContext = { ...opContext, sessionId };
this.transports.set(sessionId, transport);
this.servers.set(sessionId, currentServer);
this.sessions.set(sessionId, {
id: sessionId,
createdAt: new Date(),
lastAccessedAt: new Date(),
activeRequests: 0,
});
logger.info(`MCP Session created: ${sessionId}`, sessionContext);
},
});
transport.onclose = () => {
const sessionId = transport.sessionId;
if (sessionId) {
const closeContext = { ...opContext, sessionId };
this.closeSession(sessionId, closeContext).catch((err) => logger.error(`Error during transport.onclose cleanup for session ${sessionId}`, err, closeContext));
}
};
await server.connect(transport);
logger.debug("Server connected, handling initial request.", opContext);
const mockReq = {
headers,
method: "POST",
url: this.options.mcpHttpEndpointPath,
};
await transport.handleRequest(mockReq, mockRes, body);
const responseHeaders = convertNodeHeadersToWebHeaders(mockRes.getHeaders());
if (transport.sessionId) {
responseHeaders.set("Mcp-Session-Id", transport.sessionId);
}
const webStream = Readable.toWeb(mockRes);
return {
type: "stream",
headers: responseHeaders,
statusCode: mockRes.statusCode,
stream: webStream,
sessionId: transport.sessionId,
};
}
catch (error) {
logger.error("Failed to initialize stateful session. Cleaning up orphaned resources.", error instanceof Error ? error : undefined, { ...opContext, error: String(error) });
const sessionInitialized = transport?.sessionId && this.transports.has(transport.sessionId);
if (!sessionInitialized) {
(async () => {
await ErrorHandler.tryCatch(async () => {
if (transport)
await transport.close();
if (server)
await server.close();
}, {
operation: "initializeAndHandle.cleanupOrphaned",
context: opContext,
});
})();
}
throw ErrorHandler.handleError(error, {
operation: opContext.operation,
context: opContext,
rethrow: true,
});
}
}
/**
* Handles a subsequent request for an existing stateful session.
*/
async handleRequest(headers, body, context, sessionId) {
if (!sessionId) {
throw new McpError(BaseErrorCode.INVALID_INPUT, "Session ID is required for stateful requests.", context);
}
const sessionContext = {
...context,
sessionId,
operation: "StatefulTransportManager.handleRequest",
};
const transport = this.transports.get(sessionId);
const session = this.sessions.get(sessionId);
if (!transport || !session) {
logger.warning(`Request for non-existent session: ${sessionId}`, sessionContext);
return {
type: "buffered",
headers: new Headers({ "Content-Type": "application/json" }),
statusCode: 404,
body: {
jsonrpc: "2.0",
error: { code: -32601, message: "Session not found" },
},
};
}
session.lastAccessedAt = new Date();
session.activeRequests += 1;
logger.debug(`Incremented activeRequests for session ${sessionId}. Count: ${session.activeRequests}`, sessionContext);
try {
const mockReq = {
headers,
method: "POST",
url: this.options.mcpHttpEndpointPath,
};
const mockRes = new HonoStreamResponse();
await transport.handleRequest(mockReq, mockRes, body);
const responseHeaders = convertNodeHeadersToWebHeaders(mockRes.getHeaders());
const webStream = Readable.toWeb(mockRes);
return {
type: "stream",
headers: responseHeaders,
statusCode: mockRes.statusCode,
stream: webStream,
sessionId: transport.sessionId,
};
}
catch (error) {
throw ErrorHandler.handleError(error, {
operation: sessionContext.operation,
context: sessionContext,
rethrow: true,
});
}
finally {
session.activeRequests -= 1;
session.lastAccessedAt = new Date();
logger.debug(`Decremented activeRequests for session ${sessionId}. Count: ${session.activeRequests}`, sessionContext);
}
}
/**
* Handles a request to explicitly delete a session.
*/
async handleDeleteRequest(sessionId, context) {
const sessionContext = {
...context,
sessionId,
operation: "StatefulTransportManager.handleDeleteRequest",
};
logger.info(`Attempting to delete session: ${sessionId}`, sessionContext);
if (!this.transports.has(sessionId)) {
logger.warning(`Attempted to delete non-existent session: ${sessionId}`, sessionContext);
throw new McpError(BaseErrorCode.NOT_FOUND, "Session not found or expired.", sessionContext);
}
await this.closeSession(sessionId, sessionContext);
return {
type: "buffered",
headers: new Headers({ "Content-Type": "application/json" }),
statusCode: 200,
body: { status: "session_closed", sessionId },
};
}
/**
* Retrieves information about a specific session.
*/
getSession(sessionId) {
return this.sessions.get(sessionId);
}
/**
* Gracefully shuts down the manager, closing all active sessions.
*/
async shutdown() {
const context = requestContextService.createRequestContext({
operation: "StatefulTransportManager.shutdown",
});
logger.info("Shutting down stateful transport manager...", context);
clearInterval(this.garbageCollector);
logger.debug("Garbage collector stopped.", context);
const sessionIds = Array.from(this.transports.keys());
if (sessionIds.length > 0) {
logger.info(`Closing ${sessionIds.length} active sessions.`, context);
const closePromises = sessionIds.map((sessionId) => this.closeSession(sessionId, context));
await Promise.all(closePromises);
}
this.transports.clear();
this.sessions.clear();
this.servers.clear();
logger.info("All active sessions closed and manager shut down.", context);
}
/**
* Closes a single session and releases its associated resources.
*/
async closeSession(sessionId, context) {
const sessionContext = {
...context,
sessionId,
operation: "StatefulTransportManager.closeSession",
};
logger.debug(`Closing session: ${sessionId}`, sessionContext);
const transport = this.transports.get(sessionId);
const server = this.servers.get(sessionId);
await ErrorHandler.tryCatch(async () => {
if (transport)
await transport.close();
if (server)
await server.close();
}, { operation: "closeSession.cleanup", context: sessionContext });
this.transports.delete(sessionId);
this.servers.delete(sessionId);
this.sessions.delete(sessionId);
logger.info(`MCP Session closed and resources released: ${sessionId}`, sessionContext);
}
/**
* Periodically runs to find and clean up stale, inactive sessions.
*/
async cleanupStaleSessions() {
const context = requestContextService.createRequestContext({
operation: "StatefulTransportManager.cleanupStaleSessions",
});
logger.debug("Running stale session cleanup...", context);
const now = Date.now();
const STALE_TIMEOUT_MS = this.options.staleSessionTimeoutMs;
const staleSessionIds = [];
for (const [sessionId, session] of this.sessions.entries()) {
if (now - session.lastAccessedAt.getTime() > STALE_TIMEOUT_MS) {
if (session.activeRequests > 0) {
logger.info(`Session ${sessionId} is stale but has ${session.activeRequests} active requests. Skipping cleanup.`, { ...context, sessionId });
continue;
}
staleSessionIds.push(sessionId);
}
}
if (staleSessionIds.length > 0) {
logger.info(`Found ${staleSessionIds.length} stale sessions. Closing concurrently.`, context);
const closePromises = staleSessionIds.map((sessionId) => this.closeSession(sessionId, context).catch((err) => {
logger.error(`Error during concurrent stale session cleanup for ${sessionId}`, err, { ...context, sessionId });
}));
await Promise.all(closePromises);
logger.info(`Stale session cleanup complete. Closed ${staleSessionIds.length} sessions.`, context);
}
else {
logger.debug("No stale sessions found.", context);
}
}
}
//# sourceMappingURL=statefulTransportManager.js.map