Merge pull request #6153 from rh-id/fix/mem-state-db-read-backstop

fix(models): add defensive row backstop to display-path state.db message read

api/models.py

@@ -7346,6 +7346,7 @@ def get_state_db_session_messages(
     profile=None,
     since_timestamp=None,
     include_inactive: bool = False,
+    limit=None,
 ) -> list:
     """Read messages for a Hermes session from state.db.
 
@@ -7362,6 +7363,14 @@ def get_state_db_session_messages(
     preserving the caller's normal merge/window logic.  Full-history callers must
     leave it unset.
 
+    ``limit`` is an optional defensive row cap (applied after ORDER BY as a SQL
+    LIMIT). It is a BACKSTOP against a pathological/huge state.db materializing
+    unbounded rows into a Python list, NOT a semantic window: the display path
+    counts visible rows post-reconciliation, so a true window LIMIT here would
+    corrupt the sidecar/state.db merge (see _state_db_since_timestamp_for_limited_display,
+    which deliberately does NOT SQL-LIMIT raw rows for that reason). Callers that
+    need the full history for model-context reconstruction leave this unset.
+
     When the messages table exposes an ``active`` column, inactive rows are
     compacted/archived history and are intentionally excluded by default. WebUI
     reconciliation feeds this reader straight into the next model context; pulling
@@ -7462,14 +7471,42 @@ def get_state_db_session_messages(
             active_clause = ""
             if 'active' in available and not include_inactive:
                 active_clause = " AND (active IS NULL OR active != 0)"
-            cur.execute(f"""
-                SELECT {', '.join(selected)}, session_id
-                FROM messages
-                WHERE session_id IN ({placeholders})
-                {since_clause}
-                {active_clause}
-                ORDER BY timestamp ASC, id ASC
-            """, params)
+            # Defensive row cap (backstop only — see docstring). Applied as a
+            # SQL LIMIT bound parameter (?) so the tail (newest) rows are
+            # retained and a pathological state.db can't materialize unbounded
+            # rows. None = unchanged full-history read for model-context callers.
+            limit_clause = ""
+            if limit is not None:
+                try:
+                    limit_int = max(1, int(limit))
+                except (TypeError, ValueError):
+                    limit_int = None
+                if limit_int is not None:
+                    # The query orders ASC (oldest first); to keep the NEWEST
+                    # rows under the cap, take a descending-ordered subquery and
+                    # re-sort ascending — a plain LIMIT would keep the oldest.
+                    limit_clause = " ORDER BY timestamp DESC, id DESC LIMIT ?"
+                    params.append(limit_int)
+            if limit_clause:
+                cur.execute(f"""
+                    SELECT * FROM (
+                        SELECT {', '.join(selected)}, session_id
+                        FROM messages
+                        WHERE session_id IN ({placeholders})
+                        {since_clause}
+                        {active_clause}
+                        {limit_clause}
+                    ) ORDER BY timestamp ASC, id ASC
+                """, params)
+            else:
+                cur.execute(f"""
+                    SELECT {', '.join(selected)}, session_id
+                    FROM messages
+                    WHERE session_id IN ({placeholders})
+                    {since_clause}
+                    {active_clause}
+                    ORDER BY timestamp ASC, id ASC
+                """, params)
             msgs = []
             for row in cur.fetchall():
                 msg = {

api/routes.py

@@ -8446,6 +8446,42 @@ def _message_window_for_display(messages, msg_limit=None, msg_before=None, expan
 
 
 _LIMITED_TOOL_CONTENT_MAX_CHARS = 4096
+# Defensive row backstop for the GET /api/session display path's state.db read.
+# This is NOT a semantic window (the display window counts visible rows
+# post-reconciliation via _message_window_for_display); it is a safety net so a
+# pathological/huge state.db cannot materialize unbounded rows into memory on the
+# display path. Legitimate sessions stay far below this; the compressed-session
+# case where _state_db_since_timestamp_for_limited_display bails (and would
+# otherwise full-scan) is the main beneficiary. Generous on purpose: no real
+# conversation approaches it, and the existing since_timestamp optimization
+# already handles the common tail-load case. The full-history model-context
+# callers (reconciliation, new-turn context) do NOT use this cap.
+_STATE_DB_DISPLAY_ROW_BACKSTOP = 50000
+
+
+def _state_db_backstop_limit_for_display(session, msg_before) -> int | None:
+    """Return the row backstop to apply to the display path's state.db read, or
+    ``None`` for an uncapped (full-history) read.
+
+    The backstop is a defensive net against a pathological/huge state.db, NOT a
+    semantic window. It is applied ONLY on provably-safe reads where no
+    ``truncation_boundary`` prefix is required for the merge:
+    ``merge_session_messages_append_only`` needs the rows at/around the session's
+    ``truncation_boundary`` to reconcile correctly, and a newest-N-only SQL cap
+    would drop those boundary rows for a >N-row session and corrupt the merge
+    (silently losing the preserved prefix). So this mirrors the same conditions
+    ``_state_db_since_timestamp_for_limited_display`` uses to decide a read is
+    boundary-free: not ``msg_before`` paging, and no ``truncation_watermark`` /
+    ``truncation_boundary``. Extracted for direct test coverage.
+    """
+    has_boundary_prefix = (
+        msg_before is not None
+        or getattr(session, "truncation_watermark", None) not in (None, "")
+        or getattr(session, "truncation_boundary", None) not in (None, "")
+    )
+    return None if has_boundary_prefix else _STATE_DB_DISPLAY_ROW_BACKSTOP
+
+
 _LIMITED_TOOL_CONTENT_NOTICE = (
     "\n\n[Tool output truncated in paginated session response; "
     "load the full transcript to inspect the complete result.]"
@@ -12429,6 +12465,14 @@ def handle_get(handler, parsed) -> bool:
                 _state_db_reader_kwargs = {"profile": _session_profile}
                 if state_db_since_timestamp is not None:
                     _state_db_reader_kwargs["since_timestamp"] = state_db_since_timestamp
+                # Apply the display-path row backstop ONLY on provably-safe
+                # reads where no truncation_boundary prefix is required for the
+                # merge — see _state_db_backstop_limit_for_display. Compressed
+                # sessions and msg_before paging need their full prefix rows for
+                # correct reconciliation, so those stay uncapped.
+                _backstop = _state_db_backstop_limit_for_display(s, msg_before)
+                if _backstop is not None:
+                    _state_db_reader_kwargs["limit"] = _backstop
                 state_db_messages = get_state_db_session_messages(
                     sid,
                     **_state_db_reader_kwargs,

tests/test_core_data_loss_cases.py

@@ -499,7 +499,7 @@ def test_core_a_route_full_session_load_does_not_resurrect_deleted_turns(tmp_pat
             truncation_boundary=101.0,
         )
         s.save()
-        routes.get_state_db_session_messages = lambda sid, profile=None: list(state)
+        routes.get_state_db_session_messages = lambda sid, profile=None, since_timestamp=None, include_inactive=False, limit=None: list(state)
         routes._session_visible_to_active_profile = lambda profile, handler: True
         routes._clear_stale_stream_state = lambda s: None
         routes._resolve_effective_session_model_for_display = lambda s: getattr(s, "model", None)

tests/test_session_lineage_full_transcript.py

@@ -226,7 +226,7 @@ def test_webui_continuation_session_opens_with_snapshot_parent_messages(monkeypa
     monkeypatch.setattr(routes, "_clear_stale_stream_state", lambda s: None)
     monkeypatch.setattr(routes, "_lookup_cli_session_metadata", lambda sid: {})
     monkeypatch.setattr(routes, "_is_messaging_session_record", lambda s: False)
-    monkeypatch.setattr(routes, "get_state_db_session_messages", lambda sid, profile=None: [])
+    monkeypatch.setattr(routes, "get_state_db_session_messages", lambda sid, profile=None, since_timestamp=None, include_inactive=False, limit=None: [])
     monkeypatch.setattr(routes.Session, "load", lambda sid: parent if sid == "parent-webui" else None)
     monkeypatch.setattr(routes, "_resolve_effective_session_model_for_display", lambda s: getattr(s, "model", None))
     monkeypatch.setattr(routes, "_resolve_effective_session_model_provider_for_display", lambda s: None)

tests/test_state_db_backstop_boundary_gate.py

@@ -0,0 +1,68 @@
+"""Tests: the display-path state.db row backstop is GATED on the absence of a
+truncation_boundary prefix.
+
+Regression (reviewer): ``_STATE_DB_DISPLAY_ROW_BACKSTOP`` was applied
+unconditionally to the ``GET /api/session`` display path's state.db read. But
+``merge_session_messages_append_only`` needs the rows at/around the session's
+``truncation_boundary`` (the preserved prefix) to reconcile correctly. A
+newest-N-only SQL cap drops those boundary rows for a >N-row session and
+corrupts the merge (silently losing the preserved prefix), and also prevents
+older-page (``msg_before``) requests from reaching them.
+
+The fix gates the backstop via ``_state_db_backstop_limit_for_display`` on the
+same conditions ``_state_db_since_timestamp_for_limited_display`` uses to decide
+a read is boundary-free: NOT ``msg_before`` paging, and NO
+``truncation_watermark`` / ``truncation_boundary``. Those cases stay on the full
+(uncapped) read. The helper is the unit under test (driving the full handler
+end-to-end is brittle across its many dependencies).
+"""
+from __future__ import annotations
+
+from api.routes import _STATE_DB_DISPLAY_ROW_BACKSTOP, _state_db_backstop_limit_for_display
+
+
+class _StubSession:
+    """Minimal object exposing the boundary attributes the gate checks."""
+    def __init__(self, **attrs):
+        self.truncation_watermark = attrs.get("truncation_watermark", None)
+        self.truncation_boundary = attrs.get("truncation_boundary", None)
+
+
+def test_backstop_applied_when_no_boundary_prefix():
+    """An uncompressed, initial-tail session (no truncation_watermark /
+    truncation_boundary, no msg_before) gets the defensive row backstop."""
+    assert _state_db_backstop_limit_for_display(_StubSession(), msg_before=None) == _STATE_DB_DISPLAY_ROW_BACKSTOP
+
+
+def test_backstop_skipped_when_truncation_boundary_set():
+    """Regression: a compressed session (truncation_boundary set) needs its
+    preserved-prefix rows for the merge, so the backstop must NOT cap the read."""
+    stub = _StubSession(truncation_boundary="some-boundary-marker")
+    assert _state_db_backstop_limit_for_display(stub, msg_before=None) is None
+
+
+def test_backstop_skipped_when_truncation_watermark_set():
+    """Same gate for truncation_watermark — the merge needs the prefix rows."""
+    stub = _StubSession(truncation_watermark=12345)
+    assert _state_db_backstop_limit_for_display(stub, msg_before=None) is None
+
+
+def test_backstop_skipped_when_msg_before_paging():
+    """Older-page (msg_before) requests need to reach the boundary rows, so the
+    backstop must not cap them either."""
+    stub = _StubSession()
+    assert _state_db_backstop_limit_for_display(stub, msg_before=500) is None
+
+
+def test_backstop_boundary_check_handles_empty_string_as_unset_or_set():
+    """The gate treats empty-string boundary markers as 'not set' (matches the
+    existing since_timestamp helper's `in (None, "")` check). A real marker
+    (non-empty) triggers the skip."""
+    # Empty string = treated as unset → backstop applies.
+    assert _state_db_backstop_limit_for_display(
+        _StubSession(truncation_boundary=""), msg_before=None
+    ) == _STATE_DB_DISPLAY_ROW_BACKSTOP
+    # Non-empty string = boundary present → backstop skipped.
+    assert _state_db_backstop_limit_for_display(
+        _StubSession(truncation_watermark="wm"), msg_before=None
+    ) is None

tests/test_state_db_read_backstop.py

@@ -0,0 +1,119 @@
+"""Tests: ``get_state_db_session_messages`` defensive ``limit`` backstop.
+
+The display path (``GET /api/session``) reads state.db rows that feed a
+reconciliation merge then a visible-row window. The existing
+``since_timestamp`` optimization bounds the common uncompressed tail load but
+deliberately bails on compressed sessions (``truncation_watermark`` /
+``truncation_boundary``) and ``msg_before`` paging — in those cases the read
+used to full-scan with no SQL ``LIMIT``. A pathological/huge state.db could
+then materialize unbounded rows into memory.
+
+``get_state_db_session_messages`` now accepts an opt-in ``limit`` (applied as a
+SQL ``LIMIT`` that keeps the NEWEST rows). It is a defensive backstop, NOT a
+semantic window — full-history model-context callers leave it unset. These
+tests verify the newest-row retention and that the default (no limit) is
+unchanged.
+"""
+import sqlite3
+
+import api.models as models
+from api.models import get_state_db_session_messages
+
+
+def _make_state_db(path, n_rows=100):
+    conn = sqlite3.connect(str(path))
+    conn.execute(
+        "CREATE TABLE sessions (id TEXT PRIMARY KEY)"
+    )
+    conn.execute(
+        """
+        CREATE TABLE messages (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            session_id TEXT NOT NULL,
+            role TEXT,
+            content TEXT,
+            timestamp REAL,
+            active INTEGER DEFAULT 1
+        )
+        """
+    )
+    conn.execute("INSERT INTO sessions (id) VALUES ('s1')")
+    conn.executemany(
+        "INSERT INTO messages (session_id, role, content, timestamp) VALUES (?, ?, ?, ?)",
+        [("s1", "user" if i % 2 == 0 else "assistant", f"msg{i}", float(i)) for i in range(n_rows)],
+    )
+    conn.commit()
+    conn.close()
+
+
+def test_no_limit_returns_full_history_oldest_first(tmp_path, monkeypatch):
+    """Default (no limit) — the model-context caller contract — returns every
+    row, oldest-first. Unchanged by this fix."""
+    db = tmp_path / "state.db"
+    _make_state_db(db, n_rows=100)
+    monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
+
+    msgs = get_state_db_session_messages("s1")
+    assert len(msgs) == 100
+    assert msgs[0]["content"] == "msg0"
+    assert msgs[-1]["content"] == "msg99"
+
+
+def test_limit_keeps_newest_rows_resorted_oldest_first(tmp_path, monkeypatch):
+    """limit=N keeps the NEWEST N rows (the query orders ASC for oldest-first,
+    so the cap takes a DESC subquery then re-sorts), so a tail window is
+    retained rather than the oldest rows."""
+    db = tmp_path / "state.db"
+    _make_state_db(db, n_rows=100)
+    monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
+
+    msgs = get_state_db_session_messages("s1", limit=10)
+    assert len(msgs) == 10
+    # Newest 10 retained (msg90..msg99), re-sorted oldest-first.
+    assert [m["content"] for m in msgs] == [f"msg{i}" for i in range(90, 100)]
+
+
+def test_limit_above_row_count_returns_all(tmp_path, monkeypatch):
+    """A limit larger than the row count returns everything (no truncation)."""
+    db = tmp_path / "state.db"
+    _make_state_db(db, n_rows=50)
+    monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
+
+    msgs = get_state_db_session_messages("s1", limit=10000)
+    assert len(msgs) == 50
+
+
+def test_limit_zero_or_negative_clamps_to_one_newest(tmp_path, monkeypatch):
+    """A non-positive limit clamps to 1 and returns the single newest row."""
+    db = tmp_path / "state.db"
+    _make_state_db(db, n_rows=20)
+    monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
+
+    for bad in (0, -5):
+        msgs = get_state_db_session_messages("s1", limit=bad)
+        assert len(msgs) == 1
+        assert msgs[0]["content"] == "msg19"  # newest
+
+
+def test_limit_invalid_falls_back_to_unbounded(tmp_path, monkeypatch):
+    """A non-numeric limit is ignored (full read) rather than raising."""
+    db = tmp_path / "state.db"
+    _make_state_db(db, n_rows=30)
+    monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
+
+    msgs = get_state_db_session_messages("s1", limit="not-a-number")
+    assert len(msgs) == 30
+
+
+def test_limit_composes_with_since_timestamp(tmp_path, monkeypatch):
+    """limit and since_timestamp compose: the floor filters rows, then the cap
+    keeps the newest of the remainder."""
+    db = tmp_path / "state.db"
+    _make_state_db(db, n_rows=100)
+    monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
+
+    # since_timestamp=50.0 → rows with timestamp >= 50 (msg50..msg99 = 50 rows),
+    # then limit=10 keeps the newest 10 (msg90..msg99).
+    msgs = get_state_db_session_messages("s1", since_timestamp=50.0, limit=10)
+    assert len(msgs) == 10
+    assert [m["content"] for m in msgs] == [f"msg{i}" for i in range(90, 100)]