agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/task/__init__.py

@@ -0,0 +1,27 @@
+"""Task module for AgentEnv."""
+
+from .task import Task
+from .store import (
+    TaskStore, TaskQuery, get_task_store, set_task_store, reset_task_store,
+    TaskInstance, TaskInstanceStore, get_task_instance_store, set_task_instance_store, reset_task_instance_store,
+    TaskStepResult, TaskStepStatus, record_step_complete, record_task_cancelled, record_task_failure,
+)
+
+__all__ = [
+    "Task",
+    "TaskStore",
+    "TaskQuery",
+    "get_task_store",
+    "set_task_store",
+    "reset_task_store",
+    "TaskInstance",
+    "TaskInstanceStore",
+    "TaskStepResult",
+    "TaskStepStatus",
+    "get_task_instance_store",
+    "set_task_instance_store",
+    "reset_task_instance_store",
+    "record_step_complete",
+    "record_task_cancelled",
+    "record_task_failure",
+]

agent_env/task/interrupts.py

@@ -0,0 +1,109 @@
+"""Ctrl-C and SIGTERM for a command that tears down what its runs deployed.
+
+``Interrupts`` counts both signals from when it is entered until it exits: inside its event loops and in the
+code between them, so none is lost at a hand-off. Inside ``gather``, the first signal cancels the runs that
+haven't begun their teardown and a second cancels the rest, stopping the teardown. It takes over the process's
+handlers, so it has an effect only in the main thread, and a signal the process inherited as ignored stays
+ignored.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import signal
+import threading
+import weakref
+from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator
+from contextlib import contextmanager
+from typing import Any
+
+SIGNALS = (signal.SIGINT, signal.SIGTERM)
+
+_tearing_down: weakref.WeakSet[asyncio.Task] = weakref.WeakSet()
+
+
+def tearing_down() -> None:
+    """Spare the current task from the first signal: it is tearing down what it deployed."""
+    task = asyncio.current_task()
+    if task is not None:
+        _tearing_down.add(task)
+
+
+class Interrupts:
+    def __init__(self) -> None:
+        self.signum: int | None = None  # the first signal
+        self.count = 0
+        self._previous: dict[int, Any] = {}
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._on_signal: Callable[[int], None] | None = None
+
+    def __enter__(self) -> Interrupts:
+        if threading.current_thread() is threading.main_thread():
+            for signum in SIGNALS:
+                if signal.getsignal(signum) is not signal.SIG_IGN:
+                    self._previous[signum] = signal.signal(signum, self._record)
+        return self
+
+    def __exit__(self, *exc) -> None:
+        for signum, previous in self._previous.items():
+            signal.signal(signum, signal.SIG_DFL if previous is None else previous)
+
+    @property
+    def reason(self) -> str:
+        return "cancelled by Ctrl-C" if self.signum in (None, signal.SIGINT) else (
+            f"cancelled by {signal.Signals(self.signum).name}")
+
+    def run(self, main: Coroutine) -> Any:
+        """``asyncio.run(main)``, taking the handlers back from the loop once it closes."""
+        try:
+            return asyncio.run(main)
+        finally:
+            for signum in self._previous:
+                signal.signal(signum, self._record)
+
+    async def gather(self, runs: Iterable[Awaitable], on_signal: Callable[[int], None]) -> list[asyncio.Future]:
+        """Run ``runs`` at once until each is done, and return them. Each signal calls ``on_signal`` with the
+        count so far; the first cancels the runs not tearing down, and every later one cancels them all."""
+        futures = [asyncio.ensure_future(run) for run in runs]
+
+        def cancel(count: int) -> None:
+            on_signal(count)
+            for future in futures:
+                if count > 1 or future not in _tearing_down:
+                    future.cancel()
+
+        with self._in_loop(cancel):
+            await asyncio.gather(*futures, return_exceptions=True)
+        return futures
+
+    async def stopping(self, runs: Iterable[Awaitable]) -> list:
+        """Await ``runs`` at once; any signal cancels them all."""
+        futures = [asyncio.ensure_future(run) for run in runs]
+        with self._in_loop(lambda count: [future.cancel() for future in futures]):
+            return list(await asyncio.gather(*futures))
+
+    @contextmanager
+    def _in_loop(self, on_signal: Callable[[int], None]) -> Iterator[None]:
+        """Pass each signal to ``on_signal`` in the running loop. The loop's own handlers wake it from any
+        thread, and they stay until it closes, counting what comes after, rather than dropping a signal that
+        is already on its way."""
+        loop = asyncio.get_running_loop()
+        self._loop, self._on_signal = loop, on_signal
+        for signum in self._previous:
+            try:
+                loop.add_signal_handler(signum, self._record, signum)
+            except (NotImplementedError, RuntimeError):
+                pass  # the process's handler reaches the loop instead
+        for count in range(1, self.count + 1):
+            loop.call_soon(on_signal, count)
+        try:
+            yield
+        finally:
+            self._loop = self._on_signal = None
+
+    def _record(self, signum: int, frame=None) -> None:
+        self.count += 1
+        if self.signum is None:
+            self.signum = signum
+        if self._on_signal is not None:
+            self._loop.call_soon_threadsafe(self._on_signal, self.count)

agent_env/task/registry.py

@@ -0,0 +1,27 @@
+"""Task registry for type-based deserialization."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from agent_env.task.task import Task
+
+
+def _get_type(cls: type["Task"]) -> str:
+    return cls.type
+
+
+_registry: dict[str, type["Task"]] | None = None
+
+
+def get_task_registry() -> dict[str, type["Task"]]:
+    """Get the task registry, lazily building it to avoid circular imports."""
+    global _registry
+    if _registry is None:
+        from agent_env.task.task import Task
+
+        _registry = {
+            _get_type(Task): Task,
+        }
+    return _registry

