feat(task): a run's gateway and service-db env overrides travel with it to any runner

src/agent_env/cli/task/run.py

@@ -445,16 +445,9 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int,
     if k < 1:
         raise click.BadParameter("must be at least 1", param_hint="'--k'")
 
-    from agent_env.config import get_config
     from agent_env.task import Task
     from agent_env.task_step.context import TaskStepContext
 
-    if gateway_env_id:
-        get_config().default_gateway_env_id = gateway_env_id
-        click.echo(f"Override default_gateway_env_id: {gateway_env_id}")
-    if service_db_env_id:
-        get_config().default_service_db_env_id = service_db_env_id
-        click.echo(f"Override default_service_db_env_id: {service_db_env_id}")
     if output_dir is None:
         output_dir = f"/tmp/taskid={fs_safe(task_id)}-run{uuid4().hex[:6]}"
     if context_json:
@@ -486,6 +479,10 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int,
         initial_context.metadata.setdefault("user_overrides", {})["env_state_type"] = env_state_type
     if env_state_instance_id:
         initial_context.metadata.setdefault("user_overrides", {})["env_state_instance_id"] = env_state_instance_id
+    if gateway_env_id:
+        initial_context.metadata.setdefault("user_overrides", {})["gateway_env_id"] = gateway_env_id
+    if service_db_env_id:
+        initial_context.metadata.setdefault("user_overrides", {})["service_db_env_id"] = service_db_env_id
     initial_context.metadata["task_id"] = task_id
 
     click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...")

src/agent_env/config/__init__.py

