ci(slow tier): keep the computer tests' resolved packages between modules and runs

.github/scripts/prune_registry.py

@@ -0,0 +1,98 @@
+"""Keep only the images of a local registry's repository that a run used, before the registry's storage is cached.
+
+``prune`` drops each tag the registry's log has no manifest request for, then each manifest no remaining tag reaches,
+following an image index to the manifests it lists. The registry's ``garbage-collect`` then frees the layers only those
+manifests used; with ``--delete-untagged`` it would also delete a kept index's untagged platform manifests. ``links``,
+run after it, drops the repository's links to layers it freed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+
+_V2 = Path("docker") / "registry" / "v2"
+
+
+def main(argv: list[str]) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    pruning = commands.add_parser("prune", help="drop the tags and manifests the run didn't use")
+    linking = commands.add_parser("links", help="drop links to layers garbage-collect freed")
+    for command in (pruning, linking):
+        command.add_argument("storage", type=Path, help="the registry's storage root")
+        command.add_argument("repository")
+    pruning.add_argument("log", type=Path, help="the registry's log of the run")
+    args = parser.parse_args(argv)
+    if args.command == "prune":
+        kept, dropped = prune(args.storage, args.repository, used_tags(args.log.read_text(errors="replace"),
+                                                                        args.repository))
+        print(f"kept {kept} tags of {args.repository}, dropped {dropped}")
+    else:
+        print(f"dropped {drop_freed_links(args.storage, args.repository)} links to freed layers")
+    return 0
+
+
+def used_tags(log: str, repository: str) -> set[str]:
+    """The tags the log names a request for a manifest of ``repository`` by."""
+    return set(re.findall(rf"/v2/{re.escape(repository)}/manifests/([A-Za-z0-9_][A-Za-z0-9_.-]{{0,127}})(?=[\s\"?])",
+                          log))
+
+
+def prune(storage: Path, repository: str, used: set[str]) -> tuple[int, int]:
+    """Drop the tags of ``repository`` not in ``used``, then the manifests no remaining tag reaches; how many tags it
+    kept and dropped."""
+    manifests = storage / _V2 / "repositories" / repository / "_manifests"
+    tags = _children(manifests / "tags")
+    for tag in tags:
+        if tag.name not in used:
+            shutil.rmtree(tag)
+    kept = [tag for tag in tags if tag.name in used]
+    reached = _reached(storage, [(tag / "current" / "link").read_text().strip() for tag in kept])
+    for revision in _children(manifests / "revisions" / "sha256"):
+        if f"sha256:{revision.name}" not in reached:
+            shutil.rmtree(revision)
+    return len(kept), len(tags) - len(kept)
+
+
+def drop_freed_links(storage: Path, repository: str) -> int:
+    """Drop each link of ``repository`` to a layer that's no longer stored; how many."""
+    dropped = 0
+    for layer in _children(storage / _V2 / "repositories" / repository / "_layers" / "sha256"):
+        if not _blob(storage, f"sha256:{layer.name}").is_file():
+            shutil.rmtree(layer)
+            dropped += 1
+    return dropped
+
+
+def _reached(storage: Path, digests: list[str]) -> set[str]:
+    """``digests`` and every manifest an image index among them lists, recursively."""
+    reached, pending = set(), list(digests)
+    while pending:
+        digest = pending.pop()
+        if digest in reached:
+            continue
+        reached.add(digest)
+        try:
+            manifest = json.loads(_blob(storage, digest).read_bytes())
+        except (FileNotFoundError, ValueError):
+            continue
+        pending += [listed["digest"] for listed in manifest.get("manifests", []) if "digest" in listed]
+    return reached
+
+
+def _children(folder: Path) -> list[Path]:
+    return [child for child in folder.iterdir() if child.is_dir()] if folder.is_dir() else []
+
+
+def _blob(storage: Path, digest: str) -> Path:
+    hex_digest = digest.removeprefix("sha256:")
+    return storage / _V2 / "blobs" / "sha256" / hex_digest[:2] / hex_digest / "data"
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

.github/workflows/local-backends.yml

@@ -169,17 +169,21 @@ jobs:
       - name: Pull the base images from mirror.gcr.io
         run: python3 .github/scripts/mirror_base_images.py src tst
       # A computer image or disk is tagged with a key of everything it is built from, and the composer reuses one the
-      # image store already holds for its key. The registry keeps its storage on the runner, restored from the computer
-      # images main saved latest this week, so a run reuses them rather than building them again. A shard that runs no
-      # computer tests misses.
+      # image store already holds for its key. That key includes the lock of the packages the image installs, which the
+      # composer resolves from Ubuntu's archive, PyPI and download URLs unless the store holds it. The registry keeps
+      # its storage on the runner, and the computer tests share their locks and files in the packages folder
+      # (tst/util/package_cache.py); both are restored from what main saved latest this week, so a run neither resolves
+      # nor builds again what hasn't changed. A shard that runs no computer tests misses.
       - name: Week of the computer images' cache
         id: week
         run: echo "week=$(date -u +%G-W%V)" >> "$GITHUB_OUTPUT"
-      - name: Restore the computer images
+      - name: Restore the computer images and packages
         id: computers
         uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
         with:
-          path: ${{ runner.temp }}/registry
+          path: |
+            ${{ runner.temp }}/registry
+            ${{ runner.temp }}/packages
           key: computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}
           restore-keys: computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}-
       - name: Start the local registry
@@ -195,35 +199,49 @@ jobs:
         id: tests
         env:
           AGENTENV_TEST_SHARD: ${{ matrix.shard }}/12
+          AGENTENV_TEST_PACKAGES: ${{ runner.temp }}/packages
         run: |
           .venv/bin/python -m pytest -m int_test_slow -rA --timeout=1200 --junitxml=junit-integration-slow.xml \
             tst/integration --ignore=tst/integration/env/gateway/gateway_test.py
-      # Only the computer images are kept, and only main saves them, once its tests pass, so pull requests reuse main's
-      # without filling the repository's cache. The entry is named by the image keys the registry holds, so a run that
-      # built any new image saves a new one. The garbage collection frees what only the deleted repositories
-      # referenced; with --delete-untagged it would also delete the untagged platform manifests of an image index, and
-      # their layers.
-      - name: Keep only the computer images
+      # Only the computer images this run used are kept, with the packages, and only main saves them, once its tests
+      # pass, so pull requests reuse main's without filling the repository's cache. An image is used when the registry's
+      # log has a request for it by its tag (prune_registry.py), so what a week's earlier runs built and this one no
+      # longer asks for is dropped. The garbage collection frees what only the dropped tags and repositories referenced;
+      # with --delete-untagged it would also delete the untagged platform manifests of an image index, and their layers.
+      # The entry is named by the image keys the registry holds and the locks the packages hold, so a run that built a
+      # new image or resolved a new lock saves a new one.
+      - name: Keep only the computer images this run used
         id: kept
         if: always()
         run: |
           repositories="$RUNNER_TEMP/registry/docker/registry/v2/repositories"
           if [ ! -d "$repositories/agentenv-computers" ]; then echo "computers=false" >> "$GITHUB_OUTPUT"; exit 0; fi
+          docker logs registry > "$RUNNER_TEMP/registry.log" 2>&1
           docker stop registry
           sudo find "$repositories" -mindepth 1 -maxdepth 1 ! -name agentenv-computers -exec rm -rf {} +
+          sudo python3 .github/scripts/prune_registry.py prune "$RUNNER_TEMP/registry" agentenv-computers \
+            "$RUNNER_TEMP/registry.log"
           docker run --rm --volume "$RUNNER_TEMP/registry:/var/lib/registry" mirror.gcr.io/library/registry:2 \
             garbage-collect /etc/docker/registry/config.yml > /dev/null
-          du -sh "$RUNNER_TEMP/registry"
-          tags=$(sudo ls "$repositories/agentenv-computers/_manifests/tags" | sort | sha256sum | cut -c1-16)
-          echo "key=computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}-$tags" >> "$GITHUB_OUTPUT"
+          sudo python3 .github/scripts/prune_registry.py links "$RUNNER_TEMP/registry" agentenv-computers
+          mkdir -p "$RUNNER_TEMP/packages"
+          du -sh "$RUNNER_TEMP/registry" "$RUNNER_TEMP/packages"
+          if [ -z "$(sudo ls "$repositories/agentenv-computers/_manifests/tags")" ]; then
+            echo "computers=false" >> "$GITHUB_OUTPUT"; exit 0
+          fi
+          key=$({ sudo ls "$repositories/agentenv-computers/_manifests/tags"
+                  ls "$RUNNER_TEMP/packages/agentenv-packages/locks" 2>/dev/null; } | sort | sha256sum | cut -c1-16)
+          echo "key=computers-${{ matrix.shard }}-${{ steps.week.outputs.week }}-$key" >> "$GITHUB_OUTPUT"
           echo "computers=true" >> "$GITHUB_OUTPUT"
-      - name: Save the computer images
+      - name: Save the computer images and packages
         if: >-
           always() && github.event_name == 'push' && steps.tests.outcome == 'success'
           && steps.kept.outputs.computers == 'true' && steps.kept.outputs.key != steps.computers.outputs.cache-matched-key
         uses: actions/cache/save@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
         with:
-          path: ${{ runner.temp }}/registry
+          path: |
+            ${{ runner.temp }}/registry
+            ${{ runner.temp }}/packages
           key: ${{ steps.kept.outputs.key }}
       - name: Skips must be declared capability gaps
         if: always()

AGENTS.md

@@ -389,6 +389,11 @@ orchestrators call it the same way. What callers rely on:
   `tst/integration/slow_durations.json`. The unit and installer tiers are split the same way. When
   the shards' times drift apart, rewrite the table from all of a `main` run's `junit-*` artifacts with
   `.github/scripts/slow_test_durations.py`.
