diff --git a/api/routes.py b/api/routes.py index 373fccbc..f954e8ea 100644 --- a/api/routes.py +++ b/api/routes.py @@ -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: diff --git a/api/run_journal.py b/api/run_journal.py index e339eb1c..8be73af3 100644 --- a/api/run_journal.py +++ b/api/run_journal.py @@ -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:]: diff --git a/tests/test_issue4812_session_sse_stream.py b/tests/test_issue4812_session_sse_stream.py index b0865c93..60dbec33 100644 --- a/tests/test_issue4812_session_sse_stream.py +++ b/tests/test_issue4812_session_sse_stream.py @@ -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