diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index e1e05526..e193771d 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -3,9 +3,12 @@ import logging import os +import re import shlex from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional +import yaml + logger = logging.getLogger(__name__) # Off by default: a bake builds and pushes a multi-GB servicedb image, which is the right @@ -40,12 +43,31 @@ def _snapshot_after_load_default() -> bool: ) from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.env.envs.website import WebsiteEnv -from agent_env.env.envs.playwright_mcp import PLAYWRIGHT_MCP_RECORDER_PORT +from agent_env.env.envs.service_db import DATABASE_SERVICE_NAME +from agent_env.providers.env_providers.constants import DOCKER_COMPOSE_PATH +from agent_env.providers.env_state.local_postgres import INIT_SCRIPT_MOUNT_PATH from agent_env.store.ids import derive_id +from agent_env.store.image_store.oci_registry_credentials import image_reference_problem from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER from agent_env.attribution import Attribution +# The servicedb's image line in a deploy's compose file, as the local Postgres state renders it. +_SERVICE_DB_IMAGE = re.compile(rf"^( {re.escape(DATABASE_SERVICE_NAME)}:\n image: ).*$", re.MULTILINE) + + +def snapshot_compose(compose: str, db_image: str) -> str: + """``compose``, the file a deploy wrote, with its servicedb started from ``db_image`` and without the schema init + script, whose work a snapshot already holds. Every other service stays as the deploy wrote it.""" + if problem := image_reference_problem(db_image): + raise ValueError(problem) + swapped, count = _SERVICE_DB_IMAGE.subn(lambda match: match.group(1) + db_image, compose) + if count != 1: + raise RuntimeError(f"{DOCKER_COMPOSE_PATH} names a {DATABASE_SERVICE_NAME} image {count} times, not once, so " + "the snapshot can't replace it") + return swapped.replace(f" volumes:\n - ./init-schemas.sql:{INIT_SCRIPT_MOUNT_PATH}\n", "") + + class MultiEnv(Env): type: ClassVar[str] = "multi" description = "Multi environment combining multiple MCP servers and websites, deployed through the environment provider its env_provider_type names" @@ -577,17 +599,19 @@ async def load_artifact(self, artifact): raise ValueError(f"{type(self).__name__} '{self.id}' cannot load artifact '{getattr(artifact, 'id', '?')}' of type '{getattr(artifact, 'type', '?')}' — expected a EnvironmentArtifact or EnvironmentUniverseArtifact") async def _load_from_snapshot(self, snapshot) -> None: - """Swap the servicedb container with a pre-loaded snapshot image. Only supported when both the current deploy and the snapshot has env_state_type set to local_postgres""" + """Swap the servicedb container with a pre-loaded snapshot image. Only supported when both the current deploy and the snapshot has env_state_type set to local_postgres. + + The compose file the deploy wrote is edited in place, not rendered again, so every other service, the website + browser included, keeps the image and name it was deployed with whatever config says now.""" import asyncio from agent_env.artifact import Artifact, DockerImageArtifact from agent_env.env.env import Env from agent_env.env.envs.service_db import DB_USER from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT - from agent_env.providers import EnvironmentGatewayProvider, MCPServerConfig, WebsiteConfig - 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 import EnvironmentGatewayProvider + from agent_env.providers.env_providers.constants import GATEWAY_APP_DIR from agent_env.providers.env_state import LocalPostgresStateProvider - from agent_env.config import gateway_env_id, get_config, service_db_env_id, website_browser_env_id + from agent_env.config import get_config, service_db_env_id if self._sandbox is None: raise RuntimeError("Environment not deployed - call deploy() first") @@ -599,44 +623,16 @@ async def _load_from_snapshot(self, snapshot) -> None: logger.info(f"Loading snapshot image {db_image.image_name} into sandbox...") await self._sandbox.load_docker_images([db_image]) - # 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(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 - + logger.info("Pointing docker-compose.yml's servicedb at the snapshot image...") + # The write adds back the newline cat returns, so the file doesn't grow a blank line each reload. + deployed = (await self._sandbox.exec_script(f"cat {DOCKER_COMPOSE_PATH}")).removesuffix("\n") + compose = snapshot_compose(deployed, db_image.image_name) gw = EnvironmentGatewayProvider() gw._sandbox = self._sandbox - gw._state_provider = LocalPostgresStateProvider(service_db_config=service_db_config) + # The state store is the deploy's local Postgres; the changelog triggers below run through it. + gw._state_provider = LocalPostgresStateProvider(service_db_config=Env.get(service_db_env_id(get_config())).to_config()) gw._state_instance = LocalPostgresStateProvider.default_instance() - mcp_servers = [MCPServerConfig(image=env.docker_image_artifact.image_name, environment_name=env.environment_name) for env in self.mcp_server_envs] - website_configs = [ - WebsiteConfig(backend_image=env.backend_docker_image_artifact.image_name, frontend_image=env.frontend_docker_image_artifact.image_name, environment_name=env.environment_name) - for env in self.website_envs - ] - if website_configs: - website_browser_env = Env.get(website_browser_env_id(config)) - mcp_servers = list(mcp_servers) + [MCPServerConfig(image=website_browser_env.docker_image_artifact.image_name, environment_name=website_browser_env.environment_name, - card_port=PLAYWRIGHT_MCP_RECORDER_PORT)] - - compose_content = gw.create_docker_compose( - mcp_servers=mcp_servers, - gateway_image=gateway_env.docker_image_artifact.image_name, - gateway_port=AGENT_ENV_GATEWAY_MCP_PORT, website_configs=website_configs or None, - gateway_mode=self._gateway_mode, - state_provider=gw._state_provider, # local-only, set above - state_instance=gw._state_instance, - host_port=self._sandbox.host_port, - host_ips=self._sandbox.host_ips, - extra_hosts=self._sandbox.extra_hosts, - mcp_server_name=self._mcp_server_name or self.name, - ) - compose_content = compose_content.replace( - " volumes:\n - ./init-schemas.sql:/docker-entrypoint-initdb.d/init-schemas.sql\n", "", - ) - await gw._write_file(self._sandbox, DOCKER_COMPOSE_PATH, compose_content) + await gw._write_file(self._sandbox, DOCKER_COMPOSE_PATH, compose) # Remove old servicedb container and its anonymous volume, then start fresh from snapshot image logger.info("Recreating servicedb container with snapshot image...") @@ -660,13 +656,15 @@ async def _load_from_snapshot(self, snapshot) -> None: else: raise RuntimeError("servicedb not healthy after 60s") - # Recreate non-servicedb services so they reconnect to the new servicedb. + # Recreate every other service the compose file runs so it reconnects to the new servicedb. # Uses `up -d --force-recreate` which respects depends_on ordering # (MCP servers start before gateway), unlike `restart` which is unordered. environment_names = list(dict.fromkeys( [env.environment_name for env in self.mcp_server_envs] + [env.environment_name for env in self.website_envs] )) - svc_list = " ".join(map(shlex.quote, environment_names + [GATEWAY_SERVICE_NAME, PGWEB_SERVICE_NAME, DB_MCP_SERVICE_NAME])) + # BaseLoader keeps every name a string: safe_load reads a service named `on` as True and `123` as an int. + svc_list = " ".join(map(shlex.quote, (name for name in yaml.load(compose, Loader=yaml.BaseLoader)["services"] + if name != DATABASE_SERVICE_NAME))) logger.info(f"Recreating services: {svc_list}") await gw._compose_up(self._sandbox, f"--force-recreate {svc_list}") await gw._wait_for_gateway(self._sandbox, self._sandbox.host_port(AGENT_ENV_GATEWAY_MCP_PORT)) diff --git a/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py b/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py index 3ee8f8da..26a5fd4f 100644 --- a/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py +++ b/tst/unit/env/envs/test_multi_env_snapshot_reload_name.py @@ -1,4 +1,5 @@ -"""The snapshot reload regenerates the compose; it must reuse the name and host IPs the running gateway already has.""" +"""The snapshot reload edits the compose file the deploy wrote: it starts the servicedb from the snapshot and leaves +every other service as deployed, whatever config names now.""" from __future__ import annotations @@ -6,81 +7,130 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +import yaml from agent_env.artifact import DockerImageArtifact -from agent_env.env.envs.multi_env import MultiEnv +from agent_env.env.envs.multi_env import MultiEnv, snapshot_compose +from agent_env.env.envs.service_db import ServiceDBConfig from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT -from agent_env.providers.env_providers import EnvironmentGatewayProvider +from agent_env.providers import EnvironmentGatewayProvider, MCPServerConfig, WebsiteConfig from agent_env.providers.env_providers.constants import DOCKER_COMPOSE_PATH +from agent_env.providers.env_state import LocalPostgresStateProvider class _Stop(Exception): - """Abort right after the compose is rendered; the rest needs a real sandbox.""" + """Abort once the gateway is waited on; the rest needs a real sandbox.""" -@pytest.mark.asyncio -@pytest.mark.parametrize("declared, deployed_as, expected", [("crm", None, "crm"), (None, "env4821", "env4821"), ("crm", "crm", "crm")]) -async def test_snapshot_reload_keeps_the_env_name(declared, deployed_as, expected): - env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name=declared) - env._mcp_server_name = deployed_as - env._sandbox = MagicMock(mode="vm", load_docker_images=AsyncMock(), host_ips=("127.0.0.1",)) - db_image = MagicMock(spec=DockerImageArtifact, image_name="snap:1") - snapshot = MagicMock(db_image_artifact_id="snap", db_image_artifact_version=1) - gateway_env, service_db = MagicMock(), MagicMock() +def _deployed_compose(browser="website-browser"): + """The compose file a deploy with a website writes, its browser named ``browser``.""" + return EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack"), + MCPServerConfig(image="browser:v1", environment_name=browser, card_port=18766)], + website_configs=[WebsiteConfig(backend_image="shop-be", frontend_image="shop-fe", environment_name="shop")], + gateway_image="gateway:v38", gateway_port=AGENT_ENV_GATEWAY_MCP_PORT, + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig( + db_web_image="db-web:v1", db_mcp_image="db-mcp:v1")), + state_instance=LocalPostgresStateProvider.default_instance(), mcp_server_name="crm") - with patch("agent_env.artifact.Artifact.get", return_value=db_image), \ - patch("agent_env.env.env.Env.get", side_effect=lambda env_id, *a, **k: gateway_env if env_id == "gw-id" else service_db), \ - patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id", default_service_db_env_id="db-id")), \ - patch("agent_env.providers.env_state.LocalPostgresStateProvider"), \ - patch.object(EnvironmentGatewayProvider, "create_docker_compose", side_effect=_Stop) as compose: - with pytest.raises(_Stop): - await env._load_from_snapshot(snapshot) - assert compose.call_args.kwargs["mcp_server_name"] == expected - assert compose.call_args.kwargs["host_ips"] == ("127.0.0.1",) - assert compose.call_args.kwargs["host_port"] == env._sandbox.host_port +def _env(compose): + env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name="crm") + env.mcp_server_envs = [MagicMock(environment_name="slack")] + env.website_envs = [MagicMock(environment_name="shop")] + env._sandbox = MagicMock(mode="vm", load_docker_images=AsyncMock(), host_ips=("127.0.0.1",), + host_port=lambda port: port + 1000, + exec_script=AsyncMock(side_effect=lambda script, **k: compose + "\n" if script.startswith("cat ") else ""), + exec_with_output=AsyncMock(return_value=(0, "", ""))) + return env -@pytest.mark.asyncio -async def test_snapshot_reload_waits_for_the_gateway_where_it_is_published(): - env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name="crm") - env._sandbox = MagicMock(mode="vm", load_docker_images=AsyncMock(), exec_script=AsyncMock(return_value=""), - host_ips=("127.0.0.1",), host_port=lambda port: port + 1000) +async def _reload(env): db_image = MagicMock(spec=DockerImageArtifact, image_name="snap:1") snapshot = MagicMock(db_image_artifact_id="snap", db_image_artifact_version=1) - with patch("agent_env.artifact.Artifact.get", return_value=db_image), \ - patch("agent_env.env.env.Env.get", return_value=MagicMock()), \ - patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id", default_service_db_env_id="db-id")), \ + patch("agent_env.env.env.Env.get", return_value=MagicMock()) as get, \ + patch("agent_env.config.get_config", return_value=MagicMock(default_service_db_env_id="db-id")), \ patch("agent_env.providers.env_state.LocalPostgresStateProvider"), \ - patch.object(EnvironmentGatewayProvider, "create_docker_compose", return_value="services:\n"), \ + patch.object(EnvironmentGatewayProvider, "_compose_up", AsyncMock()) as up, \ patch.object(EnvironmentGatewayProvider, "_wait_for_gateway", side_effect=_Stop) as wait: with pytest.raises(_Stop): await env._load_from_snapshot(snapshot) + return get, up, wait + + +def _written(env): + (written,) = [argv[2] for argv in map(shlex.split, (call.args[0] for call in env._sandbox.exec_script.await_args_list)) + if argv[:1] == ["printf"]] + return written + + +@pytest.mark.asyncio +@pytest.mark.parametrize("browser", ["website-browser", "browseruse-mcp"]) +async def test_the_reload_starts_the_servicedb_from_the_snapshot_and_keeps_every_other_service_as_deployed(browser): + deployed = _deployed_compose(browser) + env = _env(deployed) + + get, _, _ = await _reload(env) + + assert _written(env) == snapshot_compose(deployed, "snap:1") # byte for byte, with no blank line added + before, after = yaml.safe_load(deployed)["services"], yaml.safe_load(_written(env))["services"] + assert after["servicedb"]["image"] == "snap:1" and "volumes" not in after["servicedb"] + assert {name: service for name, service in after.items() if name != "servicedb"} == { + name: service for name, service in before.items() if name != "servicedb"} + assert [call.args[0] for call in get.call_args_list] == ["db-id"] # the store it connects to; no image is read + + +@pytest.mark.asyncio +async def test_the_reload_recreates_every_service_but_the_servicedb_the_browser_included(): + env = _env(_deployed_compose("browseruse-mcp")) + + _, up, wait = await _reload(env) + (args,) = [call.args[1] for call in up.await_args_list] + assert set(shlex.split(args.removeprefix("--force-recreate "))) == { + "slack", "browseruse-mcp", "shop-website-backend", "shop-website-frontend", "gateway", "pgweb", "db-mcp"} assert wait.call_args.args[1] == AGENT_ENV_GATEWAY_MCP_PORT + 1000 @pytest.mark.asyncio -async def test_snapshot_reload_writes_its_compose_as_one_argument_and_names_each_service_as_one(): - env = MultiEnv(id="crm-suite", version=1, mcp_server_envs=[], name="crm") - env.mcp_server_envs = [MagicMock(environment_name="a b")] - env._sandbox = MagicMock(mode="vm", load_docker_images=AsyncMock(), exec_script=AsyncMock(return_value=""), - host_ips=("127.0.0.1",), host_port=lambda port: port) - db_image = MagicMock(spec=DockerImageArtifact, image_name="snap:1") - snapshot = MagicMock(db_image_artifact_id="snap", db_image_artifact_version=1) - compose = "services:\nCOMPOSE_EOF\ntouch pwned\n" +@pytest.mark.parametrize("name", ["on", "123", "null"]) +async def test_a_service_whose_name_yaml_reads_as_another_type_is_recreated_by_its_name(name): + env = _env(_deployed_compose().replace("slack", name)) - with patch("agent_env.artifact.Artifact.get", return_value=db_image), \ - patch("agent_env.env.env.Env.get", return_value=MagicMock()), \ - patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id", default_service_db_env_id="db-id")), \ - patch("agent_env.providers.env_state.LocalPostgresStateProvider"), \ - patch.object(EnvironmentGatewayProvider, "create_docker_compose", return_value=compose), \ - patch.object(EnvironmentGatewayProvider, "_wait_for_gateway", side_effect=_Stop): - with pytest.raises(_Stop): - await env._load_from_snapshot(snapshot) + _, up, _ = await _reload(env) + + (args,) = [call.args[1] for call in up.await_args_list] + assert name in shlex.split(args.removeprefix("--force-recreate ")) + +@pytest.mark.asyncio +async def test_the_reload_writes_its_compose_as_one_argument(): + env = _env(_deployed_compose().replace(" slack:\n", " slack:\n labels:\n note: \"$(touch pwned) 'q' `x`\"\n", 1)) + + await _reload(env) + + assert "$(touch pwned) 'q' `x`" in _written(env) scripts = [call.args[0] for call in env._sandbox.exec_script.await_args_list] - assert ["printf", "%s\\n", compose, ">", DOCKER_COMPOSE_PATH] in map(shlex.split, scripts) + assert any(argv[:2] == ["printf", "%s\\n"] and argv[3:] == [">", DOCKER_COMPOSE_PATH] + for argv in map(shlex.split, scripts)) assert not any("<<" in script for script in scripts) - assert any(script.endswith("--force-recreate 'a b' gateway pgweb db-mcp") for script in scripts) + + +def test_a_second_reload_swaps_the_image_again(): + once = snapshot_compose(_deployed_compose(), "snap:1") + assert yaml.safe_load(snapshot_compose(once, "snap:2"))["services"]["servicedb"]["image"] == "snap:2" + + +@pytest.mark.parametrize("compose, problem", [ + ("services:\n slack:\n image: mcp-slack\n", "names a servicedb image 0 times"), + ("services:\n servicedb:\n image: a\n servicedb:\n image: b\n", "names a servicedb image 2 times"), +]) +def test_a_compose_file_without_exactly_one_servicedb_image_is_refused(compose, problem): + with pytest.raises(RuntimeError, match=problem): + snapshot_compose(compose, "snap:1") + + +def test_a_snapshot_image_that_would_change_what_the_file_says_is_refused(): + with pytest.raises(ValueError): + snapshot_compose(_deployed_compose(), "snap:1\n evil:\n image: x")