agent-ui/server/services/torch-server.ts

/**
 * Torch Server
 * WebSocket server for multi-browser MCP control synchronization.
 * Only one browser can have the "torch" and connect to MCP at a time.
 */

import { PORT_TORCH } from '../config'

// Client metadata
interface TorchClient {
  ws: any
  id: string
  userAgent: string
  hostname: string
  connectedAt: Date
}

// Connected clients
const clients = new Map<any, TorchClient>()
let clientIdCounter = 1

// Torch state - who has control
let torchHolderId: string | null = null

function generateClientId(): string {
  return `client_${clientIdCounter++}_${Date.now().toString(36)}`
}

function broadcast(message: string) {
  for (const [ws] of clients) {
    try {
      ws.send(message)
    } catch {
      clients.delete(ws)
    }
  }
}

function broadcastTorchState() {
  const clientList = Array.from(clients.values()).map(c => ({
    id: c.id,
    userAgent: c.userAgent,
    hostname: c.hostname,
    connectedAt: c.connectedAt.toISOString(),
    hasTorch: c.id === torchHolderId
  }))

  const message = JSON.stringify({
    type: 'torch-update',
    holderId: torchHolderId,
    clients: clientList
  })

  broadcast(message)
}

function handleMessage(ws: any, data: any) {
  const client = clients.get(ws)
  if (!client) return

  switch (data.type) {
    case 'register': {
      client.userAgent = data.userAgent || 'Unknown'
      client.hostname = data.hostname || 'Unknown'

      // Auto-assign torch if no one has it
      const shouldHaveTorch = !torchHolderId
      if (shouldHaveTorch) {
        torchHolderId = client.id
      }

      ws.send(JSON.stringify({
        type: 'registered',
        id: client.id,
        hasTorch: shouldHaveTorch
      }))

      console.log(`[Torch] Registered: ${client.id} (torch: ${shouldHaveTorch})`)
      broadcastTorchState()
      break
    }

    case 'request': {
      const previousHolder = torchHolderId
      torchHolderId = client.id

      ws.send(JSON.stringify({ type: 'granted' }))
      console.log(`[Torch] Transferred: ${previousHolder} → ${client.id}`)
      broadcastTorchState()
      break
    }

    case 'release': {
      if (torchHolderId === client.id) {
        torchHolderId = null
        ws.send(JSON.stringify({ type: 'released' }))
        console.log(`[Torch] Released by: ${client.id}`)
        broadcastTorchState()
      }
      break
    }
  }
}

export function startTorchServer() {
  const server = Bun.serve({
    port: PORT_TORCH,
    fetch(req, server) {
      const url = new URL(req.url)

      // CORS headers
      const corsHeaders = {
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Allow-Methods': 'GET, OPTIONS',
        'Access-Control-Allow-Headers': 'Content-Type'
      }

      if (req.method === 'OPTIONS') {
        return new Response(null, { headers: corsHeaders })
      }

      // Health check
      if (url.pathname === '/health') {
        return Response.json({
          status: 'ok',
          clients: clients.size,
          torchHolder: torchHolderId
        }, { headers: corsHeaders })
      }

      // WebSocket upgrade
      const upgrade = req.headers.get('upgrade')
      if (upgrade?.toLowerCase() === 'websocket') {
        const success = server.upgrade(req)
        if (success) return undefined
        return new Response('WebSocket upgrade failed', { status: 400 })
      }

      return new Response('Torch WebSocket Server', { status: 200 })
    },
    websocket: {
      open(ws) {
        const id = generateClientId()
        clients.set(ws, {
          ws,
          id,
          userAgent: 'Unknown',
          hostname: 'Unknown',
          connectedAt: new Date()
        })
        console.log(`[Torch] Client connected: ${id} (${clients.size} total)`)
      },
      message(ws, message) {
        try {
          const data = JSON.parse(message.toString())
          handleMessage(ws, data)
        } catch (e) {
          console.error('[Torch] Invalid message:', e)
        }
      },
      close(ws) {
        const client = clients.get(ws)
        if (client) {
          console.log(`[Torch] Client disconnected: ${client.id}`)

          // If this client had the torch, release it
          if (torchHolderId === client.id) {
            torchHolderId = null

            // Auto-assign to next client
            const nextClient = clients.values().next().value
            if (nextClient && nextClient.ws !== ws) {
              torchHolderId = nextClient.id
              console.log(`[Torch] Auto-assigned to: ${nextClient.id}`)
            }
          }

          clients.delete(ws)
          broadcastTorchState()
        }
      }
    }
  })

  console.log(`[Torch] WebSocket server on port ${PORT_TORCH}`)
  return server
}

export function stopTorchServer() {
  clients.clear()
  torchHolderId = null
}