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 e9c92fb03..4cd2b4a30 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -7,6 +7,7 @@ import json import logging import re +import secrets import shlex import time from collections import Counter @@ -72,6 +73,9 @@ _TRAJECTORY_READ_ATTEMPTS = 3 +# Random bytes in the path a deploy serves its SQL consoles under (24 URL-safe characters). +_CONSOLE_PATH_BYTES = 18 + # How much of a service's output a failed deploy's error carries: enough for the end of a traceback. _LOG_TAIL_LINES = 30 _LOG_LINE_CHARS = 300 @@ -874,13 +878,14 @@ async def _stand_up(_spec): logger.warning(f"{name} container not found") gateway_url = sandbox.tunnel_urls.get(gateway_port) + console = f"/{self._state_provider.console_path}" if sidecar_names and self._state_provider.console_path else "" db_web_base_url = sandbox.tunnel_urls.get(DB_WEB_PORT) if PGWEB_SERVICE_NAME in sidecar_names else None - db_web_url = f"{db_web_base_url}/" if db_web_base_url else None + db_web_url = f"{db_web_base_url}{console}/" if db_web_base_url else None db_mcp_base_url = sandbox.tunnel_urls.get(DB_MCP_PORT) if DB_MCP_SERVICE_NAME in sidecar_names else None - db_mcp_url = f"{db_mcp_base_url}/mcp" if db_mcp_base_url else None - logger.info(f"pgweb URL: {db_web_url}") + db_mcp_url = f"{db_mcp_base_url}{console}/mcp" if db_mcp_base_url else None + logger.info(self._masked(f"pgweb URL: {db_web_url}")) if db_mcp_url: - logger.info(f"DB MCP endpoint: {db_mcp_url}") + logger.info(self._masked(f"DB MCP endpoint: {db_mcp_url}")) # Wait for the tunnel to serve the env card, which the deploy records card = await self._wait_for_tunnel(gateway_url, timeout=120) @@ -1040,7 +1045,8 @@ async def _deploy_via_containers( card_read_at_utc = datetime.now(timezone.utc).isoformat() await self._refuse_tool_name_clashes(gateway_url) - db_web_url = pgweb_sb.tunnel_urls[DB_WEB_PORT] + "/" if pgweb_sb else None + console = f"/{self._state_provider.console_path}" if pgweb_sb and self._state_provider.console_path else "" + db_web_url = f"{pgweb_sb.tunnel_urls[DB_WEB_PORT]}{console}/" if pgweb_sb else None db_mcp_url = db_mcp_sb.tunnel_urls[DB_MCP_CONTAINER_PORT] + "/mcp" if db_mcp_sb else None return DeployedGateway( gateway_url=gateway_url, @@ -1179,9 +1185,14 @@ async def _log_tail(self, sandbox: VmSandbox, container_id: str) -> str | None: except Exception as e: logger.warning(f"Can't read container {container_id}'s logs: {e}") return None - lines = [line[:_LOG_LINE_CHARS] for line in logs.splitlines() if line.strip()] if exit_code == 0 else [] + lines = [line[:_LOG_LINE_CHARS] for line in self._masked(logs).splitlines() if line.strip()] if exit_code == 0 else [] return "\n".join(f" {line}" for line in lines) or None + def _masked(self, text: str) -> str: + """``text`` with the consoles' path masked: pgweb prints it as it starts, and both log the paths they serve.""" + path = getattr(self._state_provider, "console_path", None) + return text.replace(path, "***") if path else text + async def _log_service(self, sandbox: VmSandbox, compose_service: str) -> bool: """Log a compose service's container output; False when it has no container. Diagnostics only, so a lookup or ``docker logs`` that fails is reported as a warning, never raised into the deploy.""" @@ -1196,9 +1207,9 @@ async def _log_service(self, sandbox: VmSandbox, compose_service: str) -> bool: if exit_code != 0: logger.warning(f"Skipping {compose_service} logs: docker logs exited {exit_code}: {stderr.strip()[-500:]}") return True - logger.info(f" stdout:\n{logs}") + logger.info(f" stdout:\n{self._masked(logs)}") if stderr: - logger.info(f" stderr:\n{stderr}") + logger.info(f" stderr:\n{self._masked(stderr)}") return True async def _get_container_id(self, sandbox: VmSandbox, compose_service: str) -> str | None: @@ -1365,6 +1376,8 @@ async def _deploy_gateway( if state_instance is not None else LocalPostgresStateProvider() ) + if isinstance(self._state_provider, LocalPostgresStateProvider): + self._state_provider.console_path = secrets.token_urlsafe(_CONSOLE_PATH_BYTES) deploy_kwargs = {"attribution": dict(attribution or {})} if not self.in_vm(sandbox_provider): if problem := self.sandbox_problem(sandbox_provider, websites=bool(website_configs)): @@ -1544,7 +1557,8 @@ def _compose_name_problems(mcp_servers: list[MCPServerConfig], website_configs: # Compose `environment:` keys whose values are secrets and must be masked before the # rendered compose is logged. Matches the key to the left of the first '=' on an # ` - KEY=VALUE` env line. -_SECRET_ENV_KEY_RE = re.compile(r"^(\s*-\s*)([A-Za-z0-9_]*(?:SECRET|PASSWORD|TOKEN|API_KEY)[A-Za-z0-9_]*)=.*$") +_SECRET_ENV_KEY_RE = re.compile( + r"^(\s*-\s*)([A-Za-z0-9_]*(?:SECRET|PASSWORD|TOKEN|API_KEY|URL_PREFIX|DB_MCP_PATH)[A-Za-z0-9_]*)=.*$") def _compose_ps_rows(text: str) -> list[dict]: diff --git a/src/agent_env/providers/env_state/local_postgres.py b/src/agent_env/providers/env_state/local_postgres.py index 154d1e120..372e6de8f 100644 --- a/src/agent_env/providers/env_state/local_postgres.py +++ b/src/agent_env/providers/env_state/local_postgres.py @@ -13,6 +13,7 @@ from __future__ import annotations +import json import logging from dataclasses import dataclass from typing import Optional, TYPE_CHECKING @@ -55,6 +56,14 @@ _STOCK_POSTGRES_IMAGE = "public.ecr.aws/docker/library/postgres:16-alpine" +# postgres-mcp takes no option for its MCP path, so db-mcp starts through this, which sets the path from DB_MCP_PATH +# first. It needs no change to the image, so every db-mcp image already promoted runs it. +_DB_MCP_ENTRYPOINT = [ + "python", "-c", + "import asyncio, os; from postgres_mcp import server; " + "server.mcp.settings.streamable_http_path = os.environ['DB_MCP_PATH']; asyncio.run(server.main())", +] + @dataclass class ContainerSpec: @@ -108,6 +117,9 @@ def __init__(self, service_db_config: "ServiceDBConfig | None" = None): super().__init__() self._service_db_config = service_db_config self._service_db_env: "ServiceDBEnv | None" = None + # The path pgweb and db-mcp answer under. A gateway deploy draws one, so that only the URLs it records reach + # consoles whose ports are published, and on a VM tunnelled; None serves them at their roots. + self.console_path: str | None = None def deploy_state_context( self, *, ttl_seconds: int = DEFAULT_STATE_TTL_SECONDS, name_hint: str | None = None @@ -278,6 +290,7 @@ def render_sidecar_containers( # sidecar_specs for why: Disconnect tore down pgweb's shared # session and 500'd the changelog export on Save & Next). " - PGWEB_LOCK_SESSION=true", + *([f" - PGWEB_URL_PREFIX={self.console_path}"] if self.console_path else []), " ports:", *(f' - "{spec}"' for spec in port_bindings(host_ips, publish(DB_WEB_PORT), DB_WEB_PORT)), " networks:", @@ -294,6 +307,8 @@ def render_sidecar_containers( " condition: service_healthy", " environment:", f" - DATABASE_URI={self._pg_url(DATABASE_SERVICE_NAME, scheme='postgresql')}", + *([f" - DB_MCP_PATH=/{self.console_path}/mcp", + f" entrypoint: {json.dumps(_DB_MCP_ENTRYPOINT)}"] if self.console_path else []), " command: --transport=streamable-http --streamable-http-host=0.0.0.0 --access-mode=unrestricted", " ports:", *(f' - "{spec}"' for spec in port_bindings(host_ips, publish(DB_MCP_PORT), DB_MCP_CONTAINER_PORT)), @@ -360,7 +375,8 @@ def sidecar_specs(self, *, instance: EnvStateInstance | None = None) -> list[Con # session — breaking both the Raw DB view AND the hub's changelog # export (it queries the same pgweb `/api/query`), so Save & Next # 500'd. Locking the session removes the button and the failure. - env={"PGWEB_DATABASE_URL": pgweb_url, "PGWEB_LOCK_SESSION": "true"}, + env={"PGWEB_DATABASE_URL": pgweb_url, "PGWEB_LOCK_SESSION": "true", + **({"PGWEB_URL_PREFIX": self.console_path} if self.console_path else {})}, )) elif cfg.db_web_image: logger.warning(f"pgweb image {cfg.db_web_image} is not in the image store; skipping pgweb provisioning") diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index e1ae345bc..1727310d9 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -538,12 +538,14 @@ def _is_safe_member(name: str) -> bool: def _env_on_host(deployed: DeployedSandboxEnv, sandbox: Any) -> dict: - """The env's id, address and MCP endpoint as a script on its host reaches them.""" - url, mcp_url = deployed.environment_url, deployed.mcp_url + """The env's id, address and MCP endpoint as a script on its host reaches them, and its DB MCP endpoint when it + has one, which serves under a path of the deploy's own that only the recorded URL carries.""" + url, mcp_url, db_mcp_url = deployed.environment_url, deployed.mcp_url, getattr(deployed, "db_mcp_url", None) return { "id": deployed.env_id, "url": url and sandbox.url_on_host(url), "mcp_url": mcp_url and sandbox.url_on_host(mcp_url), + **({"db_mcp_url": sandbox.url_on_host(db_mcp_url)} if db_mcp_url else {}), } diff --git a/tst/integration/env/gateway/gateway_test.py b/tst/integration/env/gateway/gateway_test.py index ab4057715..de7263d46 100644 --- a/tst/integration/env/gateway/gateway_test.py +++ b/tst/integration/env/gateway/gateway_test.py @@ -14,6 +14,7 @@ import httpx import pytest from typing import NamedTuple +from urllib.parse import urlsplit from mcp import ClientSession from mcp.client.streamable_http import streamable_http_client @@ -1375,6 +1376,12 @@ async def assert_website_tools(session, tools_result, tool_names): logger.info(f"browser_snapshot result: {content_text[:500]}") assert "Slack Workspace" in content_text, f"Expected 'Slack Workspace' in snapshot, got {content_text}" + # The SQL consoles share the browser's network but answer only under the deploy's own path. + for url in ("http://pgweb:8081/api/query?query=SELECT%201", "http://db-mcp:8000/mcp"): + result = await session.call_tool("browser_navigate", {"url": url}) + content_text = "".join(c.text for c in result.content if hasattr(c, "text")) + assert "not found" in content_text.lower(), f"Expected {url} refused to the browser, got {content_text}" + # 4. browser_close then snapshot returns empty/blank page await session.call_tool("browser_close", {}) result = await session.call_tool("browser_snapshot", {}) @@ -1429,6 +1436,20 @@ async def assert_website_tools_via_step(client, gateway_url, tools, tool_names): assert data["ok"] is True assert data["channels"] == [] + # The operator reaches pgweb and db-mcp at the URLs the deploy recorded, and nowhere else on their hosts. + async with httpx.AsyncClient() as client: + queried = await client.post(f"{result.db_web_url}api/query", data={"query": "SELECT 1 AS one"}, timeout=30) + assert queried.status_code == 200 and "one" in queried.text, f"pgweb at {result.db_web_url}: {queried.text}" + web_origin = urlsplit(result.db_web_url)._replace(path="").geturl() + for path in ("/", "/api/query?query=SELECT%201"): + assert (await client.get(f"{web_origin}{path}", timeout=30)).status_code == 404 + mcp_origin = urlsplit(result.db_mcp_url)._replace(path="").geturl() + assert (await client.post(f"{mcp_origin}/mcp", json={}, timeout=30)).status_code == 404 + async with streamable_http_client(result.db_mcp_url) as (read_stream, write_stream, _): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + assert "execute_sql" in [t.name for t in (await session.list_tools()).tools] + await slack_website_env.load_environment_artifact(slack_service_artifact) # Verify channels populated after loading diff --git a/tst/unit/env/envs/service_db/service_db_test.py b/tst/unit/env/envs/service_db/service_db_test.py index df3a3b945..cb0e2ac2a 100644 --- a/tst/unit/env/envs/service_db/service_db_test.py +++ b/tst/unit/env/envs/service_db/service_db_test.py @@ -2,6 +2,8 @@ from unittest.mock import MagicMock +import yaml + from agent_env.env.envs.service_db import ( DB_NAME, DB_PASSWORD, @@ -298,3 +300,17 @@ def test_rendered_sidecar_service_names_matches_creation_gating(): rendered = "\n".join(both.render_sidecar_containers(["svc"])) for name in both.rendered_sidecar_service_names(): assert f" {name}:" in rendered + + +def test_the_sql_consoles_answer_only_under_the_path_a_deploy_gives_them(): + """pgweb and db-mcp are published, and on a VM tunnelled, so a deploy serves them under a path of its own that only + the URLs it records carry; with none, as an exported bundle renders them, they answer at their roots.""" + provider = LocalPostgresStateProvider(service_db_config=ServiceDBConfig(db_web_image="pgweb:1", db_mcp_image="dbmcp:1")) + at_root = "\n".join(provider.render_sidecar_containers(["svc"])) + provider.console_path = "Kp3vX8nQ2wR7tY5mZ1aB9cDe" + services = yaml.safe_load("services:\n" + "\n".join(provider.render_sidecar_containers(["svc"])))["services"] + + assert "PGWEB_URL_PREFIX=Kp3vX8nQ2wR7tY5mZ1aB9cDe" in services[PGWEB_SERVICE_NAME]["environment"] + assert "DB_MCP_PATH=/Kp3vX8nQ2wR7tY5mZ1aB9cDe/mcp" in services[DB_MCP_SERVICE_NAME]["environment"] + assert services[DB_MCP_SERVICE_NAME]["entrypoint"][:2] == ["python", "-c"] + assert not any(word in at_root for word in ("PGWEB_URL_PREFIX", "DB_MCP_PATH", "entrypoint")) diff --git a/tst/unit/providers/compose_host_port_publish_test.py b/tst/unit/providers/compose_host_port_publish_test.py index c463bac02..1ce98507f 100644 --- a/tst/unit/providers/compose_host_port_publish_test.py +++ b/tst/unit/providers/compose_host_port_publish_test.py @@ -149,7 +149,7 @@ def stub_postgres(monkeypatch): monkeypatch.setattr( LocalPostgresStateProvider, "_pg_url", lambda self, *a, **k: "postgres://stub", ) - return LocalPostgresStateProvider.__new__(LocalPostgresStateProvider) + return LocalPostgresStateProvider() def _sidecar(name: str, host_port: int, container_port: int = 8000) -> SidecarConfig: diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 9941bdb5b..b72756f11 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -765,6 +765,30 @@ async def test_the_deploys_cpu_reaches_every_container(): assert aux and all(c["cpu"] == 2.0 and "memory" not in c for c in aux) +@pytest.mark.asyncio +async def test_a_container_deploy_records_pgweb_under_its_path_and_db_mcp_at_its_root(): + """In containers pgweb takes the deploy's path from the env it starts with; db-mcp, whose path only its compose + entrypoint can set, keeps /mcp there.""" + provider = ModalSandboxProvider() + gp = EnvironmentGatewayProvider() + calls: list[dict] = [] + + async def fake_create_container(**kwargs): + calls.append(kwargs) + if kwargs.get("is_tcp"): + return _FakeSandbox(f"db-{kwargs['port']}", kwargs["port"], f"tcp://db.modal.host:{kwargs['port']}") + return _FakeSandbox(f"sb-{kwargs['image_name']}", kwargs["port"], f"https://{kwargs['image_name']}.modal.host") + + with _slack_in_containers(gp, provider, fake_create_container) as args: + gp._state_provider.console_path = "Kp3vX8nQ" + deployed = await gp._deploy_via_containers(sandbox_provider=provider, **args) + + pgweb, db_mcp = (next(c["image_name"] for c in calls if c["image_name"].endswith(name)) for name in ("pgweb", "db-mcp")) + assert next(c for c in calls if c["image_name"] == pgweb)["env"]["PGWEB_URL_PREFIX"] == "Kp3vX8nQ" + assert (deployed.db_web_url, deployed.db_mcp_url) == (f"https://{pgweb}.modal.host/Kp3vX8nQ/", + f"https://{db_mcp}.modal.host/mcp") + + @pytest.mark.asyncio async def test_container_mode_injects_the_registered_environment_name(): """As in compose mode, the MCP server gets its registered name, which the SDK serves as its card name.""" @@ -1187,6 +1211,9 @@ def test_redact_compose_secrets_masks_only_secret_env_lines(): assert ' - "18768:8000"' in redacted # the un-redacted compose still carries the real secret (only the log copy is masked) assert "s3cr3t-token-value" in compose + # and so do the SQL consoles' path, which their URLs carry + consoles = " environment:\n - PGWEB_URL_PREFIX=Kp3vX8nQ\n - DB_MCP_PATH=/Kp3vX8nQ/mcp" + assert "Kp3vX8nQ" not in _redact_compose_secrets(consoles) @pytest.mark.asyncio @@ -1758,6 +1785,37 @@ class _Stop(Exception): assert not any("<<" in script for script in scripts) +@pytest.mark.asyncio +async def test_a_callers_vm_without_tunnels_for_the_consoles_deploys_without_their_urls(): + from agent_env.env.gateway import GatewayMode + + gp = EnvironmentGatewayProvider() + gp._state_provider = MagicMock(spec=LocalPostgresStateProvider, console_path="Kp3vX8nQ2wR7tY5mZ1aB9cDe", + prepare=AsyncMock(), store_images_to_load=MagicMock(return_value=[]), + rendered_sidecar_service_names=MagicMock(return_value=["pgweb", "db-mcp"])) + gp._state_instance = MagicMock(instance_id="esi-1") + vm = AsyncMock(ON_THIS_MACHINE=False, tunnel_urls={18765: "https://gw.example"}, host_ips=(), extra_hosts={}, + host_port=lambda port: port, exec_with_output=AsyncMock(return_value=(0, GATEWAY_SERVICE_NAME, ""))) + + with patch("agent_env.env.env.Env.get", return_value=MagicMock(docker_image_artifact=_loadable("gw"))), \ + patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id")), \ + patch.object(gp, "_build_local_store", AsyncMock(return_value=MagicMock(init_sql=""))), \ + patch.object(gp, "create_docker_compose", return_value="services: {}\n"), \ + patch.object(gp, "_compose_up", AsyncMock()), \ + patch.object(gp, "_wait_for_gateway", AsyncMock(return_value=True)), \ + patch.object(gp, "_log_service", AsyncMock(return_value=True)), \ + patch.object(gp, "_wait_for_tunnel", AsyncMock(return_value=_CARD)), \ + patch.object(gp, "_refuse_tool_name_clashes", AsyncMock()), \ + patch.object(env_gateway_provider.asyncio, "sleep", AsyncMock()): + deployed = await gp._deploy_via_vm( + sandbox_provider=_vm_provider(), mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], + mcp_server_images=[_loadable("mcp-slack")], gateway_port=18765, website_configs=None, website_images=None, + gateway_mode=GatewayMode.PERFORMANCE, ttl_seconds=60, disk_size_gb=10, existing_sandbox=vm, + ) + + assert (deployed.gateway_url, deployed.db_web_url, deployed.db_mcp_url) == ("https://gw.example", None, None) + + _PROBE_CARD = {"name": "env1234", "children_environments": [ {"name": "a", "capabilities": {"tools": [{"name": "a_list"}, {"name": "a_get"}]}}, {"name": "b", "capabilities": {"tools": [{"name": "b_send"}]}}, diff --git a/tst/unit/providers/env_state/test_env_state_provider.py b/tst/unit/providers/env_state/test_env_state_provider.py index 995075473..5a7dedbaf 100644 --- a/tst/unit/providers/env_state/test_env_state_provider.py +++ b/tst/unit/providers/env_state/test_env_state_provider.py @@ -269,15 +269,21 @@ def test_sidecar_specs_skip_images_the_image_store_does_not_own(image_store, cap # Both in the image store → both sidecars, pointed at the instance's host. The config is held # on the provider instance, not threaded through sidecar_specs. - both = LocalPostgresStateProvider( + provider = LocalPostgresStateProvider( service_db_config=ServiceDBConfig(db_web_image=f"{_REPOSITORIES}/pgweb:v1", db_mcp_image=f"{_REPOSITORIES}/db-mcp:v1") - ).sidecar_specs(instance=instance) + ) + both = provider.sidecar_specs(instance=instance) assert {s.name for s in both} == {"pgweb", "db-mcp"} assert all("[fdaa::1]" in next(iter(s.env.values())) for s in both) # pgweb runs locked to its single connection (Connect/Disconnect hidden) so a # contributor can't tear down the shared session the changelog export relies on. pgweb_spec = next(s for s in both if s.name == "pgweb") assert pgweb_spec.env["PGWEB_LOCK_SESSION"] == "true" + assert "PGWEB_URL_PREFIX" not in pgweb_spec.env + # A deploy's path, which the URL it records carries, is the only place pgweb answers. + provider.console_path = "Kp3vX8nQ2wR7tY5mZ1aB9cDe" + pgweb_spec = next(s for s in provider.sidecar_specs(instance=instance) if s.name == "pgweb") + assert pgweb_spec.env["PGWEB_URL_PREFIX"] == "Kp3vX8nQ2wR7tY5mZ1aB9cDe" # pgweb from Docker Hub + unset db-mcp → neither, and a warning for the configured one only. none = LocalPostgresStateProvider( diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index de5865b58..9c26000e4 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -315,6 +315,15 @@ def test_an_envs_address_and_mcp_endpoint_come_from_its_card_whatever_deployed_i assert _run_code_module._env_on_host(deployed, sandbox) == expected +def test_a_gateway_envs_db_mcp_endpoint_keeps_the_path_its_deploy_serves_it_under(): + deployed = DeployedGatewayEnv(env_id="gw", env_version=1, sandbox_id="sb", gateway_url="https://abc.modal.host", + db_mcp_url="https://db.modal.host/Kp3vX8nQ/mcp", environment_card_url=_CARD_URL, + environment_card={"name": "gw"}) + sandbox = _ContainerEnvSandbox({18765: "https://abc.modal.host", 18767: "https://db.modal.host"}) + + assert _run_code_module._env_on_host(deployed, sandbox)["db_mcp_url"] == "http://127.0.0.1:18767/Kp3vX8nQ/mcp" + + def _written_input(sandbox): """The input.json the step staged, through whichever writer its target uses.""" for call in sandbox.write_host_file.await_args_list: