claude-code/packages/remote-control-server/src/transport/acp-ws-handler.ts

import type { WSContext } from 'hono/ws'
import { randomUUID } from 'node:crypto'
import { getAcpEventBus } from './event-bus'
import type { SessionEvent } from './event-bus'
import {
  storeCreateEnvironment,
  storeGetEnvironment,
  storeMarkAcpAgentOffline,
  storeMarkAcpAgentOnline,
  storeUpdateEnvironment,
} from '../store'
import { config } from '../config'
import { log, error as logError } from '../logger'

// Per-connection state
interface AcpConnectionEntry {
  agentId: string | null // Set after register message
  channelGroupId: string
  unsub: (() => void) | null
  keepalive: ReturnType<typeof setInterval> | null
  ws: WSContext
  openTime: number
  lastClientActivity: number
  capabilities: Record<string, unknown> | null
}

const connections = new Map<string, AcpConnectionEntry>() // key: wsId

const SERVER_KEEPALIVE_INTERVAL_MS = config.wsKeepaliveInterval * 1000
const CLIENT_ACTIVITY_TIMEOUT_MS = SERVER_KEEPALIVE_INTERVAL_MS * 3

/** Send a JSON message to a WS connection (NDJSON format) */
function sendToWs(ws: WSContext, msg: object): void {
  if (ws.readyState !== 1) return
  try {
    ws.send(JSON.stringify(msg) + '\n')
  } catch (err) {
    logError('[ACP-WS] send error:', err)
  }
}

/** Called from onOpen — initializes connection tracking */
export function handleAcpWsOpen(ws: WSContext, wsId: string): void {
  log(`[ACP-WS] Connection opened: wsId=${wsId}`)

  const keepalive = setInterval(() => {
    const entry = connections.get(wsId)
    if (!entry || entry.ws.readyState !== 1) {
      clearInterval(keepalive)
      return
    }
    const silenceMs = Date.now() - entry.lastClientActivity
    if (silenceMs > CLIENT_ACTIVITY_TIMEOUT_MS) {
      log(
        `[ACP-WS] Client inactive for ${Math.round(silenceMs / 1000)}s, closing dead connection`,
      )
      try {
        entry.ws.close(1000, 'client inactive')
      } catch {
        clearInterval(keepalive)
      }
      return
    }
    sendToWs(entry.ws, { type: 'keep_alive' })
  }, SERVER_KEEPALIVE_INTERVAL_MS)

  connections.set(wsId, {
    agentId: null,
    channelGroupId: '',
    unsub: null,
    keepalive,
    ws,
    openTime: Date.now(),
    lastClientActivity: Date.now(),
    capabilities: null,
  })
}

/** Handle register message — legacy WS-only registration (still supported) */
function handleRegister(wsId: string, msg: Record<string, unknown>): void {
  const entry = connections.get(wsId)
  if (!entry) return

  if (entry.agentId) {
    sendToWs(entry.ws, { type: 'error', message: 'Already registered' })
    return
  }

  const agentName = (msg.agent_name as string) || 'unknown'
  const capabilities = msg.capabilities as Record<string, unknown> | undefined
  const channelGroupId =
    (msg.channel_group_id as string) ||
    `group_${randomUUID().replace(/-/g, '').slice(0, 12)}`
  const acpLinkVersion = (msg.acp_link_version as string) || null
  const maxSessions =
    typeof msg.max_sessions === 'number' ? msg.max_sessions : 1

  // Create EnvironmentRecord with workerType="acp"
  const secret = config.apiKeys[0] || ''
  const record = storeCreateEnvironment({
    secret,
    machineName: agentName,
    workerType: 'acp',
    bridgeId: channelGroupId,
    maxSessions,
    capabilities: capabilities || undefined,
  } as Parameters<typeof storeCreateEnvironment>[0])

  // Store ACP-specific metadata via environment update
  storeUpdateEnvironment(record.id, {
    status: 'active',
  } as Parameters<typeof storeUpdateEnvironment>[1])

  entry.agentId = record.id
  entry.channelGroupId = channelGroupId
  entry.capabilities = capabilities || null

  // Subscribe to channel group EventBus — broadcast events to this WS
  const bus = getAcpEventBus(channelGroupId)
  const unsub = bus.subscribe((event: SessionEvent) => {
    if (entry.ws.readyState !== 1) return
    if (event.direction !== 'outbound') return
    // Forward outbound events as raw ACP messages
    sendToWs(entry.ws, event.payload as object)
  })
  entry.unsub = unsub

  log(
    `[ACP-WS] Agent registered (legacy WS): agentId=${record.id} channelGroup=${channelGroupId} name=${agentName}`,
  )
  sendToWs(entry.ws, {
    type: 'registered',
    agent_id: record.id,
    channel_group_id: channelGroupId,
  })
}

