diff --git a/api/models.py b/api/models.py index fa57cf68..ff5944e1 100644 --- a/api/models.py +++ b/api/models.py @@ -8598,6 +8598,144 @@ def get_state_db_session_message_keys_before_timestamp( return None +def get_state_db_regeneration_tail_snapshot( + sid, + floor, + *, + profile=None, +): + """Return prefix proof + bounded tail from ONE read transaction (#6826 r3). + + The regeneration guard must not suffer TOCTOU: the prefix summary, the + ordered prefix keys, and the bounded tail must come from the same SQLite + snapshot, otherwise a row inserted between the proof and the data read + can be silently omitted while the proof still authorizes the bounded path. + + Returns ``None`` when a stable single-connection snapshot cannot be + obtained (callers must fall back to the full read). Otherwise returns:: + + { + "prefix": {"count": N, "null_timestamp_count": M}, + "prefix_keys": [visible-key, ...], # rows < floor, db order + "tail": [message-dict, ...], # rows >= floor (bounded) + "tail_keys": [visible-key, ...], # rows >= floor, db order + } + """ + try: + import sqlite3 + except ImportError: + return None + + if not sid: + return None + try: + floor_ts = float(floor) + 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 {"prefix": {"count": 0, "null_timestamp_count": 0}, "prefix_keys": [], "tail": [], "tail_keys": []} + + try: + with closing(open_state_db_readonly(db_path)) as conn: + conn.row_factory = sqlite3.Row + cur = conn.cursor() + cur.execute("BEGIN") + cur.execute("PRAGMA table_info(messages)") + available = {str(row['name']) for row in cur.fetchall()} + if not {'session_id', 'role', 'content', 'timestamp'}.issubset(available): + cur.execute("ROLLBACK") + return None + active_clause = "" + if 'active' in available: + active_clause = " AND (active IS NULL OR active != 0)" + # 1) prefix summary (rows with timestamp < floor) + 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} + """, + (floor_ts, str(sid)), + ) + row = cur.fetchone() + prefix = { + "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) + 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) + try: + cur.execute(prefix_key_sql, prefix_params) + except Exception: + cur.execute("ROLLBACK") + return None + prefix_keys = [ + _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, + "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: + 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" + ) + cur.execute(tail_sql, (str(sid), floor_ts)) + tail_rows = [dict(r) 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) + for r in tail_rows + ] + cur.execute("COMMIT") + return { + "prefix": prefix, + "prefix_keys": prefix_keys, + "tail": tail_rows, + "tail_keys": tail_keys, + } + except Exception: + return None + + def get_state_db_session_summary(sid, *, profile=None) -> dict: """Return a cheap message count/timestamp summary for one state.db session.""" try: diff --git a/api/session_ops.py b/api/session_ops.py index 0b9554c3..6a19beb9 100644 --- a/api/session_ops.py +++ b/api/session_ops.py @@ -254,35 +254,34 @@ def _sidecar_regeneration_read_floor(session): return min(timestamps[-_REGENERATION_SIDECAR_ANCHOR_BUDGET:]) -def _bounded_tail_read_is_safe(session, read_floor) -> bool: - """Prove the skipped state.db prefix is represented identically in the - sidecar; otherwise the caller MUST fall back to the full read (#6826 r3). +def _bounded_tail_snapshot_if_safe(session, read_floor): + """Return the bounded tail rows ONLY when it is provably identical to the + full read; otherwise None (caller must fall back to the full read). - The bounded ``since_timestamp`` read silently omits every state.db row - older than the floor. That is only safe when the skipped prefix is - provably identical in the already-loaded sidecar — same count AND same - ordered visible identity — and when a compression anchor (if any) is - covered by the floor. Any mismatch, missing database, or uncertainty - returns False (full-read fallback), so the #6611 regeneration authority - never operates on an unreconciled view. + #6826 r3: the skipped state.db prefix (rows older than the floor) must be + represented identically in the sidecar — same count AND same ordered + visible identity — and the bounded tail must not repeat any skipped key + (occurrence-count collision: a new tail turn repeating an older prompt + would be mistaken for the old sidecar duplicate and dropped). The prefix + proof and the tail data come from ONE read transaction (no TOCTOU). + + Any mismatch, missing database, or uncertainty returns None, so the #6611 + regeneration authority never operates on an unreconciled view. """ sid = getattr(session, "session_id", None) if not sid: - return False + return None profile = getattr(session, "profile", None) from api.models import ( _session_message_visible_key, - get_state_db_session_message_keys_before_timestamp, - get_state_db_session_message_prefix_summary, + get_state_db_regeneration_tail_snapshot, ) - prefix = get_state_db_session_message_prefix_summary(sid, read_floor, profile=profile) - if prefix is None: - return False # cannot prove the prefix → full read - # Compression-anchor coverage: auto-compressed sessions keep context rows - # around the anchor; if the anchor predates the floor the bounded read can - # drop compacted-tail context rows (display may still match). The anchor - # is a dict {"role","text","ts","attachments"} — compare its timestamp. + snap = get_state_db_regeneration_tail_snapshot(sid, read_floor, profile=profile) + if snap is None: + return None # cannot obtain a stable single-snapshot → full read + # Compression-anchor coverage: if the anchor predates the floor the bounded + # read can drop compacted-tail context rows (display may still match). anchor = getattr(session, "compression_anchor_message_key", None) if isinstance(anchor, dict): try: @@ -290,16 +289,12 @@ def _bounded_tail_read_is_safe(session, read_floor) -> bool: except (TypeError, ValueError): anchor_ts = None if anchor_ts is None or anchor_ts < read_floor: - return False + return None + prefix = snap["prefix"] if prefix.get("count") == 0 and prefix.get("null_timestamp_count") == 0: # Empty skipped prefix: the bounded read already covers every row. - return True + return snap["tail"] # Non-empty skipped prefix: prove identical ordered visible identity. - db_keys = get_state_db_session_message_keys_before_timestamp( - sid, read_floor, profile=profile - ) - if db_keys is None: - return False sidecar_keys = [] for message in getattr(session, "messages", None) or []: if not isinstance(message, dict): @@ -311,11 +306,18 @@ def _bounded_tail_read_is_safe(session, read_floor) -> bool: if ts is not None and ts < read_floor: key = _session_message_visible_key(message) if key is None: - return False + return None sidecar_keys.append(key) - if list(db_keys) != sidecar_keys: - return False # mismatch → full read - return True + if list(snap["prefix_keys"]) != sidecar_keys: + return None # mismatch → full read + # Occurrence-count collision (#6826 r3 #1): if any bounded-tail key ALSO + # occurs in the skipped prefix, the reconciler may drop the repeated tail + # row — fall back conservatively. + prefix_key_set = set(snap["prefix_keys"]) + for key in snap["tail_keys"]: + if key in prefix_key_set: + return None + return snap["tail"] def regeneration_state(session, *, use_sidecar=False): @@ -329,25 +331,29 @@ def regeneration_state(session, *, use_sidecar=False): and gateway apply) is preserved on the fast path. The bounded tail is only trusted when - :func:`_bounded_tail_read_is_safe` proves the skipped state.db prefix is - identical in the sidecar (count + ordered visible identity + compression - anchor coverage); otherwise the read falls back to the full transcript. + :func:`_bounded_tail_snapshot_if_safe` proves the skipped state.db prefix + is identical in the sidecar (count + ordered visible identity + no + occurrence collision + compression anchor coverage), and the tail rows + come from the SAME single read transaction as the proof (no TOCTOU); + otherwise the read falls back to the full transcript. """ from api.models import ( get_state_db_session_messages, reconciled_state_db_messages_for_session, ) - state_reader_kwargs = {} + bounded_tail = None if use_sidecar: read_floor = _sidecar_regeneration_read_floor(session) - if read_floor is not None and _bounded_tail_read_is_safe(session, read_floor): - state_reader_kwargs["since_timestamp"] = read_floor - state_messages = get_state_db_session_messages( - getattr(session, "session_id", None), - profile=getattr(session, "profile", None), - **state_reader_kwargs, - ) + if read_floor is not None: + bounded_tail = _bounded_tail_snapshot_if_safe(session, read_floor) + if bounded_tail is not None: + state_messages = bounded_tail + else: + state_messages = get_state_db_session_messages( + getattr(session, "session_id", None), + profile=getattr(session, "profile", None), + ) return ( reconciled_state_db_messages_for_session( session, diff --git a/tests/test_issue6611_regeneration_authority.py b/tests/test_issue6611_regeneration_authority.py index bd5a023c..738bbc03 100644 --- a/tests/test_issue6611_regeneration_authority.py +++ b/tests/test_issue6611_regeneration_authority.py @@ -514,30 +514,32 @@ def test_regeneration_sidecar_path_avoids_full_transcript_refetch(monkeypatch): ) # The bounded guard must prove the skipped prefix (rows < floor) is - # identical in the sidecar: the fake state.db carries the same 300 prefix - # rows, so simulate the prefix summary + keys from the sidecar itself. + # identical in the sidecar and return the tail from ONE snapshot. The fake + # state.db carries the same 300 prefix rows as the sidecar. from api.models import _session_message_visible_key floor = min(row["timestamp"] for row in session.messages[-200:]) prefix_rows = [r for r in session.messages if (r.get("timestamp") or 0) < floor] + tail_rows = [r for r in state_rows if (r.get("timestamp") or 0) >= floor] monkeypatch.setattr( models_api, - "get_state_db_session_message_prefix_summary", - lambda *_a, **_k: {"count": len(prefix_rows), "null_timestamp_count": 0}, - ) - monkeypatch.setattr( - models_api, - "get_state_db_session_message_keys_before_timestamp", - lambda *_a, **_k: [_session_message_visible_key(r) for r in prefix_rows], + "get_state_db_regeneration_tail_snapshot", + lambda *_a, **_k: { + "prefix": {"count": len(prefix_rows), "null_timestamp_count": 0}, + "prefix_keys": [_session_message_visible_key(r) for r in prefix_rows], + "tail": [dict(r) for r in tail_rows], + "tail_keys": [_session_message_visible_key(r) for r in tail_rows], + }, ) full_rows, full_context = regeneration_state(session) assert calls[-1].get("since_timestamp") is None # full read + # provably effective: the fast path does NOT consult the full reader — the + # bounded tail comes from the single snapshot + calls_before = len(calls) fast_rows, fast_context = regeneration_state(session, use_sidecar=True) - # provably effective: the fast path requests only the bounded tip window - assert calls[-1].get("since_timestamp") is not None - assert calls[-1]["since_timestamp"] > 0 + assert len(calls) == calls_before, "bounded path must not touch the full reader" # authority preserved: identical reconciled transcript on both paths, # including the state.db-only gateway tail row. @@ -566,25 +568,33 @@ def _sidecar_session(rows=3, *, anchor_key=None, anchor_ts=None): return s +def _snap(prefix_count=0, prefix_keys=None, tail=None, tail_keys=None): + """Factory for a fake get_state_db_regeneration_tail_snapshot result.""" + tail = tail if tail is not None else [] + return { + "prefix": {"count": prefix_count, "null_timestamp_count": 0}, + "prefix_keys": prefix_keys or [], + "tail": tail, + "tail_keys": tail_keys if tail_keys is not None else [], + } + + def test_bounded_guard_empty_prefix_allows_tail_read(monkeypatch): from api import session_ops session = _sidecar_session(3) monkeypatch.setattr( - "api.models.get_state_db_session_message_prefix_summary", - lambda *_a, **_k: {"count": 0, "null_timestamp_count": 0}, + "api.models.get_state_db_regeneration_tail_snapshot", + lambda *_a, **_k: _snap(tail=[{"role": "user", "content": "t", "timestamp": 300.0}]), ) seen = {} - monkeypatch.setattr( - "api.models.get_state_db_session_message_keys_before_timestamp", - lambda *_a, **_k: [], - ) monkeypatch.setattr( "api.models.get_state_db_session_messages", lambda *_a, **_k: _capture(seen, _k), ) session_ops.regeneration_state(session, use_sidecar=True) - assert seen.get("since_timestamp") is not None, "empty skipped prefix must take the bounded tail read" + # bounded path: full reader NOT consulted (tail comes from the snapshot) + assert seen == {}, "empty skipped prefix must take the bounded tail from the snapshot" def test_bounded_guard_unrepresented_prefix_row_falls_back(monkeypatch): @@ -593,12 +603,11 @@ def test_bounded_guard_unrepresented_prefix_row_falls_back(monkeypatch): session = _sidecar_session(3) # state.db has one row older than the floor that the sidecar does NOT carry monkeypatch.setattr( - "api.models.get_state_db_session_message_prefix_summary", - lambda *_a, **_k: {"count": 1, "null_timestamp_count": 0}, - ) - monkeypatch.setattr( - "api.models.get_state_db_session_message_keys_before_timestamp", - lambda *_a, **_k: [("user", "recoverable-only-in-db")], + "api.models.get_state_db_regeneration_tail_snapshot", + lambda *_a, **_k: _snap( + prefix_count=1, + prefix_keys=[("user", "recoverable-only-in-db")], + ), ) seen = {} monkeypatch.setattr( @@ -606,7 +615,33 @@ def test_bounded_guard_unrepresented_prefix_row_falls_back(monkeypatch): lambda *_a, **_k: _capture(seen, _k), ) session_ops.regeneration_state(session, use_sidecar=True) - assert "since_timestamp" not in seen, "unrepresented prefix row must force the full read (no since_timestamp)" + assert seen != {}, "unrepresented prefix row must force the full read (full reader consulted)" + + +def test_bounded_guard_repeated_tail_key_falls_back(monkeypatch): + """#6826 r3 #1: a bounded-tail key that ALSO occurs in the skipped prefix + must force the full read (occurrence-count collision would drop the + repeated tail row).""" + from api import session_ops + + session = _sidecar_session(3) + repeated_key = ("user", "m0") # m0 lives below the floor AND repeats in the tail + monkeypatch.setattr( + "api.models.get_state_db_regeneration_tail_snapshot", + lambda *_a, **_k: _snap( + prefix_count=1, + prefix_keys=[repeated_key], + tail=[{"role": "user", "content": "m0", "timestamp": 300.0}], + tail_keys=[repeated_key], + ), + ) + seen = {} + monkeypatch.setattr( + "api.models.get_state_db_session_messages", + lambda *_a, **_k: _capture(seen, _k), + ) + session_ops.regeneration_state(session, use_sidecar=True) + assert seen != {}, "repeated tail key must force the full read" def test_bounded_guard_compression_anchor_below_floor_falls_back(monkeypatch): @@ -616,12 +651,8 @@ def test_bounded_guard_compression_anchor_below_floor_falls_back(monkeypatch): # anchor timestamp predates the floor → bounded read cannot cover it session.compression_anchor_message_key = {"role": "user", "text": "m0", "ts": 100.0} monkeypatch.setattr( - "api.models.get_state_db_session_message_prefix_summary", - lambda *_a, **_k: {"count": 0, "null_timestamp_count": 0}, - ) - monkeypatch.setattr( - "api.models.get_state_db_session_message_keys_before_timestamp", - lambda *_a, **_k: [], + "api.models.get_state_db_regeneration_tail_snapshot", + lambda *_a, **_k: _snap(), ) # shrink the anchor budget window so the floor (300.0) sits above the anchor monkeypatch.setattr(session_ops, "_REGENERATION_SIDECAR_ANCHOR_BUDGET", 1) @@ -631,7 +662,7 @@ def test_bounded_guard_compression_anchor_below_floor_falls_back(monkeypatch): lambda *_a, **_k: _capture(seen, _k), ) session_ops.regeneration_state(session, use_sidecar=True) - assert "since_timestamp" not in seen, "compression anchor below floor must force the full read" + assert seen != {}, "compression anchor below floor must force the full read" def test_bounded_guard_compression_anchor_covered_allows(monkeypatch): @@ -641,12 +672,8 @@ def test_bounded_guard_compression_anchor_covered_allows(monkeypatch): session.compression_anchor_message_key = {"role": "user", "text": "m0", "ts": 100.0} # budget 200 covers all 3 rows → floor = min(all) = 100 → anchor at floor is covered monkeypatch.setattr( - "api.models.get_state_db_session_message_prefix_summary", - lambda *_a, **_k: {"count": 0, "null_timestamp_count": 0}, - ) - monkeypatch.setattr( - "api.models.get_state_db_session_message_keys_before_timestamp", - lambda *_a, **_k: [], + "api.models.get_state_db_regeneration_tail_snapshot", + lambda *_a, **_k: _snap(tail=[{"role": "user", "content": "t", "timestamp": 300.0}]), ) seen = {} monkeypatch.setattr( @@ -654,9 +681,67 @@ def test_bounded_guard_compression_anchor_covered_allows(monkeypatch): lambda *_a, **_k: _capture(seen, _k), ) session_ops.regeneration_state(session, use_sidecar=True) - assert seen.get("since_timestamp") is not None, "anchor covered by floor must keep the bounded tail read" + assert seen == {}, "anchor covered by floor must keep the bounded tail read" + + +def test_bounded_guard_missing_snapshot_falls_back(monkeypatch): + """#6826 r3 #2 (TOCTOU): if a stable single-connection snapshot cannot be + obtained, the guard must refuse the bounded path entirely.""" + from api import session_ops + + session = _sidecar_session(3) + monkeypatch.setattr( + "api.models.get_state_db_regeneration_tail_snapshot", + lambda *_a, **_k: None, + ) + seen = {} + monkeypatch.setattr( + "api.models.get_state_db_session_messages", + lambda *_a, **_k: _capture(seen, _k), + ) + session_ops.regeneration_state(session, use_sidecar=True) + assert seen != {}, "missing snapshot must force the full read (no TOCTOU window)" def _capture(seen, kwargs): seen.update(kwargs) return [] + + +def test_regeneration_tail_snapshot_reads_real_sqlite(monkeypatch, tmp_path): + """Real-SQLite: the single-connection snapshot returns prefix proof + tail + from one read, including the repeated-prompt collision shape from the + round-3 review (a prompt that lives both below and above the floor).""" + import sqlite3 + + from api import models + + db = tmp_path / "state.db" + conn = sqlite3.connect(db) + conn.execute( + "CREATE TABLE messages (id INTEGER PRIMARY KEY, session_id TEXT, role TEXT, " + "content TEXT, timestamp REAL, tool_calls TEXT, active INTEGER)" + ) + rows = [ + ("s1", "user", "prompt", 100.0), + ("s1", "assistant", "a1", 200.0), + ("s1", "user", "prompt", 300.0), # repeated prompt above the floor + ("s1", "assistant", "a2", 400.0), + ] + for sid, role, content, ts in rows: + conn.execute( + "INSERT INTO messages (session_id, role, content, timestamp, active) VALUES (?,?,?,?,1)", + (sid, role, content, ts), + ) + conn.commit() + conn.close() + + monkeypatch.setattr(models, "_active_state_db_path", lambda: db) + snap = models.get_state_db_regeneration_tail_snapshot("s1", 250.0) + assert snap is not None + assert snap["prefix"]["count"] == 2, "rows < floor must be counted" + assert len(snap["prefix_keys"]) == 2 + assert len(snap["tail"]) == 2, "rows >= floor must be the bounded tail" + # 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]