diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index 30262d78..c30406bf 100644 --- a/src/agent_env/cli/task/run.py +++ b/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'}...") diff --git a/src/agent_env/config/__init__.py b/src/agent_env/config/__init__.py index 4f4c9f3f..a34446a6 100644 --- a/src/agent_env/config/__init__.py +++ b/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", diff --git a/src/agent_env/config/runtime.py b/src/agent_env/config/runtime.py index 6c7c5c52..c833de38 100644 --- a/src/agent_env/config/runtime.py +++ b/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. diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 859366e8..b6cfdb1c 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/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 diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index fc532b42..55dc0ed0 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/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] diff --git a/src/agent_env/providers/env_state/local_postgres.py b/src/agent_env/providers/env_state/local_postgres.py index 52e600c3..2e5d5cc4 100644 --- a/src/agent_env/providers/env_state/local_postgres.py +++ b/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 diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 8a73496d..4439aa14 100644 --- a/src/agent_env/task/task.py +++ b/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, diff --git a/tst/unit/config/test_run_env_ids.py b/tst/unit/config/test_run_env_ids.py new file mode 100644 index 00000000..ea7b04ef --- /dev/null +++ b/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