From 52726aa6c77b795d7fae4f5abc999169e082c213 Mon Sep 17 00:00:00 2001 From: Brandon Leung Date: Tue, 28 Jul 2026 21:53:40 -0400 Subject: [PATCH] feat(tasks): route selected sandbox products to the ai-gateway --- posthog/settings/temporal.py | 7 + products/tasks/backend/constants.py | 2 + .../tasks/backend/logic/services/agentsh.py | 143 +++++++++++++++--- .../backend/logic/services/modal_sandbox.py | 20 ++- .../services/tests/test_modal_sandbox.py | 25 +++ .../activities/get_sandbox_for_repository.py | 3 + .../activities/provision_sandbox.py | 14 +- .../tests/test_provision_sandbox.py | 73 ++++++++- .../temporal/process_task/tests/test_utils.py | 36 +++++ .../backend/temporal/process_task/utils.py | 16 ++ products/tasks/backend/tests/test_agentsh.py | 119 ++++++++++++++- products/tasks/backend/tests/test_models.py | 9 +- 12 files changed, 432 insertions(+), 35 deletions(-) diff --git a/posthog/settings/temporal.py b/posthog/settings/temporal.py index 20b11717fb28..b124b3926da4 100644 --- a/posthog/settings/temporal.py +++ b/posthog/settings/temporal.py @@ -49,6 +49,13 @@ ) # When not set: defaults to "docker" in DEBUG mode, "modal" in production SANDBOX_API_URL: str | None = get_from_env("SANDBOX_API_URL", None, optional=True) SANDBOX_LLM_GATEWAY_URL: str | None = get_from_env("SANDBOX_LLM_GATEWAY_URL", None, optional=True) +# The Go ai-gateway runs on its own host (ai-gateway.*, vs the Python gateway.*), so the +# base URL is what selects it: no product slug on the path, attribution as one +# X-PostHog-Properties blob. SANDBOX_AI_GATEWAY_PRODUCTS limits the switch to named +# ai_product values so one product migrates without moving every other sandbox caller. +# Both must be set; clearing either rolls back to the Python gateway. +SANDBOX_AI_GATEWAY_URL: str | None = get_from_env("SANDBOX_AI_GATEWAY_URL", None, optional=True) +SANDBOX_AI_GATEWAY_PRODUCTS: str | None = get_from_env("SANDBOX_AI_GATEWAY_PRODUCTS", None, optional=True) SANDBOX_MCP_URL: str | None = get_from_env("SANDBOX_MCP_URL", None, optional=True) # OTLP destinations for agent-server run telemetry (PostHog Logs/APM). diff --git a/products/tasks/backend/constants.py b/products/tasks/backend/constants.py index bcd9d0737352..ef1d280c4d8a 100644 --- a/products/tasks/backend/constants.py +++ b/products/tasks/backend/constants.py @@ -368,6 +368,8 @@ def vm_sandbox_allowed_origins(*, distinct_id: str, organization_id: str) -> set "GITHUB_TOKEN", "GH_TOKEN", "LLM_GATEWAY_URL", + "AI_GATEWAY_URL", + "AI_GATEWAY_PRODUCTS", "POSTHOG_RESUME_RUN_ID", "POSTHOG_AGENT_OTEL_LOGS_URL", "POSTHOG_AGENT_OTEL_LOGS_TOKEN", diff --git a/products/tasks/backend/logic/services/agentsh.py b/products/tasks/backend/logic/services/agentsh.py index 91f5bf3a8641..069d1ce7d4d0 100644 --- a/products/tasks/backend/logic/services/agentsh.py +++ b/products/tasks/backend/logic/services/agentsh.py @@ -1,13 +1,19 @@ import shlex +import logging +import ipaddress from pathlib import Path from urllib.parse import urlparse from django.conf import settings +from django.core.exceptions import ValidationError +from django.core.validators import DomainNameValidator import yaml from products.tasks.backend.constants import SANDBOX_AGENT_LAUNCH_UNSET_ENV_VARS +logger = logging.getLogger(__name__) + AGENTSH_DAEMON_PORT = 18080 SESSION_ID_FILE = "/tmp/agentsh-session-id" ENV_FILE = "/tmp/agent-env" @@ -36,6 +42,8 @@ def read_gh_guard_script() -> bytes: "api.anthropic.com", "gateway.us.posthog.com", "gateway.eu.posthog.com", + "ai-gateway.us.posthog.com", + "ai-gateway.eu.posthog.com", ] @@ -69,28 +77,120 @@ def _port_from_url(url: str | None) -> int | None: return None -# Sandbox-host URLs in `.env` for local dev. Their hostnames and (non-standard) -# ports feed the DEBUG-only network rule below — e.g. llm-gateway on 3308, MCP -# wrangler on 8787 — so locally-hosted services pass the agentsh syscall-layer -# firewall. These settings are only read into the DEBUG rule; in prod they have -# no effect even if defined. -_DEBUG_SANDBOX_URL_SETTINGS = ( +# Sandbox-host URLs from deployment env or `.env`. Each one is handed to the +# sandbox as a URL it must call, so its hostname joins the enforced allow rule +# in every environment (dev's ai-gateway.dev.posthog.dev is outside +# *.posthog.com, so the static infrastructure list alone can't cover it). +# Non-standard ports (llm-gateway on 3308, MCP wrangler on 8787) feed the +# DEBUG-only rule; the enforced rule stays on cloud-routing ports. +_SANDBOX_URL_SETTINGS = ( "SANDBOX_API_URL", "SANDBOX_LLM_GATEWAY_URL", + "SANDBOX_AI_GATEWAY_URL", "SANDBOX_MCP_URL", +) + +# Sandbox-host URLs that stay out of the enforced rule: telemetry export is not +# required for the agent to run, so its host is admitted in DEBUG only and prod +# reaches its collector through `INFRASTRUCTURE_DOMAINS` or not at all. +_DEBUG_ONLY_URL_SETTINGS = ( "SANDBOX_AGENT_OTEL_LOGS_URL", "SANDBOX_AGENT_OTEL_TRACES_URL", ) +_LOOPBACK_ALIASES = ("localhost", "host.docker.internal") + +# The allow rule is an enforcement gate: a malformed value must narrow it, never +# widen it, so `*` (wildcard syntax at both layers) and any name Django rejects +# are dropped. The validator accepts a rooted FQDN, which neither layer matches +# against an unrooted request host, so a trailing dot is rejected too. +_validate_domain_name = DomainNameValidator() + + +def _is_policy_hostname(hostname: str) -> bool: + if hostname.endswith("."): + return False + try: + _validate_domain_name(hostname) + except ValidationError: + return False + return True + + +def _is_ip_literal(hostname: str) -> bool: + try: + ipaddress.ip_address(hostname) + except ValueError: + return False + return True + + +def _is_loopback(hostname: str) -> bool: + if hostname in _LOOPBACK_ALIASES: + return True + try: + return ipaddress.ip_address(hostname).is_loopback + except ValueError: + return False + + +def sandbox_url_setting_domains() -> list[str]: + """Hostnames parsed from the `SANDBOX_*_URL` settings that are usable on + the enforced allow rule. Loopback hosts are skipped silently (agentsh + allows loopback by CIDR and Modal rejects the aliases). Any other + set-but-unusable value is logged: the URL still reaches the sandbox, so + silent exclusion here would reproduce the injected-but-blocked failure + this function exists to prevent. + """ + domains: list[str] = [] + for setting_name in _SANDBOX_URL_SETTINGS: + value = getattr(settings, setting_name, None) + if not value: + continue + hostname = _hostname_from_url(value) + if hostname and _is_loopback(hostname): + continue + if not hostname or _is_ip_literal(hostname) or not _is_policy_hostname(hostname): + logger.warning( + "Sandbox URL setting %s yields no usable policy hostname; the URL will be handed to the " + "sandbox but its host will not be admitted by the network policy", + setting_name, + ) + continue + if not getattr(settings, "DEBUG", False): + port = _port_from_url(value) + if port not in (None, 443, 80, 22): + logger.warning( + "Sandbox URL setting %s uses port %d, which the enforced network policy does not " + "admit outside DEBUG; connections from the sandbox will be denied", + setting_name, + port, + ) + if hostname not in domains: + domains.append(hostname) + return domains + + +def enforced_egress_domains() -> list[str]: + """The full non-DEBUG egress source set: baked-in infrastructure plus + settings-derived hosts. Both enforcement layers (the agentsh allow rule + and Modal's outbound allowlist) build from this one assembly so a new + source cannot land in one layer and miss the other. + """ + domains = list(INFRASTRUCTURE_DOMAINS) + for domain in sandbox_url_setting_domains(): + if domain not in domains: + domains.append(domain) + return domains + def _get_debug_only_domains() -> list[str]: """Hostnames added ONLY when DEBUG is on: dev loopback aliases plus any - sandbox URL hosts parsed from `SANDBOX_*_URL` settings. Kept separate from - the prod-safe `INFRASTRUCTURE_DOMAINS` set so a stray dev hostname can't - accidentally widen prod's allowlist. + sandbox URL hosts parsed from `SANDBOX_*_URL` settings, here paired with + the dev ports those services listen on. """ domains: list[str] = ["localhost", "host.docker.internal"] - for setting_name in _DEBUG_SANDBOX_URL_SETTINGS: + for setting_name in _SANDBOX_URL_SETTINGS + _DEBUG_ONLY_URL_SETTINGS: hostname = _hostname_from_url(getattr(settings, setting_name, None)) if hostname and hostname not in domains: domains.append(hostname) @@ -106,7 +206,7 @@ def _get_debug_only_ports() -> list[int]: syscall layer even when their hostname is allowed. """ ports: list[int] = [8000, 8010] - for setting_name in _DEBUG_SANDBOX_URL_SETTINGS: + for setting_name in _SANDBOX_URL_SETTINGS + _DEBUG_ONLY_URL_SETTINGS: port = _port_from_url(getattr(settings, setting_name, None)) if port is not None and port not in ports: ports.append(port) @@ -318,13 +418,13 @@ def generate_config_yaml(*, enable_ptrace: bool = True, full_trace: bool = True) def generate_policy_yaml(allowed_domains: list[str] | None = None) -> str: """Generate agentsh policy YAML. - When allowed_domains is set, only those domains (plus infrastructure) are - reachable and everything else is denied. When None, all network traffic - is allowed (audit-only mode). + When allowed_domains is set, only those domains (plus infrastructure and + settings-derived sandbox hosts) are reachable and everything else is + denied. When None, all network traffic is allowed (audit-only mode). """ if allowed_domains is not None: prod_domains = list(allowed_domains) - for domain in INFRASTRUCTURE_DOMAINS: + for domain in enforced_egress_domains(): if domain not in prod_domains: prod_domains.append(domain) @@ -339,9 +439,8 @@ def generate_policy_yaml(allowed_domains: list[str] | None = None) -> str: "cidrs": ["169.254.169.254/32", "fd00:ec2::254/128"], "decision": "deny", }, - # Prod-safe allow rule: only the caller-provided domains plus our - # baked-in infrastructure domains, and only the cloud-routing ports - # (443, 80, 22). This rule is identical in every environment. + # Enforced allow rule: caller-provided domains plus the shared + # egress source set, on cloud-routing ports only. { "name": "allow-domains", "domains": prod_domains, @@ -349,9 +448,9 @@ def generate_policy_yaml(allowed_domains: list[str] | None = None) -> str: "decision": "allow", }, ] - # DEBUG-only additions live in their own rule so a stray dev hostname - # or port can't widen the prod allowlist by accident. Append after the - # prod rule, before default-deny. + # DEBUG-only additions (loopback aliases and non-standard dev ports) + # live in their own rule so they can't widen the prod allowlist by + # accident. Append after the prod rule, before default-deny. if getattr(settings, "DEBUG", False): network_rules.append( { @@ -428,6 +527,8 @@ def generate_policy_yaml(allowed_domains: list[str] | None = None) -> str: "JWT_PUBLIC_KEY", "GITHUB_TOKEN", "LLM_GATEWAY_URL", + "AI_GATEWAY_URL", + "AI_GATEWAY_PRODUCTS", "IS_SANDBOX", "PYTHONPATH", ], diff --git a/products/tasks/backend/logic/services/modal_sandbox.py b/products/tasks/backend/logic/services/modal_sandbox.py index 69ca7cfdaf7b..3ae0c2634ed6 100644 --- a/products/tasks/backend/logic/services/modal_sandbox.py +++ b/products/tasks/backend/logic/services/modal_sandbox.py @@ -136,6 +136,21 @@ "mcp.posthog.com", ) + +def _session_init_probe_hosts() -> list[str]: + """Hosts the startup-failure egress probe checks. Both gateway settings + are included: routed products call SANDBOX_AI_GATEWAY_URL, everything + else SANDBOX_LLM_GATEWAY_URL, and a block on either is this probe's + reason to exist. + """ + hosts = list(SESSION_INIT_PROBE_HOSTS) + for setting_name in ("SANDBOX_LLM_GATEWAY_URL", "SANDBOX_AI_GATEWAY_URL"): + gateway_host = _hostname_from_url(getattr(settings, setting_name, None)) + if gateway_host and gateway_host not in hosts: + hosts.insert(0, gateway_host) + return hosts + + # Modal region mapping based on cloud deployment MODAL_REGION_BY_DEPLOYMENT: dict[str | None, str] = { "EU": "eu-west", @@ -1055,10 +1070,7 @@ def _diagnose_startup_failure(self, allowed_domains: list[str] | None) -> dict[s return diagnostics def _probe_session_init_egress(self) -> str: - hosts = list(SESSION_INIT_PROBE_HOSTS) - gateway_host = _hostname_from_url(getattr(settings, "SANDBOX_LLM_GATEWAY_URL", None)) - if gateway_host and gateway_host not in hosts: - hosts.insert(0, gateway_host) + hosts = _session_init_probe_hosts() checks = "; ".join( f"printf '%s ' {shlex.quote(host)}; " f"curl -sS --max-time 3 -o /dev/null -w 'http_code=%{{http_code}}\\n' https://{host}/ 2>/dev/null || echo FAILED" diff --git a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py index 6dd6a863a370..05622c358ef5 100644 --- a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py +++ b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py @@ -8,6 +8,8 @@ import pytest from unittest.mock import MagicMock, patch +from django.test import override_settings + from modal.exception import ( ConnectionError as ModalConnectionError, ServiceError as ModalServiceError, @@ -39,6 +41,7 @@ _image_ref_cache, _merge_runtime_dependency_specs, _resource_create_kwargs, + _session_init_probe_hosts, ) from products.tasks.backend.logic.services.sandbox import ( AgentServerResult, @@ -1204,3 +1207,25 @@ def test_create_directory_snapshot_overrides_modal_default_timeout(self, mock_sa mock_sandbox._sandbox.snapshot_directory.assert_called_once_with( "/tmp/workspace", timeout=DIRECTORY_SNAPSHOT_TIMEOUT_SECONDS, ttl=None ) + + +class TestSessionInitProbeHosts: + @override_settings( + SANDBOX_LLM_GATEWAY_URL="https://gateway.dev.posthog.dev", + SANDBOX_AI_GATEWAY_URL="https://ai-gateway.dev.posthog.dev", + ) + def test_includes_both_configured_gateway_hosts(self): + # Routed products call the ai-gateway during session init; if the probe + # omits its host, a blocked ai-gateway diagnoses as "no egress block + # detected" (the exact failure class this probe exists to name). + hosts = _session_init_probe_hosts() + assert "gateway.dev.posthog.dev" in hosts + assert "ai-gateway.dev.posthog.dev" in hosts + + @override_settings( + SANDBOX_LLM_GATEWAY_URL="https://gateway.us.posthog.com", + SANDBOX_AI_GATEWAY_URL=None, + ) + def test_deduplicates_against_static_hosts_and_skips_unset(self): + hosts = _session_init_probe_hosts() + assert hosts.count("gateway.us.posthog.com") == 1 diff --git a/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py b/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py index 3ca704025039..b0bece33d385 100644 --- a/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py +++ b/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py @@ -34,6 +34,7 @@ from products.tasks.backend.temporal.oauth import create_oauth_access_token_for_run from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution from products.tasks.backend.temporal.process_task.utils import ( + ai_gateway_env_vars, get_git_identity_env_vars, get_sandbox_api_url, get_sandbox_github_token, @@ -240,6 +241,8 @@ def get_sandbox_for_repository(input: GetSandboxForRepositoryInput) -> GetSandbo if settings.SANDBOX_LLM_GATEWAY_URL: environment_variables["LLM_GATEWAY_URL"] = settings.SANDBOX_LLM_GATEWAY_URL + environment_variables.update(ai_gateway_env_vars()) + environment_variables.update(get_git_identity_env_vars(task, ctx.state)) run_state = parse_run_state(ctx.state) diff --git a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py index 2c4ec913182f..cbd33513fda8 100644 --- a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py @@ -18,7 +18,7 @@ OAuthTokenError, TaskNotFoundError, ) -from products.tasks.backend.logic.services.agentsh import INFRASTRUCTURE_DOMAINS, _get_debug_only_domains +from products.tasks.backend.logic.services.agentsh import _get_debug_only_domains, enforced_egress_domains from products.tasks.backend.logic.services.connection_token import ( SANDBOX_JWT_STATE_KID_KEY, get_primary_sandbox_jwt_kid, @@ -40,6 +40,7 @@ set_git_remote_token, ) from products.tasks.backend.temporal.process_task.utils import ( + ai_gateway_env_vars, get_git_identity_env_vars, get_readonly_github_token, get_sandbox_api_url, @@ -158,12 +159,13 @@ def _to_modal_domain_allowlist(allowed_domains: list[str]) -> list[str]: """Translate the agentsh allowlist into Modal's outbound_domain_allowlist. Modal fences the whole sandbox and supports `*.` wildcards that match the - apex and any subdomain, so union in the infra (and local tunnel) domains the - agent needs, drop loopback aliases Modal rejects as invalid domains, and - collapse entries already covered by a wildcard. + apex and any subdomain, so union in the shared egress source set (infra + plus settings-derived sandbox hosts) the agent needs, drop loopback + aliases Modal rejects as invalid domains, and collapse entries already + covered by a wildcard. """ domains = list(allowed_domains) - extra = list(INFRASTRUCTURE_DOMAINS) + extra = enforced_egress_domains() if settings.DEBUG: extra += _get_debug_only_domains() for domain in extra: @@ -339,6 +341,8 @@ def _build_environment_variables( if settings.SANDBOX_LLM_GATEWAY_URL: environment_variables["LLM_GATEWAY_URL"] = settings.SANDBOX_LLM_GATEWAY_URL + environment_variables.update(ai_gateway_env_vars()) + if settings.DEBUG: # Local eval runs pin models per unit; the agent's overload rescue would silently switch a # session to the fallback model mid-run, breaking prompt-cache sharing (model is part of diff --git a/products/tasks/backend/temporal/process_task/tests/test_provision_sandbox.py b/products/tasks/backend/temporal/process_task/tests/test_provision_sandbox.py index 0806f9d08e1a..7c0ffaf2c9d3 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/tests/test_provision_sandbox.py @@ -104,7 +104,16 @@ def test_build_sandbox_tags_drops_none_values(): assert all(isinstance(value, str) for value in tags.values()) -@override_settings(DEBUG=False) +# All four SANDBOX_*_URL settings are pinned: they feed the enforced allowlist +# outside DEBUG, so a developer's environment value (an ngrok SANDBOX_API_URL) +# would otherwise leak into these exact-equality expectations. +@override_settings( + DEBUG=False, + SANDBOX_API_URL=None, + SANDBOX_LLM_GATEWAY_URL=None, + SANDBOX_AI_GATEWAY_URL=None, + SANDBOX_MCP_URL=None, +) @pytest.mark.parametrize( "allowed_domains, expected", [ @@ -134,6 +143,68 @@ def test_to_modal_domain_allowlist_resolves_exact_list(allowed_domains, expected assert _to_modal_domain_allowlist(allowed_domains) == expected +@override_settings(DEBUG=False, SANDBOX_AI_GATEWAY_URL="https://ai-gateway.dev.posthog.dev") +def test_to_modal_domain_allowlist_admits_configured_gateway_host(): + # Modal fences egress independently of agentsh, so the settings-derived + # gateway host must clear this layer too; dev's host is outside + # *.posthog.com and nothing else admits it. + assert "ai-gateway.dev.posthog.dev" in _to_modal_domain_allowlist([]) + + +@override_settings(DEBUG=False, SANDBOX_LLM_GATEWAY_URL="http://127.0.0.1:3308") +def test_to_modal_domain_allowlist_drops_loopback_ip_settings_host(): + # 127.0.0.1 contains dots, so the fqdn filter alone would pass it into + # Modal's outbound_domain_allowlist, which rejects non-domain entries; + # the loopback exclusion in sandbox_url_setting_domains is the only + # defense on this path. + assert "127.0.0.1" not in _to_modal_domain_allowlist([]) + + +# A charset-only check admits empty labels and per-label hyphens, putting a host +# in Modal's allowlist its matcher never fires on: the sandbox boots, then its +# gateway calls are silently denied. +@pytest.mark.parametrize( + "hostname", + [ + "ai-gateway.dev..posthog.dev", + "ai-gateway.-dev.posthog.dev", + "ai-gateway.dev-.posthog.dev", + "-ai-gateway.dev.posthog.dev", + # Valid DNS, but no layer matches a rooted name against an unrooted host. + "ai-gateway.dev.posthog.dev.", + # Wildcard syntax at both layers; never widen to a whole suffix. + "*.posthog.dev", + ], +) +def test_to_modal_domain_allowlist_rejects_malformed_settings_host(hostname): + with override_settings(DEBUG=False, SANDBOX_AI_GATEWAY_URL=f"https://{hostname}"): + assert hostname not in _to_modal_domain_allowlist([]) + + +@override_settings(DEBUG=False, SANDBOX_AI_GATEWAY_URL="https://ai-gateway.dev.posthog.dev") +def test_to_modal_domain_allowlist_still_admits_hyphenated_host(): + # Guards against over-rejecting: hyphens inside a label are legal. + assert "ai-gateway.dev.posthog.dev" in _to_modal_domain_allowlist([]) + + +@patch(f"{_PROVISION}.get_git_identity_env_vars", return_value={}) +@patch(f"{_PROVISION}.get_sandbox_jwt_public_key", return_value="pub") +@patch(f"{_PROVISION}.get_sandbox_api_url", return_value="https://api.example") +@override_settings( + SANDBOX_AI_GATEWAY_URL="https://ai-gateway.us.posthog.com", + SANDBOX_AI_GATEWAY_PRODUCTS="signals_scout", +) +def test_build_environment_variables_injects_ai_gateway_pair(_api, _jwt, _git): + # Pins this site's wiring of the shared helper: the conjunction itself is + # tested in test_utils.py, but deleting the update() call here would merge + # green without this assertion and Modal-provisioned sandboxes would + # silently stay on the legacy gateway. + env = _build_environment_variables(_context(), MagicMock(), "", "access-token") + + assert env["AI_GATEWAY_URL"] == "https://ai-gateway.us.posthog.com" + assert env["AI_GATEWAY_PRODUCTS"] == "signals_scout" + + @patch(f"{_PROVISION}.emit_agent_log") @patch(f"{_PROVISION}.Sandbox.get_by_id") @pytest.mark.parametrize( diff --git a/products/tasks/backend/temporal/process_task/tests/test_utils.py b/products/tasks/backend/temporal/process_task/tests/test_utils.py index bfe800dbac81..b83d4be3fd36 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_utils.py +++ b/products/tasks/backend/temporal/process_task/tests/test_utils.py @@ -1069,6 +1069,42 @@ def test_collision_detection_is_case_insensitive(self): assert get_relayed_mcp_server_names(task_run, {"grafana"}) == ["Playwright", "internal-cli"] +@patch( + "products.tasks.backend.logic.services.connection_token.get_sandbox_jwt_public_key", + return_value="test-jwt-key", +) +@patch( + "products.tasks.backend.temporal.process_task.utils.get_sandbox_api_url", + return_value="https://us.posthog.com", +) +class TestBuildSandboxEnvironmentVariablesGateway(TestCase): + def _build(self): + return build_sandbox_environment_variables(github_token=None, access_token="tok", team_id=1) + + @override_settings( + SANDBOX_AI_GATEWAY_URL="https://ai-gateway.us.posthog.com", + SANDBOX_AI_GATEWAY_PRODUCTS="signals_scout,signals_research", + ) + def test_both_settings_inject_both_vars(self, _api, _jwt): + env = self._build() + self.assertEqual(env["AI_GATEWAY_URL"], "https://ai-gateway.us.posthog.com") + self.assertEqual(env["AI_GATEWAY_PRODUCTS"], "signals_scout,signals_research") + + @override_settings(SANDBOX_AI_GATEWAY_URL="https://ai-gateway.us.posthog.com", SANDBOX_AI_GATEWAY_PRODUCTS=None) + def test_url_without_products_injects_neither(self, _api, _jwt): + # A URL with no product allowlist would route every sandbox caller, so a + # half-config must inject nothing. + env = self._build() + self.assertNotIn("AI_GATEWAY_URL", env) + self.assertNotIn("AI_GATEWAY_PRODUCTS", env) + + @override_settings(SANDBOX_AI_GATEWAY_URL=None, SANDBOX_AI_GATEWAY_PRODUCTS="signals_scout") + def test_products_without_url_injects_neither(self, _api, _jwt): + env = self._build() + self.assertNotIn("AI_GATEWAY_URL", env) + self.assertNotIn("AI_GATEWAY_PRODUCTS", env) + + class TestBuildSandboxEnvironmentVariables(SimpleTestCase): @patch( "products.tasks.backend.logic.services.connection_token.get_sandbox_jwt_public_key", diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index bc880c99cd0b..3f909fd482f1 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -1194,6 +1194,8 @@ def build_sandbox_environment_variables( if settings.SANDBOX_LLM_GATEWAY_URL: env_vars["LLM_GATEWAY_URL"] = settings.SANDBOX_LLM_GATEWAY_URL + env_vars.update(ai_gateway_env_vars()) + if otel_telemetry_enabled: env_vars.update(get_sandbox_otel_env_vars()) @@ -1218,6 +1220,20 @@ def get_sandbox_otel_env_vars() -> dict[str, str]: return env_vars +def ai_gateway_env_vars() -> dict[str, str]: + """Env vars routing listed products to the Go ai-gateway, shared by every + injection site so the both-or-nothing guard cannot drift per site. Both + settings or nothing: a URL with no product allowlist would route every + sandbox caller, and a product list with no URL has nowhere to go. + """ + if settings.SANDBOX_AI_GATEWAY_URL and settings.SANDBOX_AI_GATEWAY_PRODUCTS: + return { + "AI_GATEWAY_URL": settings.SANDBOX_AI_GATEWAY_URL, + "AI_GATEWAY_PRODUCTS": settings.SANDBOX_AI_GATEWAY_PRODUCTS, + } + return {} + + def get_pr_authorship_mode(task: Task, state: dict[str, Any] | None = None) -> PrAuthorshipMode: """Return the effective PR authorship mode for a run. diff --git a/products/tasks/backend/tests/test_agentsh.py b/products/tasks/backend/tests/test_agentsh.py index a5bf7e0bbb1e..298dc0192913 100644 --- a/products/tasks/backend/tests/test_agentsh.py +++ b/products/tasks/backend/tests/test_agentsh.py @@ -1,5 +1,6 @@ import os import shlex +import logging import tempfile import subprocess from pathlib import Path @@ -59,6 +60,13 @@ def test_outputs_valid_yaml(self): class TestGeneratePolicyYaml(TestCase): + def setUp(self): + super().setUp() + # LOGGING uses disable_existing_loggers, so any test that re-applies + # logging config disables this module's import-time logger and the + # assertLogs below see nothing; re-enable it. + logging.getLogger("products.tasks.backend.logic.services.agentsh").disabled = False + def test_allows_commands(self): policy = yaml.safe_load(generate_policy_yaml(["example.com"])) allow_rule = next(rule for rule in policy["command_rules"] if rule["name"] == "allow-all-commands") @@ -132,6 +140,22 @@ def test_env_policy_allows_posthog_vars(self): policy = yaml.safe_load(generate_policy_yaml([])) self.assertIn("POSTHOG_*", policy["env_policy"]["allow"]) + def test_env_policy_allows_gateway_selection_vars(self): + # AI_GATEWAY_URL picks the Go gateway and AI_GATEWAY_PRODUCTS scopes it. If + # the firewall strips either, the sandbox silently stays on the Python + # gateway, so a migrated product keeps billing under its old tag. + policy = yaml.safe_load(generate_policy_yaml([])) + for key in ("LLM_GATEWAY_URL", "AI_GATEWAY_URL", "AI_GATEWAY_PRODUCTS"): + self.assertIn(key, policy["env_policy"]["allow"]) + + def test_go_gateway_hosts_reachable(self): + # The Go gateway is a different hostname from the Python one; without an + # allow-domains entry every sandbox model call is denied at the syscall layer. + policy = yaml.safe_load(generate_policy_yaml([])) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertIn("ai-gateway.us.posthog.com", allow_rule["domains"]) + self.assertIn("ai-gateway.eu.posthog.com", allow_rule["domains"]) + @override_settings(DEBUG=True) def test_debug_mode_adds_dev_ports(self): policy = yaml.safe_load(generate_policy_yaml([])) @@ -176,14 +200,103 @@ def test_debug_mode_adds_sandbox_hosts_to_debug_rule(self): @override_settings(DEBUG=False, SANDBOX_LLM_GATEWAY_URL="http://example.local:3308") def test_non_debug_mode_omits_debug_rule_entirely(self): - # Outside DEBUG the debug rule should not exist, and the prod rule - # must not absorb any sandbox URL hostnames or non-cloud ports. + # Outside DEBUG the debug rule should not exist and the prod rule keeps + # cloud-routing ports only. The hostname itself still joins the prod + # rule: the sandbox is configured to call it. policy = yaml.safe_load(generate_policy_yaml([])) rule_names = [rule["name"] for rule in policy["network_rules"]] self.assertNotIn("allow-debug-domains", rule_names) allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") self.assertEqual(sorted(allow_rule["ports"]), [22, 80, 443]) - self.assertNotIn("example.local", allow_rule["domains"]) + self.assertIn("example.local", allow_rule["domains"]) + + @override_settings(DEBUG=False, SANDBOX_AI_GATEWAY_URL="https://ai-gateway.dev.posthog.dev") + def test_configured_gateway_host_reachable_outside_debug(self): + # Dev's gateway host is outside *.posthog.com, so only the + # settings-derived entry admits it; without one, every routed model + # call in a restricted dev sandbox is denied at the syscall layer. + policy = yaml.safe_load(generate_policy_yaml([])) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertIn("ai-gateway.dev.posthog.dev", allow_rule["domains"]) + + # One case per member of _SANDBOX_URL_SETTINGS: the enforced-rule path is + # shared, so each entry needs its own tripping input or its deletion + # merges green. + @parameterized.expand( + [ + ("SANDBOX_API_URL", "api.sandbox.example.dev"), + ("SANDBOX_LLM_GATEWAY_URL", "llm-gw.sandbox.example.dev"), + ("SANDBOX_AI_GATEWAY_URL", "ai-gw.sandbox.example.dev"), + ("SANDBOX_MCP_URL", "mcp.sandbox.example.dev"), + ] + ) + def test_each_sandbox_url_setting_reaches_enforced_rule(self, setting_name, host): + with override_settings(DEBUG=False, **{setting_name: f"https://{host}"}): + policy = yaml.safe_load(generate_policy_yaml([])) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertIn(host, allow_rule["domains"]) + + @override_settings(DEBUG=False, SANDBOX_LLM_GATEWAY_URL="http://localhost:3308") + def test_loopback_sandbox_hosts_stay_off_prod_rule(self): + # Loopback is already allowed by CIDR; the alias would be noise in the + # enforced domain rule. + policy = yaml.safe_load(generate_policy_yaml([])) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertNotIn("localhost", allow_rule["domains"]) + + @parameterized.expand( + [ + ("bare_wildcard", "https://*", "*"), + ("wildcard_subdomain", "https://*.evil.example", "*.evil.example"), + ] + ) + def test_wildcard_settings_hosts_rejected_from_enforced_rule(self, _name, url, parsed_host): + # `*` is match-everything syntax at both enforcement layers, so a + # malformed setting value must narrow the policy, never widen it. + with override_settings(DEBUG=False, SANDBOX_AI_GATEWAY_URL=url): + with self.assertLogs("products.tasks.backend.logic.services.agentsh", level="WARNING"): + policy = yaml.safe_load(generate_policy_yaml([])) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertNotIn(parsed_host, allow_rule["domains"]) + + @override_settings(DEBUG=False, SANDBOX_LLM_GATEWAY_URL="http://10.0.5.3:3308") + def test_ip_literal_settings_hosts_rejected_from_enforced_rule(self): + # Both layers match DNS names; an IP literal would be inert on the + # agentsh rule and rejected by Modal, so it is excluded with a warning. + with self.assertLogs("products.tasks.backend.logic.services.agentsh", level="WARNING"): + policy = yaml.safe_load(generate_policy_yaml([])) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertNotIn("10.0.5.3", allow_rule["domains"]) + + @override_settings(DEBUG=False, SANDBOX_AI_GATEWAY_URL="ai-gateway.dev.posthog.dev") + def test_schemeless_setting_warns_and_admits_nothing(self): + # A scheme-less value passes the injection sites' truthiness gate but + # parses to no hostname, so the URL reaches the sandbox while its host + # is never admitted; the warning is the only backend-side signal. + with self.assertLogs("products.tasks.backend.logic.services.agentsh", level="WARNING") as logs: + policy = yaml.safe_load(generate_policy_yaml([])) + self.assertTrue(any("SANDBOX_AI_GATEWAY_URL" in line for line in logs.output)) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertNotIn("ai-gateway.dev.posthog.dev", allow_rule["domains"]) + + @override_settings(DEBUG=False, SANDBOX_AI_GATEWAY_URL="https://ai-gw.example.dev:8443") + def test_non_cloud_port_outside_debug_warns_but_keeps_host(self): + # The enforced rule stays on cloud ports, so a hosted URL on another + # port is admitted by hostname yet denied on connect; the warning keeps + # that from surfacing only as opaque denied connections in the sandbox. + with self.assertLogs("products.tasks.backend.logic.services.agentsh", level="WARNING") as logs: + policy = yaml.safe_load(generate_policy_yaml([])) + self.assertTrue(any("port 8443" in line for line in logs.output)) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertIn("ai-gw.example.dev", allow_rule["domains"]) + + @override_settings(DEBUG=False, SANDBOX_LLM_GATEWAY_URL="http://degraded-host.example:abc") + def test_malformed_port_outside_debug_keeps_hostname_on_enforced_rule(self): + # Documented degrade contract: hostname kept, port dropped. The host + # still reaches the enforced rule so the sandbox can connect on 443/80. + policy = yaml.safe_load(generate_policy_yaml([])) + allow_rule = next(rule for rule in policy["network_rules"] if rule["name"] == "allow-domains") + self.assertIn("degraded-host.example", allow_rule["domains"]) @override_settings( DEBUG=True, diff --git a/products/tasks/backend/tests/test_models.py b/products/tasks/backend/tests/test_models.py index 75fb54b58335..bdf75ff51ac1 100644 --- a/products/tasks/backend/tests/test_models.py +++ b/products/tasks/backend/tests/test_models.py @@ -1512,10 +1512,17 @@ def test_filter_user_sandbox_env_vars_drops_reserved_and_blocked(self): "NODE_OPTIONS": "--import=evil", "LD_PRELOAD": "/tmp/evil.so", "GITHUB_TOKEN": "stolen", + # Forging either would redirect the agent's model calls to an + # attacker host, so both must be reserved. + "AI_GATEWAY_URL": "https://evil.example.com", + "AI_GATEWAY_PRODUCTS": "signals_scout", } ) self.assertEqual(safe, {"SAFE_VAR": "ok"}) - self.assertEqual(sorted(skipped), ["GITHUB_TOKEN", "LD_PRELOAD", "NODE_OPTIONS"]) + self.assertEqual( + sorted(skipped), + ["AI_GATEWAY_PRODUCTS", "AI_GATEWAY_URL", "GITHUB_TOKEN", "LD_PRELOAD", "NODE_OPTIONS"], + ) @parameterized.expand( [