agent_env/task/step_journal.py

@@ -0,0 +1,99 @@
+"""The step journal's replay rules: how one step's diff is applied, in what order, and what
+that reconstructs. Pure — the store owns the collection, the sequence, and the CAS."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Iterable
+
+from agent_env.task_step.context import TaskStepContext
+from agent_env.task_step.context_ops import ContextUpdateOps
+
+_SEED_STEP_ID = "__seed__"
+# The scheduler's own record (retry audit), written by ``undo_steps_sync`` in the same CAS as
+# the undo and replayed right after the seed, so a retry's bookkeeping survives a later undo.
+_SCHEDULER_STEP_ID = "__scheduler__"
+_RESERVED_STEP_IDS = frozenset({_SEED_STEP_ID, _SCHEDULER_STEP_ID})
+
+def _get_path(doc: dict, dotted: str):
+    """Read a dotted path (e.g. 'context.metadata.x') from a doc, or None."""
+    cur: object = doc
+    for part in dotted.split("."):
+        if not isinstance(cur, dict):
+            return None
+        cur = cur.get(part)
+    return cur
+
+
+def _union(existing: list, items: list, *, move_to_end: bool = False, path: str = "") -> list:
+    """Union by value (deployed envs by instance), order-stable. ``move_to_end`` is the re-record case: that step's
+    completion moves to the end of the replay order, so its items must move with it."""
+    key = _item_key(path)
+    fresh = [x for i, x in enumerate(items) if key(x) not in [key(y) for y in items[:i]]]
+    fresh_keys = [key(x) for x in fresh]
+    if move_to_end:
+        return [x for x in existing if key(x) not in fresh_keys] + fresh
+    existing_keys = [key(x) for x in existing]
+    return list(existing) + [x for x, k in zip(fresh, fresh_keys) if k not in existing_keys]
+
+
+def _set_path(doc: dict, dotted: str, value) -> None:
+    cur = doc
+    parts = dotted.split(".")
+    for part in parts[:-1]:
+        nxt = cur.get(part)
+        if not isinstance(nxt, dict):
+            nxt = {}
+            cur[part] = nxt
+        cur = nxt
+    cur[parts[-1]] = value
+
+
+def _unset_path(doc: dict, dotted: str) -> None:
+    cur = doc
+    parts = dotted.split(".")
+    for part in parts[:-1]:
+        cur = cur.get(part)
+        if not isinstance(cur, dict):
+            return
+    cur.pop(parts[-1], None)
+
+
+def _apply_forward(doc: dict, ops: ContextUpdateOps, *, move_to_end: bool = False) -> None:
+    """Apply one journaled diff in place with the completion write's semantics."""
+    for path, value in ops.sets.items():
+        _set_path(doc, path, value)
+    for path in ops.unsets:
+        _unset_path(doc, path)
+    for path, items in ops.add_to_sets.items():
+        _set_path(doc, path, _union(_get_path(doc, path) or [], items, move_to_end=move_to_end, path=path))
+
+
+def commit_ordered(entries: list[dict], completed_steps: list[dict]) -> list[dict]:
+    """Entries in commit order: seed first, then the scheduler's record, then by position in
+    ``completed_steps``, which the CAS appends and so outranks ``seq``. A step with no
+    completion never landed; drop it."""
+    pos = {c.get("step_id"): i for i, c in enumerate(completed_steps)}
+    pos[_SCHEDULER_STEP_ID] = -1
+    pos[_SEED_STEP_ID] = -2
+    kept = [e for e in entries if e["step_id"] in pos]
+    return sorted(kept, key=lambda e: pos[e["step_id"]])
+
+
+def replay_context(entries: list[dict], rerecorded_steps: Iterable[str]) -> dict:
+    """The context ``entries`` produce over an empty one; order them with ``commit_ordered``.
+    ``rerecorded_steps`` is the instance's own list, the only record of which steps re-recorded:
+    a flag on the entry would be a pre-CAS read and can disagree with the write that landed."""
+    flagged = set(rerecorded_steps)
+    doc = {"context": dataclasses.asdict(TaskStepContext())}
+    for e in entries:
+        ops = ContextUpdateOps.from_journal_dict(e["ops"])
+        _apply_forward(doc, ops, move_to_end=e["step_id"] in flagged)
+    return doc["context"]
+
+
+def _item_key(path: str):
+    """How a context list's items are told apart: a deployed env by its instance, so two kernels' copies stay one entry."""
+    if path == "context.deployed_envs":
+        return lambda item: (item.get("instance_id") if isinstance(item, dict) else None) or item
+    return lambda item: item