fix(env): a snapshot reload edits the deployed compose file and restarts every service, the browser included

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))

tst/unit/env/envs/test_multi_env_snapshot_reload_name.py

@@ -1,86 +1,136 @@
-"""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
 
 import shlex
 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")