From 7f9a1beaf2dfe1c22fc001bab340b730513f9dac Mon Sep 17 00:00:00 2001 From: ran Date: Tue, 8 Sep 2026 16:34:37 +0200 Subject: [PATCH] feat: support resumable background runs over AG-UI Adds opt-in background execution to the AG-UI interface. Foreground runs are untouched and remain the default. A client opts in per run with forwardedProps.agnoBackground.enabled. The interface starts the run detached, then streams it by tailing Agno's buffered indexed event stream rather than the live producer. Every emitted AG-UI event carries a resume cursor in metadata.agnoBackground: the index of the Agno event it came from plus its ordinal within that event's expansion. A client that drops the connection reconnects with that cursor, receives the buffered remainder exactly once in canonical order, and continues live to completion. Ordinary messages, tool calls, tool results, state snapshots and deltas, errors and terminal events all survive the round trip. Message ids are minted deterministically for background runs so a replay reproduces the ids the first connection saw; foreground keeps random ids. Both agents and teams are covered. Support is probed at runtime rather than by version, and a request the entity cannot honor is refused with a run error instead of being silently downgraded or re-run. A cursor that points past a trimmed buffer is refused rather than answered with a gap. Consumes Agno's existing detached execution, event buffer, replay and tailing without modifying them. Tests cover replay ordering and exactly-once delivery deterministically, plus a real uvicorn server with an actual mid-stream disconnect and reconnect over HTTP, for both agents and teams. --- cookbook/05_agent_os/16_agui/README.md | 65 + cookbook/05_agent_os/16_agui/TEST_LOG.md | 111 +- .../05_agent_os/16_agui/background_run.py | 62 + .../agno/os/interfaces/agui/background.py | 675 +++++++ libs/agno/agno/os/interfaces/agui/handlers.py | 42 +- libs/agno/agno/os/interfaces/agui/router.py | 166 +- libs/agno/agno/os/interfaces/agui/state.py | 29 +- .../test_agui_background_reconnect.py | 407 ++++ .../os/interfaces/test_agui_background.py | 1785 +++++++++++++++++ 9 files changed, 3310 insertions(+), 32 deletions(-) create mode 100644 cookbook/05_agent_os/16_agui/background_run.py create mode 100644 libs/agno/agno/os/interfaces/agui/background.py create mode 100644 libs/agno/tests/integration/os/interfaces/test_agui_background_reconnect.py create mode 100644 libs/agno/tests/unit/os/interfaces/test_agui_background.py diff --git a/cookbook/05_agent_os/16_agui/README.md b/cookbook/05_agent_os/16_agui/README.md index 0b4837e4c30..b665e8b4a31 100644 --- a/cookbook/05_agent_os/16_agui/README.md +++ b/cookbook/05_agent_os/16_agui/README.md @@ -23,6 +23,7 @@ default empty prefix those routes are `POST /agui` and `GET /status`. | `human_in_the_loop.py` | Pause and resume a real backend tool that uses `requires_confirmation`. | | `research_team.py` | Stream a coordinated Team and its member activity over AG-UI. | | `multiple_instances.py` | Mount two independent AG-UI interfaces on one AgentOS. | +| `background_run.py` | Keep a run alive across a client disconnect and resume it from a cursor. | | `openui/` | Render an Agent as streaming charts, follow-ups, and validated forms with OpenUI. | ## Prerequisites @@ -51,6 +52,7 @@ Start one example at a time; every standalone server uses port 7777: .venvs/demo/bin/python cookbook/05_agent_os/16_agui/human_in_the_loop.py .venvs/demo/bin/python cookbook/05_agent_os/16_agui/research_team.py .venvs/demo/bin/python cookbook/05_agent_os/16_agui/multiple_instances.py +.venvs/demo/bin/python cookbook/05_agent_os/16_agui/background_run.py ``` The [`openui/`](openui/) example includes its own React client. Follow its @@ -71,6 +73,7 @@ endpoint: | `human_in_the_loop.py` | `http://localhost:7777/human-in-the-loop/agui` | `http://localhost:7777/human-in-the-loop/status` | | `research_team.py` | `http://localhost:7777/research-team/agui` | `http://localhost:7777/research-team/status` | | `multiple_instances.py` | `http://localhost:7777/chat/agui` and `http://localhost:7777/analyst/agui` | `/chat/status` and `/analyst/status` | +| `background_run.py` | `http://localhost:7777/background/agui` | `http://localhost:7777/background/status` | | `openui/server.py` | `http://localhost:7777/agui` | `http://localhost:7777/status` | The old all-in-one showcase is intentionally gone: starting the file you are @@ -149,3 +152,65 @@ snapshot. Media belongs in the latest user message as an AG-UI image, audio, video, or document content part. The adapter converts URL or base64 data sources into Agno media objects before calling the Gemini agent in `agent_with_media.py`. + +## Background runs and reconnection + +By default a run streams inline: closing the connection ends the run. A client +can instead ask for a background run, which executes detached from the request +that started it and buffers its events for replay. + +Opt in per request through `forwardedProps`: + +```json +"forwardedProps": {"agnoBackground": {"enabled": true}} +``` + +Every event of a background run then carries its resume cursor: + +```json +"metadata": {"agnoBackground": {"eventIndex": 12, "subIndex": 0}} +``` + +To reconnect, send the same `runId` again with the last cursor received. The +server replays whatever the client missed, once and in order, and then resumes +live streaming until the run finishes: + +```json +"forwardedProps": { + "agnoBackground": {"enabled": true, "lastEventIndex": 12, "lastSubIndex": 0} +} +``` + +A reconnection has to reproduce the events the first connection would have +sent, which shows up in how state is reported. A request that sends a `state` +is bracketed by a `STATE_SNAPSHOT` when the run starts and another when it +finishes, with no `STATE_DELTA` events in between; a request that sends none +gets neither, exactly as it would streaming inline. A client that follows +shared state along the run, as it can with `shared_state.py`, therefore sees +the state of a background run only as those two snapshots, and has to keep +sending `state` the same way on every connection. + +The buffer a background run replays from is finite. A run long enough to +overflow it keeps executing, and a connection that sends no resume position +still gets a well-formed stream starting from wherever the buffer now begins, +with a warning logged naming that event. A reconnection to such a run is refused with a +`RUN_ERROR`: the replay is rebuilt from what the buffer still holds, so a +message whose opening was trimmed would be opened again under a new identifier +while the client's own copy of it was never closed. + +A client that never sees the `agnoBackground` metadata cannot resume the run. +That covers both a server too old to offer background runs and a server that +has them but cannot apply them to this particular agent; either way the run +still streams normally on the one connection. + +Background execution needs a database, a readable run history so the server can +tell whose run a reconnection is naming, and an agent or team that executes in +this process, so a remote entity cannot use it. What happens +then depends on the request. A first connection runs in the foreground instead, +streaming normally without the ability to resume. A request that carries a +resume position is refused with a `RUN_ERROR` rather than downgraded, because +running it in the foreground would execute the whole run a second time. + +Continuing a paused run is not a resume: it starts a new leg, so it always +takes the foreground path and any resume position still being echoed alongside +it is ignored rather than refused. diff --git a/cookbook/05_agent_os/16_agui/TEST_LOG.md b/cookbook/05_agent_os/16_agui/TEST_LOG.md index 772fd71dfe6..84e1eb2b050 100644 --- a/cookbook/05_agent_os/16_agui/TEST_LOG.md +++ b/cookbook/05_agent_os/16_agui/TEST_LOG.md @@ -6,6 +6,13 @@ Tested on 2026-07-24 against Agno source commit The OpenUI addition was tested on 2026-08-18 against Agno source commit `32e5fb9c2203fa98de19ca72750133a57a075899`. +`background_run.py` was re-tested on 2026-09-04 and last re-confirmed on +2026-09-08 against the tree of the commit that carries this entry, on a machine +with no `OPENAI_API_KEY`. That run is scoped accordingly and its entry says what +it could not reach. The entry was rewritten from scratch because the previous +one, written against `8f76f52f41b4366dce9b6def7f4f687ede6c5229`, described +paused-run continuation behavior that has since changed. + Each checked-in server was first booted on its default port 7777. The sweep asserted `GET /health`, `GET /config`, every mounted AG-UI status route, and a clean shutdown. Capability-specific POST tests then used @@ -165,6 +172,96 @@ points, and both streams closed with `RUN_FINISHED`. --- +### background_run.py + +**Status:** PASS (routing, cursor stamping, and resume gating only; no model +call, so this is not a full-capability PASS like the entries above) + +**Test mode:** LIVE SERVER, NO MODEL KEY + +**Description:** Booted the checked-in server on its default port 7777 with no +`OPENAI_API_KEY` in the environment, and with `PYTHONPATH` pointed at this +worktree's `libs/agno` so the AG-UI interface under test is this tree's and not +another checkout's editable install. Checked `/health`, `/config`, and +`/background/status`, then sent ten `POST /background/agui` requests: a +background opt-in, two reconnections with the same `runId` from different +resume positions, a reconnection naming a `runId` no run uses, a second +background run seeded on a different thread plus a cross-thread reconnection to +it, a resume position with `enabled` set to `false`, two paused-run +continuations echoing a resume position, and a plain foreground run. + +**Result:** Health returned `ok`; config returned OS `agui-background-os`, +agent `agui-background-agent`, model `gpt-5.6-luna`, database +`agui-background-db`, and one AG-UI interface at route `/background`; +`/background/status` returned `available`. + +The background opt-in, sent as `forwardedProps.agnoBackground` set to the +boolean `true` shorthand rather than the `{"enabled": true}` form the README +documents, streamed five events. Both spellings are accepted. `RUN_STARTED` and `STATE_SNAPSHOT` came at cursors +`{"eventIndex": -1, "subIndex": 0}` and `{"eventIndex": -1, "subIndex": 1}`, +the snapshot carrying the submitted state. Two buffered Agno events with no +AG-UI handler arrived as `RAW` at cursors 0 and 1, wrapping `RunStarted` and +`ModelRequestStarted`. The stream ended at cursor 2 with a `RUN_ERROR` reading +`OPENAI_API_KEY not set. Please set the OPENAI_API_KEY environment variable.` + +Reconnecting with the same `runId` and `lastEventIndex` 1 replayed exactly one +event, the `RUN_ERROR` at cursor 2. Reconnecting with the same `runId` at +`lastEventIndex` -1 and `lastSubIndex` 0 replayed four events, everything from +the `STATE_SNAPSHOT` at cursor -1/1 onward, dropping only the `RUN_STARTED` at +-1/0. The cursor filter therefore discriminates within one event index as well +as across indices. + +A reconnection sending `lastEventIndex` 1 with `runId` `agui-bg-run-other`, +which no run uses, was refused with `Run agui-bg-run-other not found in this +session`, and that refusal was itself stamped at cursor 2, one past the +position the client sent. A second background run `agui-bg-run-2` was then +started on thread `agui-bg-thread-2`; reconnecting to it from thread +`agui-bg-thread-1` with `lastEventIndex` 0 was refused the same way, with +`Run agui-bg-run-2 not found in this session` stamped at cursor 1. Every +refusal observed carried a resume marker one event index past the client's own. + +A resume position sent as +`{"enabled": false, "lastEventIndex": 1, "lastSubIndex": 0}` was refused before +anything ran, with `A resume position was sent with background execution +disabled` stamped at cursor 2. + +A paused-run continuation, meaning a request carrying a trailing AG-UI tool +message, was not refused. With `agnoBackground` carrying `lastEventIndex` 1 the +server logged the warning `Background execution does not apply to a paused-run +continuation; continuing in the foreground` and took the foreground +continuation path, emitting `RUN_STARTED`, `STATE_SNAPSHOT`, and then +`RUN_ERROR` reading `No paused run matching the provided tool results found in +session agui-bg-thread-1`. None of those three events carried any +`agnoBackground` metadata. The same continuation with `agnoBackground` set to +`true` and no resume position behaved identically. That terminal error is the +continuation finding no paused run to resume in this key-less environment, not +a background refusal. + +A plain foreground run with empty `forwardedProps` produced the same five-event +shape as the opt-in run, but with no `agnoBackground` key on any event, so the +resume marker appears only on background responses. + +The server shut down cleanly on SIGINT. + +**Changed since the previous entry:** the previous version of this entry +recorded a paused-run continuation carrying a resume position as being refused +with `Background execution does not apply to a paused-run continuation` and +starting nothing. At this commit it is not refused. The router decides the +continuation downgrade before it decides the resume gate, so that sentence is +now only a server-side warning and the continuation proceeds in the foreground. + +**Not verified:** everything that needs a model call. No assistant text was +produced, so no long-running stream was disconnected and resumed mid-run, no +`TEXT_MESSAGE_*` or tool-call events were ever buffered or replayed, the +terminal `STATE_SNAPSHOT` and the absence of mid-run `STATE_DELTA` events were +not observed, and the buffer-overflow replay refusal was never reached. Those +behaviors are read from `agno/os/interfaces/agui/background.py` rather than +witnessed here. A genuinely paused run was also never reached, so the +foreground continuation path was observed only as far as its "no paused run" +rejection. + +--- + ### openui/server.py and frontend **Status:** PASS @@ -190,11 +287,13 @@ TypeScript compilation, and the Vite production build passed. ## Validation -- All 9 standalone files booted, exposed their expected `/status` route, and +- All 10 standalone files booted, exposed their expected `/status` route, and shut down cleanly. -- All 9 files completed a real capability-specific AG-UI POST flow. -- Recursive pattern validation checked exactly 9 Python files with 0 - violations. +- 9 of the 10 completed a real capability-specific AG-UI POST flow. The + exception is `background_run.py`, whose entry records what was verified + without a model call. +- Recursive pattern validation checked all 11 Python files in the folder, the + 10 standalone servers plus `openui/server.py`, with 0 violations. - Targeted Ruff format and check passed. - Python compilation, banned-model, stale-route, scope, Unicode/emoji, non-PASS status, and `git diff --check` gates passed. @@ -204,5 +303,5 @@ TypeScript compilation, and the Vite production build passed. check. Its frontend passed five tests, TypeScript compilation, and a production build. - Repository-wide Ruff, agnoctl mypy, and cookbook pattern checks passed. The - core Agno mypy step reported 27 existing errors in six files outside this - integration's diff. + core Agno mypy step now reports no issues across 1034 source files, so the 27 + pre-existing errors this section previously recorded are no longer present. diff --git a/cookbook/05_agent_os/16_agui/background_run.py b/cookbook/05_agent_os/16_agui/background_run.py new file mode 100644 index 00000000000..afe2d18a808 --- /dev/null +++ b/cookbook/05_agent_os/16_agui/background_run.py @@ -0,0 +1,62 @@ +""" +Resume a Background AG-UI Run +============================= + +Run an agent detached from the request that started it, so the run keeps going +after the client disconnects and a reconnecting client picks up exactly where +it left off. + +A client opts in per request with forwardedProps.agnoBackground. Every event of +such a run carries its resume cursor under metadata.agnoBackground. + +Resuming is addressed by run id, so the client has to choose the runId when it +starts the run and send that same runId back with the last cursor it received. +A reconnect that carries a different runId names a run the session does not +have, and is refused rather than resumed. + +Prerequisites: OPENAI_API_KEY +Run: .venvs/demo/bin/python cookbook/05_agent_os/16_agui/background_run.py +Try: POST a background run at http://localhost:7777/background/agui +""" + +from agno.agent import Agent +from agno.db.sqlite import SqliteDb +from agno.models.openai import OpenAIResponses +from agno.os import AgentOS +from agno.os.interfaces.agui import AGUI + +# --------------------------------------------------------------------------- +# Create Background-Capable Agent +# --------------------------------------------------------------------------- + +# Detached execution persists run status, so a database is required. +db = SqliteDb( + id="agui-background-db", + db_file="tmp/agui_background.db", +) + +long_form_agent = Agent( + id="agui-background-agent", + name="AG-UI Background Agent", + model=OpenAIResponses(id="gpt-5.6-luna"), + db=db, + instructions=[ + "Answer at length so the stream lasts long enough to disconnect from.", + "Write at least six paragraphs.", + ], +) + +agent_os = AgentOS( + id="agui-background-os", + description="AgentOS serving resumable background AG-UI runs.", + agents=[long_form_agent], + interfaces=[AGUI(agent=long_form_agent, prefix="/background")], +) +app = agent_os.get_app() + +# --------------------------------------------------------------------------- +# Run Background Server +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + agent_os.serve(app=app) diff --git a/libs/agno/agno/os/interfaces/agui/background.py b/libs/agno/agno/os/interfaces/agui/background.py new file mode 100644 index 00000000000..8e0f0e92962 --- /dev/null +++ b/libs/agno/agno/os/interfaces/agui/background.py @@ -0,0 +1,675 @@ +"""Resumable background runs for the AG-UI interface. + +A foreground AG-UI run streams straight out of the entity: when the client goes +away the run goes with it. A background run instead hands execution to Agno's +detached background streamer, which keeps executing after the client +disconnects and appends every event to the process event stream under a +monotonic index. This module is the bridge: it starts (or attaches to) such a +run and translates the buffered Agno events into AG-UI events. + +Because every AG-UI event is derived from a buffered Agno event, each one can +be addressed by the position of its source event plus its position within the +events that source event produced. That pair is the resume cursor: a client +echoes the last pair it saw and receives everything after it, exactly once, in +the same order it would have seen on an uninterrupted connection. + +Reproducing the same AG-UI events on a later connection requires the +translation to be a pure function of the buffered events, which drives two +choices here. ``StreamState`` is seeded with a namespace so message ids are +derived rather than random. And the translator works from its own copy of the +session state, so a run that was sent one is bracketed by a snapshot at each +end and emits no deltas in between, while one that was not is left alone. A +client that resumes has to keep sending state the same way it did, or the shape +of the run's last events shifts under it. + +Three limits are inherited rather than introduced here. The event stream's +index is monotonic but not gapless. Its buffer is finite, so a run long enough +to be trimmed replays from wherever the buffer now starts, and a client whose +own next event has been trimmed away is told to start over rather than handed a +renumbered stream. And a custom event whose class declares its own fields +reaches every background connection as its base class, so its own fields are +gone and its AG-UI name is the base name, because the buffer stores the wire +form rather than the object. +""" + +import asyncio +import contextlib +import copy +import json +import uuid +from typing import Any, AsyncIterator, Dict, List, Optional, Set, Tuple, Union + +from ag_ui.core import ( + BaseEvent, + EventType, + RawEvent, + RunAgentInput, + RunErrorEvent, + RunStartedEvent, + StateSnapshotEvent, +) + +from agno.agent import Agent, RemoteAgent +from agno.os.event_streams import get_event_stream +from agno.os.event_streams.base import BaseEventStream +from agno.os.interfaces.agui.handlers import is_completion_event, process_completion, process_event +from agno.os.interfaces.agui.input import ( + extract_context, + extract_media, + extract_user_input, + parse_client_tools, + validate_state, +) +from agno.os.interfaces.agui.state import StreamState +from agno.run.agent import RunCompletedEvent, RunEvent +from agno.run.agent import RunErrorEvent as AgnoRunErrorEvent +from agno.run.base import BaseRunOutputEvent, RunStatus +from agno.run.team import team_run_output_event_from_dict +from agno.team.remote import RemoteTeam +from agno.team.team import Team +from agno.utils.log import log_debug, log_error, log_warning + +# The forwarded_props key a client opts in with, and the event metadata key the +# server answers on. A client that sees the metadata knows the server honored +# the request and that reconnecting is safe; a server without this module +# simply never emits it, so the client stays on a single connection. +BACKGROUND_KEY = "agnoBackground" + +# Cursor for the events emitted before the first buffered event exists. +_PREFIX_INDEX = -1 + +Cursor = Tuple[int, int] + +# The events that end a run for a client. Everything else in a terminal group +# is either closing a span the client may already have been sent, or, for a run +# that paused, opening the ones its pending tool calls need. +_RUN_TERMINALS = (EventType.RUN_FINISHED, EventType.RUN_ERROR) + +# Detached drains hold the only strong reference to the producer generator for +# the life of the run; without it a garbage collection between two client +# connections would close the run the client is about to come back for. +_DRAIN_TASKS: set = set() + +# Run ids this process is between registering and starting. Two requests naming +# the same new run id would otherwise both pass the started-probe and start two +# producers into one buffer. This narrows that window to the gap between the +# probe and this set, which no await crosses. Across processes nothing +# arbitrates it, so two replicas handed the same new run id at the same moment +# can still both start one. +_STARTING_RUNS: Set[str] = set() + +# Runs this process started, and the entity, session and user each belongs to. +# The detached streamer writes the run row from inside its own task, so a client +# that reconnects immediately can arrive before that write lands. This record +# answers that without widening what a caller may attach to, and is written +# before the run is registered so nothing can wait past it. The entry goes when +# the drain finishes, so the map holds only runs still in flight here. +_STARTED_RUNS: Dict[str, Tuple[str, str, Optional[str]]] = {} + +# A background run persists its row before registering with the event stream, +# but this module registers the run first so that no event can be missed +# between starting and tailing. A reconnect landing inside that window would +# find no row yet, so the ownership read is retried briefly before refusing. +_OWNERSHIP_ATTEMPTS = 3 +_OWNERSHIP_RETRY_SECONDS = 0.1 + +# How long a second connection waits for the first one's registration. Longer +# than the ownership budget: this is waiting on another request's work rather +# than on a write that has already been issued. +_REGISTRATION_ATTEMPTS = 20 +_REGISTRATION_RETRY_SECONDS = 0.05 + + +class _OwnershipCheckFailed(Exception): + """The run's ownership could not be determined, as opposed to being denied.""" + + +def background_requested(run_input: RunAgentInput) -> bool: + """Whether this request asked for a resumable background run.""" + return _background_props(run_input) is not None + + +def background_cursor(run_input: RunAgentInput) -> Optional[Cursor]: + """The last cursor the client received, or None on a first connection. + + Read whatever the request carries, including when it also says background + is disabled: a resume position is a claim to be continuing an existing run, + and the route has to see it to refuse rather than run the whole thing again. + + Raises: + ValueError: the request carried a resume position that cannot be read. + Ignoring it would silently replay the whole run as if the client + had never connected. + """ + props = _background_props(run_input, opted_in_only=False) + if not isinstance(props, dict) or "lastEventIndex" not in props: + return None + event_index = props.get("lastEventIndex") + sub_index = props.get("lastSubIndex", 0) + if not _is_index(event_index) or not _is_index(sub_index): + raise ValueError(f"Unreadable background resume position: {props!r}") + return (int(event_index), int(sub_index)) # type: ignore[arg-type] + + +def supports_background(entity: Any) -> bool: + """Whether this entity can run detached in this process. + + Remote entities execute in another process, so their events never reach + this process's event stream and there is nothing here to replay. Detached + execution also needs a database to persist run status, and a readable run + history, without which there is no way to tell whose run a reconnection is + naming. + """ + if isinstance(entity, (RemoteAgent, RemoteTeam)): + return False + if getattr(entity, "db", None) is None: + return False + return callable(getattr(entity, "aget_run_output", None)) + + +async def run_entity_background( + entity: Union[Agent, Team], + run_input: RunAgentInput, + user_id: Optional[str] = None, +) -> AsyncIterator[BaseEvent]: + """Run an Agent or Team detached, streaming its buffered events as AG-UI events. + + Starts the run when this process has never seen it, and attaches to the + existing stream otherwise. Either way the events come from the buffer, so a + first connection and a reconnection produce the same sequence. + """ + run_id = run_input.run_id or str(uuid.uuid4()) + event_stream = get_event_stream() + cursor: Optional[Cursor] = None + highest: Optional[Cursor] = None + + try: + cursor = background_cursor(run_input) + session_state = validate_state(run_input.state, run_input.thread_id) + already_started = await event_stream.get_run_status(run_id) is not None + + starting_here = run_id in _STARTING_RUNS + + if already_started or starting_here or cursor is not None: + # Every path that attaches to a run rather than starting one goes + # through the same check: the run id is client-supplied, so without + # it a caller could name any run and read its events. + try: + owned = await _caller_owns_run(entity, run_id, run_input.thread_id, user_id) + except _OwnershipCheckFailed: + yield _stamp_refusal(f"Could not verify run {run_id}; try again", cursor) + return + if not owned: + yield _stamp_refusal(f"Run {run_id} not found in this session", cursor) + return + if starting_here and not already_started: + # Another request in this process is between registering this + # run and starting it. Attaching is what a second connection + # wants anyway, so wait rather than refuse. + if not await _wait_for_registration(event_stream, run_id): + yield _stamp_refusal(f"Run {run_id} did not start", cursor) + return + elif not already_started: + # The run exists but its events do not, so there is nothing to + # resume from. Saying so beats a silent empty stream. + yield _stamp_refusal(f"Run {run_id} is no longer available for replay", cursor) + return + else: + _STARTING_RUNS.add(run_id) + _STARTED_RUNS[run_id] = (_entity_key(entity), run_input.thread_id, user_id) + handed_over = False + try: + await event_stream.register_run(run_id, RunStatus.pending) + await _start_detached_run( + entity, run_input, run_id=run_id, user_id=user_id, session_state=session_state + ) + handed_over = True + finally: + _STARTING_RUNS.discard(run_id) + if not handed_over: + # Includes cancellation: a run registered but never handed + # over is one no producer will ever finish, and a tail + # attached to it would wait out the stream's idle recheck. + _STARTED_RUNS.pop(run_id, None) + await _abandon_run(event_stream, run_id) + + async for event in _stream_buffered_run( + event_stream=event_stream, + run_id=run_id, + thread_id=run_input.thread_id, + session_state=session_state, + cursor=cursor, + ): + delivered = _cursor_of(event) + if delivered is not None and (highest is None or delivered > highest): + highest = delivered + yield event + + except Exception as e: + log_error(f"Background AG-UI run {run_id} failed", exc_info=True) + yield _stamp_refusal(str(e)[:200], highest if highest is not None else cursor) + + +def _entity_key(entity: Any) -> str: + """A stable name for the entity a run was started on.""" + return str(getattr(entity, "id", None) or f"object:{id(entity)}") + + +def _is_index(value: Any) -> bool: + # bool is an int subclass, and a JSON true reaching here as index 1 would + # silently resume from the wrong place. + return isinstance(value, int) and not isinstance(value, bool) + + +def _stamp_refusal(message: str, after: Optional[Cursor]) -> BaseEvent: + """An error that ends the stream, positioned past everything before it. + + Stamped like any other event: the marker is what tells a client the server + honors background runs, and a client that filters by cursor would drop an + unmarked one. ``after`` is the last position the client is known to hold, + which is what it has already been sent on this connection when there is + one, and the resume position it asked from otherwise. + """ + event_index = (after[0] + 1) if after is not None else _PREFIX_INDEX + return _stamp(RunErrorEvent(type=EventType.RUN_ERROR, message=message), event_index, 0) + + +def background_error_event(message: str, after: Optional[Cursor]) -> BaseEvent: + """A stamped terminal error, for callers outside this module's own stream.""" + return _stamp_refusal(message, after) + + +def background_cursor_of(event: BaseEvent) -> Optional[Cursor]: + """The position a background event was delivered at, if it carries one.""" + return _cursor_of(event) + + +def _cursor_of(event: BaseEvent) -> Optional[Cursor]: + marker = (event.metadata or {}).get(BACKGROUND_KEY) + if not isinstance(marker, dict): + return None + event_index, sub_index = marker.get("eventIndex"), marker.get("subIndex") + if not _is_index(event_index) or not _is_index(sub_index): + return None + return (int(event_index), int(sub_index)) # type: ignore[arg-type] + + +def _background_props(run_input: RunAgentInput, *, opted_in_only: bool = True) -> Optional[Union[Dict[str, Any], bool]]: + """The background payload, or None when the request carries none.""" + forwarded = getattr(run_input, "forwarded_props", None) + if not isinstance(forwarded, dict): + return None + props = forwarded.get(BACKGROUND_KEY) + if props is True: + return props + if isinstance(props, dict) and (props.get("enabled", True) or not opted_in_only): + return props + return None + + +async def _wait_for_registration(event_stream: BaseEventStream, run_id: str) -> bool: + for _ in range(_REGISTRATION_ATTEMPTS): + if await event_stream.get_run_status(run_id) is not None: + return True + await asyncio.sleep(_REGISTRATION_RETRY_SECONDS) + log_warning(f"Background AG-UI run {run_id} never registered while starting") + return False + + +async def _caller_owns_run(entity: Any, run_id: str, thread_id: str, user_id: Optional[str]) -> bool: + """Whether ``run_id`` belongs to the session the caller is authorized for. + + The thread id is checked before streaming starts, but the run id is + separately client-supplied: without this an authorized caller could name + any run id and read another session's events back out of the buffer. + """ + if _STARTED_RUNS.get(run_id) == (_entity_key(entity), thread_id, user_id): + return True + reader = getattr(entity, "aget_run_output", None) + if not callable(reader): + log_warning(f"Cannot verify ownership of run {run_id}; refusing to attach") + return False + for attempt in range(_OWNERSHIP_ATTEMPTS): + try: + if await reader(run_id, session_id=thread_id, user_id=user_id) is not None: + return True + except Exception as e: + # A storage failure is not a denial. Reporting it as one sends the + # client off to restart a run that is still executing, so it is + # retried like a missing row and only then reported as its own + # kind of answer. + log_error(f"Ownership check failed for run {run_id}: {e}") + if attempt + 1 == _OWNERSHIP_ATTEMPTS: + raise _OwnershipCheckFailed(str(e)) from e + if attempt + 1 < _OWNERSHIP_ATTEMPTS: + await asyncio.sleep(_OWNERSHIP_RETRY_SECONDS) + return False + + +async def _start_detached_run( + entity: Union[Agent, Team], + run_input: RunAgentInput, + *, + run_id: str, + user_id: Optional[str], + session_state: Optional[Dict[str, Any]], +) -> None: + """Hand the run to Agno's detached background streamer. + + The returned generator is drained and discarded: its SSE strings are a + convenience for the connection that started the run, while the events this + interface streams always come from the buffer so that the starting + connection and every later one agree. + """ + from agno.run.base import RunContext + + event_stream = get_event_stream() + messages = run_input.messages or [] + images, audio, videos, files = extract_media(messages) + client_tools = parse_client_tools(run_input.tools) or None + ui_deps = extract_context(run_input.context) + + run_context = RunContext( + run_id=run_id, + session_id=run_input.thread_id, + user_id=user_id, + client_tools=client_tools, + dependencies=ui_deps, + session_state=session_state, + ) + + run_kwargs: Dict[str, Any] = {"run_context": run_context} + if ui_deps: + run_kwargs["add_dependencies_to_context"] = True + + producer = entity.arun( # type: ignore[call-overload] + input=extract_user_input(messages), + stream=True, + stream_events=True, + background=True, + session_id=run_input.thread_id, + user_id=user_id, + run_id=run_id, + images=images or None, + audio=audio or None, + videos=videos or None, + files=files or None, + **run_kwargs, + ) + + async def _drain() -> None: + try: + async for _ in producer: + pass + except asyncio.CancelledError: + # Whatever cancelled this, nothing else is going to end the run for + # the clients attached to it. + with contextlib.suppress(Exception): + await asyncio.shield(event_stream.complete_run(run_id, RunStatus.error)) + raise + except Exception: + log_error(f"Background AG-UI run {run_id} failed", exc_info=True) + # The producer is gone, so nothing else will end the run for the + # clients attached to it. + try: + await event_stream.complete_run(run_id, RunStatus.error) + except Exception: + log_error(f"Failed to mark background AG-UI run {run_id} as errored", exc_info=True) + + task = asyncio.create_task(_drain()) + _DRAIN_TASKS.add(task) + + def _forget(finished: "asyncio.Task[None]") -> None: + _DRAIN_TASKS.discard(finished) + _STARTED_RUNS.pop(run_id, None) + + task.add_done_callback(_forget) + log_debug(f"Started detached AG-UI run {run_id}") + + +async def _abandon_run(event_stream: BaseEventStream, run_id: str) -> None: + """Finish and forget a run that was registered but never handed over. + + The terminal comes first so anything already attached ends rather than + waiting out the stream's idle recheck. Forgetting it afterwards frees the + id to be tried again, instead of answering "already ended" for the life of + the process. + """ + try: + await event_stream.complete_run(run_id, RunStatus.error) + except Exception: + log_error(f"Failed to mark background AG-UI run {run_id} as errored", exc_info=True) + try: + await event_stream.cleanup_run(run_id) + except Exception: + log_error(f"Failed to drop the registration of background AG-UI run {run_id}", exc_info=True) + + +async def _stream_buffered_run( + *, + event_stream: BaseEventStream, + run_id: str, + thread_id: str, + session_state: Optional[Dict[str, Any]], + cursor: Optional[Cursor], +) -> AsyncIterator[BaseEvent]: + """Translate a run's buffered and live Agno events into AG-UI events.""" + # The state is copied rather than shared: a background run mutates its own + # session state inside the detached task, and a translation that read those + # live mutations would emit deltas a later replay could not reproduce. It is + # otherwise the foreground treatment, so a request that sends no state gets + # no snapshots invented for it. A client that resumes must keep sending + # state the same way, or the shape of the run's last events shifts under it. + run_state: Optional[Dict[str, Any]] = copy.deepcopy(session_state) if session_state is not None else None + state = StreamState(thread_id=thread_id, run_id=run_id, run_state=run_state) + state.id_namespace = f"{thread_id}/{run_id}" + # A replay can begin in the middle of a run, where a tool call completes + # whose start has been trimmed away. Ending a span that was never opened is + # invalid rather than merely incomplete. + state.require_started_tool_calls = True + if run_state is not None: + state.set_state_snapshot(run_state) + + terminal_chunk: Optional[BaseRunOutputEvent] = None + terminal_index = _PREFIX_INDEX + last_index = _PREFIX_INDEX + first_index: Optional[int] = None + opened = False + + tail = event_stream.tail(run_id, last_event_index=None) + try: + async for event_index, sse_data in tail: + if event_index <= _PREFIX_INDEX: + # The positions at or below this one belong to the events the + # run emits before its first buffered one. + log_warning(f"Skipping a background AG-UI frame for run {run_id} at index {event_index}") + continue + if first_index is None: + first_index = event_index + if first_index > 0: + log_warning( + f"Background AG-UI run {run_id} replays from event {first_index}; " + "earlier events have been trimmed from the buffer" + ) + if cursor is not None: + # Any trim at all defeats a resume. The translation is + # rebuilt from whatever the buffer still holds, so a + # message whose start was trimmed is opened again under + # a new id while the client's own copy of it is never + # closed. Say so rather than corrupt the stream. + yield _stamp_refusal( + f"Run {run_id} cannot be resumed from that position; " + "the events after it are no longer buffered", + cursor, + ) + return + last_index = max(last_index, event_index) + + if not opened: + opened = True + for event in _after(cursor, _PREFIX_INDEX, _opening_events(thread_id, run_id, run_state)): + yield event + + payload = _parse_frame(sse_data) + if payload is None: + continue + chunk = _event_from_payload(payload) + if chunk is None: + # The foreground translation turns an event it has no handler + # for into a raw event rather than dropping it; match that. + for event in _after(cursor, event_index, [_raw_event(payload)]): + yield event + continue + if is_completion_event(chunk): + # Held rather than emitted, and the last one wins: a team run + # buffers each member's completion before its own, so the first + # is not the end of anything the client is watching. + terminal_chunk, terminal_index = chunk, event_index + continue + state.set_id_seed(event_index) + for event in _after(cursor, event_index, process_event(chunk, state)): + yield event + finally: + closer = getattr(tail, "aclose", None) + if callable(closer): + await closer() + + if first_index is None and cursor is not None: + # Nothing at all came back, so there is no telling this from a buffer + # that lost everything the client is waiting for. + yield _stamp_refusal(f"Run {run_id} has no buffered events to resume from", cursor) + return + + if not opened: + for event in _after(cursor, _PREFIX_INDEX, _opening_events(thread_id, run_id, run_state)): + yield event + + status = await _run_status(event_stream, run_id) + if terminal_chunk is not None and _contradicts(terminal_chunk, status): + # A team run buffers its members' completions, so a held one can look + # like a finished run while the run itself died. The status wins. + log_warning(f"Background AG-UI run {run_id} ended with status {status.value if status else None}") + terminal_chunk = None + + if terminal_chunk is not None: + final_chunk: BaseRunOutputEvent = terminal_chunk + # Normally the terminal is the last thing buffered and keeps its own + # index. When something follows it, which a team run does because its + # members complete before it, the group goes after everything instead: + # reusing an index already emitted would hand two events one address. + final_index = terminal_index if terminal_index >= last_index else last_index + 1 + else: + # The tail ended without a terminal event: the producer died, or the + # stream's record of the run expired. The run's own status is the only + # honest account of how it ended, and anything short of completed is + # not a success to report as one. + if status == RunStatus.completed: + final_chunk = RunCompletedEvent() + else: + named = status.value.lower() if status is not None else "unknown" + final_chunk = AgnoRunErrorEvent(content=f"Run ended without a result, last known status {named}") + final_index = last_index + 1 + + state.set_id_seed(final_index) + for event in _closing_events(cursor, final_index, process_completion(final_chunk, state)): + yield event + + +def _opening_events(thread_id: str, run_id: str, run_state: Optional[Dict[str, Any]]) -> List[BaseEvent]: + events: List[BaseEvent] = [RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=run_id)] + if run_state is not None: + events.append(StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=copy.deepcopy(run_state))) + return events + + +def _raw_event(payload: Dict[str, Any]) -> BaseEvent: + return RawEvent(type=EventType.RAW, event=dict(payload), source="agno") + + +def _closing_events(cursor: Optional[Cursor], final_index: int, events: List[BaseEvent]) -> List[BaseEvent]: + """The run's last events, filtered by the cursor but never left unfinished. + + Span-closing events the client already received are dropped like any other. + The event that ends the run is not: a stream that closes without one leaves + a client holding a connection that will never say anything again. So when + the filter would take the terminal as well, only the terminal is re-issued, + past whatever position the client sent. + """ + delivered = _after(cursor, final_index, events) + if any(event.type in _RUN_TERMINALS for event in delivered): + return delivered + reissue_index = max(final_index, cursor[0] if cursor is not None else _PREFIX_INDEX) + 1 + return delivered + [ + _stamp(event, reissue_index, sub_index) + for sub_index, event in enumerate(event for event in events if event.type in _RUN_TERMINALS) + ] + + +def _contradicts(terminal_chunk: BaseRunOutputEvent, status: Optional[RunStatus]) -> bool: + """Whether the run's status says it ended worse than its last event claims.""" + if status not in (RunStatus.error, RunStatus.cancelled): + return False + event: Any = getattr(terminal_chunk, "event", None) + event_value = str(event.value if hasattr(event, "value") else event) + return event_value.removeprefix("Team") in (RunEvent.run_completed.value, RunEvent.run_paused.value) + + +async def _run_status(event_stream: BaseEventStream, run_id: str) -> Optional[RunStatus]: + try: + return await event_stream.get_run_status(run_id) + except Exception: + log_error(f"Could not read the status of background AG-UI run {run_id}", exc_info=True) + return None + + +def _after(cursor: Optional[Cursor], event_index: int, events: List[BaseEvent]) -> List[BaseEvent]: + """Stamp each event with its cursor and drop the ones already delivered.""" + out: List[BaseEvent] = [] + for sub_index, event in enumerate(events): + if cursor is not None and (event_index, sub_index) <= cursor: + continue + out.append(_stamp(event, event_index, sub_index)) + return out + + +def _stamp(event: BaseEvent, event_index: int, sub_index: int) -> BaseEvent: + metadata = dict(event.metadata) if event.metadata else {} + metadata[BACKGROUND_KEY] = {"eventIndex": event_index, "subIndex": sub_index} + event.metadata = metadata + return event + + +def _parse_frame(sse_data: str) -> Optional[Dict[str, Any]]: + """Read the event payload out of an SSE frame. + + The buffer hands back wire frames rather than objects on durable event + streams, so the frame is the only shape both stream implementations share. + """ + # Split on newlines only: the formatter writes non-ASCII through unescaped, + # and str.splitlines would also break on characters that are legal inside a + # JSON string, such as the line and paragraph separators. + data_lines = [line[5:].strip() for line in sse_data.split("\n") if line.startswith("data:")] + if not data_lines: + log_warning("Skipping an AG-UI background frame with no data") + return None + try: + payload = json.loads("\n".join(data_lines)) + except json.JSONDecodeError as e: + log_warning(f"Skipping an unreadable AG-UI background frame: {e}") + return None + if not isinstance(payload, dict) or not payload.get("event"): + log_warning("Skipping an AG-UI background frame that names no event") + return None + # Added when the event was written to the buffer rather than by the event + # itself, so it is dropped here and neither translation path sees it. + payload.pop("event_index", None) + return payload + + +def _event_from_payload(payload: Dict[str, Any]) -> Optional[BaseRunOutputEvent]: + """Rebuild the Agno event a frame was formatted from, or None if unknown.""" + try: + return team_run_output_event_from_dict(dict(payload)) + except Exception as e: + log_debug(f"No Agno event type for background frame {payload.get('event')}: {e}") + return None diff --git a/libs/agno/agno/os/interfaces/agui/handlers.py b/libs/agno/agno/os/interfaces/agui/handlers.py index 98a1165c282..85fb7c776c4 100644 --- a/libs/agno/agno/os/interfaces/agui/handlers.py +++ b/libs/agno/agno/os/interfaces/agui/handlers.py @@ -1,7 +1,6 @@ import copy import json -import uuid -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Union from ag_ui.core import ( BaseEvent, @@ -71,24 +70,32 @@ def _extract_team_response_chunk_content(response: TeamRunContentEvent) -> str: return main_content + members_response -def _format_reasoning_step(step: Optional[ReasoningStep], step_number: int = 0) -> str: - """Format a ReasoningStep as text for REASONING_MESSAGE_CONTENT.""" +def _format_reasoning_step(step: Optional[Union[ReasoningStep, Dict[str, Any]]], step_number: int = 0) -> str: + """Format a reasoning step as text for REASONING_MESSAGE_CONTENT. + + A step reaches a replayed stream as the plain mapping it was serialized to + rather than as a ReasoningStep, so both shapes are read the same way. + """ if step is None: return "" + + def field(name: str) -> Any: + return step.get(name) if isinstance(step, dict) else getattr(step, name, None) + parts: List[str] = [] - title = step.title or "Thinking" + title = field("title") or "Thinking" if step_number > 0: parts.append(f"## Step {step_number}: {title}") else: parts.append(f"## {title}") - if step.reasoning: - parts.append(step.reasoning) - if step.action: - parts.append(f"Action: {step.action}") - if step.result: - parts.append(f"Result: {step.result}") - if step.confidence is not None: - parts.append(f"Confidence: {step.confidence}") + if field("reasoning"): + parts.append(str(field("reasoning"))) + if field("action"): + parts.append(f"Action: {field('action')}") + if field("result"): + parts.append(f"Result: {field('result')}") + if field("confidence") is not None: + parts.append(f"Confidence: {field('confidence')}") return "\n".join(parts) + "\n\n" if parts else "" @@ -178,7 +185,7 @@ def on_tool_call_started(chunk: BaseRunOutputEvent, state: StreamState) -> List[ # Create empty parent message if none exists (AG-UI protocol requirement) if not parent_message_id: - parent_message_id = str(uuid.uuid4()) + parent_message_id = state.new_message_id() events.append( TextMessageStartEvent( type=EventType.TEXT_MESSAGE_START, @@ -219,6 +226,11 @@ def on_tool_call_completed(chunk: BaseRunOutputEvent, state: StreamState) -> Lis if tool.tool_call_id in state.ended_tool_call_ids: return events + if state.require_started_tool_calls and tool.tool_call_id not in state.active_tool_call_ids: + # This stream never saw the call start, so it has no span to end and no + # call for a result to belong to. Both are dropped together. + return events + events.append(ToolCallEndEvent(type=EventType.TOOL_CALL_END, tool_call_id=tool.tool_call_id)) state.end_tool_call(tool.tool_call_id) @@ -388,7 +400,7 @@ def on_run_completed(chunk: BaseRunOutputEvent, state: StreamState) -> List[Base paused_tools.append(req.tool_execution) if paused_tools: - assistant_message_id = str(uuid.uuid4()) + assistant_message_id = state.new_message_id() events.append( TextMessageStartEvent( type=EventType.TEXT_MESSAGE_START, diff --git a/libs/agno/agno/os/interfaces/agui/router.py b/libs/agno/agno/os/interfaces/agui/router.py index f2221f59b5c..7459cf72009 100644 --- a/libs/agno/agno/os/interfaces/agui/router.py +++ b/libs/agno/agno/os/interfaces/agui/router.py @@ -1,3 +1,5 @@ +import asyncio +import contextlib import copy import uuid from typing import AsyncIterator, Optional, Union @@ -21,6 +23,14 @@ from fastapi.responses import StreamingResponse from agno.agent import Agent, RemoteAgent +from agno.os.interfaces.agui.background import ( + background_cursor, + background_cursor_of, + background_error_event, + background_requested, + run_entity_background, + supports_background, +) from agno.os.interfaces.agui.input import ( extract_context, extract_media, @@ -134,6 +144,91 @@ async def run_entity( yield RunErrorEvent(type=EventType.RUN_ERROR, message=str(e)) +_SSE_HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Methods": "POST, GET, OPTIONS", + "Access-Control-Allow-Headers": "*", +} + + +# Idle window before an SSE comment goes out to hold the connection open. +_KEEPALIVE_INTERVAL_SECONDS = 15.0 +_KEEPALIVE_QUEUE_SIZE = 64 + + +async def _with_keepalives(events: AsyncIterator[BaseEvent], encoder: EventEncoder) -> AsyncIterator[str]: + """Encode events, emitting an SSE comment through any long silence.""" + # Bounded so a client reading slower than a replay produces cannot make the + # whole encoded run pile up in memory. + pending: asyncio.Queue = asyncio.Queue(maxsize=_KEEPALIVE_QUEUE_SIZE) + delivered: Optional[tuple] = None + failure: Optional[str] = None + done = object() + + async def _pump() -> None: + nonlocal delivered + try: + async for event in events: + position = background_cursor_of(event) + if position is not None and (delivered is None or position > delivered): + delivered = position + await pending.put(encoder.encode(event)) + except asyncio.CancelledError: + raise + except Exception as e: + nonlocal failure + log_error("AG-UI background stream failed", exc_info=True) + # Held rather than queued: a full queue would drop it, and a client + # that never hears why the stream ended is the thing this is for. + failure = encoder.encode(background_error_event(str(e)[:200], delivered)) + finally: + # Never awaited: a client that has gone away leaves nobody draining + # the queue, and a blocking put here would hang the cancellation + # that is trying to clean this task up. + with contextlib.suppress(asyncio.QueueFull): + pending.put_nowait(done) + + pump = asyncio.create_task(_pump()) + try: + while True: + try: + frame = await asyncio.wait_for(pending.get(), timeout=_KEEPALIVE_INTERVAL_SECONDS) + except asyncio.TimeoutError: + if pump.done() and pending.empty(): + # The queue was full when the pump finished, so its + # end-of-stream signal had nowhere to go. Holding the + # connection open on keepalives forever is worse than + # noticing here. + break + yield ": keepalive\n\n" + continue + if frame is done: + break + yield frame + finally: + pump.cancel() + with contextlib.suppress(BaseException): + await pump + if failure is not None: + yield failure + + +def _refusal_response(encoder: EventEncoder, message: str, after: Optional[tuple]) -> StreamingResponse: + """Answer a resume request the server cannot honor, without re-running it. + + Stamped past the client's position like every other background event: a + client that filters by cursor would otherwise drop the refusal and keep + reconnecting into it. + """ + + async def event_generator(): + yield encoder.encode(background_error_event(message, after)) + + return StreamingResponse(event_generator(), media_type="text/event-stream", headers=_SSE_HEADERS) + + def attach_routes( router: APIRouter, agent: Optional[Union[Agent, RemoteAgent]] = None, team: Optional[Union[Team, RemoteTeam]] = None ) -> APIRouter: @@ -159,20 +254,73 @@ async def run_agent_agui(request: Request, run_input: RunAgentInput): is_admin=caller_is_admin(request), ) + # A background run is opt-in per request: without it the run streams + # inline and dies with the connection, exactly as it always has. + run_in_background = background_requested(run_input) + continuing = bool(extract_tool_messages(run_input.messages or [])) + resuming_is_unreadable = False + try: + resuming = background_cursor(run_input) is not None + except ValueError: + # Unreadable, so it cannot be honored, but it is still a claim to be + # continuing a run rather than starting one. + resuming = True + resuming_is_unreadable = True + if resuming and not run_in_background and not continuing: + # A resume position with the feature switched off would otherwise + # fall through to a foreground run and execute the whole thing again. + return _refusal_response( + encoder, + "A resume position was sent with background execution disabled", + None if resuming_is_unreadable else background_cursor(run_input), + ) + + declined: Optional[str] = None + if run_in_background and continuing: + # Continuing a paused run starts a new leg rather than resuming the + # buffered one, so it takes the foreground continuation path and any + # resume position the client is still echoing does not apply to it. + # Decided first, so a continuation is never refused for carrying one. + log_warning( + "Background execution does not apply to a paused-run continuation; continuing in the foreground" + ) + run_in_background = False + elif run_in_background and not supports_background(entity): + declined = ( + f"Background execution is unavailable for '{getattr(entity, 'id', None)}': it needs a database, " + "an agent or team that runs in this process, and a readable run history" + ) + + if declined is not None: + run_in_background = False + # A request carrying a resume position is asking to continue a run + # that already exists. Quietly running it in the foreground would + # execute the whole thing a second time instead. + if resuming: + log_warning(f"{declined}; refusing to resume rather than running the whole run again") + return _refusal_response( + encoder, declined, background_cursor(run_input) if not resuming_is_unreadable else None + ) + log_warning(f"{declined}; running in the foreground (the run will not survive a disconnect)") + async def event_generator(): - async for event in run_entity(entity, run_input, user_id=user_id): # type: ignore - yield encoder.encode(event) + if not run_in_background: + async for event in run_entity(entity, run_input, user_id=user_id): # type: ignore[arg-type] + yield encoder.encode(event) + return + # A background run can sit waiting for a slot before it says + # anything, and a client whose connection a proxy closes in that + # silence has never seen a cursor to reconnect with. + async for frame in _with_keepalives( + run_entity_background(entity, run_input, user_id=user_id), # type: ignore[arg-type] + encoder, + ): + yield frame return StreamingResponse( event_generator(), media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "Access-Control-Allow-Origin": "*", - "Access-Control-Allow-Methods": "POST, GET, OPTIONS", - "Access-Control-Allow-Headers": "*", - }, + headers=_SSE_HEADERS, ) @router.get("/status") diff --git a/libs/agno/agno/os/interfaces/agui/state.py b/libs/agno/agno/os/interfaces/agui/state.py index 746a03a9ef4..f89bbb910a9 100644 --- a/libs/agno/agno/os/interfaces/agui/state.py +++ b/libs/agno/agno/os/interfaces/agui/state.py @@ -42,8 +42,33 @@ class StreamState: run_id: str = "" run_state: Optional[Dict[str, Any]] = None + # Deterministic message ids. Left unset the ids are random, which is all a + # single-pass stream needs. A replayable stream sets a namespace so that + # re-translating the same source events mints the same ids and a client + # that reconnects mid-message keeps receiving events for the message it + # already opened. The namespace identifies the run, so every connection to + # one run mints the same ids while two different runs never collide. + id_namespace: Optional[str] = None + # Set by a replayable stream, which can begin in the middle of a run and so + # must not report the end of a tool call whose start it never saw. + require_started_tool_calls: bool = False + _id_seed: str = field(default="", repr=False) + _id_counter: int = field(default=0, repr=False) + + def set_id_seed(self, seed: Any) -> None: + """Anchor the next ids to the position of the source event being translated.""" + self._id_seed = str(seed) + self._id_counter = 0 + + def new_message_id(self) -> str: + if self.id_namespace is None: + return str(uuid.uuid4()) + minted = self._id_counter + self._id_counter += 1 + return str(uuid.uuid5(uuid.NAMESPACE_URL, f"{self.id_namespace}/{self._id_seed}/{minted}")) + def open_text_message(self) -> str: - self.text_message_id = str(uuid.uuid4()) + self.text_message_id = self.new_message_id() self.text_message_open = True return self.text_message_id @@ -72,7 +97,7 @@ def clear_pending_tool_calls_parent_id(self) -> None: self.pending_tool_calls_parent_id = "" def start_reasoning(self) -> str: - self.reasoning_message_id = str(uuid.uuid4()) + self.reasoning_message_id = self.new_message_id() self.reasoning_step_count = 0 return self.reasoning_message_id diff --git a/libs/agno/tests/integration/os/interfaces/test_agui_background_reconnect.py b/libs/agno/tests/integration/os/interfaces/test_agui_background_reconnect.py new file mode 100644 index 00000000000..099a79ac8b6 --- /dev/null +++ b/libs/agno/tests/integration/os/interfaces/test_agui_background_reconnect.py @@ -0,0 +1,407 @@ +"""A real disconnect and reconnect against a live AG-UI server. + +Everything here goes over a socket: a uvicorn process serves AgentOS, an HTTP +client reads part of a background run's event stream and then drops the +connection mid-run, and a second request picks the run up from the cursor it +last saw. The point is to prove that the run keeps going with nobody attached +and that the two connections together see exactly the stream one uninterrupted +connection would have seen. +""" + +import asyncio +import json +import socket +import tempfile +import threading +import time +from pathlib import Path +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple + +import httpx +import pytest + +pytest.importorskip("ag_ui", reason="ag_ui not installed") + +import uvicorn + +from agno.agent import Agent +from agno.db.sqlite import SqliteDb +from agno.models.base import Model +from agno.models.message import MessageMetrics +from agno.models.response import ModelResponse +from agno.os.app import AgentOS +from agno.os.event_streams import get_event_stream, set_event_stream +from agno.os.event_streams.in_memory import InMemoryEventStream +from agno.os.interfaces.agui import AGUI +from agno.os.interfaces.agui import background as background_module +from agno.os.managers import EventsBuffer, SSESubscriberManager +from agno.run.base import RunStatus +from agno.team import Team + +CHUNKS = ["The ", "quick ", "brown ", "fox ", "jumps ", "over ", "the ", "lazy ", "dog."] +# Long enough that a client can disconnect part-way through with the run still +# producing, short enough to keep the test quick. +CHUNK_DELAY_SECONDS = 0.12 + +SERVER_START_TIMEOUT_SECONDS = 30.0 +SERVER_START_POLL_SECONDS = 0.05 +SERVER_STOP_TIMEOUT_SECONDS = 10.0 + +# Long enough for the whole answer to dribble out on a loaded machine, short +# enough that a run which never ends fails the test rather than hanging it. +RUN_COMPLETION_TIMEOUT_SECONDS = 60.0 +STATUS_POLL_SECONDS = 0.02 + +TERMINAL_STATUSES = (RunStatus.completed, RunStatus.error, RunStatus.cancelled, RunStatus.paused) + + +class SlowModel(Model): + """An offline model that dribbles out a fixed answer.""" + + def __init__(self, chunks: List[str]): + super().__init__(id="slow-test-model", name="slow-test-model", provider="test") + self.instructions = None + self._chunks = chunks + + def _response(self, text: str) -> ModelResponse: + return ModelResponse(content=text, role="assistant", response_usage=MessageMetrics()) + + def get_instructions_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + def get_system_message_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + async def aget_instructions_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + async def aget_system_message_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + def parse_args(self, *args: Any, **kwargs: Any) -> Dict[str, Any]: + return {} + + def invoke(self, *args: Any, **kwargs: Any) -> ModelResponse: + return self._response("".join(self._chunks)) + + async def ainvoke(self, *args: Any, **kwargs: Any) -> ModelResponse: + return self._response("".join(self._chunks)) + + def invoke_stream(self, *args: Any, **kwargs: Any) -> Iterator[ModelResponse]: + for chunk in self._chunks: + yield self._response(chunk) + + async def ainvoke_stream(self, *args: Any, **kwargs: Any) -> AsyncIterator[ModelResponse]: + for chunk in self._chunks: + await asyncio.sleep(CHUNK_DELAY_SECONDS) + yield self._response(chunk) + + def _parse_provider_response(self, response: Any, **kwargs: Any) -> ModelResponse: + return self._response("") + + def _parse_provider_response_delta(self, response: Any) -> ModelResponse: + return self._response("") + + +@pytest.fixture +def isolated_run_state() -> Iterator[None]: + """Empty the buffered events and the started-run record between tests. + + Both are process globals that the server thread reads, so a run id one + test used is otherwise still answerable in the next. Hand-picking a + distinct id per test hides that rather than fixing it, and hides it least + reliably in exactly the tests that name an id the server never saw. + + The drain task set is deliberately left alone: it holds the only strong + reference keeping a live background run from being collected. + """ + original = get_event_stream() + set_event_stream( + InMemoryEventStream(events_buffer=EventsBuffer(), subscriber_manager=SSESubscriberManager()), + ) + background_module._STARTED_RUNS.clear() + background_module._STARTING_RUNS.clear() + try: + yield + finally: + set_event_stream(original) + background_module._STARTED_RUNS.clear() + background_module._STARTING_RUNS.clear() + + +@pytest.fixture +def agui_server(isolated_run_state: None) -> Iterator[str]: + """A uvicorn server exposing an agent and a team over AG-UI. + + ``isolated_run_state`` is required rather than merely useful, and the + order is the whole point: the server thread reads the event stream out of + a process global, so the fresh one has to be installed before the thread + starts and restored only after it has stopped. Requesting it here is what + pins it either side of this fixture. A test that took the two + independently could be handed them the other way round, and would then run + against whichever stream the previous test left behind. + """ + with tempfile.TemporaryDirectory() as directory: + db = SqliteDb(db_file=str(Path(directory) / "agui-background.db")) + agent = Agent(name="bg-agent", id="bg-agent", model=SlowModel(CHUNKS), db=db) + member = Agent(name="member", id="member", model=SlowModel(CHUNKS), db=db) + team = Team(name="bg-team", id="bg-team", members=[member], model=SlowModel(CHUNKS), db=db) + + agent_os = AgentOS( + agents=[agent, member], + teams=[team], + interfaces=[AGUI(agent=agent, prefix="/agent"), AGUI(team=team, prefix="/team")], + ) + + # The listening socket is bound here and handed to uvicorn still open, + # so nothing else on the machine can claim the port in between. + listener = socket.socket() + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("127.0.0.1", 0)) + port = int(listener.getsockname()[1]) + + config = uvicorn.Config(agent_os.get_app(), host="127.0.0.1", port=port, log_level="error") + server = uvicorn.Server(config) + thread = threading.Thread(target=server.run, kwargs={"sockets": [listener]}, daemon=True) + thread.start() + try: + deadline = time.monotonic() + SERVER_START_TIMEOUT_SECONDS + while not server.started: + if not thread.is_alive(): + raise RuntimeError("AG-UI test server thread exited before the server started") + if time.monotonic() > deadline: + raise RuntimeError(f"AG-UI test server did not start within {SERVER_START_TIMEOUT_SECONDS}s") + time.sleep(SERVER_START_POLL_SECONDS) + yield f"http://127.0.0.1:{port}" + finally: + server.should_exit = True + thread.join(timeout=SERVER_STOP_TIMEOUT_SECONDS) + listener.close() + + +def request_body( + *, + thread_id: str, + run_id: str, + cursor: Optional[Tuple[int, int]] = None, +) -> Dict[str, Any]: + background: Dict[str, Any] = {"enabled": True} + if cursor is not None: + background["lastEventIndex"] = cursor[0] + background["lastSubIndex"] = cursor[1] + return { + "threadId": thread_id, + "runId": run_id, + "state": None, + "messages": [{"id": "m1", "role": "user", "content": "describe a fox"}], + "tools": [], + "context": [], + "forwardedProps": {"agnoBackground": background}, + } + + +def cursor_of(event: Dict[str, Any]) -> Tuple[int, int]: + marker = event["metadata"]["agnoBackground"] + return marker["eventIndex"], marker["subIndex"] + + +def content_of(events: List[Dict[str, Any]]) -> List[str]: + """The answer text deltas, in the order they arrived.""" + return [event["delta"] for event in events if event["type"] == "TEXT_MESSAGE_CONTENT"] + + +def identity(event: Dict[str, Any]) -> Dict[str, Any]: + """The parts of an event that must be identical across a reconnect. + + Every field the server sends is derived from the buffered Agno event, down + to the raw payload a RAW event carries, so the whole event is compared. The + exception is the protocol's optional timestamp: it is wall clock, so a leg + that carried one could never replay equal. + """ + return {key: value for key, value in event.items() if key != "timestamp"} + + +async def wait_for_terminal_status(run_id: str) -> RunStatus: + """Wait for the run to end, and report how. + + The status lives in the process-global event stream the server thread + writes to, so this reads the very record the server keeps rather than + inferring from a sleep that a run had time to finish. + """ + stream = get_event_stream() + loop = asyncio.get_running_loop() + deadline = loop.time() + RUN_COMPLETION_TIMEOUT_SECONDS + while True: + status = await stream.get_run_status(run_id) + if status in TERMINAL_STATUSES: + return status + assert loop.time() < deadline, f"run {run_id} never reached a terminal status (last seen {status})" + await asyncio.sleep(STATUS_POLL_SECONDS) + + +async def read_events( + client: httpx.AsyncClient, + url: str, + body: Dict[str, Any], + *, + stop_after_content: Optional[int] = None, +) -> List[Dict[str, Any]]: + """Read AG-UI events, optionally dropping the connection part-way through. + + Stopping is counted in answer text rather than in events, so a leg that is + cut short always holds content a later leg must not repeat. + """ + events: List[Dict[str, Any]] = [] + content_seen = 0 + async with client.stream("POST", url, json=body) as response: + assert response.status_code == 200 + async for line in response.aiter_lines(): + if not line.startswith("data:"): + continue + event = json.loads(line[5:].strip()) + events.append(event) + if event["type"] == "TEXT_MESSAGE_CONTENT": + content_seen += 1 + if stop_after_content is not None and content_seen >= stop_after_content: + break + return events + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ["agent", "team"]) +async def test_disconnect_and_reconnect_delivers_every_event_once(agui_server: str, entity: str): + url = f"{agui_server}/{entity}/agui" + thread_id = f"{entity}-interrupted" + run_id = f"{entity}-interrupted-run" + + async with httpx.AsyncClient(timeout=60.0) as client: + first_leg = await read_events( + client, url, request_body(thread_id=thread_id, run_id=run_id), stop_after_content=3 + ) + assert content_of(first_leg) == CHUNKS[:3] + + # Nothing is attached for a moment: the run has to survive on its own. + await asyncio.sleep(CHUNK_DELAY_SECONDS * 3) + + resumed = await read_events( + client, url, request_body(thread_id=thread_id, run_id=run_id, cursor=cursor_of(first_leg[-1])) + ) + # The whole run from the top, which is what an uninterrupted client + # would have received. + whole_run = await read_events(client, url, request_body(thread_id=thread_id, run_id=run_id)) + + assert resumed[-1]["type"] == "RUN_FINISHED" + assert content_of(resumed) == CHUNKS[3:] + assert [identity(event) for event in first_leg + resumed] == [identity(event) for event in whole_run] + + cursors = [cursor_of(event) for event in first_leg + resumed] + assert cursors == sorted(cursors) + assert len(set(cursors)) == len(cursors) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ["agent", "team"]) +async def test_run_completes_after_the_client_goes_away(agui_server: str, entity: str): + """The run reaches its end with nobody attached, and the client picks it up after. + + Waiting on the run's recorded status rather than on a sleep is what tells + a run that carried on from one that was merely still going when the second + connection arrived: the run is known to be over before that request goes + out, so the events it receives can only have been buffered while nothing + was reading them. + """ + url = f"{agui_server}/{entity}/agui" + thread_id = f"{entity}-detached" + run_id = f"{entity}-detached-run" + + async with httpx.AsyncClient(timeout=60.0) as client: + first_leg = await read_events( + client, url, request_body(thread_id=thread_id, run_id=run_id), stop_after_content=2 + ) + assert content_of(first_leg) == CHUNKS[:2] + + assert await wait_for_terminal_status(run_id) == RunStatus.completed + + resumed = await read_events( + client, url, request_body(thread_id=thread_id, run_id=run_id, cursor=cursor_of(first_leg[-1])) + ) + + assert resumed[-1]["type"] == "RUN_FINISHED" + # Picking up where the first leg stopped, instead of answering again from + # the top, is what tells a run that kept going from one that started over. + assert content_of(resumed) == CHUNKS[2:] + assert content_of(first_leg + resumed) == CHUNKS + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ["agent", "team"]) +async def test_resuming_from_a_position_already_passed_replays_only_what_follows(agui_server: str, entity: str): + url = f"{agui_server}/{entity}/agui" + thread_id = f"{entity}-rewind" + run_id = f"{entity}-rewind-run" + + async with httpx.AsyncClient(timeout=60.0) as client: + whole_run = await read_events(client, url, request_body(thread_id=thread_id, run_id=run_id)) + assert whole_run[-1]["type"] == "RUN_FINISHED" + + # A position the finished run went past long ago. + pivot = [index for index, event in enumerate(whole_run) if event["type"] == "TEXT_MESSAGE_CONTENT"][2] + rest = await read_events( + client, url, request_body(thread_id=thread_id, run_id=run_id, cursor=cursor_of(whole_run[pivot])) + ) + + assert content_of(rest) == CHUNKS[3:] + assert rest[-1]["type"] == "RUN_FINISHED" + assert [identity(event) for event in whole_run[: pivot + 1] + rest] == [identity(event) for event in whole_run] + + +@pytest.mark.asyncio +async def test_reconnecting_to_another_threads_run_is_refused(agui_server: str): + url = f"{agui_server}/agent/agui" + run_id = "owned-run" + + async with httpx.AsyncClient(timeout=60.0) as client: + owned = await read_events(client, url, request_body(thread_id="owner", run_id=run_id)) + assert owned[-1]["type"] == "RUN_FINISHED" + + intruder = await read_events( + client, url, request_body(thread_id="intruder", run_id=run_id, cursor=cursor_of(owned[2])) + ) + + # A lone RUN_ERROR: no RUN_STARTED, so no run began for the intruding + # thread, and none of the owner's answer was handed over. + assert [event["type"] for event in intruder] == ["RUN_ERROR"] + assert intruder[0]["message"] == f"Run {run_id} not found in this session" + + +@pytest.mark.asyncio +async def test_resuming_a_run_the_server_never_saw_is_refused(agui_server: str): + url = f"{agui_server}/agent/agui" + run_id = "never-started-run" + + async with httpx.AsyncClient(timeout=60.0) as client: + refused = await read_events(client, url, request_body(thread_id="ghost", run_id=run_id, cursor=(4, 0))) + assert [event["type"] for event in refused] == ["RUN_ERROR"] + assert refused[0]["message"] == f"Run {run_id} not found in this session" + + # The refusal ran nothing, so the same id is still free to start fresh. + fresh = await read_events(client, url, request_body(thread_id="ghost", run_id=run_id)) + + assert fresh[0]["type"] == "RUN_STARTED" + assert cursor_of(fresh[0]) == (-1, 0) + assert content_of(fresh) == CHUNKS + assert fresh[-1]["type"] == "RUN_FINISHED" + + +@pytest.mark.asyncio +async def test_foreground_run_carries_no_background_marker(agui_server: str): + url = f"{agui_server}/agent/agui" + body = request_body(thread_id="foreground", run_id="foreground-run") + body["forwardedProps"] = {} + async with httpx.AsyncClient(timeout=60.0) as client: + events = await read_events(client, url, body) + + assert events[0]["type"] == "RUN_STARTED" + assert content_of(events) == CHUNKS + assert events[-1]["type"] == "RUN_FINISHED" + assert all("agnoBackground" not in (event.get("metadata") or {}) for event in events) diff --git a/libs/agno/tests/unit/os/interfaces/test_agui_background.py b/libs/agno/tests/unit/os/interfaces/test_agui_background.py new file mode 100644 index 00000000000..c418f0bf3ff --- /dev/null +++ b/libs/agno/tests/unit/os/interfaces/test_agui_background.py @@ -0,0 +1,1785 @@ +"""Resumable background runs over the AG-UI interface. + +The contract under test: a background AG-UI run keeps executing after the +client goes away, and a reconnecting client receives every AG-UI event exactly +once, in the same order, with the same payloads it would have seen on an +uninterrupted connection. The last section drives the real route, where the +decision to run detached at all is made. +""" + +import asyncio +import contextlib +import json +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, AsyncIterator, Callable, Dict, Iterable, Iterator, List, Optional, Tuple + +import pytest + +pytest.importorskip("ag_ui", reason="ag_ui not installed") + +from ag_ui.core import BaseEvent, EventType +from fastapi.testclient import TestClient + +from agno.agent import Agent, RemoteAgent +from agno.db.sqlite import SqliteDb +from agno.models.base import Model +from agno.models.message import MessageMetrics +from agno.models.response import ModelResponse, ToolExecution +from agno.os.app import AgentOS +from agno.os.event_streams import get_event_stream, set_event_stream +from agno.os.event_streams.base import BaseEventStream +from agno.os.event_streams.in_memory import InMemoryEventStream +from agno.os.interfaces.agui import AGUI +from agno.os.interfaces.agui import background as background_module +from agno.os.interfaces.agui.background import ( + background_cursor, + background_requested, + run_entity_background, + supports_background, +) +from agno.os.managers import EventsBuffer, SSESubscriberManager +from agno.os.utils import format_sse_event_with_index +from agno.reasoning.step import ReasoningStep +from agno.run.agent import ( + ReasoningCompletedEvent, + ReasoningStartedEvent, + ReasoningStepEvent, + RunCompletedEvent, + RunContentEvent, + RunErrorEvent, + RunStartedEvent, + ToolCallCompletedEvent, + ToolCallStartedEvent, +) +from agno.run.base import RunStatus +from agno.run.team import RunCompletedEvent as TeamRunCompletedEvent +from agno.run.team import RunContentEvent as TeamRunContentEvent +from agno.team.remote import RemoteTeam + +# --------------------------------------------------------------------------- +# Fixtures and doubles +# --------------------------------------------------------------------------- + +TERMINAL_STATUSES = (RunStatus.completed, RunStatus.error, RunStatus.cancelled, RunStatus.paused) + +# Long enough that a loaded machine still gets there, short enough that a +# condition which never arrives fails the test instead of hanging the suite. +CONDITION_TIMEOUT_SECONDS = 10.0 +# Real time between polls: a zero sleep hands control back without letting the +# clock move, so a budget spent in polls is no budget at all. +POLL_SECONDS = 0.005 + + +def fresh_event_stream(max_events_per_run: int = 1000) -> InMemoryEventStream: + """Install an empty event stream so a run id can be reused from scratch.""" + stream = InMemoryEventStream( + events_buffer=EventsBuffer(max_events_per_run=max_events_per_run), + subscriber_manager=SSESubscriberManager(), + ) + set_event_stream(stream) + return stream + + +@pytest.fixture(autouse=True) +def isolated_event_stream(): + """Each test gets its own buffer and module state, so run ids never collide.""" + original = get_event_stream() + # The drain task set is deliberately not touched: it holds the only strong + # reference keeping a live background run from being collected. + background_module._STARTED_RUNS.clear() + background_module._STARTING_RUNS.clear() + try: + yield fresh_event_stream() + finally: + set_event_stream(original) + background_module._STARTED_RUNS.clear() + background_module._STARTING_RUNS.clear() + + +class FakeRunInput: + def __init__( + self, + *, + thread_id: str = "thread-1", + run_id: str = "run-1", + forwarded_props: Optional[Dict[str, Any]] = None, + state: Any = None, + messages: Optional[List[Any]] = None, + tools: Optional[List[Any]] = None, + context: Optional[List[Any]] = None, + ): + self.thread_id = thread_id + self.run_id = run_id + self.forwarded_props = forwarded_props + self.state = state + self.messages = messages if messages is not None else [_user_message("hello")] + self.tools = tools or [] + self.context = context or [] + + +@dataclass +class _FakeMessage: + id: str + role: str + content: str + tool_call_id: Optional[str] = None + tool_calls: Optional[List[Any]] = None + name: Optional[str] = None + + +def _user_message(text: str) -> _FakeMessage: + return _FakeMessage(id="m1", role="user", content=text) + + +class ScriptedEntity: + """Stands in for an Agent or Team running detached in the background. + + ``arun(background=True, stream=True)`` mirrors what Agno's background + producer does: it appends each event to the event stream (which owns index + assignment) and yields the SSE string for the originating connection. + + ``final_status`` is the status the producer leaves behind, so a script that + carries no terminal event plus an error status reproduces a producer that + died mid run. + """ + + def __init__( + self, + events: List[Any], + *, + db: Any = "db", + pause_after: Optional[int] = None, + final_status: RunStatus = RunStatus.completed, + ): + self.events = events + self.db = db + self.final_status = final_status + self.arun_kwargs: Dict[str, Any] = {} + self.arun_calls = 0 + self._pause_after = pause_after + self._gate = asyncio.Event() + self.started = asyncio.Event() + + def release(self) -> None: + self._gate.set() + + def arun(self, **kwargs): + self.arun_calls += 1 + self.arun_kwargs = kwargs + run_id = kwargs["run_id"] + + async def _produce() -> AsyncIterator[str]: + stream = get_event_stream() + await stream.register_run(run_id, RunStatus.pending) + await stream.set_run_status(run_id, RunStatus.running) + self.started.set() + try: + for position, event in enumerate(self.events): + if self._pause_after is not None and position == self._pause_after: + await self._gate.wait() + index = await stream.add_event(run_id, event) + yield format_sse_event_with_index(event, event_index=index, run_id=run_id) + finally: + await stream.complete_run(run_id, self.final_status) + + return _produce() + + async def aget_run_output(self, run_id: str, session_id: Optional[str] = None, user_id: Optional[str] = None): + return object() + + +class ForeignRunEntity(ScriptedEntity): + """A run the caller's session does not own.""" + + async def aget_run_output(self, run_id: str, session_id: Optional[str] = None, user_id: Optional[str] = None): + return None + + +class UnreadableRunEntity(ScriptedEntity): + """An entity whose run rows cannot be read: a storage failure, not a denial.""" + + async def aget_run_output(self, run_id: str, session_id: Optional[str] = None, user_id: Optional[str] = None): + raise RuntimeError("run store unavailable") + + +class NoRunOutputEntity(ScriptedEntity): + """An entity that cannot read its own run rows back, so ownership is unverifiable.""" + + aget_run_output = None + + +class UnstartableEntity(ScriptedEntity): + """An entity that refuses to start, after the run has already been registered.""" + + def arun(self, **kwargs): + self.arun_calls += 1 + raise RuntimeError("no capacity for another run") + + +class DyingProducerEntity(ScriptedEntity): + """A producer that raises part way through without ending its own run. + + Unlike ``ScriptedEntity`` it leaves no terminal status behind, so the only + thing that can end the run for an attached client is the drain's own + failure path. + """ + + def arun(self, **kwargs): + self.arun_calls += 1 + self.arun_kwargs = kwargs + run_id = kwargs["run_id"] + + async def _produce() -> AsyncIterator[str]: + stream = get_event_stream() + await stream.register_run(run_id, RunStatus.pending) + await stream.set_run_status(run_id, RunStatus.running) + self.started.set() + for event in self.events: + index = await stream.add_event(run_id, event) + yield format_sse_event_with_index(event, event_index=index, run_id=run_id) + raise RuntimeError("the producer died mid run") + + return _produce() + + +class UnmappedEvent: + """A buffered event whose wire name no Agno event class claims.""" + + event = "AnEventTypeAgnoDoesNotKnow" + + def __init__(self, payload: str): + self.payload = payload + + def to_dict(self) -> Dict[str, Any]: + return {"event": self.event, "payload": self.payload} + + +class _DelegatingEventStream(BaseEventStream): + """Forwards the stream interface to a wrapped one so subclasses override just one. + + reopen_run is left to the base class, which no background run reaches. + """ + + def __init__(self, inner: BaseEventStream): + self._inner = inner + + async def register_run(self, run_id: str, status: RunStatus = RunStatus.pending) -> None: + await self._inner.register_run(run_id, status) + + async def set_run_status(self, run_id: str, status: RunStatus, generation: Optional[int] = None) -> None: + await self._inner.set_run_status(run_id, status, generation) + + async def get_run_status(self, run_id: str) -> Optional[RunStatus]: + return await self._inner.get_run_status(run_id) + + async def complete_run(self, run_id: str, status: RunStatus, generation: Optional[int] = None) -> None: + await self._inner.complete_run(run_id, status, generation) + + async def begin_attempt(self, run_id: str, generation: int) -> None: + await self._inner.begin_attempt(run_id, generation) + + async def cleanup_run(self, run_id: str) -> None: + await self._inner.cleanup_run(run_id) + + async def reset_run_events(self, run_id: str, generation: Optional[int] = None) -> None: + await self._inner.reset_run_events(run_id, generation) + + async def add_event(self, run_id: str, event: Any, generation: Optional[int] = None) -> int: + return await self._inner.add_event(run_id, event, generation) + + async def replay(self, run_id: str, last_event_index: Optional[int] = None) -> List[Tuple[int, Any]]: + return await self._inner.replay(run_id, last_event_index) + + async def get_last_index(self, run_id: str) -> int: + return await self._inner.get_last_index(run_id) + + async def get_event_count(self, run_id: str) -> int: + return await self._inner.get_event_count(run_id) + + def tail(self, run_id: str, last_event_index: Optional[int] = None) -> AsyncIterator[Tuple[int, str]]: + return self._inner.tail(run_id, last_event_index) + + +class SseOnlyEventStream(_DelegatingEventStream): + """A durable-shaped stream: nothing but SSE strings crosses the boundary. + + Redis-backed streams behave this way. ``tail`` is built here from this + class's own ``replay`` plus a status poll instead of borrowing the + in-memory live tail, so a translation that secretly depended on the + in-memory stream handing back event objects fails against it. + """ + + # How long the tail may sit idle on a run that is still going. Reaching it + # raises rather than returning: a tail that closes quietly leaves the + # translation reporting a truncated run as a whole one, which is exactly + # the failure this double exists to catch. + _IDLE_TIMEOUT_SECONDS = CONDITION_TIMEOUT_SECONDS + + def __init__(self, inner: BaseEventStream): + super().__init__(inner) + self.replay_calls = 0 + self.tail_calls = 0 + + async def replay(self, run_id: str, last_event_index: Optional[int] = None) -> List[Tuple[int, Any]]: + self.replay_calls += 1 + return [ + (index, format_sse_event_with_index(event, event_index=index, run_id=run_id)) + for index, event in await self._inner.replay(run_id, last_event_index) + ] + + async def tail(self, run_id: str, last_event_index: Optional[int] = None) -> AsyncIterator[Tuple[int, str]]: + self.tail_calls += 1 + loop = asyncio.get_running_loop() + last = last_event_index if last_event_index is not None else -1 + idle_since = loop.time() + while True: + delivered = False + for index, frame in await self.replay(run_id, last): + delivered = True + last = max(last, index) + yield index, frame + status = await self.get_run_status(run_id) + if status in TERMINAL_STATUSES: + # Whatever landed between that replay and the status read is + # still owed to this tail before it may close. + for index, frame in await self.replay(run_id, last): + last = max(last, index) + yield index, frame + return + if status is None: + # The registration is gone, so no producer will write again. + return + if delivered: + idle_since = loop.time() + elif loop.time() - idle_since > self._IDLE_TIMEOUT_SECONDS: + raise AssertionError( + f"the durable tail for run {run_id} produced nothing for " + f"{self._IDLE_TIMEOUT_SECONDS}s while its status was still {status}" + ) + await asyncio.sleep(POLL_SECONDS) + + +class GatedRegistrationStream(_DelegatingEventStream): + """Holds the first connection inside ``register_run``. + + A connection claims a new run id before it registers it, and only a second + connection landing inside that window takes the wait-for-registration + path. Nothing in that window suspends on its own, so two connections + started together never interleave there: this gate opens the window and + counts the status reads that prove the second one went through it. + """ + + def __init__(self, inner: BaseEventStream): + super().__init__(inner) + self.registering = asyncio.Event() + self.release = asyncio.Event() + self.status_reads = 0 + + async def get_run_status(self, run_id: str) -> Optional[RunStatus]: + self.status_reads += 1 + return await self._inner.get_run_status(run_id) + + async def register_run(self, run_id: str, status: RunStatus = RunStatus.pending) -> None: + self.registering.set() + await self.release.wait() + await self._inner.register_run(run_id, status) + + +class InjectedFrameStream(_DelegatingEventStream): + """Splices extra frames into an otherwise real tail. + + Each injected frame carries the index of the event it precedes, so a frame + the parser skips must not shift the cursor of anything that follows it. + """ + + def __init__(self, inner: BaseEventStream, frames_before: Dict[int, List[str]]): + super().__init__(inner) + self._frames_before = frames_before + + async def tail(self, run_id: str, last_event_index: Optional[int] = None) -> AsyncIterator[Tuple[int, str]]: + async for index, frame in self._inner.tail(run_id, last_event_index): + for injected in self._frames_before.get(index, []): + yield index, injected + yield index, frame + + +class BreakingTailStream(_DelegatingEventStream): + """A tail that fails after handing over some of the run. + + The client keeps what it was given, so the error that ends the connection + has to be positioned relative to that rather than to the start of the run. + """ + + def __init__(self, inner: BaseEventStream, fail_after: int): + super().__init__(inner) + self._fail_after = fail_after + + async def tail(self, run_id: str, last_event_index: Optional[int] = None) -> AsyncIterator[Tuple[int, str]]: + handed_over = 0 + async for index, frame in self._inner.tail(run_id, last_event_index): + yield index, frame + handed_over += 1 + if handed_over >= self._fail_after: + raise RuntimeError("the event stream broke mid run") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def background_props(cursor: Optional[Tuple[int, int]] = None) -> Dict[str, Any]: + props: Dict[str, Any] = {"agnoBackground": {"enabled": True}} + if cursor is not None: + props["agnoBackground"]["lastEventIndex"] = cursor[0] + props["agnoBackground"]["lastSubIndex"] = cursor[1] + return props + + +def cursor_of(event: BaseEvent) -> Tuple[int, int]: + marker = (event.metadata or {}).get("agnoBackground") + assert marker is not None, f"{event.type} carries no background cursor: {event!r}" + return marker["eventIndex"], marker["subIndex"] + + +# Wall-clock fields that never replay equal. ``created_at`` is nested rather +# than top level: a raw event and a run error both carry the source event's +# wire payload, so two separately built scripts differ across a second +# boundary unless it is stripped at every depth. +_VOLATILE_FIELDS = ("timestamp", "created_at") + + +def _without_volatile_fields(value: Any) -> Any: + if isinstance(value, dict): + return {key: _without_volatile_fields(item) for key, item in value.items() if key not in _VOLATILE_FIELDS} + if isinstance(value, list): + return [_without_volatile_fields(item) for item in value] + return value + + +def fingerprint(event: BaseEvent) -> str: + """Payload identity, ignoring wall-clock fields that never replay equal.""" + data = _without_volatile_fields(event.model_dump(by_alias=True, exclude_none=True)) + return json.dumps(data, sort_keys=True, default=str) + + +def fingerprints(events: Iterable[BaseEvent]) -> List[str]: + return [fingerprint(event) for event in events] + + +def assert_spans_are_well_formed(events: List[BaseEvent]) -> None: + """Every content event belongs to a span that was opened and is later closed.""" + open_messages: List[str] = [] + open_tool_calls: List[str] = [] + for event in events: + if event.type == EventType.TEXT_MESSAGE_START: + assert event.message_id not in open_messages, f"a second start for message {event.message_id}" + open_messages.append(event.message_id) + elif event.type == EventType.TEXT_MESSAGE_CONTENT: + assert event.message_id in open_messages, f"content for an unopened message {event.message_id}" + elif event.type == EventType.TEXT_MESSAGE_END: + assert event.message_id in open_messages, f"end for an unopened message {event.message_id}" + open_messages.remove(event.message_id) + elif event.type == EventType.TOOL_CALL_START: + assert event.parent_message_id, "a tool call must parent to a message" + assert event.tool_call_id not in open_tool_calls, f"a second start for call {event.tool_call_id}" + open_tool_calls.append(event.tool_call_id) + elif event.type in (EventType.TOOL_CALL_ARGS, EventType.TOOL_CALL_END): + assert event.tool_call_id in open_tool_calls, f"args or end for an unstarted call {event.tool_call_id}" + if event.type == EventType.TOOL_CALL_END: + open_tool_calls.remove(event.tool_call_id) + assert open_messages == [] + assert open_tool_calls == [] + + +async def wait_until(predicate: Callable[[], bool], timeout: float = CONDITION_TIMEOUT_SECONDS) -> None: + """Wait for a predicate to hold, giving the loop real time to get there.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while not predicate(): + assert loop.time() < deadline, "the awaited condition never became true" + await asyncio.sleep(POLL_SECONDS) + + +async def collect(stream: AsyncIterator[BaseEvent], stop_after: Optional[int] = None) -> List[BaseEvent]: + out: List[BaseEvent] = [] + async for event in stream: + out.append(event) + if stop_after is not None and len(out) >= stop_after: + await stream.aclose() + break + return out + + +TOOL = ToolExecution(tool_call_id="call-1", tool_name="add", tool_args={"a": 1, "b": 2}, result="3") + + +def reasoning_script() -> List[Any]: + return [ + ReasoningStartedEvent(), + ReasoningStepEvent( + content=ReasoningStep(title="weigh it up", reasoning="because of the numbers"), + reasoning_content="because of the numbers", + ), + ReasoningCompletedEvent(), + RunContentEvent(content="four"), + RunCompletedEvent(content="four"), + ] + + +def agent_script() -> List[Any]: + return [ + RunStartedEvent(), + RunContentEvent(content="Hel"), + ToolCallStartedEvent(tool=TOOL), + ToolCallCompletedEvent(tool=TOOL), + RunContentEvent(content="lo"), + RunCompletedEvent(content="Hello", session_state={"counter": 1}), + ] + + +# A buffer this small keeps only the last three events of ``agent_script``, +# which drops the tool call's start and keeps the event completing it. Sized +# one larger the start survives and the trimming exercises nothing. +TRIMMED_TO_AFTER_THE_TOOL_CALL_START = 3 + + +def team_script() -> List[Any]: + return [ + TeamRunContentEvent(content="team says"), + TeamRunCompletedEvent(content="team says hi", session_state={"counter": 2}), + ] + + +def member_completion_script() -> List[Any]: + """A team run whose member completes and whose own terminal never arrives. + + The member's completion is an agent-level ``RunCompleted``, which is what + makes the last completion in the buffer say the run succeeded whatever the + run itself did. + """ + return [TeamRunContentEvent(content="team says"), RunCompletedEvent(content="member done")] + + +def script_with_events_after_the_terminal() -> List[Any]: + """A team run: the member completes, then the leader keeps talking.""" + return [ + TeamRunContentEvent(content="team says"), + RunCompletedEvent(content="member done"), + TeamRunContentEvent(content=" and stops"), + ] + + +def dying_script() -> List[Any]: + """A run that stops mid flight: no terminal event ever reaches the buffer.""" + return [RunStartedEvent(), RunContentEvent(content="Hel")] + + +# --------------------------------------------------------------------------- +# Opt-in parsing +# --------------------------------------------------------------------------- + + +class TestOptIn: + def test_absent_forwarded_props_is_foreground(self): + assert background_requested(FakeRunInput(forwarded_props=None)) is False + + def test_unrelated_forwarded_props_is_foreground(self): + assert background_requested(FakeRunInput(forwarded_props={"user_id": "u1"})) is False + + def test_nested_enabled_flag_opts_in(self): + assert background_requested(FakeRunInput(forwarded_props=background_props())) is True + + def test_bare_true_opts_in(self): + assert background_requested(FakeRunInput(forwarded_props={"agnoBackground": True})) is True + + def test_explicit_false_stays_foreground(self): + assert background_requested(FakeRunInput(forwarded_props={"agnoBackground": {"enabled": False}})) is False + + def test_cursor_is_none_on_a_first_connection(self): + assert background_cursor(FakeRunInput(forwarded_props=background_props())) is None + + def test_cursor_round_trips(self): + assert background_cursor(FakeRunInput(forwarded_props=background_props((4, 1)))) == (4, 1) + + def test_the_very_first_cursor_is_readable(self): + assert background_cursor(FakeRunInput(forwarded_props=background_props((0, 0)))) == (0, 0) + + def test_a_missing_sub_index_means_the_first_event_of_that_group(self): + props = {"agnoBackground": {"enabled": True, "lastEventIndex": 3}} + assert background_cursor(FakeRunInput(forwarded_props=props)) == (3, 0) + + @pytest.mark.parametrize( + "position", + [ + {"lastEventIndex": "4"}, + {"lastEventIndex": 4.5}, + {"lastEventIndex": None}, + # A JSON true is an int subclass in Python, so without a guard it + # would resume from index 1 rather than be rejected. + {"lastEventIndex": True}, + {"lastEventIndex": 4, "lastSubIndex": "1"}, + {"lastEventIndex": 4, "lastSubIndex": False}, + ], + ) + def test_an_unreadable_resume_position_is_rejected(self, position): + props = {"agnoBackground": {"enabled": True, **position}} + with pytest.raises(ValueError, match="Unreadable background resume position"): + background_cursor(FakeRunInput(forwarded_props=props)) + + +# --------------------------------------------------------------------------- +# Which entities may run detached +# --------------------------------------------------------------------------- + + +class TestSupportsBackground: + def test_a_remote_agent_is_refused(self): + """A remote agent's events land in another process, so none of them replay here.""" + assert supports_background(RemoteAgent(base_url="http://localhost:1", agent_id="remote-agent")) is False + + def test_a_remote_team_is_refused(self): + assert supports_background(RemoteTeam(base_url="http://localhost:1", team_id="remote-team")) is False + + def test_an_entity_without_a_database_is_refused(self): + """Detached execution needs somewhere to persist run status.""" + assert supports_background(ScriptedEntity(agent_script(), db=None)) is False + + def test_an_entity_that_cannot_read_run_output_is_refused(self): + """Without a run reader the ownership of a resumed run cannot be checked.""" + assert supports_background(NoRunOutputEntity(agent_script())) is False + + def test_an_in_process_entity_with_a_database_is_accepted(self): + assert supports_background(ScriptedEntity(agent_script())) is True + + +# --------------------------------------------------------------------------- +# Starting a background run +# --------------------------------------------------------------------------- + + +class TestBackgroundStart: + @pytest.mark.asyncio + async def test_detached_execution_is_requested(self): + entity = ScriptedEntity(agent_script()) + run_input = FakeRunInput(forwarded_props=background_props()) + + await collect(run_entity_background(entity, run_input)) + + assert entity.arun_kwargs.get("background") is True + assert entity.arun_kwargs.get("stream") is True + assert entity.arun_kwargs.get("stream_events") is True + assert entity.arun_kwargs.get("run_id") == "run-1" + assert entity.arun_kwargs.get("session_id") == "thread-1" + + @pytest.mark.asyncio + async def test_canonical_agent_sequence(self): + entity = ScriptedEntity(agent_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + types = [event.type for event in events] + assert types[0] == EventType.RUN_STARTED + assert types[-1] == EventType.RUN_FINISHED + assert EventType.TEXT_MESSAGE_START in types + assert EventType.TOOL_CALL_START in types + assert EventType.TOOL_CALL_ARGS in types + assert EventType.TOOL_CALL_END in types + assert EventType.TOOL_CALL_RESULT in types + assert_spans_are_well_formed(events) + + @pytest.mark.asyncio + async def test_every_event_carries_a_monotonic_cursor(self): + entity = ScriptedEntity(agent_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + cursors = [cursor_of(event) for event in events] + assert cursors == sorted(cursors) + assert len(set(cursors)) == len(cursors) + + @pytest.mark.asyncio + async def test_tool_call_result_payload_survives(self): + entity = ScriptedEntity(agent_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + results = [event for event in events if event.type == EventType.TOOL_CALL_RESULT] + assert [(event.tool_call_id, event.content) for event in results] == [("call-1", "3")] + + @pytest.mark.asyncio + async def test_team_background_run(self): + entity = ScriptedEntity(team_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + contents = [event.delta for event in events if event.type == EventType.TEXT_MESSAGE_CONTENT] + assert contents == ["team says"] + assert events[-1].type == EventType.RUN_FINISHED + + @pytest.mark.asyncio + async def test_state_snapshots_bracket_the_run(self): + entity = ScriptedEntity(agent_script()) + run_input = FakeRunInput(forwarded_props=background_props(), state={"counter": 0}) + events = await collect(run_entity_background(entity, run_input)) + + snapshots = [event.snapshot for event in events if event.type == EventType.STATE_SNAPSHOT] + assert snapshots == [{"counter": 0}, {"counter": 1}] + + @pytest.mark.asyncio + async def test_state_is_reported_the_way_a_foreground_run_reports_it(self): + """A background run invents no state for a request that sent none. + + A request that sends state is bracketed by an opening and a closing + snapshot; one that sends none gets neither, which is what the same run + would do streaming inline. A client that keeps sending state the same + way therefore sees the same event positions on every connection. + """ + with_state = await collect( + run_entity_background( + ScriptedEntity(agent_script()), + FakeRunInput(forwarded_props=background_props(), state={"counter": 0}), + ) + ) + + fresh_event_stream() + resumed_with_state = await collect( + run_entity_background( + ScriptedEntity(agent_script()), + FakeRunInput(forwarded_props=background_props(), state={"counter": 0}), + ) + ) + + fresh_event_stream() + without_state = await collect( + run_entity_background(ScriptedEntity(agent_script()), FakeRunInput(forwarded_props=background_props())) + ) + + assert [event.snapshot for event in with_state if event.type == EventType.STATE_SNAPSHOT] == [ + {"counter": 0}, + {"counter": 1}, + ] + assert [event.type for event in without_state if event.type == EventType.STATE_SNAPSHOT] == [] + assert fingerprints(with_state) == fingerprints(resumed_with_state) + assert with_state[-1].type == EventType.RUN_FINISHED + assert without_state[-1].type == EventType.RUN_FINISHED + + @pytest.mark.asyncio + async def test_error_terminal_is_preserved(self): + entity = ScriptedEntity([RunStartedEvent(), RunErrorEvent(content="boom")]) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert events[-1].type == EventType.RUN_ERROR + assert events[-1].message == "boom" + + +# --------------------------------------------------------------------------- +# Translating buffered frames +# --------------------------------------------------------------------------- + + +UNREADABLE_FRAMES = [ + "event: RunContent\n\n", + "event: RunContent\ndata: {not json at all}\n\n", + 'event: RunContent\ndata: {"content": "names no event"}\n\n', +] + +# U+2028 is a line terminator to str.splitlines but a legal character inside a +# JSON string, so a frame carrying one must not be split on it. +LINE_SEPARATOR_TEXT = "before\u2028after" + + +class TestFrameTranslation: + @pytest.mark.asyncio + async def test_a_reasoning_step_survives_the_trip_through_the_buffer(self): + """A reasoning step is serialized to a plain mapping on its way into the + buffer, and reading it back as one is what keeps the run alive.""" + entity = ScriptedEntity(reasoning_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + contents = [event.delta for event in events if event.type == EventType.REASONING_MESSAGE_CONTENT] + assert contents == ["## Step 1: weigh it up\nbecause of the numbers\n\n"] + assert events[-1].type == EventType.RUN_FINISHED + assert_spans_are_well_formed(events) + + @pytest.mark.asyncio + async def test_an_unmapped_event_arrives_as_a_raw_event(self): + """An event with no Agno class is forwarded raw, as the foreground path does. + + Dropping it would lose a custom event on every reconnect while the same + run had delivered it on a first connection. The payload is the event's + own, down to the field: the index the buffer stamps onto the frame is + the buffer's bookkeeping and belongs in the resume marker rather than + inside an event a client is asked to interpret. + """ + entity = ScriptedEntity([UnmappedEvent("keep me"), RunCompletedEvent(content="done")]) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.RAW, + EventType.RUN_FINISHED, + ] + raw = events[1] + assert raw.source == "agno" + assert raw.event == {"event": UnmappedEvent.event, "payload": "keep me", "run_id": "run-1"} + assert cursor_of(raw) == (0, 0) + + @pytest.mark.asyncio + async def test_unreadable_frames_are_skipped_without_ending_the_stream(self): + """A frame with no data, bad JSON, or no event name is dropped, not fatal.""" + baseline = await collect( + run_entity_background(ScriptedEntity(agent_script()), FakeRunInput(forwarded_props=background_props())) + ) + + inner = fresh_event_stream() + set_event_stream(InjectedFrameStream(inner, {index: UNREADABLE_FRAMES for index in range(len(agent_script()))})) + events = await collect( + run_entity_background(ScriptedEntity(agent_script()), FakeRunInput(forwarded_props=background_props())) + ) + + assert fingerprints(events) == fingerprints(baseline) + assert [cursor_of(event) for event in events] == [cursor_of(event) for event in baseline] + + @pytest.mark.asyncio + async def test_a_unicode_line_separator_inside_a_frame_still_decodes(self): + entity = ScriptedEntity( + [RunContentEvent(content=LINE_SEPARATOR_TEXT), RunCompletedEvent(content=LINE_SEPARATOR_TEXT)] + ) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + deltas = [event.delta for event in events if event.type == EventType.TEXT_MESSAGE_CONTENT] + assert deltas == [LINE_SEPARATOR_TEXT] + assert events[-1].type == EventType.RUN_FINISHED + + +# --------------------------------------------------------------------------- +# Terminal events +# --------------------------------------------------------------------------- + + +class TestTerminalEvents: + @pytest.mark.asyncio + @pytest.mark.parametrize("status", [RunStatus.error, RunStatus.cancelled, RunStatus.paused]) + async def test_a_tail_that_ends_badly_reports_the_runs_own_status(self, status): + """No terminal event in the buffer: the run's status decides how it ended. + + Synthesizing a RUN_FINISHED here would tell the client a run succeeded + that in fact died, was cancelled, or is sitting paused with nothing in + the buffer to say so. + """ + entity = ScriptedEntity(dying_script(), final_status=status) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.RAW, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.RUN_ERROR, + ] + assert events[-1].message == f"Run ended without a result, last known status {status.value.lower()}" + assert_spans_are_well_formed(events) + + @pytest.mark.asyncio + @pytest.mark.parametrize("status", [RunStatus.error, RunStatus.cancelled]) + async def test_a_held_member_completion_does_not_report_a_failed_run_as_a_success(self, status): + """A completion belonging to a member may not end the run happily. + + A team run buffers each member's completion before its own, so the + last completion the buffer holds can be a member's while the run + itself died or was cancelled. The run's status decides how the client + is told it ended. + """ + entity = ScriptedEntity(member_completion_script(), final_status=status) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.RUN_ERROR, + ] + assert events[-1].message == f"Run ended without a result, last known status {status.value.lower()}" + assert [event.delta for event in events if event.type == EventType.TEXT_MESSAGE_CONTENT] == ["team says"] + assert_spans_are_well_formed(events) + + @pytest.mark.asyncio + async def test_a_held_member_completion_ends_the_run_when_the_status_agrees(self): + """The status overrides a held completion only when it contradicts one.""" + entity = ScriptedEntity(member_completion_script(), final_status=RunStatus.completed) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.RUN_FINISHED, + ] + assert_spans_are_well_formed(events) + + @pytest.mark.asyncio + async def test_the_terminal_group_lands_after_everything_the_buffer_holds(self): + """A terminal with events buffered behind it is still emitted last. + + Its own index is already behind theirs, and keeping it would hand two + AG-UI events one address, so the whole terminal group moves past the + highest index the tail saw. + """ + entity = ScriptedEntity(script_with_events_after_the_terminal()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.RUN_FINISHED, + ] + assert [event.delta for event in events if event.type == EventType.TEXT_MESSAGE_CONTENT] == [ + "team says", + " and stops", + ] + terminal_group = [cursor_of(event) for event in events[-2:]] + assert terminal_group == [(3, 0), (3, 1)] + assert min(terminal_group) > max(cursor_of(event) for event in events[:-2]) + assert_spans_are_well_formed(events) + + @pytest.mark.asyncio + async def test_a_producer_that_dies_mid_run_still_ends_the_clients_stream(self): + """A run whose producer raises is marked errored, so nobody waits on it. + + The buffer holds no terminal event and the producer is gone, so + without the drain reporting the failure the client would sit through + the stream's idle recheck and then be told the run succeeded. + """ + entity = DyingProducerEntity(dying_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.RAW, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.RUN_ERROR, + ] + assert events[-1].message == f"Run ended without a result, last known status {RunStatus.error.value.lower()}" + assert await get_event_stream().get_run_status("run-1") == RunStatus.error + assert_spans_are_well_formed(events) + + @pytest.mark.asyncio + async def test_a_mid_stream_failure_ends_above_everything_already_delivered(self): + """A stream that breaks part way through still terminates, past what it sent. + + The client keeps what it received, so the error ending the connection + has to sit above all of it: one that did not would be dropped by a + client filtering by cursor, which would then reconnect into it. + """ + set_event_stream(BreakingTailStream(get_event_stream(), fail_after=2)) + entity = ScriptedEntity(agent_script()) + + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [ + EventType.RUN_STARTED, + EventType.RAW, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.RUN_ERROR, + ] + assert events[-1].message == "the event stream broke mid run" + delivered = [cursor_of(event) for event in events[:-1]] + assert delivered == [(-1, 0), (0, 0), (1, 0), (1, 1)] + assert cursor_of(events[-1]) > max(delivered) + + @pytest.mark.asyncio + async def test_a_synthesized_terminal_is_never_filtered_out_by_a_cursor(self): + """A client echoing the terminal cursor still receives a terminal event. + + Filtering the run's end away would leave that client holding a + connection that will never say anything again, so the terminal is + re-issued past whatever position the client sent. Only the terminal: + the span-closing events beside it in that group were already + delivered, and re-issuing one would end a message this leg never + opened. + """ + entity = ScriptedEntity(dying_script(), final_status=RunStatus.error) + baseline = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + assert baseline[-2].type == EventType.TEXT_MESSAGE_END + held = cursor_of(baseline[-1]) + + resumed = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props(held)))) + + assert [event.type for event in resumed] == [EventType.RUN_ERROR] + assert resumed[-1].message == baseline[-1].message + assert cursor_of(resumed[-1]) > held + assert_spans_are_well_formed(resumed) + assert entity.arun_calls == 1 + + @pytest.mark.asyncio + async def test_a_run_that_cannot_be_started_still_terminates_the_client(self): + """Starting failed after registration, so the client is owed a terminal event.""" + entity = UnstartableEntity(agent_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert events[-1].message == "no capacity for another run" + # The registration is gone afterwards, so the same run id can be tried + # again rather than answering "already ended" for the rest of the + # process. A tail that attached first was ended before it was dropped. + assert await get_event_stream().get_run_status("run-1") is None + + @pytest.mark.asyncio + async def test_a_run_id_whose_start_failed_can_be_started_again(self): + """A failed start must not poison the run id for every later attempt.""" + await collect( + run_entity_background(UnstartableEntity(agent_script()), FakeRunInput(forwarded_props=background_props())) + ) + + retried = ScriptedEntity(agent_script()) + events = await collect(run_entity_background(retried, FakeRunInput(forwarded_props=background_props()))) + + assert retried.arun_calls == 1 + assert events[0].type == EventType.RUN_STARTED + assert events[-1].type == EventType.RUN_FINISHED + + +# --------------------------------------------------------------------------- +# Disconnect and reconnect +# --------------------------------------------------------------------------- + + +class TestReconnect: + @pytest.mark.asyncio + async def test_reconnect_delivers_each_event_exactly_once_in_order(self): + baseline_entity = ScriptedEntity(agent_script()) + baseline = await collect( + run_entity_background(baseline_entity, FakeRunInput(forwarded_props=background_props())) + ) + + fresh_event_stream() + entity = ScriptedEntity(agent_script()) + first_leg = await collect( + run_entity_background(entity, FakeRunInput(forwarded_props=background_props())), + stop_after=4, + ) + resumed = await collect( + run_entity_background(entity, FakeRunInput(forwarded_props=background_props(cursor_of(first_leg[-1])))) + ) + + assert fingerprints(first_leg + resumed) == fingerprints(baseline) + + @pytest.mark.asyncio + async def test_reconnect_does_not_restart_the_run(self): + entity = ScriptedEntity(agent_script()) + first_leg = await collect( + run_entity_background(entity, FakeRunInput(forwarded_props=background_props())), stop_after=3 + ) + await collect( + run_entity_background(entity, FakeRunInput(forwarded_props=background_props(cursor_of(first_leg[-1])))) + ) + + assert entity.arun_calls == 1 + + @pytest.mark.asyncio + async def test_run_continues_while_no_client_is_attached(self): + entity = ScriptedEntity(agent_script(), pause_after=3) + stream = run_entity_background(entity, FakeRunInput(forwarded_props=background_props())) + first_leg = await collect(stream, stop_after=2) + + entity.release() + resumed = await collect( + run_entity_background(entity, FakeRunInput(forwarded_props=background_props(cursor_of(first_leg[-1])))) + ) + + assert entity.arun_calls == 1 + assert resumed[-1].type == EventType.RUN_FINISHED + + @pytest.mark.asyncio + async def test_reconnect_after_completion_replays_the_remainder(self): + entity = ScriptedEntity(agent_script()) + baseline = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + resumed = await collect( + run_entity_background(entity, FakeRunInput(forwarded_props=background_props(cursor_of(baseline[2])))) + ) + + assert fingerprints(resumed) == fingerprints(baseline[3:]) + + @pytest.mark.asyncio + async def test_replay_reproduces_identical_message_ids(self): + entity = ScriptedEntity(agent_script()) + baseline = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + full_replay = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert fingerprints(full_replay) == fingerprints(baseline) + + @pytest.mark.asyncio + async def test_durable_stream_replay_path(self): + """A stream that only ever hands back SSE strings replays identically.""" + durable = SseOnlyEventStream(get_event_stream()) + set_event_stream(durable) + entity = ScriptedEntity(agent_script()) + baseline = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + resumed = await collect( + run_entity_background(entity, FakeRunInput(forwarded_props=background_props(cursor_of(baseline[1])))) + ) + + # Asserted of each leg rather than only of the pair: two legs that + # agree can still both be truncated, which is what a tail that gives + # up quietly would produce. + assert baseline[-1].type == EventType.RUN_FINISHED + assert resumed[-1].type == EventType.RUN_FINISHED + assert fingerprints(resumed) == fingerprints(baseline[2:]) + # Both legs really came through this stream's own replay-and-poll tail, + # so the assertions above say something about the durable path. + assert durable.tail_calls == 2 + assert durable.replay_calls > 0 + + @pytest.mark.asyncio + async def test_a_second_connection_naming_a_started_run_does_not_start_it_again(self): + """A second client naming a run id this process already started attaches to it. + + Both connections are launched together, but nothing between the + started-probe and the registration suspends, so the first reaches the + buffer before the second is scheduled and the second finds a started + run rather than a starting one. The window where it finds a starting + one is held open deliberately in the test below. + """ + baseline = await collect( + run_entity_background(ScriptedEntity(agent_script()), FakeRunInput(forwarded_props=background_props())) + ) + + fresh_event_stream() + entity = ScriptedEntity(agent_script()) + first, second = await asyncio.gather( + collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))), + collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))), + ) + + assert entity.arun_calls == 1 + assert fingerprints(first) == fingerprints(baseline) + assert fingerprints(second) == fingerprints(baseline) + + @pytest.mark.asyncio + async def test_a_connection_arriving_while_the_run_registers_waits_and_attaches(self): + """A second connection inside the start window waits rather than being turned away. + + The first connection is held inside ``register_run``, so the second + arrives while the run id is claimed but not yet registered: the state + the wait-for-registration guard exists for. Attaching is what that + connection wanted anyway, so it must end up with the same stream, not + a refusal and not a second producer. + """ + gated = GatedRegistrationStream(get_event_stream()) + set_event_stream(gated) + entity = ScriptedEntity(agent_script()) + + first = asyncio.create_task( + collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + ) + await asyncio.wait_for(gated.registering.wait(), timeout=CONDITION_TIMEOUT_SECONDS) + assert "run-1" in background_module._STARTING_RUNS + + reads_before = gated.status_reads + second = asyncio.create_task( + collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + ) + # The second connection's own started-probe, and then a poll that can + # only come from the wait: the guard is reached, not merely available. + await wait_until(lambda: gated.status_reads >= reads_before + 2) + gated.release.set() + + first_events, second_events = await asyncio.gather(first, second) + + assert entity.arun_calls == 1 + assert first_events[0].type == EventType.RUN_STARTED + assert first_events[-1].type == EventType.RUN_FINISHED + assert fingerprints(second_events) == fingerprints(first_events) + assert_spans_are_well_formed(second_events) + + @pytest.mark.asyncio + async def test_a_connection_whose_run_never_registers_is_refused(self): + """A wait that runs out ends the client instead of attaching it to nothing. + + The first connection is held inside ``register_run`` for good, so the + second one spends its whole budget on a run that never appears. + Falling through would leave it tailing a run id the event stream has + never heard of, silent until something else times it out. + """ + gated = GatedRegistrationStream(get_event_stream()) + set_event_stream(gated) + entity = ScriptedEntity(agent_script()) + + held = asyncio.create_task( + collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + ) + await asyncio.wait_for(gated.registering.wait(), timeout=CONDITION_TIMEOUT_SECONDS) + + refused = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert [event.type for event in refused] == [EventType.RUN_ERROR] + assert refused[0].message == "Run run-1 did not start" + assert cursor_of(refused[0]) == (-1, 0) + assert entity.arun_calls == 0 + + held.cancel() + gated.release.set() + with contextlib.suppress(BaseException): + await held + + @pytest.mark.asyncio + async def test_attaching_from_another_session_without_a_resume_position_is_refused(self): + """A started run is not readable by a second request that merely names its id. + + The run id is client-supplied, so a caller carrying no resume position + at all still has to be checked against the session before the buffer + is handed over. + """ + owner = ScriptedEntity(agent_script()) + owned = await collect(run_entity_background(owner, FakeRunInput(forwarded_props=background_props()))) + assert owned[-1].type == EventType.RUN_FINISHED + + intruder = ForeignRunEntity(agent_script()) + events = await collect( + run_entity_background(intruder, FakeRunInput(thread_id="other-thread", forwarded_props=background_props())) + ) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert events[0].message == "Run run-1 not found in this session" + assert cursor_of(events[0]) == (-1, 0) + assert intruder.arun_calls == 0 + + @pytest.mark.asyncio + async def test_a_trimmed_buffer_still_yields_a_valid_sequence(self): + """When the buffer has dropped the start of a run, what is left is still coherent. + + The buffer is finite, so a long enough run replays from wherever it now + begins. Here the trimmed events include the start of the tool call, + while the event completing that call survives: closing a span this + connection never opened would be invalid rather than merely + incomplete, so the call is dropped whole, result and all. + """ + fresh_event_stream(max_events_per_run=TRIMMED_TO_AFTER_THE_TOOL_CALL_START) + entity = ScriptedEntity(agent_script()) + await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + replayed = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + assert entity.arun_calls == 1 + assert [event.type for event in replayed] == [ + EventType.RUN_STARTED, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.RUN_FINISHED, + ] + assert [event.delta for event in replayed if event.type == EventType.TEXT_MESSAGE_CONTENT] == ["lo"] + assert TOOL.tool_call_id not in [getattr(event, "tool_call_id", None) for event in replayed] + buffered_indices = [cursor_of(event)[0] for event in replayed if cursor_of(event)[0] >= 0] + assert min(buffered_indices) == 4 + assert_spans_are_well_formed(replayed) + + @pytest.mark.asyncio + @pytest.mark.parametrize("held", [(0, 0), (1, 0), (2, 0)]) + async def test_resuming_against_a_trimmed_buffer_is_refused(self, held): + """A trimmed buffer turns away a resume from any position at all. + + The translation is rebuilt from whatever the buffer still holds, so a + message or a tool call whose source event is gone is opened again under + a new identifier while the client's own copy of it is never closed. + Held positions inside a group are covered too: one buffered event can + produce several AG-UI events, and a client can hold the first of a + group and still be owed its siblings. + """ + fresh_event_stream(max_events_per_run=TRIMMED_TO_AFTER_THE_TOOL_CALL_START) + entity = ScriptedEntity(agent_script()) + await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + refused = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props(held)))) + + assert [event.type for event in refused] == [EventType.RUN_ERROR] + assert refused[0].message == ( + "Run run-1 cannot be resumed from that position; the events after it are no longer buffered" + ) + # Stamped past the position the client sent, so a client filtering by + # cursor cannot drop the refusal and reconnect into it forever. + assert cursor_of(refused[0]) == (held[0] + 1, 0) + assert entity.arun_calls == 1 + + @pytest.mark.asyncio + async def test_a_resume_against_a_trimmed_buffer_is_refused(self): + """A trimmed buffer cannot answer a resume, wherever the client stands. + + The translation is rebuilt from whatever the buffer still holds, so a + message whose opening was trimmed is opened again under a new id while + the client's own copy of it is never closed. That is worse than saying + the run cannot be resumed. + """ + fresh_event_stream(max_events_per_run=TRIMMED_TO_AFTER_THE_TOOL_CALL_START) + entity = ScriptedEntity(agent_script()) + # The first connection tails live, so it sees every event before the + # buffer drops any. Trimming only shows on a connection that replays. + await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + replayed = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + buffered = [cursor_of(event)[0] for event in replayed if cursor_of(event)[0] >= 0] + assert min(buffered) > 0, "the front of the run must have been trimmed for this to test anything" + held = cursor_of(replayed[1]) + + resumed = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props(held)))) + + assert [event.type for event in resumed] == [EventType.RUN_ERROR] + assert "no longer buffered" in resumed[0].message + assert cursor_of(resumed[0]) > held + assert all(cursor_of(event) > held for event in resumed) + assert entity.arun_calls == 1 + + +# --------------------------------------------------------------------------- +# Refusals +# --------------------------------------------------------------------------- + + +class TestRefusals: + @pytest.mark.asyncio + async def test_attaching_to_a_run_outside_the_session_is_refused(self): + entity = ScriptedEntity(agent_script()) + await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + intruder = ForeignRunEntity(agent_script()) + events = await collect( + run_entity_background( + intruder, + FakeRunInput(thread_id="other-thread", forwarded_props=background_props((0, 0))), + ) + ) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert events[-1].message == "Run run-1 not found in this session" + assert intruder.arun_calls == 0 + + @pytest.mark.asyncio + async def test_an_ownership_check_that_raises_does_not_read_as_a_denial(self): + """A storage failure must not send a client off to restart a run that is still live. + + The process did start this run, but its record of that is keyed by the + entity, session and user the run belongs to, and the reconnect matches + none of them. So the answer has to come from storage, which is what + fails here. + """ + entity = ScriptedEntity(agent_script()) + await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props()))) + + unreadable = UnreadableRunEntity(agent_script()) + events = await collect( + run_entity_background( + unreadable, FakeRunInput(thread_id="other-thread", forwarded_props=background_props((0, 0))) + ) + ) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert events[-1].message == "Could not verify run run-1; try again" + assert events[-1].message != "Run run-1 not found in this session" + assert unreadable.arun_calls == 0 + + @pytest.mark.asyncio + async def test_attaching_when_ownership_cannot_be_checked_at_all_is_refused(self): + """An entity with no way to read its run rows may not attach to a buffered run. + + The ownership check is what stops a caller naming any run id and + reading another session's events back out of the buffer. An entity + that cannot answer the question leaves no check to pass, so the answer + is no rather than a shrug. + """ + owner = ScriptedEntity(agent_script()) + owned = await collect(run_entity_background(owner, FakeRunInput(forwarded_props=background_props()))) + assert owned[-1].type == EventType.RUN_FINISHED + + blind = NoRunOutputEntity(agent_script()) + events = await collect( + run_entity_background(blind, FakeRunInput(forwarded_props=background_props(cursor_of(owned[1])))) + ) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert events[0].message == "Run run-1 not found in this session" + assert blind.arun_calls == 0 + + @pytest.mark.asyncio + async def test_an_unreadable_resume_position_is_refused_rather_than_replayed(self): + """A resume position that cannot be read must not restart the whole run.""" + entity = ScriptedEntity(agent_script()) + props = {"agnoBackground": {"enabled": True, "lastEventIndex": "4"}} + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=props))) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert "Unreadable background resume position" in events[-1].message + assert entity.arun_calls == 0 + + @pytest.mark.asyncio + async def test_a_json_true_resume_position_is_refused(self): + """A JSON true is an int in Python, so it would otherwise resume from index 1.""" + entity = ScriptedEntity(agent_script()) + props = {"agnoBackground": {"enabled": True, "lastEventIndex": True}} + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=props))) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert "Unreadable background resume position" in events[-1].message + assert entity.arun_calls == 0 + + @pytest.mark.asyncio + async def test_attaching_when_the_stream_state_is_gone_is_refused(self): + """The run row outlived its events, so there is nothing left to resume from. + + Saying so beats stalling on an empty stream and then reporting success. + """ + entity = ScriptedEntity(agent_script()) + events = await collect(run_entity_background(entity, FakeRunInput(forwarded_props=background_props((3, 0))))) + + assert [event.type for event in events] == [EventType.RUN_ERROR] + assert events[-1].message == "Run run-1 is no longer available for replay" + assert entity.arun_calls == 0 + + +# --------------------------------------------------------------------------- +# The route +# --------------------------------------------------------------------------- + +ANSWER_CHUNKS = ["Hel", "lo"] + +# The route quotes the entity's id when it declines to run it detached. +DETACHABLE_AGENT_ID = "background-agent" +INLINE_ONLY_AGENT_ID = "no-db-agent" +DETACHABLE_PATH = "/detachable/agui" +INLINE_ONLY_PATH = "/inline-only/agui" + + +class EchoModel(Model): + """An offline model that answers with a fixed set of chunks. + + ``invocations`` is what tells a request the route refused from one it + quietly ran again: a refusal must reach no model at all. + """ + + def __init__(self) -> None: + super().__init__(id="echo-test-model", name="echo-test-model", provider="test") + self.instructions = None + self.invocations = 0 + + def __deepcopy__(self, memo: dict) -> "EchoModel": + # Shared rather than copied, so a run reached through the route still + # counts against the instance the test holds. + return self + + def _response(self, text: str) -> ModelResponse: + return ModelResponse(content=text, role="assistant", response_usage=MessageMetrics()) + + def get_instructions_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + def get_system_message_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + async def aget_instructions_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + async def aget_system_message_for_model(self, *args: Any, **kwargs: Any) -> None: + return None + + def parse_args(self, *args: Any, **kwargs: Any) -> Dict[str, Any]: + return {} + + def invoke(self, *args: Any, **kwargs: Any) -> ModelResponse: + self.invocations += 1 + return self._response("".join(ANSWER_CHUNKS)) + + async def ainvoke(self, *args: Any, **kwargs: Any) -> ModelResponse: + self.invocations += 1 + return self._response("".join(ANSWER_CHUNKS)) + + def invoke_stream(self, *args: Any, **kwargs: Any) -> Iterator[ModelResponse]: + self.invocations += 1 + for chunk in ANSWER_CHUNKS: + yield self._response(chunk) + + async def ainvoke_stream(self, *args: Any, **kwargs: Any) -> AsyncIterator[ModelResponse]: + self.invocations += 1 + for chunk in ANSWER_CHUNKS: + yield self._response(chunk) + + def _parse_provider_response(self, response: Any, **kwargs: Any) -> ModelResponse: + return self._response("") + + def _parse_provider_response_delta(self, response: Any) -> ModelResponse: + return self._response("") + + +@dataclass +class AguiRoutes: + """A live AG-UI route pair: one entity that can run detached, one that cannot.""" + + client: TestClient + detachable_model: EchoModel + inline_only_model: EchoModel + + +@pytest.fixture +def agui_routes() -> Iterator[AguiRoutes]: + with tempfile.TemporaryDirectory() as directory: + db = SqliteDb(db_file=str(Path(directory) / "agui-route.db")) + detachable_model = EchoModel() + inline_only_model = EchoModel() + # A database and an in-process entity are what background execution + # needs, so the second agent differs from the first only in having no + # database to persist a detached run's status to. + detachable = Agent(id=DETACHABLE_AGENT_ID, name=DETACHABLE_AGENT_ID, model=detachable_model, db=db) + inline_only = Agent(id=INLINE_ONLY_AGENT_ID, name=INLINE_ONLY_AGENT_ID, model=inline_only_model) + + agent_os = AgentOS( + agents=[detachable, inline_only], + interfaces=[ + AGUI(agent=detachable, prefix="/detachable"), + AGUI(agent=inline_only, prefix="/inline-only"), + ], + ) + with TestClient(agent_os.get_app()) as client: + yield AguiRoutes(client=client, detachable_model=detachable_model, inline_only_model=inline_only_model) + + +def route_body( + *, + thread_id: str, + run_id: str, + background: Optional[Dict[str, Any]] = None, + tool_result_for: Optional[str] = None, +) -> Dict[str, Any]: + messages: List[Dict[str, Any]] = [{"id": "m1", "role": "user", "content": "say hello"}] + if tool_result_for is not None: + # A trailing tool message is what marks a request as the continuation + # of a run that paused for the client to execute a tool. + messages.append({"id": "t1", "role": "tool", "content": "{}", "toolCallId": tool_result_for}) + return { + "threadId": thread_id, + "runId": run_id, + "state": None, + "messages": messages, + "tools": [], + "context": [], + "forwardedProps": {} if background is None else {"agnoBackground": background}, + } + + +def post_events(routes: AguiRoutes, path: str, body: Dict[str, Any]) -> List[Dict[str, Any]]: + """The AG-UI events one request produced, in the order they were written.""" + response = routes.client.post(path, json=body) + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + return [json.loads(line[5:].strip()) for line in response.text.splitlines() if line.startswith("data:")] + + +def resume_marker(event: Dict[str, Any]) -> Optional[Dict[str, Any]]: + return (event.get("metadata") or {}).get("agnoBackground") + + +def route_cursor(event: Dict[str, Any]) -> Tuple[int, int]: + marker = resume_marker(event) + assert marker is not None, f"{event['type']} carries no resume marker: {event!r}" + return marker["eventIndex"], marker["subIndex"] + + +def identity(event: Dict[str, Any]) -> Dict[str, Any]: + """An event's payload, minus the protocol's optional wall-clock timestamp.""" + return {key: value for key, value in event.items() if key != "timestamp"} + + +class TestRoute: + def test_a_background_run_streams_end_to_end(self, agui_routes: AguiRoutes): + """The route runs the entity detached and serves it back out of the buffer. + + Every event carries a resume marker, which is how a client learns this + connection can be dropped and picked up again, and the answer arrives + in one piece with the model reached exactly once. + """ + events = post_events( + agui_routes, + DETACHABLE_PATH, + route_body(thread_id="detached", run_id="detached-run", background={"enabled": True}), + ) + + # Raw events pass through untranslated and their number is the model's + # business, not the route's, so the answer is read around them. + assert [event["type"] for event in events if event["type"] != "RAW"] == [ + "RUN_STARTED", + "TEXT_MESSAGE_START", + "TEXT_MESSAGE_CONTENT", + "TEXT_MESSAGE_CONTENT", + "TEXT_MESSAGE_END", + "RUN_FINISHED", + ] + assert [event["delta"] for event in events if event["type"] == "TEXT_MESSAGE_CONTENT"] == ANSWER_CHUNKS + assert agui_routes.detachable_model.invocations == 1 + + cursors = [route_cursor(event) for event in events] + assert cursors == sorted(cursors) + assert len(set(cursors)) == len(cursors) + assert cursors[0] == (-1, 0) + assert cursors[-1] == max(cursors) + + def test_a_background_run_is_picked_up_from_a_cursor_rather_than_run_again(self, agui_routes: AguiRoutes): + """A second request through the route resumes rather than answering twice. + + The remainder has to match the first connection's own tail event for + event, ids included, or a client stitching the two legs together sees + a message it never opened. + """ + body = route_body(thread_id="resumed", run_id="resumed-run", background={"enabled": True}) + whole = post_events(agui_routes, DETACHABLE_PATH, body) + assert whole[-1]["type"] == "RUN_FINISHED" + + pivot = next(index for index, event in enumerate(whole) if event["type"] == "TEXT_MESSAGE_CONTENT") + event_index, sub_index = route_cursor(whole[pivot]) + rest = post_events( + agui_routes, + DETACHABLE_PATH, + route_body( + thread_id="resumed", + run_id="resumed-run", + background={"enabled": True, "lastEventIndex": event_index, "lastSubIndex": sub_index}, + ), + ) + + assert [identity(event) for event in rest] == [identity(event) for event in whole[pivot + 1 :]] + assert agui_routes.detachable_model.invocations == 1 + + def test_a_resume_position_with_background_disabled_is_refused(self, agui_routes: AguiRoutes): + """A resume position is a claim to be continuing a run, switch or no switch. + + Falling through to a foreground run would answer the whole question a + second time, which is the one thing the client asking to resume was + trying to avoid. + """ + events = post_events( + agui_routes, + DETACHABLE_PATH, + route_body( + thread_id="disabled", + run_id="disabled-run", + background={"enabled": False, "lastEventIndex": 3, "lastSubIndex": 1}, + ), + ) + + assert [event["type"] for event in events] == ["RUN_ERROR"] + assert events[0]["message"] == "A resume position was sent with background execution disabled" + assert resume_marker(events[0]) == {"eventIndex": 4, "subIndex": 0} + assert agui_routes.detachable_model.invocations == 0 + + def test_a_resume_position_to_an_entity_that_cannot_run_detached_is_refused(self, agui_routes: AguiRoutes): + """Declining background execution must not turn a resume into a rerun.""" + events = post_events( + agui_routes, + INLINE_ONLY_PATH, + route_body( + thread_id="unsupported", + run_id="unsupported-run", + background={"enabled": True, "lastEventIndex": 3, "lastSubIndex": 1}, + ), + ) + + assert [event["type"] for event in events] == ["RUN_ERROR"] + assert events[0]["message"] == ( + f"Background execution is unavailable for '{INLINE_ONLY_AGENT_ID}': it needs a database, " + "an agent or team that runs in this process, and a readable run history" + ) + assert resume_marker(events[0]) == {"eventIndex": 4, "subIndex": 0} + assert agui_routes.inline_only_model.invocations == 0 + + def test_a_first_connection_to_such_an_entity_runs_in_the_foreground(self, agui_routes: AguiRoutes): + """With nothing to resume, the request is served inline rather than refused. + + The run then carries no resume marker at all, which is how a client + learns this connection is the only one it gets. + """ + events = post_events( + agui_routes, + INLINE_ONLY_PATH, + route_body(thread_id="inline", run_id="inline-run", background={"enabled": True}), + ) + + # Raw events pass through untranslated and their number is the model's + # business, not the route's, so the answer is read around them. + assert [event["type"] for event in events if event["type"] != "RAW"] == [ + "RUN_STARTED", + "TEXT_MESSAGE_START", + "TEXT_MESSAGE_CONTENT", + "TEXT_MESSAGE_CONTENT", + "TEXT_MESSAGE_END", + "RUN_FINISHED", + ] + assert [event["delta"] for event in events if event["type"] == "TEXT_MESSAGE_CONTENT"] == ANSWER_CHUNKS + assert all(resume_marker(event) is None for event in events) + assert agui_routes.inline_only_model.invocations == 1 + + def test_continuing_a_paused_run_while_echoing_a_resume_position_is_not_refused(self, agui_routes: AguiRoutes): + """A continuation is a new leg, so the position it still echoes does not apply. + + Deciding the continuation first is what keeps a client that + answers a tool call, and happens to still be sending the last cursor + it saw, from being turned away instead of continued. The session here + holds no paused run, so the continuation fails on its way through the + foreground path, which is the point: it got there. + """ + events = post_events( + agui_routes, + DETACHABLE_PATH, + route_body( + thread_id="continuation", + run_id="continuation-run", + background={"enabled": True, "lastEventIndex": 3, "lastSubIndex": 1}, + tool_result_for="call-1", + ), + ) + + assert [event["type"] for event in events] == ["RUN_STARTED", "RUN_ERROR"] + assert events[1]["message"] == "Session continuation not found" + # The foreground path stamps nothing, so a missing marker on the + # RUN_STARTED is itself the evidence the run was never refused: a + # refusal is a lone stamped RUN_ERROR with no run beginning at all. + assert all(resume_marker(event) is None for event in events) + + @pytest.mark.parametrize( + "background, message, marker", + [ + ( + {"enabled": True, "lastEventIndex": "4"}, + "Unreadable background resume position: {'enabled': True, 'lastEventIndex': '4'}", + {"eventIndex": -1, "subIndex": 0}, + ), + ( + {"enabled": True, "lastEventIndex": 2, "lastSubIndex": 0}, + "Run ghost-run not found in this session", + {"eventIndex": 3, "subIndex": 0}, + ), + ], + ) + def test_every_error_the_background_path_emits_carries_a_resume_marker( + self, agui_routes: AguiRoutes, background: Dict[str, Any], message: str, marker: Dict[str, int] + ): + """A client filtering by cursor would drop an unmarked error and reconnect into it. + + So a refusal is positioned past whatever the client last held, exactly + like the events it is refusing to send. + """ + events = post_events( + agui_routes, + DETACHABLE_PATH, + route_body(thread_id="ghost", run_id="ghost-run", background=background), + ) + + assert [event["type"] for event in events] == ["RUN_ERROR"] + assert events[0]["message"] == message + assert resume_marker(events[0]) == marker + assert agui_routes.detachable_model.invocations == 0