From 59bfaf9698abdacf0101d998f36dea32e6c32c00 Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Fri, 15 Aug 2025 12:48:28 -0700
Subject: [PATCH] Fix to indicate ws disconnect on web app & save interrupted
 research

- A regression had stopped indicating to user that the websocket
connection had broken. Now the interrupt has some visual indication.

- Websocket disconnects from client didn't trigger the partial
research to be saved. Now we use an interrupt signal to save partial
research before closing task.
---
 src/interface/web/app/chat/page.tsx | 28 ++++++++++++++++++++++++++++
 src/khoj/routers/api_chat.py        |  5 ++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

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():