bash: clean_env option — run task commands in the image's environment

agentix/runtime/server/worker/client.py

@@ -20,7 +20,12 @@
 from typing import Any, Protocol
 
 from agentix.runtime.server.worker.invoker import CallableInvoker
-from agentix.runtime.shared.env import AGENTIX_ADDED_PATH, BUNDLE_RUNTIME_BIN, BUNDLE_RUNTIME_PATH_ENTRIES
+from agentix.runtime.shared.env import (
+    AGENTIX_ADDED_PATH,
+    AGENTIX_SAVED_PREFIX,
+    BUNDLE_RUNTIME_BIN,
+    BUNDLE_RUNTIME_PATH_ENTRIES,
+)
 from agentix.runtime.shared.framing import read_frame, write_frame
 from agentix.runtime.shared.models import RemoteError, RemoteRequest, RemoteResponse
 
@@ -106,11 +111,20 @@ def _prepend_recorded_path_entries(env: dict[str, str], name: str, entries: Iter
 
 
 def _clean_worker_env(runtime_bin_dir: Path | None) -> dict[str, str]:
-    env = {
-        key: value
-        for key, value in os.environ.items()
-        if key not in _STRIPPED_ENV and not any(key.startswith(prefix) for prefix in _STRIPPED_ENV_PREFIXES)
-    }
+    env: dict[str, str] = {}
+    saved: dict[str, str] = {}
+    for key, value in os.environ.items():
+        if key in _STRIPPED_ENV or any(key.startswith(prefix) for prefix in _STRIPPED_ENV_PREFIXES):
+            # Stripped for the worker's own interpreter hygiene, but the task
+            # image legitimately owns these — record them so
+            # `get_env_without_agentix` can restore them for task subprocesses.
+            saved[f"{AGENTIX_SAVED_PREFIX}{key}"] = value
+            continue
+        env[key] = value
+    # The freshly-observed value is authoritative: overwrite any stale
+    # AGENTIX_SAVED_* snapshot inherited from an outer layer, so a nested
+    # clean env restores this spawn's image value, not an older one.
+    env.update(saved)
     # Build PATH from: the venv's bin (`runtime_bin_dir`), the bundle's
     # symlink-join (`BUNDLE_RUNTIME_BIN`), then the parent environment's
     # PATH. Inside the bundle runtime tree the first two are siblings and both

agentix/runtime/shared/env.py

@@ -117,6 +117,14 @@
 
 _TRACKING_PREFIX = "AGENTIX_ADDED_"
 
+AGENTIX_SAVED_PREFIX = "AGENTIX_SAVED_"
+"""Prefix for vars the worker spawn STRIPPED (rather than prepended to).
+
+The worker removes interpreter-hostile vars (``PYTHONPATH``, ``LD_PRELOAD``,
+...) from its own environment, but the task image legitimately owns them —
+so the spawn records each one under ``AGENTIX_SAVED_<NAME>`` and
+:func:`get_env_without_agentix` restores them for task subprocesses."""
+
 
 def _split_path(value: str) -> list[str]:
     return value.split(os.pathsep) if value else []
@@ -144,8 +152,11 @@ def get_env_without_agentix(
     """Return an environment for user subprocesses without Agentix-added paths.
 
     The helper only subtracts entries that Agentix explicitly recorded in
-    `AGENTIX_ADDED_*`. It intentionally does not remove arbitrary `/nix`
-    paths, because a task image may itself be Nix-based.
+    `AGENTIX_ADDED_*` (it intentionally does not remove arbitrary `/nix`
+    paths, because a task image may itself be Nix-based), and restores vars
+    the worker spawn stripped and recorded under `AGENTIX_SAVED_*`. A live
+    value wins over its saved snapshot — presence means something set it
+    after the spawn, which is more recent intent than the pre-strip original.
     """
 
     env = dict(os.environ if base is None else base)
@@ -164,6 +175,14 @@ def get_env_without_agentix(
     for name in tracking_vars:
         env.pop(name, None)
 
+    saved_vars = [name for name in env if name.startswith(AGENTIX_SAVED_PREFIX)]
+    for saved_var in saved_vars:
+        target = saved_var.removeprefix(AGENTIX_SAVED_PREFIX)
+        if target and target not in env:
+            env[target] = env[saved_var]
+    for name in saved_vars:
+        env.pop(name, None)
+
     if extra:
         env.update(extra)
     return env
@@ -172,6 +191,7 @@ def get_env_without_agentix(
 __all__ = [
     "AGENTIX_ADDED_LD_LIBRARY_PATH",
     "AGENTIX_ADDED_PATH",
+    "AGENTIX_SAVED_PREFIX",
     "BIND_HOST_ENV",
     "BIND_PORT_ENV",
     "BUNDLE_NIX_ROOT",

plugins/runtime-basic/agentix/bash/__init__.py

@@ -33,20 +33,37 @@
 
 from pydantic import Field
 
-from agentix.runtime.shared.env import BUNDLE_RUNTIME_BASH
+from agentix.runtime.shared.env import BUNDLE_RUNTIME_BASH, get_env_without_agentix
 
 
-def _clean_env(extra: dict[str, str] | None) -> dict[str, str]:
-    """Build a subprocess env: inherited runtime env + caller overrides."""
-    env = dict(os.environ)
+def _clean_env(extra: dict[str, str] | None, clean: bool = False) -> dict[str, str]:
+    """Build a subprocess env, then apply caller overrides last.
+
+    Default: the worker's own environment — right for runtime tools (the
+    claude CLI, bundled rg, ...) that live under /nix. ``clean=True``: the
+    task IMAGE's environment (bundle path additions subtracted, stripped vars
+    restored) — right for task commands that must resolve the image's own
+    toolchain instead of the bundle venv.
+    """
+    env = get_env_without_agentix() if clean else dict(os.environ)
     if extra:
         env.update(extra)
     return env
 
 
-def _shell_executable(executable: str | None, env: dict[str, str]) -> str:
+def _shell_executable(executable: str | None, env: dict[str, str], clean: bool = False) -> str:
     if executable:
         return shutil.which(executable, path=env.get("PATH")) or executable
+    if clean:
+        # A clean-env command belongs to the image's world entirely — its
+        # restored vars (LD_PRELOAD, ...) target the image's own bash, not
+        # the Nix one. Fall back to the bundled bash only when the image has
+        # none at all.
+        image_bash = shutil.which("bash", path=env.get("PATH"))
+        if image_bash:
+            return image_bash
+        if os.access("/bin/bash", os.X_OK):
+            return "/bin/bash"
     if os.access(BUNDLE_RUNTIME_BASH, os.X_OK):
         return BUNDLE_RUNTIME_BASH
     return shutil.which("bash", path=env.get("PATH")) or "/bin/bash"
@@ -145,16 +162,23 @@ async def run(
     timeout: float | None = None,
     max_output: int = 10 * 1024 * 1024,
     executable: str | None = None,
+    clean_env: bool = False,
 ) -> BashResult:
-    """Run a shell command in the sandbox and return its captured output."""
-    sub_env = _clean_env(env)
+    """Run a shell command in the sandbox and return its captured output.
+
+    ``clean_env=True`` runs the command in the task image's environment
+    (Agentix's recorded path additions subtracted, stripped vars restored)
+    instead of the bundle-polluted worker environment; ``env`` overrides win
+    last either way.
+    """
+    sub_env = _clean_env(env, clean=clean_env)
     proc = await asyncio.create_subprocess_shell(
         command,
         stdout=asyncio.subprocess.PIPE,
         stderr=asyncio.subprocess.PIPE,
         cwd=cwd,
         env=sub_env,
-        executable=_shell_executable(executable, sub_env),
+        executable=_shell_executable(executable, sub_env, clean=clean_env),
     )
     assert proc.stdout is not None and proc.stderr is not None
     stdout_task = asyncio.create_task(_read_capped(proc.stdout, max_output))
@@ -185,20 +209,22 @@ async def run_stream(
     env: dict[str, str] | None = None,
     timeout: float | None = None,
     executable: str | None = None,
+    clean_env: bool = False,
 ) -> AsyncIterator[BashEvent]:
     """Run a shell command, yielding events as the subprocess emits them.
 
     Terminates with a single `BashExit` event on normal completion or
     a single `BashError` event on timeout / wire-level failure.
+    ``clean_env`` selects the task image's environment, as in `run`.
     """
-    sub_env = _clean_env(env)
+    sub_env = _clean_env(env, clean=clean_env)
     proc = await asyncio.create_subprocess_shell(
         command,
         stdout=asyncio.subprocess.PIPE,
         stderr=asyncio.subprocess.PIPE,
         cwd=cwd,
         env=sub_env,
-        executable=_shell_executable(executable, sub_env),
+        executable=_shell_executable(executable, sub_env, clean=clean_env),
     )
 
     async def _pump(stream, tag, queue):

plugins/runtime-basic/tests/test_primitives.py

@@ -114,3 +114,42 @@ async def test_bash_run_stream_can_use_explicit_executable(tmp_path: Path):
 
     assert [event.data for event in events if isinstance(event, bash.BashStdout)] == ["zsh"]
     assert [event.exit_code for event in events if isinstance(event, bash.BashExit)] == [0]
+
+
+@pytest.mark.asyncio
+async def test_bash_run_clean_env_runs_in_the_image_environment(monkeypatch):
+    """clean_env=True subtracts the bundle's recorded PATH additions and
+    restores stripped vars, so the command resolves the image's toolchain."""
+    from agentix.runtime.shared.env import AGENTIX_ADDED_PATH
+
+    real_path = os.environ.get("PATH", "/usr/bin")
+    monkeypatch.setenv("PATH", os.pathsep.join(["/agentix-injected/bin", real_path]))
+    monkeypatch.setenv(AGENTIX_ADDED_PATH, "/agentix-injected/bin")
+    monkeypatch.setenv("AGENTIX_SAVED_PYTHONPATH", "/testbed/src")
+    # restore is live-wins: an ambient PYTHONPATH (conda, IDE runners) would
+    # legitimately shadow the snapshot and flake this test — clear it.
+    monkeypatch.delenv("PYTHONPATH", raising=False)
+
+    polluted = await bash.run(command="printf %s \"$PATH\"")
+    assert polluted.stdout.startswith("/agentix-injected/bin")
+
+    clean = await bash.run(command="printf %s \"$PATH:$PYTHONPATH\"", clean_env=True)
+    assert "/agentix-injected/bin" not in clean.stdout
+    assert clean.stdout.endswith("/testbed/src")
+
+
+def test_shell_executable_prefers_image_bash_in_clean_mode(tmp_path):
+    """A clean-env command's restored vars (LD_PRELOAD, ...) target the
+    image's own bash; the bundled Nix bash is only the last resort."""
+    fakebin = tmp_path / "bin"
+    fakebin.mkdir()
+    image_bash = fakebin / "bash"
+    image_bash.write_text("#!/bin/sh\n")
+    image_bash.chmod(0o755)
+
+    resolved = bash._shell_executable(None, {"PATH": str(fakebin)}, clean=True)
+    assert resolved == str(image_bash)
+
+    # without clean mode the bundled-bash preference is unchanged (falls back
+    # to PATH lookup here because /nix/runtime/bin/bash does not exist)
+    assert bash._shell_executable(None, {"PATH": str(fakebin)}) == str(image_bash)

tests/runtime/server/worker/test_env.py

@@ -93,3 +93,45 @@ def test_clean_worker_env_injects_recorded_runtime_build_paths(monkeypatch: pyte
         "/task/lib/pkgconfig",
     ]
     assert env["CMAKE_PREFIX_PATH"] == "/nix/runtime"
+
+
+def test_clean_worker_env_records_stripped_vars(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Stripping without recording makes the image env unrecoverable for task
+    subprocesses; every stripped var must land in AGENTIX_SAVED_*."""
+    monkeypatch.setenv("PYTHONPATH", "/testbed/src")
+    monkeypatch.setenv("LD_PRELOAD", "/usr/lib/libfoo.so")
+    monkeypatch.setenv("NIX_CFLAGS_COMPILE", "-O2")
+
+    env = _clean_worker_env(Path("/runtime/venv/bin"))
+
+    assert "PYTHONPATH" not in env
+    assert env["AGENTIX_SAVED_PYTHONPATH"] == "/testbed/src"
+    assert env["AGENTIX_SAVED_LD_PRELOAD"] == "/usr/lib/libfoo.so"
+    assert env["AGENTIX_SAVED_NIX_CFLAGS_COMPILE"] == "-O2"
+
+
+def test_clean_worker_env_records_nothing_when_nothing_stripped(monkeypatch: pytest.MonkeyPatch) -> None:
+    # The runner's own environment may legitimately carry strippable vars
+    # (NIX_*, SSL_CERT_FILE on Nix-based CI) — clear them all so this test
+    # asserts the production rule, not the runner's environment.
+    for key in list(os.environ):
+        if key in {"LD_PRELOAD", "PYTHONPATH", "PYTHONHOME", "LOCALE_ARCHIVE", "SSL_CERT_FILE"} or key.startswith(
+            ("NIX_", "FONTCONFIG_", "AGENTIX_SAVED_")
+        ):
+            monkeypatch.delenv(key, raising=False)
+
+    env = _clean_worker_env(Path("/runtime/venv/bin"))
+
+    assert not any(key.startswith("AGENTIX_SAVED_") for key in env)
+
+
+def test_clean_worker_env_fresh_strip_overwrites_inherited_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Nested spawn: the freshly-observed live value is authoritative — an
+    AGENTIX_SAVED_* snapshot inherited from an outer layer must not shadow it."""
+    monkeypatch.setenv("AGENTIX_SAVED_PYTHONPATH", "/stale/outer")
+    monkeypatch.setenv("PYTHONPATH", "/fresh/inner")
+
+    env = _clean_worker_env(Path("/runtime/venv/bin"))
+
+    assert env["AGENTIX_SAVED_PYTHONPATH"] == "/fresh/inner"
+    assert "PYTHONPATH" not in env

tests/runtime/test_env.py

@@ -93,3 +93,56 @@ def test_bash_clean_env_inherits_agentix_runtime_env(monkeypatch) -> None:
 
     assert env["PATH"].split(os.pathsep) == ["/nix/runtime/bin", "/task/bin"]
     assert env[AGENTIX_ADDED_PATH] == "/nix/runtime/bin"
+
+
+def test_get_env_without_agentix_restores_saved_vars() -> None:
+    """Vars the worker spawn STRIPPED (PYTHONPATH, LD_PRELOAD, ...) are
+    recorded under AGENTIX_SAVED_*; the clean env restores them so task
+    commands see the image's original environment."""
+    env = get_env_without_agentix(
+        base={
+            "AGENTIX_SAVED_PYTHONPATH": "/testbed/src",
+            "AGENTIX_SAVED_LD_PRELOAD": "/usr/lib/libfoo.so",
+            "PATH": "/usr/bin",
+        }
+    )
+
+    assert env["PYTHONPATH"] == "/testbed/src"
+    assert env["LD_PRELOAD"] == "/usr/lib/libfoo.so"
+    assert "AGENTIX_SAVED_PYTHONPATH" not in env
+    assert "AGENTIX_SAVED_LD_PRELOAD" not in env
+
+
+def test_get_env_without_agentix_does_not_clobber_live_values_on_restore() -> None:
+    """A live value is more recent intent than the pre-strip snapshot."""
+    env = get_env_without_agentix(
+        base={
+            "PYTHONPATH": "/set/after/spawn",
+            "AGENTIX_SAVED_PYTHONPATH": "/original",
+        }
+    )
+
+    assert env["PYTHONPATH"] == "/set/after/spawn"
+    assert "AGENTIX_SAVED_PYTHONPATH" not in env
+
+
+def test_get_env_without_agentix_extra_wins_over_restored() -> None:
+    env = get_env_without_agentix(
+        {"PYTHONPATH": "/caller"},
+        base={"AGENTIX_SAVED_PYTHONPATH": "/original"},
+    )
+
+    assert env["PYTHONPATH"] == "/caller"
+
+
+def test_bash_clean_env_opt_in_builds_the_image_environment(monkeypatch) -> None:
+    monkeypatch.setenv("PATH", os.pathsep.join(["/nix/runtime/bin", "/task/bin"]))
+    monkeypatch.setenv(AGENTIX_ADDED_PATH, "/nix/runtime/bin")
+    monkeypatch.setenv("AGENTIX_SAVED_PYTHONPATH", "/testbed/src")
+
+    env = _clean_env({"EXTRA": "1"}, clean=True)
+
+    assert env["PATH"] == "/task/bin"
+    assert env["PYTHONPATH"] == "/testbed/src"
+    assert AGENTIX_ADDED_PATH not in env
+    assert env["EXTRA"] == "1"