From ef0634135c79676280487e1f46811be35189ea46 Mon Sep 17 00:00:00 2001 From: Alishahryar1 Date: Wed, 28 Jan 2026 22:16:44 -0800 Subject: [PATCH] fixed /stop command to make sure it kills everything --- providers/claude_cli.py | 11 ++++++-- providers/message_queue.py | 57 +++++++++++++++++++++++++++++--------- server.py | 18 +++++++++++- 3 files changed, 70 insertions(+), 16 deletions(-) diff --git a/providers/claude_cli.py b/providers/claude_cli.py index e09e0ae6..eb780723 100644 --- a/providers/claude_cli.py +++ b/providers/claude_cli.py @@ -281,9 +281,16 @@ class CLISession: async def stop(self): if self.process and self.process.returncode is None: try: + logger.info(f"Stopping Claude CLI process {self.process.pid}") self.process.terminate() - await self.process.wait() + try: + await asyncio.wait_for(self.process.wait(), timeout=5.0) + except asyncio.TimeoutError: + logger.warning(f"Process {self.process.pid} did not terminate, killing...") + self.process.kill() + await self.process.wait() return True - except: + except Exception as e: + logger.error(f"Error stopping session process: {e}") return False return False diff --git a/providers/message_queue.py b/providers/message_queue.py index f35419c2..8303f4d8 100644 --- a/providers/message_queue.py +++ b/providers/message_queue.py @@ -7,7 +7,7 @@ Messages are processed one-by-one in order per session. import asyncio import logging -from typing import Callable, Awaitable, Dict, Optional, NamedTuple +from typing import Callable, Awaitable, Dict, Optional, NamedTuple, List, Any from dataclasses import dataclass logger = logging.getLogger(__name__) @@ -20,7 +20,7 @@ class QueuedMessage: chat_id: int msg_id: int reply_msg_id: int # The status message to update - event: any # Original Telegram event for context + event: Any # Original Telegram event for context class SessionQueue: @@ -31,6 +31,7 @@ class SessionQueue: self.queue: asyncio.Queue[QueuedMessage] = asyncio.Queue() self.is_processing = False self.current_task: Optional[asyncio.Task] = None + self.current_message: Optional[QueuedMessage] = None class MessageQueueManager: @@ -91,7 +92,8 @@ class MessageQueueManager: sq.is_processing = True # Process outside the lock - await self._process_message(session_id, message, processor) + sq = self._queues[session_id] + sq.current_task = asyncio.create_task(self._process_message(session_id, message, processor)) return False async def _process_message( @@ -101,11 +103,19 @@ class MessageQueueManager: processor: Callable[[str, QueuedMessage], Awaitable[None]], ) -> None: """Process a single message and then check the queue.""" + sq = self._queues.get(session_id) + if sq: + sq.current_message = message try: await processor(session_id, message) + except asyncio.CancelledError: + logger.info(f"Task for session {session_id} was cancelled") + raise except Exception as e: logger.error(f"Error processing message for session {session_id}: {e}") finally: + if sq: + sq.current_message = None # Check if there are more messages in the queue await self._process_next(session_id, processor) @@ -136,7 +146,7 @@ class MessageQueueManager: return # Process next message (outside lock) - await self._process_message(session_id, next_msg, processor) + sq.current_task = asyncio.create_task(self._process_message(session_id, next_msg, processor)) def get_queue_size(self, session_id: str) -> int: """Get the number of messages waiting in a session's queue.""" @@ -144,26 +154,47 @@ class MessageQueueManager: return 0 return self._queues[session_id].queue.qsize() - def cancel_session(self, session_id: str) -> int: + def cancel_session(self, session_id: str) -> List[QueuedMessage]: """ - Cancel all queued messages for a session. + Cancel all queued messages for a session and the running task. Returns: - Number of messages that were cancelled + List of messages that were cancelled (including the current one if any) """ if session_id not in self._queues: - return 0 + return [] sq = self._queues[session_id] - cancelled = 0 + cancelled_messages = [] + # 1. Cancel running task + if sq.current_task and not sq.current_task.done(): + sq.current_task.cancel() + if sq.current_message: + cancelled_messages.append(sq.current_message) + + # 2. Clear queue while not sq.queue.empty(): try: - sq.queue.get_nowait() - cancelled += 1 + msg = sq.queue.get_nowait() + cancelled_messages.append(msg) except asyncio.QueueEmpty: break sq.is_processing = False - logger.info(f"Cancelled {cancelled} queued messages for session {session_id}") - return cancelled + logger.info(f"Cancelled {len(cancelled_messages)} messages for session {session_id}") + return cancelled_messages + + async def cancel_all(self) -> List[QueuedMessage]: + """ + Cancel everything in all sessions. + + Returns: + List of all cancelled messages across all sessions. + """ + async with self._lock: + all_cancelled = [] + session_ids = list(self._queues.keys()) + for sid in session_ids: + all_cancelled.extend(self.cancel_session(sid)) + return all_cancelled diff --git a/server.py b/server.py index d9dea196..d05c1884 100644 --- a/server.py +++ b/server.py @@ -529,7 +529,7 @@ def register_bot_handlers(client: "TelegramClient"): except asyncio.CancelledError: logger.info(f"BOT: Task cancelled for session {captured_session_id or temp_session_id}") message_parts.append(("error", "Task was cancelled")) - await update_bot_ui("⏹ **Cancelled**", force=True) + await update_bot_ui("❌ **Failed**", force=True) except Exception as e: import traceback logger.error(f"Bot task failed: {e}\n{traceback.format_exc()}") @@ -555,8 +555,24 @@ def register_bot_handlers(client: "TelegramClient"): # 1. Handle Commands if event.text == "/stop": + # 1. Cancel all queued and running messages in the message queue + cancelled_msgs = await message_queue.cancel_all() + + # 2. Stop all CLI sessions (subprocesses) await cli_session_manager.stop_all() + + # 3. Inform the user and update status of cancelled messages await event.reply("⏹ **All Claude sessions stopped.**") + + for msg in cancelled_msgs: + try: + # We might not be able to edit if it's too old or already done, + # but we try to mark them as failed as requested. + status_msg = await client.get_messages(msg.chat_id, ids=msg.reply_msg_id) + if status_msg: + await status_msg.edit("❌ **Failed** (Stopped by user)", parse_mode="markdown") + except Exception as e: + logger.debug(f"Could not update status for cancelled msg {msg.msg_id}: {e}") return if event.text == "/stats":