Erosika
beab8b6f27
`_flush_session` discarded the observation flags when it rebuilt an evicted SDK session, and the cached path returned none, so recall fell back to the config snapshot. Both paths now return and store the flags. A deferred `save()` on a session the cap evicted puts it back in the cache, or flushes it inline when a newer object owns the key. `save()` and `stop_async_writer()` share the writer lock, and the writer drains its queue after the join, so a put that raced shutdown is written. The trim after a flush runs under the cache lock. The shutdown join takes the remaining budget instead of a fixed ten seconds. The injection audit file is created owner-only, and `logging: "false"` reads as off. The desktop passes `<provider>:<user id>` so a basic-auth alice and an OIDC alice are two peers. When a gateway platform supplies no user id, the peer notice and tool error no longer recommend peerName, which would merge every user of that gateway onto one peer. README documents `injection.sessionStart`, `logging`, and what a dashboard login does to peer resolution.
407 lines
22 KiB
Python
407 lines
22 KiB
Python
"""WebSocket transport for the tui_gateway JSON-RPC server: reuses :func:`tui_gateway.server.dispatch`
|
|
verbatim so every RPC, slash command, approval flow and agent event takes the same handlers as Ink over
|
|
stdio. Wire protocol is identical to stdio (newline-delimited JSON-RPC both ways; ``gateway.ready`` right
|
|
after accept). Mount as ``@app.websocket("/api/ws") async def ws(ws): await handle_ws(ws)``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import json
|
|
import logging
|
|
import socket
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
from tui_gateway import server
|
|
from agent.message_sanitization import _sanitize_surrogates
|
|
from tui_gateway.event_replay import replay_epoch
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
# Scale-to-zero: tell the (separate) gateway process a dashboard/desktop/TUI client is attached via
|
|
# the mtime of a marker file it reads in its idle predicate (gateway/scale_to_zero.py). Clients ping
|
|
# every 15s; one write per 5s per process is plenty.
|
|
_DASHBOARD_CLIENT_TOUCH_MIN_INTERVAL_S = 5.0
|
|
_dashboard_client_touched_at = 0.0
|
|
_dashboard_client_touch_lock = threading.Lock()
|
|
|
|
|
|
def _note_dashboard_client_activity(*, force: bool = False) -> None:
|
|
"""Refresh the dashboard-client liveness marker (throttled, best-effort)."""
|
|
global _dashboard_client_touched_at
|
|
now = time.monotonic()
|
|
with _dashboard_client_touch_lock:
|
|
if not force and now - _dashboard_client_touched_at < _DASHBOARD_CLIENT_TOUCH_MIN_INTERVAL_S:
|
|
return
|
|
_dashboard_client_touched_at = now
|
|
try:
|
|
from gateway.scale_to_zero import touch_dashboard_client_heartbeat
|
|
touch_dashboard_client_heartbeat()
|
|
except Exception: # noqa: BLE001 - liveness garnish must never break the WS
|
|
_log.debug("dashboard client heartbeat touch failed", exc_info=True)
|
|
|
|
|
|
def _sanitize_ws_text(text: str) -> str:
|
|
"""Return *text* that can be UTF-8 encoded for a WebSocket frame.
|
|
|
|
``json.dumps(..., ensure_ascii=False)`` happily emits lone UTF-16 surrogates; Starlette's
|
|
``send_text`` then raises ``UnicodeEncodeError``, which used to latch the whole connection
|
|
closed. Same U+FFFD replacement every other Hermes transport applies.
|
|
|
|
See #97288.
|
|
"""
|
|
return _sanitize_surrogates(text) if text else text
|
|
|
|
|
|
# Max seconds a pool-dispatched handler blocks waiting for the loop to flush a WS frame before we
|
|
# give up waiting (the transport is NOT marked dead).
|
|
_WS_WRITE_TIMEOUT_S = 10.0
|
|
# Max seconds one send_text may await the socket once it is actually running on the loop. A healthy
|
|
# socket returns from send_text without waiting (the frame lands in the transport buffer); only kernel
|
|
# backpressure parks it, so a GIL/loop stall cannot start this clock. Deliberately 3x the worker wait
|
|
# above and under the client's 45s heartbeat deadline (apps/shared json-rpc-gateway): a peer that
|
|
# cannot drain ~48 KiB in 30s is gone, and closing here starts its reconnect instead of leaving every
|
|
# later frame and RPC reply parked behind the writer lock (#106369).
|
|
_WS_SEND_DEADLINE_S = 30.0
|
|
_WS_LOG_PAYLOAD_PREVIEW = 240
|
|
|
|
# Per-token streaming frames are coalesced: buffered and flushed as a batch on a short timer instead
|
|
# of waking the loop once per token (each wakeup competes with the agent turn for the GIL). Keep this
|
|
# set to genuinely high-frequency, display-only events — anything a client must see promptly
|
|
# (tool/approval/status/completion) is non-streaming and flushes the buffer ahead of itself, so
|
|
# ordering is preserved. _TOKEN_COALESCE_S: max buffer wait (~30 fps; imperceptible).
|
|
_STREAMING_EVENT_TYPES = frozenset({"message.delta", "reasoning.delta", "thinking.delta"})
|
|
_TOKEN_COALESCE_S = 0.033
|
|
|
|
# starlette stays optional at import time; fall back to a generic sentinel.
|
|
try:
|
|
from starlette.websockets import WebSocketDisconnect as _WebSocketDisconnect
|
|
except ImportError: # pragma: no cover - starlette is a required install path
|
|
_WebSocketDisconnect = Exception # type: ignore[assignment]
|
|
|
|
|
|
class WSTransport:
|
|
"""Per-connection WS transport. ``write`` is safe from any thread *other than* the loop thread owning the
|
|
socket (pool workers marshal onto the loop and block on the future); from the loop thread itself it would
|
|
deadlock, so it detects that and fires-and-forgets. Loop-thread callers needing completion use ``write_async``."""
|
|
|
|
def __init__(self, ws: Any, loop: asyncio.AbstractEventLoop, *, peer: str = "unknown",
|
|
auth_identity: dict | None = None) -> None:
|
|
self._ws = ws
|
|
self._loop = loop
|
|
self._peer = peer
|
|
#: Server-verified identity from the WS-upgrade credential, stamped by ``web_server_chat._ws_auth_reason``; None
|
|
#: for legacy-token/stdio. RPC params can never populate it: sole identity authority for browser controllers
|
|
#: and for the ``user_id`` the agent is built with (``server._session_auth_user_id``).
|
|
self.auth_identity = auth_identity
|
|
self._closed = False
|
|
# Token-coalescing buffer. The lock guards the buffer + "armed" flag against worker threads
|
|
# calling write(); the timer handle is only ever touched on the loop thread.
|
|
self._token_lock = threading.Lock()
|
|
self._pending_tokens: list[str] = []
|
|
self._token_flush_handle: asyncio.TimerHandle | None = None
|
|
self._token_flush_armed = False
|
|
# Socket writes need an async boundary: several batches can queue on the loop during a stall.
|
|
self._send_lock = asyncio.Lock()
|
|
|
|
def write(self, obj: dict) -> bool:
|
|
if self._closed:
|
|
return False
|
|
line = json.dumps(obj, ensure_ascii=False)
|
|
try:
|
|
on_loop = asyncio.get_running_loop() is self._loop
|
|
except RuntimeError:
|
|
on_loop = False
|
|
# Streamed token: buffer it and arm the flush timer; the worker returns immediately.
|
|
# call_soon_threadsafe is safe from a worker or the loop.
|
|
params = obj.get("params") if isinstance(obj, dict) else None
|
|
if isinstance(params, dict) and params.get("type") in _STREAMING_EVENT_TYPES:
|
|
with self._token_lock:
|
|
self._pending_tokens.append(line)
|
|
if not self._token_flush_armed:
|
|
self._token_flush_armed = True
|
|
self._loop.call_soon_threadsafe(self._arm_token_flush)
|
|
return not self._closed
|
|
# Non-streaming frame: append behind any buffered tokens and flush the whole batch NOW so it
|
|
# can never overtake them. The send is scheduled INSIDE the lock so wire order matches buffer
|
|
# order even if the coalesce timer fires on the loop at the same moment.
|
|
from agent.async_utils import safe_schedule_threadsafe
|
|
with self._token_lock:
|
|
self._pending_tokens.append(line)
|
|
batch, self._pending_tokens = self._pending_tokens, []
|
|
if on_loop:
|
|
self._loop.create_task(self._safe_send_many(batch))
|
|
return True
|
|
fut = safe_schedule_threadsafe(self._safe_send_many(batch), self._loop)
|
|
if fut is None:
|
|
self._closed = True
|
|
return False
|
|
try:
|
|
fut.result(timeout=_WS_WRITE_TIMEOUT_S)
|
|
return not self._closed
|
|
except concurrent.futures.TimeoutError: # builtin TimeoutError on 3.11+
|
|
# The loop is stalled (GIL-heavy turn, delegation), NOT the socket dead: the send is already
|
|
# scheduled and flushes once the loop breathes. Latching _closed here permanently silenced
|
|
# live windows after one slow write; _safe_send_many latches on a real error.
|
|
_log.warning("ws write slow (loop stalled >%ss) peer=%s — frame left in flight", _WS_WRITE_TIMEOUT_S, self._peer)
|
|
return not self._closed
|
|
except Exception as exc:
|
|
self._closed = True
|
|
_log.warning("ws write failed peer=%s error_type=%s error=%s", self._peer, type(exc).__name__, exc)
|
|
return False
|
|
|
|
def _arm_token_flush(self) -> None: # loop thread
|
|
if not self._closed:
|
|
self._token_flush_handle = self._loop.call_later(_TOKEN_COALESCE_S, self._flush_tokens)
|
|
|
|
def _flush_tokens(self) -> None:
|
|
"""Timer callback (loop thread): send buffered tokens as one batch, scheduled under the lock so
|
|
wire order is fixed relative to a concurrent ``write``."""
|
|
with self._token_lock:
|
|
self._token_flush_handle = None
|
|
self._token_flush_armed = False
|
|
batch, self._pending_tokens = self._pending_tokens, []
|
|
if batch and not self._closed:
|
|
self._loop.create_task(self._safe_send_many(batch))
|
|
|
|
@property
|
|
def closed(self) -> bool:
|
|
return self._closed
|
|
|
|
async def write_async(self, obj: dict) -> bool:
|
|
"""Send from the owning loop; awaits until the frame is on the wire. Buffered tokens are flushed
|
|
ahead of it in the SAME batch so nothing slips between."""
|
|
if self._closed:
|
|
return False
|
|
with self._token_lock:
|
|
batch, self._pending_tokens = self._pending_tokens, []
|
|
batch.append(json.dumps(obj, ensure_ascii=False))
|
|
await self._safe_send_many(batch)
|
|
return not self._closed
|
|
|
|
async def _safe_send_many(self, lines: list[str]) -> None:
|
|
"""Send one indivisible batch of pre-serialized frames in wire order."""
|
|
async with self._send_lock:
|
|
if self._closed:
|
|
return
|
|
for line in lines:
|
|
if self._closed:
|
|
return
|
|
payload = _sanitize_ws_text(line)
|
|
try:
|
|
await asyncio.wait_for(self._ws.send_text(payload), timeout=_WS_SEND_DEADLINE_S)
|
|
except asyncio.TimeoutError:
|
|
# The loop is responsive (the timer fired) but the socket never drained: unlike the
|
|
# loop-stall wait in write(), this is a dead peer. Latch under the writer lock so queued
|
|
# batches bail, and close the socket so handle_ws's read loop ends and its teardown
|
|
# (session detach/reap, client reconnect) runs. See #106369.
|
|
self._closed = True
|
|
_log.warning("ws send deadline exceeded (socket stalled, loop responsive) peer=%s deadline=%ss — closing",
|
|
self._peer, _WS_SEND_DEADLINE_S)
|
|
self._loop.create_task(self._close_stalled_socket())
|
|
return
|
|
except UnicodeEncodeError as exc:
|
|
# A single illegal UTF-8 frame (lone surrogate) must not tear down the socket.
|
|
_log.warning("ws send skipped invalid utf-8 frame peer=%s error=%s", self._peer, exc)
|
|
continue
|
|
except Exception as exc:
|
|
# Latch while holding the writer lock so queued batches observe the failure first.
|
|
self._closed = True
|
|
_log.warning("ws send failed peer=%s error_type=%s error=%s", self._peer, type(exc).__name__, exc)
|
|
return
|
|
|
|
def close(self) -> None: # loop thread (handle_ws finally), so the TimerHandle is safe
|
|
self._closed = True
|
|
if self._token_flush_handle is not None:
|
|
self._token_flush_handle.cancel()
|
|
self._token_flush_handle = None
|
|
|
|
async def _close_stalled_socket(self) -> None:
|
|
"""Close the peer socket after a send deadline so ``handle_ws``'s ``receive_text`` unblocks and its
|
|
disconnect teardown runs. The server library bounds this (websockets ``close_timeout`` → abort)."""
|
|
try:
|
|
await self._ws.close(code=1011)
|
|
except Exception as exc: # noqa: BLE001 - the peer is already gone; teardown is what matters
|
|
_log.debug("ws close after send deadline failed peer=%s error=%s", self._peer, exc)
|
|
|
|
|
|
def _ws_peer_label(ws: Any) -> str:
|
|
"""``host:port`` when available, else a stable placeholder."""
|
|
client = getattr(ws, "client", None)
|
|
if client is None:
|
|
return "unknown"
|
|
host, port = getattr(client, "host", None) or "unknown", getattr(client, "port", None)
|
|
return f"{host}:{port}" if port is not None else host
|
|
|
|
|
|
def _disable_nagle(ws: Any) -> None:
|
|
"""Disable Nagle + enable TCP keepalive on the raw socket (best-effort). Without TCP_NODELAY the kernel
|
|
coalesces small per-token frames, so a burst after the model's think-pause lands in one tick and no
|
|
client-side smoothing can recover the cadence. Without keepalive a silently-dropped client (SSH tunnel
|
|
reset, sleep) leaves the leg half-open forever: receive_text() blocks and the disconnect teardown never runs."""
|
|
try:
|
|
scope = getattr(ws, "scope", None) or {}
|
|
transport = (scope.get("extensions") or {}).get("transport") or getattr(ws, "transport", None)
|
|
sock = transport.get_extra_info("socket") if transport is not None else None
|
|
if sock is not None:
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
|
|
if hasattr(socket, "TCP_KEEPIDLE"): # Linux
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30)
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 10)
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 3)
|
|
elif hasattr(socket, "TCP_KEEPALIVE"): # macOS idle seconds
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, 30)
|
|
except Exception as exc: # pragma: no cover - best-effort tuning
|
|
_log.debug("ws TCP_NODELAY skip: %s", exc)
|
|
|
|
|
|
class _SendFailed(Exception):
|
|
"""Raised by handle_ws._reply when a reply could not be written: ends the read loop."""
|
|
|
|
|
|
async def handle_ws(ws: Any, *, auth_identity: dict | None = None, subprotocol: str | None = None) -> None:
|
|
"""Run one WebSocket session. Wire-compatible with ``tui_gateway.entry``. *auth_identity* is the server-minted
|
|
``{user_id, provider}`` recorded at WS-upgrade auth, stored as ``WSTransport.auth_identity`` (the only identity
|
|
authority for browser-controller registration); callers that omit it (harnesses, embedded TUI child) get None."""
|
|
peer, transport = _ws_peer_label(ws), None
|
|
messages = parse_errors = dispatch_crashes = send_failures = 0
|
|
disconnect_reason = "not_connected"
|
|
|
|
async def _reply(frame: dict, reason: str, msg: str, *args: Any) -> None:
|
|
"""write_async; on failure record *reason*, log *msg* and end the read loop."""
|
|
nonlocal disconnect_reason, send_failures
|
|
if not await transport.write_async(frame):
|
|
disconnect_reason = reason
|
|
send_failures += 1
|
|
_log.warning(msg, *args)
|
|
raise _SendFailed
|
|
|
|
def _error(code: int, message: str, req_id: Any) -> dict:
|
|
return {"jsonrpc": "2.0", "error": {"code": code, "message": message}, "id": req_id}
|
|
|
|
try:
|
|
await (ws.accept(subprotocol=subprotocol) if subprotocol else ws.accept())
|
|
disconnect_reason = "connected"
|
|
# Mark the client attached before the (possibly slow) ready/skin setup so scale-to-zero sees it.
|
|
_note_dashboard_client_activity(force=True)
|
|
_disable_nagle(ws)
|
|
_log.info("ws accepted peer=%s", peer)
|
|
transport = WSTransport(ws, asyncio.get_running_loop(), peer=peer, auth_identity=auth_identity)
|
|
# resolve_skin() is sync I/O + CPU; pooled so the read loop can drain the frontend's initial RPC burst.
|
|
skin_payload = await asyncio.to_thread(server.resolve_skin)
|
|
# change_events: this backend broadcasts pet/cron/sessions.changed, so clients can demote legacy
|
|
# polls to backstops. replay_epoch lets reconnecting clients detect a backend restart and reset
|
|
# their per-session seq watermarks (event_replay).
|
|
ready_ok = await transport.write_async({
|
|
"jsonrpc": "2.0", "method": "event",
|
|
"params": {"type": "gateway.ready", "payload": {
|
|
"skin": skin_payload, "change_events": True, "heartbeat": True, "replay_epoch": replay_epoch(),
|
|
}},
|
|
})
|
|
if ready_ok:
|
|
# Live-apply skins Hermes activates mid-conversation, and track this peer for session-less
|
|
# global broadcasts write_json can't route.
|
|
server._ensure_skin_watcher()
|
|
server.register_live_transport(transport)
|
|
# Cross-backend liveness: a heartbeat row lets the startup orphan sweep tell "live but idle
|
|
# backend" from "truly orphaned". Idempotent and once-per-process, like the orphan sweep (the
|
|
# desktop app and web dashboard reach the agent via this sidecar, not entry.main()).
|
|
for start, what in (
|
|
(server._start_backend_heartbeat_refresher, "backend heartbeat refresher start"),
|
|
(server._schedule_startup_orphan_sweep, "startup orphan sweep scheduling"),
|
|
):
|
|
try:
|
|
start()
|
|
except Exception:
|
|
_log.warning("%s failed", what, exc_info=True)
|
|
if not ready_ok:
|
|
disconnect_reason = "ready_send_failed"
|
|
send_failures += 1
|
|
_log.error("ws ready frame send failed peer=%s", peer)
|
|
return
|
|
|
|
while True:
|
|
try:
|
|
raw = await ws.receive_text()
|
|
_note_dashboard_client_activity()
|
|
except _WebSocketDisconnect as exc:
|
|
disconnect_reason = f"client_disconnect(code={getattr(exc, 'code', None)},reason={getattr(exc, 'reason', None)})"
|
|
break
|
|
except Exception:
|
|
disconnect_reason = "receive_failed"
|
|
_log.exception("ws receive failed peer=%s", peer)
|
|
break
|
|
line = raw.strip()
|
|
if not line:
|
|
continue
|
|
messages += 1
|
|
try:
|
|
req = json.loads(line)
|
|
except json.JSONDecodeError as exc:
|
|
parse_errors += 1
|
|
_log.warning("ws parse error peer=%s index=%d error=%s payload=%r", peer, messages, exc, line[:_WS_LOG_PAYLOAD_PREVIEW])
|
|
await _reply(_error(-32700, "parse error", None), "send_failed_after_parse_error",
|
|
"ws parse-error reply send failed peer=%s", peer)
|
|
continue
|
|
req_id = req.get("id") if isinstance(req, dict) else None
|
|
req_method = req.get("method") if isinstance(req, dict) else None
|
|
if req_method == "gateway.ping":
|
|
await _reply({"jsonrpc": "2.0", "result": {"ok": True}, "id": req_id}, "send_failed_after_heartbeat",
|
|
"ws heartbeat reply send failed peer=%s id=%s", peer, req_id)
|
|
continue
|
|
# dispatch() may schedule long handlers on the pool; it returns None then and the worker
|
|
# writes the response itself via transport.write (a separate thread, so that is the safe
|
|
# path). Inline handlers return the response dict, written here from the loop.
|
|
try:
|
|
resp = await asyncio.to_thread(server.dispatch, req, transport)
|
|
except Exception:
|
|
dispatch_crashes += 1
|
|
_log.exception("ws dispatch crash peer=%s id=%s method=%s", peer, req_id, req_method)
|
|
await _reply(_error(-32603, "internal error", req_id), "send_failed_after_dispatch_crash",
|
|
"ws dispatch-crash reply send failed peer=%s id=%s method=%s", peer, req_id, req_method)
|
|
continue
|
|
if resp is not None:
|
|
await _reply(resp, "send_failed_after_response",
|
|
"ws response send failed peer=%s id=%s method=%s", peer, req_id, req_method)
|
|
except _SendFailed:
|
|
pass
|
|
finally:
|
|
reaped_sessions = detached_sessions = 0
|
|
if transport is not None:
|
|
server.unregister_live_transport(transport)
|
|
# Owner-safely park browser controllers this transport registered (a same-identity reconnect may
|
|
# deliver a terminal result for in-flight work). Offloaded: disconnect takes the controller's
|
|
# send_lock, which a worker-thread dispatch may hold while blocking on THIS loop to transmit.
|
|
try:
|
|
from gateway.browser_control_broker import get_browser_control_broker
|
|
await asyncio.to_thread(get_browser_control_broker().disconnect_owner, transport)
|
|
except Exception:
|
|
_log.exception("ws browser-controller disconnect failed peer=%s", peer)
|
|
transport.close()
|
|
try:
|
|
await asyncio.to_thread(server._release_wake_for_transport, transport)
|
|
except Exception:
|
|
_log.exception("ws wake-word teardown failed peer=%s", peer)
|
|
# The single WS-disconnect teardown path: reap sessions this transport owned (close_on_disconnect
|
|
# sidecars) or detach the rest to the drop sentinel so later emits don't hit a closed socket; detached
|
|
# ones go to the grace-windowed orphan reaper (a quick resume cancels it). Offloaded: worker.close()
|
|
# blocks (terminate + waits) plus a sync DB write, which inline would freeze the loop for every peer.
|
|
try:
|
|
reaped_sessions, detached_sessions = await asyncio.to_thread(
|
|
server._close_sessions_for_transport, transport, end_reason="ws_disconnect"
|
|
)
|
|
except Exception:
|
|
_log.exception("ws transport teardown failed peer=%s", peer)
|
|
try:
|
|
await ws.close()
|
|
except Exception as exc:
|
|
_log.debug("ws close failed peer=%s error=%s", peer, exc)
|
|
_log.info(
|
|
"ws closed peer=%s reason=%s messages=%d parse_errors=%d "
|
|
"dispatch_crashes=%d send_failures=%d reaped_sessions=%d detached_sessions=%d",
|
|
peer, disconnect_reason, messages, parse_errors, dispatch_crashes, send_failures, reaped_sessions, detached_sessions,
|
|
)
|