From 1eea672cbe38e5b59177eba01e95562f17ae537f Mon Sep 17 00:00:00 2001 From: Viswa Marepalli Date: Sat, 31 Oct 2026 01:39:10 +0000 Subject: [PATCH] feat(agents): add native Pydantic AI harness support --- agents/src/agents/__init__.py | 28 ++ agents/src/agents/harness/__init__.py | 58 ++++ agents/src/agents/harness/errors.py | 32 ++ agents/src/agents/harness/events.py | 37 +++ agents/src/agents/harness/harness.py | 438 ++++++++++++++++++++++++++ agents/src/agents/harness/session.py | 154 +++++++++ agents/src/agents/harness/store.py | 348 ++++++++++++++++++++ agents/src/agents/harness/tools.py | 98 ++++++ agents/src/agents/harness/types.py | 224 +++++++++++++ agents/tests/harness/test_harness.py | 416 ++++++++++++++++++++++++ agents/tests/test_public_api.py | 2 + 11 files changed, 1835 insertions(+) create mode 100644 agents/src/agents/harness/__init__.py create mode 100644 agents/src/agents/harness/errors.py create mode 100644 agents/src/agents/harness/events.py create mode 100644 agents/src/agents/harness/harness.py create mode 100644 agents/src/agents/harness/session.py create mode 100644 agents/src/agents/harness/store.py create mode 100644 agents/src/agents/harness/tools.py create mode 100644 agents/src/agents/harness/types.py create mode 100644 agents/tests/harness/test_harness.py diff --git a/agents/src/agents/__init__.py b/agents/src/agents/__init__.py index fa41bff..cf1495d 100644 --- a/agents/src/agents/__init__.py +++ b/agents/src/agents/__init__.py @@ -5,6 +5,7 @@ from .core import * from .dynamic_agents import * from .fibers import * +from .harness import * from .lifecycle import * from .observability import * from .queue import * @@ -34,6 +35,7 @@ "AnyToolPart", "AppendEvent", "AppendResult", + "Approval", "CallableMetadata", "CancelledRun", "ChatMessageOptions", @@ -68,6 +70,11 @@ "FiberStatus", "Fibers", "FilePart", + "Harness", + "HarnessEvent", + "HarnessInput", + "HarnessSession", + "HarnessSessions", "InvalidCronExpressionError", "JSONValue", "JobContext", @@ -90,12 +97,19 @@ "NonRetryableError", "Observability", "ObservabilityEvent", + "OperationEvent", + "OperationNotFoundError", + "OperationResult", + "OperationStatus", + "PendingOperation", "PendingRun", "Queue", "QueueItem", "ReadonlyConnectionError", "ReasoningPart", + "Receipt", "RecentHistoryResult", + "Replay", "Reschedule", "RetryOptions", "RouteAddress", @@ -108,10 +122,13 @@ "ScheduleType", "Scheduler", "Session", + "SessionBusyError", "SessionChangeEvent", "SessionChangeListener", "SessionMessage", + "SessionNotFoundError", "SessionRowStat", + "SessionSnapshot", "Sessions", "Source", "SourceDocumentPart", @@ -147,17 +164,26 @@ "TaskStep", "TaskStepAttempt", "Tasks", + "TextEvent", "TextPart", + "ToolAnswer", "ToolApproval", "ToolApprovalRequestedPart", "ToolApprovalRespondedPart", + "ToolCallEvent", "ToolInputAvailablePart", "ToolInputStreamingPart", + "ToolOutput", "ToolOutputAvailablePart", "ToolOutputDeniedPart", "ToolOutputErrorPart", "ToolPart", + "ToolRequest", + "ToolResultEvent", + "ToolSpec", "ToolState", + "Turn", + "TurnResult", "UIMessage", "UIMessagePart", "UnknownPart", @@ -165,6 +191,7 @@ "WaitingRun", "WebSocketHandlers", "WebSockets", + "WhenBusy", "build_agent_path", "build_agent_url", "callable", @@ -176,4 +203,5 @@ "route_sub_agent_request", "subscribe", "task", + "tool", ) diff --git a/agents/src/agents/harness/__init__.py b/agents/src/agents/harness/__init__.py new file mode 100644 index 0000000..4f18d06 --- /dev/null +++ b/agents/src/agents/harness/__init__.py @@ -0,0 +1,58 @@ +"""Harnesses: agent loops hosted in a Durable Object (upstream ``agents/harness``).""" + +from .errors import OperationNotFoundError, SessionBusyError, SessionNotFoundError +from .harness import Harness +from .session import HarnessSession, HarnessSessions +from .tools import tool +from .types import ( + Approval, + HarnessEvent, + HarnessInput, + OperationEvent, + OperationResult, + OperationStatus, + PendingOperation, + Receipt, + Replay, + SessionSnapshot, + TextEvent, + ToolAnswer, + ToolCallEvent, + ToolOutput, + ToolRequest, + ToolResultEvent, + ToolSpec, + Turn, + TurnResult, + WhenBusy, +) + +__all__ = ( + "Approval", + "Harness", + "HarnessEvent", + "HarnessInput", + "HarnessSession", + "HarnessSessions", + "OperationEvent", + "OperationNotFoundError", + "OperationResult", + "OperationStatus", + "PendingOperation", + "Receipt", + "Replay", + "SessionBusyError", + "SessionNotFoundError", + "SessionSnapshot", + "TextEvent", + "ToolAnswer", + "ToolCallEvent", + "ToolOutput", + "ToolRequest", + "ToolResultEvent", + "ToolSpec", + "Turn", + "TurnResult", + "WhenBusy", + "tool", +) diff --git a/agents/src/agents/harness/errors.py b/agents/src/agents/harness/errors.py new file mode 100644 index 0000000..d669375 --- /dev/null +++ b/agents/src/agents/harness/errors.py @@ -0,0 +1,32 @@ +"""Errors raised by harnesses.""" + +from ..core.errors import AgentsException + +__all__ = ("OperationNotFoundError", "SessionBusyError", "SessionNotFoundError") + + +class SessionBusyError(AgentsException): + """A session has open work and the call refuses to queue behind it.""" + + def __init__(self, session: str) -> None: + super().__init__(f"Harness session {session!r} has an operation open") + self.session = session + + +class SessionNotFoundError(AgentsException): + """No session has this id.""" + + def __init__(self, session: str) -> None: + super().__init__(f"Harness session {session!r} doesn't exist") + self.session = session + + +class OperationNotFoundError(AgentsException): + """No operation in the session has this id.""" + + def __init__(self, session: str, operation_id: str) -> None: + super().__init__( + f"Harness session {session!r} has no operation {operation_id!r}" + ) + self.session = session + self.operation_id = operation_id diff --git a/agents/src/agents/harness/events.py b/agents/src/agents/harness/events.py new file mode 100644 index 0000000..efe43ed --- /dev/null +++ b/agents/src/agents/harness/events.py @@ -0,0 +1,37 @@ +"""Fanning each session's events out to its watchers.""" + +import asyncio + +from .types import HarnessEvent + +__all__ = ("EventHub",) + + +class EventHub: + """In-memory watchers per session; events aren't stored.""" + + __slots__ = ("_watchers",) + + def __init__(self) -> None: + self._watchers: dict[str, set[asyncio.Queue[HarnessEvent]]] = {} + + def subscribe(self, session: str) -> "asyncio.Queue[HarnessEvent]": + """Return a queue that receives the session's events from now on.""" + queue: asyncio.Queue[HarnessEvent] = asyncio.Queue() + self._watchers.setdefault(session, set()).add(queue) + return queue + + def unsubscribe(self, session: str, queue: "asyncio.Queue[HarnessEvent]") -> None: + """Stop sending the session's events to ``queue``.""" + watchers = self._watchers.get(session) + if watchers is None: + return + + watchers.discard(queue) + if not watchers: + del self._watchers[session] + + def publish(self, session: str, event: HarnessEvent) -> None: + """Send ``event`` to every watcher of the session.""" + for queue in self._watchers.get(session, ()): + queue.put_nowait(event) diff --git a/agents/src/agents/harness/harness.py b/agents/src/agents/harness/harness.py new file mode 100644 index 0000000..7fe03d8 --- /dev/null +++ b/agents/src/agents/harness/harness.py @@ -0,0 +1,438 @@ +"""The base class for harnesses: an agent loop hosted in a Durable Object. + +Python's form of upstream's draft ``AgentHarness`` interface +(``agents/harness/*``). Upstream harnesses each write their own queue, +dedupe, wake logic, and events; here the base class owns them and a subclass +implements one turn. Durability follows upstream's ThinkHarness: the +operations table is the record of open work, and one Lifecycle job per +session (``wake:``) brings the object back to finish it. +""" + +import asyncio +import logging +import uuid +from collections.abc import Sequence +from datetime import UTC, datetime, timedelta +from functools import partial +from typing import Any + +from ..core.platform_errors import is_platform_failure +from ..core.timing import to_seconds +from ..core.types import Duration +from ..lifecycle.capability import LifecycleCapability +from ..lifecycle.types import JobContext, JobOutcome, MemoryLimitContext, Reschedule +from .errors import ( + OperationNotFoundError, + SessionBusyError, + SessionNotFoundError, +) +from .events import EventHub +from .session import HarnessSession, HarnessSessions +from .store import HarnessStore, OperationRow +from .tools import collect_tools +from .types import ( + HarnessEvent, + HarnessInput, + OperationEvent, + OperationResult, + OperationStatus, + PendingOperation, + Receipt, + ToolRequest, + ToolSpec, + Turn, + TurnResult, + WhenBusy, +) + +__all__ = ("Harness",) + +_log = logging.getLogger("agents.harness") + +_WAKE = "wake" +_ROOT = "" + + +class Harness(LifecycleCapability): + """Runs an agent loop in a Durable Object, one turn at a time per session. + + Install it like any capability: ``Lifecycle(self).use(harness)`` on a + plain Durable Object (forward ``alarm`` to the Lifecycle), or + ``self.use(harness)`` on an `Agent`. The base class stores sessions and + operations, queues and dedupes submissions, publishes events, and resumes + turns a crash interrupted. Subclasses implement `run_turn`. + + Parameters + ---------- + id + The capability id; also keys this harness's rows, so two harnesses + on one object need different ids. + target + The object whose ``@tool`` methods are the tools (usually the host). + max_restarts + How many times a crash-interrupted turn is resumed before its + operation settles ``unanswered`` with reason ``"interrupted"``. + heartbeat + How often a session with open work is checked; the longest a turn + waits to resume after a crash when no request wakes the object. + + Raises + ------ + ValueError + If ``max_restarts`` is negative or ``heartbeat`` isn't positive. + """ + + def __init__( + self, + id: str, + *, + target: object | None = None, + max_restarts: int = 5, + heartbeat: Duration = 30, + ) -> None: + if max_restarts < 0: + raise ValueError("max_restarts must be 0 or more") + if to_seconds(heartbeat) <= 0: + raise ValueError("heartbeat must be positive") + + super().__init__(id) + self.target = target + """The object the tools were collected from.""" + self.sessions = HarnessSessions(self) + """Creating, forking, and listing sessions.""" + + self._tools = collect_tools(target) if target is not None else [] + self._max_restarts = max_restarts + self._heartbeat = timedelta(seconds=to_seconds(heartbeat)) + self._wake_prefix = f"{_WAKE}:{id}:" + self._store_instance: HarnessStore | None = None + self._events = EventHub() + self._drives: dict[str, asyncio.Task[None]] = {} + self._waiters: dict[tuple[str, str], list[asyncio.Future[None]]] = {} + + @property + def tools(self) -> Sequence[ToolSpec]: + """The target's ``@tool`` methods, bound.""" + return self._tools + + @property + def _store(self) -> HarnessStore: + if self._store_instance is None: + self._store_instance = HarnessStore(self.lifecycle.sql, self.capability_id) + + return self._store_instance + + async def run_turn(self, turn: Turn) -> TurnResult: + """Run one turn and return the session's new history. + + Call ``turn.save`` after each step so a resumed turn (``checkpoint`` + set) can skip finished work, and ``turn.emit`` to stream events. + Raising settles the operation ``unanswered`` with the error as reason. + """ + raise NotImplementedError + + def session(self, id: str = _ROOT) -> HarnessSession: + """Return a handle on one session (default: the root session).""" + return HarnessSession(self, id) + + async def prompt(self, input: HarnessInput) -> OperationResult: + """Submit to the root session and wait for the result.""" + return await self.session().prompt(input) + + async def on_start(self) -> None: + """Create the tables and wake every session with open work.""" + self._store.ensure_tables() + self._store.ensure_session(_ROOT) + + for session in self._store.sessions_with_open_work(): + await self._wake(session) + + async def on_job(self, context: JobContext) -> JobOutcome: + """Run a session's wake job: resume its open work, or complete.""" + if context.job.fn != _WAKE: + return None + + session = context.job.id.removeprefix(self._wake_prefix) + heartbeat = Reschedule(at=datetime.now(UTC) + self._heartbeat) + if self._is_driving(session): + return heartbeat + if not self._store.open(session): + return None + + self.lifecycle.track_alarm_work(self._ensure_drive(session)) + return heartbeat + + async def on_memory_limit(self, context: MemoryLimitContext) -> None: + """Settle running turns once the breaker seals; they'd only run out again.""" + if not context.sealed: + return + + for row in self._store.running(): + self._settle(row.session, row.id, "unanswered", reason="out_of_memory") + + async def dispose(self) -> None: + """Stop the drives; their turns resume on the next wake.""" + for drive in self._drives.values(): + drive.cancel() + + async def _submit( + self, + session: str, + input: HarnessInput, + *, + operation_id: str | None, + when_busy: WhenBusy, + ) -> Receipt: + await self.lifecycle.ready() + operation_id = operation_id or str(uuid.uuid4()) + + duplicate = self._store.get(session, operation_id) is not None + if not duplicate and when_busy == "reject" and self._store.open(session): + raise SessionBusyError(session) + + accepted = self.lifecycle.storage.transactionSync( + partial(self._insert, session, operation_id, input) + ) + if accepted: + self._events.publish( + session, OperationEvent(operation_id=operation_id, status="queued") + ) + + await self._wake(session) + self._ensure_drive(session) + return Receipt(session=session, operation_id=operation_id, accepted=accepted) + + def _insert(self, session: str, operation_id: str, input: HarnessInput) -> bool: + self._store.ensure_session(session) + return self._store.insert(session, operation_id, input) + + async def _wait(self, session: str, operation_id: str) -> OperationResult: + await self.lifecycle.ready() + + while True: + row = self._store.get(session, operation_id) + if row is None: + raise OperationNotFoundError(session, operation_id) + if not row.is_open: + return row.result() + + waiter = asyncio.get_running_loop().create_future() + self._waiters.setdefault((session, operation_id), []).append(waiter) + await waiter + + async def _pending(self, session: str) -> list[PendingOperation]: + await self.lifecycle.ready() + + return [row.pending() for row in self._store.open(session)] + + async def _abort(self, session: str, operation_id: str | None) -> None: + await self.lifecycle.ready() + + rows = [ + row + for row in self._store.open(session) + if operation_id is None or row.id == operation_id + ] + for row in rows: + self._settle(session, row.id, "aborted") + + drive = self._drives.get(session) + if drive is not None and any(row.status == "running" for row in rows): + drive.cancel() + await asyncio.wait([drive]) + + if self._store.open(session): + self._ensure_drive(session) + + async def _retry(self, session: str, operation_id: str) -> None: + await self.lifecycle.ready() + + row = self._store.get(session, operation_id) + if row is None: + raise OperationNotFoundError(session, operation_id) + if row.status != "unanswered": + raise ValueError( + f"Only unanswered operations can be retried; {operation_id!r} " + f"is {row.status}" + ) + + self._store.reopen(session, operation_id) + self._events.publish( + session, OperationEvent(operation_id=operation_id, status="queued") + ) + + await self._wake(session) + self._ensure_drive(session) + + async def _reset(self, session: str) -> None: + await self.lifecycle.ready() + + if self._store.open(session): + raise SessionBusyError(session) + + self._store.set_history(session, None) + + async def _create_session(self) -> HarnessSession: + await self.lifecycle.ready() + + id = str(uuid.uuid4()) + self._store.ensure_session(id) + return HarnessSession(self, id) + + async def _fork_session(self, source: str) -> HarnessSession: + await self.lifecycle.ready() + + if not self._store.has_session(source): + raise SessionNotFoundError(source) + + id = str(uuid.uuid4()) + self._store.fork_session(source, id) + return HarnessSession(self, id) + + async def _session_ids(self) -> list[str]: + await self.lifecycle.ready() + + return self._store.session_ids() + + def _subscribe(self, session: str) -> "asyncio.Queue[HarnessEvent]": + return self._events.subscribe(session) + + def _unsubscribe(self, session: str, queue: "asyncio.Queue[HarnessEvent]") -> None: + self._events.unsubscribe(session, queue) + + async def _wake(self, session: str) -> None: + await self.lifecycle.jobs.push( + id=self._wake_prefix + session, + fn=_WAKE, + time=datetime.now(UTC), + singleflight=True, + recovery_loop=True, + ) + + def _is_driving(self, session: str) -> bool: + drive = self._drives.get(session) + return drive is not None and not drive.done() + + def _ensure_drive(self, session: str) -> asyncio.Task[None]: + drive = self._drives.get(session) + if drive is not None and not drive.done(): + return drive + + drive = asyncio.create_task(self._drive(session)) + self._drives[session] = drive + drive.add_done_callback(partial(self._drive_done, session)) + return drive + + def _drive_done(self, session: str, drive: asyncio.Task[None]) -> None: + if self._drives.get(session) is drive: + del self._drives[session] + + if not drive.cancelled() and (error := drive.exception()) is not None: + _log.error("Harness drive for %r failed", session, exc_info=error) + + async def _drive(self, session: str) -> None: + while rows := self._store.open(session): + turn = self._start_turn(rows[0]) + if turn is not None: + await self._run(turn) + + def _start_turn(self, row: OperationRow) -> Turn | None: + interrupted = row.status == "running" + if interrupted and row.attempts >= self._max_restarts: + self._settle(row.session, row.id, "unanswered", reason="interrupted") + return None + + attempt = row.attempts + 1 if interrupted else row.attempts + self._store.start(row.session, row.id, attempts=attempt) + if not interrupted: + self._events.publish( + row.session, OperationEvent(operation_id=row.id, status="running") + ) + + return Turn( + session=row.session, + operation_id=row.id, + input=row.input, + history=self._store.history(row.session), + checkpoint=row.checkpoint, + attempt=attempt, + tools=self._tools, + emit=partial(self._events.publish, row.session), + save=partial(self._store.save_checkpoint, row.session, row.id), + ) + + async def _run(self, turn: Turn) -> None: + try: + result = await self.lifecycle.run_in_host_context( + partial(self.run_turn, turn) + ) + except Exception as error: + if is_platform_failure(error): + raise + + _log.exception("Harness turn %r failed", turn.operation_id) + self._settle( + turn.session, + turn.operation_id, + "unanswered", + reason=str(error) or type(error).__name__, + ) + return + + self._settle( + turn.session, + turn.operation_id, + result.status, + text=result.text, + requests=result.requests, + history=result.history, + ) + + def _settle( + self, + session: str, + operation_id: str, + status: OperationStatus, + *, + text: str | None = None, + requests: Sequence[ToolRequest] = (), + reason: str | None = None, + history: bytes | None = None, + ) -> None: + settled = self.lifecycle.storage.transactionSync( + partial( + self._write_settlement, + session, + operation_id, + status, + text=text, + requests=requests, + reason=reason, + history=history, + ) + ) + if not settled: + return + + self._events.publish( + session, OperationEvent(operation_id=operation_id, status=status) + ) + for waiter in self._waiters.pop((session, operation_id), ()): + if not waiter.done(): + waiter.set_result(None) + + def _write_settlement( + self, + session: str, + operation_id: str, + status: OperationStatus, + *, + history: bytes | None, + **result: Any, + ) -> bool: + if not self._store.settle(session, operation_id, status, **result): + return False + + if history is not None: + self._store.set_history(session, history) + return True diff --git a/agents/src/agents/harness/session.py b/agents/src/agents/harness/session.py new file mode 100644 index 0000000..f8d6302 --- /dev/null +++ b/agents/src/agents/harness/session.py @@ -0,0 +1,154 @@ +"""Handles on a harness's sessions, matching upstream's ``AgentHarness`` names.""" + +from collections.abc import AsyncIterator, Sequence +from typing import TYPE_CHECKING + +from .types import ( + HarnessEvent, + HarnessInput, + OperationResult, + PendingOperation, + Receipt, + SessionSnapshot, + WhenBusy, +) + +if TYPE_CHECKING: + from .harness import Harness + +__all__ = ("HarnessSession", "HarnessSessions") + + +class HarnessSession: + """One conversation: its history and its operations. No I/O until called. + + Parameters + ---------- + harness + The owning harness. + id + The session id (``""`` is the root session). + """ + + __slots__ = ("_harness", "id") + + def __init__(self, harness: "Harness", id: str) -> None: + self._harness = harness + self.id = id + + def __repr__(self) -> str: + """Return ````.""" + return f"" + + async def submit( + self, + input: HarnessInput, + *, + operation_id: str | None = None, + when_busy: WhenBusy = "queue", + ) -> Receipt: + """Durably queue a prompt or tool answers; returns before the turn runs. + + Parameters + ---------- + input + A prompt, or answers to a ``needs_input`` result's requests. + operation_id + Dedupes retries: submitting an existing id stores nothing new and + returns a receipt with ``accepted=False``. Generated when omitted. + when_busy + ``"queue"`` behind open work, or ``"reject"`` it. + + Raises + ------ + SessionBusyError + If ``when_busy="reject"`` and the session has open work. + """ + return await self._harness._submit( + self.id, input, operation_id=operation_id, when_busy=when_busy + ) + + async def wait(self, operation_id: str) -> OperationResult: + """Wait for an operation to settle and return its result. + + Raises + ------ + OperationNotFoundError + If the session has no such operation. + """ + return await self._harness._wait(self.id, operation_id) + + async def prompt( + self, input: HarnessInput, *, operation_id: str | None = None + ) -> OperationResult: + """Submit and wait for the result.""" + receipt = await self.submit(input, operation_id=operation_id) + + return await self.wait(receipt.operation_id) + + async def pending(self) -> Sequence[PendingOperation]: + """Return the open operations, oldest first.""" + return await self._harness._pending(self.id) + + async def abort(self, operation_id: str | None = None) -> None: + """Stop one open operation, or all of them; queued work behind runs.""" + await self._harness._abort(self.id, operation_id) + + async def retry(self, operation_id: str) -> None: + """Run an ``unanswered`` operation again, from its last checkpoint. + + Raises + ------ + OperationNotFoundError + If the session has no such operation. + ValueError + If the operation isn't ``unanswered``. + """ + await self._harness._retry(self.id, operation_id) + + async def reset(self) -> None: + """Clear the session's history. + + Raises + ------ + SessionBusyError + If the session has open work. + """ + await self._harness._reset(self.id) + + async def watch(self) -> AsyncIterator[HarnessEvent]: + """Yield a `SessionSnapshot` of open operations, then every event.""" + queue = self._harness._subscribe(self.id) + try: + yield SessionSnapshot(operations=await self.pending()) + while True: + yield await queue.get() + finally: + self._harness._unsubscribe(self.id, queue) + + +class HarnessSessions: + """Creating, forking, and listing a harness's sessions.""" + + __slots__ = ("_harness",) + + def __init__(self, harness: "Harness") -> None: + self._harness = harness + + async def create(self) -> HarnessSession: + """Create an empty session.""" + return await self._harness._create_session() + + async def fork(self, id: str) -> HarnessSession: + """Create a session with a copy of session ``id``'s history. + + Raises + ------ + SessionNotFoundError + If session ``id`` doesn't exist. + """ + return await self._harness._fork_session(id) + + async def list(self) -> Sequence[str]: + """Return every session id, oldest first.""" + return await self._harness._session_ids() diff --git a/agents/src/agents/harness/store.py b/agents/src/agents/harness/store.py new file mode 100644 index 0000000..d923c37 --- /dev/null +++ b/agents/src/agents/harness/store.py @@ -0,0 +1,348 @@ +"""SQL storage for harnesses: sessions and their operations. + +Rows are keyed by the harness's capability id, so several harnesses can share +one object. History and checkpoints are opaque bytes the subclass owns. +""" + +import json +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from ..core.sql import Sql +from ..core.timing import now_ms +from ..core.types import JSONValue +from .types import ( + Approval, + HarnessInput, + OperationResult, + OperationStatus, + PendingOperation, + ToolAnswer, + ToolOutput, + ToolRequest, +) + +__all__ = ("HarnessStore", "OperationRow") + +_SESSIONS = """CREATE TABLE IF NOT EXISTS cf_agents_harness_sessions ( + harness TEXT NOT NULL, + id TEXT NOT NULL, + history BLOB, + created_at INTEGER NOT NULL, + PRIMARY KEY (harness, id) +) WITHOUT ROWID""" +_OPERATIONS = """CREATE TABLE IF NOT EXISTS cf_agents_harness_operations ( + harness TEXT NOT NULL, + session TEXT NOT NULL, + id TEXT NOT NULL, + seq INTEGER NOT NULL, + input TEXT NOT NULL, + status TEXT NOT NULL, + checkpoint BLOB, + attempts INTEGER NOT NULL DEFAULT 0, + result TEXT, + reason TEXT, + created_at INTEGER NOT NULL, + PRIMARY KEY (harness, session, id) +) WITHOUT ROWID""" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class OperationRow: + """One stored operation, parsed.""" + + session: str + id: str + input: HarnessInput + status: OperationStatus + checkpoint: bytes | None + attempts: int + text: str | None + requests: Sequence[ToolRequest] + reason: str | None + + @classmethod + def from_sql(cls, row: dict[str, Any]) -> "OperationRow": + """Parse a row of the operations table.""" + result = json.loads(row["result"]) if row["result"] is not None else {} + + return cls( + session=row["session"], + id=row["id"], + input=_input_from_json(json.loads(row["input"])), + status=row["status"], + checkpoint=row["checkpoint"], + attempts=row["attempts"], + text=result.get("text"), + requests=[ToolRequest(**r) for r in result.get("requests", ())], + reason=row["reason"], + ) + + @property + def is_open(self) -> bool: + """Whether the operation is queued or running.""" + return self.status in ("queued", "running") + + def result(self) -> OperationResult: + """Return the settled operation's result.""" + return OperationResult( + session=self.session, + operation_id=self.id, + status=self.status, + text=self.text, + requests=self.requests, + reason=self.reason, + ) + + def pending(self) -> PendingOperation: + """Return the open operation's summary.""" + return PendingOperation( + session=self.session, + operation_id=self.id, + status="running" if self.status == "running" else "queued", + ) + + +class HarnessStore: + """Every SQL statement a harness runs; callers wrap multi-step writes.""" + + __slots__ = ("harness", "sql") + + def __init__(self, sql: Sql, harness: str) -> None: + self.sql = sql + self.harness = harness + + def ensure_tables(self) -> None: + """Create both tables (idempotent).""" + self.sql(_SESSIONS) + self.sql(_OPERATIONS) + + def ensure_session(self, session: str) -> None: + """Create the session if it doesn't exist.""" + self.sql( + "INSERT OR IGNORE INTO cf_agents_harness_sessions" + " (harness, id, created_at) VALUES (?, ?, ?)", + self.harness, + session, + now_ms(), + ) + + def has_session(self, session: str) -> bool: + """Return whether the session exists.""" + return bool( + self.sql( + "SELECT id FROM cf_agents_harness_sessions" + " WHERE harness = ? AND id = ?", + self.harness, + session, + ) + ) + + def session_ids(self) -> list[str]: + """Return every session id, oldest first.""" + rows = self.sql( + "SELECT id FROM cf_agents_harness_sessions WHERE harness = ?" + " ORDER BY created_at, id", + self.harness, + ) + return [row["id"] for row in rows] + + def fork_session(self, source: str, session: str) -> None: + """Create ``session`` with a copy of ``source``'s history.""" + self.sql( + "INSERT INTO cf_agents_harness_sessions (harness, id, history, created_at)" + " SELECT harness, ?, history, ? FROM cf_agents_harness_sessions" + " WHERE harness = ? AND id = ?", + session, + now_ms(), + self.harness, + source, + ) + + def history(self, session: str) -> bytes | None: + """Return the session's history, or ``None`` if it has none.""" + rows = self.sql( + "SELECT history FROM cf_agents_harness_sessions" + " WHERE harness = ? AND id = ?", + self.harness, + session, + ) + return rows[0]["history"] if rows else None + + def set_history(self, session: str, history: bytes | None) -> None: + """Replace the session's history.""" + self.sql( + "UPDATE cf_agents_harness_sessions SET history = ?" + " WHERE harness = ? AND id = ?", + history, + self.harness, + session, + ) + + def insert(self, session: str, id: str, input: HarnessInput) -> bool: + """Queue an operation at the end of the session; ``False`` if ``id`` exists.""" + rows = self.sql( + "INSERT INTO cf_agents_harness_operations" + " (harness, session, id, seq, input, status, created_at)" + " VALUES (?, ?, ?, (SELECT COALESCE(MAX(seq), 0) + 1" + " FROM cf_agents_harness_operations WHERE harness = ? AND session = ?)," + " ?, 'queued', ?)" + " ON CONFLICT DO NOTHING RETURNING id", + self.harness, + session, + id, + self.harness, + session, + json.dumps(_input_to_json(input)), + now_ms(), + ) + return bool(rows) + + def get(self, session: str, id: str) -> OperationRow | None: + """Return one operation, or ``None``.""" + rows = self.sql( + "SELECT * FROM cf_agents_harness_operations" + " WHERE harness = ? AND session = ? AND id = ?", + self.harness, + session, + id, + ) + return OperationRow.from_sql(rows[0]) if rows else None + + def open(self, session: str) -> list[OperationRow]: + """Return the session's open operations, oldest first.""" + rows = self.sql( + "SELECT * FROM cf_agents_harness_operations" + " WHERE harness = ? AND session = ? AND status IN ('queued', 'running')" + " ORDER BY seq", + self.harness, + session, + ) + return [OperationRow.from_sql(row) for row in rows] + + def running(self) -> list[OperationRow]: + """Return every running operation, in any session.""" + rows = self.sql( + "SELECT * FROM cf_agents_harness_operations" + " WHERE harness = ? AND status = 'running'", + self.harness, + ) + return [OperationRow.from_sql(row) for row in rows] + + def sessions_with_open_work(self) -> list[str]: + """Return the sessions that have open operations.""" + rows = self.sql( + "SELECT DISTINCT session FROM cf_agents_harness_operations" + " WHERE harness = ? AND status IN ('queued', 'running')", + self.harness, + ) + return [row["session"] for row in rows] + + def start(self, session: str, id: str, *, attempts: int) -> None: + """Mark the operation running with this many interruptions so far.""" + self.sql( + "UPDATE cf_agents_harness_operations SET status = 'running', attempts = ?" + " WHERE harness = ? AND session = ? AND id = ?", + attempts, + self.harness, + session, + id, + ) + + def save_checkpoint(self, session: str, id: str, checkpoint: bytes) -> None: + """Replace a running operation's checkpoint.""" + self.sql( + "UPDATE cf_agents_harness_operations SET checkpoint = ?" + " WHERE harness = ? AND session = ? AND id = ? AND status = 'running'", + checkpoint, + self.harness, + session, + id, + ) + + def settle( + self, + session: str, + id: str, + status: OperationStatus, + *, + text: str | None, + requests: Sequence[ToolRequest], + reason: str | None, + ) -> bool: + """Settle an open operation; ``False`` if it was already settled. + + An ``unanswered`` operation keeps its checkpoint, so ``retry`` resumes. + """ + result: dict[str, JSONValue] = { + "text": text, + "requests": [_request_to_json(request) for request in requests], + } + + rows = self.sql( + "UPDATE cf_agents_harness_operations" + " SET status = ?, result = ?, reason = ?," + " checkpoint = CASE WHEN ? = 'unanswered' THEN checkpoint END" + " WHERE harness = ? AND session = ? AND id = ?" + " AND status IN ('queued', 'running') RETURNING id", + status, + json.dumps(result), + reason, + status, + self.harness, + session, + id, + ) + return bool(rows) + + def reopen(self, session: str, id: str) -> None: + """Queue an unanswered operation again, resuming from its checkpoint.""" + self.sql( + "UPDATE cf_agents_harness_operations" + " SET status = 'queued', attempts = 0, result = NULL, reason = NULL" + " WHERE harness = ? AND session = ? AND id = ?", + self.harness, + session, + id, + ) + + +def _input_to_json(input: HarnessInput) -> JSONValue: + if isinstance(input, str): + return {"prompt": input} + return {"answers": [_answer_to_json(answer) for answer in input]} + + +def _input_from_json(data: dict[str, Any]) -> HarnessInput: + if "prompt" in data: + return data["prompt"] + return [_answer_from_json(answer) for answer in data["answers"]] + + +def _answer_to_json(answer: ToolAnswer) -> JSONValue: + if isinstance(answer, Approval): + return { + "type": "approval", + "id": answer.id, + "approved": answer.approved, + "message": answer.message, + } + return {"type": "output", "id": answer.id, "value": answer.value} + + +def _answer_from_json(data: dict[str, Any]) -> ToolAnswer: + if data["type"] == "approval": + return Approval( + id=data["id"], approved=data["approved"], message=data["message"] + ) + return ToolOutput(id=data["id"], value=data["value"]) + + +def _request_to_json(request: ToolRequest) -> JSONValue: + return { + "id": request.id, + "name": request.name, + "args": request.args, + "kind": request.kind, + } diff --git a/agents/src/agents/harness/tools.py b/agents/src/agents/harness/tools.py new file mode 100644 index 0000000..434c7c5 --- /dev/null +++ b/agents/src/agents/harness/tools.py @@ -0,0 +1,98 @@ +"""``@tool``: tools as ``async`` methods on the host. + +The decorator only marks the function, so the method stays callable as +usual; ``Harness(target=obj)`` collects the marked methods bound to ``obj``. +""" + +import inspect +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any, overload + +from .types import Replay, ToolSpec + +__all__ = ("collect_tools", "tool") + +_MARKER = "__agents_tool__" + + +@dataclass(slots=True, frozen=True) +class _ToolOptions: + requires_approval: bool + replay: Replay + + +@overload +def tool[F: Callable[..., Awaitable[Any]]](fn: F, /) -> F: ... +@overload +def tool[F: Callable[..., Awaitable[Any]]]( + *, requires_approval: bool = False, replay: Replay = "unsafe" +) -> Callable[[F], F]: ... +def tool[F: Callable[..., Awaitable[Any]]]( + fn: F | None = None, + /, + *, + requires_approval: bool = False, + replay: Replay = "unsafe", +) -> F | Callable[[F], F]: + """Make an ``async`` method a tool the harness gives the model. + + The method's name, signature, and docstring describe the tool. + + Parameters + ---------- + requires_approval + Stop the turn with a ``needs_input`` result before each call, until + the client sends an `Approval`. + replay + ``"safe"`` lets a crash-interrupted call run again on resume; + ``"unsafe"`` (the default) reports it to the model as not run. + + Raises + ------ + TypeError + If the function isn't ``async``. + ValueError + If ``replay`` isn't ``"safe"`` or ``"unsafe"``. + + Examples + -------- + >>> class Notes(Agent): + ... @tool(replay="safe") + ... async def search(self, query: str) -> list[str]: + ... '''Search the notes.''' + """ + if replay not in ("safe", "unsafe"): + raise ValueError(f"replay must be 'safe' or 'unsafe', not {replay!r}") + + options = _ToolOptions(requires_approval, replay) + + def mark(fn: F) -> F: + _require_async(fn) + setattr(fn, _MARKER, options) + return fn + + return mark if fn is None else mark(fn) + + +def _require_async(fn: object) -> None: + if not inspect.iscoroutinefunction(fn): + raise TypeError(f"Tools must be async functions: {fn!r}") + + +def collect_tools(target: object) -> list[ToolSpec]: + """Return ``target``'s ``@tool`` methods, bound (the nearest class wins).""" + nearest: dict[str, object] = {} + for cls in reversed(type(target).__mro__): + nearest.update(vars(cls)) + + return [ + ToolSpec( + fn=getattr(target, name), + name=name, + requires_approval=options.requires_approval, + replay=options.replay, + ) + for name, value in nearest.items() + if isinstance(options := getattr(value, _MARKER, None), _ToolOptions) + ] diff --git a/agents/src/agents/harness/types.py b/agents/src/agents/harness/types.py new file mode 100644 index 0000000..635243f --- /dev/null +++ b/agents/src/agents/harness/types.py @@ -0,0 +1,224 @@ +"""Types for harnesses: inputs, turns, operation results, and events.""" + +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass +from typing import Any, Literal + +from ..core.types import JSONValue + +__all__ = ( + "Approval", + "HarnessEvent", + "HarnessInput", + "OperationEvent", + "OperationResult", + "OperationStatus", + "PendingOperation", + "Receipt", + "Replay", + "SessionSnapshot", + "TextEvent", + "ToolAnswer", + "ToolCallEvent", + "ToolOutput", + "ToolRequest", + "ToolResultEvent", + "ToolSpec", + "Turn", + "TurnResult", + "WhenBusy", +) + +type OperationStatus = Literal[ + "queued", "running", "done", "needs_input", "unanswered", "aborted" +] +"""``queued`` and ``running`` are open; the rest are settled. ``needs_input`` +means the turn stopped on tool calls the client must answer with a new +operation.""" + +type WhenBusy = Literal["queue", "reject"] +"""What ``submit`` does when the session has open work: queue behind it, or +raise `SessionBusyError`.""" + +type Replay = Literal["safe", "unsafe"] +"""Whether a tool may run again after a crash interrupted it.""" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class ToolSpec: + """A ``@tool`` method bound to its host, as a harness receives it.""" + + fn: Callable[..., Awaitable[Any]] + name: str + requires_approval: bool = False + replay: Replay = "unsafe" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class Approval: + """The client's decision on a tool call that needs approval.""" + + id: str + approved: bool + message: str | None = None + + +@dataclass(slots=True, frozen=True, kw_only=True) +class ToolOutput: + """The result of a tool call the client ran itself.""" + + id: str + value: JSONValue + + +type ToolAnswer = Approval | ToolOutput +"""An answer to one of a ``needs_input`` result's requests.""" + +type HarnessInput = str | Sequence[ToolAnswer] +"""What an operation carries: a user prompt, or answers to tool requests.""" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class ToolRequest: + """A tool call the turn stopped on: approve it, or run it and send the output.""" + + id: str + name: str + args: dict[str, JSONValue] + kind: Literal["approval", "call"] + + +@dataclass(slots=True, frozen=True, kw_only=True) +class Receipt: + """What ``submit`` returns once the operation is stored. + + ``accepted`` is ``False`` when an operation with this id already existed. + """ + + session: str + operation_id: str + accepted: bool + + +@dataclass(slots=True, frozen=True, kw_only=True) +class OperationResult: + """A settled operation.""" + + session: str + operation_id: str + status: OperationStatus + text: str | None = None + requests: Sequence[ToolRequest] = () + reason: str | None = None + + +@dataclass(slots=True, frozen=True, kw_only=True) +class PendingOperation: + """An open operation.""" + + session: str + operation_id: str + status: Literal["queued", "running"] + + +@dataclass(slots=True, frozen=True, kw_only=True) +class SessionSnapshot: + """The first event ``watch`` yields: the session's open operations.""" + + operations: Sequence[PendingOperation] + type: Literal["snapshot"] = "snapshot" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class OperationEvent: + """An operation changed status.""" + + operation_id: str + status: OperationStatus + type: Literal["operation"] = "operation" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class TextEvent: + """Model text streamed during a turn.""" + + operation_id: str + delta: str + type: Literal["text"] = "text" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class ToolCallEvent: + """A tool started running.""" + + operation_id: str + tool_call_id: str + name: str + args: dict[str, JSONValue] + type: Literal["tool-call"] = "tool-call" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class ToolResultEvent: + """A tool finished.""" + + operation_id: str + tool_call_id: str + result: JSONValue + is_error: bool = False + type: Literal["tool-result"] = "tool-result" + + +type HarnessEvent = ( + SessionSnapshot | OperationEvent | TextEvent | ToolCallEvent | ToolResultEvent +) +"""What ``watch`` yields.""" + + +@dataclass(slots=True, frozen=True, kw_only=True) +class Turn: + """One turn for ``Harness.run_turn`` to run. + + Parameters + ---------- + session + The session id. + operation_id + The operation this turn answers. + input + The operation's prompt or tool answers. + history + The session's history before this turn, as the subclass last returned + it; ``None`` for a new session. + checkpoint + The last ``save``d bytes when resuming after a crash; ``None`` on a + first run. + attempt + 0 on a first run, then how many times the turn was interrupted. + tools + The host's ``@tool`` methods, bound. + emit + Publishes an event to the session's watchers. + save + Durably replaces the turn's checkpoint. + """ + + session: str + operation_id: str + input: HarnessInput + history: bytes | None + checkpoint: bytes | None + attempt: int + tools: Sequence[ToolSpec] + emit: Callable[[HarnessEvent], None] + save: Callable[[bytes], None] + + +@dataclass(slots=True, frozen=True, kw_only=True) +class TurnResult: + """What ``Harness.run_turn`` returns; ``history`` replaces the session's.""" + + history: bytes + status: Literal["done", "needs_input"] = "done" + text: str | None = None + requests: Sequence[ToolRequest] = () diff --git a/agents/tests/harness/test_harness.py b/agents/tests/harness/test_harness.py new file mode 100644 index 0000000..2aeaaa0 --- /dev/null +++ b/agents/tests/harness/test_harness.py @@ -0,0 +1,416 @@ +import asyncio +from collections.abc import Awaitable, Callable + +import pytest +from fake_runtime import FakeCtx, Host + +from agents.harness import ( + Approval, + Harness, + OperationEvent, + OperationNotFoundError, + SessionBusyError, + SessionNotFoundError, + SessionSnapshot, + TextEvent, + ToolRequest, + Turn, + TurnResult, + tool, +) +from agents.lifecycle import Lifecycle, MemoryLimitContext + +type Script = Callable[[Turn], Awaitable[TurnResult]] + + +async def echo(turn: Turn) -> TurnResult: + """Append the prompt to the history and answer with it.""" + turn.emit(TextEvent(operation_id=turn.operation_id, delta=str(turn.input))) + + history = (turn.history or b"") + str(turn.input).encode() + b";" + return TurnResult(history=history, text=f"echo {turn.input}") + + +class ScriptedHarness(Harness): + def __init__( + self, + script: Script = echo, + *, + id: str = "test-harness", + target: object = None, + max_restarts: int = 5, + ) -> None: + super().__init__(id, target=target, max_restarts=max_restarts) + + self.script = script + self.turns: list[Turn] = [] + + async def run_turn(self, turn: Turn) -> TurnResult: + self.turns.append(turn) + return await self.script(turn) + + +def setup( + script: Script = echo, host: Host | None = None, *, max_restarts: int = 5 +) -> tuple[Lifecycle, ScriptedHarness, Host]: + host = host or Host() + harness = ScriptedHarness(script, max_restarts=max_restarts) + return Lifecycle(host).use(harness), harness, host + + +def restarted(host: Host) -> Host: + """A new host on the same storage: the object after its isolate died.""" + new = Host() + new.ctx = FakeCtx(storage=host.ctx.storage) + return new + + +class Blocker: + """A script that saves a checkpoint, then blocks until released.""" + + def __init__(self) -> None: + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def __call__(self, turn: Turn) -> TurnResult: + turn.save(b"checkpoint:" + str(turn.input).encode()) + self.started.set() + await self.release.wait() + return await echo(turn) + + +def test_prompt_runs_a_turn_and_keeps_history() -> None: + _, harness, _ = setup() + + async def scenario() -> None: + first = await harness.prompt("hi") + second = await harness.prompt("again") + + assert (first.status, first.text) == ("done", "echo hi") + assert second.text == "echo again" + assert harness.turns[1].history == b"hi;" + assert harness.turns[0].attempt == 0 + + asyncio.run(scenario()) + + +def test_submit_dedupes_by_operation_id() -> None: + _, harness, _ = setup() + + async def scenario() -> None: + session = harness.session() + first = await session.submit("hi", operation_id="op-1") + second = await session.submit("hi", operation_id="op-1") + await session.wait("op-1") + + assert (first.accepted, second.accepted) == (True, False) + assert len(harness.turns) == 1 + + asyncio.run(scenario()) + + +def test_operations_run_one_at_a_time_in_order() -> None: + blocker = Blocker() + _, harness, _ = setup(blocker) + + async def scenario() -> None: + session = harness.session() + await session.submit("a", operation_id="a") + await session.submit("b", operation_id="b") + await blocker.started.wait() + assert [op.status for op in await session.pending()] == ["running", "queued"] + + blocker.release.set() + await session.wait("b") + assert [turn.input for turn in harness.turns] == ["a", "b"] + + asyncio.run(scenario()) + + +def test_reject_when_busy() -> None: + blocker = Blocker() + _, harness, _ = setup(blocker) + + async def scenario() -> None: + session = harness.session() + await session.submit("a", operation_id="a") + with pytest.raises(SessionBusyError): + await session.submit("b", when_busy="reject") + + receipt = await session.submit("a", operation_id="a", when_busy="reject") + assert not receipt.accepted + + blocker.release.set() + await session.wait("a") + + asyncio.run(scenario()) + + +def test_wait_for_an_unknown_operation_raises() -> None: + _, harness, _ = setup() + + async def scenario() -> None: + with pytest.raises(OperationNotFoundError): + await harness.session().wait("nope") + + asyncio.run(scenario()) + + +def test_failing_turn_settles_unanswered_and_retry_resumes() -> None: + failures = [RuntimeError("model down")] + + async def flaky(turn: Turn) -> TurnResult: + turn.save(b"step-1") + if failures: + raise failures.pop() + + return await echo(turn) + + _, harness, _ = setup(flaky) + + async def scenario() -> None: + session = harness.session() + failed = await session.prompt("hi", operation_id="op") + assert (failed.status, failed.reason) == ("unanswered", "model down") + + await session.retry("op") + result = await session.wait("op") + assert result.status == "done" + assert harness.turns[1].checkpoint == b"step-1" + + with pytest.raises(ValueError, match="unanswered"): + await session.retry("op") + + asyncio.run(scenario()) + + +def test_needs_input_then_answers_as_a_new_operation() -> None: + request = ToolRequest(id="call-1", name="delete", args={"id": 3}, kind="approval") + + async def approval(turn: Turn) -> TurnResult: + if isinstance(turn.input, str): + return TurnResult( + history=b"asked", status="needs_input", requests=[request] + ) + + return TurnResult(history=b"done", text=f"answered {turn.input[0].id}") + + _, harness, _ = setup(approval) + + async def scenario() -> None: + session = harness.session() + asked = await session.prompt("delete note 3") + assert (asked.status, asked.requests) == ("needs_input", [request]) + + answered = await session.prompt([Approval(id="call-1", approved=True)]) + assert answered.text == "answered call-1" + assert harness.turns[1].input == [Approval(id="call-1", approved=True)] + + asyncio.run(scenario()) + + +def test_abort_cancels_the_running_turn_and_runs_the_next() -> None: + blocker = Blocker() + _, harness, _ = setup(blocker) + + async def scenario() -> None: + session = harness.session() + await session.submit("a", operation_id="a") + await session.submit("b", operation_id="b") + await blocker.started.wait() + await session.abort("a") + assert (await session.wait("a")).status == "aborted" + + blocker.release.set() + assert (await session.wait("b")).status == "done" + assert harness.lifecycle.storage.sql.exec( + "SELECT history FROM cf_agents_harness_sessions WHERE id = ''" + ).toArray() == [{"history": b"b;"}] + + asyncio.run(scenario()) + + +def test_wake_job_resumes_a_turn_from_its_checkpoint_after_a_crash() -> None: + blocker = Blocker() + _, dead, host = setup(blocker) + + async def scenario() -> None: + await dead.session().submit("hi", operation_id="op") + await blocker.started.wait() + + lifecycle, harness, _ = setup(host=restarted(host)) + await lifecycle.start() + await lifecycle.alarm() + result = await harness.session().wait("op") + + assert result.text == "echo hi" + assert harness.turns[0].checkpoint == b"checkpoint:hi" + assert harness.turns[0].attempt == 1 + await dead.dispose() + + asyncio.run(scenario()) + + +def test_turn_settles_unanswered_after_max_restarts() -> None: + blocker = Blocker() + _, dead, host = setup(blocker) + + async def scenario() -> None: + await dead.session().submit("hi", operation_id="op") + await blocker.started.wait() + host.ctx.storage.sql.exec( + "UPDATE cf_agents_harness_operations SET attempts = 2" + ) + + lifecycle, harness, _ = setup(host=restarted(host), max_restarts=2) + await lifecycle.start() + await lifecycle.alarm() + result = await harness.session().wait("op") + + assert (result.status, result.reason) == ("unanswered", "interrupted") + assert harness.turns == [] + await dead.dispose() + + asyncio.run(scenario()) + + +def test_platform_failure_leaves_the_turn_running_for_recovery() -> None: + async def evicted(turn: Turn) -> TurnResult: + turn.save(b"step-1") + raise RuntimeError("Durable Object reset because its code was updated.") + + _, harness, _ = setup(evicted) + + async def scenario() -> None: + await harness.session().submit("hi", operation_id="op") + await asyncio.sleep(0) + + [pending] = await harness.session().pending() + assert pending.status == "running" + + asyncio.run(scenario()) + + +def test_two_harnesses_share_one_object() -> None: + host = Host() + first = ScriptedHarness(id="first") + second = ScriptedHarness(id="second") + Lifecycle(host).use(first).use(second) + + async def scenario() -> None: + assert (await first.prompt("a")).text == "echo a" + assert (await second.prompt("b")).text == "echo b" + + asyncio.run(scenario()) + + +def test_sealed_memory_limit_settles_running_turns() -> None: + blocker = Blocker() + _, harness, _ = setup(blocker) + + async def scenario() -> None: + await harness.session().submit("hi", operation_id="op") + await blocker.started.wait() + await harness.on_memory_limit(MemoryLimitContext(sealed=True)) + result = await harness.session().wait("op") + + assert (result.status, result.reason) == ("unanswered", "out_of_memory") + await harness.dispose() + + asyncio.run(scenario()) + + +def test_wake_job_completes_once_no_work_is_open() -> None: + lifecycle, harness, _ = setup() + + async def scenario() -> None: + await harness.prompt("hi") + await lifecycle.alarm() + + assert harness.lifecycle.jobs.list() == [] + + asyncio.run(scenario()) + + +def test_fork_copies_history_and_reset_clears_it() -> None: + _, harness, _ = setup() + + async def scenario() -> None: + await harness.prompt("hi") + fork = await harness.sessions.fork("") + await fork.prompt("more") + assert harness.turns[-1].history == b"hi;" + + await harness.session().reset() + await harness.prompt("fresh") + assert harness.turns[-1].history is None + + assert await harness.sessions.list() == ["", fork.id] + with pytest.raises(SessionNotFoundError): + await harness.sessions.fork("missing") + + asyncio.run(scenario()) + + +def test_watch_yields_a_snapshot_then_events() -> None: + _, harness, _ = setup() + + async def scenario() -> None: + events = harness.session().watch() + assert await anext(events) == SessionSnapshot(operations=[]) + + await harness.session().prompt("hi", operation_id="op") + seen = [await anext(events) for _ in range(4)] + assert seen == [ + OperationEvent(operation_id="op", status="queued"), + OperationEvent(operation_id="op", status="running"), + TextEvent(operation_id="op", delta="hi"), + OperationEvent(operation_id="op", status="done"), + ] + await events.aclose() + + asyncio.run(scenario()) + + +class Notes: + @tool(replay="safe") + async def search(self, query: str) -> list[str]: + return [query] + + @tool(requires_approval=True) + async def delete(self, note_id: int) -> str: + return f"deleted {note_id}" + + async def helper(self) -> None: ... + + +class Overridden(Notes): + async def delete(self, note_id: int) -> str: + return "" + + +def test_tools_are_the_targets_marked_methods_bound() -> None: + harness = ScriptedHarness(target=Notes()) + specs = {spec.name: spec for spec in harness.tools} + + assert set(specs) == {"search", "delete"} + assert (specs["search"].replay, specs["search"].requires_approval) == ( + "safe", + False, + ) + assert specs["delete"].requires_approval + assert asyncio.run(specs["search"].fn("x")) == ["x"] + + assert [spec.name for spec in ScriptedHarness(target=Overridden()).tools] == [ + "search" + ] + + +def test_tool_rejects_sync_functions_and_bad_replay() -> None: + with pytest.raises(TypeError, match="async"): + + @tool + def sync(self) -> None: ... + + with pytest.raises(ValueError, match="replay"): + tool(replay="sometimes") diff --git a/agents/tests/test_public_api.py b/agents/tests/test_public_api.py index 1e5db9b..1dfa0dd 100644 --- a/agents/tests/test_public_api.py +++ b/agents/tests/test_public_api.py @@ -5,6 +5,7 @@ core, dynamic_agents, fibers, + harness, lifecycle, observability, queue, @@ -22,6 +23,7 @@ core, dynamic_agents, fibers, + harness, lifecycle, observability, queue,