fix(release): keep the protocol's extras when the release pins it

.github/scripts/clean_install.py

@@ -21,7 +21,7 @@
 import zipfile
 from pathlib import Path
 
-from shipped_bundles import BUNDLES, bundle_roots
+from shipped_bundles import BUNDLES, bundle_roots, pin_pyproject
 
 REPO = Path(__file__).resolve().parents[2]
 PYPI = "https://pypi.org/simple"
@@ -46,10 +46,17 @@ def main() -> int:
     venv = work / "venv"
     env = _hermetic(work, venv / "bin")
 
-    _step("Build both distributions as the release does (sdist, then the wheel from it)")
-    for package in ("packages/agentenv-protocol", "."):
-        _run([uv, "tool", "run", "--python", sys.executable, "--from", "build", "pyproject-build",
-              "--outdir", str(work / "dist"), package], env=env, cwd=REPO)
+    _step("Build both distributions as the release does (sdist, then the wheel from it), the framework pinned to the "
+          "protocol as the release pins it")
+    pyproject = REPO / "pyproject.toml"
+    unpinned = pyproject.read_text()
+    try:
+        pin_pyproject()
+        for package in ("packages/agentenv-protocol", "."):
+            _run([uv, "tool", "run", "--python", sys.executable, "--from", "build", "pyproject-build",
+                  "--outdir", str(work / "dist"), package], env=env, cwd=REPO)
+    finally:
+        pyproject.write_text(unpinned)
     protocol = _one(work / "dist", "agentenv_framework_protocol-*.whl")
     framework = _one(work / "dist", "agentenv_framework-*.whl")
 

.github/scripts/shipped_bundles.py

@@ -1,8 +1,12 @@
 """The bundles agentenv-framework ships, from its ``agent_env.bundles`` entry points: ``<name> = "<package>"`` ships the
 folder ``src/<package as a path>/<name>/``.
 
-``pin`` pins agentenv-framework-protocol, keeping its extras, in each shipped env's and artifact's requirements.txt to
-the version ``packages/agentenv-protocol`` is at, which a release ships them with. An agent's, and an env's whose image
+``pin-pyproject`` pins agentenv-framework-protocol, keeping its extras, in the framework's own requirement on it in
+``pyproject.toml``, to the version ``packages/agentenv-protocol`` is at, which the release's wheel is built with. Its
+extras are what core imports from the protocol, so a pin that dropped them would ship a wheel that can't import.
+
+``pin`` pins the protocol the same way in each shipped env's and artifact's requirements.txt, which a release ships them
+with. An agent's, and an env's whose image
 a release publishes (``PUBLISHED_ENVS``), keep their floor: a pin is one of an image's sources, so pinning it would
 build and publish a new image every release, and with the floor the image is built only when its folder changes. Every
 release job that reads those folders runs it first, so an image's tag, its build and the wheel all see the same sources.
@@ -25,6 +29,9 @@
 PUBLISHED_ENVS = ("src/agent_env/examples/mail-mcp/envs/mail",)
 # A requirement on the protocol, with its extras, and not on a package whose name starts with the protocol's.
 _PROTOCOL_REQUIREMENT = re.compile(rf"^{PROTOCOL}(\[[^\]]*\])?(?![\w.-])")
+# The quoted requirement in pyproject.toml, with its extras: the directory name in the sdist include list and the uv
+# workspace members, and the bare name keying [tool.uv.sources], aren't requirements.
+_QUOTED_REQUIREMENT = re.compile(rf'"{PROTOCOL}(\[[^\]"]*\])?(?![\w./-])[^"]*"')
 
 
 def bundle_roots(repo: Path = REPO) -> dict[str, Path]:
@@ -60,6 +67,19 @@ def pin(repo: Path = REPO) -> list[Path]:
     return pinned
 
 
+def pin_pyproject(repo: Path = REPO) -> str:
+    """Pin the framework's one quoted requirement on the protocol in ``pyproject.toml``, keeping its extras, and return
+    the pinned requirement."""
+    path = repo / "pyproject.toml"
+    text = path.read_text()
+    found = _QUOTED_REQUIREMENT.findall(text)
+    if len(found) != 1:
+        raise ValueError(f"{path} holds {len(found)} quoted {PROTOCOL} requirements, not the one a release pins")
+    pinned = f"{PROTOCOL}{found[0]}=={protocol_version(repo)}"
+    path.write_text(_QUOTED_REQUIREMENT.sub(f'"{pinned}"', text))
+    return pinned
+
+
 def _pinned(line: str, version: str) -> str:
     if not (match := _PROTOCOL_REQUIREMENT.match(line)):
         return line
@@ -69,8 +89,12 @@ def _pinned(line: str, version: str) -> str:
 
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_subparsers(dest="command", required=True).add_parser("pin")
-    parser.parse_args()
+    commands = parser.add_subparsers(dest="command", required=True)
+    commands.add_parser("pin")
+    commands.add_parser("pin-pyproject")
+    if parser.parse_args().command == "pin-pyproject":
+        print(f"pinned {pin_pyproject()} in pyproject.toml")
+        return 0
     version = protocol_version()
     for path in pin():
         print(f"pinned {PROTOCOL}=={version} in {path.relative_to(REPO).as_posix()}")

.github/workflows/publish-pypi.yml

@@ -421,11 +421,8 @@ jobs:
           echo "record=$RECORD" >> "$GITHUB_OUTPUT"
       - name: Build agentenv-framework, pinned to this release's agentenv-framework-protocol
         run: |
-          protocol=$(sed -n 's/^version = "\(.*\)"/\1/p' packages/agentenv-protocol/pyproject.toml)
-          # Only the quoted requirement: the directory name also appears in the sdist include list and the uv
-          # workspace members, and the bare name keys [tool.uv.sources]; none of those may be rewritten.
-          sed -i "s/\"agentenv-framework-protocol[^\"]*\"/\"agentenv-framework-protocol==${protocol}\"/" pyproject.toml
-          grep -q "\"agentenv-framework-protocol==${protocol}\"" pyproject.toml
+          # The framework's requirement on the protocol keeps its extras, which core imports from.
+          python3 .github/scripts/shipped_bundles.py pin-pyproject
           # Each shipped bundle's images install the protocol this release ships with, keeping its extras.
           python3 .github/scripts/shipped_bundles.py pin
           uv build --no-sources --out-dir dist

tst/unit/test_release_infra_images.py

@@ -22,7 +22,8 @@
 _RUNS_DEPENDENCIES = re.compile(r"\buv (sync|run)\b|infra_images\.py")
 # What reads the shipped bundles' folders: an image's tag, its build, and the wheel.
 _READS_SHIPPED_SOURCES = re.compile(r"infra_images\.py (plan|write)|docker build\b|uv build --no-sources --out-dir dist")
-_PIN = "python3 .github/scripts/shipped_bundles.py pin"
+# The pin command, and not pin-pyproject, which pins the framework's own requirement.
+_PIN = re.compile(r"python3 \.github/scripts/shipped_bundles\.py pin(?![\w-])")
 
 
 def _runs(job: str) -> list[str]:
@@ -122,7 +123,8 @@ def test_every_job_that_reads_the_shipped_bundles_pins_the_protocol_first():
 
     assert set(readers) == {"infra-plan", "infra-build", "agent-build", "infra-record", "framework"}
     for job, runs in readers.items():
-        assert runs.count(_PIN) == 1 and runs.index(_PIN) < _READS_SHIPPED_SOURCES.search(runs).start(), job
+        pins = [pin.start() for pin in _PIN.finditer(runs)]
+        assert len(pins) == 1 and pins[0] < _READS_SHIPPED_SOURCES.search(runs).start(), job
 
 
 def test_the_pin_changes_no_published_images_sources():

tst/unit/test_release_protocol_pins.py

@@ -1,12 +1,13 @@
 """The files ``.github/workflows/publish-pypi.yml`` rewrites to pin agentenv-framework-protocol on a release, the
-pyproject with its sed and the shipped bundles' env and artifact requirements files with ``shipped_bundles.py pin``: each
-must hold exactly one protocol requirement, or the release ships an unpinned one. An agent's, and an env's whose image a
-release publishes, keep their floor."""
+pyproject with ``shipped_bundles.py pin-pyproject`` and the shipped bundles' env and artifact requirements files with
+``shipped_bundles.py pin``: each must hold exactly one protocol requirement, or the release ships an unpinned one, and
+the pin keeps its extras. An agent's, and an env's whose image a release publishes, keep their floor."""
 
 from __future__ import annotations
 
 import importlib.util
 import re
+import shutil
 from pathlib import Path
 
 import pytest
@@ -24,6 +25,35 @@ def test_the_pyproject_holds_the_one_quoted_protocol_requirement_the_release_pin
     assert len(re.findall(r'"agentenv-framework-protocol[^"]*"', (REPO / "pyproject.toml").read_text())) == 1
 
 
+def test_the_release_pins_the_pyproject_keeping_the_extras_core_imports_from_the_protocol(tmp_path):
+    """The wheels 0.9.1363 to 0.9.1368 shipped a pin without [computer], so a bare install couldn't import av."""
+    shutil.copy(REPO / "pyproject.toml", tmp_path / "pyproject.toml")
+    (tmp_path / "packages/agentenv-protocol").mkdir(parents=True)
+    shutil.copy(REPO / "packages/agentenv-protocol/pyproject.toml", tmp_path / "packages/agentenv-protocol")
+
+    assert shipped_bundles.pin_pyproject(tmp_path) == (
+        f"agentenv-framework-protocol[computer]=={shipped_bundles.protocol_version()}")
+    assert "python3 .github/scripts/shipped_bundles.py pin-pyproject" in WORKFLOW
+    assert 'sed -i "s/\\"agentenv-framework-protocol' not in WORKFLOW
+
+
+def test_pin_pyproject_rewrites_only_the_quoted_requirement_and_again_changes_nothing(tmp_path):
+    repo = _repo(tmp_path, {})
+    text = ('[project]\nname = "x"\ndependencies = [\n    "agentenv-framework-protocol[computer]>=0.1.0",  # why\n'
+            '    "agentenv-framework-protocol-tools==1",\n]\n\n[tool.hatch.build.targets.sdist]\n'
+            'include = ["/src", "/packages/agentenv-protocol"]\n\n[tool.uv.sources]\n'
+            'agentenv-framework-protocol = { workspace = true }\n')
+    (repo / "pyproject.toml").write_text(text)
+    pinned = text.replace(">=0.1.0", "==1.2.3")
+
+    assert shipped_bundles.pin_pyproject(repo) == "agentenv-framework-protocol[computer]==1.2.3"
+    assert (repo / "pyproject.toml").read_text() == pinned
+    assert shipped_bundles.pin_pyproject(repo) == "agentenv-framework-protocol[computer]==1.2.3"
+    assert (repo / "pyproject.toml").read_text() == pinned
+    (repo / "pyproject.toml").write_text(text.replace("[computer]>=0.1.0", ""))
+    assert shipped_bundles.pin_pyproject(repo) == "agentenv-framework-protocol==1.2.3"
+
+
 def _repo(tmp_path: Path, requirements: dict[str, str]) -> Path:
     """A checkout whose protocol is at 1.2.3, shipping the agents bundle with a ``<kind>/<name>`` folder per entry of
     ``requirements``."""