fix(regeneration): snapshot reuses canonical projection+ordering; WAL data_version guard (#6826 r4)

This commit is contained in:
webtecnica authored and n committed 2026-08-26 00:24:07 +00:00
1 parent b33aff9d09
commit 7abd4773ca
2 files changed
+204 -61

No files matched your search

+80 -61
View File
@@ -8178,6 +8178,41 @@ def _state_db_active_rows_digest(rows) -> str:
return digest.hexdigest() if stamped else ''
def _project_state_db_message(row, available, id_col, optional):
"""Authoritative state.db row → WebUI message projection (#6826 r4).
Shared by ``get_state_db_session_messages`` and the regeneration
single-snapshot helper so the bounded tail can never drift from the
canonical reader: JSON-decode tool_calls/reasoning payloads, omit
empty fields, keep durable row id private (``_state_db_row_id`` only for
real Agent api_content replays), and apply ``tool_name → name``.
"""
msg = {
'role': row['role'],
'content': row['content'],
'timestamp': row['timestamp'],
}
for col in optional:
if col not in row.keys():
continue
value = row[col]
if value in (None, ''):
continue
if col in {'tool_calls', 'reasoning_details', 'codex_reasoning_items', 'codex_message_items'}:
value = _json_loads_if_string(value)
msg[col] = value
if (
id_col
and row['id'] is not None
and isinstance(msg.get('api_content'), str)
and msg['api_content']
):
msg['_state_db_row_id'] = row['id']
if msg.get('role') == 'tool' and msg.get('tool_name') and not msg.get('name'):
msg['name'] = msg['tool_name']
return msg
@overload
def get_state_db_session_messages(
sid,
@@ -8418,38 +8453,9 @@ def get_state_db_session_messages(
msgs = []
for row in rows:
msg = {
'role': row['role'],
'content': row['content'],
'timestamp': row['timestamp'],
}
# ``id`` is the durable SQLite row identity, not the WebUI's
# session-local stable message id. Keep it in a private
# provenance field so duplicate reconciliation can align a
# state.db sidecar without changing the existing ``id`` key
# used by WebUI transcript merge/dedup logic.
for col in optional:
if col not in row.keys():
continue
value = row[col]
if value in (None, ''):
continue
if col in {'tool_calls', 'reasoning_details', 'codex_reasoning_items', 'codex_message_items'}:
value = _json_loads_if_string(value)
msg[col] = value
# Keep durable provenance only alongside a real Agent replay
# sidecar. Ordinary state.db rows must retain their historic
# shape and must not acquire internal bookkeeping fields.
if (
id_col
and row['id'] is not None
and isinstance(msg.get('api_content'), str)
and msg['api_content']
):
msg['_state_db_row_id'] = row['id']
if msg.get('role') == 'tool' and msg.get('tool_name') and not msg.get('name'):
msg['name'] = msg['tool_name']
msgs.append(msg)
msgs.append(
_project_state_db_message(row, available, bool(id_col), optional)
)
except Exception:
return _state_db_session_messages_result([], None, with_revision=with_revision)
return _state_db_session_messages_result(msgs, revision, with_revision=with_revision)
@@ -8645,6 +8651,8 @@ def get_state_db_regeneration_tail_snapshot(
with closing(open_state_db_readonly(db_path)) as conn:
conn.row_factory = sqlite3.Row
cur = conn.cursor()
cur.execute("PRAGMA data_version")
dv_before = cur.fetchone()[0]
cur.execute("BEGIN")
cur.execute("PRAGMA table_info(messages)")
available = {str(row['name']) for row in cur.fetchall()}
@@ -8654,6 +8662,9 @@ def get_state_db_regeneration_tail_snapshot(
active_clause = ""
if 'active' in available:
active_clause = " AND (active IS NULL OR active != 0)"
# durable order: id ASC when present (matches the canonical reader),
# else timestamp ASC
durable_order = 'id' if 'id' in available else 'timestamp'
# 1) prefix summary (rows with timestamp < floor)
cur.execute(
f"""
@@ -8670,28 +8681,19 @@ def get_state_db_regeneration_tail_snapshot(
"count": int(row["count"]) if row else 0,
"null_timestamp_count": int(row["null_timestamp_count"]) if row else 0,
}
# 2) ordered prefix keys (rows < floor)
# 2) ordered prefix keys (rows < floor) in durable order
prefix_key_cols = "COALESCE(role,'') AS role, COALESCE(content,'') AS content"
if 'tool_calls' in available:
prefix_key_cols += ", tool_calls"
if 'api_content' in available:
prefix_key_cols += ", api_content"
if 'id' in available:
prefix_key_sql = (
f"SELECT {prefix_key_cols} FROM messages "
"WHERE session_id = ? AND timestamp IS NOT NULL AND timestamp < ? "
f"{active_clause} ORDER BY timestamp ASC, id ASC"
)
prefix_params = (str(sid), floor_ts)
else:
prefix_key_sql = (
f"SELECT {prefix_key_cols} FROM messages "
"WHERE session_id = ? AND timestamp IS NOT NULL AND timestamp < ? "
f"{active_clause} ORDER BY timestamp ASC"
)
prefix_params = (str(sid), floor_ts)
prefix_key_sql = (
f"SELECT {prefix_key_cols} FROM messages "
"WHERE session_id = ? AND timestamp IS NOT NULL AND timestamp < ? "
f"{active_clause} ORDER BY {durable_order} ASC"
)
try:
cur.execute(prefix_key_sql, prefix_params)
cur.execute(prefix_key_sql, (str(sid), floor_ts))
except Exception:
cur.execute("ROLLBACK")
return None
@@ -8699,33 +8701,50 @@ def get_state_db_regeneration_tail_snapshot(
_session_message_visible_key({
"role": r["role"],
"content": r["content"],
"tool_calls": _json_loads_if_string(r["tool_calls"]) if "tool_calls" in r.keys() else None,
"tool_calls": _json_loads_if_string(r["tool_calls"]) if "tool_calls" in r.keys() and r["tool_calls"] is not None else None,
"api_content": r["api_content"] if "api_content" in r.keys() else None,
}, normalize_workspace_prefix=True)
for r in cur.fetchall()
]
# 3) bounded tail (rows >= floor) with the columns the reconciler uses
tail_select = ['role', 'content', 'timestamp']
for col in ('tool_call_id', 'tool_calls', 'tool_name', 'reasoning', 'api_content', 'active', 'id'):
if col in available:
# 3) bounded tail (rows >= floor) with the canonical projection
optional = [
'tool_call_id', 'tool_calls', 'tool_name', 'reasoning',
'reasoning_details', 'codex_reasoning_items', 'reasoning_content',
'codex_message_items', 'api_content',
]
tail_select = ['id', 'role', 'content', 'timestamp'] if 'id' in available else ['role', 'content', 'timestamp']
for col in optional + (['active'] if 'active' in available else []):
if col in available and col not in tail_select:
tail_select.append(col)
tail_sql = (
f"SELECT {', '.join(tail_select)} FROM messages "
f"WHERE session_id = ? AND (timestamp IS NULL OR timestamp >= ?) {active_clause} "
"ORDER BY timestamp ASC, id ASC"
f"ORDER BY {durable_order} ASC"
)
cur.execute(tail_sql, (str(sid), floor_ts))
tail_rows = [dict(r) for r in cur.fetchall()]
tail_rows = [
_project_state_db_message(r, available, 'id' in available, optional)
for r in cur.fetchall()
]
tail_keys = [
_session_message_visible_key({
"role": r.get("role"),
"content": r.get("content"),
"tool_calls": _json_loads_if_string(r["tool_calls"]) if r.get("tool_calls") is not None else None,
"api_content": r.get("api_content"),
}, normalize_workspace_prefix=True)
_session_message_visible_key(
{
"role": r.get("role"),
"content": r.get("content"),
"tool_calls": r.get("tool_calls"),
"api_content": r.get("api_content"),
},
normalize_workspace_prefix=True,
)
for r in tail_rows
]
cur.execute("COMMIT")
cur.execute("PRAGMA data_version")
dv_after = cur.fetchone()[0]
if dv_before != dv_after:
# a concurrent WAL commit landed during the proof → stale
# snapshot; refuse the bounded path (TOCTOU).
return None
return {
"prefix": prefix,
"prefix_keys": prefix_keys,
@@ -745,3 +745,127 @@ def test_regeneration_tail_snapshot_reads_real_sqlite(monkeypatch, tmp_path):
# the repeated prompt key appears in BOTH the skipped prefix and the tail —
# exactly the occurrence-count collision the guard must refuse
assert snap["prefix_keys"][0] == snap["tail_keys"][0]
def _make_state_db(tmp_path, rows, with_id=True):
"""Create a real state.db with the given rows (list of dicts)."""
import sqlite3
db = tmp_path / "state.db"
conn = sqlite3.connect(db)
cols = ["session_id TEXT", "role TEXT", "content TEXT", "timestamp REAL",
"tool_calls TEXT", "tool_name TEXT", "reasoning TEXT", "active INTEGER"]
if with_id:
cols.append("id INTEGER PRIMARY KEY")
conn.execute(f"CREATE TABLE messages ({', '.join(cols)})")
for i, r in enumerate(rows):
conn.execute(
"INSERT INTO messages (session_id, role, content, timestamp, tool_calls, tool_name, reasoning, active) "
"VALUES (?,?,?,?,?,?,?,1)",
("s1", r.get("role"), r.get("content"), r.get("timestamp"),
r.get("tool_calls"), r.get("tool_name"), r.get("reasoning")),
)
conn.commit()
conn.close()
return db
def test_snapshot_tail_uses_canonical_projection(monkeypatch, tmp_path):
"""#6826 r4 #1: the snapshot tail must use the SAME projection as the
canonical reader — JSON tool_calls decoded, tool_name → name, no raw
id/active exposure (tool-call/gateway-row sessions)."""
from api import models
db = _make_state_db(tmp_path, [
{"role": "user", "content": "p", "timestamp": 100.0},
{"role": "tool", "content": "result", "timestamp": 200.0,
"tool_calls": '[{"name": "x", "arguments": {"a": 1}}]',
"tool_name": "read_file", "reasoning": "r1"},
{"role": "assistant", "content": "a", "timestamp": 300.0},
])
monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
snap = models.get_state_db_regeneration_tail_snapshot("s1", 50.0)
assert snap is not None
tool = [m for m in snap["tail"] if m.get("role") == "tool"][0]
assert isinstance(tool["tool_calls"], list), "tool_calls must be JSON-decoded"
assert tool["name"] == "read_file", "tool_name → name must be applied"
assert "id" not in tool and "active" not in tool, "raw id/active must not leak"
assert "reasoning" in tool and tool["reasoning"] == "r1"
def test_snapshot_uses_durable_id_order(monkeypatch, tmp_path):
"""#6826 r4 #2: durable id ASC order (not timestamp) — a later-id user at
ts=600 followed by its assistant at ts=550 must stay user → assistant."""
from api import models
db = _make_state_db(tmp_path, [
{"role": "user", "content": "later user", "timestamp": 600.0},
{"role": "assistant", "content": "assistant 550", "timestamp": 550.0},
])
monkeypatch.setattr(models, "_active_state_db_path", lambda: db)
snap = models.get_state_db_regeneration_tail_snapshot("s1", 50.0)
assert snap is not None
roles = [m.get("role") for m in snap["tail"]]
assert roles == ["user", "assistant"], f"durable id order violated: {roles}"
class _FakeCursor:
def __init__(self, conn):
self._c = conn._real.cursor()
self._conn = conn
self._dv = None
def execute(self, sql, *args):
s = sql.strip()
if s.startswith("PRAGMA data_version"):
self._conn._dv_calls += 1
self._dv = 42 if self._conn._dv_calls == 1 else 43 # changed on 2nd read
return self
return self._c.execute(sql, *args)
def fetchone(self):
if self._dv is not None:
v, self._dv = self._dv, None
return (v,)
return self._c.fetchone()
def fetchall(self):
return self._c.fetchall()
class _FakeConn:
def __init__(self, real):
self._real = real
self._dv_calls = 0
self.row_factory = None
def cursor(self):
return _FakeCursor(self)
def close(self):
self._real.close()
def __enter__(self):
return self
def __exit__(self, *a):
self.close()
def test_snapshot_refuses_when_wal_data_version_changes(monkeypatch, tmp_path):
"""#6826 r4 #3 (WAL TOCTOU): if PRAGMA data_version changes during the
proof (a concurrent WAL commit), the snapshot must return None so the
caller falls back to the full read."""
import sqlite3
from api import models
db = _make_state_db(tmp_path, [
{"role": "user", "content": "p", "timestamp": 100.0},
{"role": "assistant", "content": "a", "timestamp": 200.0},
])
real_conn = sqlite3.connect(db)
fake = _FakeConn(real_conn)
monkeypatch.setattr(models, "open_state_db_readonly", lambda _p: fake)
snap = models.get_state_db_regeneration_tail_snapshot("s1", 50.0)
assert snap is None, "changed data_version must refuse the bounded snapshot"