dorfteich/apps/api/src/mcp/mcp.controller.ts

import {
  Controller,
  Delete,
  Get,
  HttpException,
  HttpStatus,
  NotFoundException,
  Post,
  Req,
  Res,
  UnauthorizedException,
} from '@nestjs/common';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import type { Request, Response } from 'express';

import { Public } from '../auth/auth.guard';
import { RateLimitService } from '../rate-limit/rate-limit.service';
import { ApiTokensService } from '../public-api/api-tokens.service';
import { InstanceSettingsService } from '../settings/instance-settings.service';
import { McpService } from './mcp.service';

/** Same budget as the REST surface — one shared runaway stop per token. */
const RATE_LIMIT = { limit: 120, windowSeconds: 60 };

/**
 * The MCP Streamable-HTTP endpoint (issue #105) at `/api/mcp` — under the
 * `/api/` path on purpose: every existing reverse proxy (stages,
 * self-hosters, the bundled Caddyfile) already routes it to the api, so no
 * deployment changes anywhere. Stateless: each POST builds a fresh server +
 * transport pair (no session store), which keeps api replicas trivial;
 * GET/DELETE (SSE resumption, session teardown) answer 405 accordingly.
 *
 * Gating mirrors the REST guard with the MCP switches: instance
 * `mcp.enabled` off → 404 (feature invisible), then PAT bearer auth and the
 * per-token rate limit. Pond opt-in (`mcpEnabled`) is enforced inside every
 * tool (McpService).
 */
@Controller('api/mcp')
@Public()
export class McpController {
  constructor(
    private readonly mcp: McpService,
    private readonly tokens: ApiTokensService,
    private readonly settings: InstanceSettingsService,
    private readonly rateLimits: RateLimitService,
  ) {}

  @Post()
  async handle(@Req() request: Request, @Res() response: Response): Promise<void> {
    if (!(await this.settings.get('mcp.enabled'))) throw new NotFoundException();

    const header = request.headers.authorization ?? '';
    const raw = header.startsWith('Bearer ') ? header.slice('Bearer '.length).trim() : '';
    const validated = raw ? await this.tokens.validate(raw) : null;
    if (!validated) throw new UnauthorizedException();

    const limited = await this.rateLimits.hit(
      'mcp',
      `token:${validated.token.id}`,
      RATE_LIMIT.limit,
      RATE_LIMIT.windowSeconds,
    );
    if (!limited.allowed) {
      response.setHeader('Retry-After', String(limited.retryAfterSeconds));
      throw new HttpException({ code: 'rate_limited' }, HttpStatus.TOO_MANY_REQUESTS);
    }

    const server = this.mcp.buildServer(validated.user, validated.token);
    const transport = new StreamableHTTPServerTransport({
      // Stateless mode: no session ids, every request stands alone.
      sessionIdGenerator: undefined,
    });
    response.on('close', () => {
      void transport.close();
      void server.close();
    });
    await server.connect(transport);
    // The body was parsed by the global express.json middleware.
    await transport.handleRequest(request, response, request.body);
  }

  @Get()
  async methodNotAllowed(@Res() response: Response): Promise<void> {
    if (!(await this.settings.get('mcp.enabled'))) throw new NotFoundException();
    // Stateless server: no SSE stream to resume, no session to delete.
    response.status(HttpStatus.METHOD_NOT_ALLOWED).json({
      jsonrpc: '2.0',
      error: { code: -32000, message: 'Method not allowed — stateless transport, POST only' },
      id: null,
    });
  }

  @Delete()
  async deleteNotAllowed(@Res() response: Response): Promise<void> {
    await this.methodNotAllowed(response);
  }
}