TREK/server/tests/unit/nest/discovery.controller.test.ts

import { ALL_SCOPES } from '../../../src/mcp/scopes';
import type { AddonsService } from '../../../src/nest/addons/addons.service';
import { ConsentCoopMiddleware } from '../../../src/nest/platform/consent-coop.middleware';
import { DiscoveryMetadataService } from '../../../src/nest/platform/discovery-metadata.service';
import { DiscoveryController, McpResourceDiscoveryController } from '../../../src/nest/platform/discovery.controller';
import {
  MCP_METADATA_MIDDLEWARE,
  createMcpMetadataMiddleware,
  mcpMetadataMiddlewareProvider,
} from '../../../src/nest/platform/mcp-metadata.middleware';
import { PlatformModule } from '../../../src/nest/platform/platform.module';
import type { MikroORM } from '@mikro-orm/core';
import { mcpAuthMetadataRouter } from '@modelcontextprotocol/sdk/server/auth/router';
import { NotFoundException } from '@nestjs/common';

import { describe, it, expect, vi, beforeEach } from 'vitest';

const h = vi.hoisted(() => ({
  getMcpSafeUrl: vi.fn(() => 'https://trek.example.test'),
  // The SDK factory returns a tagged middleware so tests can identify delegation
  // without running real router wiring.
  metaRouter: vi.fn(),
}));

vi.mock('../../../src/app-config', async (importOriginal) => {
  const actual = await importOriginal<typeof import('../../../src/app-config')>();
  return { ...actual, getMcpSafeUrl: h.getMcpSafeUrl };
});
vi.mock('@modelcontextprotocol/sdk/server/auth/router', () => ({
  mcpAuthMetadataRouter: vi.fn(() => h.metaRouter),
}));

function makeRes() {
  const res = {
    statusCode: 200,
    body: undefined as unknown,
    headers: {} as Record<string, string>,
    ended: false,
    status: vi.fn(function (this: typeof res, c: number) {
      this.statusCode = c;
      return this;
    }),
    json: vi.fn(function (this: typeof res, b: unknown) {
      this.body = b;
      return this;
    }),
    end: vi.fn(function (this: typeof res) {
      this.ended = true;
      return this;
    }),
    setHeader: vi.fn(function (this: typeof res, k: string, v: string) {
      this.headers[k] = v;
      return this;
    }),
  };
  return res;
}

function addons(enabled: boolean) {
  return { isAddonEnabled: vi.fn(() => enabled) } as unknown as AddonsService;
}

/**
 * createMcpMetadataMiddleware wraps the addon check in withRequestContext
 * (Plan 3a Task 4: this pathless pre-init mount runs before
 * @mikro-orm/nestjs's per-request EM fork, so once isAddonEnabled became
 * repository-backed it needed its own context). The mocked AddonsService
 * above never touches the ORM, so this fake only needs to satisfy
 * RequestContext.create's em.fork()/em.name contract.
 */
function fakeOrm(): MikroORM {
  const em = { name: 'default', fork: () => ({ name: 'default' }) };
  return { em } as unknown as MikroORM;
}

beforeEach(() => {
  vi.clearAllMocks();
  h.getMcpSafeUrl.mockReturnValue('https://trek.example.test');
});

describe('DiscoveryMetadataService', () => {
  it('DISC-001: builds the RFC 8414 AS metadata from the live safe URL', () => {
    const meta = new DiscoveryMetadataService().getOAuthMetadata();
    expect(meta).toEqual({
      issuer: 'https://trek.example.test',
      authorization_endpoint: 'https://trek.example.test/oauth/authorize',
      token_endpoint: 'https://trek.example.test/oauth/token',
      revocation_endpoint: 'https://trek.example.test/oauth/revoke',
      registration_endpoint: 'https://trek.example.test/oauth/register',
      response_types_supported: ['code'],
      grant_types_supported: ['authorization_code', 'refresh_token', 'client_credentials'],
      code_challenge_methods_supported: ['S256'],
      token_endpoint_auth_methods_supported: ['client_secret_post', 'none'],
      scopes_supported: ALL_SCOPES,
    });
  });

  it('DISC-002: trims trailing slashes off the configured base URL', () => {
    h.getMcpSafeUrl.mockReturnValue('https://trek.example.test///');
    expect(new DiscoveryMetadataService().getOAuthMetadata().issuer).toBe('https://trek.example.test');
  });

  it('DISC-003: caches the metadata after the first lazy build', () => {
    const svc = new DiscoveryMetadataService();
    const first = svc.getOAuthMetadata();
    expect(svc.getOAuthMetadata()).toBe(first);
    expect(h.getMcpSafeUrl).toHaveBeenCalledTimes(1);
  });

  it('DISC-004: builds the SDK metadata router once, pointing at the /mcp resource', () => {
    const svc = new DiscoveryMetadataService();
    expect(svc.getMetaRouter()).toBe(h.metaRouter);
    expect(svc.getMetaRouter()).toBe(h.metaRouter);
    expect(mcpAuthMetadataRouter).toHaveBeenCalledTimes(1);
    const opts = vi.mocked(mcpAuthMetadataRouter).mock.calls[0][0];
    expect(opts.resourceServerUrl.href).toBe('https://trek.example.test/mcp');
    expect(opts.resourceName).toBe('TREK MCP');
    expect(opts.scopesSupported).toEqual(ALL_SCOPES);
  });
});

