import { afterEach, describe, expect, it, vi } from 'vitest' import type { KnowledgeService } from '../knowledge/knowledge-service' import { KnowledgeMcpGateway } from './knowledge-mcp-gateway' const firstLibraryId = '11111111-1111-4111-8111-111111111111' const secondLibraryId = '22222222-2222-4222-8222-222222222222' function createService() { const searchHybridMany = vi.fn( async (libraryIds: readonly string[]) => libraryIds.map((knowledgeBaseId, index) => ({ knowledgeBaseId, result: { document: { id: `33333333-3333-4333-8333-33333333333${index}`, title: `文档 ${index}` }, source: { displayName: `来源 ${index}`, location: `/private/${index}` }, chunk: { location: `第 ${index + 1} 段` }, snippet: `匹配 ${index}`, rank: index + 1, retrieval: { channels: ['fts'] as const, evidenceIds: [] } } })) ) const service = { database: { listKnowledgeBases: () => [ { id: firstLibraryId, name: '一号知识库' }, { id: secondLibraryId, name: '二号知识库' } ] }, searchHybridMany } as unknown as KnowledgeService return { service, searchHybridMany } } const gateways: KnowledgeMcpGateway[] = [] afterEach(async () => { await Promise.all(gateways.splice(0).map((gateway) => gateway.dispose())) }) describe('KnowledgeMcpGateway', () => { it('keeps scope server-side, strips markup, bounds model arguments, and drains references', async () => { const { service, searchHybridMany } = createService() const gateway = new KnowledgeMcpGateway(service) gateways.push(gateway) const token = gateway.grant( 'request-1', [secondLibraryId], new AbortController().signal ) expect(token).toMatch(/^[A-Za-z0-9_-]{40,}$/u) const references = await gateway.search(token!, { query: ' 要找什么 ', limit: 1 }) expect(searchHybridMany).toHaveBeenCalledWith( [secondLibraryId], '要找什么', 1, expect.any(AbortSignal) ) expect(references).toEqual([ expect.objectContaining({ libraryId: secondLibraryId, libraryName: '二号知识库', snippet: '匹配 0' }) ]) expect(gateway.drainReferences(token)).toEqual(references) expect(gateway.drainReferences(token)).toEqual([]) await expect( gateway.search(token!, { query: 'x', limit: 9, libraryIds: [firstLibraryId] }) ).rejects.toThrow() }) it('creates no capability for empty scope and rejects revoked, aborted, and expired capabilities', async () => { const { service } = createService() let now = 1_000 const gateway = new KnowledgeMcpGateway(service, { capabilityTtlMs: 10, now: () => now }) gateways.push(gateway) expect( gateway.grant('empty', [], new AbortController().signal) ).toBeUndefined() const revoked = gateway.grant( 'revoked', [firstLibraryId], new AbortController().signal )! gateway.revoke(revoked) await expect( gateway.search(revoked, { query: 'x' }) ).rejects.toThrow('unavailable or expired') const abortController = new AbortController() const aborted = gateway.grant( 'aborted', [firstLibraryId], abortController.signal )! abortController.abort() await expect( gateway.search(aborted, { query: 'x' }) ).rejects.toThrow('unavailable or expired') const expired = gateway.grant( 'expired', [firstLibraryId], new AbortController().signal )! now += 11 await expect( gateway.search(expired, { query: 'x' }) ).rejects.toThrow('unavailable or expired') }) it('binds a POST-only authenticated endpoint and rejects oversized bodies', async () => { const { service } = createService() const gateway = new KnowledgeMcpGateway(service, { maximumBodyBytes: 32 }) gateways.push(gateway) await gateway.start() const endpoint = gateway.getEndpoint()! const token = gateway.grant( 'http', [firstLibraryId], new AbortController().signal )! const getResponse = await fetch(endpoint) expect(getResponse.status).toBe(405) expect(getResponse.headers.get('access-control-allow-origin')).toBeNull() const unauthorized = await fetch(endpoint, { method: 'POST', headers: { authorization: `Bearer ${token}x` }, body: '{}' }) expect(unauthorized.status).toBe(401) const oversized = await fetch(endpoint, { method: 'POST', headers: { authorization: `Bearer ${token}` }, body: JSON.stringify({ value: 'x'.repeat(100) }) }) expect(oversized.status).toBe(413) }) })