goodbuddy/src/main/agent/knowledge-mcp-gateway.test.ts

import { mkdtemp, rm } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { createServer, type Server } from 'node:http'
import { afterEach, describe, expect, it, vi } from 'vitest'
import { z } from 'zod'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { Server as McpProtocolServer } from '@modelcontextprotocol/sdk/server/index.js'
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
import {
  LATEST_PROTOCOL_VERSION,
  ListToolsRequestSchema,
  isInitializeRequest,
  type Tool
} from '@modelcontextprotocol/sdk/types.js'
import type { KnowledgeService } from '../knowledge/knowledge-service'
import { AssistantDatabase } from '../assistant/assistant-database'
import {
  browserToolNames,
  KnowledgeMcpGateway,
  type MagicNotesDatabase
} 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: {
            id: `44444444-4444-4444-8444-44444444444${index}`,
            location: `第 ${index + 1} 段`
          },
          snippet: `<mark>匹配</mark> ${index}`,
          rank: index + 1,
          retrieval: {
            score: 0.5,
            channels: ['fts'] as const,
            lexicalRank: 1,
            evidenceIds: []
          }
        }
      }))
  )
  const service = {
    database: {
      listKnowledgeBases: () => [
        {
          id: firstLibraryId,
          name: '一号知识库',
          description: '不应暴露'
        },
        {
          id: secondLibraryId,
          name: '二号知识库',
          description: '已授权知识'
        }
      ]
    },
    searchHybridMany
  } as unknown as KnowledgeService
  return { service, searchHybridMany }
}

function customMcpServer(
  url: string,
  id = '00000000-0000-4000-8000-000000000092'
) {
  return {
    id,
    name: 'Paged MCP',
    description: '',
    enabled: true,
    allowDynamicTools: true,
    assignments: ['opencode' as const],
    secretConfigured: false,
    transport: 'http' as const,
    url
  }
}

async function startToolUpstream(
  listTools: (
    cursor: string | undefined
  ) =>
    | { tools: Tool[]; nextCursor?: string }
    | Promise<{ tools: Tool[]; nextCursor?: string }>
): Promise<{
  url: string
  notifyToolsChanged: () => Promise<void>
}> {
  const sessions = new Map<
    string,
    {
      protocol: McpProtocolServer
      transport: StreamableHTTPServerTransport
    }
  >()
  const server = createServer(async (request, response) => {
    let body: unknown
    if (request.method === 'POST') {
      const chunks: Buffer[] = []
      for await (const chunk of request) {
        chunks.push(Buffer.from(chunk))
      }
      body = JSON.parse(Buffer.concat(chunks).toString('utf8'))
    }
    const sessionId = request.headers['mcp-session-id']
    let session =
      typeof sessionId === 'string'
        ? sessions.get(sessionId)
        : undefined
    if (!session) {
      if (
        request.method !== 'POST' ||
        !isInitializeRequest(body)
      ) {
        response.writeHead(404)
        response.end()
        return
      }
      const protocol = new McpProtocolServer(
        { name: 'tool-upstream', version: '1.0.0' },
        { capabilities: { tools: { listChanged: true } } }
      )
      protocol.setRequestHandler(
        ListToolsRequestSchema,
        async (requestValue) =>
          listTools(requestValue.params?.cursor)
      )
      const transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => crypto.randomUUID(),
        onsessioninitialized: (idValue) => {
          sessions.set(idValue, { protocol, transport })
        },
        onsessionclosed: (idValue) => {
          sessions.delete(idValue)
        }
      })
      session = { protocol, transport }
      await protocol.connect(transport)
    }
    await session.transport.handleRequest(request, response, body)
  })
  httpServers.push(server)
  await new Promise<void>((resolve, reject) => {
    server.once('error', reject)
    server.listen(0, '127.0.0.1', resolve)
  })
  const address = server.address()
  if (!address || typeof address === 'string') {
    throw new Error('tool upstream did not bind')
  }
  return {
    url: `http://127.0.0.1:${address.port}/mcp`,
    notifyToolsChanged: async () => {
      await Promise.all(
        [...sessions.values()].map(({ protocol }) =>
          protocol.sendToolListChanged()
        )
      )
    }
  }
}

