stack/apps/gateway/src/mcp/mcp.controller.ts

import type { IncomingMessage, ServerResponse } from 'node:http';
import { Logger } from '@nestjs/common';
import { fromNodeHeaders } from 'better-auth/node';
import type { Auth } from '@mosaic/auth';
import type { NestFastifyApplication } from '@nestjs/platform-fastify';
import type { McpService } from './mcp.service.js';
import { AUTH } from '../auth/auth.tokens.js';

/**
 * Mounts the MCP streamable HTTP transport endpoint at /mcp on the Fastify instance.
 *
 * This follows the same low-level Fastify hook pattern used by the auth controller,
 * bypassing NestJS routing to directly delegate to the MCP SDK transport handlers.
 *
 * Endpoint: POST /mcp (and GET /mcp for SSE stream reconnect)
 * Auth: Requires a valid BetterAuth session (cookie or Authorization header).
 * Session: Stateful — each initialized client gets a session ID via Mcp-Session-Id header.
 */
export function mountMcpHandler(app: NestFastifyApplication, mcpService: McpService): void {
  const auth = app.get<Auth>(AUTH);
  const logger = new Logger('McpController');
  const fastify = app.getHttpAdapter().getInstance();

  fastify.addHook(
    'onRequest',
    (
      req: { raw: IncomingMessage; url: string; method: string },
      reply: { raw: ServerResponse; hijack: () => void },
      done: () => void,
    ) => {
      if (!req.url.startsWith('/mcp')) {
        done();
        return;
      }

      reply.hijack();

      handleMcpRequest(req, reply, auth, mcpService, logger).catch((err: unknown) => {
        logger.error(
          `MCP request handler error: ${err instanceof Error ? err.message : String(err)}`,
        );
        if (!reply.raw.headersSent) {
          reply.raw.writeHead(500, { 'Content-Type': 'application/json' });
        }
        if (!reply.raw.writableEnded) {
          reply.raw.end(JSON.stringify({ error: 'Internal server error' }));
        }
      });
    },
  );
}

async function handleMcpRequest(
  req: { raw: IncomingMessage; url: string; method: string },
  reply: { raw: ServerResponse; hijack: () => void },
  auth: Auth,
  mcpService: McpService,
  logger: Logger,
): Promise<void> {
  // ─── Authentication ─────────────────────────────────────────────────────
  const headers = fromNodeHeaders(req.raw.headers);
  const result = await auth.api.getSession({ headers });

  if (!result) {
    reply.raw.writeHead(401, { 'Content-Type': 'application/json' });
    reply.raw.end(JSON.stringify({ error: 'Unauthorized: valid session required' }));
    return;
  }

  const userId = result.user.id;

  // ─── Session routing ─────────────────────────────────────────────────────
  const sessionId = req.raw.headers['mcp-session-id'];

  if (typeof sessionId === 'string' && sessionId.length > 0) {
    // Existing session request
    const transport = mcpService.getSession(sessionId);
    if (!transport) {
      logger.warn(`MCP session not found: ${sessionId}`);
      reply.raw.writeHead(404, { 'Content-Type': 'application/json' });
      reply.raw.end(JSON.stringify({ error: 'Session not found' }));
      return;
    }

    await transport.handleRequest(req.raw, reply.raw);
    return;
  }

  // ─── Initialize new session ───────────────────────────────────────────────
  // Only POST requests can initialize a new session (must be initialize message)
  if (req.method !== 'POST') {
    reply.raw.writeHead(400, { 'Content-Type': 'application/json' });
    reply.raw.end(
      JSON.stringify({
        error: 'New session must be established via POST with initialize message',
      }),
    );
    return;
  }

  // Parse body to verify this is an initialize request before creating a session
  let body: unknown;
  try {
    body = await readRequestBody(req.raw);
  } catch (err) {
    logger.warn(
      `Failed to parse MCP request body: ${err instanceof Error ? err.message : String(err)}`,
    );
    reply.raw.writeHead(400, { 'Content-Type': 'application/json' });
    reply.raw.end(JSON.stringify({ error: 'Invalid request body' }));
    return;
  }

  // Create new session and handle this initializing request
  const { transport } = mcpService.createSession(userId);
  logger.log(`New MCP session created for user ${userId}`);

  await transport.handleRequest(req.raw, reply.raw, body);
}

/**
 * Reads and parses the JSON body from a Node.js IncomingMessage.
 */
function readRequestBody(req: IncomingMessage): Promise<unknown> {
  return new Promise((resolve, reject) => {
    const chunks: Buffer[] = [];
    req.on('data', (chunk: Buffer) => chunks.push(chunk));
    req.on('end', () => {
      const raw = Buffer.concat(chunks).toString('utf8');
      if (!raw) {
        resolve(undefined);
        return;
      }
      try {
        resolve(JSON.parse(raw));
      } catch (err) {
        reject(err);
      }
    });
    req.on('error', reject);
  });
}