From 341edca856624492461bf395aed33ac4fb6689a8 Mon Sep 17 00:00:00 2001 From: pufit Date: Mon, 6 Jul 2026 10:41:18 -0400 Subject: [PATCH] Make DB writes rollback-safe; heal leaked transactions A write task abandoned mid-transaction leaves the shared aiosqlite connection pinned to a stale WAL snapshot; after any external commit, every write fails instantly with SQLITE_BUSY_SNAPSHOT ('database is locked', busy_timeout not consulted) until ROLLBACK - which nothing ever called. Wedged a production daemon for 10h on 2026-07-06. - _atomic(): commit on success, rollback on any error incl. cancellation - _write(): single-statement writes serialized under the write lock - _heal_leaked_txn(): recover from leaked transactions on any write - migrate all 52 unlocked execute+commit sites; fix bare commits in checkpoint()/vacuum() - regression tests incl. the poisoned-connection reproduction --- nerve/db/audit.py | 8 +- nerve/db/base.py | 144 ++++++++++++++++++---- nerve/db/cron.py | 21 ++-- nerve/db/files.py | 6 +- nerve/db/maintenance.py | 20 +-- nerve/db/mcp.py | 6 +- nerve/db/messages.py | 12 +- nerve/db/notifications.py | 14 +-- nerve/db/plans.py | 6 +- nerve/db/sessions.py | 65 +++++----- nerve/db/skills.py | 15 +-- nerve/db/sources.py | 46 +++---- nerve/db/task_statuses.py | 6 +- nerve/db/tasks.py | 12 +- nerve/db/usage.py | 6 +- nerve/db/wakeups.py | 14 +-- tests/test_db.py | 5 + tests/test_db_atomicity.py | 243 +++++++++++++++++++++++++++++++++++++ 18 files changed, 475 insertions(+), 174 deletions(-) create mode 100644 tests/test_db_atomicity.py diff --git a/nerve/db/audit.py b/nerve/db/audit.py index 4b13969..53d5670 100644 --- a/nerve/db/audit.py +++ b/nerve/db/audit.py @@ -19,14 +19,12 @@ async def log_audit( ) -> int: """Write an entry to the memU audit log.""" now = datetime.now(timezone.utc).isoformat() - async with self.db.execute( + result = await self._write( "INSERT INTO memu_audit_log (timestamp, action, target_type, target_id, source, details) " "VALUES (?, ?, ?, ?, ?, ?)", (now, action, target_type, target_id, source, json.dumps(details) if details else None), - ) as cursor: - log_id = cursor.lastrowid - await self.db.commit() - return log_id + ) + return result.lastrowid async def get_audit_logs( self, diff --git a/nerve/db/base.py b/nerve/db/base.py index 7124ef5..c79d153 100644 --- a/nerve/db/base.py +++ b/nerve/db/base.py @@ -11,7 +11,7 @@ import logging from contextlib import asynccontextmanager from pathlib import Path -from typing import AsyncIterator +from typing import AsyncIterator, NamedTuple import aiosqlite @@ -39,6 +39,18 @@ SCHEMA_VERSION = max(v for v, _ in discover_migrations()) if discover_migrations() else 0 +class WriteResult(NamedTuple): + """Outcome of a single-statement write executed via :meth:`Database._write`. + + Mirrors the two cursor attributes writer call sites actually use. The + live cursor is closed before the commit, so callers can never fetch from + a cursor whose transaction has already ended. + """ + + lastrowid: int | None + rowcount: int + + # Connection pragmas applied on every ``connect()``. These mirror the tuning # memU already uses for its own SQLite connections (see # ``nerve/memory/memu_bridge.py``) — the primary operational DB is where the @@ -141,16 +153,98 @@ def db(self) -> aiosqlite.Connection: raise RuntimeError("Database not connected. Call connect() first.") return self._db + # -- Write path --------------------------------------------------------- + # + # All writes on the single shared connection MUST go through ``_atomic()`` + # (multi-statement transactions) or ``_write()`` (single statements). + # Both serialize under ``_write_lock`` and guarantee the connection is + # never left inside an open transaction — on success they COMMIT, on any + # error (including ``asyncio.CancelledError``) they ROLLBACK. + # + # Why this is load-bearing (production outage, 2026-07-06): a write task + # was abandoned mid-transaction, leaving the shared connection inside an + # open transaction pinned to a WAL read snapshot. A second process then + # committed to the same DB file, making that snapshot stale. From that + # moment every write on the shared connection failed *instantly* with + # "database is locked" (SQLITE_BUSY_SNAPSHOT — the busy handler is + # deliberately not invoked for snapshot conflicts, so ``busy_timeout`` + # does not apply), reads silently served the frozen snapshot, and nothing + # ever called ROLLBACK — wedging the daemon for 10 hours until a restart. + # ``_heal_leaked_txn()`` is the belt-and-suspenders guard that recovers + # from that state even if some future code path leaks a transaction. + @asynccontextmanager async def _atomic(self) -> AsyncIterator[None]: - """Acquire write lock for multi-statement transactions. + """Serialize a multi-statement write and make it a real transaction. + + Once a coroutine begins a multi-statement write, no other coroutine + can interleave writes before the commit. The body's statements are + committed on success and rolled back on any exception — including + task cancellation — so a failed body neither half-commits nor leaves + the shared connection inside an open (poisoned) transaction. + + Statements inside the body must use ``self.db.execute(...)`` directly + (never ``_write()``, which would deadlock on the non-reentrant lock). + """ + async with self._write_lock: + await self._heal_leaked_txn() + try: + yield + # Shield so a cancellation arriving mid-commit cannot abandon + # a half-finished transaction: the inner task runs to + # completion on aiosqlite's worker thread regardless. + await asyncio.shield(self.db.commit()) + except BaseException: + await self._rollback_quietly() + raise - Ensures that once a coroutine begins a multi-statement write, - no other coroutine can interleave writes before the commit. + async def _write(self, sql: str, params: tuple | list = ()) -> WriteResult: + """Execute one write statement and commit, under the write lock. + + The single-statement counterpart of :meth:`_atomic` with the same + guarantees: serialized against all other writers (so its commit can + never flush someone else's in-flight transaction) and commit-or- + rollback semantics (so an error or cancellation can never leave the + connection mid-transaction). """ async with self._write_lock: - yield - await self.db.commit() + await self._heal_leaked_txn() + try: + cursor = await self.db.execute(sql, params) + result = WriteResult(cursor.lastrowid, cursor.rowcount) + await cursor.close() + await asyncio.shield(self.db.commit()) + return result + except BaseException: + await self._rollback_quietly() + raise + + async def _heal_leaked_txn(self) -> None: + """Roll back a leaked open transaction on the shared connection. + + Called under ``_write_lock``. Every legitimate transaction commits or + rolls back before releasing the lock, so ``in_transaction`` being true + here means some code path abandoned a transaction (see the write-path + comment above for the outage this causes). Recover loudly. + """ + if self.db.in_transaction: + logger.error( + "Leaked open transaction detected on the shared connection — " + "rolling back to prevent a wedged write path " + "(SQLITE_BUSY_SNAPSHOT poisoning)", + ) + await self._rollback_quietly() + + async def _rollback_quietly(self) -> None: + """Best-effort ROLLBACK that never raises and survives cancellation.""" + try: + await asyncio.shield(self.db.rollback()) + except asyncio.CancelledError: + # The shielded rollback still runs to completion on the worker + # thread; re-raise so the caller's cancellation proceeds. + raise + except Exception: + logger.exception("Rollback failed on the shared connection") async def _check_fts_integrity(self) -> None: """FTS integrity check — runs every startup. @@ -167,27 +261,29 @@ async def _check_fts_integrity(self) -> None: "FTS index mismatch: %d tasks vs %d FTS entries — reseeding", task_count, fts_count, ) - await self.db.execute("DELETE FROM tasks_fts") # Read content from disk files (source of truth) instead of seeding # empty. Task file_path values are relative to the workspace root, # which is NOT the DB directory (~/.nerve) — fall back to it only # when no workspace was provided. workspace = (self.workspace or self.db_path.parent).expanduser() - async with self.db.execute("SELECT id, title, file_path FROM tasks") as cur: - rows = await cur.fetchall() - for row in rows: - content = "" - try: - fp = workspace / row["file_path"] - if fp.exists(): - content = await asyncio.to_thread( - fp.read_text, encoding="utf-8", - ) - except Exception as e: - logger.warning("Failed to read %s for FTS reseed: %s", row["file_path"], e) - await self.db.execute( - "INSERT INTO tasks_fts (task_id, title, content) VALUES (?, ?, ?)", - (row["id"], row["title"], content), - ) - await self.db.commit() + async with self._atomic(): + await self.db.execute("DELETE FROM tasks_fts") + async with self.db.execute( + "SELECT id, title, file_path FROM tasks", + ) as cur: + rows = await cur.fetchall() + for row in rows: + content = "" + try: + fp = workspace / row["file_path"] + if fp.exists(): + content = await asyncio.to_thread( + fp.read_text, encoding="utf-8", + ) + except Exception as e: + logger.warning("Failed to read %s for FTS reseed: %s", row["file_path"], e) + await self.db.execute( + "INSERT INTO tasks_fts (task_id, title, content) VALUES (?, ?, ?)", + (row["id"], row["title"], content), + ) logger.info("FTS reseeded with %d tasks (content from disk)", task_count) diff --git a/nerve/db/cron.py b/nerve/db/cron.py index 6a9d204..538210b 100644 --- a/nerve/db/cron.py +++ b/nerve/db/cron.py @@ -9,12 +9,10 @@ class CronStore: """Mixin providing cron job logging operations.""" async def log_cron_start(self, job_id: str) -> int: - async with self.db.execute( + result = await self._write( "INSERT INTO cron_logs (job_id) VALUES (?)", (job_id,) - ) as cursor: - log_id = cursor.lastrowid - await self.db.commit() - return log_id + ) + return result.lastrowid async def set_cron_log_session(self, log_id: int, session_id: str) -> None: """Link a run log to its session while the run is still in flight. @@ -22,11 +20,10 @@ async def set_cron_log_session(self, log_id: int, session_id: str) -> None: Written at run start so the UI can open the chat of a running cron; log_cron_finish keeps it as a fallback for older code paths. """ - await self.db.execute( + await self._write( "UPDATE cron_logs SET session_id = ? WHERE id = ?", (session_id, log_id), ) - await self.db.commit() async def log_cron_finish( self, @@ -37,12 +34,11 @@ async def log_cron_finish( session_id: str | None = None, ) -> None: now = utc_now_iso() - await self.db.execute( + await self._write( "UPDATE cron_logs SET finished_at = ?, status = ?, output = ?, error = ?, " "session_id = COALESCE(?, session_id) WHERE id = ?", (now, status, output, error, session_id, log_id), ) - await self.db.commit() async def get_cron_logs( self, job_id: str | None = None, limit: int = 50, offset: int = 0, @@ -122,11 +118,8 @@ async def cleanup_old_cron_logs(self, days: int = 14) -> int: thousands of rows and the unfiltered "latest N" query (used by the diagnostics endpoint) degrades to a full-table scan + memory sort. """ - cursor = await self.db.execute( + result = await self._write( "DELETE FROM cron_logs WHERE started_at < datetime('now', ? || ' days')", (f"-{days}",), ) - deleted = cursor.rowcount or 0 - await cursor.close() - await self.db.commit() - return deleted + return result.rowcount or 0 diff --git a/nerve/db/files.py b/nerve/db/files.py index 66df7b5..6f2e26e 100644 --- a/nerve/db/files.py +++ b/nerve/db/files.py @@ -20,12 +20,11 @@ async def save_uploaded_file( disk_path: str, ) -> None: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT INTO uploaded_files (id, session_id, filename, media_type, file_type, file_size, disk_path, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", (file_id, session_id, filename, media_type, file_type, file_size, disk_path, now), ) - await self.db.commit() async def get_uploaded_file(self, file_id: str) -> dict | None: async with self.db.execute( @@ -51,9 +50,8 @@ async def delete_uploaded_files(self, session_id: str) -> list[str]: (session_id,), ) as cursor: paths = [row[0] async for row in cursor] - await self.db.execute( + await self._write( "DELETE FROM uploaded_files WHERE session_id = ?", (session_id,), ) - await self.db.commit() return paths diff --git a/nerve/db/maintenance.py b/nerve/db/maintenance.py index d68d6a7..7bb4d3a 100644 --- a/nerve/db/maintenance.py +++ b/nerve/db/maintenance.py @@ -165,20 +165,24 @@ async def checkpoint(self) -> None: """Truncate the WAL after a prune pass (best-effort). Frees the WAL file; does not shrink the main DB (see :meth:`vacuum`). + Runs under the write lock: an unlocked bare ``commit()`` here could + prematurely flush another coroutine's in-flight transaction. """ - await self.db.commit() - await self.db.execute("PRAGMA wal_checkpoint(TRUNCATE)") + async with self._write_lock: + await self._heal_leaked_txn() + await self.db.execute("PRAGMA wal_checkpoint(TRUNCATE)") async def vacuum(self) -> None: """Rewrite the DB file to reclaim freelist pages (shrinks the file). - Takes a write lock and cannot run inside a transaction, so it is an - explicit operator step, never on the background loop. Run with the - daemon stopped to avoid lock contention. + Cannot run inside a transaction (``_heal_leaked_txn`` guarantees the + connection is clean) and autocommits on completion. Serialized under + the write lock; still an explicit operator step, never on the + background loop. Run with the daemon stopped to avoid lock contention. """ - await self.db.commit() - await self.db.execute("VACUUM") - await self.db.commit() + async with self._write_lock: + await self._heal_leaked_txn() + await self.db.execute("VACUUM") async def run_retention( self, diff --git a/nerve/db/mcp.py b/nerve/db/mcp.py index ea20041..e2f9363 100644 --- a/nerve/db/mcp.py +++ b/nerve/db/mcp.py @@ -17,7 +17,7 @@ async def upsert_mcp_server( ) -> None: """Register or update an MCP server entry.""" now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT INTO mcp_servers (name, type, enabled, tool_count, first_seen_at, last_seen_at) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(name) DO UPDATE SET @@ -29,7 +29,6 @@ async def upsert_mcp_server( last_seen_at = excluded.last_seen_at""", (name, server_type, enabled, tool_count, now, now), ) - await self.db.commit() async def record_mcp_tool_usage( self, @@ -41,13 +40,12 @@ async def record_mcp_tool_usage( error: str | None = None, ) -> None: """Log an MCP tool invocation.""" - await self.db.execute( + await self._write( """INSERT INTO mcp_tool_usage (server_name, tool_name, session_id, duration_ms, success, error) VALUES (?, ?, ?, ?, ?, ?)""", (server_name, tool_name, session_id, duration_ms, success, error), ) - await self.db.commit() async def get_mcp_server_stats(self) -> list[dict]: """List all MCP servers with aggregated usage stats.""" diff --git a/nerve/db/messages.py b/nerve/db/messages.py index 11fe6b4..a1d062f 100644 --- a/nerve/db/messages.py +++ b/nerve/db/messages.py @@ -160,11 +160,10 @@ async def merge_tool_result_into_call( break if not updated: return None - await self.db.execute( + await self._write( "UPDATE messages SET blocks = ? WHERE id = ?", (json.dumps(blocks), msg_id), ) - await self.db.commit() return msg_id async def merge_workflow_into_call( @@ -211,11 +210,10 @@ async def merge_workflow_into_call( updated = True break if updated: - await self.db.execute( + await self._write( "UPDATE messages SET blocks = ? WHERE id = ?", (json.dumps(blocks), row["id"]), ) - await self.db.commit() return row["id"] return None @@ -242,13 +240,12 @@ async def save_file_snapshot( Uses INSERT OR IGNORE so only the first touch per session+file is stored. """ now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT OR IGNORE INTO session_file_snapshots (session_id, file_path, original_content, created_at) VALUES (?, ?, ?, ?)""", (session_id, file_path, content, now), ) - await self.db.commit() async def get_file_snapshot( self, session_id: str, file_path: str, @@ -272,11 +269,10 @@ async def get_session_snapshots(self, session_id: str) -> list[dict]: async def delete_session_snapshots(self, session_id: str) -> None: """Delete all file snapshots for a session.""" - await self.db.execute( + await self._write( "DELETE FROM session_file_snapshots WHERE session_id = ?", (session_id,), ) - await self.db.commit() async def count_messages(self, session_id: str) -> int: async with self.db.execute( diff --git a/nerve/db/notifications.py b/nerve/db/notifications.py index 3415ddc..6204bc8 100644 --- a/nerve/db/notifications.py +++ b/nerve/db/notifications.py @@ -38,7 +38,7 @@ async def create_notification( is persisted for audit but never fanned out to channels. """ now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT INTO notifications (id, session_id, type, title, body, priority, status, options, expires_at, metadata, created_at, target_kind, target_id) @@ -49,7 +49,6 @@ async def create_notification( expires_at, json.dumps(metadata or {}), now, target_kind, target_id), ) - await self.db.commit() return {"id": notification_id, "session_id": session_id, "type": type} async def get_notification(self, notification_id: str) -> dict | None: @@ -131,11 +130,10 @@ async def dismiss_notification(self, notification_id: str) -> bool: async def dismiss_all_notifications(self) -> int: """Dismiss all pending non-question notifications. Returns count dismissed.""" - cursor = await self.db.execute( + result = await self._write( "UPDATE notifications SET status = 'dismissed' WHERE status = 'pending' AND type = 'notify'", ) - await self.db.commit() - return cursor.rowcount + return result.rowcount async def expire_due_notifications(self) -> list[dict]: """Flip pending rows past their expiry to ``expired``. @@ -287,10 +285,9 @@ async def update_notification(self, notification_id: str, **fields) -> None: sets = ", ".join(f"{k} = ?" for k in fields) vals = list(fields.values()) vals.append(notification_id) - await self.db.execute( + await self._write( f"UPDATE notifications SET {sets} WHERE id = ?", tuple(vals), ) - await self.db.commit() # ------------------------------------------------------------------ # # Notification silences (deterministic suppression rules) # @@ -312,13 +309,12 @@ async def create_silence( permanent. Returns the freshly-inserted row. """ now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT INTO notification_silences (id, pattern, action, reason, created_by, created_at, expires_at) VALUES (?, ?, ?, ?, ?, ?, ?)""", (silence_id, pattern, action, reason, created_by, now, expires_at), ) - await self.db.commit() row = await self.get_silence(silence_id) return row or {"id": silence_id, "pattern": pattern} diff --git a/nerve/db/plans.py b/nerve/db/plans.py index a59d597..a2150d0 100644 --- a/nerve/db/plans.py +++ b/nerve/db/plans.py @@ -20,12 +20,11 @@ async def create_plan( plan_type: str = "generic", ) -> dict: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT INTO plans (id, task_id, session_id, content, model, version, parent_plan_id, plan_type, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", (plan_id, task_id, session_id, content, model, version, parent_plan_id, plan_type, now), ) - await self.db.commit() return {"id": plan_id, "task_id": task_id, "version": version, "plan_type": plan_type} async def get_plan(self, plan_id: str) -> dict | None: @@ -66,10 +65,9 @@ async def update_plan(self, plan_id: str, **fields) -> None: sets = ", ".join(f"{k} = ?" for k in fields) vals = list(fields.values()) vals.append(plan_id) - await self.db.execute( + await self._write( f"UPDATE plans SET {sets} WHERE id = ?", tuple(vals), ) - await self.db.commit() async def get_plans_for_task(self, task_id: str) -> list[dict]: async with self.db.execute( diff --git a/nerve/db/sessions.py b/nerve/db/sessions.py index ba29804..d951e6f 100644 --- a/nerve/db/sessions.py +++ b/nerve/db/sessions.py @@ -20,7 +20,7 @@ async def create_session( forked_from_message: str | None = None, ) -> dict: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT OR IGNORE INTO sessions (id, title, source, metadata, status, parent_session_id, forked_from_message, created_at, updated_at) @@ -29,7 +29,6 @@ async def create_session( json.dumps(metadata or {}), status, parent_session_id, forked_from_message, now, now), ) - await self.db.commit() return { "id": session_id, "title": title or session_id, "source": source, "status": status, @@ -77,16 +76,14 @@ async def search_sessions(self, query: str, limit: int = 100) -> list[dict]: async def touch_session(self, session_id: str) -> None: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( "UPDATE sessions SET updated_at = ? WHERE id = ?", (now, session_id) ) - await self.db.commit() async def update_session_title(self, session_id: str, title: str) -> None: - await self.db.execute( + await self._write( "UPDATE sessions SET title = ? WHERE id = ?", (title, session_id) ) - await self.db.commit() async def delete_session(self, session_id: str) -> None: async with self._atomic(): @@ -102,23 +99,33 @@ async def update_session_metadata(self, session_id: str, metadata: dict) -> None Also syncs dedicated columns (sdk_session_id, connected_at) for backward compatibility with callers that still use the metadata blob. + Both updates run in one transaction (the column sync is inlined — + calling :meth:`update_session_fields` here would deadlock on the + non-reentrant write lock). """ - await self.db.execute( - "UPDATE sessions SET metadata = ? WHERE id = ?", - (json.dumps(metadata), session_id), - ) - # Sync dedicated columns from metadata col_sync: dict[str, str | None] = {} if "sdk_session_id" in metadata: col_sync["sdk_session_id"] = metadata["sdk_session_id"] if "connected_at" in metadata: col_sync["connected_at"] = metadata["connected_at"] - if col_sync: - await self.update_session_fields(session_id, col_sync) - await self.db.commit() + built = self._build_session_fields_update(session_id, col_sync) + async with self._atomic(): + await self.db.execute( + "UPDATE sessions SET metadata = ? WHERE id = ?", + (json.dumps(metadata), session_id), + ) + if built is not None: + await self.db.execute(*built) - async def update_session_fields(self, session_id: str, fields: dict) -> None: - """Update specific session columns atomically. Merges, doesn't replace.""" + @staticmethod + def _build_session_fields_update( + session_id: str, fields: dict, + ) -> tuple[str, tuple] | None: + """Build the (sql, params) for a whitelisted session-columns UPDATE. + + Returns ``None`` when no allowed field is present. Shared by + :meth:`update_session_fields` and :meth:`update_session_metadata`. + """ allowed = { "status", "sdk_session_id", "connected_at", "last_activity_at", "archived_at", "title", "message_count", "total_cost_usd", @@ -132,15 +139,21 @@ async def update_session_fields(self, session_id: str, fields: dict) -> None: set_clauses.append(f"{key} = ?") params.append(value) if not set_clauses: - return + return None set_clauses.append("updated_at = ?") params.append(datetime.now(timezone.utc).isoformat()) params.append(session_id) - await self.db.execute( + return ( f"UPDATE sessions SET {', '.join(set_clauses)} WHERE id = ?", tuple(params), ) - await self.db.commit() + + async def update_session_fields(self, session_id: str, fields: dict) -> None: + """Update specific session columns atomically. Merges, doesn't replace.""" + built = self._build_session_fields_update(session_id, fields) + if built is None: + return + await self._write(*built) async def get_sessions_with_metadata_key(self, key: str) -> list[dict]: """Find sessions whose metadata JSON contains a specific key. @@ -171,13 +184,11 @@ async def log_session_event( ) -> int: """Log a session lifecycle event.""" now = datetime.now(timezone.utc).isoformat() - async with self.db.execute( + result = await self._write( "INSERT INTO session_events (session_id, event_type, details, created_at) VALUES (?, ?, ?, ?)", (session_id, event_type, json.dumps(details) if details else None, now), - ) as cursor: - event_id = cursor.lastrowid - await self.db.commit() - return event_id + ) + return result.lastrowid async def get_session_events( self, session_id: str, limit: int = 50, @@ -209,11 +220,10 @@ async def get_channel_session(self, channel_key: str) -> dict | None: async def set_channel_session(self, channel_key: str, session_id: str) -> None: """Persist a channel-to-session mapping.""" now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( "INSERT OR REPLACE INTO channel_sessions (channel_key, session_id, updated_at) VALUES (?, ?, ?)", (channel_key, session_id, now), ) - await self.db.commit() # --- Session cleanup queries (V3) --- @@ -284,11 +294,10 @@ async def get_oldest_sessions( async def increment_message_count(self, session_id: str) -> None: """Atomically increment the message counter for a session.""" - await self.db.execute( + await self._write( "UPDATE sessions SET message_count = COALESCE(message_count, 0) + 1 WHERE id = ?", (session_id,), ) - await self.db.commit() async def get_sessions_needing_memorization(self) -> list[dict]: """Find non-archived sessions that have un-memorized messages. diff --git a/nerve/db/skills.py b/nerve/db/skills.py index 2529519..7371df1 100644 --- a/nerve/db/skills.py +++ b/nerve/db/skills.py @@ -23,7 +23,7 @@ async def upsert_skill( ) -> None: """Insert or update a skill in the registry.""" now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( """INSERT INTO skills (id, name, description, version, enabled, user_invocable, model_invocable, allowed_tools, metadata, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) @@ -38,7 +38,6 @@ async def upsert_skill( json.dumps(allowed_tools) if allowed_tools else None, json.dumps(metadata or {}), now, now), ) - await self.db.commit() async def get_skill_row(self, skill_id: str) -> dict | None: """Get a single skill record.""" @@ -57,18 +56,17 @@ async def list_skills(self) -> list[dict]: async def delete_skill_row(self, skill_id: str) -> None: """Remove a skill and its usage records.""" - await self.db.execute("DELETE FROM skill_usage WHERE skill_id = ?", (skill_id,)) - await self.db.execute("DELETE FROM skills WHERE id = ?", (skill_id,)) - await self.db.commit() + async with self._atomic(): + await self.db.execute("DELETE FROM skill_usage WHERE skill_id = ?", (skill_id,)) + await self.db.execute("DELETE FROM skills WHERE id = ?", (skill_id,)) async def update_skill_enabled(self, skill_id: str, enabled: bool) -> None: """Toggle a skill's enabled state.""" now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( "UPDATE skills SET enabled = ?, updated_at = ? WHERE id = ?", (enabled, now, skill_id), ) - await self.db.commit() async def record_skill_usage( self, @@ -80,12 +78,11 @@ async def record_skill_usage( error: str | None = None, ) -> None: """Log a skill invocation.""" - await self.db.execute( + await self._write( """INSERT INTO skill_usage (skill_id, session_id, invoked_by, duration_ms, success, error) VALUES (?, ?, ?, ?, ?, ?)""", (skill_id, session_id, invoked_by, duration_ms, success, error), ) - await self.db.commit() async def get_skill_usage(self, skill_id: str, limit: int = 50) -> list[dict]: """Get usage history for a skill.""" diff --git a/nerve/db/sources.py b/nerve/db/sources.py index ecfb40d..be64e94 100644 --- a/nerve/db/sources.py +++ b/nerve/db/sources.py @@ -20,11 +20,10 @@ async def get_sync_cursor(self, source: str) -> str | None: async def set_sync_cursor(self, source: str, cursor_value: str) -> None: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( "INSERT OR REPLACE INTO sync_cursors (source, cursor, updated_at) VALUES (?, ?, ?)", (source, cursor_value, now), ) - await self.db.commit() # --- Source run log operations --- @@ -38,14 +37,12 @@ async def log_source_run( ) -> int: """Log a source run with stats.""" now = datetime.now(timezone.utc).isoformat() - async with self.db.execute( + result = await self._write( "INSERT INTO source_run_log (source, ran_at, records_fetched, records_processed, error, session_id) " "VALUES (?, ?, ?, ?, ?, ?)", (source, now, records_fetched, records_processed, error, session_id), - ) as cursor: - log_id = cursor.lastrowid - await self.db.commit() - return log_id + ) + return result.lastrowid async def get_last_source_run(self, source: str) -> dict | None: """Get the most recent source run entry.""" @@ -222,11 +219,10 @@ async def update_source_messages_session( if not ids: return placeholders = ",".join("?" for _ in ids) - await self.db.execute( + await self._write( f"UPDATE source_messages SET run_session_id = ? WHERE source = ? AND id IN ({placeholders})", (session_id, source, *ids), ) - await self.db.commit() # Normalize ISO 8601 timestamps for consistent sorting across sources. # Different sources use different suffixes: "+00:00", "Z", or none. @@ -307,31 +303,20 @@ async def get_source_messages_storage(self) -> dict[str, dict]: async def delete_source_messages(self, source: str | None = None) -> int: """Purge source messages. If source is None, purge all. Returns count deleted.""" if source: - async with self.db.execute( - "SELECT COUNT(*) FROM source_messages WHERE source = ?", (source,) - ) as cursor: - count = (await cursor.fetchone())[0] - await self.db.execute("DELETE FROM source_messages WHERE source = ?", (source,)) + result = await self._write( + "DELETE FROM source_messages WHERE source = ?", (source,), + ) else: - async with self.db.execute("SELECT COUNT(*) FROM source_messages") as cursor: - count = (await cursor.fetchone())[0] - await self.db.execute("DELETE FROM source_messages") - await self.db.commit() - return count + result = await self._write("DELETE FROM source_messages") + return result.rowcount or 0 async def cleanup_expired_messages(self) -> int: """Delete source messages past their TTL. Returns count deleted.""" now = datetime.now(timezone.utc).isoformat() - async with self.db.execute( - "SELECT COUNT(*) FROM source_messages WHERE expires_at < ?", (now,) - ) as cursor: - count = (await cursor.fetchone())[0] - if count > 0: - await self.db.execute( - "DELETE FROM source_messages WHERE expires_at < ?", (now,) - ) - await self.db.commit() - return count + result = await self._write( + "DELETE FROM source_messages WHERE expires_at < ?", (now,) + ) + return result.rowcount or 0 # --- Consumer cursors --- @@ -379,7 +364,7 @@ async def set_consumer_cursor( """Advance cursor and refresh TTL. Links to session_id for UI tracking.""" now = datetime.now(timezone.utc) expires = (now + timedelta(days=ttl_days)).isoformat() - await self.db.execute( + await self._write( """INSERT INTO consumer_cursors (consumer, source, cursor_seq, session_id, updated_at, expires_at) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(consumer, source) DO UPDATE SET @@ -387,7 +372,6 @@ async def set_consumer_cursor( updated_at=excluded.updated_at, expires_at=excluded.expires_at""", (consumer, source, cursor_seq, session_id, now.isoformat(), expires), ) - await self.db.commit() async def list_consumer_cursors(self, consumer: str | None = None) -> list[dict]: """List active (non-expired) consumer cursors with unread counts.""" diff --git a/nerve/db/task_statuses.py b/nerve/db/task_statuses.py index 21802d4..018781a 100644 --- a/nerve/db/task_statuses.py +++ b/nerve/db/task_statuses.py @@ -133,16 +133,14 @@ async def update_task_status_def( if not sets: return params.append(name) - await self.db.execute( + await self._write( f"UPDATE task_statuses SET {', '.join(sets)} WHERE name = ?", tuple(params), ) - await self.db.commit() async def delete_task_status_def(self, name: str) -> None: """Delete a status definition. Caller enforces protection rules.""" - await self.db.execute("DELETE FROM task_statuses WHERE name = ?", (name,)) - await self.db.commit() + await self._write("DELETE FROM task_statuses WHERE name = ?", (name,)) async def count_tasks_with_status(self, name: str) -> int: """Count tasks currently set to the given status.""" diff --git a/nerve/db/tasks.py b/nerve/db/tasks.py index ad3d160..fd25218 100644 --- a/nerve/db/tasks.py +++ b/nerve/db/tasks.py @@ -133,19 +133,17 @@ async def count_tasks( async def update_task_status(self, task_id: str, status: str) -> None: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( "UPDATE tasks SET status = ?, updated_at = ? WHERE id = ?", (status, now, task_id), ) - await self.db.commit() async def update_task_tags(self, task_id: str, tags: str) -> None: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( "UPDATE tasks SET tags = ?, updated_at = ? WHERE id = ?", (tags, now, task_id), ) - await self.db.commit() # ── FTS query building ─────────────────────────────────────────────── @@ -374,18 +372,16 @@ async def find_tasks_by_source_url( async def rebuild_fts(self) -> None: """Clear the FTS index. Caller must re-populate via upsert_task().""" - await self.db.execute("DELETE FROM tasks_fts") - await self.db.commit() + await self._write("DELETE FROM tasks_fts") async def update_task_escalation( self, task_id: str, level: int, reminded_at: str | None = None ) -> None: now = datetime.now(timezone.utc).isoformat() - await self.db.execute( + await self._write( "UPDATE tasks SET escalation_level = ?, last_reminded_at = ?, updated_at = ? WHERE id = ?", (level, reminded_at or now, now, task_id), ) - await self.db.commit() async def get_task_health_stats(self) -> dict: """Get task and FTS counts for diagnostics (replaces raw sqlite3 calls).""" diff --git a/nerve/db/usage.py b/nerve/db/usage.py index a5defa8..10fcf0e 100644 --- a/nerve/db/usage.py +++ b/nerve/db/usage.py @@ -78,7 +78,7 @@ async def record_turn_usage( API returns the split; otherwise both default to 0 and only the aggregate is preserved. """ - await self.db.execute( + await self._write( "INSERT INTO session_usage " "(session_id, input_tokens, output_tokens, " "cache_creation_input_tokens, cache_read_input_tokens, " @@ -96,7 +96,6 @@ async def record_turn_usage( web_search_requests, web_fetch_requests, ), ) - await self.db.commit() async def get_session_usage_totals(self, session_id: str) -> dict: """Aggregate token usage for a session.""" @@ -284,10 +283,9 @@ async def get_usage_summary(self, days: int = 7) -> dict: async def delete_session_usage(self, session_id: str) -> None: """Delete all usage records for a session (cascade on delete).""" - await self.db.execute( + await self._write( "DELETE FROM session_usage WHERE session_id = ?", (session_id,), ) - await self.db.commit() def extract_cache_ttl_split(usage: dict) -> tuple[int, int]: diff --git a/nerve/db/wakeups.py b/nerve/db/wakeups.py index 07bc66a..645aa14 100644 --- a/nerve/db/wakeups.py +++ b/nerve/db/wakeups.py @@ -62,27 +62,21 @@ async def claim_wakeup(self, wakeup_id: int) -> bool: Returns ``True`` only for the caller that actually flipped the row, so overlapping sweeps can never fire the same wakeup twice. """ - cursor = await self.db.execute( + result = await self._write( "UPDATE session_wakeups SET status = 'fired' " "WHERE id = ? AND status = 'pending'", (wakeup_id,), ) - claimed = (cursor.rowcount or 0) == 1 - await cursor.close() - await self.db.commit() - return claimed + return (result.rowcount or 0) == 1 async def cancel_wakeups_for_session(self, session_id: str) -> int: """Delete all pending wakeups for a session. Returns rows removed.""" - cursor = await self.db.execute( + result = await self._write( "DELETE FROM session_wakeups " "WHERE session_id = ? AND status = 'pending'", (session_id,), ) - deleted = cursor.rowcount or 0 - await cursor.close() - await self.db.commit() - return deleted + return result.rowcount or 0 async def list_pending_wakeups(self, session_id: str | None = None) -> list[dict]: """List pending wakeups, optionally scoped to one session.""" diff --git a/tests/test_db.py b/tests/test_db.py index fea50d6..5adc07c 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -1420,6 +1420,10 @@ async def test_backfill_matches_runs_and_persistent(self, db: Database): "UPDATE cron_logs SET started_at = '2026-01-10 12:00:01' WHERE id = ?", (iso_log,), ) + # Raw fixture writes on the shared connection must commit themselves: + # leaving the txn open would (correctly) be rolled back by the next + # store method's leaked-transaction guard. + await db.db.commit() # Persistent session + a log row with no per-run candidate. await db.create_session("cron:pers-job", source="cron") @@ -1432,6 +1436,7 @@ async def test_backfill_matches_runs_and_persistent(self, db: Database): "UPDATE cron_logs SET started_at = '2026-01-05 00:00:00' WHERE id = ?", (far_log,), ) + await db.db.commit() # Source-runner log — no sessions at all. src_log = await db.log_cron_start("source:gmail") diff --git a/tests/test_db_atomicity.py b/tests/test_db_atomicity.py new file mode 100644 index 0000000..0395cfc --- /dev/null +++ b/tests/test_db_atomicity.py @@ -0,0 +1,243 @@ +"""Transaction-safety tests for the shared-connection write path. + +Covers the failure class behind the 2026-07-06 production wedge: an +abandoned open transaction on the shared aiosqlite connection pins a stale +WAL read snapshot, and once any other connection commits, every subsequent +write on the shared connection fails instantly with "database is locked" +(SQLITE_BUSY_SNAPSHOT — ``busy_timeout`` deliberately does not apply). +Nothing ever called ROLLBACK, so the daemon stayed wedged until restart. + +The write layer now guarantees: + * ``_atomic()`` / ``_write()`` commit on success, roll back on ANY error + (including task cancellation) — the connection is never left inside an + open transaction. + * ``_heal_leaked_txn()`` recovers from a leaked transaction produced by + any code path that bypasses the helpers. + * Every writer serializes under ``_write_lock``, so a single-statement + commit can never flush another coroutine's in-flight transaction. +""" + +import asyncio +import sqlite3 + +import pytest + +from nerve.db import Database + + +def _external_commit(db_path) -> None: + """Commit a write from a second, independent connection. + + Simulates the external process (CLI command, helper script, backup) + whose commit advances the WAL past the shared connection's snapshot. + """ + ext = sqlite3.connect(str(db_path), timeout=5) + try: + ext.execute( + "INSERT OR REPLACE INTO sync_cursors (source, cursor, updated_at) " + "VALUES ('external-writer', 'x', 'now')" + ) + ext.commit() + finally: + ext.close() + + +async def _poison_connection(db: Database) -> None: + """Leave the shared connection inside an open txn with a pinned snapshot. + + Equivalent to a write task abandoned between the implicit BEGIN and the + commit (the incident's entry path). + """ + await db.db.execute("BEGIN") + cur = await db.db.execute("SELECT COUNT(*) FROM sync_cursors") + await cur.fetchone() + await cur.close() + assert db.db.in_transaction + + +@pytest.mark.asyncio +class TestPoisonedConnectionRecovery: + """The incident reproduction: leaked txn + external commit.""" + + async def test_write_heals_leaked_transaction(self, db: Database, caplog): + await _poison_connection(db) + _external_commit(db.db_path) + + # Pre-fix behavior: this raised OperationalError("database is locked") + # in ~0ms, forever, despite busy_timeout=10s. The heal guard must + # roll back the leaked txn and let the write proceed. + await db.set_sync_cursor("healed", "42") + + assert not db.db.in_transaction + assert await db.get_sync_cursor("healed") == "42" + assert "Leaked open transaction" in caplog.text + + async def test_atomic_heals_leaked_transaction(self, db: Database): + await _poison_connection(db) + _external_commit(db.db_path) + + # Multi-statement path must recover the same way. + await db.upsert_task("t-heal", "tasks/t-heal.md", "heal test") + + assert not db.db.in_transaction + task = await db.get_task("t-heal") + assert task is not None and task["title"] == "heal test" + + async def test_stale_snapshot_write_fails_without_heal(self, db: Database): + """Documents the raw failure mode the guard protects against. + + A write attempted directly on the poisoned connection (bypassing + the helpers) fails instantly — proving the wedge is real and that + busy_timeout cannot save it. ROLLBACK is the only cure. + """ + await _poison_connection(db) + _external_commit(db.db_path) + + with pytest.raises(sqlite3.OperationalError, match="database is locked"): + await db.db.execute( + "INSERT OR REPLACE INTO sync_cursors (source, cursor, updated_at) " + "VALUES ('doomed', 'x', 'now')" + ) + + await db.db.rollback() # the cure + await db.set_sync_cursor("recovered", "1") + assert await db.get_sync_cursor("recovered") == "1" + + +@pytest.mark.asyncio +class TestAtomicRollback: + """_atomic() must discard partial writes on error and stay clean.""" + + async def test_body_error_rolls_back_partial_writes(self, db: Database): + with pytest.raises(RuntimeError, match="boom"): + async with db._atomic(): + await db.db.execute( + "INSERT INTO sync_cursors (source, cursor, updated_at) " + "VALUES ('partial', 'x', 'now')" + ) + raise RuntimeError("boom") + + assert not db.db.in_transaction + assert await db.get_sync_cursor("partial") is None + # Connection remains fully usable. + await db.set_sync_cursor("after-error", "7") + assert await db.get_sync_cursor("after-error") == "7" + + async def test_failed_statement_rolls_back(self, db: Database): + with pytest.raises(sqlite3.OperationalError): + await db._write("INSERT INTO no_such_table VALUES (1)") + assert not db.db.in_transaction + await db.set_sync_cursor("after-bad-sql", "3") + assert await db.get_sync_cursor("after-bad-sql") == "3" + + async def test_cancellation_mid_transaction_rolls_back(self, db: Database): + """A cancelled write task must not leave an open transaction. + + This is the incident's most likely entry path: the SDK cancels an + in-flight tool task between the implicit BEGIN and the commit. + """ + entered = asyncio.Event() + release = asyncio.Event() # never set — the task parks here + + async def body(): + async with db._atomic(): + await db.db.execute( + "INSERT INTO sync_cursors (source, cursor, updated_at) " + "VALUES ('cancelled', 'x', 'now')" + ) + entered.set() + await release.wait() + + task = asyncio.create_task(body()) + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert not db.db.in_transaction + assert await db.get_sync_cursor("cancelled") is None + await db.set_sync_cursor("after-cancel", "9") + assert await db.get_sync_cursor("after-cancel") == "9" + + +@pytest.mark.asyncio +class TestWriteSerialization: + """All writers hold the lock: no premature commit, no interleaving.""" + + async def test_single_writer_cannot_flush_inflight_transaction( + self, db: Database, + ): + entered = asyncio.Event() + proceed = asyncio.Event() + + async def failing_atomic(): + try: + async with db._atomic(): + await db.db.execute( + "INSERT INTO sync_cursors (source, cursor, updated_at) " + "VALUES ('inflight', 'x', 'now')" + ) + entered.set() + await proceed.wait() + raise RuntimeError("abort transaction") + except RuntimeError: + pass + + atomic_task = asyncio.create_task(failing_atomic()) + await asyncio.wait_for(entered.wait(), timeout=5) + + # A single-statement writer fired mid-transaction must block on the + # lock (pre-fix it ran immediately and its commit() flushed the + # in-flight 'inflight' row). + solo_task = asyncio.create_task(db.set_sync_cursor("solo", "v")) + await asyncio.sleep(0.05) + assert not solo_task.done(), "unlocked writer interleaved with _atomic" + + proceed.set() + await atomic_task + await asyncio.wait_for(solo_task, timeout=5) + + # The solo write landed; the aborted transaction's row did not. + assert await db.get_sync_cursor("solo") == "v" + assert await db.get_sync_cursor("inflight") is None + + async def test_concurrent_writers_smoke(self, db: Database): + async def single(i: int): + await db.set_sync_cursor(f"src-{i}", str(i)) + + async def multi(i: int): + await db.upsert_task(f"task-{i}", f"tasks/task-{i}.md", f"Task {i}") + + await asyncio.gather( + *(single(i) for i in range(25)), + *(multi(i) for i in range(25)), + ) + + assert not db.db.in_transaction + for i in range(25): + assert await db.get_sync_cursor(f"src-{i}") == str(i) + assert await db.count_tasks(status="all") == 25 + + +@pytest.mark.asyncio +class TestWriteResultPlumbing: + """lastrowid/rowcount survive the migration to _write().""" + + async def test_lastrowid_via_log_session_event(self, db: Database): + await db.create_session("s1") + first = await db.log_session_event("s1", "created") + second = await db.log_session_event("s1", "connected") + assert isinstance(first, int) and isinstance(second, int) + assert second > first + + async def test_rowcount_via_claim_wakeup(self, db: Database): + await db.create_session("s2") + wid = await db.add_wakeup("s2", "wake", "2999-01-01T00:00:00+00:00") + assert await db.claim_wakeup(wid) is True + assert await db.claim_wakeup(wid) is False # already fired + + async def test_rowcount_via_cancel_wakeups(self, db: Database): + await db.create_session("s3") + await db.add_wakeup("s3", "wake", "2999-01-01T00:00:00+00:00") + assert await db.cancel_wakeups_for_session("s3") == 1 + assert await db.cancel_wakeups_for_session("s3") == 0