shikiru-cli 0.2.0__py3-none-any.whl

shikiru/engine/quota_waits.py

@@ -0,0 +1,102 @@
+"""Which windows and role panes of a task are in a quota wait, from authoritative current facts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from shikiru.domain.ladder import pick_rung
+from shikiru.domain.quota_waits import current_quota_wait_events
+from shikiru.domain.scheduling import WaitReason
+from shikiru.domain.steps import task_acting_role
+from shikiru.domain.tasks import Role, Status
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from datetime import datetime
+
+    from shikiru.clock import Clock
+    from shikiru.config.settings import Settings
+    from shikiru.domain.quota_waits import QuotaAction
+    from shikiru.domain.tasks import Task
+    from shikiru.engine.quota_schedule import QuotaSchedule
+    from shikiru.store.events import EventLog
+    from shikiru.store.quotas import AgentBlock, QuotaStore
+
+_QUOTA_ROLES = (Role.LEADER, Role.IMPLEMENTER)
+_TURN_STATUSES = (Status.RUNNING, Status.INTEGRATING)
+
+
+@dataclass(frozen=True)
+class RoleQuotaWait:
+    """A role whose agent is quota blocked, with the engine's next action on it when known."""
+
+    action: QuotaAction | None = None
+    at: datetime | None = None
+
+
+@dataclass(frozen=True)
+class TaskQuotaWait:
+    """A task in a quota wait and the roles whose agents the wait is about.
+
+    A queued task waits when any required role is blocked, whichever role acts next; a running
+    task waits in the turn of its acting role only. `roles` can be empty while the persisted
+    wait reason still names quota, as the task list also shows.
+    """
+
+    roles: Mapping[Role, RoleQuotaWait]
+
+
+class QuotaWaits:
+    """Resolves a task's quota wait with the same rules as the task list's quota marker."""
+
+    def __init__(
+        self,
+        *,
+        quotas: QuotaStore,
+        events: EventLog,
+        schedule: QuotaSchedule,
+        clock: Clock,
+        settings: Settings,
+    ) -> None:
+        """Read blocks from `quotas`, turn facts from `events`, next actions from `schedule`."""
+        self._quotas = quotas
+        self._events = events
+        self._schedule = schedule
+        self._clock = clock
+        self._settings = settings
+
+    def resolve(self, task: Task) -> TaskQuotaWait | None:
+        """Return `task`'s quota wait, or none when it is not waiting on quota."""
+        if task.status is Status.QUEUED:
+            return self._queued(task)
+        if task.status in _TURN_STATUSES and task.step is not None:
+            return self._running(task)
+        return None
+
+    def _queued(self, task: Task) -> TaskQuotaWait | None:
+        if task.wait_reason != WaitReason.QUOTA.value:
+            return None
+        blocks = {block.agent: block for block in self._quotas.active_blocks(self._clock.now())}
+        roles = {
+            role: self._role_wait(blocks[agent])
+            for role in _QUOTA_ROLES
+            if (agent := pick_rung(task, role, self._settings).agent) in blocks
+        }
+        return TaskQuotaWait(roles)
+
+    def _running(self, task: Task) -> TaskQuotaWait | None:
+        role = task_acting_role(task)
+        agent = pick_rung(task, role, self._settings).agent
+        if current_quota_wait_events(task, role, agent, self._events.latest) is None:
+            return None
+        block = self._quotas.block_for(agent)
+        if block is None or block.until <= self._clock.now():
+            return None
+        return TaskQuotaWait({role: self._role_wait(block)})
+
+    def _role_wait(self, block: AgentBlock) -> RoleQuotaWait:
+        scheduled = self._schedule.get(block.agent)
+        if scheduled is None or scheduled.blocked_at != block.blocked_at:
+            return RoleQuotaWait()
+        return RoleQuotaWait(scheduled.action, scheduled.at)

shikiru/engine/reconcile.py