+- The computer tests compose in stores of their own, but share the package locks and files a compose
+  resolves (`tst/util/package_cache.py`) through `AGENTENV_TEST_PACKAGES`, a folder for the run unless
+  it's set. The slow tier caches that folder with the local registry's computer images, so a run
+  resolves from Ubuntu's archive, PyPI or a download URL only what changed since `main` last saved
+  them.
 - `tst/manual/` and `tst/benchmarks/` run only by hand, as each module's docstring says; no CI job
   collects them.
 - `tst/formal/` holds TLA+ models, run only by hand with each model's `check.sh`; the `AGENTS.md`

tst/conftest.py

@@ -12,7 +12,7 @@
 from agent_env.store.document_store import document_store as document_store_module
 from agent_env.store.object_store.local import grant_server as grant_server_module
 from agent_env.store.routing import disable_namespace_routing
-from tst.util import shards
+from tst.util import package_cache, shards
 
 
 def _reset_store_singletons() -> None:
@@ -27,6 +27,7 @@ def _reset_store_singletons() -> None:
 
 
 _STATE_DIR = pytest.StashKey[str]()
+_PACKAGES_DIR = pytest.StashKey[str]()
 _state_home = pytest.MonkeyPatch()
 
 
@@ -36,11 +37,17 @@ def pytest_configure(config):
     (the capability checks behind skip marks)."""
     config.stash[_STATE_DIR] = tempfile.mkdtemp(prefix="agent-env-state-")
     _state_home.setenv("XDG_STATE_HOME", config.stash[_STATE_DIR])
+    # The computer tests' resolved packages, shared by their stores for the run unless CI names a folder it keeps.
+    if not os.environ.get(package_cache.ENV_VAR):
+        config.stash[_PACKAGES_DIR] = tempfile.mkdtemp(prefix="agent-env-packages-")
+        _state_home.setenv(package_cache.ENV_VAR, config.stash[_PACKAGES_DIR])
 
 
 def pytest_unconfigure(config):
     _state_home.undo()
     shutil.rmtree(config.stash[_STATE_DIR], ignore_errors=True)
+    if packages := config.stash.get(_PACKAGES_DIR, None):
+        shutil.rmtree(packages, ignore_errors=True)
 
 
 @pytest.hookimpl(trylast=True)

tst/integration/cli/run_bundle_browser_use_test.py

@@ -27,6 +27,7 @@
 from agent_env.task_step.task_steps.browser.connect_cdp import BrowserConnectCdpTaskStep
 from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep
 from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep
+from tst.util import package_cache
 from tst.util.fake_openai import FakeOpenAI, Text, ToolCall
 
 pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow]
@@ -69,11 +70,13 @@ def state(monkeypatch, tmp_path, endpoint):
         monkeypatch.delenv(name, raising=False)
     configure()
     reset_artifact_store()
+    package_cache.seed()
     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:
+        package_cache.keep()
         logging.disable(logging.NOTSET)
         grant_server(None, None).close()
         reset_artifact_store()

tst/integration/computer/conftest.py

@@ -2,7 +2,8 @@
 shipped desktop example's computer env, the cua-driver test bundle's, the chromium test bundle's, the website-browser
 test bundle's, the TLS clients test bundle's and agentenv's chromium deployed alone, each composed once for a module's
 tests;
-``state``, local stores for a test that runs a bundle; ``docker``, which runs the docker CLI for them; and
+``state``, local stores for a test that runs a bundle; each of these stores starts with the packages the run has resolved
+(``tst.util.package_cache``); ``docker``, which runs the docker CLI for them; and
 ``cdp_clients``, Playwright and Puppeteer installed to connect to a CDP address."""
 
 import logging
@@ -27,6 +28,7 @@
 from agent_env.store.object_store.local.grant_server import grant_server
 from agent_env.store.routing import namespace_routing
 from agent_env.task.store import get_task_instance_store
+from tst.util import package_cache
 from tst.util.published_image import pull_published_else_build
 
 BASE = Path(agent_env.__file__).parent / "computers" / "artifacts" / "ubuntu-24.04"
@@ -118,9 +120,13 @@ def _composed(root: Path, bundle, env_id: str | None = None):
     patch.setenv("AGENT_ENV_OBJECT_STORE", "local")
     configure()
     reset_artifact_store()
+    package_cache.seed()
     with namespace_routing():
         plan = plan_bundle(resolve_bundle(bundle), installed=installed_bundles)
-        materialize(plan, build_here=frozenset(w.id for w in plan.writes if isinstance(w.source, BuiltImage)))
+        try:
+            materialize(plan, build_here=frozenset(w.id for w in plan.writes if isinstance(w.source, BuiltImage)))
+        finally:
+            package_cache.keep()
         yield Env.get(env_id or f"{bundle.id_root}/desktop")
     reset_artifact_store()
     reset_config()
@@ -142,11 +148,13 @@ def state(monkeypatch, tmp_path):
     monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(sandboxes))
     configure()
     reset_artifact_store()
+    package_cache.seed()
     logging.disable(logging.CRITICAL)
     try:
         yield tmp_path
         assert not sandboxes.exists() or not any(sandboxes.iterdir()), "a run left a sandbox work folder"
     finally:
+        package_cache.keep()
         logging.disable(logging.NOTSET)
         grant_server(None, None).close()
         reset_artifact_store()