From beb794123618c997e82791316df643fc61347665 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Mon, 24 Aug 2026 19:53:52 -0700 Subject: [PATCH] fix(tui-gateway): make WS reconnect replay actually deliver events (follow-up to #94219) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The #94219 replay was a production no-op: the server returned full JSON-RPC envelopes from session.events.since while the client's replay loop dispatches only elements with a top-level 'type' — every replayed event was silently skipped. Each side's tests validated its own assumption, so both suites stayed green. - server: events_since() now returns bare event objects (the frame's params), the exact shape the live dispatch path consumes; ring stores params directly; cross-language contract test added on both sides. - client: live frames racing an in-flight replay are parked and flushed seq-gated afterward — no double dispatch of deltas, no gap-skip from a watermark advanced past the replay window. - restart poisoning: seq counters are in-process, so a backend restart reset them while clients kept high watermarks (replay forever empty, truncated=false). New replay_epoch advertised in gateway.ready and echoed by session.events.since; the client clears watermarks on epoch change. - methods_session no longer reaches into event_replay privates (is_truncated() accessor). Live repro: pre-fix, 3 stamped frames -> 0 dispatchable by the client gate; post-fix 3/3. Tests: 16 py (replay+ws), 8 vitest, tsc clean, ruff clean. --- .../src/json-rpc-gateway-replay.test.ts | 139 ++++++++++++++++++ apps/shared/src/json-rpc-gateway.ts | 124 +++++++++++++++- tests/test_tui_gateway_event_replay.py | 29 +++- tui_gateway/entry.py | 9 +- tui_gateway/event_replay.py | 42 +++++- tui_gateway/methods_session.py | 16 +- tui_gateway/ws.py | 5 + 7 files changed, 339 insertions(+), 25 deletions(-) diff --git a/apps/shared/src/json-rpc-gateway-replay.test.ts b/apps/shared/src/json-rpc-gateway-replay.test.ts index dc5dd559bf..a587573830 100644 --- a/apps/shared/src/json-rpc-gateway-replay.test.ts +++ b/apps/shared/src/json-rpc-gateway-replay.test.ts @@ -194,4 +194,143 @@ describe('JsonRpcGatewayClient event-seq tracking + replay resume', () => { expect(client.getSeqWatermarks().s1).toBe(10) client.close() }) + + it('rejects envelope-shaped replay elements (the #94219 server-shape bug)', async () => { + const client = makeClient() + const seen: string[] = [] + client.on('message.delta', () => seen.push('delta')) + + const first = client.connect('ws://x') + let sock = sockets[sockets.length - 1] + sock.open() + await first + sock.serverFrame({ jsonrpc: '2.0', method: 'event', params: { type: 'message.delta', session_id: 's1', seq: 1 } }) + expect(seen).toEqual(['delta']) // the pre-drop live frame + + client.invalidate('drop') + const second = client.connect('ws://x') + sock = sockets[sockets.length - 1] + sock.open() + await second + + await vi.waitFor(() => { + expect(sock.lastRequest().method).toBe('session.events.since') + }) + const req = sock.lastRequest() + // Pre-fix servers returned FULL JSON-RPC envelopes. The client must not + // dispatch those blindly — and this documents why the server now sends + // bare event objects. + sock.serverFrame({ + jsonrpc: '2.0', + id: req.id, + result: { + events: [{ jsonrpc: '2.0', method: 'event', params: { type: 'message.delta', session_id: 's1', seq: 2 } }], + latest_seq: 2, + truncated: false, + count: 1 + } + }) + await Promise.resolve() + // Envelope-shaped replay elements must add nothing beyond the live frame. + expect(seen).toEqual(['delta']) + client.close() + }) + + it('holds live frames racing the replay fetch — no double dispatch, no skipped gap', async () => { + const client = makeClient() + const seen: number[] = [] + client.on('message.delta', e => seen.push((e as unknown as { seq: number }).seq)) + + const first = client.connect('ws://x') + let sock = sockets[sockets.length - 1] + sock.open() + await first + sock.serverFrame({ jsonrpc: '2.0', method: 'event', params: { type: 'message.delta', session_id: 's1', seq: 2 } }) + expect(seen).toEqual([2]) // pre-drop live frame dispatches normally + + client.invalidate('drop') + const second = client.connect('ws://x') + sock = sockets[sockets.length - 1] + sock.open() + await second + + await vi.waitFor(() => { + expect(sock.lastRequest().method).toBe('session.events.since') + }) + + // LIVE frames 5 and 6 arrive while the replay (which carries 3,4,5) is + // still in flight. They must be parked, not dispatched ahead of the gap. + sock.serverFrame({ jsonrpc: '2.0', method: 'event', params: { type: 'message.delta', session_id: 's1', seq: 5 } }) + sock.serverFrame({ jsonrpc: '2.0', method: 'event', params: { type: 'message.delta', session_id: 's1', seq: 6 } }) + expect(seen).toEqual([2]) // still only the pre-drop frame — 5/6 parked + + const req = sock.lastRequest() + sock.serverFrame({ + jsonrpc: '2.0', + id: req.id, + result: { + events: [ + { type: 'message.delta', session_id: 's1', seq: 3 }, + { type: 'message.delta', session_id: 's1', seq: 4 }, + { type: 'message.delta', session_id: 's1', seq: 5 } + ], + latest_seq: 5, + truncated: false, + count: 3 + } + }) + + // In-order, exactly once: replayed 3,4,5 then the parked live 6 — + // the parked duplicate of 5 is seq-gated out. + await vi.waitFor(() => { + expect(seen).toEqual([2, 3, 4, 5, 6]) + }) + expect(client.getSeqWatermarks().s1).toBe(6) + client.close() + }) + + it('clears stale watermarks when the backend epoch changes (restart poisoning)', async () => { + const client = makeClient() + + const first = client.connect('ws://x') + let sock = sockets[sockets.length - 1] + sock.open() + await first + // Learn epoch A and a high watermark. + sock.serverFrame({ + jsonrpc: '2.0', + method: 'event', + params: { type: 'gateway.ready', payload: { replay_epoch: 'epoch-A' } } + }) + sock.serverFrame({ jsonrpc: '2.0', method: 'event', params: { type: 'message.delta', session_id: 's1', seq: 97 } }) + expect(client.getSeqWatermarks()).toEqual({ s1: 97 }) + + // Backend restarts: reconnect, replay under a NEW epoch returns nothing + // (fresh process, empty ring) — pre-fix the client kept watermark 97 and + // silently believed it missed nothing, forever. + client.invalidate('drop') + const second = client.connect('ws://x') + sock = sockets[sockets.length - 1] + sock.open() + await second + + await vi.waitFor(() => { + expect(sock.lastRequest().method).toBe('session.events.since') + }) + const req = sock.lastRequest() + sock.serverFrame({ + jsonrpc: '2.0', + id: req.id, + result: { events: [], latest_seq: 0, truncated: false, count: 0, epoch: 'epoch-B' } + }) + + await vi.waitFor(() => { + expect(client.getSeqWatermarks()).toEqual({}) + }) + + // New-epoch events build fresh watermarks from scratch. + sock.serverFrame({ jsonrpc: '2.0', method: 'event', params: { type: 'message.delta', session_id: 's1', seq: 3 } }) + expect(client.getSeqWatermarks()).toEqual({ s1: 3 }) + client.close() + }) }) diff --git a/apps/shared/src/json-rpc-gateway.ts b/apps/shared/src/json-rpc-gateway.ts index 72d9d70138..93d8b2de31 100644 --- a/apps/shared/src/json-rpc-gateway.ts +++ b/apps/shared/src/json-rpc-gateway.ts @@ -111,6 +111,22 @@ export class JsonRpcGatewayClient { private lastSeenSeq = new Map() /** Set while a post-reconnect replay fetch is in flight (dedup guard). */ private replayInFlight = false + /** + * While a replay fetch is in flight, live seq'd frames for the sessions + * being replayed are parked here instead of dispatching immediately. + * Without this hold, a live frame racing the replay response is dispatched + * twice (once live, once when the replay returns the same seq) or, worse, + * advances the watermark so the gap events the replay carries get skipped. + */ + private replayHold: Map | null = null + /** + * Server process identity for the replay contract (from gateway.ready / + * session.events.since). Seq counters are in-process on the backend, so a + * restart resets them while we still hold high watermarks — without this + * check events_since(sid, 97) returns [] + truncated=false forever and we + * silently believe nothing was missed. + */ + private replayEpoch: string | null = null private readonly eventHandlers = new Map void>>() private readonly stateHandlers = new Set<(state: ConnectionState) => void>() private readonly options: Required> & @@ -446,12 +462,30 @@ export class JsonRpcGatewayClient { } if (frame.method === 'event' && frame.params?.type) { - if (frame.params.type === 'gateway.ready' && this.gatewayReadyAdvertisesHeartbeat(frame.params.payload)) { - const socket = this.socket + if (frame.params.type === 'gateway.ready') { + if (this.gatewayReadyAdvertisesHeartbeat(frame.params.payload)) { + const socket = this.socket - if (socket) { - this.startHeartbeat(socket) + if (socket) { + this.startHeartbeat(socket) + } } + + const epoch = (frame.params.payload as { replay_epoch?: unknown } | undefined)?.replay_epoch + + if (typeof epoch === 'string' && epoch) { + this.adoptReplayEpoch(epoch) + } + } + + const sid = frame.params.session_id + const seqValue = (frame.params as { seq?: unknown }).seq + + if (this.replayHold && sid && typeof seqValue === 'number' && this.replayHold.has(sid)) { + // Replay in flight for this session: park the frame; flushReplayHold + // dispatches it after the replayed gap, gated on seq. + this.replayHold.get(sid)?.push(frame.params) + return } this.recordSeq(frame.params) @@ -496,6 +530,16 @@ export class JsonRpcGatewayClient { } this.replayInFlight = true + // Park live frames for the sessions we're about to replay so a frame + // racing the replay response can't dispatch ahead of (or duplicate) the + // gap events. Sessions without watermarks are unaffected. + const hold = new Map() + + for (const sid of this.lastSeenSeq.keys()) { + hold.set(sid, []) + } + + this.replayHold = hold try { const entries = Object.entries(this.getSeqWatermarks()) @@ -516,22 +560,90 @@ export class JsonRpcGatewayClient { continue } + const epoch = (result.value as { epoch?: unknown }).epoch + + if (typeof epoch === 'string' && epoch && this.replayEpoch && epoch !== this.replayEpoch) { + // Backend restarted: its seq numbering reset, so our watermarks — + // and this replay window — are meaningless. Drop them and start + // fresh under the new epoch. + this.adoptReplayEpoch(epoch) + continue + } + + if (typeof epoch === 'string' && epoch && !this.replayEpoch) { + this.replayEpoch = epoch + } + for (const event of result.value.events) { if (!event?.type) { continue } - this.recordSeq(event as GatewayEvent) - this.dispatchEvent(event as GatewayEvent) + this.dispatchIfNewer(event as GatewayEvent) } } } catch { // Replay is an optimization over lossy-reconnect; never surface errors. } finally { + this.flushReplayHold() this.replayInFlight = false } } + /** + * Dispatch an event only when its seq advances the session watermark. + * Seq-less events always dispatch (no ordering contract to violate). + */ + private dispatchIfNewer(event: GatewayEvent): void { + const sid = event.session_id + const seq = (event as { seq?: unknown }).seq + + if (sid && typeof seq === 'number' && Number.isFinite(seq)) { + const prev = this.lastSeenSeq.get(sid) ?? 0 + + if (seq <= prev) { + return + } + + this.lastSeenSeq.set(sid, seq) + } + + this.dispatchEvent(event) + } + + /** + * Record the server's replay epoch; on change (backend restart) the old + * seq watermarks describe a numbering that no longer exists — clear them + * so the next reconnect doesn't silently believe it missed nothing. + */ + private adoptReplayEpoch(epoch: string): void { + if (this.replayEpoch === epoch) { + return + } + + if (this.replayEpoch !== null) { + this.lastSeenSeq.clear() + } + + this.replayEpoch = epoch + } + + /** Release frames parked during a replay fetch, seq-gated against dupes. */ + private flushReplayHold(): void { + const hold = this.replayHold + this.replayHold = null + + if (!hold) { + return + } + + for (const parked of hold.values()) { + for (const event of parked) { + this.dispatchIfNewer(event) + } + } + } + private gatewayReadyAdvertisesHeartbeat(payload: unknown): boolean { return Boolean(payload && typeof payload === 'object' && (payload as { heartbeat?: unknown }).heartbeat === true) } diff --git a/tests/test_tui_gateway_event_replay.py b/tests/test_tui_gateway_event_replay.py index 0fa5e0f061..5e431c29dd 100644 --- a/tests/test_tui_gateway_event_replay.py +++ b/tests/test_tui_gateway_event_replay.py @@ -60,12 +60,30 @@ def test_events_since_returns_only_newer_frames_in_order(): event_replay._stamp_event(f) got = events_since("s1", 3) - assert [f["params"]["seq"] for f in got] == [4, 5] - assert events_since("s1", 0) == frames + assert [e["seq"] for e in got] == [4, 5] + assert events_since("s1", 0) == [f["params"] for f in frames] assert events_since("s1", 99) == [] assert latest_seq("s1") == 5 +def test_events_since_returns_client_dispatchable_event_objects(): + """Cross-language contract: the client's replay loop dispatches an element + only when it has a TOP-LEVEL ``type`` (json-rpc-gateway.ts fetchReplay: + ``if (!event?.type) continue``). Returning full JSON-RPC envelopes here + makes every replayed event silently droppable — the original #94219 bug. + """ + event_replay._stamp_event(_frame("s1")) + (event,) = events_since("s1", 0) + + # Bare event object, not an envelope. + assert event["type"] == "message.delta" + assert event["session_id"] == "s1" + assert event["seq"] == 1 + assert "jsonrpc" not in event + assert "method" not in event + assert "params" not in event + + def test_unknown_session_returns_empty(): assert events_since("nope", 0) == [] assert latest_seq("nope") == 0 @@ -129,7 +147,8 @@ def test_truncation_detection_semantics(): assert oldest > 1 # eviction happened # Client saw everything up to just before the buffer → NOT truncated. - last_seen_full = oldest - 1 - assert not (last_seen_full + 1 < oldest) + assert not event_replay.is_truncated("s1", oldest - 1) # Client saw seq 5, buffer starts later → truncated. - assert (5 + 1 < oldest) + assert event_replay.is_truncated("s1", 5) + # Unknown session: nothing evicted, nothing truncated. + assert not event_replay.is_truncated("nope", 0) diff --git a/tui_gateway/entry.py b/tui_gateway/entry.py index a7e7b3d48c..ea02073d2b 100644 --- a/tui_gateway/entry.py +++ b/tui_gateway/entry.py @@ -20,6 +20,7 @@ import traceback from tui_gateway._stdin_recovery import handle_spurious_eof from tui_gateway import server +from tui_gateway.event_replay import replay_epoch from tui_gateway.server import _CRASH_LOG, dispatch, resolve_skin, write_json from tui_gateway.transport import TeeTransport @@ -447,7 +448,13 @@ def main(): "params": { "type": "gateway.ready", # change_events: see tui_gateway/ws.py — clients demote legacy polls. - "payload": {"skin": resolve_skin(), "change_events": True}, + # replay_epoch: restart detection for the WS replay contract (the + # stdio TUI ignores it). + "payload": { + "skin": resolve_skin(), + "change_events": True, + "replay_epoch": replay_epoch(), + }, }, }): _log_exit("startup write failed (broken stdout pipe before first event)") diff --git a/tui_gateway/event_replay.py b/tui_gateway/event_replay.py index 88949de81a..0ec4e2a86b 100644 --- a/tui_gateway/event_replay.py +++ b/tui_gateway/event_replay.py @@ -19,8 +19,17 @@ Design constraints honored: from __future__ import annotations import threading +import uuid from collections import OrderedDict, deque +# Process identity for the replay contract. Seq counters live in-process, so +# a gateway restart silently resets them to 1 while clients still hold high +# watermarks — events_since(sid, 97) then returns [] with truncated=False and +# the client believes it missed nothing (and its stale watermark makes every +# future replay empty too). The epoch lets clients detect the restart and +# reset their watermarks. +_REPLAY_EPOCH = uuid.uuid4().hex + # Replay ring per session. A long turn emits ~hundreds of token events; this # covers several minutes of streaming plus all control events. _REPLAY_BUFFER_MAX = 512 @@ -28,11 +37,18 @@ _REPLAY_BUFFER_MAX = 512 _REPLAY_SESSIONS_MAX = 64 _replay_lock = threading.Lock() -# sid -> OrderedDict-ish deque of (seq, frame_params_dict_without_seq) +# sid -> deque of (seq, event_object) where event_object is the frame's +# ``params`` dict (bare event: type/session_id/seq/payload) — the exact shape +# the client's dispatch path consumes. _replay_buffers: "OrderedDict[str, deque]" = OrderedDict() _replay_next_seq: dict[str, int] = {} +def replay_epoch() -> str: + """Opaque token identifying this server process's seq numbering.""" + return _REPLAY_EPOCH + + def _stamp_event(obj: dict) -> None: """Stamp one outgoing event frame (mutates obj in place) and record it.""" if obj.get("method") != "event": @@ -56,16 +72,34 @@ def _stamp_event(obj: dict) -> None: while len(_replay_buffers) > _REPLAY_SESSIONS_MAX: _oldest_sid, _oldest_buf = _replay_buffers.popitem(last=False) _replay_next_seq.pop(_oldest_sid, None) - buf.append((seq, obj)) + buf.append((seq, params)) def events_since(sid: str, last_seen: int) -> list[dict]: - """Return recorded event FRAMES with seq > last_seen for *sid*, in order.""" + """Return recorded EVENT OBJECTS with seq > last_seen for *sid*, in order. + + Shape contract: each element is the frame's ``params`` dict — a bare event + object with top-level ``type`` / ``session_id`` / ``seq`` — because that is + exactly what the client's dispatch path consumes. Returning the full + JSON-RPC envelope here would make every replayed event fail the client's + ``event.type`` gate and be silently dropped. + """ with _replay_lock: buf = _replay_buffers.get(sid or "") if not buf: return [] - return [frame for seq, frame in buf if seq > last_seen] + return [event for seq, event in buf if seq > last_seen] + + +def is_truncated(sid: str, last_seen: int) -> bool: + """True when events between *last_seen* and the ring's oldest retained + seq were evicted — the client must refetch history instead of trusting + the replay to be gap-free.""" + with _replay_lock: + buf = _replay_buffers.get(sid or "") + if not buf: + return False + return last_seen + 1 < buf[0][0] def latest_seq(sid: str) -> int: diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index e277c40c25..b82c4df717 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -3655,18 +3655,16 @@ def _(rid, params: dict) -> dict: from tui_gateway import event_replay frames = event_replay.events_since(sid, last_seen) - # Truncated when the buffer's OLDEST retained seq is past last_seen+1 — - # i.e. events between last_seen and the buffer start were evicted. - truncated = False - with event_replay._replay_lock: - buf = event_replay._replay_buffers.get(sid) - if buf and last_seen + 1 < buf[0][0]: - truncated = True return _ok(rid, { - "events": [f for f in frames], + "events": frames, "latest_seq": event_replay.latest_seq(sid), - "truncated": truncated, + "truncated": event_replay.is_truncated(sid, last_seen), "count": len(frames), + # Restart detection: seq counters are in-process, so after a gateway + # restart a client's old high watermark would silently match nothing. + # Clients compare this against the epoch they learned at gateway.ready + # and reset watermarks on mismatch. + "epoch": event_replay.replay_epoch(), }) diff --git a/tui_gateway/ws.py b/tui_gateway/ws.py index be324ae55f..ce0cbc7021 100644 --- a/tui_gateway/ws.py +++ b/tui_gateway/ws.py @@ -33,6 +33,7 @@ import time from typing import Any from tui_gateway import server +from tui_gateway.event_replay import replay_epoch _log = logging.getLogger(__name__) @@ -378,6 +379,10 @@ async def handle_ws( "skin": skin_payload, "change_events": True, "heartbeat": True, + # Replay-contract process identity: lets reconnecting + # clients detect a backend restart and reset their + # per-session seq watermarks (see event_replay). + "replay_epoch": replay_epoch(), }, }, }