@@ -0,0 +1,91 @@
+"""Engine restart recovery."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from shikiru.domain.tasks import Status
+from shikiru.paths import task_root
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+    from pathlib import Path
+
+    from shikiru.domain.tasks import Task
+    from shikiru.engine.fault_reporter import FaultReporter
+    from shikiru.engine.layout_guard import LayoutGuard
+    from shikiru.engine.turn_processes import TurnProcesses
+    from shikiru.store.tasks import TaskStore
+
+_IN_FLIGHT = (Status.RUNNING, Status.INTEGRATING)
+_TERMINAL = (Status.DONE, Status.CANCELLED)
+
+
+class Reconciler:
+    """Restore windows, end orphaned turn processes and re-run in-flight steps at engine start."""
+
+    def __init__(
+        self,
+        *,
+        tasks: TaskStore,
+        layout: LayoutGuard,
+        processes: TurnProcesses,
+        resume_task: Callable[[str], Awaitable[None]],
+        finish_task: Callable[[Task], Awaitable[Path | None]],
+        restore_escalation: Callable[[Task], Awaitable[None]],
+        restore_pull_request_closes: Callable[[], None],
+        faults: FaultReporter,
+        project_root: Path,
+    ) -> None:
+        """Resume in-flight tasks, finish terminal ones, restore escalations of escalated ones.
+
+        A fault in one step or task is recorded through `faults`; the rest goes on.
+        """
+        self._tasks = tasks
+        self._layout = layout
+        self._processes = processes
+        self._resume_task = resume_task
+        self._finish_task = finish_task
+        self._restore_escalation = restore_escalation
+        self._restore_pull_request_closes = restore_pull_request_closes
+        self._faults = faults
+        self._project_root = project_root
+        self._orphans_ended = False
+
+    async def reconcile(self) -> None:
+        """Bring a restarted engine back in line with the persisted state.
+
+        Pending commands need nothing here: the engine loop processes them once this has
+        completed. A fault in one step or one task is recorded and never stops the others.
+        Ending the orphaned turn processes is the exception: a fault there aborts the reconcile,
+        which is retried before anything starts, so no flow runs beside a live orphan and none
+        ends a turn the engine started itself.
+        """
+        active = self._tasks.list_active()
+        terminal = [
+            task
+            for task in self._tasks.list_all()
+            if task.status in _TERMINAL and task_root(self._project_root, task.id).exists()
+        ]
+        async with self._faults.isolate("reconcile layout"):
+            await self._layout.restore(active)
+        if not self._orphans_ended:
+            await self._processes.end_orphans()
+            self._orphans_ended = True
+        for task in terminal:
+            async with self._faults.isolate("reconcile cleanup", task_id=task.id):
+                await self._finish_task(task)
+        async with self._faults.isolate("reconcile pull request closes"):
+            self._restore_pull_request_closes()
+        await self.restore_escalations()
+        for task in active:
+            if task.status in _IN_FLIGHT and not task.held:
+                async with self._faults.isolate("reconcile resume", task_id=task.id):
+                    await self._resume_task(task.id)
+
+    async def restore_escalations(self) -> None:
+        """Give every escalated task an open escalation with its packet on disk."""
+        for task in self._tasks.list_active():
+            if task.status is Status.ESCALATED:
+                async with self._faults.isolate("escalation restore", task_id=task.id):
+                    await self._restore_escalation(task)

shikiru/engine/recovery.py

@@ -0,0 +1,284 @@
+"""Recovery of a turn that ended without a report: retries, quota waits and rung moves."""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from shikiru.agents.contracts import FailureKind
+from shikiru.domain.events import EventKind
+from shikiru.domain.ladder import move_up, pick_rung
+from shikiru.domain.recovery import (
+    QUOTA_MARGIN,
+    Backoff,
+    Escalate,
+    Failed,
+    NextRung,
+    RecoveryPolicy,
+    ResumeTurn,
+    RetryTurn,
+    TurnCounters,
+    WaitForQuota,
+    decide,
+    quota_block_until,
+    with_known_quota_reset,
+)
+from shikiru.domain.steps import Transition
+from shikiru.domain.tasks import EscalationReason
+from shikiru.engine.events import emit_turn_failed
+from shikiru.engine.turn_context import ORDINARY_TURN, NextTurn
+
+if TYPE_CHECKING:
+    from datetime import datetime
+
+    from shikiru.clock import Clock
+    from shikiru.config.settings import Settings
+    from shikiru.domain.recovery import Action, TurnResult
+    from shikiru.domain.tasks import Role, Step, Task
+    from shikiru.engine.quota_reconciler import QuotaReconciler
+    from shikiru.engine.transitions import TransitionWriter
+    from shikiru.engine.turn_execution import TurnExecution
+    from shikiru.store.events import EventLog
+    from shikiru.store.quotas import QuotaStore
+    from shikiru.store.tasks import TaskStore
+
+QUOTA_BLOCK_REASON = "quota"
+_QUOTA_WAIT_EXHAUSTED_MESSAGE = "quota wait exceeded"
+_CONTINUE = NextTurn(continue_turn=True)
+
+
+@dataclass(frozen=True)
+class RecoveryStep:
+    """What the flow does after a failed turn: run again, re-decide on `execution` or escalate."""
+
+    task: Task
+    counters: TurnCounters
+    execution: TurnExecution | None = None
+    escalation: EscalationReason | None = None
+    next_turn: NextTurn = ORDINARY_TURN
+
+
+class TurnRecovery:
+    """Decides and carries out the recovery of a failed turn; escalation stays with the flow."""
+
+    def __init__(
+        self,
+        *,
+        tasks: TaskStore,
+        events: EventLog,
+        quota: QuotaStore,
+        reconciler: QuotaReconciler,
+        clock: Clock,
+        writer: TransitionWriter,
+        settings: Settings,
+    ) -> None:
+        """Recover through the ladder, the quota blocks and the clock."""
+        self._tasks = tasks
+        self._events = events
+        self._quota = quota
+        self._reconciler = reconciler
+        self._clock = clock
+        self._writer = writer
+        self._settings = settings
+
+    async def recover(
+        self,
+        task: Task,
+        step: Step,
+        role: Role,
+        execution: TurnExecution,
+        counters: TurnCounters,
+    ) -> RecoveryStep:
+        """Apply the recovery `domain.recovery.decide` chooses for `execution`'s failed result."""
+        execution = self._with_known_reset(execution)
+        result = execution.result
+        emit_turn_failed(
+            self._events,
+            task,
+            step,
+            result,
+            rung=execution.rung,
+            run_id=execution.run_id,
+            turn=execution.turn_dir.name,
+            role=execution.role,
+        )
+        action = decide(
+            result,
+            counters,
+            self._policy(task, role, quota_wait_exhausted=execution.quota_wait_exhausted),
+        )
+        if isinstance(result, Failed) and result.kind is FailureKind.QUOTA:
+            self._block_quota(task, role, result.resets_at)
+        return await self._apply(
+            task, step, role, execution=execution, counters=counters, action=action
+        )
+
+    def _with_known_reset(self, execution: TurnExecution) -> TurnExecution:
+        windows = self._quota.windows_for(execution.rung.agent)
+        result = with_known_quota_reset(execution.result, windows, self._clock.now())
+        return dataclasses.replace(execution, result=result)
+
+    def _policy(self, task: Task, role: Role, *, quota_wait_exhausted: bool) -> RecoveryPolicy:
+        current = move_up(task, role, self._settings)
+        different = move_up(
+            task, role, self._settings, different_agent=True, blocked_agents=self._blocked_agents()
+        )
+        retries = self._settings.orchestration.retries
+        return RecoveryPolicy(
+            now=self._clock.now(),
+            max_quota_wait=timedelta(seconds=self._settings.orchestration.timeouts.max_quota_wait),
+            has_next_rung=current is not None,
+            has_next_rung_different_agent=different is not None,
+            unavailable_initial_seconds=retries.unavailable_initial_seconds,
+            unavailable_multiplier=retries.unavailable_multiplier,
+            unavailable_max_seconds=retries.unavailable_max_seconds,
+            unavailable_max_retries=retries.unavailable_max_retries,
+            quota_wait_exhausted=quota_wait_exhausted,
+        )
+
+    async def _apply(
+        self,
+        task: Task,
+        step: Step,
+        role: Role,
+        *,
+        execution: TurnExecution,
+        counters: TurnCounters,
+        action: Action,
+    ) -> RecoveryStep:
+        if isinstance(action, NextRung):
+            return await self._next_rung(task, step, role, execution.result, action)
+        if isinstance(action, Escalate):
+            return RecoveryStep(task, TurnCounters(), escalation=action.reason)
+        return await self._retry(task, role, execution, counters, action)
+
+    async def _retry(
+        self,
+        task: Task,
+        role: Role,
+        execution: TurnExecution,
+        counters: TurnCounters,
+        action: WaitForQuota | Backoff | ResumeTurn | RetryTurn,
+    ) -> RecoveryStep:
+        if isinstance(action, RetryTurn):
+            return RecoveryStep(
+                task, dataclasses.replace(counters, other_retries=counters.other_retries + 1)
+            )
+        if isinstance(action, Backoff):
+            await self._clock.sleep(action.delay.total_seconds())
+            retried = dataclasses.replace(
+                counters, unavailable_retries=counters.unavailable_retries + 1
+            )
+            return RecoveryStep(self._tasks.require(task.id), retried)
+        if isinstance(action, ResumeTurn):
+            resumed = dataclasses.replace(counters, stall_recoveries=counters.stall_recoveries + 1)
+            return RecoveryStep(self._tasks.require(task.id), resumed, next_turn=_CONTINUE)
+        terminal = await self._wait_for_quota(task, role)
+        if terminal is None:
+            return RecoveryStep(self._tasks.require(task.id), counters, next_turn=_CONTINUE)
+        exhausted = dataclasses.replace(execution, result=terminal, quota_wait_exhausted=True)
+        return RecoveryStep(self._tasks.require(task.id), counters, execution=exhausted)
+
+    async def _next_rung(
+        self, task: Task, step: Step, role: Role, result: TurnResult, action: NextRung
+    ) -> RecoveryStep:
+        updated = await self._move_rung(task, step, role, action.different_agent, result)
+        if updated is not None:
+            return RecoveryStep(
+                updated, TurnCounters(), next_turn=NextTurn(handoff=action.different_agent)
+            )
+        reason = (
+            EscalationReason.QUOTA_EXHAUSTED
+            if action.different_agent
+            else EscalationReason.LADDER_EXHAUSTED
+        )
+        return RecoveryStep(task, TurnCounters(), escalation=reason)
+
+    async def _move_rung(
+        self,
+        task: Task,
+        step: Step,
+        role: Role,
+        different_agent: bool,
+        result: TurnResult,
+    ) -> Task | None:
+        move = move_up(
+            task,
+            role,
+            self._settings,
+            different_agent=different_agent,
+            blocked_agents=self._blocked_agents() if different_agent else frozenset(),
+        )
+        if move is None:
+            return None
+        old_rung = pick_rung(task, role, self._settings)
+        transition = dataclasses.replace(
+            Transition.stay(task, step=step, status=task.status),
+            counters=move.counters,
+            rung_index=move.rung_index,
+            cleared_overrides=move.cleared_overrides,
+        )
+        updated = await self._writer.apply(task, transition, None)
+        new_rung = pick_rung(updated, role, self._settings)
+        self._events.append(
+            EventKind.RUNG_CHANGED,
+            task_id=task.id,
+            message=f"{task.id} {role.value} moved from {old_rung.model} to {new_rung.model}",
+            data={
+                "role": role.value,
+                "from": old_rung.model,
+                "to": new_rung.model,
+                "different_agent": different_agent,
+                "result": type(result).__name__,
+            },
+        )
+        return updated
+
+    def _blocked_agents(self) -> frozenset[str]:
+        return frozenset(block.agent for block in self._quota.active_blocks(self._clock.now()))
+
+    def _block_quota(self, task: Task, role: Role, resets_at: datetime | None) -> None:
+        agent = pick_rung(task, role, self._settings).agent
+        now = self._clock.now()
+        until = quota_block_until(resets_at, now)
+        self._quota.block(agent, until, QUOTA_BLOCK_REASON, now)
+        self._events.append(
+            EventKind.QUOTA_BLOCKED,
+            task_id=task.id,
+            message=f"{agent} quota blocked until {until.isoformat()}",
+            data={"agent": agent, "until": until.isoformat(), "reason": QUOTA_BLOCK_REASON},
+        )
+
+    async def _wait_for_quota(self, task: Task, role: Role) -> Failed | None:
+        agent = pick_rung(task, role, self._settings).agent
+        deadline = self._clock.now() + timedelta(
+            seconds=self._settings.orchestration.timeouts.max_quota_wait
+        )
+        while True:
+            current = self._quota.block_for(agent)
+            now = self._clock.now()
+            if current is None:
+                return None
+            if current.until > now:
+                if now >= deadline:
+                    return Failed(
+                        kind=FailureKind.QUOTA,
+                        resets_at=current.until - QUOTA_MARGIN,
+                        message=_QUOTA_WAIT_EXHAUSTED_MESSAGE,
+                    )
+                remaining = min(current.until, deadline) - now
+                delay = min(remaining.total_seconds(), self._reconciler.interval_seconds)
+                self._reconciler.expect_check(agent, now + timedelta(seconds=delay))
+                await self._clock.sleep(delay)
+                await self._reconciler.refresh(agent)
+                continue
+            self._quota.clear_block(agent)
+            self._events.append(
+                EventKind.QUOTA_CLEARED,
+                task_id=task.id,
+                message=f"{agent} quota block cleared",
+                data={"agent": agent},
+            )
+            return None

