Trek_CN/server/src/mcp/index.ts

import { Request, Response } from 'express';
import { randomUUID } from 'crypto';
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp';
import { User } from '../types';
import { verifyMcpToken, verifyJwtToken } from '../services/authService';
import { getUserByAccessToken } from '../services/oauthService';
import { isAddonEnabled } from '../services/adminService';
import { ADDON_IDS } from '../addons';
import { registerResources } from './resources';
import { registerTools } from './tools';
import { McpSession, sessions, revokeUserSessions, revokeUserSessionsForClient } from './sessionManager';

export { revokeUserSessions, revokeUserSessionsForClient };

const SESSION_TTL_MS = 60 * 60 * 1000; // 1 hour
const sessionParsed = Number.parseInt(process.env.MCP_MAX_SESSION_PER_USER ?? "");
const MAX_SESSIONS_PER_USER = Number.isFinite(sessionParsed) && sessionParsed > 0 ? sessionParsed : 5;
const RATE_LIMIT_WINDOW_MS = 60 * 1000; // 1 minute
const parsed = Number.parseInt(process.env.MCP_RATE_LIMIT ?? "");
const RATE_LIMIT_MAX = Number.isFinite(parsed) && parsed > 0 ? parsed : 60; // requests per minute per user

interface RateLimitEntry {
  count: number;
  windowStart: number;
}
const rateLimitMap = new Map<number, RateLimitEntry>();

function isRateLimited(userId: number): boolean {
  const now = Date.now();
  const entry = rateLimitMap.get(userId);
  if (!entry || now - entry.windowStart > RATE_LIMIT_WINDOW_MS) {
    rateLimitMap.set(userId, { count: 1, windowStart: now });
    return false;
  }
  entry.count += 1;
  return entry.count > RATE_LIMIT_MAX;
}

function countSessionsForUser(userId: number): number {
  const cutoff = Date.now() - SESSION_TTL_MS;
  let count = 0;
  for (const session of sessions.values()) {
    if (session.userId === userId && session.lastActivity >= cutoff) count++;
  }
  return count;
}

const sessionSweepInterval = setInterval(() => {
  const cutoff = Date.now() - SESSION_TTL_MS;
  let cleaned = 0;
  for (const [sid, session] of sessions) {
    if (session.lastActivity < cutoff) {
      try { session.server.close(); } catch { /* ignore */ }
      try { session.transport.close(); } catch { /* ignore */ }
      sessions.delete(sid);
      cleaned++;
    }
  }
  const rateCutoff = Date.now() - RATE_LIMIT_WINDOW_MS;
  for (const [uid, entry] of rateLimitMap) {
    if (entry.windowStart < rateCutoff) rateLimitMap.delete(uid);
  }
  if (cleaned > 0 || sessions.size > 0) {
    console.log(`[MCP] Session sweep: cleaned ${cleaned}, active ${sessions.size}`);
  }
}, 10 * 60 * 1000); // sweep every 10 minutes

// Prevent the interval from keeping the process alive if nothing else is running
sessionSweepInterval.unref();

interface VerifyTokenResult {
  user: User;
  /** null = full access (static token or JWT); string[] = OAuth 2.1 scoped access */
  scopes: string[] | null;
  /** OAuth client_id when authenticated via OAuth 2.1; null otherwise */
  clientId: string | null;
  isStaticToken: boolean;
}

function verifyToken(authHeader: string | undefined): VerifyTokenResult | null {
  if (!authHeader) return null;
  // M8: strictly require "Bearer" scheme (RFC 6750)
  const spaceIdx = authHeader.indexOf(' ');
  if (spaceIdx === -1) return null;
  const scheme = authHeader.slice(0, spaceIdx);
  const token  = authHeader.slice(spaceIdx + 1);
  if (scheme.toLowerCase() !== 'bearer' || !token) return null;

  // OAuth 2.1 access token (trekoa_...)
  if (token.startsWith('trekoa_')) {
    const result = getUserByAccessToken(token);
    if (!result) return null;
    return { user: result.user, scopes: result.scopes, clientId: result.clientId, isStaticToken: false };
  }

  // Long-lived static MCP token (trek_...) — full access + deprecation notice
  if (token.startsWith('trek_')) {
    const user = verifyMcpToken(token);
    if (!user) return null;
    return { user, scopes: null, clientId: null, isStaticToken: true };
  }

  // Short-lived JWT (TREK web session used directly) — full access, no notice
  const user = verifyJwtToken(token);
  if (!user) return null;
  return { user, scopes: null, clientId: null, isStaticToken: false };
}

export async function mcpHandler(req: Request, res: Response): Promise<void> {
  if (!isAddonEnabled(ADDON_IDS.MCP)) {
    res.status(403).json({ error: 'MCP is not enabled' });
    return;
  }

  const tokenResult = verifyToken(req.headers['authorization']);
  if (!tokenResult) {
    res.status(401).json({ error: 'Access token required' });
    return;
  }
  const { user, scopes, clientId, isStaticToken } = tokenResult;

  if (isRateLimited(user.id)) {
    res.status(429).json({ error: 'Too many requests. Please slow down.' });
    return;
  }

  const sessionId = req.headers['mcp-session-id'] as string | undefined;

  // Resume an existing session
  if (sessionId) {
    const session = sessions.get(sessionId);
    if (!session) {
      res.status(404).json({ error: 'Session not found' });
      return;
    }
    if (session.userId !== user.id) {
      res.status(403).json({ error: 'Session belongs to a different user' });
      return;
    }
    session.lastActivity = Date.now();
    try {
      await session.transport.handleRequest(req, res, req.body);
    } catch (err) {
      console.error('[MCP] transport.handleRequest error:', err);
      if (!res.headersSent) {
        res.status(500).json({ error: 'Internal MCP error' });
      }
    }
    return;
  }

  // Only POST can initialize a new session
  if (req.method !== 'POST') {
    res.status(400).json({ error: 'Missing mcp-session-id header' });
    return;
  }

  if (countSessionsForUser(user.id) >= MAX_SESSIONS_PER_USER) {
    res.status(429).json({ error: 'Session limit reached. Close an existing session before opening a new one.' });
    return;
  }

  // Create a new per-user MCP server and session
  const server = new McpServer({
    name: 'TREK MCP',
    version: '1.0.0',
    capabilities: {
      resources: { listChanged: true },
      tools: { listChanged: true },
      prompts: { listChanged: true },
    },
  });
  registerResources(server, user.id, scopes);
  registerTools(server, user.id, scopes, isStaticToken);

  const transport = new StreamableHTTPServerTransport({
    sessionIdGenerator: () => randomUUID(),
    onsessioninitialized: (sid) => {
      sessions.set(sid, { server, transport, userId: user.id, scopes, clientId, isStaticToken, lastActivity: Date.now() });
      const authMethod = isStaticToken ? 'static-token' : scopes ? `oauth(${scopes.join(',')})` : 'jwt';
      console.log(`[MCP] Session ${sid} created for user ${user.id} [${authMethod}]. Active sessions: ${sessions.size}`);
    },
    onsessionclosed: (sid) => {
      sessions.delete(sid);
    },
  });

  try {
    await server.connect(transport);
    await transport.handleRequest(req, res, req.body);
  } catch (err) {
    console.error('[MCP] transport.handleRequest error:', err);
    if (!res.headersSent) {
      res.status(500).json({ error: 'Internal MCP error', detail: String(err) });
    }
  }
}

/** Close all active MCP sessions (call during graceful shutdown). */
export function closeMcpSessions(): void {
  clearInterval(sessionSweepInterval);
  for (const [, session] of sessions) {
    try { session.server.close(); } catch { /* ignore */ }
    try { session.transport.close(); } catch { /* ignore */ }
  }
  sessions.clear();
  rateLimitMap.clear();
}