Vscode + Chrome Browser inside Modal Sandbox

gistfile1.txt

import asyncio
import secrets
from enum import Enum

import modal
from pydantic import BaseModel

from sweagent.utils.log import get_logger

logger = get_logger("browser", emoji="🌐")

### Constants ###
BROWSER_STARTUP_SLEEP = 3  # seconds
VSCODE_STARTUP_SLEEP = 3  # seconds
VSCODE_NGINX_PORT = 8080
VSCODE_PORT = 8081
BROWSER_NGINX_PORT = 8082
BROWSER_VNC_PORT = 6080
BROWSER_CDP_PORT = 9222

### Classes ###


class ServiceStatus(str, Enum):
    starting = "starting"
    ready = "ready"
    stopped = "stopped"
    error = "error"


class ServiceStatusResponse(BaseModel):
    status: ServiceStatus
    url: str | None = None
    token: str | None = None


### Functions ###


async def start_browser(sandbox: modal.Sandbox) -> ServiceStatusResponse:
    """Start Chromium browser that will be shared between agent and user."""
    status = await check_browser_status(sandbox)
    if status.status == ServiceStatus.stopped:
        logger.info(f"Starting browser server for sandbox {sandbox.object_id}")
        token = secrets.token_urlsafe(32)
        await sandbox.exec.aio("bash", "-lc", get_start_browser_cmd(token))
        await asyncio.sleep(BROWSER_STARTUP_SLEEP)

        vnc_url = await get_tunnel_url(sandbox, BROWSER_NGINX_PORT)
        logger.info(f"Browser VNC server available at {vnc_url}")
        return ServiceStatusResponse(status=ServiceStatus.ready, url=vnc_url, token=token)

    logger.info(
        f"Browser server already running for sandbox {sandbox.object_id} at {status.url} with status {status.status}"
    )
    return status


async def start_vscode(sandbox: modal.Sandbox) -> ServiceStatusResponse:
    """Start VSCode server that will be shared between agent and user."""
    status = await check_vscode_status(sandbox)
    if status.status == ServiceStatus.stopped:
        logger.info(f"Starting VSCode server for sandbox {sandbox.object_id}")
        token = secrets.token_urlsafe(32)
        await sandbox.exec.aio("bash", "-lc", get_start_vscode_cmd(token))
        await asyncio.sleep(VSCODE_STARTUP_SLEEP)
        tunnel = await get_tunnel_url(sandbox, VSCODE_NGINX_PORT)
        logger.info(f"VSCode server available at {tunnel}")
        return ServiceStatusResponse(status=ServiceStatus.ready, url=tunnel, token=token)

    logger.info(
        f"VSCode server already running for sandbox {sandbox.object_id} at {status.url} with status {status.status}"
    )
    return status


async def read_nginx_token(sandbox: modal.Sandbox, service: str) -> str:
    """Read the token from the nginx config file."""
    cmd = f"sed -n -E 's/.*add_header Set-Cookie \"{service}_token=([^;\"]+).*/\\1/p' /etc/nginx/conf.d/{service}.conf | head -1"
    result = await sandbox.exec.aio("bash", "-c", cmd)
    await result.wait.aio()
    output = await result.stdout.read.aio()
    return output.strip()


async def check_vscode_status(sandbox: modal.Sandbox) -> ServiceStatusResponse:
    try:
        cmd = f"curl -s -o /dev/null -w '%{{http_code}}' http://127.0.0.1:{VSCODE_PORT}"
        check_login = await sandbox.exec.aio("bash", "-c", cmd)
        await check_login.wait.aio()
        login_status = await check_login.stdout.read.aio()
        if login_status.strip() == "302":
            tunnel_url = await get_tunnel_url(sandbox, VSCODE_NGINX_PORT)
            return ServiceStatusResponse(status=ServiceStatus.ready, url=tunnel_url)

        return ServiceStatusResponse(status=ServiceStatus.stopped)
    except Exception as e:
        logger.error(f"Error checking VSCode status: {e}")
        return ServiceStatusResponse(status=ServiceStatus.error)


async def check_browser_status(sandbox: modal.Sandbox) -> ServiceStatusResponse:
    try:
        token = await read_nginx_token(sandbox, "browser")
        vnc_url = await get_tunnel_url(sandbox, BROWSER_NGINX_PORT)

        # Check if VNC server is responding
        check_vnc = await sandbox.exec.aio(
            "bash", "-c", "curl -s -o /dev/null -w '%{http_code}' http://127.0.0.1:6080/vnc.html"
        )
        await check_vnc.wait.aio()
        vnc_status = await check_vnc.stdout.read.aio()
        service_status = ServiceStatus.ready if vnc_status.strip() == "200" else ServiceStatus.stopped
        return ServiceStatusResponse(status=service_status, token=token, url=vnc_url)
    except Exception as e:
        logger.error(f"Error checking browser status: {e}")
        return ServiceStatusResponse(status=ServiceStatus.error)


async def get_tunnel_url(sandbox: modal.Sandbox, port: int) -> str | None:
    """Get the URL of a tunnel for a given port."""
    tunnels = await sandbox.tunnels.aio()
    if port in tunnels:
        return tunnels[port].url
    return None