/** Handle identify message — binds WS to an existing agent registered via REST */
function handleIdentify(wsId: string, msg: Record<string, unknown>): void {
  const entry = connections.get(wsId)
  if (!entry) return

  if (entry.agentId) {
    sendToWs(entry.ws, { type: 'error', message: 'Already identified' })
    return
  }

  const agentId = msg.agent_id as string
  if (!agentId) {
    sendToWs(entry.ws, { type: 'error', message: 'Missing agent_id' })
    return
  }

  // Look up the environment record (created via REST registration)
  const record = storeGetEnvironment(agentId)
  if (!record || record.workerType !== 'acp') {
    sendToWs(entry.ws, { type: 'error', message: 'Agent not found' })
    return
  }

  // Update status to active
  storeMarkAcpAgentOnline(agentId)

  const channelGroupId =
    record.bridgeId || `group_${randomUUID().replace(/-/g, '').slice(0, 12)}`

  entry.agentId = record.id
  entry.channelGroupId = channelGroupId
  entry.capabilities = record.capabilities || null

  // Subscribe to channel group EventBus — broadcast events to this WS
  const bus = getAcpEventBus(channelGroupId)
  const unsub = bus.subscribe((event: SessionEvent) => {
    if (entry.ws.readyState !== 1) return
    if (event.direction !== 'outbound') return
    sendToWs(entry.ws, event.payload as object)
  })
  entry.unsub = unsub

  log(
    `[ACP-WS] Agent identified (REST+WS): agentId=${record.id} channelGroup=${channelGroupId}`,
  )
  sendToWs(entry.ws, {
    type: 'identified',
    agent_id: record.id,
    channel_group_id: channelGroupId,
  })
}

/** Called from onMessage — processes NDJSON lines */
export function handleAcpWsMessage(
  ws: WSContext,
  wsId: string,
  data: string,
): void {
  const entry = connections.get(wsId)
  if (!entry) return

  entry.lastClientActivity = Date.now()

  const lines = data.split('\n').filter(l => l.trim())
  for (const line of lines) {
    let msg: Record<string, unknown>
    try {
      msg = JSON.parse(line)
    } catch {
      logError('[ACP-WS] parse error:', line)
      continue
    }

    // Handle keepalive
    if (msg.type === 'keep_alive') {
      // Update last activity timestamp (only if registered)
      if (entry.agentId) {
        storeUpdateEnvironment(entry.agentId, {
          lastPollAt: new Date(),
        } as Parameters<typeof storeUpdateEnvironment>[1])
      }
      continue
    }

    // Handle registration (legacy WS-only)
    if (msg.type === 'register') {
      handleRegister(wsId, msg)
      continue
    }

    // Handle identify (REST registration + WS binding)
    if (msg.type === 'identify') {
      handleIdentify(wsId, msg)
      continue
    }

    // Not registered yet — reject
    if (!entry.agentId) {
      sendToWs(entry.ws, {
        type: 'error',
        message: 'Not registered. Send register message first.',
      })
      continue
    }

    // Update agent activity
    storeUpdateEnvironment(entry.agentId, {
      lastPollAt: new Date(),
    } as Parameters<typeof storeUpdateEnvironment>[1])

    // Pass-through: publish to channel group EventBus as inbound
    const bus = getAcpEventBus(entry.channelGroupId)
    bus.publish({
      id: randomUUID(),
      sessionId: entry.channelGroupId,
      type: (msg.type as string) || 'acp_message',
      payload: msg,
      direction: 'inbound',
    })
  }
}

/** Called from onClose — marks agent offline and cleans up */
export function handleAcpWsClose(
  ws: WSContext,
  wsId: string,
  code?: number,
  reason?: string,
): void {
  const entry = connections.get(wsId)
  if (!entry) return

  const duration = Math.round((Date.now() - entry.openTime) / 1000)
  log(
    `[ACP-WS] Connection closed: wsId=${wsId} agentId=${entry.agentId} code=${code ?? 'none'} reason=${reason || '(none)'} duration=${duration}s`,
  )

  if (entry.unsub) {
    entry.unsub()
  }
  if (entry.keepalive) {
    clearInterval(entry.keepalive)
  }

  // Mark agent as offline (don't delete record — allow reconnect)
  if (entry.agentId) {
    storeMarkAcpAgentOffline(entry.agentId)

    // Notify all relay connections that this agent is gone
    if (entry.channelGroupId) {
      const bus = getAcpEventBus(entry.channelGroupId)
      bus.publish({
        id: randomUUID(),
        sessionId: entry.channelGroupId,
        type: 'agent_disconnect',
        payload: { agentId: entry.agentId },
        direction: 'inbound',
      })
    }
  }

  connections.delete(wsId)
}

/** Find an active ACP connection by agent ID */
export function findAcpConnectionByAgentId(
  agentId: string,
): AcpConnectionEntry | null {
  for (const entry of connections.values()) {
    if (entry.agentId === agentId && entry.ws.readyState === 1) {
      return entry
    }
  }
  return null
}

/** Send a JSON message directly to an agent's WebSocket connection */
export function sendToAgentWs(agentId: string, msg: object): boolean {
  const entry = findAcpConnectionByAgentId(agentId)
  if (!entry) return false
  sendToWs(entry.ws, msg)
  return true
}

/** Gracefully close all ACP WebSocket connections */
export function closeAllAcpConnections(): void {
  if (connections.size === 0) return

  log(`[ACP-WS] Gracefully closing ${connections.size} ACP connection(s)...`)
  for (const [wsId, entry] of connections) {
    try {
      if (entry.unsub) entry.unsub()
      if (entry.keepalive) clearInterval(entry.keepalive)
      if (entry.ws.readyState === 1) {
        entry.ws.close(1001, 'server_shutdown')
      }
      if (entry.agentId) {
        storeMarkAcpAgentOffline(entry.agentId)
      }
    } catch {
      // ignore errors during shutdown
    }
  }
  connections.clear()
  log('[ACP-WS] All connections closed')
}