Fix WebCRT sessions killed by StrictMode ws_ticket reuse.

Allow short-lived tickets multiple uses, stop deleting PTY on transient websocket errors, and avoid marking connect complete after a dropped wait loop.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
oliver 2026-08-02 18:49:29 +08:00
parent b42f4eb588
commit 9314f11ce7
4 changed files with 62 additions and 17 deletions

View file

@ -440,7 +440,8 @@ async def websocket_session(websocket: WebSocket, session_id: str) -> None:
} }
) )
except Exception: except Exception:
break # Client dropped during connect wait (StrictMode remount etc.) — keep PTY.
return
remaining = deadline - time.time() remaining = deadline - time.time()
if remaining <= 0: if remaining <= 0:
await websocket.send_json( await websocket.send_json(
@ -455,11 +456,19 @@ async def websocket_session(websocket: WebSocket, session_id: str) -> None:
lambda t=slice_timeout: wait_session_ready(session_id, timeout=t), lambda t=slice_timeout: wait_session_ready(session_id, timeout=t),
) )
sess = get_session(session_id) or cur sess = get_session(session_id) or cur
if sess is None or sess.closed or sess.state == "closed":
return
if sess.state == "connecting":
# wait_session_ready should not return while still connecting.
continue
break break
except HTTPException as exc: except HTTPException as exc:
if exc.status_code == 504: if exc.status_code == 504:
# Slice timeout while still connecting — keep polling with progress. # Slice timeout while still connecting — keep polling with progress.
continue continue
if exc.status_code == 404:
# Session deleted while waiting.
return
await websocket.send_json( await websocket.send_json(
{"type": "status", "state": "error", "message": str(exc.detail)} {"type": "status", "state": "error", "message": str(exc.detail)}
) )
@ -472,6 +481,21 @@ async def websocket_session(websocket: WebSocket, session_id: str) -> None:
await websocket.close(code=4502) await websocket.close(code=4502)
return return
# Session may have been deleted while the previous wait loop was exiting.
if get_session(session_id) is None or sess.closed or sess.state in {"closed", "error"}:
if sess.state == "error":
try:
await websocket.send_json(
{
"type": "status",
"state": "error",
"message": sess.connect_error or "connect_failed",
}
)
except Exception:
pass
return
await websocket.send_json( await websocket.send_json(
{ {
"type": "status", "type": "status",

View file

@ -158,7 +158,8 @@ def find_ssh_session_for_ne(ne_id: str) -> WebcrtSession | None:
def wait_session_ready(session_id: str, *, timeout: float = 120.0) -> WebcrtSession: def wait_session_ready(session_id: str, *, timeout: float = 120.0) -> WebcrtSession:
"""Block until async connect finishes (ready or error). Used by tests and WS.""" """Block until async connect finishes (ready or error). Used by tests and WS."""
deadline = time.time() + max(1.0, float(timeout)) # Honor short slice timeouts from the WS wait loop (do not clamp to 1s).
deadline = time.time() + max(0.05, float(timeout))
while time.time() < deadline: while time.time() < deadline:
sess = get_session(session_id) sess = get_session(session_id)
if sess is None: if sess is None:
@ -167,7 +168,9 @@ def wait_session_ready(session_id: str, *, timeout: float = 120.0) -> WebcrtSess
return sess return sess
if sess.state == "error": if sess.state == "error":
raise HTTPException(status_code=502, detail=sess.connect_error or "connect_failed") raise HTTPException(status_code=502, detail=sess.connect_error or "connect_failed")
sess._ready_event.wait(timeout=0.25) if sess.closed or sess.state == "closed":
raise HTTPException(status_code=404, detail="webcrt_session_not_found")
sess._ready_event.wait(timeout=min(0.25, max(0.01, deadline - time.time())))
raise HTTPException(status_code=504, detail="connect_timeout") raise HTTPException(status_code=504, detail="connect_timeout")

View file

@ -1,4 +1,8 @@
"""Short-lived WebCRT WebSocket tickets (avoid putting JWT in query strings).""" """Short-lived WebCRT WebSocket tickets (avoid putting JWT in query strings).
Tickets allow a small number of uses so React StrictMode remounts (dev) can
open a second WebSocket with the same URL without racing a one-shot consume.
"""
from __future__ import annotations from __future__ import annotations
@ -8,8 +12,11 @@ import time
from dataclasses import dataclass from dataclasses import dataclass
_TTL_SEC = 90 _TTL_SEC = 90
# React 18 StrictMode mounts effects twice in DEV; keep a spare use.
_MAX_USES = 3
_lock = threading.Lock() _lock = threading.Lock()
_tickets: dict[str, tuple[float, str, frozenset[str]]] = {} # ticket -> (expires_at, user_id, scopes, uses_remaining)
_tickets: dict[str, tuple[float, str, frozenset[str], int]] = {}
@dataclass(frozen=True) @dataclass(frozen=True)
@ -23,7 +30,7 @@ def issue_ws_ticket(*, user_id: str, scopes: frozenset[str], ttl_sec: int = _TTL
exp = time.time() + max(15, int(ttl_sec)) exp = time.time() + max(15, int(ttl_sec))
with _lock: with _lock:
_purge_locked() _purge_locked()
_tickets[tid] = (exp, str(user_id), frozenset(scopes)) _tickets[tid] = (exp, str(user_id), frozenset(scopes), int(_MAX_USES))
return tid, max(15, int(ttl_sec)) return tid, max(15, int(ttl_sec))
@ -33,17 +40,23 @@ def consume_ws_ticket(ticket: str) -> TicketInfo | None:
return None return None
with _lock: with _lock:
_purge_locked() _purge_locked()
row = _tickets.pop(raw, None) row = _tickets.get(raw)
if row is None: if row is None:
return None return None
exp, user_id, scopes = row exp, user_id, scopes, uses_left = row
if exp < time.time(): if exp < time.time() or uses_left <= 0:
_tickets.pop(raw, None)
return None return None
uses_left -= 1
if uses_left <= 0:
_tickets.pop(raw, None)
else:
_tickets[raw] = (exp, user_id, scopes, uses_left)
return TicketInfo(user_id=user_id, scopes=scopes) return TicketInfo(user_id=user_id, scopes=scopes)
def _purge_locked() -> None: def _purge_locked() -> None:
now = time.time() now = time.time()
dead = [k for k, (exp, _, _) in _tickets.items() if exp < now] dead = [k for k, (exp, _, _, _) in _tickets.items() if exp < now]
for k in dead: for k in dead:
_tickets.pop(k, None) _tickets.pop(k, None)

View file

@ -2341,14 +2341,19 @@ export function WebcrtPage() {
void openTarget(tab.target, { force: true }); void openTarget(tab.target, { force: true });
return; return;
} }
// Transient browser WS failures (StrictMode remount / ticket race) must
// NOT delete the device PTY — keep session_id for reconnect/reattach.
const transientWs =
/^(websocket_error|websocket_closed)/i.test(String(message || "")) ||
String(message || "").toLowerCase().includes("websocket");
updateTab(tab.key, { updateTab(tab.key, {
status: "error", status: transientWs ? "closed" : "error",
sessionId: "", sessionId: transientWs ? tab.sessionId : "",
wsUrl: "", wsUrl: transientWs ? tab.wsUrl : "",
connectPhase: undefined, connectPhase: undefined,
errorMessage: webcrtErrorMessage(errMsg, t), errorMessage: webcrtErrorMessage(errMsg, t),
}); });
if (sid) { if (sid && !transientWs) {
void closeWebcrtSession(sid).catch(() => undefined); void closeWebcrtSession(sid).catch(() => undefined);
} }
if (isInventoryTarget(tab.target) && isSshAuthFailure(errMsg)) { if (isInventoryTarget(tab.target) && isSshAuthFailure(errMsg)) {