Add force-stop for in-flight runs

agent_core/core/impl/action/cancellation.py

@@ -0,0 +1,165 @@
+# -*- coding: utf-8 -*-
+"""
+core.impl.action.cancellation
+
+Per-session registry of kill handles for force-stopping a run.
+
+Cancelling a turn's asyncio task aborts LLM calls and async actions, but it
+cannot reach real OS work already in flight: a shell command spawned by
+``run_shell`` (blocking a pool thread in ``communicate()``) or the python
+child of a sandboxed action (spawned inside a ProcessPoolExecutor worker).
+This module is the one place such work is registered so a user stop can
+kill it.
+
+Two mechanisms, one kill call:
+
+- ``register_process`` / ``unregister_process``: in-process registry of
+  ``subprocess.Popen`` handles, used by actions running in the main process
+  (thread-pool actions like ``run_shell``).
+- ``mark_subprocess`` / ``unmark_subprocess``: pid marker FILES under the
+  system temp dir, used by code running in a DIFFERENT process (the
+  sandboxed-action pool worker) where no in-memory registry can be shared.
+
+``kill_session_processes(session_id)`` kills both kinds, entire process
+trees included, and is safe to call at any time (missing/exited processes
+are ignored). It is blocking (taskkill / killpg) — call it from a worker
+thread, not the event loop.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import tempfile
+import threading
+from pathlib import Path
+from typing import Dict
+
+from agent_core.utils.logger import logger
+
+_lock = threading.Lock()
+# session_id -> {pid: Popen}. Popen handles registered by in-process actions.
+_procs: Dict[str, Dict[int, subprocess.Popen]] = {}
+
+
+def _marker_dir(session_id: str) -> Path:
+    return Path(tempfile.gettempdir()) / "craftbot_run_cancel" / session_id
+
+
+# ─────────────────────── In-process Popen registry ───────────────────────
+
+
+def register_process(session_id: str, proc: subprocess.Popen) -> None:
+    """Register a live child process as killable when this session is stopped."""
+    if not session_id or proc is None or proc.pid is None:
+        return
+    with _lock:
+        _procs.setdefault(session_id, {})[proc.pid] = proc
+
+
+def unregister_process(session_id: str, proc: subprocess.Popen) -> None:
+    """Remove a child process from the kill set (it finished normally)."""
+    if not session_id or proc is None or proc.pid is None:
+        return
+    with _lock:
+        session = _procs.get(session_id)
+        if session:
+            session.pop(proc.pid, None)
+            if not session:
+                _procs.pop(session_id, None)
+
+
+# ─────────────────────── Cross-process pid markers ───────────────────────
+
+
+def mark_subprocess(session_id: str, pid: int) -> None:
+    """Record a child pid from ANOTHER process (e.g. a pool worker).
+
+    The main process cannot hold the Popen handle, so the pid is written as
+    a marker file that ``kill_session_processes`` scans.
+    """
+    if not session_id or not pid:
+        return
+    try:
+        d = _marker_dir(session_id)
+        d.mkdir(parents=True, exist_ok=True)
+        (d / f"{pid}.pid").write_text(str(pid), encoding="utf-8")
+    except Exception:
+        pass  # markers are best-effort; never fail the action over them
+
+
+def unmark_subprocess(session_id: str, pid: int) -> None:
+    """Remove a pid marker (the child exited normally)."""
+    if not session_id or not pid:
+        return
+    try:
+        (_marker_dir(session_id) / f"{pid}.pid").unlink(missing_ok=True)
+    except Exception:
+        pass
+
+
+# ─────────────────────── Kill ───────────────────────
+
+
+def _kill_tree(pid: int) -> None:
+    """Kill a process and its descendants. Missing processes are fine."""
+    try:
+        if os.name == "nt":
+            subprocess.run(
+                ["taskkill", "/F", "/T", "/PID", str(pid)],
+                capture_output=True,
+                timeout=10,
+                creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
+            )
+        else:
+            import signal
+
+            try:
+                os.killpg(pid, signal.SIGKILL)
+            except (ProcessLookupError, PermissionError, OSError):
+                os.kill(pid, signal.SIGKILL)
+    except Exception as e:
+        logger.debug(f"[CANCEL] Kill of pid {pid} failed (likely already gone): {e}")
+
+
+def kill_session_processes(session_id: str) -> int:
+    """Force-kill every process registered/marked for a session.
+
+    Returns the number of kill targets attempted. Blocking — run in a
+    worker thread.
+    """
+    if not session_id:
+        return 0
+
+    with _lock:
+        handles = list(_procs.pop(session_id, {}).values())
+
+    killed = 0
+    for proc in handles:
+        if proc.poll() is None:
+            _kill_tree(proc.pid)
+            killed += 1
+            try:
+                proc.wait(timeout=5)
+            except Exception:
+                pass
+
+    # Cross-process markers (sandboxed action children).
+    try:
+        d = _marker_dir(session_id)
+        if d.is_dir():
+            for marker in d.glob("*.pid"):
+                try:
+                    _kill_tree(int(marker.stem))
+                    killed += 1
+                except ValueError:
+                    pass
+                marker.unlink(missing_ok=True)
+    except Exception as e:
+        logger.debug(f"[CANCEL] Marker sweep failed for {session_id}: {e}")
+
+    if killed:
+        logger.info(
+            f"[CANCEL] Force-killed {killed} process tree(s) for session {session_id}"
+        )
+    return killed

agent_core/core/impl/action/executor.py

@@ -391,16 +391,35 @@ def _atomic_action_venv_process(
                 encoding="utf-8",
             )
 
-            proc = subprocess.run(
+            # Popen (not subprocess.run) so the child's pid can be marked in
+            # the cross-process cancel registry: this function runs in a pool
+            # WORKER process, and a user force-stop issued in the main
+            # process kills marked pids by scanning the marker files.
+            from agent_core.core.impl.action.cancellation import (
+                mark_subprocess,
+                unmark_subprocess,
+            )
+
+            cancel_session_id = (input_data or {}).get("_session_id") or ""
+            proc = subprocess.Popen(
                 [str(python_bin), str(action_file)],
-                capture_output=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
                 text=True,
-                timeout=timeout,
             )
+            mark_subprocess(cancel_session_id, proc.pid)
+            try:
+                stdout, stderr = proc.communicate(timeout=timeout)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                proc.communicate()
+                raise
+            finally:
+                unmark_subprocess(cancel_session_id, proc.pid)
 
             return {
-                "stdout": proc.stdout.strip(),
-                "stderr": proc.stderr.strip(),
+                "stdout": (stdout or "").strip(),
+                "stderr": (stderr or "").strip(),
                 "returncode": proc.returncode,
             }
 
@@ -471,20 +490,35 @@ def _atomic_action_internal_subprocess(
         )
 
         try:
-            proc = subprocess.run(
+            from agent_core.core.impl.action.cancellation import (
+                mark_subprocess,
+                unmark_subprocess,
+            )
+
+            cancel_session_id = (input_data or {}).get("_session_id") or ""
+            popen = subprocess.Popen(
                 [python_bin, str(action_file)],
-                capture_output=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
                 text=True,
-                timeout=timeout,
             )
-
-            if proc.returncode != 0:
-                err = (
-                    proc.stderr.strip() or f"Action exited with code {proc.returncode}"
+            mark_subprocess(cancel_session_id, popen.pid)
+            try:
+                proc_stdout, proc_stderr = popen.communicate(timeout=timeout)
+            except subprocess.TimeoutExpired:
+                popen.kill()
+                popen.communicate()
+                raise
+            finally:
+                unmark_subprocess(cancel_session_id, popen.pid)
+
+            if popen.returncode != 0:
+                err = (proc_stderr or "").strip() or (
+                    f"Action exited with code {popen.returncode}"
                 )
                 return {"status": "error", "message": err}
 
-            stdout = proc.stdout.strip()
+            stdout = (proc_stdout or "").strip()
             if not stdout:
                 return {"status": "success", "output": ""}
 

agent_core/core/impl/event_stream/event_stream.py

@@ -154,11 +154,11 @@ def _append_summarization_notice(
         self, *, folded_events: int, folded_tokens: int, summary: str | None
     ) -> None:
         """Append a SYSTEM event announcing that summarization ran, so the UI
