Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 23 additions & 9 deletions src/agent_env/providers/env_providers/env_gateway_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import json
import logging
import re
import secrets
import shlex
import time
from collections import Counter
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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."""
Expand All @@ -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:
Expand Down Expand Up @@ -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)):
Expand Down Expand Up @@ -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]:
Expand Down
18 changes: 17 additions & 1 deletion src/agent_env/providers/env_state/local_postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import Optional, TYPE_CHECKING
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:",
Expand All @@ -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)),
Expand Down Expand Up @@ -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")
Expand Down
6 changes: 4 additions & 2 deletions src/agent_env/task_step/task_steps/run_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 {}),
}


Expand Down
21 changes: 21 additions & 0 deletions tst/integration/env/gateway/gateway_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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", {})
Expand Down Expand Up @@ -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
Expand Down
16 changes: 16 additions & 0 deletions tst/unit/env/envs/service_db/service_db_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

from unittest.mock import MagicMock

import yaml

from agent_env.env.envs.service_db import (
DB_NAME,
DB_PASSWORD,
Expand Down Expand Up @@ -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"))
2 changes: 1 addition & 1 deletion tst/unit/providers/compose_host_port_publish_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
58 changes: 58 additions & 0 deletions tst/unit/providers/env_providers/env_gateway_provider_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"}]}},
Expand Down
10 changes: 8 additions & 2 deletions tst/unit/providers/env_state/test_env_state_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Loading
Loading