diff --git a/src/agent_env/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index 238db5aa..c6aa58dd 100644 --- a/src/agent_env/env/envs/mcp_server.py +++ b/src/agent_env/env/envs/mcp_server.py @@ -97,12 +97,13 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: @environment_card when it isn't set. Raises BundleError listing every problem.""" return ctx.accept_env(data, cls, named_by="image") - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: # In container mode the server runs in its own container, so loads stage there, as for a MultiEnv child. return await deploy_through_provider( self, environment_name=self.environment_name, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, + artifact_id=artifact_id, artifact_version=artifact_version, ) def deploy_refusal(self, **options) -> str | None: diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index 040db384..d50dabd1 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -118,11 +118,12 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: ctx.refuse(problems) return fields - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: deployed_env = await deploy_through_provider( self, environment_name=None, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, + artifact_id=artifact_id, artifact_version=artifact_version, ) self._gateway_mode = gateway_mode self._mcp_server_name = deployed_env.mcp_server_name diff --git a/src/agent_env/env/envs/website.py b/src/agent_env/env/envs/website.py index 5d693bb4..46e2b2a7 100644 --- a/src/agent_env/env/envs/website.py +++ b/src/agent_env/env/envs/website.py @@ -116,11 +116,12 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: @environment_card when it isn't set. Raises BundleError listing every problem.""" return ctx.accept_env(data, cls, named_by="backend_image") - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: return await deploy_through_provider( self, environment_name=self.environment_name, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, + artifact_id=artifact_id, artifact_version=artifact_version, ) def deploy_refusal(self, **options) -> str | None: 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 6b164b56..5eb7d900 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -521,6 +521,8 @@ async def deploy( cpu: float | None = None, memory_mb: int | None = None, attribution: Optional[Attribution] = None, + artifact_id: str | None = None, + artifact_version: int | None = None, ) -> DeployedGatewayEnv: """A gateway in front of the env, with the state store it acquires (local Postgres by default); returns the record. diff --git a/src/agent_env/providers/env_providers/env_server_provider.py b/src/agent_env/providers/env_providers/env_server_provider.py index b0ee71c5..193a3dec 100644 --- a/src/agent_env/providers/env_providers/env_server_provider.py +++ b/src/agent_env/providers/env_providers/env_server_provider.py @@ -36,6 +36,8 @@ async def deploy( cpu: float | None = None, memory_mb: int | None = None, attribution: Optional[Attribution] = None, + artifact_id: str | None = None, + artifact_version: int | None = None, ) -> DeployedSandboxEnv: """The env's MCP server in its own container; returns the record once it serves its card and its tools.""" from agent_env.env.envs.multi_env import MultiEnv diff --git a/tst/unit/env/envs/test_mcp_server_plugin_provider.py b/tst/unit/env/envs/test_mcp_server_plugin_provider.py index 86cbdd70..4a157460 100644 --- a/tst/unit/env/envs/test_mcp_server_plugin_provider.py +++ b/tst/unit/env/envs/test_mcp_server_plugin_provider.py @@ -113,7 +113,8 @@ async def test_a_plugin_provider_gets_every_option_and_its_record_is_registered( [(sandbox_provider, options)] = SEEN["deploys"] assert sandbox_provider == "SANDBOXES" assert options == {"ttl_seconds": 60, "disk_size_gb": 10, "gateway_mode": GatewayMode.CONSISTENT, "cpu": 2.0, "memory_mb": None, - "env_state_type": None, "env_state_instance_id": None, "attribution": {"team": "t"}} + "env_state_type": None, "env_state_instance_id": None, "attribution": {"team": "t"}, + "artifact_id": None, "artifact_version": None} assert (record.instance_id, env._instance_id, env._deployed, env._sandbox, env._gateway_url) == ("inst-1", "inst-1", record, None, None) diff --git a/tst/unit/env/envs/test_multi_env_plugin_provider.py b/tst/unit/env/envs/test_multi_env_plugin_provider.py index 37b3cbef..4fa9b4a5 100644 --- a/tst/unit/env/envs/test_multi_env_plugin_provider.py +++ b/tst/unit/env/envs/test_multi_env_plugin_provider.py @@ -166,7 +166,8 @@ def register(record, ttl_seconds): [deployed] = context.deployed_envs [(deployed_env, options)] = SEEN["deploys"] assert deployed_env is env and set(options) == {"ttl_seconds", "disk_size_gb", "gateway_mode", "cpu", "memory_mb", - "env_state_type", "env_state_instance_id", "attribution"} + "env_state_type", "env_state_instance_id", "attribution", + "artifact_id", "artifact_version"} assert (deployed.instance_id, deployed.mcp_url, env._sandbox, env._deployed) == ("inst-7", f"{_URL}/mcp", None, deployed) for child in [*env.mcp_server_envs, *env.website_envs]: assert (child._env_provider, child._deployed, child._sandbox) == (env._env_provider, deployed, None) @@ -583,3 +584,28 @@ async def test_a_deploy_onto_the_state_instance_the_env_s_deployment_holds_is_re pytest.raises(ValueError, match="already has a deployment, instance 'inst-1', on env state instance 'esi-1'"): await env.deploy(env_state_instance_id="esi-1") assert (sandboxes.called, "deploys" in SEEN) == (False, False) + + +@pytest.mark.parametrize("make_env", [ + pytest.param(lambda: _multi(), id="multi"), + pytest.param(lambda: MCPServerEnv(id="mcp-slack", version=1, environment_name="slack", + docker_image_artifact=MagicMock(image_name="mcp-slack"), + env_provider_type="plugin_pods"), id="mcp_server"), + pytest.param(lambda: _site("shop", "plugin_pods"), id="website"), +]) +@pytest.mark.asyncio +async def test_every_built_in_env_forwards_the_universe_artifact_to_its_provider(make_env): + """A warm pool keys on env *and* universe, so the provider has to be told which universe + the run will load. All three envs carry it the same way, so all three are checked: a value + dropped in one would otherwise go unnoticed. Absent, both arrive as None — the pool then + treats the env as unpooled rather than guessing, and every existing caller is unaffected.""" + SEEN.clear() + SEEN["deploys"] = [] + await _deploy(make_env(), artifact_id="hg4_real", artifact_version=5) + [(_, options)] = SEEN["deploys"] + assert (options["artifact_id"], options["artifact_version"]) == ("hg4_real", 5) + + SEEN["deploys"] = [] + await _deploy(make_env()) + [(_, options)] = SEEN["deploys"] + assert (options["artifact_id"], options["artifact_version"]) == (None, None)