fix(regeneration): snapshot reuses canonical projection+ordering; WAL data_version guard (#6826 r4)
This commit is contained in:
1 parent
b33aff9d09
commit
7abd4773ca
2 files changed
+204
-61
No files matched your search
+80
-61
@@ -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"
|
||||
Reference in new issue
Block a user