feat(run_code): give a script on an env's host the env's address on any provider, and stage each run in a dir of its own

src/agent_env/providers/sandbox_providers/sandbox.py

@@ -169,6 +169,18 @@ def url_from_sandbox(self, url: str) -> str:
         """``url``, of a service this process reaches, as this sandbox reaches it."""
         return url
 
+    def url_on_host(self, url: str) -> str:
+        """``url``, of a service this sandbox publishes, as a process on its host (or in it, for a container) reaches
+        it: one under a tunnel goes to the host port the tunnel's port is published on, and any other is left alone."""
+        address = "127.0.0.1" if not self.host_ips or "127.0.0.1" in self.host_ips else self.host_ips[0]
+        address = f"[{address}]" if ":" in address else address
+        # The longest first, so a tunnel at a path under another's takes its own URLs.
+        for port, tunnel in sorted(self.tunnel_urls.items(), key=lambda item: -len(item[1])):
+            base = tunnel.rstrip("/")
+            if url == base or url.startswith(f"{base}/"):
+                return f"http://{address}:{self.host_port(port)}{url[len(base):]}"
+        return url
+
     _VM_READY_TIMEOUT = 1200      # wait_for_vm wall-clock budget (s)
     _VM_READY_POLL_INTERVAL = 30  # sparse polling (s)
 

src/agent_env/task_step/task_steps/run_code.py

@@ -2,17 +2,20 @@
 
 The script defines ``run(input) -> JSON``, where ``input = {"args", "results"}``.
 The result is stored at ``context.metadata["script_results"][result_id]``.
+A script on an env's host also gets ``input["env"] = {"id", "url", "mcp_url"}``: the env's
+address, where its card is at ``/.well-known/agent-env.json``, and its MCP endpoint, as that
+host reaches them (``Sandbox.url_on_host``), since a host can't always call its own public URL.
 
 ``script_artifact_id`` is either a single-file ``FileArtifact`` or a
 ``FileArtifactUniverse`` whose ``.py`` members are staged side by side so the entry
 file can import its siblings; non-``.py`` members are skipped. ``script_file`` names
 the entry, defaulting to the sole ``.py`` member or to ``run.py``.
 
