open-multi-agent/tests/mcp-tools.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest'
import type { ToolUseContext } from '../src/types.js'
import { ToolRegistry } from '../src/tool/framework.js'

const listToolsMock = vi.fn()
const callToolMock = vi.fn()
const connectMock = vi.fn()
const clientCloseMock = vi.fn()
const transportCloseMock = vi.fn()

class MockClient {
  async connect(
    transport: unknown,
    _options?: { timeout?: number },
  ): Promise<void> {
    connectMock(transport)
  }

  async listTools(
    params?: { cursor?: string },
    options?: { timeout?: number },
  ): Promise<{
    tools: Array<{
      name: string
      description: string
      inputSchema?: Record<string, unknown>
    }>
    nextCursor?: string
  }> {
    return listToolsMock(params, options)
  }

  async callTool(
    request: { name: string; arguments: Record<string, unknown> },
    resultSchema?: unknown,
    options?: { timeout?: number },
  ): Promise<{
    content?: Array<Record<string, unknown>>
    structuredContent?: unknown
    isError?: boolean
    toolResult?: unknown
  }> {
    return callToolMock(request, resultSchema, options)
  }

  async close(): Promise<void> {
    clientCloseMock()
  }
}

class MockStdioTransport {
  readonly config: unknown

  constructor(config: unknown) {
    this.config = config
  }

  async close(): Promise<void> {
    transportCloseMock()
  }
}

vi.mock('@modelcontextprotocol/sdk/client/index.js', () => ({
  Client: MockClient,
}))

vi.mock('@modelcontextprotocol/sdk/client/stdio.js', () => ({
  StdioClientTransport: MockStdioTransport,
}))

const context: ToolUseContext = {
  agent: { name: 'test-agent', role: 'tester', model: 'test-model' },
}

beforeEach(() => {
  vi.clearAllMocks()
})

describe('connectMCPTools', () => {
  it('connects, discovers tools, and executes MCP calls', async () => {
    listToolsMock.mockResolvedValue({
      tools: [
        {
          name: 'search_issues',
          description: 'Search repository issues.',
          inputSchema: {
            type: 'object',
            properties: { q: { type: 'string' } },
            required: ['q'],
          },
        },
      ],
    })
    callToolMock.mockResolvedValue({
      content: [{ type: 'text', text: 'found 2 issues' }],
      isError: false,
    })

    const { connectMCPTools } = await import('../src/tool/mcp.js')
    const connected = await connectMCPTools({
      command: 'npx',
      args: ['-y', 'mock-mcp-server'],
      env: { GITHUB_TOKEN: 'token' },
      namePrefix: 'github',
    })

    expect(connectMock).toHaveBeenCalledTimes(1)
    expect(connected.tools).toHaveLength(1)
    expect(connected.tools[0].name).toBe('github_search_issues')

    const registry = new ToolRegistry()
    registry.register(connected.tools[0])
    const defs = registry.toToolDefs()
    expect(defs[0].inputSchema).toMatchObject({
      type: 'object',
      properties: { q: { type: 'string' } },
      required: ['q'],
    })

    const result = await connected.tools[0].execute({ q: 'bug' }, context)
    expect(callToolMock).toHaveBeenCalledWith(
      {
        name: 'search_issues',
        arguments: { q: 'bug' },
      },
      undefined,
      expect.objectContaining({ timeout: expect.any(Number) }),
    )
    expect(result.isError).toBe(false)
    expect(result.data).toContain('found 2 issues')

    await connected.disconnect()
    expect(clientCloseMock).toHaveBeenCalledTimes(1)
    expect(transportCloseMock).not.toHaveBeenCalled()
  })

  it('aggregates paginated listTools results', async () => {
    listToolsMock.mockImplementation(
      async (params?: { cursor?: string }) => {
        if (params?.cursor === 'c1') {
          return {
            tools: [
              { name: 'b', description: 'B', inputSchema: { type: 'object' } },
            ],
          }
        }
        return {
          tools: [
            { name: 'a', description: 'A', inputSchema: { type: 'object' } },
          ],
          nextCursor: 'c1',
        }
      },
    )

    callToolMock.mockResolvedValue({ content: [{ type: 'text', text: 'ok' }] })

    const { connectMCPTools } = await import('../src/tool/mcp.js')
    const connected = await connectMCPTools({
      command: 'npx',
      args: ['-y', 'mock-mcp-server'],
    })

    expect(listToolsMock).toHaveBeenCalledTimes(2)
    expect(listToolsMock.mock.calls[1][0]).toEqual({ cursor: 'c1' })
    expect(connected.tools).toHaveLength(2)
    expect(connected.tools.map((t) => t.name)).toEqual(['a', 'b'])
  })

  it('serializes non-text MCP content blocks', async () => {
    listToolsMock.mockResolvedValue({
      tools: [{ name: 'pic', description: 'Pic', inputSchema: { type: 'object' } }],
    })
    callToolMock.mockResolvedValue({
      content: [
        {
          type: 'image',
          data: 'AAA',
          mimeType: 'image/png',
        },
      ],
      isError: false,
    })

    const { connectMCPTools } = await import('../src/tool/mcp.js')
    const connected = await connectMCPTools({ command: 'npx', args: ['x'] })
    const result = await connected.tools[0].execute({}, context)
    expect(result.data).toContain('image')
    expect(result.data).toContain('base64 length=3')
  })

  it('marks tool result as error when MCP returns isError', async () => {
    listToolsMock.mockResolvedValue({
      tools: [{ name: 'danger', description: 'Dangerous op.', inputSchema: {} }],
    })
    callToolMock.mockResolvedValue({
      content: [{ type: 'text', text: 'permission denied' }],
      isError: true,
    })

    const { connectMCPTools } = await import('../src/tool/mcp.js')
    const connected = await connectMCPTools({
      command: 'npx',
      args: ['-y', 'mock-mcp-server'],
    })

    const result = await connected.tools[0].execute({}, context)
    expect(result.isError).toBe(true)
    expect(result.data).toContain('permission denied')
  })
})