diff --git a/src/interface/web/app/chat/page.tsx b/src/interface/web/app/chat/page.tsx index 796290cf..c34e0185 100644 --- a/src/interface/web/app/chat/page.tsx +++ b/src/interface/web/app/chat/page.tsx @@ -260,6 +260,34 @@ export default function Chat() { if (idleTimerRef.current) { clearTimeout(idleTimerRef.current); } + // Mark any in-progress streamed message as completed so UI updates (stop spinner, show send icon) + setMessages((prev) => { + if (!prev || prev.length === 0) return prev; + const newMessages = [...prev]; + const last = newMessages[newMessages.length - 1]; + if (last && !last.completed) { + last.completed = true; + } + return newMessages; + }); + // Reset processing state so ChatInputArea send button reappears + setProcessQuerySignal(false); + setQueryToProcess(""); + }, + onError: (event) => { + console.error("WebSocket error", event); + // Perform same cleanup as onClose to avoid stuck UI + setMessages((prev) => { + if (!prev || prev.length === 0) return prev; + const newMessages = [...prev]; + const last = newMessages[newMessages.length - 1]; + if (last && !last.completed) { + last.completed = true; + } + return newMessages; + }); + setProcessQuerySignal(false); + setQueryToProcess(""); }, }); diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py index c647b9a7..411cb73a 100644 --- a/src/khoj/routers/api_chat.py +++ b/src/khoj/routers/api_chat.py @@ -786,6 +786,9 @@ async def event_generator( if interrupt_query == ChatEvent.END_EVENT.value: cancellation_event.set() logger.debug(f"Chat cancelled by user {user} via interrupt queue.") + elif interrupt_query == ChatEvent.INTERRUPT.value: + cancellation_event.set() + logger.debug("Chat interrupted.") else: # Pass the interrupt query to child tasks logger.info(f"Continuing chat with the new instruction: {interrupt_query}") @@ -1556,7 +1559,7 @@ async def chat_ws( except WebSocketDisconnect: logger.info(f"WebSocket disconnected for user {websocket.scope['user'].object.id}") if current_task and not current_task.done(): - current_task.cancel() + interrupt_queue.put_nowait(ChatEvent.INTERRUPT.value) except Exception as e: logger.error(f"Error in websocket chat: {e}", exc_info=True) if current_task and not current_task.done():