-Target: with ``env_id`` set, the script runs on that env's VM host (no agent, no
-LLM key); with ``sandbox_name`` set, on a ``deploy_sandbox`` sandbox (its VM host, or its
-container for a container-mode sandbox);
-otherwise it runs inside the agent container (``agent_name``), where it can see the
-agent's workspace.
+Target: with ``env_id`` set, the script runs on that env's VM host, or in its container for an
+env in a container sandbox, whose image needs python3 (no agent, no LLM key); with
+``sandbox_name`` set, on a ``deploy_sandbox`` sandbox (its VM host, or its container for a
+container-mode sandbox); otherwise it runs inside the agent container (``agent_name``), where
+it can see the agent's workspace.
 """
 
 from __future__ import annotations
@@ -28,7 +31,7 @@
 
 from agent_env.artifact import Artifact, FileArtifact
 from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse
-from agent_env.env.env import require_sandbox
+from agent_env.env.env import DeployedSandboxEnv, require_sandbox
 from agent_env.task_step.context import TaskStepContext
 from agent_env.entity_refs import EntityRef
 from agent_env.task_step.task_step import TaskStep, TaskStepDependency
@@ -164,9 +167,9 @@ def preflight(self) -> list[str]:
     async def execute(self, context: TaskStepContext) -> TaskStepContext:
         script_file, files = await asyncio.to_thread(self._load_script)
         target = await self._resolve_target(context)
-        run_dir = f"{_RUN_DIR}/{self.id}"
+        run_dir = target.run_dir(self.id)
 
-        await self._stage(target, run_dir, files, self._build_input(context))
+        await self._stage(target, run_dir, files, self._build_input(context, target))
         await self._run(target, run_dir, script_file)
         result = await self._read_output(target, run_dir)
 
@@ -289,12 +292,15 @@ def _load_script(self) -> tuple[str, dict[str, bytes]]:
             )
         return (script_file, files)
 
-    def _build_input(self, context: TaskStepContext) -> dict:
-        """Build the {args, results} input passed to the script."""
-        return {
+    def _build_input(self, context: TaskStepContext, target: _ExecTarget) -> dict:
+        """Build the {args, results} input passed to the script, and its env's endpoints on an env's host."""
+        script_input = {
             "args": self.args,
             "results": context.metadata.get("script_results", {}),
         }
+        if self.env_id is not None:
+            script_input["env"] = _env_on_host(self._deployed_env(context), target.sandbox)
+        return script_input
 
     def _target_kind(self) -> str:
         if self.env_id is not None:
@@ -304,7 +310,7 @@ def _target_kind(self) -> str:
         return "agent"
 
     async def _resolve_target(self, context: TaskStepContext) -> _ExecTarget:
-        """Pick where the script runs: the env's VM host, a deployed sandbox's host, or the agent container."""
+        """Pick where the script runs: the env's host, a deployed sandbox's host, or the agent container."""
         if self.env_id is not None:
             return await self._host_target(context)
         if self.sandbox_name is not None:
@@ -363,12 +369,8 @@ async def _read_output(self, target: _ExecTarget, run_dir: str) -> Any:
 
     # --- run targets ---
 
-    async def _host_target(self, context: TaskStepContext) -> _ExecTarget:
-        from agent_env.providers.sandbox_providers.sandbox_provider import (
-            build_sandbox_provider,
-            get_env_sandbox_provider,
-        )
-
+    def _deployed_env(self, context: TaskStepContext) -> DeployedSandboxEnv:
+        """The record of ``env_id``'s deploy, which must be in one of agent-env's sandboxes."""
         deployed = next(
             (d for d in context.deployed_envs if d.env_id == self.env_id), None
         )
@@ -379,14 +381,25 @@ async def _host_target(self, context: TaskStepContext) -> _ExecTarget:
                 f"Env '{self.env_id}' not found in context.deployed_envs "
                 "(or missing sandbox_id)"
             )
+        return deployed
+
+    async def _host_target(self, context: TaskStepContext) -> _ExecTarget:
+        from agent_env.providers.sandbox_providers.sandbox import VmSandbox
+        from agent_env.providers.sandbox_providers.sandbox_provider import (
+            build_sandbox_provider,
+            get_env_sandbox_provider,
+        )
+
+        deployed = self._deployed_env(context)
         # env sandbox provider, honoring a custom sandbox_type
         provider = (
             build_sandbox_provider(deployed.sandbox_type)
             if deployed.sandbox_type
             else get_env_sandbox_provider()
         )
         sandbox = await provider.get_sandbox(deployed.sandbox_id)
-        return _HostTarget(sandbox)
+        # An env in a container sandbox has no host to run on: the script runs in the container, beside its server.
+        return _HostTarget(sandbox) if isinstance(sandbox, VmSandbox) else _AgentTarget(sandbox, prefix=())
 
     async def _sandbox_target(self, context: TaskStepContext) -> _ExecTarget:
         from agent_env.providers.sandbox_providers.sandbox import VmSandbox
@@ -480,7 +493,20 @@ def _is_safe_member(name: str) -> bool:
     return ".." not in PurePosixPath(name).parts
 
 
+def _env_on_host(deployed: DeployedSandboxEnv, sandbox: Any) -> dict:
+    """The env's id, address and MCP endpoint as a script on its host reaches them."""
+    url, mcp_url = deployed.environment_url, deployed.mcp_url
+    return {
+        "id": deployed.env_id,
+        "url": url and sandbox.url_on_host(url),
+        "mcp_url": mcp_url and sandbox.url_on_host(mcp_url),
+    }
+
+
 class _ExecTarget(Protocol):
+    sandbox: Any
+
+    def run_dir(self, step_id: str) -> str: ...
     async def reset_dir(self, path: str) -> None: ...
     async def write_file(self, path: str, content: str) -> None: ...
     async def run(self, *command: str) -> tuple[int, str, str]: ...
@@ -492,6 +518,10 @@ class _HostTarget:
 
     sandbox: Any
 
