Make auth handoff use same browser lease

Author: claude

ax_browser_broker/api.py

@@ -61,6 +61,7 @@ class LeaseRequest(BaseModel):
     owner: str = "unknown"
     ttl_seconds: int = Field(default=14400, ge=60, le=14400)
     identity_id: str | None = None
+    headed: bool = False
 
 
 class LeaseIdRequest(BaseModel):
@@ -152,7 +153,7 @@ class AuthRequest(BaseModel):
     reason: str = "login_required"
     identity_id: str | None = None
     profile: str | None = None
-    mode: str = "vnc"
+    mode: str = "same_lease"
     ttl_seconds: int = Field(default=900, ge=60, le=14400)
     control_ttl_seconds: int = Field(default=900, ge=60, le=3600)
     wait_until: str = "domcontentloaded"
@@ -165,8 +166,8 @@ def normalize_legacy_profile(self) -> "AuthRequest":
                 raise ValueError("profile and identity_id must match when both are provided")
             self.identity_id = self.profile
         normalized_mode = self.mode.strip().lower().replace("-", "_")
-        if normalized_mode not in {"lease_control", "vnc"}:
-            raise ValueError("mode must be lease_control or vnc")
+        if normalized_mode not in {"lease_control", "same_lease", "vnc"}:
+            raise ValueError("mode must be lease_control, same_lease, or vnc")
         self.mode = normalized_mode
         return self
 
@@ -284,6 +285,76 @@ async def _legacy_vnc_auth_request(request: AuthRequest) -> dict[str, Any]:
     return result
 
 
+async def _same_lease_auth_request(request: AuthRequest) -> dict[str, Any]:
+    if request.identity_id:
+        require_identity(request.identity_id)
+    active_lease_id = _active_identity_lease_id(request.identity_id) if request.identity_id else None
+    lease_obj: dict[str, Any]
+    created_for_auth = False
+    navigation: dict[str, Any] | None = None
+    if active_lease_id:
+        lease_data = require_lease(active_lease_id)
+        if not lease_data.headed:
+            raise HTTPException(
+                status_code=409,
+                detail=(
+                    f"Identity is already leased in a non-visual slot: {request.identity_id}. "
+                    "Release that lease before starting a login handoff."
+                ),
+            )
+        lease_obj = lease_data.__dict__
+    else:
+        lease_obj = await create_lease(
+            LeaseRequest(
+                owner=request.owner,
+                ttl_seconds=request.ttl_seconds,
+                identity_id=request.identity_id,
+                headed=True,
+            )
+        )
+        created_for_auth = True
+        try:
+            navigation = await browser_navigate(
+                NavigateRequest(lease_id=str(lease_obj["lease_id"]), url=request.url, wait_until=request.wait_until)
+            )
+        except Exception:
+            await release_lease(str(lease_obj["lease_id"]))
+            raise
+    auth_data = create_auth_request(
+        request.owner,
+        request.url,
+        request.reason,
+        request.identity_id,
+        mode="same_lease",
+        lease_id=str(lease_obj["lease_id"]),
+        slot=str(lease_obj.get("name") or ""),
+        cdp=str(lease_obj.get("cdp") or ""),
+    )
+    auth_data["lease_id"] = str(lease_obj["lease_id"])
+    auth_data["lease"] = lease_obj
+    auth_data["auth_url"] = auth_data["portal_url"]
+    auth_data["handoff_url"] = auth_data["portal_url"]
+    auth_data["created_lease"] = created_for_auth
+    if navigation:
+        auth_data["navigation"] = navigation
+    _safe_record_event(
+        source=request.owner,
+        event_type="auth",
+        message="Same-lease auth handoff request created",
+        lease_id=str(lease_obj["lease_id"]),
+        url=request.url,
+        tags=["auth", "same-lease", request.reason],
+        data={
+            "token": auth_data.get("token"),
+            "status": auth_data.get("status"),
+            "identity_id": request.identity_id,
+            "slot": lease_obj.get("name"),
+            "created_lease": created_for_auth,
+        },
+    )
+    return auth_data
+
+
 async def _active_identity_control_response(request: AuthRequest) -> dict[str, Any] | None:
     if not request.identity_id:
         return None
