From 65ad367f2605f2697fb62f61d996f9019c15fb12 Mon Sep 17 00:00:00 2001 From: Huang Date: Fri, 17 Jul 2026 17:57:30 +0800 Subject: [PATCH 1/6] fix(storage): unify SQLite + graph writes behind one process-level lock Route every SQL write and its KnowledgeGraph mutation through a single shared asyncio.Lock so they form one crash-consistent critical section, closing the split-lock gap where interrupted consolidation left graph orphans. Fixes #36 --- src/hebb/agents/consolidation_agent.py | 260 ++++++---- src/hebb/api.py | 2 +- src/hebb/graph/knowledge_graph.py | 16 +- src/hebb/scheduler/manager.py | 28 +- src/hebb/server/app.py | 5 +- src/hebb/server/routers/admin.py | 28 +- src/hebb/storage/factory.py | 27 +- src/hebb/storage/partition_store.py | 61 ++- src/hebb/storage/purge.py | 30 +- src/hebb/storage/sqlite_store.py | 447 ++++++++++++------ tests/conftest.py | 24 +- .../agents/test_consolidation_agent.py | 30 +- .../scheduler/test_scheduler_manager.py | 4 +- tests/integration/storage/test_purge.py | 8 +- tests/unit/test_audit_consolidation.py | 19 +- tests/unit/test_audit_embedding.py | 3 +- tests/unit/test_audit_storage.py | 108 +++++ 17 files changed, 787 insertions(+), 313 deletions(-) diff --git a/src/hebb/agents/consolidation_agent.py b/src/hebb/agents/consolidation_agent.py index e9e1c67..b1902bb 100644 --- a/src/hebb/agents/consolidation_agent.py +++ b/src/hebb/agents/consolidation_agent.py @@ -160,29 +160,54 @@ async def consolidate_memory(self, memory: Memory) -> ConsolidationResult: result.target_partition = "kept" return result - # Step 5: Write consolidated memory to target partition + # Step 5: Write consolidated memory to target partition. + # embed() runs outside the lock (CPU/IO-bound, no contention), + # then the SQL create + KG update + source delete + KG save + # happen inside one critical section with a single transaction + # (Issue #36 — crash-consistent across store + graph). embedding = await self.embedder.embed(consolidated_content) - new_memory = await self.memory_store.create( - data=MemoryCreate( - content=consolidated_content, - partition_id=target_partition, - importance_score=importance, - tags=tags, - metadata=memory.metadata, - source="consolidation", - ), - embedding=embedding, - ) - - # Step 6: Update knowledge graph with the new consolidated memory. - self.kg.update_from_tags(tags, new_memory.id) - - # Step 7: Delete the source memory from every store (SQL + graph). - # purge_memory strips the source id from the graph too — a no-op - # unless the source was itself graphed (in-partition - # re-consolidation) — and its kg.save() persists both the new tags - # and the removal in one write. - await purge_memory(self.memory_store, self.kg, memory.id) + create_impl = getattr(self.memory_store, "_create_impl", None) + delete_impl = getattr(self.memory_store, "_delete_impl", None) + begin = getattr(self.memory_store, "_begin", None) + if callable(create_impl) and callable(delete_impl) and callable(begin): + async with self.kg.lock: + await begin() + try: + new_memory = await create_impl( + data=MemoryCreate( + content=consolidated_content, + partition_id=target_partition, + importance_score=importance, + tags=tags, + metadata=memory.metadata, + source="consolidation", + ), + embedding=embedding, + skip_tx=True, + ) + self.kg.update_from_tags(tags, new_memory.id) + self.kg.remove_memory_from_tags(memory.id) + await delete_impl(memory.id, skip_tx=True) + await self.memory_store.db.commit() + except BaseException: + await self.memory_store.db.rollback() + raise + self.kg.save() + else: + # Other backends (PG pool): no shared lock. + new_memory = await self.memory_store.create( + data=MemoryCreate( + content=consolidated_content, + partition_id=target_partition, + importance_score=importance, + tags=tags, + metadata=memory.metadata, + source="consolidation", + ), + embedding=embedding, + ) + self.kg.update_from_tags(tags, new_memory.id) + await purge_memory(self.memory_store, self.kg, memory.id) result.target_partition = target_partition result.new_memory_id = new_memory.id @@ -363,21 +388,38 @@ async def _consolidate_session_chunk( if resolution == "update" and conflict_id in recalled_ids: await self._apply_conflict_update(conflict_id, content) - # Create consolidated memory + # Create consolidated memory — embed() runs lock-free, then the + # SQL create + KG update happen in one critical section (Issue #36). embedding = await self.embedder.embed(content) - new_memory = await self.memory_store.create( - data=MemoryCreate( - content=content, - partition_id=target_partition, - importance_score=importance, - tags=tags, - metadata=MemoryMetadata.model_validate(session_meta), - source="consolidation", - ), - embedding=embedding, - ) - async with kg_lock: - self.kg.update_from_tags(tags, new_memory.id) + create_impl = getattr(self.memory_store, "_create_impl", None) + if callable(create_impl): + async with kg_lock: + new_memory = await create_impl( + data=MemoryCreate( + content=content, + partition_id=target_partition, + importance_score=importance, + tags=tags, + metadata=MemoryMetadata.model_validate(session_meta), + source="consolidation", + ), + embedding=embedding, + ) + self.kg.update_from_tags(tags, new_memory.id) + else: + new_memory = await self.memory_store.create( + data=MemoryCreate( + content=content, + partition_id=target_partition, + importance_score=importance, + tags=tags, + metadata=MemoryMetadata.model_validate(session_meta), + source="consolidation", + ), + embedding=embedding, + ) + async with kg_lock: + self.kg.update_from_tags(tags, new_memory.id) results.append( ConsolidationResult( @@ -395,18 +437,23 @@ async def _consolidate_session_chunk( # the sources: that is silent data loss with zero replacement. Keep # them in the working partition for the next consolidation pass. if results: + delete_impl = getattr(self.memory_store, "_delete_impl", None) + begin = getattr(self.memory_store, "_begin", None) async with kg_lock: - for m in memories: - await self.memory_store.delete(m.id) - # Strip the source id from the graph (no-op unless the - # source was itself graphed, i.e. in-partition - # re-consolidation). - self.kg.remove_memory_from_tags(m.id) - # Persist the graph immediately after the SQL deletes, under - # the same lock, to shrink the crash window between "source - # row deleted" and "graph reference removed". The caller's - # end-of-batch save() is now a redundant safety net rather - # than the only persistence point. + if callable(begin) and callable(delete_impl): + await begin() + try: + for m in memories: + await delete_impl(m.id, skip_tx=True) + self.kg.remove_memory_from_tags(m.id) + await self.memory_store.db.commit() + except BaseException: + await self.memory_store.db.rollback() + raise + else: + for m in memories: + await self.memory_store.delete(m.id) + self.kg.remove_memory_from_tags(m.id) self.kg.save() elif self._drain_empty and "memories" in decision and not decision.get("memories"): # Well-formed empty result: the model read the whole conversation @@ -427,9 +474,22 @@ async def _consolidate_session_chunk( preview, ) async with kg_lock: - for m in memories: - await self.memory_store.delete(m.id) - self.kg.remove_memory_from_tags(m.id) + delete_impl = getattr(self.memory_store, "_delete_impl", None) + begin = getattr(self.memory_store, "_begin", None) + if callable(begin) and callable(delete_impl): + await begin() + try: + for m in memories: + await delete_impl(m.id, skip_tx=True) + self.kg.remove_memory_from_tags(m.id) + await self.memory_store.db.commit() + except BaseException: + await self.memory_store.db.rollback() + raise + else: + for m in memories: + await self.memory_store.delete(m.id) + self.kg.remove_memory_from_tags(m.id) self.kg.save() # Report the drained sources so the run summary counts them as # processed (progress) rather than reporting "0 ok" every pass. @@ -471,9 +531,10 @@ async def _consolidate_session_chunk( async def _apply_conflict_update(self, conflict_id: str, content: str) -> None: """Apply a conflict ``update`` resolution to an existing memory. - Updates the memory's content AND re-embeds it so the stored vec0 vector - matches the new text — otherwise vector recall keeps scoring the stale - embedding against the updated content. + Updates the memory's content AND re-embeds it in a **single + transaction** so the stored vec0 vector always matches the text — + otherwise vector recall keeps scoring the stale embedding against + the updated content (Issue #36). Args: conflict_id: Id of the recalled memory being updated. Must be an id @@ -486,12 +547,19 @@ async def _apply_conflict_update(self, conflict_id: str, content: str) -> None: if not content.strip(): # Never overwrite a real memory with empty/whitespace content. return - updated = await self.memory_store.update(conflict_id, MemoryUpdate(content=content)) - if updated is None: - return - # Re-embed so the vector matches the updated text. embedding = await self.embedder.embed(content) - await self.memory_store.update_embedding(conflict_id, embedding) + # Use the composite impl that updates content + embedding in one + # transaction. Only available on SQLiteMemoryStore; fall back to + # the two-step approach for other backends. + impl = getattr(self.memory_store, "_update_content_and_embedding_impl", None) + if callable(impl): + async with self.kg.lock: + await impl(conflict_id, MemoryUpdate(content=content), embedding) + else: + updated = await self.memory_store.update(conflict_id, MemoryUpdate(content=content)) + if updated is None: + return + await self.memory_store.update_embedding(conflict_id, embedding) async def consolidate_batch( self, @@ -637,10 +705,14 @@ async def _consolidate_one(self, memory: Memory, kg_lock: asyncio.Lock) -> Conso await self._apply_conflict_update(conflict_id, consolidated_content) result.conflicts_resolved += 1 elif resolution == "discard": + delete_impl = getattr(self.memory_store, "_delete_impl", None) async with kg_lock: self.kg.remove_memory_from_tags(memory.id) + if callable(delete_impl): + await delete_impl(memory.id) + else: + await self.memory_store.delete(memory.id) self.kg.save() - await self.memory_store.delete(memory.id) result.success = True result.target_partition = "discarded" return result @@ -660,28 +732,54 @@ async def _consolidate_one(self, memory: Memory, kg_lock: asyncio.Lock) -> Conso return result embedding = await self.embedder.embed(consolidated_content) - new_memory = await self.memory_store.create( - data=MemoryCreate( - content=consolidated_content, - partition_id=target_partition, - importance_score=importance, - tags=tags, - metadata=memory.metadata, - source="consolidation", - ), - embedding=embedding, - ) - - async with kg_lock: - self.kg.update_from_tags(tags, new_memory.id) - # Strip the source id under the same lock (no-op unless the - # source was graphed via in-partition re-consolidation). - self.kg.remove_memory_from_tags(memory.id) - await self.memory_store.delete(memory.id) - # Persist the graph immediately after the SQL delete, under the - # same lock, to shrink the crash window. The end-of-batch save() - # is now a redundant safety net. - self.kg.save() + create_impl = getattr(self.memory_store, "_create_impl", None) + delete_impl = getattr(self.memory_store, "_delete_impl", None) + begin = getattr(self.memory_store, "_begin", None) + if callable(create_impl) and callable(delete_impl) and callable(begin): + # SQLite: shared lock + single BEGIN IMMEDIATE wrapping both the + # SQL create and SQL delete so they commit as one atomic unit + # (Issue #36 — crash-consistent across store + graph). + async with kg_lock: + await begin() + try: + new_memory = await create_impl( + data=MemoryCreate( + content=consolidated_content, + partition_id=target_partition, + importance_score=importance, + tags=tags, + metadata=memory.metadata, + source="consolidation", + ), + embedding=embedding, + skip_tx=True, + ) + self.kg.update_from_tags(tags, new_memory.id) + self.kg.remove_memory_from_tags(memory.id) + await delete_impl(memory.id, skip_tx=True) + await self.memory_store.db.commit() + except BaseException: + await self.memory_store.db.rollback() + raise + self.kg.save() + else: + # Other backends (PG pool): no shared lock. + new_memory = await self.memory_store.create( + data=MemoryCreate( + content=consolidated_content, + partition_id=target_partition, + importance_score=importance, + tags=tags, + metadata=memory.metadata, + source="consolidation", + ), + embedding=embedding, + ) + async with kg_lock: + self.kg.update_from_tags(tags, new_memory.id) + self.kg.remove_memory_from_tags(memory.id) + await self.memory_store.delete(memory.id) + self.kg.save() result.target_partition = target_partition result.new_memory_id = new_memory.id diff --git a/src/hebb/api.py b/src/hebb/api.py index bda629f..639fe05 100644 --- a/src/hebb/api.py +++ b/src/hebb/api.py @@ -224,7 +224,7 @@ async def _async_start(self) -> None: raise StorageError(f"Failed to seed default partitions: {exc}") from exc kg_path_value = self.settings.kg_path - self._knowledge_graph = KnowledgeGraph(Path(kg_path_value)) + self._knowledge_graph = KnowledgeGraph(Path(kg_path_value), lock=ctx.write_lock) self._searcher = MemorySearcher(store=ctx.memory_store, embedder=embedder, graph=self._knowledge_graph) def close(self) -> None: diff --git a/src/hebb/graph/knowledge_graph.py b/src/hebb/graph/knowledge_graph.py index 9f2a1fa..b41c04c 100644 --- a/src/hebb/graph/knowledge_graph.py +++ b/src/hebb/graph/knowledge_graph.py @@ -26,15 +26,17 @@ class KnowledgeGraph: """Tag-based knowledge graph with file-backed persistence.""" - def __init__(self, path: Path) -> None: + def __init__(self, path: Path, lock: asyncio.Lock | None = None) -> None: self.path = path self.graph: nx.Graph = nx.Graph() - # Single shared lock for read-modify-write + save sequences. Both the - # consolidation and forgetting paths mutate the in-memory graph and - # persist it; serialising through this one lock keeps concurrent - # background tasks from interleaving a mutation with another task's - # save (which would write out a partially-updated graph). - self.lock = asyncio.Lock() + # Shared process-level write lock (Issue #36). When provided by the + # factory, this is the *same* ``asyncio.Lock`` instance used by + # ``SQLiteMemoryStore`` and ``SQLitePartitionStore``, so a SQL write + # and its corresponding graph mutation share a single critical section + # and cannot interleave with another coroutine's half-finished + # transaction. When ``None`` (tests, SDK, PG backend) a standalone + # lock is created. + self.lock = lock or asyncio.Lock() self._load() def _load(self) -> None: diff --git a/src/hebb/scheduler/manager.py b/src/hebb/scheduler/manager.py index 5da3169..0272e47 100644 --- a/src/hebb/scheduler/manager.py +++ b/src/hebb/scheduler/manager.py @@ -336,14 +336,30 @@ async def _run_forgetting(self) -> None: await self.memory_store.update_expiry(mid, iso) # Batch-delete expired memories, then strip them from the graph and - # persist the graph once — all under the graph's shared lock so a - # concurrent consolidation can't interleave a save (C1). + # persist the graph once — all under the shared write lock so a + # concurrent consolidation can't interleave (Issue #36, C1). + # Multiple deletes are wrapped in a single BEGIN IMMEDIATE … COMMIT + # so the batch is atomic: either all sources are removed or none. if to_delete: + delete_impl = getattr(self.memory_store, "_delete_impl", None) + begin = getattr(self.memory_store, "_begin", None) async with self.knowledge_graph.lock: - for mid in to_delete: - # save=False: the graph is persisted once after the loop. - await purge_memory(self.memory_store, self.knowledge_graph, mid, save=False) - total_deleted += 1 + if callable(begin) and callable(delete_impl): + await begin() + try: + for mid in to_delete: + await delete_impl(mid, skip_tx=True) + self.knowledge_graph.remove_memory_from_tags(mid) + total_deleted += 1 + await self.memory_store.db.commit() + except BaseException: + await self.memory_store.db.rollback() + raise + else: + for mid in to_delete: + await self.memory_store.delete(mid) + self.knowledge_graph.remove_memory_from_tags(mid) + total_deleted += 1 self.knowledge_graph.save() logger.info("Forgetting job complete: %d memories deleted", total_deleted) diff --git a/src/hebb/server/app.py b/src/hebb/server/app.py index 0bfcbf7..fee43e8 100644 --- a/src/hebb/server/app.py +++ b/src/hebb/server/app.py @@ -79,8 +79,9 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # Ensure default partitions await ctx.partition_store.ensure_defaults() - # Knowledge graph - kg = KnowledgeGraph(Path(settings.kg_path)) # kg_path already resolved to absolute + # Knowledge graph — shares the process-level write lock (Issue #36) so + # SQL mutations and graph mutations cannot interleave. + kg = KnowledgeGraph(Path(settings.kg_path), lock=ctx.write_lock) # kg_path already resolved to absolute app.state.knowledge_graph = kg # Optional cross-encoder reranker (None when settings.rerank_enabled=False). diff --git a/src/hebb/server/routers/admin.py b/src/hebb/server/routers/admin.py index e4058ea..1bf1653 100644 --- a/src/hebb/server/routers/admin.py +++ b/src/hebb/server/routers/admin.py @@ -173,6 +173,8 @@ async def trigger_forgetting( partitions_swept += 1 memories = await memory_store.get_by_partition(partition.id) scanned += len(memories) + # Collect deletes to batch under the KG lock (Issue #36). + to_delete: list[str] = [] for memory in memories: # Recompute expiry from current access state so strengthened # (recently recalled) memories survive, then persist for visibility. @@ -186,10 +188,28 @@ async def trigger_forgetting( ) await memory_store.update_expiry(memory.id, expires_at.isoformat()) if expires_at < now: - await purge_memory(memory_store, kg, memory.id, save=False) - deleted += 1 - if deleted > 0: - kg.save() + to_delete.append(memory.id) + if to_delete: + delete_impl = getattr(memory_store, "_delete_impl", None) + begin = getattr(memory_store, "_begin", None) + async with kg.lock: + if callable(begin) and callable(delete_impl): + await begin() + try: + for mid in to_delete: + await delete_impl(mid, skip_tx=True) + kg.remove_memory_from_tags(mid) + deleted += 1 + await memory_store.db.commit() + except BaseException: + await memory_store.db.rollback() + raise + else: + for mid in to_delete: + await memory_store.delete(mid) + kg.remove_memory_from_tags(mid) + deleted += 1 + kg.save() except Exception as exc: # Record the failed sweep (mirrors the scheduled job) before surfacing # the error, so a manual run that errors mid-way still leaves a record. diff --git a/src/hebb/storage/factory.py b/src/hebb/storage/factory.py index 9527b8c..60b65da 100644 --- a/src/hebb/storage/factory.py +++ b/src/hebb/storage/factory.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import logging from collections.abc import Callable, Coroutine from dataclasses import dataclass @@ -19,6 +20,9 @@ class StorageContext: memory_store: Any # MemoryStore protocol partition_store: Any # PartitionStore protocol close: Callable[[], Coroutine[Any, Any, None]] # async cleanup + # Shared process-level write lock for SQLite backend (Issue #36). + # ``None`` for PostgreSQL (pool-based, no shared lock needed). + write_lock: asyncio.Lock | None = None async def create_stores(settings: Settings) -> StorageContext: @@ -44,23 +48,30 @@ async def _create_sqlite(settings: Settings) -> StorageContext: from hebb.storage.sqlite_store import SQLiteMemoryStore db = await get_connection(settings.db_path, load_vec=settings.embedding_enabled) - # Resilient pragmas on the single shared connection. WAL is already set in - # get_connection; reaffirm it and add a busy_timeout so a concurrent writer - # (e.g. a background sweep) waits for the lock instead of failing fast with - # "database is locked". Serialization itself is handled by the store's - # in-process write lock (INT-2), not by re-architecting to a pool. + # Unified process-level write lock shared across SQLiteMemoryStore, + # SQLitePartitionStore, and KnowledgeGraph so that every SQL mutation + # and its corresponding graph mutation form a single crash-consistent + # critical section (Issue #36, audit C1). The lock prevents + # interleaving between background tasks (consolidation / forgetting) + # and request handlers that share the same connection. + shared_lock = asyncio.Lock() await db.execute("PRAGMA journal_mode=WAL") await db.execute("PRAGMA busy_timeout=5000") await initialize_schema(db, settings.embedding_dim, create_vec_table=settings.embedding_enabled) - memory_store = SQLiteMemoryStore(db) - partition_store = SQLitePartitionStore(db) + memory_store = SQLiteMemoryStore(db, write_lock=shared_lock) + partition_store = SQLitePartitionStore(db, write_lock=shared_lock) async def close() -> None: await db.close() logger.info("Storage backend: SQLite (%s)", settings.db_path) - return StorageContext(memory_store=memory_store, partition_store=partition_store, close=close) + return StorageContext( + memory_store=memory_store, + partition_store=partition_store, + close=close, + write_lock=shared_lock, + ) async def _create_postgresql(settings: Settings) -> StorageContext: diff --git a/src/hebb/storage/partition_store.py b/src/hebb/storage/partition_store.py index 4c294ce..2c29e78 100644 --- a/src/hebb/storage/partition_store.py +++ b/src/hebb/storage/partition_store.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio from datetime import datetime, timezone from typing import Any @@ -34,17 +35,31 @@ def _row_to_partition(row: aiosqlite.Row) -> Partition: class SQLitePartitionStore: """SQLite-backed partition store.""" - def __init__(self, db: aiosqlite.Connection) -> None: + def __init__(self, db: aiosqlite.Connection, write_lock: asyncio.Lock | None = None) -> None: self.db = db + # Shared process-level write lock (Issue #36). Same instance as + # ``SQLiteMemoryStore._write_lock`` and ``KnowledgeGraph.lock`` so + # partition mutations are serialised with all other writes. + self._write_lock = write_lock or asyncio.Lock() + + async def _begin(self) -> None: + """Open an explicit immediate transaction (same pattern as SQLiteMemoryStore).""" + await self.db.execute("BEGIN IMMEDIATE") async def create(self, data: PartitionCreate, is_system: bool = False) -> Partition: now = _now_iso() - await self.db.execute( - """INSERT INTO partitions (id, name, description, enabled, is_system, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?)""", - (data.id, data.name, data.description, int(data.enabled), int(is_system), now, now), - ) - await self.db.commit() + async with self._write_lock: + await self._begin() + try: + await self.db.execute( + """INSERT INTO partitions (id, name, description, enabled, is_system, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?)""", + (data.id, data.name, data.description, int(data.enabled), int(is_system), now, now), + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise return await self.get(data.id) # type: ignore[return-value] async def get(self, partition_id: str) -> Partition | None: @@ -110,11 +125,17 @@ async def update(self, partition_id: str, data: PartitionUpdate) -> Partition | params.append(_now_iso()) params.append(partition_id) - await self.db.execute( - f"UPDATE partitions SET {', '.join(updates)} WHERE id = ?", - params, - ) - await self.db.commit() + async with self._write_lock: + await self._begin() + try: + await self.db.execute( + f"UPDATE partitions SET {', '.join(updates)} WHERE id = ?", + params, + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise return await self.get(partition_id) async def delete(self, partition_id: str) -> bool: @@ -123,11 +144,17 @@ async def delete(self, partition_id: str) -> bool: if not existing or existing.is_system: return False - await self.db.execute( - "DELETE FROM partitions WHERE id = ?", - (partition_id,), - ) - await self.db.commit() + async with self._write_lock: + await self._begin() + try: + await self.db.execute( + "DELETE FROM partitions WHERE id = ?", + (partition_id,), + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise return True async def ensure_defaults(self) -> None: diff --git a/src/hebb/storage/purge.py b/src/hebb/storage/purge.py index d700ab8..cc4f106 100644 --- a/src/hebb/storage/purge.py +++ b/src/hebb/storage/purge.py @@ -18,14 +18,20 @@ async def purge_memory( - store: MemoryStore, - kg: KnowledgeGraph, + store: "MemoryStore", + kg: "KnowledgeGraph", memory_id: str, *, save: bool = True, ) -> bool: """Delete a memory from every store: SQL rows, vec0, FTS5, and the graph. + When the store exposes a ``_write_lock`` (SQLite backend, Issue #36), the + SQL delete and the graph mutation are performed inside a single critical + section so they form one crash-consistent unit. For backends without a + shared lock (PostgreSQL pool), the SQL delete runs independently and the + graph mutation is guarded by ``kg.lock``. + Args: store: Memory store; deletes ``memories`` + ``memory_embeddings`` + ``memory_fts``. @@ -41,8 +47,20 @@ async def purge_memory( Graph cleanup runs regardless of the row's existence, so a pre-existing orphan referencing ``memory_id`` is still swept. """ - deleted = await store.delete(memory_id) - kg.remove_memory_from_tags(memory_id) - if save: - kg.save() + # SQLite backend: _write_lock is the shared process-level lock (= kg.lock). + # Acquire it once so the SQL delete + KG mutation form one critical section. + shared_lock = getattr(store, "_write_lock", None) + if shared_lock is not None and hasattr(store, "_delete_impl"): + async with shared_lock: + deleted = await store._delete_impl(memory_id) # type: ignore[attr-defined] + kg.remove_memory_from_tags(memory_id) + if save: + kg.save() + else: + # PostgreSQL pool backend: no shared lock; KG guarded by its own lock. + deleted = await store.delete(memory_id) + async with kg.lock: + kg.remove_memory_from_tags(memory_id) + if save: + kg.save() return deleted diff --git a/src/hebb/storage/sqlite_store.py b/src/hebb/storage/sqlite_store.py index 1c071b0..2df4294 100644 --- a/src/hebb/storage/sqlite_store.py +++ b/src/hebb/storage/sqlite_store.py @@ -59,13 +59,18 @@ def _row_to_memory(row: aiosqlite.Row) -> Memory: class SQLiteMemoryStore: """SQLite-backed memory store with sqlite-vec for vector search.""" - def __init__(self, db: aiosqlite.Connection) -> None: + def __init__(self, db: aiosqlite.Connection, write_lock: asyncio.Lock | None = None) -> None: self.db = db - # All callers share one store instance and one connection. A single - # process-wide write lock serializes every multi-statement mutation so - # no coroutine can commit another coroutine's half-finished transaction - # (audit write F1/F2). Reads are not locked. - self._write_lock = asyncio.Lock() + # Process-level write lock shared with KnowledgeGraph and + # PartitionStore (Issue #36). When provided by the factory, a single + # ``asyncio.Lock`` serialises every SQL mutation *and* its + # corresponding graph mutation so the two form one crash-consistent + # unit. When ``None`` (tests, standalone usage) a standalone lock is + # created. The ``_*_impl`` methods skip lock acquisition so callers + # that already hold the lock can compose multiple operations in one + # critical section without deadlocking on ``asyncio.Lock`` + # non-reentrancy. + self._write_lock = write_lock or asyncio.Lock() # Declared vec0 embedding width, read lazily from the schema and cached. # ``None`` means the dimension is unknown (e.g. the regular-table # fallback when vec0 is unavailable), in which case the guard is skipped. @@ -105,6 +110,19 @@ async def _get_vec_dim(self) -> int | None: return self._vec_dim async def create(self, data: MemoryCreate, embedding: _List[float] | None = None) -> Memory: + async with self._write_lock: + return await self._create_impl(data, embedding) + + async def _create_impl( + self, data: MemoryCreate, embedding: _List[float] | None = None, *, skip_tx: bool = False + ) -> Memory: + """``create`` without lock acquisition — for callers already holding ``_write_lock``. + + Args: + skip_tx: When ``True``, skip the internal ``BEGIN IMMEDIATE … COMMIT`` + boundary. The caller is responsible for wrapping this call in an + explicit transaction so multiple SQL operations form one atomic unit. + """ memory_id = str(uuid.uuid4()) now = _now_iso() vec_bytes: bytes | None = None @@ -114,42 +132,44 @@ async def create(self, data: MemoryCreate, embedding: _List[float] | None = None raise EmbeddingDimensionError(f"embedding length {len(embedding)} does not match vec0 width {expected}") vec_bytes = np.array(embedding, dtype=np.float32).tobytes() - async with self._write_lock: + if not skip_tx: await self._begin() - try: - await self.db.execute( - """INSERT INTO memories - (id, partition_id, content, importance_score, tags, metadata, source, - created_at, updated_at, last_accessed_at, access_count) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0)""", - ( - memory_id, - data.partition_id, - data.content, - data.importance_score, - json.dumps(data.tags), - json.dumps(data.metadata.model_dump(exclude_none=True)), - data.source, - now, - now, - now, - ), - ) - if vec_bytes is not None: - await self.db.execute( - "INSERT INTO memory_embeddings (memory_id, partition_id, embedding) VALUES (?, ?, ?)", - (memory_id, data.partition_id, vec_bytes), - ) - # FTS5 index for keyword search — partition_id is UNINDEXED but - # carried so retrieval can filter by partition inside MATCH. + try: + await self.db.execute( + """INSERT INTO memories + (id, partition_id, content, importance_score, tags, metadata, source, + created_at, updated_at, last_accessed_at, access_count) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0)""", + ( + memory_id, + data.partition_id, + data.content, + data.importance_score, + json.dumps(data.tags), + json.dumps(data.metadata.model_dump(exclude_none=True)), + data.source, + now, + now, + now, + ), + ) + if vec_bytes is not None: await self.db.execute( - "INSERT INTO memory_fts (memory_id, partition_id, content) VALUES (?, ?, ?)", - (memory_id, data.partition_id, data.content), + "INSERT INTO memory_embeddings (memory_id, partition_id, embedding) VALUES (?, ?, ?)", + (memory_id, data.partition_id, vec_bytes), ) + # FTS5 index for keyword search — partition_id is UNINDEXED but + # carried so retrieval can filter by partition inside MATCH. + await self.db.execute( + "INSERT INTO memory_fts (memory_id, partition_id, content) VALUES (?, ?, ?)", + (memory_id, data.partition_id, data.content), + ) + if not skip_tx: await self.db.commit() - except BaseException: + except BaseException: + if not skip_tx: await self.db.rollback() - raise + raise return await self.get(memory_id) # type: ignore[return-value] async def get(self, memory_id: str) -> Memory | None: @@ -198,6 +218,11 @@ async def list( return memories, total async def update(self, memory_id: str, data: MemoryUpdate) -> Memory | None: + async with self._write_lock: + return await self._update_impl(memory_id, data) + + async def _update_impl(self, memory_id: str, data: MemoryUpdate) -> Memory | None: + """``update`` without lock acquisition — for callers already holding ``_write_lock``.""" existing = await self.get(memory_id) if not existing: return None @@ -224,41 +249,53 @@ async def update(self, memory_id: str, data: MemoryUpdate) -> Memory | None: params.append(_now_iso()) params.append(memory_id) - async with self._write_lock: - await self._begin() - try: + await self._begin() + try: + await self.db.execute( + f"UPDATE memories SET {', '.join(updates)} WHERE id = ?", + params, + ) + # Sync FTS5 if content changed + if data.content is not None: + await self.db.execute("DELETE FROM memory_fts WHERE memory_id = ?", (memory_id,)) await self.db.execute( - f"UPDATE memories SET {', '.join(updates)} WHERE id = ?", - params, + "INSERT INTO memory_fts (memory_id, partition_id, content) VALUES (?, ?, ?)", + (memory_id, existing.partition_id, data.content), ) - # Sync FTS5 if content changed - if data.content is not None: - await self.db.execute("DELETE FROM memory_fts WHERE memory_id = ?", (memory_id,)) - await self.db.execute( - "INSERT INTO memory_fts (memory_id, partition_id, content) VALUES (?, ?, ?)", - (memory_id, existing.partition_id, data.content), - ) - await self.db.commit() - except BaseException: - await self.db.rollback() - raise + await self.db.commit() + except BaseException: + await self.db.rollback() + raise return await self.get(memory_id) async def delete(self, memory_id: str) -> bool: async with self._write_lock: + return await self._delete_impl(memory_id) + + async def _delete_impl(self, memory_id: str, *, skip_tx: bool = False) -> bool: + """``delete`` without lock acquisition — for callers already holding ``_write_lock``. + + Args: + skip_tx: When ``True``, skip the internal ``BEGIN IMMEDIATE … COMMIT`` + boundary. The caller is responsible for wrapping this call in an + explicit transaction so multiple SQL operations form one atomic unit. + """ + if not skip_tx: await self._begin() - try: - cursor = await self.db.execute( - "DELETE FROM memories WHERE id = ?", - (memory_id,), - ) - deleted = cursor.rowcount > 0 - await self.db.execute("DELETE FROM memory_embeddings WHERE memory_id = ?", (memory_id,)) - await self.db.execute("DELETE FROM memory_fts WHERE memory_id = ?", (memory_id,)) + try: + cursor = await self.db.execute( + "DELETE FROM memories WHERE id = ?", + (memory_id,), + ) + deleted = cursor.rowcount > 0 + await self.db.execute("DELETE FROM memory_embeddings WHERE memory_id = ?", (memory_id,)) + await self.db.execute("DELETE FROM memory_fts WHERE memory_id = ?", (memory_id,)) + if not skip_tx: await self.db.commit() - except BaseException: + except BaseException: + if not skip_tx: await self.db.rollback() - raise + raise return deleted async def search_by_vector( @@ -542,53 +579,61 @@ def _sort_key(m: Memory) -> int: return results async def delete_expired(self) -> _List[str]: - now = _now_iso() async with self._write_lock: - await self._begin() - try: - # Get expired memory IDs first (for embedding/fts cleanup) - cursor = await self.db.execute( - "SELECT id FROM memories WHERE expires_at IS NOT NULL AND expires_at < ?", - (now,), - ) - rows = await cursor.fetchall() - expired_ids = [row["id"] for row in rows] + return await self._delete_expired_impl() - if not expired_ids: - await self.db.commit() - return [] + async def _delete_expired_impl(self) -> _List[str]: + """``delete_expired`` without lock acquisition — for callers already holding ``_write_lock``.""" + now = _now_iso() + await self._begin() + try: + # Get expired memory IDs first (for embedding/fts cleanup) + cursor = await self.db.execute( + "SELECT id FROM memories WHERE expires_at IS NOT NULL AND expires_at < ?", + (now,), + ) + rows = await cursor.fetchall() + expired_ids = [row["id"] for row in rows] - placeholders = ",".join("?" * len(expired_ids)) - await self.db.execute( - f"DELETE FROM memory_embeddings WHERE memory_id IN ({placeholders})", - expired_ids, - ) - await self.db.execute( - f"DELETE FROM memory_fts WHERE memory_id IN ({placeholders})", - expired_ids, - ) - await self.db.execute( - f"DELETE FROM memories WHERE id IN ({placeholders})", - expired_ids, - ) + if not expired_ids: await self.db.commit() - except BaseException: - await self.db.rollback() - raise + return [] + + placeholders = ",".join("?" * len(expired_ids)) + await self.db.execute( + f"DELETE FROM memory_embeddings WHERE memory_id IN ({placeholders})", + expired_ids, + ) + await self.db.execute( + f"DELETE FROM memory_fts WHERE memory_id IN ({placeholders})", + expired_ids, + ) + await self.db.execute( + f"DELETE FROM memories WHERE id IN ({placeholders})", + expired_ids, + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise return expired_ids async def update_access(self, memory_id: str) -> None: async with self._write_lock: - await self._begin() - try: - await self.db.execute( - "UPDATE memories SET access_count = access_count + 1, last_accessed_at = ? WHERE id = ?", - (_now_iso(), memory_id), - ) - await self.db.commit() - except BaseException: - await self.db.rollback() - raise + await self._update_access_impl(memory_id) + + async def _update_access_impl(self, memory_id: str) -> None: + """``update_access`` without lock acquisition — for callers already holding ``_write_lock``.""" + await self._begin() + try: + await self.db.execute( + "UPDATE memories SET access_count = access_count + 1, last_accessed_at = ? WHERE id = ?", + (_now_iso(), memory_id), + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise async def update_access_batch(self, memory_ids: _List[str]) -> None: """Bump access_count + last_accessed_at for several memories in one @@ -598,18 +643,24 @@ async def update_access_batch(self, memory_ids: _List[str]) -> None: """ if not memory_ids: return - placeholders = ",".join("?" * len(memory_ids)) async with self._write_lock: - await self._begin() - try: - await self.db.execute( - f"UPDATE memories SET access_count = access_count + 1, last_accessed_at = ? WHERE id IN ({placeholders})", - (_now_iso(), *memory_ids), - ) - await self.db.commit() - except BaseException: - await self.db.rollback() - raise + await self._update_access_batch_impl(memory_ids) + + async def _update_access_batch_impl(self, memory_ids: _List[str]) -> None: + """``update_access_batch`` without lock acquisition.""" + if not memory_ids: + return + placeholders = ",".join("?" * len(memory_ids)) + await self._begin() + try: + await self.db.execute( + f"UPDATE memories SET access_count = access_count + 1, last_accessed_at = ? WHERE id IN ({placeholders})", + (_now_iso(), *memory_ids), + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise async def update_embedding(self, memory_id: str, embedding: _List[float]) -> None: if not embedding: @@ -617,40 +668,136 @@ async def update_embedding(self, memory_id: str, embedding: _List[float]) -> Non expected = await self._get_vec_dim() if expected is not None and len(embedding) != expected: raise EmbeddingDimensionError(f"embedding length {len(embedding)} does not match vec0 width {expected}") - vec_bytes = np.array(embedding, dtype=np.float32).tobytes() async with self._write_lock: + await self._update_embedding_impl(memory_id, embedding, _pre_validated=True) + + async def _update_embedding_impl( + self, memory_id: str, embedding: _List[float], *, _pre_validated: bool = False + ) -> None: + """``update_embedding`` without lock acquisition — for callers already holding ``_write_lock``. + + Args: + memory_id: Memory to re-embed. + embedding: New embedding vector. + _pre_validated: When ``True``, dimension validation is skipped + (caller already validated). Avoids double-checking in + composite impl methods. + """ + if not embedding: + return + if not _pre_validated: + expected = await self._get_vec_dim() + if expected is not None and len(embedding) != expected: + raise EmbeddingDimensionError(f"embedding length {len(embedding)} does not match vec0 width {expected}") + vec_bytes = np.array(embedding, dtype=np.float32).tobytes() + await self._begin() + try: + # Read partition_id from the memories row so the vec0 index keeps + # the metadata column in sync with the parent record. + cursor = await self.db.execute("SELECT partition_id FROM memories WHERE id = ?", (memory_id,)) + row = await cursor.fetchone() + if row is None: + await self.db.commit() + return + partition_id = row["partition_id"] + await self.db.execute("DELETE FROM memory_embeddings WHERE memory_id = ?", (memory_id,)) + await self.db.execute( + "INSERT INTO memory_embeddings (memory_id, partition_id, embedding) VALUES (?, ?, ?)", + (memory_id, partition_id, vec_bytes), + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise + + async def _update_content_and_embedding_impl( + self, + memory_id: str, + data: MemoryUpdate, + embedding: _List[float], + *, + skip_tx: bool = False, + ) -> Memory | None: + """Update content and re-embed in a **single transaction** (Issue #36). + + Previously ``_apply_conflict_update`` called ``update()`` then + ``update_embedding()`` as two independent transactions — if the + second failed, the memory's content was already changed but the + vector still matched the old text. This method wraps both writes + in one ``BEGIN IMMEDIATE … COMMIT`` so they either succeed as a + unit or roll back together. + + Must be called while already holding ``_write_lock``. + + Args: + memory_id: Memory to update. + data: Update payload (only ``content`` is applied here). + embedding: New embedding matching the updated content. + skip_tx: When ``True``, skip the internal ``BEGIN IMMEDIATE … COMMIT`` + boundary. The caller manages the transaction. + + Returns: + The updated ``Memory``, or ``None`` if the id doesn't exist. + """ + existing = await self.get(memory_id) + if not existing: + return None + + # When the embedder produces an empty vector (e.g. NoopEmbedder in + # tests), skip the embedding write — mirrors ``update_embedding``'s + # ``if not embedding: return`` guard. + skip_embedding = not embedding + if not skip_embedding: + expected = await self._get_vec_dim() + if expected is not None and len(embedding) != expected: + raise EmbeddingDimensionError(f"embedding length {len(embedding)} does not match vec0 width {expected}") + vec_bytes = np.array(embedding, dtype=np.float32).tobytes() + + now = _now_iso() + if not skip_tx: await self._begin() - try: - # Read partition_id from the memories row so the vec0 index keeps - # the metadata column in sync with the parent record. - cursor = await self.db.execute("SELECT partition_id FROM memories WHERE id = ?", (memory_id,)) - row = await cursor.fetchone() - if row is None: - await self.db.commit() - return - partition_id = row["partition_id"] + try: + await self.db.execute( + "UPDATE memories SET content = ?, updated_at = ? WHERE id = ?", + (data.content, now, memory_id), + ) + # Sync FTS5 index to the new content. + await self.db.execute("DELETE FROM memory_fts WHERE memory_id = ?", (memory_id,)) + await self.db.execute( + "INSERT INTO memory_fts (memory_id, partition_id, content) VALUES (?, ?, ?)", + (memory_id, existing.partition_id, data.content), + ) + if not skip_embedding: + # Replace the embedding vector. await self.db.execute("DELETE FROM memory_embeddings WHERE memory_id = ?", (memory_id,)) await self.db.execute( "INSERT INTO memory_embeddings (memory_id, partition_id, embedding) VALUES (?, ?, ?)", - (memory_id, partition_id, vec_bytes), + (memory_id, existing.partition_id, vec_bytes), ) + if not skip_tx: await self.db.commit() - except BaseException: + except BaseException: + if not skip_tx: await self.db.rollback() - raise + raise + return await self.get(memory_id) async def update_expiry(self, memory_id: str, expires_at: str) -> None: async with self._write_lock: - await self._begin() - try: - await self.db.execute( - "UPDATE memories SET expires_at = ? WHERE id = ?", - (expires_at, memory_id), - ) - await self.db.commit() - except BaseException: - await self.db.rollback() - raise + await self._update_expiry_impl(memory_id, expires_at) + + async def _update_expiry_impl(self, memory_id: str, expires_at: str) -> None: + """``update_expiry`` without lock acquisition — for callers already holding ``_write_lock``.""" + await self._begin() + try: + await self.db.execute( + "UPDATE memories SET expires_at = ? WHERE id = ?", + (expires_at, memory_id), + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise async def update_expiry_batch(self, items: _List[tuple[str, str]]) -> None: """Set ``expires_at`` for several memories under one lock + transaction. @@ -670,15 +817,21 @@ async def update_expiry_batch(self, items: _List[tuple[str, str]]) -> None: """ if not items: return - params = [(expires_at, memory_id) for memory_id, expires_at in items] async with self._write_lock: - await self._begin() - try: - await self.db.executemany( - "UPDATE memories SET expires_at = ? WHERE id = ?", - params, - ) - await self.db.commit() - except BaseException: - await self.db.rollback() - raise + await self._update_expiry_batch_impl(items) + + async def _update_expiry_batch_impl(self, items: _List[tuple[str, str]]) -> None: + """``update_expiry_batch`` without lock acquisition.""" + if not items: + return + params = [(expires_at, memory_id) for memory_id, expires_at in items] + await self._begin() + try: + await self.db.executemany( + "UPDATE memories SET expires_at = ? WHERE id = ?", + params, + ) + await self.db.commit() + except BaseException: + await self.db.rollback() + raise diff --git a/tests/conftest.py b/tests/conftest.py index c4bea15..f315561 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio from pathlib import Path from unittest.mock import AsyncMock @@ -17,6 +18,7 @@ from hebb.agents.llm_client import LLMClient from hebb.config.settings import Settings from hebb.embedding.local import NoopEmbedder +from hebb.graph.knowledge_graph import KnowledgeGraph from hebb.storage.migrations import get_connection, initialize_schema from hebb.storage.partition_store import SQLitePartitionStore from hebb.storage.sqlite_store import SQLiteMemoryStore @@ -42,8 +44,14 @@ def settings(tmp_path: Path) -> Settings: ) +@pytest.fixture +def shared_lock() -> asyncio.Lock: + """Process-level write lock shared between store and KG (Issue #36).""" + return asyncio.Lock() + + @pytest_asyncio.fixture -async def db(settings: Settings): +async def db(settings: Settings, shared_lock: asyncio.Lock): conn = await get_connection(settings.db_path) await initialize_schema(conn, settings.embedding_dim) yield conn @@ -51,12 +59,18 @@ async def db(settings: Settings): @pytest_asyncio.fixture -async def memory_store(db) -> SQLiteMemoryStore: - return SQLiteMemoryStore(db) +async def memory_store(db, shared_lock: asyncio.Lock) -> SQLiteMemoryStore: + return SQLiteMemoryStore(db, write_lock=shared_lock) @pytest_asyncio.fixture -async def partition_store(db) -> SQLitePartitionStore: - store = SQLitePartitionStore(db) +async def partition_store(db, shared_lock: asyncio.Lock) -> SQLitePartitionStore: + store = SQLitePartitionStore(db, write_lock=shared_lock) await store.ensure_defaults() return store + + +@pytest_asyncio.fixture +async def knowledge_graph(tmp_path: Path, shared_lock: asyncio.Lock) -> KnowledgeGraph: + """Knowledge graph sharing the process-level write lock (Issue #36).""" + return KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) diff --git a/tests/integration/agents/test_consolidation_agent.py b/tests/integration/agents/test_consolidation_agent.py index 27ca354..61f3bb8 100644 --- a/tests/integration/agents/test_consolidation_agent.py +++ b/tests/integration/agents/test_consolidation_agent.py @@ -2,6 +2,8 @@ from __future__ import annotations +import asyncio + import pytest from hebb.agents.consolidation_agent import ConsolidationAgent @@ -13,9 +15,9 @@ class TestConsolidationAgent: @pytest.mark.asyncio - async def test_consolidate_memory(self, mock_llm, memory_store, partition_store, noop_embedder, tmp_path): + async def test_consolidate_memory(self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path): """Consolidation should move memory from hebb to target partition.""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) # Create a memory in hebb mem = await memory_store.create( @@ -71,12 +73,12 @@ async def test_consolidate_memory(self, mock_llm, memory_store, partition_store, @pytest.mark.asyncio async def test_session_consolidation_preserves_turn_and_timestamp( - self, mock_llm, memory_store, partition_store, noop_embedder, tmp_path + self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path ): """Session consolidation must carry the source turns' turn span and earliest timestamp onto the consolidated memory (so turn-window expansion and temporal_boost still work on consolidated output).""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1 = await memory_store.create( MemoryCreate( @@ -169,12 +171,12 @@ def _agent(self, mock_llm, memory_store, partition_store, noop_embedder, kg, set @pytest.mark.asyncio async def test_session_wellformed_empty_drains_sources( - self, mock_llm, memory_store, partition_store, noop_embedder, tmp_path + self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path ): """A well-formed empty result ({"memories": []}) = the LLM deliberately judged the turns low-value, so the sources are drained from the inbox (default behavior) — otherwise they replay to the LLM forever.""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1, m2 = await self._session_pair(memory_store) mock_llm.complete_json.side_effect = [ {"queries": []}, # RecallAgent @@ -194,13 +196,13 @@ async def test_session_wellformed_empty_drains_sources( @pytest.mark.asyncio async def test_session_wellformed_empty_kept_when_flag_off( - self, mock_llm, memory_store, partition_store, noop_embedder, tmp_path + self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path ): """With consolidation_drain_empty_sources=False, a well-formed empty result keeps the sources (legacy behavior, opt-out).""" from hebb.config.settings import Settings - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1, m2 = await self._session_pair(memory_store) mock_llm.complete_json.side_effect = [ {"queries": []}, @@ -219,12 +221,12 @@ async def test_session_wellformed_empty_kept_when_flag_off( @pytest.mark.asyncio async def test_session_parse_failure_keeps_sources( - self, mock_llm, memory_store, partition_store, noop_embedder, tmp_path + self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path ): """A garbled/unparseable LLM response (complete_json -> {}) must NEVER drain — it has no "memories" key, so the sources are kept for retry. Guards against deleting data on a transient/parse failure (F4).""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1, m2 = await self._session_pair(memory_store) mock_llm.complete_json.side_effect = [ {"queries": []}, @@ -239,9 +241,9 @@ async def test_session_parse_failure_keeps_sources( assert await memory_store.get(m2.id) is not None @pytest.mark.asyncio - async def test_consolidate_batch_empty(self, mock_llm, memory_store, partition_store, noop_embedder, tmp_path): + async def test_consolidate_batch_empty(self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path): """Batch consolidation with no hebb memories returns empty list.""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) recall_agent = RecallAgent( llm=mock_llm, @@ -261,10 +263,10 @@ async def test_consolidate_batch_empty(self, mock_llm, memory_store, partition_s @pytest.mark.asyncio async def test_consolidate_handles_conflict_discard( - self, mock_llm, memory_store, partition_store, noop_embedder, tmp_path + self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path ): """When LLM says discard, the memory is deleted without creating a new one.""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) mem = await memory_store.create( MemoryCreate(content="redundant info", partition_id="mem_hippocampus"), diff --git a/tests/integration/scheduler/test_scheduler_manager.py b/tests/integration/scheduler/test_scheduler_manager.py index 00d4ad9..c968003 100644 --- a/tests/integration/scheduler/test_scheduler_manager.py +++ b/tests/integration/scheduler/test_scheduler_manager.py @@ -12,8 +12,8 @@ class TestSchedulerManager: @pytest.fixture - def scheduler(self, settings, memory_store, partition_store, tmp_path): - kg = KnowledgeGraph(tmp_path / "kg.json") + def scheduler(self, settings, memory_store, partition_store, shared_lock, tmp_path): + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) embedder = NoopEmbedder(384) return SchedulerManager( settings=settings, diff --git a/tests/integration/storage/test_purge.py b/tests/integration/storage/test_purge.py index 1fbf63d..d02ce4b 100644 --- a/tests/integration/storage/test_purge.py +++ b/tests/integration/storage/test_purge.py @@ -15,8 +15,8 @@ @pytest.mark.asyncio -async def test_purge_removes_sql_row_and_graph_node(memory_store, tmp_path): - kg = KnowledgeGraph(tmp_path / "kg.json") +async def test_purge_removes_sql_row_and_graph_node(memory_store, shared_lock, tmp_path): + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) mem = await memory_store.create( MemoryCreate(content="graphed memory", partition_id="mem_hippocampus") ) @@ -32,10 +32,10 @@ async def test_purge_removes_sql_row_and_graph_node(memory_store, tmp_path): @pytest.mark.asyncio -async def test_purge_cleans_graph_even_when_row_absent(memory_store, tmp_path): +async def test_purge_cleans_graph_even_when_row_absent(memory_store, shared_lock, tmp_path): """Graph cleanup runs regardless of the row's existence, so a stale reference is swept rather than left dangling.""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) kg.update_from_tags(["ghost"], "missing-id") assert kg.get_tag("ghost") is not None diff --git a/tests/unit/test_audit_consolidation.py b/tests/unit/test_audit_consolidation.py index 80b023b..baf61bf 100644 --- a/tests/unit/test_audit_consolidation.py +++ b/tests/unit/test_audit_consolidation.py @@ -119,7 +119,9 @@ def test_no_match_returns_empty(self, tmp_path: Path) -> None: def _make_agent( mock_llm, memory_store, partition_store, embedder, tmp_path ) -> tuple[ConsolidationAgent, KnowledgeGraph]: - kg = KnowledgeGraph(tmp_path / "kg.json") + # Reuse the store's shared lock so KG and store share one critical section (Issue #36). + shared_lock = getattr(memory_store, "_write_lock", None) + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) recall_agent = RecallAgent( llm=mock_llm, searcher=MemorySearcher(store=memory_store, embedder=embedder), @@ -207,15 +209,16 @@ async def test_conflict_update_reembeds_and_scopes_to_recalled( MemoryCreate(content="Actually user prefers coffee now", partition_id="mem_hippocampus"), ) - # Spy on update_embedding to confirm the re-embed fires for the right id. + # Spy on _update_content_and_embedding_impl to confirm the re-embed + # fires for the right id (this is the single-transaction path, Issue #36). reembedded: list[str] = [] - original_update_embedding = memory_store.update_embedding + original_impl = memory_store._update_content_and_embedding_impl - async def _spy(memory_id: str, embedding: list[float]) -> None: + async def _spy(memory_id: str, data, embedding: list[float]) -> None: reembedded.append(memory_id) - await original_update_embedding(memory_id, embedding) + await original_impl(memory_id, data, embedding) - memory_store.update_embedding = _spy # type: ignore[method-assign] + memory_store._update_content_and_embedding_impl = _spy # type: ignore[method-assign] mock_llm.complete_json.side_effect = [ # RecallAgent.recall -> query generation (returns the existing mem) @@ -250,11 +253,11 @@ async def _spy(memory_id: str, embedding: list[float]) -> None: class TestConsolidationLLMModelGuard: @pytest.mark.asyncio async def test_run_consolidation_skips_without_llm_model( - self, settings, memory_store, partition_store, noop_embedder, tmp_path + self, settings, memory_store, partition_store, noop_embedder, shared_lock, tmp_path ) -> None: """No llm_model -> run_consolidation returns [] without calling litellm.""" settings.llm_model = None - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) # Seed a memory so an unguarded path would actually try to consolidate. await memory_store.create( diff --git a/tests/unit/test_audit_embedding.py b/tests/unit/test_audit_embedding.py index 88ab134..9f9ca51 100644 --- a/tests/unit/test_audit_embedding.py +++ b/tests/unit/test_audit_embedding.py @@ -271,6 +271,7 @@ async def ensure_defaults(self) -> None: class FakeCtx: memory_store = object() partition_store = FakePartitionStore() + write_lock = None # SQLite backend would pass a shared lock (Issue #36) async def close(self) -> None: return None @@ -281,7 +282,7 @@ async def fake_create_stores(settings: Settings) -> Any: return FakeCtx() class FakeKG: - def __init__(self, path: Any) -> None: + def __init__(self, path: Any, lock: Any = None) -> None: pass def save(self) -> None: diff --git a/tests/unit/test_audit_storage.py b/tests/unit/test_audit_storage.py index 4cc673a..1b38c5f 100644 --- a/tests/unit/test_audit_storage.py +++ b/tests/unit/test_audit_storage.py @@ -9,6 +9,8 @@ - Concurrent ``create`` calls on one shared store do not interleave their multi-statement transactions (INT-2). - ``update_expiry_batch`` applies all updates atomically (INT-6). +- Concurrent write + consolidation + forgetting under the shared process-level + lock produces no orphans and no lost rows (Issue #36, C1). """ from __future__ import annotations @@ -17,6 +19,7 @@ import pytest +from hebb.graph.knowledge_graph import KnowledgeGraph from hebb.models.memory import MemoryCreate from hebb.storage.migrations import get_connection, initialize_schema from hebb.storage.sqlite_store import EmbeddingDimensionError, SQLiteMemoryStore @@ -126,3 +129,108 @@ async def test_batch_sets_all_expiries(self, store: SQLiteMemoryStore) -> None: @pytest.mark.asyncio async def test_empty_batch_is_noop(self, store: SQLiteMemoryStore) -> None: await store.update_expiry_batch([]) # must not raise + + +class TestConcurrentWriteConsolidateForget: + """Concurrent write + consolidation + forgetting under the shared process-level + write lock (Issue #36, acceptance criterion 4). + + Drives three tasks concurrently against a single SQLiteMemoryStore + + KnowledgeGraph that share one ``asyncio.Lock``: + + * Task A — interactive API write: ``store.create()`` + KG ``update_from_tags`` + ``save``. + * Task B — consolidation simulation: ``_create_impl`` + ``update_from_tags`` + + ``_delete_impl`` + ``remove_memory_from_tags`` + ``save`` (all under the shared lock). + * Task C — forgetting simulation: ``_delete_impl`` + ``remove_memory_from_tags`` + + ``save`` (all under the shared lock). + + Asserts zero KG→SQL orphans and zero lost rows after all tasks finish. + """ + + @pytest.fixture + async def shared_store_and_kg( + self, tmp_path, embedding_dim: int + ) -> tuple[SQLiteMemoryStore, KnowledgeGraph, asyncio.Lock]: + lock = asyncio.Lock() + db = await get_connection(str(tmp_path / "stress.db")) + await initialize_schema(db, embedding_dim) + store = SQLiteMemoryStore(db, write_lock=lock) + kg = KnowledgeGraph(tmp_path / "stress_kg.json", lock=lock) + yield store, kg, lock + await db.close() + + @pytest.mark.asyncio + async def test_no_orphans_after_concurrent_operations( + self, shared_store_and_kg + ) -> None: + store, kg, lock = shared_store_and_kg + + # Pre-populate: 10 memories with tags, half flagged for "forgetting". + initial_ids: list[str] = [] + for i in range(10): + m = await store.create( + MemoryCreate(partition_id="p1", content=f"seed-{i}", tags=[f"tag-{i}"]), + ) + initial_ids.append(m.id) + async with lock: + kg.update_from_tags([f"tag-{i}"], m.id) + kg.save() + + # IDs to be "forgotten" (first 5). + forget_ids = initial_ids[:5] + + async def api_write() -> None: + """Simulate REST API writes interleaving with background tasks.""" + for i in range(10): + m = await store.create( + MemoryCreate(partition_id="p1", content=f"api-{i}", tags=[f"api-tag-{i}"]), + ) + async with lock: + kg.update_from_tags([f"api-tag-{i}"], m.id) + kg.save() + + async def consolidation_sim() -> None: + """Simulate a consolidation step: create new + delete source (atomic).""" + for i in range(5): + source_id = initial_ids[5 + i] # second half as consolidation sources + async with lock: + new_mem = await store._create_impl( + MemoryCreate(partition_id="p2", content=f"consolidated-{i}", tags=[f"con-tag-{i}"]), + ) + kg.update_from_tags([f"con-tag-{i}"], new_mem.id) + kg.remove_memory_from_tags(source_id) + await store._delete_impl(source_id) + kg.save() + + async def forgetting_sim() -> None: + """Simulate a forgetting sweep: delete expired memories (atomic).""" + async with lock: + for mid in forget_ids: + await store._delete_impl(mid) + kg.remove_memory_from_tags(mid) + kg.save() + + # Run all three concurrently. + await asyncio.gather(api_write(), consolidation_sim(), forgetting_sim()) + + # 1. No KG→SQL orphan: every memory_id referenced in the KG must exist in SQL. + referenced: set[str] = set() + for node_id in kg.graph.nodes: + referenced.update(kg.graph.nodes[node_id].get("memory_ids", [])) + for mid in referenced: + assert await store.get(mid) is not None, f"KG references non-existent memory {mid} (orphan)" + + # 2. No lost rows: verify expected counts. + # - api_write created 10 new memories in p1. + # - consolidation_sim created 5 new memories in p2, deleted 5 sources from p1. + # - forgetting_sim deleted 5 memories from p1. + # Initial: 10 in p1. After: 10 (api) + 0 (5 deleted by consolidation, 5 by forget) in p1, + # 5 in p2. + p1_mems, p1_count = await store.list(partition_id="p1") + p2_mems, p2_count = await store.list(partition_id="p2") + assert p1_count == 10, f"Expected 10 in p1, got {p1_count}" + assert p2_count == 5, f"Expected 5 in p2, got {p2_count}" + + # 3. KG file can be loaded cleanly. + kg2 = KnowledgeGraph(kg.path) + assert kg2.graph.number_of_nodes() > 0 From 0b5e8db9ca7263b1d04fc688e0c70dd810a070cb Mon Sep 17 00:00:00 2001 From: Huang Date: Mon, 20 Jul 2026 17:48:39 +0800 Subject: [PATCH 2/6] =?UTF-8?q?fix(review):=20apply=20PR=20#52=20bot=20rev?= =?UTF-8?q?iew=20=E2=80=94=20post-commit=20KG=20mutation,=20guards,=20type?= =?UTF-8?q?s,=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the 22 review comments raised on PR #52 by gemini-code-assist, Copilot, and coderabbitai (topics 1–7 of the modification plan). No behavior change beyond the safety / typing / doc gaps the bots flagged. Data integrity (Issue #36; Gemini + CodeRabbit) * consolidation_agent.consolidate_memory, _consolidate_session_chunk, _consolidate_one (incl. discard), scheduler.manager scheduled forgetting, server.routers.admin manual forgetting, storage.purge_memory: move KnowledgeGraph mutations (update_from_tags / remove_memory_from_tags) and counters to AFTER db.commit() — still under the shared lock — so a SQL rollback can no longer diverge the in-memory NetworkX graph. The narrow commit → save() crash window is acknowledged; KG→SQL orphans left by a crash there are self-healed by reconcile() (KG→SQL is self-healing; SQL→KG lost forward refs are not, so the post-commit ordering is safer). Lock discipline (Issue #36; Copilot) * storage.purge_memory: explicitly acquire kg.lock even when the store's _write_lock differs, and handle the non-reentrant same-lock case (no second async with) so SDK / test wiring cannot race the graph mutation. PG fallback (CodeRabbit) * consolidation_agent non-SQLite fast-path: guard self.kg.update_from_tags with self.kg.lock so concurrent forgetting / consolidation cannot mutate the NetworkX graph unserialized. Typing (CodeRabbit + Pylance red flags) * consolidation_agent / scheduler.manager / routers.admin / purge: replace the getattr(store, "_create_impl", None) + callable() duck-type probe with isinstance(self.memory_store, SQLiteMemoryStore). Pylance / mypy now narrow .db / ._create_impl / ._delete_impl / ._begin / ._update_content_and_embedding_impl to the concrete store type, removing the 20 "await … not awaitable" / "MemoryStore has no attribute db" red flags the probe caused — no behavior change (same SQLite vs. PG split, same lock composition). * tests/conftest + the integration/audit regression tests: complete strict type annotations (fixture params + -> None / AsyncIterator[...]) so mypy is clean across the PR's test footprint. Drop two `# type: ignore` comments mypy 2.1 now reports as unused. Tests (CodeRabbit) * tests/unit/test_audit_storage.test_no_orphans_after_concurrent_operations: rewrite to exercise one transaction per operation (BEGIN + skip_tx=True + commit/rollback once) under the shared lock — previously the simulated consolidation/forgetting committed each _impl independently and released the lock between SQL and graph updates, so the test passed even when the split-lock / partial-commit failure modes regressed. * Add test_rollback_on_injected_failure_in_consolidation: inject a failure mid-transaction and assert both SQLite and the reloaded on-disk KG remain unchanged after rollback. Documentation (CodeRabbit pre-merge check, 55% -> 100% docstring coverage on the PR-touched files) * Every public callable in the PR-touched source files (consolidation_agent, scheduler.manager, server.app, server.routers.admin, storage.factory, storage.partition_store, storage.sqlite_store) now carries a docstring with Args/Returns/Raises where the signature warrants it. PR-touched coverage 70% -> 100%; src/hebb overall 68.44% -> 72.52%. CHANGELOG: add ## [Unreleased] section summarizing the above. Verified: * ruff check src/ tests/ → All checks passed * mypy src/hebb → Success: no issues (132 source files) * pytest tests/ → 779 passed, 1 skipped Refs #36 (PR #52). --- CHANGELOG.md | 66 ++++++ src/hebb/agents/consolidation_agent.py | 108 ++++++---- src/hebb/scheduler/manager.py | 26 ++- src/hebb/server/app.py | 6 + src/hebb/server/routers/admin.py | 42 ++-- src/hebb/storage/factory.py | 2 + src/hebb/storage/partition_store.py | 37 ++++ src/hebb/storage/purge.py | 40 +++- src/hebb/storage/sqlite_store.py | 58 +++++ tests/conftest.py | 12 +- .../agents/test_consolidation_agent.py | 103 +++++++-- .../scheduler/test_scheduler_manager.py | 15 +- tests/integration/storage/test_purge.py | 16 +- tests/unit/test_audit_consolidation.py | 15 +- tests/unit/test_audit_storage.py | 204 ++++++++++++++++-- 15 files changed, 624 insertions(+), 126 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 53bd004..7e8c4e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,72 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 3. Merge to main — publish.yml ships to PyPI on the pyproject.toml change and tags the release. --> +## [Unreleased] + +### Fixed + +- **In-memory knowledge-graph divergence on SQL rollback**: across every + composite write path (`consolidate_memory`, `_consolidate_session_chunk`, + `_consolidate_one`'s `discard` branch, scheduled forgetting, manual + forgetting, and `purge_memory`), `KnowledgeGraph.update_from_tags` / + `remove_memory_from_tags` now run **after** `db.commit()` succeeds — still under + the shared lock. A SQL rollback therefore can never leave the NetworkX graph + pointing at memories that still exist (a divergence `reconcile()` cannot + self-heal). The narrow commit → `kg.save()` crash window may still leave a + KG→SQL orphan that `reconcile()` sweeps on the next start; this is the safer + trade-off because KG→SQL orphans are self-healing and SQL→KG lost forward refs + are not. (Issue #36; Gemini / CodeRabbit review.) +- **`purge_memory` lock discipline**: acquire `kg.lock` explicitly even when the + store exposes a different `_write_lock`, and handle the non-reentrant + same-lock case (no second `async with`) so SDK / test wiring cannot race the + graph mutation. (Issue #36; Copilot review.) +- **PG fallback graph mutation**: guard `self.kg.update_from_tags` with + `self.kg.lock` on the non-SQLite consolidation fast-path so concurrent + forgetting / consolidation cannot mutate the NetworkX graph unserialized. + (CodeRabbit review.) + +### Changed + +- **Composite fast-path typing**: replace the `getattr(store, "_create_impl", + None)` + `callable()` duck-type probe with `isinstance(self.memory_store, + SQLiteMemoryStore)` across `consolidation_agent`, `scheduler.manager`, + `server.routers.admin`, and `storage.purge`. Pylance / mypy now narrow + `.db` / `._create_impl` / `._delete_impl` / `._begin` / + `._update_content_and_embedding_impl` to the concrete store type, removing + the 20 `await … not awaitable` and `MemoryStore has no attribute db` red + flags the previous duck-type probe caused (no behavior change: same SQLite + vs. PG split, same lock composition). + +### Tests + +- **`test_audit_storage` concurrent regression**: rewrite + `test_no_orphans_after_concurrent_operations` to exercise one transaction per + operation (`BEGIN` + `skip_tx=True` + commit/rollback once) under the shared + lock — previously the simulated consolidation / forgetting committed each + `_impl` independently and released the lock between SQL and graph updates, so + the test passed even when the split-lock / partial-commit failure modes + regressed. Add a new `test_rollback_on_injected_failure_in_consolidation` + that injects a failure mid-transaction and asserts both SQLite and the + reloaded on-disk KnowledgeGraph remain unchanged after rollback. +- **Strict typing across audit + integration tests**: complete type annotations + (fixture parameters + `-> None` / `AsyncIterator[...]` returns) for + `tests/conftest.py`, `tests/integration/agents/test_consolidation_agent.py`, + `tests/integration/scheduler/test_scheduler_manager.py`, + `tests/integration/storage/test_purge.py`, + `tests/unit/test_audit_consolidation.py`, and + `tests/unit/test_audit_storage.py` so `mypy` strict is clean across the PR's + test footprint. +- **Removed two `type: ignore` comments** that mypy 2.1 now reports as unused. + +### Documentation + +- **Docstring coverage**: every public callable in the PR-touched source files + (`consolidation_agent`, `scheduler.manager`, `server.app`, `server.routers.admin`, + `storage.factory`, `storage.partition_store`, `storage.sqlite_store`) now + carries a docstring with `Args`/`Returns`/`Raises` where the signature + warrants it. PR-touched file coverage 70% → 100%; `src/hebb` overall + 68.44% → 72.52%. + ## [0.3.0] - 2026-06-29 ### Added diff --git a/src/hebb/agents/consolidation_agent.py b/src/hebb/agents/consolidation_agent.py index b1902bb..c531dc6 100644 --- a/src/hebb/agents/consolidation_agent.py +++ b/src/hebb/agents/consolidation_agent.py @@ -21,12 +21,19 @@ from hebb.models.memory import Memory, MemoryCreate, MemoryMetadata, MemoryUpdate from hebb.storage.base import MemoryStore, PartitionStore from hebb.storage.purge import purge_memory +from hebb.storage.sqlite_store import SQLiteMemoryStore logger = logging.getLogger(__name__) @dataclass class ConsolidationResult: + """Outcome of consolidating one memory or session batch. + + Fields are populated for both success and failure so the run summary can + report progress, errors, and the produced memory's id alongside the source. + """ + original_memory_id: str target_partition: str = "" new_memory_id: str = "" @@ -166,14 +173,18 @@ async def consolidate_memory(self, memory: Memory) -> ConsolidationResult: # happen inside one critical section with a single transaction # (Issue #36 — crash-consistent across store + graph). embedding = await self.embedder.embed(consolidated_content) - create_impl = getattr(self.memory_store, "_create_impl", None) - delete_impl = getattr(self.memory_store, "_delete_impl", None) - begin = getattr(self.memory_store, "_begin", None) - if callable(create_impl) and callable(delete_impl) and callable(begin): + if isinstance(self.memory_store, SQLiteMemoryStore): + # SQLite: shared lock + single BEGIN wrapping both the create and + # source delete (Issue #36). KG mutations run AFTER the commit, so a + # SQL rollback can never diverge the in-memory graph (NetworkX has no + # transaction support; diverging would not be self-healing — + # reconcile() can drop a dead ref but not restore a lost forward ref). + # The narrow commit → save() crash window may leave a KG→SQL orphan + # that reconcile() sweeps on the next start. async with self.kg.lock: - await begin() + await self.memory_store._begin() try: - new_memory = await create_impl( + new_memory = await self.memory_store._create_impl( data=MemoryCreate( content=consolidated_content, partition_id=target_partition, @@ -185,16 +196,20 @@ async def consolidate_memory(self, memory: Memory) -> ConsolidationResult: embedding=embedding, skip_tx=True, ) - self.kg.update_from_tags(tags, new_memory.id) - self.kg.remove_memory_from_tags(memory.id) - await delete_impl(memory.id, skip_tx=True) + await self.memory_store._delete_impl(memory.id, skip_tx=True) await self.memory_store.db.commit() except BaseException: await self.memory_store.db.rollback() raise + # Commit succeeded — mutate the in-memory graph afterwards to + # avoid rollback divergence (Issue #36; Gemini review). + self.kg.update_from_tags(tags, new_memory.id) + self.kg.remove_memory_from_tags(memory.id) self.kg.save() else: - # Other backends (PG pool): no shared lock. + # Other backends (PG pool): no shared lock; kg.lock still guards the + # in-memory graph mutation. purge_memory acquires its own lock, so it + # is left outside this critical section to avoid double-acquire. new_memory = await self.memory_store.create( data=MemoryCreate( content=consolidated_content, @@ -206,7 +221,8 @@ async def consolidate_memory(self, memory: Memory) -> ConsolidationResult: ), embedding=embedding, ) - self.kg.update_from_tags(tags, new_memory.id) + async with self.kg.lock: + self.kg.update_from_tags(tags, new_memory.id) await purge_memory(self.memory_store, self.kg, memory.id) result.target_partition = target_partition @@ -391,10 +407,9 @@ async def _consolidate_session_chunk( # Create consolidated memory — embed() runs lock-free, then the # SQL create + KG update happen in one critical section (Issue #36). embedding = await self.embedder.embed(content) - create_impl = getattr(self.memory_store, "_create_impl", None) - if callable(create_impl): + if isinstance(self.memory_store, SQLiteMemoryStore): async with kg_lock: - new_memory = await create_impl( + new_memory = await self.memory_store._create_impl( data=MemoryCreate( content=content, partition_id=target_partition, @@ -437,19 +452,20 @@ async def _consolidate_session_chunk( # the sources: that is silent data loss with zero replacement. Keep # them in the working partition for the next consolidation pass. if results: - delete_impl = getattr(self.memory_store, "_delete_impl", None) - begin = getattr(self.memory_store, "_begin", None) async with kg_lock: - if callable(begin) and callable(delete_impl): - await begin() + if isinstance(self.memory_store, SQLiteMemoryStore): + await self.memory_store._begin() try: for m in memories: - await delete_impl(m.id, skip_tx=True) - self.kg.remove_memory_from_tags(m.id) + await self.memory_store._delete_impl(m.id, skip_tx=True) await self.memory_store.db.commit() except BaseException: await self.memory_store.db.rollback() raise + # Commit succeeded — mutate the in-memory graph afterwards + # to avoid rollback divergence (Issue #36; Gemini review). + for m in memories: + self.kg.remove_memory_from_tags(m.id) else: for m in memories: await self.memory_store.delete(m.id) @@ -474,18 +490,19 @@ async def _consolidate_session_chunk( preview, ) async with kg_lock: - delete_impl = getattr(self.memory_store, "_delete_impl", None) - begin = getattr(self.memory_store, "_begin", None) - if callable(begin) and callable(delete_impl): - await begin() + if isinstance(self.memory_store, SQLiteMemoryStore): + await self.memory_store._begin() try: for m in memories: - await delete_impl(m.id, skip_tx=True) - self.kg.remove_memory_from_tags(m.id) + await self.memory_store._delete_impl(m.id, skip_tx=True) await self.memory_store.db.commit() except BaseException: await self.memory_store.db.rollback() raise + # Commit succeeded — mutate the in-memory graph afterwards + # to avoid rollback divergence (Issue #36; Gemini review). + for m in memories: + self.kg.remove_memory_from_tags(m.id) else: for m in memories: await self.memory_store.delete(m.id) @@ -551,10 +568,11 @@ async def _apply_conflict_update(self, conflict_id: str, content: str) -> None: # Use the composite impl that updates content + embedding in one # transaction. Only available on SQLiteMemoryStore; fall back to # the two-step approach for other backends. - impl = getattr(self.memory_store, "_update_content_and_embedding_impl", None) - if callable(impl): + if isinstance(self.memory_store, SQLiteMemoryStore): async with self.kg.lock: - await impl(conflict_id, MemoryUpdate(content=content), embedding) + await self.memory_store._update_content_and_embedding_impl( + conflict_id, MemoryUpdate(content=content), embedding + ) else: updated = await self.memory_store.update(conflict_id, MemoryUpdate(content=content)) if updated is None: @@ -705,13 +723,14 @@ async def _consolidate_one(self, memory: Memory, kg_lock: asyncio.Lock) -> Conso await self._apply_conflict_update(conflict_id, consolidated_content) result.conflicts_resolved += 1 elif resolution == "discard": - delete_impl = getattr(self.memory_store, "_delete_impl", None) async with kg_lock: - self.kg.remove_memory_from_tags(memory.id) - if callable(delete_impl): - await delete_impl(memory.id) + # Delete the SQL row first; mutate the graph only once the + # delete has succeeded (a failed delete must not strip tags). + if isinstance(self.memory_store, SQLiteMemoryStore): + await self.memory_store._delete_impl(memory.id) else: await self.memory_store.delete(memory.id) + self.kg.remove_memory_from_tags(memory.id) self.kg.save() result.success = True result.target_partition = "discarded" @@ -732,17 +751,16 @@ async def _consolidate_one(self, memory: Memory, kg_lock: asyncio.Lock) -> Conso return result embedding = await self.embedder.embed(consolidated_content) - create_impl = getattr(self.memory_store, "_create_impl", None) - delete_impl = getattr(self.memory_store, "_delete_impl", None) - begin = getattr(self.memory_store, "_begin", None) - if callable(create_impl) and callable(delete_impl) and callable(begin): - # SQLite: shared lock + single BEGIN IMMEDIATE wrapping both the - # SQL create and SQL delete so they commit as one atomic unit - # (Issue #36 — crash-consistent across store + graph). + if isinstance(self.memory_store, SQLiteMemoryStore): + # SQLite: shared lock + single BEGIN wrapping both the create and + # source delete (Issue #36). KG mutations run AFTER the commit, so a + # SQL rollback can never diverge the in-memory graph; commit → + # save() crash orphan is self-healed by reconcile() (KG→SQL, not + # SQL→KG — reconcile can drop dead refs but cannot restore lost ones). async with kg_lock: - await begin() + await self.memory_store._begin() try: - new_memory = await create_impl( + new_memory = await self.memory_store._create_impl( data=MemoryCreate( content=consolidated_content, partition_id=target_partition, @@ -754,13 +772,13 @@ async def _consolidate_one(self, memory: Memory, kg_lock: asyncio.Lock) -> Conso embedding=embedding, skip_tx=True, ) - self.kg.update_from_tags(tags, new_memory.id) - self.kg.remove_memory_from_tags(memory.id) - await delete_impl(memory.id, skip_tx=True) + await self.memory_store._delete_impl(memory.id, skip_tx=True) await self.memory_store.db.commit() except BaseException: await self.memory_store.db.rollback() raise + self.kg.update_from_tags(tags, new_memory.id) + self.kg.remove_memory_from_tags(memory.id) self.kg.save() else: # Other backends (PG pool): no shared lock. diff --git a/src/hebb/scheduler/manager.py b/src/hebb/scheduler/manager.py index 0272e47..b65ae0e 100644 --- a/src/hebb/scheduler/manager.py +++ b/src/hebb/scheduler/manager.py @@ -35,7 +35,7 @@ ) from hebb.server.forgetting_tracker import record_run as record_forget_run from hebb.storage.base import MemoryStore, PartitionStore -from hebb.storage.purge import purge_memory +from hebb.storage.sqlite_store import SQLiteMemoryStore from hebb.upgrade import state as upgrade_state from hebb.upgrade.checker import run_check as run_upgrade_check from hebb.upgrade.helper import spawn_detached @@ -69,6 +69,11 @@ def __init__( self._consolidation_lock = asyncio.Lock() def start(self) -> None: + """Register the daily consolidation and periodic forgetting jobs and start the scheduler. + + Idempotent: reuses the existing job ids via ``replace_existing=True`` so a + second ``start()`` reschedules rather than duplicating jobs. + """ hour, minute = self._parse_consolidation_time() self.scheduler.add_job( func=self._run_consolidation, @@ -107,6 +112,7 @@ def start(self) -> None: ) def shutdown(self) -> None: + """Stop the scheduler without waiting for in-flight jobs (idempotent).""" self.scheduler.shutdown(wait=False) logger.info("Scheduler stopped") @@ -341,20 +347,24 @@ async def _run_forgetting(self) -> None: # Multiple deletes are wrapped in a single BEGIN IMMEDIATE … COMMIT # so the batch is atomic: either all sources are removed or none. if to_delete: - delete_impl = getattr(self.memory_store, "_delete_impl", None) - begin = getattr(self.memory_store, "_begin", None) async with self.knowledge_graph.lock: - if callable(begin) and callable(delete_impl): - await begin() + if isinstance(self.memory_store, SQLiteMemoryStore): + await self.memory_store._begin() try: for mid in to_delete: - await delete_impl(mid, skip_tx=True) - self.knowledge_graph.remove_memory_from_tags(mid) - total_deleted += 1 + await self.memory_store._delete_impl(mid, skip_tx=True) await self.memory_store.db.commit() except BaseException: await self.memory_store.db.rollback() raise + # Commit succeeded — mutate the in-memory graph and bump the + # counter afterwards so a SQL rollback can never diverge the + # graph or inflate total_deleted (Issue #36; Gemini review). + # The narrow commit → save() crash window may leave a KG→SQL + # orphan that reconcile() sweeps on the next start. + for mid in to_delete: + self.knowledge_graph.remove_memory_from_tags(mid) + total_deleted += 1 else: for mid in to_delete: await self.memory_store.delete(mid) diff --git a/src/hebb/server/app.py b/src/hebb/server/app.py index fee43e8..1cecf79 100644 --- a/src/hebb/server/app.py +++ b/src/hebb/server/app.py @@ -39,6 +39,12 @@ class NoCacheStaticFiles(StaticFiles): """ async def get_response(self, path: str, scope: Scope) -> Response: + """Like ``StaticFiles.get_response`` but force ``Cache-Control: no-cache``. + + The Web Console ships hashed asset names but the HTML entry points must + always be reloaded after an upgrade, so we override the default cache + header here. + """ response = await super().get_response(path, scope) response.headers["Cache-Control"] = "no-cache" return response diff --git a/src/hebb/server/routers/admin.py b/src/hebb/server/routers/admin.py index 1bf1653..4376539 100644 --- a/src/hebb/server/routers/admin.py +++ b/src/hebb/server/routers/admin.py @@ -32,7 +32,7 @@ ) from hebb.server.forgetting_tracker import record_run as record_forget_run from hebb.storage.base import MemoryStore, PartitionStore -from hebb.storage.purge import purge_memory +from hebb.storage.sqlite_store import SQLiteMemoryStore logger = logging.getLogger(__name__) @@ -49,11 +49,17 @@ async def trigger_consolidation( settings: Settings = Depends(get_settings), scheduler: SchedulerManager = Depends(get_scheduler), ) -> dict[str, Any]: - # Optional body: {"partition_ids": [...], "keep_partition": bool}. - # When ``partition_ids`` is given, each is consolidated in turn; with - # ``keep_partition`` the consolidated memories stay in their source - # partition (per-scenario benches). Empty body → default global - # HIPPOCAMPUS→long-term consolidation (production / LoCoMo). + """Manually trigger a consolidation pass and return the per-memory results. + + Optional body (``{"partition_ids": [...], "keep_partition": bool}``) scopes + the run; with ``keep_partition`` the consolidated memories stay in their + source partition (per-scenario benches). An empty body triggers the default + HIPPOCAMPUS→long-term consolidation (production flow). + + The call is routed through the scheduler's lock-protected + ``run_consolidation`` (INT-7) so a manual ``/consolidate`` cannot overlap + the daily cron job (which would double-delete / double-graph sources). + """ body = body or {} partition_ids = body.get("partition_ids") or None keep_partition = bool(body.get("keep_partition", False)) @@ -190,20 +196,24 @@ async def trigger_forgetting( if expires_at < now: to_delete.append(memory.id) if to_delete: - delete_impl = getattr(memory_store, "_delete_impl", None) - begin = getattr(memory_store, "_begin", None) async with kg.lock: - if callable(begin) and callable(delete_impl): - await begin() + if isinstance(memory_store, SQLiteMemoryStore): + await memory_store._begin() try: for mid in to_delete: - await delete_impl(mid, skip_tx=True) - kg.remove_memory_from_tags(mid) - deleted += 1 + await memory_store._delete_impl(mid, skip_tx=True) await memory_store.db.commit() except BaseException: await memory_store.db.rollback() raise + # Commit succeeded — mutate the in-memory graph and bump + # deleted afterwards so a SQL rollback can never diverge the + # graph or inflate the count (Issue #36; Gemini review). + # The narrow commit → save() crash window may leave a KG→SQL + # orphan that reconcile() sweeps on the next start. + for mid in to_delete: + kg.remove_memory_from_tags(mid) + deleted += 1 else: for mid in to_delete: await memory_store.delete(mid) @@ -346,6 +356,12 @@ async def get_stats( kg: KnowledgeGraph = Depends(get_knowledge_graph), scheduler: SchedulerManager = Depends(get_scheduler), ) -> dict[str, Any]: + """Return an aggregate health snapshot for the Web Console dashboard. + + Reports per-partition memory counts, the total memory count, an export of + the knowledge graph (tags + edges), and the scheduler's run status so the + console can render progress and signal scheduler health. + """ partitions = await partition_store.list() partition_stats = [ {"id": p.id, "name": p.name, "memory_count": p.memory_count, "enabled": p.enabled} for p in partitions diff --git a/src/hebb/storage/factory.py b/src/hebb/storage/factory.py index 60b65da..cac3856 100644 --- a/src/hebb/storage/factory.py +++ b/src/hebb/storage/factory.py @@ -63,6 +63,7 @@ async def _create_sqlite(settings: Settings) -> StorageContext: partition_store = SQLitePartitionStore(db, write_lock=shared_lock) async def close() -> None: + """Close the SQLite connection pool backing the storage context.""" await db.close() logger.info("Storage backend: SQLite (%s)", settings.db_path) @@ -98,6 +99,7 @@ async def _create_postgresql(settings: Settings) -> StorageContext: partition_store = PGPartitionStore(pool) async def close() -> None: + """Close the PostgreSQL connection pool backing the storage context.""" await pool.close() logger.info("Storage backend: PostgreSQL (pool %d-%d)", settings.pg_pool_min, settings.pg_pool_max) diff --git a/src/hebb/storage/partition_store.py b/src/hebb/storage/partition_store.py index 2c29e78..3081c7b 100644 --- a/src/hebb/storage/partition_store.py +++ b/src/hebb/storage/partition_store.py @@ -47,6 +47,16 @@ async def _begin(self) -> None: await self.db.execute("BEGIN IMMEDIATE") async def create(self, data: PartitionCreate, is_system: bool = False) -> Partition: + """Insert a new partition row under the write lock and return the created row. + + Args: + data: Partition fields to insert. + is_system: When true the row is flagged as a non-deletable system partition + (defaults + per-scenario fixtures). + + Returns: + The freshly created ``Partition`` (with memory_count populated). + """ now = _now_iso() async with self._write_lock: await self._begin() @@ -63,6 +73,12 @@ async def create(self, data: PartitionCreate, is_system: bool = False) -> Partit return await self.get(data.id) # type: ignore[return-value] async def get(self, partition_id: str) -> Partition | None: + """Fetch a single partition by id with its current memory count. + + Returns: + The ``Partition`` (with ``memory_count`` populated) or ``None`` if no + row matches ``partition_id``. + """ cursor = await self.db.execute( "SELECT * FROM partitions WHERE id = ?", (partition_id,), @@ -84,6 +100,11 @@ async def get(self, partition_id: str) -> Partition | None: return partition async def list(self) -> list[Partition]: + """List every partition ordered so ``mem_hippocampus`` sorts first, then by creation time. + + Each row carries its current ``memory_count`` (a separate COUNT query per + row; fine for the small partition cardinality we run with). + """ cursor = await self.db.execute( "SELECT * FROM partitions ORDER BY (id = 'mem_hippocampus') DESC, created_at", ) @@ -102,6 +123,16 @@ async def list(self) -> list[Partition]: return partitions async def update(self, partition_id: str, data: PartitionUpdate) -> Partition | None: + """Apply the non-None fields of ``data`` to the partition in one transaction. + + Args: + partition_id: The partition to patch. + data: Only the provided (non-None) fields are written. + + Returns: + The updated ``Partition``, the existing row unchanged if ``data`` had + nothing to apply, or ``None`` if no row matches ``partition_id``. + """ existing = await self.get(partition_id) if not existing: return None @@ -139,6 +170,12 @@ async def update(self, partition_id: str, data: PartitionUpdate) -> Partition | return await self.get(partition_id) async def delete(self, partition_id: str) -> bool: + """Delete a non-system partition row by id in one transaction. + + Returns: + ``True`` when a non-system row was deleted, ``False`` when the id is + missing or refers to a system partition (those are protected). + """ # Cannot delete system partitions existing = await self.get(partition_id) if not existing or existing.is_system: diff --git a/src/hebb/storage/purge.py b/src/hebb/storage/purge.py index cc4f106..df997d7 100644 --- a/src/hebb/storage/purge.py +++ b/src/hebb/storage/purge.py @@ -18,8 +18,8 @@ async def purge_memory( - store: "MemoryStore", - kg: "KnowledgeGraph", + store: MemoryStore, + kg: KnowledgeGraph, memory_id: str, *, save: bool = True, @@ -47,15 +47,37 @@ async def purge_memory( Graph cleanup runs regardless of the row's existence, so a pre-existing orphan referencing ``memory_id`` is still swept. """ - # SQLite backend: _write_lock is the shared process-level lock (= kg.lock). - # Acquire it once so the SQL delete + KG mutation form one critical section. + # SQLite backend exposes a process-level ``_write_lock`` that, in the + # production wiring (api.py / app.py), is the same object as ``kg.lock``. + # Acquire ``kg.lock`` explicitly so the SQL delete and the graph mutation + # always form one critical section — even when an SDK / test caller wires a + # different lock into the store. ``asyncio.Lock`` is non-reentrant, so when + # the two locks are identical we take it only once (a second ``async with`` + # on the same object would deadlock). shared_lock = getattr(store, "_write_lock", None) if shared_lock is not None and hasattr(store, "_delete_impl"): - async with shared_lock: - deleted = await store._delete_impl(memory_id) # type: ignore[attr-defined] - kg.remove_memory_from_tags(memory_id) - if save: - kg.save() + # ``_delete_impl`` already COMMITTED before we touch the graph, so the SQL + # rollback path can never diverge the in-memory graph. The narrow commit + # → save() crash window may leave a KG→SQL orphan that ``reconcile()`` + # sweeps on the next start (KG→SQL orphans are self-healing; SQL→KG lost + # forward refs are not, so ordering the graph mutation after the commit is + # the safer trade-off here — Issue #36; Gemini/Copilot review). + if shared_lock is kg.lock: + async with shared_lock: + deleted = await store._delete_impl(memory_id) + kg.remove_memory_from_tags(memory_id) + if save: + kg.save() + else: + # SDK / test wiring where the store's write lock differs from kg.lock: + # acquire both so a concurrent KG operation cannot race the mutation. + # Always acquire in the same documented order (store lock → kg lock) + # to avoid lock-ordering deadlocks with other paths in this codebase. + async with shared_lock, kg.lock: + deleted = await store._delete_impl(memory_id) + kg.remove_memory_from_tags(memory_id) + if save: + kg.save() else: # PostgreSQL pool backend: no shared lock; KG guarded by its own lock. deleted = await store.delete(memory_id) diff --git a/src/hebb/storage/sqlite_store.py b/src/hebb/storage/sqlite_store.py index 2df4294..ebf2ef4 100644 --- a/src/hebb/storage/sqlite_store.py +++ b/src/hebb/storage/sqlite_store.py @@ -110,6 +110,17 @@ async def _get_vec_dim(self) -> int | None: return self._vec_dim async def create(self, data: MemoryCreate, embedding: _List[float] | None = None) -> Memory: + """Insert a memory row + (optional) embedding + FTS5 index in one transaction. + + Args: + data: Memory fields to insert. + embedding: Optional embedding vector; length must match the declared + vec0 width or an ``EmbeddingDimensionError`` is raised (rolled back). + + Returns: + The freshly created ``Memory`` (re-read by id so timestamps + defaults + reflect what landed in the row). + """ async with self._write_lock: return await self._create_impl(data, embedding) @@ -173,6 +184,7 @@ async def _create_impl( return await self.get(memory_id) # type: ignore[return-value] async def get(self, memory_id: str) -> Memory | None: + """Fetch a single memory by id, or ``None`` if no row matches.""" cursor = await self.db.execute( "SELECT * FROM memories WHERE id = ?", (memory_id,), @@ -187,6 +199,19 @@ async def list( offset: int = 0, limit: int = 50, ) -> tuple[_List[Memory], int]: + """List memories optionally filtered by partition and tags, newest-first. + + Args: + partition_id: When provided, restricts the result to one partition. + tags: When provided, keeps only memories whose tags intersect the set + (post-query JSON filter — not a SQL predicate, so the returned + ``total`` is the surviving count, not the pre-filter count). + offset: Result offset (LIMIT/OFFSET on the SQL side). + limit: Maximum rows to return. + + Returns: + A ``(rows, total)`` pair sorted by ``created_at`` descending. + """ conditions: _List[str] = [] params: _List[Any] = [] @@ -218,6 +243,12 @@ async def list( return memories, total async def update(self, memory_id: str, data: MemoryUpdate) -> Memory | None: + """Apply the non-None ``data`` fields in one transaction under the write lock. + + Syncs the FTS5 index when ``content`` changed. Returns the updated memory + or ``None`` if no row matches ``memory_id`` (or if ``data`` carried no + fields to apply, in which case the existing row is returned unchanged). + """ async with self._write_lock: return await self._update_impl(memory_id, data) @@ -269,6 +300,13 @@ async def _update_impl(self, memory_id: str, data: MemoryUpdate) -> Memory | Non return await self.get(memory_id) async def delete(self, memory_id: str) -> bool: + """Delete a memory row + its embedding row + its FTS5 row in one transaction. + + Returns ``True`` if the base row existed (and was deleted), ``False`` + if no row matched ``memory_id``. The embedding/FTS cleanup runs + unconditionally so a stale auxiliary row is swept even when the base row + was already gone. + """ async with self._write_lock: return await self._delete_impl(memory_id) @@ -507,6 +545,7 @@ async def _get_many(self, memory_ids: _List[str]) -> dict[str, Memory]: return {row["id"]: _row_to_memory(row) for row in rows} async def get_by_partition(self, partition_id: str) -> _List[Memory]: + """Fetch every memory in a partition, newest-first (no pagination, no count).""" cursor = await self.db.execute( "SELECT * FROM memories WHERE partition_id = ? ORDER BY created_at DESC", (partition_id,), @@ -579,6 +618,12 @@ def _sort_key(m: Memory) -> int: return results async def delete_expired(self) -> _List[str]: + """Delete every memory whose ``expires_at`` has passed, in one transaction. + + Returns the list of deleted memory ids (so callers can sync the + knowledge graph / run trackers afterwards). Embedding and FTS5 rows are + cleaned up alongside the base row. + """ async with self._write_lock: return await self._delete_expired_impl() @@ -619,6 +664,11 @@ async def _delete_expired_impl(self) -> _List[str]: return expired_ids async def update_access(self, memory_id: str) -> None: + """Bump ``access_count`` and ``last_accessed_at`` for one memory in one transaction. + + No-op (still commits) if no row matches ``memory_id``; callers that need + to know should ``get()`` first. + """ async with self._write_lock: await self._update_access_impl(memory_id) @@ -663,6 +713,13 @@ async def _update_access_batch_impl(self, memory_ids: _List[str]) -> None: raise async def update_embedding(self, memory_id: str, embedding: _List[float]) -> None: + """Overwrite a memory's embedding in one transaction. + + The embedding's length must match the declared vec0 width, otherwise an + ``EmbeddingDimensionError`` is raised and nothing is written. An empty + embedding is silently skipped (the storage invariant is "0-dim if unset", + and a call with no values is treated as no-op). + """ if not embedding: return expected = await self._get_vec_dim() @@ -783,6 +840,7 @@ async def _update_content_and_embedding_impl( return await self.get(memory_id) async def update_expiry(self, memory_id: str, expires_at: str) -> None: + """Set ``expires_at`` (ISO timestamp) for one memory in one transaction.""" async with self._write_lock: await self._update_expiry_impl(memory_id, expires_at) diff --git a/tests/conftest.py b/tests/conftest.py index f315561..a567aa1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,9 +9,11 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncIterator from pathlib import Path from unittest.mock import AsyncMock +import aiosqlite import pytest import pytest_asyncio @@ -51,7 +53,7 @@ def shared_lock() -> asyncio.Lock: @pytest_asyncio.fixture -async def db(settings: Settings, shared_lock: asyncio.Lock): +async def db(settings: Settings, shared_lock: asyncio.Lock) -> AsyncIterator[aiosqlite.Connection]: conn = await get_connection(settings.db_path) await initialize_schema(conn, settings.embedding_dim) yield conn @@ -59,12 +61,16 @@ async def db(settings: Settings, shared_lock: asyncio.Lock): @pytest_asyncio.fixture -async def memory_store(db, shared_lock: asyncio.Lock) -> SQLiteMemoryStore: +async def memory_store( + db: aiosqlite.Connection, shared_lock: asyncio.Lock +) -> SQLiteMemoryStore: return SQLiteMemoryStore(db, write_lock=shared_lock) @pytest_asyncio.fixture -async def partition_store(db, shared_lock: asyncio.Lock) -> SQLitePartitionStore: +async def partition_store( + db: aiosqlite.Connection, shared_lock: asyncio.Lock +) -> SQLitePartitionStore: store = SQLitePartitionStore(db, write_lock=shared_lock) await store.ensure_defaults() return store diff --git a/tests/integration/agents/test_consolidation_agent.py b/tests/integration/agents/test_consolidation_agent.py index 61f3bb8..984d550 100644 --- a/tests/integration/agents/test_consolidation_agent.py +++ b/tests/integration/agents/test_consolidation_agent.py @@ -3,19 +3,33 @@ from __future__ import annotations import asyncio +from pathlib import Path import pytest from hebb.agents.consolidation_agent import ConsolidationAgent +from hebb.agents.llm_client import LLMClient from hebb.agents.recall_agent import RecallAgent +from hebb.config.settings import Settings +from hebb.embedding.local import NoopEmbedder from hebb.graph.knowledge_graph import KnowledgeGraph -from hebb.models.memory import MemoryCreate, MemoryMetadata +from hebb.models.memory import Memory, MemoryCreate, MemoryMetadata from hebb.retrieval.searcher import MemorySearcher +from hebb.storage.partition_store import SQLitePartitionStore +from hebb.storage.sqlite_store import SQLiteMemoryStore class TestConsolidationAgent: @pytest.mark.asyncio - async def test_consolidate_memory(self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path): + async def test_consolidate_memory( + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> None: """Consolidation should move memory from hebb to target partition.""" kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) @@ -25,7 +39,7 @@ async def test_consolidate_memory(self, mock_llm, memory_store, partition_store, ) # Mock LLM: recall returns no queries, consolidation returns classification - mock_llm.complete_json.side_effect = [ + mock_llm.complete_json.side_effect = [ # type: ignore[attr-defined] # RecallAgent.recall -> generate queries {"queries": ["user preference"]}, # ConsolidationAgent -> classify memory @@ -73,8 +87,14 @@ async def test_consolidate_memory(self, mock_llm, memory_store, partition_store, @pytest.mark.asyncio async def test_session_consolidation_preserves_turn_and_timestamp( - self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path - ): + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> None: """Session consolidation must carry the source turns' turn span and earliest timestamp onto the consolidated memory (so turn-window expansion and temporal_boost still work on consolidated output).""" @@ -99,7 +119,7 @@ async def test_session_consolidation_preserves_turn_and_timestamp( ) ) - mock_llm.complete_json.side_effect = [ + mock_llm.complete_json.side_effect = [ # type: ignore[attr-defined] {"queries": ["sweden book"]}, # RecallAgent { "memories": [ @@ -137,7 +157,7 @@ async def test_session_consolidation_preserves_turn_and_timestamp( assert meta["turn_pair"] == [0, 1] assert meta["timestamp"] == "2023-08-01T10:00:00+00:00" - async def _session_pair(self, memory_store): + async def _session_pair(self, memory_store: SQLiteMemoryStore) -> tuple[Memory, Memory]: m1 = await memory_store.create( MemoryCreate( content="turn one", @@ -154,7 +174,15 @@ async def _session_pair(self, memory_store): ) return m1, m2 - def _agent(self, mock_llm, memory_store, partition_store, noop_embedder, kg, settings=None): + def _agent( + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + kg: KnowledgeGraph, + settings: Settings | None = None, + ) -> ConsolidationAgent: recall_agent = RecallAgent( llm=mock_llm, searcher=MemorySearcher(store=memory_store, embedder=noop_embedder), @@ -171,14 +199,20 @@ def _agent(self, mock_llm, memory_store, partition_store, noop_embedder, kg, set @pytest.mark.asyncio async def test_session_wellformed_empty_drains_sources( - self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path - ): + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> None: """A well-formed empty result ({"memories": []}) = the LLM deliberately judged the turns low-value, so the sources are drained from the inbox (default behavior) — otherwise they replay to the LLM forever.""" kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1, m2 = await self._session_pair(memory_store) - mock_llm.complete_json.side_effect = [ + mock_llm.complete_json.side_effect = [ # type: ignore[attr-defined] {"queries": []}, # RecallAgent {"memories": []}, # well-formed empty: nothing worth keeping ] @@ -196,15 +230,22 @@ async def test_session_wellformed_empty_drains_sources( @pytest.mark.asyncio async def test_session_wellformed_empty_kept_when_flag_off( - self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path - ): + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> None: """With consolidation_drain_empty_sources=False, a well-formed empty result keeps the sources (legacy behavior, opt-out).""" from hebb.config.settings import Settings + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1, m2 = await self._session_pair(memory_store) - mock_llm.complete_json.side_effect = [ + mock_llm.complete_json.side_effect = [ # type: ignore[attr-defined] {"queries": []}, {"memories": []}, ] @@ -221,14 +262,20 @@ async def test_session_wellformed_empty_kept_when_flag_off( @pytest.mark.asyncio async def test_session_parse_failure_keeps_sources( - self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path - ): + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> None: """A garbled/unparseable LLM response (complete_json -> {}) must NEVER drain — it has no "memories" key, so the sources are kept for retry. Guards against deleting data on a transient/parse failure (F4).""" kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1, m2 = await self._session_pair(memory_store) - mock_llm.complete_json.side_effect = [ + mock_llm.complete_json.side_effect = [ # type: ignore[attr-defined] {"queries": []}, {}, # parse failure: _parse_json returns {} on unparseable output ] @@ -241,7 +288,15 @@ async def test_session_parse_failure_keeps_sources( assert await memory_store.get(m2.id) is not None @pytest.mark.asyncio - async def test_consolidate_batch_empty(self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path): + async def test_consolidate_batch_empty( + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> None: """Batch consolidation with no hebb memories returns empty list.""" kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) @@ -263,8 +318,14 @@ async def test_consolidate_batch_empty(self, mock_llm, memory_store, partition_s @pytest.mark.asyncio async def test_consolidate_handles_conflict_discard( - self, mock_llm, memory_store, partition_store, noop_embedder, shared_lock, tmp_path - ): + self, + mock_llm: LLMClient, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> None: """When LLM says discard, the memory is deleted without creating a new one.""" kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) @@ -272,7 +333,7 @@ async def test_consolidate_handles_conflict_discard( MemoryCreate(content="redundant info", partition_id="mem_hippocampus"), ) - mock_llm.complete_json.side_effect = [ + mock_llm.complete_json.side_effect = [ # type: ignore[attr-defined] {"queries": []}, { "target_partition": "mem_semantic", diff --git a/tests/integration/scheduler/test_scheduler_manager.py b/tests/integration/scheduler/test_scheduler_manager.py index c968003..2538f33 100644 --- a/tests/integration/scheduler/test_scheduler_manager.py +++ b/tests/integration/scheduler/test_scheduler_manager.py @@ -2,17 +2,30 @@ from __future__ import annotations +import asyncio +from pathlib import Path + import pytest from apscheduler.triggers.cron import CronTrigger +from hebb.config.settings import Settings from hebb.embedding.local import NoopEmbedder from hebb.graph.knowledge_graph import KnowledgeGraph from hebb.scheduler.manager import SchedulerManager +from hebb.storage.partition_store import SQLitePartitionStore +from hebb.storage.sqlite_store import SQLiteMemoryStore class TestSchedulerManager: @pytest.fixture - def scheduler(self, settings, memory_store, partition_store, shared_lock, tmp_path): + def scheduler( + self, + settings: Settings, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + shared_lock: asyncio.Lock, + tmp_path: Path, + ) -> SchedulerManager: kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) embedder = NoopEmbedder(384) return SchedulerManager( diff --git a/tests/integration/storage/test_purge.py b/tests/integration/storage/test_purge.py index d02ce4b..1946ab6 100644 --- a/tests/integration/storage/test_purge.py +++ b/tests/integration/storage/test_purge.py @@ -7,15 +7,23 @@ from __future__ import annotations +import asyncio +from pathlib import Path + import pytest from hebb.graph.knowledge_graph import KnowledgeGraph from hebb.models.memory import MemoryCreate from hebb.storage.purge import purge_memory +from hebb.storage.sqlite_store import SQLiteMemoryStore @pytest.mark.asyncio -async def test_purge_removes_sql_row_and_graph_node(memory_store, shared_lock, tmp_path): +async def test_purge_removes_sql_row_and_graph_node( + memory_store: SQLiteMemoryStore, + shared_lock: asyncio.Lock, + tmp_path: Path, +) -> None: kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) mem = await memory_store.create( MemoryCreate(content="graphed memory", partition_id="mem_hippocampus") @@ -32,7 +40,11 @@ async def test_purge_removes_sql_row_and_graph_node(memory_store, shared_lock, t @pytest.mark.asyncio -async def test_purge_cleans_graph_even_when_row_absent(memory_store, shared_lock, tmp_path): +async def test_purge_cleans_graph_even_when_row_absent( + memory_store: SQLiteMemoryStore, + shared_lock: asyncio.Lock, + tmp_path: Path, +) -> None: """Graph cleanup runs regardless of the row's existence, so a stale reference is swept rather than left dangling.""" kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) diff --git a/tests/unit/test_audit_consolidation.py b/tests/unit/test_audit_consolidation.py index baf61bf..7f8ad9b 100644 --- a/tests/unit/test_audit_consolidation.py +++ b/tests/unit/test_audit_consolidation.py @@ -13,6 +13,7 @@ from __future__ import annotations +import asyncio from datetime import datetime, timedelta, timezone from pathlib import Path @@ -20,6 +21,8 @@ from hebb.agents.consolidation_agent import ConsolidationAgent from hebb.agents.recall_agent import RecallAgent +from hebb.config.settings import Settings +from hebb.embedding.local import NoopEmbedder from hebb.graph.knowledge_graph import KnowledgeGraph from hebb.models.memory import Memory, MemoryCreate from hebb.retrieval.searcher import MemorySearcher @@ -29,6 +32,8 @@ compute_expires_at, forget_idle_days, ) +from hebb.storage.partition_store import SQLitePartitionStore +from hebb.storage.sqlite_store import SQLiteMemoryStore # Representative parameter set (the global user-partition defaults). _PARAMS = dict(half_life_days=60.0, k_importance=2.0, k_access=1.5, threshold=0.3) @@ -218,7 +223,7 @@ async def _spy(memory_id: str, data, embedding: list[float]) -> None: reembedded.append(memory_id) await original_impl(memory_id, data, embedding) - memory_store._update_content_and_embedding_impl = _spy # type: ignore[method-assign] + memory_store._update_content_and_embedding_impl = _spy mock_llm.complete_json.side_effect = [ # RecallAgent.recall -> query generation (returns the existing mem) @@ -253,7 +258,13 @@ async def _spy(memory_id: str, data, embedding: list[float]) -> None: class TestConsolidationLLMModelGuard: @pytest.mark.asyncio async def test_run_consolidation_skips_without_llm_model( - self, settings, memory_store, partition_store, noop_embedder, shared_lock, tmp_path + self, + settings: Settings, + memory_store: SQLiteMemoryStore, + partition_store: SQLitePartitionStore, + noop_embedder: NoopEmbedder, + shared_lock: asyncio.Lock, + tmp_path: Path, ) -> None: """No llm_model -> run_consolidation returns [] without calling litellm.""" settings.llm_model = None diff --git a/tests/unit/test_audit_storage.py b/tests/unit/test_audit_storage.py index 1b38c5f..53e46a5 100644 --- a/tests/unit/test_audit_storage.py +++ b/tests/unit/test_audit_storage.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncIterator import pytest @@ -31,7 +32,7 @@ def embedding_dim() -> int: @pytest.fixture -async def store(tmp_path, embedding_dim: int) -> SQLiteMemoryStore: +async def store(tmp_path, embedding_dim: int) -> AsyncIterator[SQLiteMemoryStore]: db = await get_connection(str(tmp_path / "audit.db")) await initialize_schema(db, embedding_dim) yield SQLiteMemoryStore(db) @@ -41,7 +42,7 @@ async def store(tmp_path, embedding_dim: int) -> SQLiteMemoryStore: async def _count_memories(store: SQLiteMemoryStore) -> int: cursor = await store.db.execute("SELECT count(*) FROM memories") row = await cursor.fetchone() - return int(row[0]) if row else 0 + return int(row[0]) if row is not None else 0 class TestCreateAtomicity: @@ -52,7 +53,7 @@ async def test_create_rolls_back_on_fts_error( """A failure on the FTS insert must roll back the memories row too.""" real_execute = store.db.execute - async def flaky_execute(sql, *args, **kwargs): # type: ignore[no-untyped-def] + async def flaky_execute(sql: str, *args: object, **kwargs: object): if "memory_fts" in sql and sql.strip().upper().startswith("INSERT"): raise RuntimeError("forced FTS failure") return await real_execute(sql, *args, **kwargs) @@ -106,6 +107,7 @@ async def test_concurrent_creates_do_not_interleave( # FTS index stayed in lockstep with the base table. cursor = await store.db.execute("SELECT count(*) FROM memory_fts") row = await cursor.fetchone() + assert row is not None assert int(row[0]) == n @@ -150,7 +152,7 @@ class TestConcurrentWriteConsolidateForget: @pytest.fixture async def shared_store_and_kg( self, tmp_path, embedding_dim: int - ) -> tuple[SQLiteMemoryStore, KnowledgeGraph, asyncio.Lock]: + ) -> AsyncIterator[tuple[SQLiteMemoryStore, KnowledgeGraph, asyncio.Lock]]: lock = asyncio.Lock() db = await get_connection(str(tmp_path / "stress.db")) await initialize_schema(db, embedding_dim) @@ -161,52 +163,99 @@ async def shared_store_and_kg( @pytest.mark.asyncio async def test_no_orphans_after_concurrent_operations( - self, shared_store_and_kg + self, + shared_store_and_kg: tuple[SQLiteMemoryStore, KnowledgeGraph, asyncio.Lock], ) -> None: store, kg, lock = shared_store_and_kg # Pre-populate: 10 memories with tags, half flagged for "forgetting". initial_ids: list[str] = [] for i in range(10): - m = await store.create( - MemoryCreate(partition_id="p1", content=f"seed-{i}", tags=[f"tag-{i}"]), - ) - initial_ids.append(m.id) async with lock: + await store._begin() + try: + m = await store._create_impl( + MemoryCreate( + partition_id="p1", content=f"seed-{i}", tags=[f"tag-{i}"] + ), + skip_tx=True, + ) + await store.db.commit() + except BaseException: + await store.db.rollback() + raise kg.update_from_tags([f"tag-{i}"], m.id) kg.save() + initial_ids.append(m.id) # IDs to be "forgotten" (first 5). forget_ids = initial_ids[:5] async def api_write() -> None: - """Simulate REST API writes interleaving with background tasks.""" + """Simulate REST API writes interleaving with background tasks. + + Uses the composite create path (``_create_impl`` + explicit BEGIN) + under the shared lock, with the KG mutation applied AFTER the commit + succeeds — mirroring the production ordering the review enforced + (Issue #36; Gemini review of in-memory graph divergence on rollback). + """ for i in range(10): - m = await store.create( - MemoryCreate(partition_id="p1", content=f"api-{i}", tags=[f"api-tag-{i}"]), - ) async with lock: + await store._begin() + try: + m = await store._create_impl( + MemoryCreate( + partition_id="p1", + content=f"api-{i}", + tags=[f"api-tag-{i}"], + ), + skip_tx=True, + ) + await store.db.commit() + except BaseException: + await store.db.rollback() + raise kg.update_from_tags([f"api-tag-{i}"], m.id) kg.save() async def consolidation_sim() -> None: - """Simulate a consolidation step: create new + delete source (atomic).""" + """Simulate a consolidation step: create new + delete source in one + SQL transaction under the shared lock, then sync the KG afterwards.""" for i in range(5): source_id = initial_ids[5 + i] # second half as consolidation sources async with lock: - new_mem = await store._create_impl( - MemoryCreate(partition_id="p2", content=f"consolidated-{i}", tags=[f"con-tag-{i}"]), - ) + await store._begin() + try: + new_mem = await store._create_impl( + MemoryCreate( + partition_id="p2", + content=f"consolidated-{i}", + tags=[f"con-tag-{i}"], + ), + skip_tx=True, + ) + await store._delete_impl(source_id, skip_tx=True) + await store.db.commit() + except BaseException: + await store.db.rollback() + raise kg.update_from_tags([f"con-tag-{i}"], new_mem.id) kg.remove_memory_from_tags(source_id) - await store._delete_impl(source_id) kg.save() async def forgetting_sim() -> None: - """Simulate a forgetting sweep: delete expired memories (atomic).""" + """Simulate a forgetting sweep: delete expired memories in one SQL + transaction under the shared lock, then strip KG refs afterwards.""" async with lock: + await store._begin() + try: + for mid in forget_ids: + await store._delete_impl(mid, skip_tx=True) + await store.db.commit() + except BaseException: + await store.db.rollback() + raise for mid in forget_ids: - await store._delete_impl(mid) kg.remove_memory_from_tags(mid) kg.save() @@ -226,11 +275,122 @@ async def forgetting_sim() -> None: # - forgetting_sim deleted 5 memories from p1. # Initial: 10 in p1. After: 10 (api) + 0 (5 deleted by consolidation, 5 by forget) in p1, # 5 in p2. - p1_mems, p1_count = await store.list(partition_id="p1") - p2_mems, p2_count = await store.list(partition_id="p2") + _, p1_count = await store.list(partition_id="p1") + _, p2_count = await store.list(partition_id="p2") assert p1_count == 10, f"Expected 10 in p1, got {p1_count}" assert p2_count == 5, f"Expected 5 in p2, got {p2_count}" # 3. KG file can be loaded cleanly. kg2 = KnowledgeGraph(kg.path) assert kg2.graph.number_of_nodes() > 0 + + @pytest.mark.asyncio + async def test_rollback_on_injected_failure_in_consolidation( + self, + shared_store_and_kg: tuple[SQLiteMemoryStore, KnowledgeGraph, asyncio.Lock], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """An injected SQL failure mid-consolidation must leave both the SQL store + AND the in-memory KG unchanged — no half-applied rows, no half-stripped KG + refs (Issue #36; Gemini review of in-memory graph divergence on rollback). + + Reproduces the exact ordering the consolidated fast-path now uses: + ``_begin`` → ``_create_impl(skip_tx=True)`` → ``_delete_impl(skip_tx=True)`` → + ``commit`` → KG mutations, with KG mutations running only after the commit + succeeds. If the delete fails, ``db.rollback()`` reverts SQL and the graph + must not have been mutated, so a reloaded graph matches the pre-failure + snapshot exactly. + """ + store, kg, lock = shared_store_and_kg + + # Seed one target memory (will be deleted during consolidation) tagged so + # we can assert the tag survives the rollback. + async with lock: + await store._begin() + try: + target = await store._create_impl( + MemoryCreate(partition_id="p1", content="target", tags=["kept-tag"]), + skip_tx=True, + ) + await store.db.commit() + except BaseException: + await store.db.rollback() + raise + kg.update_from_tags(["kept-tag"], target.id) + kg.save() + + # Snapshot SQL row count and on-disk KG state before injection. + rows_before = await _count_memories(store) + kg_before_path = kg.path + kg_before_text = kg_before_path.read_text() if kg_before_path.exists() else "" + + # Inject a failure on the second SQL call inside the consolidation path + # (the source delete) so that the create has already executed but commit + # never happens. + original_delete = store._delete_impl + delete_calls = 0 + + async def flaky_delete(memory_id: str, *, skip_tx: bool = False) -> bool: + nonlocal delete_calls + delete_calls += 1 + if delete_calls == 1: + raise RuntimeError("injected delete failure") + return await original_delete(memory_id, skip_tx=skip_tx) + + monkeypatch.setattr(store, "_delete_impl", flaky_delete) + + # Drive the same critical section the production path uses. The lock is + # acquired explicitly so the test exercises the real ordering. + with pytest.raises(RuntimeError, match="injected delete failure"): + async with lock: + await store._begin() + try: + await store._create_impl( + MemoryCreate( + partition_id="p2", + content="consolidated", + tags=["new-tag"], + ), + skip_tx=True, + ) + # This raises — the create must NOT commit. + await store._delete_impl(target.id, skip_tx=True) + await store.db.commit() + except BaseException: + await store.db.rollback() + raise + # These lines must NOT execute when the delete raises. + kg.update_from_tags(["new-tag"], "") + kg.remove_memory_from_tags(target.id) + kg.save() + + # 1. SQL row count unchanged: the rollback reverted the create. + assert await _count_memories(store) == rows_before, ( + "rollback failed: rows committed despite the injected failure" + ) + # 2. The seeded target row still exists. + assert await store.get(target.id) is not None, "target row was lost despite rollback" + + # 3. KG on disk unchanged: save() never ran after the commit failed. + assert ( + kg.path.read_text() if kg.path.exists() else "" + ) == kg_before_text, "KG was persisted despite the rollback" + + # 4. In-memory KG state unchanged: the tag still references the target. + assert kg.get_tag("kept-tag") is not None, "KG lost the pre-failure tag" + referenced: list[str] = [] + for node_id in kg.graph.nodes: + referenced.extend(kg.graph.nodes[node_id].get("memory_ids", [])) + assert target.id in referenced, "KG dropped the target ref on rollback" + + # 5. Reloaded KG from disk matches the in-memory graph — divergence would + # not be self-healing, so the assertion catches it before the orphan + # can be persisted by a later save. + reloaded = KnowledgeGraph(kg.path) + for node_id in kg.graph.nodes: + assert node_id in reloaded.graph.nodes, f"node {node_id} missing from reloaded KG" + in_mem_refs = set(kg.graph.nodes[node_id].get("memory_ids", [])) + on_disk_refs = set(reloaded.graph.nodes[node_id].get("memory_ids", [])) + assert in_mem_refs == on_disk_refs, ( + f"in-memory KG diverged from on-disk KG for node {node_id}" + ) From e4eef52687f9fd94cafdb65818e037171e61c109 Mon Sep 17 00:00:00 2001 From: Huang Date: Mon, 27 Jul 2026 11:35:08 +0800 Subject: [PATCH 3/6] fix(storage): add partition_id to sqlite-vec fallback table schema --- src/hebb/storage/migrations.py | 48 ++++++++++++++++++++++++++++++++-- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/src/hebb/storage/migrations.py b/src/hebb/storage/migrations.py index c1428ad..03ddfc5 100644 --- a/src/hebb/storage/migrations.py +++ b/src/hebb/storage/migrations.py @@ -219,6 +219,40 @@ async def _ensure_vec_table(db: aiosqlite.Connection, embedding_dim: int) -> Non await db.execute(_VEC_CREATE_SQL.format(dim=dim)) +async def _ensure_fallback_partition_column(db: aiosqlite.Connection) -> None: + """Ensure the non-vec0 fallback ``memory_embeddings`` table carries ``partition_id``. + + Only acts on the regular (BLOB) fallback table — an existing vec0 virtual + table is migrated by :func:`_ensure_vec_table`. Writes always supply + ``partition_id`` (the partition-aware vector-KNN contract), so a fallback + table created before that contract needs an additive ``ALTER`` rather than + a destructive recreate, else every insert fails with + ``no such column: partition_id`` in vec0-unavailable environments. + + Args: + db: Open database connection. + """ + cursor = await db.execute( + "SELECT sql FROM sqlite_master WHERE name = 'memory_embeddings'" + ) + row = await cursor.fetchone() + if not row or not row[0] or "USING vec0" in row[0]: + # vec0 virtual table (or absent) — schema maintained by _ensure_vec_table. + return + + info_cursor = await db.execute("PRAGMA table_info(memory_embeddings)") + columns = {r[1] for r in await info_cursor.fetchall()} + if "partition_id" not in columns: + await db.execute( + "ALTER TABLE memory_embeddings " + "ADD COLUMN partition_id TEXT NOT NULL DEFAULT 'default'" + ) + await db.execute( + "CREATE INDEX IF NOT EXISTS idx_memory_embeddings_partition_id " + "ON memory_embeddings(partition_id)" + ) + + async def initialize_schema( db: aiosqlite.Connection, embedding_dim: int = 384, *, create_vec_table: bool = True ) -> None: @@ -267,11 +301,21 @@ async def initialize_schema( raise except Exception as e: logger.warning("Could not create vec0 table: %s. %s", e, _VEC_FIX_HINT) - # Fallback: regular table so embedding INSERT/DELETE still works + # Fallback: regular table so embedding INSERT/DELETE still works. + # Schema mirrors the vec0 contract (memory_id, partition_id, + # embedding) so partition-aware writes don't fail in + # vec0-unavailable environments. await db.execute( - "CREATE TABLE IF NOT EXISTS memory_embeddings (memory_id TEXT PRIMARY KEY, embedding BLOB)" + "CREATE TABLE IF NOT EXISTS memory_embeddings (" + "memory_id TEXT PRIMARY KEY, " + "partition_id TEXT NOT NULL DEFAULT 'default', " + "embedding BLOB)" ) + # Additive migration for fallback tables created before the partition_id + # write contract (vec0 tables are skipped inside). Idempotent. + await _ensure_fallback_partition_column(db) + # FTS5 full-text search table for keyword matching. # `porter unicode61` chains Porter stemming on top of unicode61 so that # plural/inflected forms match the query stem ("researched" matches From 29ab376710ccfaf1996df8221b2f8df8146ffa68 Mon Sep 17 00:00:00 2001 From: Huang Date: Mon, 3 Aug 2026 13:46:42 +0800 Subject: [PATCH 4/6] fix(storage): enforce dim guard on vec0-unavailable fallback; pin mcp<2 The macOS-only test failures (test_dim_mismatch_raises_and_writes_nothing, test_dimension_mismatch_raises_without_opt_in, test_dimension_mismatch_drops_with_opt_in) come from a pre-existing latent gap: the embedding-dim guard is skipped on the BLOB fallback path used when sqlite-vec can't load (macOS CI), because the fallback table has no inline width to read from sqlite_master. Adding partition_id to the fallback schema made the probe INSERT succeed on the 3-column fallback and return early, bypassing the mismatch check entirely. Persist the configured embedding_dim in a new schema_meta table so both the write-path guard (SQLiteMemoryStore._get_vec_dim) and the migration mismatch check (_ensure_vec_table) can read the existing width even on the fallback path. _ensure_vec_table now also creates the BLOB fallback itself when vec0 is unavailable, so direct calls (the audit-write-api tests) work without initialize_schema's except handler. Verified on both vec0-enabled and vec0-disabled (macOS CI sim) paths. Also pin mcp>=1.0.0,<2: upstream mcp 2.0.0 (2026-07-28) removed mcp.server.fastmcp and the v1 client API (mcp.ClientSession / mcp.client.stdio.stdio_client) that hebb.mcp.server and its e2e test import, breaking e2e on all platforms. Pinning unblocks CI on main and every PR with no code change; migrating to standalone fastmcp / mcp 2.0 MCPServer is a separate follow-up. Co-Authored-By: Claude --- pyproject.toml | 6 +- src/hebb/storage/migrations.py | 171 +++++++++++++++++++++---------- src/hebb/storage/sqlite_store.py | 16 ++- 3 files changed, 138 insertions(+), 55 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1d623b4..5d484a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,11 @@ dependencies = [ "httpx>=0.27.0", "PyJWT>=2.8.0", "networkx>=3.1", - "mcp>=1.0.0", + # mcp 2.0 (2026-07-28) removed ``mcp.server.fastmcp`` (FastMCP -> MCPServer) + # and the v1 client API (``mcp.ClientSession`` / ``mcp.client.stdio.stdio_client``) + # that hebb.mcp.server and its e2e test depend on. Pin <2 until the migration + # to the standalone ``fastmcp`` package (or mcp 2.0) lands. + "mcp>=1.0.0,<2", # Repairs malformed LLM JSON (doubled braces, truncation, trailing commas). # Pure-Python, MIT, no mandatory transitive deps. "json-repair>=0.30.0", diff --git a/src/hebb/storage/migrations.py b/src/hebb/storage/migrations.py index 03ddfc5..f59743d 100644 --- a/src/hebb/storage/migrations.py +++ b/src/hebb/storage/migrations.py @@ -78,6 +78,27 @@ def _load_vec(conn: Any) -> None: "tokenize='porter unicode61')" ) +# Regular (BLOB) fallback table used when the vec0 extension can't be loaded +# (e.g. macOS python.org builds without loadable-extension support). Mirrors +# the vec0 contract (memory_id, partition_id, embedding) so partition-aware +# writes don't fail in vec0-unavailable environments. The BLOB column has no +# fixed width, so the configured width is recorded in ``schema_meta`` (see +# ``_read_meta_dim``) to preserve the dim guard on this path. +_FALLBACK_CREATE_SQL = ( + "CREATE TABLE IF NOT EXISTS memory_embeddings (" + "memory_id TEXT PRIMARY KEY, " + "partition_id TEXT NOT NULL DEFAULT 'default', " + "embedding BLOB)" +) + +# Key/value row recording the embedding width the current ``memory_embeddings`` +# table was built at. vec0 declares the width inline (``float[N]``) so this is +# only authoritative for the BLOB fallback, but it is written in both cases so +# the write-path dim guard and the migration mismatch check still work on the +# fallback path (where ``sqlite_master`` has no width to read). +_META_CREATE_SQL = "CREATE TABLE IF NOT EXISTS schema_meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)" +_META_DIM_KEY = "embedding_dim" + async def _ensure_fts_table(db: aiosqlite.Connection) -> None: """Create FTS5 with partition_id, force-rebuild if old schema present.""" @@ -138,8 +159,33 @@ async def _vec_row_count(db: aiosqlite.Connection) -> int: return int(row[0]) if row and row[0] is not None else 0 +async def _read_meta_dim(db: aiosqlite.Connection) -> int | None: + """Embedding width recorded in ``schema_meta``, or ``None`` if unset. + + Authoritative for the BLOB fallback table (which has no inline width); + informational for vec0 (which declares its width inline). Returns ``None`` + when the meta table is absent or the key is unset — e.g. first init, or a + DB created before this metadata existed. + """ + try: + cursor = await db.execute("SELECT value FROM schema_meta WHERE key = ?", (_META_DIM_KEY,)) + row = await cursor.fetchone() + except Exception: + return None + return int(row[0]) if row and row[0] is not None else None + + +async def _write_meta_dim(db: aiosqlite.Connection, dim: int) -> None: + """Record (upsert) the embedding width the current table is built at.""" + await db.execute(_META_CREATE_SQL) + await db.execute( + "INSERT INTO schema_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (_META_DIM_KEY, str(dim)), + ) + + async def _ensure_vec_table(db: aiosqlite.Connection, embedding_dim: int) -> None: - """Create or recreate vec0 table if dimension or schema changed. + """Create or recreate the vec0 table, or its BLOB fallback, at ``embedding_dim``. Schema invariant: ``(memory_id, partition_id, embedding)``. The ``partition_id`` metadata column lets retrieval push the partition @@ -153,6 +199,12 @@ async def _ensure_vec_table(db: aiosqlite.Connection, embedding_dim: int) -> Non First-time creation and the additive ``partition_id`` migration (empty or correctly-dimensioned tables) are unaffected. + The existing width is read from the vec0 declaration (``float[N]``) when the + extension is available, or from ``schema_meta`` for the BLOB fallback used + when vec0 can't be loaded (macOS python.org builds). Recording the width in + ``schema_meta`` keeps the write-path dim guard and the migration mismatch + check working on the fallback path too. + Args: db: Open database connection. embedding_dim: Target embedding width for the configured model. @@ -164,59 +216,78 @@ async def _ensure_vec_table(db: aiosqlite.Connection, embedding_dim: int) -> Non import numpy as np dim = int(embedding_dim) - # Fast path for fresh DB - try: - await db.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS memory_embeddings USING vec0(" - f"memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding float[{dim}])" - ) - except Exception: - pass - # Probe both the dimension AND the partition_id column. A pre-refactor - # table (no partition_id) succeeds at CREATE IF NOT EXISTS silently; - # this insert raises iff the column is missing or the dim mismatches. - probe = np.zeros(dim, dtype=np.float32).tobytes() - try: - await db.execute( - "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) " - "VALUES ('__schema_probe__', '__probe__', ?)", - (probe,), - ) - await db.execute("DELETE FROM memory_embeddings WHERE memory_id = '__schema_probe__'") - return - except Exception: - pass - - # Probe failed. Distinguish a benign schema gap (missing partition_id, or an - # empty table) from a TRUE dimension mismatch on data we'd be destroying. + # The width the existing table was built at. vec0 declares it inline + # (``float[N]`` in sqlite_master); the BLOB fallback (vec0 unavailable) + # records it in ``schema_meta``. ``None`` means first-time creation. declared = await _declared_vec_dim(db) + existing_dim = declared if declared is not None else await _read_meta_dim(db) row_count = await _vec_row_count(db) - true_dim_mismatch = declared is not None and declared != dim + true_dim_mismatch = existing_dim is not None and existing_dim != dim + # A true mismatch on a POPULATED table is operator-actionable: refuse before + # touching anything so we never silently drop vectors. First-time creation, + # an empty-table dim change, and the partition_id migration fall through. if true_dim_mismatch and row_count > 0 and os.getenv(_ALLOW_EMBED_DROP_ENV) != "1": raise EmbeddingDimensionMismatchError( - f"memory_embeddings has {row_count} vectors of width {declared}, but the " + f"memory_embeddings has {row_count} vectors of width {existing_dim}, but the " f"configured embedding dimension is {dim}. Dropping the table would lose " f"every embedding. Run `hebb memory reembed` to rebuild vectors at the new " f"dimension, or set {_ALLOW_EMBED_DROP_ENV}=1 to drop and rebuild empty." ) + # --- vec0 path: try the virtual table. --- + try: + if true_dim_mismatch: + # Empty-table dim change or opt-in destructive rebuild at ``dim``. + await db.execute("DROP TABLE IF EXISTS memory_embeddings") + await db.execute( + f"CREATE VIRTUAL TABLE IF NOT EXISTS memory_embeddings USING vec0(" + f"memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding float[{dim}])" + ) + except Exception: + # vec0 unavailable (extension not loaded) -> BLOB fallback below. If we + # already dropped for a mismatch, the fallback recreates the table empty. + pass + else: + # Probe the partition_id column by inserting a dim-width vector. A + # pre-refactor table (no partition_id) accepts CREATE IF NOT EXISTS + # silently but rejects this insert; we then recreate cleanly. On a + # matching-dim table with the column, the probe round-trips and we're + # done. (The probe double-checks the dim on the vec0 path only — the + # BLOB fallback accepts any width, so dim enforcement there relies on + # the ``schema_meta`` mismatch check above.) + probe = np.zeros(dim, dtype=np.float32).tobytes() + try: + await db.execute( + "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) " + "VALUES ('__schema_probe__', '__probe__', ?)", + (probe,), + ) + await db.execute("DELETE FROM memory_embeddings WHERE memory_id = '__schema_probe__'") + except Exception: + await db.execute("DROP TABLE IF EXISTS memory_embeddings") + await db.execute(_VEC_CREATE_SQL.format(dim=dim)) + await _write_meta_dim(db, dim) + return + + # --- Fallback path: vec0 unavailable -> regular BLOB table at ``dim``. --- if true_dim_mismatch and row_count > 0: logger.warning( "Dropping %d embeddings: dim %d -> %d (%s=1 opt-in)", row_count, - declared, + existing_dim, dim, _ALLOW_EMBED_DROP_ENV, ) - else: + elif true_dim_mismatch: logger.warning( "Recreating vec0 table (dim=%d, partition_id column) — any existing embeddings will be lost", dim, ) await db.execute("DROP TABLE IF EXISTS memory_embeddings") - await db.execute(_VEC_CREATE_SQL.format(dim=dim)) + await db.execute(_FALLBACK_CREATE_SQL) + await _write_meta_dim(db, dim) async def _ensure_fallback_partition_column(db: aiosqlite.Connection) -> None: @@ -232,9 +303,7 @@ async def _ensure_fallback_partition_column(db: aiosqlite.Connection) -> None: Args: db: Open database connection. """ - cursor = await db.execute( - "SELECT sql FROM sqlite_master WHERE name = 'memory_embeddings'" - ) + cursor = await db.execute("SELECT sql FROM sqlite_master WHERE name = 'memory_embeddings'") row = await cursor.fetchone() if not row or not row[0] or "USING vec0" in row[0]: # vec0 virtual table (or absent) — schema maintained by _ensure_vec_table. @@ -243,14 +312,8 @@ async def _ensure_fallback_partition_column(db: aiosqlite.Connection) -> None: info_cursor = await db.execute("PRAGMA table_info(memory_embeddings)") columns = {r[1] for r in await info_cursor.fetchall()} if "partition_id" not in columns: - await db.execute( - "ALTER TABLE memory_embeddings " - "ADD COLUMN partition_id TEXT NOT NULL DEFAULT 'default'" - ) - await db.execute( - "CREATE INDEX IF NOT EXISTS idx_memory_embeddings_partition_id " - "ON memory_embeddings(partition_id)" - ) + await db.execute("ALTER TABLE memory_embeddings ADD COLUMN partition_id TEXT NOT NULL DEFAULT 'default'") + await db.execute("CREATE INDEX IF NOT EXISTS idx_memory_embeddings_partition_id ON memory_embeddings(partition_id)") async def initialize_schema( @@ -286,6 +349,14 @@ async def initialize_schema( CREATE INDEX IF NOT EXISTS idx_memories_created_at ON memories(created_at); CREATE INDEX IF NOT EXISTS idx_memories_last_accessed ON memories(last_accessed_at); CREATE INDEX IF NOT EXISTS idx_memories_expires_at ON memories(expires_at); + + -- Records the embedding width the memory_embeddings table was built at, + -- so the dim guard and migration mismatch check work on the vec0-unavailable + -- BLOB fallback (which has no inline width to read from sqlite_master). + CREATE TABLE IF NOT EXISTS schema_meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); """) # sqlite-vec virtual table (cannot be created via executescript) @@ -300,17 +371,13 @@ async def initialize_schema( # raise happens before any DROP. (write F7 / embedding F7) raise except Exception as e: + # Defensive net: ``_ensure_vec_table`` already builds the BLOB + # fallback itself when vec0 is unavailable, so this only fires on a + # genuinely unexpected failure. Ensure a writable table exists and + # record the configured width so the dim guard stays consistent. logger.warning("Could not create vec0 table: %s. %s", e, _VEC_FIX_HINT) - # Fallback: regular table so embedding INSERT/DELETE still works. - # Schema mirrors the vec0 contract (memory_id, partition_id, - # embedding) so partition-aware writes don't fail in - # vec0-unavailable environments. - await db.execute( - "CREATE TABLE IF NOT EXISTS memory_embeddings (" - "memory_id TEXT PRIMARY KEY, " - "partition_id TEXT NOT NULL DEFAULT 'default', " - "embedding BLOB)" - ) + await db.execute(_FALLBACK_CREATE_SQL) + await _write_meta_dim(db, embedding_dim) # Additive migration for fallback tables created before the partition_id # write contract (vec0 tables are skipped inside). Idempotent. diff --git a/src/hebb/storage/sqlite_store.py b/src/hebb/storage/sqlite_store.py index ebf2ef4..f43b218 100644 --- a/src/hebb/storage/sqlite_store.py +++ b/src/hebb/storage/sqlite_store.py @@ -92,8 +92,10 @@ async def _get_vec_dim(self) -> int | None: Returns: The integer dimension parsed from the ``memory_embeddings`` CREATE - statement, or ``None`` when the table is absent or is the - regular-BLOB fallback (which has no fixed width). + statement. When the table is the regular-BLOB fallback (vec0 + unavailable — no ``float[N]`` to read), the width is read from the + ``schema_meta`` table written at init, so the dim guard still fires + on the fallback path. ``None`` only when both are absent. """ if self._vec_dim_loaded: return self._vec_dim @@ -106,6 +108,16 @@ async def _get_vec_dim(self) -> int | None: match = _VEC_DIM_RE.search(row[0]) if match: self._vec_dim = int(match.group(1)) + if self._vec_dim is None: + # BLOB fallback (vec0 unavailable): the width isn't inline, so read + # the value ``initialize_schema`` recorded in ``schema_meta``. + try: + cursor = await self.db.execute("SELECT value FROM schema_meta WHERE key = 'embedding_dim'") + meta_row = await cursor.fetchone() + except Exception: + meta_row = None + if meta_row and meta_row[0] is not None: + self._vec_dim = int(meta_row[0]) self._vec_dim_loaded = True return self._vec_dim From 1a3157a85386d7c8ab0cf0eb2368c89fa1aebc83 Mon Sep 17 00:00:00 2001 From: Huang Date: Mon, 3 Aug 2026 14:14:22 +0800 Subject: [PATCH 5/6] fix(storage): gate fallback DROP on dim mismatch; add same-dim re-init guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses CodeRabbit review on the dim-guard fix: the vec0-unavailable fallback path unconditionally executed DROP TABLE before _FALLBACK_CREATE_SQL. On builds where CREATE VIRTUAL TABLE IF NOT EXISTS short-circuits (the BLOB table already exists), the probe returns early so the DROP is never reached; but on builds that resolve the vec0 module before the existence check, the CREATE raises and control falls through to the fallback — where an unconditional DROP would wipe a same-dim populated table. Gate the DROP on true_dim_mismatch so a same-dim (or first-time) table is never dropped on this path; _FALLBACK_CREATE_SQL is IF NOT EXISTS, so skipping the DROP is a no-op when the table exists and a clean create when it doesn't. Add test_same_dim_reinit_preserves_embeddings as an invariant guard: a populated table must survive a same-dim re-init on both the vec0 and fallback paths. Verified on vec0-enabled and vec0-disabled (macOS CI sim). Finding 2 from the review (probe-except DROP for a populated legacy vec0 table missing the partition_id column) is pre-existing — the original code took the same path, and vec0 virtual tables cannot be ALTERed to add the column. Changing it would alter the upgrade behavior for legacy installs, so it is left for a separate follow-up rather than bundled into this fix. Co-Authored-By: Claude --- src/hebb/storage/migrations.py | 37 +++++++++++++++++---------- tests/unit/test_audit_write_api.py | 41 +++++++++++++++++++++++++----- 2 files changed, 58 insertions(+), 20 deletions(-) diff --git a/src/hebb/storage/migrations.py b/src/hebb/storage/migrations.py index f59743d..55067a1 100644 --- a/src/hebb/storage/migrations.py +++ b/src/hebb/storage/migrations.py @@ -272,20 +272,29 @@ async def _ensure_vec_table(db: aiosqlite.Connection, embedding_dim: int) -> Non return # --- Fallback path: vec0 unavailable -> regular BLOB table at ``dim``. --- - if true_dim_mismatch and row_count > 0: - logger.warning( - "Dropping %d embeddings: dim %d -> %d (%s=1 opt-in)", - row_count, - existing_dim, - dim, - _ALLOW_EMBED_DROP_ENV, - ) - elif true_dim_mismatch: - logger.warning( - "Recreating vec0 table (dim=%d, partition_id column) — any existing embeddings will be lost", - dim, - ) - await db.execute("DROP TABLE IF EXISTS memory_embeddings") + if true_dim_mismatch: + # Opt-in destructive rebuild, or an empty-table dim change. Gate the + # DROP on a true mismatch so a same-dim (or first-time) table is never + # wiped here: depending on the SQLite build, the preceding + # ``CREATE VIRTUAL TABLE IF NOT EXISTS`` may raise (module resolved + # before the existence check) even when the BLOB table already exists, + # falling through to this branch. ``_FALLBACK_CREATE_SQL`` is + # ``IF NOT EXISTS``, so skipping the DROP is a no-op when the table + # exists and a clean create when it doesn't. + if row_count > 0: + logger.warning( + "Dropping %d embeddings: dim %d -> %d (%s=1 opt-in)", + row_count, + existing_dim, + dim, + _ALLOW_EMBED_DROP_ENV, + ) + else: + logger.warning( + "Recreating vec0 table (dim=%d, partition_id column) — any existing embeddings will be lost", + dim, + ) + await db.execute("DROP TABLE IF EXISTS memory_embeddings") await db.execute(_FALLBACK_CREATE_SQL) await _write_meta_dim(db, dim) diff --git a/tests/unit/test_audit_write_api.py b/tests/unit/test_audit_write_api.py index 826f357..6bb36a9 100644 --- a/tests/unit/test_audit_write_api.py +++ b/tests/unit/test_audit_write_api.py @@ -193,9 +193,7 @@ async def test_update_rejects_empty_content( @pytest.mark.asyncio -async def test_update_omitted_content_is_no_change( - memory_store: SQLiteMemoryStore, noop_embedder: Any -) -> None: +async def test_update_omitted_content_is_no_change(memory_store: SQLiteMemoryStore, noop_embedder: Any) -> None: mem = await memory_store.create(MemoryCreate(content="original", partition_id="mem_hippocampus")) updated = await memories_router.update_memory( memory_id=mem.id, @@ -218,9 +216,7 @@ def _settings_with_strengthening(enabled: bool = True) -> Settings: def _fake_searcher(memory_ids: list[str]) -> AsyncMock: searcher = AsyncMock() - results = [ - type("R", (), {"memory": type("M", (), {"id": mid})()})() for mid in memory_ids - ] + results = [type("R", (), {"memory": type("M", (), {"id": mid})()})() for mid in memory_ids] searcher.search.return_value = SearchResponse.model_construct(results=results, related=[]) return searcher @@ -351,3 +347,36 @@ async def test_first_time_creation_unaffected(tmp_path: Any) -> None: ) finally: await conn.close() + + +@pytest.mark.asyncio +async def test_same_dim_reinit_preserves_embeddings(tmp_path: Any) -> None: + """Re-initializing at the SAME dim must not drop existing embeddings. + + Invariant guard for the vec0-unavailable fallback path: the BLOB fallback + table must survive a same-dim re-init (e.g. a service restart) regardless of + SQLite build behavior. When ``CREATE VIRTUAL TABLE IF NOT EXISTS`` short- + circuits (table already exists), the probe returns early; when a build + resolves the vec0 module before the existence check and raises, control + falls through to the fallback — which must still leave a same-dim table + intact. Either way no embedding may be lost. + """ + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path) + try: + await initialize_schema(conn, embedding_dim=384) + await _populate_one_vector(conn, 384) + assert await _vec_row_count_proxy(conn) == 1 + + # Simulate a service restart: re-init at the same configured dim. + await initialize_schema(conn, embedding_dim=384) + assert await _vec_row_count_proxy(conn) == 1, "same-dim re-init wiped existing embeddings on the fallback path" + finally: + await conn.close() + + +async def _vec_row_count_proxy(conn: Any) -> int: + """Count rows in ``memory_embeddings`` directly on the connection.""" + cursor = await conn.execute("SELECT count(*) FROM memory_embeddings") + row = await cursor.fetchone() + return int(row[0]) if row and row[0] is not None else 0 From ae9fbcab72fb6ed439de0223fd45205438127d3c Mon Sep 17 00:00:00 2001 From: Huang Date: Wed, 2 Sep 2026 15:14:24 +0800 Subject: [PATCH 6/6] fix(storage): fail closed on unsafe vector migrations Prevent schema inspection and rebuild failures from silently relabeling or dropping populated embedding tables. Infer legacy fallback dimensions, require explicit destructive opt-in, and preserve the prior table through transactional rebuild failures. --- src/hebb/storage/migrations.py | 393 +++++++++++++++++++++-------- tests/unit/test_audit_write_api.py | 275 +++++++++++++++++++- 2 files changed, 552 insertions(+), 116 deletions(-) diff --git a/src/hebb/storage/migrations.py b/src/hebb/storage/migrations.py index 55067a1..d213734 100644 --- a/src/hebb/storage/migrations.py +++ b/src/hebb/storage/migrations.py @@ -21,14 +21,23 @@ # inside the CREATE VIRTUAL TABLE SQL recorded in ``sqlite_master``. _VEC_DIM_RE = re.compile(r"embedding\s+float\[(\d+)\]", re.IGNORECASE) -# Opt-in escape hatch: when set to "1", a true embedding-dimension mismatch on a -# populated table is resolved by dropping the table (losing every vector) instead -# of raising. Off by default so a config change never silently wipes embeddings. +# Opt-in escape hatch: when set to "1", an incompatible populated embeddings +# table may be dropped and rebuilt empty. Off by default so migrations never +# silently wipe embeddings. _ALLOW_EMBED_DROP_ENV = "HEBB_ALLOW_EMBED_DROP" -class EmbeddingDimensionMismatchError(RuntimeError): - """Raised when an existing populated vec0 table's width != the configured dim. +class EmbeddingSchemaMigrationError(RuntimeError): + """Raised when migrating ``memory_embeddings`` would discard stored vectors. + + The existing table is left unchanged. Operators can rebuild embeddings with + ``hebb memory reembed`` or explicitly allow an empty rebuild by setting + ``HEBB_ALLOW_EMBED_DROP=1``. + """ + + +class EmbeddingDimensionMismatchError(EmbeddingSchemaMigrationError): + """Raised when a populated embeddings table's width != the configured dim. Dropping the table would discard every stored embedding, so we refuse and point the operator at ``hebb memory reembed`` (or the explicit opt-in env @@ -44,6 +53,17 @@ async def get_connection(db_path: str, *, load_vec: bool = True) -> aiosqlite.Co This means ``aiosqlite`` (which imports ``sqlite3``) transparently uses the patched module, and ``conn.enable_load_extension`` is always available when ``pysqlite3-binary`` is installed. + + Args: + db_path: Filesystem path to the SQLite database. + load_vec: Whether to attempt loading the sqlite-vec extension. When + false or unavailable, schema initialization uses a BLOB fallback. + + Returns: + Configured asynchronous SQLite connection. + + Raises: + Exception: Opening or configuring the SQLite database failed. """ db = await aiosqlite.connect(db_path) db.row_factory = aiosqlite.Row @@ -150,15 +170,62 @@ async def _declared_vec_dim(db: aiosqlite.Connection) -> int | None: async def _vec_row_count(db: aiosqlite.Connection) -> int: - """Number of stored embeddings, or 0 if the table is missing/unreadable.""" + """Return the number of stored embeddings. + + Args: + db: Open database connection with an existing embeddings table. + + Returns: + Number of rows in ``memory_embeddings``. + + Raises: + EmbeddingSchemaMigrationError: The table exists but cannot be read. An + unreadable table must never be treated as empty during migration. + """ try: cursor = await db.execute("SELECT count(*) FROM memory_embeddings") row = await cursor.fetchone() - except Exception: - return 0 + except Exception as e: + raise EmbeddingSchemaMigrationError( + "Could not determine whether memory_embeddings contains data; " + "refusing to rebuild it because stored embeddings may be lost. " + f"Restore sqlite-vec access and retry, or set {_ALLOW_EMBED_DROP_ENV}=1 " + "to explicitly allow a destructive rebuild." + ) from e return int(row[0]) if row and row[0] is not None else 0 +async def _embedding_table_schema(db: aiosqlite.Connection) -> tuple[str | None, set[str]]: + """Read the embeddings table declaration and columns without mutating it. + + Args: + db: Open database connection. + + Returns: + A pair of the ``sqlite_master`` CREATE statement and declared column + names. The CREATE statement is ``None`` when the table does not exist. + + Raises: + EmbeddingSchemaMigrationError: The existing table cannot be inspected. + """ + try: + cursor = await db.execute("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'memory_embeddings'") + row = await cursor.fetchone() + if not row or not row[0]: + return None, set() + + create_sql = str(row[0]) + info_cursor = await db.execute("PRAGMA table_info(memory_embeddings)") + columns = {str(info_row[1]) for info_row in await info_cursor.fetchall()} + except Exception as e: + raise EmbeddingSchemaMigrationError( + "Could not inspect the existing memory_embeddings schema; refusing " + "to rebuild it because stored embeddings may be lost. Restore " + "sqlite-vec access and retry." + ) from e + return create_sql, columns + + async def _read_meta_dim(db: aiosqlite.Connection) -> int | None: """Embedding width recorded in ``schema_meta``, or ``None`` if unset. @@ -166,13 +233,94 @@ async def _read_meta_dim(db: aiosqlite.Connection) -> int | None: informational for vec0 (which declares its width inline). Returns ``None`` when the meta table is absent or the key is unset — e.g. first init, or a DB created before this metadata existed. + + Args: + db: Open database connection. + + Returns: + Recorded positive embedding dimension, or ``None`` only when the table + or metadata key does not exist. + + Raises: + EmbeddingSchemaMigrationError: Metadata exists but cannot be read or is + not a positive integer. """ try: + table_cursor = await db.execute("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'schema_meta'") + if await table_cursor.fetchone() is None: + return None cursor = await db.execute("SELECT value FROM schema_meta WHERE key = ?", (_META_DIM_KEY,)) row = await cursor.fetchone() - except Exception: + if row is None: + return None + if row[0] is None: + raise ValueError("embedding dimension metadata is NULL") + dim = int(row[0]) + if dim <= 0: + raise ValueError("embedding dimension must be positive") + except Exception as e: + raise EmbeddingSchemaMigrationError( + "Could not read valid embedding-dimension metadata from schema_meta; " + "refusing to continue because relabeling existing vectors may corrupt " + "the index. Repair the metadata table and retry." + ) from e + return dim + + +async def _infer_fallback_dim(db: aiosqlite.Connection) -> int | None: + """Infer float32 width from all rows in a legacy BLOB fallback table. + + Args: + db: Open database connection with a regular ``memory_embeddings`` table. + + Returns: + The uniform float32 dimension, or ``None`` when the table is empty. + + Raises: + EmbeddingSchemaMigrationError: Stored values are NULL, not BLOBs, have + inconsistent lengths, have an invalid float32 byte width, or cannot + be inspected. + """ + try: + cursor = await db.execute( + "SELECT count(*), " + "sum(CASE WHEN embedding IS NULL THEN 1 ELSE 0 END), " + "sum(CASE WHEN embedding IS NOT NULL AND typeof(embedding) != 'blob' THEN 1 ELSE 0 END), " + "count(DISTINCT length(embedding)), min(length(embedding)), max(length(embedding)) " + "FROM memory_embeddings" + ) + row = await cursor.fetchone() + except Exception as e: + raise EmbeddingSchemaMigrationError( + "Could not infer the dimension of legacy fallback embeddings; " + "refusing to write new dimension metadata or rebuild the table." + ) from e + + if row is None: + raise EmbeddingSchemaMigrationError( + "Could not infer the dimension of legacy fallback embeddings because " + "the aggregate query returned no result." + ) + row_count = int(row[0]) if row[0] is not None else 0 + if row_count == 0: return None - return int(row[0]) if row and row[0] is not None else None + + null_count = int(row[1] or 0) + non_blob_count = int(row[2] or 0) + distinct_lengths = int(row[3] or 0) + min_bytes = int(row[4]) if row[4] is not None else 0 + max_bytes = int(row[5]) if row[5] is not None else 0 + if null_count or non_blob_count or distinct_lengths != 1 or min_bytes != max_bytes: + raise EmbeddingSchemaMigrationError( + "Legacy fallback embeddings do not have one consistent non-NULL BLOB " + "width; refusing to infer or overwrite embedding-dimension metadata." + ) + if min_bytes <= 0 or min_bytes % 4 != 0: + raise EmbeddingSchemaMigrationError( + f"Legacy fallback embeddings have invalid byte width {min_bytes}; " + "float32 vector bytes must be non-empty and divisible by 4." + ) + return min_bytes // 4 async def _write_meta_dim(db: aiosqlite.Connection, dim: int) -> None: @@ -196,8 +344,8 @@ async def _ensure_vec_table(db: aiosqlite.Connection, embedding_dim: int) -> Non a fatal, recoverable condition: dropping it would lose every vector, so we raise :class:`EmbeddingDimensionMismatchError` (run ``hebb memory reembed`` or set ``HEBB_ALLOW_EMBED_DROP=1`` to opt into the destructive rebuild). - First-time creation and the additive ``partition_id`` migration (empty or - correctly-dimensioned tables) are unaffected. + First-time creation and empty-table schema upgrades remain automatic. A + regular BLOB fallback missing ``partition_id`` is migrated additively. The existing width is read from the vec0 declaration (``float[N]``) when the extension is available, or from ``schema_meta`` for the BLOB fallback used @@ -210,93 +358,121 @@ async def _ensure_vec_table(db: aiosqlite.Connection, embedding_dim: int) -> Non embedding_dim: Target embedding width for the configured model. Raises: - EmbeddingDimensionMismatchError: Populated table with a different width - and no ``HEBB_ALLOW_EMBED_DROP=1`` opt-in. + EmbeddingSchemaMigrationError: A populated table requires destructive + migration and no ``HEBB_ALLOW_EMBED_DROP=1`` opt-in was provided. """ - import numpy as np - dim = int(embedding_dim) + allow_drop = os.getenv(_ALLOW_EMBED_DROP_ENV) == "1" + create_sql, columns = await _embedding_table_schema(db) + table_exists = create_sql is not None + is_vec = bool(create_sql and re.search(r"\bUSING\s+vec0\b", create_sql, re.IGNORECASE)) + declared = await _declared_vec_dim(db) if is_vec else None + validation_error: EmbeddingSchemaMigrationError | None = None + try: + existing_dim = declared if declared is not None else await _read_meta_dim(db) + except EmbeddingSchemaMigrationError as e: + if not allow_drop: + raise + existing_dim = None + validation_error = e + + # A regular fallback with the core columns is already usable. A missing + # partition_id is handled additively by _ensure_fallback_partition_column, + # so repeated initialization must never replace a populated fallback merely + # because vec0 happens to be unavailable on this connection. + required_columns = {"memory_id", "embedding"} + fallback_compatible = not is_vec and required_columns <= columns + vec_compatible = is_vec and "partition_id" in columns and declared is not None + if table_exists and fallback_compatible and existing_dim is None and validation_error is None: + try: + existing_dim = await _infer_fallback_dim(db) + except EmbeddingSchemaMigrationError as e: + if not allow_drop: + raise + validation_error = e + dim_mismatch = existing_dim is not None and existing_dim != dim + + incompatibility = ( + f"embedding dimension could not be validated: {validation_error}" if validation_error is not None else None + ) + if table_exists: + if incompatibility is not None: + pass + elif not required_columns <= columns: + missing = ", ".join(sorted(required_columns - columns)) + incompatibility = f"missing required column(s): {missing}" + elif is_vec and "partition_id" not in columns: + incompatibility = "missing partition_id column" + elif is_vec and declared is None: + incompatibility = "embedding dimension is not declared" + elif dim_mismatch: + incompatibility = f"embedding dimension is {existing_dim}, configured dimension is {dim}" + + if table_exists and incompatibility is None and (fallback_compatible or vec_compatible): + await _write_meta_dim(db, dim) + return - # The width the existing table was built at. vec0 declares it inline - # (``float[N]`` in sqlite_master); the BLOB fallback (vec0 unavailable) - # records it in ``schema_meta``. ``None`` means first-time creation. - declared = await _declared_vec_dim(db) - existing_dim = declared if declared is not None else await _read_meta_dim(db) - row_count = await _vec_row_count(db) - true_dim_mismatch = existing_dim is not None and existing_dim != dim - - # A true mismatch on a POPULATED table is operator-actionable: refuse before - # touching anything so we never silently drop vectors. First-time creation, - # an empty-table dim change, and the partition_id migration fall through. - if true_dim_mismatch and row_count > 0 and os.getenv(_ALLOW_EMBED_DROP_ENV) != "1": - raise EmbeddingDimensionMismatchError( - f"memory_embeddings has {row_count} vectors of width {existing_dim}, but the " - f"configured embedding dimension is {dim}. Dropping the table would lose " - f"every embedding. Run `hebb memory reembed` to rebuild vectors at the new " - f"dimension, or set {_ALLOW_EMBED_DROP_ENV}=1 to drop and rebuild empty." + row_count: int | None = 0 + if table_exists: + try: + row_count = await _vec_row_count(db) + except EmbeddingSchemaMigrationError: + if not allow_drop: + raise + row_count = None + + if table_exists and row_count is not None and row_count > 0 and not allow_drop: + if dim_mismatch: + raise EmbeddingDimensionMismatchError( + f"memory_embeddings has {row_count} vectors of width {existing_dim}, but the " + f"configured embedding dimension is {dim}. Dropping the table would lose " + f"every embedding. Run `hebb memory reembed` to rebuild vectors at the new " + f"dimension, or set {_ALLOW_EMBED_DROP_ENV}=1 to drop and rebuild empty." + ) + raise EmbeddingSchemaMigrationError( + f"memory_embeddings has {row_count} stored vector(s), but its schema is " + f"incompatible ({incompatibility}). Rebuilding it would lose every embedding. " + "Run `hebb memory reembed` to rebuild the index, or set " + f"{_ALLOW_EMBED_DROP_ENV}=1 to explicitly drop and rebuild empty." ) - # --- vec0 path: try the virtual table. --- try: - if true_dim_mismatch: - # Empty-table dim change or opt-in destructive rebuild at ``dim``. - await db.execute("DROP TABLE IF EXISTS memory_embeddings") - await db.execute( - f"CREATE VIRTUAL TABLE IF NOT EXISTS memory_embeddings USING vec0(" - f"memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding float[{dim}])" - ) - except Exception: - # vec0 unavailable (extension not loaded) -> BLOB fallback below. If we - # already dropped for a mismatch, the fallback recreates the table empty. - pass - else: - # Probe the partition_id column by inserting a dim-width vector. A - # pre-refactor table (no partition_id) accepts CREATE IF NOT EXISTS - # silently but rejects this insert; we then recreate cleanly. On a - # matching-dim table with the column, the probe round-trips and we're - # done. (The probe double-checks the dim on the vec0 path only — the - # BLOB fallback accepts any width, so dim enforcement there relies on - # the ``schema_meta`` mismatch check above.) - probe = np.zeros(dim, dtype=np.float32).tobytes() + await db.execute("SAVEPOINT rebuild_memory_embeddings") + except Exception as e: + raise EmbeddingSchemaMigrationError( + "Could not start the memory_embeddings rebuild; the existing table was left unchanged." + ) from e + try: + if table_exists: + if row_count is None: + logger.warning( + "Dropping unreadable memory_embeddings table (%s=1 opt-in)", + _ALLOW_EMBED_DROP_ENV, + ) + elif row_count > 0: + logger.warning( + "Dropping %d embeddings with incompatible schema (%s; %s=1 opt-in)", + row_count, + incompatibility, + _ALLOW_EMBED_DROP_ENV, + ) + await db.execute("DROP TABLE memory_embeddings") + try: - await db.execute( - "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) " - "VALUES ('__schema_probe__', '__probe__', ?)", - (probe,), - ) - await db.execute("DELETE FROM memory_embeddings WHERE memory_id = '__schema_probe__'") - except Exception: - await db.execute("DROP TABLE IF EXISTS memory_embeddings") await db.execute(_VEC_CREATE_SQL.format(dim=dim)) + except Exception: + await db.execute(_FALLBACK_CREATE_SQL) await _write_meta_dim(db, dim) - return - - # --- Fallback path: vec0 unavailable -> regular BLOB table at ``dim``. --- - if true_dim_mismatch: - # Opt-in destructive rebuild, or an empty-table dim change. Gate the - # DROP on a true mismatch so a same-dim (or first-time) table is never - # wiped here: depending on the SQLite build, the preceding - # ``CREATE VIRTUAL TABLE IF NOT EXISTS`` may raise (module resolved - # before the existence check) even when the BLOB table already exists, - # falling through to this branch. ``_FALLBACK_CREATE_SQL`` is - # ``IF NOT EXISTS``, so skipping the DROP is a no-op when the table - # exists and a clean create when it doesn't. - if row_count > 0: - logger.warning( - "Dropping %d embeddings: dim %d -> %d (%s=1 opt-in)", - row_count, - existing_dim, - dim, - _ALLOW_EMBED_DROP_ENV, - ) - else: - logger.warning( - "Recreating vec0 table (dim=%d, partition_id column) — any existing embeddings will be lost", - dim, - ) - await db.execute("DROP TABLE IF EXISTS memory_embeddings") - await db.execute(_FALLBACK_CREATE_SQL) - await _write_meta_dim(db, dim) + except BaseException as e: + await db.execute("ROLLBACK TO SAVEPOINT rebuild_memory_embeddings") + await db.execute("RELEASE SAVEPOINT rebuild_memory_embeddings") + if not isinstance(e, Exception): + raise + raise EmbeddingSchemaMigrationError( + "Could not rebuild memory_embeddings; the migration was rolled back and the previous table was preserved." + ) from e + else: + await db.execute("RELEASE SAVEPOINT rebuild_memory_embeddings") async def _ensure_fallback_partition_column(db: aiosqlite.Connection) -> None: @@ -328,7 +504,21 @@ async def _ensure_fallback_partition_column(db: aiosqlite.Connection) -> None: async def initialize_schema( db: aiosqlite.Connection, embedding_dim: int = 384, *, create_vec_table: bool = True ) -> None: - """Create all tables idempotently.""" + """Create all database tables idempotently. + + Args: + db: Open database connection. + embedding_dim: Width of vectors produced by the configured embedder. + create_vec_table: Whether to initialize the vector index or fallback. + + Returns: + None. + + Raises: + EmbeddingSchemaMigrationError: A populated embeddings table requires a + destructive schema migration without explicit operator opt-in. + Exception: A required SQLite schema operation failed. + """ await db.executescript(""" CREATE TABLE IF NOT EXISTS partitions ( id TEXT PRIMARY KEY, @@ -370,23 +560,10 @@ async def initialize_schema( # sqlite-vec virtual table (cannot be created via executescript) if create_vec_table: - try: - # Check if vec0 table exists with a different dimension - await _ensure_vec_table(db, embedding_dim) - except EmbeddingDimensionMismatchError: - # A populated table with a mismatched width is operator-actionable, - # not a missing-extension fallback — surface it instead of silently - # degrading to a dimensionless BLOB table. No data was dropped: the - # raise happens before any DROP. (write F7 / embedding F7) - raise - except Exception as e: - # Defensive net: ``_ensure_vec_table`` already builds the BLOB - # fallback itself when vec0 is unavailable, so this only fires on a - # genuinely unexpected failure. Ensure a writable table exists and - # record the configured width so the dim guard stays consistent. - logger.warning("Could not create vec0 table: %s. %s", e, _VEC_FIX_HINT) - await db.execute(_FALLBACK_CREATE_SQL) - await _write_meta_dim(db, embedding_dim) + # _ensure_vec_table handles vec0 unavailability itself. Any other error + # must fail closed: treating it as an unavailable extension could retain + # old BLOBs while relabeling them with the newly configured dimension. + await _ensure_vec_table(db, embedding_dim) # Additive migration for fallback tables created before the partition_id # write contract (vec0 tables are skipped inside). Idempotent. diff --git a/tests/unit/test_audit_write_api.py b/tests/unit/test_audit_write_api.py index 6bb36a9..817b6de 100644 --- a/tests/unit/test_audit_write_api.py +++ b/tests/unit/test_audit_write_api.py @@ -28,6 +28,7 @@ from hebb.server.routers import search as search_router from hebb.storage.migrations import ( EmbeddingDimensionMismatchError, + EmbeddingSchemaMigrationError, _ensure_vec_table, get_connection, initialize_schema, @@ -354,27 +355,285 @@ async def test_same_dim_reinit_preserves_embeddings(tmp_path: Any) -> None: """Re-initializing at the SAME dim must not drop existing embeddings. Invariant guard for the vec0-unavailable fallback path: the BLOB fallback - table must survive a same-dim re-init (e.g. a service restart) regardless of - SQLite build behavior. When ``CREATE VIRTUAL TABLE IF NOT EXISTS`` short- - circuits (table already exists), the probe returns early; when a build - resolves the vec0 module before the existence check and raises, control - falls through to the fallback — which must still leave a same-dim table - intact. Either way no embedding may be lost. + table must survive a same-dim re-init (e.g. a service restart). The second + connection also has extension loading disabled, so this exercises the real + fallback path rather than relying on host-specific sqlite-vec behavior. """ db_path = str(tmp_path / "hebb.db") - conn = await get_connection(db_path) + conn = await get_connection(db_path, load_vec=False) try: await initialize_schema(conn, embedding_dim=384) + schema_cursor = await conn.execute("SELECT sql FROM sqlite_master WHERE name = 'memory_embeddings'") + schema_row = await schema_cursor.fetchone() + assert schema_row is not None + assert "USING vec0" not in schema_row[0] + await _populate_one_vector(conn, 384) assert await _vec_row_count_proxy(conn) == 1 + finally: + await conn.close() - # Simulate a service restart: re-init at the same configured dim. + # Simulate a service restart on another connection where vec0 is still + # unavailable: initialization must reuse the populated BLOB fallback. + conn = await get_connection(db_path, load_vec=False) + try: await initialize_schema(conn, embedding_dim=384) assert await _vec_row_count_proxy(conn) == 1, "same-dim re-init wiped existing embeddings on the fallback path" finally: await conn.close() +@pytest.mark.asyncio +async def test_populated_legacy_fallback_adds_partition_without_data_loss(tmp_path: Any) -> None: + """Fallback schema migration is additive even when vec0 is unavailable.""" + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path, load_vec=False) + try: + await conn.execute("CREATE TABLE memory_embeddings (memory_id TEXT PRIMARY KEY, embedding BLOB)") + embedding = np.zeros(384, dtype=np.float32).tobytes() + await conn.execute( + "INSERT INTO memory_embeddings(memory_id, embedding) VALUES ('legacy', ?)", + (embedding,), + ) + await conn.commit() + + await initialize_schema(conn, embedding_dim=384) + + assert await _vec_row_count_proxy(conn) == 1 + info_cursor = await conn.execute("PRAGMA table_info(memory_embeddings)") + columns = {row[1] for row in await info_cursor.fetchall()} + assert "partition_id" in columns + row_cursor = await conn.execute("SELECT memory_id, partition_id, embedding FROM memory_embeddings") + row = await row_cursor.fetchone() + assert tuple(row) == ("legacy", "default", embedding) + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_populated_legacy_fallback_infers_dimension_before_writing_meta(tmp_path: Any, monkeypatch: Any) -> None: + """Missing fallback metadata must not relabel existing vectors at a new width.""" + monkeypatch.delenv("HEBB_ALLOW_EMBED_DROP", raising=False) + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path, load_vec=False) + try: + await conn.execute( + "CREATE TABLE memory_embeddings (memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding BLOB)" + ) + embedding = np.zeros(384, dtype=np.float32).tobytes() + await conn.execute( + "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) VALUES ('legacy', 'default', ?)", + (embedding,), + ) + await conn.commit() + + with pytest.raises(EmbeddingDimensionMismatchError, match="384"): + await initialize_schema(conn, embedding_dim=512) + + assert await _vec_row_count_proxy(conn) == 1 + meta_cursor = await conn.execute("SELECT value FROM schema_meta WHERE key = 'embedding_dim'") + assert await meta_cursor.fetchone() is None + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_fallback_metadata_read_error_fails_closed(tmp_path: Any) -> None: + """A broken metadata table must not be treated as an absent metadata key.""" + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path, load_vec=False) + try: + await conn.execute( + "CREATE TABLE memory_embeddings (memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding BLOB)" + ) + await conn.execute("CREATE TABLE schema_meta (key TEXT PRIMARY KEY)") + embedding = np.zeros(384, dtype=np.float32).tobytes() + await conn.execute( + "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) VALUES ('legacy', 'default', ?)", + (embedding,), + ) + await conn.commit() + + with pytest.raises(EmbeddingSchemaMigrationError, match="metadata"): + await _ensure_vec_table(conn, embedding_dim=384) + + assert await _vec_row_count_proxy(conn) == 1 + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_embedding_schema_inspection_error_preserves_data_and_meta(tmp_path: Any, monkeypatch: Any) -> None: + """A schema-query failure must not fall through to fallback initialization.""" + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path, load_vec=False) + try: + await initialize_schema(conn, embedding_dim=384) + await _populate_one_vector(conn, 384) + original_execute = conn.execute + failed = False + + async def fail_embedding_schema_query(sql: str, *args: Any) -> Any: + nonlocal failed + if not failed and "sqlite_master" in sql and "memory_embeddings" in sql: + failed = True + raise RuntimeError("injected schema inspection failure") + return await original_execute(sql, *args) + + monkeypatch.setattr(conn, "execute", fail_embedding_schema_query) + + with pytest.raises(EmbeddingSchemaMigrationError, match="inspect"): + await initialize_schema(conn, embedding_dim=512) + + assert await _vec_row_count_proxy(conn) == 1 + meta_cursor = await conn.execute("SELECT value FROM schema_meta WHERE key = 'embedding_dim'") + meta_row = await meta_cursor.fetchone() + assert meta_row is not None + assert meta_row[0] == "384" + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_opt_in_rebuild_rolls_back_when_metadata_write_fails(tmp_path: Any, monkeypatch: Any) -> None: + """A failed destructive rebuild restores the original populated table.""" + monkeypatch.setenv("HEBB_ALLOW_EMBED_DROP", "1") + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path, load_vec=False) + try: + await conn.execute( + "CREATE TABLE memory_embeddings (memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding BLOB)" + ) + await conn.execute("CREATE TABLE schema_meta (key TEXT PRIMARY KEY)") + embedding = np.zeros(384, dtype=np.float32).tobytes() + await conn.execute( + "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) VALUES ('legacy', 'default', ?)", + (embedding,), + ) + await conn.commit() + + with pytest.raises(EmbeddingSchemaMigrationError, match="rolled back"): + await _ensure_vec_table(conn, embedding_dim=512) + + assert await _vec_row_count_proxy(conn) == 1 + row_cursor = await conn.execute("SELECT memory_id, partition_id, embedding FROM memory_embeddings") + row = await row_cursor.fetchone() + assert tuple(row) == ("legacy", "default", embedding) + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_opt_in_savepoint_failure_preserves_data_and_meta(tmp_path: Any, monkeypatch: Any) -> None: + """Failure to start a destructive rebuild must leave the old index intact.""" + monkeypatch.setenv("HEBB_ALLOW_EMBED_DROP", "1") + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path, load_vec=False) + try: + await initialize_schema(conn, embedding_dim=384) + await _populate_one_vector(conn, 384) + original_execute = conn.execute + failed = False + + async def fail_rebuild_savepoint(sql: str, *args: Any) -> Any: + nonlocal failed + if not failed and sql == "SAVEPOINT rebuild_memory_embeddings": + failed = True + raise RuntimeError("injected savepoint failure") + return await original_execute(sql, *args) + + monkeypatch.setattr(conn, "execute", fail_rebuild_savepoint) + + with pytest.raises(EmbeddingSchemaMigrationError, match="start"): + await initialize_schema(conn, embedding_dim=512) + + assert await _vec_row_count_proxy(conn) == 1 + meta_cursor = await conn.execute("SELECT value FROM schema_meta WHERE key = 'embedding_dim'") + meta_row = await meta_cursor.fetchone() + assert meta_row is not None + assert meta_row[0] == "384" + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_inconsistent_legacy_fallback_rebuilds_with_opt_in(tmp_path: Any, monkeypatch: Any) -> None: + """Explicit opt-in permits rebuilding fallback vectors of unknown width.""" + monkeypatch.setenv("HEBB_ALLOW_EMBED_DROP", "1") + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path, load_vec=False) + try: + await conn.execute( + "CREATE TABLE memory_embeddings (memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding BLOB)" + ) + await conn.executemany( + "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) VALUES (?, 'default', ?)", + [ + ("legacy-384", np.zeros(384, dtype=np.float32).tobytes()), + ("legacy-512", np.zeros(512, dtype=np.float32).tobytes()), + ], + ) + await conn.commit() + + await initialize_schema(conn, embedding_dim=768) + + assert await _vec_row_count_proxy(conn) == 0 + meta_cursor = await conn.execute("SELECT value FROM schema_meta WHERE key = 'embedding_dim'") + meta_row = await meta_cursor.fetchone() + assert meta_row is not None + assert meta_row[0] == "768" + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_populated_legacy_vec_table_without_partition_is_preserved(tmp_path: Any, monkeypatch: Any) -> None: + """A failed legacy-schema probe must never silently discard vectors.""" + monkeypatch.delenv("HEBB_ALLOW_EMBED_DROP", raising=False) + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path) + try: + await conn.execute( + "CREATE VIRTUAL TABLE memory_embeddings USING vec0(memory_id TEXT PRIMARY KEY, embedding float[384])" + ) + embedding = np.zeros(384, dtype=np.float32).tobytes() + await conn.execute( + "INSERT INTO memory_embeddings(memory_id, embedding) VALUES ('legacy', ?)", + (embedding,), + ) + await conn.commit() + + with pytest.raises(EmbeddingSchemaMigrationError, match="partition_id"): + await initialize_schema(conn, embedding_dim=384) + + assert await _vec_row_count_proxy(conn) == 1 + schema_cursor = await conn.execute("SELECT sql FROM sqlite_master WHERE name = 'memory_embeddings'") + schema_row = await schema_cursor.fetchone() + assert schema_row is not None + assert "partition_id" not in schema_row[0] + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_empty_legacy_vec_table_without_partition_is_rebuilt(tmp_path: Any) -> None: + """An incompatible vec0 schema may be rebuilt when it contains no data.""" + db_path = str(tmp_path / "hebb.db") + conn = await get_connection(db_path) + try: + await conn.execute( + "CREATE VIRTUAL TABLE memory_embeddings USING vec0(memory_id TEXT PRIMARY KEY, embedding float[384])" + ) + + await initialize_schema(conn, embedding_dim=384) + + info_cursor = await conn.execute("PRAGMA table_info(memory_embeddings)") + columns = {row[1] for row in await info_cursor.fetchall()} + assert "partition_id" in columns + finally: + await conn.close() + + async def _vec_row_count_proxy(conn: Any) -> int: """Count rows in ``memory_embeddings`` directly on the connection.""" cursor = await conn.execute("SELECT count(*) FROM memory_embeddings")