fix: add bot run cancellation and timeout support (#1919)

* fix: add bot run cancellation and timeout support (fixes #1917)

- Add MESSAGE_ABORT event type to gateway protocols
- Enhance BotSessionManager with InterruptController integration
- Add configurable run_timeout parameter (default 300s)
- Implement cancel_run() method for user-initiated cancellation
- Add /stop command for cancelling active runs
- Add BotRunTimeout exception for timeout handling
- Maintain backward compatibility
- Follow protocol-driven architecture from AGENTS.md

Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com>

* fix: resolve critical bugs in bot run cancellation and timeout support

Fixes three critical issues identified by code review:

1. Fix parameter name mismatch: Agent.chat() expects 'cancel_token', not 'interrupt_controller'
   - Cooperative cancellation is now functional

2. Wire /stop command to all bot adapters (telegram, discord, slack, whatsapp)
   - Users can now actually cancel runs via /stop command

3. Prevent BotRunTimeout from being queued in DLQ to avoid infinite retry loops
   - Timeout exceptions are excluded from replay to prevent loops

Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com>

---------

Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com>
Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: 3 people

src/praisonai-agents/praisonaiagents/gateway/protocols.py

@@ -53,6 +53,7 @@ class EventType(str, Enum):
     # Message events
     MESSAGE = "message"
     MESSAGE_ACK = "message_ack"
+    MESSAGE_ABORT = "message_abort"
     TYPING = "typing"
 
     # Streaming events (relayed from agent's StreamEventEmitter)

src/praisonai/praisonai/bots/_commands.py

@@ -76,11 +76,57 @@ def format_help(
     return "\n".join(lines)
 
 
-async def handle_stop_command(
+def handle_stop_command(session_manager, user_id: str) -> str:
+    """Handle a /stop command to cancel an active run.
+    
+    This function works with both the legacy BotSessionManager approach
+    and the newer SessionRunControl approach for maximum compatibility.
+    
+    Args:
+        session_manager: BotSessionManager instance or SessionRunControl
+        user_id: User ID to cancel run for
+        
+    Returns:
+        Response message indicating success or failure
+    """
+    # Handle SessionRunControl (newer approach)
+    if hasattr(session_manager, 'stop'):
+        try:
+            import asyncio
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # We're in an async context, need to use create_task
+                task = asyncio.create_task(session_manager.stop(user_id))
+                # This is a synchronous function, so we can't await
+                # Return a message indicating async operation is happening
+                return "⏳ Cancellation requested..."
+            else:
+                # No event loop, run synchronously
+                stopped = asyncio.run(session_manager.stop(user_id))
+                if stopped:
+                    return "✅ Current task cancelled. Send a new message to start fresh."
+                else:
+                    return "ℹ️ No active task to cancel."
+        except Exception as e:
+            return f"❌ Error stopping task: {e}"
+    
+    # Handle BotSessionManager (legacy approach)
+    elif hasattr(session_manager, 'cancel_run'):
+        was_cancelled = session_manager.cancel_run(user_id, "user_stop_command")
+        if was_cancelled:
+            return "✅ Current task cancelled. Send a new message to start fresh."
+        else:
+            return "ℹ️ No active task to cancel."
+    
+    else:
+        return "❌ Stop command not available (run control not enabled)"
+
+
+async def handle_stop_command_async(
     user_id: str,
     run_control: Optional["SessionRunControl"] = None,
 ) -> str:
-    """Handle /stop command to cancel current agent task.
+    """Async version of handle_stop_command for SessionRunControl.
     
     Args:
         user_id: User identifier

src/praisonai/praisonai/bots/_session.py

@@ -26,6 +26,11 @@
 logger = logging.getLogger(__name__)
 
 
+class BotRunTimeout(Exception):
+    """Exception raised when a bot run times out."""
+    pass
+
+
 class BotSessionManager:
     """Lightweight per-user session store for bot agents.
 
@@ -66,6 +71,7 @@ def __init__(
         identity_resolver: Optional[Any] = None,
         ingress_journal: Optional[Any] = None,
         run_control: Optional[Any] = None,
+        run_timeout: float = 300.0,  # 5 minutes default timeout
     ) -> None:
         self._histories: Dict[str, List[Dict[str, Any]]] = {}
         self._locks: Dict[str, asyncio.Lock] = {}
@@ -89,6 +95,9 @@ def __init__(
         self._last_journal_key = None  # Store key for delayed completion
         # Run control for in-flight message handling
         self._run_control = run_control
+        # Run timeout and active run tracking for cancellation support
+        self._run_timeout = run_timeout
+        self._active_runs: Dict[str, Any] = {}  # user_id -> InterruptController
 
     def _storage_key(self, user_id: str) -> str:
         """Resolve a raw platform user id to the in-memory/store key.
@@ -275,29 +284,67 @@ async def chat(
                     async with agent_lock:
                         saved_history = agent.chat_history
                         agent.chat_history = user_history
+                        
+                        # Create interrupt controller for this run and register it
+                        try:
+                            from praisonaiagents.agent.interrupt import InterruptController
+                        except ImportError:
+                            # Fallback if InterruptController is not available
+                            InterruptController = None
+                        
+                        controller = InterruptController() if InterruptController else None
+                        storage_key = self._storage_key(user_id)
+                        if controller:
+                            self._active_runs[storage_key] = controller
+                        
                         try:
                             # Choose streaming vs non-streaming path based on callback
                             if stream_callback:
-                                # Streaming path: use agent.astart() with stream callback
-                                response = await agent.astart(prompt, stream_callback=stream_callback)
-                                # Handle AutonomyResult when autonomy is enabled in caller mode
-                                if hasattr(response, 'output'):
-                                    response = response.output
+                                # Streaming path: use agent.astart() with stream callback and timeout
+                                try:
+                                    # astart is async so we can use asyncio.wait_for directly 
+                                    response = await asyncio.wait_for(
+                                        agent.astart(prompt, stream_callback=stream_callback),
+                                        timeout=self._run_timeout if self._run_timeout > 0 else None,
+                                    )
+                                    # Handle AutonomyResult when autonomy is enabled in caller mode
+                                    if hasattr(response, 'output'):
+                                        response = response.output
+                                except asyncio.TimeoutError:
+                                    if controller:
+                                        controller.request("run timeout")
+                                    raise BotRunTimeout(f"Agent run timed out after {self._run_timeout}s")
                             else:
-                                # Legacy non-streaming path: use agent.chat() in executor
-                                # Copy current task's contextvars (incl. SessionContext)
-                                # into the worker thread so tools the agent invokes can
-                                # read platform/user metadata.
+                                # Legacy non-streaming path: use agent.chat() in executor with cancel_token and timeout
                                 import contextvars
+                                import inspect
+                                from functools import partial
                                 _ctx = contextvars.copy_context()
-                                response = await loop.run_in_executor(
-                                    None, _ctx.run, agent.chat, prompt
-                                )
+                                
+                                # Create agent.chat call with cancel_token if supported
+                                # Use inspect.signature for safer parameter checking
+                                _chat_params = inspect.signature(agent.chat).parameters if (controller and hasattr(agent, 'chat')) else {}
+                                if controller and 'cancel_token' in _chat_params:
+                                    chat_call = partial(agent.chat, prompt, cancel_token=controller)
+                                else:
+                                    chat_call = partial(agent.chat, prompt)
+                                
+                                # Run with timeout and interruption support
+                                try:
+                                    response = await asyncio.wait_for(
+                                        loop.run_in_executor(None, _ctx.run, chat_call),
+                                        timeout=self._run_timeout if self._run_timeout > 0 else None,
+                                    )
+                                except asyncio.TimeoutError:
+                                    if controller:
+                                        controller.request("run timeout")
+                                    raise BotRunTimeout(f"Agent run timed out after {self._run_timeout}s")
                             # Capture updated history before restoring caller's.
                             updated_history = agent.chat_history
                         except Exception as exc:
                             # N4: persist the failed inbound message before bubbling.
-                            if self._dlq is not None:
+                            # Skip DLQ for timeout exceptions to prevent infinite retry loops
+                            if self._dlq is not None and not isinstance(exc, BotRunTimeout):
                                 try:
                                     await loop.run_in_executor(
                                         None,
@@ -319,6 +366,9 @@ async def chat(
                             raise
                         finally:
                             agent.chat_history = saved_history
+                            # Clean up active run tracking
+                            if controller:
+                                self._active_runs.pop(storage_key, None)
 
                     # Persist outside the agent_lock — it's per-user and the agent
                     # is no longer touched.
@@ -532,6 +582,27 @@ def reset(self, user_id: str) -> bool:
 
         return existed
 
+    def cancel_run(self, user_id: str, reason: str = "user_cancel") -> bool:
+        """Cancel an active run for a user.
+        
+        Args:
+            user_id: Raw platform user id
+            reason: Reason for cancellation
+            
+        Returns:
+            bool: True if there was an active run to cancel, False otherwise
+        """
+        storage_key = self._storage_key(user_id)
+        controller = self._active_runs.get(storage_key)
+        if controller:
+            controller.request(reason)
+            return True
+        return False
+
+    def get_active_runs(self) -> List[str]:
+        """Get list of user IDs with active runs."""
+        return list(self._active_runs.keys())
+
     def _add_mirror_entry_sync(self, user_id: str, entry: dict) -> bool:
         """Thread-safe method to add a mirror entry to user's history.
         

src/praisonai/praisonai/bots/discord.py

@@ -24,7 +24,7 @@
     MessageType,
 )
 
-from ._commands import format_status, format_help
+from ._commands import format_status, format_help, handle_stop_command
 from ._session import BotSessionManager
 from ._debounce import InboundDebouncer
 from ._ack import AckReactor
@@ -200,6 +200,11 @@ async def on_message(message):
                 elif command == "help":
                     await message.reply(self._format_help())
                     return
+                elif command == "stop":
+                    user_id = str(message.author.id)
+                    response = handle_stop_command(self._session, user_id)
+                    await message.reply(response)
+                    return
                 elif command and command in self._command_handlers:
                     handler = self._command_handlers[command]
                     try:

src/praisonai/praisonai/bots/slack.py

@@ -26,7 +26,7 @@
 )
 
 from .media import split_media_from_output, is_audio_file
-from ._commands import format_status, format_help
+from ._commands import format_status, format_help, handle_stop_command
 from ._session import BotSessionManager
 from ._debounce import InboundDebouncer
 from ._ack import AckReactor
@@ -213,6 +213,11 @@ async def handle_message(event, say):
             elif text == "/help":
                 await say(text=self._format_help(), thread_ts=event.get("ts"))
                 return
+            elif text == "/stop":
+                user_id = event.get("user", "unknown")
+                response = handle_stop_command(self._session, user_id)
+                await say(text=response, thread_ts=event.get("ts"))
+                return
 
             for handler in self._message_handlers:
                 try:

src/praisonai/praisonai/bots/telegram.py

@@ -27,7 +27,7 @@
 )
 
 from .media import split_media_from_output, is_audio_file
-from ._commands import format_status, format_help
+from ._commands import format_status, format_help, handle_stop_command
 from ._session import BotSessionManager
 from ._debounce import InboundDebouncer
 from ._ack import AckReactor
@@ -398,9 +398,20 @@ async def handle_help(update: Update, context: ContextTypes.DEFAULT_TYPE):
                 return
             await update.message.reply_text(self._format_help())
 
+        async def handle_stop(update: Update, context: ContextTypes.DEFAULT_TYPE):
+            if not update.message:
+                return
+            message = await process_inbound_telegram_message(update, self)
+            if not message:
+                return
+            user_id = message.sender.user_id if message.sender else "unknown"
+            response = handle_stop_command(self._session, user_id)
+            await update.message.reply_text(response)
+        
         self._application.add_handler(CommandHandler("status", handle_status))
         self._application.add_handler(CommandHandler("new", handle_new))
         self._application.add_handler(CommandHandler("help", handle_help))
+        self._application.add_handler(CommandHandler("stop", handle_stop))
 
         for command in self._command_handlers:
             self._application.add_handler(CommandHandler(command, handle_command))

src/praisonai/praisonai/bots/whatsapp.py

@@ -43,7 +43,7 @@
     MessageType,
 )
 
-from ._commands import format_status, format_help
+from ._commands import format_status, format_help, handle_stop_command
 from ._session import BotSessionManager
 from ._rate_limit import RateLimiter
 from ._ack import AckReactor
@@ -175,13 +175,18 @@ async def _help(msg):
             extra = {
                 name: info.get("description", "")
                 for name, info in self._command_info.items()
-                if name not in ("status", "new", "help")
+                if name not in ("status", "new", "help", "stop")
             }
             return format_help(self._agent, "whatsapp", extra or None)
 
+        async def _stop(msg):
+            user_id = msg.sender.user_id if msg.sender else "unknown"
+            return handle_stop_command(self._session_mgr, user_id)
+
         self.register_command("status", _status, description="Show bot status and info")
         self.register_command("new", _new, description="Reset conversation session")
         self.register_command("help", _help, description="Show this help message")
+        self.register_command("stop", _stop, description="Cancel the current agent run")
 
     # ── Properties ──────────────────────────────────────────────────