@@ -369,6 +440,11 @@ async def _active_identity_control_response(request: AuthRequest) -> dict[str, A
 
 
 async def _open_auth_lease_control(request: AuthRequest) -> dict[str, Any]:
+    if request.mode == "same_lease":
+        if not request.identity_id:
+            legacy_request = request.model_copy(update={"mode": "vnc"})
+            return await _legacy_vnc_auth_request(legacy_request)
+        return await _same_lease_auth_request(request)
     if request.mode == "vnc":
         return await _legacy_vnc_auth_request(request)
     active_response = await _active_identity_control_response(request)
@@ -636,6 +712,18 @@ async def _verify_live_auth_browser(cdp: str | None, target_url: str, host: str
         return {"ok": False, "checked": False, "host": host, "url": target_url, "error": str(error)}
 
 
+async def _verify_live_lease_page_state(lease_id: str, target_url: str, host: str | None) -> dict[str, Any]:
+    if not lease_id:
+        return {"ok": False, "checked": False, "error": "missing lease_id"}
+    try:
+        if target_url:
+            await browser_navigate(NavigateRequest(lease_id=lease_id, url=target_url))
+        snapshot = await browser_snapshot(LeaseIdRequest(lease_id=lease_id))
+        return _assess_auth_page_state(snapshot, host)
+    except Exception as error:
+        return {"ok": False, "checked": False, "host": host, "url": target_url, "error": str(error)}
+
+
 def _record_browser_failure(request: LeaseIdRequest, action: str, error: Exception, data: dict[str, Any] | None = None) -> None:
     _safe_record_event(
         source="broker-api",
@@ -2223,6 +2311,11 @@ async def openbrowser_docs(_auth: str = Depends(require_openbrowser_api_key)) ->
         "dashboard": "/openbrowser",
         "base_url": _openbrowser_base_url(),
         "endpoints": _openbrowser_endpoint_catalog(),
+        "agent_guidance": {
+            "quickstart": docs("quickstart"),
+            "auth": docs("auth"),
+            "routing": docs("routing"),
+        },
         "identities": {
             "generic": "omit identity_id for a neutral non-account browser",
             "configured": {
@@ -2467,7 +2560,7 @@ async def openbrowser_telemetry_summary(
 @app.post("/lease")
 async def create_lease(request: LeaseRequest) -> dict[str, Any]:
     try:
-        lease_obj = lease(request.owner, request.ttl_seconds, request.identity_id)
+        lease_obj = lease(request.owner, request.ttl_seconds, request.identity_id, headed=request.headed)
         result = lease_obj.__dict__
         _safe_record_event(
             source=lease_obj.owner,
@@ -4474,6 +4567,44 @@ async def auth_complete(token: str) -> dict[str, Any]:
         pending_target_url = str(pending_request.get("url") or "")
         pending_host = urllib.parse.urlsplit(pending_target_url).hostname if pending_target_url else None
         pending_vnc = pending_request.get("vnc") or {}
+        if pending_request.get("mode") == "same_lease":
+            live_verification = await _verify_live_lease_page_state(
+                str(pending_request.get("lease_id") or ""),
+                pending_target_url,
+                pending_host,
+            )
+            if not live_verification.get("ok"):
+                response = dict(pending_request)
+                response["completion_blocked"] = True
+                response["auth_verified"] = False
+                response["live_page_verification"] = live_verification
+                _safe_record_event(
+                    source=str(pending_request.get("owner", "unknown")),
+                    event_type="auth",
+                    message="Same-lease auth completion blocked because live tab still appears signed out",
+                    severity="warning",
+                    lease_id=str(pending_request.get("lease_id") or ""),
+                    url=pending_target_url,
+                    tags=["auth", "complete", "blocked", "same-lease", "live-page-signed-out"],
+                    data={"token": token, "identity_id": pending_identity_id, "verification": live_verification},
+                )
+                return response
+            request = complete_auth_request(token)
+            request["vnc_stop"] = stop_auth_vnc(token, missing_ok=True)
+            request["auth_verified"] = True
+            request["live_page_verification"] = live_verification
+            request["lease_id"] = str(pending_request.get("lease_id") or "")
+            request["handoff_url"] = request.get("portal_url")
+            _safe_record_event(
+                source=str(request.get("owner", "unknown")),
+                event_type="auth",
+                message="Same-lease auth request completed",
+                lease_id=str(request.get("lease_id") or ""),
+                url=str(request.get("url") or ""),
+                tags=["auth", "complete", "same-lease"],
+                data={"token": token, "status": request.get("status")},
+            )
+            return request
         if pending_identity_id and pending_vnc.get("cdp"):
             live_verification = await _verify_live_auth_browser(
                 str(pending_vnc.get("cdp")),

ax_browser_broker/auth.py

@@ -145,11 +145,14 @@ def create_auth_request(
     reason: str = "login_required",
     identity_id: str | None = None,
     mode: str = "vnc",
+    lease_id: str | None = None,
+    slot: str | None = None,
+    cdp: str | None = None,
 ) -> dict[str, Any]:
     now = int(time.time())
     normalized_mode = mode.strip().lower().replace("-", "_")
-    if normalized_mode not in {"lease_control", "vnc"}:
-        raise AuthError("mode must be lease_control or vnc")
+    if normalized_mode not in {"lease_control", "same_lease", "vnc"}:
+        raise AuthError("mode must be lease_control, same_lease, or vnc")
     if identity_id:
         require_identity(identity_id)
     token = secrets.token_urlsafe(24)
@@ -167,6 +170,12 @@ def create_auth_request(
     }
     if identity_id:
         request["identity_id"] = identity_id
+    if lease_id:
+        request["lease_id"] = lease_id
+    if slot:
+        request["slot"] = slot
+    if cdp:
+        request["cdp"] = cdp
     with locked_auth_state() as state:
         gc_auth_requests(state)
         state["requests"][token] = request
@@ -605,7 +614,10 @@ def start_auth_vnc(token: str, websocket_port: int = 6081, vnc_port: int = 5901)
     os.chmod(password_file, 0o600)
 
     try:
-        if request.get("identity_id"):
+        if request.get("mode") == "same_lease":
+            vnc_state = _start_lease_vnc(request, websocket_port, vnc_port, password_file, log_path)
+            display = str(vnc_state["display"])
+        elif request.get("identity_id"):
             vnc_state = _start_identity_auth_vnc(request, websocket_port, vnc_port, password_file, log_path)
             display = str(vnc_state["display"])
         else:
@@ -743,6 +755,80 @@ def _close_chrome_via_cdp(cdp_port: int) -> bool:
         return False
 
 
+def _start_lease_vnc(
+    request: dict[str, Any],
+    websocket_port: int,
+    vnc_port: int,
+    password_file: Path,
+    log_path: Path,
+) -> dict[str, Any]:
+    from .pool import require_lease
+
+    lease_id = str(request.get("lease_id") or "")
+    if not lease_id:
+        raise AuthError("same_lease auth request is missing lease_id")
+    lease = require_lease(lease_id)
+    display_file = BROWSER_POOL_DIR / "state" / f"{lease.name}.display"
+    try:
+        display = display_file.read_text(encoding="utf-8").strip()
+    except OSError as error:
+        raise AuthError(f"Headed display is not available for lease slot {lease.name}") from error
+    if not display:
+        raise AuthError(f"Headed display is empty for lease slot {lease.name}")
+    x11vnc = shutil.which("x11vnc")
+    websockify = shutil.which("websockify")
+    if not x11vnc or not websockify:
+        raise AuthError("x11vnc or websockify is missing")
+    env = os.environ.copy()
+    env["DISPLAY"] = display
+    with log_path.open("ab") as log:
+        x11vnc_proc = subprocess.Popen(
+            [
+                x11vnc,
+                "-display",
+                display,
+                "-rfbport",
+                str(vnc_port),
+                "-localhost",
+                "-forever",
+                "-shared",
+                "-passwdfile",
+                str(password_file),
+            ],
+            stdout=log,
+            stderr=log,
+            env=env,
+            start_new_session=True,
+        )
+        websockify_proc = subprocess.Popen(
+            [
+                websockify,
+                "--web=/usr/share/novnc",
+                f"127.0.0.1:{websocket_port}",
+                f"127.0.0.1:{vnc_port}",
+            ],
+            stdout=log,
+            stderr=log,
+            env=env,
+            start_new_session=True,
+        )
+    return {
+        "mode": "same-lease",
+        "lease_id": lease_id,
+        "slot": lease.name,
+        "identity_id": lease.identity_id,
+        "display": display,
+        "x11vnc_pid": x11vnc_proc.pid,
+        "websockify_pid": websockify_proc.pid,
+        "websocket_port": websocket_port,
+        "vnc_port": vnc_port,
+        "password_file": str(password_file),
+        "started_at": int(time.time()),
+        "profile_dir": lease.profile_dir,
+        "cdp": lease.cdp,
+    }
+
+
 def stop_auth_vnc(token: str, missing_ok: bool = False) -> dict[str, Any]:
     try:
         request = get_auth_request(token)

ax_browser_broker/browser_pool/bin/launch_chrome.sh

@@ -11,8 +11,12 @@ PROFILE_DIR="${BROWSER_POOL_DIR}/profiles/${NAME}"
 LOG="${BROWSER_POOL_DIR}/logs/${NAME}.log"
 MAINTENANCE_FILE="${BROWSER_POOL_DIR}/state/maintenance/${NAME}.json"
 PROXY_PID_FILE="${BROWSER_POOL_DIR}/state/${NAME}.proxy.pid"
+XVFB_PID_FILE="${BROWSER_POOL_DIR}/state/${NAME}.xvfb.pid"
+DISPLAY_FILE="${BROWSER_POOL_DIR}/state/${NAME}.display"
 PROXY_ARGS=()
 SYNC_ARGS=()
+HEADLESS_ARGS=(--headless)
+DISPLAY_ARGS=()
 CHROME_LANG="en-US"
 
 find_chrome_bin() {
@@ -88,6 +92,14 @@ if [[ -f "$PROXY_PID_FILE" ]]; then
   fi
   rm -f "$PROXY_PID_FILE"
 fi
+if [[ -f "$XVFB_PID_FILE" ]]; then
+  old_xvfb_pid="$(cat "$XVFB_PID_FILE" 2>/dev/null || true)"
+  if [[ -n "$old_xvfb_pid" ]]; then
+    kill "$old_xvfb_pid" 2>/dev/null || true
+  fi
+  rm -f "$XVFB_PID_FILE"
+fi
+rm -f "$DISPLAY_FILE"
 sleep 1
 
 if [[ -n "${PROXY_REF:-}" ]]; then
@@ -142,12 +154,29 @@ if ! CHROME_BIN="$(find_chrome_bin)"; then
   exit 1
 fi
 
+if [[ "${CHROME_HEADLESS:-1}" == "0" ]]; then
+  XVFB_BIN="${OPENBROWSER_XVFB_BIN:-$(command -v Xvfb || true)}"
+  if [[ -z "$XVFB_BIN" ]]; then
+    echo "$(date -Is) ${NAME} headed mode requested but Xvfb is missing" >>"$LOG"
+    exit 1
+  fi
+  DISPLAY_NUM="${CHROME_DISPLAY_NUM:-$((PORT - 9000))}"
+  DISPLAY=":${DISPLAY_NUM}"
+  "$XVFB_BIN" "$DISPLAY" -screen 0 "${CHROME_WINDOW_SIZE:-1280x800x24}" -nolisten tcp >>"$LOG" 2>&1 &
+  XVFB_PID=$!
+  echo "$XVFB_PID" > "$XVFB_PID_FILE"
+  echo "$DISPLAY" > "$DISPLAY_FILE"
+  export DISPLAY
+  HEADLESS_ARGS=()
+  DISPLAY_ARGS=(--ozone-platform=x11)
+fi
+
 (
   if [[ -w "/proc/${BASHPID}/oom_score_adj" ]]; then
     echo 300 >"/proc/${BASHPID}/oom_score_adj" || true
   fi
   exec "$CHROME_BIN" \
-    --headless \
+    "${HEADLESS_ARGS[@]}" \
     --user-data-dir="$PROFILE_DIR" \
     --no-sandbox \
     --disable-gpu \
@@ -169,6 +198,7 @@ fi
     --window-size=1280,800 \
     --window-position=0,0 \
     --remote-allow-origins='*' \
+    "${DISPLAY_ARGS[@]}" \
     "${PROXY_ARGS[@]}"
 ) >"$LOG" 2>&1 &
 CHROME_PID=$!