Skip to content

Commit c251e40

Browse files
ochafikclaude
ochafik
and
claude
committed
fix(examples): improve server-utils with stateful sessions and security
Address critical issues in the shared server utility: 1. Stateful sessions: Use sessionIdGenerator + onsessioninitialized to persist StreamableHTTPServerTransport across requests 2. Unified session store: Single Map for both transport types with proper type discrimination 3. Error handling: Try/catch on all endpoints with JSON-RPC errors 4. DNS rebinding protection: Use SDK's createMcpExpressApp helper 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 6d3ff41 commit c251e40

File tree

1 file changed

+113
-70
lines changed

1 file changed

+113
-70
lines changed

‎examples/shared/server-utils.ts‎

Lines changed: 113 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,36 @@
11
/**
22
* Shared utilities for running MCP servers with multiple transports.
33
*
4-
* This module provides a unified way to start MCP servers supporting:
5-
* - stdio transport (for local CLI tools)
6-
* - Streamable HTTP transport (current spec)
7-
* - Legacy SSE transport (deprecated, for backwards compatibility)
4+
* Supports:
5+
* - stdio transport (--stdio flag)
6+
* - Streamable HTTP transport (/mcp) - stateful sessions
7+
* - Legacy SSE transport (/sse, /messages) - backwards compatibility
88
*/
99

1010
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
11+
import { createMcpExpressApp } from "@modelcontextprotocol/sdk/server/express.js";
1112
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
1213
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
1314
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
14-
import cors from "cors";
15-
import express, { type Request, type Response } from "express";
15+
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
16+
import { randomUUID } from "node:crypto";
17+
import type { Request, Response } from "express";
1618

1719
export interface ServerOptions {
18-
/** Port to listen on for HTTP mode. Defaults to 3001 or PORT env variable. */
20+
/** Port to listen on. Defaults to PORT env var or 3001. */
1921
port?: number;
20-
/** Server name for logging. Defaults to "MCP Server". */
22+
/** Server name for logging. */
2123
name?: string;
2224
}
2325