function testTool(name: string): Tool {
  return {
    name,
    description: name,
    inputSchema: { type: 'object' }
  }
}

const gateways: KnowledgeMcpGateway[] = []
const databases: AssistantDatabase[] = []
const temporaryDirectories: string[] = []
const httpServers: Server[] = []

afterEach(async () => {
  await Promise.all(gateways.splice(0).map((gateway) => gateway.dispose()))
  for (const database of databases.splice(0)) {
    database.close()
  }
  await Promise.all(
    temporaryDirectories
      .splice(0)
      .map((directory) => rm(directory, { recursive: true, force: true }))
  )
  await Promise.all(
    httpServers.splice(0).map(
      (server) =>
        new Promise<void>((resolve) => server.close(() => resolve()))
    )
  )
})

describe('KnowledgeMcpGateway', () => {
  it('exposes the assigned browser through a request-scoped conversation', async () => {
    const { service } = createService()
    const browserService = {
      getOrigin: vi.fn(() => undefined),
      navigate: vi.fn(async () => ({
        url: 'https://example.com/',
        origin: 'https://example.com'
      })),
      snapshot: vi.fn(),
      click: vi.fn(),
      type: vi.fn(),
      select: vi.fn(),
      back: vi.fn(),
      screenshot: vi.fn(),
      releaseConversation: vi.fn(async () => undefined)
    }
    const gateway = new KnowledgeMcpGateway(service, {
      browserService
    })
    gateways.push(gateway)
    await gateway.start()
    const token = gateway.grant(
      'browser-request',
      [],
      new AbortController().signal,
      'none',
      undefined,
      'browser-conversation'
    )!
    expect(gateway.getAvailableToolNames(token)).toEqual(browserToolNames)
    const client = new Client({
      name: 'browser-loopback-test',
      version: '1.0.0'
    })
    await client.connect(
      new StreamableHTTPClientTransport(
        new URL(gateway.getEndpoint()!),
        {
          requestInit: {
            headers: { Authorization: `Bearer ${token}` }
          }
        }
      )
    )
    try {
      const listed = await client.listTools()
      expect(listed.tools.map((tool) => tool.name)).toEqual(
        browserToolNames
      )
      await expect(
        client.callTool({
          name: 'browser_navigate',
          arguments: { url: 'https://example.com' }
        })
      ).resolves.toMatchObject({
        content: [
          {
            type: 'text',
            text: JSON.stringify({
              url: 'https://example.com/',
              origin: 'https://example.com'
            })
          }
        ]
      })
      expect(browserService.navigate).toHaveBeenCalledWith(
        'browser-conversation',
        'https://example.com',
        expect.any(AbortSignal)
      )
    } finally {
      await client.close()
    }
  })

  it('exposes GoodBuddy config reads in Ask and apply only in Execute', async () => {
    const { service } = createService()
    const configService = {
      getCapabilities: vi.fn(() => ({ server: 'goodbuddy_config' })),
      getSnapshot: vi.fn(async () => ({ application: {}, skills: [], mcpServers: [] })),
      plan: vi.fn(async () => ({ planId: 'plan' })),
      apply: vi.fn(async () => ({ status: 'applied' })),
      revokeRequest: vi.fn()
    }
    const gateway = new KnowledgeMcpGateway(service, {
      configService: configService as never
    })
    gateways.push(gateway)
    const readToken = gateway.grant(
      'config-read',
      [],
      new AbortController().signal,
      'none',
      { access: 'read', workspacePath: process.cwd() }
    )!
    const authorizeApply = vi.fn(async () => true)
    const writeToken = gateway.grant(
      'config-write',
      [],
      new AbortController().signal,
      'none',
      {
        access: 'write',
        workspacePath: process.cwd(),
        authorizeApply
      }
    )!

    expect(gateway.getAvailableToolNames(readToken)).toEqual([
      'goodbuddy_config_capabilities',
      'goodbuddy_config_get',
      'goodbuddy_config_plan'
    ])
    expect(gateway.getAvailableToolNames(writeToken)).toEqual([
      'goodbuddy_config_capabilities',
      'goodbuddy_config_get',
      'goodbuddy_config_plan',
      'goodbuddy_config_apply'
    ])
    await gateway.callGoodBuddyConfigTool(
      readToken,
      'goodbuddy_config_capabilities',
      {}
    )
    expect(configService.getCapabilities).toHaveBeenCalledWith({})
    await expect(
      gateway.callGoodBuddyConfigTool(
        readToken,
        'goodbuddy_config_apply',
        { planId: crypto.randomUUID() }
      )
    ).rejects.toThrow('unavailable')
    await gateway.callGoodBuddyConfigTool(
      writeToken,
      'goodbuddy_config_apply',
      { planId: crypto.randomUUID() }
    )
    expect(configService.apply).toHaveBeenCalledWith(
      'config-write',
      expect.any(Object),
      expect.any(AbortSignal),
      authorizeApply
    )
    gateway.revoke(writeToken)
    expect(configService.revokeRequest).toHaveBeenCalledWith('config-write')
  })

  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)
    expect(gateway.getAvailableToolNames(token!)).toEqual([
      'knowledge_list',
      'knowledge_search'
    ])
    expect(gateway.listLibraries(token!)).toEqual([
      {
        id: secondLibraryId,
        name: '二号知识库',
        description: '已授权知识'
      }
    ])
    expect(() =>
      gateway.listLibraries(token!, {
        libraryIds: [firstLibraryId]
      })
    ).toThrow()
    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: '二号知识库',
        chunkId: '44444444-4444-4444-8444-444444444440',
        score: 0.5,
        snippet: '匹配 0'
      })
    ])
    expect(references[0]?.sourceLocation).toBeUndefined()
    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('grants bounded global Magic Notes search without a knowledge scope', () => {
    const { service } = createService()
    const searchMagicNotes = vi.fn(() => [
      {
        noteId: '00000000-0000-4000-8000-000000000701',
        noteTitle: '发布计划',
        entryId: '00000000-0000-4000-8000-000000000702',
        content: '核对构建产物',
        updatedAt: '2026-08-10T00:00:00.000Z'
      }
    ])
    const gateway = new KnowledgeMcpGateway(service, {
      magicNotesDatabase: {
        listMagicNotes: vi.fn(() => []),
        getMagicNote: vi.fn(() => {
          throw new Error('not used')
        }),
        getMagicNoteEntry: vi.fn(() => {
          throw new Error('not used')
        }),
        searchMagicNotes,
        createMagicNote: vi.fn(() => {
          throw new Error('not used')
        }),
        updateMagicNote: vi.fn(() => {
          throw new Error('not used')
        }),
        deleteMagicNote: vi.fn(),
        createMagicNoteEntry: vi.fn(() => {
          throw new Error('not used')
        }),
        updateMagicNoteEntry: vi.fn(() => {
          throw new Error('not used')
        }),
        deleteMagicNoteEntry: vi.fn(() => {
          throw new Error('not used')
        })
      } satisfies MagicNotesDatabase
    })
    gateways.push(gateway)
    const token = gateway.grant(
      'notes',
      [],
      new AbortController().signal,
      'read'
    )!

    expect(gateway.getAvailableToolNames(token)).toEqual([
      'note_list',
      'note_get',
      'note_search'
    ])
    expect(
      gateway.searchMagicNotes(token, {
        query: '  发布  ',
        limit: 3
      })
    ).toEqual([
      expect.objectContaining({
        noteTitle: '发布计划',
        content: '核对构建产物'
      })
    ])
    expect(searchMagicNotes).toHaveBeenCalledWith('发布', 3)
    expect(() =>
      gateway.searchMagicNotes(token, {
        query: '发布',
        noteIds: ['not-allowed']
      })
    ).toThrow()
  })

  it('keeps Ask read-only and supports revision-safe Magic Notes CRUD in Execute', async () => {
    const { service } = createService()
    const directory = await mkdtemp(join(tmpdir(), 'goodbuddy-note-mcp-'))
    temporaryDirectories.push(directory)
    const database = new AssistantDatabase(
      join(directory, 'assistant.sqlite')
    )
    databases.push(database)
    database.initialize('C:\\Workspace')
    const gateway = new KnowledgeMcpGateway(service, {
      magicNotesDatabase: database
    })
    gateways.push(gateway)
    const readToken = gateway.grant(
      'notes-read',
      [],
      new AbortController().signal,
      'read'
    )!
    const writeToken = gateway.grant(
      'notes-write',
      [],
      new AbortController().signal,
      'write'
    )!

    expect(gateway.getAvailableToolNames(readToken)).toEqual([
      'note_list',
      'note_get',
      'note_search'
    ])
    expect(gateway.getAvailableToolNames(writeToken)).toEqual([
      'note_list',
      'note_get',
      'note_search',
      'note_create',
      'note_update',
      'note_entry_create',
      'note_entry_update',
      'note_entry_delete',
      'note_delete'
    ])
    expect(() =>
      gateway.createMagicNote(readToken, { title: '不允许创建' })
    ).toThrow('unavailable')

    const created = gateway.createMagicNote(writeToken, {
      title: '发布计划',
      content: '核对构建产物'
    })
    expect(gateway.listMagicNotes(readToken)).toEqual([
      expect.objectContaining({
        id: created.id,
        title: '发布计划',
        revision: 1,
        entryCount: 1
      })
    ])
    expect(created.entries[0]?.content).toBe('核对构建产物')
    const withEntry = gateway.createMagicNoteEntry(writeToken, {
      noteId: created.id,
      content: '通知发布负责人'
    })
    const entry = withEntry.entries[1]!
    expect(entry.content).toBe('通知发布负责人')

    const updatedEntry = gateway.updateMagicNoteEntry(writeToken, {
      entryId: entry.id,
      content: '核对六个平台构建产物',
      expectedRevision: entry.revision
    })
    expect(updatedEntry.entries[1]?.content).toBe(
      '核对六个平台构建产物'
    )
    expect(() =>
      gateway.deleteMagicNoteEntry(writeToken, {
        entryId: entry.id,
        expectedRevision: entry.revision
      })
    ).toThrow('已被更新')

    const withoutEntry = gateway.deleteMagicNoteEntry(writeToken, {
      entryId: entry.id,
      expectedRevision: updatedEntry.entries[1]!.revision
    })
    expect(withoutEntry.entries).toEqual([
      expect.objectContaining({ content: '核对构建产物' })
    ])
    expect(
      gateway.deleteMagicNote(writeToken, {
        noteId: created.id,
        expectedRevision: withoutEntry.revision
      })
    ).toEqual({ deleted: true, noteId: created.id })
    expect(() =>
      gateway.getMagicNote(readToken, { noteId: created.id })
    ).toThrow('笔记不存在')
  })

  it('binds an authenticated MCP 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(401)
    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)
  })

  it('proxies custom MCP through a request-scoped loopback token without exposing the upstream credential', async () => {
    const upstreamAuthorizations: Array<string | undefined> = []
    const upstream = createServer(async (request, response) => {
      upstreamAuthorizations.push(request.headers.authorization)
      if (request.method !== 'POST') {
        response.writeHead(405)
        response.end()
        return
      }
      const chunks: Buffer[] = []
      for await (const chunk of request) {
        chunks.push(Buffer.from(chunk))
      }
      const body = JSON.parse(
        Buffer.concat(chunks).toString('utf8')
      ) as unknown
      const mcp = new McpServer({
        name: 'private-upstream',
        version: '1.0.0'
      })
      mcp.registerTool(
        'echo_private',
        {
          description: 'Echo through the private server',
          inputSchema: {
            value: z.string().max(100)
          }
        },
        async ({ value }) => ({
          content: [{ type: 'text', text: `upstream:${value}` }]
        })
      )
      const transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: undefined
      })
      await mcp.connect(transport)
      await transport.handleRequest(request, response, body)
      await Promise.allSettled([transport.close(), mcp.close()])
    })
    httpServers.push(upstream)
    await new Promise<void>((resolve, reject) => {
      upstream.once('error', reject)
      upstream.listen(0, '127.0.0.1', resolve)
    })
    const address = upstream.address()
    if (!address || typeof address === 'string') {
      throw new Error('upstream did not bind')
    }

    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    await gateway.start()
    const controller = new AbortController()
    const token = gateway.grantCustomMcp(
      'custom-request',
      [
        {
          id: '00000000-0000-4000-8000-000000000091',
          name: 'Private MCP',
          description: '',
          enabled: true,
          allowDynamicTools: true,
          assignments: ['opencode'],
          secretConfigured: true,
          secret: 'upstream-secret',
          transport: 'http',
          url: `http://127.0.0.1:${address.port}/mcp`
        }
      ],
      controller.signal
    )!
    const client = new Client({
      name: 'loopback-test-client',
      version: '1.0.0'
    })
    await client.connect(
      new StreamableHTTPClientTransport(
        new URL(gateway.getEndpoint()!),
        {
          requestInit: {
            headers: {
              Authorization: `Bearer ${token}`
            }
          }
        }
      )
    )
    try {
      const listed = await client.listTools()
      expect(listed.tools).toEqual([
        expect.objectContaining({
          name: expect.stringMatching(
            /^mcp_[a-f0-9]{8}_[a-f0-9]{8}_echo_private$/u
          ),
          description: expect.stringContaining('Private MCP')
        })
      ])
      expect(JSON.stringify(listed)).not.toContain('upstream-secret')
      expect(JSON.stringify(listed)).not.toContain(
        `127.0.0.1:${address.port}`
      )
      const result = await client.callTool({
        name: listed.tools[0]!.name,
        arguments: { value: 'hello' }
      })
      expect(result).toMatchObject({
        content: [{ type: 'text', text: 'upstream:hello' }]
      })
      expect(upstreamAuthorizations).toContain(
        'Bearer upstream-secret'
      )
    } finally {
      gateway.revoke(token)
      await client.close()
    }
  })

  it('loads every tools/list page before exposing custom MCP tools', async () => {
    const listTools = vi.fn((cursor: string | undefined) =>
      cursor !== undefined
        ? { tools: [testTool('second')] }
        : {
            tools: [testTool('first')],
            nextCursor: 'page-2'
          }
    )
    const upstream = await startToolUpstream(listTools)
    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    const token = gateway.grantCustomMcp(
      'paged-tools',
      [customMcpServer(upstream.url)],
      new AbortController().signal
    )!

    const tools = await gateway.prepareCustomMcpTools(token)

    expect(tools.map((tool) => tool.name)).toEqual([
      expect.stringMatching(/_first$/u),
      expect.stringMatching(/_second$/u)
    ])
    expect(listTools).toHaveBeenNthCalledWith(1, undefined)
    expect(listTools).toHaveBeenNthCalledWith(2, 'page-2')
  })

  it('continues tools/list pagination with an empty cursor', async () => {
    const listTools = vi.fn((cursor: string | undefined) =>
      cursor === undefined
        ? {
            tools: [testTool('first')],
            nextCursor: ''
          }
        : { tools: [testTool('second')] }
    )
    const upstream = await startToolUpstream(listTools)
    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    const token = gateway.grantCustomMcp(
      'empty-cursor-tools',
      [customMcpServer(upstream.url)],
      new AbortController().signal
    )!

    const tools = await gateway.prepareCustomMcpTools(token)

    expect(tools.map((tool) => tool.name)).toEqual([
      expect.stringMatching(/_first$/u),
      expect.stringMatching(/_second$/u)
    ])
    expect(listTools).toHaveBeenNthCalledWith(2, '')
  })

  it('explicitly requests task execution for required tools from earlier pages', async () => {
    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    const server = customMcpServer('http://127.0.0.1:1/mcp')
    const token = gateway.grantCustomMcp(
      'required-task-tool',
      [server],
      new AbortController().signal
    )!
    const callToolStream = vi.fn(
      async function* () {
        yield {
          type: 'result' as const,
          result: {
            content: [{ type: 'text' as const, text: 'done' }],
            structuredContent: { value: 'done' }
          }
        }
      }
    )
    const outputValidator = vi.fn(() => ({
      valid: true as const,
      data: { value: 'done' },
      errorMessage: undefined
    }))
    const callRequiredTool = (
      gateway as unknown as {
        callCustomMcpTool(
          capabilityToken: string,
          binding: unknown,
          input: Record<string, unknown>,
          signal: AbortSignal
        ): Promise<unknown>
      }
    ).callCustomMcpTool.bind(gateway)

    await expect(
      callRequiredTool(
        token,
        {
          client: {
            experimental: {
              tasks: {
                callToolStream,
                cancelTask: vi.fn()
              }
            }
          },
          server,
          originalName: 'required-first-page',
          taskSupport: 'required',
          outputValidator,
          exposedTool: testTool('required-first-page')
        },
        {},
        new AbortController().signal
      )
    ).resolves.toMatchObject({
      content: [{ type: 'text', text: 'done' }]
    })
    expect(callToolStream).toHaveBeenCalledWith(
      {
        name: 'required-first-page',
        arguments: {}
      },
      undefined,
      expect.objectContaining({ task: {} })
    )
    expect(outputValidator).toHaveBeenCalledWith({ value: 'done' })
  })

  it('rejects cyclic cursors and custom MCP tool counts over 100', async () => {
    const cyclicUpstream = await startToolUpstream(
      (cursor: string | undefined) => ({
        tools: [testTool(cursor ? 'second' : 'first')],
        nextCursor: 'cycle'
      })
    )
    const excessiveUpstream = await startToolUpstream(() => ({
      tools: Array.from({ length: 101 }, (_, index) =>
        testTool(`tool-${index}`)
      )
    }))
    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    const cycleToken = gateway.grantCustomMcp(
      'cursor-cycle',
      [customMcpServer(cyclicUpstream.url)],
      new AbortController().signal
    )!
    const excessiveToken = gateway.grantCustomMcp(
      'excessive-tools',
      [
        customMcpServer(
          excessiveUpstream.url,
          '00000000-0000-4000-8000-000000000093'
        )
      ],
      new AbortController().signal
    )!

    await expect(
      gateway.prepareCustomMcpTools(cycleToken)
    ).rejects.toMatchObject({
      cause: expect.objectContaining({
        message: expect.stringContaining('分页游标发生循环')
      })
    })
    await expect(
      gateway.prepareCustomMcpTools(excessiveToken)
    ).rejects.toMatchObject({
      cause: expect.objectContaining({
        message: expect.stringContaining('工具数量超过安全限制')
      })
    })
  })

  it('releases rejected initialize attempts before enforcing the session limit', async () => {
    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    await gateway.start()
    const endpoint = gateway.getEndpoint()!
    const token = gateway.grant(
      'initialize-retry',
      [firstLibraryId],
      new AbortController().signal
    )!
    const initialize = {
      jsonrpc: '2.0',
      id: 1,
      method: 'initialize',
      params: {
        protocolVersion: LATEST_PROTOCOL_VERSION,
        capabilities: {},
        clientInfo: {
          name: 'initialize-retry-fixture',
          version: '1.0.0'
        }
      }
    }

    for (let attempt = 0; attempt < 8; attempt += 1) {
      const rejected = await fetch(endpoint, {
        method: 'POST',
        headers: {
          accept: 'application/json',
          authorization: `Bearer ${token}`,
          'content-type': 'application/json'
        },
        body: JSON.stringify(initialize)
      })
      expect(rejected.status).toBe(406)
    }

    const accepted = await fetch(endpoint, {
      method: 'POST',
      headers: {
        accept: 'application/json, text/event-stream',
        authorization: `Bearer ${token}`,
        'content-type': 'application/json'
      },
      body: JSON.stringify(initialize)
    })
    expect(accepted.status).toBe(200)
    expect(accepted.headers.get('mcp-session-id')).toEqual(
      expect.any(String)
    )
  })

  it('publishes upstream tool changes downstream after a successful refresh', async () => {
    let tools = [testTool('before')]
    const listTools = vi.fn(async () => ({ tools }))
    const upstream = await startToolUpstream(listTools)
    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    await gateway.start()
    const token = gateway.grantCustomMcp(
      'dynamic-tools',
      [customMcpServer(upstream.url)],
      new AbortController().signal
    )!
    const listChanged = vi.fn()
    const client = new Client(
      { name: 'dynamic-client', version: '1.0.0' },
      {
        listChanged: {
          tools: {
            autoRefresh: false,
            debounceMs: 0,
            onChanged: listChanged
          }
        }
      }
    )
    await client.connect(
      new StreamableHTTPClientTransport(
        new URL(gateway.getEndpoint()!),
        {
          requestInit: {
            headers: { Authorization: `Bearer ${token}` }
          }
        }
      )
    )
    try {
      const initial = await client.listTools()
      expect(initial.tools[0]?.name).toMatch(/_before$/u)
      await new Promise((resolve) => setTimeout(resolve, 25))
      tools = [testTool('after')]

      await upstream.notifyToolsChanged()

      await vi.waitFor(() => {
        expect(listChanged).toHaveBeenCalledWith(null, null)
      })
      const updated = await client.listTools()
      expect(updated.tools.map((tool) => tool.name)).toEqual([
        expect.stringMatching(/_after$/u)
      ])
    } finally {
      await client.close()
    }
  })

  it('does not publish a downstream change when dynamic refresh fails', async () => {
    let failRefresh = false
    const listTools = vi.fn(async () => {
      if (failRefresh) {
        throw new Error('refresh failed')
      }
      return { tools: [testTool('stable')] }
    })
    const upstream = await startToolUpstream(listTools)
    const { service } = createService()
    const gateway = new KnowledgeMcpGateway(service)
    gateways.push(gateway)
    await gateway.start()
    const token = gateway.grantCustomMcp(
      'failed-refresh',
      [customMcpServer(upstream.url)],
      new AbortController().signal
    )!
    const listChanged = vi.fn()
    const client = new Client(
      { name: 'failed-refresh-client', version: '1.0.0' },
      {
        listChanged: {
          tools: {
            autoRefresh: false,
            debounceMs: 0,
            onChanged: listChanged
          }
        }
      }
    )
    await client.connect(
      new StreamableHTTPClientTransport(
        new URL(gateway.getEndpoint()!),
        {
          requestInit: {
            headers: { Authorization: `Bearer ${token}` }
          }
        }
      )
    )
    try {
      await client.listTools()
      await new Promise((resolve) => setTimeout(resolve, 25))
      failRefresh = true

      await upstream.notifyToolsChanged()

      await vi.waitFor(() => {
        expect(listTools).toHaveBeenCalledTimes(2)
      })
      await new Promise((resolve) => setTimeout(resolve, 25))
      expect(listChanged).not.toHaveBeenCalled()
    } finally {
      await client.close()
    }
  })
})