openpencil/tests/engine/mcp/server/index.test.ts

/* eslint-disable max-lines -- MCP server scenarios share lifecycle and transport fixtures */

import { describe, expect, test, beforeEach, afterEach } from 'bun:test'
import { mkdir, stat } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'

import { Client, StreamableHTTPClientTransport } from '@modelcontextprotocol/client'

import { SceneGraph } from '@open-pencil/scene-graph'

import { startServer, type ServerHandle } from '#mcp/server'
import { createToolDescriptors, getMCPToolDefinitions } from '#mcp/tool/manifest'
import { parseDiscoveryInfo } from '#mcp/transport/discovery'

import { expectDefined } from '#tests/helpers/assert'
import {
  connectMockBrowser,
  readHealth,
  waitForBrowserRegistration,
  type MockBrowser
} from '#tests/helpers/mcp/server'

const isUnix = process.platform !== 'win32'
const SOCKET_DIR = join(tmpdir(), `openpencil-test-server-${process.pid}`)
const TEST_MCP_ROOT = join(tmpdir(), 'open-pencil-mcp-root')
const TEST_AUTH_TOKEN = 'test-auth-token'
const TEST_CLIENT_AUTH_TOKEN = 'test-client-token'
let testCounter = 0

// ---------------------------------------------------------------------------
// Test client: starts server with ephemeral TCP port, connects mock browser + MCP client
// ---------------------------------------------------------------------------

function testSocketPath(): string | null {
  if (!isUnix) return null
  return join(SOCKET_DIR, `mcp-test-${process.pid}-${++testCounter}.sock`)
}

async function createTestClient(disabledTools: string[] = []) {
  if (isUnix) await mkdir(SOCKET_DIR, { recursive: true })
  const authToken = TEST_CLIENT_AUTH_TOKEN
  const handle = await startServer({
    httpPort: 0,
    withTcp: true,
    socketPath: testSocketPath(),
    authToken,
    disabledTools,
    enableEval: false,
    mcpRoot: null
  })

  const httpPort = handle.httpPort
  if (!httpPort) {
    await handle.close()
    throw new Error('TCP listener not started')
  }

  const graph = new SceneGraph()
  let browser: MockBrowser | undefined
  let client: Client | undefined

  try {
    browser = await connectMockBrowser(httpPort, graph, authToken)
    await waitForBrowserRegistration(httpPort)

    client = new Client({ name: 'test-client', version: '0.0.0' })
    const transport = new StreamableHTTPClientTransport(
      new URL(`http://127.0.0.1:${httpPort}/mcp`),
      {
        requestInit: { headers: { Authorization: `Bearer ${authToken}` } }
      }
    )
    await client.connect(transport)
  } catch (e) {
    await client?.close().catch(() => undefined)
    browser?.close()
    await handle.close()
    throw e
  }

  const safeClient = client
  const safeBrowser = browser
  return {
    client: safeClient,
    graph,
    handle,
    close: async () => {
      const errors: unknown[] = []
      try {
        await safeClient?.close()
      } catch (e) {
        errors.push(e)
      }
      try {
        safeBrowser?.close()
      } catch (e) {
        errors.push(e)
      }
      try {
        await handle.close()
      } catch (e) {
        errors.push(e)
      }
      if (errors.length > 0) throw errors[0]
    }
  }
}

function parseResult(result: { content: { type: string; text?: string }[] }): unknown {
  const textContent = result.content.find((c) => c.type === 'text')
  return textContent?.text ? JSON.parse(textContent.text) : null
}

// ---------------------------------------------------------------------------
// MCP tool + session tests
// ---------------------------------------------------------------------------

test('MCP exposure exclusions apply to both descriptors and registered tools', async () => {
  const def = getMCPToolDefinitions().find((tool) => tool.name === 'get_page_tree')
  if (!def) throw new Error('Missing tool fixture')
  const previous = def.exposure
  let context: Awaited<ReturnType<typeof createTestClient>> | undefined
  try {
    def.exposure = { ...previous, mcp: false }
    expect(createToolDescriptors(false).some((tool) => tool.name === def.name)).toBe(false)
    context = await createTestClient()
    const { tools } = await context.client.listTools()
    expect(tools.some((tool) => tool.name === def.name)).toBe(false)
  } finally {
    def.exposure = previous
    await context?.close()
  }
})

