openpencil/tests/helpers/mcp/server.ts

import { type RequestOptions, request as httpRequest } from 'node:http'

import * as v from 'valibot'
import { WebSocket } from 'ws'

import {
  type SceneGraph,
  ALL_TOOLS,
  FigmaAPI,
  computeAllLayouts,
  executeRPCCommand
} from '@open-pencil/core'
import type { ToolDescriptor } from '@open-pencil/mcp/tools'
import { randomHex } from '@open-pencil/scene-graph/random'

const ToolDescriptorShape = v.looseObject({ name: v.string(), enabled: v.boolean() })

export const HealthResponseSchema = v.object({
  status: v.string(),
  version: v.string(),
  authRequired: v.boolean(),
  tools: v.optional(v.array(v.custom<ToolDescriptor>((value) => v.is(ToolDescriptorShape, value)))),
  discoveryPath: v.optional(v.string())
})

export type HealthResponse = v.InferOutput<typeof HealthResponseSchema>

/** Read a JSON response body or file, checked against `schema`. */
export async function readJSON<T>(
  source: { text(): Promise<string> },
  schema: v.GenericSchema<unknown, T>
): Promise<T> {
  return v.parse(v.pipe(v.string(), v.parseJson(), schema), await source.text())
}

export function readHealth(response: Response): Promise<HealthResponse> {
  return readJSON(response, HealthResponseSchema)
}

export interface MockBrowserRequest {
  command: string
  args?: unknown
}

export interface MockBrowser {
  ws: WebSocket
  graph: SceneGraph
  requests: MockBrowserRequest[]
  close: () => void
}

/** A `register` broadcast: the browser bridge announcing its token. */
export const RegisterMessage = v.object({
  type: v.string(),
  token: v.optional(v.nullable(v.string()))
})

/** A tool response relayed through the browser bridge. */
export const ResponseMessage = v.object({
  type: v.string(),
  id: v.string(),
  ok: v.optional(v.boolean()),
  result: v.optional(v.looseObject({ name: v.optional(v.string()) })),
  error: v.optional(v.string())
})

/** Read the next WebSocket JSON message with a timeout, checked against `schema`. */
export function readWsJSON<T>(
  ws: WebSocket,
  schema: v.GenericSchema<unknown, T>,
  timeoutMs = 1000
): Promise<T> {
  return new Promise((resolve, reject) => {
    const cleanup = () => {
      clearTimeout(timer)
      ws.off('message', onMessage)
      ws.off('error', onError)
    }
    const onMessage = (raw: WebSocket.RawData) => {
      cleanup()
      try {
        resolve(v.parse(v.pipe(v.string(), v.parseJson(), schema), raw.toString()))
      } catch (error) {
        reject(error)
      }
    }
    const onError = (error: Error) => {
      cleanup()
      reject(error)
    }
    const timer = setTimeout(() => {
      cleanup()
      reject(new Error('Timed out waiting for WebSocket message'))
    }, timeoutMs)
    ws.on('message', onMessage)
    ws.on('error', onError)
  })
}

/** Open a WebSocket connection, optionally authenticating as a stdio bridge client. */
export function openWs(url: string, authToken?: string | null): Promise<WebSocket> {
  const ws = new WebSocket(url)
  return new Promise((resolve, reject) => {
    ws.once('open', () => {
      if (authToken) {
        // Authenticate as a stdio bridge client (not the browser app).
        // The "auth" message type validates the token and adds the client
        // to authenticatedClients without replacing the registered browser.
        ws.send(JSON.stringify({ type: 'auth', token: authToken }))
      }
      resolve(ws)
    })
    ws.once('error', reject)
  })
}

/** Read the next WebSocket JSON message, skipping any 'register' broadcasts. */
export async function readNextResponse<T>(
  ws: WebSocket,
  schema: v.GenericSchema<unknown, T>,
  timeoutMs = 5000
): Promise<T> {
  const start = Date.now()
  for (let i = 0; ; i++) {
    const remaining = timeoutMs - (Date.now() - start)
    if (remaining <= 0) throw new Error(`Timed out after reading ${i} register messages`)
    const msg = await readWsJSON(ws, v.looseObject({ type: v.string() }), remaining)
    if (msg.type !== 'register') return v.parse(schema, msg)
  }
}

/** Low-level HTTP request via Node's http module, with timeout. */
export function nodeHttpRequest(
  opts: RequestOptions,
  bodyJSON?: string,
  timeoutMs = 5_000
): Promise<{ status: number; data: unknown }> {
  return new Promise((resolve, reject) => {
    const req = httpRequest(opts, (res) => {
      const chunks: Buffer[] = []
      res.on('data', (chunk: Buffer) => chunks.push(chunk))
      res.on('end', () => {
        clearTimeout(timeout)
        const raw = Buffer.concat(chunks).toString('utf-8')
        let data: unknown
        try {
          data = JSON.parse(raw)
        } catch {
          data = raw
        }
        resolve({ status: res.statusCode ?? 200, data })
      })
      res.on('error', (err) => {
        clearTimeout(timeout)
        reject(err)
      })
    })
    const timeout = setTimeout(() => {
      req.destroy(new Error(`nodeHttpRequest timed out after ${timeoutMs / 1000}s`))
    }, timeoutMs)
    req.on('error', (err) => {
      clearTimeout(timeout)
      reject(err)
    })
    // Connection-level timeout: destroys the socket if the server stalls
    // during connection or mid-request. The response-level timeout above
    // handles slow responses; this handles hung connections (e.g., socket
    // exists but no one is listening).
    req.setTimeout(timeoutMs, () => {
      req.destroy(new Error(`nodeHttpRequest connection timed out after ${timeoutMs / 1000}s`))
    })
    if (bodyJSON) req.write(bodyJSON)
    req.end()
  })
}

