shikiru-cli 0.2.0__py3-none-any.whl

shikiru/engine/turn_loop.py

@@ -0,0 +1,230 @@
+"""Runs a step's turns and their recoveries until one completes or the task stops."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from shikiru.domain.recovery import Completed
+from shikiru.domain.steps import ACTING_ROLE, Transition
+from shikiru.engine.events import emit_nudge, emit_turn_started, emit_turn_stopped
+from shikiru.engine.turn_context import ORDINARY_TURN, NextTurn
+from shikiru.engine.turn_execution import TurnExecution
+from shikiru.engine.turns import TurnHeld
+from shikiru.errors import InvariantError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from shikiru.domain.feedback import FeedbackInput
+    from shikiru.domain.recovery import TurnCounters
+    from shikiru.domain.reports import Report, ReportKind
+    from shikiru.domain.tasks import Role, Step, Task
+    from shikiru.engine.escalation_filing import EscalationFiler
+    from shikiru.engine.recovery import TurnRecovery
+    from shikiru.engine.sessions import SessionManager
+    from shikiru.engine.turn_context import TurnContextPreparer
+    from shikiru.engine.turns import TurnRunner
+    from shikiru.prompts.renderer import PromptName
+    from shikiru.store.events import EventLog
+    from shikiru.store.runs import RunStore
+    from shikiru.store.tasks import TaskStore
+
+
+@dataclass(frozen=True)
+class TurnOverrides:
+    """What an integration turn changes about a step's ordinary turn."""
+
+    role: Role | None = None
+    template: PromptName | None = None
+    kind: ReportKind | None = None
+    conflict_context: str | None = None
+    feedback: FeedbackInput | None = None
+
+
+NO_OVERRIDES = TurnOverrides()
+
+
+@dataclass(frozen=True)
+class TurnsCompleted:
+    """The loop ended with a turn that produced `report`."""
+
+    task: Task
+    counters: TurnCounters
+    execution: TurnExecution
+    report: Report
+
+
+@dataclass(frozen=True)
+class TurnsStopped:
+    """The loop ended because the task was held or escalated."""
+
+    task: Task
+    counters: TurnCounters
+
+
+TurnLoopResult = TurnsCompleted | TurnsStopped
+
+
+def _acting_role(step: Step, role_override: Role | None) -> Role:
+    return role_override if role_override is not None else ACTING_ROLE[step]
+
+
+class TurnLoop:
+    """Runs a step's turn, recovering failed turns until one completes or the task stops."""
+
+    def __init__(
+        self,
+        *,
+        tasks: TaskStore,
+        events: EventLog,
+        sessions: SessionManager,
+        turns: TurnRunner,
+        contexts: TurnContextPreparer,
+        recovery: TurnRecovery,
+        escalations: EscalationFiler,
+        runs: RunStore,
+    ) -> None:
+        """Run turns through `turns`, recover through `recovery`, escalate through `escalations`."""
+        self._tasks = tasks
+        self._events = events
+        self._sessions = sessions
+        self._turns = turns
+        self._contexts = contexts
+        self._recovery = recovery
+        self._escalations = escalations
+        self._runs = runs
+        self._resumed: set[str] = set()
+
+    def expect_resume(self, task_id: str) -> None:
+        """Let `task_id`'s next turn consume a report its interrupted turn already submitted."""
+        self._resumed.add(task_id)
+
+    def forget_resume(self, task_id: str) -> None:
+        """Drop the pending resume of `task_id`."""
+        self._resumed.discard(task_id)
+
+    async def run(
+        self,
+        task: Task,
+        step: Step,
+        counters: TurnCounters,
+        overrides: TurnOverrides = NO_OVERRIDES,
+    ) -> TurnLoopResult:
+        """Run `step`'s turn and its recoveries until a report, a hold or an escalation."""
+        current = self._tasks.require(task.id)
+        if current.held:
+            return TurnsStopped(current, counters)
+        current_step = step
+        execution = await self._execute(current, current_step, overrides, ORDINARY_TURN)
+        while execution is not None:
+            result = execution.result
+            if isinstance(result, Completed):
+                return TurnsCompleted(current, counters, execution, result.report)
+            role = _acting_role(current_step, overrides.role)
+            recovery = await self._recovery.recover(
+                current, current_step, role, execution, counters
+            )
+            current, counters = recovery.task, recovery.counters
+            if recovery.escalation is not None:
+                transition = Transition.escalate(current, recovery.escalation, step=current_step)
+                escalated = await self._escalations.escalate(
+                    current, transition, None, result=result, role=role
+                )
+                return TurnsStopped(escalated, counters)
+            if recovery.execution is not None:
+                execution = recovery.execution
+                continue
+            if current.step is None:
+                message = f"{current.id} recovery left the task with no step"
+                raise InvariantError(message)
+            current_step = current.step
+            current = self._tasks.require(current.id)
+            if current.held:
+                return TurnsStopped(current, counters)
+            execution = await self._execute(current, current_step, overrides, recovery.next_turn)
+        return TurnsStopped(self._tasks.require(current.id), counters)
+
+    async def _execute(
+        self, task: Task, step: Step, overrides: TurnOverrides, next_turn: NextTurn
+    ) -> TurnExecution | None:
+        role = _acting_role(step, overrides.role)
+        selection = self._sessions.session_for(task, role)
+        turn_request = await self._contexts.prepare(
+            task,
+            step,
+            session_handoff=selection.handoff,
+            next_turn=next_turn,
+            template_override=overrides.template,
+            kind_override=overrides.kind,
+            conflict_context=overrides.conflict_context,
+            feedback=overrides.feedback,
+        )
+        usage_before = selection.run.usage
+        resumed = task.id in self._resumed
+        self._resumed.discard(task.id)
+        completed = self._turns.completed_turn(task.id, role, turn_request) if resumed else None
+        if completed is not None:
+            result, turn_dir = completed
+            return TurnExecution(
+                result=result,
+                task_id=task.id,
+                role=role,
+                turn_dir=turn_dir,
+                rung=selection.rung,
+                run_id=selection.run.id,
+                usage_before=usage_before,
+                usage_after=usage_before,
+                kind=turn_request.kind,
+            )
+        turn_name: str | None = None
+
+        def on_start(turn_dir: Path) -> None:
+            nonlocal turn_name
+            turn_name = turn_dir.name
+            emit_turn_started(
+                self._events,
+                task,
+                step,
+                role,
+                selection.rung,
+                kind=turn_request.kind,
+                run_id=selection.run.id,
+                turn=turn_name,
+            )
+
+        def on_nudge() -> None:
+            emit_nudge(self._events, task, step, role, turn=turn_name)
+
+        try:
+            result, turn_dir, selection = await self._turns.run(
+                task, role, selection, turn_request, on_start=on_start, on_nudge=on_nudge
+            )
+        except BaseException as stop:
+            emit_turn_stopped(
+                self._events,
+                task,
+                step,
+                role,
+                rung=selection.rung,
+                run_id=selection.run.id,
+                turn=turn_name,
+            )
+            if isinstance(stop, TurnHeld):
+                return None
+            raise
+        live_run = self._runs.get_live(task.id, role)
+        if live_run is None:
+            message = f"{task.id} {role.value} run disappeared after its turn"
+            raise InvariantError(message)
+        return TurnExecution(
+            result=result,
+            task_id=task.id,
+            role=role,
+            turn_dir=turn_dir,
+            rung=selection.rung,
+            run_id=selection.run.id,
+            usage_before=usage_before,
+            usage_after=live_run.usage,
+            kind=turn_request.kind,
+        )