describe('startServer option validation', () => {
  test('rejects an app wait timeout that is not a safe non-negative integer', async () => {
    await expect(startServer({ appWaitTimeoutMs: -1 })).rejects.toThrow(RangeError)
    await expect(startServer({ appWaitTimeoutMs: 1.5 })).rejects.toThrow(
      'appWaitTimeoutMs must be a safe integer'
    )
  })
})

describe('MCP server', () => {
  let client: Client
  let graph: SceneGraph
  let cleanup: (() => Promise<void>) | null = null

  beforeEach(async () => {
    try {
      const ctx = await createTestClient()
      client = ctx.client
      graph = ctx.graph
      cleanup = ctx.close
    } catch (e) {
      if (cleanup) await cleanup().catch(() => undefined)
      throw e
    }
  })

  afterEach(async () => {
    if (cleanup) await cleanup()
    cleanup = null
  })

  test('lists all registered tools', async () => {
    const { tools } = await client.listTools()
    const expectedNames = createToolDescriptors(false)
      .filter((tool) => tool.availability === 'default')
      .map((tool) => tool.name)
      .sort()
    expect(tools.map((tool) => tool.name).sort()).toEqual(expectedNames)
  })

  test('describes effects, capabilities, and runtime availability independently', () => {
    const descriptors = createToolDescriptors(true)
    const byName = new Map(descriptors.map((tool) => [tool.name, tool] as const))
    expect(byName.get('get_page_tree')?.effect).toBe('read')
    expect(byName.get('switch_page')?.effect).toBe('read')
    expect(byName.get('viewport_set')?.effect).toBe('read')
    expect(byName.get('export_image')?.effect).toBe('read')
    expect(byName.get('save_file')?.effect).toBe('write')
    expect(byName.get('open_file')?.effect).toBe('read')
    expect(byName.get('open_file')?.capabilities).toEqual(['filesystem:read', 'document:read'])
    expect(byName.get('close_file')?.effect).toBe('write')
    expect(byName.get('update_node')?.effect).toBe('write')
    expect(byName.get('new_document')?.capabilities).toEqual(['document:write', 'filesystem:write'])
    expect(byName.get('activate_document')?.effect).toBe('read')
    expect(byName.get('undo')?.capabilities).toEqual(['document:write'])
    expect(byName.get('redo')?.effect).toBe('write')
    expect(byName.get('get_settings')?.capabilities).toEqual(['settings:read'])
    expect(byName.get('update_settings')?.capabilities).toEqual(['settings:write'])
    expect(byName.get('eval')?.availability).toBe('eval')
    expect(byName.get('eval')?.capabilities).toContain('code:execute')
  })

  test('omits tools disabled in the generic registration filter', async () => {
    if (cleanup) await cleanup()
    cleanup = null

    const ctx = await createTestClient(['create_shape', 'list_documents'])
    client = ctx.client
    graph = ctx.graph
    cleanup = ctx.close

    const { tools } = await client.listTools()
    const names = tools.map((tool) => tool.name)
    expect(names).not.toContain('create_shape')
    expect(names).not.toContain('list_documents')
    expect(names).toContain('get_page_tree')

    const healthResponse = await fetch(`http://127.0.0.1:${ctx.handle.httpPort}/health`, {
      headers: { Authorization: `Bearer ${TEST_CLIENT_AUTH_TOKEN}` }
    })
    const health = await readHealth(healthResponse)
    const descriptors = health.tools ?? []
    expect(descriptors.find((tool) => tool.name === 'create_shape')?.enabled).toBe(false)
    expect(descriptors.find((tool) => tool.name === 'list_documents')?.enabled).toBe(false)
    expect(descriptors.find((tool) => tool.name === 'get_page_tree')?.enabled).toBe(true)
  })

  test('tools expose standard MCP effect annotations', async () => {
    const { tools } = await client.listTools()
    const byName = new Map(tools.map((tool) => [tool.name, tool] as const))
    expect(byName.get('get_page_tree')?.annotations?.readOnlyHint).toBe(true)
    expect(byName.get('get_page_tree')?.annotations?.destructiveHint).toBe(false)
    expect(byName.get('update_node')?.annotations?.readOnlyHint).toBe(false)
    expect(byName.get('update_node')?.annotations?.destructiveHint).toBe(true)
  })

  test('tools have descriptions and input schemas', async () => {
    const { tools } = await client.listTools()
    for (const tool of tools) {
      expect(tool.description).toBeTruthy()
      expect(tool.inputSchema).toBeDefined()
    }
  })

  test('create_shape creates a node on the live canvas', async () => {
    const result = await client.callTool({
      name: 'create_shape',
      arguments: { type: 'FRAME', x: 0, y: 0, width: 200, height: 100, name: 'Test' }
    })
    expect(result.isError).not.toBe(true)
    const data = parseResult(result) as { id: string; name: string; type: string }
    expect(data.type).toBe('FRAME')
    expect(data.name).toBe('Test')

    const node = graph.getNode(data.id)
    expect(node).toBeDefined()
    expect(node?.name).toBe('Test')
  })

  test('set_fill validates and applies color', async () => {
    const create = await client.callTool({
      name: 'create_shape',
      arguments: { type: 'RECTANGLE', x: 0, y: 0, width: 50, height: 50 }
    })
    const { id } = parseResult(create) as { id: string }

    const fill = await client.callTool({
      name: 'set_fill',
      arguments: { id, color: '#00ff00' }
    })
    expect(fill.isError).not.toBe(true)
  })

  test('get_page_tree returns page structure', async () => {
    await client.callTool({
      name: 'create_shape',
      arguments: { type: 'FRAME', x: 0, y: 0, width: 100, height: 100, name: 'F1' }
    })
    const result = await client.callTool({ name: 'get_page_tree', arguments: {} })
    expect(result.isError).not.toBe(true)
    const data = parseResult(result) as { children: { name: string }[] }
    expect(data.children.some((c) => c.name === 'F1')).toBe(true)
  })

  test('delete_node removes a node', async () => {
    const create = await client.callTool({
      name: 'create_shape',
      arguments: { type: 'RECTANGLE', x: 0, y: 0, width: 50, height: 50 }
    })
    const { id } = parseResult(create) as { id: string }

    await client.callTool({ name: 'delete_node', arguments: { id } })

    const get = await client.callTool({ name: 'get_node', arguments: { id } })
    const data = parseResult(get) as { error?: string }
    expect(data.error).toContain('not found')
  })

  test('find_nodes filters by type', async () => {
    await client.callTool({
      name: 'create_shape',
      arguments: { type: 'FRAME', x: 0, y: 0, width: 100, height: 100 }
    })
    await client.callTool({
      name: 'create_shape',
      arguments: { type: 'RECTANGLE', x: 0, y: 0, width: 50, height: 50 }
    })
    await client.callTool({
      name: 'create_shape',
      arguments: { type: 'FRAME', x: 0, y: 0, width: 100, height: 100 }
    })
    const result = await client.callTool({ name: 'find_nodes', arguments: { type: 'FRAME' } })
    const data = parseResult(result) as { count: number }
    expect(data.count).toBe(2)
  })

  test('get_codegen_prompt returns prompt text', async () => {
    const result = await client.callTool({ name: 'get_codegen_prompt', arguments: {} })
    expect(result.isError).not.toBe(true)
    const data = parseResult(result) as { prompt: string }
    expect(data.prompt.length).toBeGreaterThan(100)
  })
})

