From f8924f252102738dfafddc658ced22a121bdec16 Mon Sep 17 00:00:00 2001 From: Debanjum Date: Thu, 24 Jul 2025 17:16:26 -0500 Subject: [PATCH] Avoid duplicate chat turn save if chat cancelled during final response Save to conversation in normal flow should only be done if interrupt wasn't triggered. Saving conversations on interrupt is handled completely by the disconnect monitor since the improvements to interrupt. This abort is handled correctly for steps before final response. But not if interrupt occurs while final response is being sent. This changes checks for cancellation after final response send attempt and avoids duplicate chat turn save. --- src/khoj/routers/api_chat.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py index d792eadc..13cc0325 100644 --- a/src/khoj/routers/api_chat.py +++ b/src/khoj/routers/api_chat.py @@ -1425,6 +1425,13 @@ async def event_generator( logger.warning(f"Error during streaming. Stopping send: {e}") break + # Check if the user has disconnected + if cancellation_event.is_set(): + logger.debug(f"Stopping LLM response to user {user} on {common.client} client.") + # Cancel the disconnect monitor task if it is still running + await cancel_disconnect_monitor() + return + # Save conversation once finish streaming asyncio.create_task( save_to_conversation_log( @@ -1450,16 +1457,16 @@ async def event_generator( ) # Signal end of LLM response after the loop finishes - if not cancellation_event.is_set(): - async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""): - yield result - # Send Usage Metadata once llm interactions are complete - if tracer.get("usage"): - async for event in send_event(ChatEvent.USAGE, tracer.get("usage")): - yield event - async for result in send_event(ChatEvent.END_RESPONSE, ""): - yield result - logger.debug("Finished streaming response") + async for result in send_event(ChatEvent.END_LLM_RESPONSE, ""): + yield result + + # Send Usage Metadata once llm interactions are complete + if tracer.get("usage"): + async for event in send_event(ChatEvent.USAGE, tracer.get("usage")): + yield event + async for result in send_event(ChatEvent.END_RESPONSE, ""): + yield result + logger.debug("Finished streaming response") # Cancel the disconnect monitor task if it is still running await cancel_disconnect_monitor()