@animalabs/agent-framework 0.21.0 → 0.22.0

package/src/code-execution/runtime-py.ts

@@ -37,6 +37,10 @@
  *                {op:"wake", id, exec_id, line, payload}
  *                {op:"exec_result", id, stdout, stderr, return_code, tail?}
  *
+ * At a deadline the host sends cancel and then SIGINT, repeated until the
+ * script reports (not on Windows): the cancel stops a script waiting on an
+ * await, the signal also one stuck in blocking code (see _on_sigint).
+ *
  * IMPORTANT: the python source below must not contain backticks or the
  * sequence dollar+brace (TS template literal syntax). String.raw preserves
  * backslashes, so \n inside python string literals is fine.
@@ -49,6 +53,7 @@ import inspect
 import io
 import json
 import os
+import signal
 import sys
 import threading
 import traceback
@@ -169,6 +174,90 @@ _next_wake_id = 0
 _current_exec_id = None
 _current_exec_task = None
 
+# How the host stops a script (its time limit, or an abort): ONE per-script
+# state, so the script is stopped exactly once, by whichever path gets there
+# first. Every other path then does nothing to that script.
+#   PENDING   -> nothing asked yet.
+#   REQUESTED -> asked while the script was outside its own code (waiting on
+#                an await, writing a protocol line, or not started yet). A
+#                cancellation is queued on the loop, or delivered as the
+#                script starts.
+#   DELIVERED -> KeyboardInterrupt raised in its code, or its task cancelled.
+#                Final: whatever the script does next (except/finally, awaits
+#                in its cleanup) runs to completion or to the host's kill.
+# Transitions happen only in the main thread: in the SIGINT handler (between
+# bytecodes, so it never cancels anything itself -- see _on_sigint) and in
+# loop callbacks. Neither can interleave a transition of the other halfway
+# in a way that matters: the handler only moves PENDING to REQUESTED or, with
+# the script's own frame on the stack, to DELIVERED; a loop callback never
+# runs while the script's frame is on the stack.
+#
+# Where the signal (or the cancel op) can land, and what happens:
+#   1. script's own code, running or blocking -> KeyboardInterrupt there
+#   2. script waiting on an await (loop idle) -> cancellation queued; it wakes
+#      the loop and lands at that await
+#   3. main() mid-dispatch, between the done() check of a future the script
+#      awaits and its set_result (reply behind the cancellation) -> queued;
+#      set_result completes; the script gets CancelledError instead of it
+#   4. the script writing a protocol line (send) -> queued; the line stays
+#      whole; the cancellation lands at the script's next await
+#   5. reply ahead of the queued cancellation, script then awaits -> the
+#      cancellation lands at that await
+#   6. reply ahead, script then blocks -> the next signal (the host repeats
+#      it) raises KeyboardInterrupt there; the queued cancellation then does
+#      nothing, so cleanup that awaits is not cut short
+#   7. reply ahead, script finishes before the queued cancellation runs ->
+#      it does nothing, and never reaches the next script
+#   8. script not started yet (signal, or the cancel op) -> delivered as
+#      CancelledError when it starts, so it still reports
+#   9. cancel op before the signal -> delivered; the signal and its repeats
+#      do nothing, so cleanup after a caught CancelledError runs on
+#  10. cancel op or repeated signal after delivery -> nothing
+#  11. script finished, reporting its result -> nothing
+#  12. between scripts -> nothing
+#  13. a previous script installed an eager task factory -> no effect: the
+#      runtime builds its script task directly, and the state is in place
+#      before the script's first step
+_PENDING = "pending"
+_REQUESTED = "requested"
+_DELIVERED = "delivered"
+
+
+class _ScriptStop:
+    __slots__ = ("task", "running", "state", "interrupt")
+
+    def __init__(self):
+        self.task = None
+        self.running = False  # the script's own code may be on the stack
+        self.state = _PENDING
+        self.interrupt = None  # the KeyboardInterrupt raised in its code
+
+
+_stop = None  # the running script's _ScriptStop; None between scripts
+
+
+def _deliver_cancel(st):
+    # Loop context only: the cancel op in main()'s dispatch, or queued by the
+    # SIGINT handler. Never from inside the handler, where a cancellation
+    # could land between main()'s done() check of a future the script awaits
+    # and its set_result.
+    if st is None or st is not _stop or st.state == _DELIVERED:
+        return
+    if not st.running:
+        st.state = _REQUESTED  # not started yet: delivered as it starts
+        return
+    st.state = _DELIVERED
+    st.task.cancel()
+
+
+def _resolve(fut, value):
+    # Defense in depth: the script may have been cancelled since the lookup.
+    if fut is not None and not fut.done():
+        try:
+            fut.set_result(value)
+        except asyncio.InvalidStateError:
+            pass
+
 
 def _make_tool_fn(tool_name, py_name):
     async def tool_fn(args=None):