// ---------------------------------------------------------------------------
// mcpRoot tests
// ---------------------------------------------------------------------------

describe('MCP server with mcpRoot', () => {
  async function withMCPRootServer(
    mcpRoot: string | null,
    fn: (client: Client, browser: MockBrowser, graph: SceneGraph) => Promise<void>
  ) {
    if (isUnix) await mkdir(SOCKET_DIR, { recursive: true })
    if (mcpRoot) await mkdir(mcpRoot, { recursive: true })
    const handle = await startServer({
      httpPort: 0,
      withTcp: true,
      socketPath: testSocketPath(),
      authToken: TEST_AUTH_TOKEN,
      enableEval: false,
      mcpRoot
    })

    let browser: MockBrowser | null = null
    let client: Client | null = null

    try {
      const httpPort = handle.httpPort
      if (!httpPort) throw new Error('withTcp: true did not produce an HTTP port')

      const graph = new SceneGraph()
      browser = await connectMockBrowser(httpPort, graph, TEST_AUTH_TOKEN)
      await waitForBrowserRegistration(httpPort)

      client = new Client({ name: 'test-root', version: '0.0.0' })
      const transport = new StreamableHTTPClientTransport(
        new URL(`http://127.0.0.1:${httpPort}/mcp`),
        { requestInit: { headers: { Authorization: `Bearer ${TEST_AUTH_TOKEN}` } } }
      )
      await client.connect(transport)

      await fn(client, browser, graph)
    } finally {
      await client?.close().catch(() => undefined)
      browser?.close()
      await handle.close()
    }
  }

  test('registers open_file and new_document tools when mcpRoot is set', async () => {
    await withMCPRootServer(TEST_MCP_ROOT, async (client) => {
      const { tools } = await client.listTools()
      const byName = new Map(tools.map((tool) => [tool.name, tool] as const))
      expect(byName.has('open_file')).toBe(true)
      expect(byName.get('open_file')?.annotations?.readOnlyHint).toBe(true)
      expect(byName.has('new_document')).toBe(true)
    })
  })

  test('registers close_file as destructive and forwards its unsaved choice and target', async () => {
    await withMCPRootServer(TEST_MCP_ROOT, async (client, browser) => {
      const { tools } = await client.listTools()
      const closeFile = tools.find((tool) => tool.name === 'close_file')
      expect(closeFile?.annotations?.destructiveHint).toBe(true)

      const result = await client.callTool({
        name: 'close_file',
        arguments: { document_id: 'doc-1', unsaved: 'discard' }
      })
      expect(result.isError).not.toBe(true)
      expect(browser.requests.find((item) => item.command === 'close_file')?.args).toEqual({
        document_id: 'doc-1',
        unsaved: 'discard'
      })
    })
  })

  test('close_file keeps a save path inside mcpRoot', async () => {
    await withMCPRootServer(TEST_MCP_ROOT, async (client, browser) => {
      const outside = await client.callTool({
        name: 'close_file',
        arguments: { unsaved: 'save', path: '/etc/escape.fig' }
      })
      expect(outside.isError).toBe(true)
      expect(browser.requests.some((item) => item.command === 'close_file')).toBe(false)
    })
  })

  test('forwards document activation and history steps with their target', async () => {
    await withMCPRootServer(TEST_MCP_ROOT, async (client, browser) => {
      const activated = await client.callTool({
        name: 'activate_document',
        arguments: { document_id: 'doc-2', page_id: '0:5' }
      })
      expect(activated.isError).not.toBe(true)
      expect(browser.requests.find((item) => item.command === 'activate_document')?.args).toEqual({
        document_id: 'doc-2',
        page_id: '0:5'
      })

      const undone = await client.callTool({ name: 'undo', arguments: { document_id: 'doc-2' } })
      expect(undone.isError).not.toBe(true)
      expect(JSON.stringify(undone.content)).toContain('Agent: mock')
      expect(browser.requests.find((item) => item.command === 'undo')?.args).toEqual({
        document_id: 'doc-2'
      })
    })
  })

  test('forwards settings updates and returns the resulting settings', async () => {
    await withMCPRootServer(TEST_MCP_ROOT, async (client, browser) => {
      const settings = { appearance: { theme: 'light' } }
      const result = await client.callTool({ name: 'update_settings', arguments: { settings } })
      expect(result.isError).not.toBe(true)
      expect(browser.requests.find((item) => item.command === 'update_settings')?.args).toEqual({
        settings
      })
      expect(JSON.stringify(result.content)).toContain('light')
    })
  })

  test('save_file accepts an explicit path inside mcpRoot', async () => {
    await withMCPRootServer(TEST_MCP_ROOT, async (client, browser) => {
      await mkdir(join(TEST_MCP_ROOT, 'unicode'), { recursive: true })
      const savePath = join(TEST_MCP_ROOT, 'unicode', 'пример.fig')
      const result = await client.callTool({
        name: 'save_file',
        arguments: { path: savePath }
      })

      expect(result.isError).not.toBe(true)
      const request = browser.requests.find((item) => item.command === 'save_file')
      // The server sends the canonical (realpath-resolved) path to the browser
      // to prevent TOCTOU races. On macOS, /var -> /private/var.
      const { realpath } = await import('node:fs/promises')
      const { dirname, basename } = await import('node:path')
      const canonicalPath = join(await realpath(dirname(savePath)), basename(savePath))
      expect(request?.args).toEqual({ path: canonicalPath })
    })
  })

  test('save_file rejects paths outside mcpRoot', async () => {
    await withMCPRootServer(TEST_MCP_ROOT, async (client, browser) => {
      const result = await client.callTool({
        name: 'save_file',
        arguments: { path: join(join(TEST_MCP_ROOT, '..'), 'outside.fig') }
      })

      expect(result.isError).toBe(true)
      expect(browser.requests.some((item) => item.command === 'save_file')).toBe(false)
    })
  })

  test('does not register open_file when mcpRoot is null', async () => {
    await withMCPRootServer(null, async (client) => {
      const { tools } = await client.listTools()
      const names = tools.map((t) => t.name)
      expect(names).not.toContain('open_file')
      expect(names).not.toContain('new_document')
    })
  })
})