26+
type Transport = StreamableHTTPServerTransport | SSEServerTransport;
27+
2428
/**
25-
* Starts an MCP server with support for stdio and HTTP transports.
26-
*
27-
* Transport is selected based on command line arguments:
28-
* - `--stdio`: Uses stdio transport for local process communication
29-
* - Otherwise: Starts HTTP server with Streamable HTTP and legacy SSE support
29+
* Starts an MCP server with stdio and HTTP transports.
3030
*
31-
* @param server - The MCP server instance to start
32-
* @param options - Optional configuration
31+
* HTTP mode provides:
32+
* - /mcp (GET/POST/DELETE): Streamable HTTP with stateful sessions
33+
* - /sse (GET) + /messages (POST): Legacy SSE for older clients
3334
*/
3435
export async function startServer(
3536
server: McpServer,
@@ -40,74 +41,116 @@ export async function startServer(
4041
const name = options.name ?? "MCP Server";
4142

4243
if (process.argv.includes("--stdio")) {
43-
const transport = new StdioServerTransport();
44-
await server.connect(transport);
44+
await server.connect(new StdioServerTransport());
4545
console.error(`${name} running in stdio mode`);
46-
} else {
47-
const app = express();
48-
app.use(cors());
49-
app.use(express.json());
50-
51-
// Streamable HTTP transport (current spec) - handles GET, POST, DELETE
52-
app.all("/mcp", async (req: Request, res: Response) => {
53-
try {
54-
const transport = new StreamableHTTPServerTransport({
55-
sessionIdGenerator: undefined,
56-
enableJsonResponse: true,
57-
});
58-
res.on("close", () => {
59-
transport.close();
46+
return;
47+
}
48+
49+
// Unified session store for both transport types
50+
const sessions = new Map<string, Transport>();
51+
52+
// Express with DNS rebinding protection
53+
const app = createMcpExpressApp();
54+
55+
// Streamable HTTP (stateful)
56+
app.all("/mcp", async (req: Request, res: Response) => {
57+
try {
58+
const sessionId = req.headers["mcp-session-id"] as string | undefined;
59+
let transport = sessionId
60+
? (sessions.get(sessionId) as StreamableHTTPServerTransport | undefined)
61+
: undefined;
62+
63+
// Session exists but wrong transport type
64+
if (sessionId && sessions.has(sessionId) && !transport) {
65+
return res.status(400).json({
66+
jsonrpc: "2.0",
67+
error: { code: -32000, message: "Session uses different transport" },
68+
id: null,
6069
});
70+
}
6171

62-
await server.connect(transport);
63-
await transport.handleRequest(req, res, req.body);
64-
} catch (error) {
65-
console.error("Error handling MCP request:", error);
66-
if (!res.headersSent) {
67-
res.status(500).json({
72+
// New session requires initialize request
73+
if (!transport) {
74+
if (req.method !== "POST" || !isInitializeRequest(req.body)) {
75+
return res.status(400).json({
6876
jsonrpc: "2.0",
69-
error: { code: -32603, message: "Internal server error" },
77+
error: { code: -32000, message: "Bad request: not initialized" },
7078
id: null,
7179
});
7280
}
81+
82+
transport = new StreamableHTTPServerTransport({
83+
sessionIdGenerator: () => randomUUID(),
84+
onsessioninitialized: (id) => {
85+
sessions.set(id, transport!);
86+
},
87+
});
88+
const t = transport;
89+
t.onclose = () => {
90+
if (t.sessionId) sessions.delete(t.sessionId);
91+
};
92+
await server.connect(transport);
7393
}
74-
});
7594

76-
// Legacy SSE transport (deprecated) - for backwards compatibility
77-
const sseTransports = new Map<string, SSEServerTransport>();
95+
await transport.handleRequest(req, res, req.body);
96+
} catch (error) {
97+
console.error("MCP error:", error);
98+
if (!res.headersSent) {
99+
res.status(500).json({
100+
jsonrpc: "2.0",
101+
error: { code: -32603, message: "Internal server error" },
102+
id: null,
103+
});
104+
}
105+
}
106+
});
78107

79-
app.get("/sse", async (_req: Request, res: Response) => {
108+
// Legacy SSE
109+
app.get("/sse", async (_req: Request, res: Response) => {
110+
try {
80111
const transport = new SSEServerTransport("/messages", res);
81-
sseTransports.set(transport.sessionId, transport);
82-
res.on("close", () => {
83-
sseTransports.delete(transport.sessionId);
84-
});
112+
sessions.set(transport.sessionId, transport);
113+
res.on("close", () => sessions.delete(transport.sessionId));
85114
await server.connect(transport);
86-
});
115+
} catch (error) {
116+
console.error("SSE error:", error);
117+
if (!res.headersSent) res.status(500).end();
118+
}
119+
});
87120

88-
app.post("/messages", async (req: Request, res: Response) => {
89-
const sessionId = req.query.sessionId as string;
90-
const transport = sseTransports.get(sessionId);
91-
if (!transport) {
92-
res.status(404).json({ error: "Session not found" });
93-
return;
121+
app.post("/messages", async (req: Request, res: Response) => {
122+
try {
123+
const transport = sessions.get(req.query.sessionId as string);
124+
if (!(transport instanceof SSEServerTransport)) {
125+
return res.status(404).json({
126+
jsonrpc: "2.0",
127+
error: { code: -32001, message: "Session not found" },
128+
id: null,
129+
});
94130
}
95131
await transport.handlePostMessage(req, res, req.body);
96-
});
97-
98-
const httpServer = app.listen(port, () => {
99-
console.log(`${name} listening on http://localhost:${port}/mcp`);
100-
});
101-
102-
const shutdown = () => {
103-
console.log("\nShutting down...");
104-
httpServer.close(() => {
105-
console.log("Server closed");
106-
process.exit(0);
107-
});
108-
};
109-
110-
process.on("SIGINT", shutdown);
111-
process.on("SIGTERM", shutdown);
112-
}
132+
} catch (error) {
133+
console.error("Message error:", error);
134+
if (!res.headersSent) {
135+
res.status(500).json({
136+
jsonrpc: "2.0",
137+
error: { code: -32603, message: "Internal server error" },
138+
id: null,
139+
});
140+
}
141+
}
142+
});
143+
144+
const httpServer = app.listen(port, () => {
145+
console.log(`${name} listening on http://localhost:${port}/mcp`);
146+
});
147+
148+
const shutdown = () => {
149+
console.log("\nShutting down...");
150+
sessions.forEach((t) => t.close().catch(() => {}));
151+
httpServer.close(() => process.exit(0));
152+
};
153+
154+
process.on("SIGINT", shutdown);
155+
process.on("SIGTERM", shutdown);
113156
}

0 commit comments

Comments
 (0)