mcpctl/src/mcplocal/tests/plugin-gate.test.ts

/**
 * Gate Plugin Tests — verify the gate plugin produces identical behavior
 * to the legacy hardcoded gate in router.ts when wired via setPlugin().
 */
import { describe, it, expect, vi } from 'vitest';
import { McpRouter } from '../src/router.js';
import type { UpstreamConnection, JsonRpcRequest, JsonRpcResponse } from '../src/types.js';
import type { McpdClient } from '../src/http/mcpd-client.js';
import { ProviderRegistry } from '../src/providers/registry.js';
import type { LlmProvider, CompletionResult } from '../src/providers/types.js';
import { createGatePlugin } from '../src/proxymodel/plugins/gate.js';
import { LLMProviderAdapter } from '../src/proxymodel/llm-adapter.js';
import { MemoryCache } from '../src/proxymodel/cache.js';

function mockUpstream(
  name: string,
  opts: { tools?: Array<{ name: string; description?: string }> } = {},
): UpstreamConnection {
  return {
    name,
    isAlive: vi.fn(() => true),
    close: vi.fn(async () => {}),
    onNotification: vi.fn(),
    send: vi.fn(async (req: JsonRpcRequest): Promise<JsonRpcResponse> => {
      if (req.method === 'tools/list') {
        return { jsonrpc: '2.0', id: req.id, result: { tools: opts.tools ?? [] } };
      }
      if (req.method === 'tools/call') {
        return {
          jsonrpc: '2.0',
          id: req.id,
          result: { content: [{ type: 'text', text: `Called ${(req.params as Record<string, unknown>)?.name}` }] },
        };
      }
      if (req.method === 'resources/list') {
        return { jsonrpc: '2.0', id: req.id, result: { resources: [] } };
      }
      if (req.method === 'prompts/list') {
        return { jsonrpc: '2.0', id: req.id, result: { prompts: [] } };
      }
      return { jsonrpc: '2.0', id: req.id, error: { code: -32601, message: 'Not found' } };
    }),
  } as UpstreamConnection;
}

function mockMcpdClient(prompts: Array<{ name: string; priority: number; summary: string | null; chapters: string[] | null; content: string }> = []): McpdClient {
  return {
    get: vi.fn(async (path: string) => {
      if (path.includes('/prompts/visible')) {
        return prompts.map((p) => ({ ...p, type: 'prompt' }));
      }
      return [];
    }),
    post: vi.fn(async () => ({})),
    put: vi.fn(async () => ({})),
    delete: vi.fn(async () => {}),
    forward: vi.fn(async () => ({ status: 200, body: {} })),
    withHeaders: vi.fn(function (this: McpdClient) { return this; }),
  } as unknown as McpdClient;
}

const samplePrompts = [
  { name: 'common-mistakes', priority: 10, summary: 'Critical safety rules everyone must follow', chapters: null, content: 'NEVER do X. ALWAYS do Y.' },
  { name: 'zigbee-pairing', priority: 7, summary: 'How to pair Zigbee devices with the hub', chapters: ['Setup', 'Troubleshooting'], content: 'Step 1: Put device in pairing mode...' },
  { name: 'mqtt-config', priority: 5, summary: 'MQTT broker configuration guide', chapters: ['Broker Setup', 'Authentication'], content: 'Configure the MQTT broker at...' },
  { name: 'security-policy', priority: 8, summary: 'Security policies for production deployments', chapters: ['Network', 'Auth'], content: 'All connections must use TLS...' },
];