@@ -279,8 +368,8 @@ def handle_init(msg):
         SCRIPT_GLOBALS.pop("wake_agent", None)
 
 
-async def _run_script(exec_id, code):
-    global _current_exec_id, _current_exec_task
+async def _run_script(exec_id, code, st):
+    global _current_exec_id, _current_exec_task, _stop
     if BACKGROUND:
         out = LogTee(LOG_PATH, TAIL_CHARS)
         err = out  # interleave, terminal-style; tail is shared
@@ -293,17 +382,38 @@ async def _run_script(exec_id, code):
     return_code = 0
     try:
         compiled = compile(code, "<script>", "exec", flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT)
-        result = eval(compiled, SCRIPT_GLOBALS)
-        if inspect.iscoroutine(result):
-            await result
+        st.running = True
+        try:
+            if st.state == _REQUESTED:
+                # Stopped before it started (row 8).
+                st.state = _DELIVERED
+                raise asyncio.CancelledError()
+            result = eval(compiled, SCRIPT_GLOBALS)
+            if inspect.iscoroutine(result):
+                await result
+        finally:
+            st.running = False
     except asyncio.CancelledError:
         err.write("\nKeyboardInterrupt: script cancelled by host\n")
         return_code = 1
-    except BaseException:
-        traceback.print_exc(file=err)
+    except BaseException as exc:
+        if exc is st.interrupt:
+            # Where the script was when it was stopped, without the runtime's frames.
+            frames = [
+                f for f in traceback.extract_tb(exc.__traceback__)
+                if f.filename != _RUN_SCRIPT_CODE.co_filename
+            ]
+            err.write("Traceback (most recent call last):\n")
+            err.write("".join(traceback.format_list(frames)))
+            err.write("KeyboardInterrupt: script interrupted by host\n")
+        else:
+            traceback.print_exc(file=err)
         return_code = 1
     finally:
         sys.stdout, sys.stderr, sys.stdin = old_out, old_err, old_in
+        st.interrupt = None
+        if _stop is st:
+            _stop = None
         _current_exec_id = None
         _current_exec_task = None
         for fut in list(_pending_tool_futures.values()):
@@ -343,8 +453,47 @@ async def _run_script(exec_id, code):
     })
 
 
+def _dispatch(msg):
+    """Handle one protocol message from the host; False means stop reading."""
+    global _current_exec_id, _current_exec_task, _stop
+    op = msg.get("op")
+    if op == "init":
+        handle_init(msg)
+    elif op == "exec":
+        if _current_exec_task is not None and not _current_exec_task.done():
+            send({
+                "op": "exec_result",
+                "id": msg.get("id"),
+                "stdout": "",
+                "stderr": "RuntimeError: another script is already running in this interpreter",
+                "return_code": 1,
+            })
+            return True
+        # Everything _on_sigint reads is in place before the script can run
+        # (row 13). The task is built directly, not through the loop's task
+        # factory: a script may have installed one (interpreter state
+        # persists), and an eager one would run the next script inside this
+        # call, before the assignments below.
+        st = _ScriptStop()
+        _stop = st
+        _current_exec_id = msg.get("id")
+        _current_exec_task = asyncio.Task(
+            _run_script(msg.get("id"), msg.get("code") or "", st),
+            loop=asyncio.get_running_loop(),
+        )
+        st.task = _current_exec_task
+    elif op == "tool_result":
+        _resolve(_pending_tool_futures.get(msg.get("id")), str(msg.get("result", "")))
+    elif op == "wake_ack":
+        _resolve(_pending_wake_futures.get(msg.get("id")), msg.get("error") or None)
+    elif op == "cancel":
+        _deliver_cancel(_stop)
+    elif op == "exit":
+        return False
+    return True
+
+
 async def main():
-    global _current_exec_id, _current_exec_task
     loop = asyncio.get_running_loop()
     queue = asyncio.Queue()
 
