diff --git a/agentix/runtime/server/worker/client.py b/agentix/runtime/server/worker/client.py index 200d899..aacadee 100644 --- a/agentix/runtime/server/worker/client.py +++ b/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 diff --git a/agentix/runtime/shared/env.py b/agentix/runtime/shared/env.py index 37116c6..4a9acc5 100644 --- a/agentix/runtime/shared/env.py +++ b/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_`` 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", diff --git a/plugins/runtime-basic/agentix/bash/__init__.py b/plugins/runtime-basic/agentix/bash/__init__.py index d6537ca..d1d88e8 100644 --- a/plugins/runtime-basic/agentix/bash/__init__.py +++ b/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): diff --git a/plugins/runtime-basic/tests/test_primitives.py b/plugins/runtime-basic/tests/test_primitives.py index c14aa3d..e32aa06 100644 --- a/plugins/runtime-basic/tests/test_primitives.py +++ b/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) diff --git a/tests/runtime/server/worker/test_env.py b/tests/runtime/server/worker/test_env.py index 17a6fe0..40b9ca6 100644 --- a/tests/runtime/server/worker/test_env.py +++ b/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 diff --git a/tests/runtime/test_env.py b/tests/runtime/test_env.py index 3d41a17..991cf36 100644 --- a/tests/runtime/test_env.py +++ b/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"