diff --git a/CLAUDE.md b/CLAUDE.md index f7692ed7..a5571843 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/architecture/api-reference.md b/docs/architecture/api-reference.md index 8e853733..4d03bb49 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,48 @@ 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. + +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" +``` + +!!! 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. @@ -360,18 +401,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` @@ -429,6 +492,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/cli/commands.md b/docs/cli/commands.md index cc44aa75..32c935f2 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 f0bee242..c97fe17c 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/docs/guide/configuration.md b/docs/guide/configuration.md index f18a4ca6..082d8351 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/docs/guide/status.md b/docs/guide/status.md index 02d6b556..1203b8bf 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/docs/guide/tasks.md b/docs/guide/tasks.md index f64a5dcb..beb10d7a 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/pyproject.toml b/pyproject.toml index ccfa6fc6..d29512d3 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/agent_guide.py b/src/agentomatic/cli/agent_guide.py index e65c5c85..345a4056 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/cli/commands.py b/src/agentomatic/cli/commands.py index f9b87fa3..7ed98427 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/core/agent_card.py b/src/agentomatic/core/agent_card.py new file mode 100644 index 00000000..0a3f813d --- /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 6911815f..e6a3d7a3 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( @@ -1810,23 +1814,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 7fc26d19..ea39789d 100644 --- a/src/agentomatic/core/router_factory.py +++ b/src/agentomatic/core/router_factory.py @@ -2,20 +2,23 @@ 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, Query, Request from fastapi.responses import StreamingResponse 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 +from agentomatic.streaming import get_replay_buffer, new_stream_id, sse_frame # --------------------------------------------------------------------------- # Request / Response Models @@ -162,7 +165,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 @@ -733,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", "") @@ -767,6 +798,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, owner=stream_owner) + return sse_frame(body, event_id=sequence) async def event_stream(): """Yield SSE frames from graph or node execution.""" @@ -785,30 +830,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: @@ -849,7 +894,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( @@ -863,12 +915,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 @@ -889,6 +947,66 @@ 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 = Query(default=0, ge=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() + 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, owner=owner) + 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 @@ -1153,26 +1271,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, - "capabilities": { - "streaming": True, - "chat": True, - "invoke": True, - "a2a": True, - }, - "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", - }, - "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. @@ -1202,6 +1306,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 +1415,90 @@ 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 = 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 + 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, + } + ) + # ``_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, replay_record), 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 + # 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 + 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/core/status.py b/src/agentomatic/core/status.py index 5302c7f0..c5b4f8e9 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/protocols/decorators.py b/src/agentomatic/protocols/decorators.py index 16b27aa8..5cb98718 100644 --- a/src/agentomatic/protocols/decorators.py +++ b/src/agentomatic/protocols/decorators.py @@ -65,14 +65,41 @@ def create_streaming_response( generator: Any, agent_name: str = "", media_type: str = "text/event-stream", + *, + stream_id: str | None = None, + owner: str = "", ) -> 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. 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. + """ + 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, owner=owner) + headers["X-Stream-Id"] = stream_id + + return StreamingResponse(body, media_type=media_type, headers=headers) diff --git a/src/agentomatic/stacks/defaults.py b/src/agentomatic/stacks/defaults.py index eea60c9b..68c648ad 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/src/agentomatic/streaming/__init__.py b/src/agentomatic/streaming/__init__.py new file mode 100644 index 00000000..e6e86e9b --- /dev/null +++ b/src/agentomatic/streaming/__init__.py @@ -0,0 +1,33 @@ +"""Shared plumbing for Server-Sent Events streams.""" + +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, + set_replay_buffer, + sse_frame, +) + +__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", + "set_replay_buffer", + "sse_frame", +] diff --git a/src/agentomatic/streaming/replay.py b/src/agentomatic/streaming/replay.py new file mode 100644 index 00000000..a998f287 --- /dev/null +++ b/src/agentomatic/streaming/replay.py @@ -0,0 +1,427 @@ +"""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 os +import sys +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 + +#: 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. + + 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, + owner: 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. + owner: Ownership tag required to read the frames back. + 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, owner=owner) + 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. + size: Measured in-memory footprint of ``data``, used for budgeting. + """ + + __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. + + 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, + 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._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, *, 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. + """ + frame = StreamFrame(sequence=0, data=data) + 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()) + 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) + + 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) + self._owners.pop(stream_id, None) + + 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 [] + 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, *, 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 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: + """Discard a stream's retained frames. + + Args: + stream_id: Stream to forget. + """ + async with self._lock: + 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. + + Returns: + The shared :class:`StreamReplayBuffer`. + """ + global _DEFAULT_BUFFER + if _DEFAULT_BUFFER is None: + _DEFAULT_BUFFER = buffer_from_env() + 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 0eceff93..fb4e92c4 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 new_stream_id, numbered_stream from agentomatic.studio.adapters import resolve_adapter from agentomatic.studio.models import ( StudioAgentInfo, @@ -350,8 +351,12 @@ 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, + owner=f"studio:{name}", ), media_type="text/event-stream", headers={ @@ -528,10 +533,20 @@ 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( - _stream(), + numbered_stream(_stream(), stream_id=resume_stream_id, owner=f"studio:{name}"), 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": resume_stream_id, + }, ) @router.get( diff --git a/src/agentomatic/tasks/__init__.py b/src/agentomatic/tasks/__init__.py index 931f5e05..f7ed1936 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 00000000..f3a90693 --- /dev/null +++ b/src/agentomatic/tasks/event_log.py @@ -0,0 +1,473 @@ +"""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 +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 + +#: 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 + +#: 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.""" + + @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 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__( + self, + *, + 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: + task_id = event.task_id + bucket = self._events.setdefault(task_id, deque()) + bucket.append(event) + 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``.""" + 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) + 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) + 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._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): + """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 11830831..ec18a66e 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 @@ -240,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 @@ -285,6 +304,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 +529,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 641c5b91..26d519f2 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 da1a4f1a..57159c10 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 00000000..da75e35b --- /dev/null +++ b/tests/test_a2a_streaming.py @@ -0,0 +1,246 @@ +# 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"] + + +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 + + +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_cli_init_project_root.py b/tests/test_cli_init_project_root.py new file mode 100644 index 00000000..6cae891e --- /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 00000000..a3eb004b --- /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 55bea01c..ba82fc70 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. @@ -320,6 +322,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 +344,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/tests/test_stream_replay.py b/tests/test_stream_replay.py new file mode 100644 index 00000000..5b8a962c --- /dev/null +++ b/tests/test_stream_replay.py @@ -0,0 +1,602 @@ +# 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, + buffer_from_env, + 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() + + +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"] + + +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) + + +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 + + +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 + + +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 diff --git a/tests/test_task_event_resume.py b/tests/test_task_event_resume.py new file mode 100644 index 00000000..c7c1c7c2 --- /dev/null +++ b/tests/test_task_event_resume.py @@ -0,0 +1,529 @@ +# 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, + estimate_event_size, + event_log_from_env, + 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"])) + + +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) diff --git a/uv.lock b/uv.lock index 1db6c700..bf51ff98 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" },