modelcontextprotocol/example-remote-server
Watch
1
0
Fork
You've already forked example-remote-server
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

WIP shttp changes

Browse source
This commit is contained in:
Jerome 2025-07-01 15:10:55 +01:00
commit c2119ed9f9
No known key found for this signature in database
4 changed files with 88 additions and 58 deletions
Download patch file Download diff file Expand all files Collapse all files

46
src/handlers/shttp.ts
View file

@ -1,11 +1,10 @@
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { Request, Response } from "express";
import { getFirstShttpTransport, getShttpTransport, isLive, startServerListeningToRedis } from "../services/redisTransport.js";
import { getFirstShttpTransport, getShttpTransport, isLive, RedisTransport, startServerListeningToRedis } from "../services/redisTransport.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
import { randomUUID } from "crypto";
import { createMcpServer } from "../services/mcp.js";
import getRawBody from "raw-body";
declare module "express-serve-static-core" {
@ -22,7 +21,8 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
console.log('Received MCP request:', JSON.stringify(req.body, null, 2));
console.log('Request headers:', JSON.stringify(req.headers, null, 2));
let transport: StreamableHTTPServerTransport | undefined = undefined;
let shttpTransport: StreamableHTTPServerTransport | undefined = undefined;
let redisTransport: RedisTransport | undefined = undefined;
try {
// Check for existing session ID
const sessionId = req.headers['mcp-session-id'] as string | undefined;
@ -31,7 +31,11 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
if (sessionId && await isLive(sessionId)) {
console.log('Session is live, reusing existing transport for session:', sessionId);
// Reuse existing transport
transport = await getShttpTransport(sessionId)
// const { shttpTransport, redisTransport } = await getShttpTransport(sessionId);
// destructuring to get both transports
({ shttpTransport, redisTransport } = await getShttpTransport(sessionId));
console.log('Created transport for session:', sessionId);
console.log('Retrieved transport from Redis for session:', sessionId);
} else if (!sessionId && isInitializeRequest(req.body)) {
console.log('New initialization request detected, creating new session');
@ -45,20 +49,9 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
startServerListeningToRedis(server.server, sessionId)
console.log('Started server listening to Redis for session:', sessionId);
transport = await getFirstShttpTransport(sessionId);
({ shttpTransport, redisTransport } = await getFirstShttpTransport(sessionId));
console.log('Retrieved first transport for session:', sessionId);
console.log('Transport object:', transport.constructor.name, 'sessionId:', transport.sessionId);
// Connect the transport to the MCP server BEFORE handling the request
console.log('Connecting transport to MCP server...');
await server.server.connect(transport);
console.log('Transport connected successfully');
console.log('Handling initialization request...');
await transport.handleRequest(req, res, req.body);
console.log('Initialization request handled successfully');
console.log('=== handleStreamableHTTP END (initialization) ===');
return; // Already handled
console.log('Transport object:', shttpTransport.constructor.name, 'sessionId:', shttpTransport.sessionId);
} else {
console.log('Invalid request - no session ID and not an initialization request');
console.log('Session ID present:', !!sessionId);
@ -79,9 +72,9 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
// Handle the request with existing transport - no need to reconnect
console.log('Handling request with existing transport for session:', sessionId);
console.log('Transport object:', transport.constructor.name, 'sessionId:', transport.sessionId);
console.log('Request body:', JSON.stringify(req.body, null, 2));
await transport.handleRequest(req, res, req.body);
console.log('Transport object:', shttpTransport.constructor.name, 'sessionId:', shttpTransport.sessionId);
console.log('Request body:', JSON.stringify(req.body, null, 2));
await shttpTransport.handleRequest(req, res, req.body);
console.log('Request handled successfully');
} catch (error) {
console.error('=== ERROR in handleStreamableHTTP ===');
@ -104,12 +97,13 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
console.log('Response headers already sent, cannot send error response');
}
} finally {
// if (transport) {
// console.log('Closing transport in finally block');
// // Close transports because they are ephemeral in this setup.
// transport.close();
// console.log('Transport closed');
// }
// Set up cleanup when response is complete
res.on('finish', async () => {
console.log('HTTP response finished, closing transport');
await shttpTransport?.close();
// await redisTransport?.close();
console.log('Transport closed after response');
});
console.log('=== handleStreamableHTTP END ===');
}
}

15
src/index.ts
View file