+    def run_dir(self, step_id: str) -> str:
+        # The sandbox's own: runs on the local provider share this machine's /tmp, and teardown removes it there.
+        return f"{self.sandbox.scoped_name(_RUN_DIR)}/{step_id}"
+
     async def reset_dir(self, path: str) -> None:
         await self.sandbox.exec_script(
             f"rm -rf {shlex.quote(path)} && mkdir -p {shlex.quote(path)}"
@@ -507,11 +537,14 @@ async def run(self, *command: str) -> tuple[int, str, str]:
 
 @dataclass
 class _AgentTarget:
-    """Runs the script inside the agent container."""
+    """Runs the script inside a container: an agent's, a container sandbox's, or an env's in one."""
 
     sandbox: Any
     prefix: tuple[str, ...]
 
+    def run_dir(self, step_id: str) -> str:
+        return f"{_RUN_DIR}/{step_id}"
+
     async def reset_dir(self, path: str) -> None:
         await self.run("rm", "-rf", path)
         await self.run("mkdir", "-p", path)

tst/integration/cli/run_bundle_built_envs_test.py

@@ -8,7 +8,8 @@
 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; an output a bundle's task wrote is written under the folder
-that imports the task. Needs a Docker daemon and the local registry."""
+that imports the task. A script run on an env's host reaches the env's gateway and MCP endpoint at the addresses its
+input gives. Needs a Docker daemon and the local registry."""
 
 import hashlib
 import json
@@ -574,6 +575,47 @@ def _store_echo_server() -> None:
     MCPServerEnv.put(id="echo-server", docker_image_artifact=image, environment_name="echo")
 
 
+# Fetches the env's card at its address and opens an MCP session at its endpoint, both as its input gives them: on the
+# local provider, the host port the sandbox mapped, not the port the env serves on.
+_REACH = """\
+import json, urllib.parse, urllib.request
+
+def run(inp):
+    env = inp["env"]
+    address = urllib.parse.urlsplit(env["url"])
+    assert (address.hostname, address.port != 18765) == ("127.0.0.1", True), env
+    with urllib.request.urlopen(env["url"] + "/.well-known/agent-env.json", timeout=30) as response:
+        card = json.load(response)
+    initialize = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {
+        "protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "reach", "version": "1"}}}
+    request = urllib.request.Request(env["mcp_url"], data=json.dumps(initialize).encode(), headers={
+        "Content-Type": "application/json", "Accept": "application/json, text/event-stream"})
+    with urllib.request.urlopen(request, timeout=30) as response:
+        assert response.status == 200, response.status
+    return {"card": card["name"]}
+"""
+
+
+@pytest.mark.parametrize("env_provider_type", ["gateway", "server"])
+def test_a_script_on_an_envs_host_reaches_it_and_its_mcp_endpoint_at_the_addresses_its_input_gives(
+    state, env_provider_type
+):
+    root = state / "reach"
+    shutil.copytree(ECHO / "envs", root / "envs")
+    (root / "envs" / "echo" / "env.toml").write_text(f'env_provider_type = "{env_provider_type}"\n')
+    (root / "artifacts" / "reach").mkdir(parents=True)
+    (root / "artifacts" / "reach" / "run.py").write_text(_REACH)
+    (root / "tasks").mkdir()
+    (root / "tasks" / "reach.json").write_text(json.dumps([
+        {"id": "env", "type": "deploy_env", "env_id": "echo", "sandbox_type": "local"},
+        {"id": "reach", "type": "run_code", "env_id": "echo", "script_artifact_id": "reach", "depends_on": ["env"]}]))
+
+    result = _run(root)
+
+    assert result.exit_code == 0, result.output
+    assert "tasks/reach.json v1: unscored" in result.output, result.output
+
+
 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",

tst/unit/providers/sandbox_providers/test_url_rewrites.py

@@ -7,8 +7,10 @@
 
 from agent_env.config import reset_config
 from agent_env.env.env import DeployedEnv, DeployedSandboxEnv
-from agent_env.providers.sandbox_providers.local_sandbox import LocalSandboxProvider, host_url_for
+from agent_env.providers.sandbox_providers import local_sandbox as local_sandbox_module
+from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider, host_url_for
 from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider
+from agent_env.providers.sandbox_providers.sandbox import Sandbox
 from agent_env.providers.sandbox_providers.sandbox_provider import (
     SandboxProvider,
     all_sandbox_url_rewrites,
@@ -210,4 +212,62 @@ def test_choose_mcp_url_matrix(env_type, agent_type, rewritten):
 
 def test_an_env_outside_our_sandboxes_keeps_its_own_mcp_url():
     hosted = DeployedEnv(env_id="e", env_version=1, env_provider_type="hosted", mcp_url=f"https://{_INTERNAL}/mcp")
-    assert _choose_mcp_url(_record("modal"), hosted) == f"https://{_INTERNAL}/mcp"
+    assert _choose_mcp_url(_record("modal"), hosted) == f"https://{_INTERNAL}/mcp"
+
+# --- a sandbox's own URLs, as its host reaches them ------------------------------
+
+
+class _Published(Sandbox):
+    """A sandbox's tunnels, the addresses it publishes on and its host ports, as a provider reports them."""
+
+    type = "published"
+    mode = "vm"
+    vnc_url = None
+
+    def __init__(self, tunnel_urls, host_ips=(), port_map=None):
+        self.sandbox_id = "sb"
+        self.tunnel_urls = tunnel_urls
+        self.host_ips = host_ips
+        self._port_map = port_map or {}
+
+    def host_port(self, port):
+        return self._port_map.get(port, port)
+
+    async def terminate(self):  # pragma: no cover - never provisioned here
+        pass
+
+
+_PROXY = {80: "https://proxy.example/sb/vm-1", 18765: "https://proxy.example/sb/vm-1/port/18765/"}
+
+
+@pytest.mark.parametrize("tunnels,url,expected", [
+    ({18765: "https://abc.modal.host"}, "https://abc.modal.host", "http://127.0.0.1:18765"),
+    ({18765: "https://abc.modal.host"}, "https://abc.modal.host/mcp", "http://127.0.0.1:18765/mcp"),
+    # a proxy's path prefix, with another tunnel at a path above it
+    (_PROXY, "https://proxy.example/sb/vm-1/port/18765/mcp", "http://127.0.0.1:18765/mcp"),
+    (_PROXY, "https://proxy.example/sb/vm-1/web/", "http://127.0.0.1:80/web/"),
+    # a host that only starts like the tunnel's, and one elsewhere
+    ({18765: "https://abc.modal.host"}, "https://abc.modal.hostile/mcp", "https://abc.modal.hostile/mcp"),
+    ({18765: "https://abc.modal.host"}, "https://elsewhere.example/mcp", "https://elsewhere.example/mcp"),
+])
+def test_a_url_under_a_sandboxs_tunnel_is_reached_at_the_port_it_publishes(tunnels, url, expected):
+    assert _Published(tunnels).url_on_host(url) == expected
+
+
+@pytest.mark.parametrize("host_ips,expected", [
+    ((), "http://127.0.0.1:23456/mcp"),
+    (("127.0.0.1", "172.17.0.1"), "http://127.0.0.1:23456/mcp"),
+    (("10.0.0.5",), "http://10.0.0.5:23456/mcp"),
+    (("fd00::5",), "http://[fd00::5]:23456/mcp"),
+])
+def test_a_sandbox_is_reached_on_the_address_it_publishes_on(host_ips, expected):
+    sandbox = _Published({18765: "https://abc.modal.host"}, host_ips, port_map={18765: 23456})
+
+    assert sandbox.url_on_host("https://abc.modal.host/mcp") == expected
+
+
+def test_a_local_sandbox_is_reached_at_the_host_port_it_mapped(tmp_path, monkeypatch):
+    monkeypatch.setattr(local_sandbox_module, "_host_ips", lambda: ("127.0.0.1", "172.17.0.1"))
+    sandbox = LocalSandbox(work_dir=tmp_path, port_map={18765: 23456})
+
+    assert sandbox.url_on_host(f"{sandbox.tunnel_urls[18765]}/mcp") == "http://127.0.0.1:23456/mcp"