TREK/server/tests/integration/mcp.test.ts

/**
 * MCP integration tests.
 * Covers MCP-001 to MCP-013.
 *
 * The MCP endpoint uses JWT auth and server-sent events / streaming HTTP.
 * Tests cover authentication, session management, rate limiting, and API token auth.
 */
import { describe, it, expect, vi, beforeAll, beforeEach, afterAll } from 'vitest';
import request from 'supertest';
import type { Application } from 'express';
import type { INestApplication } from '@nestjs/common';

const { testDb, dbMock } = vi.hoisted(() => {
  const Database = require('better-sqlite3');
  const db = new Database(':memory:');
  db.exec('PRAGMA journal_mode = WAL');
  db.exec('PRAGMA foreign_keys = ON');
  db.exec('PRAGMA busy_timeout = 5000');
  const mock = {
    db,
    closeDb: () => {},
    reinitialize: () => {},
    getPlaceWithTags: (placeId: number) => {
      const place: any = db.prepare(`SELECT p.*, c.name as category_name, c.color as category_color, c.icon as category_icon FROM places p LEFT JOIN categories c ON p.category_id = c.id WHERE p.id = ?`).get(placeId);
      if (!place) return null;
      const tags = db.prepare(`SELECT t.* FROM tags t JOIN place_tags pt ON t.id = pt.tag_id WHERE pt.place_id = ?`).all(placeId);
      return { ...place, category: place.category_id ? { id: place.category_id, name: place.category_name, color: place.category_color, icon: place.category_icon } : null, tags };
    },
    canAccessTrip: (tripId: any, userId: number) =>
      db.prepare(`SELECT t.id, t.user_id FROM trips t LEFT JOIN trip_members m ON m.trip_id = t.id AND m.user_id = ? WHERE t.id = ? AND (t.user_id = ? OR m.user_id IS NOT NULL)`).get(userId, tripId, userId),
    isOwner: (tripId: any, userId: number) =>
      !!db.prepare('SELECT id FROM trips WHERE id = ? AND user_id = ?').get(tripId, userId),
  };
  return { testDb: db, dbMock: mock };
});

vi.mock('../../src/db/database', () => dbMock);
vi.mock('../../src/config', () => ({
  JWT_SECRET: 'test-jwt-secret-for-trek-testing-only',
  ENCRYPTION_KEY: 'a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6a7b8c9d0e1f2a3b4c5d6a7b8c9d0e1f2',
  updateJwtSecret: () => {},
  SESSION_DURATION: '24h',
  SESSION_DURATION_MS: 86400000,
  SESSION_DURATION_SECONDS: 86400,
  DEFAULT_LANGUAGE: 'en',
}));
vi.mock('../../src/websocket', () => ({ broadcast: vi.fn(), broadcastToUser: vi.fn() }));

import { buildApp } from '../../src/bootstrap';
import { createTables } from '../../src/db/schema';
import { runMigrations } from '../../src/db/migrations';
import { resetTestDb, resetRateLimits } from '../helpers/test-db';
import { createUser } from '../helpers/factories';
import { generateToken } from '../helpers/auth';
import { createMcpToken } from '../helpers/factories';
import { closeMcpSessions } from '../../src/mcp/index';

let nestApp: INestApplication;
let app: Application;

beforeAll(async () => {
  createTables(testDb);
  runMigrations(testDb);
  nestApp = await buildApp();
  app = nestApp.getHttpAdapter().getInstance();
});

beforeEach(() => {
  resetTestDb(testDb);
  resetRateLimits(nestApp);
});

afterAll(async () => {
  closeMcpSessions();
  await nestApp.close();
  testDb.close();
});

describe('MCP authentication', () => {
  // MCP handler checks if the 'mcp' addon is enabled first (403 if not),
  // then checks auth (401). In test DB the addon may be disabled.

  it('MCP-001 — POST /mcp without auth returns 403 (addon disabled before auth check)', async () => {
    const res = await request(app)
      .post('/mcp')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1 });
    // MCP handler checks addon enabled before verifying auth; addon is disabled in test DB
    expect(res.status).toBe(403);
  });

  it('MCP-001 — GET /mcp without auth returns 403 (addon disabled)', async () => {
    const res = await request(app).get('/mcp');
    expect(res.status).toBe(403);
  });

  it('MCP-001 — DELETE /mcp without auth returns 403 (addon disabled)', async () => {
    const res = await request(app)
      .delete('/mcp')
      .set('Mcp-Session-Id', 'fake-session-id');
    expect(res.status).toBe(403);
  });
});

describe('MCP session init', () => {
  it('MCP-002 — POST /mcp with valid JWT passes auth check (may fail if addon disabled)', async () => {
    const { user } = createUser(testDb);
    const token = generateToken(user.id);

    // Enable MCP addon in test DB
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const res = await request(app)
      .post('/mcp')
      .set('Authorization', `Bearer ${token}`)
      .set('Accept', 'application/json, text/event-stream')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1, params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1' } } });
    // Valid JWT + enabled addon → auth passes; SDK returns 200 with session headers
    expect(res.status).toBe(200);
  });

  it('MCP-003 — DELETE /mcp with unknown session returns 404', async () => {
    const { user } = createUser(testDb);
    const token = generateToken(user.id);

    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const res = await request(app)
      .delete('/mcp')
      .set('Authorization', `Bearer ${token}`)
      .set('Mcp-Session-Id', 'nonexistent-session-id');
    expect(res.status).toBe(404);
  });

  it('MCP-004 — POST /mcp with invalid JWT returns 401 (when addon enabled)', async () => {
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const res = await request(app)
      .post('/mcp')
      .set('Authorization', 'Bearer invalid.jwt.token')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1 });
    expect(res.status).toBe(401);
  });
});

