feat(sandbox): configure the Modal environment for modal and modal_vm

src/agent_env/providers/sandbox_providers/modal_image_build.py

@@ -24,19 +24,20 @@
 
 logger = logging.getLogger(__name__)
 
-# The id of each image built from a build context, by Modal token, app and source digest, for every provider in the
-# process to reuse; one build of each runs at a time, so deploys of one image at once share its build.
-_built_image_ids: dict[tuple[str, str, str], str] = {}
-_build_locks: weakref.WeakValueDictionary[tuple[asyncio.AbstractEventLoop, tuple[str, str, str]], asyncio.Lock] = (
+# The id of each image built from a build context, by Modal token, environment, app and source digest, for every
+# provider in the process to reuse; one build of each runs at a time, so deploys of one image at once share its build.
+_built_image_ids: dict[tuple[str, str, str, str], str] = {}
+_build_locks: weakref.WeakValueDictionary[tuple[asyncio.AbstractEventLoop, tuple[str, str, str, str]], asyncio.Lock] = (
     weakref.WeakValueDictionary()
 )
 
 
 async def context_image(image: DockerImageArtifact, *, client: modal.Client, app: modal.App, token_id: str,
-                        app_name: str, stale: str | None = None) -> tuple[modal.Image, str]:
+                        app_name: str, environment_name: str | None,
+                        stale: str | None = None) -> tuple[modal.Image, str]:
     """The Modal image built from ``image``'s build context in ``app``, and its id: built now unless one built from the
     same sources is, other than ``stale``, an id Modal no longer holds."""
-    key = (token_id, app_name, image.source_digest or image.build_context_object_url)
+    key = (token_id, environment_name or "", app_name, image.source_digest or image.build_context_object_url)
     async with _build_locks.setdefault((asyncio.get_running_loop(), key), asyncio.Lock()):
         image_id = _built_image_ids.get(key)
         if image_id is None or image_id == stale:

src/agent_env/providers/sandbox_providers/modal_sandbox.py

@@ -234,9 +234,11 @@ def __init__(
         app_name: Optional[str] = None,
         gpu: Optional[str] = None,
         ecr_pull_secret_name: Optional[str] = None,
+        environment_name: Optional[str] = None,
     ):
         self._app_name = _resolve_app_base_name(app_name)
         self._ecr_pull_secret_name = ecr_pull_secret_name
+        self._environment_name = environment_name
         self._client: Optional[modal.Client] = None
         # Optional Modal GPU spec (e.g. "H100", "A100", "H100:2"). When set, create_container
         # provisions the sandbox with the V1 ``modal.Sandbox.create`` factory (which accepts
