wehub-resource-sync
b4fbd6fe9f
Deploy Site / deploy-vercel (push) Has been skipped
Deploy Site / deploy-docs (push) Has been skipped
Build Skills Index / build-index (push) Has been skipped
CI / Deny unrelated histories (push) Has been skipped
CI / Detect affected areas (push) Successful in 27m35s
CI / OSV scan (push) Failing after 4s
CI / Build&Test Docker image (push) Successful in 9s
CI / Supply-chain scan (push) Has been skipped
CI / Lint Docker scripts (push) Failing after 5m13s
CI / Check contributors (push) Failing after 12m8s
CI / Docs Site (push) Failing after 12m8s
CI / TypeScript (push) Failing after 12m8s
CI / Python lints (push) Failing after 12m9s
CI / Python tests (push) Failing after 12m9s
CI / Check uv.lock (push) Failing after 23m22s
CI / CI timing report (push) Has been cancelled
Build Skills Index / trigger-deploy (push) Has been cancelled
CI / All required checks pass (push) Has been cancelled
467 lines
19 KiB
Python
467 lines
19 KiB
Python
"""WebSocket transport for the tui_gateway JSON-RPC server.
|
|
|
|
Reuses :func:`tui_gateway.server.dispatch` verbatim so every RPC method, every
|
|
slash command, every approval/clarify/sudo flow, and every agent event flows
|
|
through the same handlers whether the client is Ink over stdio or an iOS /
|
|
web client over WebSocket.
|
|
|
|
Wire protocol
|
|
-------------
|
|
Identical to stdio: newline-delimited JSON-RPC in both directions. The server
|
|
emits a ``gateway.ready`` event immediately after connection accept, then
|
|
echoes responses/events for inbound requests. No framing differences.
|
|
|
|
Mounting
|
|
--------
|
|
from fastapi import WebSocket
|
|
from tui_gateway.ws import handle_ws
|
|
|
|
@app.websocket("/api/ws")
|
|
async def ws(ws: WebSocket):
|
|
await handle_ws(ws)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import json
|
|
import logging
|
|
import socket
|
|
import threading
|
|
from typing import Any
|
|
|
|
from tui_gateway import server
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
# Max seconds a pool-dispatched handler will block waiting for the event loop
|
|
# to flush a WS frame before we mark the transport dead. Protects handler
|
|
# threads from a wedged socket.
|
|
_WS_WRITE_TIMEOUT_S = 10.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 event loop once per token. A model reply
|
|
# emits hundreds of these in a burst, and each one is a loop wakeup competing
|
|
# with the agent turn for the GIL — coalescing cuts that churn (CF-2). The task
|
|
# that introduced this called them "agent.token"/"agent.thinking"; in this
|
|
# codebase the per-token frames are the ``*.delta`` stream events below. Keep
|
|
# this set to genuinely high-frequency, display-only events — anything a client
|
|
# must see promptly (tool/approval/status/completion frames) is non-streaming
|
|
# and flushes the buffer ahead of itself, so ordering is preserved.
|
|
_STREAMING_EVENT_TYPES = frozenset({
|
|
"message.delta",
|
|
"reasoning.delta",
|
|
"thinking.delta",
|
|
})
|
|
# Max time a streamed token waits in the buffer before flush (~30 fps). Short
|
|
# enough to stay imperceptible to the live token cadence.
|
|
_TOKEN_COALESCE_S = 0.033
|
|
|
|
# Keep starlette optional at import time; handle_ws uses the real class when
|
|
# it's available and falls back to a generic Exception sentinel otherwise.
|
|
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 to call from any thread *other than* the event loop
|
|
thread that owns the socket. Pool workers (the only real caller) run in
|
|
their own threads, so marshalling onto the loop via
|
|
:func:`asyncio.run_coroutine_threadsafe` + ``future.result()`` is correct
|
|
and deadlock-free there.
|
|
|
|
When called from the loop thread itself (e.g. by ``handle_ws`` for an
|
|
inline response) the same call would deadlock: we'd schedule work onto
|
|
the loop we're currently blocking. We detect that case and fire-and-
|
|
forget instead. Callers that need to know when the bytes are on the wire
|
|
should use :meth:`write_async` from the loop thread.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
ws: Any,
|
|
loop: asyncio.AbstractEventLoop,
|
|
*,
|
|
peer: str = "unknown",
|
|
) -> None:
|
|
self._ws = ws
|
|
self._loop = loop
|
|
self._peer = peer
|
|
self._closed = False
|
|
# Token-coalescing buffer (CF-2). Streamed token frames land here and a
|
|
# short timer flushes the batch. The lock guards the buffer + the
|
|
# "armed" flag against the worker threads that call 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
|
|
|
|
@staticmethod
|
|
def _is_streaming_frame(obj: dict) -> bool:
|
|
"""True for high-frequency per-token frames eligible for coalescing."""
|
|
params = obj.get("params") if isinstance(obj, dict) else None
|
|
if not isinstance(params, dict):
|
|
return False
|
|
return params.get("type") in _STREAMING_EVENT_TYPES
|
|
|
|
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
|
|
|
|
# Coalesce streamed token frames: buffer this frame and arm a short
|
|
# flush timer instead of waking the loop right now. Cheap and
|
|
# non-blocking — the worker returns immediately. Ordering is preserved
|
|
# because every non-streaming frame (below) drains the buffer ahead of
|
|
# itself.
|
|
if self._is_streaming_frame(obj):
|
|
with self._token_lock:
|
|
self._pending_tokens.append(line)
|
|
if not self._token_flush_armed:
|
|
self._token_flush_armed = True
|
|
# call_soon_threadsafe arms the call_later timer on the loop
|
|
# thread and is safe to call from a worker or the loop.
|
|
self._loop.call_soon_threadsafe(self._arm_token_flush)
|
|
return not self._closed
|
|
|
|
# Non-streaming frame (RPC response, control frame, non-token event):
|
|
# append it behind any buffered tokens and flush the whole batch NOW so
|
|
# it can never overtake the tokens that preceded it. The send is
|
|
# scheduled INSIDE the lock so the on-the-wire order matches the 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:
|
|
# Fire-and-forget — don't block the loop waiting on itself.
|
|
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 event loop is stalled (GIL-heavy agent turn, delegation
|
|
# running N children), NOT the socket dead. The send coroutine is
|
|
# already scheduled and will flush once the loop breathes — latching
|
|
# _closed here permanently silenced live windows after one slow
|
|
# write (the "subagent window shows zero streaming" bug). Unblock
|
|
# the worker thread and keep the transport alive; _safe_send_many
|
|
# latches on a real socket error when the frame actually fails.
|
|
_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:
|
|
"""Arm the coalesce timer. Runs on the loop thread (call_soon_threadsafe)."""
|
|
if self._closed:
|
|
return
|
|
self._token_flush_handle = self._loop.call_later(
|
|
_TOKEN_COALESCE_S, self._flush_tokens
|
|
)
|
|
|
|
def _flush_tokens(self) -> None:
|
|
"""Send buffered tokens as one batch. Runs on the loop thread (timer).
|
|
|
|
The send is scheduled under the lock so its wire order is fixed relative
|
|
to a concurrent non-streaming flush in :meth:`write`.
|
|
"""
|
|
with self._token_lock:
|
|
self._token_flush_handle = None
|
|
self._token_flush_armed = False
|
|
if not self._pending_tokens or self._closed:
|
|
self._pending_tokens = []
|
|
return
|
|
batch = self._pending_tokens
|
|
self._pending_tokens = []
|
|
self._loop.create_task(self._safe_send_many(batch))
|
|
|
|
async def write_async(self, obj: dict) -> bool:
|
|
"""Send from the owning event loop. Awaits until the frame is on the wire."""
|
|
if self._closed:
|
|
return False
|
|
# Flush any buffered streamed tokens ahead of this frame (RPC response /
|
|
# control frame) so it can't overtake the tokens that preceded it.
|
|
with self._token_lock:
|
|
pending = self._pending_tokens
|
|
self._pending_tokens = []
|
|
if pending:
|
|
await self._safe_send_many(pending)
|
|
await self._safe_send(json.dumps(obj, ensure_ascii=False))
|
|
return not self._closed
|
|
|
|
async def _safe_send(self, line: str) -> None:
|
|
try:
|
|
await self._ws.send_text(line)
|
|
except Exception as exc:
|
|
self._closed = True
|
|
_log.warning(
|
|
"ws send failed peer=%s error_type=%s error=%s",
|
|
self._peer, type(exc).__name__, exc,
|
|
)
|
|
|
|
async def _safe_send_many(self, lines: list[str]) -> None:
|
|
"""Send a batch of pre-serialized frames in order on the loop thread."""
|
|
try:
|
|
for line in lines:
|
|
await self._ws.send_text(line)
|
|
except Exception as exc:
|
|
self._closed = True
|
|
_log.warning(
|
|
"ws send failed peer=%s error_type=%s error=%s",
|
|
self._peer, type(exc).__name__, exc,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self._closed = True
|
|
# Cancel any pending coalesce flush. close() runs on the loop thread
|
|
# (the handle_ws finally), so touching the TimerHandle here is safe.
|
|
handle = self._token_flush_handle
|
|
if handle is not None:
|
|
handle.cancel()
|
|
self._token_flush_handle = None
|
|
|
|
|
|
def _ws_peer_label(ws: Any) -> str:
|
|
"""Return ``host:port`` when available, else a stable placeholder."""
|
|
client = getattr(ws, "client", None)
|
|
if client is None:
|
|
return "unknown"
|
|
host = getattr(client, "host", None) or "unknown"
|
|
port = getattr(client, "port", None)
|
|
return f"{host}:{port}" if port is not None else host
|
|
|
|
|
|
def _disable_nagle(ws: Any) -> None:
|
|
"""Disable Nagle so streamed JSON-RPC frames go out individually.
|
|
|
|
Without it the kernel coalesces the small per-token frames, so a burst after
|
|
the model's think-pause lands on the client in one tick and no client-side
|
|
smoothing can recover the cadence. GUI/WS only; chat platforms don't hit
|
|
this path. Best-effort — skip silently if the socket isn't reachable.
|
|
"""
|
|
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)
|
|
except Exception as exc: # pragma: no cover - best-effort tuning
|
|
_log.debug("ws TCP_NODELAY skip: %s", exc)
|
|
|
|
|
|
async def handle_ws(ws: Any) -> None:
|
|
"""Run one WebSocket session. Wire-compatible with ``tui_gateway.entry``."""
|
|
peer = _ws_peer_label(ws)
|
|
transport: WSTransport | None = None
|
|
messages = 0
|
|
parse_errors = 0
|
|
dispatch_crashes = 0
|
|
send_failures = 0
|
|
disconnect_reason = "not_connected"
|
|
|
|
try:
|
|
await ws.accept()
|
|
disconnect_reason = "connected"
|
|
# Push small streamed frames out immediately instead of letting Nagle
|
|
# batch them — keeps the live token cadence intact for GUI clients.
|
|
_disable_nagle(ws)
|
|
_log.info("ws accepted peer=%s", peer)
|
|
|
|
transport = WSTransport(ws, asyncio.get_running_loop(), peer=peer)
|
|
|
|
# The desktop app and dashboard chat reach the agent through this WS
|
|
# sidecar, NOT through tui_gateway.entry.main() (the stdio TUI path that
|
|
# spawns the background MCP discovery thread). Without starting it here,
|
|
# discovery never runs in this process: _make_agent only *waits* on the
|
|
# thread (wait_for_mcp_discovery), which no-ops when it was never
|
|
# created, so the agent snapshots an MCP-less tool list and the only way
|
|
# to surface MCP tools is a manual /reload-mcp. Start it once per
|
|
# process here (idempotent, config-gated) before gateway.ready so the
|
|
# first agent build can pick up already-spawning servers. (#38945)
|
|
from hermes_cli.mcp_startup import start_background_mcp_discovery
|
|
|
|
start_background_mcp_discovery(
|
|
logger=_log,
|
|
thread_name="tui-ws-mcp-discovery",
|
|
)
|
|
|
|
ready_ok = await transport.write_async(
|
|
{
|
|
"jsonrpc": "2.0",
|
|
"method": "event",
|
|
"params": {
|
|
"type": "gateway.ready",
|
|
"payload": {"skin": server.resolve_skin()},
|
|
},
|
|
}
|
|
)
|
|
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()
|
|
except _WebSocketDisconnect as exc:
|
|
disconnect_reason = (
|
|
"client_disconnect("
|
|
f"code={getattr(exc, 'code', None)},"
|
|
f"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],
|
|
)
|
|
ok = await transport.write_async(
|
|
{
|
|
"jsonrpc": "2.0",
|
|
"error": {"code": -32700, "message": "parse error"},
|
|
"id": None,
|
|
}
|
|
)
|
|
if not ok:
|
|
disconnect_reason = "send_failed_after_parse_error"
|
|
send_failures += 1
|
|
_log.warning("ws parse-error reply send failed peer=%s", peer)
|
|
break
|
|
continue
|
|
|
|
# dispatch() may schedule long handlers on the pool; it returns
|
|
# None in that case and the worker writes the response itself via
|
|
# the transport we pass in (a separate thread, so transport.write
|
|
# is the safe path there). For inline handlers it returns the
|
|
# response dict, which we write here from the loop.
|
|
req_id = req.get("id") if isinstance(req, dict) else None
|
|
req_method = req.get("method") if isinstance(req, dict) else None
|
|
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,
|
|
)
|
|
ok = await transport.write_async(
|
|
{
|
|
"jsonrpc": "2.0",
|
|
"error": {"code": -32603, "message": "internal error"},
|
|
"id": req_id if req_id is not None else None,
|
|
}
|
|
)
|
|
if not ok:
|
|
disconnect_reason = "send_failed_after_dispatch_crash"
|
|
send_failures += 1
|
|
_log.warning(
|
|
"ws dispatch-crash reply send failed peer=%s id=%s method=%s",
|
|
peer,
|
|
req_id,
|
|
req_method,
|
|
)
|
|
break
|
|
continue
|
|
if resp is not None and not await transport.write_async(resp):
|
|
disconnect_reason = "send_failed_after_response"
|
|
send_failures += 1
|
|
_log.warning(
|
|
"ws response send failed peer=%s id=%s method=%s",
|
|
peer,
|
|
req_id,
|
|
req_method,
|
|
)
|
|
break
|
|
finally:
|
|
reaped_sessions = 0
|
|
detached_sessions = 0
|
|
if transport is not None:
|
|
transport.close()
|
|
|
|
# Reap sessions this transport owned (close_on_disconnect sidecar
|
|
# sessions) or detach the rest to the drop sentinel so later emits
|
|
# don't crash into a closed socket or fall through to desktop stdout
|
|
# logs. Detached sessions are handed to the grace-windowed WS-orphan
|
|
# reaper inside _close_sessions_for_transport (a quick reconnect /
|
|
# session.resume cancels it). This is the single WS-disconnect
|
|
# teardown path.
|
|
#
|
|
# Offloaded: _close_session_by_id does a blocking worker.close()
|
|
# (terminate + waits) plus a synchronous DB write — inline that
|
|
# would freeze the uvicorn event loop for every other live
|
|
# connection.
|
|
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,
|
|
)
|