Fix stop generating: actually abort stream instead of draining in background

When user clicks Stop, worker_registry.stop_conversation() now marks the
conversation as explicitly stopped. The CancelledError and GeneratorExit
handlers in _execute_chat_completions_stream check this flag — if the
conversation was user-stopped, the stream is closed immediately instead
of being drained to completion in a background task.

The background drain is preserved for passive disconnects (browser close,
network drop) where the agent should continue processing so the response
is saved for when the user returns.

Author: Josh-XT

agixt/WorkerRegistry.py

@@ -21,6 +21,7 @@ def __init__(self):
         self._conversation_tasks: Dict[str, asyncio.Task] = (
             {}
         )  # conversation_id -> task
+        self._stopped_conversations: Set[str] = set()  # explicitly stopped by user
         self._lock = threading.Lock()
 
     def register_conversation(
@@ -76,8 +77,15 @@ def unregister_conversation(self, conversation_id: str) -> bool:
             if conversation_id in self._conversation_tasks:
                 del self._conversation_tasks[conversation_id]
 
+            self._stopped_conversations.discard(conversation_id)
+
             return removed
 
+    def is_stopped(self, conversation_id: str) -> bool:
+        """Check if a conversation was explicitly stopped by the user."""
+        with self._lock:
+            return conversation_id in self._stopped_conversations
+
     def get_conversation_info(self, conversation_id: str) -> Optional[Dict]:
         """
         Get information about an active conversation
@@ -147,6 +155,10 @@ async def stop_conversation(
                 )
                 return False
 
+            # Mark as explicitly stopped so the stream handler knows not to
+            # continue processing in the background.
+            self._stopped_conversations.add(conversation_id)
+
             # Get the task if it exists
             task = self._conversation_tasks.get(conversation_id)
 

agixt/XT.py

@@ -4988,6 +4988,20 @@ async def _next_event(iterator):
                 yield f"data: {json.dumps(tts_end_chunk)}\n\n"
 
         except asyncio.CancelledError:
+            # Check if user explicitly stopped this conversation
+            if worker_registry.is_stopped(conversation_id):
+                logging.info(
+                    f"[_execute_chat_completions_stream] Conversation {conversation_id} "
+                    f"stopped by user — aborting stream."
+                )
+                # Close the stream iterator without draining
+                if stream_iter is not None:
+                    try:
+                        await stream_iter.aclose()
+                    except Exception:
+                        pass
+                raise
+
             logging.warning(
                 f"[_execute_chat_completions_stream] CancelledError for conversation {conversation_id}. "
                 f"Detaching run_stream to continue in background."
@@ -5032,6 +5046,19 @@ async def _drain_stream_background(stream, pending, conv_id, _SEND):
                 )
             raise
         except GeneratorExit:
+            # Check if user explicitly stopped this conversation
+            if worker_registry.is_stopped(conversation_id):
+                logging.info(
+                    f"[_execute_chat_completions_stream] GeneratorExit: conversation {conversation_id} "
+                    f"stopped by user — aborting stream."
+                )
+                if stream_iter is not None:
+                    try:
+                        await stream_iter.aclose()
+                    except Exception:
+                        pass
+                return
+
             logging.warning(
                 f"[_execute_chat_completions_stream] GeneratorExit for conversation {conversation_id}. "
                 f"Detaching run_stream to continue in background."