f4aed937f0- Add structured logger utility with Cloud Run trace context support - Uses AsyncLocalStorage for request context propagation - Replace all console.log/error/warn statements with structured logger - Add logging middleware to capture request/response metadata - Support Google Cloud Logging severity levels and trace correlation 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
94 lines
3.2 KiB
TypeScript
94 lines
3.2 KiB
TypeScript
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
|
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
|
|
import { Request, Response } from "express";
|
|
import { getShttpTransport, isSessionOwnedBy, redisRelayToMcpServer, ServerRedisTransport, setSessionOwner, shutdownSession } from "../services/redisTransport.js";
|
|
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
|
|
import { randomUUID } from "crypto";
|
|
import { createMcpServer } from "../services/mcp.js";
|
|
import { logger } from "../utils/logger.js";
|
|
|
|
|
|
declare module "express-serve-static-core" {
|
|
interface Request {
|
|
/**
|
|
* Information about the validated access token, if the `requireBearerAuth` middleware was used.
|
|
*/
|
|
auth?: AuthInfo;
|
|
}
|
|
}
|
|
|
|
function getUserIdFromAuth(auth?: AuthInfo): string | null {
|
|
return auth?.extra?.userId as string || null;
|
|
}
|
|
|
|
export async function handleStreamableHTTP(req: Request, res: Response) {
|
|
let shttpTransport: StreamableHTTPServerTransport | undefined = undefined;
|
|
|
|
res.on('finish', async () => {
|
|
await shttpTransport?.close();
|
|
});
|
|
|
|
const onsessionclosed = async (sessionId: string) => {
|
|
await shutdownSession(sessionId);
|
|
}
|
|
|
|
try {
|
|
// Check for existing session ID
|
|
const sessionId = req.headers['mcp-session-id'] as string | undefined;
|
|
const userId = getUserIdFromAuth(req.auth);
|
|
|
|
// if no userid, return 401, we shouldn't get here ideally
|
|
if (!userId) {
|
|
res.status(401)
|
|
return;
|
|
}
|
|
|
|
const isGetRequest = req.method === 'GET';
|
|
|
|
// incorrect session for the authed user, return 401
|
|
if (sessionId) {
|
|
if (!(await isSessionOwnedBy(sessionId, userId))) {
|
|
res.status(401)
|
|
return;
|
|
}
|
|
// Reuse existing transport for owned session
|
|
shttpTransport = await getShttpTransport(sessionId, onsessionclosed, isGetRequest);
|
|
} else if (isInitializeRequest(req.body)) {
|
|
// New initialization request - use JSON response mode
|
|
const onsessioninitialized = async (sessionId: string) => {
|
|
const { server, cleanup: mcpCleanup } = createMcpServer();
|
|
|
|
const serverRedisTransport = new ServerRedisTransport(sessionId);
|
|
serverRedisTransport.onclose = mcpCleanup;
|
|
await server.connect(serverRedisTransport)
|
|
|
|
// Set session ownership
|
|
await setSessionOwner(sessionId, userId);
|
|
}
|
|
|
|
const sessionId = randomUUID();
|
|
shttpTransport = new StreamableHTTPServerTransport({
|
|
sessionIdGenerator: () => sessionId,
|
|
onsessionclosed,
|
|
onsessioninitialized,
|
|
});
|
|
shttpTransport.onclose = await redisRelayToMcpServer(sessionId, shttpTransport);
|
|
} else {
|
|
// Invalid request - no session ID and not initialization request
|
|
res.status(400)
|
|
return;
|
|
}
|
|
// Handle the request with existing transport - no need to reconnect
|
|
await shttpTransport.handleRequest(req, res, req.body);
|
|
} catch (error) {
|
|
logger.error('Error handling MCP request', error as Error, {
|
|
sessionId: req.headers['mcp-session-id'] as string | undefined,
|
|
method: req.method,
|
|
userId: getUserIdFromAuth(req.auth)
|
|
});
|
|
|
|
if (!res.headersSent) {
|
|
res.status(500)
|
|
}
|
|
}
|
|
}
|