@newsails/veil-studio 1.0.0

package/server/utils/veil-ws.js

@@ -0,0 +1,207 @@
+'use strict'
+
+const WebSocket = require('ws')
+const { VEIL_WS_URL, VEIL_SECRET } = require('./config.js')
+const { getIO } = require('./socket-io.js')
+const boardState = require('./board-state.js')
+const { initSessionStates } = require('./session-watcher.js')
+
+let _activeStreams = null
+let _setElementStreamingFn = null
+
+// Sessions that received iteration.start but not yet session.stream done
+// Used to detect sse=false sessions (agent_message tool) and emit chat:done fallback
+const _typingSessions = new Set()
+
+function setActiveStreamsRef(ref) {
+  _activeStreams = ref
+}
+
+function setElementStreamingRef(fn) {
+  _setElementStreamingFn = fn
+}
+
+let _ws = null
+let _connected = false
+let _retryDelay = 2000
+const MAX_RETRY_DELAY = 60000
+
+function isVeilConnected() {
+  return _connected
+}
+
+function connectVeilWS() {
+  const wsUrl = `${VEIL_WS_URL}/ws`
+  const headers = {}
+  if (VEIL_SECRET) headers['X-Veil-Secret'] = VEIL_SECRET
+
+  try {
+    _ws = new WebSocket(wsUrl, { headers })
+  } catch (err) {
+    console.warn(`[veil-ws] Could not connect to ${wsUrl}: ${err.message}`)
+    _scheduleReconnect()
+    return
+  }
+
+  _ws.on('open', () => {
+    _connected = true
+    _retryDelay = 2000
+    console.log(`[veil-ws] Connected to ${wsUrl}`)
+    try {
+      getIO().emit('veil:status', { connected: true })
+    } catch (_) {}
+    initSessionStates().catch(e => console.warn('[veil-ws] initSessionStates:', e.message))
+  })
+
+  _ws.on('message', (data) => {
+    try {
+      const event = JSON.parse(data.toString())
+      try { getIO().emit('veil:event', event) } catch (_) {}
+
+      // Update element state from task/session events
+      _handleVeilEvent(event)
+    } catch (err) {
+      console.error('[veil-ws] Failed to parse message:', err.message)
+    }
+  })
+
+  _ws.on('close', () => {
+    _connected = false
+    console.log('[veil-ws] Disconnected. Reconnecting...')
+    try { getIO().emit('veil:status', { connected: false }) } catch (_) {}
+    _scheduleReconnect()
+  })
+
+  _ws.on('error', (err) => {
+    console.error('[veil-ws] Error:', err.message)
+    // close event will fire next
+  })
+}
+
+function _scheduleReconnect() {
+  setTimeout(() => {
+    connectVeilWS()
+    _retryDelay = Math.min(_retryDelay * 2, MAX_RETRY_DELAY)
+  }, _retryDelay)
+}
+
+function _isExternalSession(sessionId) {
+  return !_activeStreams || !_activeStreams.has(sessionId)
+}
+
+function _handleVeilEvent(event) {
+  if (!event || !event.type) return
+
+  const io = getIO()
+
+  // ── session.stream — mirrors SSE stream for every session in real-time ──────
+  if (event.type === 'session.stream' && event.sessionId) {
+    if (!_isExternalSession(event.sessionId)) return  // Studio's chat:send owns this
+
+    const { sessionId, eventType, data } = event
+
+    if (eventType === 'error') {
+      try { io.to('board').emit('chat:error', { sessionId, error: data?.error ?? 'Unknown error' }) } catch (_) {}
+    } else if (eventType === 'done') {
+      _typingSessions.delete(sessionId)
+      try {
+        io.to('board').emit('chat:stream', { sessionId, eventType: 'done', data })
+        io.to('board').emit('chat:done', { sessionId })
+      } catch (_) {}
+      // Clear element streaming indicator
+      if (_setElementStreamingFn) _setElementStreamingFn(sessionId, false)
+    } else {
+      // inference.chunk, inference.tool, message (assistant/tool)
+      try { io.to('board').emit('chat:stream', { sessionId, eventType, data }) } catch (_) {}
+    }
+    return
+  }
+
+  // ── chat.user_message — user message before AI starts ───────────────────────
+  if (event.type === 'chat.user_message' && event.sessionId) {
+    if (!_isExternalSession(event.sessionId)) return
+    try {
+      io.to('board').emit('chat:user-message', {
+        sessionId:  event.sessionId,
+        agentName:  event.agentName || '',
+        message:    event.event?.content ?? '',
+      })
+    } catch (_) {}
+    return
+  }
+
+  // ── chat.event — iteration start → typing indicator + element streaming ─────
+  if (event.type === 'chat.event' && event.sessionId) {
+    if (event.event?.type === 'iteration.start' && _isExternalSession(event.sessionId)) {
+      _typingSessions.add(event.sessionId)
+      try { io.to('board').emit('chat:typing', { sessionId: event.sessionId }) } catch (_) {}
+      if (_setElementStreamingFn) _setElementStreamingFn(event.sessionId, true)
+    }
+    return
+  }
+
+  // ── chat.inference_tool — tool indicator for non-SSE sessions ───────────────
+  // Equivalent of session.stream eventType inference.tool for HTTP/agent_message paths
+  if (event.type === 'chat.inference_tool' && event.sessionId) {
+    if (_isExternalSession(event.sessionId)) {
+      try {
+        io.to('board').emit('chat:stream', {
+          sessionId: event.sessionId,
+          eventType: 'inference.tool',
+          data: { name: event.event?.name ?? '' },
+        })
+      } catch (_) {}
+    }
+    return
+  }
+
+  // ── chat.message — sse=false final message (agent_message tool invocation) ────
+  // Emitted instead of session.stream events when the session was called with sse=false
+  if (event.type === 'chat.message' && event.sessionId) {
+    if (_isExternalSession(event.sessionId)) {
+      try {
+        io.to('board').emit('chat:stream', {
+          sessionId: event.sessionId,
+          eventType: 'message',
+          data: event.event,
+        })
+      } catch (_) {}
+    }
+    return
+  }
+
+  // ── chat.response — fallback: stop streaming if session.stream done never fired ─
+  if (event.type === 'chat.response' && event.sessionId) {
+    const sessionId = event.sessionId
+    if (_isExternalSession(sessionId) && _typingSessions.has(sessionId)) {
+      _typingSessions.delete(sessionId)
+      try {
+        io.to('board').emit('chat:done', { sessionId })
+      } catch (_) {}
+      if (_setElementStreamingFn) _setElementStreamingFn(sessionId, false)
+    }
+    return
+  }
+
+  // ── task.status — update task element state on board ────────────────────────
+  if (event.type === 'task.status' && event.taskId && event.event) {
+    const status = event.event.status
+    if (!status) return
+    const stateMap = {
+      running:   'running',
+      finished:  'completed',
+      failed:    'error',
+      cancelled: 'default',
+    }
+    const elementState = stateMap[status]
+    if (!elementState) return
+
+    for (const el of boardState.getElements()) {
+      if (el.type === 'task' && el.data && el.data.taskId === event.taskId) {
+        boardState.updateElementState(el.id, elementState)
+      }
+    }
+  }
+}
+
+module.exports = { connectVeilWS, isVeilConnected, setActiveStreamsRef, setElementStreamingRef }