Add UNAVAILABLE_WITHOUT_CONTAINER_NAME feature flag

libs/python/computer-server/computer_server/main.py

@@ -40,6 +40,38 @@
 # Authentication session TTL (in seconds). Override via env var CUA_AUTH_TTL_SECONDS. Default: 60s
 AUTH_SESSION_TTL_SECONDS: int = int(os.environ.get("CUA_AUTH_TTL_SECONDS", "60"))
 
+# Status code returned when UNAVAILABLE_WITHOUT_CONTAINER_NAME is set and CONTAINER_NAME is missing.
+DEFAULT_UNAVAILABLE_STATUS_CODE: int = 503
+
+
+def _parse_bool_env(name: str) -> bool:
+    return os.environ.get(name, "").lower().strip() in ("1", "true", "yes", "y", "on")
+
+
+def _unavailable_status_code() -> Optional[int]:
+    """Return the HTTP status code to use when CONTAINER_NAME is required but unset.
+
+    When ``UNAVAILABLE_WITHOUT_CONTAINER_NAME`` is truthy and ``CONTAINER_NAME`` is not
+    set, the server should reject requests with the configured status code
+    (``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE``, default 503) rather
+    than passing through to local dev mode. Returns ``None`` when the server should
+    proceed normally (either because ``CONTAINER_NAME`` is set, or because the
+    unavailable-without-container flag is not enabled).
+    """
+    if os.environ.get("CONTAINER_NAME"):
+        return None
+    if not _parse_bool_env("UNAVAILABLE_WITHOUT_CONTAINER_NAME"):
+        return None
+    try:
+        return int(
+            os.environ.get(
+                "UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE",
+                str(DEFAULT_UNAVAILABLE_STATUS_CODE),
+            )
+        )
+    except ValueError:
+        return DEFAULT_UNAVAILABLE_STATUS_CODE
+
 try:
     from cua_agent import ComputerAgent
 
@@ -76,6 +108,77 @@
     redirect_slashes=False,
 )
 
