Guard state db prefix keys with tool calls

Co-Authored-By: Paperclip <noreply@paperclip.ing>

api/models.py

@@ -5171,13 +5171,14 @@ def get_state_db_session_message_keys_before_timestamp(
     before_timestamp,
     *,
     profile=None,
-) -> list[tuple[str, str]] | None:
-    """Return sorted raw ``(role, content)`` keys before ``before_timestamp``.
+) -> list[tuple] | None:
+    """Return visible-identity keys before ``before_timestamp`` in DB order.
 
     Missing timestamps are intentionally excluded because the bounded reader
     keeps them with ``timestamp IS NULL OR timestamp >= ?``.  The caller uses
     this as a conservative prefix-identity guard before taking the optimized
-    tail-read path.
+    tail-read path, so schemas that cannot prove the merge-visible identity
+    force a full read.
     """
     try:
         import sqlite3
@@ -5206,21 +5207,30 @@ def get_state_db_session_message_keys_before_timestamp(
             cur = conn.cursor()
             cur.execute("PRAGMA table_info(messages)")
             available = {str(row['name']) for row in cur.fetchall()}
-            if not {'session_id', 'role', 'content', 'timestamp'}.issubset(available):
+            if not {'id', 'session_id', 'role', 'content', 'timestamp', 'tool_calls'}.issubset(available):
                 return None
             cur.execute(
                 """
-                SELECT COALESCE(role, '') AS role, COALESCE(content, '') AS content
+                SELECT
+                    COALESCE(role, '') AS role,
+                    COALESCE(content, '') AS content,
+                    tool_calls
                 FROM messages
                 WHERE session_id = ? AND timestamp IS NOT NULL AND timestamp < ?
                 ORDER BY timestamp ASC, id ASC
                 """,
                 (str(sid), before_ts),
             )
-            return sorted(
-                (str(row['role'] or ''), str(row['content'] or ''))
+            return [
+                _session_message_visible_key(
+                    {
+                        "role": row["role"],
+                        "content": row["content"],
+                        "tool_calls": _json_loads_if_string(row["tool_calls"]),
+                    }
+                )
                 for row in cur.fetchall()
-            )
+            ]
     except Exception:
         return None
 

api/routes.py

@@ -6132,14 +6132,11 @@ 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_keys = sorted(
-        (
-            str(msg.get("role") or ""),
-            str(msg.get("content") or ""),
-        )
+    sidecar_before_keys = [
+        _session_message_visible_key(msg)
         for msg, ts in zip(sidecar_messages, sidecar_timestamps, strict=True)
         if ts < floor
-    )
+    ]
     state_before_keys = get_state_db_session_message_keys_before_timestamp(
         getattr(session, "session_id", None),
         floor,

tests/test_webui_state_db_reconciliation.py

@@ -666,6 +666,69 @@ def test_msg_limit_session_load_bails_when_prefloor_key_counts_mask_offset_chang
     assert len(full_all_messages) == len(sidecar_messages) + 1
 
 
+def test_msg_limit_session_load_bails_when_prefloor_tool_calls_mask_offset_change(monkeypatch, tmp_path):
+    import api.routes as routes
+
+    sid = "webui_reconcile_limited_tail_tool_calls_bail"
+    sidecar_tool_calls = [{"id": "call_sidecar", "function": {"name": "terminal", "arguments": "{}"}}]
+    state_tool_calls = [{"id": "call_state", "function": {"name": "terminal", "arguments": "{}"}}]
+    sidecar_messages = [
+        {"role": "user", "content": f"sidecar {idx}", "timestamp": float(idx)}
+        for idx in range(500)
+    ]
+    sidecar_messages[100] = {
+        "role": "assistant",
+        "content": "",
+        "timestamp": 100.0,
+        "tool_calls": sidecar_tool_calls,
+    }
+    state_messages = [
+        dict(msg, tool_calls=json.dumps(msg["tool_calls"]))
+        if msg.get("tool_calls")
+        else dict(msg)
+        for msg in sidecar_messages
+    ]
+    state_messages[100] = {
+        "role": "assistant",
+        "content": "",
+        "timestamp": 100.0,
+        "tool_calls": json.dumps(state_tool_calls),
+    }
+    session = _install_test_session(monkeypatch, tmp_path, sid, sidecar_messages)
+    _make_state_db(tmp_path / "state.db", sid, state_messages)
+
+    real_reader = routes.get_state_db_session_messages
+    full_state_messages = real_reader(sid)
+    full_all_messages = routes._limited_webui_messages_for_display(
+        session,
+        full_state_messages,
+    )
+    expected_window, expected_offset = routes._message_window_for_display(
+        full_all_messages,
+        msg_limit=30,
+    )
+    captured = {}
+
+    def wrapped_reader(*args, **kwargs):
+        captured["since_timestamp"] = kwargs.get("since_timestamp")
+        return real_reader(*args, **kwargs)
+
+    monkeypatch.setattr(routes, "get_state_db_session_messages", wrapped_reader)
+
+    handler = _GetHandler(
+        f"/api/session?session_id={sid}&messages=1&resolve_model=0&msg_limit=30"
+    )
+    routes.handle_get(handler, urlparse(handler.path))
+
+    assert handler.status == 200
+    assert captured["since_timestamp"] is None
+    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
+    assert len(full_all_messages) == len(sidecar_messages) + 1
+
+
 def test_msg_limit_session_load_bails_when_truncation_boundary_is_set(monkeypatch, tmp_path):
     import api.routes as routes