fix: abort in-flight tool body on cooperative interrupt (fixes #4596)

Thread the InterruptController's cancellation Event into the tool
invocation context (AgentState.cancel_event / InvocationContext.cancel_event)
so a running tool body can observe an interrupt and abort promptly instead
of running until its own timeout. The built-in shell tool now polls the
event while its subprocess runs and kills the process group on cancel,
returning an `interrupted` outcome. Fully backward compatible: a tool that
ignores the event behaves exactly as before (no new hard deps, lazy imports).

Co-authored-by: MervinPraison <[email protected]>

src/praisonai-agents/praisonaiagents/agent/interrupt.py

@@ -31,7 +31,12 @@ class InterruptControllerProtocol(Protocol):
     def reason(self) -> Optional[str]:
         """Get interrupt reason if set."""
         ...
-    
+
+    @property
+    def event(self) -> "threading.Event":
+        """Get the underlying cancellation Event for in-tool observation."""
+        ...
+
     def check(self) -> None:
         """Check for interrupt and raise if set."""
         ...
@@ -126,6 +131,16 @@ class InterruptController:
         with self._lock:
             return self._reason
 
+    @property
+    def event(self) -> threading.Event:
+        """Expose the underlying cancellation Event.
+
+        Lets a running tool body observe the interrupt cooperatively (e.g. wait
+        on it with a timeout, or poll ``is_set()``) so an in-flight subprocess/
+        request can be aborted rather than left running until its own timeout.
+        """
+        return self._flag
+
     def check(self) -> None:
         """Check for cancellation and raise if requested.
         

src/praisonai-agents/praisonaiagents/agent/tool_execution.py

@@ -469,6 +469,12 @@ class ToolExecutionMixin:
         state_metadata = {'agent_name': self.name}
         if idempotency_key:
             state_metadata['idempotency_key'] = idempotency_key
+        # Thread the interrupt controller's cancellation Event into the tool
+        # context so a running tool body can observe it and abort promptly
+        # (e.g. shell tool killing its subprocess group). A tool that ignores
+        # this behaves exactly as before.
+        _controller = getattr(self, 'interrupt_controller', None)
+        cancel_event = getattr(_controller, 'event', None) if _controller is not None else None
         state = AgentState(
             agent_id=self.name,
             run_id=getattr(self, '_current_run_id', 'unknown'),
@@ -477,6 +483,7 @@ class ToolExecutionMixin:
             memory=getattr(self, '_memory_instance', None),
             learn_manager=getattr(getattr(self, '_memory_instance', None), 'learn', None),
             metadata=state_metadata,
+            cancel_event=cancel_event,
         )
         
         # Route through user-supplied tool middleware (Agent(hooks=[...])) when
@@ -496,6 +503,7 @@ class ToolExecutionMixin:
                 session_id=getattr(self, '_session_id', None) or 'default',
                 tool_name=function_name,
                 metadata={"tool_call_id": tool_call_id},
+                cancel_event=cancel_event,
             ),
         )
 

src/praisonai-agents/praisonaiagents/hooks/middleware.py

@@ -46,6 +46,7 @@ class InvocationContext:
     tool_name: Optional[str] = None
     model_name: Optional[str] = None
     metadata: Dict[str, Any] = field(default_factory=dict)
+    cancel_event: Optional[Any] = None
 
 
 @dataclass

src/praisonai-agents/praisonaiagents/tools/injected.py

@@ -63,7 +63,8 @@ class AgentState:
     learn_manager: Optional[Any] = None
     previous_tool_results: List[Any] = field(default_factory=list)
     metadata: Dict[str, Any] = field(default_factory=dict)
-    
+    cancel_event: Optional[Any] = None
+
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary for Injected[dict] usage."""
         return {
@@ -74,6 +75,7 @@ class AgentState:
             'last_agent_message': self.last_agent_message,
             'previous_tool_results': self.previous_tool_results,
             'metadata': self.metadata,
+            'cancel_event': self.cancel_event,
         }
 
 

src/praisonai-agents/praisonaiagents/tools/shell_tools.py

@@ -33,6 +33,12 @@ class _StreamDrainTimeout(Exception):
     process.
     """
 
+
+class _CommandCancelled(Exception):
+    """Raised when a cooperative cancellation signal fires while the child
+    process is still running, so the caller kills the process group and reports
+    an ``interrupted`` outcome instead of waiting for the timeout."""
+
 class ShellTools:
     """Tools for executing shell commands safely."""
     
@@ -125,11 +131,17 @@ class ShellTools:
                 text=True
             )
             
+            # Observe the agent's cooperative cancellation signal (when running
+            # inside an agent tool call) so an interrupt (/stop, Ctrl-C) aborts
+            # this subprocess promptly instead of waiting for the timeout. When
+            # run standalone the event is None and behaviour is unchanged.
+            cancel_event = self._current_cancel_event()
+
             try:
                 # Read stdout/stderr incrementally so a progress channel (when
                 # active) can surface output live while the command runs. Falls
                 # back to fully-buffered behaviour when no sink is listening.
-                stdout, stderr = self._communicate_streaming(process, timeout)
+                stdout, stderr = self._communicate_streaming(process, timeout, cancel_event)
 
                 # Handle over-budget output. Rather than permanently discarding
                 # the middle (often where the real error lives), spill the full
@@ -154,6 +166,20 @@ class ShellTools:
                     result['stderr_artifact'] = stderr_path
                 return result
             
+            except _CommandCancelled:
+                # A cooperative interrupt (/stop, Ctrl-C) fired while the child
+                # was still running. Kill the whole process group so descendants
+                # do not keep mutating state, then report an interrupted outcome.
+                self._kill_process_tree(process)
+                return {
+                    'stdout': '',
+                    'stderr': 'Command interrupted by user',
+                    'exit_code': -1,
+                    'success': False,
+                    'interrupted': True,
+                    'execution_time': time.time() - start_time
+                }
+
             except _StreamDrainTimeout:
                 # The direct child already exited; only an inherited pipe kept a
                 # reader blocked. Report a timeout but do NOT kill — the direct
@@ -169,20 +195,7 @@ class ShellTools:
 
             except subprocess.TimeoutExpired:
                 # Kill process on timeout
-                try:
-                    import psutil
-                    parent = psutil.Process(process.pid)
-                    children = parent.children(recursive=True)
-                    for child in children:
-                        child.kill()
-                    parent.kill()
-                except ImportError:
-                    # Fallback: kill without psutil
-                    process.kill()
-                except psutil.NoSuchProcess:
-                    # Process already gone — nothing to kill.
-                    pass
-                
+                self._kill_process_tree(process)
                 return {
                     'stdout': '',
                     'stderr': f'Command timed out after {timeout} seconds',
@@ -202,7 +215,48 @@ class ShellTools:
                 'execution_time': 0
             }
     
-    def _communicate_streaming(self, process, timeout):
+    def _current_cancel_event(self):
+        """Return the agent's cooperative cancellation Event, if any.
+
+        Sourced from the injected AgentState set by the agent tool-execution
+        loop. Returns ``None`` when the tool runs outside an agent (standalone),
+        keeping behaviour and imports lazy with zero overhead.
+        """
+        try:
+            from .injected import get_current_state
+            state = get_current_state()
+        except Exception:
+            return None
+        if state is None:
+            return None
+        return getattr(state, 'cancel_event', None)
+
+    def _kill_process_tree(self, process):
+        """Kill a child process and its descendants (best-effort).
+
+        Uses psutil to reach the whole tree when available, falling back to
+        ``process.kill()`` otherwise. Safe to call if the process already exited.
+        """
+        try:
+            import psutil
+            parent = psutil.Process(process.pid)
+            children = parent.children(recursive=True)
+            for child in children:
+                try:
+                    child.kill()
+                except psutil.NoSuchProcess:
+                    pass
+            parent.kill()
+        except ImportError:
+            try:
+                process.kill()
+            except Exception:
+                pass
+        except psutil.NoSuchProcess:
+            # Process already gone — nothing to kill.
+            pass
+
+    def _communicate_streaming(self, process, timeout, cancel_event=None):
         """Read stdout/stderr line-buffered, emitting live progress.
 
         Reads both streams concurrently in background threads so a line on
@@ -273,16 +327,38 @@ class ShellTools:
         t_err.start()
 
         deadline = None if timeout is None else time.monotonic() + timeout
-        try:
-            process.wait(timeout=timeout)
-        except subprocess.TimeoutExpired:
-            # The direct child is still running. Disable any further emission and
-            # re-raise TimeoutExpired so the caller's timeout path KILLS the
-            # process before returning — a recorded read error must not short-
-            # circuit that cleanup (otherwise the still-running child leaks with
-            # its pipes open).
-            stop_emitting.set()
-            raise
+        if cancel_event is None:
+            try:
+                process.wait(timeout=timeout)
+            except subprocess.TimeoutExpired:
+                # The direct child is still running. Disable any further emission
+                # and re-raise TimeoutExpired so the caller's timeout path KILLS
+                # the process before returning — a recorded read error must not
+                # short-circuit that cleanup (otherwise the still-running child
+                # leaks with its pipes open).
+                stop_emitting.set()
+                raise
+        else:
+            # Poll in short slices so a cooperative interrupt aborts the child
+            # promptly (within the poll interval) rather than only on timeout.
+            _POLL = 0.1
+            while True:
+                if cancel_event.is_set():
+                    stop_emitting.set()
+                    raise _CommandCancelled()
+                if deadline is not None:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        stop_emitting.set()
+                        raise subprocess.TimeoutExpired(process.args, timeout)
+                    wait_slice = min(_POLL, remaining)
+                else:
+                    wait_slice = _POLL
+                try:
+                    process.wait(timeout=wait_slice)
+                    break
+                except subprocess.TimeoutExpired:
+                    continue
 
         # Wait for the readers to drain the pipe buffers before returning,
         # preserving subprocess.communicate()'s "all captured output" semantics.

src/praisonai-agents/tests/unit/tools/test_shell_interrupt.py

@@ -0,0 +1,69 @@
+"""Tests for cooperative interrupt aborting an in-flight shell subprocess.
+
+Verifies that a running shell command is terminated promptly when the agent's
+cancellation Event fires (issue #4596), instead of running until its timeout.
+"""
+
+import os
+import sys
+import threading
+import time
+
+import pytest
+
+from praisonaiagents.tools.shell_tools import ShellTools
+from praisonaiagents.tools.injected import AgentState, with_injection_context
+from praisonaiagents.agent.interrupt import InterruptController
+
+
+@pytest.fixture(autouse=True)
+def _auto_approve(monkeypatch):
+    # execute_command is approval-gated (critical); auto-approve for the test.
+    monkeypatch.setenv("PRAISONAI_AUTO_APPROVE", "true")
+
+
+def _state_with(cancel_event):
+    return AgentState(
+        agent_id="test",
+        run_id="run",
+        session_id="sess",
+        cancel_event=cancel_event,
+    )
+
+
+def test_interrupt_aborts_running_command_within_grace():
+    """A long sleep is interrupted well before its timeout and reports interrupted."""
+    controller = InterruptController()
+    tools = ShellTools()
+
+    result_holder = {}
+
+    def run():
+        state = _state_with(controller.event)
+        with with_injection_context(state):
+            result_holder["result"] = tools.execute_command(
+                f"{sys.executable} -c \"import time; time.sleep(30)\"",
+                timeout=30,
+            )
+
+    worker = threading.Thread(target=run)
+    start = time.time()
+    worker.start()
+    time.sleep(0.5)  # let the subprocess start
+    controller.request("user_stop")
+    worker.join(timeout=10)
+    elapsed = time.time() - start
+
+    assert not worker.is_alive(), "shell command was not aborted"
+    assert elapsed < 10, f"interrupt took too long: {elapsed:.1f}s"
+    result = result_holder["result"]
+    assert result.get("interrupted") is True
+    assert result.get("success") is False
+
+
+def test_no_cancel_event_preserves_normal_behaviour():
+    """Without a cancel_event a quick command still returns normally."""
+    tools = ShellTools()
+    result = tools.execute_command(f"{sys.executable} -c \"print('hi')\"", timeout=10)
+    assert result.get("success") is True
+    assert "hi" in result.get("stdout", "")