shared/execution.py

"""Cooperative cancellation and deadlines for one process's active graph work."""

from __future__ import annotations

import threading
import time
from contextvars import ContextVar
from functools import wraps
from uuid import UUID


class ExecutionStopped(RuntimeError):
    def __init__(self, reason: str):
        self.reason = reason
        self.code = "research_timeout" if reason == "deadline" else "research_cancelled"
        super().__init__(
            "Research took too long. Try a more focused question."
            if reason == "deadline"
            else "Research was stopped."
        )


class Execution:
    def __init__(self, seconds: float):
        self.deadline = time.monotonic() + seconds
        self.reason = ""
        self.stopped = threading.Event()
        self.idle = threading.Event()
        self.idle.set()
        self._lock = threading.RLock()
        self._workers = 0

    def stop(self, reason: str = "cancelled") -> None:
        with self._lock:
            if not self.reason:
                self.reason = reason
            self.stopped.set()

    def check(self) -> None:
        if time.monotonic() >= self.deadline:
            self.stop("deadline")
        if self.stopped.is_set():
            raise ExecutionStopped(self.reason)

    def enter(self) -> None:
        with self._lock:
            self.check()
            self._workers += 1
            self.idle.clear()

    def leave(self) -> None:
        with self._lock:
            self._workers -= 1
            if not self._workers:
                self.idle.set()


execution: ContextVar[Execution | None] = ContextVar("precedent_execution", default=None)
_active: dict[tuple[str, UUID], Execution] = {}
_registry_lock = threading.Lock()


def register(owner: str, ident: UUID, control: Execution) -> bool:
    with _registry_lock:
        key = (owner, ident)
        if key in _active:
            return False
        _active[key] = control
        return True


def unregister(owner: str, ident: UUID, control: Execution) -> None:
    with _registry_lock:
        if _active.get((owner, ident)) is control:
            del _active[(owner, ident)]


def stop_request(owner: str, ident: UUID) -> bool:
    with _registry_lock:
        control = _active.get((owner, ident))
        if control is None:
            return False
        control.stop()
        return True


def check_execution() -> None:
    control = execution.get()
    if control is not None:
        control.check()


def remaining_seconds(maximum: float) -> float:
    check_execution()
    control = execution.get()
    return min(maximum, max(0.01, control.deadline - time.monotonic())) if control else maximum


def pause(seconds: float) -> None:
    control = execution.get()
    if control:
        control.stopped.wait(remaining_seconds(seconds))
        control.check()
    else:
        time.sleep(seconds)


def guarded_node(node):
    @wraps(node)
    def guarded(*args, **kwargs):
        control = execution.get()
        if control is None:
            return node(*args, **kwargs)
        control.enter()
        try:
            result = node(*args, **kwargs)
            control.check()
            return result
        finally:
            control.leave()

    return guarded