playwriter/playwriter/src/cdp-relay.ts

import { Hono } from 'hono'
import { serve } from '@hono/node-server'
import { createNodeWebSocket } from '@hono/node-ws'
import type { WSContext } from 'hono/ws'
import type { Protocol } from './cdp-types.js'
import type { CDPCommand, CDPResponseBase, CDPEventBase, CDPEventFor, RelayServerEvents } from './cdp-types.js'
import type { ExtensionMessage, ExtensionEventMessage } from './protocol.js'
import chalk from 'chalk'
import { EventEmitter } from 'node:events'
import { VERSION } from './utils.js'

type ConnectedTarget = {
  sessionId: string
  targetId: string
  targetInfo: Protocol.Target.TargetInfo
}



type PlaywrightClient = {
  id: string
  ws: WSContext
}


export type RelayServer = {
  close(): void
  on<K extends keyof RelayServerEvents>(event: K, listener: RelayServerEvents[K]): void
  off<K extends keyof RelayServerEvents>(event: K, listener: RelayServerEvents[K]): void
}

export async function startPlayWriterCDPRelayServer({ port = 19988, host = '127.0.0.1', token, logger }: { port?: number; host?: string; token?: string; logger?: { log(...args: any[]): void; error(...args: any[]): void } } = {}): Promise<RelayServer> {
  const emitter = new EventEmitter()
  const connectedTargets = new Map<string, ConnectedTarget>()

  const playwrightClients = new Map<string, PlaywrightClient>()
  let extensionWs: WSContext | null = null

  const extensionPendingRequests = new Map<number, {
    resolve: (result: any) => void
    reject: (error: Error) => void
  }>()
  let extensionMessageId = 0
  let extensionPingInterval: ReturnType<typeof setInterval> | null = null

  function startExtensionPing() {
    if (extensionPingInterval) {
      clearInterval(extensionPingInterval)
    }
    extensionPingInterval = setInterval(() => {
      extensionWs?.send(JSON.stringify({ method: 'ping' }))
    }, 5000)
  }

  function stopExtensionPing() {
    if (extensionPingInterval) {
      clearInterval(extensionPingInterval)
      extensionPingInterval = null
    }
  }

  function logCdpMessage({
    direction,
    clientId,
    method,
    sessionId,
    params,
    id,
    source
  }: {
    direction: 'to-playwright' | 'from-playwright' | 'from-extension'
    clientId?: string
    method: string
    sessionId?: string
    params?: any
    id?: number
    source?: 'extension' | 'server'
  }) {
    const noisyEvents = [
      'Network.requestWillBeSentExtraInfo',
      'Network.responseReceived',
      'Network.responseReceivedExtraInfo',
      'Network.dataReceived',
      'Network.requestWillBeSent',
      'Network.loadingFinished'
    ]

    if (noisyEvents.includes(method)) {
      return
    }

    const details: string[] = []

    if (id !== undefined) {
      details.push(`id=${id}`)
    }

    if (sessionId) {
      details.push(`sessionId=${sessionId}`)
    }

    if (params) {
      if (params.targetId) {
        details.push(`targetId=${params.targetId}`)
      }
      if (params.targetInfo?.targetId) {
        details.push(`targetId=${params.targetInfo.targetId}`)
      }
      if (params.sessionId && params.sessionId !== sessionId) {
        details.push(`sessionId=${params.sessionId}`)
      }
    }

    const detailsStr = details.length > 0 ? ` ${chalk.gray(details.join(', '))}` : ''

    if (direction === 'from-playwright') {
      const clientLabel = clientId ? chalk.blue(`[${clientId}]`) : ''
      logger?.log(chalk.cyan('← Playwright'), clientLabel + ':', method + detailsStr)
    } else if (direction === 'from-extension') {
      logger?.log(chalk.yellow('← Extension:'), method + detailsStr)
    } else if (direction === 'to-playwright') {
      const color = source === 'server' ? chalk.magenta : chalk.green
      const sourceLabel = source === 'server' ? chalk.gray(' (server-generated)') : ''
      const clientLabel = clientId ? chalk.blue(`[${clientId}]`) : chalk.blue('[ALL]')
      logger?.log(color('→ Playwright'), clientLabel + ':', method + detailsStr + sourceLabel)
    }
  }

  function sendToPlaywright({
    message,
    clientId,
    source = 'extension'
  }: {
    message: CDPResponseBase | CDPEventBase
    clientId?: string
    source?: 'extension' | 'server'
  }) {
    const messageToSend = source === 'server' && 'method' in message
      ? { ...message, __serverGenerated: true }
      : message

    if ('method' in message) {
      logCdpMessage({
        direction: 'to-playwright',
        clientId,
        method: message.method,
        sessionId: 'sessionId' in message ? message.sessionId : undefined,
        params: 'params' in message ? message.params : undefined,
        source
      })
    }

    const messageStr = JSON.stringify(messageToSend)

    if (clientId) {
      const client = playwrightClients.get(clientId)
      if (client) {
        client.ws.send(messageStr)
      }
    } else {
      for (const client of playwrightClients.values()) {
        client.ws.send(messageStr)
      }
    }
  }

  async function sendToExtension({ method, params, timeout = 30000 }: { method: string; params?: any; timeout?: number }) {
    if (!extensionWs) {
      throw new Error('Extension not connected')
    }

    const id = ++extensionMessageId
    const message = { id, method, params }

    extensionWs.send(JSON.stringify(message))

    return new Promise((resolve, reject) => {
      const timeoutId = setTimeout(() => {
        extensionPendingRequests.delete(id)
        reject(new Error(`Extension request timeout after ${timeout}ms: ${method}`))
      }, timeout)

      extensionPendingRequests.set(id, {
        resolve: (result) => {
          clearTimeout(timeoutId)
          resolve(result)
        },
        reject: (error) => {
          clearTimeout(timeoutId)
          reject(error)
        }
      })
    })
  }

  // Auto-create initial tab when PLAYWRITER_AUTO_ENABLE is set and no targets exist.
  // This allows Playwright to connect and immediately have a page to work with.
  async function maybeAutoCreateInitialTab(): Promise<void> {
    if (!process.env.PLAYWRITER_AUTO_ENABLE) {
      return
    }
    if (!extensionWs) {
      return
    }
    if (connectedTargets.size > 0) {
      return
    }

    try {
      logger?.log(chalk.blue('Auto-creating initial tab for Playwright client'))
      const result = await sendToExtension({ method: 'createInitialTab', timeout: 10000 }) as {
        success: boolean
        tabId: number
        sessionId: string
        targetInfo: Protocol.Target.TargetInfo
      }
      if (result.success && result.sessionId && result.targetInfo) {
        connectedTargets.set(result.sessionId, {
          sessionId: result.sessionId,
          targetId: result.targetInfo.targetId,
          targetInfo: result.targetInfo
        })
        logger?.log(chalk.blue(`Auto-created tab, now have ${connectedTargets.size} targets`))
      }
    } catch (e) {
      logger?.error('Failed to auto-create initial tab:', e)
    }
  }

  async function routeCdpCommand({ method, params, sessionId }: { method: string; params: any; sessionId?: string }) {
    switch (method) {
      case 'Browser.getVersion': {
        return {
          protocolVersion: '1.3',
          product: 'Chrome/Extension-Bridge',
          revision: '1.0.0',
          userAgent: 'CDP-Bridge-Server/1.0.0',
          jsVersion: 'V8'
        } satisfies Protocol.Browser.GetVersionResponse
      }

      case 'Browser.setDownloadBehavior': {
        return {}
      }

      // Target.setAutoAttach is a CDP command Playwright sends on first connection.
      // We use it as the hook to auto-create an initial tab. If Playwright changes
      // its initialization sequence in the future, this could be moved to a different command.
      case 'Target.setAutoAttach': {
        if (sessionId) {
          break
        }
        await maybeAutoCreateInitialTab()
        return {}
      }

      case 'Target.setDiscoverTargets': {
        return {}
      }

      case 'Target.attachToTarget': {
        const targetId = params?.targetId
        if (!targetId) {
          throw new Error('targetId is required for Target.attachToTarget')
        }

        for (const target of connectedTargets.values()) {
          if (target.targetId === targetId) {
            return { sessionId: target.sessionId } satisfies Protocol.Target.AttachToTargetResponse
          }
        }

        throw new Error(`Target ${targetId} not found in connected targets`)
      }

      case 'Target.getTargetInfo': {
        const targetId = params?.targetId

        if (targetId) {
          for (const target of connectedTargets.values()) {
            if (target.targetId === targetId) {
              return { targetInfo: target.targetInfo }
            }
          }
        }

        if (sessionId) {
          const target = connectedTargets.get(sessionId)
          if (target) {
            return { targetInfo: target.targetInfo }
          }
        }

        const firstTarget = Array.from(connectedTargets.values())[0]
        return { targetInfo: firstTarget?.targetInfo }
      }

      case 'Target.getTargets': {
        return {
          targetInfos: Array.from(connectedTargets.values()).map((t) => ({
            ...t.targetInfo,
            attached: true
          }))
        }
      }

      case 'Target.createTarget': {
        return await sendToExtension({
          method: 'forwardCDPCommand',
          params: { method, params }
        })
      }

      case 'Target.closeTarget': {
        return await sendToExtension({
          method: 'forwardCDPCommand',
          params: { method, params }
        })
      }

      case 'Runtime.enable': {
        if (!sessionId) {
          break
        }

        const contextCreatedPromise = new Promise<void>((resolve) => {
          const handler = ({ event }: { event: CDPEventBase }) => {
            if (event.method === 'Runtime.executionContextCreated' && event.sessionId === sessionId) {
              const params = event.params as Protocol.Runtime.ExecutionContextCreatedEvent | undefined
              if (params?.context?.auxData?.isDefault === true) {
                clearTimeout(timeout)
                emitter.off('cdp:event', handler)
                resolve()
              }
            }
          }
          const timeout = setTimeout(() => {
            emitter.off('cdp:event', handler)
            logger?.log(chalk.yellow(`IMPORTANT: Runtime.enable timed out waiting for main frame executionContextCreated (sessionId: ${sessionId}). This may cause pages to not be visible immediately.`))
            resolve()
          }, 3000)
          emitter.on('cdp:event', handler)
        })

        const result = await sendToExtension({
          method: 'forwardCDPCommand',
          params: { sessionId, method, params }
        })

        await contextCreatedPromise

        return result
      }
    }

    return await sendToExtension({
      method: 'forwardCDPCommand',
      params: { sessionId, method, params }
    })
  }

  const app = new Hono()
  const { injectWebSocket, upgradeWebSocket } = createNodeWebSocket({ app })

  app.get('/', (c) => {
    return c.text('OK')
  })

  app.get('/version', (c) => {
    return c.json({ version: VERSION })
  })

  app.get('/extension/status', (c) => {
    return c.json({ connected: extensionWs !== null })
  })

  app.post('/mcp-log', async (c) => {
    try {
      const { level, args } = await c.req.json()
      const logFn = (logger as any)?.[level] || logger?.log
      const prefix = chalk.red(`[MCP] [${level.toUpperCase()}]`)
      logFn?.(prefix, ...args)
      return c.json({ ok: true })
    } catch {
      return c.json({ ok: false }, 400)
    }
  })

  app.get('/cdp/:clientId?', (c, next) => {
    if (token) {
      const url = new URL(c.req.url, 'http://localhost')
      const providedToken = url.searchParams.get('token')
      if (providedToken !== token) {
        return c.text('Unauthorized', 401)
      }
    }
    return next()
  }, upgradeWebSocket((c) => {
    const clientId = c.req.param('clientId') || 'default'

    return {
      async onOpen(_event, ws) {
        if (playwrightClients.has(clientId)) {
          logger?.log(chalk.red(`Rejecting duplicate client ID: ${clientId}`))
          ws.close(1000, 'Client ID already connected')
          return
        }

        // Add client first so it can receive Target.attachedToTarget events
        playwrightClients.set(clientId, { id: clientId, ws })
        logger?.log(chalk.green(`Playwright client connected: ${clientId} (${playwrightClients.size} total) (extension? ${!!extensionWs}) (${connectedTargets.size} pages)`))
      },

      async onMessage(event, ws) {
        let message: CDPCommand

        try {
          message = JSON.parse(event.data.toString())
        } catch {
          return
        }

        const { id, sessionId, method, params } = message

        logCdpMessage({
          direction: 'from-playwright',
          clientId,
          method,
          sessionId,
          id
        })

        emitter.emit('cdp:command', { clientId, command: message })

        if (!extensionWs) {
          sendToPlaywright({
            message: {
              id,
              sessionId,
              error: { message: 'Extension not connected' }
            },
            clientId
          })
          return
        }

        try {
          const result: any = await routeCdpCommand({ method, params, sessionId })

          if (method === 'Target.setAutoAttach' && !sessionId) {
            for (const target of connectedTargets.values()) {
              const attachedPayload = {
                method: 'Target.attachedToTarget',
                params: {
                  sessionId: target.sessionId,
                  targetInfo: {
                    ...target.targetInfo,
                    attached: true
                  },
                  waitingForDebugger: false
                }
              } satisfies CDPEventFor<'Target.attachedToTarget'>
              logger?.log(chalk.magenta('[Server] Target.attachedToTarget full payload:'), JSON.stringify(attachedPayload))
              sendToPlaywright({
                message: attachedPayload,
                clientId,
                source: 'server'
              })
            }
          }

          if (method === 'Target.setDiscoverTargets' && (params as any)?.discover) {
            for (const target of connectedTargets.values()) {
              const targetCreatedPayload = {
                method: 'Target.targetCreated',
                params: {
                  targetInfo: {
                    ...target.targetInfo,
                    attached: true
                  }
                }
              } satisfies CDPEventFor<'Target.targetCreated'>
              logger?.log(chalk.magenta('[Server] Target.targetCreated full payload:'), JSON.stringify(targetCreatedPayload))
              sendToPlaywright({
                message: targetCreatedPayload,
                clientId,
                source: 'server'
              })
            }
          }

          if (method === 'Target.attachToTarget' && result?.sessionId) {
            const targetId = params?.targetId
            const target = Array.from(connectedTargets.values()).find(t => t.targetId === targetId)
            if (target) {
              const attachedPayload = {
                method: 'Target.attachedToTarget',
                params: {
                  sessionId: result.sessionId,
                  targetInfo: {
                    ...target.targetInfo,
                    attached: true
                  },
                  waitingForDebugger: false
                }
              } satisfies CDPEventFor<'Target.attachedToTarget'>
              logger?.log(chalk.magenta('[Server] Target.attachedToTarget (from attachToTarget) payload:'), JSON.stringify(attachedPayload))
              sendToPlaywright({
                message: attachedPayload,
                clientId,
                source: 'server'
              })
            }
          }

          const response: CDPResponseBase = { id, sessionId, result }
          sendToPlaywright({ message: response, clientId })
          emitter.emit('cdp:response', { clientId, response, command: message })
        } catch (e) {
          logger?.error('Error handling CDP command:', method, params, e)
          const errorResponse: CDPResponseBase = {
            id,
            sessionId,
            error: { message: (e as Error).message }
          }
          sendToPlaywright({ message: errorResponse, clientId })
          emitter.emit('cdp:response', { clientId, response: errorResponse, command: message })
        }
      },

      onClose() {
        playwrightClients.delete(clientId)
        logger?.log(chalk.yellow(`Playwright client disconnected: ${clientId} (${playwrightClients.size} remaining)`))
      },

      onError(event) {
        logger?.error(`Playwright WebSocket error [${clientId}]:`, event)
      }
    }
  }))

  app.get('/extension', upgradeWebSocket(() => {
    return {
      onOpen(_event, ws) {
        if (extensionWs) {
          logger?.log(chalk.yellow('Closing existing extension connection to replace with new one'))
          extensionWs.close(4001, 'Extension Replaced')

          // Clear state from the old connection to prevent leaks
          connectedTargets.clear()
          for (const pending of extensionPendingRequests.values()) {
            pending.reject(new Error('Extension connection replaced'))
          }
          extensionPendingRequests.clear()

          for (const client of playwrightClients.values()) {
            client.ws.close(1000, 'Extension Replaced')
          }
          playwrightClients.clear()
        }

        extensionWs = ws
        startExtensionPing()
        logger?.log('Extension connected with clean state')
      },

      async onMessage(event, ws) {
        let message: ExtensionMessage

        try {
          message = JSON.parse(event.data.toString())
        } catch {
          ws.close(1000, 'Invalid JSON')
          return
        }

        if (message.id !== undefined) {
          const pending = extensionPendingRequests.get(message.id)
          if (!pending) {
            logger?.log('Unexpected response with id:', message.id)
            return
          }

          extensionPendingRequests.delete(message.id)

          if (message.error) {
            pending.reject(new Error(message.error))
          } else {
            pending.resolve(message.result)
          }
        } else if (message.method === 'pong') {
          // Keep-alive response, nothing to do
        } else if (message.method === 'log') {
          const { level, args } = message.params
          const logFn = (logger as any)?.[level] || logger?.log
          const prefix = chalk.yellow(`[Extension] [${level.toUpperCase()}]`)
          logFn?.(prefix, ...args)
        } else {
          const extensionEvent = message as ExtensionEventMessage

          if (extensionEvent.method !== 'forwardCDPEvent') {
            return
          }

          const { method, params, sessionId } = extensionEvent.params

          logCdpMessage({
            direction: 'from-extension',
            method,
            sessionId,
            params
          })

          const cdpEvent: CDPEventBase = { method, sessionId, params }
          emitter.emit('cdp:event', { event: cdpEvent, sessionId })

          if (method === 'Target.attachedToTarget') {
            const targetParams = params as Protocol.Target.AttachedToTargetEvent

            logger?.log(chalk.yellow('[Extension] Target.attachedToTarget full payload:'), JSON.stringify({ method, params: targetParams, sessionId }))

            // Check if we already sent this target to clients (e.g., from Target.setAutoAttach response)
            const alreadyConnected = connectedTargets.has(targetParams.sessionId)

            // Always update our local state with latest target info
            connectedTargets.set(targetParams.sessionId, {
              sessionId: targetParams.sessionId,
              targetId: targetParams.targetInfo.targetId,
              targetInfo: targetParams.targetInfo
            })

            // Only forward to Playwright if this is a new target to avoid duplicates
            if (!alreadyConnected) {
              sendToPlaywright({
                message: {
                  method: 'Target.attachedToTarget',
                  params: targetParams
                } as CDPEventBase,
                source: 'extension'
              })
            }
          } else if (method === 'Target.detachedFromTarget') {
            const detachParams = params as Protocol.Target.DetachedFromTargetEvent
            connectedTargets.delete(detachParams.sessionId)

            sendToPlaywright({
              message: {
                method: 'Target.detachedFromTarget',
                params: detachParams
              } as CDPEventBase,
              source: 'extension'
            })
          } else if (method === 'Target.targetCrashed') {
            const crashParams = params as Protocol.Target.TargetCrashedEvent
            for (const [sid, target] of connectedTargets.entries()) {
              if (target.targetId === crashParams.targetId) {
                connectedTargets.delete(sid)
                logger?.log(chalk.red('[Server] Target crashed, removing:'), crashParams.targetId)
                break
              }
            }

            sendToPlaywright({
              message: {
                method: 'Target.targetCrashed',
                params: crashParams
              } as CDPEventBase,
              source: 'extension'
            })
          } else if (method === 'Target.targetInfoChanged') {
            const infoParams = params as Protocol.Target.TargetInfoChangedEvent
            for (const target of connectedTargets.values()) {
              if (target.targetId === infoParams.targetInfo.targetId) {
                target.targetInfo = infoParams.targetInfo
                break
              }
            }

            sendToPlaywright({
              message: {
                method: 'Target.targetInfoChanged',
                params: infoParams
              } as CDPEventBase,
              source: 'extension'
            })
          } else if (method === 'Page.frameNavigated') {
            const frameParams = params as Protocol.Page.FrameNavigatedEvent
            if (!frameParams.frame.parentId && sessionId) {
              const target = connectedTargets.get(sessionId)
              if (target) {
                target.targetInfo = {
                  ...target.targetInfo,
                  url: frameParams.frame.url,
                  title: frameParams.frame.name || target.targetInfo.title,
                }
                logger?.log(chalk.magenta('[Server] Updated target URL from Page.frameNavigated:'), frameParams.frame.url)
              }
            }

            sendToPlaywright({
              message: {
                sessionId,
                method,
                params
              } as CDPEventBase,
              source: 'extension'
            })
          } else if (method === 'Page.navigatedWithinDocument') {
            const navParams = params as Protocol.Page.NavigatedWithinDocumentEvent
            if (sessionId) {
              const target = connectedTargets.get(sessionId)
              if (target) {
                target.targetInfo = {
                  ...target.targetInfo,
                  url: navParams.url,
                }
                logger?.log(chalk.magenta('[Server] Updated target URL from Page.navigatedWithinDocument:'), navParams.url)
              }
            }

            sendToPlaywright({
              message: {
                sessionId,
                method,
                params
              } as CDPEventBase,
              source: 'extension'
            })
          } else {
            sendToPlaywright({
              message: {
                sessionId,
                method,
                params
              } as CDPEventBase,
              source: 'extension'
            })
          }
        }
      },

      onClose(event, ws) {
        logger?.log(`Extension disconnected: code=${event.code} reason=${event.reason || 'none'}`)
        stopExtensionPing()

        // If this is an old connection closing after we've already established a new one,
        // don't clear the global state
        if (extensionWs && extensionWs !== ws) {
           logger?.log('Old extension connection closed, keeping new one active')
           return
        }

        for (const pending of extensionPendingRequests.values()) {
          pending.reject(new Error('Extension connection closed'))
        }
        extensionPendingRequests.clear()

        extensionWs = null
        connectedTargets.clear()

        for (const client of playwrightClients.values()) {
          client.ws.close(1000, 'Extension disconnected')
        }
        playwrightClients.clear()
      },

      onError(event) {
        logger?.error('Extension WebSocket error:', event)
      }
    }
  }))

  const server = serve({ fetch: app.fetch, port, hostname: host })
  injectWebSocket(server)

  const wsHost = `ws://${host}:${port}`
  const cdpEndpoint = `${wsHost}/cdp`
  const extensionEndpoint = `${wsHost}/extension`

  logger?.log('CDP relay server started')
  logger?.log('Host:', host)
  logger?.log('Port:', port)
  logger?.log('Extension endpoint:', extensionEndpoint)
  logger?.log('CDP endpoint:', cdpEndpoint)

  return {
    close() {
      for (const client of playwrightClients.values()) {
        client.ws.close(1000, 'Server stopped')
      }
      playwrightClients.clear()
      extensionWs?.close(1000, 'Server stopped')
      server.close()
      emitter.removeAllListeners()
    },
    on<K extends keyof RelayServerEvents>(event: K, listener: RelayServerEvents[K]) {
      emitter.on(event, listener as (...args: unknown[]) => void)
    },
    off<K extends keyof RelayServerEvents>(event: K, listener: RelayServerEvents[K]) {
      emitter.off(event, listener as (...args: unknown[]) => void)
    }
  }
}