From f15f178b5ed92bea30541ea45cb69b06e834fa2c Mon Sep 17 00:00:00 2001 From: Frank Song Date: Fri, 26 Jun 2026 18:01:14 +0800 Subject: [PATCH] Guard state db prefix keys with tool calls Co-Authored-By: Paperclip --- api/models.py | 26 ++++++--- api/routes.py | 9 +-- tests/test_webui_state_db_reconciliation.py | 63 +++++++++++++++++++++ 3 files changed, 84 insertions(+), 14 deletions(-) diff --git a/api/models.py b/api/models.py index 3f49c8f4..485a1a42 100644 --- a/api/models.py +++ b/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 diff --git a/api/routes.py b/api/routes.py index fd56d792..187122b6 100644 --- a/api/routes.py +++ b/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, diff --git a/tests/test_webui_state_db_reconciliation.py b/tests/test_webui_state_db_reconciliation.py index b9b22061..441e1a31 100644 --- a/tests/test_webui_state_db_reconciliation.py +++ b/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