shikiru/engine/retention.py

@@ -0,0 +1,102 @@
+"""Prunes finished history so a long-running engine stays bounded."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import shutil
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from shikiru.domain.tasks import Status
+from shikiru.paths import run_dir, task_root
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from datetime import datetime
+    from pathlib import Path
+
+    from shikiru.clock import Clock
+    from shikiru.config.settings import RetentionConfig
+    from shikiru.store.notification_cursor import NotificationCursor
+    from shikiru.store.retention import StateRetention
+    from shikiru.store.tasks import TaskStore
+
+
+class Retention:
+    """Removes long-finished tasks' run directories and events and old commands, periodically.
+
+    One pass runs every `interval_seconds`, also removing old engine events. A task's files and
+    events go only when the task is `done` or `cancelled` past the horizon and nothing unresolved
+    still needs them: no open escalation, no worktree kept for its unintegrated work, no hold, no
+    pull request of a cancelled task left open. Run directories go before their events, so an
+    interrupted pass is simply repeated.
+    """
+
+    def __init__(
+        self,
+        *,
+        store: StateRetention,
+        tasks: TaskStore,
+        cursor: NotificationCursor,
+        project_root: Path,
+        has_open_pull_request: Callable[[str], bool],
+        clock: Clock,
+        settings: RetentionConfig,
+    ) -> None:
+        """Prune `store` and the run directories under `project_root` on `clock`'s schedule."""
+        self._store = store
+        self._tasks = tasks
+        self._cursor = cursor
+        self._project_root = project_root
+        self._has_open_pull_request = has_open_pull_request
+        self._clock = clock
+        self._settings = settings
+        self._next_run: datetime | None = None
+        self._logger = logging.getLogger(__name__)
+
+    async def run_if_due(self) -> None:
+        """Run one pruning pass when `interval_seconds` have passed since the last one."""
+        now = self._clock.now()
+        if self._next_run is not None and now < self._next_run:
+            return
+        await self._prune(now)
+        self._next_run = now + timedelta(seconds=self._settings.interval_seconds)
+
+    async def _prune(self, now: datetime) -> None:
+        delivered_up_to = self._cursor.read()
+        if delivered_up_to is None:
+            return
+        horizon = now - timedelta(days=self._settings.finished_task_days)
+        settled = [
+            task_id
+            for task_id in self._store.settled_task_ids(horizon)
+            if self._releasable(task_id)
+        ]
+        await asyncio.to_thread(_remove_run_dirs, [run_dir(self._project_root, t) for t in settled])
+        events = self._store.delete_task_events(settled, delivered_up_to=delivered_up_to)
+        events += self._store.delete_engine_events(
+            now - timedelta(days=self._settings.event_days), delivered_up_to=delivered_up_to
+        )
+        commands = self._store.delete_finished_commands(
+            now - timedelta(days=self._settings.command_days)
+        )
+        if settled or events or commands:
+            self._logger.info(
+                "retention pruned %d tasks, %d events, %d commands", len(settled), events, commands
+            )
+
+    def _releasable(self, task_id: str) -> bool:
+        task = self._tasks.get(task_id)
+        return (
+            task is not None
+            and not task.held
+            and not task_root(self._project_root, task_id).exists()
+            and not (task.status is Status.CANCELLED and self._has_open_pull_request(task_id))
+        )
+
+
+def _remove_run_dirs(directories: list[Path]) -> None:
+    for directory in directories:
+        if directory.exists():
+            shutil.rmtree(directory)