@@ -18,9 +18,12 @@
 from agent_env.config.runtime import (
     Config,
     configure,
+    gateway_env_id,
     get_config,
     get_runner,
     reset_config,
+    run_env_ids,
+    service_db_env_id,
     set_document_store,
     set_image_store,
     set_object_store,
@@ -39,9 +42,12 @@
     "load_impl",
     "Config",
     "configure",
+    "gateway_env_id",
     "get_config",
     "get_runner",
     "reset_config",
+    "run_env_ids",
+    "service_db_env_id",
     "set_document_store",
     "set_image_store",
     "set_object_store",

src/agent_env/config/runtime.py

@@ -28,9 +28,12 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import os
 import threading
+from collections.abc import Iterator
+from contextvars import ContextVar
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Mapping, Optional
@@ -1028,6 +1031,31 @@ def get_runner():
     return get_config().get_runner()
 
 
+_RUN_ENV_IDS: ContextVar[Mapping[str, str]] = ContextVar("agent_env_run_env_ids", default={})
+
+
+@contextlib.contextmanager
+def run_env_ids(user_overrides: Mapping[str, Any]) -> Iterator[None]:
+    """Within a run, the ``gateway_env_id`` and ``service_db_env_id`` its ``user_overrides`` name win over the
+    configured defaults, in whichever process runs it."""
+    token = _RUN_ENV_IDS.set(
+        {key: user_overrides[key] for key in ("gateway_env_id", "service_db_env_id") if user_overrides.get(key)})
+    try:
+        yield
+    finally:
+        _RUN_ENV_IDS.reset(token)
+
+
+def gateway_env_id(config: Config) -> str:
+    """The gateway env a deploy uses: the run's override, else ``config.default_gateway_env_id``."""
+    return _RUN_ENV_IDS.get().get("gateway_env_id") or config.default_gateway_env_id
+
+
+def service_db_env_id(config: Config) -> str:
+    """The service-db env a deploy uses: the run's override, else ``config.default_service_db_env_id``."""
+    return _RUN_ENV_IDS.get().get("service_db_env_id") or config.default_service_db_env_id
+
+
 def get_config() -> Config:
     """The process-wide Config, built on first use.
 

src/agent_env/env/envs/multi_env.py

@@ -574,7 +574,7 @@ async def _load_from_snapshot(self, snapshot) -> None:
         from agent_env.env.envs.service_db import DATABASE_SERVICE_NAME, DB_MCP_SERVICE_NAME, PGWEB_SERVICE_NAME
         from agent_env.providers.env_providers.constants import DOCKER_COMPOSE_PATH, GATEWAY_APP_DIR, GATEWAY_SERVICE_NAME
         from agent_env.providers.env_state import LocalPostgresStateProvider
-        from agent_env.config import get_config
+        from agent_env.config import gateway_env_id, get_config, service_db_env_id
 
         if self._sandbox is None:
             raise RuntimeError("Environment not deployed - call deploy() first")
@@ -589,8 +589,8 @@ async def _load_from_snapshot(self, snapshot) -> None:
         # Regenerate docker-compose.yml with the snapshot image (no init-schemas volume needed)
         logger.info("Regenerating docker-compose.yml with snapshot image...")
         config = get_config()
-        gateway_env = Env.get(config.default_gateway_env_id)
-        service_db = Env.get(config.default_service_db_env_id)
+        gateway_env = Env.get(gateway_env_id(config))
+        service_db = Env.get(service_db_env_id(config))
         service_db_config = service_db.to_config()
         service_db_config.db_image = db_image.image_name
 

src/agent_env/providers/env_providers/env_gateway_provider.py

@@ -573,7 +573,7 @@ async def deploy(
 
         A MultiEnv's gateway fronts its MCP servers and websites, a WebsiteEnv's its website, and any other env's its one MCP server.
         """
-        from agent_env.config import get_config
+        from agent_env.config import get_config, service_db_env_id
         from agent_env.env.env import Env
         from agent_env.providers.env_state import LOCAL_POSTGRES_STATE_TYPE, acquire_state_for_deploy
 
@@ -582,7 +582,7 @@ async def deploy(
             env_state_type=env_state_type, ttl_seconds=ttl_seconds, name_hint=env.id, env_state_instance_id=env_state_instance_id,
         )
         if state_instance is None or state_instance.state_type == LOCAL_POSTGRES_STATE_TYPE:
-            Env.get(get_config().default_service_db_env_id)  # the local store runs from it: check it before any sandbox
+            Env.get(service_db_env_id(get_config()))  # the local store runs from it: check it before any sandbox
         deployed = await self._launch(
             sandbox_provider, env.id, state_instance=state_instance, mcp_servers=topology.mcp_servers,
             mcp_server_images=topology.mcp_server_images, website_configs=topology.website_configs,
@@ -706,14 +706,14 @@ async def _deploy_via_vm(
         mcp_server_name: str | None = None,
     ) -> DeployedGateway:
         from agent_env.env.env import Env
-        from agent_env.config import get_config
+        from agent_env.config import gateway_env_id, get_config
 
         website_configs = website_configs or []
         website_images = website_images or []
         sidecars = sidecars or []
 
         config = get_config()
-        gateway_env = Env.get(config.default_gateway_env_id)
+        gateway_env = Env.get(gateway_env_id(config))
 
         # Auto-include website_browser MCP server when websites are present
         if website_configs:
@@ -910,10 +910,9 @@ async def _deploy_via_containers(
         mcp_server_name: str | None = None,
     ) -> DeployedGateway:
         from agent_env.env.env import Env
-        from agent_env.config import get_config
+        from agent_env.config import gateway_env_id, get_config
 
-        config = get_config()
-        gateway_env = Env.get(config.default_gateway_env_id)
+        gateway_env = Env.get(gateway_env_id(get_config()))
         # Each image is checked, and prepared (one the provider can't pull, built from its build context), before
         # anything is created, so an image the provider can't run, or a failed build, leaves nothing.
         images = [gateway_env.docker_image_artifact, *mcp_server_images]

src/agent_env/providers/env_state/local_postgres.py

@@ -120,9 +120,9 @@ def _default_service_db_env(self) -> "ServiceDBEnv":
         ``service_db_config`` off the same resolve unless one was injected at construction."""
         if self._service_db_env is None:
             from agent_env.env.env import Env
-            from agent_env.config import get_config
+            from agent_env.config import get_config, service_db_env_id
 
-            self._service_db_env = Env.get(get_config().default_service_db_env_id)
+            self._service_db_env = Env.get(service_db_env_id(get_config()))
             if self._service_db_config is None:
                 self._service_db_config = self._service_db_env.to_config()
         return self._service_db_env

src/agent_env/task/task.py

@@ -15,6 +15,7 @@
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Callable, ClassVar, Optional
 
+from agent_env.config import run_env_ids
 from agent_env.plugins import _registration
 from agent_env.store.document_store import UnknownDocumentTypeError
 from agent_env.store.routing import run_scope
@@ -415,7 +416,7 @@ async def run(
         context.metadata['failed_steps']. The run is scoped to this task's id, so under namespace
         routing an ``@local`` task's records stay in the local stores.
         """
-        with run_scope(self.id):
+        with run_scope(self.id), run_env_ids((context.metadata.get("user_overrides") or {}) if context else {}):
             return await self._run(
                 on_step_start, on_step_complete, agent_model, agent_artifact_id,
                 start_step, end_step, context, instance_id,

tst/unit/config/test_run_env_ids.py

@@ -0,0 +1,39 @@
+"""A run's ``gateway_env_id`` / ``service_db_env_id`` overrides name the infra envs its deploys use, only within it."""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from agent_env.config import gateway_env_id, get_config, service_db_env_id, set_document_store
+from agent_env.store import LocalSqliteDocumentStore
+from agent_env.task.task import Task
+from agent_env.task_step.context import TaskStepContext
+
+
+class _Probe:
+    type = "probe"
+    id, depends_on, retry_config, fail_task_on_error = "probe", None, None, True
+
+    def __init__(self, both_running: asyncio.Barrier):
+        self.both_running, self.seen = both_running, None
+
+    async def execute(self, context: TaskStepContext) -> TaskStepContext:
+        await self.both_running.wait()
+        self.seen = (gateway_env_id(get_config()), service_db_env_id(get_config()))
+        return context
+
+
+@pytest.mark.asyncio
+async def test_concurrent_runs_each_use_their_own_infra_envs_and_only_within_the_run(tmp_path):
+    set_document_store(LocalSqliteDocumentStore(str(tmp_path / "runs.db")))
+    both_running = asyncio.Barrier(2)
+    probes = {name: _Probe(both_running) for name in ("a", "b")}
+
+    await asyncio.gather(*(Task(id=f"t-{name}", version=1, steps=[probe]).run(context=TaskStepContext(metadata={
+        "user_overrides": {"gateway_env_id": f"gw-{name}", "service_db_env_id": f"db-{name}"}}))
+        for name, probe in probes.items()))
+
+    assert {name: probe.seen for name, probe in probes.items()} == {"a": ("gw-a", "db-a"), "b": ("gw-b", "db-b")}
+    assert gateway_env_id(get_config()) == get_config().default_gateway_env_id