From fa247df4d99a4e758cfb0fb755998c8e5c082a54 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 19:50:48 -0700 Subject: [PATCH 1/6] feat(bundle): import a stored task or eval as a bundle folder, with its references pinned agent-env bundle init --from-task ID [--version N] writes a task a store holds as a bundle that agent-env run runs: tasks/.json is the stored steps read through the current step models, with every store reference pinned at the version the store holds now, so the folder keeps running what the store ran. An id a step of the task writes stays unpinned, since the run writes it again. --from-eval writes an eval's tasks and evals/.toml. Values that look like secrets are left out: env-var entries whose names look secret or whose values hold a URL's credentials are dropped, and a URL's credentials elsewhere are masked. The README records what was imported, pinned, floating and left out. Nothing is written to a store; the folder is written beside its destination, checked as agent-env run checks a bundle, and moved into place, so a failed import leaves nothing. The parser's id-root derivation and segment check, the planner's store getters and config's secret-key predicate are now public, so the importer reuses them. Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 2 +- README.md | 2 +- src/agent_env/bundle/_toml.py | 60 ++++ src/agent_env/bundle/importer.py | 326 ++++++++++++++++++ src/agent_env/bundle/parse.py | 58 ++-- src/agent_env/bundle/plan.py | 10 +- src/agent_env/cli/__init__.py | 2 + src/agent_env/cli/bundle.py | 48 +++ src/agent_env/config/describe.py | 6 +- tst/integration/cli/bundle_init_test.py | 94 +++++ .../cli/run_bundle_built_envs_test.py | 35 +- tst/unit/bundle/importer_test.py | 316 +++++++++++++++++ tst/unit/bundle/plan_test.py | 6 +- tst/unit/bundle/toml_writer_test.py | 47 +++ tst/unit/cli/plugins_test.py | 1 + 15 files changed, 977 insertions(+), 36 deletions(-) create mode 100644 src/agent_env/bundle/_toml.py create mode 100644 src/agent_env/bundle/importer.py create mode 100644 src/agent_env/cli/bundle.py create mode 100644 tst/integration/cli/bundle_init_test.py create mode 100644 tst/unit/bundle/importer_test.py create mode 100644 tst/unit/bundle/toml_writer_test.py 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/_toml.py b/src/agent_env/bundle/_toml.py new file mode 100644 index 00000000..6af74ce7 --- /dev/null +++ b/src/agent_env/bundle/_toml.py @@ -0,0 +1,60 @@ +"""TOML for the tomls a bundle import writes, since ``tomllib`` only reads: strings, numbers, booleans, lists and +inline tables, and tables of those under named top-level keys.""" + +from __future__ import annotations + +import math +import re +from collections.abc import Collection +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], *, tables: Collection[str] = ()) -> str: + """``document`` as TOML. A value under a key in ``tables`` is written as a ``[table]`` after the other keys; + every other dict is an inline table. Raises ValueError for a value TOML can't hold, such as None.""" + lines = [f"{_key(key)} = {_value(value)}" for key, value in document.items() if key not in tables] + for key in (key for key in document if key in tables): + lines += ["", f"[{_key(key)}]", *(f"{_key(k)} = {_value(v)}" for k, v in document[key].items())] + return "\n".join(lines).lstrip("\n") + "\n" if lines else "" + + +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, since 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..fa01e059 --- /dev/null +++ b/src/agent_env/bundle/importer.py @@ -0,0 +1,326 @@ +"""Import a stored task, or an eval and its tasks, as a bundle folder that ``agent-env run`` runs. + +Each task is written as ``task get`` prints it: the stored steps, read through the current step models. Every store +reference a step makes is pinned at the version the store holds now, so the folder runs what the store held when it +was imported; an id a step of the task writes stays unpinned, since the run writes it again. Values that look like +secrets are dropped. Nothing is written to a store, and the folder is written beside its destination, checked the way +``agent-env run`` checks a bundle before it reads a store, then moved into place, so a failed import leaves nothing. +""" + +from __future__ import annotations + +import json +import os +import re +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 typing import Any + +from agent_env.bundle import _toml +from agent_env.bundle._fs import fold, show +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.config.describe import is_secret_key +from agent_env.entity_refs import EntityKind, RefRole, ref_sites +from agent_env.eval import Eval +from agent_env.store.base import NotFoundError +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.verifiers.rubrics_verifier import RubricsVerifierTaskStep + +# The step fields that map environment variable names to values, where a credential is most often set. +_ENV_VAR_MAPS = ("env_vars", "default_env_vars") +# 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 +# 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. +_URL_CREDENTIALS = re.compile(r"(?P[a-zA-Z][\w+.-]*://)(?:[^/@\s]*:[^/@\s]*|[^/@\s:]{16,})@") + + +@dataclass(frozen=True) +class Pinned: + """A store reference written with a version: the one the store held 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 BundleImport: + """What an import wrote, and how it treated each reference and secret.""" + + root: Path + files: tuple[str, ...] + tasks: tuple[ImportedTask, ...] + eval: tuple[str, str, int] | None # (path, id, version) of the eval, when one was imported + 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 sections(self) -> list[tuple[str, list[str]]]: + """The import's report, as titled lists of lines; a section with nothing to say is left out.""" + sections = [ + ("Tasks", [f"{task.path}: task {task.id} v{task.version}" for task in self.tasks]), + ("Eval", [f"{self.eval[0]}: eval {self.eval[1]} v{self.eval[2]}"] if self.eval else []), + ("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]), + ] + return [(title, lines) for title, lines in sections if lines] + + +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``, a folder that is + empty or isn't there yet in one that is. Raises BundleError listing every problem, writing nothing.""" + 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 every task it names as a bundle at + ``dest``, as ``import_task`` writes a task. Raises BundleError listing every problem, writing nothing.""" + importer = _Importer(Path(dest)) + stored = Eval.get(eval_id, version) + names, seen = [], {} + 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[task.id, task.version] = importer.add_task(task) + names.append(seen[task.id, task.version]) + 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 = dest.expanduser() + self.files: dict[str, str] = {} + self.taken: dict[str, dict[str, str]] = {"tasks": {}, "evals": {}} + self.tasks: list[ImportedTask] = [] + self.eval: tuple[str, str, int] | 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.latest: dict[tuple[EntityKind, str], 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() + classes = [registry[step["type"]] for step in steps] + outputs = {(site.ref.kind, _nfc(site.value)) + for step, cls in zip(steps, classes) for site in ref_sites(cls.entity_refs or (), step) + if site.ref.role is RefRole.OUTPUT and isinstance(site.value, str)} + for step, cls in zip(steps, classes): + where = f"{path}: step {step.get('id')!r}" + self._drop_secrets(step, where) + if cls.entity_refs is None: + self.undeclared.append(f"{where} ({cls.type})") + for site in ref_sites(cls.entity_refs or (), step): + if site.ref.role is RefRole.OUTPUT or not (isinstance(site.value, str) and site.value): + continue + if (site.ref.kind, _nfc(site.value)) in outputs: + continue # written by a step of this task, which the bundle names under the task + self._pin(site, f"{where}: {site.path}") + 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 = (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(self.dest, tuple(files), tuple(self.tasks), self.eval, tuple(self.pinned), + tuple(self.floating), tuple(self.dropped), tuple(self.undeclared), + tuple(self.duplicates)) + self.files["README.md"] = _readme(description, imported) + staging = self.dest.parent / f".{self.dest.name}.import-{uuid.uuid4().hex[:8]}" + 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") + check_bundle(resolve_bundle(parse_bundle(staging))) + emptied = self.dest.is_dir() + if emptied: + self.dest.rmdir() # empty, as _check_destination found it + try: + os.replace(staging, self.dest) + except OSError: + if emptied: + self.dest.mkdir(exist_ok=True) + raise + finally: + shutil.rmtree(staging, ignore_errors=True) + return imported + + def _pin(self, site: Any, where: str) -> None: + kind, entity_id = site.ref.kind, site.value + if site.version is not None: + self.pinned.append(Pinned(where, kind, entity_id, site.version, stored=True)) + elif site.version_key is not None: + latest = self._latest(kind, entity_id) + if isinstance(latest, int): + site.rewrite(entity_id, latest) + self.pinned.append(Pinned(where, kind, entity_id, latest, stored=False)) + else: + self.problems.append(f"{where}: {latest}") + elif kind is not EntityKind.ENV: # an env field with no version names an env the task deployed + self.floating.append(Floating(where, kind, entity_id, "the field takes no version, so a run reads the " + "latest")) + + def _latest(self, kind: EntityKind, entity_id: str) -> int | str: + """The version of ``entity_id`` the store holds now, or why it can't be read.""" + if (kind, entity_id) not in self.latest: + try: + self.latest[kind, entity_id] = STORE_GETTERS[kind](entity_id, None).version + except NotFoundError: + self.latest[kind, entity_id] = f"there is no {kind.value} {entity_id!r} in the store" + except (ValueError, KeyError, TypeError) as e: + self.latest[kind, entity_id] = f"{kind.value} {entity_id!r} can't be read ({type(e).__name__}: {e})" + return self.latest[kind, entity_id] + + 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 _drop_secrets(self, value: Any, where: str, path: str = "") -> None: + """Drop each variable of an env-var map whose name looks like a secret or whose value holds a URL's + credentials, and mask the credentials of a URL anywhere else, in place.""" + 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) + if secret or (isinstance(setting, str) and _URL_CREDENTIALS.search(setting)): + del item[name] + why = "its name looks like a secret" if secret else "its value holds a URL's credentials" + self.dropped.append(Dropped(f"{where}: {at}.{name}", f"dropped: {why}")) + elif isinstance(item, str) and _URL_CREDENTIALS.search(item): + value[key] = _URL_CREDENTIALS.sub(r"\g***@", item) + self.dropped.append(Dropped(f"{where}: {at}", "masked: a URL's credentials")) + else: + self._drop_secrets(item, where, at) + + def _name(self, kind: str, entity_id: str) -> str: + """A name for ``entity_id`` in the bundle's ``kind`` folder: its last segment, made valid, and numbered when + another entry of the kind took it, as names compare 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][fold(name)] = name + return name + + +def _readme(description: str, imported: BundleImport) -> str: + """The bundle's README: what it was imported from, its first line, 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.", ""] + for title, items in imported.sections(): + lines += [f"## {title}", "", *(f"- {item}" for item in items), ""] + return "\n".join(lines).rstrip("\n") + "\n" + + +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 becomes ``-``, and dots that would hide it or + spaces at its edges go. Empty when no letter or number is left.""" + text = unicodedata.normalize("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 _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..379b7bc4 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,36 @@ def _rel(self, path) -> str: return show(text) +def root_problems(root: Path) -> list[str]: + """Why a bundle at ``root``, a folder that needn't exist yet in one that does, would have no id; empty when it + would.""" + 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 of a bundle at the canonical path ``root``, and each part of the path that can't be one + segment of an id.""" + 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 +568,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/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/bundle.py b/src/agent_env/cli/bundle.py new file mode 100644 index 00000000..bdf05b67 --- /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 + + +@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") + 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[1]} v{imported.eval[2]} 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.sections(): + if title in ("Tasks", "Eval"): + continue + 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/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..99043206 --- /dev/null +++ b/tst/integration/cli/bundle_init_test.py @@ -0,0 +1,94 @@ +"""``agent-env bundle init`` on a task and an eval held in a store under bare ids, as a registry holds them, 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. Needs `bash` and the usual shell tools on PATH.""" + +import json +import logging +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.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.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()}) + steps = json.loads((HELLO / "tasks" / "hello.json").read_text()) + Task.put(id="greet", steps=[build_step(step) for step in 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 _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_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..89cb93a8 100644 --- a/tst/integration/cli/run_bundle_built_envs_test.py +++ b/tst/integration/cli/run_bundle_built_envs_test.py @@ -6,25 +6,30 @@ 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. 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.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.routing import namespace_routing +from agent_env.task import Task pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -553,3 +558,29 @@ 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 + + +def test_a_stored_mcp_task_imports_as_a_bundle_that_runs_through_the_gateway(state): + echo = REPO / "src" / "agent_env" / "examples" / "echo-mcp" + 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") + 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 diff --git a/tst/unit/bundle/importer_test.py b/tst/unit/bundle/importer_test.py new file mode 100644 index 00000000..df6d00bb --- /dev/null +++ b/tst/unit/bundle/importer_test.py @@ -0,0 +1,316 @@ +"""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 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 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 + + +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_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 == ("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"] + + +@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(OSError): + import_task("triage", store / "dest") + + assert _left(store) == before and (store / "dest").is_dir() + + +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/toml_writer_test.py b/tst/unit/bundle/toml_writer_test.py new file mode 100644 index 00000000..b87973e3 --- /dev/null +++ b/tst/unit/bundle/toml_writer_test.py @@ -0,0 +1,47 @@ +"""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, tables=("metadata",))) == DOCUMENT + + +def test_a_table_key_is_written_as_a_table_after_the_other_keys(): + text = dumps({"metadata": {"default_model": "m"}, "image": "x"}, tables=("metadata",)) + + assert text == 'image = "x"\n\n[metadata]\ndefault_model = "m"\n' + + +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", From eddda2c2a0819d94d0d01b1a4d2a1803d706cf16 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 20:48:47 -0700 Subject: [PATCH 2/6] fix(bundle): import review fixes, a linear upstream walk, and one-line errors when the write fails - An @local output the stored task writes is renamed to its last segment, and the steps reading it follow. - `bundle init .` takes the folder's absolute path. - An env ref without a version floats unless the task deploys it; an existing pin is read at its version. - A step's upstream steps are worked out on demand and in linear time, rather than for every step up front in cubic time; a 3,000-step task took ~400 s to resolve and now takes under a second. - A folder the import can't write in, a destination something else wrote meanwhile, and a failed move each give one line rather than a traceback; SIGTERM cleans up like Ctrl-C. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/bundle/importer.py | 119 ++++++++++++++++++--------- src/agent_env/bundle/resolve.py | 51 +++++++----- src/agent_env/cli/_plugin_changes.py | 16 +--- src/agent_env/cli/bundle.py | 6 +- src/agent_env/cli/utils.py | 16 ++++ tst/unit/bundle/importer_test.py | 92 ++++++++++++++++++++- tst/unit/bundle/resolve_test.py | 33 ++++++++ 7 files changed, 255 insertions(+), 78 deletions(-) diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py index fa01e059..774b2852 100644 --- a/src/agent_env/bundle/importer.py +++ b/src/agent_env/bundle/importer.py @@ -9,6 +9,7 @@ from __future__ import annotations +import contextlib import json import os import re @@ -22,7 +23,7 @@ from typing import Any from agent_env.bundle import _toml -from agent_env.bundle._fs import fold, show +from agent_env.bundle._fs import fold, os_reason, show 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 @@ -31,6 +32,7 @@ from agent_env.entity_refs import EntityKind, RefRole, 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 @@ -145,7 +147,7 @@ def import_eval(eval_id: str, dest: Path | str, *, version: int | None = None) - class _Importer: def __init__(self, dest: Path): - self.dest = dest.expanduser() + 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, dict[str, str]] = {"tasks": {}, "evals": {}} self.tasks: list[ImportedTask] = [] @@ -156,7 +158,7 @@ def __init__(self, dest: Path): self.undeclared: list[str] = [] self.duplicates: list[str] = [] self.problems: list[str] = [] - self.latest: dict[tuple[EntityKind, str], int | str] = {} + self.read: dict[tuple[EntityKind, str, int | None], int | str] = {} _check_destination(self.dest) def add_task(self, task: Task) -> str: @@ -165,21 +167,27 @@ def add_task(self, task: Task) -> str: path = f"tasks/{name}.json" steps = task.to_dict()["steps"] registry = get_task_step_registry() - classes = [registry[step["type"]] for step in steps] - outputs = {(site.ref.kind, _nfc(site.value)) - for step, cls in zip(steps, classes) for site in ref_sites(cls.entity_refs or (), step) - if site.ref.role is RefRole.OUTPUT and isinstance(site.value, str)} - for step, cls in zip(steps, classes): + found = [(step, cls, list(ref_sites(cls.entity_refs or (), step))) + for step in steps for cls in (registry[step["type"]],)] + sites = [site for _, _, step_sites in found for site in step_sites + if isinstance(site.value, str) and site.value] + outputs = self._outputs(sites) + deployed = {_nfc(site.value) for site in sites if site.ref.kind is EntityKind.ENV + and site.ref.role is RefRole.INPUT and site.version_key is not None} + for step, cls, step_sites in found: where = f"{path}: step {step.get('id')!r}" self._drop_secrets(step, where) if cls.entity_refs is None: self.undeclared.append(f"{where} ({cls.type})") - for site in ref_sites(cls.entity_refs or (), step): + for site in step_sites: if site.ref.role is RefRole.OUTPUT or not (isinstance(site.value, str) and site.value): continue - if (site.ref.kind, _nfc(site.value)) in outputs: - continue # written by a step of this task, which the bundle names under the task - self._pin(site, f"{where}: {site.path}") + output = outputs.get((site.ref.kind, _nfc(site.value))) + if output is not None: # written by a step of this task, which the bundle names under the task + if output != site.value: + site.rewrite(output) + continue + 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))) @@ -201,48 +209,81 @@ def write(self, description: str) -> BundleImport: self.files["README.md"] = _readme(description, imported) staging = self.dest.parent / f".{self.dest.name}.import-{uuid.uuid4().hex[:8]}" 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") - check_bundle(resolve_bundle(parse_bundle(staging))) - emptied = self.dest.is_dir() - if emptied: - self.dest.rmdir() # empty, as _check_destination found it try: - os.replace(staging, self.dest) - except OSError: - if emptied: - self.dest.mkdir(exist_ok=True) - raise + 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 _pin(self, site: Any, where: str) -> None: + def _move(self, staging: Path) -> None: + """Move ``staging`` to the destination, which is put back if it was an empty folder and the move fails.""" + 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 OSError as e: + if emptied: + with contextlib.suppress(OSError): + self.dest.mkdir(exist_ok=True) + _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 + + @staticmethod + def _outputs(sites: list[Any]) -> dict[tuple[EntityKind, str], str]: + """The name each id a step of the task writes keeps. One a bundle named under its own task when it wrote the + task takes its last segment, so this bundle names it under its task rather than writing more versions of the + other's; the step that writes it is rewritten to the name.""" + outputs = {} + for site in sites: + if site.ref.role is RefRole.OUTPUT: + name = site.value.rsplit("/", 1)[-1] if site.value.startswith(LOCAL_PREFIX) else site.value + outputs[site.ref.kind, _nfc(site.value)] = name + if name != site.value: + site.rewrite(name) + return outputs + + def _pin(self, site: Any, where: str, deployed: set[str]) -> None: kind, entity_id = site.ref.kind, site.value if site.version is not None: - self.pinned.append(Pinned(where, kind, entity_id, site.version, stored=True)) + read = self._read(kind, entity_id, site.version) + if isinstance(read, int): + self.pinned.append(Pinned(where, kind, entity_id, site.version, stored=True)) + else: + self.problems.append(f"{where}: {read}") elif site.version_key is not None: - latest = self._latest(kind, entity_id) - if isinstance(latest, int): - site.rewrite(entity_id, latest) - self.pinned.append(Pinned(where, kind, entity_id, latest, stored=False)) + read = self._read(kind, entity_id) + if isinstance(read, int): + site.rewrite(entity_id, read) + self.pinned.append(Pinned(where, kind, entity_id, read, stored=False)) else: - self.problems.append(f"{where}: {latest}") - elif kind is not EntityKind.ENV: # an env field with no version names an env the task deployed + self.problems.append(f"{where}: {read}") + 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: # an env field with no version names an env a step deploys, if one does + 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 _latest(self, kind: EntityKind, entity_id: str) -> int | str: - """The version of ``entity_id`` the store holds now, or why it can't be read.""" - if (kind, entity_id) not in self.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.latest[kind, entity_id] = STORE_GETTERS[kind](entity_id, None).version + self.read[kind, entity_id, version] = STORE_GETTERS[kind](entity_id, version).version except NotFoundError: - self.latest[kind, entity_id] = f"there is no {kind.value} {entity_id!r} in the store" + 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.latest[kind, entity_id] = f"{kind.value} {entity_id!r} can't be read ({type(e).__name__}: {e})" - return self.latest[kind, entity_id] + 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.""" diff --git a/src/agent_env/bundle/resolve.py b/src/agent_env/bundle/resolve.py index 35ca5fe8..560477ef 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,32 @@ 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 indexes of the steps step ``index`` runs after: those its edges name or, when it has none, every step + before it, followed transitively; the rule the task scheduler runs steps by. Linear in the task's size, as + the steps below the highest edge-less step reached are all upstream.""" + seen: set[int] = set() + below = 0 + 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/_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 index bdf05b67..10e5002e 100644 --- a/src/agent_env/cli/bundle.py +++ b/src/agent_env/cli/bundle.py @@ -7,6 +7,7 @@ import click from agent_env.bundle.importer import BundleImport, import_eval, import_task +from agent_env.cli.utils import interruptible @click.group() @@ -28,8 +29,9 @@ def init(folder: Path, task_id: str | None, eval_id: str | None, version: int | writes nothing.""" if (task_id is None) == (eval_id is None): raise click.UsageError("give one of --from-task or --from-eval") - imported = (import_task(task_id, folder, version=version) if task_id is not None - else import_eval(eval_id, folder, version=version)) + 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) 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/tst/unit/bundle/importer_test.py b/tst/unit/bundle/importer_test.py index df6d00bb..765972ef 100644 --- a/tst/unit/bundle/importer_test.py +++ b/tst/unit/bundle/importer_test.py @@ -3,6 +3,8 @@ 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 @@ -148,6 +150,49 @@ def test_the_folder_materializes_back_to_the_stored_steps(store): 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_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]]) @@ -242,12 +287,57 @@ def fail(src, dst): raise OSError("disk gone") monkeypatch.setattr(importer_module.os, "replace", fail) - with pytest.raises(OSError): + 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() +@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.""" 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'), From 5a3c2bf3671fca628b1ff88774630b0b295437a3 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 20:57:01 -0700 Subject: [PATCH 3/6] fix(bundle): number a renamed output whose name is taken, and put an empty destination back on an interrupt - An @local output that takes its last segment is numbered (-2, -3, ...) when another output or a store reference of its kind in the task already has that name, and the steps reading it follow. - An interrupt between emptying the destination and moving the import into place puts the empty folder back. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/bundle/importer.py | 35 ++++++++++++++++++++++++-------- tst/unit/bundle/importer_test.py | 31 ++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 9 deletions(-) diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py index 774b2852..94aacccf 100644 --- a/src/agent_env/bundle/importer.py +++ b/src/agent_env/bundle/importer.py @@ -222,17 +222,20 @@ def write(self, description: str) -> BundleImport: return imported def _move(self, staging: Path) -> None: - """Move ``staging`` to the destination, which is put back if it was an empty folder and the move fails.""" + """Move ``staging`` to the destination, which is put back if it was an empty folder and 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 OSError as e: + 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 @@ -240,14 +243,28 @@ def _move(self, staging: Path) -> None: def _outputs(sites: list[Any]) -> dict[tuple[EntityKind, str], str]: """The name each id a step of the task writes keeps. One a bundle named under its own task when it wrote the task takes its last segment, so this bundle names it under its task rather than writing more versions of the - other's; the step that writes it is rewritten to the name.""" - outputs = {} + other's, numbered when another output or a store reference of its kind has it; the step that writes it is + rewritten to the name.""" + written = {(site.ref.kind, _nfc(site.value)) for site in sites if site.ref.role is RefRole.OUTPUT} + taken = {(site.ref.kind, _nfc(site.value)) for site in sites + if (site.ref.kind, _nfc(site.value)) not in written or not site.value.startswith(LOCAL_PREFIX)} + outputs: dict[tuple[EntityKind, str], str] = {} for site in sites: - if site.ref.role is RefRole.OUTPUT: - name = site.value.rsplit("/", 1)[-1] if site.value.startswith(LOCAL_PREFIX) else site.value - outputs[site.ref.kind, _nfc(site.value)] = name - if name != site.value: - site.rewrite(name) + key = (site.ref.kind, _nfc(site.value)) + if site.ref.role is not RefRole.OUTPUT or key in outputs: + 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))) + outputs[key] = name + for site in sites: + if site.ref.role is RefRole.OUTPUT and outputs[site.ref.kind, _nfc(site.value)] != site.value: + site.rewrite(outputs[site.ref.kind, _nfc(site.value)]) return outputs def _pin(self, site: Any, where: str, deployed: set[str]) -> None: diff --git a/tst/unit/bundle/importer_test.py b/tst/unit/bundle/importer_test.py index 765972ef..a61cf901 100644 --- a/tst/unit/bundle/importer_test.py +++ b/tst/unit/bundle/importer_test.py @@ -161,6 +161,23 @@ def test_an_id_a_bundle_named_under_its_task_is_renamed_so_the_import_names_it_u 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"}]) @@ -294,6 +311,20 @@ def fail(src, dst): 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" From bc5e2d0eb0aa15ab88977829a9f5d68328b08215 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 21:51:30 -0700 Subject: [PATCH 4/6] feat(bundle): leave out values shaped like known credentials, wherever they sit in a step An env var whose value is an API key (sk-/sk_), a GitHub, AWS, Slack, Google or Hugging Face token, a JWT or a private key is dropped like one with a secret-looking name, and any other string has such a value masked, as a URL's credentials already are. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/bundle/importer.py | 42 ++++++++++++++++++++++++----- tst/unit/bundle/importer_test.py | 45 ++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 7 deletions(-) diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py index 94aacccf..3941551d 100644 --- a/src/agent_env/bundle/importer.py +++ b/src/agent_env/bundle/importer.py @@ -44,6 +44,18 @@ _MAX_NAME_BYTES = 200 # 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. _URL_CREDENTIALS = re.compile(r"(?P[a-zA-Z][\w+.-]*://)(?:[^/@\s]*:[^/@\s]*|[^/@\s:]{16,})@") +# Credentials by the shape their issuers give them, wherever they sit. +_TOKENS = ( + ("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)), +) @dataclass(frozen=True) @@ -315,21 +327,24 @@ def _default_agents(self, step: dict, where: str) -> None: f"the step {what}, so a run deploys the configured default")) def _drop_secrets(self, value: Any, where: str, path: str = "") -> None: - """Drop each variable of an env-var map whose name looks like a secret or whose value holds a URL's - credentials, and mask the credentials of a URL anywhere else, in place.""" + """Drop each variable of an env-var map whose name looks like a secret or whose value holds a credential, and + mask each credential anywhere else, in place.""" 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) - if secret or (isinstance(setting, str) and _URL_CREDENTIALS.search(setting)): + 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 "its value holds a URL's credentials" + why = "its name looks like a secret" if secret else f"its value holds {found[0]}" self.dropped.append(Dropped(f"{where}: {at}.{name}", f"dropped: {why}")) - elif isinstance(item, str) and _URL_CREDENTIALS.search(item): - value[key] = _URL_CREDENTIALS.sub(r"\g***@", item) - self.dropped.append(Dropped(f"{where}: {at}", "masked: a URL's credentials")) + elif isinstance(item, str): + masked, found = _masked(item) + if found: + value[key] = masked + self.dropped.append(Dropped(f"{where}: {at}", f"masked: {', '.join(found)}")) else: self._drop_secrets(item, where, at) @@ -380,5 +395,18 @@ def _valid_name(text: str) -> str: return name if any(char.isalnum() for char in name) and segment_problem(name) is None else "" +def _masked(text: str) -> tuple[str, list[str]]: + """``text`` with each credential in it replaced by ``***``, and what kinds it held.""" + found = [] + if _URL_CREDENTIALS.search(text): + text = _URL_CREDENTIALS.sub(r"\g***@", text) + found.append("a URL's credentials") + for what, pattern in _TOKENS: + if pattern.search(text): + text = pattern.sub("***", text) + found.append(what) + return text, found + + def _nfc(text: str) -> str: return unicodedata.normalize("NFC", text) diff --git a/tst/unit/bundle/importer_test.py b/tst/unit/bundle/importer_test.py index a61cf901..d123774b 100644 --- a/tst/unit/bundle/importer_test.py +++ b/tst/unit/bundle/importer_test.py @@ -128,6 +128,51 @@ def test_env_vars_that_look_like_secrets_are_left_out_and_a_urls_credentials_mas 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() From 4b58ec5950f9918f771abb30fc6b87813d154631 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 22:30:47 -0700 Subject: [PATCH 5/6] refactor(bundle): a secrets module, outputs named in one pass, and integration tests for versions, re-imports and secrets - The secret scrubbing moves to bundle/_secrets.py as one function, scrub(), so the env and agent tomls an import writes next can share it. - Outputs are named by a function that only names them, and one loop renames each output and every read of it. - A pin the task already holds and a pin to the latest share one store read. - The envs a task deploys are the ones its deploy_env steps name, as preflight reads them. - The imported eval is an ImportedEval, and the CLI prints the report without filtering titles. - Docstrings are one line. - New tests: `--version`, an eval naming a task the store doesn't hold, a task a bundle ran imported into another folder and run there, secrets left out of the task a run writes, and an output a bundle's task wrote written under the importing folder. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/bundle/_secrets.py | 62 +++++ src/agent_env/bundle/_toml.py | 8 +- src/agent_env/bundle/importer.py | 217 +++++++----------- src/agent_env/bundle/parse.py | 6 +- src/agent_env/bundle/resolve.py | 6 +- src/agent_env/cli/bundle.py | 8 +- tst/integration/cli/bundle_init_test.py | 67 +++++- .../cli/run_bundle_built_envs_test.py | 44 +++- tst/unit/bundle/importer_test.py | 25 +- 9 files changed, 272 insertions(+), 171 deletions(-) create mode 100644 src/agent_env/bundle/_secrets.py diff --git a/src/agent_env/bundle/_secrets.py b/src/agent_env/bundle/_secrets.py new file mode 100644 index 00000000..9194cfbf --- /dev/null +++ b/src/agent_env/bundle/_secrets.py @@ -0,0 +1,62 @@ +"""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") +# 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. +_URL_CREDENTIALS = re.compile(r"(?P[a-zA-Z][\w+.-]*://)(?:[^/@\s]*:[^/@\s]*|[^/@\s:]{16,})@") +# Credentials by the shape their issuers give them, wherever they sit. +_TOKENS = ( + ("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 = [] + if _URL_CREDENTIALS.search(text): + text = _URL_CREDENTIALS.sub(r"\g***@", text) + found.append("a URL's credentials") + for what, pattern in _TOKENS: + if pattern.search(text): + text = pattern.sub("***", text) + found.append(what) + return text, found diff --git a/src/agent_env/bundle/_toml.py b/src/agent_env/bundle/_toml.py index 6af74ce7..def8ad93 100644 --- a/src/agent_env/bundle/_toml.py +++ b/src/agent_env/bundle/_toml.py @@ -1,5 +1,4 @@ -"""TOML for the tomls a bundle import writes, since ``tomllib`` only reads: strings, numbers, booleans, lists and -inline tables, and tables of those under named top-level keys.""" +"""A TOML writer for the tomls a bundle import writes, since ``tomllib`` only reads.""" from __future__ import annotations @@ -13,8 +12,7 @@ def dumps(document: dict[str, Any], *, tables: Collection[str] = ()) -> str: - """``document`` as TOML. A value under a key in ``tables`` is written as a ``[table]`` after the other keys; - every other dict is an inline table. Raises ValueError for a value TOML can't hold, such as None.""" + """``document`` as TOML, keys in ``tables`` as ``[table]``s after the rest; ValueError for what TOML can't hold.""" lines = [f"{_key(key)} = {_value(value)}" for key, value in document.items() if key not in tables] for key in (key for key in document if key in tables): lines += ["", f"[{_key(key)}]", *(f"{_key(k)} = {_value(v)}" for k, v in document[key].items())] @@ -26,7 +24,7 @@ def _key(key: str) -> str: def _value(value: Any, *, top: bool = True) -> str: - """``value`` as TOML; a list of several items spans lines only at the top, since an inline table can't.""" + """``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): diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py index 3941551d..bc6b8534 100644 --- a/src/agent_env/bundle/importer.py +++ b/src/agent_env/bundle/importer.py @@ -1,18 +1,10 @@ -"""Import a stored task, or an eval and its tasks, as a bundle folder that ``agent-env run`` runs. - -Each task is written as ``task get`` prints it: the stored steps, read through the current step models. Every store -reference a step makes is pinned at the version the store holds now, so the folder runs what the store held when it -was imported; an id a step of the task writes stays unpinned, since the run writes it again. Values that look like -secrets are dropped. Nothing is written to a store, and the folder is written beside its destination, checked the way -``agent-env run`` checks a bundle before it reads a store, then moved into place, so a failed import leaves nothing. -""" +"""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 re import shutil import unicodedata import uuid @@ -20,47 +12,31 @@ from datetime import datetime, timezone from importlib.metadata import version as package_version from pathlib import Path -from typing import Any 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.config.describe import is_secret_key -from agent_env.entity_refs import EntityKind, RefRole, ref_sites +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 -# The step fields that map environment variable names to values, where a credential is most often set. -_ENV_VAR_MAPS = ("env_vars", "default_env_vars") # 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 -# 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. -_URL_CREDENTIALS = re.compile(r"(?P[a-zA-Z][\w+.-]*://)(?:[^/@\s]*:[^/@\s]*|[^/@\s:]{16,})@") -# Credentials by the shape their issuers give them, wherever they sit. -_TOKENS = ( - ("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)), -) @dataclass(frozen=True) class Pinned: - """A store reference written with a version: the one the store held at import, or the task's own pin.""" + """A store reference written with a version: the store's at import, or the task's own pin.""" where: str kind: EntityKind @@ -95,6 +71,13 @@ class ImportedTask: 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.""" @@ -102,31 +85,27 @@ class BundleImport: root: Path files: tuple[str, ...] tasks: tuple[ImportedTask, ...] - eval: tuple[str, str, int] | None # (path, id, version) of the eval, when one was imported + 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 sections(self) -> list[tuple[str, list[str]]]: - """The import's report, as titled lists of lines; a section with nothing to say is left out.""" - sections = [ - ("Tasks", [f"{task.path}: task {task.id} v{task.version}" for task in self.tasks]), - ("Eval", [f"{self.eval[0]}: eval {self.eval[1]} v{self.eval[2]}"] if self.eval else []), + 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]), - ] - return [(title, lines) for title, lines in sections if lines] + ]) 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``, a folder that is - empty or isn't there yet in one that is. Raises BundleError listing every problem, writing nothing.""" + """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) @@ -134,8 +113,7 @@ def import_task(task_id: str, dest: Path | str, *, version: int | None = None) - 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 every task it names as a bundle at - ``dest``, as ``import_task`` writes a task. Raises BundleError listing every problem, writing nothing.""" + """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 = [], {} @@ -163,7 +141,7 @@ def __init__(self, dest: Path): self.files: dict[str, str] = {} self.taken: dict[str, dict[str, str]] = {"tasks": {}, "evals": {}} self.tasks: list[ImportedTask] = [] - self.eval: tuple[str, str, int] | None = None + self.eval: ImportedEval | None = None self.pinned: list[Pinned] = [] self.floating: list[Floating] = [] self.dropped: list[Dropped] = [] @@ -179,27 +157,26 @@ def add_task(self, task: Task) -> str: path = f"tasks/{name}.json" steps = task.to_dict()["steps"] registry = get_task_step_registry() - found = [(step, cls, list(ref_sites(cls.entity_refs or (), step))) - for step in steps for cls in (registry[step["type"]],)] - sites = [site for _, _, step_sites in found for site in step_sites - if isinstance(site.value, str) and site.value] - outputs = self._outputs(sites) - deployed = {_nfc(site.value) for site in sites if site.ref.kind is EntityKind.ENV - and site.ref.role is RefRole.INPUT and site.version_key is not None} + 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._drop_secrets(step, where) + 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 site.ref.role is RefRole.OUTPUT or not (isinstance(site.value, str) and site.value): - continue - output = outputs.get((site.ref.kind, _nfc(site.value))) - if output is not None: # written by a step of this task, which the bundle names under the task - if output != site.value: - site.rewrite(output) - continue - self._pin(site, f"{where}: {site.path}", deployed) + 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))) @@ -208,7 +185,7 @@ def add_task(self, task: Task) -> str: 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 = (path, stored.id, stored.version) + 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.""" @@ -234,8 +211,7 @@ def write(self, description: str) -> BundleImport: return imported def _move(self, staging: Path) -> None: - """Move ``staging`` to the destination, which is put back if it was an empty folder and the move fails or is - interrupted.""" + """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(): @@ -251,53 +227,20 @@ def _move(self, staging: Path) -> None: _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 - @staticmethod - def _outputs(sites: list[Any]) -> dict[tuple[EntityKind, str], str]: - """The name each id a step of the task writes keeps. One a bundle named under its own task when it wrote the - task takes its last segment, so this bundle names it under its task rather than writing more versions of the - other's, numbered when another output or a store reference of its kind has it; the step that writes it is - rewritten to the name.""" - written = {(site.ref.kind, _nfc(site.value)) for site in sites if site.ref.role is RefRole.OUTPUT} - taken = {(site.ref.kind, _nfc(site.value)) for site in sites - if (site.ref.kind, _nfc(site.value)) not in written or not site.value.startswith(LOCAL_PREFIX)} - outputs: dict[tuple[EntityKind, str], str] = {} - for site in sites: - key = (site.ref.kind, _nfc(site.value)) - if site.ref.role is not RefRole.OUTPUT or key in outputs: - 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))) - outputs[key] = name - for site in sites: - if site.ref.role is RefRole.OUTPUT and outputs[site.ref.kind, _nfc(site.value)] != site.value: - site.rewrite(outputs[site.ref.kind, _nfc(site.value)]) - return outputs - - def _pin(self, site: Any, where: str, deployed: set[str]) -> 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: + if site.version is not None or site.version_key is not None: read = self._read(kind, entity_id, site.version) - if isinstance(read, int): - self.pinned.append(Pinned(where, kind, entity_id, site.version, stored=True)) - else: + if isinstance(read, str): self.problems.append(f"{where}: {read}") - elif site.version_key is not None: - read = self._read(kind, entity_id) - if isinstance(read, int): + return + if site.version is None: site.rewrite(entity_id, read) - self.pinned.append(Pinned(where, kind, entity_id, read, stored=False)) - else: - self.problems.append(f"{where}: {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: # an env field with no version names an env a step deploys, if one does + 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")) @@ -326,31 +269,8 @@ def _default_agents(self, step: dict, where: str) -> None: 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 _drop_secrets(self, value: Any, where: str, path: str = "") -> None: - """Drop each variable of an env-var map whose name looks like a secret or whose value holds a credential, and - mask each credential anywhere else, in place.""" - 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]}" - self.dropped.append(Dropped(f"{where}: {at}.{name}", f"dropped: {why}")) - elif isinstance(item, str): - masked, found = _masked(item) - if found: - value[key] = masked - self.dropped.append(Dropped(f"{where}: {at}", f"masked: {', '.join(found)}")) - else: - self._drop_secrets(item, where, at) - def _name(self, kind: str, entity_id: str) -> str: - """A name for ``entity_id`` in the bundle's ``kind`` folder: its last segment, made valid, and numbered when - another entry of the kind took it, as names compare caselessly.""" + """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]: @@ -360,17 +280,44 @@ def _name(self, kind: str, entity_id: str) -> str: 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: what it was imported from, its first line, then the import's report.""" + """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.", ""] - for title, items in imported.sections(): + 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()): @@ -385,8 +332,7 @@ def _check_destination(dest: Path) -> None: def _valid_name(text: str) -> str: - """``text`` as one bundle name: each character an id can't hold becomes ``-``, and dots that would hide it or - spaces at its edges go. Empty when no letter or number is left.""" + """``text`` as one bundle name, each character an id can't hold made ``-``; empty with no letter or digit left.""" text = unicodedata.normalize("NFC", text) name = "".join(char if char.isalnum() or unicodedata.category(char).startswith("M") or char in " ._-~@()+,&" else "-" for char in text).strip(" ").lstrip(".") @@ -395,17 +341,8 @@ def _valid_name(text: str) -> str: return name if any(char.isalnum() for char in name) and segment_problem(name) is None else "" -def _masked(text: str) -> tuple[str, list[str]]: - """``text`` with each credential in it replaced by ``***``, and what kinds it held.""" - found = [] - if _URL_CREDENTIALS.search(text): - text = _URL_CREDENTIALS.sub(r"\g***@", text) - found.append("a URL's credentials") - for what, pattern in _TOKENS: - if pattern.search(text): - text = pattern.sub("***", text) - found.append(what) - return text, found +def _key(site: RefSite) -> tuple[EntityKind, str]: + return site.ref.kind, _nfc(site.value) def _nfc(text: str) -> str: diff --git a/src/agent_env/bundle/parse.py b/src/agent_env/bundle/parse.py index 379b7bc4..71a1f654 100644 --- a/src/agent_env/bundle/parse.py +++ b/src/agent_env/bundle/parse.py @@ -525,15 +525,13 @@ def _rel(self, path) -> str: def root_problems(root: Path) -> list[str]: - """Why a bundle at ``root``, a folder that needn't exist yet in one that does, would have no id; empty when it - would.""" + """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 of a bundle at the canonical path ``root``, and each part of the path that can't be one - segment of an id.""" + """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) diff --git a/src/agent_env/bundle/resolve.py b/src/agent_env/bundle/resolve.py index 560477ef..53899ece 100644 --- a/src/agent_env/bundle/resolve.py +++ b/src/agent_env/bundle/resolve.py @@ -466,11 +466,9 @@ def _step_edges(steps: list[dict], depends: list[list[str] | None]) -> list[list def _upstream(edges: list[list[int] | None], index: int) -> set[int]: - """The indexes of the steps step ``index`` runs after: those its edges name or, when it has none, every step - before it, followed transitively; the rule the task scheduler runs steps by. Linear in the task's size, as - the steps below the highest edge-less step reached are all upstream.""" + """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 + below = 0 # each step under it is already seen, which keeps the walk linear frontier = [index] while frontier: current = frontier.pop() diff --git a/src/agent_env/cli/bundle.py b/src/agent_env/cli/bundle.py index 10e5002e..11d509d9 100644 --- a/src/agent_env/cli/bundle.py +++ b/src/agent_env/cli/bundle.py @@ -36,14 +36,12 @@ def init(folder: Path, task_id: str | None, eval_id: str | None, version: int | def _report(imported: BundleImport, folder: Path) -> None: - what = f"eval {imported.eval[1]} v{imported.eval[2]} and its tasks" if imported.eval else ( - f"task {imported.tasks[0].id} v{imported.tasks[0].version}") + 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.sections(): - if title in ("Tasks", "Eval"): - continue + for title, lines in imported.report(): click.echo(f"{title}:") for line in lines: click.echo(f" {line}") diff --git a/tst/integration/cli/bundle_init_test.py b/tst/integration/cli/bundle_init_test.py index 99043206..84ad1304 100644 --- a/tst/integration/cli/bundle_init_test.py +++ b/tst/integration/cli/bundle_init_test.py @@ -1,9 +1,11 @@ -"""``agent-env bundle init`` on a task and an eval held in a store under bare ids, as a registry holds them, 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. Needs `bash` and the usual shell tools on PATH.""" +"""``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 @@ -12,11 +14,13 @@ 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 @@ -39,8 +43,7 @@ def registry(monkeypatch, tmp_path): reset_artifact_store() greeting = HELLO / "artifacts" / "greeting" FileArtifactUniverse.put_bundled("greeting", files={path.name: path for path in greeting.iterdir()}) - steps = json.loads((HELLO / "tasks" / "hello.json").read_text()) - Task.put(id="greet", steps=[build_step(step) for step in steps]) + 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: @@ -52,6 +55,10 @@ def registry(monkeypatch, tmp_path): 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 @@ -73,6 +80,56 @@ def test_an_imported_task_runs_reruns_writing_nothing_and_keeps_its_pins_when_th 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") diff --git a/tst/integration/cli/run_bundle_built_envs_test.py b/tst/integration/cli/run_bundle_built_envs_test.py index 89cb93a8..26048db3 100644 --- a/tst/integration/cli/run_bundle_built_envs_test.py +++ b/tst/integration/cli/run_bundle_built_envs_test.py @@ -7,7 +7,8 @@ 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, and a task held under bare ids, as a registry -holds it, imports as a bundle that runs through the gateway. Needs a Docker daemon and the local 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 @@ -19,7 +20,7 @@ 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 @@ -28,6 +29,7 @@ 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 @@ -560,15 +562,22 @@ def test_the_shipped_echo_mcp_example_runs_by_name_and_a_rerun_writes_nothing(st assert "tasks/smoke.json v1: passed (echoes: 1)" in second.output, second.output -def test_a_stored_mcp_task_imports_as_a_bundle_that_runs_through_the_gateway(state): - echo = REPO / "src" / "agent_env" / "examples" / "echo-mcp" +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) + 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()) + 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", @@ -584,3 +593,26 @@ def test_a_stored_mcp_task_imports_as_a_bundle_that_runs_through_the_gateway(sta 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 index d123774b..b15bcb1e 100644 --- a/tst/unit/bundle/importer_test.py +++ b/tst/unit/bundle/importer_test.py @@ -17,7 +17,7 @@ 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 import_eval, import_task +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 @@ -269,11 +269,32 @@ def test_an_eval_imports_with_each_task_it_names_once_under_names_taken_from_the 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 == ("evals/regression.toml", "nightly/regression", 1) + 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), From dbcd58d362ab2f19098e6379f5addff612aef51e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 15 Oct 2026 08:06:32 -0700 Subject: [PATCH 6/6] refactor(bundle): one credentials table, keyword BundleImport, and no unused TOML tables - The URL and token patterns are one table of (what, pattern, mask), each masked in one subn pass. - `_toml.dumps` drops its unused `tables` parameter; every dict is an inline table. - `BundleImport` is built with keyword arguments, and name and eval bookkeeping that only checks membership uses sets. - `_valid_name` uses the importer's own `_nfc`. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/bundle/_secrets.py | 34 ++++++++++++++--------------- src/agent_env/bundle/_toml.py | 10 +++------ src/agent_env/bundle/importer.py | 19 ++++++++-------- tst/unit/bundle/toml_writer_test.py | 8 +------ 4 files changed, 30 insertions(+), 41 deletions(-) diff --git a/src/agent_env/bundle/_secrets.py b/src/agent_env/bundle/_secrets.py index 9194cfbf..22f1590c 100644 --- a/src/agent_env/bundle/_secrets.py +++ b/src/agent_env/bundle/_secrets.py @@ -9,19 +9,20 @@ # The fields that map environment variable names to values, where a credential is most often set. _ENV_VAR_MAPS = ("env_vars", "default_env_vars") -# 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. -_URL_CREDENTIALS = re.compile(r"(?P[a-zA-Z][\w+.-]*://)(?:[^/@\s]*:[^/@\s]*|[^/@\s:]{16,})@") -# Credentials by the shape their issuers give them, wherever they sit. -_TOKENS = ( - ("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,}")), +# 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)), + r"(?:.*?-----END (?:[A-Z]+ )*PRIVATE KEY-----|.*)", re.DOTALL), "***"), ) @@ -52,11 +53,8 @@ def scrub(value: Any, path: str = "") -> list[tuple[str, str]]: def _masked(text: str) -> tuple[str, list[str]]: """``text`` with each credential replaced by ``***``, and the kinds it held.""" found = [] - if _URL_CREDENTIALS.search(text): - text = _URL_CREDENTIALS.sub(r"\g***@", text) - found.append("a URL's credentials") - for what, pattern in _TOKENS: - if pattern.search(text): - text = pattern.sub("***", text) + 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 index def8ad93..d06a33a5 100644 --- a/src/agent_env/bundle/_toml.py +++ b/src/agent_env/bundle/_toml.py @@ -4,19 +4,15 @@ import math import re -from collections.abc import Collection 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], *, tables: Collection[str] = ()) -> str: - """``document`` as TOML, keys in ``tables`` as ``[table]``s after the rest; ValueError for what TOML can't hold.""" - lines = [f"{_key(key)} = {_value(value)}" for key, value in document.items() if key not in tables] - for key in (key for key in document if key in tables): - lines += ["", f"[{_key(key)}]", *(f"{_key(k)} = {_value(v)}" for k, v in document[key].items())] - return "\n".join(lines).lstrip("\n") + "\n" if lines else "" +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: diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py index bc6b8534..48ec9686 100644 --- a/src/agent_env/bundle/importer.py +++ b/src/agent_env/bundle/importer.py @@ -116,7 +116,7 @@ def import_eval(eval_id: str, dest: Path | str, *, version: int | None = None) - """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 = [], {} + names, seen = [], set() for entry in stored.tasks: try: task = Task.get(entry.task_id, entry.task_version) @@ -128,8 +128,8 @@ def import_eval(eval_id: str, dest: Path | str, *, version: int | None = None) - if (task.id, task.version) in seen: importer.duplicates.append(f"task {task.id} v{task.version}") continue - seen[task.id, task.version] = importer.add_task(task) - names.append(seen[task.id, task.version]) + 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.") @@ -139,7 +139,7 @@ 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, dict[str, str]] = {"tasks": {}, "evals": {}} + 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] = [] @@ -192,9 +192,10 @@ def write(self, description: str) -> BundleImport: if self.problems: raise BundleError(self.problems) files = sorted([*self.files, "README.md"]) - imported = BundleImport(self.dest, tuple(files), tuple(self.tasks), self.eval, tuple(self.pinned), - tuple(self.floating), tuple(self.dropped), tuple(self.undeclared), - tuple(self.duplicates)) + 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: @@ -276,7 +277,7 @@ def _name(self, kind: str, entity_id: str) -> str: while fold(name) in self.taken[kind]: n += 1 name = f"{base}-{n}" - self.taken[kind][fold(name)] = name + self.taken[kind].add(fold(name)) return name @@ -333,7 +334,7 @@ def _check_destination(dest: Path) -> None: 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 = unicodedata.normalize("NFC", text) + 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: diff --git a/tst/unit/bundle/toml_writer_test.py b/tst/unit/bundle/toml_writer_test.py index b87973e3..6d379ab9 100644 --- a/tst/unit/bundle/toml_writer_test.py +++ b/tst/unit/bundle/toml_writer_test.py @@ -23,13 +23,7 @@ def test_what_it_writes_reads_back_the_same(): - assert tomllib.loads(dumps(DOCUMENT, tables=("metadata",))) == DOCUMENT - - -def test_a_table_key_is_written_as_a_table_after_the_other_keys(): - text = dumps({"metadata": {"default_model": "m"}, "image": "x"}, tables=("metadata",)) - - assert text == 'image = "x"\n\n[metadata]\ndefault_model = "m"\n' + assert tomllib.loads(dumps(DOCUMENT)) == DOCUMENT def test_a_list_of_several_spans_lines_only_at_the_top():