feat: message abort and web socket support (#222)

* feat: message abort and web socket support

* fix(web): chat end

* fix: lint

* fix: lint

apps/web/src/composables/api/useChat.ts

@@ -2,3 +2,4 @@ export * from './useChat.types'
 export * from './useChat.chat-api'
 export * from './useChat.message-api'
 export * from './useChat.content'
+export * from './useChat.ws'

apps/web/src/composables/api/useChat.types.ts

@@ -47,7 +47,7 @@ export interface StreamEvent {
     | 'reasoning_start' | 'reasoning_delta' | 'reasoning_end'
     | 'tool_call_start' | 'tool_call_end'
     | 'attachment_delta'
-    | 'agent_start' | 'agent_end'
+    | 'agent_start' | 'agent_end' | 'agent_abort'
     | 'processing_started' | 'processing_completed' | 'processing_failed'
     | 'error'
   delta?: string

apps/web/src/composables/api/useChat.ws.ts

@@ -0,0 +1,139 @@
+import { client } from '@memoh/sdk/client'
+import type { StreamEvent, MessageStreamEvent, ChatAttachment, StreamEventHandler } from './useChat.types'
+
+export interface WSClientMessage {
+  type: 'message' | 'abort'
+  text?: string
+  attachments?: ChatAttachment[]
+}
+
+export interface ChatWebSocket {
+  send: (msg: WSClientMessage) => void
+  abort: () => void
+  close: () => void
+  readonly connected: boolean
+  onOpen: (() => void) | null
+  onClose: (() => void) | null
+}
+
+function resolveWebSocketUrl(botId: string): string {
+  const baseUrl = String(client.getConfig().baseUrl || '').trim()
+  const path = `/bots/${encodeURIComponent(botId)}/web/ws`
+
+  if (!baseUrl || baseUrl.startsWith('/')) {
+    const loc = window.location
+    const proto = loc.protocol === 'https:' ? 'wss:' : 'ws:'
+    const base = baseUrl || '/api'
+    return `${proto}//${loc.host}${base.replace(/\/+$/, '')}${path}`
+  }
+
+  try {
+    const url = new URL(path, baseUrl)
+    url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:'
+    return url.toString()
+  } catch {
+    const loc = window.location
+    const proto = loc.protocol === 'https:' ? 'wss:' : 'ws:'
+    return `${proto}//${loc.host}/api${path}`
+  }
+}
+
+export function connectWebSocket(
+  botId: string,
+  onStreamEvent: StreamEventHandler,
+  onMessageEvent?: (event: MessageStreamEvent) => void,
+): ChatWebSocket {
+  const id = botId.trim()
+  if (!id) throw new Error('bot id is required')
+
+  const wsUrl = resolveWebSocketUrl(id)
+  const token = localStorage.getItem('token') ?? ''
+  const url = token ? `${wsUrl}?token=${encodeURIComponent(token)}` : wsUrl
+
+  let ws: WebSocket | null = null
+  let isConnected = false
+  let closed = false
+  let reconnectTimer: ReturnType<typeof setTimeout> | null = null
+  let reconnectDelay = 1000
+
+  const handle: ChatWebSocket = {
+    send(msg: WSClientMessage) {
+      if (ws && ws.readyState === WebSocket.OPEN) {
+        ws.send(JSON.stringify(msg))
+      }
+    },
+    abort() {
+      if (ws && ws.readyState === WebSocket.OPEN) {
+        ws.send(JSON.stringify({ type: 'abort' }))
+      }
+    },
+    close() {
+      closed = true
+      if (reconnectTimer) {
+        clearTimeout(reconnectTimer)
+        reconnectTimer = null
+      }
+      if (ws) {
+        ws.close()
+        ws = null
+      }
+      isConnected = false
+    },
+    get connected() {
+      return isConnected
+    },
+    onOpen: null,
+    onClose: null,
+  }
+
+  function connect() {
+    if (closed) return
+    ws = new WebSocket(url)
+
+    ws.onopen = () => {
+      isConnected = true
+      reconnectDelay = 1000
+      handle.onOpen?.()
+    }
+
+    ws.onclose = () => {
+      isConnected = false
+      handle.onClose?.()
+      scheduleReconnect()
+    }
+
+    ws.onerror = () => {
+      // onerror is always followed by onclose; reconnect handled there.
+    }
+
+    ws.onmessage = (event) => {
+      if (typeof event.data !== 'string') return
+      try {
+        const parsed = JSON.parse(event.data)
+        if (!parsed || typeof parsed !== 'object') return
+
+        const eventType = String(parsed.type ?? '').trim()
+        if (eventType === 'message_created' && onMessageEvent) {
+          onMessageEvent(parsed as MessageStreamEvent)
+          return
+        }
+
+        onStreamEvent(parsed as StreamEvent)
+      } catch {
+        // Ignore unparsable messages.
+      }
+    }
+  }
+
+  function scheduleReconnect() {
+    if (closed) return
+    reconnectTimer = setTimeout(() => {
+      reconnectTimer = null
+      connect()
+    }, reconnectDelay)
+    reconnectDelay = Math.min(reconnectDelay * 1.5, 10000)
+  }
+
+  connect()
+  return handle
+}

apps/web/src/store/chat-list.ts

@@ -21,7 +21,9 @@ import {
   sendLocalChannelMessage,
   streamLocalChannel,
   streamMessageEvents,
+  connectWebSocket,
   type ChatAttachment,
+  type ChatWebSocket,
 } from '@/composables/api/useChat'
 
 // ---- Message model (blocks-based, aligned with main branch) ----
@@ -103,6 +105,7 @@ export const useChatStore = defineStore('chat', () => {
   let pendingAssistantStream: PendingAssistantStream | null = null
   const messageEventsStream = useRetryingStream()
   const localStream = useRetryingStream()
+  let activeWs: ChatWebSocket | null = null
 
   const participantChats = computed(() =>
     chats.value.filter((c) => (c.access_mode ?? 'participant') === 'participant'),
@@ -123,6 +126,7 @@ export const useChatStore = defineStore('chat', () => {
     } else {
       stopMessageEvents()
       stopLocalStream()
+      stopWebSocket()
       rejectPendingAssistantStream(new Error('Bot stream stopped'))
       messageEventsSince = ''
       chats.value = []
@@ -332,6 +336,9 @@ export const useChatStore = defineStore('chat', () => {
   // ---- Abort ----
 
   function abort() {
+    if (activeWs) {
+      activeWs.abort()
+    }
     abortFn?.()
     abortFn = null
     for (const msg of messages) {
@@ -356,6 +363,25 @@ export const useChatStore = defineStore('chat', () => {
     localStream.stop()
   }
 
+  function stopWebSocket() {
+    if (activeWs) {
+      activeWs.close()
+      activeWs = null
+    }
+  }
+
+  function startWebSocket(targetBotId: string) {
+    const bid = targetBotId.trim()
+    stopWebSocket()
+    if (!bid) return
+
+    activeWs = connectWebSocket(
+      bid,
+      handleLocalStreamEvent,
+      (e) => handleStreamEvent(bid, e),
+    )
+  }
+
   function pushAssistantBlock(session: PendingAssistantStream, block: ContentBlock): number {
     session.assistantMsg.blocks.push(block)
     return session.assistantMsg.blocks.length - 1
@@ -587,8 +613,11 @@ export const useChatStore = defineStore('chat', () => {
         rejectPendingAssistantStream(new Error(message))
         break
       }
-      case 'agent_start':
+      case 'agent_abort':
       case 'agent_end':
+        resolvePendingAssistantStream()
+        break
+      case 'agent_start':
       default: {
         const fallback = extractFallbackText(event)
         if (fallback) {
@@ -750,6 +779,7 @@ export const useChatStore = defineStore('chat', () => {
     loadingChats.value = true
     stopMessageEvents()
     stopLocalStream()
+    stopWebSocket()
     try {
       const bid = await ensureBot()
       if (!bid) {
@@ -772,6 +802,7 @@ export const useChatStore = defineStore('chat', () => {
         : visible[0]!.id
       chatId.value = activeChatId
       await loadMessages(bid, activeChatId)
+      startWebSocket(bid)
       startMessageEvents(bid)
       startLocalStream(bid)
     } finally {
@@ -900,10 +931,17 @@ export const useChatStore = defineStore('chat', () => {
       abortFn = () => {
         const abortError = new Error('aborted')
         abortError.name = 'AbortError'
+        if (activeWs) {
+          activeWs.abort()
+        }
         rejectPendingAssistantStream(abortError)
       }
 
-      await sendLocalChannelMessage(bid, trimmed, attachments)
+      if (activeWs?.connected) {
+        activeWs.send({ type: 'message', text: trimmed, attachments })
+      } else {
+        await sendLocalChannelMessage(bid, trimmed, attachments)
+      }
       await completion
 
       assistantMsg.streaming = false