function setupPluginRouter(opts: {
  gated?: boolean;
  prompts?: typeof samplePrompts;
  withLlm?: boolean;
  llmResponse?: string;
  /** When set, the heavy provider's complete() rejects — exercises the loud fallback. */
  llmError?: boolean;
  byteBudget?: number;
} = {}): { router: McpRouter; mcpdClient: McpdClient } {
  const router = new McpRouter();
  const prompts = opts.prompts ?? samplePrompts;
  const mcpdClient = mockMcpdClient(prompts);
  router.setPromptConfig(mcpdClient, 'test-project');

  let providerRegistry: ProviderRegistry | null = null;
  if (opts.withLlm) {
    providerRegistry = new ProviderRegistry();
    const mockProvider: LlmProvider = {
      name: 'mock-heavy',
      complete: opts.llmError
        ? vi.fn().mockRejectedValue(new Error('upstream 400: model not found'))
        : vi.fn().mockResolvedValue({
          content: opts.llmResponse ?? '{ "selectedNames": ["zigbee-pairing"], "reasoning": "User is working with zigbee" }',
          toolCalls: [],
          usage: { promptTokens: 100, completionTokens: 50, totalTokens: 150 },
          finishReason: 'stop',
        } satisfies CompletionResult),
      listModels: vi.fn().mockResolvedValue([]),
      isAvailable: vi.fn().mockResolvedValue(true),
    };
    providerRegistry.register(mockProvider);
    providerRegistry.assignTier(mockProvider.name, 'heavy');
  }

  // Wire the gate PLUGIN instead of legacy setGateConfig
  const gatePlugin = createGatePlugin({
    gated: opts.gated !== false,
    providerRegistry,
    byteBudget: opts.byteBudget,
  });
  router.setPlugin(gatePlugin);

  // Wire proxymodel services (needed for plugin context)
  const llmAdapter = providerRegistry ? new LLMProviderAdapter(providerRegistry) : {
    complete: async () => '',
    available: () => false,
  };
  router.setProxyModel('default', llmAdapter, new MemoryCache());

  return { router, mcpdClient };
}

