diff --git a/AGENTS.md b/AGENTS.md index eb67487c..dba3d8ca 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -104,7 +104,7 @@ and deserialized through a registry. | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns `(run_id, instance_id)`; `LocalRunner` is built in. | | `explorer/` | Optional local web UI: `agent-env up`, needs the `explorer` extra, binds loopback `:8234`. | -| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. | +| `cli/` | Click CLI with the groups `a2a-agent`, `artifact`, `bundle`, `config`, `env`, `eval`, `plugin`, `run`, `task`, `up`. | | `examples/` | The bundles agent-env ships (`hello`, and `echo-mcp`, which needs Docker), registered under `agent_env.bundles` in `pyproject.toml`. Tests run them, but they are examples for users, not fixtures. | ## Extension points diff --git a/README.md b/README.md index abb409ef..d010c268 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ uv tool install agentenv-framework agent-env run hello ``` -`agent-env run hello` runs the built-in hello task, which needs no Docker, model or configuration. `agent-env run echo-mcp` needs Docker: it builds an MCP server from a Dockerfile, deploys it through the gateway and checks it by calling its tool. `agent-env run` lists both. To try hello without installing anything, run it with uvx: +`agent-env run hello` runs the built-in hello task, which needs no Docker, model or configuration. `agent-env run echo-mcp` needs Docker: it builds an MCP server from a Dockerfile, deploys it through the gateway and checks it by calling its tool. `agent-env run` lists both. `agent-env bundle init --from-task ` writes a task a store holds as a bundle folder, with each reference it makes pinned at the version the store holds, so you can run and edit it locally. To try hello without installing anything, run it with uvx: ```bash uvx --from agentenv-framework agent-env run hello diff --git a/src/agent_env/bundle/_secrets.py b/src/agent_env/bundle/_secrets.py new file mode 100644 index 00000000..22f1590c --- /dev/null +++ b/src/agent_env/bundle/_secrets.py @@ -0,0 +1,60 @@ +"""Secrets a bundle import leaves out: env vars that name or hold one, and credentials anywhere else in a value.""" + +from __future__ import annotations + +import re +from typing import Any + +from agent_env.config.describe import is_secret_key + +# The fields that map environment variable names to values, where a credential is most often set. +_ENV_VAR_MAPS = ("env_vars", "default_env_vars") +# Each credential a value can hold, how to find it and what masks it: a URL's user info when it holds a password or is +# long enough to be a token rather than a user name such as git, then values in the shape their issuers give them. +_CREDENTIALS = ( + ("a URL's credentials", re.compile(r"(?P[a-zA-Z][\w+.-]*://)(?:[^/@\s]*:[^/@\s]*|[^/@\s:]{16,})@"), + r"\g***@"), + ("an API key", re.compile(r"\bsk[-_](?=[\w-]*\d)[\w-]{20,}"), "***"), + ("a GitHub token", re.compile(r"\b(?:gh[pousr]_[A-Za-z0-9]{36,}|github_pat_\w{22,})"), "***"), + ("an AWS access key", re.compile(r"\b(?:AKIA|ASIA)[0-9A-Z]{16}\b"), "***"), + ("a Slack token", re.compile(r"\bxox[abposr]-[A-Za-z0-9-]{10,}"), "***"), + ("a Google API key", re.compile(r"\bAIza[\w-]{35}"), "***"), + ("a Hugging Face token", re.compile(r"\bhf_[A-Za-z0-9]{30,}"), "***"), + ("a JWT", re.compile(r"\beyJ[\w-]{8,}\.eyJ[\w-]{8,}\.[\w-]{8,}"), "***"), + ("a private key", re.compile(r"-----BEGIN (?:[A-Z]+ )*PRIVATE KEY-----" + r"(?:.*?-----END (?:[A-Z]+ )*PRIVATE KEY-----|.*)", re.DOTALL), "***"), +) + + +def scrub(value: Any, path: str = "") -> list[tuple[str, str]]: + """Drop secret env vars from ``value`` and mask credentials elsewhere in it, in place; each one's path and why.""" + removed = [] + items = value.items() if isinstance(value, dict) else enumerate(value) if isinstance(value, list) else () + for key, item in list(items): + at = f"{path}.{key}" if path and isinstance(key, str) else f"{path}[{key}]" if path else str(key) + if isinstance(value, dict) and key in _ENV_VAR_MAPS and isinstance(item, dict): + for name, setting in list(item.items()): + secret = is_secret_key(name) + found = _masked(setting)[1] if isinstance(setting, str) and not secret else [] + if secret or found: + del item[name] + why = "its name looks like a secret" if secret else f"its value holds {found[0]}" + removed.append((f"{at}.{name}", f"dropped: {why}")) + elif isinstance(item, str): + masked, found = _masked(item) + if found: + value[key] = masked + removed.append((at, f"masked: {', '.join(found)}")) + else: + removed += scrub(item, at) + return removed + + +def _masked(text: str) -> tuple[str, list[str]]: + """``text`` with each credential replaced by ``***``, and the kinds it held.""" + found = [] + for what, pattern, mask in _CREDENTIALS: + text, count = pattern.subn(mask, text) + if count: + found.append(what) + return text, found diff --git a/src/agent_env/bundle/_toml.py b/src/agent_env/bundle/_toml.py new file mode 100644 index 00000000..d06a33a5 --- /dev/null +++ b/src/agent_env/bundle/_toml.py @@ -0,0 +1,54 @@ +"""A TOML writer for the tomls a bundle import writes, since ``tomllib`` only reads.""" + +from __future__ import annotations + +import math +import re +from typing import Any + +_BARE_KEY = re.compile(r"[A-Za-z0-9_-]+") +_ESCAPES = {'"': '\\"', "\\": "\\\\", "\b": "\\b", "\t": "\\t", "\n": "\\n", "\f": "\\f", "\r": "\\r"} + + +def dumps(document: dict[str, Any]) -> str: + """``document`` as TOML, each dict an inline table; ValueError for a value TOML can't hold.""" + return "".join(f"{_key(key)} = {_value(value)}\n" for key, value in document.items()) + + +def _key(key: str) -> str: + return key if _BARE_KEY.fullmatch(key) else _string(key) + + +def _value(value: Any, *, top: bool = True) -> str: + """``value`` as TOML; a list of several items spans lines only at the top, as an inline table can't.""" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, int): + return str(value) + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError(f"{value!r} is not a TOML number") + return repr(value) + if isinstance(value, str): + return _string(value) + if isinstance(value, (list, tuple)): + items = [_value(item, top=False) for item in value] + if top and len(items) > 1: + return "[\n" + "".join(f" {item},\n" for item in items) + "]" + return f"[{', '.join(items)}]" + if isinstance(value, dict): + pairs = ", ".join(f"{_key(k)} = {_value(v, top=False)}" for k, v in value.items()) + return f"{{ {pairs} }}" if pairs else "{}" + raise ValueError(f"TOML can't hold {value!r}") + + +def _string(text: str) -> str: + out = [] + for char in text: + if char in _ESCAPES: + out.append(_ESCAPES[char]) + elif ord(char) < 0x20 or ord(char) == 0x7F: + out.append(f"\\u{ord(char):04x}") + else: + out.append(char) + return '"' + "".join(out) + '"' diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py new file mode 100644 index 00000000..48ec9686 --- /dev/null +++ b/src/agent_env/bundle/importer.py @@ -0,0 +1,350 @@ +"""Write a stored task, or an eval and its tasks, as a bundle folder that runs what the store held, writing no store.""" + +from __future__ import annotations + +import contextlib +import json +import os +import shutil +import unicodedata +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from importlib.metadata import version as package_version +from pathlib import Path + +from agent_env.bundle import _toml +from agent_env.bundle._fs import fold, os_reason, show +from agent_env.bundle._secrets import scrub +from agent_env.bundle.parse import BundleError, parse_bundle, root_problems, segment_problem +from agent_env.bundle.plan import STORE_GETTERS, check_bundle +from agent_env.bundle.resolve import resolve_bundle +from agent_env.config import get_config +from agent_env.entity_refs import EntityKind, RefRole, RefSite, ref_sites +from agent_env.eval import Eval +from agent_env.store.base import NotFoundError +from agent_env.store.ids import LOCAL_PREFIX +from agent_env.task import Task +from agent_env.task_step.registry import get_task_step_registry +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep +from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep + +# How long a name taken from an id may be, in bytes; ids are capped far above it, and a longer name helps nobody. +_MAX_NAME_BYTES = 200 + + +@dataclass(frozen=True) +class Pinned: + """A store reference written with a version: the store's at import, or the task's own pin.""" + + where: str + kind: EntityKind + id: str + version: int + stored: bool # the task already pinned it + + +@dataclass(frozen=True) +class Floating: + """A store reference written without a version, so a run reads whatever version is current then.""" + + where: str + kind: EntityKind + id: str + why: str + + +@dataclass(frozen=True) +class Dropped: + """A value left out of the folder, or masked in it, because it looks like a secret.""" + + where: str + why: str + + +@dataclass(frozen=True) +class ImportedTask: + path: str # in the bundle: tasks/.json + id: str + version: int + steps: int + + +@dataclass(frozen=True) +class ImportedEval: + path: str # in the bundle: evals/.toml + id: str + version: int + + +@dataclass(frozen=True) +class BundleImport: + """What an import wrote, and how it treated each reference and secret.""" + + root: Path + files: tuple[str, ...] + tasks: tuple[ImportedTask, ...] + eval: ImportedEval | None + pinned: tuple[Pinned, ...] + floating: tuple[Floating, ...] + dropped: tuple[Dropped, ...] + undeclared: tuple[str, ...] # the steps whose type declares no references, so any they make float + duplicates: tuple[str, ...] # the eval's entries that named a task it already named + + def report(self) -> list[tuple[str, list[str]]]: + """How each reference and secret was written, as titled lists of lines; an empty section is left out.""" + return _sections([ + ("Pinned", [f"{pin.kind.value} {pin.id} v{pin.version}{' (as stored)' if pin.stored else ''}: " + f"{pin.where}" for pin in self.pinned]), + ("Floating", [f"{ref.kind.value} {ref.id}: {ref.where}: {ref.why}" for ref in self.floating] + + [f"{step}: its type declares no references, so any it makes float" for step in self.undeclared]), + ("Left out", [f"{drop.where}: {drop.why}" for drop in self.dropped] + + [f"{entry}, which the eval names more than once" for entry in self.duplicates]), + ]) + + +def import_task(task_id: str, dest: Path | str, *, version: int | None = None) -> BundleImport: + """Write the stored task ``task_id``, at ``version`` or its latest, as a bundle at ``dest``.""" + importer = _Importer(Path(dest)) + task = Task.get(task_id, version) + importer.add_task(task) + return importer.write(f"Imported from the store's task {task.id} v{task.version}.") + + +def import_eval(eval_id: str, dest: Path | str, *, version: int | None = None) -> BundleImport: + """Write the stored eval ``eval_id``, at ``version`` or its latest, and the tasks it names as a bundle.""" + importer = _Importer(Path(dest)) + stored = Eval.get(eval_id, version) + names, seen = [], set() + for entry in stored.tasks: + try: + task = Task.get(entry.task_id, entry.task_version) + except NotFoundError: + at = f" v{entry.task_version}" if entry.task_version is not None else "" + importer.problems.append(f"eval {stored.id} v{stored.version} names task {entry.task_id}{at}, which isn't " + "in the store") + continue + if (task.id, task.version) in seen: + importer.duplicates.append(f"task {task.id} v{task.version}") + continue + seen.add((task.id, task.version)) + names.append(importer.add_task(task)) + if names: + importer.add_eval(stored, names) + return importer.write(f"Imported from the store's eval {stored.id} v{stored.version} and the tasks it names.") + + +class _Importer: + def __init__(self, dest: Path): + self.dest = Path(os.path.abspath(dest.expanduser())) # so ``.`` has a name, and a sibling to stage in + self.files: dict[str, str] = {} + self.taken: dict[str, set[str]] = {"tasks": set(), "evals": set()} # folded names, per folder + self.tasks: list[ImportedTask] = [] + self.eval: ImportedEval | None = None + self.pinned: list[Pinned] = [] + self.floating: list[Floating] = [] + self.dropped: list[Dropped] = [] + self.undeclared: list[str] = [] + self.duplicates: list[str] = [] + self.problems: list[str] = [] + self.read: dict[tuple[EntityKind, str, int | None], int | str] = {} + _check_destination(self.dest) + + def add_task(self, task: Task) -> str: + """Write ``task`` under a name taken from its id, and return the name.""" + name = self._name("tasks", task.id) + path = f"tasks/{name}.json" + steps = task.to_dict()["steps"] + registry = get_task_step_registry() + found = [] + for step in steps: + cls = registry[step["type"]] + found.append((step, cls, [site for site in ref_sites(cls.entity_refs or (), step) + if isinstance(site.value, str) and site.value])) + sites = [site for _, _, step_sites in found for site in step_sites] + outputs = _output_names(sites) + deployed = {_nfc(site.value) for _, cls, step_sites in found if issubclass(cls, DeployEnvTaskStep) + for site in step_sites if site.ref.kind is EntityKind.ENV} + for site in sites: # each output, and each step that reads it, takes the output's name + if outputs.get(_key(site), site.value) != site.value: + site.rewrite(outputs[_key(site)]) + for step, cls, step_sites in found: + where = f"{path}: step {step.get('id')!r}" + self.dropped += [Dropped(f"{where}: {at}", why) for at, why in scrub(step)] + if cls.entity_refs is None: + self.undeclared.append(f"{where} ({cls.type})") + for site in step_sites: + if _key(site) not in outputs: # the run writes an output again, so it takes no pin + self._pin(site, f"{where}: {site.path}", deployed) + self._default_agents(step, where) + self.files[path] = json.dumps(steps, indent=2, ensure_ascii=False, allow_nan=False) + "\n" + self.tasks.append(ImportedTask(path, task.id, task.version, len(steps))) + return name + + def add_eval(self, stored: Eval, names: list[str]) -> None: + path = f"evals/{self._name('evals', stored.id)}.toml" + self.files[path] = _toml.dumps({"tasks": names}) + self.eval = ImportedEval(path, stored.id, stored.version) + + def write(self, description: str) -> BundleImport: + """Write the folder, with ``description`` as its README's first line, and say what was written.""" + if self.problems: + raise BundleError(self.problems) + files = sorted([*self.files, "README.md"]) + imported = BundleImport(root=self.dest, files=tuple(files), tasks=tuple(self.tasks), eval=self.eval, + pinned=tuple(self.pinned), floating=tuple(self.floating), + dropped=tuple(self.dropped), undeclared=tuple(self.undeclared), + duplicates=tuple(self.duplicates)) + self.files["README.md"] = _readme(description, imported) + staging = self.dest.parent / f".{self.dest.name}.import-{uuid.uuid4().hex[:8]}" + try: + try: + for path, text in self.files.items(): + (staging / path).parent.mkdir(parents=True, exist_ok=True) + (staging / path).write_text(text, encoding="utf-8") + except OSError as e: + raise BundleError([f"{show(self.dest.parent)}: can't write the import there: {os_reason(e)}"]) from None + check_bundle(resolve_bundle(parse_bundle(staging))) + self._move(staging) + finally: + shutil.rmtree(staging, ignore_errors=True) + return imported + + def _move(self, staging: Path) -> None: + """Move ``staging`` to the destination, putting an emptied one back if the move fails or is interrupted.""" + emptied = False + try: + if self.dest.is_dir(): + self.dest.rmdir() # empty, as _check_destination found it + emptied = True + os.replace(staging, self.dest) + except BaseException as e: + if emptied: + with contextlib.suppress(OSError): + self.dest.mkdir(exist_ok=True) + if not isinstance(e, OSError): + raise + _check_destination(self.dest) # something else wrote it meanwhile + raise BundleError([f"{show(self.dest)}: can't move the import into place: {os_reason(e)}"]) from None + + def _pin(self, site: RefSite, where: str, deployed: set[str]) -> None: + kind, entity_id = site.ref.kind, site.value + if site.version is not None or site.version_key is not None: + read = self._read(kind, entity_id, site.version) + if isinstance(read, str): + self.problems.append(f"{where}: {read}") + return + if site.version is None: + site.rewrite(entity_id, read) + self.pinned.append(Pinned(where, kind, entity_id, read, stored=site.version is not None)) + elif kind is not EntityKind.ENV: + self.floating.append(Floating(where, kind, entity_id, "the field takes no version, so a run reads the " + "latest")) + elif _nfc(entity_id) not in deployed: + self.floating.append(Floating(where, kind, entity_id, "the field takes no version and no step of the task " + "deploys it, so a run reads the latest")) + + def _read(self, kind: EntityKind, entity_id: str, version: int | None = None) -> int | str: + """The version of ``entity_id`` the store holds, at ``version`` or its latest, or why it can't be read.""" + if (kind, entity_id, version) not in self.read: + try: + self.read[kind, entity_id, version] = STORE_GETTERS[kind](entity_id, version).version + except NotFoundError: + at = f" v{version}" if version is not None else "" + self.read[kind, entity_id, version] = f"there is no {kind.value} {entity_id!r}{at} in the store" + except (ValueError, KeyError, TypeError) as e: + self.read[kind, entity_id, version] = (f"{kind.value} {entity_id!r} can't be read " + f"({type(e).__name__}: {e})") + return self.read[kind, entity_id, version] + + def _default_agents(self, step: dict, where: str) -> None: + """A step that names no agent deploys the configured default, which no field pins.""" + if step["type"] == DeployAgentTaskStep.type and step.get("a2a_agent_id") is None: + field, what = "a2a_agent_id", "names no agent" + elif (step["type"] == RubricsVerifierTaskStep.type and step.get("use_agent_judge") + and step.get("agent_name") is None and step.get("judge_a2a_agent_id") is None): + field, what = "judge_a2a_agent_id", "names no judge agent" + else: + return + self.floating.append(Floating(f"{where}: {field}", EntityKind.AGENT, get_config().get_default_a2a_agent_id(), + f"the step {what}, so a run deploys the configured default")) + + def _name(self, kind: str, entity_id: str) -> str: + """A name in the ``kind`` folder: the id's last segment made valid, numbered when taken, caselessly.""" + base = _valid_name(entity_id.rsplit("/", 1)[-1]) or kind.removesuffix("s") + name, n = base, 1 + while fold(name) in self.taken[kind]: + n += 1 + name = f"{base}-{n}" + self.taken[kind].add(fold(name)) + return name + + +def _output_names(sites: list[RefSite]) -> dict[tuple[EntityKind, str], str]: + """The name each output takes: as written, or an ``@local`` id's last segment, numbered when it is taken.""" + written = {_key(site) for site in sites if site.ref.role is RefRole.OUTPUT} + taken = {_key(site) for site in sites if _key(site) not in written or not site.value.startswith(LOCAL_PREFIX)} + names: dict[tuple[EntityKind, str], str] = {} + for site in sites: + if site.ref.role is not RefRole.OUTPUT or _key(site) in names: + continue + name = site.value + if name.startswith(LOCAL_PREFIX): + base = name = name.rsplit("/", 1)[-1] + n = 1 + while (site.ref.kind, _nfc(name)) in taken: + n += 1 + name = f"{base}-{n}" + taken.add((site.ref.kind, _nfc(name))) + names[_key(site)] = name + return names + + +def _readme(description: str, imported: BundleImport) -> str: + """The bundle's README: its description, where it came from, what it holds, then the import's report.""" + stamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC") + lines = [description, "", f"Imported with agent-env {package_version('agentenv-framework')} on {stamp}. Each store " + "reference is pinned at the version the store held then, unless it floats, so `agent-env run` runs this " + "folder as the store ran it.", ""] + held = _sections([("Tasks", [f"{task.path}: task {task.id} v{task.version}" for task in imported.tasks]), + ("Eval", [f"{imported.eval.path}: eval {imported.eval.id} v{imported.eval.version}"] + if imported.eval else [])]) + for title, items in [*held, *imported.report()]: + lines += [f"## {title}", "", *(f"- {item}" for item in items), ""] + return "\n".join(lines).rstrip("\n") + "\n" + + +def _sections(sections: list[tuple[str, list[str]]]) -> list[tuple[str, list[str]]]: + return [(title, lines) for title, lines in sections if lines] + + +def _check_destination(dest: Path) -> None: + """Refuse a destination that holds anything, isn't a folder, or would give the bundle no id.""" + if dest.is_symlink() or (dest.exists() and not dest.is_dir()): + raise BundleError([f"{show(dest)}: not a folder; import into a new or empty folder"]) + if dest.is_dir() and any(dest.iterdir()): + raise BundleError([f"{show(dest)}: not empty; import into a new or empty folder"]) + if not dest.parent.is_dir(): + raise BundleError([f"{show(dest.parent)}: no such folder"]) + problems = root_problems(dest) + if problems: + raise BundleError(problems) + + +def _valid_name(text: str) -> str: + """``text`` as one bundle name, each character an id can't hold made ``-``; empty with no letter or digit left.""" + text = _nfc(text) + name = "".join(char if char.isalnum() or unicodedata.category(char).startswith("M") or char in " ._-~@()+,&" + else "-" for char in text).strip(" ").lstrip(".") + while len(name.encode("utf-8")) > _MAX_NAME_BYTES: + name = name[:-1].rstrip(" ") + return name if any(char.isalnum() for char in name) and segment_problem(name) is None else "" + + +def _key(site: RefSite) -> tuple[EntityKind, str]: + return site.ref.kind, _nfc(site.value) + + +def _nfc(text: str) -> str: + return unicodedata.normalize("NFC", text) diff --git a/src/agent_env/bundle/parse.py b/src/agent_env/bundle/parse.py index 592cb110..71a1f654 100644 --- a/src/agent_env/bundle/parse.py +++ b/src/agent_env/bundle/parse.py @@ -152,25 +152,11 @@ def parse(self) -> Bundle: ) def _derive_id_root(self) -> str: - user_home = home() - home_relative = user_home is not None and self.root.is_relative_to(user_home) - if home_relative: - parts = ["~", *self.root.relative_to(user_home).parts] - above = user_home - else: - parts = list(self.root.parts[1:]) - above = Path(self.root.anchor) - if parts[0] == "~": - self._problem(above / "~", "a folder named '~' at the filesystem root would read as home; " - "move the bundle") - self.root_ok = False - for part in parts[1:] if home_relative else parts: - above = above / part - reason = _refused(part) - if reason: - self._problem(above, f"move the bundle, or rename the folder: {reason}") - self.root_ok = False - return LOCAL_PREFIX + "/".join(unicodedata.normalize("NFC", part) for part in parts) + id_root, problems = _id_root(self.root) + for path, message in problems: + self._problem(path, message) + self.root_ok = False + return id_root # Walking the tree @@ -381,7 +367,7 @@ def _entity_id(self, kind: BundleKind, name: str, path: Path, declared: Any, dec if declared is not None: self._check_declared(declared, declared_in or path, artifact=kind.store == "artifact") return declared - reason = _refused(name) + reason = segment_problem(name) if reason: self._skip(path, f"rename it: {reason}") # The root and the name are each validated, so only the length is left to check. @@ -538,6 +524,34 @@ def _rel(self, path) -> str: return show(text) +def root_problems(root: Path) -> list[str]: + """Why a bundle at ``root``, which needn't exist yet, would have no id; empty when it would have one.""" + parent = _canonical(Path(root).expanduser().parent) + return [f"{show(path)}: {message}" for path, message in _id_root(parent / Path(root).name)[1]] + + +def _id_root(root: Path) -> tuple[str, list[tuple[Path, str]]]: + """The ``@local`` id root for the canonical path ``root``, and each part of it that can't be an id segment.""" + problems = [] + user_home = home() + home_relative = user_home is not None and root.is_relative_to(user_home) + if home_relative: + parts = ["~", *root.relative_to(user_home).parts] + above = user_home + else: + parts = list(root.parts[1:]) + above = Path(root.anchor) + if parts[0] == "~": + problems.append((above / "~", "a folder named '~' at the filesystem root would read as home; " + "move the bundle")) + for part in parts[1:] if home_relative else parts: + above = above / part + reason = segment_problem(part) + if reason: + problems.append((above, f"move the bundle, or rename the folder: {reason}")) + return LOCAL_PREFIX + "/".join(unicodedata.normalize("NFC", part) for part in parts), problems + + def _canonical(root: Path) -> Path: try: resolved = root.expanduser().resolve(strict=True) @@ -552,7 +566,7 @@ def _canonical(root: Path) -> Path: return on_disk(resolved) -def _refused(segment: str) -> str | None: +def segment_problem(segment: str) -> str | None: """Why ``segment`` can't be one segment of an id, or None; the reason names the segment.""" try: segment.encode("utf-8") diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index fd8ef766..961bf28e 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -65,8 +65,8 @@ _Key = tuple[str, str] # (store, id) -_GETTERS = {EntityKind.ENV: Env.get, EntityKind.AGENT: A2AAgent.get, EntityKind.ARTIFACT: Artifact.get, - EntityKind.TASK: Task.get} +STORE_GETTERS = {EntityKind.ENV: Env.get, EntityKind.AGENT: A2AAgent.get, EntityKind.ARTIFACT: Artifact.get, + EntityKind.TASK: Task.get} # The artifact types written from their folders, and what each one's write reads there, which the ledger hashes. _FOLDER_WALKS = {FileArtifact: one_file, FileArtifactUniverse: entry_files, EnvironmentArtifact: environment_files, @@ -320,7 +320,7 @@ def _store_ref_problem(self, ref: Reference, read: dict) -> str | None: lookup = (ref.kind, ref.id, ref.version) if lookup not in read: try: - read[lookup] = _GETTERS[ref.kind](ref.id, ref.version) + read[lookup] = STORE_GETTERS[ref.kind](ref.id, ref.version) except NotFoundError: read[lookup] = None except (ValueError, KeyError, TypeError) as e: @@ -590,8 +590,8 @@ def _environment_name(self, ref: Reference, store_latest: dict[tuple[EntityKind, if store_latest is None: return None try: - entity = _GETTERS[ref.kind](ref.id, ref.version if ref.version is not None - else store_latest.get((ref.kind, ref.id))) + entity = STORE_GETTERS[ref.kind](ref.id, ref.version if ref.version is not None + else store_latest.get((ref.kind, ref.id))) except (NotFoundError, ValueError, KeyError, TypeError): return None # reported by the store ref check return getattr(entity, "environment_name", None) diff --git a/src/agent_env/bundle/resolve.py b/src/agent_env/bundle/resolve.py index 35ca5fe8..53899ece 100644 --- a/src/agent_env/bundle/resolve.py +++ b/src/agent_env/bundle/resolve.py @@ -9,8 +9,9 @@ from __future__ import annotations import copy +import functools import unicodedata -from collections.abc import Iterator +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -197,7 +198,7 @@ def _task(self, entry: BundleEntry) -> ResolvedEntry: elif self._derived_outputs(entry, step, cls): declared.append((index, step, cls)) outputs = self._outputs(entry, declared) - upstream = _upstream(steps, self._depends_on(entry, steps)) + edges = _step_edges(steps, self._depends_on(entry, steps)) references: list[Reference] = [] for index, step, cls in declared: where = f"step {step.get('id')!r}: " @@ -210,8 +211,8 @@ def _task(self, entry: BundleEntry) -> ResolvedEntry: if not (isinstance(site.value, str) and site.value): self._problem(entry, f"{where}{site.path} must be an id, not {site.value!r}") continue - self._resolve(entry, site, f"{where}{site.path}", site.value, site.version, outputs, upstream[index], - references) + self._resolve(entry, site, f"{where}{site.path}", site.value, site.version, outputs, + functools.partial(_upstream, edges, index), references) return ResolvedEntry(entry, steps, tuple(references), tuple(outputs.values())) def _image_folder(self, entry: BundleEntry) -> bool: @@ -374,12 +375,12 @@ def _depends_on(self, entry: BundleEntry, steps: list[dict]) -> list[list[str] | return depends def _resolve(self, entry: BundleEntry, site: RefSite, where: str, name: str, version: int | None, - outputs: dict[tuple[EntityKind, str], Output], upstream: set[int], + outputs: dict[tuple[EntityKind, str], Output], upstream: Callable[[], set[int]], references: list[Reference]) -> None: kind, expected, key = site.ref.kind, site.ref.artifact_type or site.ref.env_type, _nfc(name) output = outputs.get((kind, key)) if output is not None: - if output.step not in upstream: + if output.step not in upstream(): self._problem(entry, f"{where}: {name!r} is written by step {output.step_id!r}, which this step " "doesn't depend on") elif version is not None: @@ -455,26 +456,30 @@ def _nfc(name: str) -> str: return unicodedata.normalize("NFC", name) -def _upstream(steps: list[dict], depends: list[list[str] | None]) -> list[set[int]]: - """The indexes of the steps each step runs after: the step ids ``depends`` holds for it or, when it has - none, every step before it, followed transitively; the rule the task scheduler runs steps by.""" +def _step_edges(steps: list[dict], depends: list[list[str] | None]) -> list[list[int] | None]: + """The indexes of the steps each step's ``depends`` names; None for a step that names none.""" index_of: dict[str, int] = {} for index, step in enumerate(steps): if isinstance(step.get("id"), str): index_of.setdefault(step["id"], index) - direct = [set(range(index)) if ids is None else {index_of[dep_id] for dep_id in ids} - for index, ids in enumerate(depends)] - upstream = [] - for index in range(len(steps)): - seen: set[int] = set() - frontier = list(direct[index]) - while frontier: - current = frontier.pop() - if current not in seen: - seen.add(current) - frontier.extend(direct[current]) - upstream.append(seen) - return upstream + return [None if ids is None else [index_of[dep_id] for dep_id in ids] for ids in depends] + + +def _upstream(edges: list[list[int] | None], index: int) -> set[int]: + """The steps step ``index`` runs after: those its edges name, or all before it when it names none, transitively.""" + seen: set[int] = set() + below = 0 # each step under it is already seen, which keeps the walk linear + frontier = [index] + while frontier: + current = frontier.pop() + named = edges[current] + if named is None: + named, below = range(below, current), max(below, current) + for step in named: + if step not in seen: + seen.add(step) + frontier.append(step) + return seen def _same_type(artifact_type: str, expected: str) -> bool: diff --git a/src/agent_env/cli/__init__.py b/src/agent_env/cli/__init__.py index aad88a4a..a2b4f10b 100644 --- a/src/agent_env/cli/__init__.py +++ b/src/agent_env/cli/__init__.py @@ -11,6 +11,7 @@ from .a2a_agent import a2a_agent from .artifact import artifact +from .bundle import bundle from .config import config from .env import env from .eval import eval @@ -78,6 +79,7 @@ def cli(ctx: click.Context, verbose: bool): cli.add_command(a2a_agent) cli.add_command(env) cli.add_command(artifact) +cli.add_command(bundle) cli.add_command(config) cli.add_command(eval) cli.add_command(plugin) diff --git a/src/agent_env/cli/_plugin_changes.py b/src/agent_env/cli/_plugin_changes.py index 518b4eff..ecc00107 100644 --- a/src/agent_env/cli/_plugin_changes.py +++ b/src/agent_env/cli/_plugin_changes.py @@ -28,6 +28,7 @@ from agent_env.cli import _installers from agent_env.cli._installers import Environment, InstallerError, Plan, ToolReceipt, normalize +from agent_env.cli.utils import interruptible from agent_env.config import get_config from agent_env.config.loader import interpolate from agent_env.config.paths import state_root @@ -501,7 +502,7 @@ def _uninstall(plan: Plan, saved: _Rollback, targets: list[str], also: list[str] def _apply(saved: _Rollback, change: Callable[[], list[str]], tools: Tools, *, keep: bool = False) -> int: """Make the change, and put the environment back as it was unless it took effect.""" - with _interruptible(): + with interruptible(): # so SIGTERM reaches the rollback try: problems = change() except KeyboardInterrupt: @@ -553,19 +554,6 @@ def _lock_path(env: Environment) -> Path: return state_root() / "locks" / f"plugin-{digest}.lock" -@contextmanager -def _interruptible() -> Iterator[None]: - """SIGTERM interrupts like Ctrl-C, so both reach the rollback.""" - def interrupt(signum, frame): - raise KeyboardInterrupt - - previous = signal.signal(signal.SIGTERM, interrupt) - try: - yield - finally: - signal.signal(signal.SIGTERM, previous) - - @contextmanager def _uninterrupted() -> Iterator[None]: """Ctrl-C and SIGTERM wait until the restore finishes; stopping halfway would leave the diff --git a/src/agent_env/cli/bundle.py b/src/agent_env/cli/bundle.py new file mode 100644 index 00000000..11d509d9 --- /dev/null +++ b/src/agent_env/cli/bundle.py @@ -0,0 +1,48 @@ +"""``agent-env bundle``: start a bundle folder from what a store holds.""" + +from __future__ import annotations + +from pathlib import Path + +import click + +from agent_env.bundle.importer import BundleImport, import_eval, import_task +from agent_env.cli.utils import interruptible + + +@click.group() +def bundle(): + """Create bundle folders, which ``agent-env run`` runs.""" + + +@bundle.command("init") +@click.argument("folder", type=click.Path(file_okay=False, path_type=Path)) +@click.option("--from-task", "task_id", default=None, help="Import this stored task, by id.") +@click.option("--from-eval", "eval_id", default=None, help="Import this stored eval and every task it names, by id.") +@click.option("--version", type=int, default=None, help="The version to import (defaults to the latest).") +def init(folder: Path, task_id: str | None, eval_id: str | None, version: int | None): + """Write a stored task, or an eval and its tasks, as a bundle at FOLDER, which must be new or empty. + + Each task is written as the store holds it, and every store reference it makes is pinned at the version the + store holds now, so the folder runs what the store ran; what a step writes stays unpinned, since the run writes + it again. Values that look like secrets are left out. Nothing is written to a store, and an import that fails + writes nothing.""" + if (task_id is None) == (eval_id is None): + raise click.UsageError("give one of --from-task or --from-eval") + with interruptible(): + imported = (import_task(task_id, folder, version=version) if task_id is not None + else import_eval(eval_id, folder, version=version)) + _report(imported, folder) + + +def _report(imported: BundleImport, folder: Path) -> None: + what = (f"eval {imported.eval.id} v{imported.eval.version} and its tasks" if imported.eval + else f"task {imported.tasks[0].id} v{imported.tasks[0].version}") + click.echo(f"Imported {what} into {folder}:") + for path in imported.files: + click.echo(f" {path}") + for title, lines in imported.report(): + click.echo(f"{title}:") + for line in lines: + click.echo(f" {line}") + click.echo(f"Run it with: agent-env run {folder}") diff --git a/src/agent_env/cli/utils.py b/src/agent_env/cli/utils.py index 7df4225e..21a0b72d 100644 --- a/src/agent_env/cli/utils.py +++ b/src/agent_env/cli/utils.py @@ -1,4 +1,7 @@ import asyncio +import signal +from collections.abc import Iterator +from contextlib import contextmanager import click @@ -11,6 +14,19 @@ from agent_env.utils.docker_build import DEFAULT_BUILD_PLATFORM +@contextmanager +def interruptible() -> Iterator[None]: + """SIGTERM interrupts like Ctrl-C, raising KeyboardInterrupt, so cleanup runs for both.""" + def interrupt(signum, frame): + raise KeyboardInterrupt + + previous = signal.signal(signal.SIGTERM, interrupt) + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) + + def build_platform_option(f): """Shared `--platform` option for CLI commands that run `docker build` locally. diff --git a/src/agent_env/config/describe.py b/src/agent_env/config/describe.py index d0b49a97..1445a0b3 100644 --- a/src/agent_env/config/describe.py +++ b/src/agent_env/config/describe.py @@ -193,7 +193,7 @@ class ConfigReport: plugins: PluginsReport = field(default_factory=PluginsReport) -def _is_secret_key(key: str) -> bool: +def is_secret_key(key: str) -> bool: return any(part in key.lower() for part in _SECRET_KEY_PARTS) @@ -214,7 +214,7 @@ def mask_value(key: str, value: Any, *, secret_scope: bool = False) -> Any: ref, sep, default = value.partition("?") return f"{ref}?{mask_value(key, default, secret_scope=secret_scope)}" if sep else value # Before the string check: a secret is no less a secret for being written as a number. - if secret_scope or _is_secret_key(key): + if secret_scope or is_secret_key(key): return MASK if not isinstance(value, str): return value @@ -230,7 +230,7 @@ def _mask_table(table: Any, *, key: str = "", secret_scope: bool = False) -> Any consulted the leaf's own key would print both. """ if isinstance(table, dict): - return {k: _mask_table(v, key=k, secret_scope=secret_scope or _is_secret_key(k)) + return {k: _mask_table(v, key=k, secret_scope=secret_scope or is_secret_key(k)) for k, v in table.items()} if isinstance(table, list): # A list carries no keys of its own — its items stay in the enclosing key's scope. diff --git a/tst/integration/cli/bundle_init_test.py b/tst/integration/cli/bundle_init_test.py new file mode 100644 index 00000000..84ad1304 --- /dev/null +++ b/tst/integration/cli/bundle_init_test.py @@ -0,0 +1,151 @@ +"""``agent-env bundle init`` on tasks and an eval held in a store, with no infra: the local stores and a local VM-mode +sandbox. The imported folder runs, reruns writing nothing, and keeps running what the store held at import after the +store moves on; an older version imports as asked, a task a bundle ran imports into another folder and runs there, and +secrets left out never reach the run. Needs `bash` and the usual shell tools on PATH.""" + +import json +import logging +import shutil +from pathlib import Path + +import pytest +from click.testing import CliRunner + +import agent_env +from agent_env.artifact import FileArtifactUniverse +from agent_env.artifact.store import reset_artifact_store +from agent_env.bundle import parse_bundle +from agent_env.bundle.resolve import build_step +from agent_env.cli import cli +from agent_env.config import configure, reset_config +from agent_env.eval import Eval +from agent_env.eval.eval import EvalTask +from agent_env.store.routing import namespace_routing +from agent_env.task import Task + +pytestmark = pytest.mark.integration + +HELLO = Path(agent_env.__file__).parent / "examples" / "hello" + + +@pytest.fixture +def registry(monkeypatch, tmp_path): + """hello's task and greeting, put under bare ids in the configured local store. Whatever ran, no sandbox work + folder may be left behind.""" + sandboxes = tmp_path / "sandboxes" + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local") + monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes)) + configure() + reset_artifact_store() + greeting = HELLO / "artifacts" / "greeting" + FileArtifactUniverse.put_bundled("greeting", files={path.name: path for path in greeting.iterdir()}) + Task.put(id="greet", steps=[build_step(step) for step in _hello_steps()]) + Eval.put(id="smoke", tasks=[EvalTask("greet")]) + logging.disable(logging.CRITICAL) # pytest's live logging would take CliRunner's stdout + try: + yield tmp_path + assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a run left a sandbox work folder" + finally: + logging.disable(logging.NOTSET) + reset_artifact_store() + reset_config() + + +def _hello_steps(): + return json.loads((HELLO / "tasks" / "hello.json").read_text()) + + +def _cli(*args): + result = CliRunner().invoke(cli, [*map(str, args)]) + assert result.exit_code == 0, result.output + return result.output + + +def test_an_imported_task_runs_reruns_writing_nothing_and_keeps_its_pins_when_the_store_moves_on(registry): + imported = _cli("bundle", "init", registry / "greet", "--from-task", "greet") + + assert " artifact greeting v1: tasks/greet.json: step 'load': artifacts[0].id" in imported.splitlines() + first = _cli("run", registry / "greet") + assert "tasks/greet.json v1: passed (hello: 1)" in first + second = _cli("run", registry / "greet") + assert "tasks/greet.json: v1, unchanged" in second + assert "tasks/greet.json v1: passed (hello: 1)" in second + + (registry / "bye.txt").write_text("bye\n") + FileArtifactUniverse.put_bundled("greeting", files={"hello.txt": registry / "bye.txt"}) + assert "tasks/greet.json v1: passed (hello: 1)" in _cli("run", registry / "greet") + + +def test_an_older_version_imports_as_it_was_stored(registry): + steps = _hello_steps() + steps[2]["criteria"][0]["expected"] = "goodbye" + Task.put(id="greet", steps=[build_step(step) for step in steps]) + + _cli("bundle", "init", registry / "old", "--from-task", "greet", "--version", 1) + _cli("bundle", "init", registry / "new", "--from-task", "greet") + + assert (registry / "old" / "README.md").read_text().splitlines()[0] == "Imported from the store's task greet v1." + assert "tasks/greet.json v1: passed (hello: 1)" in _cli("run", registry / "old") + assert "tasks/greet.json v1: scored below 1" in _cli("run", registry / "new") + + +def test_a_task_a_bundle_ran_imports_into_another_folder_and_runs_there(registry): + shutil.copytree(HELLO, registry / "first") + assert "tasks/hello.json v1: passed (hello: 1)" in _cli("run", registry / "first") + first = parse_bundle(registry / "first").id_root + + imported = _cli("bundle", "init", registry / "second", "--from-task", f"{first}/hello") + + assert f" artifact {first}/greeting v1: tasks/hello.json: step 'load': artifacts[0].id" in imported.splitlines() + assert "tasks/hello.json v1: passed (hello: 1)" in _cli("run", registry / "second") + assert "tasks/hello.json: v1, unchanged" in _cli("run", registry / "second") + + +def test_secrets_left_out_of_an_import_stay_out_of_the_task_its_run_writes(registry): + key = "AK" + "IA" + "Q4ZX7" * 3 + "W" # an AWS access key's shape, joined so no scanner takes it for one + url = "https://u:" + "pw-1234@example.com/x" + steps = _hello_steps() + steps[0]["env_vars"] = {"API_TOKEN": "a-stored-token", "NOTE": key, "MODE": "kept"} + steps[2]["criteria"].append({"type": "bash_cmd_succeeds", "criterion": "fetches", "bash_cmd": f"echo {url}"}) + Task.put(id="guarded", steps=[build_step(step) for step in steps]) + + imported = _cli("bundle", "init", registry / "guarded", "--from-task", "guarded").splitlines() + ran = _cli("run", registry / "guarded") + + at = imported.index("Left out:") + assert imported[at + 1:at + 4] == [ + " tasks/guarded.json: step 'box': env_vars.API_TOKEN: dropped: its name looks like a secret", + " tasks/guarded.json: step 'box': env_vars.NOTE: dropped: its value holds an AWS access key", + " tasks/guarded.json: step 'hello': criteria[2].bash_cmd: masked: a URL's credentials"] + assert "tasks/guarded.json v1: passed (hello: 1)" in ran + with namespace_routing(): + run_steps = Task.get(f"{parse_bundle(registry / 'guarded').id_root}/guarded").to_dict()["steps"] + assert run_steps[0]["env_vars"] == {"MODE": "kept"} + assert run_steps[2]["criteria"][2]["bash_cmd"] == "echo https://***@example.com/x" + written = "".join(path.read_text() for path in (registry / "guarded").rglob("*") if path.is_file()) + assert not any(secret in written for secret in ("a-stored-token", key, "pw-1234")) + + +def test_an_imported_eval_runs_its_tasks(registry): + imported = _cli("bundle", "init", registry / "nightly", "--from-eval", "smoke") + + assert imported.splitlines()[0] == f"Imported eval smoke v1 and its tasks into {registry / 'nightly'}:" + output = _cli("run", registry / "nightly") + assert "tasks/greet.json v1: passed (hello: 1)" in output + assert "evals/smoke.toml v1: 1 of 1 passed" in output + + +def test_an_import_that_fails_leaves_nothing_beside_its_destination(registry): + Task.put(id="broken", steps=[build_step({"id": "load", "type": "load_artifact", "artifact_id": "gone", + "sandbox_name": "box", "destination_path": "/app"})]) + before = sorted(path.name for path in registry.iterdir()) + + result = CliRunner().invoke(cli, ["bundle", "init", str(registry / "broken"), "--from-task", "broken"]) + + assert result.exit_code == 1 + assert "there is no artifact 'gone' in the store" in result.output + assert sorted(path.name for path in registry.iterdir()) == before diff --git a/tst/integration/cli/run_bundle_built_envs_test.py b/tst/integration/cli/run_bundle_built_envs_test.py index 0ce46d1c..26048db3 100644 --- a/tst/integration/cli/run_bundle_built_envs_test.py +++ b/tst/integration/cli/run_bundle_built_envs_test.py @@ -6,25 +6,32 @@ when a child is. Environments and a universe the bundle writes load into an MCP server deployed through the gateway, and an edit to their files loads anew. A universe composed every way a bundle composes one seeds each server of a multi of bundle and store servers, and a store universe shared across servers seeds the one server it's loaded into. The -echo-mcp example agent-env ships runs by name, and a rerun writes nothing. Needs a Docker daemon and the local -registry.""" +echo-mcp example agent-env ships runs by name, and a rerun writes nothing, and a task held under bare ids, as a registry +holds it, imports as a bundle that runs through the gateway; an output a bundle's task wrote is written under the folder +that imports the task. Needs a Docker daemon and the local registry.""" import hashlib import json import logging import shutil +import subprocess from pathlib import Path import pytest from click.testing import CliRunner -from agent_env.artifact import EnvironmentArtifact, EnvironmentUniverseArtifact, FileArtifact +from agent_env.artifact import Artifact, EnvironmentArtifact, EnvironmentUniverseArtifact, FileArtifact +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.artifact.store import reset_artifact_store from agent_env.bundle import parse_bundle +from agent_env.bundle.resolve import build_step from agent_env.cli import cli from agent_env.config import configure, reset_config from agent_env.env import Env +from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.store.base import NotFoundError from agent_env.store.routing import namespace_routing +from agent_env.task import Task pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -553,3 +560,59 @@ def test_the_shipped_echo_mcp_example_runs_by_name_and_a_rerun_writes_nothing(st for entry in ("envs/echo (Dockerfile image)", "envs/echo", "artifacts/check_echo", "tasks/smoke.json"): assert f"{entry}: v1, unchanged" in second.output, second.output assert "tasks/smoke.json v1: passed (echoes: 1)" in second.output, second.output + + +ECHO = REPO / "src" / "agent_env" / "examples" / "echo-mcp" + + +def _store_echo_server() -> None: + """The echo example's MCP server, built and put under bare ids, as a registry holds it.""" + tag = "bundle-import-echo:test" + build = subprocess.run(["docker", "build", "-t", tag, str(ECHO / "envs" / "echo")], capture_output=True, text=True) + assert build.returncode == 0, build.stderr[-2000:] + image = DockerImageArtifact.put("echo-server-image", description="echo server", image_name=tag) + MCPServerEnv.put(id="echo-server", docker_image_artifact=image, environment_name="echo") + + +def test_a_stored_mcp_task_imports_as_a_bundle_that_runs_through_the_gateway(state): + _store_echo_server() + FileArtifact.put_bytes("echo-check", description="calls echo", filename="verify.py", + content=(ECHO / "artifacts" / "check_echo" / "verify.py").read_bytes()) + Task.put(id="echo-smoke", steps=[build_step(step) for step in [ + {"id": "env", "type": "deploy_env", "env_id": "echo-server", "sandbox_type": "local"}, + {"id": "echoes", "type": "env_outcome_verifier", "env_id": "echo-server", "file_artifact_id": "echo-check", + "verifier_id": "echoes"}]]) + + imported = CliRunner().invoke(cli, ["bundle", "init", str(state / "echo"), "--from-task", "echo-smoke"]) + + assert imported.exit_code == 0, imported.output + assert " env echo-server v1: tasks/echo-smoke.json: step 'env': env_id" in imported.output.splitlines() + assert " artifact echo-check v1: tasks/echo-smoke.json: step 'echoes': file_artifact_id" in imported.output + first, second = _run(state / "echo"), _run(state / "echo") + assert first.exit_code == 0, first.output + assert "tasks/echo-smoke.json v1: passed (echoes: 1)" in first.output, first.output + assert "tasks/echo-smoke.json: v1, unchanged" in second.output, second.output + assert "tasks/echo-smoke.json v1: passed (echoes: 1)" in second.output, second.output + + +def test_an_output_a_bundles_task_wrote_is_written_under_the_folder_that_imports_it(state): + _store_echo_server() + old = "@local/~/old/echo-cli" + with namespace_routing(): + Task.put(id=old, steps=[build_step(step) for step in [ + {"id": "env", "type": "deploy_env", "env_id": "echo-server", "sandbox_type": "local"}, + {"id": "cli", "type": "build_mcp_cli", "env_id": "echo-server", "command_name": "echo", + "cli_artifact_id": f"{old}/echo-cli", "depends_on": ["env"]}]]) + + imported = CliRunner().invoke(cli, ["bundle", "init", str(state / "again"), "--from-task", old]) + + assert imported.exit_code == 0, imported.output + steps = json.loads((state / "again" / "tasks" / "echo-cli.json").read_text()) + assert steps[1]["cli_artifact_id"] == "echo-cli" + ran = _run(state / "again") + assert ran.exit_code == 0, ran.output + with namespace_routing(): + written = Artifact.get(f"{parse_bundle(state / 'again').id_root}/echo-cli/echo-cli") + assert written.type == "cli" + with pytest.raises(NotFoundError): + Artifact.get(f"{old}/echo-cli") diff --git a/tst/unit/bundle/importer_test.py b/tst/unit/bundle/importer_test.py new file mode 100644 index 00000000..b15bcb1e --- /dev/null +++ b/tst/unit/bundle/importer_test.py @@ -0,0 +1,503 @@ +"""Importing a stored task, or an eval and its tasks, as a bundle: each step as the store holds it, every store +reference pinned at the version the store holds now, an id a step of the task writes left unpinned, secrets left +out, and a folder that plans and materializes back to the stored steps. A failed import writes nothing.""" + +import json +import os +import signal +import tomllib +from typing import ClassVar + +import pytest +from click.testing import CliRunner + +from agent_env.a2a_agent import A2AAgent +from agent_env.artifact import FileArtifact +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.artifact.store import get_artifact_store +from agent_env.bundle import BundleError, parse_bundle +from agent_env.bundle import importer as importer_module +from agent_env.bundle.importer import ImportedEval, import_eval, import_task +from agent_env.bundle.materialize import materialize +from agent_env.bundle.plan import plan_bundle +from agent_env.bundle.resolve import resolve_bundle +from agent_env.cli import cli +from agent_env.config import get_config +from agent_env.config.runtime import Config +from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.eval import Eval +from agent_env.eval.eval import EvalTask +from agent_env.store import Filter +from agent_env.store.routing import namespace_routing +from agent_env.task import Task +from agent_env.task_step.task_step import TaskStep + +STEPS = [ + {"id": "env", "type": "deploy_env", "env_id": "crm-env"}, + {"id": "load", "type": "load_artifact", "env_id": "crm-env", "artifacts": [{"id": "crm-data"}]}, + {"id": "agent", "type": "deploy_agent", "env_ids": ["crm-env"], "a2a_agent_id": "solver", + "env_vars": {"LITELLM_API_KEY": "sk-1", "MODE": "fast", "DB": "postgres://admin:pw@db/crm"}}, + {"id": "helper", "type": "deploy_agent", "agent_name": "helper", "env_ids": []}, + {"id": "judge", "type": "rubrics_verifier", "verifier_id": "judge", "prompt_id": "agent", + "judge_a2a_agent_id": "solver", + "criteria": [{"id": "c1", "description": "cites https://user:hunter2@wiki.example.com/x"}, + {"id": "c2", "description": "clones ssh://git@github.com/org/repo"}]}, + {"id": "cli", "type": "build_mcp_cli", "env_id": "crm-env", "command_name": "crm", "cli_artifact_id": "crm-cli"}, + {"id": "skills", "type": "add_skills", "agent_name": "default-agent", "cli_artifact_ids": ["crm-cli"], + "depends_on": [{"task_step_id": "cli"}]}, +] + + +def _image(id): + return get_artifact_store().put_document(DockerImageArtifact( + id=id, description=id, image_name=f"registry.example.com/{id}:v1", tar_gz_s3_url=f"file:///{id}.tar.gz")) + + +@pytest.fixture +def store(local_stores, tmp_path, monkeypatch): + """A registry of bare ids in the configured local store, holding the entities STEPS names: crm-env at v2.""" + monkeypatch.setenv("HOME", str(tmp_path)) + MCPServerEnv.put(id="crm-env", docker_image_artifact=_image("crm-image"), environment_name="crm") + MCPServerEnv.put(id="crm-env", docker_image_artifact=_image("crm-image"), environment_name="crm") + A2AAgent.put(id="solver", docker_image_artifact=_image("solver-image")) + (tmp_path / "crm.json").write_text("{}") + FileArtifact.put("crm-data", description="crm seed", file_path=str(tmp_path / "crm.json")) + _task("triage", STEPS) + return tmp_path + + +def _task(task_id, steps): + from agent_env.bundle.resolve import build_step + return Task.put(id=task_id, steps=[build_step(dict(step)) for step in steps]) + + +def _steps(root, name): + return json.loads((root / "tasks" / f"{name}.json").read_text()) + + +def _left(root): + """Everything an import left beside its destination, staging folders included.""" + return sorted(path.name for path in root.iterdir()) + + +def test_a_task_imports_as_its_stored_steps_with_each_store_reference_pinned_at_its_version_now(store): + imported = import_task("triage", store / "triage") + steps = _steps(store / "triage", "triage") + stored = Task.get("triage").to_dict()["steps"] + + assert imported.files == ("README.md", "tasks/triage.json") + assert [(step["id"], step["type"]) for step in steps] == [(step["id"], step["type"]) for step in stored] + assert steps[0]["env_version"] == 2 and stored[0]["env_version"] is None + assert steps[1]["artifacts"] == [{"id": "crm-data", "version": 1}] + assert steps[2]["a2a_agent_version"] == 1 + assert {k: v for k, v in steps[4].items() if k != "criteria"} == { + k: v for k, v in stored[4].items() if k != "criteria"} + assert sorted((pin.kind.value, pin.id, pin.version, pin.stored) for pin in imported.pinned) == [ + ("agent", "solver", 1, False), ("artifact", "crm-data", 1, False), ("env", "crm-env", 2, False)] + + +def test_an_id_a_step_of_the_task_writes_stays_unpinned_and_unchanged(store): + steps = _steps(import_task("triage", store / "triage").root, "triage") + + assert steps[5]["cli_artifact_id"] == "crm-cli" + assert steps[6]["cli_artifact_ids"] == ["crm-cli"] + + +def test_a_reference_that_takes_no_version_and_a_step_that_names_no_agent_are_reported_as_floating(store): + imported = import_task("triage", store / "triage") + + assert sorted((ref.where, ref.kind.value, ref.id) for ref in imported.floating) == [ + ("tasks/triage.json: step 'helper': a2a_agent_id", "agent", "a2a-default"), + ("tasks/triage.json: step 'judge': judge_a2a_agent_id", "agent", "solver"), + ] + + +def test_env_vars_that_look_like_secrets_are_left_out_and_a_urls_credentials_masked_elsewhere(store): + imported = import_task("triage", store / "triage") + steps = _steps(imported.root, "triage") + + assert steps[2]["env_vars"] == {"MODE": "fast"} + assert [c["description"] for c in steps[4]["criteria"]] == [ + "cites https://***@wiki.example.com/x", "clones ssh://git@github.com/org/repo"] + assert [(drop.where, drop.why) for drop in imported.dropped] == [ + ("tasks/triage.json: step 'agent': env_vars.LITELLM_API_KEY", "dropped: its name looks like a secret"), + ("tasks/triage.json: step 'agent': env_vars.DB", "dropped: its value holds a URL's credentials"), + ("tasks/triage.json: step 'judge': criteria[0].description", "masked: a URL's credentials"), + ] + written = (imported.root / "README.md").read_text() + (imported.root / "tasks/triage.json").read_text() + assert "sk-1" not in written and "hunter2" not in written and "admin:pw" not in written + + +# Each in the shape its issuer gives it, joined at run time so no scanner takes it for a real one. +TOKENS = { + "an API key": "sk" + "-proj-" + "a1B2" * 6, + "a GitHub token": "gh" + "p_" + "a1B2c3" * 6, + "an AWS access key": "AK" + "IA" + "Q4ZX7" * 3 + "W", + "a Slack token": "xo" + "xb-" + "1234-5678-abcd", + "a Google API key": "AI" + "za" + "Sy" + "a1B2c" * 6 + "d3E", + "a Hugging Face token": "hf" + "_" + "a1B2c3" * 6, + "a JWT": "ey" + "JhbGciOiJIUzI1NiJ9." + "ey" + "JzdWIiOiIxMjM0In0." + "c2lnbmF0dXJlMTIz", + "a private key": "-----BEGIN " + "RSA PRIVATE KEY-----\nMIIBOgIBAAJBAK\n-----END RSA PRIVATE KEY-----", +} + + +@pytest.mark.parametrize("what", TOKENS) +def test_a_credential_by_its_shape_is_dropped_from_an_env_var_and_masked_anywhere_else(store, what): + token = TOKENS[what] + _task("creds", [STEPS[0], STEPS[1], {**STEPS[2], "env_vars": {"SETTING": token}}, + {**STEPS[4], "criteria": [{"id": "c1", "description": f"logs in with {token} today"}]}]) + + imported = import_task("creds", store / "creds") + steps = _steps(imported.root, "creds") + + assert steps[2]["env_vars"] == {} + assert steps[3]["criteria"][0]["description"] == "logs in with *** today" + assert [(drop.where, drop.why) for drop in imported.dropped] == [ + ("tasks/creds.json: step 'agent': env_vars.SETTING", f"dropped: its value holds {what}"), + ("tasks/creds.json: step 'judge': criteria[0].description", f"masked: {what}"), + ] + assert not any(token in path.read_text() for path in imported.root.rglob("*") if path.is_file()) + + +def test_text_that_only_resembles_a_credential_is_kept(store): + kept = ["sk-learn-is-a-python-library", "max_tokens=4096", "eyJ is how a JWT starts", "ssh://git@github.com/o/r", + "commit 5a3c2bf0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f5a6"] + _task("kept", [STEPS[0], STEPS[1], {**STEPS[2], "env_vars": {f"V{i}": text for i, text in enumerate(kept)}}, + {**STEPS[4], "criteria": [{"id": f"c{i}", "description": text} for i, text in enumerate(kept)]}]) + + imported = import_task("kept", store / "kept") + steps = _steps(imported.root, "kept") + + assert list(steps[2]["env_vars"].values()) == kept + assert [criterion["description"] for criterion in steps[3]["criteria"]] == kept + assert imported.dropped == () + + +def test_the_readme_says_where_the_folder_came_from_and_how_each_reference_was_written(store): + readme = (import_task("triage", store / "triage").root / "README.md").read_text() + + assert readme.splitlines()[0] == "Imported from the store's task triage v1." + assert "- env crm-env v2: tasks/triage.json: step 'env': env_id" in readme + assert "## Floating" in readme and "## Left out" in readme + assert parse_bundle(store / "triage").description == "Imported from the store's task triage v1." + + +def test_the_folder_materializes_back_to_the_stored_steps(store): + imported = import_task("triage", store / "triage") + bundle = parse_bundle(imported.root) + with namespace_routing(): + materialization = materialize(plan_bundle(resolve_bundle(bundle))) + written = Task.get(f"{bundle.id_root}/triage", materialization.version_of("task", f"{bundle.id_root}/triage")) + expected = _steps(imported.root, "triage") + output = f"{bundle.id_root}/triage/crm-cli" + expected[5]["cli_artifact_id"], expected[6]["cli_artifact_ids"] = output, [output] + + assert written.to_dict()["steps"] == expected + + +def test_an_id_a_bundle_named_under_its_task_is_renamed_so_the_import_names_it_under_its_own(store): + old = "@local/~/old/triage" + with namespace_routing(): + _task(old, [STEPS[0], {**STEPS[5], "cli_artifact_id": f"{old}/crm-cli"}, + {**STEPS[6], "cli_artifact_ids": [f"{old}/crm-cli"]}]) + imported = import_task(old, store / "again") + steps = _steps(imported.root, "triage") + + assert (steps[1]["cli_artifact_id"], steps[2]["cli_artifact_ids"]) == ("crm-cli", ["crm-cli"]) + + +def test_a_renamed_output_is_numbered_when_another_output_or_a_store_reference_has_its_name(store): + old = "@local/~/old/triage" + with namespace_routing(): + _task(old, [STEPS[0], STEPS[1], + {**STEPS[5], "cli_artifact_id": f"{old}/crm-cli"}, + {**STEPS[5], "id": "cli2", "cli_artifact_id": "@local/~/other/triage/crm-cli"}, + {**STEPS[5], "id": "cli3", "cli_artifact_id": f"{old}/crm-data"}, + {**STEPS[6], "cli_artifact_ids": [f"{old}/crm-cli", "@local/~/other/triage/crm-cli"], + "depends_on": [{"task_step_id": "cli"}, {"task_step_id": "cli2"}]}]) + imported = import_task(old, store / "again") + steps = _steps(imported.root, "triage") + + assert [step.get("cli_artifact_id") for step in steps[2:5]] == ["crm-cli", "crm-cli-2", "crm-data-2"] + assert steps[5]["cli_artifact_ids"] == ["crm-cli", "crm-cli-2"] + assert steps[1]["artifacts"] == [{"id": "crm-data", "version": 1}] + + +def test_an_env_no_step_of_the_task_deploys_floats(store): + _task("live", [{"id": "agent", "type": "deploy_agent", "env_ids": ["crm-env"], "a2a_agent_id": "solver"}]) + + imported = import_task("live", store / "live") + + assert [(ref.kind.value, ref.id, ref.where) for ref in imported.floating] == [ + ("env", "crm-env", "tasks/live.json: step 'agent': env_ids[0]")] + + +def test_a_pin_the_task_already_has_is_kept_when_the_store_holds_it_and_refused_when_it_doesnt(store): + _task("pinned", [{"id": "env", "type": "deploy_env", "env_id": "crm-env", "env_version": 1}]) + _task("gone", [{"id": "env", "type": "deploy_env", "env_id": "crm-env", "env_version": 9}]) + + kept = import_task("pinned", store / "pinned") + with pytest.raises(BundleError) as e: + import_task("gone", store / "gone") + + assert [(pin.id, pin.version, pin.stored) for pin in kept.pinned] == [("crm-env", 1, True)] + assert e.value.problems == ("tasks/gone.json: step 'env': env_id: there is no env 'crm-env' v9 in the store",) + + +def test_the_current_folder_imports_when_it_is_empty(store, monkeypatch): + (store / "here").mkdir() + monkeypatch.chdir(store / "here") + before = _left(store) + + result = CliRunner().invoke(cli, ["bundle", "init", ".", "--from-task", "triage"]) + + assert result.exit_code == 0, result.output + assert (store / "here" / "tasks" / "triage.json").exists() and _left(store) == before + + +def test_an_eval_imports_with_each_task_it_names_once_under_names_taken_from_their_ids(store): + _task("team/run", [STEPS[0]]) + _task("other/run", [STEPS[0]]) + _task("RUN", [STEPS[0]]) + Eval.put(id="nightly/regression", tasks=[EvalTask("team/run"), EvalTask("other/run"), EvalTask("RUN"), + EvalTask("team/run"), EvalTask("triage", 1)]) + + imported = import_eval("nightly/regression", store / "nightly") + + assert imported.files == ("README.md", "evals/regression.toml", "tasks/RUN-3.json", "tasks/run-2.json", + "tasks/run.json", "tasks/triage.json") + assert tomllib.loads((imported.root / "evals/regression.toml").read_text()) == { + "tasks": ["run", "run-2", "RUN-3", "triage"]} + assert imported.duplicates == ("task team/run v1",) + assert imported.eval == ImportedEval("evals/regression.toml", "nightly/regression", 1) + plan = plan_bundle(resolve_bundle(parse_bundle(imported.root))) + assert [entry.entry.name for entry in plan.evals] == ["regression"] + + +def test_an_eval_naming_a_task_the_store_doesnt_hold_is_refused_and_writes_nothing(store): + Eval.put(id="broken", tasks=[EvalTask("triage"), EvalTask("gone"), EvalTask("triage", 7)]) + before = _left(store) + + with pytest.raises(BundleError) as caught: + import_eval("broken", store / "broken") + + assert caught.value.problems == ("eval broken v1 names task gone, which isn't in the store", + "eval broken v1 names task triage v7, which isn't in the store") + assert _left(store) == before + + +def test_an_older_version_imports_when_asked_for(store): + _task("triage", [STEPS[0]]) + + old, new = import_task("triage", store / "old", version=1), import_task("triage", store / "new") + + assert [(task.version, task.steps) for task in old.tasks + new.tasks] == [(1, len(STEPS)), (2, 1)] + assert (old.root / "README.md").read_text().splitlines()[0] == "Imported from the store's task triage v1." + + +@pytest.mark.parametrize("task_id, name", [ + ("crm#triage", "crm-triage"), ("..hidden", "hidden"), (" edge ", "edge"), ("org/", "task"), ("😀", "task"), + ("a" * 300, "a" * 200), +]) +def test_a_name_is_the_ids_last_segment_made_valid(store, task_id, name): + _task(task_id, [STEPS[0]]) + + assert import_task(task_id, store / "named").tasks[0].path == f"tasks/{name}.json" + + +@pytest.mark.parametrize("prepare, problem", [ + (lambda dest: (dest.mkdir(), (dest / "x").write_text("x")), "not empty; import into a new or empty folder"), + (lambda dest: dest.write_text("x"), "not a folder; import into a new or empty folder"), +]) +def test_a_destination_that_holds_anything_is_refused(store, prepare, problem): + dest = store / "dest" + prepare(dest) + + with pytest.raises(BundleError) as e: + import_task("triage", dest) + + assert e.value.problems == (f"{dest}: {problem}",) + + +def test_a_destination_whose_folder_gives_no_id_or_has_no_parent_is_refused(store): + with pytest.raises(BundleError, match="rename the folder"): + import_task("triage", store / "bad​name") + with pytest.raises(BundleError, match="no such folder"): + import_task("triage", store / "missing" / "dest") + assert not (store / "missing").exists() + + +def test_an_empty_destination_is_filled(store): + (store / "dest").mkdir() + + assert import_task("triage", store / "dest").files == ("README.md", "tasks/triage.json") + + +def test_a_reference_the_store_doesnt_hold_is_refused_with_every_problem_and_nothing_is_written(store): + _task("broken", [{"id": "env", "type": "deploy_env", "env_id": "gone"}, + {"id": "load", "type": "load_artifact", "env_id": "gone", "artifacts": [{"id": "lost"}]}]) + before = _left(store) + + with pytest.raises(BundleError) as e: + import_task("broken", store / "broken") + + assert e.value.problems == ( + "tasks/broken.json: step 'env': env_id: there is no env 'gone' in the store", + "tasks/broken.json: step 'load': artifacts[0].id: there is no artifact 'lost' in the store", + ) + assert _left(store) == before + + +def test_a_folder_that_fails_its_check_is_not_moved_into_place_and_its_staging_is_removed(store, monkeypatch): + (store / "dest").mkdir() + before = _left(store) + + def refuse(resolved): + raise BundleError(["refused"]) + + monkeypatch.setattr(importer_module, "check_bundle", refuse) + with pytest.raises(BundleError): + import_task("triage", store / "dest") + + assert _left(store) == before and not any((store / "dest").iterdir()) + + +def test_an_empty_destination_is_put_back_when_the_move_fails(store, monkeypatch): + (store / "dest").mkdir() + before = _left(store) + + def fail(src, dst): + raise OSError("disk gone") + + monkeypatch.setattr(importer_module.os, "replace", fail) + with pytest.raises(BundleError) as caught: + import_task("triage", store / "dest") + + assert caught.value.problems == (f"{store / 'dest'}: can't move the import into place: disk gone",) + assert _left(store) == before and (store / "dest").is_dir() + + +def test_an_empty_destination_is_put_back_when_an_interrupt_stops_the_move(store, monkeypatch): + (store / "dest").mkdir() + before = _left(store) + + def interrupt(src, dst): + raise KeyboardInterrupt + + monkeypatch.setattr(importer_module.os, "replace", interrupt) + with pytest.raises(KeyboardInterrupt): + import_task("triage", store / "dest") + + assert _left(store) == before and (store / "dest").is_dir() + + +@pytest.mark.parametrize("empty", [False, True]) +def test_a_destination_something_else_writes_during_the_import_is_left_as_it_wrote_it(store, monkeypatch, empty): + dest = store / "dest" + if empty: + dest.mkdir() + + def land(resolved): + dest.mkdir(exist_ok=True) + (dest / "theirs.txt").write_text("theirs") + + monkeypatch.setattr(importer_module, "check_bundle", land) + with pytest.raises(BundleError) as caught: + import_task("triage", dest) + + assert caught.value.problems == (f"{dest}: not empty; import into a new or empty folder",) + assert sorted(path.name for path in dest.iterdir()) == ["theirs.txt"] + assert not [name for name in _left(store) if ".import-" in name] + + +@pytest.mark.skipif(not hasattr(os, "geteuid") or os.geteuid() == 0, reason="root writes anywhere") +def test_a_folder_the_import_cant_write_in_is_refused_in_one_line(store): + (store / "ro").mkdir() + (store / "ro").chmod(0o555) + try: + with pytest.raises(BundleError) as caught: + import_task("triage", store / "ro" / "dest") + left = list((store / "ro").iterdir()) + finally: + (store / "ro").chmod(0o755) + + assert caught.value.problems == (f"{store / 'ro'}: can't write the import there: permission denied",) + assert left == [] + + +def test_sigterm_during_an_import_leaves_nothing_beside_its_destination(store, monkeypatch): + before = _left(store) + monkeypatch.setattr(importer_module, "check_bundle", lambda resolved: os.kill(os.getpid(), signal.SIGTERM)) + + result = CliRunner().invoke(cli, ["bundle", "init", str(store / "dest"), "--from-task", "triage"]) + + assert result.exit_code == 1 and "Aborted!" in result.output + assert _left(store) == before + + +class _Undeclared(TaskStep): + """A plugin step type that declares no references.""" + + type: ClassVar[str] = "undeclared_import_test" + + def __init__(self, target=None, **base): + super().__init__(**base) + self.target = target + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data), target=data.get("target")) + + def to_dict(self): + return {**super().to_dict(), "target": self.target} + + async def execute(self, context): + return context + + +def test_a_step_whose_type_declares_no_references_is_copied_and_reported(store, monkeypatch): + steps = {**Config().task_step_registry(), _Undeclared.type: _Undeclared} + monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) + _task("plugin", [{"id": "p", "type": _Undeclared.type, "target": "crm-env"}]) + + imported = import_task("plugin", store / "plugin") + + assert _steps(imported.root, "plugin")[0]["target"] == "crm-env" + assert imported.undeclared == (f"tasks/plugin.json: step 'p' ({_Undeclared.type})",) + + +def test_the_cli_imports_and_says_what_it_pinned_left_out_and_floats(store): + result = CliRunner().invoke(cli, ["bundle", "init", str(store / "triage"), "--from-task", "triage"]) + + assert result.exit_code == 0, result.output + lines = result.output.splitlines() + assert lines[0] == f"Imported task triage v1 into {store / 'triage'}:" + assert lines[1:3] == [" README.md", " tasks/triage.json"] + assert "Pinned:" in lines and "Floating:" in lines and "Left out:" in lines + assert " env crm-env v2: tasks/triage.json: step 'env': env_id" in lines + assert lines[-1] == f"Run it with: agent-env run {store / 'triage'}" + + +@pytest.mark.parametrize("args", [[], ["--from-task", "triage", "--from-eval", "e"]]) +def test_the_cli_takes_exactly_one_source(store, args): + result = CliRunner().invoke(cli, ["bundle", "init", str(store / "x"), *args]) + + assert result.exit_code == 2 + assert "give one of --from-task or --from-eval" in result.output + + +def test_the_cli_says_in_one_line_when_the_store_has_no_such_task(store): + result = CliRunner().invoke(cli, ["bundle", "init", str(store / "x"), "--from-task", "nope"]) + + assert result.exit_code == 1 + assert result.output.startswith("Error: ") and "Traceback" not in result.output + assert not (store / "x").exists() + + +def test_reading_the_store_writes_nothing_to_it(store): + documents = get_config().get_document_store() + counts = {name: documents.count(name, Filter()) for name in ("tasks", "envs", "a2a_agents", "artifacts")} + + import_task("triage", store / "triage") + + assert {name: documents.count(name, Filter()) for name in counts} == counts diff --git a/tst/unit/bundle/plan_test.py b/tst/unit/bundle/plan_test.py index 3dd8d0c7..ef6b1477 100644 --- a/tst/unit/bundle/plan_test.py +++ b/tst/unit/bundle/plan_test.py @@ -243,7 +243,7 @@ def test_an_id_has_one_writer(make, tasks, files, problem): def test_store_ids_the_selection_reaches_are_read_once_and_kept(make, monkeypatch): put_env("crm") reads = [] - monkeypatch.setitem(plan_module._GETTERS, EntityKind.ENV, lambda *key: reads.append(key) or Env.get(*key)) + monkeypatch.setitem(plan_module.STORE_GETTERS, EntityKind.ENV, lambda *key: reads.append(key) or Env.get(*key)) plan = plan_bundle(make(tasks={"t": [deploy("crm")], "u": [deploy("crm")]})) assert reads == [("crm", None)] ((kind, env_id, version, where),) = [(r.kind, r.id, r.version, r.where) for r in plan.store_refs] @@ -291,7 +291,7 @@ def test_an_id_under_the_bundles_own_root_that_it_doesnt_write_is_refused_withou def read(*_): raise AssertionError("the plan read a store") - monkeypatch.setattr(plan_module, "_GETTERS", dict.fromkeys(plan_module._GETTERS, read)) + monkeypatch.setattr(plan_module, "STORE_GETTERS", dict.fromkeys(plan_module.STORE_GETTERS, read)) assert problems(make(tasks={"t": [deploy(f"{ROOT}/crm")]})) == ( f"tasks/t.json: step 'deploy': env_id: '{ROOT}/crm' isn't in this bundle", ) @@ -347,7 +347,7 @@ def test_a_check_covers_every_task_and_eval_and_reads_no_store(make, monkeypatch def read(*_): raise AssertionError("the check read a store") - monkeypatch.setattr(plan_module, "_GETTERS", dict.fromkeys(plan_module._GETTERS, read)) + monkeypatch.setattr(plan_module, "STORE_GETTERS", dict.fromkeys(plan_module.STORE_GETTERS, read)) check_bundle(make(tasks={"t": [deploy("tickets")], "u": [deploy("crm")]}, evals={"e": 'tasks = ["t"]\n'})) resolved = make(tasks={"u": [deploy("tickets", "d"), deploy("tickets", "d")], "v": [deploy("@local/~/other/crm")]}, diff --git a/tst/unit/bundle/resolve_test.py b/tst/unit/bundle/resolve_test.py index d8c7d60c..96d08bec 100644 --- a/tst/unit/bundle/resolve_test.py +++ b/tst/unit/bundle/resolve_test.py @@ -1,6 +1,7 @@ """Resolving a parsed bundle: every declared reference rewritten to the id it names, before any write.""" import json +import random import pytest @@ -253,6 +254,38 @@ def test_an_output_resolves_for_a_step_that_depends_on_its_writer(make, depends_ assert entry.config[2]["artifact_id"] == f"{ROOT}/t/tickets-cli" +def _walked(edges, index): + """The steps ``index`` runs after, by the plain transitive walk of the scheduler's rule.""" + direct = [set(range(i)) if named is None else set(named) for i, named in enumerate(edges)] + seen, frontier = set(), list(direct[index]) + while frontier: + current = frontier.pop() + if current not in seen: + seen.add(current) + frontier.extend(direct[current]) + return seen + + +def test_a_steps_upstream_is_what_the_scheduler_runs_it_after_cycles_and_later_steps_included(): + rng = random.Random(0) + for _ in range(500): + size = rng.randint(1, 10) + edges = [None if rng.random() < 0.4 else rng.sample(range(size), rng.randint(0, min(3, size))) + for _ in range(size)] + for index in range(size): + assert resolve_module._upstream(edges, index) == _walked(edges, index), (edges, index) + + +@pytest.mark.timeout(60) +def test_a_task_of_thousands_of_steps_resolves(make): + steps = [{"id": "cli", "type": "build_mcp_cli", "env_id": "tickets", "command_name": "c", + "cli_artifact_id": "tickets-cli"}] + steps += [{"id": f"d{i}", "type": "deploy_env", "env_id": "tickets"} for i in range(3000)] + steps.append({"id": "install", "type": "load_artifact", "agent_name": "solver", "artifact_id": "tickets-cli"}) + + assert resolved(make(tasks={"t": steps}), "t").config[-1]["artifact_id"] == f"{ROOT}/t/tickets-cli" + + @pytest.mark.parametrize(("depends_on", "problem"), [ ("cli", "depends_on is a list of step ids, not 'cli'"), ([7], 'a depends_on entry is a step id or {"task_step_id": ""}, not 7'), diff --git a/tst/unit/bundle/toml_writer_test.py b/tst/unit/bundle/toml_writer_test.py new file mode 100644 index 00000000..6d379ab9 --- /dev/null +++ b/tst/unit/bundle/toml_writer_test.py @@ -0,0 +1,41 @@ +"""The TOML a bundle import writes reads back as what was written.""" + +import math +import tomllib + +import pytest + +from agent_env.bundle._toml import dumps + +DOCUMENT = { + "tasks": ["a", "b c", 'q"\\x\n\x01\x7fé'], + "one": ["x"], + "none": [], + "n": 3, + "f": 1.5, + "b": True, + "image": {"artifact": "img/1", "version": 2}, + "mcp_server_envs": ["e", {"env": "crm", "version": 2}], + "nested": {"l": [1, 2, {"k": [3, 4]}]}, + "weird key": "v", + "metadata": {"default_model": "m", "min_disk_size_gb": 20, "tags": ["x", "y"]}, +} + + +def test_what_it_writes_reads_back_the_same(): + assert tomllib.loads(dumps(DOCUMENT)) == DOCUMENT + + +def test_a_list_of_several_spans_lines_only_at_the_top(): + assert dumps({"tasks": ["a", "b"], "t": {"l": ["a", "b"]}}) == ( + 'tasks = [\n "a",\n "b",\n]\nt = { l = ["a", "b"] }\n') + + +def test_an_empty_document_is_empty(): + assert dumps({}) == "" + + +@pytest.mark.parametrize("value", [None, math.nan, math.inf, object()]) +def test_a_value_toml_cant_hold_is_refused(value): + with pytest.raises(ValueError): + dumps({"k": value}) diff --git a/tst/unit/cli/plugins_test.py b/tst/unit/cli/plugins_test.py index 14019630..e2150e4b 100644 --- a/tst/unit/cli/plugins_test.py +++ b/tst/unit/cli/plugins_test.py @@ -348,6 +348,7 @@ def test_bare_install_core_commands_unchanged(): assert sorted(cli.commands) == [ "a2a-agent", "artifact", + "bundle", "config", "env", "eval",