@@ -353,7 +502,17 @@ async def main():
             loop.call_soon_threadsafe(queue.put_nowait, raw_line)
         loop.call_soon_threadsafe(queue.put_nowait, None)
 
-    threading.Thread(target=reader, daemon=True).start()
+    # The reader thread blocks SIGINT (it inherits the mask while it starts),
+    # so the signal reaches the main thread, where a blocking call in the
+    # script is waiting to be interrupted.
+    masked = hasattr(signal, "pthread_sigmask") and hasattr(signal, "SIGINT")
+    if masked:
+        previous_mask = signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGINT})
+    try:
+        threading.Thread(target=reader, daemon=True).start()
+    finally:
+        if masked:
+            signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask)
     send({"op": "ready"})
 
     while True:
@@ -369,35 +528,7 @@ async def main():
             REAL_STDERR.write("[pytc-runtime] bad protocol line\n")
             REAL_STDERR.flush()
             continue
-        op = msg.get("op")
-        if op == "init":
-            handle_init(msg)
-        elif op == "exec":
-            if _current_exec_task is not None and not _current_exec_task.done():
-                send({
-                    "op": "exec_result",
-                    "id": msg.get("id"),
-                    "stdout": "",
-                    "stderr": "RuntimeError: another script is already running in this interpreter",
-                    "return_code": 1,
-                })
-                continue
-            _current_exec_id = msg.get("id")
-            _current_exec_task = asyncio.ensure_future(
-                _run_script(msg.get("id"), msg.get("code") or "")
-            )
-        elif op == "tool_result":
-            fut = _pending_tool_futures.get(msg.get("id"))
-            if fut is not None and not fut.done():
-                fut.set_result(str(msg.get("result", "")))
-        elif op == "wake_ack":
-            fut = _pending_wake_futures.get(msg.get("id"))
-            if fut is not None and not fut.done():
-                fut.set_result(msg.get("error") or None)
-        elif op == "cancel":
-            if _current_exec_task is not None and not _current_exec_task.done():
-                _current_exec_task.cancel()
-        elif op == "exit":
+        if not _dispatch(msg):
             break
 
     # Drain: give a cancelled exec a moment to emit its exec_result.
@@ -409,6 +540,52 @@ async def main():
             pass
 
 
+# At the deadline the host sends a cancel op, which lands only when the script
+# awaits: a script blocked in time.sleep(), a busy loop or a blocking read
+# never gives the event loop control back. So the host also sends SIGINT,
+# repeated until the script reports, and this stops the script through its
+# _ScriptStop (see the table there):
+#   - With the script's own code on the stack (_run_script's frame, reached
+#     before any send frame), it raises KeyboardInterrupt there; _run_script
+#     catches it, so the output so far and the interpreter's globals survive.
+#     Checking the current task instead is not enough: asyncio runs its own
+#     code for the task between steps, and an exception raised there escapes
+#     the event loop and ends the interpreter.
+#   - Anywhere else it only queues the cancellation on the loop (which also
+#     wakes select(), otherwise resuming its wait after a signal): cancelling
+#     here could cancel a future between main()'s done() check and its
+#     set_result, and an exception mid-send would tear a protocol line.
+_RUN_SCRIPT_CODE = _run_script.__code__
+_SEND_CODE = send.__code__
+
+
+def _on_sigint(signum, frame):
+    st = _stop
+    if st is None or st.state == _DELIVERED:
+        return
+    if st.running:
+        while frame is not None:
+            code = frame.f_code
+            if code is _SEND_CODE:
+                break
+            if code is _RUN_SCRIPT_CODE:
+                st.state = _DELIVERED
+                st.interrupt = KeyboardInterrupt("script interrupted by host")
+                raise st.interrupt
+            frame = frame.f_back
+    if st.state == _PENDING:
+        st.state = _REQUESTED
+        if st.running:
+            st.task.get_loop().call_soon_threadsafe(_deliver_cancel, st)
+
+
 if __name__ == "__main__":
+    # A handler of our own also keeps asyncio.run from installing one that
+    # would cancel main() and end the interpreter.
+    if hasattr(signal, "SIGINT"):
+        try:
+            signal.signal(signal.SIGINT, _on_sigint)
+        except (ValueError, OSError):
+            pass
     asyncio.run(main())
 `;