describe('Gate Plugin via setPlugin()', () => {
  describe('initialize with gating', () => {
    it('creates gated session on initialize', async () => {
      const { router } = setupPluginRouter();

      const res = await router.route(
        { jsonrpc: '2.0', id: 1, method: 'initialize' },
        { sessionId: 's1' },
      );

      expect(res.result).toBeDefined();
      const toolsRes = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/list' },
        { sessionId: 's1' },
      );
      const tools = (toolsRes.result as { tools: Array<{ name: string }> }).tools;
      expect(tools).toHaveLength(1);
      expect(tools[0]!.name).toBe('begin_session');
    });

    it('creates ungated session when project is not gated', async () => {
      const { router } = setupPluginRouter({ gated: false });
      router.addUpstream(mockUpstream('ha', { tools: [{ name: 'get_entities' }] }));

      await router.route(
        { jsonrpc: '2.0', id: 1, method: 'initialize' },
        { sessionId: 's1' },
      );

      const toolsRes = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/list' },
        { sessionId: 's1' },
      );
      const tools = (toolsRes.result as { tools: Array<{ name: string }> }).tools;
      const names = tools.map((t) => t.name);
      expect(names).toContain('ha/get_entities');
      expect(names).toContain('read_prompts');
      expect(names).toContain('propose_prompt');
      expect(names).not.toContain('begin_session');
    });
  });

  describe('tools/list gating', () => {
    it('shows only begin_session when session is gated', async () => {
      const { router } = setupPluginRouter();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/list' },
        { sessionId: 's1' },
      );

      const tools = (res.result as { tools: Array<{ name: string }> }).tools;
      expect(tools).toHaveLength(1);
      expect(tools[0]!.name).toBe('begin_session');
    });

    it('shows all tools plus read_prompts after ungating', async () => {
      const { router } = setupPluginRouter();
      router.addUpstream(mockUpstream('ha', { tools: [{ name: 'get_entities' }] }));
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['zigbee'] } } },
        { sessionId: 's1' },
      );

      const toolsRes = await router.route(
        { jsonrpc: '2.0', id: 3, method: 'tools/list' },
        { sessionId: 's1' },
      );
      const tools = (toolsRes.result as { tools: Array<{ name: string }> }).tools;
      const names = tools.map((t) => t.name);
      expect(names).toContain('ha/get_entities');
      expect(names).toContain('propose_prompt');
      expect(names).toContain('read_prompts');
      expect(names).not.toContain('begin_session');
    });
  });

  describe('begin_session', () => {
    it('falls back with a visible warning when the LLM selector errors (no silent degrade)', async () => {
      const { router } = setupPluginRouter({ withLlm: true, llmError: true });
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['zigbee', 'pairing'] } } },
        { sessionId: 's1' },
      );

      // Fell back — the call succeeded rather than failing/hanging …
      expect(res.error).toBeUndefined();
      const text = ((res.result as { content: Array<{ text: string }> }).content[0]!.text);
      // … and the degradation is visible, not silent …
      expect(text).toContain('Smart prompt-selection unavailable');
      // … while deterministic tag matching still returned relevant prompts.
      expect(text).toContain('common-mistakes');
      expect(text).toContain('zigbee-pairing');
    });

    it('returns matched prompts with keyword matching', async () => {
      const { router } = setupPluginRouter();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['zigbee', 'pairing'] } } },
        { sessionId: 's1' },
      );

      expect(res.error).toBeUndefined();
      const text = ((res.result as { content: Array<{ text: string }> }).content[0]!.text);
      expect(text).toContain('common-mistakes');
      expect(text).toContain('NEVER do X');
      expect(text).toContain('zigbee-pairing');
      expect(text).toContain('pairing mode');
      expect(text).toContain('read_prompts');
    });

    it('includes priority 10 prompts even without matching tags', async () => {
      const { router } = setupPluginRouter();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['unrelated-keyword'] } } },
        { sessionId: 's1' },
      );

      const text = ((res.result as { content: Array<{ text: string }> }).content[0]!.text);
      expect(text).toContain('common-mistakes');
      expect(text).toContain('NEVER do X');
    });

    it('uses LLM selection when provider is available', async () => {
      const { router } = setupPluginRouter({
        withLlm: true,
        llmResponse: '{ "selectedNames": ["zigbee-pairing", "security-policy"], "reasoning": "Zigbee pairing needs security awareness" }',
      });
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['zigbee'] } } },
        { sessionId: 's1' },
      );

      const text = ((res.result as { content: Array<{ text: string }> }).content[0]!.text);
      expect(text).toContain('Zigbee pairing needs security awareness');
      expect(text).toContain('zigbee-pairing');
      expect(text).toContain('security-policy');
      expect(text).toContain('common-mistakes');
    });

    it('rejects empty tags', async () => {
      const { router } = setupPluginRouter();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: [] } } },
        { sessionId: 's1' },
      );

      expect(res.error).toBeDefined();
      expect(res.error!.code).toBe(-32602);
    });

    it('returns message when session is already ungated', async () => {
      const { router } = setupPluginRouter();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['zigbee'] } } },
        { sessionId: 's1' },
      );

      const res = await router.route(
        { jsonrpc: '2.0', id: 3, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['mqtt'] } } },
        { sessionId: 's1' },
      );

      const text = ((res.result as { content: Array<{ text: string }> }).content[0]!.text);
      expect(text).toContain('already started');
      expect(text).toContain('read_prompts');
    });

    it('accepts description and tokenizes to keywords', async () => {
      const { router } = setupPluginRouter();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { description: 'I want to pair a zigbee device with mqtt' } } },
        { sessionId: 's1' },
      );

      expect(res.error).toBeUndefined();
      const text = (res.result as { content: Array<{ text: string }> }).content[0]!.text;
      expect(text).toContain('zigbee-pairing');
      expect(text).toContain('mqtt-config');
    });
  });

  describe('read_prompts', () => {
    it('returns prompts matching keywords', async () => {
      const { router } = setupPluginRouter({ gated: false });
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'read_prompts', arguments: { tags: ['mqtt', 'broker'] } } },
        { sessionId: 's1' },
      );

      expect(res.error).toBeUndefined();
      const text = ((res.result as { content: Array<{ text: string }> }).content[0]!.text);
      expect(text).toContain('mqtt-config');
      expect(text).toContain('Configure the MQTT broker');
    });

    it('filters out already-sent prompts', async () => {
      const { router } = setupPluginRouter({ byteBudget: 80 });
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['zigbee'] } } },
        { sessionId: 's1' },
      );

      const res = await router.route(
        { jsonrpc: '2.0', id: 3, method: 'tools/call', params: { name: 'read_prompts', arguments: { tags: ['mqtt'] } } },
        { sessionId: 's1' },
      );

      const text = ((res.result as { content: Array<{ text: string }> }).content[0]!.text);
      expect(text).toContain('mqtt-config');
      expect(text).not.toContain('NEVER do X');
    });

    it('rejects empty tags', async () => {
      const { router } = setupPluginRouter({ gated: false });
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'read_prompts', arguments: { tags: [] } } },
        { sessionId: 's1' },
      );

      expect(res.error).toBeDefined();
      expect(res.error!.code).toBe(-32602);
    });
  });

  describe('gated intercept', () => {
    it('auto-ungates when gated session calls a real tool', async () => {
      const { router } = setupPluginRouter();
      const ha = mockUpstream('ha', { tools: [{ name: 'get_entities' }] });
      router.addUpstream(ha);
      await router.discoverTools();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'ha/get_entities', arguments: { domain: 'light' } } },
        { sessionId: 's1' },
      );

      expect(res.error).toBeUndefined();
      const result = res.result as { content: Array<{ type: string; text: string }> };
      expect(result.content.length).toBeGreaterThanOrEqual(1);

      const toolsRes = await router.route(
        { jsonrpc: '2.0', id: 3, method: 'tools/list' },
        { sessionId: 's1' },
      );
      const tools = (toolsRes.result as { tools: Array<{ name: string }> }).tools;
      expect(tools.map((t) => t.name)).toContain('ha/get_entities');
    });

    it('includes project context in intercepted response', async () => {
      const { router } = setupPluginRouter();
      const ha = mockUpstream('ha', { tools: [{ name: 'get_entities' }] });
      router.addUpstream(ha);
      await router.discoverTools();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'ha/get_entities', arguments: { domain: 'light' } } },
        { sessionId: 's1' },
      );

      const result = res.result as { content: Array<{ type: string; text: string }> };
      const briefing = result.content[0]!.text;
      expect(briefing).toContain('common-mistakes');
      expect(briefing).toContain('NEVER do X');
    });
  });

  describe('initialize instructions for gated projects', () => {
    it('includes gate message and prompt index in instructions', async () => {
      const { router } = setupPluginRouter();

      const res = await router.route(
        { jsonrpc: '2.0', id: 1, method: 'initialize' },
        { sessionId: 's1' },
      );

      const result = res.result as { instructions?: string };
      expect(result.instructions).toBeDefined();
      expect(result.instructions).toContain('begin_session');
      expect(result.instructions).toContain('gated session');
      expect(result.instructions).toContain('common-mistakes');
      expect(result.instructions).toContain('zigbee-pairing');
    });

    it('does not include gate message for non-gated projects', async () => {
      const { router } = setupPluginRouter({ gated: false });
      router.setInstructions('Base project instructions');

      const res = await router.route(
        { jsonrpc: '2.0', id: 1, method: 'initialize' },
        { sessionId: 's1' },
      );

      const result = res.result as { instructions?: string };
      expect(result.instructions).toBe('Base project instructions');
      expect(result.instructions).not.toContain('gated session');
    });
  });

  describe('notifications after ungating', () => {
    it('queues tools/list_changed after begin_session ungating', async () => {
      const { router } = setupPluginRouter();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['zigbee'] } } },
        { sessionId: 's1' },
      );

      const notifications = router.consumeNotifications('s1');
      expect(notifications).toHaveLength(1);
      expect(notifications[0]!.method).toBe('notifications/tools/list_changed');
    });

    it('queues tools/list_changed after gated intercept', async () => {
      const { router } = setupPluginRouter();
      const ha = mockUpstream('ha', { tools: [{ name: 'get_entities' }] });
      router.addUpstream(ha);
      await router.discoverTools();
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'ha/get_entities', arguments: {} } },
        { sessionId: 's1' },
      );

      const notifications = router.consumeNotifications('s1');
      expect(notifications).toHaveLength(1);
      expect(notifications[0]!.method).toBe('notifications/tools/list_changed');
    });
  });

  describe('response size cap', () => {
    it('truncates begin_session response over 24K chars', async () => {
      const largePrompts = [
        { name: 'huge-prompt', priority: 10, summary: 'A very large prompt', chapters: null, content: 'x'.repeat(30_000) },
      ];
      const { router } = setupPluginRouter({ prompts: largePrompts, byteBudget: 50_000 });
      await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });

      const res = await router.route(
        { jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'begin_session', arguments: { tags: ['huge'] } } },
        { sessionId: 's1' },
      );

      expect(res.error).toBeUndefined();
      const text = (res.result as { content: Array<{ text: string }> }).content[0]!.text;
      expect(text.length).toBeLessThanOrEqual(24_000 + 100);
      expect(text).toContain('[Response truncated');
    });
  });
});