/**
 * createMcpMetadataMiddleware resolves the addon gate inside an async IIFE
 * (Express cannot await a middleware), so on a /.well-known path the 404 or the
 * router delegation lands a microtask after mw() returns. Drain the microtask
 * queue — no timer is involved, so there is nothing to sleep for.
 */
async function flushMicrotasks(): Promise<void> {
  for (let i = 0; i < 10; i++) await Promise.resolve();
}

describe('createMcpMetadataMiddleware', () => {
  it('DISC-010: 404s a /.well-known path with an empty body when MCP is disabled', async () => {
    const mw = createMcpMetadataMiddleware(new DiscoveryMetadataService(), addons(false), fakeOrm());
    const res = makeRes();
    const next = vi.fn();
    mw({ path: '/.well-known/oauth-authorization-server' } as never, res as never, next);
    await flushMicrotasks();
    expect(res.statusCode).toBe(404);
    expect(res.ended).toBe(true);
    expect(next).not.toHaveBeenCalled();
    expect(h.metaRouter).not.toHaveBeenCalled();
  });

  it('DISC-011: delegates a /.well-known path to the SDK router when MCP is enabled', async () => {
    const mw = createMcpMetadataMiddleware(new DiscoveryMetadataService(), addons(true), fakeOrm());
    const next = vi.fn();
    mw({ path: '/.well-known/oauth-authorization-server' } as never, makeRes() as never, next);
    await flushMicrotasks();
    expect(h.metaRouter).toHaveBeenCalled();
  });

  it('DISC-012: delegates non-well-known paths without consulting the addon gate', () => {
    const gate = addons(false);
    const mw = createMcpMetadataMiddleware(new DiscoveryMetadataService(), gate, fakeOrm());
    mw({ path: '/api/trips' } as never, makeRes() as never, vi.fn());
    expect(gate.isAddonEnabled).not.toHaveBeenCalled();
    expect(h.metaRouter).toHaveBeenCalled();
  });

  it('DISC-013: the factory provider wires the token to the factory with its three deps (MikroORM added — Plan 3a Task 4)', () => {
    expect(mcpMetadataMiddlewareProvider.provide).toBe(MCP_METADATA_MIDDLEWARE);
    expect(mcpMetadataMiddlewareProvider.useFactory).toBe(createMcpMetadataMiddleware);
    expect(mcpMetadataMiddlewareProvider.inject.length).toBe(3);
  });
});

describe('ConsentCoopMiddleware', () => {
  it('DISC-020: relaxes COOP then continues', () => {
    const res = makeRes();
    const next = vi.fn();
    new ConsentCoopMiddleware().use({} as never, res as never, next);
    expect(res.headers['Cross-Origin-Opener-Policy']).toBe('unsafe-none');
    expect(next).toHaveBeenCalled();
  });
});