+class UnavailableWithoutContainerMiddleware:
+    """ASGI middleware that rejects all requests when CONTAINER_NAME is required but unset.
+
+    Controlled by env vars (read per-request so tests and dynamic config work):
+    - ``UNAVAILABLE_WITHOUT_CONTAINER_NAME``: if truthy and ``CONTAINER_NAME`` is unset,
+      every HTTP and WebSocket request is rejected.
+    - ``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE``: HTTP status code for
+      rejections (default 503).
+
+    When disabled (either env var not set), requests pass through unchanged, preserving
+    the original "local development mode" behavior for backwards compatibility.
+    """
+
+    _DETAIL = "Service unavailable: CONTAINER_NAME is required but not configured"
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        scope_type = scope.get("type")
+        if scope_type in ("http", "websocket"):
+            status_code = _unavailable_status_code()
+            if status_code is not None:
+                if scope_type == "http":
+                    await self._reject_http(send, status_code)
+                else:
+                    await self._reject_websocket(receive, send, status_code)
+                return
+        await self.app(scope, receive, send)
+
+    @classmethod
+    async def _reject_http(cls, send, status_code):
+        body = json.dumps({"detail": cls._DETAIL}).encode()
+        await send(
+            {
+                "type": "http.response.start",
+                "status": status_code,
+                "headers": [
+                    (b"content-type", b"application/json"),
+                    (b"content-length", str(len(body)).encode()),
+                ],
+            }
+        )
+        await send({"type": "http.response.body", "body": body})
+
+    @classmethod
+    async def _reject_websocket(cls, receive, send, status_code):
+        # Accept first so we can send a structured JSON error before closing — this
+        # preserves the existing error shape that clients already handle.
+        event = await receive()
+        if event.get("type") != "websocket.connect":
+            return
+        await send({"type": "websocket.accept"})
+        await send(
+            {
+                "type": "websocket.send",
+                "text": json.dumps(
+                    {
+                        "success": False,
+                        "error": cls._DETAIL,
+                        "status_code": status_code,
+                    }
+                ),
+            }
+        )
+        # 1008 = Policy Violation
+        await send({"type": "websocket.close", "code": 1008})
+
+
+app.add_middleware(UnavailableWithoutContainerMiddleware)
+
 # CORS configuration
 origins = ["*"]
 app.add_middleware(

libs/python/computer-server/tests/test_auth_availability.py

@@ -0,0 +1,189 @@
+"""Integration tests for UNAVAILABLE_WITHOUT_CONTAINER_NAME behavior.
+
+These tests verify two things:
+
+1. **Backwards compat** — when neither ``CONTAINER_NAME`` nor
+   ``UNAVAILABLE_WITHOUT_CONTAINER_NAME`` is set, the server continues to
+   operate in local development mode (no auth required, requests succeed).
+
+2. **New behavior** — when ``UNAVAILABLE_WITHOUT_CONTAINER_NAME`` is truthy
+   and ``CONTAINER_NAME`` is unset, requests are rejected with the status
+   code configured by
+   ``UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE`` (default 503)
+   rather than being allowed through.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+try:
+    from fastapi.testclient import TestClient
+
+    from computer_server.main import _unavailable_status_code, app
+except Exception as import_error:  # pragma: no cover - environment-dependent
+    pytest.skip(
+        f"computer_server.main unavailable in this environment: {import_error}",
+        allow_module_level=True,
+    )
+
+
+@pytest.fixture
+def clean_env(monkeypatch):
+    """Remove all env vars that influence auth availability for a clean baseline."""
+    for var in (
+        "CONTAINER_NAME",
+        "UNAVAILABLE_WITHOUT_CONTAINER_NAME",
+        "UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE",
+    ):
+        monkeypatch.delenv(var, raising=False)
+    return monkeypatch
+
+
+@pytest.fixture
+def client():
+    return TestClient(app)
+
+
+class TestUnavailableStatusCode:
+    """Unit tests for the `_unavailable_status_code` helper."""
+
+    def test_returns_none_when_both_unset(self, clean_env):
+        assert _unavailable_status_code() is None
+
+    def test_returns_none_when_container_name_set(self, clean_env):
+        clean_env.setenv("CONTAINER_NAME", "vm-abc")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        # CONTAINER_NAME being set overrides the unavailable flag.
+        assert _unavailable_status_code() is None
+
+    def test_returns_default_503_when_flag_truthy_and_container_missing(self, clean_env):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        assert _unavailable_status_code() == 503
+
+    @pytest.mark.parametrize("value", ["1", "true", "True", "YES", "y", "on"])
+    def test_accepts_various_truthy_values(self, clean_env, value):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", value)
+        assert _unavailable_status_code() == 503
+
+    @pytest.mark.parametrize("value", ["0", "false", "no", "", "random"])
+    def test_rejects_falsy_values(self, clean_env, value):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", value)
+        assert _unavailable_status_code() is None
+
+    def test_custom_status_code(self, clean_env):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "1")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "418")
+        assert _unavailable_status_code() == 418
+
+    def test_invalid_status_code_falls_back_to_503(self, clean_env):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "1")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "not-a-number")
+        assert _unavailable_status_code() == 503
+
+
+class TestCmdEndpoint:
+    """Integration tests for the POST /cmd endpoint."""
+
+    def test_backwards_compat_local_dev_allows_requests(self, clean_env, client):
+        # No CONTAINER_NAME, no availability flag — old "local dev" behavior.
+        resp = client.post("/cmd", json={"command": "version", "params": {}})
+        assert resp.status_code == 200, resp.text
+        assert "success" in resp.text
+
+    def test_unavailable_flag_rejects_with_default_503(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        resp = client.post("/cmd", json={"command": "version", "params": {}})
+        assert resp.status_code == 503
+        assert "CONTAINER_NAME" in resp.json()["detail"]
+
+    def test_unavailable_flag_with_custom_status_code(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "1")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "418")
+        resp = client.post("/cmd", json={"command": "version", "params": {}})
+        assert resp.status_code == 418
+
+    def test_container_name_set_bypasses_unavailable_flag(self, clean_env, client):
+        # CONTAINER_NAME being set means auth is required — but the unavailable
+        # flag should NOT apply. Without valid creds, this should 401, not 503.
+        clean_env.setenv("CONTAINER_NAME", "vm-xyz")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        resp = client.post("/cmd", json={"command": "version", "params": {}})
+        assert resp.status_code == 401
+
+
+class TestPtyEndpointAuthGate:
+    """Integration tests for PTY endpoints (via `_require_auth`)."""
+
+    def test_backwards_compat_local_dev_allows_access(self, clean_env, client):
+        # Use a non-existent PID — we just want to verify we get past the auth gate.
+        # If auth passes, we get 404 (PTY not found); if not, we get 401/503.
+        resp = client.get("/pty/999999")
+        assert resp.status_code == 404, resp.text
+
+    def test_unavailable_flag_rejects_with_default_503(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        resp = client.get("/pty/999999")
+        assert resp.status_code == 503
+
+    def test_unavailable_flag_with_custom_status_code(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "599")
+        resp = client.get("/pty/999999")
+        assert resp.status_code == 599
+
+
+class TestPlaywrightExecEndpoint:
+    def test_backwards_compat_local_dev_accepts_auth(self, clean_env, client):
+        # Browser manager may fail for other reasons, but it should NOT be 503/401.
+        resp = client.post("/playwright_exec", json={"command": "noop", "params": {}})
+        assert resp.status_code not in (401, 503)
+
+    def test_unavailable_flag_rejects(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        resp = client.post("/playwright_exec", json={"command": "noop", "params": {}})
+        assert resp.status_code == 503
+
+
+class TestStatusEndpointMiddlewareGating:
+    """The middleware applies uniformly — /status is reachable when the flag is off."""
+
+    def test_status_accessible_in_local_dev_mode(self, clean_env, client):
+        resp = client.get("/status")
+        assert resp.status_code == 200
+        assert resp.json()["status"] == "ok"
+
+    def test_status_rejected_by_middleware_when_flag_set(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        resp = client.get("/status")
+        assert resp.status_code == 503
+
+    def test_status_accessible_when_container_name_set(self, clean_env, client):
+        clean_env.setenv("CONTAINER_NAME", "vm-abc")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        resp = client.get("/status")
+        assert resp.status_code == 200
+
+
+class TestWebSocketEndpoint:
+    def test_backwards_compat_local_dev_allows_commands(self, clean_env, client):
+        with client.websocket_connect("/ws") as ws:
+            ws.send_json({"command": "version", "params": {}})
+            data = ws.receive_json()
+            assert data["success"] is True
+
+    def test_unavailable_flag_closes_with_error(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        with client.websocket_connect("/ws") as ws:
+            data = ws.receive_json()
+            assert data["success"] is False
+            assert data["status_code"] == 503
+            assert "CONTAINER_NAME" in data["error"]
+
+    def test_unavailable_flag_reports_custom_status_code(self, clean_env, client):
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME", "true")
+        clean_env.setenv("UNAVAILABLE_WITHOUT_CONTAINER_NAME_RESPONSE_STATUS_CODE", "599")
+        with client.websocket_connect("/ws") as ws:
+            data = ws.receive_json()
+            assert data["success"] is False
+            assert data["status_code"] == 599