agent-ui/frontend/src/services/torch.ts

import { useTorchStore } from '../stores/torch'
import { autoConnect, disconnectWebMCP } from './webmcp'
import { onTorchConnected, onTorchDisconnected } from './toolRegistry'
import { endpoints } from '../config/endpoints'

let torchWs: WebSocket | null = null
let clientId: string | null = null
let reconnectTimeout: number | null = null

/**
 * Connect to torch WebSocket server
 */
function connectToTorchServer(): Promise<void> {
  return new Promise((resolve, reject) => {
    if (torchWs?.readyState === WebSocket.OPEN) {
      resolve()
      return
    }

    console.log('[Torch] Connecting to server...')
    torchWs = new WebSocket(endpoints.torch)

    torchWs.onopen = () => {
      console.log('[Torch] Connected to server')
      const torchStore = useTorchStore()
      // Register this client with name and autoRequest
      torchWs?.send(JSON.stringify({
        type: 'register',
        userAgent: navigator.userAgent,
        hostname: window.location.hostname,
        name: torchStore.clientName || 'Anonymous',
        autoRequest: torchStore.autoRequest
      }))
      resolve()
    }

    torchWs.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data)
        handleMessage(data)
      } catch (e) {
        console.error('[Torch] Invalid message:', e)
      }
    }

    torchWs.onclose = () => {
      console.log('[Torch] Disconnected from server')
      torchWs = null

      // Reconnect after delay
      if (!reconnectTimeout) {
        reconnectTimeout = window.setTimeout(() => {
          reconnectTimeout = null
          connectToTorchServer()
        }, 2000)
      }
    }

    torchWs.onerror = (e) => {
      console.error('[Torch] WebSocket error:', e)
      reject(e)
    }
  })
}

/**
 * Handle messages from torch server
 */
async function handleMessage(data: any) {
  const torchStore = useTorchStore()

  switch (data.type) {
    case 'registered': {
      clientId = data.id
      torchStore.setClientId(data.id)
      console.log(`[Torch] Registered as ${data.id}, hasTorch: ${data.hasTorch}`)
      // Don't set torch state here — let torch-update handle the transition
      // so connectToMCP() is triggered correctly
      break
    }

    case 'torch-update': {
      const hadTorch = torchStore.hasTorch
      const hasTorchNow = data.holderId === clientId

      torchStore.setClients(data.clients)
      torchStore.setTorchState(data.holderId)

      if (hadTorch && !hasTorchNow) {
        console.log('[Torch] Lost torch, disconnecting from MCP')
        disconnectFromMCP()
      } else if (!hadTorch && hasTorchNow) {
        console.log('[Torch] Got torch, connecting to MCP')
        await connectToMCP()
      }

      // Auto-request: if no one holds the torch and we have autoRequest enabled
      if (!hasTorchNow && data.holderId === null && torchStore.autoRequest) {
        console.log('[Torch] Auto-requesting torch (no holder)')
        requestTorch()
      }
      break
    }

    case 'granted': {
      console.log('[Torch] Torch granted!')
      torchStore.setRequesting(false)
      break
    }

    case 'released': {
      console.log('[Torch] Torch released')
      break
    }
  }
}

/**
 * Request the torch
 */
export async function requestTorch(): Promise<boolean> {
  if (!torchWs || torchWs.readyState !== WebSocket.OPEN) {
    console.error('[Torch] Not connected to server')
    return false
  }

  const torchStore = useTorchStore()
  torchStore.setRequesting(true)

  torchWs.send(JSON.stringify({ type: 'request' }))
  return true
}

/**
 * Release the torch
 */
export async function releaseTorch(): Promise<boolean> {
  if (!torchWs || torchWs.readyState !== WebSocket.OPEN) {
    console.error('[Torch] Not connected to server')
    return false
  }

  const torchStore = useTorchStore()
  if (!torchStore.hasTorch) {
    console.warn('[Torch] Cannot release - do not have torch')
    return false
  }

  torchWs.send(JSON.stringify({ type: 'release' }))
  return true
}

/**
 * Transfer the torch to a specific client
 */
export async function transferTorch(targetId: string): Promise<boolean> {
  if (!torchWs || torchWs.readyState !== WebSocket.OPEN) {
    console.error('[Torch] Not connected to server')
    return false
  }

  torchWs.send(JSON.stringify({ type: 'transfer', targetId }))
  return true
}

/**
 * Update client name on server
 */
export function updateName(name: string): void {
  const torchStore = useTorchStore()
  torchStore.setClientName(name)

  if (torchWs?.readyState === WebSocket.OPEN) {
    torchWs.send(JSON.stringify({ type: 'update-name', name }))
  }
}

/**
 * Get list of connected clients
 */
export function getTorchClients() {
  const torchStore = useTorchStore()
  return torchStore.clients
}

/**
 * Get current torch status
 */
export function getTorchStatus() {
  const torchStore = useTorchStore()
  return {
    clientId: torchStore.clientId,
    hasTorch: torchStore.hasTorch,
    torchHolderId: torchStore.torchHolderId,
    clientCount: torchStore.clients.length
  }
}

/**
 * Connect to MCP (when we have the torch)
 */
async function connectToMCP(): Promise<void> {
  console.log('[Torch] Connecting to MCP...')
  const success = await autoConnect()
  if (success) {
    console.log('[Torch] Connected to MCP, activating tools')
    await onTorchConnected()
  } else {
    console.error('[Torch] Failed to connect to MCP')
  }
}

/**
 * Disconnect from MCP (when we lose the torch)
 */
function disconnectFromMCP(): void {
  console.log('[Torch] Disconnecting from MCP...')
  onTorchDisconnected()
  disconnectWebMCP()
}

/**
 * Initialize torch system
 */
export async function initTorch(): Promise<void> {
  await connectToTorchServer()

  // Cleanup on page unload
  window.addEventListener('beforeunload', () => {
    if (reconnectTimeout) {
      clearTimeout(reconnectTimeout)
    }
    torchWs?.close()
  })
}

/**
 * Cleanup torch system
 */
export function destroyTorch(): void {
  if (reconnectTimeout) {
    clearTimeout(reconnectTimeout)
    reconnectTimeout = null
  }
  torchWs?.close()
  torchWs = null
  clientId = null

  const torchStore = useTorchStore()
  torchStore.reset()
}