@@ -263,7 +265,9 @@ async def _get_app(self, app_name: str, app_tags: Optional[dict[str, str]] = Non
         app = self._apps.get(app_name)
         if app is None:
             client = await self._get_client()
-            app = await modal.App.lookup.aio(app_name, create_if_missing=True, client=client)
+            app = await modal.App.lookup.aio(
+                app_name, create_if_missing=True, client=client, environment_name=self._environment_name,
+            )
             if app_tags:
                 # App tags are the dimension Modal's billing report breaks down by; set them
                 # once per app (cached below) rather than on every sandbox.
@@ -291,6 +295,7 @@ async def _context_image(self, image: DockerImageArtifact, *,
         client = await self._get_client()
         app = await self._get_app(self._app_name, _attribution_tags({}))
         return await context_image(image, client=client, app=app, token_id=self._token_id, app_name=self._app_name,
+                                   environment_name=self._environment_name,
                                    stale=stale)
 
     async def create_container(
@@ -493,7 +498,8 @@ async def _registry_image(self, image_name: str) -> modal.Image:
         auth = None if uses_ecr_pull_secret else await asyncio.to_thread(image_store.auth, image_name)
         if uses_ecr_pull_secret:
             return modal.Image.from_aws_ecr(image_name, secret=modal.Secret.from_name(
-                self._ecr_pull_secret_name, required_keys=ECR_READER_SECRET_KEYS))
+                self._ecr_pull_secret_name, required_keys=ECR_READER_SECRET_KEYS,
+                environment_name=self._environment_name))
         if auth is not None:
             return modal.Image.from_registry(image_name, secret=modal.Secret.from_dict({
                 "REGISTRY_USERNAME": auth.username,

src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py

@@ -281,8 +281,9 @@ def supports_network_policy(cls, policy: NetworkPolicy) -> bool:
 
     """Provisions a single Modal VM sandbox per env (the VM-mode Modal backend)."""
 
-    def __init__(self, app_name: Optional[str] = None):
+    def __init__(self, app_name: Optional[str] = None, environment_name: Optional[str] = None):
         self._app_name = _resolve_app_base_name(app_name)
+        self._environment_name = environment_name
         self._client: Optional[modal.Client] = None
         # Per-attribution-scope Modal App, looked up on demand (billing aggregates per App).
         # Self-contained copy of ModalSandboxProvider's client/app handling — intentionally
@@ -300,7 +301,9 @@ async def _get_app(self, app_name: str, app_tags: Optional[dict[str, str]] = Non
         app = self._apps.get(app_name)
         if app is None:
             client = await self._get_client()
-            app = await modal.App.lookup.aio(app_name, create_if_missing=True, client=client)
+            app = await modal.App.lookup.aio(
+                app_name, create_if_missing=True, client=client, environment_name=self._environment_name,
+            )
             if app_tags:
                 try:
                     await app.set_tags.aio(app_tags)

tst/unit/providers/sandbox_providers/modal_context_build_test.py

@@ -64,8 +64,8 @@ def builds(monkeypatch):
     return builds
 
 
-def _provider() -> ModalSandboxProvider:
-    provider = ModalSandboxProvider(app_name="agent-env-test")
+def _provider(environment_name: str | None = None) -> ModalSandboxProvider:
+    provider = ModalSandboxProvider(app_name="agent-env-test", environment_name=environment_name)
     provider._get_client = AsyncMock(return_value=MagicMock())
     provider._get_app = AsyncMock(return_value="app")
     return provider
@@ -110,6 +110,17 @@ async def test_deploys_of_the_same_sources_at_once_share_one_build_and_later_one
     assert len(builds.calls) == 1
 
 
+@pytest.mark.asyncio
+async def test_each_modal_environment_builds_its_own_image_of_the_same_sources(local_stores, context, builds):
+    image = _context_image("solver-image", context)
+
+    await _provider().prepare_image(image)
+    await _provider("rl-prod").prepare_image(image)
+    await _provider("rl-prod").prepare_image(image)
+
+    assert len(builds.calls) == 2
+
+
 @pytest.mark.asyncio
 async def test_a_failed_build_names_the_image_and_the_next_deploy_builds_it_again(local_stores, context, builds):
     image = _context_image("solver-image", context)

tst/unit/providers/sandbox_providers/modal_sandbox_test.py

@@ -22,6 +22,16 @@ def test_provider_default_pull_secret_is_none():
     assert ModalSandboxProvider()._ecr_pull_secret_name is None
 
 
+@pytest.mark.asyncio
+async def test_apps_are_looked_up_in_the_configured_environment():
+    provider = ModalSandboxProvider(environment_name="rl-prod")
+    provider._get_client = AsyncMock(return_value=MagicMock())  # type: ignore[method-assign]
+    with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.App.lookup") as lookup:
+        lookup.aio = AsyncMock(return_value=MagicMock())
+        await provider._get_app("agent-env-p")
+    assert lookup.aio.call_args.kwargs["environment_name"] == "rl-prod"
+
+
 @pytest.fixture
 def _patched_provider():
     provider = ModalSandboxProvider(app_name="agent-env-test")
@@ -308,6 +318,7 @@ def mint(self, host):
     ecr_host = "123456789012.dkr.ecr.us-west-2.amazonaws.com"
     ecr = "123456789012.dkr.ecr.us-west-2.amazonaws.com/example/x:v1"
     _patched_provider._ecr_pull_secret_name = "agent-env-ecr-reader"
+    _patched_provider._environment_name = "rl-prod"
     p = lambda n: patch(f"agent_env.providers.sandbox_providers.modal_sandbox.modal.{n}")  # noqa: E731
 
     with p("Image.from_aws_ecr") as aws_ecr, p("Image.from_registry") as registry, \
@@ -340,6 +351,7 @@ def mint(self, host):
     from_name.assert_called_once_with(
         "agent-env-ecr-reader",
         required_keys=["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION"],
+        environment_name="rl-prod",
     )
     assert registry.call_args_list[0].kwargs["secret"] is not None
     assert registry.call_args_list[1].args == ("pub.example/x:v1",)

tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py

@@ -7,7 +7,7 @@
 
 from __future__ import annotations
 
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -385,3 +385,13 @@ async def record(*args):
         c[-1].split(" ", 2)[2].rsplit(" >> ", 1)[0].strip("'") for c in commands if c[-1].startswith("printf '%s' ")
     )
     assert appended == base64.b64encode(data).decode()
+
+
+@pytest.mark.asyncio
+async def test_apps_are_looked_up_in_the_configured_environment():
+    provider = ModalVmSandboxProvider(environment_name="rl-prod")
+    provider._get_client = AsyncMock(return_value=MagicMock())  # type: ignore[method-assign]
+    with patch("agent_env.providers.sandbox_providers.modal_vm_sandbox.modal.App.lookup") as lookup:
+        lookup.aio = AsyncMock(return_value=MagicMock())
+        await provider._get_app("agent-env-p")
+    assert lookup.aio.call_args.kwargs["environment_name"] == "rl-prod"

tst/unit/providers/sandbox_providers/test_config_sandbox_providers.py

@@ -187,6 +187,16 @@ def test_builtin_accepts_a_config_only_table(monkeypatch, tmp_path):
     assert provider._ecr_pull_secret_name == "my-ecr-reader"
 
 
+@pytest.mark.parametrize("name", ["modal", "modal_vm"])
+def test_modal_builtins_accept_an_environment_name(monkeypatch, tmp_path, name):
+    cfg = _write_config(tmp_path, f"""
+        [sandbox.providers.{name}.config]
+        environment_name = "rl-prod"
+    """)
+    monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg))
+    assert build_sandbox_provider(name)._environment_name == "rl-prod"
+
+
 def test_e2b_builtin_receives_interpolated_key_and_versioned_base(monkeypatch, tmp_path):
     cfg = _write_config(tmp_path, """
         [sandbox.providers.e2b.config]