def get_start_vscode_cmd(token: str) -> str:
    """Start code-server with nginx reverse proxy in the background."""
    code_server_args = [
        "--bind-addr",
        f"127.0.0.1:{VSCODE_PORT}",
        "--auth",
        "none",
        "--disable-telemetry",
        "--disable-workspace-trust",
        "/repos",
    ]
    code_server_args_str = " ".join(code_server_args)

    # Start both services inside the sandbox
    # - Write Nginx config (with token), start/reload Nginx on :8080
    # - Start code-server on localhost:8081 with no built-in auth
    return f"""
set -euo pipefail
mkdir -p /etc/nginx/conf.d
cat > /etc/nginx/conf.d/vscode.conf <<'EOF'
{get_vscode_nginx_conf()}
EOF
sed -i "s/__TOKEN__/{token}/g" /etc/nginx/conf.d/vscode.conf

# Ensure any stale instance is stopped, then start/reload
nginx -s reload >/dev/null 2>&1 || nginx

sleep 2

# Start code-server bound to localhost only; Nginx proxies external traffic
/code-server.sh {code_server_args_str} &
"""


def get_start_browser_cmd(token: str, start_url: str = "google.com") -> str:
    """Start Chromium browser with VNC and noVNC server with nginx reverse proxy."""
    chromium_args = [
        "--display=:99",
        "--remote-debugging-address=0.0.0.0",
        f"--remote-debugging-port={BROWSER_CDP_PORT}",
        "--no-sandbox",
        "--disable-web-security",
        "--disable-features=VizDisplayCompositor",
        "--disable-dev-shm-usage",
        "--no-first-run",
        "--disable-background-timer-throttling",
        "--disable-backgrounding-occluded-windows",
        "--disable-renderer-backgrounding",
        "--memory-pressure-off",
        "--max_old_space_size=256",
        "--disable-background-networking",
        "--disable-default-apps",
        "--disable-extensions",
        "--disable-sync",
        "--disable-translate",
        "--hide-scrollbars",
        "--mute-audio",
        "--no-default-browser-check",
        "--no-pings",
        "--disable-logging",
        "--disable-gpu-sandbox",
        "--window-size=1600,1200",
        start_url,
    ]

    chromium_args_str = " ".join(chromium_args)

    # Start browser, VNC, and noVNC services inside the sandbox
    return f"""
set -euo pipefail

# Setup Nginx config for browser VNC
mkdir -p /etc/nginx/conf.d
cat > /etc/nginx/conf.d/browser.conf <<'EOF'
{get_browser_nginx_conf()}
EOF
sed -i "s/__TOKEN__/{token}/g" /etc/nginx/conf.d/browser.conf

# Reload nginx with browser config
nginx -s reload >/dev/null 2>&1 || nginx

# Start Xvfb (virtual display) - taller aspect ratio
export DISPLAY=:99
Xvfb :99 -screen 0 1600x1200x24 -nolisten tcp -noreset >/dev/null 2>&1 &

# Wait for Xvfb to start
sleep 1

# Start Chromium in debug mode on the virtual display
chromium {chromium_args_str} >/dev/null 2>&1 &

# Wait for Chromium to start
sleep 1

# Start x11vnc to capture the display - optimized settings
x11vnc -display :99 -nopw -listen localhost -xkb -rfbport 5900 -shared -forever -bg -q \
    -threads -cursor arrow -cursorpos -solid

# Start websockify (noVNC) to proxy VNC over WebSocket
websockify --web=/usr/share/novnc {BROWSER_VNC_PORT} localhost:5900 >/dev/null 2>&1 &

# Wait for all services to be ready
sleep 2
"""


### nginx configs ###


def get_browser_nginx_conf() -> str:
    return r"""
    server {
        listen 8082;
        server_name _;
        absolute_redirect off;

        # /login sets cookie when token is valid then redirects to /
        location = /login {
            if ($arg_token != "__TOKEN__") { return 401; }
            add_header Set-Cookie "browser_token=__TOKEN__; Path=/; HttpOnly; SameSite=None; Secure; Partitioned" always;
            return 302 /;
        }

        # Exact root: support /?token=... by bouncing to /login, else send to noVNC
        location = / {
            if ($arg_token) { return 302 /login?$args; }
            if ($cookie_browser_token != "__TOKEN__") { return 401; }
            return 302 /vnc.html?autoconnect=true&resize=scale&quality=9&compression=9&view_only=false&shared=true&show_dot=false;
        }

        # Everything else proxies to noVNC, requires cookie
        location / {
            if ($cookie_browser_token != "__TOKEN__") { return 401; }
            proxy_pass http://127.0.0.1:6080;
            proxy_redirect http://127.0.0.1:6080 $scheme://$host;
            proxy_set_header Host $host;
            proxy_set_header X-Real-IP $remote_addr;
            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
            proxy_set_header X-Forwarded-Host $host;
            proxy_set_header X-Forwarded-Proto $scheme;
            proxy_set_header Upgrade $http_upgrade;
            proxy_set_header Connection "upgrade";
            proxy_http_version 1.1;
            proxy_read_timeout 3600s;
            proxy_send_timeout 3600s;
            proxy_buffering off;
        }
    }
    """


def get_vscode_nginx_conf() -> str:
    return r"""
server {
    listen 8080;
    server_name _;
    absolute_redirect off;

    location / {
        # Auth: allow if cookie matches OR token query matches (and set cookie)
        set $auth_ok 0;
        if ($cookie_vscode_token = "__TOKEN__") { set $auth_ok 1; }
        if ($arg_token = "__TOKEN__") {
            set $auth_ok 1;
            add_header Set-Cookie "vscode_token=__TOKEN__; Path=/; HttpOnly; SameSite=None; Secure; Partitioned" always;
        }
        if ($auth_ok = 0) { return 401; }

        proxy_pass http://127.0.0.1:8081;
        proxy_redirect http://127.0.0.1:8081 $scheme://$host;
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Host $host;
        proxy_set_header X-Forwarded-Proto $scheme;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection "upgrade";
        proxy_http_version 1.1;
        proxy_read_timeout 3600s;
        proxy_send_timeout 3600s;
        proxy_buffering off;
    }
}
"""