feat(gateway): add gateway.ping heartbeat wire contract

Additive WebSocket wire contract for a client-driven heartbeat.

The gateway.ready payload now advertises "heartbeat": True so clients can
discover the capability, and the WS read loop answers a gateway.ping request
with a {"ok": True} pong short-circuited BEFORE method dispatch (no method is
invoked). The WSTransport gains closed / last_inbound_at properties and a
mark_inbound() hook, updated on every inbound frame, for later liveness checks.

Backward-compatible in both directions: old clients never send gateway.ping,
and old servers simply never advertise the heartbeat flag. This is the first
slice of a WebSocket-recovery series; the follow-on slices consume this
contract (server-side transport rebind, TUI/desktop clients).

Receipts:
  bash scripts/run_tests.sh tests/test_tui_gateway_ws.py -q
  => 1 file, 7 tests passed, 0 failed (100%) in 0.8s; exit 0
This commit is contained in:
Eva
2026-08-19 18:04:13 +07:00
committed by Teknium
parent 9525c0e5b7
commit 9a71cb95cb
2 changed files with 79 additions and 1 deletions

View File

@@ -153,6 +153,50 @@ def test_ws_starts_mcp_discovery_before_ready(monkeypatch):
assert events == ["accept", "ready_after_0"]
def test_ws_ready_advertises_heartbeat_and_ping_is_inline(monkeypatch):
sent = []
inbound = iter(
[
json.dumps(
{
"jsonrpc": "2.0",
"id": "heartbeat-1",
"method": "gateway.ping",
"params": {},
}
)
]
)
monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)
class FakeWS:
async def accept(self):
pass
async def send_text(self, line):
sent.append(json.loads(line))
async def receive_text(self):
try:
return next(inbound)
except StopIteration:
raise ws_mod._WebSocketDisconnect()
async def close(self):
pass
asyncio.run(ws_mod.handle_ws(FakeWS()))
ready = sent[0]["params"]
assert ready["type"] == "gateway.ready"
assert ready["payload"]["heartbeat"] is True
assert sent[1] == {
"jsonrpc": "2.0",
"result": {"ok": True},
"id": "heartbeat-1",
}
def test_ws_transport_serializes_concurrent_sends():
active_sends = 0
max_active_sends = 0

View File

@@ -29,6 +29,7 @@ import json
import logging
import socket
import threading
import time
from typing import Any
from tui_gateway import server
@@ -103,6 +104,7 @@ class WSTransport:
#: browser-controller registration.
self.auth_identity = auth_identity
self._closed = False
self._last_inbound_at = time.monotonic()
# 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
@@ -116,6 +118,17 @@ class WSTransport:
# the owning loop while it recovers from a stall.
self._send_lock = asyncio.Lock()
@property
def closed(self) -> bool:
return self._closed
@property
def last_inbound_at(self) -> float:
return self._last_inbound_at
def mark_inbound(self) -> None:
self._last_inbound_at = time.monotonic()
@staticmethod
def _is_streaming_frame(obj: dict) -> bool:
"""True for high-frequency per-token frames eligible for coalescing."""
@@ -361,7 +374,11 @@ async def handle_ws(
# change_events: this backend broadcasts pet.changed /
# cron.changed / sessions.changed, so clients can demote
# their legacy polls to slow backstops.
"payload": {"skin": skin_payload, "change_events": True},
"payload": {
"skin": skin_payload,
"change_events": True,
"heartbeat": True,
},
},
}
)
@@ -404,6 +421,7 @@ async def handle_ws(
line = raw.strip()
if not line:
continue
transport.mark_inbound()
messages += 1
try:
@@ -438,6 +456,22 @@ async def handle_ws(
# 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
if req_method == "gateway.ping":
ok = await transport.write_async(
{
"jsonrpc": "2.0",
"result": {"ok": True},
"id": req_id,
}
)
if not ok:
disconnect_reason = "send_failed_after_heartbeat"
send_failures += 1
_log.warning("ws heartbeat reply send failed peer=%s id=%s", peer, req_id)
break
continue
try:
resp = await asyncio.to_thread(server.dispatch, req, transport)
except Exception: