Merge pull request #6053 from nesquena/stage/6014

Release exp: faster session reconciliation — short-circuit impossible prefix (#6014)

CHANGELOG.md

@@ -21,6 +21,10 @@
 
 - **New chats honor your agent's `worktree:` config default.** If your agent config sets a `worktree:` default, a deliberate New Chat now inherits it (matching TUI behavior). System-created sessions — boot auto-bind, profile switches, workspace binding, file/folder creation, terminal sessions — never inherit it, so simply opening the UI can't spawn a stray worktree + branch. An explicit request flag always wins over the config default, and the config value is read strictly (only a real `true` opts in; anything else is safely ignored). Thanks @ayxuerui. (#6029, #6022)
 
+### Performance
+
+- **Faster session reconciliation for large histories.** When the WebUI reconciles a session against the agent's `state.db`, it now uses a cheap prefix-count check to skip the expensive full identity-key materialization whenever the two prefixes provably can't match — falling back to the full comparison whenever a match can't be ruled out (null timestamps, equal counts, ties, or a locked/unreadable db all take the safe full path). No change to which sessions or messages are shown. Thanks @Isla-Liu. (#6014)
+
 - **Send-to-first-token latency (TTFT) is now measured.** Each streamed turn records how long it took from send to the first visible token (`ttft_ms`), surfaced in metering stats and as a `_firstTokenMs` diagnostic on the message. Measurement-only — computed once on the first token with no added latency on the streaming hot path, and existing token/usage/cost metering is unchanged. Thanks @rodboev. (#6042, #6006)
 
 - **"Approve once" no longer silently sticks for the whole session.** On local (non-gateway) deployments, choosing **Approve once** on a tool approval prompt was wrongly persisting as a session-wide approval, so the next matching guarded tool call skipped its approval card. "Once" now applies to only the single call that prompted it, and the next matching call re-prompts as intended; "Approve for session" and "Always allow" are unchanged. Thanks @rudidev08. (#6035, #6017)

api/models.py

@@ -7351,6 +7351,75 @@ def get_state_db_session_messages(
     return msgs
 
 
+def get_state_db_session_message_prefix_summary(
+    sid,
+    before_timestamp,
+    *,
+    profile=None,
+) -> dict | None:
+    """Return prefix timestamp counts, or ``None`` when they cannot be proven.
+
+    The projection intentionally avoids message content and tool-call columns so
+    callers can reject impossible prefix matches before materializing visible
+    identities. Missing databases are an authoritative empty prefix and are not
+    created by this read path.
+    """
+    try:
+        import sqlite3
+    except ImportError:
+        return None
+
+    if not sid:
+        return None
+    try:
+        before_ts = float(before_timestamp)
+    except (TypeError, ValueError):
+        return None
+
+    if isinstance(profile, str) and profile:
+        db_path = _get_profile_home(profile) / 'state.db'
+        if not db_path.exists():
+            db_path = _active_state_db_path()
+    else:
+        db_path = _active_state_db_path()
+    if not db_path.exists():
+        return {"count": 0, "null_timestamp_count": 0}
+
+    try:
+        with closing(open_state_db_readonly(db_path)) as conn:
+            conn.row_factory = sqlite3.Row
+            cur = conn.cursor()
+            cur.execute("PRAGMA table_info(messages)")
+            available = {str(row['name']) for row in cur.fetchall()}
+            if not {'session_id', 'timestamp'}.issubset(available):
+                return None
+            active_clause = ""
+            if 'active' in available:
+                active_clause = " AND (active IS NULL OR active != 0)"
+            cur.execute(
+                f"""
+                SELECT
+                    COUNT(CASE
+                        WHEN timestamp IS NOT NULL AND timestamp < ? THEN 1
+                    END) AS count,
+                    COUNT(CASE WHEN timestamp IS NULL THEN 1 END) AS null_timestamp_count
+                FROM messages
+                WHERE session_id = ?
+                {active_clause}
+                """,
+                (before_ts, str(sid)),
+            )
+            row = cur.fetchone()
+            if row is None:
+                return None
+            return {
+                "count": int(row["count"]),
+                "null_timestamp_count": int(row["null_timestamp_count"]),
+            }
+    except Exception:
+        return None
+
+
 def get_state_db_session_message_keys_before_timestamp(
     sid,
     before_timestamp,

api/routes.py

@@ -8557,6 +8557,24 @@ def _state_db_since_timestamp_for_limited_display(session, msg_limit, msg_before
         return None, sidecar_messages
 
     floor = min(sidecar_timestamps[-raw_budget:])
+    sidecar_before_count = sum(1 for ts in sidecar_timestamps if ts < floor)
+    prefix_summary = get_state_db_session_message_prefix_summary(
+        getattr(session, "session_id", None),
+        floor,
+        profile=getattr(session, "profile", None) or None,
+    )
+    if prefix_summary is None:
+        return None, sidecar_messages
+    try:
+        state_before_count = int(prefix_summary["count"])
+        null_timestamp_count = int(prefix_summary["null_timestamp_count"])
+    except (KeyError, TypeError, ValueError):
+        return None, sidecar_messages
+    if null_timestamp_count or state_before_count != sidecar_before_count:
+        return None, sidecar_messages
+    if sidecar_before_count == 0:
+        return floor, sidecar_messages
+
     sidecar_before_keys = [
         _session_message_visible_key(msg)
         for msg, ts in zip(sidecar_messages, sidecar_timestamps, strict=True)
@@ -9170,6 +9188,7 @@ from api.models import (
     get_cli_sessions,
     get_cli_session_messages,
     get_state_db_session_messages,
+    get_state_db_session_message_prefix_summary,
     get_state_db_session_message_keys_before_timestamp,
     get_state_db_session_summary,
     merge_session_messages_append_only,

tests/test_state_db_readonly_reads_models.py

@@ -15,6 +15,7 @@ write-capable handle on the agent DB.
 """
 import sqlite3
 from contextlib import closing
+from pathlib import Path
 
 import api.agent_sessions as agent_sessions
 import api.models as models
@@ -148,6 +149,72 @@ def test_get_state_db_message_keys_before_timestamp_opens_read_only(tmp_path, mo
     _assert_read_only(calls)
 
 
+def test_get_state_db_message_prefix_summary_opens_read_only(tmp_path, monkeypatch):
+    db = tmp_path / "state.db"
+    _make_state_db(db)
+    _point_models_at(monkeypatch, db)
+    calls = _record_connects(monkeypatch)
+
+    summary = models.get_state_db_session_message_prefix_summary("sess-1", 1000.5)
+
+    assert summary == {"count": 1, "null_timestamp_count": 0}
+    _assert_read_only(calls)
+
+
+def test_get_state_db_message_prefix_summary_preserves_active_filtering(tmp_path, monkeypatch):
+    db = tmp_path / "state.db"
+    conn = sqlite3.connect(str(db))
+    conn.execute(
+        "CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id TEXT, timestamp REAL, active INTEGER)"
+    )
+    conn.executemany(
+        "INSERT INTO messages (id, session_id, timestamp, active) VALUES (?, ?, ?, ?)",
+        [
+            (1, "sess-1", 10.0, 1),
+            (2, "sess-1", 20.0, 0),
+            (3, "sess-1", None, None),
+            (4, "sess-1", None, 0),
+            (5, "other", 10.0, 1),
+        ],
+    )
+    conn.commit()
+    conn.close()
+    _point_models_at(monkeypatch, db)
+
+    summary = models.get_state_db_session_message_prefix_summary("sess-1", 30.0)
+
+    assert summary == {"count": 1, "null_timestamp_count": 1}
+
+
+def test_get_state_db_message_prefix_summary_returns_none_for_inconclusive_schema(
+    tmp_path,
+    monkeypatch,
+):
+    db = tmp_path / "state.db"
+    conn = sqlite3.connect(str(db))
+    conn.execute("CREATE TABLE messages (id INTEGER PRIMARY KEY, content TEXT)")
+    conn.commit()
+    conn.close()
+    _point_models_at(monkeypatch, db)
+
+    assert models.get_state_db_session_message_prefix_summary("sess-1", 1000.5) is None
+
+
+def test_get_state_db_message_prefix_summary_missing_db_creates_no_sqlite_files(
+    tmp_path,
+    monkeypatch,
+):
+    db = tmp_path / "missing" / "state.db"
+    _point_models_at(monkeypatch, db)
+
+    summary = models.get_state_db_session_message_prefix_summary("sess-1", 1000.5)
+
+    assert summary == {"count": 0, "null_timestamp_count": 0}
+    assert not db.exists()
+    assert not Path(f"{db}-wal").exists()
+    assert not Path(f"{db}-shm").exists()
+
+
 def test_get_state_db_session_summary_opens_read_only(tmp_path, monkeypatch):
     db = tmp_path / "state.db"
     _make_state_db(db)

tests/test_webui_state_db_reconciliation.py

@@ -125,6 +125,13 @@ def _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages):
     return session
 
 
+def _large_timestamped_sidecar_messages(count=500):
+    return [
+        {"role": "user", "content": f"sidecar {idx}", "timestamp": float(idx)}
+        for idx in range(count)
+    ]
+
+
 def test_sidebar_state_db_overlay_preserves_numeric_actual_count():
     import api.models as models
 
@@ -532,7 +539,7 @@ def test_msg_limit_session_load_reads_only_recent_state_db_tail(monkeypatch, tmp
     assert messages[-1]["content"] == "external answer"
 
 
-def test_msg_limit_session_load_matches_full_read_with_null_state_db_timestamp(monkeypatch, tmp_path):
+def test_msg_limit_session_load_falls_back_with_null_state_db_timestamp(monkeypatch, tmp_path):
     import api.routes as routes
 
     sid = "webui_reconcile_limited_tail_null"
@@ -580,14 +587,257 @@ def test_msg_limit_session_load_matches_full_read_with_null_state_db_timestamp(m
     routes.handle_get(handler, urlparse(handler.path))
 
     assert handler.status == 200
-    assert captured["since_timestamp"] == 200.0
-    assert captured["row_count"] == 301
+    assert captured["since_timestamp"] is None
+    assert captured["row_count"] == 501
     session_payload = handler.response_json["session"]
     assert session_payload["messages"] == expected_window
     assert session_payload["message_count"] == len(full_all_messages)
     assert session_payload["_messages_offset"] == expected_offset
 
 
+def test_limited_state_db_prefix_missing_db_skips_visible_key_normalization(monkeypatch, tmp_path):
+    import api.models as models
+    import api.routes as routes
+
+    sid = "webui_reconcile_prefix_missing_db"
+    sidecar_messages = _large_timestamped_sidecar_messages(10_000)
+    session = _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages)
+    visible_key_calls = 0
+    real_visible_key = routes._session_message_visible_key
+
+    def counted_visible_key(message):
+        nonlocal visible_key_calls
+        visible_key_calls += 1
+        return real_visible_key(message)
+
+    monkeypatch.setattr(routes, "_session_message_visible_key", counted_visible_key)
+    monkeypatch.setattr(models, "_session_message_visible_key", counted_visible_key)
+
+    floor, returned_sidecar = routes._state_db_since_timestamp_for_limited_display(
+        session,
+        30,
+    )
+
+    assert floor is None
+    assert returned_sidecar == sidecar_messages
+    assert visible_key_calls == 0
+
+
+def test_limited_state_db_prefix_count_mismatch_skips_visible_key_normalization(monkeypatch, tmp_path):
+    import api.models as models
+    import api.routes as routes
+
+    sid = "webui_reconcile_prefix_count_mismatch"
+    sidecar_messages = _large_timestamped_sidecar_messages()
+    session = _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages)
+    _make_state_db(tmp_path / "state.db", sid, sidecar_messages[:10])
+    visible_key_calls = 0
+    real_visible_key = routes._session_message_visible_key
+
+    def counted_visible_key(message):
+        nonlocal visible_key_calls
+        visible_key_calls += 1
+        return real_visible_key(message)
+
+    monkeypatch.setattr(routes, "_session_message_visible_key", counted_visible_key)
+    monkeypatch.setattr(models, "_session_message_visible_key", counted_visible_key)
+
+    floor, returned_sidecar = routes._state_db_since_timestamp_for_limited_display(
+        session,
+        30,
+    )
+
+    assert floor is None
+    assert returned_sidecar == sidecar_messages
+    assert visible_key_calls == 0
+
+
+def test_limited_state_db_prefix_exact_match_runs_key_comparison(monkeypatch, tmp_path):
+    import api.routes as routes
+
+    sid = "webui_reconcile_prefix_exact"
+    sidecar_messages = _large_timestamped_sidecar_messages()
+    session = _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages)
+    _make_state_db(tmp_path / "state.db", sid, sidecar_messages)
+    summary_calls = []
+    key_calls = []
+    visible_key_calls = 0
+    real_summary_reader = routes.get_state_db_session_message_prefix_summary
+    real_key_reader = routes.get_state_db_session_message_keys_before_timestamp
+    real_visible_key = routes._session_message_visible_key
+
+    def prefix_summary(*args, **kwargs):
+        summary_calls.append((args, kwargs))
+        return real_summary_reader(*args, **kwargs)
+
+    def counted_key_reader(*args, **kwargs):
+        key_calls.append((args, kwargs))
+        return real_key_reader(*args, **kwargs)
+
+    def counted_visible_key(message):
+        nonlocal visible_key_calls
+        visible_key_calls += 1
+        return real_visible_key(message)
+
+    monkeypatch.setattr(
+        routes,
+        "get_state_db_session_message_prefix_summary",
+        prefix_summary,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        routes,
+        "get_state_db_session_message_keys_before_timestamp",
+        counted_key_reader,
+    )
+    monkeypatch.setattr(routes, "_session_message_visible_key", counted_visible_key)
+
+    floor, returned_sidecar = routes._state_db_since_timestamp_for_limited_display(
+        session,
+        30,
+    )
+
+    assert floor == 200.0
+    assert returned_sidecar == sidecar_messages
+    assert len(summary_calls) == 1
+    assert len(key_calls) == 1
+    assert visible_key_calls == 200
+
+
+def test_limited_state_db_prefix_equal_count_different_content_falls_back(monkeypatch, tmp_path):
+    import api.routes as routes
+
+    sid = "webui_reconcile_prefix_content_mismatch"
+    sidecar_messages = _large_timestamped_sidecar_messages()
+    state_messages = [dict(message) for message in sidecar_messages]
+    state_messages[100]["content"] = "edited state content"
+    session = _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages)
+    _make_state_db(tmp_path / "state.db", sid, state_messages)
+    summary_calls = []
+    key_calls = []
+    real_summary_reader = routes.get_state_db_session_message_prefix_summary
+    real_key_reader = routes.get_state_db_session_message_keys_before_timestamp
+
+    def prefix_summary(*args, **kwargs):
+        summary_calls.append((args, kwargs))
+        return real_summary_reader(*args, **kwargs)
+
+    def counted_key_reader(*args, **kwargs):
+        key_calls.append((args, kwargs))
+        return real_key_reader(*args, **kwargs)
+
+    monkeypatch.setattr(
+        routes,
+        "get_state_db_session_message_prefix_summary",
+        prefix_summary,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        routes,
+        "get_state_db_session_message_keys_before_timestamp",
+        counted_key_reader,
+    )
+
+    floor, returned_sidecar = routes._state_db_since_timestamp_for_limited_display(
+        session,
+        30,
+    )
+
+    assert floor is None
+    assert returned_sidecar == sidecar_messages
+    assert len(summary_calls) == 1
+    assert len(key_calls) == 1
+
+
+def test_limited_state_db_prefix_equal_empty_assistant_different_tool_calls_falls_back(
+    monkeypatch,
+    tmp_path,
+):
+    import api.routes as routes
+
+    sid = "webui_reconcile_prefix_tool_calls_mismatch"
+    sidecar_messages = _large_timestamped_sidecar_messages()
+    sidecar_messages[100] = {
+        "role": "assistant",
+        "content": "",
+        "timestamp": 100.0,
+        "tool_calls": [{"id": "sidecar-call", "function": {"name": "terminal"}}],
+    }
+    state_messages = [dict(message) for message in sidecar_messages]
+    state_messages[100] = {
+        "role": "assistant",
+        "content": "",
+        "timestamp": 100.0,
+        "tool_calls": json.dumps([{"id": "state-call", "function": {"name": "terminal"}}]),
+    }
+    session = _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages)
+    _make_state_db(tmp_path / "state.db", sid, state_messages)
+    summary_calls = []
+    key_calls = []
+    real_summary_reader = routes.get_state_db_session_message_prefix_summary
+    real_key_reader = routes.get_state_db_session_message_keys_before_timestamp
+
+    def prefix_summary(*args, **kwargs):
+        summary_calls.append((args, kwargs))
+        return real_summary_reader(*args, **kwargs)
+
+    def counted_key_reader(*args, **kwargs):
+        key_calls.append((args, kwargs))
+        return real_key_reader(*args, **kwargs)
+
+    monkeypatch.setattr(
+        routes,
+        "get_state_db_session_message_prefix_summary",
+        prefix_summary,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        routes,
+        "get_state_db_session_message_keys_before_timestamp",
+        counted_key_reader,
+    )
+
+    floor, returned_sidecar = routes._state_db_since_timestamp_for_limited_display(
+        session,
+        30,
+    )
+
+    assert floor is None
+    assert returned_sidecar == sidecar_messages
+    assert len(summary_calls) == 1
+    assert len(key_calls) == 1
+
+
+def test_limited_state_db_prefix_missing_sidecar_timestamp_preserves_full_fallback(
+    monkeypatch,
+    tmp_path,
+):
+    import api.routes as routes
+
+    sid = "webui_reconcile_prefix_missing_sidecar_timestamp"
+    sidecar_messages = _large_timestamped_sidecar_messages()
+    sidecar_messages[100]["timestamp"] = None
+    session = _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages)
+
+    def unexpected_prefix_summary(*args, **kwargs):
+        raise AssertionError("missing sidecar timestamps must fall back before state.db preflight")
+
+    monkeypatch.setattr(
+        routes,
+        "get_state_db_session_message_prefix_summary",
+        unexpected_prefix_summary,
+        raising=False,
+    )
+
+    floor, returned_sidecar = routes._state_db_since_timestamp_for_limited_display(
+        session,
+        30,
+    )
+
+    assert floor is None
+    assert returned_sidecar == sidecar_messages
+
+
 def test_msg_limit_session_load_bails_when_older_state_db_row_changes_offsets(monkeypatch, tmp_path):
     import api.routes as routes