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/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/agents/consolidation_agent.py b/src/hebb/agents/consolidation_agent.py index e9e1c67..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 = "" @@ -160,29 +167,63 @@ 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) + 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 self.memory_store._begin() + try: + new_memory = await self.memory_store._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, + ) + 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; 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, + partition_id=target_partition, + importance_score=importance, + tags=tags, + metadata=memory.metadata, + source="consolidation", + ), + embedding=embedding, + ) + 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 result.new_memory_id = new_memory.id @@ -363,21 +404,37 @@ 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) + if isinstance(self.memory_store, SQLiteMemoryStore): + async with kg_lock: + new_memory = await self.memory_store._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( @@ -396,17 +453,23 @@ async def _consolidate_session_chunk( # them in the working partition for the next consolidation pass. if results: 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 isinstance(self.memory_store, SQLiteMemoryStore): + await self.memory_store._begin() + try: + for m in memories: + 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) + 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 +490,23 @@ 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) + if isinstance(self.memory_store, SQLiteMemoryStore): + await self.memory_store._begin() + try: + for m in memories: + 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) + 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 +548,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 +564,20 @@ 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. + if isinstance(self.memory_store, SQLiteMemoryStore): + async with self.kg.lock: + 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: + return + await self.memory_store.update_embedding(conflict_id, embedding) async def consolidate_batch( self, @@ -638,9 +724,14 @@ async def _consolidate_one(self, memory: Memory, kg_lock: asyncio.Lock) -> Conso result.conflicts_resolved += 1 elif resolution == "discard": async with kg_lock: + # 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() - await self.memory_store.delete(memory.id) result.success = True result.target_partition = "discarded" return result @@ -660,28 +751,53 @@ 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() + 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 self.memory_store._begin() + try: + new_memory = await self.memory_store._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, + ) + 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. + 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..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") @@ -336,14 +342,34 @@ 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: 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 isinstance(self.memory_store, SQLiteMemoryStore): + await self.memory_store._begin() + try: + for mid in to_delete: + 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) + 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..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 @@ -79,8 +85,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..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)) @@ -173,6 +179,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 +194,32 @@ 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: + async with kg.lock: + if isinstance(memory_store, SQLiteMemoryStore): + await memory_store._begin() + try: + for mid in to_delete: + 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) + 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. @@ -326,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 9527b8c..cac3856 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,31 @@ 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: + """Close the SQLite connection pool backing the storage context.""" 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: @@ -87,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/migrations.py b/src/hebb/storage/migrations.py index c1428ad..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 @@ -78,6 +98,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.""" @@ -129,17 +170,170 @@ 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. + + 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. + + 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() + 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 + + 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: + """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 @@ -150,79 +344,181 @@ 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 + 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. 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) - # Fast path for fresh DB + 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: - 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}])" + 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 + + 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." ) - 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. - declared = await _declared_vec_dim(db) - row_count = await _vec_row_count(db) - true_dim_mismatch = declared is not None and declared != dim - - 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"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." - ) + 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") - if true_dim_mismatch and row_count > 0: - logger.warning( - "Dropping %d embeddings: dim %d -> %d (%s=1 opt-in)", - row_count, - declared, - dim, - _ALLOW_EMBED_DROP_ENV, - ) + try: + await db.execute(_VEC_CREATE_SQL.format(dim=dim)) + except Exception: + 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: - 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("RELEASE SAVEPOINT rebuild_memory_embeddings") + + +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: - """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, @@ -252,25 +548,26 @@ 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) 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: - logger.warning("Could not create vec0 table: %s. %s", e, _VEC_FIX_HINT) - # Fallback: regular table so embedding INSERT/DELETE still works - await db.execute( - "CREATE TABLE IF NOT EXISTS memory_embeddings (memory_id TEXT PRIMARY KEY, embedding BLOB)" - ) + # _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. + 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 diff --git a/src/hebb/storage/partition_store.py b/src/hebb/storage/partition_store.py index 4c294ce..3081c7b 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,20 +35,50 @@ 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: + """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() - 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: + """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,), @@ -69,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", ) @@ -87,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 @@ -110,24 +156,42 @@ 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: + """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: 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..df997d7 100644 --- a/src/hebb/storage/purge.py +++ b/src/hebb/storage/purge.py @@ -26,6 +26,12 @@ async def purge_memory( ) -> 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,42 @@ 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 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"): + # ``_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) + 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..f43b218 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. @@ -87,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 @@ -101,10 +108,44 @@ 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 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) + + 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,45 +155,48 @@ 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: + """Fetch a single memory by id, or ``None`` if no row matches.""" cursor = await self.db.execute( "SELECT * FROM memories WHERE id = ?", (memory_id,), @@ -167,6 +211,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] = [] @@ -198,6 +255,17 @@ 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) + + 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 +292,60 @@ 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: + """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) + + 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( @@ -470,6 +557,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,), @@ -542,53 +630,72 @@ def _sort_key(m: Memory) -> int: return results async def delete_expired(self) -> _List[str]: - now = _now_iso() + """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: - 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: + """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._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,59 +705,169 @@ 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: + """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() 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: + """Set ``expires_at`` (ISO timestamp) for one memory in one transaction.""" 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 +887,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..a567aa1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,15 +8,19 @@ 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 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 +46,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) -> AsyncIterator[aiosqlite.Connection]: conn = await get_connection(settings.db_path) await initialize_schema(conn, settings.embedding_dim) yield conn @@ -51,12 +61,22 @@ async def db(settings: Settings): @pytest_asyncio.fixture -async def memory_store(db) -> SQLiteMemoryStore: - return SQLiteMemoryStore(db) +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) -> SQLitePartitionStore: - store = SQLitePartitionStore(db) +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 + + +@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..984d550 100644 --- a/tests/integration/agents/test_consolidation_agent.py +++ b/tests/integration/agents/test_consolidation_agent.py @@ -2,20 +2,36 @@ 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, 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") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) # Create a memory in hebb mem = await memory_store.create( @@ -23,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 @@ -71,12 +87,18 @@ 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: 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).""" - kg = KnowledgeGraph(tmp_path / "kg.json") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) m1 = await memory_store.create( MemoryCreate( @@ -97,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": [ @@ -135,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", @@ -152,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), @@ -169,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, 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") + 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 ] @@ -194,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, 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") + + 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": []}, ] @@ -219,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, 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") + 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 ] @@ -239,9 +288,17 @@ 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: 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") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) recall_agent = RecallAgent( llm=mock_llm, @@ -261,16 +318,22 @@ 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: 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") + kg = KnowledgeGraph(tmp_path / "kg.json", lock=shared_lock) mem = await memory_store.create( 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 00d4ad9..2538f33 100644 --- a/tests/integration/scheduler/test_scheduler_manager.py +++ b/tests/integration/scheduler/test_scheduler_manager.py @@ -2,18 +2,31 @@ 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, tmp_path): - kg = KnowledgeGraph(tmp_path / "kg.json") + 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( settings=settings, diff --git a/tests/integration/storage/test_purge.py b/tests/integration/storage/test_purge.py index 1fbf63d..1946ab6 100644 --- a/tests/integration/storage/test_purge.py +++ b/tests/integration/storage/test_purge.py @@ -7,16 +7,24 @@ 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, tmp_path): - kg = KnowledgeGraph(tmp_path / "kg.json") +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,10 +40,14 @@ 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: 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") + 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..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) @@ -119,7 +124,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 +214,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 mock_llm.complete_json.side_effect = [ # RecallAgent.recall -> query generation (returns the existing mem) @@ -250,11 +258,17 @@ 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: 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 - 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..53e46a5 100644 --- a/tests/unit/test_audit_storage.py +++ b/tests/unit/test_audit_storage.py @@ -9,14 +9,18 @@ - 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 import asyncio +from collections.abc import AsyncIterator 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 @@ -28,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) @@ -38,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: @@ -49,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) @@ -103,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 @@ -126,3 +131,266 @@ 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 + ) -> AsyncIterator[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: 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): + 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. + + 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): + 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 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: + 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) + kg.save() + + async def forgetting_sim() -> None: + """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: + 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_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}" + ) diff --git a/tests/unit/test_audit_write_api.py b/tests/unit/test_audit_write_api.py index 826f357..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, @@ -193,9 +194,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 +217,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 +348,294 @@ 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). 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, 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 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") + row = await cursor.fetchone() + return int(row[0]) if row and row[0] is not None else 0