@ -1,4 +1,4 @@
import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js";
import { BearerAuthMiddlewareOptions, requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js";
import { AuthRouterOptions, mcpAuthRouter } from "@modelcontextprotocol/sdk/server/auth/router.js";
import cors from "cors";
import express from "express";
@ -9,6 +9,7 @@ import { handleFakeAuthorize, handleFakeAuthorizeRedirect } from "./handlers/fak
import { handleStreamableHTTP } from "./handlers/shttp.js";
import { handleMessage, handleSSEConnection } from "./handlers/sse.js";
import { redisClient } from "./redis.js";
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
const app = express();
@ -78,6 +79,8 @@ app.use(baseSecurityHeaders);
// Enable CORS pre-flight requests
app.options('*', cors(corsOptions));
const authProvider = new EverythingAuthProvider();
// Auth configuration
const options: AuthRouterOptions = {
provider: new EverythingAuthProvider(),
@ -89,8 +92,16 @@ const options: AuthRouterOptions = {
}
}
};
const dearerAuthMiddlewareOptions: BearerAuthMiddlewareOptions = {
// verifyAccessToken(token: string): Promise<AuthInfo>;
verifier: {
verifyAccessToken: authProvider.verifyAccessToken.bind(authProvider),
}
}
app.use(mcpAuthRouter(options));
const bearerAuth = requireBearerAuth(options);
const bearerAuth = requireBearerAuth(dearerAuthMiddlewareOptions);
// MCP routes (legacy SSE transport)
app.get("/sse", cors(corsOptions), bearerAuth, authContext, sseHeaders, handleSSEConnection);

19
src/services/mcp.ts
View file

