feat(tasks): route selected sandbox products to the ai-gateway

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

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",

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,19 +439,18 @@ 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,
                 "ports": [443, 80, 22],
                 "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",
             ],

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"

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

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)

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