describe('MCP API token auth', () => {
  it('MCP-002 — POST /mcp with valid trek_ API token authenticates successfully', async () => {
    const { user } = createUser(testDb);
    const { rawToken } = createMcpToken(testDb, user.id);
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const res = await request(app)
      .post('/mcp')
      .set('Authorization', `Bearer ${rawToken}`)
      .set('Accept', 'application/json, text/event-stream')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1, params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1' } } });
    expect(res.status).toBe(200);
  });

  it('MCP-002 — last_used_at is updated on token use', async () => {
    const { user } = createUser(testDb);
    const { rawToken, id: tokenId } = createMcpToken(testDb, user.id);
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const before = (testDb.prepare('SELECT last_used_at FROM mcp_tokens WHERE id = ?').get(tokenId) as { last_used_at: string | null }).last_used_at;

    await request(app)
      .post('/mcp')
      .set('Authorization', `Bearer ${rawToken}`)
      .set('Accept', 'application/json, text/event-stream')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1, params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1' } } });

    const after = (testDb.prepare('SELECT last_used_at FROM mcp_tokens WHERE id = ?').get(tokenId) as { last_used_at: string | null }).last_used_at;
    expect(after).not.toBeNull();
    expect(after).not.toBe(before);
  });

  it('MCP — POST /mcp with unknown trek_ token returns 401', async () => {
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const res = await request(app)
      .post('/mcp')
      .set('Authorization', 'Bearer trek_totally_fake_token_not_in_db')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1 });
    expect(res.status).toBe(401);
  });

  it('MCP — POST /mcp with no Authorization header returns 401', async () => {
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const res = await request(app)
      .post('/mcp')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1 });
    expect(res.status).toBe(401);
  });
});

describe('MCP session management', () => {
  async function createSession(userId: number): Promise<string> {
    const token = generateToken(userId);
    const res = await request(app)
      .post('/mcp')
      .set('Authorization', `Bearer ${token}`)
      .set('Accept', 'application/json, text/event-stream')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1, params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1' } } });
    expect(res.status).toBe(200);
    const sessionId = res.headers['mcp-session-id'];
    expect(sessionId).toBeTruthy();
    return sessionId as string;
  }

  it('MCP-003 — session limit of 5 per user', async () => {
    const { user } = createUser(testDb);
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    // Create 5 sessions
    for (let i = 0; i < 20; i++) {
      await createSession(user.id);
    }

    // 6th should fail
    const token = generateToken(user.id);
    const res = await request(app)
      .post('/mcp')
      .set('Authorization', `Bearer ${token}`)
      .set('Accept', 'application/json, text/event-stream')
      .send({ jsonrpc: '2.0', method: 'initialize', id: 1, params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1' } } });
    expect(res.status).toBe(429);
    expect(res.body.error).toMatch(/session limit/i);
  });

  it('MCP — session resumption with valid mcp-session-id', async () => {
    const { user } = createUser(testDb);
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();
    const sessionId = await createSession(user.id);
    const token = generateToken(user.id);

    const res = await request(app)
      .post('/mcp')
      .set('Authorization', `Bearer ${token}`)
      .set('mcp-session-id', sessionId)
      .set('Accept', 'application/json, text/event-stream')
      .send({ jsonrpc: '2.0', method: 'tools/list', id: 2, params: {} });
    expect(res.status).toBe(200);
  });

  it('MCP — session belongs to different user returns 403', async () => {
    const { user: user1 } = createUser(testDb);
    const { user: user2 } = createUser(testDb);
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();

    const sessionId = await createSession(user1.id);
    const token2 = generateToken(user2.id);

    const res = await request(app)
      .post('/mcp')
      .set('Authorization', `Bearer ${token2}`)
      .set('mcp-session-id', sessionId)
      .send({ jsonrpc: '2.0', method: 'tools/list', id: 2 });
    expect(res.status).toBe(403);
  });

  it('MCP — GET without mcp-session-id returns 400', async () => {
    const { user } = createUser(testDb);
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();
    const token = generateToken(user.id);

    const res = await request(app)
      .get('/mcp')
      .set('Authorization', `Bearer ${token}`);
    expect(res.status).toBe(400);
  });
});

describe('MCP rate limiting', () => {
  it('MCP-005 — requests below limit succeed', async () => {
    const { user } = createUser(testDb);
    testDb.prepare("UPDATE addons SET enabled = 1 WHERE id = 'mcp'").run();
    const token = generateToken(user.id);

    // Set a very low rate limit via env for this test
    const originalLimit = process.env.MCP_RATE_LIMIT;
    process.env.MCP_RATE_LIMIT = '3';

    try {
      for (let i = 0; i < 3; i++) {
        const res = await request(app)
          .post('/mcp')
          .set('Authorization', `Bearer ${token}`)
          .set('Accept', 'application/json, text/event-stream')
          .send({ jsonrpc: '2.0', method: 'initialize', id: i + 1, params: { protocolVersion: '2024-11-05', capabilities: {}, clientInfo: { name: 'test', version: '1' } } });
        // Each should pass (no rate limit hit yet since limit is read at module init,
        // but we can verify that the responses are not 429)
        expect(res.status).not.toBe(429);
      }
    } finally {
      if (originalLimit === undefined) delete process.env.MCP_RATE_LIMIT;
      else process.env.MCP_RATE_LIMIT = originalLimit;
    }
  });
});