fix(agent-proxy): survive a client that drops during agent-VM startup (#10880)

agent_ws accepts the client socket before booting the user's agent VM, and
_wait_for_vm_healthy polls for up to 120s — far past uvicorn's WebSocket
keepalive window. When the phone is dropped mid-wait ("sent 1011 keepalive
ping timeout"), the startup path's unguarded send_text/close raised
WebSocketDisconnect straight out of the handler: 24 "Exception in ASGI
application" tracebacks per day on prod agent-proxy, every one of them at the
"Agent VM is not responding" send.

The relay loop below already owned this condition (except Exception plus a
guarded final close); the startup path did not. Route every startup status,
error, and close through one client-liveness seam that treats peer-gone as the
terminal, expected end of the connection, so the handler still emits its
intended 4002/4003 close code and returns cleanly.

Verified: backend/tests/unit/test_agent_proxy_startup_client_gone.py drives the
real agent_ws against a socket that raises WebSocketDisconnect on every write —
both cases fail with the production exception on the parent commit and pass
here; the sibling agent-proxy suite stays green (8 passed).

Failure-Class: FC-peer-close-aborts-owed-write

Author: kodjima33

backend/agent-proxy/main.py

@@ -426,6 +426,31 @@ async def _wait_for_vm_healthy(vm_ip: str, auth_token: str, timeout: float = 120
     return False
 
 
+async def _send_startup_event(websocket: WebSocket, uid: str, payload: Dict[str, Any]) -> bool:
+    """Push a VM-startup status/error event to the client. False means the client is gone.
+
+    VM startup outlives the client's keepalive window (the health wait alone is 120s), so by
+    the time these events go out uvicorn may already have dropped the phone with a 1011 ping
+    timeout. A vanished client is the terminal, expected end of this connection — not an ASGI
+    error — and must not escape agent_ws, which is the only ownership the relay loop below
+    already has and this startup path did not.
+    """
+    try:
+        await websocket.send_text(json.dumps(payload))
+        return True
+    except Exception as e:
+        logger.info(f"[agent-proxy] uid={uid} client gone during VM startup: {type(e).__name__}")
+        return False
+
+
+async def _close_client(websocket: WebSocket, uid: str, code: int, reason: str) -> None:
+    """Close the client socket, tolerating a client that already went away."""
+    try:
+        await websocket.close(code=code, reason=reason)
+    except Exception as e:
+        logger.debug(f"[agent-proxy] uid={uid} close({code}) on gone client: {type(e).__name__}")
+
+
 # --------------- encryption helpers ---------------
 
 
@@ -624,34 +649,34 @@ async def agent_ws(websocket: WebSocket):
         except Exception:
             # VM not reachable — check GCE and restart/reset if needed
             logger.info(f"[agent-proxy] uid={uid} VM {vm_ip} not reachable, checking GCE...")
-            await websocket.send_text(json.dumps({"type": "status", "message": "Starting your agent VM..."}))
+            await _send_startup_event(websocket, uid, {"type": "status", "message": "Starting your agent VM..."})
             vm = await _ensure_vm_running(uid, vm, health_failed=True)
             if not vm or vm.get("status") != "ready" or not vm.get("ip"):
-                await websocket.send_text(json.dumps(await run_blocking(db_executor, _vm_unavailable_event, uid)))
-                await websocket.close(code=4002, reason="VM startup failed")
+                await _send_startup_event(websocket, uid, await run_blocking(db_executor, _vm_unavailable_event, uid))
+                await _close_client(websocket, uid, 4002, "VM startup failed")
                 return
             vm_ip = vm["ip"]
             vm_token = vm["authToken"]
             # Wait for VM to be healthy after restart
             healthy = await _wait_for_vm_healthy(vm_ip, vm_token)
             if not healthy:
-                await websocket.send_text(json.dumps({"type": "error", "message": "Agent VM is not responding"}))
-                await websocket.close(code=4003, reason="VM not healthy")
+                await _send_startup_event(websocket, uid, {"type": "error", "message": "Agent VM is not responding"})
+                await _close_client(websocket, uid, 4003, "VM not healthy")
                 return
     else:
         # No IP or not ready — must restart
-        await websocket.send_text(json.dumps({"type": "status", "message": "Starting your agent VM..."}))
+        await _send_startup_event(websocket, uid, {"type": "status", "message": "Starting your agent VM..."})
         vm = await _ensure_vm_running(uid, vm)
         if not vm or vm.get("status") != "ready" or not vm.get("ip"):
-            await websocket.send_text(json.dumps(await run_blocking(db_executor, _vm_unavailable_event, uid)))
-            await websocket.close(code=4002, reason="VM startup failed")
+            await _send_startup_event(websocket, uid, await run_blocking(db_executor, _vm_unavailable_event, uid))
+            await _close_client(websocket, uid, 4002, "VM startup failed")
             return
         vm_ip = vm["ip"]
         vm_token = vm["authToken"]
         healthy = await _wait_for_vm_healthy(vm_ip, vm_token)
         if not healthy:
-            await websocket.send_text(json.dumps({"type": "error", "message": "Agent VM is not responding"}))
-            await websocket.close(code=4003, reason="VM not healthy")
+            await _send_startup_event(websocket, uid, {"type": "error", "message": "Agent VM is not responding"})
+            await _close_client(websocket, uid, 4003, "VM not healthy")
             return
 
     vm_uri = f"ws://{vm_ip}:8080/ws?token={vm_token}"

backend/tests/unit/test_agent_proxy_startup_client_gone.py

@@ -0,0 +1,111 @@
+"""A client that vanishes during agent-VM startup must end the connection, not crash agent_ws.
+
+``agent_ws`` accepts the client socket and only then boots/waits for the user's VM —
+``_wait_for_vm_healthy`` alone polls for 120s, far past uvicorn's WebSocket keepalive
+window. When the phone is dropped mid-wait (``sent 1011 keepalive ping timeout``), the
+startup path's unguarded ``send_text``/``close`` raised ``WebSocketDisconnect`` straight
+out of the handler: 24 "Exception in ASGI application" tracebacks per day on prod
+agent-proxy, all at the "Agent VM is not responding" send. The relay loop below already
+owned this (``except Exception`` + guarded close); the startup path did not.
+"""
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from types import ModuleType
+from unittest.mock import MagicMock
+
+import firebase_admin
+import pytest
+from fastapi import WebSocketDisconnect
+from firebase_admin import firestore
+
+BACKEND_DIR = Path(__file__).resolve().parents[2]
+AGENT_PROXY_DIR = BACKEND_DIR / "agent-proxy"
+if str(BACKEND_DIR) not in sys.path:
+    sys.path.insert(0, str(BACKEND_DIR))
+# `agent-proxy/main.py` imports its siblings by bare name (`from resilience import ...`),
+# so loading it by file path also needs its own directory importable.
+if str(AGENT_PROXY_DIR) not in sys.path:
+    sys.path.insert(0, str(AGENT_PROXY_DIR))
+
+
+@pytest.fixture
+def agent_proxy(monkeypatch) -> ModuleType:
+    monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS", raising=False)
+    monkeypatch.setattr(firebase_admin, "initialize_app", MagicMock(return_value=object()))
+    monkeypatch.setattr(firestore, "client", MagicMock(return_value=object()))
+
+    spec = importlib.util.spec_from_file_location("agent_proxy_startup_client_gone_test", AGENT_PROXY_DIR / "main.py")
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class _GoneClientWebSocket:
+    """A client socket that uvicorn already dropped: every write reports the disconnect."""
+
+    def __init__(self) -> None:
+        self.headers = {"authorization": "Bearer test-token"}
+        self.accepted = False
+        self.send_attempts: list[str] = []
+        self.close_attempts: list[int] = []
+
+    async def accept(self) -> None:
+        self.accepted = True
+
+    async def send_text(self, text: str) -> None:
+        self.send_attempts.append(text)
+        raise WebSocketDisconnect(code=1006)
+
+    async def close(self, code: int = 1000, reason: str = "") -> None:
+        self.close_attempts.append(code)
+        raise WebSocketDisconnect(code=1006)
+
+
+def _stub_startup(agent_proxy: ModuleType, monkeypatch, *, ensure_result, healthy: bool) -> None:
+    """Drive agent_ws down the restart path with an authenticated uid and no live VM."""
+
+    async def direct_run_blocking(_executor, func, *args, **kwargs):
+        return func(*args, **kwargs)
+
+    async def ensure_vm_running(_uid, _vm, health_failed=False):
+        return ensure_result
+
+    async def wait_for_vm_healthy(_ip, _token, timeout=120):
+        return healthy
+
+    monkeypatch.setattr(agent_proxy, "run_blocking", direct_run_blocking)
+    monkeypatch.setattr(agent_proxy, "_verify_id_token", lambda _token: {"uid": "uid-gone"})
+    monkeypatch.setattr(
+        agent_proxy,
+        "_get_user_context",
+        lambda _uid: ({"vmName": "omi-agent-gone", "zone": "us-central1-a", "status": "stopped"}, "standard"),
+    )
+    monkeypatch.setattr(agent_proxy, "_ensure_vm_running", ensure_vm_running)
+    monkeypatch.setattr(agent_proxy, "_wait_for_vm_healthy", wait_for_vm_healthy)
+    monkeypatch.setattr(agent_proxy, "_vm_unavailable_event", lambda _uid: {"type": "error", "code": "unavailable"})
+    monkeypatch.setattr(agent_proxy, "httpx", types.SimpleNamespace(AsyncClient=MagicMock()))
+
+
+class TestAgentWsStartupSurvivesAGoneClient:
+    async def test_unavailable_vm_does_not_raise_when_the_client_is_gone(self, agent_proxy, monkeypatch):
+        _stub_startup(agent_proxy, monkeypatch, ensure_result=None, healthy=False)
+        websocket = _GoneClientWebSocket()
+
+        await agent_proxy.agent_ws(websocket)
+
+        assert websocket.accepted
+        assert websocket.close_attempts == [4002], "the connection must still be closed with its startup-failure code"
+
+    async def test_unhealthy_vm_does_not_raise_when_the_client_is_gone(self, agent_proxy, monkeypatch):
+        ready_vm = {"vmName": "omi-agent-gone", "status": "ready", "ip": "34.9.9.9", "authToken": "vm-token"}
+        _stub_startup(agent_proxy, monkeypatch, ensure_result=ready_vm, healthy=False)
+        websocket = _GoneClientWebSocket()
+
+        await agent_proxy.agent_ws(websocket)
+
+        assert [s for s in websocket.send_attempts if "not responding" in s], "the error event must still be attempted"
+        assert websocket.close_attempts == [4003], "the connection must still be closed with its not-healthy code"