shikiru/engine/turn_processes.py

@@ -0,0 +1,79 @@
+"""The `turn.pid` record of each running turn, and ending the ones a dead engine left."""
+
+from __future__ import annotations
+
+import json
+import signal
+from dataclasses import asdict
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from shikiru.agents.contracts import role_dir
+from shikiru.paths import RUNS_DIR_NAME, project_state_dir
+from shikiru.private import make_private_dir, write_private_text
+from shikiru.process import TURN_STOP_GRACE, ProcessIdentity
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from shikiru.clock import Clock
+    from shikiru.domain.tasks import Role
+    from shikiru.process import SubprocessRunner
+
+TURN_PID_NAME = "turn.pid"
+EXIT_POLL_SECONDS = 0.1
+_STAGED_SUFFIX = ".staged"
+
+
+def turn_pid_path(run_dir: Path, role: Role) -> Path:
+    """Return the file recording the process that runs `role`'s current turn."""
+    return role_dir(run_dir, role.value) / TURN_PID_NAME
+
+
+def read_turn_pid(path: Path) -> ProcessIdentity:
+    """Read the process identity recorded in a `turn.pid` file."""
+    return ProcessIdentity(**json.loads(path.read_text()))
+
+
+class TurnProcesses:
+    """Record the process running each role's turn and stop orphans after an engine restart."""
+
+    def __init__(self, *, runner: SubprocessRunner, clock: Clock, project_root: Path) -> None:
+        """Identify and signal processes through `runner`, waiting on `clock`."""
+        self._runner = runner
+        self._clock = clock
+        self._project_root = project_root
+
+    def record(self, run_dir: Path, role: Role, pid: int) -> None:
+        """Write `role`'s `turn.pid` with the identity of `pid`, or of its group once it left."""
+        path = turn_pid_path(run_dir, role)
+        identity = self._runner.identify(pid)
+        if identity is None and self._runner.group_is_running(pid):
+            identity = ProcessIdentity(pid, None)
+        if identity is None:
+            path.unlink(missing_ok=True)
+            return
+        make_private_dir(path.parent)
+        staged = path.with_suffix(_STAGED_SUFFIX)
+        write_private_text(staged, json.dumps(asdict(identity)))
+        staged.replace(path)
+
+    @staticmethod
+    def clear(run_dir: Path, role: Role) -> None:
+        """Remove `role`'s `turn.pid` once its turn process has ended."""
+        turn_pid_path(run_dir, role).unlink(missing_ok=True)
+
+    async def end_orphans(self) -> None:
+        """Stop every recorded turn process that still runs, then drop its `turn.pid`."""
+        runs_root = project_state_dir(self._project_root) / RUNS_DIR_NAME
+        for path in sorted(runs_root.glob(f"*/*/{TURN_PID_NAME}")):
+            await self._stop(read_turn_pid(path))
+            path.unlink()
+
+    async def _stop(self, identity: ProcessIdentity) -> None:
+        if not self._runner.signal_group(identity, signal.SIGTERM):
+            return
+        deadline = self._clock.now() + timedelta(seconds=TURN_STOP_GRACE)
+        while self._runner.group_is_running(identity.pid) and self._clock.now() < deadline:
+            await self._clock.sleep(EXIT_POLL_SECONDS)
+        self._runner.signal_group(identity, signal.SIGKILL)