From 302e7547db4f43945685ada6722735a9e92de233 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 10:12:31 +0000 Subject: [PATCH 1/9] fix(cli): scaffold into the real project instead of the current directory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `agentomatic init` resolved `agents/` against the current working directory. Run it from a subdirectory of a project — or from the parent it had just been created in — and it silently built a second, parallel project, leaving the real `agents/` empty while reporting success. The `--dir` flag made it worse: the bootstrap wrote `stacks/`, `main.py` and `.agentomatic-stack` into whatever directory happened to contain the target, a home directory in the worst case. - Resolve the enclosing project by walking up from the current directory, stopping at a repository root or the user's home. Markers are specific (`.agentomatic-stack`, or `agents/` alongside `stacks/` or a platform `main.py`) so a scaffold cannot escape into a lookalike ancestor. - Use a lone project found in a subdirectory (the `new` then forgot-to-cd case) and say so; name the candidates and stay put when several match. - Add `--here` to opt out of detection. - Confine `--dir` to placing the package: bootstrap only when the target actually sits inside an identifiable project. - Fix the same cwd assumption in `add connection` (failed outright from a subdirectory) and `add ingestion` (landed in `/ingestion/`). Also, so that a scaffolded project runs and deploys as advertised: - Prompt for NAME instead of hard-exiting on a bare `agentomatic init`, and reject names that are not importable package names. - Report the absolute location and verify files reached disk, rather than claiming success over an empty directory. - Bootstrap the platform `main.py` (never overwriting an existing one, and warning when it is not an Agentomatic app) — `deploy` renders `uvicorn main:app`, so a project without one built a container that could not boot. Keep the bootstrap narrow: no empty `plugins/`, `endpoints/`, `ingestion/` or `pipelines/`. - Detect the model a local OpenAI-compatible server has loaded and bake that into `stacks/local.yaml`. The hardcoded default meant any other loaded model answered the first invoke with HTTP 400. - Add `openai` (langchain-openai) to the `all` extra: the default local stack's `omlx` provider needs it, so `agentomatic[all]` scaffolded a project whose every LLM call failed on a missing client library. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- docs/cli/commands.md | 39 ++- docs/getting-started/installation.md | 24 +- pyproject.toml | 2 +- src/agentomatic/cli/commands.py | 369 ++++++++++++++++++++++++-- src/agentomatic/stacks/defaults.py | 61 ++++- tests/test_cli_init_project_root.py | 332 +++++++++++++++++++++++ tests/test_local_stack_model_probe.py | 130 +++++++++ tests/test_production_readiness.py | 11 +- uv.lock | 7 +- 9 files changed, 932 insertions(+), 43 deletions(-) create mode 100644 tests/test_cli_init_project_root.py create mode 100644 tests/test_local_stack_model_probe.py diff --git a/docs/cli/commands.md b/docs/cli/commands.md index cc44aa7..32c935f 100644 --- a/docs/cli/commands.md +++ b/docs/cli/commands.md @@ -65,18 +65,48 @@ Commands: Scaffold a new agent directory from one of the pre-built templates. ```text -Usage: agentomatic init [OPTIONS] NAME +Usage: agentomatic init [OPTIONS] [NAME] Arguments: - NAME Agent name (snake_case) [required] + NAME Agent name (snake_case) [prompted for if omitted] Options: -t, --template [basic|class|full|coordinator|pipeline|rag|chatbot|deepagent|custom|legacy_dict|plugin|endpoint|connection|ingestion|extraction|langchain] Template to use (interactive picker if omitted) - -d, --dir TEXT Agents parent directory [default: agents] + -d, --dir TEXT Parent directory for the package (default: the + enclosing project's agents/, or the template's own + directory — plugins/, endpoints/, ingestion/, + pipelines/) + --here Scaffold into the current directory, skipping + project-root detection -f, --force Overwrite existing agent directory ``` +#### Where the code lands + +`init` writes into the **enclosing project**, not blindly into the current +directory. It walks up from where you are (stopping at a repository root or +your home directory) looking for a project marker — `.agentomatic-stack`, or +an `agents/` directory alongside `stacks/` or a platform `main.py`. + +```bash +cd my_platform/notebooks +agentomatic init helper --template basic +# 📦 Project root: /path/to/my_platform (detected from notebooks/) +# 📍 Location: /path/to/my_platform/agents/helper +``` + +If you run it in the directory you *just created a project in* — the +`agentomatic new my_platform` then forgot-to-`cd` case — the lone project in a +subdirectory is used, and the choice is printed. Pass `--here` to scaffold into +the current directory instead, or `--dir` to choose the location outright. + +When no project is found, the current directory is bootstrapped into one: +`stacks/local.yaml`, `stacks/remote.yaml`, `.agentomatic-stack`, +`agents/__init__.py` and a platform `main.py` (an existing `main.py` is never +overwritten). Nothing else is created — no empty `plugins/`, `endpoints/`, +`ingestion/` or `pipelines/`; use `agentomatic new` for the full layout. + #### Templates | Template | Description | @@ -101,6 +131,9 @@ Options: #### Examples ```bash +# Fully interactive — prompts for the name, then the template +agentomatic init + # Interactive template selection (requires questionary) agentomatic init support_agent diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index f0bee24..c97fe17 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -84,19 +84,27 @@ If you prefer a lightweight install, you can select only the modules and depende !!! warning "What `all` does *not* include" - `all` covers `langgraph`, `ollama`, `metrics`, `db`, `cli`, `studio`, - `optimize`, `telemetry`, `dotenv`, `security`, `swarm` and `vector`. - - It deliberately leaves out the vendor LLM drivers — `openai`, `azure`, - `vertex` — which follow the provider-agnostic principle: you install the - SDK for the backend you actually use. It also leaves out `db-postgres` - (an alternative to `db`) and `ui` (Chainlit), which is a heavy dependency. + `all` covers `langgraph`, `langchain`, `openai`, `ollama`, `metrics`, + `db`, `cli`, `studio`, `optimize`, `telemetry`, `dotenv`, `security`, + `swarm` and `vector`. + + `openai` is in that list for the **local-first** path, not for OpenAI: + the `omlx` provider in a freshly scaffolded `stacks/local.yaml` talks to + your own OpenAI-*compatible* server (oMLX, llama.cpp, vLLM, LM Studio) + through `langchain-openai`. Without it, `agentomatic[all]` scaffolds a + project whose every LLM call fails with *"client library is not + installed"*. + + It still leaves out the hosted vendor drivers `azure` and `vertex`, which + follow the provider-agnostic principle: you install the SDK for the + backend you actually use. It also leaves out `db-postgres` (an + alternative to `db`) and `ui` (Chainlit), which is a heavy dependency. So `agentomatic ui` needs `pip install "agentomatic[ui]"` even after an `all` install. Add what you need alongside it: ```bash - pip install "agentomatic[all,openai,ui]" + pip install "agentomatic[all,ui]" ``` !!! tip "Quote the extras in zsh/bash" diff --git a/pyproject.toml b/pyproject.toml index ccfa6fc..d29512d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,7 @@ dotenv = ["python-dotenv>=1.0"] security = ["pyjwt[crypto]>=2.8", "cryptography>=42.0"] swarm = ["langgraph-swarm>=0.1"] vector = ["numpy>=1.26"] -all = ["agentomatic[langgraph,ollama,metrics,db,cli,studio,optimize,telemetry,dotenv,security,swarm,vector]"] +all = ["agentomatic[langgraph,langchain,openai,ollama,metrics,db,cli,studio,optimize,telemetry,dotenv,security,swarm,vector]"] docs = [ "mkdocs-material>=9.7.7", "mike>=2.1", diff --git a/src/agentomatic/cli/commands.py b/src/agentomatic/cli/commands.py index f9b87fa..7ed9842 100644 --- a/src/agentomatic/cli/commands.py +++ b/src/agentomatic/cli/commands.py @@ -286,25 +286,166 @@ def _served_model_ids(response: Any) -> set[str]: return {str(item["id"]) for item in entries if isinstance(item, dict) and item.get("id")} -def _ensure_project_context(parent: Path) -> list[str]: +def _is_project_root(path: Path) -> bool: + """Report whether ``path`` looks like an Agentomatic project root. + + Markers are deliberately specific. A lone ``stacks/`` or ``main.py`` says + almost nothing — plenty of unrelated projects have both — and treating + either as a marker would let a scaffold escape into an ancestor directory + that merely resembles a project. + + Args: + path: Directory to test. + + Returns: + True when the directory carries an Agentomatic project marker. + """ + # Written by ``agentomatic new`` and by the init bootstrap alike, so every + # real project has one. + if (path / ".agentomatic-stack").is_file(): + return True + if not (path / "agents").is_dir(): + return False + stacks = path / "stacks" + if stacks.is_dir() and any(stacks.glob("*.yaml")): + return True + return (path / "main.py").is_file() + + +def _is_walk_boundary(path: Path) -> bool: + """Report whether the upward search must stop *after* inspecting ``path``. + + A repository root or the user's home directory is as far as a scaffold + should ever reach; beyond that any match is a coincidence. + + Args: + path: Directory just inspected. + + Returns: + True when the walk should not continue past ``path``. + """ + if (path / ".git").exists(): + return True + try: + return path == Path.home().resolve() + except (OSError, RuntimeError): + return False + + +def find_project_root(start: Path | None = None, *, ceiling: int = 6) -> Path | None: + """Walk up from ``start`` looking for the enclosing Agentomatic project. + + ``agentomatic init`` used to resolve ``agents/`` against the current + working directory, so running it from a subdirectory of a project silently + built a *second*, parallel project instead of adding an agent to the real + one — leaving the project's own ``agents/`` empty while the CLI reported + success. + + Args: + start: Directory to start from (default: the current directory). + ceiling: Maximum number of directories to inspect while walking up. + + Returns: + The project root, or None when ``start`` is not inside a project. + """ + current = (start or Path.cwd()).resolve() + for _ in range(ceiling): + if _is_project_root(current): + return current + if _is_walk_boundary(current) or current.parent == current: + break + current = current.parent + return None + + +def find_child_projects(start: Path | None = None) -> list[Path]: + """List Agentomatic projects in the immediate subdirectories of ``start``. + + Covers the most common mistake by far: ``agentomatic new my_platform`` + followed by ``agentomatic init hello`` *without* the ``cd`` in between. + + Args: + start: Directory whose children are scanned (default: cwd). + + Returns: + Child project roots, sorted by name. + """ + base = (start or Path.cwd()).resolve() + try: + return [p for p in sorted(base.iterdir()) if p.is_dir() and _is_project_root(p)] + except OSError: + return [] + + +def find_child_project(start: Path | None = None) -> Path | None: + """Find a lone Agentomatic project in an immediate subdirectory. + + Args: + start: Directory whose children are scanned (default: cwd). + + Returns: + The sole child project root, or None when there is not exactly one. + """ + children = find_child_projects(start) + return children[0] if len(children) == 1 else None + + +def _display_path(path: Path) -> str: + """Render ``path`` relative to the current directory when that is shorter. + + Absolute paths are what make a mis-targeted scaffold obvious, but they are + noisy for the common in-project case. + + Args: + path: Path to render. + + Returns: + A path string for display. + """ + resolved = path.resolve() + try: + relative = resolved.relative_to(Path.cwd().resolve()) + except ValueError: + return str(resolved) + return str(relative) if str(relative) != "." else str(resolved) + + +def _ensure_project_context( + parent: Path, + *, + root: Path | None = None, + project_name: str = "agentomatic-app", +) -> list[str]: """Create the project files a scaffolded agent needs in order to run. ``agentomatic init `` is often the first command a user types, in a directory that has never seen ``agentomatic new``. Without a stack the agent gets no LLM, and ``train.py`` cannot resolve one either — so the agent scaffolds successfully and then fails at the first invoke. This - fills in the missing pieces (stacks, active-stack marker, package init) - without touching anything that already exists. + fills in the missing pieces (stacks, active-stack marker, package init and + the platform ``main.py``) without touching anything that already exists. + + Deliberately narrow: it does not create ``plugins/``, ``endpoints/``, + ``ingestion/`` or ``pipelines/``. Those are discovered lazily and only + clutter a directory the user did not ask to have restructured — use + ``agentomatic new`` for the full project layout. Args: parent: Directory the agent package was written into (e.g. ``agents``). + root: Project root to bootstrap. Defaults to ``parent.parent``, which + is only correct when ``parent`` is the project's own ``agents/``: + with an explicit ``--dir`` it would scatter ``stacks/``, + ``main.py`` and ``.agentomatic-stack`` into whatever happens to + contain that directory — a home directory, in the worst case. + project_name: Name baked into the generated ``main.py`` title. Returns: - Relative paths of the files that were created, in write order. + Paths of the files that were created, in write order. """ + from agentomatic.cli.project import get_project_files from agentomatic.stacks.defaults import get_default_stack_yaml - root = parent.parent + root = root if root is not None else parent.parent created: list[str] = [] wanted: dict[Path, str] = { root / "stacks" / "local.yaml": get_default_stack_yaml("local"), @@ -312,15 +453,122 @@ def _ensure_project_context(parent: Path) -> list[str]: root / ".agentomatic-stack": "local\n", parent / "__init__.py": f'"""Agentomatic {parent.name} package."""\n', } + + # ``agentomatic run`` prefers ``main:app`` and ``agentomatic deploy`` + # renders ``uvicorn main:app`` unconditionally, so a project without one + # deploys to a container that cannot boot. + main_py = root / "main.py" + if not main_py.exists(): + wanted[main_py] = get_project_files(project_name)["main.py"] + for path, content in wanted.items(): if path.exists(): continue path.parent.mkdir(parents=True, exist_ok=True) path.write_text(content) - created.append(str(path)) + created.append(_display_path(path)) return created +_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_-]*$") + + +def _validate_agent_name(name: str) -> str | None: + """Validate an agent/component name. + + A name with a path separator in it used to scaffold into a surprising + place; one starting with a digit produced a package Python cannot import. + + Args: + name: Candidate name. + + Returns: + An error message, or None when the name is usable. + """ + if not name or not name.strip(): + return "NAME cannot be empty" + if "/" in name or "\\" in name or name in {".", ".."}: + return f"NAME cannot contain path separators: {name!r} (use --dir for the location)" + if not _NAME_RE.match(name): + return ( + f"Invalid NAME {name!r} — use letters, digits, '_' and '-', " + "starting with a letter or underscore" + ) + return None + + +def _prompt_agent_name() -> str | None: + """Ask for an agent name interactively. + + Returns: + The entered name, or None when there is no TTY or the user cancels. + """ + if not sys.stdin.isatty(): + return None + try: + answer = click.prompt( + click.style("Agent name", fg="cyan", bold=True), + default="my_agent", + show_default=True, + type=str, + ) + except (click.Abort, EOFError): + return None + return str(answer).strip() or None + + +def _resolve_project_root(*, here: bool = False) -> tuple[Path, str]: + """Decide which project directory a scaffold should be written into. + + Resolution order: + + 1. ``--here`` — the current directory, no questions asked. + 2. The nearest enclosing project root, walking up from the current + directory (so ``init`` works from anywhere inside a project). + 3. A lone Agentomatic project in an immediate subdirectory — the + ``agentomatic new x`` / forgot-to-``cd`` case. + 4. The current directory, which then gets bootstrapped into a project. + + Args: + here: Skip detection and use the current directory. + + Returns: + A ``(root, note)`` pair; ``note`` is a line to echo, or "". + """ + cwd = Path.cwd() + if here: + return cwd, "" + + enclosing = find_project_root(cwd) + if enclosing is not None: + if enclosing == cwd.resolve(): + return enclosing, "" + return enclosing, click.style( + f" 📦 Project root: {_display_path(enclosing)} (detected from {cwd.name}/)", + fg="cyan", + ) + + children = find_child_projects(cwd) + if len(children) == 1: + return children[0], click.style( + f" 📦 Project root: {_display_path(children[0])} — detected in a subdirectory.\n" + f" Pass --here to scaffold into {cwd.name}/ instead.", + fg="yellow", + ) + if children: + # Guessing between them would put the agent somewhere the user has no + # reason to look, which is the failure this detection exists to stop. + names = ", ".join(p.name for p in children[:4]) + return cwd, click.style( + f" ⚠️ Several projects here ({names}) — scaffolding into " + f"{_display_path(cwd)} itself.\n" + f" cd into the one you meant, or pass --dir.", + fg="yellow", + ) + + return cwd, "" + + # Default top-level directories for non-agent templates (platform discovery paths). _TEMPLATE_DEFAULT_DIRS: dict[str, str] = { "plugin": "plugins", @@ -356,12 +604,18 @@ def _ensure_project_context(parent: Path) -> list[str]: is_flag=True, help="Scaffold a full Agentomatic project (main.py, stacks, dirs)", ) +@click.option( + "--here", + is_flag=True, + help="Scaffold into the current directory, skipping project-root detection", +) @click.option("--force", "-f", is_flag=True, help="Overwrite existing files") def init( name: str | None, target_dir: str | None, template: str | None, as_project: bool, + here: bool, force: bool, ) -> None: """Create a new agent/project from a template. @@ -403,8 +657,17 @@ def init( click.echo(" agentomatic run --studio") return + # A bare ``agentomatic init`` is the first command most people type. Hard + # exiting on a missing NAME turned that into a dead end, so ask for it. if not name: - _print_error("NAME is required (or use --project NAME)") + name = _prompt_agent_name() + if not name: + _print_error("NAME is required. Try: agentomatic init my_agent --template basic") + sys.exit(1) + + invalid = _validate_agent_name(name) + if invalid: + _print_error(invalid) sys.exit(1) # Template selection @@ -430,11 +693,29 @@ def init( sys.exit(1) # Resolve target directory: ingestion→ingestion/, plugins→plugins/, etc. - if target_dir is None: - parent = _TEMPLATE_DEFAULT_DIRS.get(template, "agents") + # + # These used to resolve against the *current* directory, so running + # ``agentomatic init`` from a subdirectory of a project — or from the + # parent it was just created in — silently built a second, parallel + # project and left the real ``agents/`` empty while reporting success. + subdir = _TEMPLATE_DEFAULT_DIRS.get(template, "agents") + if target_dir is not None: + # An explicit --dir places the package, and must never decide where + # the project itself lives: bootstrapping relative to it scattered + # stacks/, main.py and .agentomatic-stack into whatever contained it. + # With a custom layout there is nothing to bootstrap unless the + # directory sits inside a project we can actually identify. + parent_path = Path(target_dir) + project_root = find_project_root(parent_path) or find_project_root() + root_note = "" else: - parent = target_dir - target = Path(parent) / name + project_root, root_note = _resolve_project_root(here=here) + parent_path = project_root / subdir + parent = str(parent_path) + target = parent_path / name + + if root_note: + click.echo(root_note) files = get_template_files(template, name) @@ -469,16 +750,43 @@ def init( file_path.write_text(content) written += 1 + # Report what is actually on disk, not what we believe we wrote — a + # scaffold that claims success while the directory is empty is the single + # most confusing failure this command can produce. + on_disk = sorted(q for q in target.rglob("*") if q.is_file()) click.echo() + if not on_disk: + _print_error( + f"Template '{template}' produced no files in {target.resolve()} — " + "this is a bug, please report it." + ) + sys.exit(1) + logger.success(f"Created '{name}' with template '{template}'") - click.echo(f" 📍 Location: {target}") - click.echo(f" 📦 Written: {written} · Skipped: {skipped}") + click.echo(f" 📍 Location: {target.resolve()}") + click.echo(f" 📦 Written: {written} · Skipped: {skipped} · On disk: {len(on_disk)}") # A bare `agentomatic init ` should leave a runnable workspace, not # an agent with no stack to get its LLM from. - bootstrapped = _ensure_project_context(Path(parent)) - if bootstrapped: - click.echo(f" 🧩 Bootstrapped: {', '.join(bootstrapped)}") + if project_root is not None: + bootstrapped = _ensure_project_context(Path(parent), root=project_root, project_name=name) + if bootstrapped: + click.echo(f" 🧩 Bootstrapped: {', '.join(bootstrapped)}") + + # ``uv init`` (and cookiecutters) leave a main.py of their own. It is not + # ours to overwrite, but ``deploy`` renders ``uvicorn main:app`` against + # whatever is there, so an unrelated main.py deploys to a dead container. + if ( + project_root is not None + and (project_root / "main.py").is_file() + and not _has_project_main_app(project_root) + ): + _print_warning( + f"{_display_path(project_root / 'main.py')} exists but does not define an " + "Agentomatic `app`. `agentomatic run` falls back to folder discovery; " + "`agentomatic deploy` needs `main:app` — see the scaffold in " + "`agentomatic new`." + ) click.echo() if template == "plugin": @@ -604,15 +912,30 @@ def add() -> None: @add.command("connection") @click.argument("agent_name") -@click.option("--dir", "-d", "agents_dir", default="agents", help="Agents directory") +@click.option( + "--dir", + "-d", + "agents_dir", + default=None, + help="Agents directory (default: the enclosing project's agents/)", +) @click.option("--force", "-f", is_flag=True, help="Overwrite connections.py if present") -def add_connection(agent_name: str, agents_dir: str, force: bool) -> None: +def add_connection(agent_name: str, agents_dir: str | None, force: bool) -> None: """Add ``connections.py`` to an existing agent package.""" from .templates import _connections_env_example, _connections_py, _connections_readme - target = Path(agents_dir) / agent_name + # ``agents`` used to be resolved against the current directory, so this + # failed outright from any subdirectory of the project it was aimed at. + if agents_dir is not None: + base = Path(agents_dir) + else: + base = (find_project_root() or Path.cwd()) / "agents" + target = base / agent_name if not target.exists(): - _print_error(f"Agent directory not found: {target}") + _print_error( + f"Agent directory not found: {_display_path(target)}\n" + f" Create it first: agentomatic init {agent_name} --template basic" + ) sys.exit(1) files = { @@ -653,11 +976,11 @@ def add_connection(agent_name: str, agents_dir: str, force: bool) -> None: "--dir", "-d", "ingestion_dir", - default="ingestion", - help="Ingestion directory (default: ingestion)", + default=None, + help="Ingestion directory (default: the enclosing project's ingestion/)", ) @click.option("--force", "-f", is_flag=True, help="Overwrite existing files") -def add_ingestion(name: str, ingestion_dir: str, force: bool) -> None: +def add_ingestion(name: str, ingestion_dir: str | None, force: bool) -> None: """Scaffold an ingestor under ``ingestion//`` (never into agents/).""" ctx = click.get_current_context() ctx.invoke( diff --git a/src/agentomatic/stacks/defaults.py b/src/agentomatic/stacks/defaults.py index eea60c9..68c648a 100644 --- a/src/agentomatic/stacks/defaults.py +++ b/src/agentomatic/stacks/defaults.py @@ -42,17 +42,70 @@ """ -def default_local_model() -> str: +def probe_local_models(base_url: str | None = None, *, timeout: float = 1.5) -> list[str]: + """List the models an OpenAI-compatible server currently has loaded. + + Every scaffolded project used to hard-code :data:`DEFAULT_LOCAL_MODEL`, so + a local server holding anything else answered the first invoke with HTTP + 400 (unknown model) — the generated stack was wrong the moment it was + written. Asking the server what it actually serves removes that mismatch. + + The probe is best-effort and never raises: no server, a slow one, or a + non-OpenAI endpoint all yield an empty list. + + Args: + base_url: Server base URL (default: :func:`default_local_base_url`). + timeout: Per-request timeout in seconds. + + Returns: + Model identifiers, in the order the server reported them. + """ + import json + import urllib.error + import urllib.request + + url = (base_url or default_local_base_url()).rstrip("/") + "/models" + try: + # A localhost model server must not be reached through a proxy, and + # an ambient HTTP_PROXY would otherwise swallow the request. + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + with opener.open(url, timeout=timeout) as response: + payload = json.loads(response.read().decode("utf-8")) + except (urllib.error.URLError, OSError, ValueError, json.JSONDecodeError): + return [] + + entries = payload.get("data") if isinstance(payload, dict) else None + if not isinstance(entries, list): + return [] + models: list[str] = [] + for entry in entries: + identifier = entry.get("id") if isinstance(entry, dict) else None + if isinstance(identifier, str) and identifier.strip(): + models.append(identifier.strip()) + return models + + +def default_local_model(*, probe: bool = True) -> str: """Return the model name to bake into a generated local stack. - Reads ``AGENTOMATIC_LOCAL_MODEL`` then ``OMLX_MODEL`` so a machine that - already has a model loaded scaffolds projects pointing straight at it. + Resolution order: ``AGENTOMATIC_LOCAL_MODEL``, ``OMLX_MODEL``, whatever a + running local server reports, then :data:`DEFAULT_LOCAL_MODEL`. + + Args: + probe: Ask a running local server when no env var is set. Returns: A model identifier, e.g. ``"Qwen3.5-9B-MLX-4bit"``. """ model = os.getenv("AGENTOMATIC_LOCAL_MODEL") or os.getenv("OMLX_MODEL") or "" - return model.strip().removeprefix("omlx/") or DEFAULT_LOCAL_MODEL + resolved = model.strip().removeprefix("omlx/") + if resolved: + return resolved + if probe: + detected = probe_local_models() + if detected: + return detected[0] + return DEFAULT_LOCAL_MODEL def default_local_base_url() -> str: diff --git a/tests/test_cli_init_project_root.py b/tests/test_cli_init_project_root.py new file mode 100644 index 0000000..6cae891 --- /dev/null +++ b/tests/test_cli_init_project_root.py @@ -0,0 +1,332 @@ +# pyright: reportMissingParameterType=none +"""Regression tests for ``agentomatic init`` target resolution. + +``init`` used to resolve ``agents/`` against the *current* directory. Running +it from a subdirectory of a project — or from the parent it had just been +created in — silently built a second, parallel project and left the real +``agents/`` empty while reporting success. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from agentomatic.cli.commands import ( + _validate_agent_name, + cli, + find_child_project, + find_child_projects, + find_project_root, +) +from agentomatic.cli.project import scaffold_project + + +def _make_project(root: Path, name: str = "demo") -> Path: + """Scaffold a project and return its path.""" + target = root / name + scaffold_project(target, name) + return target + + +class TestProjectRootDiscovery: + def test_detects_scaffolded_project(self, tmp_path: Path) -> None: + project = _make_project(tmp_path) + assert find_project_root(project) == project.resolve() + + def test_walks_up_from_subdirectory(self, tmp_path: Path) -> None: + project = _make_project(tmp_path) + deep = project / "docs" / "nested" + deep.mkdir(parents=True) + assert find_project_root(deep) == project.resolve() + + def test_returns_none_outside_a_project(self, tmp_path: Path) -> None: + bare = tmp_path / "bare" + bare.mkdir() + assert find_project_root(bare) is None + + def test_finds_lone_child_project(self, tmp_path: Path) -> None: + project = _make_project(tmp_path) + assert find_child_project(tmp_path) == project.resolve() + + def test_ignores_ambiguous_children(self, tmp_path: Path) -> None: + _make_project(tmp_path, "one") + _make_project(tmp_path, "two") + assert find_child_project(tmp_path) is None + + +class TestInitTargetsTheRealProject: + def test_init_from_subdirectory_writes_into_project( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + project = _make_project(tmp_path) + deep = project / "docs" / "nested" + deep.mkdir(parents=True) + monkeypatch.chdir(deep) + + result = CliRunner().invoke(cli, ["init", "hello", "--template", "basic"]) + + assert result.exit_code == 0, result.output + assert (project / "agents" / "hello" / "agent.py").is_file() + assert not (deep / "agents").exists(), "scaffolded a parallel project" + + def test_init_from_parent_writes_into_child_project( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """``agentomatic new x`` then ``init`` without the ``cd`` in between.""" + project = _make_project(tmp_path) + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, ["init", "hello", "--template", "basic"]) + + assert result.exit_code == 0, result.output + assert (project / "agents" / "hello" / "agent.py").is_file() + assert not (tmp_path / "agents").exists(), "scaffolded a parallel project" + + def test_ambiguous_parent_does_not_redirect( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two candidate projects must not silently pick one — and must say so.""" + one = _make_project(tmp_path, "one") + two = _make_project(tmp_path, "two") + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, ["init", "hello", "--template", "basic"]) + + assert result.exit_code == 0, result.output + assert (tmp_path / "agents" / "hello" / "agent.py").is_file() + assert not (one / "agents" / "hello").exists() + assert not (two / "agents" / "hello").exists() + assert "Several projects here" in result.output + + def test_here_flag_overrides_detection( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + project = _make_project(tmp_path) + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, ["init", "hello", "--template", "basic", "--here"]) + + assert result.exit_code == 0, result.output + assert (tmp_path / "agents" / "hello" / "agent.py").is_file() + assert not (project / "agents" / "hello").exists() + + def test_dir_flag_still_wins(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + project = _make_project(tmp_path) + monkeypatch.chdir(project) + explicit = tmp_path / "elsewhere" + + result = CliRunner().invoke( + cli, ["init", "hello", "--template", "basic", "--dir", str(explicit)] + ) + + assert result.exit_code == 0, result.output + assert (explicit / "hello" / "agent.py").is_file() + + def test_reports_an_absolute_location( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The path printed must be unambiguous about where code landed.""" + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, ["init", "hello", "--template", "basic"]) + + assert result.exit_code == 0, result.output + assert str((tmp_path / "agents" / "hello").resolve()) in result.output + + +class TestInitBootstrapsRunnableProject: + def test_bare_directory_gets_main_py_and_stacks( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, ["init", "solo", "--template", "basic"]) + + assert result.exit_code == 0, result.output + assert (tmp_path / "agents" / "solo" / "agent.py").is_file() + # ``deploy`` renders ``uvicorn main:app``; without main.py the image + # cannot boot. + assert (tmp_path / "main.py").is_file() + assert (tmp_path / "stacks" / "local.yaml").is_file() + assert (tmp_path / ".agentomatic-stack").is_file() + + def test_existing_main_py_is_never_overwritten( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """``uv init`` leaves its own main.py; the scaffold must not clobber it.""" + sentinel = "# hand-written, do not touch\n" + (tmp_path / "main.py").write_text(sentinel) + monkeypatch.chdir(tmp_path) + + result = CliRunner().invoke(cli, ["init", "solo", "--template", "basic"]) + + assert result.exit_code == 0, result.output + assert (tmp_path / "main.py").read_text() == sentinel + + def test_no_stray_component_directories( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The bootstrap must not restructure a directory it was not given.""" + monkeypatch.chdir(tmp_path) + + CliRunner().invoke(cli, ["init", "solo", "--template", "basic"]) + + for junk in ("plugins", "endpoints", "ingestion", "pipelines"): + assert not (tmp_path / junk).exists(), f"bootstrap created {junk}/" + + +class TestNameValidation: + def test_accepts_ordinary_names(self) -> None: + for name in ("hello", "my_agent", "my-agent", "_x9"): + assert _validate_agent_name(name) is None + + def test_rejects_path_separators(self) -> None: + assert _validate_agent_name("a/b") is not None + + def test_rejects_empty_and_dots(self) -> None: + assert _validate_agent_name("") is not None + assert _validate_agent_name(".") is not None + + def test_rejects_leading_digit(self) -> None: + assert _validate_agent_name("9lives") is not None + + def test_cli_rejects_bad_name(self, tmp_path: Path) -> None: + result = CliRunner().invoke( + cli, ["init", "bad/name", "--template", "basic", "--dir", str(tmp_path)] + ) + assert result.exit_code == 1 + + +class TestExplicitDirStaysContained: + def test_dir_does_not_scatter_project_files_into_its_parent( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """``--dir`` places the package; it must not relocate the project. + + The bootstrap used to write ``stacks/``, ``main.py`` and + ``.agentomatic-stack`` into ``--dir``'s parent — a home directory when + the flag pointed at one. With a custom layout and no project in sight + there is nothing to bootstrap, so nothing is written. + """ + workspace = tmp_path / "workspace" + workspace.mkdir() + monkeypatch.chdir(workspace) + + result = CliRunner().invoke( + cli, + ["init", "hello", "--template", "basic", "--dir", str(workspace / "custom")], + ) + + assert result.exit_code == 0, result.output + assert (workspace / "custom" / "hello" / "agent.py").is_file() + for stray in (".agentomatic-stack", "stacks", "main.py"): + assert not (tmp_path / stray).exists(), f"leaked {stray} above --dir" + assert not (workspace / stray).exists(), f"scattered {stray} into cwd" + + def test_dir_inside_a_project_bootstraps_that_project( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A custom directory *within* a project still gets its stacks.""" + project = _make_project(tmp_path) + monkeypatch.chdir(project) + + result = CliRunner().invoke( + cli, + ["init", "hello", "--template", "basic", "--dir", str(project / "custom")], + ) + + assert result.exit_code == 0, result.output + assert (project / "custom" / "hello" / "agent.py").is_file() + assert (project / "stacks" / "local.yaml").is_file() + assert not (tmp_path / "stacks").exists() + + +class TestRootDetectionStaysBounded: + def test_lone_stacks_dir_is_not_a_project(self, tmp_path: Path) -> None: + """A ``stacks/`` directory alone is far too weak a signal.""" + (tmp_path / "stacks").mkdir() + (tmp_path / "stacks" / "local.yaml").write_text("name: local\n") + child = tmp_path / "child" + child.mkdir() + assert find_project_root(child) is None + + def test_lone_main_py_is_not_a_project(self, tmp_path: Path) -> None: + (tmp_path / "main.py").write_text("print('hi')\n") + child = tmp_path / "child" + child.mkdir() + assert find_project_root(child) is None + + def test_walk_stops_at_a_git_boundary(self, tmp_path: Path) -> None: + """A scaffold must not reach out of the repository it sits in.""" + outer = _make_project(tmp_path, "outer") + repo = outer / "vendored" + (repo / ".git").mkdir(parents=True) + inner = repo / "src" + inner.mkdir() + assert find_project_root(inner) is None + + def test_walk_stops_at_home(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + home = tmp_path / "home" + (home / "projects").mkdir(parents=True) + _make_project(tmp_path, "above_home") + monkeypatch.setenv("HOME", str(home)) + assert find_project_root(home / "projects") is None + + +class TestAddResolvesTheProject: + def test_add_connection_works_from_a_subdirectory( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + project = _make_project(tmp_path) + monkeypatch.chdir(project) + assert CliRunner().invoke(cli, ["init", "hello", "--template", "basic"]).exit_code == 0 + + sub = project / "notebooks" + sub.mkdir() + monkeypatch.chdir(sub) + + result = CliRunner().invoke(cli, ["add", "connection", "hello"]) + + assert result.exit_code == 0, result.output + assert (project / "agents" / "hello" / "connections.py").is_file() + + def test_add_ingestion_lands_in_the_project( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + project = _make_project(tmp_path) + sub = project / "notebooks" + sub.mkdir() + monkeypatch.chdir(sub) + + result = CliRunner().invoke(cli, ["add", "ingestion", "docs"]) + + assert result.exit_code == 0, result.output + assert (project / "ingestion" / "docs").is_dir() + assert not (sub / "ingestion").exists() + + def test_add_connection_missing_agent_is_actionable( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + project = _make_project(tmp_path) + monkeypatch.chdir(project) + + result = CliRunner().invoke(cli, ["add", "connection", "ghost"]) + + assert result.exit_code == 1 + assert "agentomatic init ghost" in result.output + + +class TestChildProjectListing: + def test_lists_every_child_project(self, tmp_path: Path) -> None: + _make_project(tmp_path, "alpha") + _make_project(tmp_path, "beta") + (tmp_path / "unrelated").mkdir() + found = {p.name for p in find_child_projects(tmp_path)} + assert found == {"alpha", "beta"} + + def test_empty_without_projects(self, tmp_path: Path) -> None: + (tmp_path / "unrelated").mkdir() + assert find_child_projects(tmp_path) == [] diff --git a/tests/test_local_stack_model_probe.py b/tests/test_local_stack_model_probe.py new file mode 100644 index 0000000..a3eb004 --- /dev/null +++ b/tests/test_local_stack_model_probe.py @@ -0,0 +1,130 @@ +# pyright: reportMissingParameterType=none +"""Tests for local-model auto-detection in scaffolded stacks. + +A generated ``stacks/local.yaml`` used to hard-code one model name, so a local +server holding anything else rejected the very first invoke with HTTP 400 +(unknown model). The scaffold now asks the server what it serves. +""" + +from __future__ import annotations + +import json +import threading +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest + +from agentomatic.stacks.defaults import ( + DEFAULT_LOCAL_MODEL, + default_local_model, + probe_local_models, +) + +SERVED_MODEL = "Llama-3.2-3B-Instruct-4bit" + + +def _make_handler(payload: bytes | None, status: int = 200) -> type[BaseHTTPRequestHandler]: + class Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 - http.server API + if payload is None: + self.send_response(500) + self.end_headers() + return + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args: object) -> None: + """Silence the default stderr access log.""" + + return Handler + + +@pytest.fixture +def model_server() -> Iterator[str]: + """Serve an OpenAI-compatible ``/v1/models`` response; yield its base URL.""" + body = json.dumps({"object": "list", "data": [{"id": SERVED_MODEL}]}).encode() + server = HTTPServer(("127.0.0.1", 0), _make_handler(body)) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}/v1" + finally: + server.shutdown() + server.server_close() + + +class TestProbeLocalModels: + def test_reports_served_models(self, model_server: str) -> None: + assert probe_local_models(model_server) == [SERVED_MODEL] + + def test_no_server_yields_empty_list(self) -> None: + # Port 1 is reserved and never listening. + assert probe_local_models("http://127.0.0.1:1/v1", timeout=0.5) == [] + + def test_non_json_body_yields_empty_list(self) -> None: + server = HTTPServer(("127.0.0.1", 0), _make_handler(b"not json")) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + assert probe_local_models(f"http://127.0.0.1:{server.server_port}/v1") == [] + finally: + server.shutdown() + server.server_close() + + def test_server_error_yields_empty_list(self) -> None: + server = HTTPServer(("127.0.0.1", 0), _make_handler(None)) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + assert probe_local_models(f"http://127.0.0.1:{server.server_port}/v1") == [] + finally: + server.shutdown() + server.server_close() + + +class TestDefaultLocalModel: + def test_env_var_wins_over_probe( + self, model_server: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OMLX_BASE_URL", model_server) + monkeypatch.setenv("AGENTOMATIC_LOCAL_MODEL", "explicit-choice") + assert default_local_model() == "explicit-choice" + + def test_probe_fills_in_when_no_env_var( + self, model_server: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("AGENTOMATIC_LOCAL_MODEL", raising=False) + monkeypatch.delenv("OMLX_MODEL", raising=False) + monkeypatch.setenv("OMLX_BASE_URL", model_server) + assert default_local_model() == SERVED_MODEL + + def test_falls_back_when_nothing_is_running(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AGENTOMATIC_LOCAL_MODEL", raising=False) + monkeypatch.delenv("OMLX_MODEL", raising=False) + monkeypatch.setenv("OMLX_BASE_URL", "http://127.0.0.1:1/v1") + assert default_local_model() == DEFAULT_LOCAL_MODEL + + def test_probe_can_be_disabled( + self, model_server: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("AGENTOMATIC_LOCAL_MODEL", raising=False) + monkeypatch.delenv("OMLX_MODEL", raising=False) + monkeypatch.setenv("OMLX_BASE_URL", model_server) + assert default_local_model(probe=False) == DEFAULT_LOCAL_MODEL + + +def test_scaffolded_stack_uses_the_served_model( + model_server: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """End to end: the YAML written to disk names the model that is loaded.""" + monkeypatch.delenv("AGENTOMATIC_LOCAL_MODEL", raising=False) + monkeypatch.delenv("OMLX_MODEL", raising=False) + monkeypatch.setenv("OMLX_BASE_URL", model_server) + + from agentomatic.stacks.defaults import get_default_stack_yaml + + assert SERVED_MODEL in get_default_stack_yaml("local") diff --git a/tests/test_production_readiness.py b/tests/test_production_readiness.py index 55bea01..de37c25 100644 --- a/tests/test_production_readiness.py +++ b/tests/test_production_readiness.py @@ -320,6 +320,11 @@ def test_all_extra_contents_match_what_the_docs_claim() -> None: documented = { "langgraph", + # ``langchain`` + ``openai`` carry langchain-openai, which the default + # scaffolded ``stacks/local.yaml`` needs: its ``omlx`` provider speaks + # to a local OpenAI-*compatible* server. + "langchain", + "openai", "ollama", "metrics", "db", @@ -337,9 +342,11 @@ def test_all_extra_contents_match_what_the_docs_claim() -> None: f"(added={included - documented}, removed={documented - included})" ) - # These are deliberately excluded — vendor SDKs (provider-agnostic + # These are deliberately excluded — hosted vendor SDKs (provider-agnostic # principle), an alternative DB driver, and the heavy Chainlit UI. - for deliberately_excluded in ("openai", "azure", "vertex", "db-postgres", "ui"): + # ``openai`` is *not* on this list: the default local stack's ``omlx`` + # provider needs langchain-openai to reach an OpenAI-compatible server. + for deliberately_excluded in ("azure", "vertex", "db-postgres", "ui"): assert deliberately_excluded in extras, f"{deliberately_excluded} extra vanished" assert deliberately_excluded not in included diff --git a/uv.lock b/uv.lock index 1db6c70..bf51ff9 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ [[package]] name = "agentomatic" -version = "1.12.0" +version = "1.13.0" source = { editable = "." } dependencies = [ { name = "click" }, @@ -31,8 +31,11 @@ all = [ { name = "cryptography" }, { name = "deepeval" }, { name = "holysheet" }, + { name = "langchain" }, + { name = "langchain-community" }, { name = "langchain-core" }, { name = "langchain-ollama" }, + { name = "langchain-openai" }, { name = "langgraph" }, { name = "langgraph-swarm" }, { name = "numpy" }, @@ -159,7 +162,7 @@ docs = [ [package.metadata] requires-dist = [ - { name = "agentomatic", extras = ["langgraph", "ollama", "metrics", "db", "cli", "studio", "optimize", "telemetry", "dotenv", "security", "swarm", "vector"], marker = "extra == 'all'" }, + { name = "agentomatic", extras = ["langgraph", "langchain", "openai", "ollama", "metrics", "db", "cli", "studio", "optimize", "telemetry", "dotenv", "security", "swarm", "vector"], marker = "extra == 'all'" }, { name = "aiosqlite", marker = "extra == 'db'", specifier = ">=0.20" }, { name = "asyncpg", marker = "extra == 'db-postgres'", specifier = ">=0.30" }, { name = "build", marker = "extra == 'dev'", specifier = ">=1.0" }, From 02b1ef9be16844b73ad1d38291588329208b9c6d Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 14:40:38 +0000 Subject: [PATCH 2/9] style(tests): format an assert stably across both pinned ruff versions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PR Quality Gate runs pre-commit over the files a PR changes, and its ruff-format hook is pinned to v0.8.6 while the repo's dev dependency resolves to ruff 0.15.17. The two wrap an assert message in opposite directions, so this pre-existing block reformatted under the pinned hook and failed the gate as soon as a PR touched the file — it is unrelated to the surrounding change and already reformats on main. Hoist the message into a local so neither version needs to wrap it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- tests/test_production_readiness.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_production_readiness.py b/tests/test_production_readiness.py index de37c25..ba82fc7 100644 --- a/tests/test_production_readiness.py +++ b/tests/test_production_readiness.py @@ -85,9 +85,11 @@ def test_slug_alias_routes_work_but_are_not_documented_twice(dual_mounted_platfo spec = app.openapi() assert "/api/v1/hello/invoke" in spec["paths"], "canonical route must be documented" - assert "/api/v1/agent-hello/invoke" not in spec["paths"], ( - "the slug alias must not be documented — it doubles the advertised surface" - ) + # Kept on its own line: ruff 0.8.6 (pinned in .pre-commit-config.yaml) and + # newer ruff wrap an assert message in opposite directions, so an inline + # message here reformats differently depending on which version runs. + alias_note = "the slug alias must not be documented — it doubles the advertised surface" + assert "/api/v1/agent-hello/invoke" not in spec["paths"], alias_note with TestClient(app) as client: # ...but it must still route, so Studio's slug-based calls keep working. From cc1f746511d0e787b3ac580527520e531c4b1aa9 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 14:57:16 +0000 Subject: [PATCH 3/9] feat(tasks): make task progress streams resumable after a dropped connection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task events existed only while somebody was listening: ``_emit`` returned early when a task had no subscribers, so everything emitted while a client was disconnected was never created, let alone retained. A dropped SSE connection could therefore only be recovered as far as the task's *current* state — every intermediate step was lost, and there was no way to ask for the gap. - Give TaskEvent a monotonic per-task ``sequence``, gapless from 1. - Add TaskEventLog: an ordered, bounded replay buffer with an in-memory default, a null implementation, and a provider registry so a shared backend can be registered for cross-worker replay (the interface ships, not a vendor client). - Record every event whether or not anyone is subscribed. - Publish the sequence as the SSE ``id:``, so a browser EventSource resumes through ``Last-Event-ID`` with no client code; ``?since=`` does the same for other clients. A new subscriber's snapshot also carries a cursor, so a client that drops immediately after it can still resume precisely. - Announce a truncated resume point instead of serving a partial history as if it were complete. - Drop a task's retained events when the task record is deleted. This lands for every target type at once — agents, pipelines, plugins, endpoints and ingestion share the one task manager, so a pipeline's per-step progress now resumes exactly like an agent's. A2A gains the matching surface: GET /a2a/tasks/{id}/events streams the same events shaped as A2A task objects, with the same resumption contract. The agent card stops hardcoding its capabilities and reports streaming/pushNotifications/stateTransitionHistory from what the deployment actually wired up — a card claiming a capability the platform lacks sends clients down a path that then 501s. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- docs/architecture/api-reference.md | 33 ++ docs/guide/tasks.md | 69 +++- src/agentomatic/core/router_factory.py | 122 ++++++- src/agentomatic/tasks/__init__.py | 14 + src/agentomatic/tasks/event_log.py | 303 ++++++++++++++++++ src/agentomatic/tasks/manager.py | 55 +++- src/agentomatic/tasks/models.py | 14 +- src/agentomatic/tasks/routes.py | 106 ++++++- tests/test_a2a_streaming.py | 170 ++++++++++ tests/test_task_event_resume.py | 423 +++++++++++++++++++++++++ 10 files changed, 1285 insertions(+), 24 deletions(-) create mode 100644 src/agentomatic/tasks/event_log.py create mode 100644 tests/test_a2a_streaming.py create mode 100644 tests/test_task_event_resume.py diff --git a/docs/architecture/api-reference.md b/docs/architecture/api-reference.md index 8e85373..12739dc 100644 --- a/docs/architecture/api-reference.md +++ b/docs/architecture/api-reference.md @@ -429,6 +429,39 @@ curl http://localhost:8000/api/v1/my_agent/a2a/tasks/task_a1b2c3d4e5f6 --- +#### `GET /a2a/tasks/{task_id}/events` + +Subscribe — or **re-subscribe** — to an A2A task's progress over +Server-Sent Events. Frames are A2A task objects (`submitted`, `working`, +`completed`, `failed`, `canceled`), and the stream ends with `data: [DONE]`. +Returns `501` without the task subsystem, `404` for an unknown task. + +Every frame carries an SSE `id:` holding the event's monotonic `sequence`, so +a client that loses the connection resumes precisely instead of restarting: + +```bash +# Subscribe +curl -N http://localhost:8000/api/v1/my_agent/a2a/tasks/task_a1b2c3d4e5f6/events + +# Re-subscribe after event 4 — replays 5 onwards, then goes live +curl -N -H "Last-Event-ID: 4" \ + http://localhost:8000/api/v1/my_agent/a2a/tasks/task_a1b2c3d4e5f6/events +curl -N "http://localhost:8000/api/v1/my_agent/a2a/tasks/task_a1b2c3d4e5f6/events?since=4" +``` + +```json +{"task_id": "task_a1b2c3d4e5f6", "sequence": 5, "status": "working", + "progress": {"percent": 60.0, "message": "step 3"}} +``` + +A reconnect whose resume point has already been evicted receives a +`{"event": "truncated"}` frame followed by the current task object, so a +client is never handed a partial history as if it were complete. See +[Tasks & Execution Modes](../guide/tasks.md#resuming-a-dropped-stream) for +retention limits and the multi-worker caveat. + +--- + #### `POST /a2a/tasks/{task_id}/cancel` Request cooperative cancellation of an in-flight A2A task. This endpoint is diff --git a/docs/guide/tasks.md b/docs/guide/tasks.md index f64a5dc..beb10d7 100644 --- a/docs/guide/tasks.md +++ b/docs/guide/tasks.md @@ -72,7 +72,7 @@ is enabled by default (`enable_tasks=True`). | `GET` | `/tasks` | List / filter tasks | | `GET` | `/tasks/{id}` | Poll status + progress | | `GET` | `/tasks/{id}/result` | Fetch the result (`409` if still pending) | -| `GET` | `/tasks/{id}/events` | Live SSE progress stream | +| `GET` | `/tasks/{id}/events` | Resumable SSE progress stream | | `POST` | `/tasks/{id}/cancel` | Request cancellation | | `DELETE` | `/tasks/{id}` | Delete a terminal record | @@ -202,10 +202,75 @@ const es = new EventSource("/api/v1/tasks/task_a1b2c3d4e5f6a7b8/events"); es.onmessage = (e) => { if (e.data === "[DONE]") { es.close(); return; } const evt = JSON.parse(e.data); - console.log(evt.status, evt.progress?.percent, evt.progress?.message); + console.log(evt.sequence, evt.status, evt.progress?.percent); }; ``` +#### Resuming a dropped stream + +Every frame carries an SSE `id:` holding the event's `sequence` — a +per-task counter that starts at 1 and never skips. That is what makes a +dropped connection recoverable rather than merely survivable. + +A browser `EventSource` needs no extra code: it reconnects on its own and +sends the last id it saw back in `Last-Event-ID`, and the server replays +only what was missed. Other clients pass the same value explicitly: + +```bash +# Resume after event 7 — events 8 onwards, then live updates. +curl -N "http://localhost:8000/api/v1/tasks/task_a1b2c3d4e5f6a7b8/events?since=7" +curl -N -H "Last-Event-ID: 7" \ + http://localhost:8000/api/v1/tasks/task_a1b2c3d4e5f6a7b8/events +``` + +A client reconnecting *without* a cursor is treated as new: it receives the +task's current snapshot (itself carrying an `id:`, so it can resume from +there) rather than the whole history. + +!!! warning "Replay is bounded, and says so" + + The event log keeps the last 512 events of each of the last 1024 tasks by + default. If you resume from a point that has already been evicted, the + stream sends a `{"event": "truncated"}` frame followed by the current + snapshot — it never serves a partial history as if it were complete. + + Tune it, or turn replay off, when constructing the manager: + + ```python + from agentomatic.tasks import InMemoryTaskEventLog, TaskManager + + manager = TaskManager( + event_log=InMemoryTaskEventLog(max_events_per_task=2048, max_tasks=256) + ) + ``` + +!!! note "One process, by default" + + The built-in event log lives in the worker that ran the task. Task + *records* are shared through the store, so a client reconnecting to + another worker still gets the correct snapshot and terminal result — but + the missed intermediate events are not replayable from there. + + For replay across workers or replicas, register a shared log. Agentomatic + ships the interface, not a vendor client: + + ```python + from agentomatic.tasks import TaskEventLog, register_event_log_provider + + class RedisEventLog(TaskEventLog): + async def append(self, event): ... + async def replay(self, task_id, *, after=0): ... + async def earliest_sequence(self, task_id): ... + async def latest_sequence(self, task_id): ... + async def drop(self, task_id): ... + + register_event_log_provider("redis", lambda **cfg: RedisEventLog(**cfg)) + ``` + +This applies to **every** task target — agents, pipelines, plugins, +endpoints and ingestion all run through the same manager, so a pipeline's +per-step progress resumes exactly like an agent's. + ### Fetching the result ```bash diff --git a/src/agentomatic/core/router_factory.py b/src/agentomatic/core/router_factory.py index 7fc26d1..be96c39 100644 --- a/src/agentomatic/core/router_factory.py +++ b/src/agentomatic/core/router_factory.py @@ -2,13 +2,14 @@ from __future__ import annotations +import asyncio import json import time import uuid from collections.abc import Callable from typing import Any, cast -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, Request from fastapi.responses import StreamingResponse from loguru import logger from pydantic import BaseModel, ConfigDict, Field @@ -1159,17 +1160,32 @@ async def get_card() -> dict[str, Any]: "description": m.description, "version": m.version, "framework": m.framework, + # Reported from what this deployment actually wired up. A card + # that claims a capability unconditionally sends A2A clients down + # code paths that then 501, which is worse than admitting the gap. "capabilities": { "streaming": True, "chat": True, "invoke": True, "a2a": True, + # Resumable task streams and push both need the task manager. + "stateTransitionHistory": task_manager is not None, + "pushNotifications": task_manager is not None, + "resumableStreams": task_manager is not None, }, "endpoints": { "invoke": f"{api_prefix}/{agent_name}/invoke", "chat": f"{api_prefix}/{agent_name}/chat", "stream": f"{api_prefix}/{agent_name}/invoke/stream", "health": f"{api_prefix}/{agent_name}/health", + **( + { + "a2a_tasks": f"{api_prefix}/{agent_name}/a2a/tasks", + "a2a_events": (f"{api_prefix}/{agent_name}/a2a/tasks/{{task_id}}/events"), + } + if task_manager is not None + else {} + ), }, "metadata": m.metadata, } @@ -1202,6 +1218,34 @@ def _a2a_view(record: Any) -> dict[str, Any]: view["error"] = record.error return view + def _a2a_event_view(evt: Any, record: Any) -> dict[str, Any]: + """Render one recorded event as an A2A-shaped task object. + + A replayed event describes the state *at that point*, so its status and + progress come from the event rather than from the task's current + snapshot — otherwise every historical frame would claim the task's + latest state. The record is consulted only for the terminal payload. + + Args: + evt: The recorded :class:`~agentomatic.tasks.models.TaskEvent`. + record: The task record, used for the result on a terminal event. + + Returns: + An A2A task object for one point in the task's history. + """ + view: dict[str, Any] = { + "task_id": evt.task_id, + "sequence": evt.sequence, + "status": _A2A_STATE.get(evt.status.value, evt.status.value), + "progress": evt.progress.model_dump() if evt.progress else None, + } + if evt.status.is_terminal and record is not None: + terminal = _a2a_view(record) + for key in ("result", "artifacts", "raw", "error"): + if key in terminal: + view[key] = terminal[key] + return view + @router.post("/a2a/tasks") async def submit_a2a_task(request: A2ATaskRequest) -> dict[str, Any]: """Submit an A2A task. @@ -1283,6 +1327,82 @@ async def get_a2a_task(task_id: str) -> dict[str, Any]: "message": "Task tracking requires a task manager", } + @router.get("/a2a/tasks/{task_id}/events") + async def stream_a2a_task(task_id: str, request: Request, since: int = 0): + """Subscribe (or re-subscribe) to an A2A task's progress over SSE. + + This is the A2A analogue of the platform's task stream: same + resumption contract — every frame carries its sequence as the SSE + ``id:``, and a reconnect replays the gap named by ``Last-Event-ID`` or + ``?since=`` — but the frames are A2A task objects rather than raw task + records, so an A2A client sees the states it expects. + """ + from agentomatic.tasks.event_log import parse_last_event_id + + if task_manager is None: + raise HTTPException(501, "A2A task streaming requires a task manager") + record = await task_manager.get(task_id) + if record is None: + raise HTTPException(404, f"Task '{task_id}' not found") + + resume_from = since or parse_last_event_id(request.headers.get("last-event-id")) + + def _frame(payload: dict[str, Any], seq: int | None = None) -> str: + body = f"data: {json.dumps(payload, default=str)}\n\n" + return f"id: {seq}\n{body}" if seq else body + + async def event_stream() -> Any: + queue = await task_manager.subscribe(task_id) + try: + resumption = await task_manager.resume(task_id, after=resume_from) + cursor = resumption.cursor + if resume_from <= 0: + cursor = await task_manager.event_log.latest_sequence(task_id) + + if resumption.truncated: + yield _frame( + { + "task_id": task_id, + "event": "truncated", + "resumed_after": resume_from, + } + ) + for evt in resumption.replay: + yield _frame( + _a2a_event_view(evt, await task_manager.get(task_id)), evt.sequence + ) + + current = await task_manager.get(task_id) + if current is not None and (resume_from <= 0 or resumption.truncated): + yield _frame(_a2a_view(current), cursor) + if current is not None and current.status.is_terminal: + yield "data: [DONE]\n\n" + return + + while True: + try: + evt = await asyncio.wait_for(queue.get(), timeout=15.0) + except TimeoutError: + yield ": keep-alive\n\n" + continue + if evt.sequence and evt.sequence <= cursor: + continue + cursor = evt.sequence or cursor + yield _frame( + _a2a_event_view(evt, await task_manager.get(task_id)), evt.sequence + ) + if evt.status.is_terminal: + yield "data: [DONE]\n\n" + return + finally: + task_manager.unsubscribe(task_id, queue) + + return StreamingResponse( + event_stream(), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + ) + @router.post("/a2a/tasks/{task_id}/cancel") async def cancel_a2a_task(task_id: str) -> dict[str, Any]: """Cancel an in-flight A2A task.""" diff --git a/src/agentomatic/tasks/__init__.py b/src/agentomatic/tasks/__init__.py index 931f5e0..f7ed193 100644 --- a/src/agentomatic/tasks/__init__.py +++ b/src/agentomatic/tasks/__init__.py @@ -35,6 +35,14 @@ make_plugin_dispatcher, make_plugin_input_validator, ) +from .event_log import ( + InMemoryTaskEventLog, + NullTaskEventLog, + TaskEventLog, + available_event_log_providers, + create_event_log, + register_event_log_provider, +) from .manager import TaskInputValidationError, TaskManager from .models import ( TargetType, @@ -61,12 +69,15 @@ __all__ = [ "BatchSubmitRequest", + "InMemoryTaskEventLog", "InMemoryTaskStore", + "NullTaskEventLog", "SQLAlchemyTaskStore", "TargetNotFoundError", "TargetType", "TaskContext", "TaskEvent", + "TaskEventLog", "TaskManager", "TaskInputValidationError", "TaskProgress", @@ -76,7 +87,9 @@ "TaskStore", "TaskSubmitRequest", "attach_execution_modes", + "available_event_log_providers", "bind_task_context", + "create_event_log", "create_task_router", "get_task_context", "install_task_progress_bridge", @@ -88,6 +101,7 @@ "make_ingestion_input_validator", "make_pipeline_dispatcher", "make_pipeline_input_validator", + "register_event_log_provider", "make_plugin_dispatcher", "make_plugin_input_validator", "report_stage", diff --git a/src/agentomatic/tasks/event_log.py b/src/agentomatic/tasks/event_log.py new file mode 100644 index 0000000..cd1207a --- /dev/null +++ b/src/agentomatic/tasks/event_log.py @@ -0,0 +1,303 @@ +"""Replay buffers for task progress events. + +A live SSE subscriber is not enough to make streaming reliable: the moment a +client's connection drops, every event emitted while it was away is gone, and +the client can only re-read the task's *current* state — never the steps it +missed. An event log keeps a bounded, ordered history per task so a +reconnecting client can resume exactly where it left off. + +The default implementation is in-process and bounded. Register a shared one +(Redis, a database, a log service) with +:func:`~agentomatic.tasks.event_log.register_event_log_provider` when tasks +must be resumable across workers or replicas — Agentomatic ships the +interface, not a vendor client. +""" + +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections import OrderedDict, deque +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from agentomatic.tasks.models import TaskEvent + +#: Events retained per task before the oldest are discarded. +DEFAULT_MAX_EVENTS_PER_TASK = 512 + +#: Tasks retained in the log before the least recently used are discarded. +DEFAULT_MAX_TASKS = 1024 + + +class TaskEventLog(ABC): + """Ordered, replayable history of the events emitted for each task.""" + + @abstractmethod + async def append(self, event: TaskEvent) -> None: + """Record one event. + + Args: + event: The event to retain. Its ``sequence`` is already assigned. + """ + + @abstractmethod + async def replay(self, task_id: str, *, after: int = 0) -> list[TaskEvent]: + """Return retained events for a task, oldest first. + + Args: + task_id: Task whose history to read. + after: Return only events with ``sequence`` greater than this. + + Returns: + The matching events in ascending sequence order. + """ + + @abstractmethod + async def earliest_sequence(self, task_id: str) -> int: + """Return the lowest sequence still retained for a task. + + A caller that asks to resume from before this point cannot be served a + complete history and must be told the stream is truncated. + + Args: + task_id: Task to inspect. + + Returns: + The lowest retained sequence, or 0 when nothing is retained. + """ + + @abstractmethod + async def latest_sequence(self, task_id: str) -> int: + """Return the highest sequence recorded for a task. + + A new subscriber uses this as its starting cursor, so it can resume + precisely even if it drops before the first live event arrives. + + Args: + task_id: Task to inspect. + + Returns: + The highest recorded sequence, or 0 when nothing is retained. + """ + + @abstractmethod + async def drop(self, task_id: str) -> None: + """Discard a task's retained history. + + Args: + task_id: Task whose history to forget. + """ + + +class InMemoryTaskEventLog(TaskEventLog): + """Bounded in-process event history. + + Retention is capped twice over — per task and across tasks — so a + long-running platform cannot accumulate history without limit. Eviction is + least-recently-appended. + + Args: + max_events_per_task: Events retained per task. + max_tasks: Number of tasks retained. + """ + + def __init__( + self, + *, + max_events_per_task: int = DEFAULT_MAX_EVENTS_PER_TASK, + max_tasks: int = DEFAULT_MAX_TASKS, + ) -> None: + self._max_events_per_task = max(1, max_events_per_task) + self._max_tasks = max(1, max_tasks) + self._events: OrderedDict[str, deque[TaskEvent]] = OrderedDict() + self._lock = asyncio.Lock() + + async def append(self, event: TaskEvent) -> None: + """Record one event, evicting the oldest history when full.""" + async with self._lock: + bucket = self._events.get(event.task_id) + if bucket is None: + bucket = deque(maxlen=self._max_events_per_task) + self._events[event.task_id] = bucket + bucket.append(event) + self._events.move_to_end(event.task_id) + while len(self._events) > self._max_tasks: + self._events.popitem(last=False) + + async def replay(self, task_id: str, *, after: int = 0) -> list[TaskEvent]: + """Return retained events for a task with sequence greater than ``after``.""" + async with self._lock: + bucket = self._events.get(task_id) + if not bucket: + return [] + return [evt for evt in bucket if evt.sequence > after] + + async def earliest_sequence(self, task_id: str) -> int: + """Return the lowest sequence still retained, or 0 when empty.""" + async with self._lock: + bucket = self._events.get(task_id) + if not bucket: + return 0 + return bucket[0].sequence + + async def latest_sequence(self, task_id: str) -> int: + """Return the highest recorded sequence, or 0 when empty.""" + async with self._lock: + bucket = self._events.get(task_id) + if not bucket: + return 0 + return bucket[-1].sequence + + async def drop(self, task_id: str) -> None: + """Discard a task's retained history.""" + async with self._lock: + self._events.pop(task_id, None) + + +class NullTaskEventLog(TaskEventLog): + """Event log that retains nothing. + + Use it to opt out of replay entirely. Reconnecting clients still receive + the task's current snapshot; they simply cannot recover missed steps. + """ + + async def append(self, event: TaskEvent) -> None: + """Discard the event.""" + + async def replay(self, task_id: str, *, after: int = 0) -> list[TaskEvent]: + """Return no history.""" + return [] + + async def earliest_sequence(self, task_id: str) -> int: + """Report that nothing is retained.""" + return 0 + + async def latest_sequence(self, task_id: str) -> int: + """Report that nothing is retained.""" + return 0 + + async def drop(self, task_id: str) -> None: + """Nothing is retained, so nothing is dropped.""" + + +#: Factories for named event-log backends, keyed by provider name. +EventLogFactory = Callable[..., TaskEventLog] + +_PROVIDERS: dict[str, EventLogFactory] = {} + + +def register_event_log_provider(name: str, factory: EventLogFactory) -> None: + """Register a named event-log backend. + + Cross-worker resumption needs a log every worker can read. Rather than + ship a client for one vendor, Agentomatic lets you register any backend:: + + register_event_log_provider("redis", lambda **cfg: MyRedisEventLog(**cfg)) + + Args: + name: Provider name used by :func:`create_event_log`. + factory: Callable returning a :class:`TaskEventLog`. + """ + _PROVIDERS[name] = factory + + +def available_event_log_providers() -> list[str]: + """Return the registered provider names, built-ins included. + + Returns: + Sorted provider names accepted by :func:`create_event_log`. + """ + return sorted({"memory", "null", *_PROVIDERS}) + + +def create_event_log(name: str = "memory", **options: Any) -> TaskEventLog: + """Build an event log by provider name. + + Args: + name: ``"memory"``, ``"null"``, or a registered provider. + **options: Passed to the provider factory. + + Returns: + A :class:`TaskEventLog`. + + Raises: + ValueError: When ``name`` is not a known provider. + """ + factory = _PROVIDERS.get(name) + if factory is not None: + return factory(**options) + if name == "memory": + return InMemoryTaskEventLog(**options) + if name == "null": + return NullTaskEventLog() + raise ValueError( + f"Unknown task event log provider '{name}'. " + f"Known: {', '.join(available_event_log_providers())}. " + "Register your own with register_event_log_provider()." + ) + + +class StreamResumption: + """The outcome of resolving a client's resume point. + + Attributes: + replay: Events the client missed, oldest first. + truncated: True when the requested resume point fell off the end of + the retained history, so ``replay`` does not cover every missed + event. Clients should treat the following snapshot as + authoritative rather than assume continuity. + cursor: Sequence the client is caught up to after ``replay``. + """ + + __slots__ = ("cursor", "replay", "truncated") + + def __init__(self, replay: list[TaskEvent], *, truncated: bool, cursor: int) -> None: + self.replay = replay + self.truncated = truncated + self.cursor = cursor + + +async def resolve_resumption(log: TaskEventLog, task_id: str, *, after: int) -> StreamResumption: + """Work out what a reconnecting client still needs. + + Args: + log: Event log to read. + task_id: Task being resumed. + after: Last sequence the client successfully processed (0 = new client). + + Returns: + A :class:`StreamResumption` describing the gap. + """ + if after <= 0: + return StreamResumption([], truncated=False, cursor=0) + replay = await log.replay(task_id, after=after) + earliest = await log.earliest_sequence(task_id) + # ``earliest > after + 1`` means the events between the client's cursor and + # the oldest retained event have already been evicted. + truncated = bool(earliest and earliest > after + 1) + cursor = replay[-1].sequence if replay else after + return StreamResumption(replay, truncated=truncated, cursor=cursor) + + +def parse_last_event_id(raw: str | None) -> int: + """Parse an SSE ``Last-Event-ID`` value into a sequence number. + + The header is client-supplied and echoed from a previous stream, so it is + treated as a hint: anything unparseable means "start from the beginning" + rather than an error. + + Args: + raw: Header value, or None. + + Returns: + The parsed sequence, or 0 when absent or malformed. + """ + if not raw: + return 0 + try: + value = int(str(raw).strip()) + except (TypeError, ValueError): + return 0 + return value if value > 0 else 0 diff --git a/src/agentomatic/tasks/manager.py b/src/agentomatic/tasks/manager.py index 1183083..72577da 100644 --- a/src/agentomatic/tasks/manager.py +++ b/src/agentomatic/tasks/manager.py @@ -28,6 +28,12 @@ from agentomatic.core.errors import client_safe_message from .context import TaskContext +from .event_log import ( + InMemoryTaskEventLog, + StreamResumption, + TaskEventLog, + resolve_resumption, +) from .models import ( TargetType, TaskEvent, @@ -65,8 +71,10 @@ def __init__( *, max_concurrency: int = 8, default_batch_concurrency: int = 4, + event_log: TaskEventLog | None = None, ) -> None: self.store = store or InMemoryTaskStore() + self.event_log: TaskEventLog = event_log or InMemoryTaskEventLog() self._semaphore = asyncio.Semaphore(max_concurrency) self._max_concurrency = max_concurrency self._default_batch_concurrency = default_batch_concurrency @@ -75,6 +83,7 @@ def __init__( self._running: dict[str, asyncio.Task[Any]] = {} self._cancel_requested: set[str] = set() self._subscribers: dict[str, list[asyncio.Queue[TaskEvent]]] = {} + self._sequences: dict[str, int] = {} # ------------------------------------------------------------------ # Registration & lifecycle @@ -285,6 +294,27 @@ def unsubscribe(self, task_id: str, queue: asyncio.Queue[TaskEvent]) -> None: if subs is not None and not subs: self._subscribers.pop(task_id, None) + async def resume(self, task_id: str, *, after: int) -> StreamResumption: + """Work out what a client reconnecting at ``after`` still needs. + + Args: + task_id: Task being resumed. + after: Last sequence the client processed (0 for a new client). + + Returns: + The events it missed, and whether the history was truncated. + """ + return await resolve_resumption(self.event_log, task_id, after=after) + + async def forget_events(self, task_id: str) -> None: + """Discard a task's retained event history. + + Args: + task_id: Task whose history to drop. + """ + self._sequences.pop(task_id, None) + await self.event_log.drop(task_id) + # ------------------------------------------------------------------ # Internal execution # ------------------------------------------------------------------ @@ -489,23 +519,38 @@ async def _finalize( async def _emit( self, record: TaskRecord, event: str, *, data: dict[str, Any] | None = None ) -> None: - """Broadcast an event to all subscribers of a task.""" - subs = self._subscribers.get(record.id) - if not subs: - return + """Record an event and broadcast it to the task's live subscribers. + + Recording happens whether or not anyone is currently listening. The + previous behaviour — build the event only when a subscriber existed — + made a dropped connection unrecoverable: the events emitted while the + client was away were never created, let alone retained. + """ + self._sequences[record.id] = self._sequences.get(record.id, 0) + 1 evt = TaskEvent( task_id=record.id, + sequence=self._sequences[record.id], event=event, status=record.status, progress=record.progress, data=data or {}, ) - for queue in list(subs): + try: + await self.event_log.append(evt) + except Exception as exc: # noqa: BLE001 - the log must never break a task + logger.warning(f"Task event log rejected an event for {record.id}: {exc}") + + for queue in list(self._subscribers.get(record.id, ())): try: queue.put_nowait(evt) except asyncio.QueueFull: # pragma: no cover - slow consumer + # The subscriber can still recover: it reconnects with its last + # sequence and replays the gap from the event log. logger.debug(f"Dropping task event for slow subscriber on {record.id}") + if record.status.is_terminal: + self._sequences.pop(record.id, None) + async def _fire_webhook(self, record: TaskRecord) -> None: """POST the final record to ``callback_url`` (best-effort).""" if not record.callback_url: diff --git a/src/agentomatic/tasks/models.py b/src/agentomatic/tasks/models.py index 641c5b9..26d519f 100644 --- a/src/agentomatic/tasks/models.py +++ b/src/agentomatic/tasks/models.py @@ -91,9 +91,21 @@ class TaskProgress(BaseModel): class TaskEvent(BaseModel): - """A single progress/status event emitted during task execution.""" + """A single progress/status event emitted during task execution. + + ``sequence`` is what makes a dropped stream recoverable: it numbers the + events of one task from 1 upwards with no gaps, so a client that + reconnects can say how far it got and be sent only what it missed. It is + the value published as the SSE ``id:`` field, which browsers echo back in + ``Last-Event-ID`` automatically. + """ task_id: str + sequence: int = Field( + default=0, + ge=0, + description="Monotonic per-task event number, starting at 1 (0 = unassigned).", + ) event: str = Field( description="Event name: queued|started|progress|succeeded|failed|cancelled." ) diff --git a/src/agentomatic/tasks/routes.py b/src/agentomatic/tasks/routes.py index da1a4f1..57159c1 100644 --- a/src/agentomatic/tasks/routes.py +++ b/src/agentomatic/tasks/routes.py @@ -7,7 +7,7 @@ GET {api_prefix}/tasks list/filter tasks GET {api_prefix}/tasks/{id} poll status + progress GET {api_prefix}/tasks/{id}/result fetch the result (409 if pending) - GET {api_prefix}/tasks/{id}/events live SSE progress stream + GET {api_prefix}/tasks/{id}/events resumable SSE progress stream POST {api_prefix}/tasks/{id}/cancel request cancellation DELETE {api_prefix}/tasks/{id} delete a terminal record """ @@ -19,10 +19,11 @@ from collections.abc import AsyncIterator from typing import Any, Literal -from fastapi import APIRouter, HTTPException, Query, Response +from fastapi import APIRouter, HTTPException, Query, Request, Response from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field +from .event_log import parse_last_event_id from .manager import TaskInputValidationError, TaskManager from .models import TargetType, TaskStatus @@ -168,32 +169,94 @@ async def delete_task(task_id: str) -> dict[str, Any]: if not record.status.is_terminal: raise HTTPException(status_code=409, detail="Cannot delete a running task") await manager.store.delete(task_id) + # The record is gone; its retained events would otherwise linger until + # evicted, and could be replayed for a task that no longer exists. + await manager.forget_events(task_id) return {"task_id": task_id, "deleted": True} - @router.get("/{task_id}/events", summary="Stream task progress (SSE)") - async def stream_events(task_id: str) -> StreamingResponse: - """Stream live progress events for a task as Server-Sent Events.""" + @router.get("/{task_id}/events", summary="Stream task progress (SSE, resumable)") + async def stream_events( + task_id: str, + request: Request, + since: int = Query( + default=0, + ge=0, + description=( + "Resume after this event sequence. Usually unnecessary — a browser " + "EventSource replays its Last-Event-ID header automatically." + ), + ), + ) -> StreamingResponse: + """Stream a task's progress as Server-Sent Events, resumably. + + Every frame carries an ``id:`` holding the event's monotonic sequence. + A client that loses the connection reconnects and is sent the events it + missed before the stream goes live again — from the ``Last-Event-ID`` + header a browser sends on its own, or from an explicit ``?since=``. + + Replay depth is bounded by the configured event log. When a resume + point has already been evicted the stream says so in a ``truncated`` + frame and follows it with the current snapshot, so a client is never + left believing it has a complete history. + """ record = await manager.get(task_id) if record is None: raise HTTPException(status_code=404, detail=f"Task '{task_id}' not found") + resume_from = since or parse_last_event_id(request.headers.get("last-event-id")) + async def event_stream() -> AsyncIterator[str]: + # Subscribe *before* replaying, so events emitted during replay are + # queued rather than lost in the gap between the two. queue = await manager.subscribe(task_id) try: - # Emit the current snapshot immediately. + resumption = await manager.resume(task_id, after=resume_from) + cursor = resumption.cursor + if resume_from <= 0: + # Read the high-water mark *before* the snapshot. An event + # landing between the two is then re-delivered from the + # live queue rather than skipped: duplicated state is + # harmless, a silently dropped event is not. + cursor = await manager.event_log.latest_sequence(task_id) + + if resumption.truncated: + yield _sse( + { + "event": "truncated", + "task_id": task_id, + "resumed_after": resume_from, + "detail": ( + "Events before this point are no longer retained; " + "treat the following snapshot as authoritative." + ), + } + ) + for evt in resumption.replay: + yield _sse(evt.model_dump(), event_id=evt.sequence) + + # A fresh subscriber (or a truncated one) needs the current + # state; a caught-up one already has it from the replay. snapshot = await manager.get(task_id) - if snapshot is not None: - yield _sse(snapshot.public_dict()) - if snapshot.status.is_terminal: - yield "data: [DONE]\n\n" - return + if snapshot is not None and (resume_from <= 0 or resumption.truncated): + # The snapshot carries the cursor it reflects, so a client + # that drops immediately after it can still resume exactly. + yield _sse(snapshot.public_dict(), event_id=cursor) + if snapshot is not None and snapshot.status.is_terminal: + yield "data: [DONE]\n\n" + return + while True: try: evt = await asyncio.wait_for(queue.get(), timeout=15.0) except TimeoutError: yield ": keep-alive\n\n" continue - yield _sse(evt.model_dump()) + # Replay and the live queue overlap by design; skip what + # the client has already been sent. + if evt.sequence and evt.sequence <= cursor: + continue + cursor = evt.sequence or cursor + yield _sse(evt.model_dump(), event_id=evt.sequence) if evt.status.is_terminal: yield "data: [DONE]\n\n" return @@ -219,6 +282,19 @@ def _links(task_id: str) -> dict[str, str]: } -def _sse(payload: dict[str, Any]) -> str: - """Format a payload as an SSE ``data:`` frame.""" - return f"data: {json.dumps(payload, default=str)}\n\n" +def _sse(payload: dict[str, Any], *, event_id: int | None = None) -> str: + """Format a payload as an SSE frame. + + Args: + payload: JSON-serialisable frame body. + event_id: Sequence to publish as the SSE ``id:`` field. Browsers echo + the last one they saw back in ``Last-Event-ID`` when they + reconnect, which is what makes the stream resumable for free. + + Returns: + The encoded frame. + """ + body = f"data: {json.dumps(payload, default=str)}\n\n" + if event_id: + return f"id: {event_id}\n{body}" + return body diff --git a/tests/test_a2a_streaming.py b/tests/test_a2a_streaming.py new file mode 100644 index 0000000..dec6274 --- /dev/null +++ b/tests/test_a2a_streaming.py @@ -0,0 +1,170 @@ +# pyright: reportMissingParameterType=none +# pyright: reportAttributeAccessIssue=none +"""Tests for the A2A task stream and the capabilities the agent card claims. + +Agentomatic exposed A2A task submission, status and cancel, but nothing an A2A +client could subscribe to: a caller that lost its connection mid-task could +only re-poll for the current state. These cover the resumable A2A stream and +the card that advertises it. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from agentomatic import AgentPlatform +from agentomatic.core.manifest import AgentManifest + +BASE = "/api/v1" + + +@dataclass +class _State: + """Minimal state for the streaming test agent.""" + + request: str = "" + output: dict[str, Any] = field(default_factory=dict) + + +async def _echo_fn(state: dict[str, Any]) -> dict[str, Any]: + """Return a canned response for the incoming query.""" + return {"response": f"echo: {state.get('current_query', '')}"} + + +@pytest.fixture +def client() -> Any: + """A platform with one agent and the task manager enabled.""" + platform = AgentPlatform( + agents_dir="/tmp/agentomatic_a2a_stream_test", + title="A2A Stream Test", + version="0.0.1", + enable_tasks=True, + ) + platform.register_agent( + manifest=AgentManifest( + name="echo", + slug="fn-echo", + description="Echo agent", + version="1.0.0", + ), + node_fn=_echo_fn, + ) + with TestClient(platform.build()) as test_client: + yield test_client + + +def _frames(body: str) -> list[dict[str, Any]]: + """Parse an SSE body into its JSON data frames.""" + out: list[dict[str, Any]] = [] + for line in body.splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + out.append(json.loads(line[len("data: ") :])) + return out + + +def _ids(body: str) -> list[int]: + """Extract SSE ``id:`` values from a body.""" + return [int(line[len("id: ") :]) for line in body.splitlines() if line.startswith("id: ")] + + +def _submit(client: Any) -> str: + """Submit an A2A task and return its id.""" + response = client.post( + f"{BASE}/echo/a2a/tasks", + json={"message": {"parts": [{"type": "text", "text": "stream me"}]}}, + ) + assert response.status_code == 200, response.text + return str(response.json()["task_id"]) + + +class TestA2AStream: + def test_stream_endpoint_exists(self, client: Any) -> None: + task_id = _submit(client) + response = client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events") + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + + def test_frames_use_a2a_task_states(self, client: Any) -> None: + """A2A clients expect submitted/working/completed, not internal names.""" + task_id = _submit(client) + body = client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events").text + + states = {frame.get("status") for frame in _frames(body)} + assert states, "no frames" + assert states <= {"submitted", "working", "completed", "failed", "canceled"} + + def test_stream_carries_event_ids(self, client: Any) -> None: + task_id = _submit(client) + body = client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events").text + assert _ids(body), "no id: lines — an A2A client cannot resume" + + def test_since_resumes_without_replaying_seen_events(self, client: Any) -> None: + task_id = _submit(client) + client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events").text + resumed = client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events?since=1").text + assert all(seq > 1 for seq in _ids(resumed)) + + def test_last_event_id_header_matches_since(self, client: Any) -> None: + task_id = _submit(client) + by_query = _ids(client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events?since=1").text) + by_header = _ids( + client.get( + f"{BASE}/echo/a2a/tasks/{task_id}/events", + headers={"Last-Event-ID": "1"}, + ).text + ) + assert by_header == by_query + + def test_terminal_stream_closes_with_done(self, client: Any) -> None: + task_id = _submit(client) + body = client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events").text + assert body.rstrip().endswith("data: [DONE]") + + def test_terminal_frame_carries_the_result(self, client: Any) -> None: + task_id = _submit(client) + body = client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events").text + terminal = [f for f in _frames(body) if f.get("status") == "completed"] + assert terminal, "no completed frame" + assert "result" in terminal[-1] + + def test_unknown_task_404s(self, client: Any) -> None: + response = client.get(f"{BASE}/echo/a2a/tasks/does-not-exist/events") + assert response.status_code == 404 + + +class TestAgentCardCapabilities: + def test_card_advertises_the_stream_endpoint(self, client: Any) -> None: + card = client.get(f"{BASE}/echo/card").json() + assert "a2a_events" in card["endpoints"] + + def test_card_reports_resumable_streams(self, client: Any) -> None: + card = client.get(f"{BASE}/echo/card").json() + assert card["capabilities"]["resumableStreams"] is True + assert card["capabilities"]["stateTransitionHistory"] is True + assert card["capabilities"]["pushNotifications"] is True + + def test_capabilities_are_not_claimed_without_a_task_manager(self) -> None: + """A card that claims a capability the deployment lacks sends clients + down a path that then 501s.""" + platform = AgentPlatform( + agents_dir="/tmp/agentomatic_a2a_stream_test_notasks", + title="No Tasks", + version="0.0.1", + enable_tasks=False, + ) + platform.register_agent( + manifest=AgentManifest(name="echo", slug="fn-echo", description="Echo"), + node_fn=_echo_fn, + ) + with TestClient(platform.build()) as test_client: + card = test_client.get(f"{BASE}/echo/card").json() + + assert card["capabilities"]["resumableStreams"] is False + assert card["capabilities"]["stateTransitionHistory"] is False + assert card["capabilities"]["pushNotifications"] is False + assert "a2a_events" not in card["endpoints"] diff --git a/tests/test_task_event_resume.py b/tests/test_task_event_resume.py new file mode 100644 index 0000000..cf0a5b8 --- /dev/null +++ b/tests/test_task_event_resume.py @@ -0,0 +1,423 @@ +# pyright: reportMissingParameterType=none +"""Tests for resumable task progress streams. + +Task events used to exist only while somebody was listening: ``_emit`` +returned early when a task had no subscribers, so everything emitted while a +client was disconnected was never created, let alone retained. A dropped +connection could therefore only be recovered as far as the task's *current* +state — every intermediate step was gone. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import Any + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from agentomatic.tasks.event_log import ( + InMemoryTaskEventLog, + NullTaskEventLog, + available_event_log_providers, + create_event_log, + parse_last_event_id, + register_event_log_provider, + resolve_resumption, +) +from agentomatic.tasks.manager import TaskManager +from agentomatic.tasks.models import TargetType, TaskEvent, TaskStatus +from agentomatic.tasks.routes import create_task_router + + +def _stepping_dispatcher(steps: int = 5, delay: float = 0.02): + """Return a dispatcher that reports ``steps`` progress events.""" + + async def dispatch(target: str, payload: Any, ctx: Any) -> dict[str, Any]: + for i in range(1, steps + 1): + await ctx.report( + percent=i * (100 / steps), message=f"step {i}", current=i, total=steps + ) + await asyncio.sleep(delay) + return {"steps": steps} + + return dispatch + + +def _build_app(manager: TaskManager) -> FastAPI: + """Mount the task router for ``manager`` on a bare app.""" + app = FastAPI() + app.include_router(create_task_router(manager), prefix="/api/v1/tasks") + return app + + +def _frames(body: str) -> list[dict[str, Any]]: + """Parse an SSE response body into its JSON data frames.""" + out: list[dict[str, Any]] = [] + for block in body.split("\n\n"): + for line in block.splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + out.append(json.loads(line[len("data: ") :])) + return out + + +def _ids(body: str) -> list[int]: + """Extract the SSE ``id:`` values from a response body.""" + return [int(line[len("id: ") :]) for line in body.splitlines() if line.startswith("id: ")] + + +# ===================================================================== +# Sequencing +# ===================================================================== + + +class TestEventSequencing: + @pytest.mark.asyncio + async def test_events_are_recorded_without_any_subscriber(self) -> None: + """The whole point: nobody listening must not mean nothing retained.""" + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(3)) + + record = await manager.submit(TargetType.AGENT, "x", input={}, mode="async") + await asyncio.sleep(0.4) + + events = await manager.event_log.replay(record.id) + assert events, "no events retained for an unobserved task" + assert events[-1].status.is_terminal + + @pytest.mark.asyncio + async def test_sequences_are_gapless_and_start_at_one(self) -> None: + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(4)) + + record = await manager.submit(TargetType.AGENT, "x", input={}, mode="async") + await asyncio.sleep(0.5) + + sequences = [e.sequence for e in await manager.event_log.replay(record.id)] + assert sequences == list(range(1, len(sequences) + 1)) + + @pytest.mark.asyncio + async def test_sequences_are_independent_per_task(self) -> None: + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(2)) + + first = await manager.submit(TargetType.AGENT, "x", input={}, mode="async") + second = await manager.submit(TargetType.AGENT, "y", input={}, mode="async") + await asyncio.sleep(0.5) + + for record in (first, second): + sequences = [e.sequence for e in await manager.event_log.replay(record.id)] + assert sequences[0] == 1, "each task numbers its own events from 1" + + @pytest.mark.asyncio + async def test_reconnect_replays_exactly_the_gap(self) -> None: + """No missed event, and nothing delivered twice.""" + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(6, delay=0.03)) + + record = await manager.submit(TargetType.AGENT, "x", input={}, mode="async") + await asyncio.sleep(0.1) + + seen = [e.sequence for e in await manager.event_log.replay(record.id)] + cursor = seen[-1] + await asyncio.sleep(0.6) + + resumed = await manager.resume(record.id, after=cursor) + replayed = [e.sequence for e in resumed.replay] + everything = [e.sequence for e in await manager.event_log.replay(record.id)] + + assert seen + replayed == everything + assert not resumed.truncated + + +# ===================================================================== +# The SSE surface +# ===================================================================== + + +class TestResumableSSE: + def test_new_subscriber_gets_a_snapshot_carrying_a_cursor(self) -> None: + """Without an id on the snapshot, a client that drops straight after it + has nothing to resume from and must restart the whole stream.""" + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(3)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={"target_type": "agent", "target": "x", "input": {}, "wait": True}, + ).json() + task_id = submitted["id"] + body = client.get(f"/api/v1/tasks/{task_id}/events").text + + ids = _ids(body) + assert ids, "no id: lines — EventSource cannot resume without them" + assert ids[-1] == asyncio.run(manager.event_log.latest_sequence(task_id)) + assert _frames(body)[-1]["id"] == task_id, "a new client gets the full record" + + def test_since_replays_exactly_the_missed_events(self) -> None: + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(4)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={"target_type": "agent", "target": "x", "input": {}, "wait": True}, + ).json() + task_id = submitted["id"] + resumed = client.get(f"/api/v1/tasks/{task_id}/events?since=2").text + + recorded = [e.sequence for e in asyncio.run(manager.event_log.replay(task_id))] + assert _ids(resumed) == [s for s in recorded if s > 2] + assert all(seq > 2 for seq in _ids(resumed)), "re-sent already-seen events" + + def test_last_event_id_header_matches_the_since_parameter(self) -> None: + """A browser EventSource resumes with this header and no query param.""" + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(4)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={"target_type": "agent", "target": "x", "input": {}, "wait": True}, + ).json() + task_id = submitted["id"] + + by_query = _ids(client.get(f"/api/v1/tasks/{task_id}/events?since=2").text) + by_header = _ids( + client.get( + f"/api/v1/tasks/{task_id}/events", + headers={"Last-Event-ID": "2"}, + ).text + ) + + assert by_header == by_query + assert by_header, "the header was ignored" + + def test_malformed_last_event_id_is_treated_as_a_new_subscriber(self) -> None: + """It must never be read as a valid cursor that silently skips events.""" + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(2)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={"target_type": "agent", "target": "x", "input": {}, "wait": True}, + ).json() + task_id = submitted["id"] + body = client.get( + f"/api/v1/tasks/{task_id}/events", + headers={"Last-Event-ID": "not-a-number"}, + ).text + fresh = client.get(f"/api/v1/tasks/{task_id}/events").text + + assert _ids(body) == _ids(fresh) + # A new subscriber's snapshot is the complete current state, so nothing + # is lost by discarding the unusable cursor. + assert _frames(body)[-1]["status"] == "succeeded" + + def test_truncated_history_is_announced(self) -> None: + """Silently serving a partial history would be the dangerous failure.""" + manager = TaskManager(event_log=InMemoryTaskEventLog(max_events_per_task=2)) + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(5)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={"target_type": "agent", "target": "x", "input": {}, "wait": True}, + ).json() + body = client.get(f"/api/v1/tasks/{submitted['id']}/events?since=1").text + + assert any(f.get("event") == "truncated" for f in _frames(body)) + + def test_terminal_task_stream_closes_with_done(self) -> None: + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(2)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={"target_type": "agent", "target": "x", "input": {}, "wait": True}, + ).json() + body = client.get(f"/api/v1/tasks/{submitted['id']}/events").text + + assert body.rstrip().endswith("data: [DONE]") + + def test_unknown_task_still_404s(self) -> None: + manager = TaskManager() + with TestClient(_build_app(manager)) as client: + assert client.get("/api/v1/tasks/nope/events").status_code == 404 + + def test_deleting_a_task_forgets_its_events(self) -> None: + manager = TaskManager() + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(2)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={"target_type": "agent", "target": "x", "input": {}, "wait": True}, + ).json() + task_id = submitted["id"] + assert client.delete(f"/api/v1/tasks/{task_id}").status_code == 200 + + assert client.get(f"/api/v1/tasks/{task_id}/events").status_code == 404 + + +# ===================================================================== +# The log itself +# ===================================================================== + + +def _event(task_id: str, sequence: int) -> TaskEvent: + """Build a minimal event for log-level tests.""" + return TaskEvent( + task_id=task_id, sequence=sequence, event="progress", status=TaskStatus.RUNNING + ) + + +class TestInMemoryEventLog: + @pytest.mark.asyncio + async def test_replay_filters_by_sequence(self) -> None: + log = InMemoryTaskEventLog() + for seq in range(1, 5): + await log.append(_event("t", seq)) + assert [e.sequence for e in await log.replay("t", after=2)] == [3, 4] + + @pytest.mark.asyncio + async def test_per_task_retention_is_bounded(self) -> None: + log = InMemoryTaskEventLog(max_events_per_task=3) + for seq in range(1, 11): + await log.append(_event("t", seq)) + assert [e.sequence for e in await log.replay("t")] == [8, 9, 10] + + @pytest.mark.asyncio + async def test_task_count_is_bounded(self) -> None: + log = InMemoryTaskEventLog(max_tasks=2) + for name in ("a", "b", "c"): + await log.append(_event(name, 1)) + assert await log.replay("a") == [], "oldest task should have been evicted" + assert await log.replay("c") + + @pytest.mark.asyncio + async def test_drop_forgets_one_task_only(self) -> None: + log = InMemoryTaskEventLog() + await log.append(_event("a", 1)) + await log.append(_event("b", 1)) + await log.drop("a") + assert await log.replay("a") == [] + assert await log.replay("b") + + @pytest.mark.asyncio + async def test_earliest_sequence_reports_retention_floor(self) -> None: + log = InMemoryTaskEventLog(max_events_per_task=2) + for seq in range(1, 6): + await log.append(_event("t", seq)) + assert await log.earliest_sequence("t") == 4 + assert await log.earliest_sequence("missing") == 0 + + +class TestResumptionResolution: + @pytest.mark.asyncio + async def test_new_client_gets_no_replay(self) -> None: + log = InMemoryTaskEventLog() + await log.append(_event("t", 1)) + resumed = await resolve_resumption(log, "t", after=0) + assert resumed.replay == [] + assert not resumed.truncated + + @pytest.mark.asyncio + async def test_evicted_gap_is_flagged_truncated(self) -> None: + log = InMemoryTaskEventLog(max_events_per_task=2) + for seq in range(1, 6): + await log.append(_event("t", seq)) + resumed = await resolve_resumption(log, "t", after=1) + assert resumed.truncated + + @pytest.mark.asyncio + async def test_contiguous_history_is_not_truncated(self) -> None: + log = InMemoryTaskEventLog() + for seq in range(1, 6): + await log.append(_event("t", seq)) + resumed = await resolve_resumption(log, "t", after=3) + assert not resumed.truncated + assert resumed.cursor == 5 + + +class TestNullEventLog: + @pytest.mark.asyncio + async def test_retains_nothing_but_never_fails(self) -> None: + manager = TaskManager(event_log=NullTaskEventLog()) + manager.register_dispatcher(TargetType.AGENT, _stepping_dispatcher(2)) + record = await manager.submit(TargetType.AGENT, "x", input={}, mode="async") + await asyncio.sleep(0.3) + assert await manager.event_log.replay(record.id) == [] + + +class TestProviderRegistry: + def test_builtin_providers_resolve(self) -> None: + assert isinstance(create_event_log("memory"), InMemoryTaskEventLog) + assert isinstance(create_event_log("null"), NullTaskEventLog) + + def test_unknown_provider_names_the_registry(self) -> None: + with pytest.raises(ValueError, match="register_event_log_provider"): + create_event_log("redis") + + def test_custom_provider_can_be_registered(self) -> None: + register_event_log_provider("test-shared", lambda **_: NullTaskEventLog()) + assert "test-shared" in available_event_log_providers() + assert isinstance(create_event_log("test-shared"), NullTaskEventLog) + + +class TestLastEventIdParsing: + @pytest.mark.parametrize( + ("raw", "expected"), + [("5", 5), (" 7 ", 7), (None, 0), ("", 0), ("abc", 0), ("-3", 0), ("0", 0)], + ) + def test_parses_defensively(self, raw: str | None, expected: int) -> None: + assert parse_last_event_id(raw) == expected + + +# ===================================================================== +# Coverage across every resource type +# ===================================================================== + + +class TestEveryTargetTypeResumes: + """Agents, pipelines, plugins, endpoints and ingestion share one task + manager, so the resumable stream must work identically for all of them — + a pipeline's per-step progress is exactly the case where losing the + connection hurts most.""" + + @pytest.mark.parametrize("target_type", list(TargetType)) + def test_stream_is_resumable(self, target_type: TargetType) -> None: + manager = TaskManager() + manager.register_dispatcher(target_type, _stepping_dispatcher(4)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={ + "target_type": target_type.value, + "target": f"{target_type.value}-x", + "input": {}, + "wait": True, + }, + ).json() + task_id = submitted["id"] + resumed = client.get(f"/api/v1/tasks/{task_id}/events?since=2").text + + recorded = [e.sequence for e in asyncio.run(manager.event_log.replay(task_id))] + assert recorded == list(range(1, len(recorded) + 1)), "gapless sequence required" + assert _ids(resumed) == [s for s in recorded if s > 2] + + @pytest.mark.parametrize("target_type", list(TargetType)) + def test_events_survive_having_no_listener(self, target_type: TargetType) -> None: + manager = TaskManager() + manager.register_dispatcher(target_type, _stepping_dispatcher(3)) + with TestClient(_build_app(manager)) as client: + submitted = client.post( + "/api/v1/tasks", + json={ + "target_type": target_type.value, + "target": f"{target_type.value}-x", + "input": {}, + "wait": True, + }, + ).json() + + # Nothing ever subscribed, yet the history is there to replay. + assert asyncio.run(manager.event_log.replay(submitted["id"])) From 75af79c039a9b4fe9b30b2577aee7656d6eafae2 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 15:10:09 +0000 Subject: [PATCH 4/9] feat(streaming): number and retain frames on every SSE surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The task stream became resumable, but the other streams still wrote frames straight to the socket and kept nothing. A client that dropped mid-answer on ``/invoke/stream`` lost the partial response entirely, with no id to reconnect against and nothing to reconnect to. Add agentomatic.streaming: one SSE encoder, a stream identity, and a bounded replay buffer with the same per-stream and cross-stream retention caps as the task event log. - ``/invoke/stream`` numbers every frame, records it, and returns the identity as ``X-Stream-Id``. - ``GET /invoke/stream/{stream_id}`` re-sends the frames already produced, honouring ``?since=`` and ``Last-Event-ID`` and announcing a truncated resume point rather than serving a partial history as a complete one. - Studio's run and resume streams are numbered through the same helper, so a reloaded debug UI is not forced to re-run an expensive agent. Replay here recovers what was *produced*; it does not restart a run, because the run is bound to the request that was cancelled. The docs say so plainly and point at the durable path — submit as a task, follow the task events — rather than implying a guarantee this cannot make. Also fix a pre-existing UP038 in router_factory that the pinned ruff 0.8.6 in .pre-commit-config.yaml flags but the repo's own ruff 0.15 no longer does. It already fails on main and blocked the gate as soon as a PR touched the file. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- docs/architecture/api-reference.md | 36 ++++ src/agentomatic/core/router_factory.py | 102 +++++++++- src/agentomatic/streaming/__init__.py | 27 +++ src/agentomatic/streaming/replay.py | 258 +++++++++++++++++++++++++ src/agentomatic/studio/router.py | 16 +- tests/test_stream_replay.py | 210 ++++++++++++++++++++ 6 files changed, 636 insertions(+), 13 deletions(-) create mode 100644 src/agentomatic/streaming/__init__.py create mode 100644 src/agentomatic/streaming/replay.py create mode 100644 tests/test_stream_replay.py diff --git a/docs/architecture/api-reference.md b/docs/architecture/api-reference.md index 12739dc..b88a8d6 100644 --- a/docs/architecture/api-reference.md +++ b/docs/architecture/api-reference.md @@ -213,10 +213,13 @@ curl -X POST http://localhost:8000/api/v1/my_agent/invoke/stream \ **SSE Response:** ``` +id: 1 data: {"response": "", "steps_taken": ["retrieve"]} +id: 2 data: {"response": "Quantum computing uses...", "steps_taken": ["retrieve", "generate"]} +id: 3 data: [DONE] ``` @@ -226,10 +229,43 @@ data: [DONE] |---|---| | `Content-Type` | `text/event-stream` | | `X-Agent` | Agent name | +| `X-Stream-Id` | Identity to replay this stream with | | `Cache-Control` | `no-cache` | --- +#### `GET /invoke/stream/{stream_id}` + +Re-send the frames a previous `/invoke/stream` already produced, using the +`X-Stream-Id` that response carried. Accepts `?since=` or `Last-Event-ID` and +returns only the frames after that point; `404` once the stream is no longer +retained. + +```bash +curl -N "http://localhost:8000/api/v1/my_agent/invoke/stream/stream_a1b2c3d4?since=2" +``` + +!!! warning "Replay is not resumption" + + This returns what was already **produced**. It does not restart the run: + the agent execution is bound to the request that was cancelled, so a + client that disconnects mid-answer loses the *generation*, and replay + recovers only the part that had been emitted. + + For a stream that keeps working across a dropped connection, submit the + work as a task and follow the task's events — that path is resumable end + to end: + + ```bash + TASK=$(curl -sX POST http://localhost:8000/api/v1/my_agent/invoke/async \ + -H 'Content-Type: application/json' -d '{"query": "..."}' | jq -r .id) + curl -N "http://localhost:8000/api/v1/tasks/$TASK/events" + ``` + + See [Tasks & Execution Modes](../guide/tasks.md#resuming-a-dropped-stream). + +--- + #### `POST /chat` Session-aware chat with automatic conversation memory. When a thread store is configured, this endpoint automatically loads history, invokes the agent, and persists messages. diff --git a/src/agentomatic/core/router_factory.py b/src/agentomatic/core/router_factory.py index be96c39..00f9538 100644 --- a/src/agentomatic/core/router_factory.py +++ b/src/agentomatic/core/router_factory.py @@ -17,6 +17,7 @@ from agentomatic.core.agent_invoke import build_invoke_state, invoke_registered_agent from agentomatic.core.errors import client_safe_detail, client_safe_message from agentomatic.langchain_adapter import dict_to_messages, json_default, to_jsonable +from agentomatic.streaming import get_replay_buffer, new_stream_id, sse_frame # --------------------------------------------------------------------------- # Request / Response Models @@ -163,7 +164,7 @@ def coerce_agent_invoke_payload( return raw_response, output, context # Explicit structured response value. - if isinstance(raw_response, (dict, list)): + if isinstance(raw_response, dict | list): jsonable_response = to_jsonable(raw_response) return _json_dumps(jsonable_response), jsonable_response, context @@ -768,6 +769,20 @@ async def invoke_stream(request: Any) -> StreamingResponse: ) stream_t0 = time.perf_counter() + stream_id = new_stream_id() + replay_buffer = get_replay_buffer() + + async def emit(payload: Any) -> str: + """Record one frame and render it with its sequence as the SSE id. + + Recording is what lets a client that drops mid-answer recover the + part that was already produced, rather than losing the exchange. + """ + body = ( + payload if isinstance(payload, str) else json.dumps(payload, default=json_default) + ) + sequence = await replay_buffer.record(stream_id, body) + return sse_frame(body, event_id=sequence) async def event_stream(): """Yield SSE frames from graph or node execution.""" @@ -786,30 +801,30 @@ async def event_stream(): state_obj = class_agent.input_to_state(input_data) graph = class_agent.graph async for event in graph.astream(state_obj): - yield f"data: {json.dumps(event, default=json_default)}\n\n" + yield await emit(event) result = class_agent.state_to_output(state_obj) if hasattr(class_agent, "_graph") and class_agent._graph is not None: class_agent._traces.append(class_agent._graph.last_trace) finally: class_agent._end_request_prompt() - yield f"data: {json.dumps(result, default=json_default)}\n\n" + yield await emit(result) collected_response = result.get("response", "") collected_output = result elif agent.graph_fn: graph = agent.graph_fn() async for event in graph.astream(state): - yield f"data: {json.dumps(event, default=json_default)}\n\n" + yield await emit(event) # Collect response for persistence if isinstance(event, dict) and "response" in event: collected_response = event["response"] collected_output = event elif agent.node_fn: result = await agent.node_fn(state) - yield f"data: {json.dumps(result, default=json_default)}\n\n" + yield await emit(result) if isinstance(result, dict): collected_response = result.get("response", "") collected_output = result - yield "data: [DONE]\n\n" + yield await emit("[DONE]") # Persist the turn after streaming completes if memory_mgr and thread_id and query and collected_response: @@ -850,7 +865,14 @@ async def event_stream(): status="suspended", error=exc.message, ) - yield f"data: {json.dumps({'status': 'suspended', 'approval_id': exc.approval_id, 'node_name': exc.node_name, 'message': exc.message})}\n\n" + yield await emit( + { + "status": "suspended", + "approval_id": exc.approval_id, + "node_name": exc.node_name, + "message": exc.message, + } + ) except Exception as exc: if log_recorder is not None: await log_recorder.record( @@ -864,12 +886,18 @@ async def event_stream(): error=str(exc), ) safe = client_safe_detail(exc, context="Agent streaming failed") - yield f"data: {json.dumps(safe)}\n\n" + yield await emit(safe) return StreamingResponse( event_stream(), media_type="text/event-stream", - headers={"X-Agent": agent_name, "Cache-Control": "no-cache"}, + headers={ + "X-Agent": agent_name, + "Cache-Control": "no-cache", + # A client that drops mid-answer needs this to ask for the + # frames it missed. + "X-Stream-Id": stream_id, + }, ) # Apply annotations dynamically @@ -890,6 +918,62 @@ async def event_stream(): summary="Invoke agent with SSE streaming", ) + @router.get("/invoke/stream/{stream_id}", summary="Replay a streamed response") + async def replay_invoke_stream( + stream_id: str, + request: Request, + since: int = 0, + ) -> StreamingResponse: + """Re-send the frames a previous ``/invoke/stream`` already produced. + + A dropped SSE connection used to lose the whole partial answer: the + frames were written straight to the socket and kept nowhere. They are + now retained under the ``X-Stream-Id`` the original response carried, + so a client can come back and collect what it missed. + + This replays what was *produced*; it does not restart the run, because + the run is bound to the request that was cancelled. For a stream that + keeps working across a disconnect, submit the work as a task + (``/invoke/async``) and follow ``/api/v1/tasks/{id}/events``, which is + resumable end to end. + """ + from agentomatic.tasks.event_log import parse_last_event_id + + buffer = get_replay_buffer() + if not await buffer.knows(stream_id): + raise HTTPException( + 404, + f"Stream '{stream_id}' is unknown or its frames are no longer retained.", + ) + + resume_from = since or parse_last_event_id(request.headers.get("last-event-id")) + frames = await buffer.replay(stream_id, after=resume_from) + earliest = await buffer.earliest_sequence(stream_id) + truncated = bool(resume_from and earliest and earliest > resume_from + 1) + + async def replay_stream() -> Any: + if truncated: + yield sse_frame( + { + "event": "truncated", + "stream_id": stream_id, + "resumed_after": resume_from, + "detail": "Earlier frames are no longer retained.", + } + ) + for frame in frames: + yield sse_frame(frame.data, event_id=frame.sequence) + + return StreamingResponse( + replay_stream(), + media_type="text/event-stream", + headers={ + "X-Agent": agent_name, + "Cache-Control": "no-cache", + "X-Stream-Id": stream_id, + }, + ) + # ── POST /invoke/async + /invoke/batch (via the task system) ── if task_manager is not None: from agentomatic.tasks.models import TargetType diff --git a/src/agentomatic/streaming/__init__.py b/src/agentomatic/streaming/__init__.py new file mode 100644 index 0000000..ec55c34 --- /dev/null +++ b/src/agentomatic/streaming/__init__.py @@ -0,0 +1,27 @@ +"""Shared plumbing for Server-Sent Events streams.""" + +from __future__ import annotations + +from .replay import ( + DEFAULT_MAX_FRAMES_PER_STREAM, + DEFAULT_MAX_STREAMS, + StreamFrame, + StreamReplayBuffer, + get_replay_buffer, + new_stream_id, + numbered_stream, + set_replay_buffer, + sse_frame, +) + +__all__ = [ + "DEFAULT_MAX_FRAMES_PER_STREAM", + "DEFAULT_MAX_STREAMS", + "StreamFrame", + "StreamReplayBuffer", + "get_replay_buffer", + "new_stream_id", + "numbered_stream", + "set_replay_buffer", + "sse_frame", +] diff --git a/src/agentomatic/streaming/replay.py b/src/agentomatic/streaming/replay.py new file mode 100644 index 0000000..dc8d1dd --- /dev/null +++ b/src/agentomatic/streaming/replay.py @@ -0,0 +1,258 @@ +"""Sequenced SSE frames and a bounded replay buffer for raw streams. + +The task subsystem has its own event log, because a task's events outlive the +connection watching them. Raw response streams — ``/invoke/stream``, Studio's +run stream — are different: the work is bound to the request, so a client that +disconnects loses the *run*, not just the view of it. + +That still leaves something worth keeping. The frames already produced are the +partial answer, and a client that reconnects can be handed them instead of +being told the whole exchange is gone. This module provides the numbering and +the bounded buffer that makes that possible, plus the one SSE encoder the +streaming endpoints share. + +For a stream that must survive a dropped connection *and* keep working, run +the agent as a task (``/invoke/async``) and follow +``/api/v1/tasks/{id}/events``, which is resumable end to end. +""" + +from __future__ import annotations + +import asyncio +import json +import time +import uuid +from collections import OrderedDict, deque +from typing import Any + +#: Frames retained per stream before the oldest are discarded. +DEFAULT_MAX_FRAMES_PER_STREAM = 512 + +#: Streams retained before the least recently used are discarded. +DEFAULT_MAX_STREAMS = 256 + + +def new_stream_id() -> str: + """Return an identifier for one response stream. + + Returns: + A short, unique stream id. + """ + return f"stream_{uuid.uuid4().hex[:16]}" + + +def sse_frame(payload: Any, *, event_id: int | None = None) -> str: + """Encode one Server-Sent Events frame. + + Args: + payload: A JSON-serialisable body, or a pre-rendered ``str``. + event_id: Sequence to publish as the SSE ``id:`` field. Browsers echo + the last id they saw back in ``Last-Event-ID`` when reconnecting. + + Returns: + The encoded frame, terminated by a blank line. + """ + from agentomatic.langchain_adapter import json_default + + body = payload if isinstance(payload, str) else json.dumps(payload, default=json_default) + frame = f"data: {body}\n\n" + if event_id: + return f"id: {event_id}\n{frame}" + return frame + + +async def numbered_stream( + source: Any, + *, + stream_id: str, + buffer: StreamReplayBuffer | None = None, +) -> Any: + """Record and number frames from a generator that already emits SSE text. + + Endpoints that build their own ``data: ...`` frames can be made replayable + without rewriting them: this wraps the generator, stores each frame's + payload against ``stream_id``, and re-emits it carrying the sequence as + the SSE ``id:``. Non-data chunks (comments, keep-alives) pass through + untouched and are not numbered. + + Args: + source: Async iterator of SSE-encoded strings. + stream_id: Identity the frames are retained under. + buffer: Replay buffer to use (default: the process-wide one). + + Yields: + The same frames, numbered. + """ + store = buffer if buffer is not None else get_replay_buffer() + async for chunk in source: + payload = _sse_payload(chunk) + if payload is None: + yield chunk + continue + sequence = await store.record(stream_id, payload) + yield sse_frame(payload, event_id=sequence) + + +def _sse_payload(chunk: str) -> str | None: + """Extract the body of a single-frame ``data:`` chunk. + + Args: + chunk: A chunk yielded by an SSE generator. + + Returns: + The payload, or None when the chunk is not a lone data frame (a + comment, a keep-alive, or a multi-line frame this must not mangle). + """ + if not isinstance(chunk, str) or not chunk.startswith("data: "): + return None + body = chunk.removeprefix("data: ").rstrip("\n") + return None if "\n" in body else body + + +class StreamFrame: + """One recorded frame of a response stream. + + Attributes: + sequence: Position in the stream, from 1, with no gaps. + data: The frame body exactly as it was sent. + timestamp: When the frame was recorded. + """ + + __slots__ = ("data", "sequence", "timestamp") + + def __init__(self, sequence: int, data: str) -> None: + self.sequence = sequence + self.data = data + self.timestamp = time.time() + + def as_dict(self) -> dict[str, Any]: + """Return the frame as a JSON-serialisable mapping. + + Returns: + The frame's sequence, body and timestamp. + """ + return {"sequence": self.sequence, "data": self.data, "timestamp": self.timestamp} + + +class StreamReplayBuffer: + """Bounded, in-process history of the frames each stream produced. + + Retention is capped per stream and across streams, so a long-running + process cannot accumulate transcripts without limit. + + Args: + max_frames_per_stream: Frames retained per stream. + max_streams: Number of streams retained. + """ + + def __init__( + self, + *, + max_frames_per_stream: int = DEFAULT_MAX_FRAMES_PER_STREAM, + max_streams: int = DEFAULT_MAX_STREAMS, + ) -> None: + self._max_frames = max(1, max_frames_per_stream) + self._max_streams = max(1, max_streams) + self._streams: OrderedDict[str, deque[StreamFrame]] = OrderedDict() + self._sequences: dict[str, int] = {} + self._lock = asyncio.Lock() + + async def record(self, stream_id: str, data: str) -> int: + """Append one frame and return the sequence assigned to it. + + Args: + stream_id: Stream the frame belongs to. + data: The frame body as sent to the client. + + Returns: + The frame's sequence number. + """ + async with self._lock: + sequence = self._sequences.get(stream_id, 0) + 1 + self._sequences[stream_id] = sequence + bucket = self._streams.get(stream_id) + if bucket is None: + bucket = deque(maxlen=self._max_frames) + self._streams[stream_id] = bucket + bucket.append(StreamFrame(sequence, data)) + self._streams.move_to_end(stream_id) + while len(self._streams) > self._max_streams: + evicted, _ = self._streams.popitem(last=False) + self._sequences.pop(evicted, None) + return sequence + + async def replay(self, stream_id: str, *, after: int = 0) -> list[StreamFrame]: + """Return retained frames with sequence greater than ``after``. + + Args: + stream_id: Stream to read. + after: Exclusive lower bound on sequence. + + Returns: + The matching frames, oldest first. + """ + async with self._lock: + bucket = self._streams.get(stream_id) + if not bucket: + return [] + return [frame for frame in bucket if frame.sequence > after] + + async def earliest_sequence(self, stream_id: str) -> int: + """Return the lowest sequence still retained, or 0 when empty. + + Args: + stream_id: Stream to inspect. + + Returns: + The retention floor. + """ + async with self._lock: + bucket = self._streams.get(stream_id) + return bucket[0].sequence if bucket else 0 + + async def knows(self, stream_id: str) -> bool: + """Report whether any frames are retained for a stream. + + Args: + stream_id: Stream to check. + + Returns: + True when the stream has retained frames. + """ + async with self._lock: + return bool(self._streams.get(stream_id)) + + async def drop(self, stream_id: str) -> None: + """Discard a stream's retained frames. + + Args: + stream_id: Stream to forget. + """ + async with self._lock: + self._streams.pop(stream_id, None) + self._sequences.pop(stream_id, None) + + +_DEFAULT_BUFFER: StreamReplayBuffer | None = None + + +def get_replay_buffer() -> StreamReplayBuffer: + """Return the process-wide replay buffer, creating it on first use. + + Returns: + The shared :class:`StreamReplayBuffer`. + """ + global _DEFAULT_BUFFER + if _DEFAULT_BUFFER is None: + _DEFAULT_BUFFER = StreamReplayBuffer() + return _DEFAULT_BUFFER + + +def set_replay_buffer(buffer: StreamReplayBuffer | None) -> None: + """Replace the process-wide replay buffer. + + Args: + buffer: The buffer to install, or None to reset to a fresh default. + """ + global _DEFAULT_BUFFER + _DEFAULT_BUFFER = buffer diff --git a/src/agentomatic/studio/router.py b/src/agentomatic/studio/router.py index 0eceff9..19560a0 100644 --- a/src/agentomatic/studio/router.py +++ b/src/agentomatic/studio/router.py @@ -15,6 +15,7 @@ from agentomatic.core.errors import client_safe_detail from agentomatic.core.schemas import SchemaValidator, load_schema_models +from agentomatic.streaming import numbered_stream from agentomatic.studio.adapters import resolve_adapter from agentomatic.studio.models import ( StudioAgentInfo, @@ -350,8 +351,11 @@ async def stream_run(name: str, request: StudioRunRequest) -> StreamingResponse: state = _build_studio_state(request, thread_id) return StreamingResponse( - tracker.execute_with_adapter( - adapter, state, run.id, thread_id, request.checkpoint_id, request.breakpoints + numbered_stream( + tracker.execute_with_adapter( + adapter, state, run.id, thread_id, request.checkpoint_id, request.breakpoints + ), + stream_id=run.id, ), media_type="text/event-stream", headers={ @@ -529,9 +533,13 @@ async def _stream() -> AsyncGenerator[str, None]: yield f"data: {error_data}\n\n" return StreamingResponse( - _stream(), + numbered_stream(_stream(), stream_id=thread_id), media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + "X-Stream-Id": thread_id, + }, ) @router.get( diff --git a/tests/test_stream_replay.py b/tests/test_stream_replay.py new file mode 100644 index 0000000..4e207e2 --- /dev/null +++ b/tests/test_stream_replay.py @@ -0,0 +1,210 @@ +# pyright: reportMissingParameterType=none +# pyright: reportAttributeAccessIssue=none +"""Tests for replaying a dropped ``/invoke/stream`` response. + +Frames used to be written straight to the socket and kept nowhere, so a +client that lost the connection mid-answer lost the partial answer with it. +They are now numbered and retained under the ``X-Stream-Id`` the response +carries, so the produced frames can be collected afterwards. + +Replay returns what was *produced*; it does not restart a run cancelled with +the request. The durable path is a task — see ``test_task_event_resume.py``. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from agentomatic import AgentPlatform +from agentomatic.core.manifest import AgentManifest +from agentomatic.streaming import ( + StreamReplayBuffer, + get_replay_buffer, + new_stream_id, + set_replay_buffer, + sse_frame, +) + +BASE = "/api/v1" + + +async def _echo_fn(state: dict[str, Any]) -> dict[str, Any]: + """Return a canned response for the incoming query.""" + return {"response": f"echo: {state.get('current_query', '')}"} + + +@pytest.fixture(autouse=True) +def _fresh_buffer(): + """Isolate each test from the process-wide buffer.""" + set_replay_buffer(StreamReplayBuffer()) + yield + set_replay_buffer(None) + + +@pytest.fixture +def client() -> Any: + """A platform exposing one streaming agent.""" + platform = AgentPlatform( + agents_dir="/tmp/agentomatic_stream_replay_test", + title="Stream Replay Test", + version="0.0.1", + ) + platform.register_agent( + manifest=AgentManifest(name="echo", slug="fn-echo", description="Echo"), + node_fn=_echo_fn, + ) + with TestClient(platform.build()) as test_client: + yield test_client + + +def _ids(body: str) -> list[int]: + """Extract SSE ``id:`` values from a body.""" + return [int(line[len("id: ") :]) for line in body.splitlines() if line.startswith("id: ")] + + +def _data(body: str) -> list[str]: + """Extract SSE ``data:`` payloads from a body.""" + return [line[len("data: ") :] for line in body.splitlines() if line.startswith("data: ")] + + +class TestInvokeStreamIsReplayable: + def test_response_advertises_a_stream_id(self, client: Any) -> None: + response = client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}) + assert response.status_code == 200 + assert response.headers.get("X-Stream-Id", "").startswith("stream_") + + def test_frames_are_numbered(self, client: Any) -> None: + response = client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}) + sequences = _ids(response.text) + assert sequences == list(range(1, len(sequences) + 1)), "gapless numbering required" + + def test_replay_returns_the_produced_frames(self, client: Any) -> None: + original = client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}) + stream_id = original.headers["X-Stream-Id"] + + replayed = client.get(f"{BASE}/echo/invoke/stream/{stream_id}") + + assert replayed.status_code == 200 + assert _data(replayed.text) == _data(original.text) + assert _ids(replayed.text) == _ids(original.text) + + def test_replay_honours_since(self, client: Any) -> None: + original = client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}) + stream_id = original.headers["X-Stream-Id"] + + replayed = client.get(f"{BASE}/echo/invoke/stream/{stream_id}?since=1") + + assert all(seq > 1 for seq in _ids(replayed.text)) + assert _ids(replayed.text) == [s for s in _ids(original.text) if s > 1] + + def test_replay_honours_last_event_id(self, client: Any) -> None: + original = client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}) + stream_id = original.headers["X-Stream-Id"] + + by_query = client.get(f"{BASE}/echo/invoke/stream/{stream_id}?since=1").text + by_header = client.get( + f"{BASE}/echo/invoke/stream/{stream_id}", + headers={"Last-Event-ID": "1"}, + ).text + + assert _ids(by_header) == _ids(by_query) + + def test_terminal_marker_is_retained(self, client: Any) -> None: + """The [DONE] marker must survive replay, or a client never stops.""" + original = client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}) + stream_id = original.headers["X-Stream-Id"] + replayed = client.get(f"{BASE}/echo/invoke/stream/{stream_id}").text + assert "[DONE]" in replayed + + def test_unknown_stream_404s(self, client: Any) -> None: + response = client.get(f"{BASE}/echo/invoke/stream/stream_deadbeef") + assert response.status_code == 404 + + def test_each_invocation_gets_its_own_stream(self, client: Any) -> None: + first = client.post(f"{BASE}/echo/invoke/stream", json={"query": "one"}) + second = client.post(f"{BASE}/echo/invoke/stream", json={"query": "two"}) + assert first.headers["X-Stream-Id"] != second.headers["X-Stream-Id"] + + replayed = client.get(f"{BASE}/echo/invoke/stream/{first.headers['X-Stream-Id']}").text + assert "one" in replayed + assert "two" not in replayed + + +class TestReplayBuffer: + @pytest.mark.asyncio + async def test_record_assigns_gapless_sequences(self) -> None: + buffer = StreamReplayBuffer() + sequences = [await buffer.record("s", f"frame-{i}") for i in range(5)] + assert sequences == [1, 2, 3, 4, 5] + + @pytest.mark.asyncio + async def test_replay_filters_by_sequence(self) -> None: + buffer = StreamReplayBuffer() + for i in range(1, 5): + await buffer.record("s", f"f{i}") + assert [f.sequence for f in await buffer.replay("s", after=2)] == [3, 4] + + @pytest.mark.asyncio + async def test_frames_per_stream_are_bounded(self) -> None: + buffer = StreamReplayBuffer(max_frames_per_stream=3) + for i in range(1, 8): + await buffer.record("s", f"f{i}") + assert [f.sequence for f in await buffer.replay("s")] == [5, 6, 7] + + @pytest.mark.asyncio + async def test_stream_count_is_bounded(self) -> None: + buffer = StreamReplayBuffer(max_streams=2) + for name in ("a", "b", "c"): + await buffer.record(name, "f") + assert not await buffer.knows("a"), "oldest stream should be evicted" + assert await buffer.knows("c") + + @pytest.mark.asyncio + async def test_drop_forgets_one_stream_only(self) -> None: + buffer = StreamReplayBuffer() + await buffer.record("a", "f") + await buffer.record("b", "f") + await buffer.drop("a") + assert not await buffer.knows("a") + assert await buffer.knows("b") + + @pytest.mark.asyncio + async def test_earliest_sequence_reports_the_retention_floor(self) -> None: + buffer = StreamReplayBuffer(max_frames_per_stream=2) + for i in range(1, 6): + await buffer.record("s", f"f{i}") + assert await buffer.earliest_sequence("s") == 4 + assert await buffer.earliest_sequence("missing") == 0 + + @pytest.mark.asyncio + async def test_streams_number_independently(self) -> None: + buffer = StreamReplayBuffer() + await buffer.record("a", "f") + assert await buffer.record("b", "f") == 1, "streams must not share a counter" + + @pytest.mark.asyncio + async def test_concurrent_records_do_not_collide(self) -> None: + buffer = StreamReplayBuffer() + sequences = await asyncio.gather(*(buffer.record("s", f"f{i}") for i in range(50))) + assert sorted(sequences) == list(range(1, 51)) + + +class TestSseEncoding: + def test_omits_id_when_unnumbered(self) -> None: + assert sse_frame({"a": 1}) == 'data: {"a": 1}\n\n' + + def test_includes_id_when_numbered(self) -> None: + assert sse_frame({"a": 1}, event_id=7).startswith("id: 7\n") + + def test_passes_through_pre_rendered_strings(self) -> None: + assert sse_frame("[DONE]") == "data: [DONE]\n\n" + + def test_stream_ids_are_unique(self) -> None: + assert len({new_stream_id() for _ in range(100)}) == 100 + + def test_default_buffer_is_shared(self) -> None: + assert get_replay_buffer() is get_replay_buffer() From 7fc8d66dfc731e0e0f35af88e95c8b1fe5e3d215 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 15:32:06 +0000 Subject: [PATCH 5/9] fix(a2a): render one agent card for discovery and the per-agent endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The card was built in two places and they had drifted. ``/.well-known/agent.json`` — the canonical A2A discovery document, and the one a conforming client actually reads — emitted a reduced card: no capabilities block at all, and only ``invoke`` and ``chat`` in endpoints. A client doing proper discovery therefore learned *less* about an agent than one that guessed the per-agent ``/card`` URL, and could not discover the streaming or A2A task endpoints at all. Both now render through core.agent_card.build_agent_card, so a capability can only be advertised in one place and the two cannot drift again. The discovery card gains framework, capabilities, metadata and the full endpoint set; its platform envelope is unchanged. The endpoint list also names what the earlier commits added — the A2A task stream and the invoke-stream replay route — so they are discoverable rather than only documented. Also let create_streaming_response opt into the same treatment: pass a ``stream_id`` and a user-built endpoint gets numbered, retained frames and an ``X-Stream-Id`` header, exactly like the built-in streams. Omitting it keeps the previous behaviour. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- docs/architecture/api-reference.md | 26 ++++++++- src/agentomatic/core/agent_card.py | 75 +++++++++++++++++++++++++ src/agentomatic/core/platform.py | 34 ++++++----- src/agentomatic/core/router_factory.py | 42 +++----------- src/agentomatic/protocols/decorators.py | 41 ++++++++++---- tests/test_a2a_streaming.py | 47 ++++++++++++++++ tests/test_stream_replay.py | 53 +++++++++++++++++ 7 files changed, 257 insertions(+), 61 deletions(-) create mode 100644 src/agentomatic/core/agent_card.py diff --git a/docs/architecture/api-reference.md b/docs/architecture/api-reference.md index b88a8d6..934784a 100644 --- a/docs/architecture/api-reference.md +++ b/docs/architecture/api-reference.md @@ -396,18 +396,40 @@ curl http://localhost:8000/api/v1/my_agent/card "streaming": true, "chat": true, "invoke": true, - "a2a": true + "a2a": true, + "stateTransitionHistory": true, + "pushNotifications": true, + "resumableStreams": true }, "endpoints": { "invoke": "/api/v1/my_agent/invoke", "chat": "/api/v1/my_agent/chat", "stream": "/api/v1/my_agent/invoke/stream", - "health": "/api/v1/my_agent/health" + "stream_replay": "/api/v1/my_agent/invoke/stream/{stream_id}", + "health": "/api/v1/my_agent/health", + "card": "/api/v1/my_agent/card", + "a2a_tasks": "/api/v1/my_agent/a2a/tasks", + "a2a_events": "/api/v1/my_agent/a2a/tasks/{task_id}/events" }, "metadata": {} } ``` +`stateTransitionHistory`, `pushNotifications` and `resumableStreams` all +depend on the task subsystem. When it is disabled they report `false` and the +`a2a_*` endpoints are omitted — the card describes what this deployment +actually serves, so a client is never sent down a path that answers `501`. + +!!! note "Discovery returns the same card" + + `GET /.well-known/agent.json` — the canonical A2A discovery document — + renders each agent through the same builder, so its cards are identical to + this one, wrapped in a platform envelope: + + ```json + {"platform": "My Platform", "version": "1.0.0", "agents": {"my_agent": { ... }}} + ``` + --- #### `POST /a2a/tasks` diff --git a/src/agentomatic/core/agent_card.py b/src/agentomatic/core/agent_card.py new file mode 100644 index 0000000..0a3f813 --- /dev/null +++ b/src/agentomatic/core/agent_card.py @@ -0,0 +1,75 @@ +"""The single source of truth for an agent's A2A card. + +The card was built in two places — ``GET /api/v1/{agent}/card`` and the +platform's ``/.well-known/agent.json`` — and they had drifted. The per-agent +card advertised capabilities and a streaming endpoint; the well-known one, the +canonical A2A discovery document, listed only ``invoke`` and ``chat`` and no +capabilities at all. A client doing proper discovery therefore learned *less* +about an agent than one that guessed the per-agent URL. + +Both now render through :func:`build_agent_card`, so a capability can only be +advertised in one place: here. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from agentomatic.core.manifest import AgentManifest + + +def build_agent_card( + manifest: AgentManifest, + agent_name: str, + api_prefix: str, + *, + supports_tasks: bool, +) -> dict[str, Any]: + """Render one agent's A2A card. + + Capabilities are reported from what the deployment actually wired up + rather than asserted unconditionally: a card that claims a capability the + platform lacks sends A2A clients down a path that then answers 501. + + Args: + manifest: The agent's manifest. + agent_name: Name the agent is mounted under. + api_prefix: API prefix the endpoints hang off (e.g. ``/api/v1``). + supports_tasks: Whether the task subsystem is configured. Resumable + streams, task history and push callbacks all depend on it. + + Returns: + The agent card. + """ + base = f"{api_prefix}/{agent_name}" + endpoints: dict[str, str] = { + "invoke": f"{base}/invoke", + "chat": f"{base}/chat", + "stream": f"{base}/invoke/stream", + "stream_replay": f"{base}/invoke/stream/{{stream_id}}", + "health": f"{base}/health", + "card": f"{base}/card", + } + if supports_tasks: + endpoints["a2a_tasks"] = f"{base}/a2a/tasks" + endpoints["a2a_events"] = f"{base}/a2a/tasks/{{task_id}}/events" + + return { + "name": manifest.slug, + "description": manifest.description, + "version": manifest.version, + "framework": manifest.framework, + "capabilities": { + "streaming": True, + "chat": True, + "invoke": True, + "a2a": True, + # These three are the task subsystem's, not the agent's. + "stateTransitionHistory": supports_tasks, + "pushNotifications": supports_tasks, + "resumableStreams": supports_tasks, + }, + "endpoints": endpoints, + "metadata": manifest.metadata, + } diff --git a/src/agentomatic/core/platform.py b/src/agentomatic/core/platform.py index 6911815..2d7148c 100644 --- a/src/agentomatic/core/platform.py +++ b/src/agentomatic/core/platform.py @@ -1810,23 +1810,29 @@ async def readiness() -> dict[str, Any]: # A2A discovery @app.get("/.well-known/agent.json", tags=["Platform"]) async def a2a_discovery() -> dict[str, Any]: - """Return A2A agent cards for all registered agents.""" - cards: dict[str, Any] = {} - for name, agent in self._registry.all().items(): - m = agent.manifest - cards[name] = { - "name": m.slug, - "description": m.description, - "version": m.version, - "endpoints": { - "invoke": f"{self.api_prefix}/{name}/invoke", - "chat": f"{self.api_prefix}/{name}/chat", - }, - } + """Return A2A agent cards for all registered agents. + + This is the canonical A2A discovery document, so it renders the + same card as ``GET {api_prefix}/{agent}/card``. It used to emit a + reduced one — no capabilities, and only ``invoke`` and ``chat`` — + which left a conforming client knowing less about an agent than + one that guessed the per-agent URL. + """ + from agentomatic.core.agent_card import build_agent_card + + supports_tasks = self._task_manager is not None return { "platform": self.title, "version": self.version, - "agents": cards, + "agents": { + name: build_agent_card( + agent.manifest, + name, + self.api_prefix, + supports_tasks=supports_tasks, + ) + for name, agent in self._registry.all().items() + }, } # Agents list diff --git a/src/agentomatic/core/router_factory.py b/src/agentomatic/core/router_factory.py index 00f9538..993118b 100644 --- a/src/agentomatic/core/router_factory.py +++ b/src/agentomatic/core/router_factory.py @@ -14,6 +14,7 @@ from loguru import logger from pydantic import BaseModel, ConfigDict, Field +from agentomatic.core.agent_card import build_agent_card from agentomatic.core.agent_invoke import build_invoke_state, invoke_registered_agent from agentomatic.core.errors import client_safe_detail, client_safe_message from agentomatic.langchain_adapter import dict_to_messages, json_default, to_jsonable @@ -1238,41 +1239,12 @@ async def get_prompts() -> dict[str, Any]: async def get_card() -> dict[str, Any]: """A2A Agent Card.""" agent = _get_agent() - m = agent.manifest - return { - "name": m.slug, - "description": m.description, - "version": m.version, - "framework": m.framework, - # Reported from what this deployment actually wired up. A card - # that claims a capability unconditionally sends A2A clients down - # code paths that then 501, which is worse than admitting the gap. - "capabilities": { - "streaming": True, - "chat": True, - "invoke": True, - "a2a": True, - # Resumable task streams and push both need the task manager. - "stateTransitionHistory": task_manager is not None, - "pushNotifications": task_manager is not None, - "resumableStreams": task_manager is not None, - }, - "endpoints": { - "invoke": f"{api_prefix}/{agent_name}/invoke", - "chat": f"{api_prefix}/{agent_name}/chat", - "stream": f"{api_prefix}/{agent_name}/invoke/stream", - "health": f"{api_prefix}/{agent_name}/health", - **( - { - "a2a_tasks": f"{api_prefix}/{agent_name}/a2a/tasks", - "a2a_events": (f"{api_prefix}/{agent_name}/a2a/tasks/{{task_id}}/events"), - } - if task_manager is not None - else {} - ), - }, - "metadata": m.metadata, - } + return build_agent_card( + agent.manifest, + agent_name, + api_prefix, + supports_tasks=task_manager is not None, + ) # ── A2A task lifecycle ──────────────────────────────────────── # Maps the unified TaskStatus to canonical A2A task states. diff --git a/src/agentomatic/protocols/decorators.py b/src/agentomatic/protocols/decorators.py index 16b27aa..eda9afc 100644 --- a/src/agentomatic/protocols/decorators.py +++ b/src/agentomatic/protocols/decorators.py @@ -65,14 +65,35 @@ def create_streaming_response( generator: Any, agent_name: str = "", media_type: str = "text/event-stream", + *, + stream_id: str | None = None, ) -> StreamingResponse: - """Create an SSE streaming response.""" - return StreamingResponse( - generator, - media_type=media_type, - headers={ - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "X-Agent": agent_name, - }, - ) + """Create an SSE streaming response. + + Args: + generator: Async iterator yielding SSE-encoded strings. + agent_name: Reported in the ``X-Agent`` header. + media_type: Response media type. + stream_id: Retain the stream's frames under this identity and number + them, so a client that drops mid-response can collect what it + missed. Get one from + :func:`~agentomatic.streaming.new_stream_id`; the value is echoed + back in ``X-Stream-Id``. Omit it to stream without retention, as + before. + + Returns: + The streaming response. + """ + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Agent": agent_name, + } + body = generator + if stream_id: + from agentomatic.streaming import numbered_stream + + body = numbered_stream(generator, stream_id=stream_id) + headers["X-Stream-Id"] = stream_id + + return StreamingResponse(body, media_type=media_type, headers=headers) diff --git a/tests/test_a2a_streaming.py b/tests/test_a2a_streaming.py index dec6274..483f932 100644 --- a/tests/test_a2a_streaming.py +++ b/tests/test_a2a_streaming.py @@ -168,3 +168,50 @@ def test_capabilities_are_not_claimed_without_a_task_manager(self) -> None: assert card["capabilities"]["stateTransitionHistory"] is False assert card["capabilities"]["pushNotifications"] is False assert "a2a_events" not in card["endpoints"] + + +class TestDiscoveryMatchesTheAgentCard: + """``/.well-known/agent.json`` is the canonical A2A discovery document. + + It used to emit a reduced card — no capabilities, only ``invoke`` and + ``chat`` — so a conforming client learned *less* about an agent than one + that guessed the per-agent URL and could not discover streaming at all. + """ + + def test_discovery_card_matches_the_per_agent_card(self, client: Any) -> None: + per_agent = client.get(f"{BASE}/echo/card").json() + discovery = client.get("/.well-known/agent.json").json()["agents"]["echo"] + assert discovery == per_agent + + def test_discovery_advertises_capabilities(self, client: Any) -> None: + discovery = client.get("/.well-known/agent.json").json()["agents"]["echo"] + assert discovery["capabilities"]["resumableStreams"] is True + + def test_discovery_advertises_the_streaming_endpoints(self, client: Any) -> None: + endpoints = client.get("/.well-known/agent.json").json()["agents"]["echo"]["endpoints"] + assert "a2a_events" in endpoints + assert "stream_replay" in endpoints + + def test_discovery_keeps_its_platform_envelope(self, client: Any) -> None: + """Unifying the card must not change the document around it.""" + body = client.get("/.well-known/agent.json").json() + assert body["platform"] == "A2A Stream Test" + assert "version" in body + assert "echo" in body["agents"] + + def test_discovery_omits_task_endpoints_without_a_task_manager(self) -> None: + platform = AgentPlatform( + agents_dir="/tmp/agentomatic_a2a_discovery_notasks", + title="No Tasks", + version="0.0.1", + enable_tasks=False, + ) + platform.register_agent( + manifest=AgentManifest(name="echo", slug="fn-echo", description="Echo"), + node_fn=_echo_fn, + ) + with TestClient(platform.build()) as test_client: + discovery = test_client.get("/.well-known/agent.json").json()["agents"]["echo"] + + assert "a2a_events" not in discovery["endpoints"] + assert discovery["capabilities"]["resumableStreams"] is False diff --git a/tests/test_stream_replay.py b/tests/test_stream_replay.py index 4e207e2..e0a65a4 100644 --- a/tests/test_stream_replay.py +++ b/tests/test_stream_replay.py @@ -208,3 +208,56 @@ def test_stream_ids_are_unique(self) -> None: def test_default_buffer_is_shared(self) -> None: assert get_replay_buffer() is get_replay_buffer() + + +class TestPublicStreamingHelper: + """``create_streaming_response`` is exported for user-built endpoints, so + it should be able to opt into the same numbering and retention.""" + + @pytest.mark.asyncio + async def test_without_a_stream_id_nothing_is_retained(self) -> None: + from agentomatic.protocols import create_streaming_response + + async def frames(): + yield "data: one\n\n" + + response = create_streaming_response(frames(), agent_name="a") + assert "X-Stream-Id" not in response.headers + body = "".join([chunk async for chunk in response.body_iterator]) + assert "id:" not in body + + @pytest.mark.asyncio + async def test_a_stream_id_numbers_and_retains(self) -> None: + from agentomatic.protocols import create_streaming_response + + async def frames(): + yield "data: one\n\n" + yield "data: two\n\n" + + stream_id = new_stream_id() + response = create_streaming_response(frames(), agent_name="a", stream_id=stream_id) + assert response.headers["X-Stream-Id"] == stream_id + + body = "".join([chunk async for chunk in response.body_iterator]) + assert "id: 1" in body + assert "id: 2" in body + + retained = await get_replay_buffer().replay(stream_id) + assert [frame.data for frame in retained] == ["one", "two"] + + @pytest.mark.asyncio + async def test_non_data_chunks_pass_through_unnumbered(self) -> None: + """Keep-alive comments must not be numbered or retained as frames.""" + from agentomatic.protocols import create_streaming_response + + async def frames(): + yield ": keep-alive\n\n" + yield "data: real\n\n" + + stream_id = new_stream_id() + response = create_streaming_response(frames(), stream_id=stream_id) + body = "".join([chunk async for chunk in response.body_iterator]) + + assert ": keep-alive" in body + retained = await get_replay_buffer().replay(stream_id) + assert [frame.data for frame in retained] == ["real"] From d576e6e320bb3f68cb8090c2bf29c57066ec5bbd Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 15:33:52 +0000 Subject: [PATCH 6/9] test(streaming): assert Studio's frames stay numbered and retained Studio builds its SSE frames inside run_tracker, so the numbering is applied by wrapping the generator at the router rather than at the yield sites. That is easy to undo by accident: a refactor that stops wrapping would drop the ids silently, with every test still green. Cover both the sequence and the retention under the run id. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- tests/test_stream_replay.py | 38 +++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/test_stream_replay.py b/tests/test_stream_replay.py index e0a65a4..f21df55 100644 --- a/tests/test_stream_replay.py +++ b/tests/test_stream_replay.py @@ -261,3 +261,41 @@ async def frames(): assert ": keep-alive" in body retained = await get_replay_buffer().replay(stream_id) assert [frame.data for frame in retained] == ["real"] + + +class TestStudioStreamsAreNumbered: + """Studio builds its frames inside ``run_tracker``, so the numbering is + applied by wrapping the generator at the router. Worth asserting directly: + a refactor that stops wrapping would silently drop the ids.""" + + @pytest.fixture + def studio_client(self) -> Any: + platform = AgentPlatform( + agents_dir="/tmp/agentomatic_studio_numbering_test", + title="Studio Numbering", + version="0.0.1", + enable_studio=True, + ) + platform.register_agent( + manifest=AgentManifest(name="echo", slug="fn-echo", description="Echo"), + node_fn=_echo_fn, + ) + with TestClient(platform.build()) as test_client: + yield test_client + + def test_run_stream_frames_carry_ids(self, studio_client: Any) -> None: + response = studio_client.post( + "/studio/agents/echo/runs/stream", json={"input": {"query": "hi"}} + ) + assert response.status_code == 200 + sequences = _ids(response.text) + assert sequences == list(range(1, len(sequences) + 1)) + + def test_run_stream_frames_are_retained(self, studio_client: Any) -> None: + response = studio_client.post( + "/studio/agents/echo/runs/stream", json={"input": {"query": "hi"}} + ) + run_id = response.headers["X-Studio-Run-Id"] + retained = asyncio.run(get_replay_buffer().replay(run_id)) + assert retained, "studio frames were not retained under the run id" + assert [frame.sequence for frame in retained] == _ids(response.text) From cb264d83861732ac7e1c0c648f4df25874fd1a12 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 15:53:47 +0000 Subject: [PATCH 7/9] fix(streaming): bound replay memory, key Studio resumes safely, size config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Production hardening of the replay work. Three defects and one gap. **Retention was not a memory bound.** Both buffers capped entries by count only: 512 frames across 256 streams, and 512 events across 1024 tasks. With frames carrying generated answers that is several hundred megabytes of entirely ordinary traffic, and a long answer per frame pushes it past a container's limit. Both now carry byte budgets alongside the counts, with per-stream and total ceilings, and the deques lost their ``maxlen`` because an implicit drop evicted an entry without its size leaving the running total — the accounting would have drifted upwards until nothing could be retained. The newest entry always survives, so a single oversized frame is still replayable rather than silently lost. **Studio keyed its resume stream on ``thread_id``** — a client-supplied path parameter. Resuming the same thread twice continued the first stream's numbering, so a client reconnecting with Last-Event-ID was served the wrong frames, and a caller could reach another thread's retained frames by naming it. Keyed on a fresh per-request id instead. **The A2A stream read the task record once per replayed event**, though the record is only needed for the terminal frame — invisible with the in-memory store, one query per frame with SQL. Read once. **Retention is now sized from the environment** like every other platform switch, so an operator can cap or disable it without editing Python: AGENTOMATIC_TASK_EVENT_LOG / _EVENTS_PER_TASK / _EVENT_TASKS / _EVENT_MAX_MB and AGENTOMATIC_STREAM_REPLAY / _FRAMES / _COUNT / _MAX_MB. A bad value logs and falls back rather than stopping the platform booting. Also confirms with tests that every new streaming route sits behind auth. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- CLAUDE.md | 4 + docs/guide/configuration.md | 8 +- src/agentomatic/cli/agent_guide.py | 4 + src/agentomatic/core/platform.py | 4 + src/agentomatic/core/router_factory.py | 18 ++- src/agentomatic/streaming/__init__.py | 6 + src/agentomatic/streaming/replay.py | 171 +++++++++++++++++++-- src/agentomatic/studio/router.py | 12 +- src/agentomatic/tasks/event_log.py | 202 +++++++++++++++++++++++-- tests/test_a2a_streaming.py | 29 ++++ tests/test_stream_replay.py | 183 ++++++++++++++++++++++ tests/test_task_event_resume.py | 106 +++++++++++++ 12 files changed, 709 insertions(+), 38 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index f7692ed..a557184 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -147,6 +147,10 @@ into the image/compose that drive the same `main.py`. | `AGENTOMATIC_LOGS_HISTORY` | Persist per-agent invoke/chat/stream history (default off) | | `AGENTOMATIC_ALLOW_LOGSLLM_ANALYSIS` | Enable LLM analysis over those logs (default off) | | `AGENTOMATIC_INGESTION_ROOT` | Confine ingestion source/output paths to this dir (default: cwd) | +| `AGENTOMATIC_TASK_EVENT_LOG` | Task-event replay backend: `memory` (default), `none` to disable, or a registered provider | +| `AGENTOMATIC_TASK_EVENTS_PER_TASK` / `AGENTOMATIC_TASK_EVENT_TASKS` / `AGENTOMATIC_TASK_EVENT_MAX_MB` | Task-event replay retention (default 512 / 1024 / 32 MB) | +| `AGENTOMATIC_STREAM_REPLAY` | Retain `/invoke/stream` frames for replay (default on; `0` disables) | +| `AGENTOMATIC_STREAM_FRAMES` / `AGENTOMATIC_STREAM_COUNT` / `AGENTOMATIC_STREAM_MAX_MB` | Stream replay retention (default 512 / 256 / 64 MB) | | `AGENTOMATIC_DEBUG_ERRORS` | Return raw exception text in API errors (default off — dev only) | | `AGENTOMATIC_OTEL_CONSOLE` | Print OpenTelemetry spans to stdout (default off) | | `AGENTOMATIC_RATE_LIMIT_TRUST_PROXY_HEADERS` | Honour X-Forwarded-For for rate-limit keys (default off) | diff --git a/docs/guide/configuration.md b/docs/guide/configuration.md index f18a4ca..082d835 100644 --- a/docs/guide/configuration.md +++ b/docs/guide/configuration.md @@ -193,6 +193,10 @@ directory, host, port, and stack values passed to `agentomatic run`. | `AGENTOMATIC_PLUGIN_AUTORELOAD` | `plugin_autoreload` | `false` | Watch the promoted artifact version and reload plugins safely | | `AGENTOMATIC_PLUGIN_AUTORELOAD_INTERVAL` | `plugin_autoreload_interval` | `5` | Seconds between artifact-pointer checks | | `AGENTOMATIC_RUNS_ROOT` | `runs_root` | `.local/runs` | Scratch directory for pipeline/task outputs | +| `AGENTOMATIC_TASK_EVENT_LOG` | task-event replay backend | `memory` | `none` disables replay; any registered provider name also works | +| `AGENTOMATIC_TASK_EVENTS_PER_TASK` / `AGENTOMATIC_TASK_EVENT_TASKS` / `AGENTOMATIC_TASK_EVENT_MAX_MB` | task-event retention | `512` / `1024` / `32` | Caps how much history `/tasks/{id}/events` can replay. The MB ceiling is what protects a container's memory limit | +| `AGENTOMATIC_STREAM_REPLAY` | `/invoke/stream` frame retention | `true` | `0` streams without retaining anything | +| `AGENTOMATIC_STREAM_FRAMES` / `AGENTOMATIC_STREAM_COUNT` / `AGENTOMATIC_STREAM_MAX_MB` | stream retention | `512` / `256` / `64` | Caps what `GET /invoke/stream/{id}` can replay | | `AGENTOMATIC_AUDIT_LOG` | `audit_log` | `""` (disabled) | JSONL op-audit sink path (non-PII metadata only) | | `AGENTOMATIC_AUDIT_HASH_KEY` | — | API key or process-local key | HMAC key for audit correlation references | | `AGENTOMATIC_CHUNK_SIZE_TOKENS` | `chunk_size_tokens` | `1200` | Default ingestion chunk size | @@ -487,7 +491,7 @@ platform.run( ??? question "My environment variables aren't being loaded" Check these common causes: - + 1. **`.env` file location**: Must be in the project root (same directory as `main.py`) 2. **Variable shape**: Generated-platform variables start with `AGENTOMATIC_` (for example, `AGENTOMATIC_LOG_LEVEL`); `PlatformSettings` variables use @@ -497,7 +501,7 @@ platform.run( ??? question "CORS errors from my frontend" Pass your frontend's origin to `cors_origins`: - + ```python platform = AgentPlatform.from_folder( "agents/", diff --git a/src/agentomatic/cli/agent_guide.py b/src/agentomatic/cli/agent_guide.py index e65c5c8..345a405 100644 --- a/src/agentomatic/cli/agent_guide.py +++ b/src/agentomatic/cli/agent_guide.py @@ -191,6 +191,10 @@ def state_to_output(self, state: HelloState) -> dict[str, Any]: | `AGENTOMATIC_LOGS_HISTORY` | Persist per-agent invoke/chat/stream history (default off) | | `AGENTOMATIC_ALLOW_LOGSLLM_ANALYSIS` | Enable LLM analysis over those logs (default off) | | `AGENTOMATIC_INGESTION_ROOT` | Confine ingestion source/output paths to this dir (default: cwd) | +| `AGENTOMATIC_TASK_EVENT_LOG` | Task-event replay backend: `memory` (default), `none` to disable, or a registered provider | +| `AGENTOMATIC_TASK_EVENTS_PER_TASK` / `AGENTOMATIC_TASK_EVENT_TASKS` / `AGENTOMATIC_TASK_EVENT_MAX_MB` | Task-event replay retention (default 512 / 1024 / 32 MB) | +| `AGENTOMATIC_STREAM_REPLAY` | Retain `/invoke/stream` frames for replay (default on; `0` disables) | +| `AGENTOMATIC_STREAM_FRAMES` / `AGENTOMATIC_STREAM_COUNT` / `AGENTOMATIC_STREAM_MAX_MB` | Stream replay retention (default 512 / 256 / 64 MB) | | `AGENTOMATIC_DEBUG_ERRORS` | Return raw exception text in API errors (default off — dev only) | | `AGENTOMATIC_OTEL_CONSOLE` | Print OpenTelemetry spans to stdout (default off) | | `AGENTOMATIC_RATE_LIMIT_TRUST_PROXY_HEADERS` | Honour X-Forwarded-For for rate-limit keys (default off) | diff --git a/src/agentomatic/core/platform.py b/src/agentomatic/core/platform.py index 2d7148c..e6a3d7a 100644 --- a/src/agentomatic/core/platform.py +++ b/src/agentomatic/core/platform.py @@ -975,12 +975,16 @@ def _build_task_manager(self) -> TaskManager: make_plugin_dispatcher, make_plugin_input_validator, ) + from agentomatic.tasks.event_log import event_log_from_env from agentomatic.tasks.manager import TaskManager from agentomatic.tasks.models import TargetType manager = TaskManager( store=self._task_store, max_concurrency=self._task_max_concurrency, + # Retention is an operational concern — sized from the + # environment, like every other platform switch. + event_log=event_log_from_env(), ) manager.register_dispatcher(TargetType.AGENT, make_agent_dispatcher(self._registry)) manager.register_input_validator( diff --git a/src/agentomatic/core/router_factory.py b/src/agentomatic/core/router_factory.py index 993118b..5012c07 100644 --- a/src/agentomatic/core/router_factory.py +++ b/src/agentomatic/core/router_factory.py @@ -1423,10 +1423,16 @@ async def event_stream() -> Any: "resumed_after": resume_from, } ) + # ``_a2a_event_view`` consults the record only for a + # terminal event's payload, so read it once rather than once + # per replayed event — that was a store round-trip per frame. + replay_record = ( + await task_manager.get(task_id) + if any(evt.status.is_terminal for evt in resumption.replay) + else None + ) for evt in resumption.replay: - yield _frame( - _a2a_event_view(evt, await task_manager.get(task_id)), evt.sequence - ) + yield _frame(_a2a_event_view(evt, replay_record), evt.sequence) current = await task_manager.get(task_id) if current is not None and (resume_from <= 0 or resumption.truncated): @@ -1444,9 +1450,11 @@ async def event_stream() -> Any: if evt.sequence and evt.sequence <= cursor: continue cursor = evt.sequence or cursor - yield _frame( - _a2a_event_view(evt, await task_manager.get(task_id)), evt.sequence + # Only a terminal event needs the record's result. + live_record = ( + await task_manager.get(task_id) if evt.status.is_terminal else None ) + yield _frame(_a2a_event_view(evt, live_record), evt.sequence) if evt.status.is_terminal: yield "data: [DONE]\n\n" return diff --git a/src/agentomatic/streaming/__init__.py b/src/agentomatic/streaming/__init__.py index ec55c34..e6e86e9 100644 --- a/src/agentomatic/streaming/__init__.py +++ b/src/agentomatic/streaming/__init__.py @@ -3,10 +3,13 @@ from __future__ import annotations from .replay import ( + DEFAULT_MAX_BYTES_PER_STREAM, + DEFAULT_MAX_BYTES_TOTAL, DEFAULT_MAX_FRAMES_PER_STREAM, DEFAULT_MAX_STREAMS, StreamFrame, StreamReplayBuffer, + buffer_from_env, get_replay_buffer, new_stream_id, numbered_stream, @@ -15,10 +18,13 @@ ) __all__ = [ + "DEFAULT_MAX_BYTES_PER_STREAM", + "DEFAULT_MAX_BYTES_TOTAL", "DEFAULT_MAX_FRAMES_PER_STREAM", "DEFAULT_MAX_STREAMS", "StreamFrame", "StreamReplayBuffer", + "buffer_from_env", "get_replay_buffer", "new_stream_id", "numbered_stream", diff --git a/src/agentomatic/streaming/replay.py b/src/agentomatic/streaming/replay.py index dc8d1dd..5fa92a8 100644 --- a/src/agentomatic/streaming/replay.py +++ b/src/agentomatic/streaming/replay.py @@ -20,6 +20,8 @@ import asyncio import json +import os +import sys import time import uuid from collections import OrderedDict, deque @@ -31,6 +33,15 @@ #: Streams retained before the least recently used are discarded. DEFAULT_MAX_STREAMS = 256 +#: Bytes retained per stream. A frame carrying a long generated answer is +#: kilobytes, so a count-only bound is not a memory bound: 512 frames across +#: 256 streams is hundreds of megabytes of perfectly ordinary traffic. +DEFAULT_MAX_BYTES_PER_STREAM = 4 * 1024 * 1024 + +#: Bytes retained across all streams — the ceiling that actually protects a +#: container's memory limit. +DEFAULT_MAX_BYTES_TOTAL = 64 * 1024 * 1024 + def new_stream_id() -> str: """Return an identifier for one response stream. @@ -116,14 +127,19 @@ class StreamFrame: sequence: Position in the stream, from 1, with no gaps. data: The frame body exactly as it was sent. timestamp: When the frame was recorded. + size: Measured in-memory footprint of ``data``, used for budgeting. """ - __slots__ = ("data", "sequence", "timestamp") + __slots__ = ("data", "sequence", "size", "timestamp") def __init__(self, sequence: int, data: str) -> None: self.sequence = sequence self.data = data self.timestamp = time.time() + # Real in-memory footprint of the payload, measured once. Character + # count would under-report multi-byte text, which is exactly the + # traffic most likely to blow a budget. + self.size = sys.getsizeof(data) def as_dict(self) -> dict[str, Any]: """Return the frame as a JSON-serialisable mapping. @@ -150,10 +166,19 @@ def __init__( *, max_frames_per_stream: int = DEFAULT_MAX_FRAMES_PER_STREAM, max_streams: int = DEFAULT_MAX_STREAMS, + max_bytes_per_stream: int = DEFAULT_MAX_BYTES_PER_STREAM, + max_bytes_total: int = DEFAULT_MAX_BYTES_TOTAL, ) -> None: self._max_frames = max(1, max_frames_per_stream) self._max_streams = max(1, max_streams) + self._max_bytes_per_stream = max(1, max_bytes_per_stream) + self._max_bytes_total = max(1, max_bytes_total) + # No ``maxlen`` on the deques: an implicit drop would evict a frame + # without its size leaving the running total, and the byte accounting + # would drift upwards until nothing could be retained at all. self._streams: OrderedDict[str, deque[StreamFrame]] = OrderedDict() + self._bytes: dict[str, int] = {} + self._total_bytes = 0 self._sequences: dict[str, int] = {} self._lock = asyncio.Lock() @@ -167,20 +192,82 @@ async def record(self, stream_id: str, data: str) -> int: Returns: The frame's sequence number. """ + frame = StreamFrame(sequence=0, data=data) async with self._lock: sequence = self._sequences.get(stream_id, 0) + 1 self._sequences[stream_id] = sequence - bucket = self._streams.get(stream_id) - if bucket is None: - bucket = deque(maxlen=self._max_frames) - self._streams[stream_id] = bucket - bucket.append(StreamFrame(sequence, data)) + frame.sequence = sequence + + bucket = self._streams.setdefault(stream_id, deque()) + bucket.append(frame) + self._bytes[stream_id] = self._bytes.get(stream_id, 0) + frame.size + self._total_bytes += frame.size self._streams.move_to_end(stream_id) - while len(self._streams) > self._max_streams: - evicted, _ = self._streams.popitem(last=False) - self._sequences.pop(evicted, None) + + self._trim_stream(stream_id) + self._trim_total(keep=stream_id) return sequence + def _trim_stream(self, stream_id: str) -> None: + """Drop this stream's oldest frames until it fits its own limits.""" + bucket = self._streams.get(stream_id) + if bucket is None: + return + while len(bucket) > 1 and len(bucket) > self._max_frames: + self._drop_oldest(stream_id) + self._trim_bytes(stream_id, self._max_bytes_per_stream) + + def _trim_bytes(self, stream_id: str, limit: int) -> None: + """Drop a stream's oldest frames until it fits ``limit`` bytes. + + The newest frame is always kept: a single frame larger than the budget + should still be delivered and replayable, not silently dropped. + + Args: + stream_id: Stream to trim. + limit: Byte ceiling for this stream. + """ + bucket = self._streams.get(stream_id) + if bucket is None: + return + while len(bucket) > 1 and self._bytes.get(stream_id, 0) > limit: + self._drop_oldest(stream_id) + + def _drop_oldest(self, stream_id: str) -> None: + """Remove a stream's oldest frame, keeping the byte totals honest.""" + bucket = self._streams[stream_id] + dropped = bucket.popleft() + self._bytes[stream_id] = self._bytes.get(stream_id, 0) - dropped.size + self._total_bytes -= dropped.size + + def _trim_total(self, *, keep: str) -> None: + """Evict least-recently-used streams until the whole buffer fits. + + Args: + keep: Stream that must survive — it is the one just written to. + """ + while len(self._streams) > 1 and ( + len(self._streams) > self._max_streams or self._total_bytes > self._max_bytes_total + ): + oldest = next(iter(self._streams)) + if oldest == keep: + # ``keep`` was just moved to the end, so it is the newest; + # reaching it means every other stream is already gone. + break + self._forget(oldest) + + # Nothing left to evict but still over budget: the surviving stream is + # the one holding the bytes, so trim its history rather than let the + # total ceiling be quietly ignored. + if self._total_bytes > self._max_bytes_total: + self._trim_bytes(keep, self._max_bytes_total) + + def _forget(self, stream_id: str) -> None: + """Remove one stream and its byte accounting.""" + self._streams.pop(stream_id, None) + self._total_bytes -= self._bytes.pop(stream_id, 0) + self._sequences.pop(stream_id, None) + async def replay(self, stream_id: str, *, after: int = 0) -> list[StreamFrame]: """Return retained frames with sequence greater than ``after``. @@ -229,13 +316,73 @@ async def drop(self, stream_id: str) -> None: stream_id: Stream to forget. """ async with self._lock: - self._streams.pop(stream_id, None) - self._sequences.pop(stream_id, None) + self._forget(stream_id) + + async def stats(self) -> dict[str, int]: + """Return current occupancy, for tests and operational visibility. + + Returns: + Stream count, retained frame count and total retained bytes. + """ + async with self._lock: + return { + "streams": len(self._streams), + "frames": sum(len(bucket) for bucket in self._streams.values()), + "bytes": self._total_bytes, + } _DEFAULT_BUFFER: StreamReplayBuffer | None = None +def _env_int(name: str, default: int) -> int: + """Read a positive integer from the environment. + + Args: + name: Environment variable to read. + default: Value used when unset, unparseable, or not positive. + + Returns: + The configured value, or ``default``. + """ + raw = os.getenv(name) + if not raw: + return default + try: + value = int(raw.strip()) + except ValueError: + return default + return value if value > 0 else default + + +def buffer_from_env() -> StreamReplayBuffer: + """Build the replay buffer a deployment's environment asks for. + + Reads ``AGENTOMATIC_STREAM_FRAMES``, ``AGENTOMATIC_STREAM_COUNT`` and + ``AGENTOMATIC_STREAM_MAX_MB``. Set ``AGENTOMATIC_STREAM_REPLAY=0`` to keep + nothing at all: streams still work, they just cannot be replayed. + + Returns: + The configured :class:`StreamReplayBuffer`. + """ + enabled = (os.getenv("AGENTOMATIC_STREAM_REPLAY") or "1").strip().lower() + if enabled in {"0", "false", "no", "off"}: + # A zero-capacity buffer keeps one frame per stream at most, which is + # the cheapest way to disable replay without special-casing callers. + return StreamReplayBuffer( + max_frames_per_stream=1, max_streams=1, max_bytes_per_stream=1, max_bytes_total=1 + ) + return StreamReplayBuffer( + max_frames_per_stream=_env_int("AGENTOMATIC_STREAM_FRAMES", DEFAULT_MAX_FRAMES_PER_STREAM), + max_streams=_env_int("AGENTOMATIC_STREAM_COUNT", DEFAULT_MAX_STREAMS), + max_bytes_total=_env_int( + "AGENTOMATIC_STREAM_MAX_MB", DEFAULT_MAX_BYTES_TOTAL // (1024 * 1024) + ) + * 1024 + * 1024, + ) + + def get_replay_buffer() -> StreamReplayBuffer: """Return the process-wide replay buffer, creating it on first use. @@ -244,7 +391,7 @@ def get_replay_buffer() -> StreamReplayBuffer: """ global _DEFAULT_BUFFER if _DEFAULT_BUFFER is None: - _DEFAULT_BUFFER = StreamReplayBuffer() + _DEFAULT_BUFFER = buffer_from_env() return _DEFAULT_BUFFER diff --git a/src/agentomatic/studio/router.py b/src/agentomatic/studio/router.py index 19560a0..033cda6 100644 --- a/src/agentomatic/studio/router.py +++ b/src/agentomatic/studio/router.py @@ -15,7 +15,7 @@ from agentomatic.core.errors import client_safe_detail from agentomatic.core.schemas import SchemaValidator, load_schema_models -from agentomatic.streaming import numbered_stream +from agentomatic.streaming import new_stream_id, numbered_stream from agentomatic.studio.adapters import resolve_adapter from agentomatic.studio.models import ( StudioAgentInfo, @@ -532,13 +532,19 @@ async def _stream() -> AsyncGenerator[str, None]: ) yield f"data: {error_data}\n\n" + # Key the retained frames on a fresh per-request id, never on + # ``thread_id``: that is a client-supplied path parameter, so reusing + # it would continue a previous resume's numbering (serving the wrong + # frames to a client reconnecting with Last-Event-ID) and let a caller + # collide with another thread's retained frames by choosing its name. + resume_stream_id = new_stream_id() return StreamingResponse( - numbered_stream(_stream(), stream_id=thread_id), + numbered_stream(_stream(), stream_id=resume_stream_id), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", - "X-Stream-Id": thread_id, + "X-Stream-Id": resume_stream_id, }, ) diff --git a/src/agentomatic/tasks/event_log.py b/src/agentomatic/tasks/event_log.py index cd1207a..f3a9069 100644 --- a/src/agentomatic/tasks/event_log.py +++ b/src/agentomatic/tasks/event_log.py @@ -16,11 +16,15 @@ from __future__ import annotations import asyncio +import json +import os from abc import ABC, abstractmethod from collections import OrderedDict, deque from collections.abc import Callable from typing import TYPE_CHECKING, Any +from loguru import logger + if TYPE_CHECKING: from agentomatic.tasks.models import TaskEvent @@ -30,6 +34,97 @@ #: Tasks retained in the log before the least recently used are discarded. DEFAULT_MAX_TASKS = 1024 +#: Bytes retained across all tasks. Most events are small, but a progress +#: payload can carry a pipeline checkpoint's ``sub_result``, so a count-only +#: bound is not a memory bound. +DEFAULT_MAX_BYTES_TOTAL = 32 * 1024 * 1024 + + +def _env_int(name: str, default: int) -> int: + """Read a positive integer from the environment. + + Args: + name: Environment variable to read. + default: Value used when unset, unparseable, or not positive. + + Returns: + The configured value, or ``default``. + """ + raw = os.getenv(name) + if not raw: + return default + try: + value = int(raw.strip()) + except ValueError: + logger.warning(f"{name}={raw!r} is not an integer — using {default}") + return default + if value <= 0: + logger.warning(f"{name}={value} must be positive — using {default}") + return default + return value + + +def event_log_from_env() -> TaskEventLog: + """Build the task event log a deployment's environment asks for. + + The platform is env-driven everywhere else, so retention has to be + tunable without editing Python — an operator sizing a container needs to + cap it, and someone who wants no replay at all needs to say so. + + Reads ``AGENTOMATIC_TASK_EVENT_LOG`` (a provider name; ``none``/``off`` + selects the null log), plus ``AGENTOMATIC_TASK_EVENTS_PER_TASK``, + ``AGENTOMATIC_TASK_EVENT_TASKS`` and ``AGENTOMATIC_TASK_EVENT_MAX_MB``. + + Returns: + The configured :class:`TaskEventLog`. + """ + provider = (os.getenv("AGENTOMATIC_TASK_EVENT_LOG") or "memory").strip().lower() + if provider in {"none", "off", "disabled", "null"}: + return NullTaskEventLog() + + options: dict[str, Any] = {} + if provider == "memory": + options = { + "max_events_per_task": _env_int( + "AGENTOMATIC_TASK_EVENTS_PER_TASK", DEFAULT_MAX_EVENTS_PER_TASK + ), + "max_tasks": _env_int("AGENTOMATIC_TASK_EVENT_TASKS", DEFAULT_MAX_TASKS), + "max_bytes_total": _env_int( + "AGENTOMATIC_TASK_EVENT_MAX_MB", DEFAULT_MAX_BYTES_TOTAL // (1024 * 1024) + ) + * 1024 + * 1024, + } + try: + return create_event_log(provider, **options) + except ValueError as exc: + # An unknown provider must not take the platform down; replay is a + # convenience, and the default still works. + logger.warning(f"{exc} Falling back to the in-memory event log.") + return InMemoryTaskEventLog() + + +def estimate_event_size(event: TaskEvent) -> int: + """Approximate an event's in-memory footprint in bytes. + + Cheap in the common case: most events carry no ``data``, so this is a + constant. When ``data`` is present it is measured, because that is the + field that can carry a large payload. + + Args: + event: Event to size. + + Returns: + An estimated byte count, never less than the fixed overhead. + """ + overhead = 512 + if not event.data: + return overhead + try: + return overhead + len(json.dumps(event.data, default=str)) + except (TypeError, ValueError): # pragma: no cover - exotic payloads + return overhead + class TaskEventLog(ABC): """Ordered, replayable history of the events emitted for each task.""" @@ -94,13 +189,16 @@ async def drop(self, task_id: str) -> None: class InMemoryTaskEventLog(TaskEventLog): """Bounded in-process event history. - Retention is capped twice over — per task and across tasks — so a - long-running platform cannot accumulate history without limit. Eviction is + Retention is capped three ways — events per task, tasks, and total bytes. + The byte ceiling is the one that protects a container's memory limit: a + progress event can carry a pipeline checkpoint's ``sub_result``, so + counting events alone bounds nothing in particular. Eviction is least-recently-appended. Args: max_events_per_task: Events retained per task. max_tasks: Number of tasks retained. + max_bytes_total: Estimated bytes retained across all tasks. """ def __init__( @@ -108,23 +206,86 @@ def __init__( *, max_events_per_task: int = DEFAULT_MAX_EVENTS_PER_TASK, max_tasks: int = DEFAULT_MAX_TASKS, + max_bytes_total: int = DEFAULT_MAX_BYTES_TOTAL, ) -> None: self._max_events_per_task = max(1, max_events_per_task) self._max_tasks = max(1, max_tasks) + self._max_bytes_total = max(1, max_bytes_total) + # No ``maxlen``: an implicit drop would evict an event without its + # size leaving the running total, and the accounting would drift up + # until nothing could be retained at all. self._events: OrderedDict[str, deque[TaskEvent]] = OrderedDict() + self._sizes: dict[str, int] = {} + self._total_bytes = 0 self._lock = asyncio.Lock() async def append(self, event: TaskEvent) -> None: """Record one event, evicting the oldest history when full.""" + size = estimate_event_size(event) async with self._lock: - bucket = self._events.get(event.task_id) - if bucket is None: - bucket = deque(maxlen=self._max_events_per_task) - self._events[event.task_id] = bucket + task_id = event.task_id + bucket = self._events.setdefault(task_id, deque()) bucket.append(event) - self._events.move_to_end(event.task_id) - while len(self._events) > self._max_tasks: - self._events.popitem(last=False) + self._sizes[task_id] = self._sizes.get(task_id, 0) + size + self._total_bytes += size + self._events.move_to_end(task_id) + + while len(bucket) > 1 and len(bucket) > self._max_events_per_task: + self._drop_oldest(task_id) + self._trim_total(keep=task_id) + + def _drop_oldest(self, task_id: str) -> None: + """Remove a task's oldest event, keeping the byte totals honest. + + ``estimate_event_size`` is deterministic, so what is subtracted here + is exactly what was added when the event was appended. + """ + dropped = estimate_event_size(self._events[task_id].popleft()) + self._sizes[task_id] = self._sizes.get(task_id, 0) - dropped + self._total_bytes -= dropped + + def _trim_bytes(self, task_id: str, limit: int) -> None: + """Drop a task's oldest events until it fits ``limit`` bytes. + + The newest event always survives: a single oversized one should still + be replayable rather than silently dropped. + + Args: + task_id: Task to trim. + limit: Byte ceiling. + """ + bucket = self._events.get(task_id) + if bucket is None: + return + while len(bucket) > 1 and self._sizes.get(task_id, 0) > limit: + self._drop_oldest(task_id) + + def _trim_total(self, *, keep: str) -> None: + """Evict least-recently-used tasks until the whole log fits. + + Args: + keep: Task that must survive — it is the one just written to. + """ + while len(self._events) > 1 and ( + len(self._events) > self._max_tasks or self._total_bytes > self._max_bytes_total + ): + oldest = next(iter(self._events)) + if oldest == keep: + # ``keep`` was just moved to the end, so reaching it means + # every other task's history is already gone. + break + self._forget(oldest) + + # Nothing left to evict but still over budget: the surviving task + # holds the bytes, so trim its history rather than quietly ignoring + # the ceiling. + if self._total_bytes > self._max_bytes_total: + self._trim_bytes(keep, self._max_bytes_total) + + def _forget(self, task_id: str) -> None: + """Remove one task's history and its byte accounting.""" + self._events.pop(task_id, None) + self._total_bytes = max(0, self._total_bytes - self._sizes.pop(task_id, 0)) async def replay(self, task_id: str, *, after: int = 0) -> list[TaskEvent]: """Return retained events for a task with sequence greater than ``after``.""" @@ -138,22 +299,31 @@ async def earliest_sequence(self, task_id: str) -> int: """Return the lowest sequence still retained, or 0 when empty.""" async with self._lock: bucket = self._events.get(task_id) - if not bucket: - return 0 - return bucket[0].sequence + return bucket[0].sequence if bucket else 0 async def latest_sequence(self, task_id: str) -> int: """Return the highest recorded sequence, or 0 when empty.""" async with self._lock: bucket = self._events.get(task_id) - if not bucket: - return 0 - return bucket[-1].sequence + return bucket[-1].sequence if bucket else 0 async def drop(self, task_id: str) -> None: """Discard a task's retained history.""" async with self._lock: - self._events.pop(task_id, None) + self._forget(task_id) + + async def stats(self) -> dict[str, int]: + """Return current occupancy, for tests and operational visibility. + + Returns: + Task count, retained event count and total estimated bytes. + """ + async with self._lock: + return { + "tasks": len(self._events), + "events": sum(len(bucket) for bucket in self._events.values()), + "bytes": self._total_bytes, + } class NullTaskEventLog(TaskEventLog): diff --git a/tests/test_a2a_streaming.py b/tests/test_a2a_streaming.py index 483f932..da75e35 100644 --- a/tests/test_a2a_streaming.py +++ b/tests/test_a2a_streaming.py @@ -215,3 +215,32 @@ def test_discovery_omits_task_endpoints_without_a_task_manager(self) -> None: assert "a2a_events" not in discovery["endpoints"] assert discovery["capabilities"]["resumableStreams"] is False + + +class TestStreamDoesNotHammerTheStore: + """Rendering an A2A frame needs the task record only for the terminal + payload. Reading it per event turned a replay into one store round-trip + per frame — invisible with the in-memory store, N queries with SQL.""" + + def test_replay_reads_the_record_at_most_once(self, client: Any, monkeypatch) -> None: + task_id = _submit(client) + # Let the task finish so there is a real history to replay. + client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events") + + from agentomatic.tasks.manager import TaskManager + + original = TaskManager.get + calls: list[str] = [] + + async def counting_get(self, tid: str): + calls.append(tid) + return await original(self, tid) + + monkeypatch.setattr(TaskManager, "get", counting_get) + body = client.get(f"{BASE}/echo/a2a/tasks/{task_id}/events").text + + frames = len(_frames(body)) + assert frames >= 1 + # One lookup to 404-check, one for the snapshot, at most one for the + # terminal payload — never one per frame. + assert len(calls) <= 3, f"{len(calls)} store reads for {frames} frames" diff --git a/tests/test_stream_replay.py b/tests/test_stream_replay.py index f21df55..f34e46c 100644 --- a/tests/test_stream_replay.py +++ b/tests/test_stream_replay.py @@ -23,6 +23,7 @@ from agentomatic.core.manifest import AgentManifest from agentomatic.streaming import ( StreamReplayBuffer, + buffer_from_env, get_replay_buffer, new_stream_id, set_replay_buffer, @@ -299,3 +300,185 @@ def test_run_stream_frames_are_retained(self, studio_client: Any) -> None: retained = asyncio.run(get_replay_buffer().replay(run_id)) assert retained, "studio frames were not retained under the run id" assert [frame.sequence for frame in retained] == _ids(response.text) + + +class TestNewRoutesRequireAuth: + """Auth is middleware with a skip-list, so a new route is protected by + default — but only until someone adds a path to that list. A replay route + that stops requiring a key hands out other callers' responses.""" + + @pytest.fixture + def secured(self) -> Any: + platform = AgentPlatform( + agents_dir="/tmp/agentomatic_stream_auth_test", + title="Secured", + version="0.0.1", + enable_auth=True, + auth_api_key="secret", + enable_tasks=True, + ) + platform.register_agent( + manifest=AgentManifest(name="echo", slug="fn-echo", description="Echo"), + node_fn=_echo_fn, + ) + with TestClient(platform.build()) as test_client: + yield test_client + + def test_stream_replay_requires_a_key(self, secured: Any) -> None: + header = {"X-API-Key": "secret"} + stream_id = secured.post( + f"{BASE}/echo/invoke/stream", json={"query": "hi"}, headers=header + ).headers["X-Stream-Id"] + + assert secured.get(f"{BASE}/echo/invoke/stream/{stream_id}").status_code == 401 + assert ( + secured.get(f"{BASE}/echo/invoke/stream/{stream_id}", headers=header).status_code + == 200 + ) + + def test_a2a_event_stream_requires_a_key(self, secured: Any) -> None: + header = {"X-API-Key": "secret"} + task_id = secured.post( + f"{BASE}/echo/a2a/tasks", + json={"message": {"content": "hi"}}, + headers=header, + ).json()["task_id"] + + assert secured.get(f"{BASE}/echo/a2a/tasks/{task_id}/events").status_code == 401 + assert ( + secured.get(f"{BASE}/echo/a2a/tasks/{task_id}/events", headers=header).status_code + == 200 + ) + + def test_task_event_stream_requires_a_key(self, secured: Any) -> None: + header = {"X-API-Key": "secret"} + task_id = secured.post( + f"{BASE}/echo/a2a/tasks", + json={"message": {"content": "hi"}}, + headers=header, + ).json()["task_id"] + + assert secured.get(f"{BASE}/tasks/{task_id}/events").status_code == 401 + + def test_stream_ids_are_not_guessable(self) -> None: + """The id is the only thing standing between a caller and a retained + response, so it must carry real entropy.""" + ids = {new_stream_id() for _ in range(500)} + assert len(ids) == 500 + # uuid4 hex[:16] — 64 bits. + assert all(len(sid.removeprefix("stream_")) == 16 for sid in ids) + + +class TestMemoryIsBoundedByBytes: + """Count limits are not memory limits. 512 frames across 256 streams is a + few hundred megabytes of perfectly ordinary traffic, and one long + generated answer per frame pushes it past a container's limit.""" + + @pytest.mark.asyncio + async def test_total_bytes_stay_under_budget(self) -> None: + buffer = StreamReplayBuffer(max_bytes_total=200_000) + big = "x" * 10_000 + for stream in range(40): + for _ in range(20): + await buffer.record(f"s{stream}", big) + + stats = await buffer.stats() + assert stats["bytes"] <= 200_000, "byte budget exceeded" + + @pytest.mark.asyncio + async def test_per_stream_bytes_stay_under_budget(self) -> None: + buffer = StreamReplayBuffer(max_bytes_per_stream=50_000) + for _ in range(200): + await buffer.record("s", "y" * 5_000) + + stats = await buffer.stats() + assert stats["bytes"] <= 50_000 + + @pytest.mark.asyncio + async def test_newest_frame_survives_even_when_oversized(self) -> None: + """A frame bigger than the whole budget must still be replayable — + dropping it would silently lose the answer it carries.""" + buffer = StreamReplayBuffer(max_bytes_per_stream=100, max_bytes_total=100) + await buffer.record("s", "z" * 50_000) + + retained = await buffer.replay("s") + assert len(retained) == 1 + assert retained[0].data == "z" * 50_000 + + @pytest.mark.asyncio + async def test_eviction_keeps_the_accounting_honest(self) -> None: + """If evicted bytes were not subtracted, the total would drift up + until the buffer refused to retain anything at all.""" + buffer = StreamReplayBuffer(max_frames_per_stream=5) + for _ in range(500): + await buffer.record("s", "a" * 1_000) + + stats = await buffer.stats() + retained = await buffer.replay("s") + assert stats["frames"] == len(retained) == 5 + assert stats["bytes"] == sum(frame.size for frame in retained) + + @pytest.mark.asyncio + async def test_dropping_a_stream_reclaims_its_bytes(self) -> None: + buffer = StreamReplayBuffer() + await buffer.record("a", "x" * 1_000) + await buffer.record("b", "x" * 1_000) + before = (await buffer.stats())["bytes"] + + await buffer.drop("a") + after = (await buffer.stats())["bytes"] + + assert after < before + assert after == sum(frame.size for frame in await buffer.replay("b")) + + @pytest.mark.asyncio + async def test_sequences_stay_gapless_through_eviction(self) -> None: + """Eviction trims history; it must never renumber what survives.""" + buffer = StreamReplayBuffer(max_frames_per_stream=4) + for _ in range(20): + await buffer.record("s", "f") + + sequences = [frame.sequence for frame in await buffer.replay("s")] + assert sequences == [17, 18, 19, 20] + + +class TestStreamReplayEnvConfiguration: + def test_defaults_when_nothing_is_set(self, monkeypatch: pytest.MonkeyPatch) -> None: + for var in ( + "AGENTOMATIC_STREAM_REPLAY", + "AGENTOMATIC_STREAM_FRAMES", + "AGENTOMATIC_STREAM_COUNT", + "AGENTOMATIC_STREAM_MAX_MB", + ): + monkeypatch.delenv(var, raising=False) + assert isinstance(buffer_from_env(), StreamReplayBuffer) + + @pytest.mark.asyncio + async def test_replay_can_be_turned_off(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AGENTOMATIC_STREAM_REPLAY", "0") + buffer = buffer_from_env() + for _ in range(20): + await buffer.record("s", "frame") + # Streaming still works; only the history is gone. + assert len(await buffer.replay("s")) <= 1 + + @pytest.mark.asyncio + async def test_frame_limit_is_read_from_the_environment( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("AGENTOMATIC_STREAM_FRAMES", "3") + buffer = buffer_from_env() + for _ in range(10): + await buffer.record("s", "frame") + assert len(await buffer.replay("s")) == 3 + + @pytest.mark.asyncio + async def test_a_bad_value_falls_back_to_the_default( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A typo in a container's env must not stop the platform booting.""" + monkeypatch.setenv("AGENTOMATIC_STREAM_FRAMES", "lots") + buffer = buffer_from_env() + for _ in range(10): + await buffer.record("s", "frame") + assert len(await buffer.replay("s")) == 10 diff --git a/tests/test_task_event_resume.py b/tests/test_task_event_resume.py index cf0a5b8..c7c1c7c 100644 --- a/tests/test_task_event_resume.py +++ b/tests/test_task_event_resume.py @@ -23,6 +23,8 @@ NullTaskEventLog, available_event_log_providers, create_event_log, + estimate_event_size, + event_log_from_env, parse_last_event_id, register_event_log_provider, resolve_resumption, @@ -421,3 +423,107 @@ def test_events_survive_having_no_listener(self, target_type: TargetType) -> Non # Nothing ever subscribed, yet the history is there to replay. assert asyncio.run(manager.event_log.replay(submitted["id"])) + + +class TestEventLogMemoryIsBounded: + """A progress event can carry a pipeline checkpoint's ``sub_result``, so + capping the number of events is not the same as capping memory.""" + + @pytest.mark.asyncio + async def test_total_bytes_stay_under_budget(self) -> None: + log = InMemoryTaskEventLog(max_bytes_total=200_000) + payload = {"blob": "x" * 10_000} + for task in range(40): + for seq in range(1, 21): + await log.append( + TaskEvent( + task_id=f"t{task}", + sequence=seq, + event="progress", + status=TaskStatus.RUNNING, + data=payload, + ) + ) + + assert (await log.stats())["bytes"] <= 200_000 + + @pytest.mark.asyncio + async def test_eviction_keeps_the_accounting_honest(self) -> None: + log = InMemoryTaskEventLog(max_events_per_task=5) + for seq in range(1, 201): + await log.append(_event("t", seq)) + + stats = await log.stats() + assert stats["events"] == 5 + assert stats["bytes"] == sum(estimate_event_size(evt) for evt in await log.replay("t")) + + @pytest.mark.asyncio + async def test_dropping_a_task_reclaims_its_bytes(self) -> None: + log = InMemoryTaskEventLog() + await log.append(_event("a", 1)) + await log.append(_event("b", 1)) + before = (await log.stats())["bytes"] + + await log.drop("a") + + assert (await log.stats())["bytes"] < before + + @pytest.mark.asyncio + async def test_sequences_stay_gapless_through_eviction(self) -> None: + log = InMemoryTaskEventLog(max_events_per_task=4) + for seq in range(1, 21): + await log.append(_event("t", seq)) + + assert [e.sequence for e in await log.replay("t")] == [17, 18, 19, 20] + + @pytest.mark.asyncio + async def test_empty_data_costs_only_the_fixed_overhead(self) -> None: + """The common case must not pay for JSON serialisation.""" + assert estimate_event_size(_event("t", 1)) == 512 + + @pytest.mark.asyncio + async def test_unserialisable_data_does_not_raise(self) -> None: + event = _event("t", 1) + event.data = {"obj": object()} + assert estimate_event_size(event) >= 512 + + +class TestEnvironmentConfiguration: + """Retention is an operational concern: an operator sizing a container has + to be able to cap or disable it without editing Python.""" + + def test_defaults_when_nothing_is_set(self, monkeypatch: pytest.MonkeyPatch) -> None: + for var in ( + "AGENTOMATIC_TASK_EVENT_LOG", + "AGENTOMATIC_TASK_EVENTS_PER_TASK", + "AGENTOMATIC_TASK_EVENT_TASKS", + "AGENTOMATIC_TASK_EVENT_MAX_MB", + ): + monkeypatch.delenv(var, raising=False) + assert isinstance(event_log_from_env(), InMemoryTaskEventLog) + + def test_replay_can_be_turned_off(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AGENTOMATIC_TASK_EVENT_LOG", "none") + assert isinstance(event_log_from_env(), NullTaskEventLog) + + @pytest.mark.asyncio + async def test_limits_are_read_from_the_environment( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("AGENTOMATIC_TASK_EVENT_LOG", "memory") + monkeypatch.setenv("AGENTOMATIC_TASK_EVENTS_PER_TASK", "3") + log = event_log_from_env() + for seq in range(1, 11): + await log.append(_event("t", seq)) + assert [e.sequence for e in await log.replay("t")] == [8, 9, 10] + + def test_a_bad_value_falls_back_instead_of_crashing( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A typo in a container's env must not stop the platform booting.""" + monkeypatch.setenv("AGENTOMATIC_TASK_EVENTS_PER_TASK", "not-a-number") + assert isinstance(event_log_from_env(), InMemoryTaskEventLog) + + def test_an_unknown_provider_falls_back(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AGENTOMATIC_TASK_EVENT_LOG", "redis-typo") + assert isinstance(event_log_from_env(), InMemoryTaskEventLog) From c4a805438f7826525c37da29fc93bce8ac417f2b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 16:03:23 +0000 Subject: [PATCH 8/9] feat(status): report replay buffer occupancy on the status dashboard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The event log and the stream replay buffer both hold memory now, sized by limits an operator is expected to tune. Without visibility there is no way to tell whether those limits fit the box: the buffers were invisible until something ran out of memory. ``/status`` gains a ``streams`` section and ``tasks.event_log``, each reporting retained entries and bytes — ``bytes`` being the number to watch against the configured ceiling. ``stats`` is deliberately not added to the TaskEventLog contract, so a registered custom backend need not implement it; the field is simply absent when it does not. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- docs/guide/status.md | 19 ++++++++++++++++++- src/agentomatic/core/status.py | 11 +++++++++++ src/agentomatic/tasks/manager.py | 12 +++++++++++- tests/test_stream_replay.py | 31 +++++++++++++++++++++++++++++++ 4 files changed, 71 insertions(+), 2 deletions(-) diff --git a/docs/guide/status.md b/docs/guide/status.md index 02d6b55..1203b8b 100644 --- a/docs/guide/status.md +++ b/docs/guide/status.md @@ -48,13 +48,30 @@ starting the platform. It polls the JSON endpoint every few seconds and renders: "by_status": { "queued": 0, "running": 1, "succeeded": 10, "failed": 1, "cancelled": 0 }, "running": 1, "max_concurrency": 8, - "supported_targets": ["agent", "endpoint", "ingestion", "pipeline", "plugin"] + "supported_targets": ["agent", "endpoint", "ingestion", "pipeline", "plugin"], + "event_log": { "tasks": 4, "events": 61, "bytes": 31232 } }, + "streams": { "streams": 2, "frames": 9, "bytes": 12408 }, "storage": { "status": "healthy" }, "generated_at": 1718900000.0 } ``` +### Replay occupancy + +`tasks.event_log` and `streams` report what the replay buffers currently +hold — the history that makes +[dropped streams resumable](tasks.md#resuming-a-dropped-stream). They hold +memory, so `bytes` is the number to watch when sizing a container: if it sits +at the configured ceiling, raise it or lower retention with +`AGENTOMATIC_TASK_EVENT_MAX_MB` / `AGENTOMATIC_STREAM_MAX_MB`, or turn replay +off entirely (`AGENTOMATIC_TASK_EVENT_LOG=none`, +`AGENTOMATIC_STREAM_REPLAY=0`). See +[Configuration](configuration.md). + +`event_log` is absent when a registered custom backend does not report +occupancy — `stats` is not part of the `TaskEventLog` contract. + The top-level `status` is `degraded` if any resource (including a configured connection) is unhealthy or storage is unhealthy, otherwise `healthy` — handy for a single uptime check. diff --git a/src/agentomatic/core/status.py b/src/agentomatic/core/status.py index 5302c7f..c5b4f8e 100644 --- a/src/agentomatic/core/status.py +++ b/src/agentomatic/core/status.py @@ -95,6 +95,16 @@ async def build_status_payload(platform: AgentPlatform) -> dict[str, Any]: except Exception as exc: # noqa: BLE001 tasks = {"enabled": True, "error": str(exc)} + # Stream replay buffer — retained response frames hold memory, so its + # occupancy belongs beside the task counts rather than being invisible. + streams: dict[str, Any] + try: + from agentomatic.streaming import get_replay_buffer + + streams = await get_replay_buffer().stats() + except Exception as exc: # noqa: BLE001 - status must never fail + streams = {"error": str(exc)} + # Connections (platform + per-agent scopes) connections: dict[str, Any] = {} try: @@ -153,6 +163,7 @@ async def build_status_payload(platform: AgentPlatform) -> dict[str, Any]: }, "resources": sections, "tasks": tasks, + "streams": streams, "storage": storage, "generated_at": time.time(), } diff --git a/src/agentomatic/tasks/manager.py b/src/agentomatic/tasks/manager.py index 72577da..ec18a66 100644 --- a/src/agentomatic/tasks/manager.py +++ b/src/agentomatic/tasks/manager.py @@ -249,13 +249,23 @@ async def stats(self) -> dict[str, Any]: the number currently executing, and the configured concurrency. """ by_status = {status.value: await self.store.count(status=status) for status in TaskStatus} - return { + stats: dict[str, Any] = { "total": await self.store.count(), "by_status": by_status, "running": len(self._running), "max_concurrency": self._max_concurrency, "supported_targets": self.supported_targets, } + # The event log holds memory, so its occupancy belongs on the status + # dashboard. ``stats`` is not part of the TaskEventLog contract — a + # registered custom backend need not implement it. + occupancy = getattr(self.event_log, "stats", None) + if callable(occupancy): + try: + stats["event_log"] = await occupancy() + except Exception as exc: # noqa: BLE001 - status must never fail + stats["event_log"] = {"error": str(exc)} + return stats # ------------------------------------------------------------------ # Cancellation diff --git a/tests/test_stream_replay.py b/tests/test_stream_replay.py index f34e46c..ae6324d 100644 --- a/tests/test_stream_replay.py +++ b/tests/test_stream_replay.py @@ -482,3 +482,34 @@ async def test_a_bad_value_falls_back_to_the_default( for _ in range(10): await buffer.record("s", "frame") assert len(await buffer.replay("s")) == 10 + + +class TestReplayOccupancyIsObservable: + """The buffers hold memory. If an operator cannot see how much, they + cannot tell whether the retention limits are sized right for the box.""" + + def test_status_reports_stream_occupancy(self, client: Any) -> None: + client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}) + streams = client.get(f"{BASE}/status").json()["streams"] + + assert streams["streams"] >= 1 + assert streams["frames"] >= 1 + assert streams["bytes"] > 0 + + def test_status_reports_event_log_occupancy(self, client: Any) -> None: + client.post(f"{BASE}/echo/a2a/tasks", json={"message": {"content": "hi"}}) + tasks = client.get(f"{BASE}/status").json()["tasks"] + + assert tasks["enabled"] is True + assert tasks["event_log"]["events"] >= 1 + + def test_status_survives_a_log_without_stats(self) -> None: + """``stats`` is not part of the TaskEventLog contract, so a registered + custom backend that omits it must not break the dashboard.""" + from agentomatic.tasks.event_log import NullTaskEventLog + from agentomatic.tasks.manager import TaskManager + + manager = TaskManager(event_log=NullTaskEventLog()) + assert not hasattr(NullTaskEventLog, "stats") + stats = asyncio.run(manager.stats()) + assert "event_log" not in stats From de013daa8dfe0c3da90a7330090b52b7d9f06614 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 9 Sep 2026 16:11:25 +0000 Subject: [PATCH 9/9] fix(security): scope retained stream replay to its agent and principal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A security review of this branch found a cross-agent authorization bypass in the replay route added earlier in the PR, and reproduced it. ``GET /{agent}/invoke/stream/{stream_id}`` resolved the id against a single process-wide buffer with no check that the retained stream belonged to the agent named in the path. Zero-trust derives its policy purely from that agent segment, so a caller permitted on a permissive agent could name a restricted agent's stream id on their own agent's route and receive the restricted agent's response frames — precisely the boundary the policy engine exists to enforce. Studio run frames shared the same flat namespace, and Studio run ids are enumerable through a path zero-trust does not cover at all. The codebase already had the right pattern two files over: the Studio run endpoint 404s when ``run.agent_name != name``. This restores it. - Frames are recorded under an ownership tag pinning them to the agent that produced them and the principal they were produced for; reads under a different tag return nothing. - The replay route answers 404 for a stream it may not read — the same answer as an id that does not exist, so the response never confirms someone else's stream is there. - Studio frames are tagged ``studio:``, so a Studio run id is not a valid argument to an agent replay route. - ``create_streaming_response`` takes an ``owner`` and its docstring now requires a fresh ``new_stream_id`` rather than suggesting one, so a third-party endpoint cannot key retention on a caller-supplied value. - ``since`` rejects negatives at the edge on both new routes, matching the constraint the task route already applied. Covered by tests that reproduce the original bypass and assert the indistinguishable 404. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Pk5bxkw7Q3j8VN1TvxeQ5V --- docs/architecture/api-reference.md | 5 ++ src/agentomatic/core/router_factory.py | 46 +++++++++++-- src/agentomatic/protocols/decorators.py | 16 +++-- src/agentomatic/streaming/replay.py | 34 ++++++++-- src/agentomatic/studio/router.py | 3 +- tests/test_stream_replay.py | 87 +++++++++++++++++++++++++ 6 files changed, 172 insertions(+), 19 deletions(-) diff --git a/docs/architecture/api-reference.md b/docs/architecture/api-reference.md index 934784a..4d03bb4 100644 --- a/docs/architecture/api-reference.md +++ b/docs/architecture/api-reference.md @@ -241,6 +241,11 @@ Re-send the frames a previous `/invoke/stream` already produced, using the returns only the frames after that point; `404` once the stream is no longer retained. +A retained stream is scoped to the agent that produced it **and** the +principal it was produced for. Replaying another agent's stream — or another +caller's — returns `404`, the same answer as an id that does not exist, so the +response never confirms that someone else's stream is there. + ```bash curl -N "http://localhost:8000/api/v1/my_agent/invoke/stream/stream_a1b2c3d4?since=2" ``` diff --git a/src/agentomatic/core/router_factory.py b/src/agentomatic/core/router_factory.py index 5012c07..ea39789 100644 --- a/src/agentomatic/core/router_factory.py +++ b/src/agentomatic/core/router_factory.py @@ -9,7 +9,7 @@ from collections.abc import Callable from typing import Any, cast -from fastapi import APIRouter, HTTPException, Request +from fastapi import APIRouter, HTTPException, Query, Request from fastapi.responses import StreamingResponse from loguru import logger from pydantic import BaseModel, ConfigDict, Field @@ -736,9 +736,37 @@ async def invoke(request: Any) -> Any: detail=client_safe_detail(exc, context="Agent invocation failed"), ) from exc - async def invoke_stream(request: Any) -> StreamingResponse: + def _stream_owner(http_request: Request | None) -> str: + """Return the ownership tag for a stream on this agent. + + The replay buffer is process-wide, so a lookup by id alone would read + across every boundary the platform enforces per agent: zero-trust + derives its policy from the agent segment of the path, so a caller + allowed on one agent could otherwise fetch another agent's retained + response by naming its stream id on their own agent's route. + + The tag pins a stream to the agent that produced it *and* to the + principal it was produced for, so neither a different agent's route + nor a different caller can read it back. + + Args: + http_request: The incoming request, when one is available. + + Returns: + An opaque ownership tag. + """ + principal = "anon" + if http_request is not None: + state = http_request.scope.get("state") or {} + principal = str( + getattr(http_request.state, "user_id", None) or state.get("user_id") or "anon" + ) + return f"agent:{agent_name}:{principal}" + + async def invoke_stream(request: Any, http_request: Request) -> StreamingResponse: """Invoke agent with SSE streaming.""" agent = _get_agent() + stream_owner = _stream_owner(http_request) state = _build_initial_state(request) thread_id = state.get("thread_id", "") query = state.get("current_query", "") @@ -782,7 +810,7 @@ async def emit(payload: Any) -> str: body = ( payload if isinstance(payload, str) else json.dumps(payload, default=json_default) ) - sequence = await replay_buffer.record(stream_id, body) + sequence = await replay_buffer.record(stream_id, body, owner=stream_owner) return sse_frame(body, event_id=sequence) async def event_stream(): @@ -923,7 +951,7 @@ async def event_stream(): async def replay_invoke_stream( stream_id: str, request: Request, - since: int = 0, + since: int = Query(default=0, ge=0), ) -> StreamingResponse: """Re-send the frames a previous ``/invoke/stream`` already produced. @@ -941,14 +969,18 @@ async def replay_invoke_stream( from agentomatic.tasks.event_log import parse_last_event_id buffer = get_replay_buffer() - if not await buffer.knows(stream_id): + owner = _stream_owner(request) + # A stream belonging to another agent or another caller reports "not + # found" rather than "forbidden": answering differently would confirm + # that the id exists. + if not await buffer.knows(stream_id, owner=owner): raise HTTPException( 404, f"Stream '{stream_id}' is unknown or its frames are no longer retained.", ) resume_from = since or parse_last_event_id(request.headers.get("last-event-id")) - frames = await buffer.replay(stream_id, after=resume_from) + frames = await buffer.replay(stream_id, after=resume_from, owner=owner) earliest = await buffer.earliest_sequence(stream_id) truncated = bool(resume_from and earliest and earliest > resume_from + 1) @@ -1384,7 +1416,7 @@ async def get_a2a_task(task_id: str) -> dict[str, Any]: } @router.get("/a2a/tasks/{task_id}/events") - async def stream_a2a_task(task_id: str, request: Request, since: int = 0): + async def stream_a2a_task(task_id: str, request: Request, since: int = Query(default=0, ge=0)): """Subscribe (or re-subscribe) to an A2A task's progress over SSE. This is the A2A analogue of the platform's task stream: same diff --git a/src/agentomatic/protocols/decorators.py b/src/agentomatic/protocols/decorators.py index eda9afc..5cb9871 100644 --- a/src/agentomatic/protocols/decorators.py +++ b/src/agentomatic/protocols/decorators.py @@ -67,6 +67,7 @@ def create_streaming_response( media_type: str = "text/event-stream", *, stream_id: str | None = None, + owner: str = "", ) -> StreamingResponse: """Create an SSE streaming response. @@ -76,10 +77,15 @@ def create_streaming_response( media_type: Response media type. stream_id: Retain the stream's frames under this identity and number them, so a client that drops mid-response can collect what it - missed. Get one from - :func:`~agentomatic.streaming.new_stream_id`; the value is echoed - back in ``X-Stream-Id``. Omit it to stream without retention, as - before. + missed. It **must** be a fresh + :func:`~agentomatic.streaming.new_stream_id` — never a + caller-supplied value such as a path parameter, which would let + one caller name another's stream. The value is echoed back in + ``X-Stream-Id``. Omit it to stream without retention, as before. + owner: Ownership tag required to read these frames back. Retention is + process-wide, so without a tag scoping the stream to its endpoint + and principal, anything that can reach a replay route could read + it by id alone. Returns: The streaming response. @@ -93,7 +99,7 @@ def create_streaming_response( if stream_id: from agentomatic.streaming import numbered_stream - body = numbered_stream(generator, stream_id=stream_id) + body = numbered_stream(generator, stream_id=stream_id, owner=owner) headers["X-Stream-Id"] = stream_id return StreamingResponse(body, media_type=media_type, headers=headers) diff --git a/src/agentomatic/streaming/replay.py b/src/agentomatic/streaming/replay.py index 5fa92a8..a998f28 100644 --- a/src/agentomatic/streaming/replay.py +++ b/src/agentomatic/streaming/replay.py @@ -76,6 +76,7 @@ async def numbered_stream( source: Any, *, stream_id: str, + owner: str = "", buffer: StreamReplayBuffer | None = None, ) -> Any: """Record and number frames from a generator that already emits SSE text. @@ -89,6 +90,7 @@ async def numbered_stream( Args: source: Async iterator of SSE-encoded strings. stream_id: Identity the frames are retained under. + owner: Ownership tag required to read the frames back. buffer: Replay buffer to use (default: the process-wide one). Yields: @@ -100,7 +102,7 @@ async def numbered_stream( if payload is None: yield chunk continue - sequence = await store.record(stream_id, payload) + sequence = await store.record(stream_id, payload, owner=owner) yield sse_frame(payload, event_id=sequence) @@ -177,17 +179,23 @@ def __init__( # without its size leaving the running total, and the byte accounting # would drift upwards until nothing could be retained at all. self._streams: OrderedDict[str, deque[StreamFrame]] = OrderedDict() + self._owners: dict[str, str] = {} self._bytes: dict[str, int] = {} self._total_bytes = 0 self._sequences: dict[str, int] = {} self._lock = asyncio.Lock() - async def record(self, stream_id: str, data: str) -> int: + async def record(self, stream_id: str, data: str, *, owner: str = "") -> int: """Append one frame and return the sequence assigned to it. Args: stream_id: Stream the frame belongs to. data: The frame body as sent to the client. + owner: Opaque tag identifying who may read this stream back — + the agent that produced it, and the principal it was produced + for. Reads with a different tag are refused. The buffer is + process-wide, so without this a lookup by id alone would cross + every authorization boundary the platform enforces per agent. Returns: The frame's sequence number. @@ -196,6 +204,7 @@ async def record(self, stream_id: str, data: str) -> int: async with self._lock: sequence = self._sequences.get(stream_id, 0) + 1 self._sequences[stream_id] = sequence + self._owners.setdefault(stream_id, owner) frame.sequence = sequence bucket = self._streams.setdefault(stream_id, deque()) @@ -267,18 +276,25 @@ def _forget(self, stream_id: str) -> None: self._streams.pop(stream_id, None) self._total_bytes -= self._bytes.pop(stream_id, 0) self._sequences.pop(stream_id, None) + self._owners.pop(stream_id, None) - async def replay(self, stream_id: str, *, after: int = 0) -> list[StreamFrame]: + async def replay( + self, stream_id: str, *, after: int = 0, owner: str | None = None + ) -> list[StreamFrame]: """Return retained frames with sequence greater than ``after``. Args: stream_id: Stream to read. after: Exclusive lower bound on sequence. + owner: When given, return nothing unless the stream was recorded + under this exact tag. Returns: The matching frames, oldest first. """ async with self._lock: + if owner is not None and self._owners.get(stream_id) != owner: + return [] bucket = self._streams.get(stream_id) if not bucket: return [] @@ -297,16 +313,22 @@ async def earliest_sequence(self, stream_id: str) -> int: bucket = self._streams.get(stream_id) return bucket[0].sequence if bucket else 0 - async def knows(self, stream_id: str) -> bool: - """Report whether any frames are retained for a stream. + async def knows(self, stream_id: str, *, owner: str | None = None) -> bool: + """Report whether a stream is retained and readable by ``owner``. + + A stream owned by someone else reports False rather than raising, so + callers answer "not found" and do not confirm that the id exists. Args: stream_id: Stream to check. + owner: When given, require this exact ownership tag. Returns: - True when the stream has retained frames. + True when the stream has retained frames the caller may read. """ async with self._lock: + if owner is not None and self._owners.get(stream_id) != owner: + return False return bool(self._streams.get(stream_id)) async def drop(self, stream_id: str) -> None: diff --git a/src/agentomatic/studio/router.py b/src/agentomatic/studio/router.py index 033cda6..fb4e92c 100644 --- a/src/agentomatic/studio/router.py +++ b/src/agentomatic/studio/router.py @@ -356,6 +356,7 @@ async def stream_run(name: str, request: StudioRunRequest) -> StreamingResponse: adapter, state, run.id, thread_id, request.checkpoint_id, request.breakpoints ), stream_id=run.id, + owner=f"studio:{name}", ), media_type="text/event-stream", headers={ @@ -539,7 +540,7 @@ async def _stream() -> AsyncGenerator[str, None]: # collide with another thread's retained frames by choosing its name. resume_stream_id = new_stream_id() return StreamingResponse( - numbered_stream(_stream(), stream_id=resume_stream_id), + numbered_stream(_stream(), stream_id=resume_stream_id, owner=f"studio:{name}"), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", diff --git a/tests/test_stream_replay.py b/tests/test_stream_replay.py index ae6324d..5b8a962 100644 --- a/tests/test_stream_replay.py +++ b/tests/test_stream_replay.py @@ -513,3 +513,90 @@ def test_status_survives_a_log_without_stats(self) -> None: assert not hasattr(NullTaskEventLog, "stats") stats = asyncio.run(manager.stats()) assert "event_log" not in stats + + +class TestReplayIsScopedToItsOwner: + """The replay buffer is process-wide, and zero-trust derives its policy + from the agent segment of the path. Without an ownership check, a caller + permitted on one agent could read another agent's retained response by + naming its stream id on their own agent's route — bypassing exactly the + boundary the platform enforces per agent.""" + + @pytest.fixture + def two_agents(self) -> Any: + async def secret(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "TOP-SECRET-SALARY-DATA"} + + platform = AgentPlatform( + agents_dir="/tmp/agentomatic_stream_owner_test", + title="Two Agents", + version="0.0.1", + enable_studio=True, + ) + platform.register_agent( + manifest=AgentManifest(name="secret_agent", slug="s", description="s"), + node_fn=secret, + ) + platform.register_agent( + manifest=AgentManifest(name="open_agent", slug="o", description="o"), + node_fn=_echo_fn, + ) + with TestClient(platform.build()) as test_client: + yield test_client + + def test_another_agents_route_cannot_replay_the_stream(self, two_agents: Any) -> None: + produced = two_agents.post(f"{BASE}/secret_agent/invoke/stream", json={"query": "salary?"}) + stream_id = produced.headers["X-Stream-Id"] + assert "TOP-SECRET-SALARY-DATA" in produced.text + + leaked = two_agents.get(f"{BASE}/open_agent/invoke/stream/{stream_id}") + + assert leaked.status_code == 404 + assert "TOP-SECRET-SALARY-DATA" not in leaked.text + + def test_the_owning_agent_can_still_replay(self, two_agents: Any) -> None: + stream_id = two_agents.post( + f"{BASE}/secret_agent/invoke/stream", json={"query": "salary?"} + ).headers["X-Stream-Id"] + + replayed = two_agents.get(f"{BASE}/secret_agent/invoke/stream/{stream_id}") + + assert replayed.status_code == 200 + assert "TOP-SECRET-SALARY-DATA" in replayed.text + + def test_a_studio_run_id_is_not_valid_on_an_agent_route(self, two_agents: Any) -> None: + """Studio run ids are enumerable, so they must live in their own + namespace rather than being accepted as agent stream ids.""" + run_id = two_agents.post( + "/studio/agents/open_agent/runs/stream", json={"input": {"query": "hi"}} + ).headers["X-Studio-Run-Id"] + + assert two_agents.get(f"{BASE}/open_agent/invoke/stream/{run_id}").status_code == 404 + + def test_a_missing_stream_and_a_forbidden_one_look_the_same(self, two_agents: Any) -> None: + """Answering differently would confirm that an id exists.""" + stream_id = two_agents.post( + f"{BASE}/secret_agent/invoke/stream", json={"query": "hi"} + ).headers["X-Stream-Id"] + + forbidden = two_agents.get(f"{BASE}/open_agent/invoke/stream/{stream_id}") + missing = two_agents.get(f"{BASE}/open_agent/invoke/stream/stream_deadbeefdeadbeef") + + assert forbidden.status_code == missing.status_code == 404 + + @pytest.mark.asyncio + async def test_buffer_refuses_reads_under_a_different_owner(self) -> None: + buffer = StreamReplayBuffer() + await buffer.record("s", "frame", owner="agent:a:alice") + + assert await buffer.replay("s", owner="agent:a:alice") + assert await buffer.replay("s", owner="agent:b:alice") == [] + assert await buffer.replay("s", owner="agent:a:bob") == [] + assert not await buffer.knows("s", owner="agent:b:alice") + + def test_negative_resume_points_are_rejected(self, client: Any) -> None: + stream_id = client.post(f"{BASE}/echo/invoke/stream", json={"query": "hi"}).headers[ + "X-Stream-Id" + ] + + assert client.get(f"{BASE}/echo/invoke/stream/{stream_id}?since=-5").status_code == 422