fix(#4812): close session SSE replay recovery gaps

This commit is contained in:
Rod Boev 2026-07-10 10:46:47 -04:00
parent da885659bd
commit 956c276362
3 changed files with 169 additions and 26 deletions

View file

@ -16203,13 +16203,20 @@ def _handle_session_run_journal_stream_for_session(handler, parsed, session_id):
if event_id:
note_sent_event_id(event_id)
def emit_session_snapshot(active_stream_id):
try:
fresh_session = get_session(session_id, metadata_only=True)
except KeyError:
fresh_session = session
_sse(handler, "session_snapshot", _session_snapshot_payload(fresh_session, active_stream_id=active_stream_id))
try:
replay_events = []
replay_ok = False
if resume_event_id:
replay = read_session_run_events(session_id, after_event_id=resume_event_id)
if replay.get("status") != "ok":
_sse(handler, "session_snapshot", _session_snapshot_payload(session, active_stream_id=active_stream_id))
emit_session_snapshot(active_stream_id)
else:
replay_ok = True
replay_events = replay.get("events") or []
@ -16231,6 +16238,8 @@ def _handle_session_run_journal_stream_for_session(handler, parsed, session_id):
reconciled = read_session_run_events(session_id, after_event_id=resume_event_id)
if reconciled.get("status") == "ok":
emit_replay(reconciled.get("events") or [], active_stream_id, replay_cutoff_seq)
else:
emit_session_snapshot(active_stream_id)
try:
while True:
try:

View file

@ -23,6 +23,7 @@ _FSYNC_MODE_EAGER = "eager"
_FSYNC_MODE_TERMINAL_ONLY = "terminal-only"
_SESSION_REPLAY_MAX_BYTES = 4 * 1024 * 1024
_SESSION_REPLAY_MAX_ROWS = 4096
_SESSION_REPLAY_READ_CHUNK_BYTES = 64 * 1024
def _default_session_dir() -> Path:
@ -152,6 +153,42 @@ def _event_created_at(event: dict, *, fallback: float = 0.0) -> float:
return fallback
def _iter_bounded_raw_jsonl_lines(path: Path, *, max_bytes: int, retained_bytes: int = 0):
line_no = 0
buffered = bytearray()
total_bytes = int(retained_bytes)
try:
with path.open("rb") as fh:
while True:
chunk = fh.read(_SESSION_REPLAY_READ_CHUNK_BYTES)
if not chunk:
if buffered:
if total_bytes + len(buffered) > max_bytes:
raise ValueError("replay_limit_bytes")
line_no += 1
total_bytes += len(buffered)
yield line_no, bytes(buffered), total_bytes
return
start = 0
while start < len(chunk):
newline = chunk.find(b"\n", start)
if newline == -1:
buffered.extend(chunk[start:])
if total_bytes + len(buffered) > max_bytes:
raise ValueError("replay_limit_bytes")
break
buffered.extend(chunk[start : newline + 1])
if total_bytes + len(buffered) > max_bytes:
raise ValueError("replay_limit_bytes")
line_no += 1
total_bytes += len(buffered)
yield line_no, bytes(buffered), total_bytes
buffered.clear()
start = newline + 1
except FileNotFoundError:
return
def append_run_event(
session_id: str,
run_id: str,
@ -339,32 +376,36 @@ def read_session_run_events(
events: list[dict] = []
expected_seq = 1
try:
with path.open(encoding="utf-8") as fh:
for raw in fh:
retained_bytes += len(raw.encode("utf-8"))
if retained_bytes > max_bytes:
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_limit_bytes", "events": []}
if not raw.strip():
continue
try:
event = json.loads(raw)
seq = int(event.get("seq")) if isinstance(event, dict) else 0
except (json.JSONDecodeError, TypeError, ValueError):
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_malformed", "events": []}
if (
seq != expected_seq
or event.get("event_id") != f"{run_id}:{seq}"
or event.get("run_id") != run_id
or event.get("session_id") != sid
):
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_noncontiguous", "events": []}
expected_seq += 1
retained_rows += 1
if retained_rows > max_rows:
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_limit_rows", "events": []}
events.append(event)
for _line_no, raw, retained_bytes in _iter_bounded_raw_jsonl_lines(
path,
max_bytes=max_bytes,
retained_bytes=retained_bytes,
):
if not raw.strip():
continue
try:
event = json.loads(raw.decode("utf-8"))
seq = int(event.get("seq")) if isinstance(event, dict) else 0
except (UnicodeDecodeError, json.JSONDecodeError, TypeError, ValueError):
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_malformed", "events": []}
if (
seq != expected_seq
or event.get("event_id") != f"{run_id}:{seq}"
or event.get("run_id") != run_id
or event.get("session_id") != sid
):
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_noncontiguous", "events": []}
expected_seq += 1
retained_rows += 1
if retained_rows > max_rows:
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_limit_rows", "events": []}
events.append(event)
except FileNotFoundError:
continue
except ValueError as exc:
if str(exc) == "replay_limit_bytes":
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "replay_limit_bytes", "events": []}
raise
created_at = min((_event_created_at(event) for event in events), default=path.stat().st_mtime)
runs.append((created_at, run_id, events))
runs.sort(key=lambda run: (run[0], run[1]))
@ -383,7 +424,7 @@ def read_session_run_events(
"events": [],
}
cursor_events = runs[cursor_index][2]
if cursor_seq is None or cursor_seq >= len(cursor_events):
if cursor_seq is None or cursor_seq > len(cursor_events):
return {"session_id": sid, "cursor_run_id": cursor_run_id, "cursor_seq": cursor_seq, "status": "cursor_event_missing", "events": []}
replay_events = [event for event in cursor_events if event["seq"] > cursor_seq]
for _created_at, _run_id, events in runs[cursor_index + 1:]:

View file

@ -214,6 +214,35 @@ def test_session_journal_rejects_bounds_and_invalid_contiguous_rows(tmp_path):
assert invalid["events"] == []
def test_session_journal_rejects_oversized_line_before_decode(tmp_path, monkeypatch):
path = tmp_path / "_run_journal" / "session_1" / "run_a.jsonl"
path.parent.mkdir(parents=True, exist_ok=True)
payload = '{"event_id":"run_a:1","seq":1,"run_id":"run_a","session_id":"session_1","event":"done","type":"done","payload":{"text":"' + ("x" * 5000) + '"}}\n'
path.write_bytes(payload.encode("utf-8"))
monkeypatch.setattr("api.run_journal.json.loads", lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("oversized line should not decode")))
limited = read_session_run_events("session_1", after_event_id="run_a:1", session_dir=tmp_path, max_bytes=1024)
assert limited["status"] == "replay_limit_bytes"
assert limited["events"] == []
def test_session_journal_accepts_line_at_cap_and_rejects_next_byte(tmp_path):
path = tmp_path / "_run_journal" / "session_1" / "run_a.jsonl"
path.parent.mkdir(parents=True, exist_ok=True)
payload = '{"event_id":"run_a:1","seq":1,"run_id":"run_a","session_id":"session_1","event":"done","type":"done","payload":{"text":"ok"}}\n'
path.write_bytes(payload.encode("utf-8"))
cap = len(payload.encode("utf-8"))
at_cap = read_session_run_events("session_1", after_event_id="run_a:1", session_dir=tmp_path, max_bytes=cap)
over = read_session_run_events("session_1", after_event_id="run_a:1", session_dir=tmp_path, max_bytes=cap - 1)
assert at_cap["status"] == "ok"
assert at_cap["events"] == []
assert over["status"] == "replay_limit_bytes"
assert over["events"] == []
def test_session_journal_requires_exact_cursor_and_contiguous_suffix(tmp_path):
append_run_event("session_1", "run_a", "token", {}, session_dir=tmp_path, seq=1)
exact_missing = read_session_run_events("session_1", after_event_id="run_a:2", session_dir=tmp_path)
@ -226,6 +255,16 @@ def test_session_journal_requires_exact_cursor_and_contiguous_suffix(tmp_path):
assert gap["events"] == []
def test_session_journal_accepts_final_cursor_and_replays_later_runs(tmp_path):
append_run_event("session_1", "run_a", "token", {}, session_dir=tmp_path, seq=1, created_at=100.0)
append_run_event("session_1", "run_b", "token", {}, session_dir=tmp_path, seq=1, created_at=200.0)
replay = read_session_run_events("session_1", after_event_id="run_a:1", session_dir=tmp_path)
assert replay["status"] == "ok"
assert [event["event_id"] for event in replay["events"]] == ["run_b:1"]
def test_session_route_prefers_last_event_id_then_query_fallback(monkeypatch):
import api.routes as routes
@ -554,6 +593,60 @@ def test_session_route_reconciles_late_attach_cutoff_once(monkeypatch):
assert stream.unsubscribed == 1
def test_session_route_emits_snapshot_when_reconciliation_fails(monkeypatch):
import api.routes as routes
class _FakeStream:
def __init__(self):
self.q = queue.Queue()
self.q.put_nowait(("token", {"text": "buffered"}, "run_active:1"))
self.q.put_nowait(("stream_end", {"status": "done"}, "run_active:2"))
def subscribe_with_snapshot(self):
return self.q, {"last_event_id": "run_active:1"}
stream = _FakeStream()
reads = []
sessions = iter(
[
SimpleNamespace(session_id="session_1", compact=lambda **_kwargs: {"session_id": "session_1", "title": "stale"}),
SimpleNamespace(session_id="session_1", compact=lambda **_kwargs: {"session_id": "session_1", "title": "fresh"}),
]
)
monkeypatch.setattr(routes, "_session_id_visible_to_request_profile", lambda *_args, **_kwargs: True)
monkeypatch.setattr(
routes,
"get_session",
lambda sid, metadata_only=False: next(sessions),
)
monkeypatch.setattr(routes, "_active_run_stream_for_session", lambda *_args, **_kwargs: "run_active")
monkeypatch.setattr(routes, "STREAMS", {"run_active": stream})
monkeypatch.setattr(
routes,
"read_session_run_events",
lambda *_args, **_kwargs: reads.append(1) or (
{"status": "ok", "events": [{"event": "token", "payload": {"text": "replayed"}, "event_id": "run_prev:1"}]}
if len(reads) == 1
else {"status": "replay_noncontiguous", "events": []}
),
)
handler = _FakeHandler()
routes._handle_session_sse_stream_for_session(
handler,
urlparse("/api/sessions/session_1/events?after_event_id=run_prev:0"),
"session_1",
)
body = handler.wfile.getvalue().decode("utf-8")
assert "event: session_snapshot\n" in body
assert '"title": "fresh"' in body
assert '"title": "stale"' not in body
assert "id: run_prev:1\n" not in body
assert body.count("id: run_active:1\n") == 0
assert "id: run_active:2\n" in body
def test_session_route_bounds_sent_event_id_deduplication(monkeypatch):
import api.routes as routes