// ---------------------------------------------------------------------------
// GAP-01: server.close() resource cleanup
// ---------------------------------------------------------------------------

describe('MCP server lifecycle', () => {
  test('close() removes the discovery file from disk', async () => {
    if (isUnix) await mkdir(SOCKET_DIR, { recursive: true })
    const { getDiscoveryPath } = await import('#mcp/transport/paths')
    const discoveryPath = await getDiscoveryPath()

    const handle = await startServer({
      httpPort: 0,
      withTcp: true,
      socketPath: testSocketPath(),
      authToken: TEST_AUTH_TOKEN,
      enableEval: false,
      mcpRoot: null
    })
    try {
      expect(await Bun.file(discoveryPath).exists()).toBe(true)
    } finally {
      await handle.close()
    }
    expect(await Bun.file(discoveryPath).exists()).toBe(false)
  })

  test('close() removes the unix socket file from disk', async () => {
    if (!isUnix) return
    await mkdir(SOCKET_DIR, { recursive: true })
    const socketPath = expectDefined(testSocketPath(), 'socket path')
    expect(socketPath).toBeTruthy()

    const handle = await startServer({
      httpPort: 0,
      withTcp: true,
      socketPath,
      authToken: TEST_AUTH_TOKEN,
      enableEval: false,
      mcpRoot: null
    })
    try {
      const info = await stat(socketPath)
      expect(info.isSocket()).toBe(true)
    } finally {
      await handle.close()
    }
    await expect(stat(socketPath)).rejects.toThrow()
  })

  test('close() removes socket file when no replacement server is listening', async () => {
    if (!isUnix) return
    await mkdir(SOCKET_DIR, { recursive: true })
    const socketPath = expectDefined(testSocketPath(), 'socket path')

    const handle1 = await startServer({
      httpPort: 0,
      withTcp: false,
      socketPath,
      authToken: TEST_AUTH_TOKEN,
      enableEval: false,
      mcpRoot: null
    })
    try {
      expect(handle1.socketPath).toBe(socketPath)
    } finally {
      await handle1.close()
    }
    await expect(stat(socketPath)).rejects.toThrow()
  })

  test('close() is idempotent and does not throw on second call', async () => {
    if (isUnix) await mkdir(SOCKET_DIR, { recursive: true })
    const handle = await startServer({
      httpPort: 0,
      withTcp: true,
      socketPath: testSocketPath(),
      authToken: TEST_AUTH_TOKEN,
      enableEval: false,
      mcpRoot: null
    })
    await handle.close()
    await expect(handle.close()).resolves.toBeUndefined()
  })

  test('close() rejects subsequent RPC requests', async () => {
    if (isUnix) await mkdir(SOCKET_DIR, { recursive: true })
    const handle = await startServer({
      httpPort: 0,
      withTcp: true,
      socketPath: testSocketPath(),
      authToken: TEST_AUTH_TOKEN,
      enableEval: false,
      mcpRoot: null
    })
    const httpPort = handle.httpPort
    if (!httpPort) {
      await handle.close()
      throw new Error('withTcp: true did not produce an HTTP port')
    }

    try {
      const healthyResp = await fetch(`http://127.0.0.1:${httpPort}/health`)
      expect(healthyResp.status).toBe(200)
    } finally {
      await handle.close()
    }

    await expect(fetch(`http://127.0.0.1:${httpPort}/health`)).rejects.toThrow()
  })
})

// ---------------------------------------------------------------------------
// GAP-02: concurrent startServer calls produce non-overlapping servers
// ---------------------------------------------------------------------------

describe('MCP server concurrent startServer', () => {
  test('two simultaneous startServer calls each get their own discovery file and both stay up', async () => {
    if (isUnix) await mkdir(SOCKET_DIR, { recursive: true })
    const { getDiscoveryPath } = await import('#mcp/transport/paths')

    // Use unique socket paths so the two servers don't fight over the
    // socket file. The test still covers discovery atomicity by writing
    // to the same shared discovery path.
    const results = await Promise.allSettled([
      startServer({
        httpPort: 0,
        withTcp: true,
        socketPath: testSocketPath(),
        authToken: 'token-a',
        enableEval: false,
        mcpRoot: null
      }),
      startServer({
        httpPort: 0,
        withTcp: true,
        socketPath: testSocketPath(),
        authToken: 'token-b',
        enableEval: false,
        mcpRoot: null
      })
    ])

    const a = results[0].status === 'fulfilled' ? results[0].value : null
    const b = results[1].status === 'fulfilled' ? results[1].value : null

    if (!a || !b) {
      if (a) await a.close()
      if (b) await b.close()
      throw new Error('Multi-server startup failed')
    }

    try {
      // Both servers are listening on their own ephemeral ports.
      expect(a.httpPort).toBeGreaterThan(0)
      expect(b.httpPort).toBeGreaterThan(0)
      expect(a.httpPort).not.toBe(b.httpPort)

      // Each responds on /health with the expected auth state.
      const aHealth = await readHealth(await fetch(`http://127.0.0.1:${a.httpPort}/health`))
      const bHealth = await readHealth(await fetch(`http://127.0.0.1:${b.httpPort}/health`))
      expect(aHealth.status).toBe('no_app')
      expect(bHealth.status).toBe('no_app')

      // Discovery file exists and was last written by one of the two servers
      // (atomic rename guarantees it's a complete file, not interleaved).
      const discoveryPath = await getDiscoveryPath()
      const file = Bun.file(discoveryPath)
      expect(await file.exists()).toBe(true)
      const info = expectDefined(parseDiscoveryInfo(await file.text()), 'discovery file')
      expect(info.pid).toBe(process.pid)
      expect(['token-a', 'token-b']).toContain(
        expectDefined(info.authToken, 'discovery auth token')
      )
    } finally {
      await a.close()
      await b.close()
    }
  }, 15000)

  test("closing one server does not delete another server's discovery file", async () => {
    if (isUnix) await mkdir(SOCKET_DIR, { recursive: true })
    const { getDiscoveryPath } = await import('#mcp/transport/paths')

    const a = await startServer({
      httpPort: 0,
      withTcp: true,
      socketPath: testSocketPath(),
      authToken: 'token-a',
      enableEval: false,
      mcpRoot: null
    })
    let b: ServerHandle | undefined
    try {
      b = await startServer({
        httpPort: 0,
        withTcp: true,
        socketPath: testSocketPath(),
        authToken: 'token-b',
        enableEval: false,
        mcpRoot: null
      })
    } catch (err) {
      await a.close()
      throw err
    }

    try {
      // Close server a — its discovery cleanup should NOT remove the file
      // because server b still owns it (different auth token).
      await a.close()

      // Server b should still be healthy and reachable.
      const bHealth = await readHealth(await fetch(`http://127.0.0.1:${b.httpPort}/health`))
      expect(bHealth.status).toBe('no_app')

      // Discovery file should still exist (owned by server b now).
      const discoveryPath = await getDiscoveryPath()
      const file = Bun.file(discoveryPath)
      expect(await file.exists()).toBe(true)
      const info = expectDefined(parseDiscoveryInfo(await file.text()), 'discovery file')
      expect(info.authToken).toBe('token-b')
    } finally {
      await b.close()
    }
  }, 15000)
})