@ -15,13 +15,16 @@ import {
SetLevelRequestSchema,
SubscribeRequestSchema,
Tool,
ToolSchema,
UnsubscribeRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import { zodToJsonSchema } from "zod-to-json-schema";
type ToolInput = z.infer<typeof ToolSchema.shape.inputSchema>;
type ToolInput = {
type: "object";
properties?: Record<string, unknown>;
required?: string[];
};
/* Input schemas for tools implemented in this server */
const EchoSchema = z.object({
@ -109,18 +112,6 @@ export const createMcpServer = () => {
}
);
// Add logging to track all incoming messages
const originalSetRequestHandler = server.setRequestHandler.bind(server);
server.setRequestHandler = function(schema: any, handler: any) {
const wrappedHandler = async (request: any) => {
console.log(`[MCP Server] Handling request: ${request.method}`);
const result = await handler(request);
console.log(`[MCP Server] Request ${request.method} completed`);
return result;
};
return originalSetRequestHandler(schema, wrappedHandler);
};
const subscriptions: Set<string> = new Set();
// Set up update interval for subscribed resources

66
src/services/redisTransport.ts
View file

@ -5,6 +5,7 @@ import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
let redisTransportCounter = 0;
function getToServerChannel(sessionId: string): string {
return `mcp:shttp:toserver:${sessionId}`;
@ -24,6 +25,7 @@ export async function isLive(sessionId: string): Promise<boolean> {
export class RedisTransport implements Transport {
private redisCleanup: (() => Promise<void>) | undefined;
private counter: number;
onclose?: (() => void) | undefined;
onerror?: ((error: Error) => void) | undefined;
@ -34,6 +36,7 @@ export class RedisTransport implements Transport {
private recvChannel: string,
private isLiveKey: string | undefined = undefined
) {
this.counter = redisTransportCounter++;
this.sendChannel = sendChannel;
this.recvChannel = recvChannel;
this.isLiveKey = isLiveKey;
@ -42,18 +45,26 @@ export class RedisTransport implements Transport {
async start(): Promise<void> {
console.log(`[RedisTransport.start] Starting transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`);
console.log(`[RedisTransport.${this.counter}.start] Starting transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`);
if (this.redisCleanup) {
throw new Error(`Redis transport already started for channels ${this.sendChannel} and ${this.recvChannel}`);
}
// Log when onmessage is set
console.log(`[RedisTransport.${this.counter}.start] onmessage handler is ${this.onmessage ? 'SET' : 'NOT SET'}`);
this.redisCleanup = await redisClient.createSubscription(
this.recvChannel,
(json) => {
console.log(`[RedisTransport] Received message on ${this.recvChannel}:`, json.substring(0, 100));
console.log(`[RedisTransport.${this.counter}] Received message on ${this.recvChannel}:`, json.substring(0, 100));
const message = JSON.parse(json);
const extra = popExtra(message);
this.onmessage?.(message, extra)
if (this.onmessage) {
console.log(`[RedisTransport.${this.counter}] Calling onmessage handler for ${this.recvChannel}`);
this.onmessage(message, extra)
} else {
console.log(`[RedisTransport.${this.counter}] WARNING: No onmessage handler for ${this.recvChannel}!`);
}
},
(error) => {
console.error(
@ -63,23 +74,27 @@ export class RedisTransport implements Transport {
this.close()
},
);
console.log(`[RedisTransport.start] Successfully subscribed to ${this.recvChannel}`);
console.log(`[RedisTransport.${this.counter}.start] Successfully subscribed to ${this.recvChannel}`);
}
send(message: JSONRPCMessage, options?: TransportSendOptions): Promise<void> {
async send(message: JSONRPCMessage, options?: TransportSendOptions): Promise<void> {
if (options) {
setOptions(message, options);
}
const messageStr = JSON.stringify(message);
console.log(`[RedisTransport.send] Publishing to ${this.sendChannel}:`, messageStr.substring(0, 100));
return redisClient.publish(this.sendChannel, messageStr);
console.log(`[RedisTransport.${this.counter}.send] Publishing to ${this.sendChannel}:`, messageStr.substring(0, 100));
console.log(`[RedisTransport.${this.counter}.send] Full message:`, messageStr);
await redisClient.publish(this.sendChannel, messageStr);
console.log(`[RedisTransport.${this.counter}.send] Published successfully to ${this.sendChannel}`);
}
async cleanup(): Promise<void> {
if (this.redisCleanup) {
console.log(`[RedisTransport.${this.counter}.cleanup] Unsubscribing from ${this.recvChannel}`);
await this.redisCleanup()
this.redisCleanup = undefined;
console.log(`[RedisTransport.${this.counter}.cleanup] Successfully unsubscribed from ${this.recvChannel}`);
}
if (this.isLiveKey) {
await redisClient.del(this.isLiveKey);
@ -87,9 +102,9 @@ export class RedisTransport implements Transport {
}
async close(): Promise<void> {
console.log(`[RedisTransport.close] Closing transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`);
console.log(`[RedisTransport.${this.counter}.close] Closing transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`);
this.onclose?.();
this.cleanup()
await this.cleanup()
}
}
@ -186,11 +201,20 @@ function relayTransports(transport1: Transport, transport2: Transport): void {
transport1.onerror?.(error);
};
// Prevent circular close calls
let closing = false;
transport1.onclose = () => {
transport2.close().catch(console.error);
if (!closing) {
closing = true;
transport2.close().catch(console.error);
}
};
transport2.onclose = () => {
transport1.close().catch(console.error);
if (!closing) {
closing = true;
transport1.close().catch(console.error);
}
};
}
@ -198,6 +222,15 @@ function relayTransports(transport1: Transport, transport2: Transport): void {
export async function startServerListeningToRedis(server: Server, sessionId: string) {
console.log(`[startServerListeningToRedis] Starting background server for session ${sessionId}`);
const serverRedisTransport = createBackgroundTaskSideRedisTransport(sessionId)
// Log all messages sent by the server
const originalSend = serverRedisTransport.send.bind(serverRedisTransport);
serverRedisTransport.send = async (message, options) => {
console.log(`[MCP Server -> Redis] Sending response:`, JSON.stringify(message).substring(0, 200));
return originalSend(message, options);
};
// The server.connect() will call start() on the transport
await server.connect(serverRedisTransport)
console.log(`[startServerListeningToRedis] Background server connected for session ${sessionId}`);
}
@ -216,7 +249,7 @@ function createBackgroundTaskSideRedisTransport(sessionId: string): RedisTranspo
);
}
export async function getFirstShttpTransport(sessionId: string): Promise<StreamableHTTPServerTransport> {
export async function getFirstShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, redisTransport: RedisTransport}> {
const transport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => sessionId,
enableJsonResponse: true, // Enable JSON response mode
@ -227,14 +260,15 @@ export async function getFirstShttpTransport(sessionId: string): Promise<Streama
// When shttpTransport closes, so does the redisTransport
relayTransports(redisTransport, transport);
await redisTransport.start()
return transport;
return { shttpTransport: transport, redisTransport };
}
export async function getShttpTransport(sessionId: string): Promise<StreamableHTTPServerTransport> {
export async function getShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, redisTransport: RedisTransport}> {
// Giving undefined here and setting the sessionId means the
// transport wont try to create a new session.
const shttpTransport = new StreamableHTTPServerTransport({
sessionIdGenerator: undefined
sessionIdGenerator: undefined,
enableJsonResponse: true, // Use JSON response mode for all requests
})
shttpTransport.sessionId = sessionId;
@ -243,5 +277,5 @@ export async function getShttpTransport(sessionId: string): Promise<StreamableHT
// When shttpTransport closes, so does the redisTransport
relayTransports(redisTransport, shttpTransport);
await redisTransport.start()
return shttpTransport;
return { shttpTransport, redisTransport };
}