describe('DiscoveryController', () => {
  function controller(enabled = true) {
    return new DiscoveryController(new DiscoveryMetadataService(), addons(enabled));
  }

  it('DISC-030: openid-configuration is the AS metadata plus userinfo_endpoint', () => {
    const res = makeRes();
    controller().openidConfiguration(res as never);
    const body = res.body as { issuer: string; userinfo_endpoint: string };
    expect(body.issuer).toBe('https://trek.example.test');
    expect(body.userinfo_endpoint).toBe('https://trek.example.test/oauth/userinfo');
  });

  it('DISC-031: flat oauth-protected-resource 404s empty when MCP is disabled', async () => {
    const res = makeRes();
    await controller(false).protectedResource(res as never);
    expect(res.statusCode).toBe(404);
    expect(res.ended).toBe(true);
  });

  it('DISC-032: flat oauth-protected-resource serves the PRM document when enabled', async () => {
    const res = makeRes();
    await controller().protectedResource(res as never);
    expect(res.body).toEqual({
      resource: 'https://trek.example.test/mcp',
      authorization_servers: ['https://trek.example.test'],
      bearer_methods_supported: ['header'],
      scopes_supported: ALL_SCOPES,
      resource_name: 'TREK MCP',
    });
  });

  it('DISC-033: the catchall answers 404 JSON, never throwing into the SPA fallback', () => {
    const res = makeRes();
    controller().wellKnownFallback(res as never);
    expect(res.statusCode).toBe(404);
    expect(res.body).toEqual({ error: 'not_found' });
  });

  it('DISC-034: /.well-known/ (trailing slash) gets the same 404 JSON', () => {
    const res = makeRes();
    controller().wellKnownRoot({ path: '/.well-known/' } as never, res as never);
    expect(res.statusCode).toBe(404);
    expect(res.body).toEqual({ error: 'not_found' });
  });

  it('DISC-035: bare /.well-known reproduces the framework 404 (SpaFallback route)', () => {
    const req = { path: '/.well-known', method: 'GET', originalUrl: '/.well-known' };
    expect(() => controller().wellKnownRoot(req as never, makeRes() as never)).toThrowError(
      new NotFoundException('Cannot GET /.well-known'),
    );
  });

  it('DISC-036: route declaration order keeps the concrete documents ahead of the catchalls', () => {
    // Nest registers routes in method-declaration order; the catchall must not
    // shadow the concrete documents.
    const names = Object.getOwnPropertyNames(DiscoveryController.prototype).filter((n) => n !== 'constructor');
    expect(names.indexOf('openidConfiguration')).toBeLessThan(names.indexOf('wellKnownRoot'));
    expect(names.indexOf('protectedResource')).toBeLessThan(names.indexOf('wellKnownRoot'));
    expect(names.indexOf('wellKnownRoot')).toBeLessThan(names.indexOf('wellKnownFallback'));
  });

  it('DISC-037: the RFC 8414 path-inserted AS metadata is served, not 404', () => {
    const res = makeRes();
    controller().authorizationServerForMcp(res as never);
    expect((res.body as { issuer: string }).issuer).toBe('https://trek.example.test');
    expect(res.statusCode).toBe(200);
  });
});

describe('McpResourceDiscoveryController', () => {
  function resourceController(enabled = true) {
    return new McpResourceDiscoveryController(new DiscoveryController(new DiscoveryMetadataService(), addons(enabled)));
  }

  it('DISC-038: serves the three documents under the resource URL, never SPA HTML', async () => {
    const oidc = makeRes();
    resourceController().openidConfiguration(oidc as never);
    expect((oidc.body as { userinfo_endpoint: string }).userinfo_endpoint).toBe(
      'https://trek.example.test/oauth/userinfo',
    );

    const as = makeRes();
    resourceController().authorizationServer(as as never);
    expect((as.body as { registration_endpoint: string }).registration_endpoint).toBe(
      'https://trek.example.test/oauth/register',
    );

    const prm = makeRes();
    await resourceController().protectedResource(prm as never);
    expect((prm.body as { resource: string }).resource).toBe('https://trek.example.test/mcp');
  });

  it('DISC-039: an unknown path under the resource URL answers 404 JSON, and the PRM follows the addon switch', async () => {
    const res = makeRes();
    resourceController().fallback(res as never);
    expect(res.statusCode).toBe(404);
    expect(res.body).toEqual({ error: 'not_found' });

    const off = makeRes();
    await resourceController(false).protectedResource(off as never);
    expect(off.statusCode).toBe(404);
    expect(off.ended).toBe(true);
  });
});

describe('PlatformModule', () => {
  it('DISC-040: applies the consent COOP middleware to the /oauth/consent prefix', () => {
    const forRoutes = vi.fn();
    const consumer = { apply: vi.fn(() => ({ forRoutes })) };
    new PlatformModule().configure(consumer as never);
    expect(consumer.apply).toHaveBeenCalledWith(ConsentCoopMiddleware);
    expect(forRoutes).toHaveBeenCalledWith('oauth/consent');
  });
});