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