/** * 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 => { 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)?.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; 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: 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('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'); }); }); });