/** Make an HTTP request via TCP to a specific port. */
export function tcpRequest(
  port: number,
  method: string,
  path: string,
  body?: Record<string, unknown>,
  headers?: Record<string, string>
): Promise<{ status: number; data: unknown }> {
  if (!port) throw new Error('tcpRequest: port not initialized — beforeAll must run first')
  const bodyJSON = body ? JSON.stringify(body) : undefined
  return nodeHttpRequest(
    {
      hostname: '127.0.0.1',
      port,
      path,
      method,
      headers: { ...(bodyJSON ? { 'Content-Type': 'application/json' } : {}), ...headers }
    },
    bodyJSON
  )
}

/** Make an HTTP request via Unix domain socket. */
export function socketRequest(
  socketPath: string,
  method: string,
  path: string,
  headers?: Record<string, string>
): Promise<{ status: number; data: unknown }> {
  return nodeHttpRequest({ socketPath, path, method, headers: headers ?? {} })
}

// App commands whose mock result does not depend on the graph or arguments.
const FIXED_MOCK_RESULTS: Partial<Record<string, object>> = {
  save_file: {},
  new_document: {},
  open_file: {},
  close_file: { closed: true },
  activate_document: { activated: true },
  undo: { applied: true, label: 'Agent: mock' },
  redo: { applied: true, label: 'Agent: mock' }
}

async function handleMockCommand(
  graph: SceneGraph,
  command: string,
  rawArgs: unknown
): Promise<unknown> {
  const args = rawArgs as { name?: string; args?: Record<string, unknown> } | undefined

  if (command === 'tool' && args?.name) {
    const def = ALL_TOOLS.find((t) => t.name === args.name)
    if (!def) throw new Error(`Unknown tool: ${args.name}`)
    const api = new FigmaAPI(graph)
    const pages = graph.getPages()
    if (pages.length === 0) throw new Error('No pages in graph')
    api.currentPage = api.wrapNode(pages[0].id)
    const result = await def.execute(api, args.args ?? {})
    if (def.mutates) computeAllLayouts(graph)
    return result
  }

  if (command === 'list_documents') {
    const pages = graph.getPages()
    const currentPage = pages[0]
    return {
      documents: [
        {
          id: 'doc-1',
          name: 'Mock document',
          active: true,
          current_page_id: currentPage?.id ?? '',
          current_page_name: currentPage?.name ?? '',
          pages: pages.map((page) => ({ id: page.id, name: page.name }))
        }
      ]
    }
  }

  const fixed = FIXED_MOCK_RESULTS[command]
  if (fixed) return fixed

  if (command === 'get_settings' || command === 'update_settings') {
    const settings = (rawArgs as { settings?: unknown } | undefined)?.settings
    return { settings: settings ?? { appearance: { theme: 'dark' } } }
  }

  return executeRPCCommand(graph, command, args ?? {})
}

export function connectMockBrowser(
  port: number,
  graph: SceneGraph,
  authToken?: string
): Promise<MockBrowser> {
  return new Promise((resolve, reject) => {
    const ws = new WebSocket(`ws://127.0.0.1:${port}`)
    const requests: MockBrowserRequest[] = []
    const token = authToken ?? 'test-token-' + randomHex(8)

    let settled = false

    ws.on('open', () => {
      ws.send(JSON.stringify({ type: 'register', token }))

      ws.on('message', async (raw) => {
        let parsed: unknown
        try {
          parsed = JSON.parse(raw.toString())
        } catch {
          return // Ignore malformed payloads
        }
        if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) return
        const msg = parsed as { type: string; id: string; command: string; args?: unknown }
        if (msg.type !== 'request') return

        try {
          requests.push({ command: msg.command, args: msg.args })
          const result = await handleMockCommand(graph, msg.command, msg.args)
          ws.send(JSON.stringify({ type: 'response', id: msg.id, ok: true, result }))
        } catch (e) {
          ws.send(
            JSON.stringify({
              type: 'response',
              id: msg.id,
              ok: false,
              error: e instanceof Error ? e.message : String(e)
            })
          )
        }
      })

      if (!settled) {
        settled = true
        // Wait for the server to confirm browser registration before resolving.
        // Without this, follow-up test traffic can start before the server marks
        // the browser as connected, causing race conditions.
        void waitForBrowserRegistration(port).then(
          () => {
            resolve({ ws, graph, requests, close: () => ws.close() })
            return undefined
          },
          (err) => {
            ws.close()
            reject(err)
          }
        )
      }
    })

    ws.on('error', (err) => {
      if (!settled) {
        settled = true
        reject(err)
      }
    })
  })
}

export async function waitForBrowserRegistration(port: number, timeoutMs = 5000): Promise<void> {
  const start = Date.now()
  let lastStatus = 'unknown'
  while (Date.now() - start < timeoutMs) {
    try {
      const resp = await fetch(`http://127.0.0.1:${port}/health`)
      const health = await readHealth(resp)
      lastStatus = health.status
      if (health.status === 'ok') return
    } catch {
      void 0
    }
    await new Promise<void>((r) => {
      setTimeout(r, 100)
    })
  }
  throw new Error(
    `Browser registration not confirmed within ${timeoutMs}ms (last status: ${lastStatus})`
  )
}