diff --git a/CLAUDE.md b/CLAUDE.md index 5f572d5..41cf0a0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -36,7 +36,7 @@ agentomatic deploy --profile minimal --stack remote # production-lean | Command | Purpose | | --- | --- | | `agentomatic new NAME` | Scaffold a full project (alias for `init --project`) | -| `agentomatic init NAME [--project] [--template basic\|full\|class\|rag\|chatbot\|coordinator\|extraction\|deepagent\|legacy_dict\|ingestion]` | Scaffold an agent (or project) | +| `agentomatic init NAME [--project] [--template basic\|full\|class\|rag\|chatbot\|coordinator\|extraction\|deepagent\|legacy_dict\|ingestion\|langchain]` | Scaffold an agent (or project) | | `agentomatic add connection\|ingestion NAME` | Add components to an existing agent | | `agentomatic run [--studio/--no-studio] [--with-ui] [--port 8000] [--ssl-certfile ...] [--require-auth-globally]` | Start the platform (prefers `main:app` when present) | | `agentomatic deploy [--profile full\|minimal] [--minimal] [--distroless] [--stack NAME]` | Generate Dockerfile/compose/.env | @@ -103,6 +103,13 @@ fully-featured, **env-driven** `app` — nothing is silently dropped. Agents mount at **`/api/v1/{agent_name}/invoke`** (no `/agents/` segment). +!!! note "`query` on the wire, `current_query` in state" + + The REST body field is **`query`** (`{"query": "..."}`); posting + `current_query` returns 422. The framework normalises it, so the dict your + `input_to_state` receives has **`current_query`** — which is why the example + above reads `data.get("current_query", "")`. `/chat` uses `content` instead. + ### Deploy profiles - **full** (default): everything on — REST API, Swagger, Studio UI, health, @@ -123,12 +130,16 @@ into the image/compose that drive the same `main.py`. | `AGENTOMATIC_ENABLE_METRICS` | Prometheus metrics (default on) | | `AGENTOMATIC_LOG_LEVEL` | Log level (default `INFO`; `WARNING` in minimal) | | `AGENTOMATIC_ENABLE_AUTH` / `AGENTOMATIC_API_KEY` | API-key auth | -| `AGENTOMATIC_ENABLE_JWT` | JWT auth (configure JWKS via stack) | +| `AGENTOMATIC_ENABLE_JWT` | JWT auth (JWKS from `AUTH__JWKS_URL` / `AUTH__ISSUER` / `AUTH__AUDIENCE`, or the active stack's `auth:` block) | | `AGENTOMATIC_REQUIRE_AUTH` | Require auth globally (implies JWT + zero-trust) | | `AGENTOMATIC_ENABLE_CONTROL_PLANE` / `AGENTOMATIC_CONTROL_TOKEN` | Control plane | | `AGENTOMATIC_ENABLE_RATE_LIMIT` | Rate limiting | | `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_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) | | `AGENTOMATIC_TITLE` | Platform title | | `AGENTOMATIC_STACK` | Active stack name | | `AGENTOMATIC_AGENTS` | Comma-separated allow-list scoping agent discovery | diff --git a/Dockerfile.distroless b/Dockerfile.distroless index 1693363..abc22d7 100644 --- a/Dockerfile.distroless +++ b/Dockerfile.distroless @@ -68,3 +68,8 @@ EXPOSE 8000 # distroless images boot through the same ``agentomatic run`` entrypoint. ENTRYPOINT ["/app/.venv/bin/python", "/app/.venv/bin/agentomatic"] CMD ["run", "--agents-dir", "agents", "--host", "0.0.0.0", "--port", "8000"] + +# No shell and no curl in this image — hit /health with the venv Python +# instead (exec form, so no shell is needed to run this CMD either). +HEALTHCHECK --interval=30s --timeout=10s --start-period=30s --retries=3 \ + CMD ["/app/.venv/bin/python", "-c", "import urllib.request as u; u.urlopen('http://localhost:8000/health', timeout=5)"] diff --git a/docs/architecture/api-reference.md b/docs/architecture/api-reference.md index a3ed20c..3f9ad52 100644 --- a/docs/architecture/api-reference.md +++ b/docs/architecture/api-reference.md @@ -376,10 +376,20 @@ Submit an A2A protocol task. | Field | Type | Required | Description | |---|---|---|---| -| `message` | `object` | ✅ | A2A message with `content` field | +| `message` | `object` | ✅ | A2A message — protocol `parts`, or a `content` / `text` string | | `metadata` | `object` | | Extra metadata | +The protocol form carries text in `parts`; the simplified `content` string is +also accepted. A message with no readable text is rejected with `422` rather +than silently running the agent on an empty query. + ```bash +# Protocol form +curl -X POST http://localhost:8000/api/v1/my_agent/a2a/tasks \ + -H "Content-Type: application/json" \ + -d '{"message": {"role": "user", "parts": [{"type": "text", "text": "Analyze this dataset"}]}}' + +# Simplified form curl -X POST http://localhost:8000/api/v1/my_agent/a2a/tasks \ -H "Content-Type: application/json" \ -d '{"message": {"content": "Analyze this dataset"}}' diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index d1f263c..5ca6054 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -68,7 +68,30 @@ If you prefer a lightweight install, you can select only the modules and depende | `studio` | `pip install agentomatic[studio]` | Agentomatic Studio visual debugger | | `optimize` | `pip install agentomatic[optimize]` | DSPy-style optimizer loop + DeepEval validation | | `telemetry` | `pip install agentomatic[telemetry]` | OpenTelemetry APM tracing exporters | -| `all` | `pip install agentomatic[all]` | Installs all components and drivers above | +| `all` | `pip install "agentomatic[all]"` | Everything except the vendor LLM drivers and the Chainlit UI — see the note below | + +!!! 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. + + 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]" + ``` + +!!! tip "Quote the extras in zsh/bash" + + Square brackets are glob syntax in most shells, so quote them: + `pip install "agentomatic[all]"`. Unquoted, zsh fails with + `no matches found`. !!! note "Combining extras" You can combine multiple extras in a single install command: diff --git a/docs/guide/control-plane.md b/docs/guide/control-plane.md index e5ef35d..f06afa0 100644 --- a/docs/guide/control-plane.md +++ b/docs/guide/control-plane.md @@ -90,6 +90,28 @@ flowchart LR - A disabled agent → only that agent's routes return `503`. - The control plane itself remains reachable so you can turn things back on. +!!! warning "Control-plane state is per-process and does not survive a restart" + + `maintenance_mode` and the set of disabled agents live in memory for the + lifetime of the process. They are **not** persisted, so any restart — a + redeploy, a crash-restart, a rolling update, or scaling to a new replica — + resets them to "maintenance off, nothing disabled". + + Two consequences worth planning around in production: + + - An agent you drained comes back **live** after a deploy. If a drain has + to outlast a restart, enforce it upstream (load balancer, ingress rule, + or by removing the agent from `AGENTOMATIC_AGENTS`) rather than relying + on the control plane alone. + - With more than one replica, a control call only affects the **replica + that served it**. Route control requests to every replica, or treat the + control plane as a per-instance debugging tool rather than a + fleet-wide switch. + + Both toggles are enforced on the hot request path (not merely reported), + and an agent mounted under both its folder name and its manifest slug is + disabled under both aliases. + ## Typical rollout flow 1. Enable maintenance mode before a risky migration. diff --git a/docs/guide/deployment.md b/docs/guide/deployment.md index aad286b..8405b39 100644 --- a/docs/guide/deployment.md +++ b/docs/guide/deployment.md @@ -118,6 +118,14 @@ agentomatic deploy --profile minimal --stack remote --distroless agentomatic deploy --minimal --stack remote # shorthand ``` +!!! note "Distroless images pin Python 3.11" + `gcr.io/distroless/python3-debian12` is Debian 12's Python **3.11**, so the + distroless build stage uses `python:3.11-slim` and installs dependencies + with `pip --target=/app/deps` (on `PYTHONPATH`) rather than into a + virtualenv — a venv's `bin/python` points at the *build* stage's + interpreter, which does not exist in the runtime image. If you customise + the generated Dockerfile, keep the two Python versions in step. + !!! warning "Swagger is always available" `--profile minimal` **never** disables `/docs`, `/redoc`, or `/openapi.json`. It only sets `AGENTOMATIC_ENABLE_STUDIO=0` and @@ -477,6 +485,41 @@ Run it: `AGENTOMATIC_TITLE`, `AGENTOMATIC_LOG_LEVEL`), so `uvicorn main:app` in the generated Dockerfile drops no functionality versus running the CLI. +!!! tip "Turning on verified JWT auth" + `AGENTOMATIC_ENABLE_JWT=1` alone gives you a middleware with nothing to + verify against. Point it at your identity provider's JWKS endpoint with + `AUTH__JWKS_URL` (plus `AUTH__ISSUER` / `AUTH__AUDIENCE`), or set the same + values in the active stack's `auth:` block — `agentomatic deploy` writes + them into the generated `.env` for you. Environment wins over the stack. + + ```bash + AGENTOMATIC_REQUIRE_AUTH=1 + AUTH__JWKS_URL=https://idp.example.com/.well-known/jwks.json + AUTH__ISSUER=https://idp.example.com/ + AUTH__AUDIENCE=agentomatic + ``` + + With `AGENTOMATIC_REQUIRE_AUTH=1` and no JWKS, the platform will not accept + unsigned tokens: it enforces with your API key if one is configured + (`AGENTOMATIC_ENABLE_AUTH=1` + `AGENTOMATIC_API_KEY`), and refuses to start + if neither is set. + +!!! warning "Rate limiting behind a proxy" + The limiter keys on the client address, so behind a reverse proxy every + request looks like it comes from the proxy — set + `AGENTOMATIC_RATE_LIMIT_TRUST_PROXY_HEADERS=1` so `X-Forwarded-For` is used + instead. Leave it **off** when the platform is exposed directly: the header + is caller-controlled, and rotating it per request would bypass the limiter. + + That flag governs Agentomatic's own reading of the header. Uvicorn's + `--proxy-headers` (on by default) separately rewrites `request.client` from + `X-Forwarded-For` for peers in `--forwarded-allow-ips` (default + `127.0.0.1`), before any middleware runs. Keep that list limited to your + real proxy — a caller connecting *from* an allowed peer address can steer + the rate-limit key regardless of the flag above. The generated + `nginx.conf` sits on that trusted hop, which is why it is the right place + to set the header. + !!! warning "Workers and in-memory state" Connection pools and per-process caches live **per worker**. Keep shared state (threads, memory, cache) in external services (Postgres, redis) so it diff --git a/docs/guide/langgraph.md b/docs/guide/langgraph.md index 8162c07..d299a58 100644 --- a/docs/guide/langgraph.md +++ b/docs/guide/langgraph.md @@ -243,6 +243,16 @@ metadata: Annotated[dict, _merge_dicts] The Studio adapter extracts state from the LangGraph checkpointer automatically: +!!! warning "Requires a checkpointer" + These endpoints read the graph's checkpointer, and nothing else writes + checkpoints. An agent whose `build_graph` returns `self.new_graph().compile()` + (the native `GraphBuilder` runtime) has no checkpointer, so `/state` returns + `{}` and `/history` returns `[]` — Studio's State and History panels stay + empty even though the thread's chat messages are persisted separately. + To get thread state and time-travel, build the graph with LangGraph's + `StateGraph` and compile it with `AgentomaticCheckpointer` as shown in + [Storage](storage.md#checkpointer-api). + ```bash # Get current thread state curl http://localhost:8000/studio/agents/my_agent/threads/thread_001/state @@ -524,7 +534,7 @@ sequenceDiagram participant ST as BaseStore (Memory/SQL) LG->>CP: aput(config, checkpoint, metadata) - CP->>CP: _ensure_json_serializable(checkpoint) + CP->>CP: encode_for_storage(checkpoint) CP->>ST: save_checkpoint(thread_id, ns, id, data) ST-->>CP: saved CP-->>LG: RunnableConfig @@ -574,7 +584,13 @@ graph = builder.compile(checkpointer=checkpointer) ``` !!! tip "Safe Serialization" - The checkpointer automatically handles non-JSON-serializable objects (datetimes, bytes, custom classes) via `_ensure_json_serializable()`. No extra configuration needed. + The checkpointer automatically handles non-JSON-serializable objects — datetimes, + bytes, custom classes, and LangChain `BaseMessage` objects (`HumanMessage`, + `AIMessage`, `ToolMessage`, ...) — via LangGraph's own `JsonPlusSerializer` + (`encode_for_storage()` / `decode_from_storage()`). Messages round-trip back + into real message objects, not stringified reprs, so `add_messages` and + `prompt | llm` chains keep working across checkpoint resumes. No extra + configuration needed. ### Replaying from Checkpoints diff --git a/docs/guide/platform-features.md b/docs/guide/platform-features.md index 2015bdf..2f3f33e 100644 --- a/docs/guide/platform-features.md +++ b/docs/guide/platform-features.md @@ -353,17 +353,31 @@ platform.register_after_node_hook(audit_output_hook) ## 🛡️ 8. Safe Checkpoint Serialization -LangGraph checkpoints can contain non-JSON-serializable Python objects (datetimes, bytes, custom classes). The `AgentomaticCheckpointer` automatically handles this via a safe JSON round-trip: +LangGraph checkpoints can contain non-JSON-serializable Python objects — datetimes, +bytes, custom classes, and (critically) LangChain `BaseMessage` objects +(`HumanMessage`, `AIMessage`, `ToolMessage`, ...) sitting in channel values like +`messages`. The `AgentomaticCheckpointer` automatically handles this using +LangGraph's own `JsonPlusSerializer`, which round-trips these objects **without +losing their type or structure**: ```python -from agentomatic.storage.checkpointer import AgentomaticCheckpointer, _ensure_json_serializable +from agentomatic.storage.checkpointer import decode_from_storage, encode_for_storage -# Objects like datetimes, bytes, and custom classes are safely converted -data = _ensure_json_serializable({"ts": datetime.now(), "raw": b"bytes"}) -# → {"ts": "2026-06-14 12:00:00", "raw": "b'bytes'"} +# Rich objects are encoded into a JSON-safe wrapper for storage... +data = encode_for_storage({"ts": datetime.now(), "raw": b"bytes"}) +# → {"__agentomatic_serde_type__": "msgpack", "__agentomatic_serde_data__": ""} + +# ...and decoded back into the *original* Python objects, not strings. +restored = decode_from_storage(data) +# → {"ts": datetime(...), "raw": b"bytes"} ``` -This happens transparently inside `aput()` — no user configuration needed. All values are round-tripped through `json.dumps(obj, default=str)` → `json.loads()`. +This happens transparently inside `aput()`/`aget_tuple()` — no user configuration +needed. Unlike a naive `json.dumps(obj, default=str)`, a `HumanMessage`/`AIMessage` +stored in `channel_values["messages"]` comes back as a real message object (with +`tool_calls`, `tool_call_id`, etc. intact) on the next `graph.ainvoke()`, so the +LangGraph `add_messages` reducer and any `prompt | llm` chain keep working across +checkpoint resumes. --- diff --git a/docs/guide/storage.md b/docs/guide/storage.md index 2207f75..e7dc61f 100644 --- a/docs/guide/storage.md +++ b/docs/guide/storage.md @@ -404,8 +404,13 @@ The checkpointer implements the full LangGraph `BaseCheckpointSaver` interface: | `list(config, ...)` | List checkpoints (sync wrapper, returns `Iterator`) | | `alist(config, ...)` | List checkpoints (async, returns `AsyncIterator`) | -!!! info "JSON Serialization" - The checkpointer automatically ensures all checkpoint and metadata values are JSON-serializable using a `default=str` fallback. Custom objects, datetimes, and bytes are converted to their string representation. +!!! info "Serialization" + Checkpoints and metadata are serialized with LangGraph's own + `JsonPlusSerializer`, so LangChain `BaseMessage` subclasses, pydantic + models, dataclasses, datetimes, and bytes round-trip as themselves — a + reloaded thread gives you back real `HumanMessage` / `AIMessage` objects, + not their `repr()`. Payloads written before this used a lossy + `default=str` fallback and are returned as-is on read (best effort). --- diff --git a/pyproject.toml b/pyproject.toml index d21fff8..2e5db27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -117,7 +117,12 @@ docs = [ ] [build-system] -requires = ["hatchling"] +# Upper-bounded deliberately: hatchling 1.30+ emits ``Metadata-Version: 2.5``, +# which current twine rejects ("'2.5' is not a valid metadata version") — and +# the release workflow publishes through pypa/gh-action-pypi-publish, which +# verifies metadata with twine before upload. Lift the bound once twine and +# PyPI accept 2.5. +requires = ["hatchling>=1.27,<1.30"] build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] diff --git a/src/__pycache__/__init__.cpython-312.pyc b/src/__pycache__/__init__.cpython-312.pyc index 72c52fd..c890669 100644 Binary files a/src/__pycache__/__init__.cpython-312.pyc and b/src/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/agentomatic/agents/base.py b/src/agentomatic/agents/base.py index 31d6550..b85b4ec 100644 --- a/src/agentomatic/agents/base.py +++ b/src/agentomatic/agents/base.py @@ -1258,7 +1258,12 @@ async def node_fn(state: dict[str, Any]) -> dict[str, Any]: """Node function adapter for registry.""" input_data = { "query": state.get("current_query", ""), - **{k: v for k, v in state.items() if k not in ("messages", "thread_id")}, + # Forward ``messages``/``thread_id`` too — see the rationale in + # ``agentomatic.core.agent_invoke._input_from_state``: they are + # required by LangChain-style agents (MessagesPlaceholder + # history and RunnableConfig thread_id) and were previously + # unreachable from ``input_to_state``. + **state, } return await self.atransform(input_data) diff --git a/src/agentomatic/agents/graph.py b/src/agentomatic/agents/graph.py index 89fe6b4..8b82ad4 100644 --- a/src/agentomatic/agents/graph.py +++ b/src/agentomatic/agents/graph.py @@ -19,6 +19,8 @@ from loguru import logger +from agentomatic.core.errors import client_safe_detail + from .types import StateT, TraceEvent # Sentinel for "end of graph" @@ -308,11 +310,18 @@ async def astream(self, state: StateT) -> AsyncGenerator[dict[str, Any], None]: import copy def _state_to_dict(s: Any) -> dict[str, Any]: + from agentomatic.langchain_adapter import to_jsonable + if hasattr(s, "model_dump"): - return s.model_dump() - if hasattr(s, "__dict__"): - return vars(s) - return dict(s) if isinstance(s, dict) else {} + raw = s.model_dump() + elif hasattr(s, "__dict__"): + raw = vars(s) + else: + raw = dict(s) if isinstance(s, dict) else {} + # Convert any LangChain BaseMessage objects in state fields (e.g. + # ``messages: list[BaseMessage]``) to plain dicts so this event + # payload is safe to json.dumps downstream (SSE / REST). + return cast(dict[str, Any], to_jsonable(raw)) errors = self.validate() if errors: @@ -380,11 +389,18 @@ def _now_iso() -> str: return datetime.now(UTC).isoformat() def _state_to_dict(s: Any) -> dict[str, Any]: + from agentomatic.langchain_adapter import to_jsonable + if hasattr(s, "model_dump"): - return s.model_dump() - if hasattr(s, "__dict__"): - return vars(s) - return dict(s) if isinstance(s, dict) else {} + raw = s.model_dump() + elif hasattr(s, "__dict__"): + raw = vars(s) + else: + raw = dict(s) if isinstance(s, dict) else {} + # Convert any LangChain BaseMessage objects in state fields (e.g. + # ``messages: list[BaseMessage]``) to plain dicts so this event + # payload is safe to json.dumps downstream (SSE / REST). + return cast(dict[str, Any], to_jsonable(raw)) yield { "event": "run_start", @@ -435,7 +451,8 @@ def _state_to_dict(s: Any) -> dict[str, Any]: "run_id": run_id, "node": current_node_name, "timestamp": _now_iso(), - "data": {"error": str(exc)}, + # Streamed to Studio clients — sanitise (full detail is logged). + "data": client_safe_detail(exc, context="Node failed"), } raise RuntimeError(f"Node '{current_node_name}' failed: {exc}") from exc diff --git a/src/agentomatic/agents/history.py b/src/agentomatic/agents/history.py index ca2ea66..d13d3bf 100644 --- a/src/agentomatic/agents/history.py +++ b/src/agentomatic/agents/history.py @@ -389,9 +389,7 @@ def _render( for key, delta in params.items(): old_v, new_v = delta["old"], delta["new"] if isinstance(old_v, list) or isinstance(new_v, list): - parts.append( - f" {key}: {len(old_v or [])} → {len(new_v or [])} items" - ) + parts.append(f" {key}: {len(old_v or [])} → {len(new_v or [])} items") else: parts.append(f" {key}: {old_v!r} → {new_v!r}") diff --git a/src/agentomatic/agents/optimizers.py b/src/agentomatic/agents/optimizers.py index e12ed59..83032a0 100644 --- a/src/agentomatic/agents/optimizers.py +++ b/src/agentomatic/agents/optimizers.py @@ -380,16 +380,13 @@ def _build_fitter(self, agent: Any, name: str) -> Any: effective_space = overrides.get("search_space") or kwargs.get("search_space") or _Space() baseline_model_params: dict[str, Any] = {} - if ( - isinstance(compiled_cfg, dict) - and getattr(effective_space, "optimize_model_params", False) + if isinstance(compiled_cfg, dict) and getattr( + effective_space, "optimize_model_params", False ): param_keys = set(getattr(effective_space, "model_param_space", {}) or {}) param_keys.add("temperature") baseline_model_params = { - k: v - for k, v in compiled_cfg.items() - if k in param_keys and v is not None + k: v for k, v in compiled_cfg.items() if k in param_keys and v is not None } # Same compounding for few-shot examples accepted by an earlier epoch. baseline_few_shot: list[dict[str, Any]] = [] diff --git a/src/agentomatic/cli/agent_guide.py b/src/agentomatic/cli/agent_guide.py index cb247cd..51605da 100644 --- a/src/agentomatic/cli/agent_guide.py +++ b/src/agentomatic/cli/agent_guide.py @@ -83,7 +83,7 @@ | Command | Purpose | | --- | --- | | `agentomatic new NAME` | Scaffold a full project (alias for `init --project`) | -| `agentomatic init NAME [--project] [--template basic\\|full\\|class\\|rag\\|chatbot\\|coordinator\\|extraction\\|deepagent\\|legacy_dict\\|ingestion]` | Scaffold an agent (or project) | +| `agentomatic init NAME [--project] [--template basic\\|full\\|class\\|rag\\|chatbot\\|coordinator\\|extraction\\|deepagent\\|legacy_dict\\|ingestion\\|langchain]` | Scaffold an agent (or project) | | `agentomatic add connection\\|ingestion NAME` | Add components to an existing agent | | `agentomatic run [--studio/--no-studio] [--with-ui] [--port 8000] [--ssl-certfile ...] [--require-auth-globally]` | Start the platform (prefers `main:app` when present) | | `agentomatic deploy [--profile full\\|minimal] [--minimal] [--distroless] [--stack NAME]` | Generate Dockerfile/compose/.env | @@ -151,6 +151,13 @@ def state_to_output(self, state: HelloState) -> dict[str, Any]: Agents mount at **`/api/v1/{{agent_name}}/invoke`** (no `/agents/` segment). +!!! note "`query` on the wire, `current_query` in state" + + The REST body field is **`query`** (`{{"query": "..."}}`); posting + `current_query` returns 422. The framework normalises it, so the dict your + `input_to_state` receives has **`current_query`** — which is why the example + above reads `data.get("current_query", "")`. `/chat` uses `content` instead. + ### Deploy profiles - **full** (default): everything on — REST API, Swagger, Studio UI, health, @@ -177,6 +184,10 @@ def state_to_output(self, state: HelloState) -> dict[str, Any]: | `AGENTOMATIC_ENABLE_RATE_LIMIT` | Rate limiting | | `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_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) | | `AGENTOMATIC_TITLE` | Platform title | | `AGENTOMATIC_STACK` | Active stack name | | `AGENTOMATIC_AGENTS` | Comma-separated allow-list scoping agent discovery | diff --git a/src/agentomatic/cli/commands.py b/src/agentomatic/cli/commands.py index d3075af..d1b7337 100644 --- a/src/agentomatic/cli/commands.py +++ b/src/agentomatic/cli/commands.py @@ -29,6 +29,7 @@ from loguru import logger from agentomatic.cli.agent_guide import WRITE_TARGETS as _AGENT_GUIDE_TARGETS +from agentomatic.cli.templates import TEMPLATES # Graceful Rich fallback try: @@ -199,25 +200,10 @@ def cli(ctx: click.Context, version: bool) -> None: @click.option( "--template", "-t", - type=click.Choice( - [ - "basic", - "class", # alias for basic (class-owned BaseGraphAgent) - "full", - "coordinator", - "pipeline", - "rag", - "chatbot", - "deepagent", - "custom", - "legacy_dict", - "plugin", - "endpoint", - "connection", - "ingestion", - "extraction", - ] - ), + # Derived from the template registry so the CLI can never drift out of + # sync with it (a hardcoded list previously omitted "langchain", making + # a shipped template unreachable from the CLI). + type=click.Choice(list(TEMPLATES)), default=None, help="Template to use (default: interactive selection)", ) @@ -254,7 +240,11 @@ def init( name = "agentomatic-app" from .project import scaffold_project - dest = Path(target_dir) if target_dir else Path(name) + # ``--dir`` is documented as the *parent* directory, and that is how + # it behaves for agents (``init foo --dir x`` → ``x/foo``). Treating it + # as the project directory itself scattered the project's 15 files + # straight into a directory the user only meant to scaffold *inside*. + dest = Path(target_dir) / name if target_dir else Path(name) result = scaffold_project(dest, name, force=force) click.echo() logger.success(f"Created project '{name}'") @@ -400,6 +390,19 @@ def init( f" 2. Prefer: [cyan]agentomatic add connection {name}[/cyan] " f"to attach to an existing agent without overwrite\n\n" ) + elif template == "deepagent": + # This template imports a third-party package that agentomatic does + # not depend on. Without it the agent scaffolds and registers fine but + # every invocation fails, so say so here rather than at the first 500. + steps = ( + f"[bold]Next steps:[/bold]\n\n" + f" 1. [cyan]pip install deepagents[/cyan] " + f"[yellow](required — this template imports it)[/yellow]\n" + f" 2. Edit [yellow]{target / edit_file}[/yellow]\n" + f" 3. Review [yellow]{target / '__init__.py'}[/yellow] " + f"(AgentManifest / card)\n" + f" 4. [cyan]agentomatic run --studio[/cyan]\n\n" + ) else: steps = ( f"[bold]Next steps:[/bold]\n\n" @@ -638,8 +641,33 @@ def run( log_level=log_level, require_auth_globally=require_auth_globally, ) + # ``uvicorn.run("main:app")`` resolves the import string against + # ``sys.path``. When agentomatic is launched as a console script + # (``uv run agentomatic run``, or any pipx/venv install), the project + # directory is NOT on sys.path — only the interpreter's script dir is + # — so importing ``main`` fails. Put the project dir on the path + # explicitly, and export it via PYTHONPATH so uvicorn's --reload + # subprocess inherits it too. + project_dir = str(Path.cwd()) + if project_dir not in sys.path: + sys.path.insert(0, project_dir) + existing_pythonpath = os.environ.get("PYTHONPATH", "") + if project_dir not in existing_pythonpath.split(os.pathsep): + os.environ["PYTHONPATH"] = ( + f"{project_dir}{os.pathsep}{existing_pythonpath}" + if existing_pythonpath + else project_dir + ) logger.info("Found main.py — starting via uvicorn main:app ...") - uvicorn.run("main:app", **run_kwargs) + # The scaffolded main.py enables the platform's LoggingMiddleware, which + # already logs every request with a correlation id; uvicorn's access log + # would duplicate each line. Set AGENTOMATIC_UVICORN_ACCESS_LOG=1 to + # keep uvicorn's own access log as well. + run_kwargs.setdefault( + "access_log", + _env_bool("AGENTOMATIC_UVICORN_ACCESS_LOG", False), + ) + uvicorn.run("main:app", app_dir=project_dir, **run_kwargs) return logger.info(f"Starting platform from {agents_dir} (plugins: {plugins_dir})...") @@ -1391,7 +1419,8 @@ def doctor(agents_dir: str) -> None: ver = getattr(mod, "__version__", "installed") checks.append((f"{pkg} [{extra}]", True, ver)) except ImportError: - checks.append((f"{pkg} [{extra}]", False, f"pip install agentomatic[{extra}]")) + # Quote the extra — unquoted brackets are glob syntax in zsh/bash. + checks.append((f"{pkg} [{extra}]", False, f'pip install "agentomatic[{extra}]"')) # Agents directory agents_path = Path(agents_dir) @@ -2023,8 +2052,22 @@ def stack_list(stacks_dir: str) -> None: @stack.command("show") @click.argument("name") @click.option("--dir", "-d", "stacks_dir", default="stacks", help="Stacks directory") -def stack_show(name: str, stacks_dir: str) -> None: - """Show the contents of a stack configuration.""" +@click.option( + "--reveal", + is_flag=True, + help="Show secret values in clear text (default: redacted).", +) +def stack_show(name: str, stacks_dir: str, reveal: bool) -> None: + """Show the contents of a stack configuration. + + Secret-looking values (api keys, passwords, tokens, and credentials + embedded in URLs) are redacted by default, because this output lands in + terminal scrollback, CI logs, and screen shares. ``${ENV_VAR}`` references + are always shown as-is — they are indirections, not secrets. Pass + ``--reveal`` to print the file verbatim. + """ + from agentomatic.stacks.redaction import redact_yaml_text + stacks_path = Path(stacks_dir) stack_file = stacks_path / f"{name}.yaml" if not stack_file.exists(): @@ -2035,6 +2078,10 @@ def stack_show(name: str, stacks_dir: str) -> None: return content = stack_file.read_text() + if not reveal: + content, redactions = redact_yaml_text(content) + if redactions: + logger.info(f"🔒 {redactions} secret value(s) redacted — use --reveal to show them") if HAS_RICH: from rich.syntax import Syntax diff --git a/src/agentomatic/cli/deploy.py b/src/agentomatic/cli/deploy.py index 9db729f..1fcc2dc 100644 --- a/src/agentomatic/cli/deploy.py +++ b/src/agentomatic/cli/deploy.py @@ -11,8 +11,8 @@ and launches the project's ``main.py`` via ``uvicorn main:app`` so the platform's ``AgentPlatform`` configuration is honoured. A ``--distroless`` variant is produced from ``gcr.io/distroless/python3-debian12`` which runs -under the built-in ``nonroot`` user (numeric UID 65532) — verified to have -execute permissions on ``/app/.venv/bin/python``. +under the built-in ``nonroot`` user (numeric UID 65532) and launches the base +image's own ``/usr/bin/python3`` with dependencies on ``PYTHONPATH``. Two deploy *profiles* select how much of the platform the image runs, driven purely through ``AGENTOMATIC_*`` env vars so both share one ``main.py`` code @@ -40,6 +40,7 @@ from typing import TYPE_CHECKING from agentomatic._version import __version__ +from agentomatic.stacks.redaction import env_example_value, redact_url_credentials if TYPE_CHECKING: from agentomatic.stacks.manager import LLMStackEntry, StackConfig @@ -303,7 +304,14 @@ def render_dockerfile_distroless( image scanners and Kubernetes ``runAsNonRoot`` policies do not have to introspect the base image. ``agentomatic`` is installed from PyPI (pinned to *version*) in the build stage; distroless has no shell, so the app is - launched by invoking ``uvicorn`` directly through the venv Python. + launched by invoking ``uvicorn`` through the base image's own interpreter. + + The build stage must match that interpreter. ``distroless/python3-debian12`` + is Debian 12's Python 3.11, so packages are built on ``python:3.11-slim`` + and installed with ``pip --target`` rather than into a virtualenv: a venv's + ``bin/python`` is a symlink to the *builder's* interpreter, which does not + exist in the runtime image, and C extensions built for another minor + version would not import even if it did. Args: version: ``agentomatic`` version to pin (``agentomatic[all]==``). @@ -321,7 +329,10 @@ def render_dockerfile_distroless( # ============================================================================= # ---- Build stage ------------------------------------------------------------ -FROM python:3.12-slim AS builder +# Must be the same Python minor version as the distroless runtime below +# (distroless/python3-debian12 is Debian 12's Python 3.11), or the compiled +# wheels installed here will not import there. +FROM python:3.11-slim AS builder ENV PYTHONDONTWRITEBYTECODE=1 \\ PYTHONUNBUFFERED=1 \\ @@ -334,11 +345,11 @@ def render_dockerfile_distroless( WORKDIR /app -RUN python -m venv /app/.venv -ENV PATH="/app/.venv/bin:$PATH" - +# ``--target`` instead of a virtualenv: the runtime stage runs the distroless +# image's own interpreter, which cannot use a venv built around a different +# Python binary. A plain directory on ``PYTHONPATH`` works with any 3.11. RUN pip install --upgrade pip \\ - && pip install "agentomatic[all]=={version}" + && pip install --target=/app/deps "agentomatic[all]=={version}" # Copy project sources into the build stage so they can be chowned + carried # into the runtime stage (distroless cannot chown at runtime — no shell). @@ -352,8 +363,7 @@ def render_dockerfile_distroless( ENV PYTHONDONTWRITEBYTECODE=1 \\ PYTHONUNBUFFERED=1 \\ - PATH="/app/.venv/bin:$PATH" \\ - VIRTUAL_ENV="/app/.venv" + PYTHONPATH="/app/deps" WORKDIR /app @@ -365,8 +375,9 @@ def render_dockerfile_distroless( EXPOSE 8000 {profile_env_block} # Distroless has no shell but can execute binaries directly. Launch uvicorn -# through the venv Python so main.py's AgentPlatform config is honoured. -ENTRYPOINT ["/app/.venv/bin/python", "-m", "uvicorn"] +# through the base image's interpreter (dependencies come from PYTHONPATH) so +# main.py's AgentPlatform config is honoured. +ENTRYPOINT ["/usr/bin/python3", "-m", "uvicorn"] CMD ["main:app", "--host", "0.0.0.0", "--port", "8000"] """ @@ -385,6 +396,7 @@ def render_docker_compose( build_context: str = ".", volume_prefix: str = ".", profile: str = "full", + distroless: bool = False, ) -> str: """Return a docker-compose file wiring the platform and optional agents. @@ -403,10 +415,23 @@ def render_docker_compose( profile: Deploy profile (``"full"`` or ``"minimal"``); ``minimal`` adds ``AGENTOMATIC_*`` env vars that disable Studio and quiet logs while keeping Swagger, health, metrics, and auth. + distroless: When ``True``, use a curl-free, shell-free healthcheck — + ``gcr.io/distroless/python3-debian12`` has neither, so the + ``curl``-based check used for the regular image would leave the + container permanently reporting "unhealthy". """ profile_env_lines = "".join( f" - {key}={value}\n" for key, value in profile_env(profile).items() ) + if distroless: + # No shell, no curl in distroless — hit /health with the base image's + # own interpreter (stdlib only, so no PYTHONPATH needed). + healthcheck_test = ( + '["CMD", "/usr/bin/python3", "-c", ' + "\"import urllib.request as u; u.urlopen('http://localhost:8000/health', timeout=5)\"]" + ) + else: + healthcheck_test = '["CMD", "curl", "-f", "http://localhost:8000/health"]' stubs = "" for name in agent_names or []: service = name.replace("_", "-").lower() @@ -456,7 +481,7 @@ def render_docker_compose( - {volume_prefix}/stacks:/app/stacks:ro - agentomatic-data:/app/data healthcheck: - test: ["CMD", "curl", "-f", "http://localhost:8000/health"] + test: {healthcheck_test} interval: 30s timeout: 10s retries: 3 @@ -560,9 +585,11 @@ def _llm_env_lines(profile: str, entry: LLMStackEntry) -> list[str]: lines.append(f"{prefix}LLM__BASE_URL={entry.base_url}") if entry.api_key: key_var = f"LLM__{entry.provider.upper()}_API_KEY" - # Show placeholder as-is (usually ``${OPENAI_API_KEY}``) so the - # operator knows which secret to inject at runtime. - lines.append(f"{prefix}{key_var}={entry.api_key}") + # A ``${OPENAI_API_KEY}``-style reference is shown as-is so the operator + # knows which secret to inject at runtime. A *literal* key is replaced + # with a placeholder: .env.example is conventionally committed, so + # copying a real key into it would publish the secret. + lines.append(f"{prefix}{key_var}={env_example_value(entry.api_key)}") return lines @@ -612,7 +639,7 @@ def render_env_example(stack: StackConfig) -> str: f"EMBEDDING__DIMENSION={stack.embedding.dimension}", "", "# --- Database ----------------------------------------------------", - f"DB__URL={stack.database.url}", + f"DB__URL={redact_url_credentials(stack.database.url)}", f"DB__POOL_SIZE={stack.database.pool_size}", f"DB__MAX_OVERFLOW={stack.database.max_overflow}", "", @@ -628,7 +655,7 @@ def render_env_example(stack: StackConfig) -> str: f"# AUTH__METHOD={stack.auth.method}", ] if stack.auth.api_key: - lines.append(f"AUTH__API_KEY={stack.auth.api_key}") + lines.append(f"AUTH__API_KEY={env_example_value(stack.auth.api_key)}") if stack.auth.jwks_url: lines.append(f"AUTH__JWKS_URL={stack.auth.jwks_url}") if stack.auth.issuer: @@ -823,6 +850,7 @@ def generate_deploy( build_context=build_context, volume_prefix=volume_prefix, profile=profile, + distroless=distroless, ) compose_path = out_path / "docker-compose.yml" compose_path.write_text(compose_content) diff --git a/src/agentomatic/cli/project.py b/src/agentomatic/cli/project.py index eb0a9f1..997a9ac 100644 --- a/src/agentomatic/cli/project.py +++ b/src/agentomatic/cli/project.py @@ -38,8 +38,12 @@ def _main_py(name: str) -> str: Returns: Rendered ``main.py`` source. """ - display = name.replace("_", " ").title() - return f'''"""{name} — Agentomatic platform entrypoint. + # Use only the final path segment: ``agentomatic new /srv/apps/my_proj`` + # would otherwise bake the whole filesystem path into the platform title, + # which is published via /openapi.json, /.well-known/agent.json and + # /studio/info — leaking the server's directory layout. + display = Path(name).name.replace("_", " ").title() + return f'''"""{display} — Agentomatic platform entrypoint. Serves an identical feature set whether launched with ``agentomatic run`` (dev) or ``uvicorn main:app`` (container). Toggle features with @@ -48,8 +52,10 @@ def _main_py(name: str) -> str: agentomatic stack use local # or remote """ + from __future__ import annotations +import importlib.metadata import os from agentomatic import AgentPlatform @@ -73,6 +79,34 @@ def _env_bool(var: str, default: bool) -> bool: return raw.strip().lower() in {{"1", "true", "yes", "on"}} +def _platform_api_error(exc: TypeError) -> RuntimeError: + """Turn an unknown-keyword ``TypeError`` into an actionable message. + + ``main.py`` is generated against the agentomatic version that scaffolded + it. If the installed wheel is older (a stale pin in ``requirements.txt``, + a cached container layer, or a Dockerfile pinning a release that predates + a new option), the call below fails with a bare + ``unexpected keyword argument`` and no hint about what to do. + + Args: + exc: The original ``TypeError`` raised by ``AgentPlatform``. + + Returns: + A ``RuntimeError`` naming the installed version and the fix. + """ + try: + installed = importlib.metadata.version("agentomatic") + except importlib.metadata.PackageNotFoundError: # pragma: no cover + installed = "unknown" + return RuntimeError( + f"agentomatic {{installed}} does not support an option this main.py uses " + f"({{exc}}). The installed version is older than the one this project was " + "generated with. Upgrade it (`pip install -U 'agentomatic[all]'`, and " + "raise the pin in requirements.txt / Dockerfile), or remove the " + "unsupported argument from create_platform()." + ) + + def create_platform() -> AgentPlatform: """Build a fully-featured platform matching ``agentomatic run`` defaults. @@ -83,38 +117,53 @@ def create_platform() -> AgentPlatform: Returns: A configured :class:`AgentPlatform` ready to :meth:`build`. + + Raises: + RuntimeError: If the installed agentomatic is too old to accept the + options below (see :func:`_platform_api_error`). """ # ``require_auth`` mirrors ``agentomatic run --require-auth-globally``: # it implies zero-trust + JWT unless those are overridden individually. require_auth = _env_bool("AGENTOMATIC_REQUIRE_AUTH", False) - return AgentPlatform.from_folder( - "agents/", - plugins_dir="plugins/", - endpoints_dir="endpoints/", - ingestion_dir="ingestion/", - stacks_dir="stacks/", - # Pipelines are auto-discovered from ../pipelines/ (sibling of agents/). - # stack="local", # or set via .agentomatic-stack / STACK env - title=os.getenv("AGENTOMATIC_TITLE", "{display} Platform"), - description="Agentomatic multi-agent platform for {name}", - log_level=os.getenv("AGENTOMATIC_LOG_LEVEL", "INFO"), - # On by default — matches `agentomatic run` (Studio) + prod observability. - enable_studio=_env_bool("AGENTOMATIC_ENABLE_STUDIO", True), - enable_metrics=_env_bool("AGENTOMATIC_ENABLE_METRICS", True), - # Opt-in hardening — drive from env / stack; no code edits needed. - enable_auth=_env_bool("AGENTOMATIC_ENABLE_AUTH", False), - auth_api_key=os.getenv("AGENTOMATIC_API_KEY", ""), - enable_jwt_auth=_env_bool("AGENTOMATIC_ENABLE_JWT", require_auth), - enable_zero_trust=_env_bool("AGENTOMATIC_ENABLE_ZERO_TRUST", require_auth), - require_auth_globally=require_auth, - # On with Studio so Control / Endpoints / Connections tabs work out of the box. - enable_control_plane=_env_bool("AGENTOMATIC_ENABLE_CONTROL_PLANE", True), - control_token=os.getenv("AGENTOMATIC_CONTROL_TOKEN", ""), - enable_rate_limit=_env_bool("AGENTOMATIC_ENABLE_RATE_LIMIT", False), - # Opt-in per-agent invocation history + optional LLM log analysis. - logs_history=_env_bool("AGENTOMATIC_LOGS_HISTORY", False), - allow_logsllm_analysis=_env_bool("AGENTOMATIC_ALLOW_LOGSLLM_ANALYSIS", False), - ) + try: + return AgentPlatform.from_folder( + "agents/", + plugins_dir="plugins/", + endpoints_dir="endpoints/", + ingestion_dir="ingestion/", + stacks_dir="stacks/", + # Pipelines are auto-discovered from ../pipelines/ (sibling of agents/). + # stack="local", # or set via .agentomatic-stack / STACK env + title=os.getenv("AGENTOMATIC_TITLE", "{display} Platform"), + description="Agentomatic multi-agent platform for {display}", + log_level=os.getenv("AGENTOMATIC_LOG_LEVEL", "INFO"), + # On by default — matches `agentomatic run` (Studio) + prod observability. + enable_studio=_env_bool("AGENTOMATIC_ENABLE_STUDIO", True), + enable_metrics=_env_bool("AGENTOMATIC_ENABLE_METRICS", True), + # Opt-in hardening — drive from env / stack; no code edits needed. + enable_auth=_env_bool("AGENTOMATIC_ENABLE_AUTH", False), + auth_api_key=os.getenv("AGENTOMATIC_API_KEY", ""), + enable_jwt_auth=_env_bool("AGENTOMATIC_ENABLE_JWT", require_auth), + enable_zero_trust=_env_bool("AGENTOMATIC_ENABLE_ZERO_TRUST", require_auth), + require_auth_globally=require_auth, + # On with Studio so Control / Endpoints / Connections tabs work out of the box. + enable_control_plane=_env_bool("AGENTOMATIC_ENABLE_CONTROL_PLANE", True), + control_token=os.getenv("AGENTOMATIC_CONTROL_TOKEN", ""), + enable_rate_limit=_env_bool("AGENTOMATIC_ENABLE_RATE_LIMIT", False), + # Only trust X-Forwarded-For for rate-limit keys behind a real proxy + # that overwrites it (e.g. a load balancer) — otherwise any caller + # can spoof the header and bypass the limiter. + rate_limit_trust_proxy_headers=_env_bool( + "AGENTOMATIC_RATE_LIMIT_TRUST_PROXY_HEADERS", False + ), + # Opt-in per-agent invocation history + optional LLM log analysis. + logs_history=_env_bool("AGENTOMATIC_LOGS_HISTORY", False), + allow_logsllm_analysis=_env_bool("AGENTOMATIC_ALLOW_LOGSLLM_ANALYSIS", False), + ) + except TypeError as exc: + if "unexpected keyword argument" not in str(exc): + raise + raise _platform_api_error(exc) from exc _platform = create_platform() @@ -226,7 +275,13 @@ def _env_example() -> str: QDRANT_API_KEY= # --- Observability --- -OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 +# Uncomment once a collector is actually reachable. Setting this with nothing +# listening makes the exporter retry every span and fill the log with +# "Transient error StatusCode.UNAVAILABLE / Failed to export traces". +# Spans are still recorded either way; this only controls shipping them. +# OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 +# Print spans to stdout instead (local debugging only — very verbose): +# AGENTOMATIC_OTEL_CONSOLE=1 """ diff --git a/src/agentomatic/cli/templates.py b/src/agentomatic/cli/templates.py index b981c31..7505dce 100644 --- a/src/agentomatic/cli/templates.py +++ b/src/agentomatic/cli/templates.py @@ -61,7 +61,10 @@ def _system_prompt(self) -> str: # Honour optimize/fit overrides via resolve_system_prompt (compiled_config, # system_prompt_override, prompt_manager) — required for train/optimize. return self.resolve_system_prompt( - default="You are a helpful RAG assistant. Ground answers in retrieved context." + default=( + "You are a helpful RAG assistant. " + "Ground answers in retrieved context." + ) ) def build_graph(self): @@ -318,7 +321,7 @@ def _config_py(name: str) -> str: def _schemas_py(name: str) -> str: title = name.replace("_", " ").title().replace(" ", "") - return f'''"""Custom schemas for {name}."""\nfrom __future__ import annotations\n\nfrom pydantic import BaseModel, Field\n\n\nclass {title}Request(BaseModel):\n """Custom request model."""\n query: str = Field(..., description="User query")\n context: dict = Field(default_factory=dict)\n\n\nclass {title}Response(BaseModel):\n """Custom response model."""\n answer: str\n confidence: float = Field(0.0, ge=0.0, le=1.0)\n sources: list[str] = Field(default_factory=list)\n''' + return f'''"""Custom schemas for {name}."""\nfrom __future__ import annotations\n\nfrom pydantic import BaseModel, Field\n\n\nclass {title}Request(BaseModel):\n """Custom request model."""\n query: str = Field(..., description="User query")\n context: dict = Field(default_factory=dict)\n\n\nclass {title}Response(BaseModel):\n """Custom response model.\n\n Field names must match what the agent\'s ``state_to_output()`` actually\n returns, or every invoke logs an output-validation warning.\n """\n response: str\n agent_type: str = ""\n confidence: float = Field(0.0, ge=0.0, le=1.0)\n sources: list[str] = Field(default_factory=list)\n''' def _tools_py(name: str) -> str: @@ -326,7 +329,7 @@ def _tools_py(name: str) -> str: def _api_py(name: str) -> str: - return f'''"""Custom API router for {name}.\n\nIf this file exports a `router`, it REPLACES the auto-generated endpoints.\nRemove this file to use auto-generated endpoints instead.\n"""\nfrom __future__ import annotations\n\nfrom fastapi import APIRouter\n\nrouter = APIRouter()\n\n\n@router.get("/status")\nasync def status() -> dict:\n """Custom status endpoint."""\n return {{"agent": "{name}", "custom_router": True}}\n''' + return f'''"""Custom API router for {name}.\n\nExporting a module-level ``router`` REPLACES *all* auto-generated endpoints\nfor this agent — /invoke, /chat, /invoke/stream, /card and /health included.\n\nThis scaffold therefore names it ``custom_router``, which the registry does\nnot pick up, so the agent keeps its auto-generated endpoints out of the box.\nRename it to ``router`` when you genuinely want to take over the agent\'s\nroutes entirely (and re-add any of the generated ones you still need).\n"""\nfrom __future__ import annotations\n\nfrom fastapi import APIRouter\n\ncustom_router = APIRouter()\n\n\n@custom_router.get("/status")\nasync def status() -> dict:\n """Custom status endpoint."""\n return {{"agent": "{name}", "custom_router": True}}\n''' def _prompts_json() -> str: @@ -430,8 +433,20 @@ def _resolve_model() -> str: @lru_cache(maxsize=1) def create_agent(): - """Create and compile the deep agent.""" - from deepagents import create_deep_agent + """Create and compile the deep agent. + + Raises: + RuntimeError: If the third-party ``deepagents`` package is missing. + """ + try: + from deepagents import create_deep_agent + except ImportError as exc: # pragma: no cover - depends on the environment + raise RuntimeError( + "The deepagent template requires the third-party 'deepagents' " + "package, which agentomatic does not install. Run " + "`pip install deepagents` (and add it to requirements.txt) " + "before invoking this agent." + ) from exc return create_deep_agent( model=_resolve_model(), @@ -459,7 +474,7 @@ def _legacy_init_py( def _legacy_graph_py(name: str) -> str: - return f'''"""LangGraph graph for {name}."""\nfrom __future__ import annotations\n\nfrom functools import lru_cache\n\nfrom langgraph.graph import END, StateGraph\n\nfrom agentomatic import BaseAgentState\n\nfrom . import nodes\n\n\ndef build_graph() -> StateGraph:\n g = StateGraph(BaseAgentState)\n g.add_node("process", nodes.process)\n g.set_entry_point("process")\n g.add_edge("process", END)\n return g\n\n\n@lru_cache(maxsize=1)\ndef get_graph():\n return build_graph().compile()\n''' + return f'''"""LangGraph graph for {name}."""\nfrom __future__ import annotations\n\nfrom functools import lru_cache\n\nfrom agentomatic import BaseAgentState\nfrom langgraph.graph import END, StateGraph\n\nfrom . import nodes\n\n\ndef build_graph() -> StateGraph:\n g = StateGraph(BaseAgentState)\n g.add_node("process", nodes.process)\n g.set_entry_point("process")\n g.add_edge("process", END)\n return g\n\n\n@lru_cache(maxsize=1)\ndef get_graph():\n return build_graph().compile()\n''' def _legacy_nodes_py(name: str) -> str: @@ -519,8 +534,11 @@ def classify(self, state: {title}State) -> {title}State: TODO: Replace keyword matching with an LLM classifier. """ query = state.request.lower() - # Add your routing logic here - state.classification = "default" + # Add your routing logic here, e.g.: + # if "invoice" in query: + # state.classification = "billing" + # return state + state.classification = "billing" if "invoice" in query else "default" return state def route(self, state: {title}State) -> {title}State: @@ -529,7 +547,7 @@ def route(self, state: {title}State) -> {title}State: tools = get_handoff_tools() if not tools: - state.output = {{"response": f"No delegation targets configured"}} + state.output = {{"response": "No delegation targets configured"}} return state # Pick the first tool as default, or match by classification @@ -855,7 +873,7 @@ async def evaluate(dataset_path: str, split: str = "all") -> None: # Aggregate scores print("\\n" + "=" * 60) - print(f"\\n Pipeline: {name}") + print("\\n Pipeline: {name}") print(f" Examples: {{len(examples)}}") print(f" Passed: {{len(examples) - failures}}") print(f" Failed: {{failures}}") @@ -866,7 +884,7 @@ async def evaluate(dataset_path: str, split: str = "all") -> None: for key in all_scores[0]: values = [s.get(key, 0.0) for s in all_scores] avg_scores[key] = sum(values) / len(values) - print(f"\\n Average scores:") + print("\\n Average scores:") for k, v in avg_scores.items(): print(f" {{k}}: {{v:.3f}}") @@ -920,7 +938,6 @@ def _pipeline_optimize_py(name: str) -> str: import argparse import asyncio import json -import time from itertools import product from pathlib import Path from typing import Any @@ -1162,18 +1179,18 @@ def main() -> None: print("Usage: python -m pipelines.{name}.run \\"your query\\"") sys.exit(1) - print(f"Running pipeline '{name}'...") + print("Running pipeline '{name}'...") print(f" Query: {{query}}\\n") result = asyncio.run(run_via_api(query)) print(f" Status: {{result.get('status', 'unknown')}}") print(f" Duration: {{result.get('duration_ms', 0):.0f}}ms") - print(f"\\n Output:") + print("\\n Output:") print(json.dumps(result.get("output", {{}}), indent=4)) if result.get("steps"): - print(f"\\n Steps:") + print("\\n Steps:") for step_name, step_data in result["steps"].items(): status = step_data.get("status", "?") dur = step_data.get("duration_ms", 0) @@ -1250,8 +1267,10 @@ def _plugin_py(name: str) -> str: from __future__ import annotations from typing import Any -from pydantic import BaseModel, Field + from agentomatic.plugins import BaseMLPlugin +from pydantic import BaseModel, Field + class {title}Input(BaseModel): """Input schema for {name}.""" @@ -1265,12 +1284,21 @@ class {title}Output(BaseModel): class {title}Plugin(BaseMLPlugin[{title}Input, {title}Output]): """Classical ML model wrapper for {name}.""" + # Without these the plugin inherits BaseMLPlugin's "default_plugin" name, + # so it mounts at /api/v1/plugins/default_plugin/* and a second scaffolded + # plugin would collide with it. + plugin_name = "{name}" + plugin_version = "1.0.0" + async def load_model(self) -> None: """Load the ML model weights into memory. This is called automatically during platform startup. """ # TODO: Load your model here (e.g., joblib.load, torch.load) self.model = "dummy_model_instance" + # Marks the plugin ready. Without it /predict answers 503 and /health + # reports the platform as "degraded". + await super().load_model() async def predict(self, inputs: {title}Input) -> {title}Output: """Run inference using the loaded model.""" @@ -1315,7 +1343,10 @@ def _plugin_train_py(name: str) -> str: import logging from pathlib import Path -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(levelname)s - %(message)s", +) logger = logging.getLogger(__name__) DATA_DIR = Path(__file__).parent @@ -1336,7 +1367,9 @@ def load_data(filepath: str | Path) -> list[dict]: return [json.loads(line) for line in f if line.strip()] -def train_eval_split(data: list[dict], eval_fraction: float = 0.2) -> tuple[list[dict], list[dict]]: +def train_eval_split( + data: list[dict], eval_fraction: float = 0.2 +) -> tuple[list[dict], list[dict]]: """Deterministic tail-split (last N% held out for eval).""" if not data: return data, [] @@ -1377,7 +1410,9 @@ def train() -> None: def _plugin_eval_py(name: str) -> str: - title = name.replace("_", "").title() + # Must match _plugin_py exactly: "ag_plugin" -> "AgPlugin" (not "Agplugin"), + # or the generated imports reference classes that were never defined. + title = name.replace("_", " ").title().replace(" ", "") return f'''"""Evaluation script for {name} ML plugin. Loads the plugin, runs it against the labelled dataset, and computes a @@ -1399,7 +1434,10 @@ def _plugin_eval_py(name: str) -> str: from .plugin import {title}Input, {title}Plugin -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(levelname)s - %(message)s", +) logger = logging.getLogger(__name__) DATA_DIR = Path(__file__).parent @@ -1431,10 +1469,18 @@ def to_label(result: str) -> int | None: return None -def weighted_score(component_scores: dict[str, float], weights: dict[str, float]) -> float: +def weighted_score( + component_scores: dict[str, float], weights: dict[str, float] +) -> float: """Return a weight-normalised composite score across components.""" total_w = sum(weights.get(k, 0.0) for k in component_scores) or 1.0 - return sum(component_scores.get(k, 0.0) * weights.get(k, 0.0) for k in component_scores) / total_w + return ( + sum( + component_scores.get(k, 0.0) * weights.get(k, 0.0) + for k in component_scores + ) + / total_w + ) async def evaluate() -> None: @@ -1444,7 +1490,9 @@ async def evaluate() -> None: await plugin.load_model() if not DATASET.exists(): - logger.error("No dataset at %s — add labelled JSONL rows before evaluating.", DATASET) + logger.error( + "No dataset at %s — add labelled JSONL rows before evaluating.", DATASET + ) raise SystemExit(1) examples = load_data(DATASET) @@ -1498,7 +1546,6 @@ def _plugin_optimize_py(name: str) -> str: from __future__ import annotations import logging -import sys logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -1520,10 +1567,14 @@ def optimize() -> None: def _plugin_predict_py(name: str) -> str: - title = name.replace("_", "").title() + # Must match _plugin_py exactly: "ag_plugin" -> "AgPlugin" (not "Agplugin"), + # or the generated imports reference classes that were never defined. + title = name.replace("_", " ").title().replace(" ", "") return f'''"""Local inference script for {name} plugin.""" import asyncio -from .plugin import {title}Plugin, {title}Input + +from .plugin import {title}Input, {title}Plugin + async def predict(text: str): plugin = {title}Plugin() @@ -1606,6 +1657,8 @@ def _train_py(name: str) -> str: AGENTOMATIC_STACK=gemini uv run python agents/{name}/train.py \\ --augment --n-examples 40 --persist --optimizer rewrite """ +# ruff: noqa: E402, I001 - the sys.path bootstrap below must run before the +# project imports, so those imports cannot sit at the top of the file. from __future__ import annotations import os @@ -1689,7 +1742,8 @@ def main(argv: list[str] | None = None) -> int: # ) # history = fit_agent(compiled, data, epochs=cli.epochs, trials=cli.trials) # scores = evaluate_agent(compiled, data.test or data.validation).scores - # generate_fit_report(compiled.fit_result, output_path=HERE / "reports" / f"train_{{AGENT}}.html", + # generate_fit_report(compiled.fit_result, + # output_path=HERE / "reports" / f"train_{{AGENT}}.html", # keras_history=history.history, eval_scores=scores) return 0 @@ -1714,6 +1768,8 @@ def _eval_py(name: str) -> str: AGENTOMATIC_STACK=gemini uv run python agents/{name}/eval.py \\ --split test --prefer-augmented --limit 3 """ +# ruff: noqa: E402, I001 - the sys.path bootstrap below must run before the +# project imports, so those imports cannot sit at the top of the file. from __future__ import annotations import os @@ -1751,7 +1807,11 @@ def main(argv: list[str] | None = None) -> int: apply_stack_defaults(stacks) # --- agent --- - llm = None if not cli.judge else get_llm_for_agent(AGENT, role="default", stack_manager=stacks) + llm = ( + None + if not cli.judge + else get_llm_for_agent(AGENT, role="default", stack_manager=stacks) + ) agent = {title}Agent(llm=llm) if cli.compiled: agent.load_compiled(cli.compiled) @@ -1821,8 +1881,6 @@ def _optimize_py(name: str) -> str: import os from pathlib import Path -from .agent import {title}Agent - from agentomatic.agents import AgentDataset from agentomatic.agents.metrics import ( CallableMetric, @@ -1832,6 +1890,8 @@ def _optimize_py(name: str) -> str: ) from agentomatic.agents.optimizers import GridSearchOptimizer, PromptFitterBridge +from .agent import {title}Agent + DATA_DIR = Path(__file__).parent COMPILED_DIR = Path("compiled") / "{name}" DEFAULT_SEARCH_SPACE = DATA_DIR / "search_space.yaml" @@ -1969,7 +2029,10 @@ def main() -> None: ) parser.add_argument( "--search-space", type=str, default=None, - help="Path to a search_space.yaml (defaults to agents/{name}/search_space.yaml)" + help=( + "Path to a search_space.yaml " + "(defaults to agents/{name}/search_space.yaml)" + ) ) args = parser.parse_args() @@ -2052,7 +2115,9 @@ def main() -> None: # Batch mode if args.input: input_path = Path(args.input) - output_path = Path(args.output) if args.output else DATA_DIR / "predictions.jsonl" + output_path = ( + Path(args.output) if args.output else DATA_DIR / "predictions.jsonl" + ) queries = [] with open(input_path) as f: @@ -2066,7 +2131,9 @@ def main() -> None: result = agent.transform(query_data) results.append({{"input": query_data, "output": result, "status": "ok"}}) except Exception as exc: - results.append({{"input": query_data, "error": str(exc), "status": "error"}}) + results.append( + {{"input": query_data, "error": str(exc), "status": "error"}} + ) print(f" [{{i}}/{{len(queries)}}] done") with open(output_path, "w") as f: @@ -2210,8 +2277,6 @@ def _endpoint_py(name: str) -> str: """ from __future__ import annotations -from pydantic import BaseModel, Field - from agentomatic.endpoints import ( AggregationStrategy, AuthType, @@ -2219,12 +2284,15 @@ def _endpoint_py(name: str) -> str: UpstreamAuthConfig, UpstreamConfig, ) +from pydantic import BaseModel, Field class {title}Request(BaseModel): """Input schema for {name}.""" - payload: dict = Field(default_factory=dict, description="Data forwarded to upstreams.") + payload: dict = Field( + default_factory=dict, description="Data forwarded to upstreams." + ) class {title}Endpoint(BaseEndpoint): @@ -2472,9 +2540,8 @@ def _ingestor_py(name: str) -> str: """ from __future__ import annotations -from pydantic import BaseModel, Field - from agentomatic.ingestion import BaseIngestor, IngestionResult +from pydantic import BaseModel, Field class {title}Request(BaseModel): @@ -2509,7 +2576,14 @@ async def ingest(self, request: {title}Request, ctx) -> IngestionResult: import pymupdf4llm # your PDF -> markdown lib from langchain_text_splitters import MarkdownTextSplitter - markdown = pymupdf4llm.to_markdown(request.source) + from agentomatic.ingestion.paths import resolve_within_root + + # ``request.source`` arrives over HTTP. Resolve it through + # resolve_within_root() (see the import above) or your ingestor + # becomes an arbitrary file read: a caller can pass + # "/etc/passwd", or "../../" out of any directory you intended. + source = resolve_within_root(request.source, description="source") + markdown = pymupdf4llm.to_markdown(str(source)) chunks = MarkdownTextSplitter().split_text(markdown) upserted = 0 @@ -2677,11 +2751,17 @@ def load(self, state: {title}State) -> {title}State: if not raw: return state candidate = Path(raw).expanduser() - if candidate.exists() and candidate.is_file() and candidate.stat().st_size < 5_000_000: + if ( + candidate.exists() + and candidate.is_file() + and candidate.stat().st_size < 5_000_000 + ): try: state.markdown = candidate.read_text(encoding="utf-8") except UnicodeDecodeError: - state.markdown = candidate.read_bytes().decode("utf-8", errors="replace") + state.markdown = candidate.read_bytes().decode( + "utf-8", errors="replace" + ) return state def extract(self, state: {title}State) -> {title}State: @@ -2826,7 +2906,12 @@ def _langchain_init_py(name: str, description: str, keywords: str) -> str: def _langchain_agent_py(name: str) -> str: title = name.replace("_", " ").title().replace(" ", "") return f'''"""LangChain-native agent: {name}. -Uses native LangChain abstractions with agentomatic integration. + +Demonstrates the full set of LangChain abstractions inside an agentomatic +class agent: ``ChatPromptTemplate`` + ``MessagesPlaceholder``, an LCEL chain +(``prompt | llm``), real ``HumanMessage``/``AIMessage`` objects, and an +explicit ``RunnableConfig`` threaded into the chain invocation so tracing/ +callbacks work the same way they would in a hand-rolled LangGraph app. """ from __future__ import annotations @@ -2834,13 +2919,19 @@ def _langchain_agent_py(name: str) -> str: from typing import Any from agentomatic.agents import BaseGraphAgent -from agentomatic.langchain_adapter import dict_to_messages, serialize_messages +from agentomatic.langchain_adapter import ( + dict_to_messages, + make_config, + serialize_messages, +) +from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder @dataclass class {title}State: request: str = "" messages: list[Any] = field(default_factory=list) + thread_id: str | None = None response: str = "" output: dict[str, Any] = field(default_factory=dict) @@ -2858,6 +2949,18 @@ def __init__(self, *, llm: Any = None, prompt_manager: Any = None) -> None: super().__init__() self.llm = llm self.prompt_manager = prompt_manager + # ChatPromptTemplate with a system message + a placeholder for the + # full conversation history — the standard LangChain chat pattern. + self.prompt_template = ChatPromptTemplate.from_messages( + [ + ("system", "{{system_message}}"), + MessagesPlaceholder("messages"), + ] + ) + # LCEL chain: ``prompt | llm``. When no llm is injected (e.g. in + # tests), ``self.chain`` stays None and ``chat()`` falls back to a + # deterministic stub response. + self.chain = self.prompt_template | self.llm if self.llm is not None else None def _system_prompt(self) -> str: return self.resolve_system_prompt( @@ -2872,7 +2975,7 @@ def build_graph(self): return g.compile() def chat(self, state: {title}State) -> {title}State: - prompt = self._system_prompt() + system_message = self._system_prompt() # Normalise REST/Studio dict messages → LangChain message objects. # Fall back to current_query/request so optimize/invoke paths work # when the payload has no prior message history. @@ -2882,21 +2985,26 @@ def chat(self, state: {title}State) -> {title}State: lc_messages = dict_to_messages({{"current_query": state.request}}) else: lc_messages = [] - if self.llm is not None: + + if self.chain is not None: + # RunnableConfig carries tracing tags / thread_id through to the + # underlying LLM call, same as a hand-rolled LangGraph node would. + config = make_config(thread_id=state.thread_id, tags=["{name}"]) try: - if lc_messages: - result = self.llm.invoke(lc_messages) - else: - result = self.llm.invoke( - f"{{prompt}}" + "\\n\\nUser: " + f"{{state.request}}" - ) + result = self.chain.invoke( + {{"system_message": system_message, "messages": lc_messages}}, + config=config, + ) text = getattr(result, "content", None) or str(result) except Exception: text = "Response to: " + f"{{state.request}}" else: - text = f"{{prompt}}" + ": Response to '" + f"{{state.request}}" + "'" + text = f"{{system_message}}" + ": Response to '" + f"{{state.request}}" + "'" + state.response = text - state.messages = serialize_messages(lc_messages) if lc_messages else state.messages + state.messages = ( + serialize_messages(lc_messages) if lc_messages else state.messages + ) state.output = {{"response": text, "agent_type": "{name}"}} return state @@ -2904,6 +3012,7 @@ def input_to_state(self, input_data: dict[str, Any]) -> {title}State: return {title}State( request=input_data.get("current_query", ""), messages=input_data.get("messages", []), + thread_id=input_data.get("thread_id"), ) def state_to_output(self, state: {title}State) -> dict[str, Any]: diff --git a/src/agentomatic/connections/custom.py b/src/agentomatic/connections/custom.py index 9bdaa3f..deaed50 100644 --- a/src/agentomatic/connections/custom.py +++ b/src/agentomatic/connections/custom.py @@ -17,6 +17,7 @@ from __future__ import annotations +import asyncio import importlib import inspect from typing import Any @@ -60,6 +61,7 @@ class CustomConnection: def __init__(self, config: CustomConnectionConfig) -> None: self.config = config self._client: Any = None + self._init_lock = asyncio.Lock() @property def name(self) -> str: @@ -71,21 +73,25 @@ async def initialize(self) -> None: if self._client is not None: return - factory = self.config.factory - if isinstance(factory, str): - factory = import_from_path(resolve_env(factory)) - if not callable(factory): - raise TypeError( - f"Custom connection '{self.name}' factory is not callable: {factory!r}" - ) - - args = resolve_env_deep(self.config.args) - kwargs = resolve_env_deep(self.config.kwargs) - result = factory(*args, **kwargs) - if inspect.isawaitable(result): - result = await result - self._client = result - logger.info(f"🔌 Custom connection '{self.name}' initialized") + async with self._init_lock: + if self._client is not None: + return + + factory = self.config.factory + if isinstance(factory, str): + factory = import_from_path(resolve_env(factory)) + if not callable(factory): + raise TypeError( + f"Custom connection '{self.name}' factory is not callable: {factory!r}" + ) + + args = resolve_env_deep(self.config.args) + kwargs = resolve_env_deep(self.config.kwargs) + result = factory(*args, **kwargs) + if inspect.isawaitable(result): + result = await result + self._client = result + logger.info(f"🔌 Custom connection '{self.name}' initialized") @property def client(self) -> Any: diff --git a/src/agentomatic/connections/database.py b/src/agentomatic/connections/database.py index 9b2156d..faa206a 100644 --- a/src/agentomatic/connections/database.py +++ b/src/agentomatic/connections/database.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any from urllib.parse import quote_plus, urlsplit, urlunsplit @@ -63,6 +64,11 @@ def __init__(self, config: DatabaseConnectionConfig) -> None: self._engine: AsyncEngine | None = None self._sessionmaker: Any = None self._resolved_url: str = "" + # Guards initialize() so two concurrent first-callers (e.g. two + # requests racing at cold start via session()'s lazy init) can't + # both pass the `self._engine is None` check and each build (and + # leak) their own engine/pool. + self._init_lock = asyncio.Lock() @property def name(self) -> str: @@ -84,38 +90,45 @@ async def initialize(self) -> None: if self._engine is not None: return - try: - from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine - except ImportError as exc: # pragma: no cover - raise ImportError( - "SQLAlchemy is required for database connections. " - "Install with: pip install 'agentomatic[db]'" - ) from exc - - cfg = self.config - url = _inject_credentials( - resolve_env(cfg.url), - resolve_env(cfg.username), - resolve_env(cfg.password), - ) - self._resolved_url = url - - engine_kwargs: dict[str, Any] = { - "echo": cfg.echo, - "pool_pre_ping": cfg.pool_pre_ping, - "connect_args": cfg.connect_args, - } - # SQLite async engines do not support pool sizing kwargs. - if not url.startswith("sqlite"): - engine_kwargs.update( - pool_size=cfg.pool_size, - max_overflow=cfg.max_overflow, - pool_timeout=cfg.pool_timeout, + async with self._init_lock: + # Re-check: another coroutine may have finished initializing + # while we were waiting for the lock. + if self._engine is not None: + return + + try: + from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + except ImportError as exc: # pragma: no cover + raise ImportError( + "SQLAlchemy is required for database connections. " + "Install with: pip install 'agentomatic[db]'" + ) from exc + + cfg = self.config + url = _inject_credentials( + resolve_env(cfg.url), + resolve_env(cfg.username), + resolve_env(cfg.password), ) + self._resolved_url = url - self._engine = create_async_engine(url, **engine_kwargs) - self._sessionmaker = async_sessionmaker(self._engine, expire_on_commit=False) - logger.info(f"🗄️ Database connection '{self.name}' initialized") + engine_kwargs: dict[str, Any] = { + "echo": cfg.echo, + "pool_pre_ping": cfg.pool_pre_ping, + "connect_args": cfg.connect_args, + } + # SQLite async engines do not support pool sizing kwargs. + if not url.startswith("sqlite"): + engine_kwargs.update( + pool_size=cfg.pool_size, + max_overflow=cfg.max_overflow, + pool_timeout=cfg.pool_timeout, + ) + + engine = create_async_engine(url, **engine_kwargs) + self._sessionmaker = async_sessionmaker(engine, expire_on_commit=False) + self._engine = engine + logger.info(f"🗄️ Database connection '{self.name}' initialized") @asynccontextmanager async def session(self) -> AsyncGenerator[AsyncSession, None]: diff --git a/src/agentomatic/connections/models.py b/src/agentomatic/connections/models.py index fa7582e..39f47bf 100644 --- a/src/agentomatic/connections/models.py +++ b/src/agentomatic/connections/models.py @@ -83,9 +83,7 @@ class HttpConnectionConfig(BaseModel): ) base_url: str = Field(..., description="Base URL of the service (supports ${ENV}).") headers: dict[str, str] = Field(default_factory=dict) - auth: UpstreamAuthConfig = Field( - default_factory=lambda: UpstreamAuthConfig.model_construct() - ) + auth: UpstreamAuthConfig = Field(default_factory=lambda: UpstreamAuthConfig.model_construct()) timeout: float = Field(30.0, gt=0) max_retries: int = Field(2, ge=0, le=10) verify_ssl: bool = Field(True) diff --git a/src/agentomatic/control/router.py b/src/agentomatic/control/router.py index 5e8ce1d..8a53905 100644 --- a/src/agentomatic/control/router.py +++ b/src/agentomatic/control/router.py @@ -11,6 +11,7 @@ from __future__ import annotations +import hmac from typing import TYPE_CHECKING, Any from fastapi import APIRouter, Header, HTTPException @@ -26,6 +27,7 @@ ToggleResponse, ) from agentomatic.control.state import ControlPlaneState +from agentomatic.core.errors import client_safe_message if TYPE_CHECKING: from agentomatic.core.platform import AgentPlatform @@ -52,9 +54,25 @@ def create_control_router( def _authorize(token: str | None) -> None: """Reject mutating requests lacking the configured control token.""" - if control_token and token != control_token: + # Bytes comparison — see the note in agentomatic.middleware.auth: a + # non-ASCII token would otherwise raise TypeError and surface as 500. + if control_token and ( + not token or not hmac.compare_digest(token.encode(), control_token.encode()) + ): raise HTTPException(status_code=401, detail="Invalid or missing control token") + def _agent_aliases(agent: Any) -> set[str]: + """Return every path segment an agent's routes are mounted under. + + Agents are mounted under both their folder name and their manifest + slug when the two differ (see ``_mount_agent_router`` in + ``core/platform.py``) — both must be tracked so enable/disable + can't be bypassed via whichever alias wasn't targeted. + """ + manifest = getattr(agent, "manifest", None) + aliases = {getattr(manifest, "name", None), getattr(manifest, "slug", None)} + return {alias for alias in aliases if alias} + def _agent_policy(agent: Any) -> tuple[bool, list[str], list[str]]: """Extract (requires_auth, roles, scopes) from an agent policy.""" policy = getattr(agent, "security_policy", None) @@ -100,7 +118,12 @@ async def _agent_health(agent: Any, *, timeout: float = 2.0) -> dict[str, Any]: except TimeoutError: return {"status": "timeout", "error": "health_check timed out"} except Exception as exc: # noqa: BLE001 - return {"status": "error", "error": str(exc)} + # These control *reads* are not token-gated, so never echo raw + # exception text from a backend health check. + return { + "status": "error", + "error": client_safe_message(exc, context="health_check failed"), + } @router.get("/agents", response_model=list[ControlAgentInfo], summary="List agents") async def list_agents() -> list[ControlAgentInfo]: @@ -266,9 +289,15 @@ async def disable_agent( ) -> ToggleResponse: """Stop routing traffic to an agent (returns 503 for its routes).""" _authorize(x_control_token) - if platform._registry.get(name) is None: + agent = platform._registry.get(name) + if agent is None: raise HTTPException(404, f"Agent '{name}' not found") - state.disable_agent(name) + # An agent is mounted under BOTH its folder name and its manifest + # slug when they differ (so Studio, which addresses agents by slug, + # doesn't 404). Disable both aliases, or draining by whichever one + # the operator didn't use leaves the other alias's routes live. + for alias in _agent_aliases(agent): + state.disable_agent(alias) logger.warning(f"🛑 Control plane: agent '{name}' disabled") return ToggleResponse(target=name, state="disabled") @@ -283,9 +312,11 @@ async def enable_agent( ) -> ToggleResponse: """Resume routing traffic to a previously disabled agent.""" _authorize(x_control_token) - if platform._registry.get(name) is None: + agent = platform._registry.get(name) + if agent is None: raise HTTPException(404, f"Agent '{name}' not found") - state.enable_agent(name) + for alias in _agent_aliases(agent): + state.enable_agent(alias) logger.info(f"✅ Control plane: agent '{name}' enabled") return ToggleResponse(target=name, state="enabled") diff --git a/src/agentomatic/core/agent_invoke.py b/src/agentomatic/core/agent_invoke.py index 1b54020..2147003 100644 --- a/src/agentomatic/core/agent_invoke.py +++ b/src/agentomatic/core/agent_invoke.py @@ -127,11 +127,18 @@ def _input_from_state(state: dict[str, Any]) -> dict[str, Any]: Dict passed to ``BaseGraphAgent.input_to_state`` / ``atransform``. """ query = state.get("current_query", state.get("query", "")) - # Drop conversation bookkeeping that agents rarely want as input fields. - skip = {"messages", "thread_id"} + # ``messages`` and ``thread_id`` are forwarded (they used to be dropped as + # "conversation bookkeeping"). Both are load-bearing for class agents built + # on LangChain: ``messages`` feeds a ``MessagesPlaceholder`` so the model + # sees prior turns, and ``thread_id`` feeds ``RunnableConfig``'s + # ``configurable.thread_id`` for checkpointing/tracing. Dropping them made + # the scaffolded template's ``input_to_state`` read values that were + # guaranteed empty on every HTTP path. ``input_to_state`` is agent-authored + # and only reads the keys it asks for, so forwarding extra keys is inert + # for agents that ignore them. payload: dict[str, Any] = { "query": query, - **{k: v for k, v in state.items() if k not in skip and k != "context"}, + **{k: v for k, v in state.items() if k != "context"}, } context = state.get("context") if isinstance(context, dict): diff --git a/src/agentomatic/core/errors.py b/src/agentomatic/core/errors.py new file mode 100644 index 0000000..ea3647c --- /dev/null +++ b/src/agentomatic/core/errors.py @@ -0,0 +1,84 @@ +"""Client-safe error reporting. + +An exception's ``str()`` routinely carries operational detail that should not +leave the process: database drivers put full DSNs (credentials included) in +their messages, HTTP clients echo request URLs with tokens in the query string, +and auth libraries name internal hosts. Interpolating ``{exc}`` straight into +an HTTP response therefore leaks secrets to anyone who can trigger a failure. + +:func:`client_safe_detail` logs the full exception server-side and returns a +sanitised payload carrying a short *error id* that correlates the response to +that log line — so operators keep full diagnostics without publishing them. + +Set ``AGENTOMATIC_DEBUG_ERRORS=1`` to include raw exception text in responses +while developing locally. +""" + +from __future__ import annotations + +import os +import uuid +from typing import Any + +from loguru import logger + +#: Opt-in flag that puts raw exception text back into HTTP responses. +DEBUG_ERRORS_ENV = "AGENTOMATIC_DEBUG_ERRORS" + + +def debug_errors_enabled() -> bool: + """Whether raw exception text may be returned to clients.""" + return os.getenv(DEBUG_ERRORS_ENV, "").strip().lower() in {"1", "true", "yes", "on"} + + +def client_safe_detail( + exc: BaseException, + *, + context: str, + extra: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Log *exc* in full and return a payload that is safe to send to a client. + + Args: + exc: The exception being handled. + context: Short description of the failing operation, e.g. + ``"Agent invocation failed"``. This is caller-authored text, so it + is always safe to return. + extra: Additional non-sensitive fields to merge into the payload. + + Returns: + ``{"error": , "error_id": , "error_type": }``, + plus ``"detail"`` with the raw message when debug errors are enabled. + """ + error_id = uuid.uuid4().hex[:12] + # Full detail (including traceback) goes to the server log only. + logger.opt(exception=exc).error(f"{context} [error_id={error_id}]") + + payload: dict[str, Any] = { + "error": context, + "error_id": error_id, + "error_type": type(exc).__name__, + } + if extra: + payload.update(extra) + if debug_errors_enabled(): + payload["detail"] = str(exc) + return payload + + +def client_safe_message( + exc: BaseException, + *, + context: str, +) -> str: + """Log *exc* in full and return a one-line message safe to hand a client. + + The string form of :func:`client_safe_detail`, for places that persist an + error into a plain ``str`` field — a background task record, a per-item + batch result — which is later served over HTTP. The correlation id ties it + back to the full server-side log entry. + """ + payload = client_safe_detail(exc, context=context) + if "detail" in payload: # debug mode: raw text is allowed through + return f"{context}: {payload['detail']} [error_id={payload['error_id']}]" + return f"{context} ({payload['error_type']}) [error_id={payload['error_id']}]" diff --git a/src/agentomatic/core/lifespan.py b/src/agentomatic/core/lifespan.py index 30ceaf4..cba2198 100644 --- a/src/agentomatic/core/lifespan.py +++ b/src/agentomatic/core/lifespan.py @@ -30,7 +30,11 @@ def configure_logging(level: str = "INFO") -> None: format=( "{time:YYYY-MM-DD HH:mm:ss.SSS} | " "{level: <8} | " - "{name}:{function}:{line} — " + # Loguru's built-in format uses " - ", and lines emitted before + # this sink is installed still use it. Matching it keeps one + # separator across the whole log (and keeps it ASCII, which is + # kinder to log-shipping regexes than an em dash). + "{name}:{function}:{line} - " "{message}" ), ) diff --git a/src/agentomatic/core/memory_manager.py b/src/agentomatic/core/memory_manager.py index 0b9ef07..6ff1194 100644 --- a/src/agentomatic/core/memory_manager.py +++ b/src/agentomatic/core/memory_manager.py @@ -36,6 +36,7 @@ HumanMessage: type[Any] = _lc_messages.HumanMessage AIMessage: type[Any] = _lc_messages.AIMessage SystemMessage: type[Any] = _lc_messages.SystemMessage + ToolMessage: type[Any] = _lc_messages.ToolMessage except ImportError: HAS_LANGCHAIN = False @@ -61,13 +62,20 @@ class _SystemMessage(_BaseMessage): def __init__(self, content: str = "", **kwargs: Any) -> None: super().__init__(content=content, type="system", **kwargs) + class _ToolMessage(_BaseMessage): + def __init__(self, content: str = "", **kwargs: Any) -> None: + self.tool_call_id = kwargs.pop("tool_call_id", "") + self.name = kwargs.pop("name", None) + super().__init__(content=content, type="tool", **kwargs) + # Rebind the public names so the rest of this module (and callers) can - # use ``HumanMessage`` / ``AIMessage`` / ``SystemMessage`` uniformly - # whether or not langchain_core is installed. + # use ``HumanMessage`` / ``AIMessage`` / ``SystemMessage`` / ``ToolMessage`` + # uniformly whether or not langchain_core is installed. BaseMessage = _BaseMessage HumanMessage = _HumanMessage AIMessage = _AIMessage SystemMessage = _SystemMessage + ToolMessage = _ToolMessage # Default summary system prompt @@ -362,17 +370,37 @@ async def _get_or_generate_summary( def _convert_to_langchain_messages( messages: list[dict[str, Any]], ) -> list[Any]: - """Convert stored message dicts to LangChain BaseMessage objects.""" + """Convert stored message dicts to LangChain BaseMessage objects. + + Preserves tool-calling metadata: an assistant turn keeps its + ``tool_calls``, and a ``tool`` turn becomes a real ``ToolMessage`` + carrying its ``tool_call_id``. Without that, a tool-using agent's + history is silently corrupted (the tool result degrades to a plain + human turn and the call/result pairing is lost), which the provider + rejects on the next turn. + """ result: list[Any] = [] for msg in messages: - role = msg.get("role", "user") + role = str(msg.get("role", "user")).lower() content = msg.get("content", "") - if role == "user": + if role in ("user", "human"): result.append(HumanMessage(content=content)) - elif role == "assistant": - result.append(AIMessage(content=content)) + elif role in ("assistant", "ai"): + tool_calls = msg.get("tool_calls") + if tool_calls: + result.append(AIMessage(content=content, tool_calls=tool_calls)) + else: + result.append(AIMessage(content=content)) elif role == "system": result.append(SystemMessage(content=content)) + elif role == "tool": + result.append( + ToolMessage( + content=content, + tool_call_id=str(msg.get("tool_call_id", "")), + name=msg.get("name"), + ) + ) else: # Default to HumanMessage for unknown roles result.append(HumanMessage(content=content)) diff --git a/src/agentomatic/core/platform.py b/src/agentomatic/core/platform.py index 36c149f..4be48d2 100644 --- a/src/agentomatic/core/platform.py +++ b/src/agentomatic/core/platform.py @@ -60,10 +60,17 @@ def _mount_agent_router(app: FastAPI, agent: RegisteredAgent, name: str, api_pre ) slug = getattr(agent, "slug", None) or getattr(getattr(agent, "manifest", None), "slug", None) if slug and slug != name: + # The slug mount is a compatibility alias for the canonical + # ``{api_prefix}/{name}`` routes above. Keep it out of the OpenAPI + # schema: documenting both copies doubled the advertised surface and + # produced a duplicate ``operationId`` for every route (FastAPI emits + # one UserWarning each, and duplicate ids break client codegen). + # The alias still routes normally at runtime. app.include_router( agent.router, prefix=f"{api_prefix}/{slug}", tags=[_agent_tag(name)], + include_in_schema=False, ) @@ -254,6 +261,7 @@ def __init__( enable_rate_limit: bool = False, rate_limit_requests: int = 100, rate_limit_window: int = 60, + rate_limit_trust_proxy_headers: bool = False, enable_metrics: bool = False, enable_feedback: bool = True, enable_telemetry: bool = True, @@ -334,6 +342,13 @@ def __init__( allow_logsllm_analysis: When ``True``, expose LLM log-analysis endpoints that score recent logs and return recommendations. """ + # Apply the requested level before anything is logged. Construction + # and ``build()`` narrate settings loading, discovery, and every mount, + # while the lifespan's ``configure_logging`` only runs at startup — so + # without this, ``log_level="WARNING"`` (what ``--profile minimal`` + # bakes into the image) still printed every INFO and DEBUG line. + configure_logging(log_level) + self.agents_dir = Path(agents_dir).resolve() self.plugins_dir = Path(plugins_dir).resolve() self.endpoints_dir = Path(endpoints_dir).resolve() @@ -362,6 +377,7 @@ def __init__( self._enable_rate_limit = enable_rate_limit self._rate_limit_requests = rate_limit_requests self._rate_limit_window = rate_limit_window + self._rate_limit_trust_proxy_headers = rate_limit_trust_proxy_headers self._enable_metrics = enable_metrics self._enable_feedback = enable_feedback self._enable_telemetry = enable_telemetry @@ -571,6 +587,21 @@ def register_agent( **kwargs: Extra keyword arguments forwarded to :class:`RegisteredAgent`. """ + # A class agent registered as ``class_instance=MyAgent()`` already owns + # a compiled graph. Studio's graph view *and* its run streaming both + # key off ``graph_fn``, so without deriving it here the agent shows an + # empty topology and streaming fails with "Agent has no graph_fn". + # (``BaseGraphAgent.as_registered_agent()`` sets it; this is the + # programmatic-registration path catching up.) + class_instance = kwargs.get("class_instance") + if graph_fn is None and class_instance is not None: + instance_graph = getattr(class_instance, "graph", None) + if instance_graph is not None: + + def graph_fn() -> Any: # noqa: D401 - adapter closure + """Return the class agent's compiled graph.""" + return class_instance.graph + agent = RegisteredAgent( manifest=manifest, node_fn=node_fn, @@ -578,6 +609,9 @@ def register_agent( **kwargs, ) self._registry._agents[manifest.name] = agent # noqa: SLF001 + # Keep the slug index in sync so ``registry.get(slug)`` resolves via the + # index (matching folder-discovered agents) rather than the fallback scan. + self._registry._index_slug(manifest.name, agent) # noqa: SLF001 logger.info(f" ✅ Programmatically registered: {manifest.name} ({manifest.slug})") # ------------------------------------------------------------------ @@ -757,6 +791,53 @@ def _resolve_database_url(self) -> str | None: return None return url + def _resolve_jwt_config(self) -> Any: + """Build a :class:`JWTConfig` from ``AUTH__*`` env vars or the stack. + + Verified JWT auth was documented as configurable "via stack", and + ``agentomatic deploy`` writes ``AUTH__JWKS_URL`` / ``AUTH__ISSUER`` / + ``AUTH__AUDIENCE`` into the generated ``.env`` — but nothing read any + of them. Only the in-process ``jwt_config=`` kwarg reached the + middleware, so a deployed container running the scaffolded ``main.py`` + had no way to turn on signature verification at all. + + Environment wins over the stack, matching the database-URL resolution + above; ``${VAR}`` placeholders in stack values are expanded. + + Returns: + A ``JWTConfig`` when a JWKS endpoint is configured, else ``None``. + """ + import os + + from agentomatic.security.jwt_auth import JWTConfig + + def _expand(value: str) -> str: + value = (value or "").strip() + if value.startswith("${") and value.endswith("}"): + value = (os.getenv(value[2:-1]) or "").strip() + return "" if "${" in value else value + + stack_auth = getattr( + getattr(getattr(self, "_stack_manager", None), "_active_stack", None), "auth", None + ) + + def _setting(env_key: str, field: str) -> str: + from_env = (os.getenv(env_key) or "").strip() + if from_env: + return from_env + return _expand(str(getattr(stack_auth, field, "") or "")) + + jwks_url = _setting("AUTH__JWKS_URL", "jwks_url") + if not jwks_url: + return None + + return JWTConfig( + enabled=True, + jwks_url=jwks_url, + issuer=_setting("AUTH__ISSUER", "issuer"), + audience=_setting("AUTH__AUDIENCE", "audience"), + ) + async def _auto_derive_store_from_connections(self) -> None: """Populate ``self._store`` from the first MEMORY connection, if any.""" from agentomatic.connections.manager import PLATFORM_SCOPE, all_managers @@ -990,6 +1071,12 @@ async def lifespan(app: FastAPI): # noqa: ARG001 for name, plugin in platform._plugin_registry.list_plugins().items(): try: await plugin.load_model() + # A subclass that overrides load_model() without calling + # super() leaves _is_loaded False, so /predict answers 503 + # and /health reports "degraded" — while this line claimed + # success. Stamp it here so a forgotten super() can't + # silently produce a permanently unusable plugin. + plugin.mark_loaded() logger.info(f" ✅ Plugin '{name}' loaded successfully") except Exception as e: logger.error(f" ❌ Failed to load plugin '{name}': {e}") @@ -1171,16 +1258,10 @@ async def lifespan(app: FastAPI): # noqa: ARG001 app.add_middleware(LoggingMiddleware) - # Auth - if self._enable_auth and self._auth_api_key: - from agentomatic.middleware.auth import AuthMiddleware - - app.add_middleware(AuthMiddleware, api_key=self._auth_api_key) - logger.info("🔒 Auth middleware enabled") - - # Zero Trust Enforcer (v0.6) — added BEFORE JWT so that (thanks to - # Starlette's reverse middleware ordering) the JWT middleware runs - # first and populates ``request.state.jwt_claims`` before enforcement. + # Zero Trust Enforcer (v0.6) — added BEFORE the authentication + # middlewares so that (thanks to Starlette's reverse middleware + # ordering) JWT *and* API-key auth both run first and record the + # caller's identity on ``request.state`` before enforcement. # Per-agent enforcement is opt-in via each agent's ``security.py`` # policy (``require_auth`` / ``allowed_roles`` / ``allowed_scopes``). if self._enable_zero_trust: @@ -1209,32 +1290,54 @@ async def lifespan(app: FastAPI): # noqa: ARG001 if self._enable_jwt_auth: from agentomatic.security.jwt_auth import JWTAuthMiddleware, JWTConfig - jwt_cfg = self._jwt_config or JWTConfig(enabled=True) + jwt_cfg = self._jwt_config or self._resolve_jwt_config() or JWTConfig(enabled=True) # Under the global auth lock, signature-disabled (dev) JWT is a # bypass: forged/unsigned tokens would authenticate EVERY request. # Refuse to start unless real verification (jwks_url) is configured # — or API-key auth guards the platform instead. Raised outside the # try below so the misconfiguration is not silently swallowed. - if self._require_auth_globally and not jwt_cfg.jwks_url and not self._enable_auth: - raise RuntimeError( - "require_auth_globally=True but JWT signature verification " - "is not configured (no jwks_url) and API-key auth is " - "disabled — this would accept forged/unsigned JWTs. Fix by " - "one of: (a) set JWTConfig.jwks_url (with issuer/audience) " - "and pass it via jwt_config=/stack; (b) enable_auth=True " - "with auth_api_key; or (c) drop require_auth_globally for " - "local dev." + if self._require_auth_globally and not jwt_cfg.jwks_url: + if not self._enable_auth: + raise RuntimeError( + "require_auth_globally=True but JWT signature verification " + "is not configured (no jwks_url) and API-key auth is " + "disabled — this would accept forged/unsigned JWTs. Fix by " + "one of: (a) set JWTConfig.jwks_url (with issuer/audience) " + "and pass it via jwt_config=/stack; (b) enable_auth=True " + "with auth_api_key; or (c) drop require_auth_globally for " + "local dev." + ) + # Remedy (b): API-key auth guards the platform. Adding the JWT + # middleware anyway is not merely redundant — under the auth + # lock it demands ``require_signature`` and raises from its own + # constructor, which Starlette runs when it builds the + # middleware stack on the FIRST REQUEST, not here. The app + # would start clean and then 500 on every route including + # /health. Enforce with the API key alone and say so. + logger.warning( + "JWT auth is enabled with require_auth_globally but no " + "jwks_url is configured — enforcing with API-key auth only. " + "Set JWTConfig.jwks_url (with issuer/audience) to verify JWTs." ) - if self._require_auth_globally: - # Enforce signature verification for the global auth lock. - jwt_cfg = jwt_cfg.model_copy(update={"require_signature": True}) + else: + if self._require_auth_globally: + # Enforce signature verification for the global auth lock. + jwt_cfg = jwt_cfg.model_copy(update={"require_signature": True}) - try: app.add_middleware(JWTAuthMiddleware, config=jwt_cfg) logger.info("🔐 JWT auth middleware enabled") - except Exception as exc: - logger.warning(f"JWT auth setup failed: {exc}") + + # API-key auth — registered after zero-trust (so it runs before it) + # and marks the request authenticated. Registering it earlier meant + # zero-trust denied every request before the key was ever checked, so + # API-key auth plus require_auth_globally could not serve a single + # request: valid key, 401 "no valid JWT claims found". + if self._enable_auth and self._auth_api_key: + from agentomatic.middleware.auth import AuthMiddleware + + app.add_middleware(AuthMiddleware, api_key=self._auth_api_key) + logger.info("🔒 Auth middleware enabled") # Rate limiting if self._enable_rate_limit: @@ -1244,6 +1347,7 @@ async def lifespan(app: FastAPI): # noqa: ARG001 RateLimitMiddleware, max_requests=self._rate_limit_requests, window_seconds=self._rate_limit_window, + trust_proxy_headers=self._rate_limit_trust_proxy_headers, ) logger.info( f"🚦 Rate limit: {self._rate_limit_requests} req/{self._rate_limit_window}s" @@ -1816,6 +1920,13 @@ def run( if ssl_certfile or ssl_keyfile: logger.info(f"🔐 HTTPS enabled (certfile={ssl_certfile}, keyfile={ssl_keyfile})") + # The platform's own LoggingMiddleware already logs every request with a + # correlation id, so leaving uvicorn's access log on doubles the line + # count for the same information — real volume (and cost) in a hosted + # log pipeline. An explicit access_log=... from the caller still wins. + if self._enable_logging and "access_log" not in run_kwargs: + run_kwargs["access_log"] = False + # uvicorn requires an import string (re-imported per worker subprocess) # for reload / multi-worker mode. Passing an app *instance* makes modern # uvicorn exit(1). Reconstruct via a module-level factory when possible. diff --git a/src/agentomatic/core/router_factory.py b/src/agentomatic/core/router_factory.py index 4d831c9..c7a165d 100644 --- a/src/agentomatic/core/router_factory.py +++ b/src/agentomatic/core/router_factory.py @@ -14,6 +14,8 @@ from pydantic import BaseModel, ConfigDict, Field 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 # --------------------------------------------------------------------------- # Request / Response Models @@ -55,9 +57,13 @@ class AgentInvokeRequest(BaseModel): query: str = Field(..., description="User query or input") user_id: str = Field("default-user", description="User identifier") context: dict[str, Any] = Field(default_factory=dict, description="Additional context") - thread_id: str | None = Field(default=None, description="Thread ID for conversation continuity") + thread_id: str | None = Field( + default=None, description="Thread ID for conversation continuity" + ) prompt_version: str = Field("v1", description="Prompt version to use") - temperature: float | None = Field(default=None, ge=0.0, le=2.0, description="Temperature override") + temperature: float | None = Field( + default=None, ge=0.0, le=2.0, description="Temperature override" + ) max_tokens: int | None = Field(default=None, ge=1, description="Max tokens override") metadata: dict[str, Any] = Field(default_factory=dict, description="Extra metadata") @@ -98,7 +104,12 @@ class AgentInvokeResponse(BaseModel): "steps_taken", "context", "metadata", - "messages", + # NOTE: "messages" is deliberately NOT filtered. Every other key here + # is surfaced elsewhere in AgentInvokeResponse (or is an echo of the + # request), but the envelope has no ``messages`` field — so filtering + # it silently DROPPED the conversation a class agent returned from + # ``state_to_output()``, which is exactly what a LangChain/chat agent + # produces. It now flows through to ``output.messages``. "current_query", "query", "user_id", @@ -111,7 +122,7 @@ class AgentInvokeResponse(BaseModel): def _json_dumps(value: Any) -> str: """Serialize *value* as JSON text (fallback to ``str`` on failure).""" try: - return json.dumps(value, ensure_ascii=False, default=str) + return json.dumps(value, ensure_ascii=False, default=json_default) except (TypeError, ValueError): return str(value) @@ -141,21 +152,24 @@ def coerce_agent_invoke_payload( text = "" if result is None else str(result) return text, None, {} - context = result.get("context", result.get("retrieved_documents", {})) + context = to_jsonable(result.get("context", result.get("retrieved_documents", {}))) raw_response = result.get("response") # Explicit string response (classic BaseAgentState / chat agents). if isinstance(raw_response, str) and raw_response: extras = {k: v for k, v in result.items() if k not in _FRAMEWORK_RESULT_KEYS} - output: Any | None = extras or None + output: Any | None = to_jsonable(extras) if extras else None return raw_response, output, context # Explicit structured response value. if isinstance(raw_response, (dict, list)): - return _json_dumps(raw_response), raw_response, context + jsonable_response = to_jsonable(raw_response) + return _json_dumps(jsonable_response), jsonable_response, context - # Class-agent ``state_to_output``: whole dict is the payload. - output = dict(result) + # Class-agent ``state_to_output``: whole dict is the payload. May contain raw + # LangChain BaseMessage objects (e.g. ``{"messages": state.messages}``) — make + # it JSON-safe before it becomes an HTTP response body. + output = to_jsonable(dict(result)) text = _human_response_text(output) or _json_dumps(output) return text, output, context if context is not None else {} @@ -205,7 +219,9 @@ class AgentChatRequest(BaseModel): class CreateThreadRequest(BaseModel): """Request to explicitly create a thread.""" - thread_id: str | None = Field(default=None, description="Custom thread ID (auto-generated if omitted)") + thread_id: str | None = Field( + default=None, description="Custom thread ID (auto-generated if omitted)" + ) user_id: str = Field("default-user", description="User identifier") title: str | None = Field(default=None, description="Thread title") metadata: dict[str, Any] = Field(default_factory=dict) @@ -225,6 +241,57 @@ class A2ATaskRequest(BaseModel): metadata: dict[str, Any] = Field(default_factory=dict) +def a2a_message_text(message: dict[str, Any]) -> str: + """Extract the user text from an A2A message, whatever shape it arrives in. + + The A2A protocol carries text in ``message.parts`` — a list of typed part + objects. Agentomatic historically read only ``message.content``, so a + spec-shaped request produced an *empty* query and the agent answered + nothing, with a 200 and no hint that the text had been dropped. + + Accepted shapes (checked in this order): + + - ``{"parts": [{"type"|"kind": "text", "text": "..."}]}`` — the protocol form + (a bare ``{"text": ...}`` part is also accepted) + - ``{"content": "..."}`` — the simplified form Agentomatic documents + - ``{"content": [ ...parts... ]}`` — content carrying parts + - ``{"text": "..."}`` + + Args: + message: The raw ``message`` object from the request body. + + Returns: + The concatenated text, or ``""`` when the message carries none. + """ + + def _from_parts(parts: Any) -> str: + if not isinstance(parts, list): + return "" + texts: list[str] = [] + for part in parts: + if isinstance(part, str): + texts.append(part) + elif isinstance(part, dict): + value = part.get("text") + if isinstance(value, str) and value: + texts.append(value) + return "\n".join(texts) + + text = _from_parts(message.get("parts")) + if text: + return text + + content = message.get("content") + if isinstance(content, str) and content: + return content + text = _from_parts(content) + if text: + return text + + value = message.get("text") + return value if isinstance(value, str) else "" + + class OptimizeInvokeRequest(BaseModel): """Optimization-specific invocation — returns full pipeline context.""" @@ -652,8 +719,10 @@ async def invoke(request: Any) -> Any: status="error", error=str(exc), ) - logger.error(f"Agent {agent_name} invocation failed: {exc}") - raise HTTPException(500, detail={"error": f"Agent invocation failed: {exc}"}) from exc + raise HTTPException( + 500, + detail=client_safe_detail(exc, context="Agent invocation failed"), + ) from exc async def invoke_stream(request: Any) -> StreamingResponse: """Invoke agent with SSE streaming.""" @@ -707,26 +776,26 @@ 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=str)}\n\n" + yield f"data: {json.dumps(event, default=json_default)}\n\n" 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=str)}\n\n" + yield f"data: {json.dumps(result, default=json_default)}\n\n" 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=str)}\n\n" + yield f"data: {json.dumps(event, default=json_default)}\n\n" # 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=str)}\n\n" + yield f"data: {json.dumps(result, default=json_default)}\n\n" if isinstance(result, dict): collected_response = result.get("response", "") collected_output = result @@ -784,7 +853,8 @@ async def event_stream(): status="error", error=str(exc), ) - yield f"data: {json.dumps({'error': str(exc)})}\n\n" + safe = client_safe_detail(exc, context="Agent streaming failed") + yield f"data: {json.dumps(safe)}\n\n" return StreamingResponse( event_stream(), @@ -874,34 +944,16 @@ async def chat(request: AgentChatRequest) -> dict[str, Any]: # ── Load conversation history ──────────────────────────── history_loaded = 0 if request.messages is not None: - # User supplied their own messages — use them directly - try: - from langchain_core.messages import AIMessage, HumanMessage, SystemMessage - - lc_messages: list[Any] = [] - for msg in request.messages: - role = msg.get("role", "user") - content_val = msg.get("content", "") - if role == "assistant": - lc_messages.append(AIMessage(content=content_val)) - elif role == "system": - lc_messages.append(SystemMessage(content=content_val)) - else: - lc_messages.append(HumanMessage(content=content_val)) - lc_messages.append(HumanMessage(content=request.content)) - state["messages"] = lc_messages - except ImportError: - # langchain_core not installed — use plain dicts - lc_messages_plain: list[dict[str, str]] = [] - for msg in request.messages: - lc_messages_plain.append( - { - "role": msg.get("role", "user"), - "content": msg.get("content", ""), - } - ) - lc_messages_plain.append({"role": "user", "content": request.content}) - state["messages"] = lc_messages_plain + # User supplied their own messages — use them directly. + # Delegate to the canonical converter in ``langchain_adapter`` so + # tool-calling metadata survives: a hand-rolled conversion here + # previously dropped ``tool_calls`` and turned a ``tool`` turn into + # a HumanMessage, breaking the call/result pairing that providers + # require on the next turn. It also degrades gracefully to plain + # role/content dicts when langchain_core isn't installed. + state["messages"] = dict_to_messages( + [*request.messages, {"role": "user", "content": request.content}] + ) history_loaded = len(request.messages) elif memory_mgr and request.include_history: try: @@ -922,7 +974,10 @@ async def chat(request: AgentChatRequest) -> dict[str, Any]: history_loaded = max(0, len(messages) - 1) except Exception as exc: logger.warning(f"History loading failed for chat: {exc}") - state["metadata"]["_history_error"] = str(exc) + # Surfaced in the 200 response metadata, so sanitise it too. + state["metadata"]["_history_error"] = client_safe_message( + exc, context="History loading failed" + ) # Run before_node hooks for hook in registry.before_node_hooks: @@ -1035,8 +1090,10 @@ async def chat(request: AgentChatRequest) -> dict[str, Any]: status="error", error=str(exc), ) - logger.error(f"Chat with {agent_name} failed: {exc}") - raise HTTPException(500, detail={"error": f"Chat failed: {exc}"}) from exc + raise HTTPException( + 500, + detail=client_safe_detail(exc, context="Chat failed"), + ) from exc # ── GET /health ─────────────────────────────────────────────── @router.get("/health") @@ -1136,7 +1193,16 @@ async def submit_a2a_task(request: A2ATaskRequest) -> dict[str, Any]: synchronous (blocking) execution for backward compatibility. """ agent = _get_agent() - query = request.message.get("content", "") + query = a2a_message_text(request.message) + if not query: + raise HTTPException( + status_code=422, + detail=( + "A2A message carries no text. Send either " + '{"message": {"parts": [{"type": "text", "text": "..."}]}} ' + 'or {"message": {"content": "..."}}.' + ), + ) payload = { "query": query, "user_id": "a2a", @@ -1179,7 +1245,11 @@ async def submit_a2a_task(request: A2ATaskRequest) -> dict[str, Any]: "output": structured_output, } except Exception as exc: - return {"task_id": task_id, "status": "failed", "error": str(exc)} + return { + "task_id": task_id, + "status": "failed", + "error": client_safe_message(exc, context="A2A task failed"), + } @router.get("/a2a/tasks/{task_id}") async def get_a2a_task(task_id: str) -> dict[str, Any]: @@ -1325,13 +1395,18 @@ async def clear_thread_messages(thread_id: str) -> dict[str, Any]: @router.get("/threads/{thread_id}/summary") async def get_thread_summary(thread_id: str) -> dict[str, Any]: """Get or generate a conversation summary for a thread.""" - if not memory_mgr: + # ``memory_mgr`` wraps the lazy store proxy and stays truthy even when + # no store is configured, so check the store itself — otherwise the + # RuntimeError surfaced as a 500 for what is a configuration issue. + if not memory_mgr or not thread_store: raise HTTPException(400, "Memory manager not configured (requires thread storage)") try: summary = await memory_mgr.get_conversation_summary(thread_id) return {"thread_id": thread_id, "summary": summary} except Exception as exc: - raise HTTPException(500, f"Failed to generate summary: {exc}") from exc + raise HTTPException( + 500, detail=client_safe_detail(exc, context="Failed to generate summary") + ) from exc # ── POST /optimize/invoke ───────────────────────────────────── @router.post("/optimize/invoke", response_model=OptimizeInvokeResponse) @@ -1412,7 +1487,9 @@ async def optimize_invoke(request: OptimizeInvokeRequest) -> OptimizeInvokeRespo raise except Exception as exc: logger.error(f"Optimize invoke for {agent_name} failed: {exc}") - raise HTTPException(500, f"Optimize invoke failed: {exc}") from exc + raise HTTPException( + 500, detail=client_safe_detail(exc, context="Optimize invoke failed") + ) from exc # ── POST /feedback ──────────────────────────────────────────── @router.post("/feedback") @@ -1478,7 +1555,11 @@ async def list_invocation_logs( "Set logs_history=True / AGENTOMATIC_LOGS_HISTORY=1." }, ) - if thread_store is None: + # NOT ``is None``: thread_store is a _LazyStoreProxy, which is never + # None. Its __bool__ reports whether a store is actually configured, + # so an identity check silently fell through and the first attribute + # access raised RuntimeError as a bare 500. + if not thread_store: raise HTTPException(400, detail={"error": "Storage backend is not configured"}) limit = max(1, min(limit, 200)) offset = max(0, offset) @@ -1522,7 +1603,11 @@ async def get_latest_log_analysis() -> dict[str, Any]: "AGENTOMATIC_ALLOW_LOGSLLM_ANALYSIS=1." }, ) - if thread_store is None: + # NOT ``is None``: thread_store is a _LazyStoreProxy, which is never + # None. Its __bool__ reports whether a store is actually configured, + # so an identity check silently fell through and the first attribute + # access raised RuntimeError as a bare 500. + if not thread_store: raise HTTPException(400, detail={"error": "Storage backend is not configured"}) analysis = await thread_store.get_latest_log_analysis(agent_name, resource_type="agent") if not analysis: @@ -1551,7 +1636,11 @@ async def analyze_invocation_logs( 400, detail={"error": "logs_history must be enabled to analyse invocation logs."}, ) - if thread_store is None: + # NOT ``is None``: thread_store is a _LazyStoreProxy, which is never + # None. Its __bool__ reports whether a store is actually configured, + # so an identity check silently fell through and the first attribute + # access raised RuntimeError as a bare 500. + if not thread_store: raise HTTPException(400, detail={"error": "Storage backend is not configured"}) opts = request or AnalyzeLogsRequest(sample_limit=20, persist=True) @@ -1569,7 +1658,9 @@ async def analyze_invocation_logs( ) except Exception as exc: # noqa: BLE001 logger.error("Log analysis failed for '{}': {}", agent_name, exc) - raise HTTPException(500, detail={"error": f"Log analysis failed: {exc}"}) from exc + raise HTTPException( + 500, detail=client_safe_detail(exc, context="Log analysis failed") + ) from exc return { "agent": agent_name, "resource": "agent", @@ -1589,7 +1680,11 @@ async def get_invocation_log(log_id: str) -> dict[str, Any]: "Set logs_history=True / AGENTOMATIC_LOGS_HISTORY=1." }, ) - if thread_store is None: + # NOT ``is None``: thread_store is a _LazyStoreProxy, which is never + # None. Its __bool__ reports whether a store is actually configured, + # so an identity check silently fell through and the first attribute + # access raised RuntimeError as a bare 500. + if not thread_store: raise HTTPException(400, detail={"error": "Storage backend is not configured"}) entry = await thread_store.get_invocation_log(log_id) if ( @@ -1608,7 +1703,11 @@ async def list_optimization_runs( offset: int = 0, ) -> dict[str, Any]: """List auditable prompt-fit / retrain runs for this agent.""" - if thread_store is None: + # NOT ``is None``: thread_store is a _LazyStoreProxy, which is never + # None. Its __bool__ reports whether a store is actually configured, + # so an identity check silently fell through and the first attribute + # access raised RuntimeError as a bare 500. + if not thread_store: raise HTTPException(400, detail={"error": "Storage backend is not configured"}) limit = max(1, min(limit, 200)) offset = max(0, offset) @@ -1689,7 +1788,7 @@ async def approve_suspended_state(thread_id: str, request: ApproveSuspendedReque except Exception as exc: raise HTTPException( status_code=500, - detail=f"Error resuming execution: {exc}", + detail=client_safe_detail(exc, context="Error resuming execution"), ) from exc # ── POST /threads/{thread_id}/reject ────────────────────────── @@ -1723,7 +1822,9 @@ async def fork_thread(thread_id: str, request: ForkThreadRequest) -> dict[str, A title=request.title, ) except Exception as exc: - raise HTTPException(500, f"Failed to fork thread: {exc}") from exc + raise HTTPException( + 500, detail=client_safe_detail(exc, context="Failed to fork thread") + ) from exc if not forked: raise HTTPException(404, f"Thread '{thread_id}' not found") return forked @@ -1738,6 +1839,8 @@ async def get_thread_lineage(thread_id: str) -> dict[str, Any]: lineage = await thread_store.get_thread_lineage(thread_id) return lineage except Exception as exc: - raise HTTPException(500, f"Failed to retrieve lineage: {exc}") from exc + raise HTTPException( + 500, detail=client_safe_detail(exc, context="Failed to retrieve lineage") + ) from exc return router diff --git a/src/agentomatic/core/schemas.py b/src/agentomatic/core/schemas.py index 2b86cf8..2231101 100644 --- a/src/agentomatic/core/schemas.py +++ b/src/agentomatic/core/schemas.py @@ -28,6 +28,7 @@ # ``{Title}`` is the agent folder name converted to CamelCase # (e.g. ``weather_bot`` -> ``WeatherBot``). + def schema_model_candidates( agent_name: str, ) -> tuple[list[str], list[str]]: diff --git a/src/agentomatic/endpoints/models.py b/src/agentomatic/endpoints/models.py index ab6f8e5..bbc8fe9 100644 --- a/src/agentomatic/endpoints/models.py +++ b/src/agentomatic/endpoints/models.py @@ -55,9 +55,7 @@ class UpstreamAuthConfig(BaseModel): # oauth2 client credentials token_url: str = Field(default="", description="OAuth2 token endpoint (supports ${ENV}).") client_id: str = Field(default="", description="OAuth2 client id (supports ${ENV}).") - client_secret: str = Field( - default="", description="OAuth2 client secret (supports ${ENV})." - ) + client_secret: str = Field(default="", description="OAuth2 client secret (supports ${ENV}).") scope: str = Field(default="", description="Optional OAuth2 scope(s), space-delimited.") audience: str = Field(default="", description="Optional OAuth2 audience claim.") token_leeway: int = Field( diff --git a/src/agentomatic/endpoints/router.py b/src/agentomatic/endpoints/router.py index 5c3227b..feb6845 100644 --- a/src/agentomatic/endpoints/router.py +++ b/src/agentomatic/endpoints/router.py @@ -9,6 +9,7 @@ from fastapi import APIRouter, HTTPException from loguru import logger +from agentomatic.core.errors import client_safe_detail from agentomatic.endpoints.base import BaseEndpoint if TYPE_CHECKING: @@ -132,7 +133,9 @@ async def call_endpoint(request: Any) -> Any: status="error", recorder=log_recorder, ) - raise HTTPException(status_code=500, detail=str(exc)) from exc + raise HTTPException( + status_code=500, detail=client_safe_detail(exc, context="Endpoint call failed") + ) from exc finally: _observe_endpoint(endpoint.endpoint_name, status, time.perf_counter() - t0) diff --git a/src/agentomatic/ingestion/builtin/markdown.py b/src/agentomatic/ingestion/builtin/markdown.py index 599a6ef..18f3828 100644 --- a/src/agentomatic/ingestion/builtin/markdown.py +++ b/src/agentomatic/ingestion/builtin/markdown.py @@ -30,6 +30,11 @@ from agentomatic.ingestion.base import BaseIngestor from agentomatic.ingestion.context import IngestionContext from agentomatic.ingestion.models import IngestionResult +from agentomatic.ingestion.paths import ( + IngestionPathError, + resolve_within_root, + safe_output_filename, +) class MarkdownIngestRequest(BaseModel): @@ -103,7 +108,16 @@ async def ingest( ``{"path": ..., "engine": ..., "size_bytes": ...}``. """ t0 = time.perf_counter() - source_path = Path(request.source).expanduser() + # Caller-supplied paths arrive over HTTP — confine them, or this + # ingestor becomes an arbitrary file read/write primitive. + try: + source_path = resolve_within_root(request.source, description="source") + except IngestionPathError as exc: + return IngestionResult( + ingestor=self.ingestor_name, + status="failed", + errors=[str(exc)], + ) if not source_path.exists(): return IngestionResult( ingestor=self.ingestor_name, @@ -134,9 +148,20 @@ async def ingest( stage="write", ) - out_dir = Path(request.output_dir).expanduser() + try: + out_dir = resolve_within_root(request.output_dir, description="output_dir") + # ``output_filename`` is a name, not a path — a traversal in it + # would escape the (otherwise confined) output directory. + out_name = safe_output_filename( + request.output_filename, default=f"{source_path.stem}.md" + ) + except IngestionPathError as exc: + return IngestionResult( + ingestor=self.ingestor_name, + status="failed", + errors=[str(exc)], + ) out_dir.mkdir(parents=True, exist_ok=True) - out_name = request.output_filename or f"{source_path.stem}.md" out_path = out_dir / out_name out_path.write_text(markdown, encoding="utf-8") diff --git a/src/agentomatic/ingestion/paths.py b/src/agentomatic/ingestion/paths.py new file mode 100644 index 0000000..ba4a208 --- /dev/null +++ b/src/agentomatic/ingestion/paths.py @@ -0,0 +1,96 @@ +"""Filesystem confinement for ingestion jobs. + +Ingestion requests carry caller-supplied paths (what to read, where to write). +Those arrive over HTTP, so treating them as trusted turns an ingestor into an +arbitrary file read/write primitive: ``source=/etc/passwd`` exfiltrates any +file the process can read, and ``output_dir``/``output_filename`` can escape +to any writable location (including overwriting code on an import path). + +Every ingestor that touches caller-supplied paths should resolve them through +:func:`resolve_within_root`, which confines them to an ingestion root. + +The root defaults to the current working directory (the project root for a +normal ``agentomatic run``), so ordinary relative paths keep working. Operators +who genuinely ingest from elsewhere set ``AGENTOMATIC_INGESTION_ROOT``. Setting +it to ``/`` restores the old unconfined behaviour and is deliberately explicit. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +#: Environment variable naming the directory ingestion paths are confined to. +INGESTION_ROOT_ENV = "AGENTOMATIC_INGESTION_ROOT" + + +class IngestionPathError(ValueError): + """Raised when a caller-supplied path escapes the ingestion root.""" + + +def ingestion_root() -> Path: + """Return the directory ingestion paths are confined to. + + Falls back to the current working directory when + ``AGENTOMATIC_INGESTION_ROOT`` is unset. + """ + configured = os.getenv(INGESTION_ROOT_ENV, "").strip() + base = Path(configured).expanduser() if configured else Path.cwd() + return base.resolve() + + +def resolve_within_root( + candidate: str | Path, + *, + root: Path | None = None, + description: str = "path", +) -> Path: + """Resolve *candidate* and require the result to sit inside *root*. + + Args: + candidate: Caller-supplied path (absolute or relative to the root). + root: Confinement root; defaults to :func:`ingestion_root`. + description: Field name used in the error message. + + Returns: + The resolved, confined :class:`~pathlib.Path`. + + Raises: + IngestionPathError: If the path escapes *root*. Symlinks are resolved + before the check, so a symlink pointing outside is rejected too. + """ + base = (root or ingestion_root()).resolve() + raw = Path(candidate).expanduser() + resolved = (raw if raw.is_absolute() else base / raw).resolve() + + if base == Path(resolved.anchor): + # Root is the filesystem root — explicitly unconfined. + return resolved + + if resolved != base and base not in resolved.parents: + raise IngestionPathError( + f"{description} {str(candidate)!r} resolves outside the ingestion root " + f"({base}). Set {INGESTION_ROOT_ENV} to widen it." + ) + return resolved + + +def safe_output_filename(name: str | None, *, default: str) -> str: + """Return a bare filename, rejecting any path separators or traversal. + + ``output_filename`` is a *name*, not a path — allowing ``../../evil`` in it + lets a caller escape an otherwise-confined output directory. + + Raises: + IngestionPathError: If *name* contains a separator or is a traversal. + """ + if not name: + return default + candidate = name.strip() + if not candidate: + return default + if candidate in {".", ".."} or os.sep in candidate or "/" in candidate: + raise IngestionPathError(f"output_filename {name!r} must be a bare filename, not a path") + if os.altsep and os.altsep in candidate: + raise IngestionPathError(f"output_filename {name!r} must be a bare filename, not a path") + return candidate diff --git a/src/agentomatic/ingestion/router.py b/src/agentomatic/ingestion/router.py index 3d51206..5b045cf 100644 --- a/src/agentomatic/ingestion/router.py +++ b/src/agentomatic/ingestion/router.py @@ -14,6 +14,8 @@ from fastapi import APIRouter, HTTPException from loguru import logger +from agentomatic.core.errors import client_safe_detail + from .context import NullIngestionContext from .registry import IngestionRegistry @@ -87,7 +89,10 @@ async def run_endpoint(request: Any, _ingestor: Any = ingestor) -> Any: status="error", recorder=log_recorder, ) - raise HTTPException(status_code=500, detail=str(exc)) from exc + raise HTTPException( + status_code=500, + detail=client_safe_detail(exc, context="Ingestion failed"), + ) from exc duration = (time.perf_counter() - t0) * 1000 logger.debug(f"Ingestor '{_ingestor.ingestor_name}' ran in {duration:.1f}ms") if log_recorder is not None: diff --git a/src/agentomatic/langchain_adapter.py b/src/agentomatic/langchain_adapter.py index 3842eca..b030701 100644 --- a/src/agentomatic/langchain_adapter.py +++ b/src/agentomatic/langchain_adapter.py @@ -264,7 +264,7 @@ def messages_to_dict( # → {"current_query": "Hi", "response": "Hello!", "messages": [...]} """ state: dict[str, Any] = dict(fallback_state or {}) - lc_list: list[dict[str, str]] = [] + lc_list: list[dict[str, Any]] = [] last_query = "" last_response = "" @@ -272,7 +272,17 @@ def messages_to_dict( content = getattr(msg, "content", "") or "" content_str = str(content) role = _msg_role(msg) - lc_list.append({"role": role, "content": content_str}) + entry: dict[str, Any] = {"role": role, "content": content_str} + tool_call_id = getattr(msg, "tool_call_id", None) + if tool_call_id: + entry["tool_call_id"] = tool_call_id + name = getattr(msg, "name", None) + if name: + entry["name"] = name + tool_calls = getattr(msg, "tool_calls", None) + if tool_calls: + entry["tool_calls"] = tool_calls + lc_list.append(entry) if role in ("user", "human"): last_query = content_str elif role in ("ai", "assistant"): @@ -292,6 +302,79 @@ def serialize_messages(messages: Sequence[Any]) -> list[dict[str, str]]: return list(messages_to_dict(messages).get("messages", [])) +def message_to_dict(msg: Any) -> dict[str, Any]: + """Convert a single LangChain ``BaseMessage`` (or role dict) to a plain JSON-safe dict. + + Preserves ``tool_call_id``, ``name``, and ``tool_calls`` when present, so tool-using + agents keep protocol fidelity when their state is serialized for REST/SSE responses. + + Raises: + TypeError: If *msg* has no ``content`` attribute and isn't already a dict + (i.e. it isn't message-shaped at all). + """ + if isinstance(msg, dict): + return msg + if not hasattr(msg, "content"): + raise TypeError(f"not a message-like object: {msg!r}") + + entry: dict[str, Any] = {"role": _msg_role(msg), "content": str(msg.content or "")} + tool_call_id = getattr(msg, "tool_call_id", None) + if tool_call_id: + entry["tool_call_id"] = tool_call_id + name = getattr(msg, "name", None) + if name: + entry["name"] = name + tool_calls = getattr(msg, "tool_calls", None) + if tool_calls: + entry["tool_calls"] = tool_calls + return entry + + +def to_jsonable(value: Any) -> Any: + """Recursively convert *value* into plain JSON-safe Python types. + + Walks dicts/lists/tuples and converts any LangChain ``BaseMessage`` found + along the way into a plain ``{"role": ..., "content": ...}`` dict (via + :func:`message_to_dict`), instead of leaving it to be stringified into a + ``repr()`` — or to blow up entirely — by a downstream JSON encoder. + + This is what a class agent's ``state_to_output()`` result passes through + on the generic REST/Studio response path, so returning raw ``BaseMessage`` + objects (e.g. ``{"messages": state.messages}``) "just works". + """ + _base_message_cls: Any + try: + from langchain_core.messages import BaseMessage as _base_message_cls + except ImportError: + _base_message_cls = None + + if _base_message_cls is not None and isinstance(value, _base_message_cls): + return message_to_dict(value) + if isinstance(value, dict): + return {k: to_jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [to_jsonable(v) for v in value] + return value + + +def json_default(obj: Any) -> Any: + """A ``json.dumps(..., default=...)`` callable that understands LangChain messages. + + Converts ``BaseMessage`` instances to plain ``{"role": ..., "content": ...}`` dicts + (instead of stringifying their ``repr()``) and falls back to ``str(obj)`` for + anything else, matching the platform's previous ``default=str`` behaviour. + """ + _base_message_cls: Any + try: + from langchain_core.messages import BaseMessage as _base_message_cls + except ImportError: + _base_message_cls = None + + if _base_message_cls is not None and isinstance(obj, _base_message_cls): + return message_to_dict(obj) + return str(obj) + + # ===================================================================== # RunnableConfig # ===================================================================== @@ -739,8 +822,12 @@ def _render_text(msg: Any) -> str: def _dict_to_lc(d: dict[str, Any]) -> Any: """Convert a dict to a LangChain message object.""" + if not isinstance(d, dict): + # Already a LangChain message (or other object) — pass through. + return d + try: - from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage except ImportError: return d @@ -750,9 +837,19 @@ def _dict_to_lc(d: dict[str, Any]) -> Any: if role in ("user", "human"): return HumanMessage(content=content) if role in ("ai", "assistant"): - return AIMessage(content=content) + tool_calls = d.get("tool_calls") + kwargs: dict[str, Any] = {} + if tool_calls: + kwargs["tool_calls"] = tool_calls + return AIMessage(content=content, **kwargs) if role == "tool": - return ToolMessage(content=content, tool_call_id=str(d.get("tool_call_id", ""))) + return ToolMessage( + content=content, + tool_call_id=str(d.get("tool_call_id", "")), + name=d.get("name"), + ) + if role == "system": + return SystemMessage(content=content) return HumanMessage(content=content) diff --git a/src/agentomatic/logs/router.py b/src/agentomatic/logs/router.py index 7aa8c41..e00c503 100644 --- a/src/agentomatic/logs/router.py +++ b/src/agentomatic/logs/router.py @@ -12,6 +12,7 @@ from loguru import logger from pydantic import BaseModel, Field +from agentomatic.core.errors import client_safe_detail from agentomatic.logs.runtime import RESOURCE_TYPES, normalize_resource_type if TYPE_CHECKING: @@ -48,7 +49,9 @@ def _require_history() -> None: "Set logs_history=True / AGENTOMATIC_LOGS_HISTORY=1." }, ) - if store is None: + # See router_factory: the store may be a lazy proxy that is never + # None — rely on its __bool__ instead of an identity check. + if not store: raise HTTPException(400, detail={"error": "Storage backend is not configured"}) def _require_analysis() -> None: @@ -163,7 +166,9 @@ async def analyze_logs(request: AnalyzeLogsRequest) -> dict[str, Any]: ) except Exception as exc: # noqa: BLE001 logger.error("Log analysis failed for {}:{}: {}", rtype, request.name, exc) - raise HTTPException(500, detail={"error": f"Log analysis failed: {exc}"}) from exc + raise HTTPException( + 500, detail=client_safe_detail(exc, context="Log analysis failed") + ) from exc return {"resource": rtype, "name": request.name, "analysis": result.to_dict()} @router.get("/{log_id}") diff --git a/src/agentomatic/middleware/auth.py b/src/agentomatic/middleware/auth.py index cfe12aa..8b099b8 100644 --- a/src/agentomatic/middleware/auth.py +++ b/src/agentomatic/middleware/auth.py @@ -6,6 +6,7 @@ from __future__ import annotations +import hmac from collections.abc import Awaitable, Callable from typing import Any @@ -23,7 +24,12 @@ "/openapi.json", "/redoc", "/", - "/studio", + # Only the static UI shell is public (like the Swagger UI shell at + # /docs) — NOT "/studio", which (via prefix matching in + # path_is_skipped) would also exempt the entire Studio debug REST API + # (/studio/agents/..., /studio/.../threads/{id}/state, etc.), letting + # an unauthenticated caller read/mutate any agent's run state. + "/studio/ui", "/status", } @@ -64,10 +70,19 @@ async def dispatch( return response key = request.headers.get(self._header) or request.query_params.get(self._query) - if not key or key != self._api_key: + # Compare as bytes: hmac.compare_digest raises TypeError on a str + # containing non-ASCII, which would turn a bad key into a 500 + # instead of a 401 (and is trivially reachable via ?api_key=…). + if not key or not hmac.compare_digest(key.encode(), self._api_key.encode()): return JSONResponse( {"detail": "Invalid or missing API key"}, status_code=401, ) + # Record the authenticated principal so downstream authorization (the + # zero-trust enforcer) can tell an API-key caller from an anonymous + # one. Without this it looked for JWT claims, found none, and denied + # a request that had just presented a valid key. + request.state.api_key_authenticated = True + request.state.auth_method = "api_key" response = await call_next(request) return response diff --git a/src/agentomatic/middleware/metrics.py b/src/agentomatic/middleware/metrics.py index c8a0bba..f492159 100644 --- a/src/agentomatic/middleware/metrics.py +++ b/src/agentomatic/middleware/metrics.py @@ -33,9 +33,7 @@ class MetricsMiddleware(BaseHTTPMiddleware): """Prometheus metrics collection per request.""" - def __init__( - self, app: Any, *, prefix: str = "agentomatic" - ) -> None: + def __init__(self, app: Any, *, prefix: str = "agentomatic") -> None: super().__init__(app) self._requests: Any | None = None self._duration: Any | None = None diff --git a/src/agentomatic/middleware/rate_limit.py b/src/agentomatic/middleware/rate_limit.py index 99c047d..69a8805 100644 --- a/src/agentomatic/middleware/rate_limit.py +++ b/src/agentomatic/middleware/rate_limit.py @@ -33,16 +33,31 @@ def __init__( *, max_requests: int = 100, window_seconds: int = 60, + trust_proxy_headers: bool = False, ) -> None: super().__init__(app) self._max = max_requests self._window = window_seconds self._hits: dict[str, list[float]] = defaultdict(list) + # X-Forwarded-For is client-controlled unless a trusted reverse proxy + # sets/overwrites it — trusting it by default lets any caller rotate + # the header per request and bypass the limiter entirely. Only honour + # it when the deployer explicitly confirms a trusted proxy is in front. + # + # This flag governs *this* middleware only. Uvicorn's own + # ``--proxy-headers`` (on by default) rewrites ``request.client`` from + # X-Forwarded-For for peers listed in ``--forwarded-allow-ips`` + # (default ``127.0.0.1``), and that rewrite happens before any of this + # runs — the original peer address is not recoverable. So a caller who + # can reach the server *from an allowed peer address* can still steer + # the key. Keep ``--forwarded-allow-ips`` limited to your real proxy. + self._trust_proxy_headers = trust_proxy_headers def _client_key(self, request: Request) -> str: - forwarded = request.headers.get("X-Forwarded-For") - if forwarded: - return forwarded.split(",")[0].strip() + if self._trust_proxy_headers: + forwarded = request.headers.get("X-Forwarded-For") + if forwarded: + return forwarded.split(",")[0].strip() return request.client.host if request.client else "unknown" async def dispatch( diff --git a/src/agentomatic/observability/telemetry.py b/src/agentomatic/observability/telemetry.py index 70840b2..59665ff 100644 --- a/src/agentomatic/observability/telemetry.py +++ b/src/agentomatic/observability/telemetry.py @@ -35,6 +35,12 @@ async def retrieve_docs(query: str) -> list[str]: from loguru import logger + +def _env_flag(name: str) -> bool: + """Return True when *name* is set to a truthy value in the environment.""" + return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"} + + F = TypeVar("F", bound=Callable[..., Any]) # --------------------------------------------------------------------------- @@ -173,8 +179,23 @@ def setup_telemetry( logger.warning("OTLP exporter packages not installed — falling back to console") enable_console = True - if enable_console or not otlp_endpoint: + # Console export is strictly opt-in. It previously defaulted on whenever no + # OTLP endpoint was configured — i.e. for most deployments — which dumped a + # full JSON span document to stdout for *every* HTTP request (thousands of + # log lines per minute, and real money in a hosted log pipeline). Spans are + # still recorded either way; this only controls printing them. + if not enable_console and _env_flag("AGENTOMATIC_OTEL_CONSOLE"): + enable_console = True + + if enable_console: provider.add_span_processor(BatchSpanProcessor(SafeConsoleSpanExporter())) + logger.info("🔭 OTEL console span export enabled") + elif not otlp_endpoint: + logger.debug( + "OTEL: no OTLP endpoint configured and console export is off — " + "spans are recorded but not exported. Set OTEL_EXPORTER_OTLP_ENDPOINT " + "to ship them, or AGENTOMATIC_OTEL_CONSOLE=1 to print them locally." + ) trace.set_tracer_provider(provider) _tracer = trace.get_tracer("agentomatic") diff --git a/src/agentomatic/optimize/dashboard.py b/src/agentomatic/optimize/dashboard.py index 0800a95..aa79100 100644 --- a/src/agentomatic/optimize/dashboard.py +++ b/src/agentomatic/optimize/dashboard.py @@ -210,9 +210,7 @@ def compose(self) -> Any: def on_mount(self) -> None: """Initialise the candidates DataTable columns.""" - table: Any = self.query_one( - "#candidates-pane", DataTable - ) + table: Any = self.query_one("#candidates-pane", DataTable) table.add_columns("Round", "Name", "Source", "Score", "Status") # ── public update entry-point ────────────────────── @@ -297,9 +295,7 @@ def _ingest_candidate( return try: - table: Any = self.query_one( - "#candidates-pane", DataTable - ) + table: Any = self.query_one("#candidates-pane", DataTable) score_str = f"{data.score:.4f}" if data.score is not None else "—" table.add_row( str(data.round_idx or self._round_idx), diff --git a/src/agentomatic/optimize/fitter.py b/src/agentomatic/optimize/fitter.py index c87d34f..911cdf2 100644 --- a/src/agentomatic/optimize/fitter.py +++ b/src/agentomatic/optimize/fitter.py @@ -118,6 +118,7 @@ def _candidate_source_rank(source: str) -> int: """Return the tie-break rank for a candidate ``source`` label.""" return _CANDIDATE_SOURCE_PRIORITY.get(source, 3) + _EARLY_STOP_PATIENCE: int = 3 """Stop if no improvement for this many consecutive rounds.""" @@ -1266,7 +1267,9 @@ def _record_round( logger.debug("APO beam update skipped: {}", exc) if hasattr(opt, "observe"): try: - getattr(opt, "observe")(dict(cand.config.model_params or {}), full_score) + getattr(opt, "observe")( + dict(cand.config.model_params or {}), full_score + ) except Exception as exc: # pragma: no cover logger.debug("Param observe skipped: {}", exc) else: diff --git a/src/agentomatic/optimize/fitter_optimizers.py b/src/agentomatic/optimize/fitter_optimizers.py index 3eabc31..2957293 100644 --- a/src/agentomatic/optimize/fitter_optimizers.py +++ b/src/agentomatic/optimize/fitter_optimizers.py @@ -598,7 +598,11 @@ async def propose( usable = [ r for r in scored if r.get("query") and r.get("response") and r.get("score", 0.0) > 0 ] - if len(usable) < self.k_examples: + # ``k_examples`` can be 0 (e.g. caller passes min(4, len(eval_results)) + # with an empty eval_results) — guard on it explicitly, since + # ``len(usable) < 0`` never fires and would otherwise fall through to + # a ZeroDivisionError below when averaging an empty subset. + if self.k_examples <= 0 or len(usable) < self.k_examples: logger.warning( "FewShotBootstrapOptimizer: only {} usable results, need {} — skipping", len(usable), diff --git a/src/agentomatic/optimize/progress.py b/src/agentomatic/optimize/progress.py index bd2d41e..36a6f41 100644 --- a/src/agentomatic/optimize/progress.py +++ b/src/agentomatic/optimize/progress.py @@ -230,9 +230,7 @@ def _on_baseline_evaluated(self, data: EventData) -> None: self._best_score = score self._scores.append(score) if self._console is not None: - self._console.print( - f" 📊 Baseline score: [bold]{score:.4f}[/]" - ) + self._console.print(f" 📊 Baseline score: [bold]{score:.4f}[/]") def _on_round_start(self, data: EventData) -> None: """Start a per-round sub-progress bar.""" @@ -306,16 +304,12 @@ def _on_round_end(self, data: EventData) -> None: spark = _make_sparkline(self._scores) first = self._scores[0] last = self._scores[-1] - self._console.print( - f" 📈 {spark} {first:.2f} → {last:.2f}" - ) + self._console.print(f" 📈 {spark} {first:.2f} → {last:.2f}") def _on_early_stop(self, data: EventData) -> None: """Log early-stop notification.""" if self._console is not None: - self._console.print( - "\n ⏹️ [yellow]Early stop triggered[/]" - ) + self._console.print("\n ⏹️ [yellow]Early stop triggered[/]") def _on_fit_complete(self, data: EventData) -> None: """Print final summary table and stop progress bars.""" @@ -337,9 +331,7 @@ def _on_sample_result(self, data: EventData) -> None: return s_score = data.sample_score if data.sample_score is not None else 0.0 query_preview = (data.query or "")[:60] - self._console.print( - f" 🔹 {s_score:.3f} {query_preview}" - ) + self._console.print(f" 🔹 {s_score:.3f} {query_preview}") def _on_rewrite_accepted(self, data: EventData) -> None: """Display rewrite acceptance info.""" diff --git a/src/agentomatic/optimize/reward.py b/src/agentomatic/optimize/reward.py index 212b997..fdb5b1d 100644 --- a/src/agentomatic/optimize/reward.py +++ b/src/agentomatic/optimize/reward.py @@ -29,9 +29,7 @@ def reward_from_eval(self, result: EvalResult) -> RewardSignal: dims = {} raw_dims = result.metadata.get("dimensions") or {} if isinstance(raw_dims, dict): - dims = { - str(k): float(v) for k, v in raw_dims.items() if isinstance(v, (int, float)) - } + dims = {str(k): float(v) for k, v in raw_dims.items() if isinstance(v, (int, float))} return RewardSignal( value=float(result.score), dimensions=dims, diff --git a/src/agentomatic/pipelines/engine.py b/src/agentomatic/pipelines/engine.py index 2d97b4f..823a9fd 100644 --- a/src/agentomatic/pipelines/engine.py +++ b/src/agentomatic/pipelines/engine.py @@ -313,7 +313,7 @@ async def _execute_steps( (each step runs after all of its declared upstreams), tie-broken by list index for determinism. """ - from .ordering import compute_execution_order + from .ordering import compute_execution_order, upstreams_of try: order = compute_execution_order(self.config.steps) @@ -325,7 +325,31 @@ async def _execute_steps( ordered_steps = [self.config.steps[i] for i in order] total_steps = len(ordered_steps) + # Steps that ended up FAILED or SKIPPED, in DAG mode — used to cascade + # a skip to any declared downstream dependent instead of letting it + # run against a missing/stale upstream output. Only applies to steps + # that actually declare ``upstreams``: legacy sequential pipelines + # (no upstreams anywhere) keep their existing "keep going" behavior + # under on_error="continue". + unsuccessful_steps: set[str] = set() for exec_pos, step_config in enumerate(ordered_steps): + blocking_upstreams = [u for u in upstreams_of(step_config) if u in unsuccessful_steps] + if blocking_upstreams: + logger.warning( + f" ⏭️ Skipping '{step_config.name}' " + f"(upstream(s) {blocking_upstreams} did not complete successfully)" + ) + pipeline_result.steps[step_config.name] = StepResult( + name=step_config.name, + status=StepStatus.SKIPPED, + error=f"Skipped: upstream(s) {blocking_upstreams} did not complete successfully", + ) + unsuccessful_steps.add(step_config.name) + await self._report_step_progress( + exec_pos + 1, total_steps, step_config.name, "skipped" + ) + continue + # Evaluate condition if present condition = getattr(step_config, "condition", None) if condition and not self._evaluate_condition(condition, ctx): @@ -360,6 +384,11 @@ async def _execute_steps( # Handle failure if result.status == StepStatus.FAILED: + # Mark unsuccessful up front so any DAG dependent of this step + # is skipped (see the ``blocking_upstreams`` check above), + # rather than running against a missing/failed upstream's + # output under a pipeline-level on_error="continue" policy. + unsuccessful_steps.add(step_config.name) on_error = getattr(step_config, "on_error", ErrorPolicy.FAIL) if on_error == ErrorPolicy.SKIP: logger.warning( diff --git a/src/agentomatic/pipelines/loader.py b/src/agentomatic/pipelines/loader.py index 27a0f88..950ae6b 100644 --- a/src/agentomatic/pipelines/loader.py +++ b/src/agentomatic/pipelines/loader.py @@ -480,6 +480,35 @@ def _parse_step(data: dict[str, Any]) -> StepConfigUnion: # --------------------------------------------------------------------------- +def _iter_pipeline_dir(pipelines_dir: Path) -> Iterator[Path]: + """Yield pipeline YAML files held by a ``pipelines/`` folder. + + Covers both layouts: + + - flat — ``pipelines/estimation.yaml`` + - per-pipeline folder — ``pipelines/estimation/pipeline.yaml``, which is + what ``agentomatic init NAME --template pipeline`` scaffolds (alongside + its ``dataset.jsonl``, ``eval.py`` and ``Makefile``). Only the flat form + used to be scanned, so a freshly scaffolded pipeline was never + discovered — ``agentomatic pipeline list`` reported none and the + Pipelines API mounted empty. + + Args: + pipelines_dir: The ``pipelines/`` directory to scan. + + Yields: + Candidate pipeline YAML paths (unresolved). + """ + for child in sorted(pipelines_dir.iterdir()): + if child.is_file() and child.suffix in _YAML_SUFFIXES: + yield child + elif child.is_dir(): + for suffix in ("yaml", "yml"): + candidate = child / f"pipeline.{suffix}" + if candidate.is_file(): + yield candidate + + def _iter_pipeline_files(directory: Path) -> Iterator[Path]: """Yield pipeline YAML files in discovery order. @@ -487,8 +516,9 @@ def _iter_pipeline_files(directory: Path) -> Iterator[Path]: - ``pipeline.yaml`` / ``pipeline.yml`` in *directory* itself. - Every ``*.yaml`` / ``*.yml`` file when *directory* is named - ``pipelines`` (flat project layout). - - Every ``*.yaml`` / ``*.yml`` file in a ``pipelines/`` subdirectory. + ``pipelines`` (flat project layout), plus ``pipeline.yaml`` inside each + of its subfolders (the scaffolded per-pipeline layout). + - The same two shapes under a ``pipelines/`` subdirectory. - ``pipeline.yaml`` / ``pipeline.yml`` inside each ``agents/*/`` folder. Args: @@ -507,16 +537,12 @@ def _iter_pipeline_files(directory: Path) -> Iterator[Path]: # 1b. Flat layout: callers often pass the pipelines/ folder itself. if directory.name == "pipelines": - for child in sorted(directory.iterdir()): - if child.is_file() and child.suffix in _YAML_SUFFIXES: - yield child + yield from _iter_pipeline_dir(directory) # 2. All YAML files under a `pipelines/` subdirectory (project root) pipelines_dir = directory / "pipelines" if pipelines_dir.is_dir(): - for child in sorted(pipelines_dir.iterdir()): - if child.is_file() and child.suffix in _YAML_SUFFIXES: - yield child + yield from _iter_pipeline_dir(pipelines_dir) # 3. pipeline.yaml inside each `agents/*/` folder agents_dir = directory / "agents" diff --git a/src/agentomatic/plugins/ml.py b/src/agentomatic/plugins/ml.py index beb80f4..0411e80 100644 --- a/src/agentomatic/plugins/ml.py +++ b/src/agentomatic/plugins/ml.py @@ -92,6 +92,18 @@ async def load_model(self) -> None: self._is_loaded = True self._loaded_at = datetime.now(UTC).isoformat() + def mark_loaded(self) -> None: + """Record that the model is loaded and ready to serve. + + Called by the platform after :meth:`load_model` returns, so a subclass + that overrides ``load_model`` without calling ``super()`` still ends up + in a usable state instead of answering 503 forever while the startup + log claims success. Idempotent — an existing ``_loaded_at`` is kept. + """ + self._is_loaded = True + if self._loaded_at is None: + self._loaded_at = datetime.now(UTC).isoformat() + def artifact_dir(self) -> Path | None: """Return the active artifact bundle directory, or ``None``. @@ -116,9 +128,8 @@ async def reload_model(self) -> dict[str, Any]: self._is_loaded = False self._loaded_at = None await self.load_model() - # Stamp even when a subclass sets ``_is_loaded`` without calling super(). - if self._loaded_at is None: - self._loaded_at = datetime.now(UTC).isoformat() + # Stamp even when a subclass overrides load_model without calling super(). + self.mark_loaded() return self.info(include_model_card=True) def info(self, *, include_model_card: bool = False) -> dict[str, Any]: diff --git a/src/agentomatic/plugins/router.py b/src/agentomatic/plugins/router.py index bbea589..1c2f140 100644 --- a/src/agentomatic/plugins/router.py +++ b/src/agentomatic/plugins/router.py @@ -8,6 +8,8 @@ from fastapi import APIRouter, HTTPException from loguru import logger +from agentomatic.core.errors import client_safe_detail + from .ml import BaseMLPlugin if TYPE_CHECKING: @@ -53,7 +55,9 @@ async def reload_plugin() -> dict[str, Any]: return await plugin.reload_model() except Exception as exc: logger.error("Reload failed for plugin '{}': {}", plugin.plugin_name, exc) - raise HTTPException(status_code=500, detail=str(exc)) from exc + raise HTTPException( + status_code=500, detail=client_safe_detail(exc, context="Plugin call failed") + ) from exc # Create the dynamic predict endpoint async def predict_endpoint(request: Any) -> Any: @@ -101,7 +105,9 @@ async def predict_endpoint(request: Any) -> Any: status="error", recorder=log_recorder, ) - raise HTTPException(status_code=500, detail=str(e)) from e + raise HTTPException( + status_code=500, detail=client_safe_detail(e, context="Plugin prediction failed") + ) from e # Dynamically adjust the signature so FastAPI extracts the correct Pydantic schemas import inspect diff --git a/src/agentomatic/protocols/decorators.py b/src/agentomatic/protocols/decorators.py index 56da331..16b27aa 100644 --- a/src/agentomatic/protocols/decorators.py +++ b/src/agentomatic/protocols/decorators.py @@ -13,6 +13,8 @@ from loguru import logger from pydantic import BaseModel, Field +from agentomatic.core.errors import client_safe_detail + class APIResponse(BaseModel): """Standard JSON response envelope.""" @@ -35,7 +37,7 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any: raise except Exception as exc: logger.error(f"Unhandled error in {fn.__name__}: {exc}") - raise HTTPException(500, detail=str(exc)) + raise HTTPException(500, detail=client_safe_detail(exc, context="Request failed")) return wrapper diff --git a/src/agentomatic/security/jwt_auth.py b/src/agentomatic/security/jwt_auth.py index f3d8c52..9fde8c6 100644 --- a/src/agentomatic/security/jwt_auth.py +++ b/src/agentomatic/security/jwt_auth.py @@ -53,7 +53,10 @@ "/openapi.json", "/redoc", "/", - "/studio", + # Only the static UI shell is public — see the identical comment in + # agentomatic.middleware.auth._SKIP_PATHS for why this must not be the + # bare "/studio" prefix. + "/studio/ui", "/status", } diff --git a/src/agentomatic/security/zero_trust.py b/src/agentomatic/security/zero_trust.py index e5de900..ecc6287 100644 --- a/src/agentomatic/security/zero_trust.py +++ b/src/agentomatic/security/zero_trust.py @@ -88,7 +88,13 @@ def verify_request( raw_claims = getattr(request.state, "jwt_claims", None) claims: dict[str, Any] = raw_claims if isinstance(raw_claims, dict) else {} - if auth_required and not claims: + # An API key is a valid way to authenticate — the API-key middleware + # runs first and marks the request. It carries no roles or scopes, + # though, so an agent policy that restricts either cannot be evaluated + # for such a caller and must fail closed rather than silently pass. + api_key_authenticated = bool(getattr(request.state, "api_key_authenticated", False)) + + if auth_required and not claims and not api_key_authenticated: self.audit_log( "request_denied", agent_name, @@ -96,6 +102,22 @@ def verify_request( ) return False, "Authentication is required but no valid JWT claims found" + if ( + api_key_authenticated + and not claims + and (policy.allowed_roles or policy.allowed_scopes) + ): + self.audit_log( + "request_denied", + agent_name, + {"reason": "api_key_cannot_satisfy_role_or_scope_policy"}, + ) + return False, ( + f"Agent '{agent_name}' restricts roles/scopes, which an API key " + "cannot carry — authenticate with a JWT that has the required " + "claims." + ) + # Prefer middleware-normalized lists; fall back to claim extraction so # Keycloak-style realm_access / scope strings are honoured. state_roles = getattr(request.state, "roles", None) diff --git a/src/agentomatic/stacks/redaction.py b/src/agentomatic/stacks/redaction.py new file mode 100644 index 0000000..322a00a --- /dev/null +++ b/src/agentomatic/stacks/redaction.py @@ -0,0 +1,129 @@ +"""Redaction helpers for displaying stack configuration. + +Stack YAML is *meant* to reference secrets indirectly (``api_key: ${OPENAI_API_KEY}``), +but nothing enforces that convention — a literal key or a database URL with +embedded credentials is perfectly valid YAML. Commands that print a stack to a +terminal therefore have to assume the file may contain real secrets, because +that output lands in scrollback, CI logs, and screen shares. + +These helpers redact secret-*looking* values while leaving ``${ENV_VAR}`` +references intact (those are safe, and hiding them would obscure the very thing +an operator is trying to verify). +""" + +from __future__ import annotations + +import re + +REDACTED = "***REDACTED***" + +#: Key names whose values are treated as secrets. +_SECRET_KEY_RE = re.compile( + r"(api[_-]?key|secret|password|passwd|token|credential|private[_-]?key|access[_-]?key)", + re.IGNORECASE, +) + +#: ``key: value`` on a YAML line (captures indent, key, value, trailing comment). +_YAML_PAIR_RE = re.compile(r"^(?P\s*(?:-\s*)?)(?P[\w.\-]+)\s*:\s*(?P.*)$") + +#: ``scheme://user:password@host`` — credentials embedded in a URL. +_URL_CREDENTIALS_RE = re.compile( + r"(?P[a-zA-Z][\w+.-]*://)(?P[^:/@\s]+):(?P[^@/\s]+)@" +) + +#: A pure ``${VAR}`` / ``$VAR`` indirection — not a secret itself. +_ENV_REF_RE = re.compile(r"^\$\{[^}]+\}$|^\$[A-Za-z_][A-Za-z0-9_]*$") + + +def _is_env_reference(value: str) -> bool: + """Whether *value* is only an environment-variable reference.""" + stripped = value.strip().strip("\"'") + return bool(_ENV_REF_RE.match(stripped)) + + +def redact_url_credentials(text: str) -> str: + """Mask the password in any ``scheme://user:password@host`` URL in *text*.""" + return _URL_CREDENTIALS_RE.sub( + lambda m: f"{m.group('scheme')}{m.group('user')}:{REDACTED}@", text + ) + + +def redact_secret_value(key: str, value: str) -> str: + """Return *value* redacted when *key* names a secret. + + ``${ENV_VAR}`` references, empty values, and YAML block/flow openers are + left alone — they carry no secret and hiding them would only obscure the + structure an operator is inspecting. + """ + bare = value.strip() + if not bare or bare in {"~", "null", "{}", "[]", "|", ">"}: + return value + if _is_env_reference(bare): + return value + if _SECRET_KEY_RE.search(key): + # Preserve any trailing comment so the file still reads naturally. + comment = "" + if " #" in value: + bare, comment = value.split(" #", 1) + comment = f" #{comment}" + return f"{REDACTED}{comment}" + return value + + +#: Placeholder written into generated ``.env.example`` files in place of a +#: literal secret. Actionable rather than merely masked, because the operator +#: is expected to fill it in. +ENV_EXAMPLE_PLACEHOLDER = "CHANGEME" + + +def env_example_value(value: str) -> str: + """Return a value safe to write into a committed ``.env.example``. + + ``${ENV_VAR}`` references pass through — showing them is the whole point + of the file. A literal secret is replaced with a placeholder: unlike a + stack file (which may be gitignored), ``.env.example`` is conventionally + committed, so writing a real key there publishes it. + """ + if not value or _is_env_reference(value): + return value + return ENV_EXAMPLE_PLACEHOLDER + + +def redact_yaml_text(text: str) -> tuple[str, int]: + """Redact secret-looking values in YAML *text*. + + Operates line-by-line so comments, ordering, and formatting survive — the + point is to show the operator their real file, minus the secrets. + + Args: + text: Raw YAML source. + + Returns: + ``(redacted_text, number_of_redactions)``. + """ + out: list[str] = [] + redactions = 0 + + for line in text.splitlines(): + stripped = line.lstrip() + if stripped.startswith("#") or not stripped: + out.append(line) + continue + + match = _YAML_PAIR_RE.match(line) + if match: + key, value = match.group("key"), match.group("value") + new_value = redact_secret_value(key, value) + if new_value != value: + redactions += 1 + out.append(f"{match.group('indent')}{key}: {new_value}") + continue + + # Even on non-secret keys, a URL may carry inline credentials. + masked = redact_url_credentials(line) + if masked != line: + redactions += 1 + out.append(masked) + + trailing_newline = "\n" if text.endswith("\n") else "" + return "\n".join(out) + trailing_newline, redactions diff --git a/src/agentomatic/storage/checkpointer.py b/src/agentomatic/storage/checkpointer.py index 85bb6ec..4c9c71e 100644 --- a/src/agentomatic/storage/checkpointer.py +++ b/src/agentomatic/storage/checkpointer.py @@ -3,8 +3,8 @@ from __future__ import annotations import asyncio +import base64 import builtins -import json from collections.abc import AsyncIterator, Iterator from typing import Any @@ -16,20 +16,42 @@ CheckpointMetadata, CheckpointTuple, ) +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from agentomatic.storage.base import BaseStore +# Encodes/decodes checkpoint payloads using LangGraph's own serializer, which +# knows how to round-trip LangChain ``BaseMessage`` subclasses, pydantic +# models, and other rich objects that show up in graph channel values. +# A naive ``json.dumps(obj, default=str)`` would stringify those objects to +# their ``repr()`` and permanently lose their structure on reload. +_SERDE = JsonPlusSerializer() -def _ensure_json_serializable(obj: Any) -> Any: - """Round-trip through JSON to guarantee all values are JSON-serializable. +_ENCODED_TYPE_KEY = "__agentomatic_serde_type__" +_ENCODED_DATA_KEY = "__agentomatic_serde_data__" - Non-serializable objects (custom classes, datetimes, bytes, etc.) are - converted to their string representation via ``default=str``. + +def encode_for_storage(obj: Any) -> dict[str, str]: + """Serialize *obj* into a JSON-safe wrapper using LangGraph's serde. + + Preserves rich objects (LangChain messages, pydantic models, dataclasses, + etc.) so they can be reconstructed exactly via :func:`decode_from_storage`. """ - try: - return json.loads(json.dumps(obj, default=str)) - except (TypeError, ValueError): - return obj + type_name, raw = _SERDE.dumps_typed(obj) + return { + _ENCODED_TYPE_KEY: type_name, + _ENCODED_DATA_KEY: base64.b64encode(raw).decode("ascii"), + } + + +def decode_from_storage(payload: Any) -> Any: + """Inverse of :func:`encode_for_storage`.""" + if isinstance(payload, dict) and _ENCODED_TYPE_KEY in payload and _ENCODED_DATA_KEY in payload: + raw = base64.b64decode(payload[_ENCODED_DATA_KEY]) + return _SERDE.loads_typed((payload[_ENCODED_TYPE_KEY], raw)) + # Back-compat: checkpoints written before this fix used naive JSON + # (via ``default=str``) and are returned as-is — best effort only. + return payload class AgentomaticCheckpointer(BaseCheckpointSaver): @@ -78,8 +100,8 @@ async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: "checkpoint_id": cp_data["checkpoint_id"], } }, - checkpoint=cp_data["checkpoint"], - metadata=cp_data["metadata"], + checkpoint=decode_from_storage(cp_data["checkpoint"]), + metadata=decode_from_storage(cp_data["metadata"]), parent_config=( { "configurable": { @@ -145,8 +167,8 @@ async def aput( checkpoint_ns=checkpoint_ns, checkpoint_id=checkpoint_id_str, parent_checkpoint_id=parent_id_str, - checkpoint=_ensure_json_serializable(dict(checkpoint)), - metadata=_ensure_json_serializable(dict(metadata)), + checkpoint=encode_for_storage(dict(checkpoint)), + metadata=encode_for_storage(dict(metadata)), ) return { @@ -225,8 +247,8 @@ async def _alist_list( "checkpoint_id": cp["checkpoint_id"], } }, - checkpoint=cp["checkpoint"], - metadata=cp["metadata"], + checkpoint=decode_from_storage(cp["checkpoint"]), + metadata=decode_from_storage(cp["metadata"]), parent_config=( { "configurable": { diff --git a/src/agentomatic/storage/memory.py b/src/agentomatic/storage/memory.py index 62b97e3..7ea0e67 100644 --- a/src/agentomatic/storage/memory.py +++ b/src/agentomatic/storage/memory.py @@ -7,6 +7,9 @@ from .base import BaseStore +_DEFAULT_CHECKPOINT_LIST_LIMIT = 1000 +"""Cap applied to list_checkpoints() when no explicit limit is given.""" + class MemoryStore(BaseStore): """Fast in-memory store — perfect for development and unit tests. @@ -80,6 +83,12 @@ async def delete_thread(self, thread_id: str) -> bool: ] for sid in orphan_ids: del self._suspended_states[sid] + # Clean up orphaned checkpoints (keyed by (thread_id, ns, checkpoint_id)). + orphan_cp_keys = [key for key in self._checkpoints if key[0] == thread_id] + for key in orphan_cp_keys: + del self._checkpoints[key] + # Clean up orphaned feedback entries. + self._feedback = [fb for fb in self._feedback if fb.get("thread_id") != thread_id] return True return False @@ -394,8 +403,8 @@ async def list_checkpoints( break if found_idx != -1: cps = cps[found_idx + 1 :] - if limit is not None: - cps = cps[:limit] + # Always cap the result set (see SQLAlchemyStore.list_checkpoints for why). + cps = cps[: limit if limit is not None else _DEFAULT_CHECKPOINT_LIST_LIMIT] return cps # ------------------------------------------------------------------ diff --git a/src/agentomatic/storage/models.py b/src/agentomatic/storage/models.py index 539d10d..81bacc0 100644 --- a/src/agentomatic/storage/models.py +++ b/src/agentomatic/storage/models.py @@ -175,7 +175,14 @@ class CheckpointModel(Base): __tablename__ = "checkpoints" - thread_id: Mapped[str] = mapped_column(String(64), primary_key=True) + # No ForeignKey to ThreadModel: LangGraph checkpoints are keyed by whatever + # ``configurable.thread_id`` the caller passes to the checkpointer, which + # is not required to correspond to an agentomatic ThreadModel row (a graph + # can be checkpointed standalone, without ever calling create_thread()). + # Orphan cleanup on thread deletion is therefore done explicitly in + # SQLAlchemyStore.delete_thread()/MemoryStore.delete_thread() instead of + # relying on a DB-level cascade. + thread_id: Mapped[str] = mapped_column(String(64), primary_key=True, index=True) checkpoint_ns: Mapped[str] = mapped_column(String(128), primary_key=True, default="") checkpoint_id: Mapped[str] = mapped_column(String(64), primary_key=True) parent_checkpoint_id: Mapped[str | None] = mapped_column(String(64), nullable=True) diff --git a/src/agentomatic/storage/sqlalchemy.py b/src/agentomatic/storage/sqlalchemy.py index f5f0de8..777ece3 100644 --- a/src/agentomatic/storage/sqlalchemy.py +++ b/src/agentomatic/storage/sqlalchemy.py @@ -25,7 +25,7 @@ from typing import Any from loguru import logger -from sqlalchemy import delete, func, select +from sqlalchemy import delete, event, func, select from sqlalchemy.ext.asyncio import ( AsyncSession, async_sessionmaker, @@ -45,6 +45,9 @@ ThreadModel, ) +_DEFAULT_CHECKPOINT_LIST_LIMIT = 1000 +"""Cap applied to list_checkpoints() when no explicit limit is given.""" + def _ensure_logs_resource_type_columns(sync_conn: Any) -> None: """Best-effort ADD COLUMN for DBs created before resource_type existed. @@ -102,6 +105,7 @@ def __init__( engine: Any = None, ) -> None: self._url = url + self._initialized = False # Reuse an existing engine (e.g. from a per-agent DatabaseConnection) # so memory shares the agent's own database + pool. if engine is not None: @@ -122,6 +126,20 @@ def __init__( self._owns_engine = True self._engine = create_async_engine(url, **engine_kwargs) + # Only when this store built the engine itself: with engine= the url + # argument is the unused default, so testing it could attach a SQLite + # pragma to someone else's non-SQLite engine. + if self._owns_engine and "sqlite" in url: + # SQLite disables foreign-key enforcement per connection unless + # explicitly turned on — without this, the ondelete="CASCADE" + # declared on CheckpointModel/FeedbackModel/SuspendedStateModel + # is silently a no-op and deleting a thread orphans their rows. + @event.listens_for(self._engine.sync_engine, "connect") + def _enable_sqlite_fk(dbapi_connection: Any, _connection_record: Any) -> None: + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + self._session_factory = async_sessionmaker( self._engine, expire_on_commit=False, @@ -135,10 +153,19 @@ def __init__( # ------------------------------------------------------------------ async def initialize(self) -> None: - """Create all database tables and ensure newer columns exist.""" + """Create all database tables and ensure newer columns exist. + + Idempotent: the platform can reach this from several places during + startup (an explicitly configured store, one derived from + ``DATABASE_URL``, and a post-connection pass), and re-running the DDL + on every one of them is wasted round trips plus duplicated log lines. + """ + if self._initialized: + return async with self._engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) await conn.run_sync(_ensure_logs_resource_type_columns) + self._initialized = True logger.info("🗄️ Database tables created/verified") async def close(self) -> None: @@ -150,6 +177,7 @@ async def close(self) -> None: """ if self._owns_engine: await self._engine.dispose() + self._initialized = False logger.info("🗄️ Database connection pool closed") async def health_check(self) -> dict[str, Any]: @@ -223,11 +251,17 @@ async def list_threads( return [t.to_dict() for t in result.scalars().all()] async def delete_thread(self, thread_id: str) -> bool: - """Delete a thread and all its messages (cascading).""" + """Delete a thread and all its messages, feedback, and checkpoints.""" async with self._session() as session: result = await session.execute(select(ThreadModel).where(ThreadModel.id == thread_id)) thread = result.scalar_one_or_none() if thread: + # CheckpointModel has no FK to ThreadModel (see models.py), so + # it isn't covered by the ondelete="CASCADE" on Feedback/ + # SuspendedState — clean it up explicitly or it's orphaned. + await session.execute( + delete(CheckpointModel).where(CheckpointModel.thread_id == thread_id) + ) await session.delete(thread) await session.commit() return True @@ -647,8 +681,11 @@ async def list_checkpoints( if before_cp: stmt = stmt.where(CheckpointModel.created_at < before_cp.created_at) - if limit is not None: - stmt = stmt.limit(limit) + # Always cap the result set — a long-running thread can accumulate + # thousands of checkpoints, and an unbounded query (e.g. a caller, + # or LangGraph's own checkpointer.alist(), passing limit=None) + # would load the entire history into memory on every call. + stmt = stmt.limit(limit if limit is not None else _DEFAULT_CHECKPOINT_LIST_LIMIT) result = await session.execute(stmt) return [c.to_dict() for c in result.scalars().all()] diff --git a/src/agentomatic/studio/adapters/langgraph.py b/src/agentomatic/studio/adapters/langgraph.py index 94089f2..adf0c33 100644 --- a/src/agentomatic/studio/adapters/langgraph.py +++ b/src/agentomatic/studio/adapters/langgraph.py @@ -14,6 +14,7 @@ from loguru import logger +from agentomatic.storage.checkpointer import decode_from_storage from agentomatic.studio.adapter import StudioAdapter from agentomatic.studio.models import ( StudioCheckpoint, @@ -99,7 +100,14 @@ def capabilities(self) -> list[str]: async def get_graph(self) -> StudioGraphTopology: if self._agent.graph_fn is None: - raise ValueError(f"LangGraph agent '{self.agent_name}' has no graph_fn") + # An agent registered with only a node_fn has no graph to draw. + # Raising here surfaced as a bare 500 from /studio/agents/{name}/graph, + # which the Studio UI calls for *every* agent. Return an empty + # topology, matching what the graph-agent adapter already does. + return StudioGraphTopology( + agent_name=self.agent_name, + metadata={"reason": "agent has no graph_fn (node_fn only)"}, + ) graph = self._agent.graph_fn() drawable = graph.get_graph() @@ -240,7 +248,10 @@ async def get_state(self, thread_id: str) -> StudioStateSnapshot | None: cps = await self._store.list_checkpoints(thread_id, "", limit=1) if cps: latest = cps[0] - state_data = latest.get("checkpoint", {}) + # Checkpoints are stored through LangGraph's serde (so + # BaseMessage objects survive), which means the raw row + # holds an encoded wrapper — decode before displaying it. + state_data = decode_from_storage(latest.get("checkpoint", {})) checkpoint_id = latest.get("checkpoint_id") except Exception as exc: logger.warning(f"Store fallback get_state failed: {exc}") @@ -295,8 +306,8 @@ async def get_history(self, thread_id: str) -> list[StudioCheckpoint]: id=cp.get("checkpoint_id", f"cp_{idx}"), thread_id=thread_id, step=idx, - state=cp.get("checkpoint", {}), - metadata=cp.get("metadata", {}), + state=decode_from_storage(cp.get("checkpoint", {})), + metadata=decode_from_storage(cp.get("metadata", {})), parent_id=cp.get("parent_checkpoint_id"), timestamp=cp.get("timestamp", _now_iso()), ) diff --git a/src/agentomatic/studio/models.py b/src/agentomatic/studio/models.py index 1344c2d..1c47534 100644 --- a/src/agentomatic/studio/models.py +++ b/src/agentomatic/studio/models.py @@ -69,7 +69,9 @@ class StudioGraphEdge(BaseModel): id: str = Field(..., description="Unique edge identifier") source: str = Field(..., description="Source node ID") target: str = Field(..., description="Target node ID") - condition: str | None = Field(default=None, description="Condition label for conditional edges") + condition: str | None = Field( + default=None, description="Condition label for conditional edges" + ) metadata: dict[str, Any] = Field(default_factory=dict) @@ -138,7 +140,9 @@ class StudioRunRequest(BaseModel): query: str = Field(..., description="User query or input text") user_id: str = Field("default-user", description="User identifier") - thread_id: str | None = Field(default=None, description="Thread ID for conversation continuity") + thread_id: str | None = Field( + default=None, description="Thread ID for conversation continuity" + ) context: dict[str, Any] = Field(default_factory=dict, description="Additional context") metadata: dict[str, Any] = Field(default_factory=dict, description="Extra metadata") prompt_version: str = Field("v1", description="Prompt version to use") @@ -218,7 +222,9 @@ class StudioStateSnapshot(BaseModel): agent_name: str = Field(..., description="Agent name") state: dict[str, Any] = Field(default_factory=dict, description="Full state dict") timestamp: str = Field(..., description="ISO-8601 snapshot timestamp") - checkpoint_id: str | None = Field(default=None, description="Checkpoint ID if backed by storage") + checkpoint_id: str | None = Field( + default=None, description="Checkpoint ID if backed by storage" + ) class StudioStateUpdate(BaseModel): diff --git a/src/agentomatic/studio/router.py b/src/agentomatic/studio/router.py index 7402ce5..201e058 100644 --- a/src/agentomatic/studio/router.py +++ b/src/agentomatic/studio/router.py @@ -13,6 +13,7 @@ from loguru import logger from pydantic import BaseModel, Field +from agentomatic.core.errors import client_safe_detail from agentomatic.core.schemas import SchemaValidator, load_schema_models from agentomatic.studio.adapters import resolve_adapter from agentomatic.studio.models import ( @@ -228,6 +229,15 @@ async def get_graph(name: str) -> StudioGraphTopology: agent_name=name, metadata={"error": "get_graph timed out"}, ) + except Exception as exc: + # Graph introspection runs user code (graph_fn touches imports and + # connections). A debug view failing to draw must not 500 — degrade + # to an empty topology carrying the reason. + logger.warning(f"Studio get_graph failed for '{name}': {exc}") + return StudioGraphTopology( + agent_name=name, + metadata=client_safe_detail(exc, context="get_graph failed"), + ) @router.get( "/agents/{name}/schemas", @@ -440,13 +450,34 @@ async def resume_execution( detail=f"Agent '{name}' does not support interrupt/resume (no graph_fn)", ) + # Resume is a LangGraph feature: it needs ``astream_events`` and + # ``Command(resume=...)``. Agentomatic's own lightweight AgentGraph has + # neither, so without this guard the call raised AttributeError *inside* + # the SSE body — surfacing a raw internal error with a 200 status. + try: + resume_graph = agent.graph_fn() + except Exception as exc: + raise HTTPException( + status_code=500, + detail=client_safe_detail(exc, context="Failed to build agent graph"), + ) from exc + if not hasattr(resume_graph, "astream_events"): + raise HTTPException( + status_code=501, + detail=( + f"Agent '{name}' does not support interrupt/resume: its graph is not " + "a LangGraph runnable (no 'astream_events'). Compile the agent with " + "LangGraph to use interrupts." + ), + ) + async def _stream() -> AsyncGenerator[str, None]: try: # Guard re-check for the closure (the route already 400s above). if agent.graph_fn is None: yield 'data: {"event": "run_error", "data": {"error": "no graph_fn"}}\n\n' return - graph = agent.graph_fn() + graph = resume_graph # already built (and validated) above config = {"configurable": {"thread_id": thread_id}} # Use LangGraph's Command to resume from interrupt @@ -467,7 +498,15 @@ async def _stream() -> AsyncGenerator[str, None]: yield 'data: {"event": "done"}\n\n' except Exception as exc: - error_data = json_mod.dumps({"event": "run_error", "data": {"error": str(exc)}}) + # Don't echo raw internal exception text to the client — it can + # carry credentials/paths. Full detail goes to the server log, + # correlated by error_id. + error_data = json_mod.dumps( + { + "event": "run_error", + "data": client_safe_detail(exc, context="Resume failed"), + } + ) yield f"data: {error_data}\n\n" return StreamingResponse( diff --git a/src/agentomatic/studio/run_tracker.py b/src/agentomatic/studio/run_tracker.py index eff06b9..695512d 100644 --- a/src/agentomatic/studio/run_tracker.py +++ b/src/agentomatic/studio/run_tracker.py @@ -10,6 +10,7 @@ from loguru import logger +from agentomatic.core.errors import client_safe_message from agentomatic.studio.models import StudioRunEvent, StudioRunInfo if TYPE_CHECKING: @@ -182,6 +183,14 @@ async def execute_with_adapter( try: last_output: dict[str, Any] = {} async for event in adapter.stream_execution(state, config, breakpoints, checkpoint_id): + # This tracker owns the run lifecycle — it already bracketed the + # stream with its own run_start/run_complete (carrying the real + # run_id, timing and output). Some adapters emit their own pair + # too (AgentGraph.astream_studio_events does), which reached the + # client as a duplicate run and made the Studio UI render every + # agent reply twice. + if event.event in {"run_start", "run_complete"}: + continue # Stamp the run_id onto adapter events event.run_id = run_id self.add_event(run_id, event) @@ -207,14 +216,17 @@ async def execute_with_adapter( except Exception as exc: duration = (time.monotonic() - start_time) * 1000 - self.fail_run(run_id, str(exc)) - logger.error(f"Studio run {run_id} failed: {exc}") + # Studio runs are reachable unauthenticated in the default + # `agentomatic run` posture, and this error is both stored on the + # run and streamed over SSE — so never put raw exception text in it. + safe = client_safe_message(exc, context="Studio run failed") + self.fail_run(run_id, safe) error_event = StudioRunEvent( event="run_error", run_id=run_id, timestamp=_now_iso(), - data={"error": str(exc), "type": type(exc).__name__}, + data={"error": safe, "type": type(exc).__name__}, duration_ms=round(duration, 2), ) self.add_event(run_id, error_event) diff --git a/src/agentomatic/studio/serve.py b/src/agentomatic/studio/serve.py index 7d62ac8..7bbfc99 100644 --- a/src/agentomatic/studio/serve.py +++ b/src/agentomatic/studio/serve.py @@ -226,7 +226,8 @@ async def studio_root_redirect() -> RedirectResponse: """Redirect bare path to trailing-slash version.""" return RedirectResponse(url=f"{prefix}/") - logger.info("🎨 Studio UI mounted at %s/", prefix) + # loguru uses {}-style formatting, not printf — "%s" would print literally. + logger.info("🎨 Studio UI mounted at {}/", prefix) def mount_studio_disabled_page( diff --git a/src/agentomatic/tasks/manager.py b/src/agentomatic/tasks/manager.py index f9f6664..94e9636 100644 --- a/src/agentomatic/tasks/manager.py +++ b/src/agentomatic/tasks/manager.py @@ -25,6 +25,8 @@ from loguru import logger +from agentomatic.core.errors import client_safe_message + from .context import TaskContext from .models import ( TargetType, @@ -84,9 +86,23 @@ async def initialize(self) -> None: await self.store.initialize() async def shutdown(self) -> None: - """Cancel in-flight tasks and close the store.""" + """Cancel in-flight tasks and close the store. + + ``asyncio.Task.cancel()`` only *schedules* ``CancelledError`` at the + task's next suspension point — the task's own cleanup (including its + ``_finalize()`` call, which persists the terminal status via + ``self.store.save``) runs later on the event loop. Closing the store + immediately after requesting cancellation (the previous behavior) + raced that save against ``store.close()`` disposing the underlying + connection, so a cancelled task's terminal status could silently + fail to persist. Gather the (now-cancelled) tasks first so their + cleanup has actually run before the store goes away. + """ + running_tasks = list(self._running.values()) for task_id in list(self._running): await self.cancel(task_id) + if running_tasks: + await asyncio.gather(*running_tasks, return_exceptions=True) await self.store.close() # ------------------------------------------------------------------ @@ -311,7 +327,16 @@ async def _run(self, record: TaskRecord, batch_concurrency: int | None) -> None: await asyncio.sleep(delay) logger.exception(f"Task {record.id} failed after {record.attempts} attempts") - await self._finalize(record, TaskStatus.FAILED, error=str(last_error)) + # Raw exception text routinely carries DSNs/paths, and this + # record is served verbatim by GET /tasks/{id}, /result and the + # task list — so sanitise before it is persisted. + await self._finalize( + record, + TaskStatus.FAILED, + error=client_safe_message(last_error, context="Task failed") + if isinstance(last_error, BaseException) + else str(last_error), + ) def _prepare_input(self, record: TaskRecord) -> Any: """Build the payload for a dispatcher invocation, injecting checkpoints. @@ -356,7 +381,10 @@ async def run_item(index: int, payload: Any) -> None: try: results[index] = await dispatcher(record.target, payload, ctx) except Exception as exc: # noqa: BLE001 - collect per-item errors - results[index] = {"error": str(exc)} + # Per-item errors are returned to the caller too. + results[index] = { + "error": client_safe_message(exc, context="Batch item failed") + } async with lock: done += 1 await ctx.report( diff --git a/src/agentomatic/tasks/progress.py b/src/agentomatic/tasks/progress.py index 13e7d47..bae9f7d 100644 --- a/src/agentomatic/tasks/progress.py +++ b/src/agentomatic/tasks/progress.py @@ -25,6 +25,12 @@ _task_ctx: ContextVar[Any] = ContextVar("agentomatic_task_ctx", default=None) _INSTALLED = False +# asyncio only holds a *weak* reference to scheduled tasks — without a strong +# reference kept somewhere, the event loop can garbage-collect a fire-and- +# forget task mid-execution, silently dropping the progress report. Keep one +# here until it completes. +_background_tasks: set[asyncio.Task[Any]] = set() + def bind_task_context(ctx: TaskContext | None) -> Any: """Bind *ctx* for the current async task; return a reset token.""" @@ -78,7 +84,7 @@ def report_stage_sync( loop = asyncio.get_running_loop() except RuntimeError: return - loop.create_task( + task = loop.create_task( report_stage( stage, percent=percent, @@ -87,6 +93,8 @@ def report_stage_sync( message=message, ) ) + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) def _wrap_dispatcher(dispatcher: Any) -> Any: diff --git a/src/agents/__pycache__/__init__.cpython-312.pyc b/src/agents/__pycache__/__init__.cpython-312.pyc index b6803c6..42c0632 100644 Binary files a/src/agents/__pycache__/__init__.cpython-312.pyc and b/src/agents/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/agents/alpha/__pycache__/__init__.cpython-312.pyc b/src/agents/alpha/__pycache__/__init__.cpython-312.pyc index fb9d381..b2dfd2c 100644 Binary files a/src/agents/alpha/__pycache__/__init__.cpython-312.pyc and b/src/agents/alpha/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/agents/alpha/__pycache__/agent.cpython-312.pyc b/src/agents/alpha/__pycache__/agent.cpython-312.pyc index d1cdd59..7857a53 100644 Binary files a/src/agents/alpha/__pycache__/agent.cpython-312.pyc and b/src/agents/alpha/__pycache__/agent.cpython-312.pyc differ diff --git a/src/agents/alpha/__pycache__/config.cpython-312.pyc b/src/agents/alpha/__pycache__/config.cpython-312.pyc index 2600ab1..b415b91 100644 Binary files a/src/agents/alpha/__pycache__/config.cpython-312.pyc and b/src/agents/alpha/__pycache__/config.cpython-312.pyc differ diff --git a/src/agents/alpha/__pycache__/schemas.cpython-312.pyc b/src/agents/alpha/__pycache__/schemas.cpython-312.pyc index 458e8ae..aa0aa5f 100644 Binary files a/src/agents/alpha/__pycache__/schemas.cpython-312.pyc and b/src/agents/alpha/__pycache__/schemas.cpython-312.pyc differ diff --git a/src/agents/beta/__pycache__/__init__.cpython-312.pyc b/src/agents/beta/__pycache__/__init__.cpython-312.pyc index 359b6aa..cb80c33 100644 Binary files a/src/agents/beta/__pycache__/__init__.cpython-312.pyc and b/src/agents/beta/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/agents/beta/__pycache__/agent.cpython-312.pyc b/src/agents/beta/__pycache__/agent.cpython-312.pyc index 671e8de..4fcedeb 100644 Binary files a/src/agents/beta/__pycache__/agent.cpython-312.pyc and b/src/agents/beta/__pycache__/agent.cpython-312.pyc differ diff --git a/src/agents/beta/__pycache__/config.cpython-312.pyc b/src/agents/beta/__pycache__/config.cpython-312.pyc index cc0cc7b..460632e 100644 Binary files a/src/agents/beta/__pycache__/config.cpython-312.pyc and b/src/agents/beta/__pycache__/config.cpython-312.pyc differ diff --git a/src/agents/beta/__pycache__/schemas.cpython-312.pyc b/src/agents/beta/__pycache__/schemas.cpython-312.pyc index b00758d..ef09225 100644 Binary files a/src/agents/beta/__pycache__/schemas.cpython-312.pyc and b/src/agents/beta/__pycache__/schemas.cpython-312.pyc differ diff --git a/src/common/__pycache__/__init__.cpython-312.pyc b/src/common/__pycache__/__init__.cpython-312.pyc index 5e87359..e0b1b33 100644 Binary files a/src/common/__pycache__/__init__.cpython-312.pyc and b/src/common/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/common/__pycache__/base_agent.cpython-312.pyc b/src/common/__pycache__/base_agent.cpython-312.pyc index bbe12b7..3f14d0f 100644 Binary files a/src/common/__pycache__/base_agent.cpython-312.pyc and b/src/common/__pycache__/base_agent.cpython-312.pyc differ diff --git a/src/common/__pycache__/llm_factory.cpython-312.pyc b/src/common/__pycache__/llm_factory.cpython-312.pyc index ba56985..7a13600 100644 Binary files a/src/common/__pycache__/llm_factory.cpython-312.pyc and b/src/common/__pycache__/llm_factory.cpython-312.pyc differ diff --git a/src/common/__pycache__/prompt_manager.cpython-312.pyc b/src/common/__pycache__/prompt_manager.cpython-312.pyc index 688beaa..1435510 100644 Binary files a/src/common/__pycache__/prompt_manager.cpython-312.pyc and b/src/common/__pycache__/prompt_manager.cpython-312.pyc differ diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..fec8126 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,66 @@ +"""Shared test helpers.""" + +from __future__ import annotations + +import contextlib +import importlib +import sys +from collections.abc import Iterator +from pathlib import Path + + +def install_plugin_package( + plugins_dir: Path | str, + package: str, + source: str, +) -> contextlib.AbstractContextManager[None]: + """Write a plugin package and make it importable, hermetically. + + The platform discovers plugins as ``..plugin`` + — see ``AgentPlatform.build``, where the prefix defaults to the plugins + directory's own name — resolved through ``sys.path``. A test that merely + writes the files therefore depends on ambient interpreter state, and + passes or fails depending on which tests ran before it. + + This puts the *parent* of ``plugins_dir`` on ``sys.path`` for the duration, + invalidates importlib's cached directory listings so the freshly written + files are visible, and evicts the modules again afterwards. + + Args: + plugins_dir: The platform's plugins directory. + package: Package name to create inside it. + source: Contents of the package's ``plugin.py``. + + Returns: + A context manager that makes the package importable while active. + """ + plugins_dir = Path(plugins_dir) + prefix = plugins_dir.name + root = plugins_dir.parent + + plugins_dir.mkdir(parents=True, exist_ok=True) + (plugins_dir / "__init__.py").write_text("", encoding="utf-8") + target = plugins_dir / package + target.mkdir(parents=True, exist_ok=True) + (target / "__init__.py").write_text("", encoding="utf-8") + (target / "plugin.py").write_text(source, encoding="utf-8") + + def _evict() -> None: + for name in [m for m in sys.modules if m == prefix or m.startswith(f"{prefix}.")]: + del sys.modules[name] + + @contextlib.contextmanager + def _importable() -> Iterator[None]: + sys.path.insert(0, str(root)) + _evict() + # The files were created after interpreter start, so importlib's cached + # directory listings would otherwise not see them. + importlib.invalidate_caches() + try: + yield + finally: + with contextlib.suppress(ValueError): + sys.path.remove(str(root)) + _evict() + + return _importable() diff --git a/tests/test_agent_endpoints.py b/tests/test_agent_endpoints.py index 152ce94..d626ddb 100644 --- a/tests/test_agent_endpoints.py +++ b/tests/test_agent_endpoints.py @@ -41,6 +41,7 @@ from __future__ import annotations import json +import time from dataclasses import dataclass, field from typing import Any @@ -974,3 +975,57 @@ async def fn(state): assert "hooked_agent" in before_calls assert "hooked_agent" in after_calls + + +class TestA2AMessageShapes: + """The A2A protocol carries text in ``message.parts``, not ``content``. + + Reading only ``content`` meant a spec-shaped request ran the agent on an + empty query and returned 200 with a meaningless result — the text was + dropped silently. + """ + + def test_protocol_parts_reach_the_agent(self, client): + r = client.post( + f"{BASE}/echo/a2a/tasks", + json={ + "message": { + "role": "user", + "parts": [{"type": "text", "text": "parts-shaped query"}], + } + }, + ) + assert r.status_code == 200, r.text + task_id = r.json()["task_id"] + + for _ in range(100): + data = client.get(f"{BASE}/echo/a2a/tasks/{task_id}").json() + if data["status"] in {"completed", "failed", "canceled"}: + break + time.sleep(0.02) + + assert data["status"] == "completed", data + assert "parts-shaped query" in str(data["result"]) + + def test_documented_content_shape_still_works(self, client): + r = client.post( + f"{BASE}/echo/a2a/tasks", + json={"message": {"content": "content-shaped query"}}, + ) + assert r.status_code == 200, r.text + + def test_bare_text_key_is_accepted(self, client): + r = client.post( + f"{BASE}/echo/a2a/tasks", + json={"message": {"text": "text-shaped query"}}, + ) + assert r.status_code == 200, r.text + + def test_message_without_any_text_is_rejected(self, client): + """Better a 422 naming the accepted shapes than a silent empty run.""" + r = client.post( + f"{BASE}/echo/a2a/tasks", + json={"message": {"role": "user", "parts": [{"type": "image"}]}}, + ) + assert r.status_code == 422, r.text + assert "parts" in r.json()["detail"] diff --git a/tests/test_agent_graph_stream.py b/tests/test_agent_graph_stream.py index ff78c43..6db275a 100644 --- a/tests/test_agent_graph_stream.py +++ b/tests/test_agent_graph_stream.py @@ -65,3 +65,40 @@ def node_a(state): assert "node_end" in event_types assert "state_update" in event_types assert "run_complete" in event_types + + +@pytest.mark.asyncio +async def test_agent_graph_astream_serializes_langchain_messages(): + """State containing raw BaseMessage objects must stream as JSON-safe dicts. + + A class-agent node commonly does ``state.messages = [HumanMessage(...), ...]``. + The streamed event payloads must be safe to ``json.dumps`` (no BaseMessage + objects left over) so SSE/Studio consumers don't crash or get stringified reprs. + """ + import json + + from langchain_core.messages import AIMessage, HumanMessage + + def node_a(state): + state["messages"] = [HumanMessage(content="hi"), AIMessage(content="hello")] + return state + + graph = AgentGraph( + nodes={"a": GraphNode(name="a", handler=node_a)}, + edges={}, + entrypoint="a", + finish="a", + ) + + events = [] + async for event in graph.astream_studio_events({"messages": []}, run_id="test_lc"): + events.append(event) + + state_update = next(e for e in events if e["event"] == "state_update") + messages = state_update["data"]["messages"] + assert messages == [ + {"role": "human", "content": "hi"}, + {"role": "ai", "content": "hello"}, + ] + # Must be JSON-serializable without a `default=` fallback. + json.dumps(state_update) diff --git a/tests/test_auth_coverage.py b/tests/test_auth_coverage.py new file mode 100644 index 0000000..eb16ef0 --- /dev/null +++ b/tests/test_auth_coverage.py @@ -0,0 +1,264 @@ +# pyright: reportMissingParameterType=none +# pyright: reportCallIssue=none +# pyright: reportArgumentType=none +# pyright: reportAttributeAccessIssue=none +"""Exhaustive auth coverage over every mounted route. + +Two auth bypasses shipped in this codebase (the Studio debug API riding a +``/studio`` skip-prefix, and the control-plane drain being escapable via an +agent's slug alias). Both were found by inspection, which cannot prove the +absence of a third. + +These tests instead enumerate *every* route the platform mounts and probe it +with auth enabled and no credentials. The set that answers anything other than +401 must exactly equal an explicit, reviewed allowlist — so a newly added route +that is accidentally public fails here rather than in production. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from agentomatic import AgentManifest, AgentPlatform + +_API_KEY = "unit-test-api-key" +_CONTROL_TOKEN = "unit-test-control-token" + +#: Routes that are public *by design*, each with the reason it must stay so. +#: Adding to this set is a security decision and should be reviewed as one. +_INTENTIONALLY_PUBLIC: dict[str, str] = { + "/": "root landing page — no data", + "/health": "liveness probe: orchestrators cannot authenticate", + "/readiness": "readiness probe: same", + "/status": "human status page — verified to carry no secrets", + "/docs": "Swagger UI shell (the API it documents is still gated)", + "/docs/oauth2-redirect": "Swagger OAuth redirect target", + "/redoc": "ReDoc shell", + "/openapi.json": "API schema — deliberately always available", + "/studio/ui": "Studio SPA shell", + "/studio/ui/static": "Studio static assets", + "/studio/ui/{filename:path}": "Studio SPA fallback route (serves index.html)", +} + + +def _build_platform(tmp_path, **overrides: Any) -> AgentPlatform: + async def echo(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok", "agent_type": "echo"} + + kwargs: dict[str, Any] = { + "agents_dir": tmp_path / "agents", + "plugins_dir": tmp_path / "plugins", + "endpoints_dir": tmp_path / "endpoints", + "enable_studio": True, + "enable_control_plane": True, + "control_token": _CONTROL_TOKEN, + "enable_auth": True, + "auth_api_key": _API_KEY, + } + kwargs.update(overrides) + platform = AgentPlatform(**kwargs) + platform.register_agent( + manifest=AgentManifest(name="echo_agent", slug="echo", description="echo"), + node_fn=echo, + ) + return platform + + +def _concrete(path: str) -> str | None: + """Substitute path params, or return None when the route can't be probed.""" + filled = path + for placeholder in ( + "{name}", + "{agent_name}", + "{thread_id}", + "{run_id}", + "{task_id}", + "{log_id}", + "{tid}", + "{filename:path}", + ): + replacement = "echo_agent" if "name" in placeholder else "probe" + filled = filled.replace(placeholder, replacement) + return None if "{" in filled else filled + + +def _probe_unauthenticated(app) -> list[tuple[str, str, int]]: + """Return ``(method, path, status)`` for routes reachable with no credentials.""" + reachable: list[tuple[str, str, int]] = [] + with TestClient(app, raise_server_exceptions=False) as client: + for route in app.routes: + template = getattr(route, "path", None) + if not template: + continue + target = _concrete(template) + if target is None: + continue + methods = getattr(route, "methods", None) or {"GET"} + for method in sorted(m for m in methods if m not in {"HEAD", "OPTIONS"}): + body = {} if method in {"POST", "PUT", "PATCH"} else None + response = client.request(method, target, json=body) + if response.status_code != 401: + reachable.append((method, template, response.status_code)) + return reachable + + +def test_api_key_auth_covers_every_route_except_the_reviewed_allowlist(tmp_path) -> None: + platform = _build_platform(tmp_path) + reachable = _probe_unauthenticated(platform.build()) + + unexpected = sorted({path for _, path, _ in reachable} - set(_INTENTIONALLY_PUBLIC)) + assert not unexpected, ( + "These routes answered without credentials while API-key auth was " + f"enabled: {unexpected}. If a route is genuinely meant to be public, " + "add it to _INTENTIONALLY_PUBLIC with a justification — that is a " + "security decision and should be reviewed as one." + ) + + +def test_jwt_auth_covers_every_route_except_the_reviewed_allowlist(tmp_path) -> None: + """The JWT middleware keeps its own skip list, so it needs its own sweep.""" + platform = _build_platform( + tmp_path, + enable_auth=False, + auth_api_key="", + enable_jwt_auth=True, + ) + reachable = _probe_unauthenticated(platform.build()) + + unexpected = sorted({path for _, path, _ in reachable} - set(_INTENTIONALLY_PUBLIC)) + assert not unexpected, ( + f"Routes reachable without a JWT while JWT auth was enabled: {unexpected}" + ) + + +def test_the_sweep_actually_probes_a_meaningful_number_of_routes(tmp_path) -> None: + """Guard the harness: a broken substitution would skip everything and + make the assertions above vacuously true. + """ + platform = _build_platform(tmp_path) + app = platform.build() + + probed = sum( + 1 + for route in app.routes + if getattr(route, "path", None) and _concrete(route.path) is not None + for m in (getattr(route, "methods", None) or {"GET"}) + if m not in {"HEAD", "OPTIONS"} + ) + assert probed >= 100, f"only {probed} routes probed — substitution likely broke" + + +def test_agent_data_routes_are_gated(tmp_path) -> None: + """Spot-check the routes that actually carry data, under both aliases.""" + platform = _build_platform(tmp_path) + with TestClient(platform.build(), raise_server_exceptions=False) as client: + for path in ( + "/api/v1/echo_agent/invoke", + "/api/v1/echo/invoke", # slug alias + "/api/v1/echo_agent/chat", + "/studio/agents", + "/studio/agents/echo_agent/config", + "/api/v1/control/agents", + ): + method = "POST" if path.endswith(("invoke", "chat")) else "GET" + response = client.request(method, path, json={"query": "x"}) + assert response.status_code == 401, f"{path} reachable without credentials" + + +@pytest.mark.parametrize( + "public_path", + # The templated SPA-fallback route is excluded here rather than skipped + # inside the test: its id contains brackets, which pytest plugins that + # render skip reasons through rich try to parse as markup. + sorted(p for p in _INTENTIONALLY_PUBLIC if "{" not in p), +) +def test_public_routes_do_not_leak_configured_secrets(public_path, tmp_path) -> None: + """Whatever is public must not carry the API key or control token.""" + platform = _build_platform(tmp_path) + with TestClient(platform.build(), raise_server_exceptions=False) as client: + body = client.get(public_path).text + + assert _API_KEY not in body + assert _CONTROL_TOKEN not in body + + +# ===================================================================== +# The skip-prefix must not be escapable +# ===================================================================== + + +async def _raw_get(app, path: str) -> tuple[int, bytes]: + """Send a raw, un-normalised path straight into the ASGI app. + + An HTTP client normally collapses ``..`` before sending, which would mask a + traversal bug. Real attackers do not (``curl --path-as-is``, raw sockets), + so the scope is constructed by hand here. + """ + chunks: list[bytes] = [] + status: int | None = None + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "GET", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "headers": [(b"host", b"testserver")], + "client": ("1.2.3.4", 1), + "server": ("testserver", 80), + "scheme": "http", + "root_path": "", + } + + async def receive() -> dict[str, Any]: + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: dict[str, Any]) -> None: + nonlocal status + if message["type"] == "http.response.start": + status = message["status"] + elif message["type"] == "http.response.body": + chunks.append(message.get("body", b"")) + + await app(scope, receive, send) + return status or 0, b"".join(chunks) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path", + [ + "/studio/ui/../agents", + "/studio/ui/..%2fagents", + "/studio/ui/../../../../etc/passwd", + "//studio/agents", + "/studio/uiadmin", + "/studio/ui-admin", + ], +) +async def test_public_studio_prefix_cannot_be_escaped(path, tmp_path) -> None: + """``/studio/ui`` is public; the debug API next to it is not. + + A traversal that starts inside the public prefix must not reach the + protected routes — neither by skipping auth and then normalising into + ``/studio/agents``, nor by reading files off disk through the SPA fallback. + """ + platform = _build_platform(tmp_path) + app = platform.build() + + async with app.router.lifespan_context(app): + status, body = await _raw_get(app, path) + + # Either the request is rejected, or it lands on the SPA shell — never on + # agent data and never on a file from outside the static directory. + assert b"root:x:" not in body, f"{path} read a file off disk" + if status == 200: + assert body.lstrip()[:9].lower() == b" same real client key, + # so this third request (with yet another spoofed header) must + # still be rate-limited. + resp = client.get("/api/test", headers={"X-Forwarded-For": "3.3.3.3"}) + assert resp.status_code == 429 + + def test_trusted_proxy_headers_opt_in_honors_forwarded_for(self): + from starlette.applications import Starlette + from starlette.responses import JSONResponse + from starlette.routing import Route + + from agentomatic.middleware.rate_limit import RateLimitMiddleware + + async def home(request): + return JSONResponse({"ok": True}) + + app = Starlette(routes=[Route("/api/test", home)]) + app.add_middleware( + RateLimitMiddleware, + max_requests=1, + window_seconds=60, + trust_proxy_headers=True, + ) + client = TestClient(app) + assert client.get("/api/test", headers={"X-Forwarded-For": "1.1.1.1"}).status_code == 200 + # Same real connection, but a distinct forwarded IP gets its own bucket. + assert client.get("/api/test", headers={"X-Forwarded-For": "2.2.2.2"}).status_code == 200 + # Second request from the same forwarded IP is over the limit. + assert client.get("/api/test", headers={"X-Forwarded-For": "1.1.1.1"}).status_code == 429 + # ───────────────────────────────────────────────────────────────────── # Metrics Middleware (without prometheus) diff --git a/tests/test_deploy_cli.py b/tests/test_deploy_cli.py index b223afd..0983015 100644 --- a/tests/test_deploy_cli.py +++ b/tests/test_deploy_cli.py @@ -51,8 +51,8 @@ def test_distroless_uses_nonroot_numeric_uid(self) -> None: content = deploy_mod.render_dockerfile_distroless() assert "gcr.io/distroless/python3-debian12:nonroot" in content assert "USER 65532:65532" in content - assert '"/app/.venv/bin/python"' in content - assert 'pip install "agentomatic[all]==' in content + assert '"/usr/bin/python3"' in content + assert 'pip install --target=/app/deps "agentomatic[all]==' in content assert '"main:app"' in content def test_copy_lines_only_include_existing(self, tmp_path: Path) -> None: @@ -100,6 +100,65 @@ def test_compose_references_distroless(self) -> None: ) assert "dockerfile: Dockerfile.distroless" in content + def test_compose_distroless_healthcheck_has_no_curl(self) -> None: + """The distroless image has no shell and no curl — a curl-based + healthcheck would leave the container permanently "unhealthy". + """ + content = deploy_mod.render_docker_compose( + stack_name="local", + dockerfile_name="Dockerfile.distroless", + distroless=True, + ) + assert "curl" not in content + assert "http://localhost:8000/health" in content + # The base image's own interpreter — a venv Python from the build + # stage does not exist in the distroless runtime. + assert "/usr/bin/python3" in content + assert "/app/.venv/bin/python" not in content + + def test_distroless_builder_python_matches_the_runtime_interpreter(self) -> None: + """The build stage must match ``distroless/python3-debian12``'s Python. + + That base image is Debian 12's Python 3.11. Building the dependencies + on ``python:3.12-slim`` produced an image that could not start at all: + the venv's ``bin/python`` symlinks to the builder's interpreter, which + is absent from the runtime stage, so the ENTRYPOINT was a dangling + symlink (``exec: "/app/.venv/bin/python": no such file or directory``). + Even resolved, cp312 wheels would not import under 3.11. + """ + content = deploy_mod.render_dockerfile_distroless() + + assert "FROM python:3.11-slim AS builder" in content + assert "python:3.12-slim" not in content + # Dependencies land in a plain directory on PYTHONPATH, not a venv + # built around an interpreter the runtime image does not have. + assert "--target=/app/deps" in content + assert 'PYTHONPATH="/app/deps"' in content + assert "python -m venv" not in content + assert ".venv" not in content + assert 'ENTRYPOINT ["/usr/bin/python3", "-m", "uvicorn"]' in content + + def test_non_distroless_image_keeps_its_venv(self) -> None: + """The regular image runs its own interpreter, so a venv is correct.""" + content = deploy_mod.render_dockerfile() + + assert "python -m venv /app/.venv" in content + assert 'PATH="/app/.venv/bin:$PATH"' in content + + def test_compose_non_distroless_still_uses_curl(self) -> None: + content = deploy_mod.render_docker_compose(stack_name="local", distroless=False) + assert '"curl"' in content + + def test_generate_deploy_distroless_compose_has_no_curl(self, tmp_path: Path) -> None: + plan = deploy_mod.generate_deploy( + out_dir=tmp_path / "out", + stack_name="local", + stacks_dir=tmp_path / "no-stacks", + distroless=True, + ) + compose = plan.files["docker-compose.yml"].read_text() + assert "curl" not in compose + # ========================================================================= # Deploy profiles — full (default) vs minimal (Studio off, Swagger stays on) diff --git a/tests/test_e2e_langchain_class_agent.py b/tests/test_e2e_langchain_class_agent.py new file mode 100644 index 0000000..fa4ba01 --- /dev/null +++ b/tests/test_e2e_langchain_class_agent.py @@ -0,0 +1,384 @@ +# pyright: reportMissingParameterType=none +# pyright: reportCallIssue=none +# pyright: reportArgumentType=none +# pyright: reportAttributeAccessIssue=none +"""End-to-end tests for class-based agents built on LangChain abstractions. + +These exercise a ``BaseGraphAgent`` subclass that uses the full LangChain +surface a real user would reach for — ``ChatPromptTemplate`` with +``MessagesPlaceholder``, an LCEL chain (``prompt | llm``), real +``HumanMessage``/``AIMessage``/``SystemMessage``/``ToolMessage`` objects in +state, ``@tool``-decorated tools, and an explicit ``RunnableConfig`` — served +through the *actual* platform HTTP stack (REST invoke / chat / SSE stream / +Studio debug API), not just called directly in-process. + +The LLM is a real LangChain ``FakeListChatModel`` runnable, so the chain +composes and executes for real (no mocking of LangChain itself) while staying +deterministic and offline. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any + +import pytest +from fastapi.testclient import TestClient +from langchain_core.language_models.fake_chat_models import FakeListChatModel +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder +from langchain_core.tools import tool + +from agentomatic import AgentManifest, AgentPlatform +from agentomatic.agents import BaseGraphAgent +from agentomatic.langchain_adapter import dict_to_messages, make_config, serialize_messages + +# ===================================================================== +# A realistic LangChain-native class agent +# ===================================================================== + + +@tool +def lookup_order(order_id: str) -> str: + """Look up an order by its id.""" + return f"order {order_id} is shipped" + + +@dataclass +class SupportState: + """State carrying real LangChain message objects.""" + + request: str = "" + thread_id: str | None = None + # NOTE: raw LangChain BaseMessage objects live here — the framework must + # serialize these safely on every outbound path. + messages: list[Any] = field(default_factory=list) + response: str = "" + used_config: dict[str, Any] = field(default_factory=dict) + + +class SupportAgent(BaseGraphAgent[SupportState]): + """Class agent using ChatPromptTemplate + MessagesPlaceholder + LCEL.""" + + agent_name = "support" + agent_description = "LangChain-native support agent" + agent_framework = "graph_agent" + + def __init__(self, *, llm: Any = None) -> None: + super().__init__() + self.llm = llm + self.tools = [lookup_order] + self.prompt_template = ChatPromptTemplate.from_messages( + [ + ("system", "{system_message}"), + MessagesPlaceholder("messages"), + ] + ) + self.chain = self.prompt_template | self.llm if self.llm is not None else None + + def build_graph(self) -> Any: + g = self.new_graph() + g.add_node("respond", self.respond) + g.set_entry_point("respond") + g.set_finish_point("respond") + return g.compile() + + def respond(self, state: SupportState) -> SupportState: + lc_messages = ( + dict_to_messages(state.messages) + if state.messages + else dict_to_messages({"current_query": state.request}) + ) + config = make_config(thread_id=state.thread_id, tags=["support"]) + # Record the RunnableConfig so a test can assert it was threaded through. + state.used_config = dict(config) + + result = self.chain.invoke( + {"system_message": "You are a support agent.", "messages": lc_messages}, + config=config, + ) + text = getattr(result, "content", None) or str(result) + + # Deliberately leave RAW BaseMessage objects in state — the framework + # must make these JSON-safe on the REST/SSE/Studio paths. + state.messages = [*lc_messages, AIMessage(content=text)] + state.response = text + return state + + def input_to_state(self, data: dict[str, Any]) -> SupportState: + return SupportState( + request=data.get("current_query", ""), + messages=data.get("messages", []) or [], + thread_id=data.get("thread_id"), + ) + + def state_to_output(self, state: SupportState) -> dict[str, Any]: + return { + "response": state.response, + "messages": state.messages, # raw BaseMessage objects on purpose + "used_config": state.used_config, + } + + +def _make_agent(replies: list[str] | None = None) -> SupportAgent: + llm = FakeListChatModel(responses=replies or ["Hello from the fake LLM"]) + return SupportAgent(llm=llm) + + +@pytest.fixture +def client(tmp_path): + agent = _make_agent() + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + title="LangChain E2E", + enable_studio=True, + ) + platform.register_agent( + manifest=AgentManifest( + name="support", + slug="support", + description="LangChain-native support agent", + ), + class_instance=agent, + ) + with TestClient(platform.build()) as c: + yield c + + +# ===================================================================== +# In-process: the LangChain abstractions really execute +# ===================================================================== + + +def test_lcel_chain_executes_and_threads_runnable_config() -> None: + """prompt | llm composes and runs, and a RunnableConfig reaches the chain.""" + agent = _make_agent(["Order is on the way"]) + state = agent.input_to_state({"current_query": "where is my order?", "thread_id": "t-1"}) + out = agent.respond(state) + + assert out.response == "Order is on the way" + # The RunnableConfig carried the thread_id through for tracing/checkpointing. + assert out.used_config["configurable"]["thread_id"] == "t-1" + assert "support" in out.used_config["tags"] + # Real message objects, not dicts/strings. + assert isinstance(out.messages[0], HumanMessage) + assert isinstance(out.messages[-1], AIMessage) + + +def test_message_placeholder_receives_full_history() -> None: + """MessagesPlaceholder("messages") must receive prior turns, not just the query.""" + agent = _make_agent(["ack"]) + state = agent.input_to_state( + { + "messages": [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"}, + ] + } + ) + out = agent.respond(state) + + types = [type(m) for m in out.messages[:3]] + assert types == [SystemMessage, HumanMessage, AIMessage] + + +def test_tool_messages_round_trip_with_tool_call_id() -> None: + """ToolMessage/AIMessage tool-call metadata survives the state round-trip.""" + agent = _make_agent(["done"]) + state = agent.input_to_state( + { + "messages": [ + {"role": "user", "content": "check order 42"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"name": "lookup_order", "args": {"order_id": "42"}, "id": "call_1"} + ], + }, + { + "role": "tool", + "content": "order 42 is shipped", + "tool_call_id": "call_1", + "name": "lookup_order", + }, + ] + } + ) + out = agent.respond(state) + + tool_msg = next(m for m in out.messages if isinstance(m, ToolMessage)) + assert tool_msg.tool_call_id == "call_1" + ai_with_calls = next( + m for m in out.messages if isinstance(m, AIMessage) and getattr(m, "tool_calls", None) + ) + assert ai_with_calls.tool_calls[0]["id"] == "call_1" + + +# ===================================================================== +# Over real HTTP: REST invoke / chat / SSE stream +# ===================================================================== + + +def test_http_invoke_serializes_raw_langchain_messages(client) -> None: + """state_to_output() returning raw BaseMessage objects must not break the + JSON response — they become plain role/content dicts. + """ + resp = client.post("/api/v1/support/invoke", json={"query": "hi"}) + assert resp.status_code == 200, resp.text + + body = resp.json() + assert body["response"] == "Hello from the fake LLM" + + messages = body["output"]["messages"] + assert messages == [ + {"role": "human", "content": "hi"}, + {"role": "ai", "content": "Hello from the fake LLM"}, + ] + # No Python reprs leaked into the payload. + assert "HumanMessage(" not in resp.text + assert "additional_kwargs" not in resp.text + + +def test_http_invoke_response_is_valid_json_end_to_end(client) -> None: + """The whole envelope must be re-parseable (no stringified objects).""" + resp = client.post("/api/v1/support/invoke", json={"query": "hi"}) + reparsed = json.loads(resp.content) + assert isinstance(reparsed["output"]["messages"], list) + assert all(isinstance(m, dict) for m in reparsed["output"]["messages"]) + + +def test_http_chat_creates_thread_and_returns_text(client) -> None: + resp = client.post("/api/v1/support/chat", json={"content": "hello"}) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["thread_id"] + assert body["response"] == "Hello from the fake LLM" + + +def test_http_sse_stream_emits_jsonable_messages(client) -> None: + """SSE frames must be valid JSON with serialized messages.""" + resp = client.post("/api/v1/support/invoke/stream", json={"query": "stream me"}) + assert resp.status_code == 200, resp.text + + frames = [ + json.loads(line[len("data: ") :]) + for line in resp.text.splitlines() + if line.startswith("data: ") and line.strip() != "data: [DONE]" + ] + assert frames, f"no SSE frames parsed from: {resp.text[:400]}" + + node_frame = next(f for f in frames if "respond" in f) + streamed = node_frame["respond"]["messages"] + assert {"role": "human", "content": "stream me"} in streamed + assert "[DONE]" in resp.text + + +def test_http_chat_preserves_tool_call_fidelity(client) -> None: + """A caller-supplied tool-calling history on /chat must reach the agent + (and come back) with ``tool_calls`` and ``tool_call_id`` intact. + + Losing either breaks the call/result pairing that OpenAI/Anthropic + require on the following turn. + """ + resp = client.post( + "/api/v1/support/chat", + json={ + "content": "check it", + "messages": [ + {"role": "user", "content": "check order 42"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"name": "lookup_order", "args": {"order_id": "42"}, "id": "call_1"} + ], + }, + { + "role": "tool", + "content": "order 42 is shipped", + "tool_call_id": "call_1", + "name": "lookup_order", + }, + ], + }, + ) + assert resp.status_code == 200, resp.text + + messages = resp.json()["output"]["messages"] + tool_entry = next(m for m in messages if m["role"] == "tool") + assert tool_entry["tool_call_id"] == "call_1" + assert tool_entry["name"] == "lookup_order" + + ai_entry = next(m for m in messages if m.get("tool_calls")) + assert ai_entry["tool_calls"][0]["id"] == "call_1" + + +def test_http_chat_forwards_history_and_thread_id_to_input_to_state(client) -> None: + """``messages`` and ``thread_id`` must actually reach ``input_to_state``. + + They used to be stripped as "conversation bookkeeping", which made the + scaffolded LangChain template's history/thread handling dead code on + every HTTP path. + """ + resp = client.post( + "/api/v1/support/chat", + json={ + "content": "second turn", + "messages": [{"role": "user", "content": "first turn"}], + }, + ) + assert resp.status_code == 200, resp.text + + body = resp.json() + # thread_id reached the agent and was threaded into the RunnableConfig. + assert body["output"]["used_config"]["configurable"]["thread_id"] == body["thread_id"] + # Prior turn was visible to the agent (so MessagesPlaceholder sees it). + contents = [m["content"] for m in body["output"]["messages"]] + assert "first turn" in contents + assert "second turn" in contents + + +# ===================================================================== +# Studio debug API works for a LangChain class agent +# ===================================================================== + + +def test_studio_lists_and_introspects_langchain_class_agent(client) -> None: + listing = client.get("/studio/agents") + assert listing.status_code == 200, listing.text + assert any(a["name"] == "support" for a in listing.json()) + + graph = client.get("/studio/agents/support/graph") + assert graph.status_code == 200, graph.text + node_ids = {n["id"] for n in graph.json()["nodes"]} + assert "respond" in node_ids + + schemas = client.get("/studio/agents/support/schemas") + assert schemas.status_code == 200, schemas.text + + +def test_openapi_is_valid_with_langchain_class_agent(client) -> None: + """A bad response_model would make /openapi.json 500 — guard against it.""" + resp = client.get("/openapi.json") + assert resp.status_code == 200, resp.text + spec = resp.json() + assert "/api/v1/support/invoke" in spec["paths"] + + +# ===================================================================== +# Serialization helper used by the scaffolded template +# ===================================================================== + + +def test_serialize_messages_matches_rest_representation() -> None: + """The helper the scaffold uses produces the same shape the REST layer emits.""" + msgs = [HumanMessage(content="hi"), AIMessage(content="yo")] + assert serialize_messages(msgs) == [ + {"role": "human", "content": "hi"}, + {"role": "ai", "content": "yo"}, + ] diff --git a/tests/test_fitter.py b/tests/test_fitter.py index 84afc77..543552e 100644 --- a/tests/test_fitter.py +++ b/tests/test_fitter.py @@ -905,6 +905,38 @@ async def test_param_search_no_changes_skips(self): # The baseline already has temperature=0.5, so no new candidates assert len(candidates) == 0 + async def test_few_shot_bootstrap_empty_eval_results_no_crash(self): + """k_examples=min(4, len(eval_results))=0 with empty results must not + raise ZeroDivisionError — it should just skip (return no candidates). + """ + opt = FewShotBootstrapOptimizer(n_candidates=5, k_examples=0) + baseline = PromptRuntimeConfig(system_prompt="test") + candidates = await opt.propose( + current_config=baseline, + eval_results=[], + dataset_sample=[], + search_space=PromptSearchSpace(), + ) + assert candidates == [] + + async def test_few_shot_bootstrap_k_zero_with_usable_results_no_crash(self): + """Same guard, but with usable (non-empty) eval_results present — + k_examples=0 must still skip cleanly rather than dividing by zero. + """ + opt = FewShotBootstrapOptimizer(n_candidates=5, k_examples=0) + baseline = PromptRuntimeConfig(system_prompt="test") + eval_results = [ + {"query": "q1", "response": "r1", "score": 0.9}, + {"query": "q2", "response": "r2", "score": 0.8}, + ] + candidates = await opt.propose( + current_config=baseline, + eval_results=eval_results, + dataset_sample=[], + search_space=PromptSearchSpace(), + ) + assert candidates == [] + # ===================================================================== # PromptFitter Tests diff --git a/tests/test_invoke_response_coercion.py b/tests/test_invoke_response_coercion.py index 5240241..2b847d4 100644 --- a/tests/test_invoke_response_coercion.py +++ b/tests/test_invoke_response_coercion.py @@ -40,3 +40,25 @@ def test_coerce_json_fallback_when_no_human_text() -> None: text, output, _ = coerce_agent_invoke_payload(payload) assert output == payload assert json.loads(text)["p50"] == 10 + + +def test_coerce_state_to_output_with_raw_langchain_messages() -> None: + """A class agent returning ``{"messages": state.messages}`` (raw BaseMessage + objects, the pattern shown in TODO.md's LangGraph example) must not crash the + REST response — messages become plain role/content dicts. + """ + from langchain_core.messages import AIMessage, HumanMessage + + result = { + "agent_type": "chatbot", + "messages": [HumanMessage(content="hi"), AIMessage(content="hello there")], + } + text, output, _ = coerce_agent_invoke_payload(result) + + assert output["messages"] == [ + {"role": "human", "content": "hi"}, + {"role": "ai", "content": "hello there"}, + ] + # The whole envelope must be safe to json-encode as a real HTTP response. + envelope = AgentInvokeResponse(response=text, output=output, agent_type="chatbot") + json.dumps(envelope.model_dump()) diff --git a/tests/test_langchain_adapter.py b/tests/test_langchain_adapter.py index 0826f4c..9b69c4a 100644 --- a/tests/test_langchain_adapter.py +++ b/tests/test_langchain_adapter.py @@ -276,6 +276,55 @@ def test_dict_to_messages_accepts_list() -> None: assert plain[0]["content"] == "Hi" +def test_dict_to_lc_system_role_becomes_system_message() -> None: + """A 'system' role dict must round-trip to a real SystemMessage, not HumanMessage.""" + from langchain_core.messages import SystemMessage + + from agentomatic.langchain_adapter import dict_to_messages + + msgs = dict_to_messages([{"role": "system", "content": "be terse"}]) + assert len(msgs) == 1 + assert isinstance(msgs[0], SystemMessage) + assert msgs[0].content == "be terse" + + +def test_messages_to_dict_preserves_tool_call_fidelity() -> None: + """tool_call_id / tool_calls / name must survive the message -> dict round-trip. + + Without this, a ToolMessage response loses its tool_call_id and can no longer + satisfy the OpenAI/Anthropic tool-call protocol on the next turn. + """ + from langchain_core.messages import AIMessage, ToolMessage + + from agentomatic.langchain_adapter import messages_to_dict + + ai_msg = AIMessage( + content="", + tool_calls=[{"name": "search", "args": {"q": "hi"}, "id": "call_1"}], + ) + tool_msg = ToolMessage(content="result", tool_call_id="call_1", name="search") + + result = messages_to_dict([ai_msg, tool_msg]) + entries = result["messages"] + + assert entries[0]["tool_calls"][0]["id"] == "call_1" + assert entries[1]["tool_call_id"] == "call_1" + assert entries[1]["name"] == "search" + + +def test_dict_to_lc_round_trips_tool_call_id() -> None: + """dict -> LangChain message -> dict must preserve tool_call_id (see BUG: dropped id).""" + from agentomatic.langchain_adapter import dict_to_messages, messages_to_dict + + original = [{"role": "tool", "content": "42", "tool_call_id": "call_9", "name": "calc"}] + lc_messages = dict_to_messages(original) + assert lc_messages[0].tool_call_id == "call_9" + + round_tripped = messages_to_dict(lc_messages)["messages"] + assert round_tripped[0]["tool_call_id"] == "call_9" + assert round_tripped[0]["name"] == "calc" + + # ===================================================================== # RunnableConfig # ===================================================================== @@ -612,3 +661,46 @@ async def astream(self, state, config=None): async for event in adapted.astream({"current_query": "test"}): events.append(event) assert len(events) == 3 + + +# ===================================================================== +# JSON-safe serialization of LangChain objects (to_jsonable / json_default) +# ===================================================================== + + +def test_message_to_dict_preserves_tool_fields() -> None: + from langchain_core.messages import AIMessage + + from agentomatic.langchain_adapter import message_to_dict + + msg = AIMessage(content="hi", tool_calls=[{"name": "x", "args": {}, "id": "c1"}]) + result = message_to_dict(msg) + assert result == {"role": "ai", "content": "hi", "tool_calls": msg.tool_calls} + + +def test_to_jsonable_converts_nested_messages() -> None: + """A class agent's raw ``state.messages`` (list[BaseMessage]) must become JSON-safe.""" + import json + + from langchain_core.messages import AIMessage, HumanMessage + + from agentomatic.langchain_adapter import to_jsonable + + payload = {"messages": [HumanMessage(content="hi"), AIMessage(content="hello")], "n": 1} + result = to_jsonable(payload) + assert result == { + "messages": [{"role": "human", "content": "hi"}, {"role": "ai", "content": "hello"}], + "n": 1, + } + json.dumps(result) # must not raise + + +def test_json_default_handles_base_message() -> None: + import json + + from langchain_core.messages import HumanMessage + + from agentomatic.langchain_adapter import json_default + + text = json.dumps({"m": HumanMessage(content="hi")}, default=json_default) + assert json.loads(text)["m"] == {"role": "human", "content": "hi"} diff --git a/tests/test_live_omlx_keras_optimize.py b/tests/test_live_omlx_keras_optimize.py index 12a5aca..950f196 100644 --- a/tests/test_live_omlx_keras_optimize.py +++ b/tests/test_live_omlx_keras_optimize.py @@ -435,9 +435,9 @@ def test_keras_fit_loss_decreases_per_epoch(mode: str) -> None: # Prompt modes discover at least the deterministic marker. assert history["banana"][-1] == 1.0, mode prompt = agent.compiled_config.get("system_prompt") or agent.system_prompt - assert ( - MARKER in prompt.lower() or agent.compiled_config.get("few_shot_examples") - ), f"{mode}: banana missing from applied config: {prompt[:120]}" + assert MARKER in prompt.lower() or agent.compiled_config.get("few_shot_examples"), ( + f"{mode}: banana missing from applied config: {prompt[:120]}" + ) @pytest.mark.asyncio diff --git a/tests/test_llm_thinking.py b/tests/test_llm_thinking.py index a91cf85..e6b5829 100644 --- a/tests/test_llm_thinking.py +++ b/tests/test_llm_thinking.py @@ -79,7 +79,39 @@ def test_strip_thinking_for_json() -> None: assert cleaned.strip().startswith("{") or "ok" in cleaned -def test_openai_compat_mirrors_enable_thinking_into_chat_template() -> None: +@pytest.fixture +def _stub_langchain_openai(monkeypatch): + """Stand in for ``langchain_openai`` so this runs without the vendor extra. + + ``langchain-openai`` lives in the optional ``openai`` extra, which + ``agentomatic[all]`` deliberately does not pull in (the platform ships no + first-party vendor connectors). Importing it unconditionally made this + test fail — rather than skip — for anyone running the suite after the + documented install. The rest of the suite stubs the module the same way. + """ + import sys + import types + + if "langchain_openai" in sys.modules: # real package present — use it + yield + return + + class _ChatOpenAI: + model_fields = {"extra_body": object()} + + def __init__(self, **kwargs): + self.kwargs = kwargs + + module = types.ModuleType("langchain_openai") + module.ChatOpenAI = _ChatOpenAI + module.AzureChatOpenAI = _ChatOpenAI + monkeypatch.setitem(sys.modules, "langchain_openai", module) + yield + + +def test_openai_compat_mirrors_enable_thinking_into_chat_template( + _stub_langchain_openai, +) -> None: """oMLX/Qwen need chat_template_kwargs.enable_thinking, not only the flat key.""" out = _openai_compat_kwargs({"extra": {"enable_thinking": False}}) body = out["extra_body"] diff --git a/tests/test_middleware_paths.py b/tests/test_middleware_paths.py index aa1d756..a40f09c 100644 --- a/tests/test_middleware_paths.py +++ b/tests/test_middleware_paths.py @@ -19,6 +19,20 @@ def test_prefix_match(self) -> None: assert path_is_skipped("/status/platform", skips) assert not path_is_skipped("/api/v1/agent/invoke", skips) - def test_jwt_defaults_include_studio(self) -> None: - assert "/studio" in _DEFAULT_SKIP_PATHS - assert path_is_skipped("/studio/agents", _DEFAULT_SKIP_PATHS) + def test_jwt_defaults_exempt_only_the_studio_ui_shell(self) -> None: + """Regression: the bare "/studio" prefix must NOT be in the default + skip set — via path_is_skipped's prefix matching, that would also + exempt the entire Studio debug REST API (/studio/agents, + /studio/.../threads/{id}/state, etc.) from JWT auth, letting an + unauthenticated caller read/mutate any agent's run state. + + Only the static UI shell ("/studio/ui") is public, matching how + "/docs" only exempts the Swagger UI shell, not the API it documents. + """ + assert "/studio" not in _DEFAULT_SKIP_PATHS + assert "/studio/ui" in _DEFAULT_SKIP_PATHS + assert path_is_skipped("/studio/ui/", _DEFAULT_SKIP_PATHS) + assert not path_is_skipped("/studio/agents", _DEFAULT_SKIP_PATHS) + assert not path_is_skipped( + "/studio/agents/hello/threads/victim/state", _DEFAULT_SKIP_PATHS + ) diff --git a/tests/test_no_unhandled_500s.py b/tests/test_no_unhandled_500s.py new file mode 100644 index 0000000..f7557c6 --- /dev/null +++ b/tests/test_no_unhandled_500s.py @@ -0,0 +1,216 @@ +# pyright: reportMissingParameterType=none +# pyright: reportCallIssue=none +# pyright: reportArgumentType=none +# pyright: reportAttributeAccessIssue=none +"""Every route must answer, not crash — swept exhaustively. + +The unit suite passed 2092 tests while a stock deployment returned a bare +``500 Internal Server Error`` on ``/api/v1/{agent}/optimization-runs``: the +route guarded a lazy store proxy with ``is None`` (never true for a proxy), so +a ``RuntimeError`` escaped the handler. Two more turned up the same way — +``/studio/agents/{name}/graph`` for an agent with no ``graph_fn``, and the +thread-summary route. + +None of those were caught by testing individual features, because each was +only reachable in a configuration no individual test happened to build. This +module instead walks *every* mounted route and asserts none of them 5xx, in the +default no-store posture that a fresh `agentomatic run` actually uses. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from agentomatic import AgentManifest, AgentPlatform +from agentomatic.storage import MemoryStore + +#: Substitutions that turn a route template into a concrete probe URL. The +#: values are deliberately non-existent ids — a missing resource is a 404, not +#: a 500, and asserting that is the point. +_PATH_PARAMS = { + "{name}": "a1", + "{agent_name}": "a1", + "{thread_id}": "no-such-thread", + "{run_id}": "no-such-run", + "{task_id}": "no-such-task", + "{log_id}": "no-such-log", + "{tid}": "no-such-thread", + "{filename:path}": "index.html", +} + +#: A superset body — routes ignore the keys they don't declare, and the ones +#: they do declare get a plausible value so we exercise the handler rather than +#: bouncing off request validation. +_PROBE_BODY: dict[str, Any] = { + "query": "hi", + "content": "hi", + "value": "x", + "enabled": False, + "updates": {}, + "message": {"content": "hi"}, + "input": {"query": "hi"}, + "message_index": 0, + "text": "hi", +} + +_CONTROL_TOKEN = "sweep-control-token" + + +@pytest.fixture(scope="module") +def swept_app(): + """A fully-featured platform in the DEFAULT posture: no store configured.""" + import tempfile + from pathlib import Path + + async def echo(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok", "agent_type": "echo"} + + tmp = Path(tempfile.mkdtemp()) + platform = AgentPlatform( + agents_dir=tmp / "agents", + plugins_dir=tmp / "plugins", + endpoints_dir=tmp / "endpoints", + enable_studio=True, + enable_control_plane=True, + control_token=_CONTROL_TOKEN, + ) + platform.register_agent( + manifest=AgentManifest(name="a1", slug="a1", description="echo"), + node_fn=echo, + ) + return platform.build() + + +def _probe_targets(app) -> list[tuple[str, str, str]]: + """Return ``(method, concrete_path, route_template)`` for every route.""" + targets: list[tuple[str, str, str]] = [] + for route in app.routes: + template = getattr(route, "path", None) + if not template: + continue + concrete = template + for placeholder, value in _PATH_PARAMS.items(): + concrete = concrete.replace(placeholder, value) + if "{" in concrete: # an unknown param we can't fill safely + continue + methods = getattr(route, "methods", None) or {"GET"} + for method in sorted(m for m in methods if m not in {"HEAD", "OPTIONS"}): + targets.append((method, concrete, template)) + return targets + + +def test_no_route_returns_an_unhandled_server_error(swept_app) -> None: + """No route may 5xx in the default configuration. + + A 4xx is fine everywhere — missing resource, unconfigured backend, bad + input. A 5xx means an exception escaped a handler. + """ + failures: list[str] = [] + + with TestClient(swept_app, raise_server_exceptions=False) as client: + for method, path, template in _probe_targets(swept_app): + body = _PROBE_BODY if method in {"POST", "PUT", "PATCH"} else None + try: + response = client.request( + method, path, json=body, headers={"X-Control-Token": _CONTROL_TOKEN} + ) + except Exception as exc: # noqa: BLE001 - a raised error is a failure too + failures.append(f"{method} {template} raised {type(exc).__name__}: {exc}") + continue + if response.status_code >= 500: + failures.append( + f"{method} {template} -> {response.status_code}: {response.text[:160]}" + ) + + assert not failures, "Routes returned a server error:\n" + "\n".join(failures) + + +def test_the_sweep_covers_a_meaningful_number_of_routes(swept_app) -> None: + """Guard the harness — a broken substitution would skip everything and make + the assertion above vacuously true. + """ + targets = _probe_targets(swept_app) + assert len(targets) >= 80, f"only {len(targets)} route/method pairs probed" + + +def test_store_dependent_routes_answer_4xx_rather_than_crashing(swept_app) -> None: + """The specific regression: these need a store, and none is configured.""" + with TestClient(swept_app, raise_server_exceptions=False) as client: + for path in ( + "/api/v1/a1/optimization-runs", + "/api/v1/a1/logs", + "/api/v1/a1/threads/no-such-thread/summary", + ): + response = client.get(path) + assert 400 <= response.status_code < 500, ( + f"{path} -> {response.status_code} (expected a 4xx): {response.text[:160]}" + ) + + +def test_studio_graph_degrades_for_an_agent_without_a_graph(swept_app) -> None: + """The Studio UI calls this for *every* agent; a node_fn-only agent has no + graph to draw, which must not be a 500. + """ + with TestClient(swept_app, raise_server_exceptions=False) as client: + response = client.get("/studio/agents/a1/graph") + + assert response.status_code == 200, response.text + assert response.json()["agent_name"] == "a1" + + +@pytest.fixture(scope="module") +def swept_app_with_store(): + """The *other* posture: a store, invocation history, and the task API on. + + The no-store sweep above cannot reach the code paths that only run once a + store exists (history reads, checkpoint lookups, thread summaries). A + container sweep in that configuration is what surfaced the A2A and + template defects, so the same surface is covered here. + """ + import tempfile + from pathlib import Path + + async def echo(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok", "agent_type": "echo"} + + tmp = Path(tempfile.mkdtemp()) + platform = AgentPlatform( + agents_dir=tmp / "agents", + plugins_dir=tmp / "plugins", + endpoints_dir=tmp / "endpoints", + enable_studio=True, + enable_control_plane=True, + control_token=_CONTROL_TOKEN, + store=MemoryStore(), + logs_history=True, + ) + platform.register_agent( + manifest=AgentManifest(name="a1", slug="a1", description="echo"), + node_fn=echo, + ) + return platform.build() + + +def test_no_route_returns_a_server_error_with_a_store_configured(swept_app_with_store) -> None: + """Same sweep, store-backed posture — the one a real deployment runs.""" + failures: list[str] = [] + + with TestClient(swept_app_with_store, raise_server_exceptions=False) as client: + for method, path, template in _probe_targets(swept_app_with_store): + body = _PROBE_BODY if method in {"POST", "PUT", "PATCH"} else None + try: + response = client.request( + method, path, json=body, headers={"X-Control-Token": _CONTROL_TOKEN} + ) + except Exception as exc: # noqa: BLE001 + failures.append(f"{method} {template} raised {type(exc).__name__}: {exc}") + continue + if response.status_code >= 500: + failures.append( + f"{method} {template} -> {response.status_code}: {response.text[:160]}" + ) + + assert not failures, "Routes returned a server error:\n" + "\n".join(failures) diff --git a/tests/test_pipeline_dag.py b/tests/test_pipeline_dag.py index 78689c7..4c2c3a7 100644 --- a/tests/test_pipeline_dag.py +++ b/tests/test_pipeline_dag.py @@ -141,6 +141,45 @@ async def test_conditions_still_skip_in_dag_order(self) -> None: assert result.steps["c"].status.value == "skipped" assert result.steps["b"].status.value == "success" + @pytest.mark.asyncio + async def test_downstream_step_skipped_when_upstream_fails_and_pipeline_continues( + self, + ) -> None: + """A step whose declared upstream FAILED must not run against a + missing/stale output just because the pipeline-level policy is + "continue" — it should be skipped, and that skip should cascade to + its own dependents too. + """ + config = PipelineConfig( + name="dag", + on_error="continue", + steps=[ + TransformStepConfig(name="a", code="raise ValueError('boom')"), + TransformStepConfig( + name="b", + code="return {'v': ctx.steps['a'].output['v'] + 1}", + upstreams=["a"], + ), + TransformStepConfig( + name="c", + code="return {'v': ctx.steps['b'].output['v'] + 1}", + upstreams=["b"], + ), + TransformStepConfig(name="d", code="return {'v': 100}"), + ], + ) + engine = PipelineEngine(config, MagicMock()) + result = await engine.run({}) + + assert result.steps["a"].status.value == "failed" + # b depends directly on the failed step a -> skipped, not run. + assert result.steps["b"].status.value == "skipped" + assert "a" in result.steps["b"].error + # c depends on b, which is now unsuccessful too -> cascaded skip. + assert result.steps["c"].status.value == "skipped" + # d has no dependency on the failed branch -> runs normally. + assert result.steps["d"].status.value == "success" + @pytest.mark.asyncio async def test_cycle_fails_pipeline_without_validate(self) -> None: config = _dag_config( diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py index e8250a0..364c730 100644 --- a/tests/test_pipelines.py +++ b/tests/test_pipelines.py @@ -986,6 +986,26 @@ def test_discover_pipelines_in_subdir(self, tmp_path): pipelines = PipelineLoader.discover_pipelines(tmp_path) assert "research" in pipelines + def test_discover_pipelines_in_a_per_pipeline_folder(self, tmp_path): + """``pipelines//pipeline.yaml`` is what the CLI scaffolds. + + ``agentomatic init NAME --template pipeline`` writes the YAML into its + own folder next to ``dataset.jsonl``, ``eval.py`` and a ``Makefile``. + Discovery only scanned flat ``pipelines/*.yaml``, so a freshly + scaffolded pipeline was invisible: ``agentomatic pipeline list`` + reported none and the Pipelines API mounted empty. + """ + from agentomatic.pipelines.loader import PipelineLoader + + folder = tmp_path / "pipelines" / "estimation" + folder.mkdir(parents=True) + (folder / "pipeline.yaml").write_text("name: estimation\nsteps:\n - agent: estimator\n") + (folder / "dataset.jsonl").write_text("{}\n") + + assert "estimation" in PipelineLoader.discover_pipelines(tmp_path) + # Also when the pipelines/ folder itself is what gets passed. + assert "estimation" in PipelineLoader.discover_pipelines(tmp_path / "pipelines") + def test_discover_pipelines_flat_pipelines_dir(self, tmp_path): """When callers pass pipelines/ itself, flat *.yaml must load.""" from agentomatic.pipelines.loader import PipelineLoader diff --git a/tests/test_platform_features.py b/tests/test_platform_features.py index a1cf626..1836a51 100644 --- a/tests/test_platform_features.py +++ b/tests/test_platform_features.py @@ -650,24 +650,31 @@ async def test_checkpoint_serialization_with_non_json_objects(store): """Verify checkpointer handles non-JSON-serializable objects (datetimes, custom classes).""" from datetime import datetime - from agentomatic.storage.checkpointer import AgentomaticCheckpointer, _ensure_json_serializable + from agentomatic.storage.checkpointer import ( + AgentomaticCheckpointer, + decode_from_storage, + encode_for_storage, + ) checkpointer = AgentomaticCheckpointer(store) - # 1. Test _ensure_json_serializable directly - assert _ensure_json_serializable({"a": 1, "b": "text"}) == {"a": 1, "b": "text"} + # 1. Test encode_for_storage / decode_from_storage round-trip directly + encoded = encode_for_storage({"a": 1, "b": "text"}) + assert isinstance(encoded, dict) + assert decode_from_storage(encoded) == {"a": 1, "b": "text"} - # 2. Test with datetime values (non-JSON-native) + # 2. Test with datetime values (non-JSON-native) — value survives as a real datetime. dt = datetime(2026, 6, 14, 12, 0, 0) - result = _ensure_json_serializable({"ts": dt, "val": 42}) - assert result["val"] == 42 - assert isinstance(result["ts"], str) # datetime converted to string + encoded_dt = encode_for_storage({"ts": dt, "val": 42}) + decoded_dt = decode_from_storage(encoded_dt) + assert decoded_dt["val"] == 42 + assert decoded_dt["ts"] == dt - # 3. Test with bytes - result_bytes = _ensure_json_serializable({"data": b"binary"}) - assert isinstance(result_bytes["data"], str) + # 3. Test with bytes — value survives as real bytes. + encoded_bytes = encode_for_storage({"data": b"binary"}) + assert decode_from_storage(encoded_bytes)["data"] == b"binary" - # 4. Test full round-trip through checkpointer + # 4. Test full round-trip through checkpointer, encoded value is JSON-safe for storage. config = { "configurable": { "thread_id": "thread_serde_test", @@ -682,8 +689,51 @@ async def test_checkpoint_serialization_with_non_json_objects(store): retrieved = await checkpointer.aget_tuple(config) assert retrieved is not None assert retrieved.checkpoint["channel_values"]["key"] == "val" - # The datetime should be stored as a string - assert isinstance(retrieved.checkpoint["ts"], str) + # The datetime should round-trip back to a real datetime, not a string. + assert retrieved.checkpoint["ts"] == dt + + +@pytest.mark.asyncio +async def test_checkpoint_serialization_preserves_langchain_messages(store): + """LangChain BaseMessage objects in channel_values must survive a checkpoint round-trip. + + A naive ``json.dumps(obj, default=str)`` would stringify HumanMessage/AIMessage + to their repr(), breaking the ``add_messages`` reducer and any chain built on + ``prompt | llm`` when the graph resumes from a checkpoint. + """ + from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage + + from agentomatic.storage.checkpointer import AgentomaticCheckpointer + + checkpointer = AgentomaticCheckpointer(store) + config = { + "configurable": { + "thread_id": "thread_lc_messages", + "checkpoint_ns": "", + "checkpoint_id": "cp_lc_1", + } + } + messages = [ + SystemMessage(content="be helpful"), + HumanMessage(content="hi there"), + AIMessage( + content="", + tool_calls=[{"name": "search", "args": {"q": "hi"}, "id": "call_1"}], + ), + ToolMessage(content="result", tool_call_id="call_1"), + ] + checkpoint = {"v": 1, "channel_values": {"messages": messages}} + metadata = {"source": "input", "step": 1} + + await checkpointer.aput(config, checkpoint, metadata, {}) + retrieved = await checkpointer.aget_tuple(config) + + assert retrieved is not None + restored = retrieved.checkpoint["channel_values"]["messages"] + assert [type(m) for m in restored] == [SystemMessage, HumanMessage, AIMessage, ToolMessage] + assert restored[1].content == "hi there" + assert restored[2].tool_calls[0]["name"] == "search" + assert restored[3].tool_call_id == "call_1" # ========================================================================= @@ -995,8 +1045,8 @@ async def test_lineage_cycle_guard_sqlalchemy(): @pytest.mark.asyncio async def test_ensure_json_serializable_nested(): - """Verify _ensure_json_serializable handles deeply nested non-serializable objects.""" - from agentomatic.storage.checkpointer import _ensure_json_serializable + """Verify the checkpoint encode/decode round-trip handles deeply nested objects.""" + from agentomatic.storage.checkpointer import decode_from_storage, encode_for_storage nested = { "level1": { @@ -1008,11 +1058,11 @@ async def test_ensure_json_serializable_nested(): }, "list_with_dt": [datetime(2026, 6, 1), "normal", 123], } - result = _ensure_json_serializable(nested) + result = decode_from_storage(encode_for_storage(nested)) assert result["level1"]["level2"]["num"] == 42 - assert isinstance(result["level1"]["level2"]["dt"], str) - assert isinstance(result["level1"]["level2"]["data"], str) - assert isinstance(result["list_with_dt"][0], str) + assert result["level1"]["level2"]["dt"] == datetime(2026, 1, 1, 12, 0) + assert result["level1"]["level2"]["data"] == b"binary_nested" + assert result["list_with_dt"][0] == datetime(2026, 6, 1) assert result["list_with_dt"][1] == "normal" @@ -1050,6 +1100,70 @@ async def test_delete_thread_cleans_up_suspended_states(store): assert state is None +@pytest.mark.asyncio +async def test_delete_thread_cleans_up_checkpoints_memory_store(store): + """Deleting a thread must not leave its LangGraph checkpoints orphaned.""" + await store.create_thread("cp_del_thread", "user_1", "agent_1") + await store.save_checkpoint( + thread_id="cp_del_thread", + checkpoint_ns="", + checkpoint_id="cp_1", + parent_checkpoint_id=None, + checkpoint={"v": 1}, + metadata={"source": "input"}, + ) + + assert await store.get_checkpoint("cp_del_thread", "", "cp_1") is not None + + await store.delete_thread("cp_del_thread") + + assert await store.get_checkpoint("cp_del_thread", "", "cp_1") is None + assert await store.list_checkpoints("cp_del_thread", "") == [] + + +@pytest.mark.asyncio +async def test_delete_thread_cleans_up_checkpoints_sqlalchemy_store(): + """Same guarantee on the SQLAlchemy backend (checkpoints have no DB-level FK).""" + db_store = SQLAlchemyStore("sqlite+aiosqlite:///:memory:") + await db_store.initialize() + try: + await db_store.create_thread("cp_del_thread_sqla", "user_1", "agent_1") + await db_store.save_checkpoint( + thread_id="cp_del_thread_sqla", + checkpoint_ns="", + checkpoint_id="cp_1", + parent_checkpoint_id=None, + checkpoint={"v": 1}, + metadata={"source": "input"}, + ) + + assert await db_store.get_checkpoint("cp_del_thread_sqla", "", "cp_1") is not None + + await db_store.delete_thread("cp_del_thread_sqla") + + assert await db_store.get_checkpoint("cp_del_thread_sqla", "", "cp_1") is None + assert await db_store.list_checkpoints("cp_del_thread_sqla", "") == [] + finally: + await db_store.close() + + +@pytest.mark.asyncio +async def test_sqlite_foreign_keys_are_enforced(): + """FK constraints (ondelete=CASCADE on feedback/suspended_state) must be + enforced on SQLite, which disables FK checks per-connection by default. + """ + from sqlalchemy import text + + db_store = SQLAlchemyStore("sqlite+aiosqlite:///:memory:") + await db_store.initialize() + try: + async with db_store._engine.connect() as conn: + result = await conn.execute(text("PRAGMA foreign_keys")) + assert result.scalar() == 1 + finally: + await db_store.close() + + @pytest.mark.asyncio async def test_create_thread_returns_consistent_shape(store): """Verify create_thread returns dict with parent_thread_id and fork_message_index.""" diff --git a/tests/test_production_readiness.py b/tests/test_production_readiness.py new file mode 100644 index 0000000..d44513b --- /dev/null +++ b/tests/test_production_readiness.py @@ -0,0 +1,547 @@ +# pyright: reportMissingParameterType=none +# pyright: reportCallIssue=none +# pyright: reportArgumentType=none +# pyright: reportAttributeAccessIssue=none +"""Regression tests for release-blocking defects found by end-to-end testing. + +Each test here corresponds to something that was observed failing (or being +unusably noisy) when the platform was actually booted and driven over HTTP, +rather than to a hypothetical from reading code. +""" + +from __future__ import annotations + +import sys +import warnings +from typing import Any + +import pytest +from conftest import install_plugin_package +from fastapi.testclient import TestClient + +from agentomatic import AgentManifest, AgentPlatform + + +async def _echo(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok", "agent_type": "echo"} + + +@pytest.fixture +def dual_mounted_platform(tmp_path): + """A platform whose agent's folder name and manifest slug differ. + + Such an agent is mounted under BOTH names so Studio (which addresses + agents by slug) does not 404. + """ + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + title="Prod Readiness", + ) + platform.register_agent( + manifest=AgentManifest(name="hello", slug="agent-hello", description="Hello"), + node_fn=_echo, + ) + return platform + + +# ===================================================================== +# OpenAPI: no duplicate operationIds from the name/slug dual mount +# ===================================================================== + + +def test_openapi_has_no_duplicate_operation_ids(dual_mounted_platform) -> None: + """Duplicate operationIds break OpenAPI client codegen. + + Mounting each agent under both its folder name and its slug previously + emitted one ``UserWarning: Duplicate Operation ID`` per route (~205 on a + small project) and produced a spec generators reject. + """ + app = dual_mounted_platform.build() + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + spec = app.openapi() + + duplicate_warnings = [w for w in caught if "Duplicate Operation ID" in str(w.message)] + assert not duplicate_warnings, ( + f"{len(duplicate_warnings)} duplicate operationId warning(s): " + f"{[str(w.message) for w in duplicate_warnings[:3]]}" + ) + + operation_ids = [ + operation["operationId"] + for path_item in spec["paths"].values() + for operation in path_item.values() + if isinstance(operation, dict) and "operationId" in operation + ] + assert len(operation_ids) == len(set(operation_ids)), "operationIds are not unique" + + +def test_slug_alias_routes_work_but_are_not_documented_twice(dual_mounted_platform) -> None: + """The slug mount is a compatibility alias: live, but not in the schema.""" + app = dual_mounted_platform.build() + 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" + ) + + with TestClient(app) as client: + # ...but it must still route, so Studio's slug-based calls keep working. + assert client.post("/api/v1/agent-hello/invoke", json={"query": "x"}).status_code == 200 + assert client.post("/api/v1/hello/invoke", json={"query": "x"}).status_code == 200 + + +# ===================================================================== +# OpenTelemetry console export is opt-in +# ===================================================================== + + +def test_otel_console_export_is_opt_in_by_default(monkeypatch) -> None: + """Console span export must not default on. + + It previously attached whenever no OTLP endpoint was configured — i.e. for + most deployments — dumping a full JSON span document to stdout for every + single HTTP request. + """ + from agentomatic.observability import telemetry + + monkeypatch.delenv("AGENTOMATIC_OTEL_CONSOLE", raising=False) + assert telemetry._env_flag("AGENTOMATIC_OTEL_CONSOLE") is False + + for truthy in ("1", "true", "TRUE", "yes", "on"): + monkeypatch.setenv("AGENTOMATIC_OTEL_CONSOLE", truthy) + assert telemetry._env_flag("AGENTOMATIC_OTEL_CONSOLE") is True, truthy + + for falsy in ("0", "false", "no", "", "off"): + monkeypatch.setenv("AGENTOMATIC_OTEL_CONSOLE", falsy) + assert telemetry._env_flag("AGENTOMATIC_OTEL_CONSOLE") is False, falsy + + +def test_otel_setup_does_not_attach_console_exporter_by_default(monkeypatch) -> None: + """Without the opt-in env var, no ConsoleSpanExporter is registered.""" + pytest.importorskip("opentelemetry.sdk") + from agentomatic.observability import telemetry + + if not telemetry.HAS_OTEL: # pragma: no cover - depends on extras + pytest.skip("OpenTelemetry not installed") + + monkeypatch.delenv("AGENTOMATIC_OTEL_CONSOLE", raising=False) + monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False) + + attached: list[Any] = [] + real_provider_cls = telemetry.TracerProvider + + class _RecordingProvider(real_provider_cls): # type: ignore[misc, valid-type] + def add_span_processor(self, processor: Any) -> None: + attached.append(processor) + super().add_span_processor(processor) + + monkeypatch.setattr(telemetry, "TracerProvider", _RecordingProvider) + telemetry.setup_telemetry(app=None, service_name="test-svc") + + assert not attached, ( + "a span processor was attached with no OTLP endpoint and console export " + "off — this is the per-request stdout span dump regression" + ) + + +# ===================================================================== +# `agentomatic run` can import the project's main.py +# ===================================================================== + + +def test_run_puts_project_dir_on_sys_path_before_uvicorn(monkeypatch, tmp_path) -> None: + """``uvicorn.run("main:app")`` resolves the import string against sys.path. + + Launched as a console script (``uv run agentomatic run``), the project + directory is not on sys.path, so importing ``main`` failed outright. The + run command must add it (and export PYTHONPATH for the --reload child). + """ + import os + import sys + + from agentomatic.cli import commands + + (tmp_path / "main.py").write_text("app = object()\n", encoding="utf-8") + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(commands, "_has_project_main_app", lambda *a, **k: True) + monkeypatch.delenv("PYTHONPATH", raising=False) + + captured: dict[str, Any] = {} + + class _FakeUvicorn: + @staticmethod + def run(target: str, **kwargs: Any) -> None: + captured["target"] = target + captured["app_dir"] = kwargs.get("app_dir") + captured["sys_path_0"] = sys.path[0] + captured["pythonpath"] = os.environ.get("PYTHONPATH", "") + + monkeypatch.setitem(sys.modules, "uvicorn", _FakeUvicorn) + + original_sys_path = list(sys.path) + try: + # Invoke the click command's underlying callback directly. + commands.run.callback( + agents_dir="agents", + plugins_dir="plugins", + endpoints_dir="endpoints", + ingestion_dir="ingestion", + stacks_dir="stacks", + host="127.0.0.1", + port=8000, + reload=False, + title=None, + log_level="INFO", + with_ui=False, + studio=True, + ssl_certfile=None, + ssl_keyfile=None, + require_auth_globally=False, + ) + finally: + sys.path[:] = original_sys_path + + project_dir = str(tmp_path) + assert captured["target"] == "main:app" + assert captured["app_dir"] == project_dir + assert captured["sys_path_0"] == project_dir + assert project_dir in captured["pythonpath"].split(os.pathsep) + + +# ===================================================================== +# Scaffold hygiene: no filesystem paths or dead collectors leaked +# ===================================================================== + + +def test_project_title_uses_only_the_directory_name() -> None: + """Regression: ``agentomatic new /srv/apps/my_proj`` baked the whole + filesystem path into the platform title, which is published via + ``/openapi.json``, ``/.well-known/agent.json`` and ``/studio/info`` — + leaking the server's directory layout to any caller. + """ + from agentomatic.cli.project import get_project_files + + main_py = get_project_files("/srv/apps/my_proj")["main.py"] + + assert '"My Proj Platform"' in main_py + assert "/srv/apps" not in main_py.split("description=")[0] + + +def test_env_example_does_not_enable_a_collector_that_is_not_there() -> None: + """An OTLP endpoint set with nothing listening makes the exporter retry + every span and flood the log with transient-failure warnings. + """ + from agentomatic.cli.project import get_project_files + + env_example = get_project_files("demo")[".env.example"] + + for line in env_example.splitlines(): + stripped = line.strip() + if stripped.startswith("OTEL_EXPORTER_OTLP_ENDPOINT="): + pytest.fail(f"OTLP endpoint enabled by default in .env.example: {stripped!r}") + # It should still be documented, just commented out. + assert "OTEL_EXPORTER_OTLP_ENDPOINT" in env_example + + +def test_full_template_agent_keeps_its_generated_endpoints() -> None: + """A module-level ``router`` in an agent's ``api.py`` REPLACES every + auto-generated endpoint (/invoke, /chat, /stream, /card, /health). The + ``full`` template shipped one, so the flagship scaffold produced an agent + with no way to invoke it. + """ + from agentomatic.cli.templates import get_template_files + + api_py = get_template_files("full", "sample_agent")["api.py"] + + # The registry keys on the exact name ``router``; anything else is inert. + assert "\nrouter = APIRouter()" not in api_py + assert "custom_router = APIRouter()" in api_py + # The pattern is still demonstrated and explained. + assert "REPLACES" in api_py + + +def test_all_extra_contents_match_what_the_docs_claim() -> None: + """``[all]`` is the advertised "recommended" install, so what it contains + must stay in sync with the documented list — a user following the docs and + then hitting a missing dependency is a release defect. + """ + import tomllib + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + with (root / "pyproject.toml").open("rb") as fh: + pyproject = tomllib.load(fh) + + extras = pyproject["project"]["optional-dependencies"] + all_spec = " ".join(extras["all"]) + included = {e.strip() for e in all_spec.split("[", 1)[1].rstrip("]\"' ").split(",")} + + documented = { + "langgraph", + "ollama", + "metrics", + "db", + "cli", + "studio", + "optimize", + "telemetry", + "dotenv", + "security", + "swarm", + "vector", + } + assert included == documented, ( + "The `all` extra changed; update docs/getting-started/installation.md " + f"(added={included - documented}, removed={documented - included})" + ) + + # These are deliberately excluded — vendor SDKs (provider-agnostic + # principle), an alternative DB driver, and the heavy Chainlit UI. + for deliberately_excluded in ("openai", "azure", "vertex", "db-postgres", "ui"): + assert deliberately_excluded in extras, f"{deliberately_excluded} extra vanished" + assert deliberately_excluded not in included + + +def test_platform_marks_plugin_loaded_even_if_subclass_forgets_super(tmp_path) -> None: + """A plugin overriding ``load_model`` without calling ``super()`` used to + stay ``_is_loaded=False`` forever: /predict answered 503 and /health went + "degraded", while startup logged "loaded successfully". The platform now + stamps the flag itself so the footgun cannot produce a dead plugin. + """ + from fastapi.testclient import TestClient + + from agentomatic import AgentPlatform + + plugins_dir = tmp_path / "plugins" + source = '''"""Plugin that overrides load_model without calling super().""" +from __future__ import annotations + +from pydantic import BaseModel + +from agentomatic.plugins import BaseMLPlugin + + +class Inp(BaseModel): + text: str + + +class Out(BaseModel): + result: str + + +class ForgetfulPlugin(BaseMLPlugin[Inp, Out]): + plugin_name = "forgetful" + + async def load_model(self) -> None: + self.model = object() # deliberately no: await super().load_model() + + async def predict(self, inputs: Inp) -> Out: + return Out(result=inputs.text.upper()) +''' + importable = install_plugin_package(plugins_dir, "forgetful", source) + + with importable: + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=plugins_dir, + endpoints_dir=tmp_path / "endpoints", + ) + with TestClient(platform.build()) as client: + assert client.get("/health").json()["status"] == "healthy" + response = client.post("/api/v1/plugins/forgetful/predict", json={"text": "hi"}) + assert response.status_code == 200, response.text + assert response.json()["result"] == "HI" + + +# ===================================================================== +# Log hygiene: one line per request, one separator, one DDL pass +# ===================================================================== + + +def test_log_format_separator_matches_loguru_default() -> None: + """Lines emitted before ``configure_logging`` installs our sink use loguru's + built-in format. Using a different separator afterwards produced two + formats in one log, which breaks log-shipping regexes. + """ + import inspect + + from agentomatic.core.lifespan import configure_logging + + source = inspect.getsource(configure_logging) + assert "{line} - " in source + assert "—" not in source, "em dash in the log format: non-ASCII and inconsistent" + + +def test_platform_run_disables_uvicorn_access_log_when_middleware_logs(tmp_path) -> None: + """The platform's LoggingMiddleware already logs every request, so leaving + uvicorn's access log on doubles the volume for the same information. + """ + from unittest.mock import patch + + from agentomatic import AgentPlatform + + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + enable_logging=True, + ) + with patch("uvicorn.run") as mock_run: + platform.run(host="127.0.0.1", port=9999) + assert mock_run.call_args.kwargs["access_log"] is False + + # An explicit choice from the caller still wins. + with patch("uvicorn.run") as mock_run: + platform.run(host="127.0.0.1", port=9999, access_log=True) + assert mock_run.call_args.kwargs["access_log"] is True + + +def test_platform_run_keeps_access_log_when_middleware_is_off(tmp_path) -> None: + """With the middleware disabled, uvicorn's access log is the only record + of requests — it must not be suppressed. + """ + from unittest.mock import patch + + from agentomatic import AgentPlatform + + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + enable_logging=False, + ) + with patch("uvicorn.run") as mock_run: + platform.run(host="127.0.0.1", port=9999) + assert "access_log" not in mock_run.call_args.kwargs + + +@pytest.mark.asyncio +async def test_sqlalchemy_store_initialize_is_idempotent(tmp_path) -> None: + """Startup can reach ``initialize()`` from several paths (configured store, + one derived from DATABASE_URL, and a post-connection pass). Re-running the + DDL each time is wasted round trips and duplicate log lines. + """ + from loguru import logger + + from agentomatic.storage.sqlalchemy import SQLAlchemyStore + + messages: list[str] = [] + sink_id = logger.add(lambda m: messages.append(m), level="INFO") + + store = SQLAlchemyStore(url=f"sqlite+aiosqlite:///{tmp_path / 'x.db'}") + try: + await store.initialize() + assert store._initialized is True + await store.initialize() + await store.initialize() + finally: + logger.remove(sink_id) + await store.close() + + created = [m for m in messages if "Database tables created/verified" in m] + assert len(created) == 1, f"DDL ran {len(created)} times, expected once" + + +def test_store_dependent_routes_fail_cleanly_without_a_store(tmp_path) -> None: + """No store is the *default*, so these routes must answer, not crash. + + ``thread_store`` is a ``_LazyStoreProxy`` — never ``None``, with a + ``__bool__`` reporting whether a store actually exists. Five routes guarded + with ``is None``, which is never true for the proxy, so the guard fell + through and the first attribute access raised RuntimeError out of the + handler as a bare 500 with no body. + """ + from fastapi.testclient import TestClient + + from agentomatic import AgentManifest, AgentPlatform + + async def echo(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok"} + + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + ) + platform.register_agent( + manifest=AgentManifest(name="a1", slug="a1", description="echo"), + node_fn=echo, + ) + + with TestClient(platform.build(), raise_server_exceptions=False) as client: + for path in ( + "/api/v1/a1/optimization-runs", + "/api/v1/a1/logs", + "/api/v1/a1/logs/analysis", + "/api/v1/a1/threads", + ): + response = client.get(path) + assert response.status_code != 500, ( + f"{path} raised an unhandled exception without a store configured: " + f"{response.text[:200]}" + ) + # And it must say something useful rather than an empty body. + assert response.text.strip(), f"{path} returned an empty body" + + +def test_log_level_applies_to_build_not_just_startup(tmp_path, monkeypatch) -> None: + """``--profile minimal`` bakes ``LOG_LEVEL=WARNING`` and must be obeyed. + + ``build()`` narrates discovery and every mount. Configuring loguru only in + the lifespan (which runs at startup, after ``build()`` returns) let all of + that INFO/DEBUG output through, so the "quieter logs" the minimal profile + advertises never materialised in a container. + """ + import io + + from agentomatic.core.lifespan import configure_logging + from agentomatic.core.platform import AgentPlatform + + agents_dir = tmp_path / "agents" + agents_dir.mkdir() + + stream = io.StringIO() + monkeypatch.setattr(sys, "stdout", stream) + try: + # Start noisy, so a build that ignores ``log_level`` is visible. + configure_logging("DEBUG") + AgentPlatform(agents_dir=agents_dir, log_level="WARNING").build() + finally: + monkeypatch.undo() + # Leave the global logger as the rest of the suite expects it. + configure_logging("INFO") + + noisy = [ + line + for line in stream.getvalue().splitlines() + if " | INFO " in line or " | DEBUG " in line + ] + assert not noisy, f"build() logged below WARNING: {noisy[:5]}" + + +def test_build_backend_is_pinned_to_publishable_metadata() -> None: + """The wheel must carry metadata PyPI's tooling accepts. + + hatchling 1.30+ emits ``Metadata-Version: 2.5``, which current twine + rejects outright ("'2.5' is not a valid metadata version"). The release + workflow publishes through ``pypa/gh-action-pypi-publish``, which verifies + metadata with twine before upload — so an unbounded ``requires`` turns + every release into a coin flip on whatever hatchling resolves that day. + """ + import tomllib + from pathlib import Path + + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + requires = tomllib.loads(pyproject.read_text())["build-system"]["requires"] + + hatchling = [r for r in requires if r.replace("_", "-").startswith("hatchling")] + assert hatchling, f"expected hatchling in build-system.requires, got {requires}" + assert "<" in hatchling[0], ( + f"build backend {hatchling[0]!r} has no upper bound — a newer hatchling " + "can emit metadata that twine and PyPI reject at publish time" + ) diff --git a/tests/test_rc_fixes.py b/tests/test_rc_fixes.py index c999ef5..59f5538 100644 --- a/tests/test_rc_fixes.py +++ b/tests/test_rc_fixes.py @@ -204,8 +204,13 @@ def test_input_from_state_flattens_context(self) -> None: assert payload["snapshot"] == {"k": 1} assert payload["query"] == "q" # top-level wins over context.query assert payload["context"] == {"snapshot": {"k": 1}, "query": "ignored"} - assert "messages" not in payload - assert "thread_id" not in payload + # ``messages`` and ``thread_id`` are forwarded to ``input_to_state``. + # They were previously dropped as "conversation bookkeeping", which + # made LangChain-style class agents unable to see prior turns (for a + # MessagesPlaceholder) or thread a RunnableConfig thread_id through — + # the scaffolded template read both and always got empty values. + assert payload["messages"] == [] + assert payload["thread_id"] == "t1" def test_sync_invoke_reads_flattened_context(self, tmp_path: Any) -> None: """REST ``/invoke`` with ``context.snapshot`` reaches ``input_to_state``.""" diff --git a/tests/test_security.py b/tests/test_security.py index 9381874..6f6a2c4 100644 --- a/tests/test_security.py +++ b/tests/test_security.py @@ -288,11 +288,19 @@ def test_blacklist_denies(self) -> None: def _mock_request( claims: dict | None = None, user_id: str | None = None, + api_key_authenticated: bool = False, ) -> MagicMock: """Build a mock Starlette Request with ``state`` attributes.""" request = MagicMock() state = MagicMock() + # ``MagicMock`` would answer truthily to *any* attribute, which would make + # every mocked request look API-key authenticated. + if api_key_authenticated: + state.api_key_authenticated = True + else: + del state.api_key_authenticated + # MagicMock auto-creates attributes, so we need to explicitly control # which attributes are present on state. if claims is not None: @@ -328,6 +336,24 @@ def test_auth_required_globally_no_claims_denied(self) -> None: assert ok is False assert "Authentication is required" in reason + def test_api_key_authentication_satisfies_the_global_auth_lock(self) -> None: + """An API key is a valid credential — it just carries no claims.""" + enforcer = ZeroTrustEnforcer(require_auth_globally=True) + request = _mock_request(api_key_authenticated=True) + + ok, reason = enforcer.verify_request(request, "agent") + assert ok is True, reason + + def test_api_key_cannot_satisfy_a_role_restricted_policy(self) -> None: + """Fail closed: a key has no roles, so a role policy is unevaluable.""" + enforcer = ZeroTrustEnforcer(require_auth_globally=True) + enforcer.register_policy("agent", AgentSecurityPolicy(allowed_roles=["admin"])) + request = _mock_request(api_key_authenticated=True) + + ok, reason = enforcer.verify_request(request, "agent") + assert ok is False + assert "API key" in reason + def test_auth_required_per_policy_no_claims_denied(self) -> None: enforcer = ZeroTrustEnforcer() enforcer.register_policy("agent", AgentSecurityPolicy(require_auth=True)) diff --git a/tests/test_security_hardening.py b/tests/test_security_hardening.py new file mode 100644 index 0000000..4007ea7 --- /dev/null +++ b/tests/test_security_hardening.py @@ -0,0 +1,739 @@ +# pyright: reportMissingParameterType=none +# pyright: reportCallIssue=none +# pyright: reportArgumentType=none +# pyright: reportAttributeAccessIssue=none +"""Security regressions: path confinement, error sanitisation, secret redaction. + +Each test here corresponds to a vulnerability that was reproduced by actually +executing it against a running platform, not inferred from reading code. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from agentomatic import AgentManifest, AgentPlatform + +# ===================================================================== +# Ingestion path confinement +# ===================================================================== + + +class _StubIngestionContext: + """Minimal ingestion context (the real one is a Protocol).""" + + cancelled = False + + async def report(self, **kwargs: Any) -> None: + return None + + +@pytest.fixture +def ingest_env(monkeypatch, tmp_path): + """An ingestion root plus a secret file safely outside it.""" + from agentomatic.ingestion.paths import INGESTION_ROOT_ENV + + root = tmp_path / "project" + root.mkdir() + (root / "legit.txt").write_text("legitimate project content", encoding="utf-8") + + outside = tmp_path / "elsewhere" + outside.mkdir() + secret = outside / "secret.txt" + secret.write_text("TOPSECRET-ABC123", encoding="utf-8") + + monkeypatch.setenv(INGESTION_ROOT_ENV, str(root)) + return {"root": root, "secret": secret, "out": root / "out"} + + +async def _ingest(**kwargs: Any): + from agentomatic.ingestion.builtin.markdown import MarkdownIngestor, MarkdownIngestRequest + + kwargs.setdefault("engine", "plain") + return await MarkdownIngestor().ingest( + MarkdownIngestRequest(**kwargs), _StubIngestionContext() + ) + + +@pytest.mark.asyncio +async def test_ingestion_rejects_absolute_path_outside_root(ingest_env) -> None: + """`source=/etc/passwd` must not exfiltrate arbitrary files.""" + result = await _ingest(source="/etc/passwd", output_dir=str(ingest_env["out"])) + assert result.status == "failed" + assert "outside the ingestion root" in result.errors[0] + + +@pytest.mark.asyncio +async def test_ingestion_rejects_reading_secret_outside_root(ingest_env) -> None: + result = await _ingest(source=str(ingest_env["secret"]), output_dir=str(ingest_env["out"])) + assert result.status == "failed" + # The secret's contents must not have been copied anywhere. + assert not list(ingest_env["out"].glob("*.md")) if ingest_env["out"].exists() else True + + +@pytest.mark.asyncio +async def test_ingestion_rejects_write_outside_root(ingest_env, tmp_path) -> None: + target = tmp_path / "pwned_dir" + result = await _ingest(source="legit.txt", output_dir=str(target)) + assert result.status == "failed" + assert not target.exists(), "attacker-chosen directory must not be created" + + +@pytest.mark.asyncio +async def test_ingestion_rejects_filename_traversal(ingest_env) -> None: + """output_filename is a name, not a path — traversal escapes the out dir.""" + result = await _ingest( + source="legit.txt", + output_dir=str(ingest_env["out"]), + output_filename="../../pwned.txt", + ) + assert result.status == "failed" + assert "bare filename" in result.errors[0] + + +@pytest.mark.asyncio +async def test_ingestion_still_works_inside_the_root(ingest_env) -> None: + """Confinement must not break legitimate in-project ingestion.""" + result = await _ingest(source="legit.txt", output_dir=str(ingest_env["out"])) + assert result.status == "succeeded", result.errors + assert (ingest_env["out"] / "legit.md").exists() + + +def test_ingestion_root_defaults_to_cwd(monkeypatch, tmp_path) -> None: + from agentomatic.ingestion.paths import INGESTION_ROOT_ENV, ingestion_root + + monkeypatch.delenv(INGESTION_ROOT_ENV, raising=False) + monkeypatch.chdir(tmp_path) + assert ingestion_root() == tmp_path.resolve() + + +def test_safe_output_filename_rejects_separators() -> None: + from agentomatic.ingestion.paths import IngestionPathError, safe_output_filename + + assert safe_output_filename(None, default="x.md") == "x.md" + assert safe_output_filename("report.md", default="x.md") == "report.md" + for bad in ("../evil", "a/b.md", "..", "."): + with pytest.raises(IngestionPathError): + safe_output_filename(bad, default="x.md") + + +# ===================================================================== +# Exception message sanitisation +# ===================================================================== + +_SECRET_DSN = "postgres://user:HUNTER2@db:5432" + + +@pytest.fixture +def leaky_client(tmp_path): + async def boom(state: dict[str, Any]) -> dict[str, Any]: + raise RuntimeError(f"connection failed: {_SECRET_DSN}") + + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + ) + platform.register_agent( + manifest=AgentManifest(name="boom", slug="boom", description="raises"), + node_fn=boom, + ) + with TestClient(platform.build(), raise_server_exceptions=False) as client: + yield client + + +@pytest.mark.parametrize( + ("path", "payload"), + [ + ("/api/v1/boom/invoke", {"query": "x"}), + ("/api/v1/boom/chat", {"content": "x"}), + ("/api/v1/boom/invoke/stream", {"query": "x"}), + ], +) +def test_exception_text_does_not_leak_to_clients(leaky_client, monkeypatch, path, payload) -> None: + """A driver/exception message routinely carries credentials — it must not + be interpolated into an HTTP response. + """ + from agentomatic.core.errors import DEBUG_ERRORS_ENV + + monkeypatch.delenv(DEBUG_ERRORS_ENV, raising=False) + response = leaky_client.post(path, json=payload) + + assert "HUNTER2" not in response.text + assert _SECRET_DSN not in response.text + # A correlation id is returned so operators can find the full server log. + assert "error_id" in response.text + + +def test_debug_errors_opt_in_restores_raw_detail(leaky_client, monkeypatch) -> None: + """Local development can opt back into raw exception text.""" + from agentomatic.core.errors import DEBUG_ERRORS_ENV + + monkeypatch.setenv(DEBUG_ERRORS_ENV, "1") + response = leaky_client.post("/api/v1/boom/invoke", json={"query": "x"}) + assert "HUNTER2" in response.text + + +def test_client_safe_detail_shape() -> None: + from agentomatic.core.errors import client_safe_detail + + payload = client_safe_detail(ValueError("secret-value"), context="Thing failed") + assert payload["error"] == "Thing failed" + assert payload["error_type"] == "ValueError" + assert len(payload["error_id"]) == 12 + assert "secret-value" not in str(payload) + + +# ===================================================================== +# Stack secret redaction +# ===================================================================== + +_STACK_YAML = """# Prod stack +name: prod +llm: + default: + provider: openai + api_key: ${OPENAI_API_KEY} + model: gpt-4o + legacy: + api_key: sk-proj-REALSECRET123 +database: + url: postgresql+asyncpg://admin:SuperSecret99@db.internal:5432/app + pool_size: 10 +""" + + +def test_redaction_masks_literal_secrets_but_keeps_env_refs() -> None: + from agentomatic.stacks.redaction import redact_yaml_text + + redacted, count = redact_yaml_text(_STACK_YAML) + + assert "sk-proj-REALSECRET123" not in redacted + assert "SuperSecret99" not in redacted + # ${ENV_VAR} indirections are not secrets and must stay visible. + assert "${OPENAI_API_KEY}" in redacted + # Non-secret structure survives so the file is still readable. + assert "provider: openai" in redacted + assert "pool_size: 10" in redacted + assert "postgresql+asyncpg://admin:" in redacted # host/user kept for debugging + assert count == 2 + + +def test_env_example_value_never_emits_a_literal_secret() -> None: + """.env.example is conventionally committed — a literal key would publish it.""" + from agentomatic.stacks.redaction import ENV_EXAMPLE_PLACEHOLDER, env_example_value + + assert env_example_value("${OPENAI_API_KEY}") == "${OPENAI_API_KEY}" + assert env_example_value("sk-proj-REALSECRET") == ENV_EXAMPLE_PLACEHOLDER + assert env_example_value("") == "" + + +def test_generated_env_example_contains_no_literal_secret(tmp_path) -> None: + """End-to-end: a stack with literal secrets must not leak into deploy output.""" + from agentomatic.cli import deploy as deploy_mod + + stacks_dir = tmp_path / "stacks" + stacks_dir.mkdir() + (stacks_dir / "prod.yaml").write_text(_STACK_YAML, encoding="utf-8") + + plan = deploy_mod.generate_deploy( + out_dir=tmp_path / "out", + stack_name="prod", + stacks_dir=stacks_dir, + ) + env_example = plan.files[".env.example"].read_text(encoding="utf-8") + + assert "sk-proj-REALSECRET123" not in env_example + assert "SuperSecret99" not in env_example + + +# ===================================================================== +# Studio resume: clear status code + sanitised errors +# ===================================================================== + + +def test_studio_resume_rejects_non_langgraph_agent_cleanly(tmp_path) -> None: + """Resume is a LangGraph feature (``astream_events`` + ``Command``). + + Agentomatic's own lightweight AgentGraph has neither, so the call used to + raise AttributeError *inside* the SSE body — returning HTTP 200 with the + raw internal message ``'AgentGraph' object has no attribute + 'astream_events'``. It must fail fast with a real status code instead. + """ + from agentomatic.agents.graph import AgentGraph, GraphNode + + def node(state: dict[str, Any]) -> dict[str, Any]: + return state + + def graph_fn() -> AgentGraph: + return AgentGraph( + nodes={"n": GraphNode(name="n", handler=node)}, + edges={}, + entrypoint="n", + finish="n", + ) + + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + enable_studio=True, + ) + + async def node_fn(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok"} + + platform.register_agent( + manifest=AgentManifest(name="plain", slug="plain", description="no langgraph"), + node_fn=node_fn, + graph_fn=graph_fn, + ) + + with TestClient(platform.build(), raise_server_exceptions=False) as client: + response = client.post( + "/studio/agents/plain/threads/does-not-exist/resume", + json={"value": "hi"}, + ) + + assert response.status_code == 501 + body = response.text + assert "astream_events" in body # actionable: names what's missing + assert "AttributeError" not in body + + +# ===================================================================== +# Async / background paths must sanitise too +# ===================================================================== + + +@pytest.fixture +def async_leaky_client(tmp_path): + """A platform whose agent fails with a credential-bearing exception.""" + + async def boom(state: dict[str, Any]) -> dict[str, Any]: + raise RuntimeError(f"connect failed: {_SECRET_DSN} at /srv/secret/config.yaml") + + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + enable_studio=True, + ) + platform.register_agent( + manifest=AgentManifest(name="boom", slug="boom", description="raises"), + node_fn=boom, + ) + with TestClient(platform.build(), raise_server_exceptions=False) as client: + yield client + + +def _await_terminal(client, task_id: str) -> dict[str, Any]: + import time + + for _ in range(60): + body = client.get(f"/api/v1/tasks/{task_id}").json() + if body.get("status") in {"succeeded", "failed", "cancelled"}: + return body + time.sleep(0.05) + raise AssertionError("task never reached a terminal state") + + +def test_background_task_record_does_not_leak_exception_text(async_leaky_client) -> None: + """The sync paths were sanitised first; the async ones serve the stored + record verbatim via /tasks/{id}, /result and the task list. + """ + submitted = async_leaky_client.post("/api/v1/boom/invoke/async", json={"query": "x"}) + task_id = submitted.json().get("id") or submitted.json().get("task_id") + record = _await_terminal(async_leaky_client, task_id) + + assert record["status"] == "failed" + assert "HUNTER2" not in str(record) + assert "/srv/secret" not in str(record) + # Still actionable: names the type and carries a correlation id. + assert "RuntimeError" in record["error"] + assert "error_id=" in record["error"] + + for path in (f"/api/v1/tasks/{task_id}", f"/api/v1/tasks/{task_id}/result", "/api/v1/tasks"): + body = async_leaky_client.get(path).text + assert "HUNTER2" not in body, f"{path} leaked the DSN" + assert "/srv/secret" not in body, f"{path} leaked a server path" + + +def test_studio_run_and_stream_do_not_leak_exception_text(async_leaky_client) -> None: + """Studio runs are reachable unauthenticated in the default `agentomatic + run` posture, and the error is both stored on the run and streamed by SSE. + """ + run = async_leaky_client.post("/studio/agents/boom/runs", json={"query": "x"}) + assert run.status_code == 200 + assert "HUNTER2" not in run.text + assert "/srv/secret" not in run.text + + stream = async_leaky_client.post("/studio/agents/boom/runs/stream", json={"query": "x"}) + assert "HUNTER2" not in stream.text + assert "/srv/secret" not in stream.text + + +def test_non_ascii_credentials_are_rejected_not_a_server_error(tmp_path) -> None: + """``hmac.compare_digest`` raises TypeError on a non-ASCII ``str``, which + turned a bad key into a 500 — trivially reachable via ``?api_key=…``. + """ + + async def echo(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok"} + + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + enable_auth=True, + auth_api_key="SECRETKEY", + ) + platform.register_agent( + manifest=AgentManifest(name="a1", slug="a1", description="echo"), + node_fn=echo, + ) + with TestClient(platform.build(), raise_server_exceptions=False) as client: + assert client.get("/api/v1/a1/health?api_key=%C3%A9vil").status_code == 401 + assert client.get("/api/v1/a1/health?api_key=wrong").status_code == 401 + assert ( + client.get("/api/v1/a1/health", headers={"X-API-Key": "SECRETKEY"}).status_code == 200 + ) + + +# ===================================================================== +# Studio must display decoded checkpoint state, not the storage wrapper +# ===================================================================== + + +@pytest.mark.asyncio +async def test_studio_state_and_history_decode_stored_checkpoints() -> None: + """Checkpoints are persisted through LangGraph's serde so BaseMessage + objects survive a round-trip. Studio reads those rows directly, so it must + decode them — otherwise the debug UI shows an opaque + ``{__agentomatic_serde_type__, __agentomatic_serde_data__}`` blob instead + of the actual state. + """ + from types import SimpleNamespace + + from langchain_core.messages import AIMessage, HumanMessage + + from agentomatic.storage.checkpointer import AgentomaticCheckpointer + from agentomatic.storage.memory import MemoryStore + from agentomatic.studio.adapters.langgraph import LangGraphAdapter + + store = MemoryStore() + await store.initialize() + checkpointer = AgentomaticCheckpointer(store) + await checkpointer.aput( + {"configurable": {"thread_id": "t1", "checkpoint_ns": "", "checkpoint_id": "c1"}}, + { + "v": 1, + "channel_values": { + "messages": [HumanMessage(content="hello"), AIMessage(content="hi back")], + "answer": "42", + }, + }, + {"source": "input"}, + {}, + ) + + agent = SimpleNamespace( + name="a1", + slug="a1", + graph_fn=None, + manifest=SimpleNamespace( + name="a1", slug="a1", description="d", version="1", framework="langgraph" + ), + ) + adapter = LangGraphAdapter(agent, store) + + snapshot = await adapter.get_state("t1") + assert "__agentomatic_serde_type__" not in str(snapshot.state) + channels = snapshot.state["channel_values"] + assert channels["answer"] == "42" + assert [m.content for m in channels["messages"]] == ["hello", "hi back"] + + history = await adapter.get_history("t1") + assert history + assert "__agentomatic_serde_type__" not in str(history[0].state) + + +class TestGlobalAuthLockStaysServable: + """``require_auth_globally`` with an API key must produce a working app. + + The build-time guard accepts that configuration — it is remedy (b) in its + own error message — but the scaffolded ``main.py`` also switches JWT auth + on when ``AGENTOMATIC_REQUIRE_AUTH`` is set. ``JWTAuthMiddleware`` then + refuses to construct without a ``jwks_url``, and Starlette builds the + middleware stack on the *first request*, not at ``build()``. The container + started clean and answered 500 to every route, ``/health`` and ``/docs`` + included. + """ + + def _app(self): + import tempfile + from pathlib import Path + + from agentomatic import AgentManifest, AgentPlatform + + async def echo(state): + return {"response": "ok", "agent_type": "echo"} + + tmp = Path(tempfile.mkdtemp()) + platform = AgentPlatform( + agents_dir=tmp / "agents", + plugins_dir=tmp / "plugins", + endpoints_dir=tmp / "endpoints", + enable_auth=True, + auth_api_key="zt-key", + enable_jwt_auth=True, + require_auth_globally=True, + ) + platform.register_agent( + manifest=AgentManifest(name="a1", slug="a1", description="echo"), + node_fn=echo, + ) + return platform.build() + + def test_every_route_still_answers(self): + from fastapi.testclient import TestClient + + with TestClient(self._app(), raise_server_exceptions=False) as client: + assert client.get("/health").status_code < 500 + assert client.get("/docs").status_code < 500 + + def test_zero_trust_accepts_an_api_key_authenticated_caller(self): + """Zero-trust ran before the API-key middleware and denied everything. + + With ``require_auth_globally`` the enforcer looked for JWT claims, + found none — the key had not been checked yet — and returned + ``zero_trust_denied`` to a caller presenting a perfectly valid key. + The configuration served no request at all. + """ + import tempfile + from pathlib import Path + + from fastapi.testclient import TestClient + + from agentomatic import AgentManifest, AgentPlatform + + async def echo(state): + return {"response": "ok", "agent_type": "echo"} + + tmp = Path(tempfile.mkdtemp()) + platform = AgentPlatform( + agents_dir=tmp / "agents", + plugins_dir=tmp / "plugins", + endpoints_dir=tmp / "endpoints", + enable_auth=True, + auth_api_key="zt-key", + enable_jwt_auth=True, + enable_zero_trust=True, + require_auth_globally=True, + ) + platform.register_agent( + manifest=AgentManifest(name="a1", slug="a1", description="echo"), + node_fn=echo, + ) + + with TestClient(platform.build(), raise_server_exceptions=False) as client: + authenticated = client.post( + "/api/v1/a1/invoke", + json={"query": "x"}, + headers={"X-API-Key": "zt-key"}, + ) + assert authenticated.status_code == 200, authenticated.text + + anonymous = client.post("/api/v1/a1/invoke", json={"query": "x"}) + assert anonymous.status_code == 401, anonymous.text + + def test_the_api_key_still_gates_agent_routes(self): + from fastapi.testclient import TestClient + + with TestClient(self._app(), raise_server_exceptions=False) as client: + unauthenticated = client.post("/api/v1/a1/invoke", json={"query": "x"}) + assert unauthenticated.status_code == 401, unauthenticated.text + + authenticated = client.post( + "/api/v1/a1/invoke", + json={"query": "x"}, + headers={"X-API-Key": "zt-key"}, + ) + assert authenticated.status_code == 200, authenticated.text + + def test_no_api_key_and_no_jwks_still_refuses_to_boot(self): + """The forged-JWT hole must stay closed.""" + import tempfile + from pathlib import Path + + import pytest + + from agentomatic import AgentPlatform + + tmp = Path(tempfile.mkdtemp()) + platform = AgentPlatform( + agents_dir=tmp / "agents", + enable_auth=False, + enable_jwt_auth=True, + require_auth_globally=True, + ) + with pytest.raises(RuntimeError, match="forged/unsigned JWTs"): + platform.build() + + +class TestJwtConfigFromEnvironmentAndStack: + """Verified JWT auth must be reachable from a container's environment. + + ``agentomatic deploy`` writes ``AUTH__JWKS_URL`` / ``AUTH__ISSUER`` / + ``AUTH__AUDIENCE`` into the generated ``.env`` and the docs said JWKS was + configurable "via stack" — but only the in-process ``jwt_config=`` kwarg + ever reached the middleware. A deployed container running the scaffolded + ``main.py`` had no way to switch signature verification on, and + ``require_auth_globally`` refused to boot without an API key. + """ + + def _platform(self, tmp_path, **kwargs): + from agentomatic import AgentPlatform + + return AgentPlatform(agents_dir=tmp_path / "agents", **kwargs) + + def test_env_vars_produce_a_verifying_config(self, tmp_path, monkeypatch): + monkeypatch.setenv("AUTH__JWKS_URL", "https://idp.test/jwks.json") + monkeypatch.setenv("AUTH__ISSUER", "https://idp.test/") + monkeypatch.setenv("AUTH__AUDIENCE", "agentomatic") + + cfg = self._platform(tmp_path)._resolve_jwt_config() + + assert cfg is not None + assert cfg.jwks_url == "https://idp.test/jwks.json" + assert cfg.issuer == "https://idp.test/" + assert cfg.audience == "agentomatic" + + def test_nothing_configured_returns_none(self, tmp_path, monkeypatch): + for var in ("AUTH__JWKS_URL", "AUTH__ISSUER", "AUTH__AUDIENCE"): + monkeypatch.delenv(var, raising=False) + + assert self._platform(tmp_path)._resolve_jwt_config() is None + + def test_stack_supplies_the_jwks_url(self, tmp_path, monkeypatch): + from types import SimpleNamespace + + for var in ("AUTH__JWKS_URL", "AUTH__ISSUER", "AUTH__AUDIENCE"): + monkeypatch.delenv(var, raising=False) + + platform = self._platform(tmp_path) + platform._stack_manager = SimpleNamespace( + _active_stack=SimpleNamespace( + auth=SimpleNamespace( + jwks_url="https://stack.test/jwks.json", + issuer="https://stack.test/", + audience="from-stack", + ) + ) + ) + + cfg = platform._resolve_jwt_config() + assert cfg is not None + assert cfg.jwks_url == "https://stack.test/jwks.json" + assert cfg.audience == "from-stack" + + def test_environment_wins_over_the_stack(self, tmp_path, monkeypatch): + from types import SimpleNamespace + + monkeypatch.setenv("AUTH__JWKS_URL", "https://env.test/jwks.json") + platform = self._platform(tmp_path) + platform._stack_manager = SimpleNamespace( + _active_stack=SimpleNamespace( + auth=SimpleNamespace( + jwks_url="https://stack.test/jwks.json", issuer="", audience="" + ) + ) + ) + + assert platform._resolve_jwt_config().jwks_url == "https://env.test/jwks.json" + + def test_unexpanded_stack_placeholder_is_not_treated_as_a_url(self, tmp_path, monkeypatch): + """``${JWKS_URL}`` with nothing in the env must not become the URL.""" + from types import SimpleNamespace + + monkeypatch.delenv("AUTH__JWKS_URL", raising=False) + monkeypatch.delenv("JWKS_URL", raising=False) + + platform = self._platform(tmp_path) + platform._stack_manager = SimpleNamespace( + _active_stack=SimpleNamespace( + auth=SimpleNamespace(jwks_url="${JWKS_URL}", issuer="", audience="") + ) + ) + + assert platform._resolve_jwt_config() is None + + def test_global_auth_lock_boots_on_a_jwks_url_alone(self, tmp_path, monkeypatch): + """No API key needed — remedy (a) from the lock's own error message.""" + from fastapi.testclient import TestClient + + monkeypatch.setenv("AUTH__JWKS_URL", "https://idp.test/jwks.json") + + platform = self._platform( + tmp_path, + plugins_dir=tmp_path / "plugins", + endpoints_dir=tmp_path / "endpoints", + enable_jwt_auth=True, + require_auth_globally=True, + ) + app = platform.build() + + with TestClient(app, raise_server_exceptions=False) as client: + # No token — rejected, not a 500 from a middleware that refused to + # construct. + assert client.get("/api/v1/agents").status_code == 401 + + +class TestRateLimitClientKey: + """What the ``trust_proxy_headers`` flag does — and does not — cover.""" + + def _middleware(self, *, trust: bool): + from agentomatic.middleware.rate_limit import RateLimitMiddleware + + return RateLimitMiddleware(app=None, trust_proxy_headers=trust) + + def _request(self, *, peer: str, forwarded: str | None): + from types import SimpleNamespace + + headers = {"X-Forwarded-For": forwarded} if forwarded else {} + return SimpleNamespace(headers=headers, client=SimpleNamespace(host=peer)) + + def test_forwarded_header_is_ignored_by_default(self): + """Otherwise any caller rotates the header and never gets limited.""" + mw = self._middleware(trust=False) + + assert mw._client_key(self._request(peer="10.0.0.1", forwarded="9.9.9.9")) == "10.0.0.1" + + def test_forwarded_header_is_used_when_a_proxy_is_declared(self): + mw = self._middleware(trust=True) + + key = mw._client_key(self._request(peer="10.0.0.1", forwarded="9.9.9.9, 10.0.0.1")) + assert key == "9.9.9.9" + + def test_key_falls_back_to_unknown_without_a_peer(self): + from types import SimpleNamespace + + mw = self._middleware(trust=False) + request = SimpleNamespace(headers={}, client=None) + + assert mw._client_key(request) == "unknown" + + def test_a_rewritten_peer_address_is_still_taken_at_face_value(self): + """Uvicorn's own --proxy-headers rewrites ``request.client`` upstream. + + By the time this middleware runs the original peer is gone, so the + flag cannot undo it. This documents the boundary: ``--forwarded-allow-ips`` + (uvicorn) is what decides whether that rewrite happens at all. + """ + mw = self._middleware(trust=False) + + # What uvicorn hands us after rewriting from X-Forwarded-For. + assert mw._client_key(self._request(peer="9.9.9.9", forwarded="9.9.9.9")) == "9.9.9.9" diff --git a/tests/test_stacks.py b/tests/test_stacks.py index a887131..c97b6c6 100644 --- a/tests/test_stacks.py +++ b/tests/test_stacks.py @@ -172,9 +172,7 @@ def test_load_from_yaml(self, tmp_path: Path) -> None: assert stack.llm["default"].provider == "ollama" assert stack.llm["default"].temperature == 0.2 - def test_load_missing_file_falls_back_to_builtin( - self, tmp_path: Path - ) -> None: + def test_load_missing_file_falls_back_to_builtin(self, tmp_path: Path) -> None: mgr = StackManager(stacks_dir=tmp_path) stack = mgr.load("local") assert stack.name == "local" diff --git a/tests/test_studio.py b/tests/test_studio.py index dedd36c..b10ebc1 100644 --- a/tests/test_studio.py +++ b/tests/test_studio.py @@ -1133,3 +1133,56 @@ async def test_stream_execution_fallback(self): history = await adapter.get_history("t1") assert len(history) == 1 assert history[0].metadata["framework"] == "langchain" + + +@pytest.mark.asyncio +async def test_studio_run_stream_emits_one_lifecycle_pair() -> None: + """The run tracker brackets each run with its own run_start/run_complete. + + Adapters may emit their own pair too (``AgentGraph.astream_studio_events`` + does), and forwarding those verbatim sent the client two runs' worth of + lifecycle events — which made the Studio UI render every agent reply twice. + Caught by driving the real UI in a browser, not by any unit test. + """ + import json as json_mod + + from agentomatic.studio.models import StudioRunEvent + from agentomatic.studio.run_tracker import RunTracker + + class _Adapter: + capabilities = ["streaming"] + + async def stream_execution(self, state, config, breakpoints, checkpoint_id): + # An adapter that emits its own lifecycle events, as the real + # graph runtime does. + for event in ( + StudioRunEvent(event="run_start", run_id="inner", timestamp="t", data={}), + StudioRunEvent(event="node_start", run_id="inner", timestamp="t", data={}), + StudioRunEvent( + event="node_end", run_id="inner", timestamp="t", data={"response": "ok"} + ), + StudioRunEvent(event="run_complete", run_id="inner", timestamp="t", data={}), + ): + yield event + + tracker = RunTracker() + run = tracker.create_run(agent_name="a1", thread_id="t1", request_data={"query": "hi"}) + + events: list[dict] = [] + async for frame in tracker.execute_with_adapter( + adapter=_Adapter(), + state={"query": "hi"}, + run_id=run.id, + thread_id="t1", + ): + for line in frame.splitlines(): + if line.startswith("data: "): + payload = line[len("data: ") :].strip() + if payload and payload.startswith("{"): + events.append(json_mod.loads(payload)) + + kinds = [e["event"] for e in events] + assert kinds.count("run_start") == 1, kinds + assert kinds.count("run_complete") == 1, kinds + # The adapter's real work still flows through. + assert "node_start" in kinds and "node_end" in kinds diff --git a/tests/test_studio_bundle_alignment.py b/tests/test_studio_bundle_alignment.py new file mode 100644 index 0000000..b707c09 --- /dev/null +++ b/tests/test_studio_bundle_alignment.py @@ -0,0 +1,297 @@ +# pyright: reportMissingParameterType=none +# pyright: reportCallIssue=none +# pyright: reportArgumentType=none +# pyright: reportAttributeAccessIssue=none +"""Backend ↔ Studio-frontend API alignment. + +The Studio UI ships as a pre-built React bundle synced from a separate repo, +so nothing in this repository's own test suite would otherwise notice if a +backend route were renamed, moved, or removed out from under it — the UI would +simply start 404-ing at runtime. + +This module parses the API paths the shipped bundle actually calls out of the +JavaScript, then asserts every one of them resolves to a real route on a fully +featured platform. It is a drift alarm for exactly the failure mode that +"the checked-in bundle is stale relative to the backend" produces. +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any + +import pytest +from fastapi.testclient import TestClient + +from agentomatic import AgentManifest, AgentPlatform + +_BUNDLE_DIR = Path(__file__).resolve().parents[1] / ("src/agentomatic/studio/static/static/js") + +# Paths the bundle builds dynamically in ways the static extractor cannot see +# (or that are intentionally probed with a fallback). Keep this list tiny and +# justified — every entry is a hole in the drift alarm. +_EXTRACTOR_BLIND_SPOTS: frozenset[str] = frozenset() + + +# ===================================================================== +# Bundle parsing +# ===================================================================== + + +def _find_bundle() -> Path | None: + """Return the shipped Studio JS bundle, if the UI assets are present.""" + if not _BUNDLE_DIR.is_dir(): + return None + bundles = sorted(_BUNDLE_DIR.glob("main.*.js")) + return bundles[0] if bundles else None + + +def _parse_concat_chain(text: str, quote_index: int) -> tuple[str, int]: + """Parse a JS string-concat chain into a path template. + + ``"/studio/agents/".concat(e,"/graph")`` → ``/studio/agents/*/graph`` + + Args: + text: Full bundle source. + quote_index: Index of the opening double quote of the chain. + + Returns: + ``(path_template, index_just_past_the_chain)``. + """ + end_quote = text.index('"', quote_index + 1) + parts: list[str] = [text[quote_index + 1 : end_quote]] + cursor = end_quote + 1 + + while text.startswith(".concat(", cursor): + cursor += len(".concat(") + depth, arg_start = 1, cursor + while depth: + char = text[cursor] + if char in "([": + depth += 1 + elif char in ")]": + depth -= 1 + cursor += 1 + inner = text[arg_start : cursor - 1] + + # Split the concat() arguments on top-level commas only. + args: list[str] = [] + depth, current = 0, "" + for char in inner: + if char in "([": + depth += 1 + elif char in ")]": + depth -= 1 + if char == "," and depth == 0: + args.append(current) + current = "" + else: + current += char + args.append(current) + + for arg in args: + arg = arg.strip() + if len(arg) >= 2 and arg.startswith('"') and arg.endswith('"'): + parts.append(arg[1:-1]) + else: + # A runtime expression — a path parameter. + parts.append("*") + + return "".join(parts), cursor + + +def extract_frontend_api_paths(bundle_source: str) -> set[str]: + """Extract the ``/api/v1`` and ``/studio`` paths the bundle requests.""" + paths: set[str] = set() + for match in re.finditer(r'(?:request|fetch)\(\s*"(?=/(?:api/v1|studio)/)', bundle_source): + quote_index = bundle_source.index('"', match.end() - 1) + try: + template, _ = _parse_concat_chain(bundle_source, quote_index) + except (ValueError, IndexError): # pragma: no cover - defensive + continue + # Drop any query string; routing only cares about the path. + paths.add(template.split("?", 1)[0].rstrip("/")) + return {p for p in paths if p} + + +# ===================================================================== +# Backend route collection +# ===================================================================== + + +async def _echo(state: dict[str, Any]) -> dict[str, Any]: + return {"response": "ok", "agent_type": "echo"} + + +_DEMO_PLUGIN_SOURCE = ''' +"""Minimal plugin so per-plugin routes (predict/model_card) are mounted.""" +from __future__ import annotations + +from pydantic import BaseModel + +from agentomatic.plugins import BaseMLPlugin + + +class DemoInput(BaseModel): + text: str = "" + + +class DemoOutput(BaseModel): + label: str = "" + + +class DemoPlugin(BaseMLPlugin[DemoInput, DemoOutput]): + plugin_name = "demo" + plugin_description = "Alignment-test plugin" + + async def predict(self, inputs: DemoInput) -> DemoOutput: + return DemoOutput(label="ok") +''' + + +@pytest.fixture(scope="module") +def backend_route_templates(tmp_path_factory) -> set[str]: + """Every route path a fully featured platform exposes. + + A plugin is discovered from disk so the per-plugin routes the Studio UI + calls (``/predict``, ``/model_card``) are actually mounted — without one, + the plugin section of the UI has nothing to align against. + """ + import importlib + import sys + + tmp_path = tmp_path_factory.mktemp("alignment") + plugins_dir = tmp_path / "plugins" + plugins_dir.mkdir(parents=True, exist_ok=True) + # The platform discovers plugins under the package prefix + # ``plugins_dir.name`` (i.e. ``plugins.demo_plugin``), so this must be a + # real package with its PARENT on sys.path. + (plugins_dir / "__init__.py").write_text("", encoding="utf-8") + (plugins_dir / "demo_plugin.py").write_text(_DEMO_PLUGIN_SOURCE, encoding="utf-8") + + # The repository has its own top-level ``plugins`` package which would + # otherwise shadow this one, so drop any cached import of it and put the + # temp parent first on sys.path. invalidate_caches() is required because + # these files were created after interpreter start. + saved_modules = { + name: sys.modules.pop(name) + for name in list(sys.modules) + if name == "plugins" or name.startswith("plugins.") + } + sys.path.insert(0, str(tmp_path)) + importlib.invalidate_caches() + try: + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=plugins_dir, + endpoints_dir=tmp_path / "endpoints", + title="Alignment", + enable_studio=True, + enable_control_plane=True, + control_token="t", + ) + platform.register_agent( + manifest=AgentManifest(name="echo_agent", slug="echo", description="Echo"), + node_fn=_echo, + ) + app = platform.build() + with TestClient(app): + routes = {r.path for r in app.routes if hasattr(r, "path")} + # Fail loudly if the demo plugin was not discovered — otherwise the + # per-plugin alignment assertions would silently have nothing to check. + assert any("/plugins/demo" in r for r in routes), ( + "demo plugin was not discovered, so per-plugin routes are missing; " + f"the alignment check would be vacuous. plugin routes seen: " + f"{sorted(r for r in routes if '/plugins' in r)}" + ) + return routes + finally: + sys.path.remove(str(tmp_path)) + for name in [n for n in sys.modules if n == "plugins" or n.startswith("plugins.")]: + del sys.modules[name] + sys.modules.update(saved_modules) + importlib.invalidate_caches() + + +def _segments(path: str) -> list[str]: + return [s for s in path.split("/") if s] + + +def _matches(frontend_path: str, backend_path: str) -> bool: + """Whether a frontend path template matches a backend route template. + + ``*`` (frontend runtime expression) and ``{param}`` (FastAPI path param) + both match any single segment. + """ + fe, be = _segments(frontend_path), _segments(backend_path) + if len(fe) != len(be): + return False + for fe_seg, be_seg in zip(fe, be, strict=True): + if fe_seg == "*" or (be_seg.startswith("{") and be_seg.endswith("}")): + continue + if fe_seg != be_seg: + return False + return True + + +# ===================================================================== +# Tests +# ===================================================================== + + +def test_studio_bundle_is_present() -> None: + """The packaged Studio UI assets must ship with the wheel.""" + assert _find_bundle() is not None, ( + f"No Studio JS bundle found under {_BUNDLE_DIR}. The Studio UI is " + "expected to be synced into the package." + ) + + +def test_extractor_finds_a_meaningful_number_of_paths() -> None: + """Guard the parser itself — a silent regex break would void this suite.""" + bundle = _find_bundle() + assert bundle is not None + paths = extract_frontend_api_paths(bundle.read_text(encoding="utf-8", errors="replace")) + assert len(paths) >= 25, f"extractor found only {len(paths)} paths — parser likely broke" + # Spot-check a few well-known calls the UI certainly makes. + assert "/studio/agents" in paths + assert "/api/v1/control/agents" in paths + + +def test_every_frontend_api_path_exists_on_the_backend(backend_route_templates) -> None: + """Every endpoint the shipped Studio UI calls must exist on the backend. + + A failure here means the checked-in UI bundle and the Python backend have + drifted: the UI will 404 at runtime against this version of the platform. + """ + bundle = _find_bundle() + assert bundle is not None + frontend_paths = extract_frontend_api_paths( + bundle.read_text(encoding="utf-8", errors="replace") + ) + + missing = sorted( + fe + for fe in frontend_paths + if fe not in _EXTRACTOR_BLIND_SPOTS + and not any(_matches(fe, be) for be in backend_route_templates) + ) + assert not missing, ( + "The Studio UI calls endpoints that do not exist on the backend " + f"(bundle/backend drift): {missing}" + ) + + +def test_studio_debug_api_paths_are_all_under_the_studio_prefix( + backend_route_templates, +) -> None: + """Sanity: the Studio calls we protect with auth really are /studio/*.""" + bundle = _find_bundle() + assert bundle is not None + paths = extract_frontend_api_paths(bundle.read_text(encoding="utf-8", errors="replace")) + studio_paths = {p for p in paths if p.startswith("/studio")} + assert studio_paths, "expected the UI to call the Studio debug API" + # None of them are the public UI shell — those are asset requests, not + # API calls, so the debug API is entirely inside the authenticated set. + assert all(not p.startswith("/studio/ui") for p in studio_paths) diff --git a/tests/test_task_progress.py b/tests/test_task_progress.py index 4335bf6..e0b6b6d 100644 --- a/tests/test_task_progress.py +++ b/tests/test_task_progress.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio from typing import Any import pytest @@ -117,3 +118,36 @@ async def report_fn(**kwargs: Any) -> None: reset_task_context(token) assert reports assert reports[0]["stage"] == "sync-stage" + + +@pytest.mark.asyncio +async def test_report_stage_sync_keeps_strong_task_reference() -> None: + """The scheduled task must be held strongly until it completes. + + asyncio only keeps a *weak* reference to a task created via + ``loop.create_task`` with no reference kept elsewhere — without an + explicit strong reference, the task can be garbage-collected mid-flight, + silently dropping the progress report. + """ + import gc + + from agentomatic.tasks.progress import _background_tasks, bind_task_context, reset_task_context + + async def slow_report_fn(**kwargs: Any) -> None: + await asyncio.sleep(0.05) + + ctx = TaskContext(task_id="t3", report_fn=slow_report_fn, is_cancelled=lambda: False) + token = bind_task_context(ctx) + try: + report_stage_sync("gc-stage") + # Force a collection pass immediately — before our fix this could + # collect the fire-and-forget task before it ever ran. + gc.collect() + assert len(_background_tasks) == 1 + pending = next(iter(_background_tasks)) + assert not pending.done() + await pending + finally: + reset_task_context(token) + # The done-callback must remove it from the tracking set afterwards. + assert len(_background_tasks) == 0 diff --git a/tests/test_tasks.py b/tests/test_tasks.py index 249ff51..e397a16 100644 --- a/tests/test_tasks.py +++ b/tests/test_tasks.py @@ -131,7 +131,12 @@ async def boom(target, payload, ctx): mgr.register_dispatcher(TargetType.AGENT, boom) rec = await mgr.submit_and_wait(TargetType.AGENT, "x", input={}) assert rec.status == TaskStatus.FAILED - assert "kaboom" in (rec.error or "") + # The record is served verbatim by GET /tasks/{id} and the task list, + # so the raw exception text (which routinely carries DSNs) must not be + # stored on it — only the type and a correlation id for the full log. + assert "kaboom" not in (rec.error or "") + assert "RuntimeError" in (rec.error or "") + assert "error_id=" in (rec.error or "") async def test_batch_progress(self): async def run(target, payload, ctx): @@ -162,6 +167,33 @@ async def slow(target, payload, ctx): assert refreshed is not None assert refreshed.status == TaskStatus.CANCELLED + async def test_shutdown_waits_for_cancelled_task_finalize(self): + """shutdown() must wait for a cancelled task's cleanup (which + persists its terminal status via the store) before closing the + store — otherwise the save can race the store's own disposal and + silently fail to persist the CANCELLED status. + """ + started = asyncio.Event() + + async def slow(target, payload, ctx): + started.set() + await asyncio.sleep(5) + return "done" + + mgr = TaskManager() + mgr.register_dispatcher(TargetType.AGENT, slow) + await mgr.submit(TargetType.AGENT, "x", input={}) + await asyncio.wait_for(started.wait(), timeout=2) + + # Before the fix, shutdown() closed the store immediately after + # requesting cancellation, without waiting for the task's own + # _finalize() to run and persist CANCELLED. + await mgr.shutdown() + + # The task must have actually finished (been cancelled) by the time + # shutdown() returns, not still be pending in the background. + assert not mgr._running + async def test_progress_reporting(self): async def run(target, payload, ctx: TaskContext): await ctx.report(current=1, total=2, message="half") diff --git a/tests/test_template_scaffold_quality.py b/tests/test_template_scaffold_quality.py new file mode 100644 index 0000000..e741137 --- /dev/null +++ b/tests/test_template_scaffold_quality.py @@ -0,0 +1,298 @@ +# pyright: reportMissingParameterType=none +# pyright: reportCallIssue=none +# pyright: reportArgumentType=none +# pyright: reportAttributeAccessIssue=none +"""Quality gate for the code `agentomatic init` scaffolds. + +Generated projects are the first thing a user sees, so a template that emits +code which does not compile — or that trips the linter the project itself +recommends — is a release defect. These tests render every template and check +the output the same way a user's own CI would. +""" + +from __future__ import annotations + +import ast +import subprocess +import sys +from pathlib import Path + +import pytest +from conftest import install_plugin_package + +from agentomatic.cli.templates import TEMPLATES, get_template_files + +# Templates render Python plus supporting files; only .py files are compiled. +_PY_SUFFIX = ".py" + +# No allowances: every template must render lint-clean as-is. The scripts that +# genuinely need a ``sys.path`` bootstrap before their project imports carry a +# file-level E402 suppression of their own, so they satisfy the linter without +# this gate having to look the other way. +_ALLOWED_RUFF_CODES: set[str] = set() + + +def _rendered(template: str) -> dict[str, str]: + return get_template_files(template, "sample_agent") + + +@pytest.mark.parametrize("template", sorted(TEMPLATES)) +def test_every_template_emits_syntactically_valid_python(template: str) -> None: + """Every generated .py file must parse.""" + files = _rendered(template) + py_files = {p: c for p, c in files.items() if p.endswith(_PY_SUFFIX)} + assert py_files, f"template {template!r} generated no Python files" + + for path, content in py_files.items(): + try: + ast.parse(content, filename=path) + except SyntaxError as exc: # pragma: no cover - failure path + pytest.fail(f"{template}/{path} is not valid Python: {exc}") + + +@pytest.mark.parametrize("template", sorted(TEMPLATES)) +def test_every_template_passes_ruff(template: str, tmp_path: Path) -> None: + """Generated code must be clean under the linter (bar allowed codes). + + A scaffold that ships lint-dirty code makes a new project fail its own + first CI run. + """ + files = _rendered(template) + for rel_path, content in files.items(): + target = tmp_path / rel_path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(content, encoding="utf-8") + + result = subprocess.run( # noqa: S603 + [ + sys.executable, + "-m", + "ruff", + "check", + "--isolated", + # Deliberately ruff's *default* 88-char line length, not the + # project's 99: a scaffolded project starts with no ruff config, + # so the defaults are what its first CI run will enforce. + "--select", + "E,F,W,I", + "--output-format", + "concise", + ".", + ], + capture_output=True, + text=True, + check=False, + # Run from inside the scaffolded dir, exactly as a user's CI would. + # From the repo root, ruff's isort resolves ``src/agentomatic`` and + # would misclassify ``agentomatic`` as a first-party import. + cwd=tmp_path, + ) + if result.returncode == 0: + return + + offending = [ + line + for line in result.stdout.splitlines() + if line.strip() + # Drop ruff's trailing summary lines ("Found N errors.", "[*] N fixable…") + and ":" in line + and not any(f" {code} " in line for code in _ALLOWED_RUFF_CODES) + ] + assert not offending, f"template {template!r} generates lint-dirty code:\n" + "\n".join( + offending + ) + + +def test_langchain_template_is_reachable_from_the_cli() -> None: + """Regression: ``langchain`` shipped in the registry but was missing from + the CLI's ``--template`` choices, so it could not actually be scaffolded. + """ + from click.types import Choice + + from agentomatic.cli.commands import init as init_cmd + + template_opt = next(p for p in init_cmd.params if p.name == "template") + assert isinstance(template_opt.type, Choice) + choices = set(template_opt.type.choices) + + assert "langchain" in choices + # The choice list is derived from the registry, so it can never drift again. + assert choices == set(TEMPLATES), ( + "CLI --template choices have drifted from the TEMPLATES registry: " + f"missing={set(TEMPLATES) - choices} extra={choices - set(TEMPLATES)}" + ) + + +def test_langchain_template_demonstrates_the_advertised_abstractions() -> None: + """The ``langchain`` template's description promises specific LangChain + abstractions — the generated code must actually use them. + """ + agent_py = _rendered("langchain")["agent.py"] + for expected in ( + "ChatPromptTemplate", + "MessagesPlaceholder", + "make_config", # builds the RunnableConfig + "self.prompt_template | self.llm", # a real LCEL chain + ): + assert expected in agent_py, f"langchain template does not use {expected!r}" + + +# ===================================================================== +# Scaffolded ML plugin must actually be usable +# ===================================================================== + + +def test_plugin_template_sets_its_own_name() -> None: + """Without ``plugin_name`` the scaffold inherits BaseMLPlugin's + ``default_plugin``, so it mounts at /api/v1/plugins/default_plugin/* and a + second scaffolded plugin silently collides with the first. + """ + plugin_py = get_template_files("plugin", "sentiment")["plugin.py"] + assert 'plugin_name = "sentiment"' in plugin_py + + +def test_plugin_template_marks_itself_loaded() -> None: + """Overriding ``load_model`` without calling super() leaves ``_is_loaded`` + False: /predict answers 503 and /health reports the platform "degraded", + while startup logs claim the plugin loaded successfully. + """ + plugin_py = get_template_files("plugin", "sentiment")["plugin.py"] + assert "await super().load_model()" in plugin_py + + +def test_scaffolded_plugin_serves_predictions_and_reports_healthy(tmp_path) -> None: + """End-to-end: render the plugin template, mount it, and call /predict.""" + from fastapi.testclient import TestClient + + from agentomatic import AgentPlatform + + plugins_dir = tmp_path / "plugins" + files = get_template_files("plugin", "sentiment") + importable = install_plugin_package(plugins_dir, "sentiment", files["plugin.py"]) + for rel, content in files.items(): + if rel == "plugin.py": + continue + path = plugins_dir / "sentiment" / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + with importable: + platform = AgentPlatform( + agents_dir=tmp_path / "agents", + plugins_dir=plugins_dir, + endpoints_dir=tmp_path / "endpoints", + ) + with TestClient(platform.build()) as client: + listed = client.get("/api/v1/plugins").json() + entries = listed if isinstance(listed, list) else listed.get("plugins", []) + assert [e.get("name") for e in entries] == ["sentiment"] + + assert client.get("/health").json()["status"] == "healthy" + + response = client.post("/api/v1/plugins/sentiment/predict", json={"text": "hi"}) + assert response.status_code == 200, response.text + + +def test_full_template_response_schema_matches_its_agent_output() -> None: + """`schemas.py` required ``answer`` while the agent returned ``response``, + so every invoke logged an output-validation warning. + """ + files = get_template_files("full", "sample_agent") + schemas_py, agent_py = files["schemas.py"], files["agent.py"] + + assert "response: str" in schemas_py + assert "answer: str" not in schemas_py + # The agent really does emit "response". + assert '"response": text' in agent_py + + +def test_scaffolded_main_explains_a_version_skew_instead_of_a_bare_typeerror( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """An older installed agentomatic must fail with actionable guidance. + + ``main.py`` is generated against the version that scaffolded it. If the + image pins an older release (a stale ``requirements.txt`` line or a + Dockerfile pin that predates a new option), ``AgentPlatform`` raises a bare + ``unexpected keyword argument`` at import time and the container dies with + no hint about the cause. + """ + import agentomatic + from agentomatic.cli.project import _main_py + + class _OldPlatform: + @staticmethod + def from_folder(*args: object, **kwargs: object) -> object: + raise TypeError( + "AgentPlatform.__init__() got an unexpected keyword argument " + "'rate_limit_trust_proxy_headers'" + ) + + monkeypatch.setattr(agentomatic, "AgentPlatform", _OldPlatform) + monkeypatch.chdir(tmp_path) + + with pytest.raises(RuntimeError) as excinfo: + exec(compile(_main_py("demo"), "main.py", "exec"), {"__name__": "main"}) + + message = str(excinfo.value) + assert "rate_limit_trust_proxy_headers" in message + assert "older than the one this project was generated with" in message + assert "requirements.txt" in message + # The original TypeError stays chained so the traceback is not lost. + assert isinstance(excinfo.value.__cause__, TypeError) + + +def test_scaffolded_main_does_not_swallow_unrelated_type_errors( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Only unknown-keyword failures are reframed; real bugs propagate as-is.""" + import agentomatic + from agentomatic.cli.project import _main_py + + class _BrokenPlatform: + @staticmethod + def from_folder(*args: object, **kwargs: object) -> object: + raise TypeError("unhashable type: 'dict'") + + monkeypatch.setattr(agentomatic, "AgentPlatform", _BrokenPlatform) + monkeypatch.chdir(tmp_path) + + with pytest.raises(TypeError, match="unhashable"): + exec(compile(_main_py("demo"), "main.py", "exec"), {"__name__": "main"}) + + +def test_deepagent_template_names_its_missing_dependency() -> None: + """The deepagent template imports a package agentomatic does not install. + + Without it the agent scaffolds, registers, and reports healthy, but every + invocation returns a sanitised 500 whose only clue is the exception type + ``ModuleNotFoundError`` — the caller cannot tell what to install. + """ + files = get_template_files("deepagent", "mydeep") + agent_py = files["agent.py"] + + assert "pip install deepagents" in agent_py + assert "except ImportError" in agent_py + # The bare import must not remain outside the guard. + guarded = agent_py.split("try:", 1)[1] + assert "from deepagents import create_deep_agent" in guarded + + +def test_deepagent_scaffold_tells_the_user_to_install_it(tmp_path: Path) -> None: + """``agentomatic init --template deepagent`` must surface the dependency.""" + result = subprocess.run( + [ + sys.executable, + "-c", + "from agentomatic.cli.commands import cli; cli()", + "init", + "mydeep", + "--template", + "deepagent", + ], + cwd=tmp_path, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + assert "pip install deepagents" in result.stdout + result.stderr diff --git a/tests/test_train_pipeline_fixes.py b/tests/test_train_pipeline_fixes.py index d9cf002..71cdce4 100644 --- a/tests/test_train_pipeline_fixes.py +++ b/tests/test_train_pipeline_fixes.py @@ -71,7 +71,9 @@ def respond(self, state: _PromptState) -> _PromptState: return state def input_to_state(self, input_data: dict[str, Any]) -> _PromptState: - return _PromptState(request=input_data.get("current_query") or input_data.get("query") or "") + return _PromptState( + request=input_data.get("current_query") or input_data.get("query") or "" + ) def state_to_output(self, state: _PromptState) -> dict[str, Any]: return state.output