-        surfaces it as a system message in the session's chat. The LLM-facing
-        `message` stays a one-liner (the summary itself already lives in
-        head_summary — repeating it in the tail would double its token cost);
-        the full summary rides on `display_message`, which only the UI reads.
-        Caller holds the lock."""
+        surfaces it as a system message in the session's chat. Both the
+        LLM-facing `message` and the UI-facing `display_message` are
+        one-liners: the summary text itself lives only in head_summary
+        (repeating it in the tail would double its token cost, and dumping
+        it into the chat drowns the conversation). Caller holds the lock."""
         line = (
             f"Summarized {folded_events} older events (~{folded_tokens} tokens) "
             "into the running head summary."
@@ -169,16 +169,13 @@ def _append_summarization_notice(
                 f"(~{folded_tokens} tokens) without a summary."
             )
             display = (
-                "Context summarization was triggered but the summary could not "
-                f"be generated. The {folded_events} oldest events "
-                f"(~{folded_tokens} tokens) were pruned to keep the context lean."
+                f"Event stream summarization failed, {folded_tokens} tokens "
+                "were pruned without a summary"
             )
         else:
             display = (
-                f"Context summarization triggered: the {folded_events} oldest "
-                f"events (~{folded_tokens} tokens) were folded into a running "
-                "summary to keep the context lean.\n\n"
-                f"**Summary of folded events:**\n\n{summary}"
+                f"Summarized event stream, {folded_tokens} tokens were folded "
+                "into summary"
             )
         ev = Event(
             message=line,

agent_core/core/impl/trigger/session_queue.py

@@ -123,6 +123,27 @@ async def pop_due_batch(self) -> List[Trigger]:
                 batch.append(heapq.heappop(self._heap))
             return [entry[2] for entry in batch]
 
+    async def purge(self, predicate) -> int:
+        """Remove queued triggers matching ``predicate`` (a Trigger -> bool).
+
+        Used by user force-stop to drop a run's pending continuation rows
+        without touching unrelated triggers (user messages, schedules).
+        Removed triggers are reported to the lifecycle listener so their
+        durable rows settle instead of rehydrating next boot. Returns the
+        number of triggers removed.
+        """
+        async with self._cv:
+            if self._closed or not self._heap:
+                return 0
+            kept = [entry for entry in self._heap if not predicate(entry[2])]
+            removed = [entry[2] for entry in self._heap if predicate(entry[2])]
+            if not removed:
+                return 0
+            self._heap = kept
+            heapq.heapify(self._heap)
+            self._notify_evicted(removed)
+            return len(removed)
+
     async def close(self) -> List[Trigger]:
         """Close the queue (session deletion) and return discarded triggers.