shikiru/engine/root_preparation.py

@@ -0,0 +1,86 @@
+"""Off-tick task root preparation under a per-tick budget."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import suppress
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Collection
+    from datetime import datetime
+
+    from shikiru.clock import Clock
+    from shikiru.domain.tasks import RepoState
+    from shikiru.workspace.roots import TaskRootBuilder
+
+
+class RootPreparation:
+    """Builds task roots on their own tasks so a slow fetch or cache copy never holds the tick.
+
+    A tick may wait `budget_seconds` in total for the roots it starts; a root that is not ready
+    by then keeps building and is collected by a later tick, which never waits for it again. At most
+    one build runs at a time and none starts once a tick's budget is spent, so a hung remote costs
+    one budget, not one per ready task and tick, and builds in one repository never overlap.
+    """
+
+    def __init__(self, roots: TaskRootBuilder, clock: Clock, *, budget_seconds: float) -> None:
+        """Build through `roots`, measuring each tick's budget on `clock`."""
+        self._roots = roots
+        self._clock = clock
+        self._budget = timedelta(seconds=budget_seconds)
+        self._expires: datetime = clock.now()
+        self._builds: dict[str, asyncio.Task[RepoState]] = {}
+
+    def begin_tick(self, queued: Collection[str]) -> None:
+        """Start a tick's budget and let go of the builds of tasks that are no longer `queued`.
+
+        A build in flight for such a task (cancelled meanwhile) is cancelled, and counts as
+        running until it has ended; a finished one is dropped.
+        """
+        self._expires = self._clock.now() + self._budget
+        for task_id, build in tuple(self._builds.items()):
+            if task_id in queued:
+                continue
+            if not build.done():
+                build.cancel()
+                continue
+            del self._builds[task_id]
+            if not build.cancelled():
+                build.exception()
+
+    async def prepare(self, task_id: str, title: str, seed_branch: str | None) -> RepoState | None:
+        """Return `task_id`'s built root, or `None` while it is not ready within this tick.
+
+        Raises what the build raised.
+        """
+        build = self._builds.get(task_id)
+        if build is None:
+            remaining = (self._expires - self._clock.now()).total_seconds()
+            if remaining <= 0 or any(not other.done() for other in self._builds.values()):
+                return None
+            build = asyncio.create_task(self._roots.build(task_id, title, seed_branch=seed_branch))
+            self._builds[task_id] = build
+            await self._wait(build, remaining)
+        if not build.done():
+            return None
+        del self._builds[task_id]
+        return build.result()
+
+    async def shutdown(self) -> None:
+        """Cancel every build in flight and wait until its processes are gone."""
+        builds = tuple(self._builds.values())
+        self._builds.clear()
+        for build in builds:
+            build.cancel()
+        await asyncio.gather(*builds, return_exceptions=True)
+
+    async def _wait(self, build: asyncio.Task[RepoState], seconds: float) -> None:
+        timer = asyncio.create_task(self._clock.sleep(seconds))
+        try:
+            await asyncio.wait({build, timer}, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            timer.cancel()
+            with suppress(asyncio.CancelledError):
+                await timer