136 lines
5.0 KiB
Python
136 lines
5.0 KiB
Python
"""Cooperative cancellation token for deep-research jobs.
|
|
|
|
The vendored GPT Researcher loops (recursive deep research, sub-topic fan-out,
|
|
multi-agent review cycles) can run for minutes. Every adapter and runner polls
|
|
a single :class:`CancellationToken` so a user cancel (or a lease takeover)
|
|
promptly unwinds the whole tree instead of burning model budget to completion.
|
|
|
|
Design:
|
|
|
|
- ``CancellationToken`` is an asyncio-friendly flag — ``cancel()`` flips it,
|
|
``raise_if_cancelled()`` is the checked-call sites sprinkle at loop boundaries.
|
|
- A token is created **per job execution** by the executor and threaded into the
|
|
``AdapterBundle``; it is never a process global.
|
|
- The executor's cancellation watcher calls ``cancel()`` when it observes
|
|
``cancel_requested`` in the DB; the running ``asyncio.Task`` sees
|
|
``CancelledError`` propagate up.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
|
|
class CancellationToken:
|
|
"""A lightweight cooperative cancel flag.
|
|
|
|
Thread-safe in the asyncio sense (single event loop): ``cancel()`` and
|
|
``is_cancelled`` only touch the flag / an asyncio event.
|
|
"""
|
|
|
|
__slots__ = ("_cancelled", "_event", "_reason")
|
|
|
|
def __init__(self) -> None:
|
|
self._cancelled: bool = False
|
|
self._reason: str = ""
|
|
# Used by ``wait()``; lazily created on first await so we don't need a
|
|
# running loop at construction time (tokens are often built in sync code).
|
|
self._event: asyncio.Event | None = None
|
|
|
|
# ── state queries ───────────────────────────────────────────────────────
|
|
|
|
@property
|
|
def is_cancelled(self) -> bool:
|
|
return self._cancelled
|
|
|
|
@property
|
|
def reason(self) -> str:
|
|
return self._reason
|
|
|
|
# ── transitions ─────────────────────────────────────────────────────────
|
|
|
|
def cancel(self, reason: str = "cancelled") -> None:
|
|
"""Flip the flag and wake any coroutine blocked on ``wait()``."""
|
|
self._cancelled = True
|
|
self._reason = reason
|
|
if self._event is not None:
|
|
self._event.set()
|
|
|
|
# ── checked-call helpers ────────────────────────────────────────────────
|
|
|
|
def raise_if_cancelled(self) -> None:
|
|
"""Raise :class:`ResearchCancelled` if the token has been cancelled.
|
|
|
|
Sprinkle at the top of every loop iteration in vendored / runner code:
|
|
search fan-out, context compression, report chunking, graph nodes.
|
|
"""
|
|
if self._cancelled:
|
|
raise ResearchCancelled(self._reason)
|
|
|
|
async def wait(self, timeout: float | None = None) -> None:
|
|
"""Block until cancelled (or ``timeout`` elapses).
|
|
|
|
Useful when a node needs to *park* until either cancellation or an
|
|
external resume signal — though for human-review interrupts the
|
|
LangGraph ``interrupt()`` mechanism is preferred.
|
|
"""
|
|
if self._cancelled:
|
|
return
|
|
if self._event is None:
|
|
self._event = asyncio.Event()
|
|
try:
|
|
await asyncio.wait_for(self._event.wait(), timeout=timeout)
|
|
except TimeoutError:
|
|
pass
|
|
|
|
# ── context-manager sugar (check on enter/exit of a scoped unit) ────────
|
|
|
|
def __enter__(self) -> CancellationToken:
|
|
self.raise_if_cancelled()
|
|
return self
|
|
|
|
def __exit__(self, *exc: object) -> None:
|
|
self.raise_if_cancelled()
|
|
|
|
|
|
class ResearchCancelled(Exception):
|
|
"""Raised when a :class:`CancellationToken` is flipped mid-run.
|
|
|
|
Deliberately **not** a subclass of ``asyncio.CancelledError``: the latter is
|
|
used by the executor to kill the task object; ``ResearchCancelled`` is the
|
|
*checked* variant that lets adapters/runners unwind cleanly and record a
|
|
``cancelled`` terminal state rather than an error.
|
|
"""
|
|
|
|
def __init__(self, reason: str = "cancelled") -> None:
|
|
super().__init__(reason)
|
|
self.reason = reason
|
|
|
|
|
|
class ResearchAwaitingInput(Exception):
|
|
"""Signal a durable human-review pause to the job executor.
|
|
|
|
This is intentionally separate from cancellation: the runner has completed
|
|
a safe checkpoint and the executor must preserve its plan, release the
|
|
worker lease, and let a later ``resume`` job continue from that checkpoint.
|
|
"""
|
|
|
|
def __init__(self, interrupt_id: str, plan: dict) -> None:
|
|
super().__init__(f"Research is awaiting input at {interrupt_id}")
|
|
self.interrupt_id = interrupt_id
|
|
self.plan = plan
|
|
|
|
|
|
def raise_if_cancelled_or_return(value: CancellationToken | None) -> None:
|
|
"""Convenience for adapters that receive an optional token."""
|
|
if value is not None:
|
|
value.raise_if_cancelled()
|
|
|
|
|
|
__all__ = [
|
|
"CancellationToken",
|
|
"ResearchAwaitingInput",
|
|
"ResearchCancelled",
|
|
"raise_if_cancelled_or_return",
|
|
]
|