fix(storage): unify SQLite + graph writes behind one process-level lock - #52
fix(storage): unify SQLite + graph writes behind one process-level lock#52Inoriac wants to merge 5 commits into
Conversation
Route every SQL write and its KnowledgeGraph mutation through a single shared asyncio.Lock so they form one crash-consistent critical section, closing the split-lock gap where interrupted consolidation left graph orphans. Fixes afx-team#36
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a shared process-level write lock across the SQLite memory store, partition store, and knowledge graph to ensure crash consistency and prevent interleaving of SQL and graph mutations. It refactors the SQLite memory store to expose non-locking implementation methods allowing atomic composition of operations. The review feedback highlights a critical consistency issue across several files where the in-memory knowledge graph is mutated before the database transaction commits. If the transaction fails, the graph is left in an inconsistent state. It is highly recommended to apply the suggested changes to perform graph mutations only after successful database commits.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if callable(create_impl) and callable(delete_impl) and callable(begin): | ||
| async with self.kg.lock: | ||
| await begin() | ||
| try: | ||
| new_memory = await create_impl( | ||
| data=MemoryCreate( | ||
| content=consolidated_content, | ||
| partition_id=target_partition, | ||
| importance_score=importance, | ||
| tags=tags, | ||
| metadata=memory.metadata, | ||
| source="consolidation", | ||
| ), | ||
| embedding=embedding, | ||
| skip_tx=True, | ||
| ) | ||
| self.kg.update_from_tags(tags, new_memory.id) | ||
| self.kg.remove_memory_from_tags(memory.id) | ||
| await delete_impl(memory.id, skip_tx=True) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise | ||
| self.kg.save() |
There was a problem hiding this comment.
If the database transaction fails or rolls back (e.g., due to a busy timeout or write error), the in-memory knowledge graph (self.kg.graph) has already been mutated by self.kg.update_from_tags and self.kg.remove_memory_from_tags. Since there is no rollback mechanism for the in-memory NetworkX graph, the graph is left in an inconsistent state, and subsequent saves will persist this corrupted state to disk.
To ensure consistency, perform all knowledge graph mutations only after the database transaction has successfully committed.
if callable(create_impl) and callable(delete_impl) and callable(begin):
async with self.kg.lock:
await begin()
try:
new_memory = await create_impl(
data=MemoryCreate(
content=consolidated_content,
partition_id=target_partition,
importance_score=importance,
tags=tags,
metadata=memory.metadata,
source="consolidation",
),
embedding=embedding,
skip_tx=True,
)
await delete_impl(memory.id, skip_tx=True)
await self.memory_store.db.commit()
except BaseException:
await self.memory_store.db.rollback()
raise
self.kg.update_from_tags(tags, new_memory.id)
self.kg.remove_memory_from_tags(memory.id)
self.kg.save()| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for m in memories: | ||
| await delete_impl(m.id, skip_tx=True) | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise |
There was a problem hiding this comment.
If the database transaction fails or rolls back, the in-memory knowledge graph will have already been mutated by self.kg.remove_memory_from_tags. Perform the graph mutations only after the database transaction has successfully committed.
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for m in memories: | |
| await delete_impl(m.id, skip_tx=True) | |
| self.kg.remove_memory_from_tags(m.id) | |
| await self.memory_store.db.commit() | |
| except BaseException: | |
| await self.memory_store.db.rollback() | |
| raise | |
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for m in memories: | |
| await delete_impl(m.id, skip_tx=True) | |
| await self.memory_store.db.commit() | |
| except BaseException: | |
| await self.memory_store.db.rollback() | |
| raise | |
| for m in memories: | |
| self.kg.remove_memory_from_tags(m.id) |
| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for m in memories: | ||
| await delete_impl(m.id, skip_tx=True) | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise |
There was a problem hiding this comment.
If the database transaction fails or rolls back, the in-memory knowledge graph will have already been mutated by self.kg.remove_memory_from_tags. Perform the graph mutations only after the database transaction has successfully committed.
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for m in memories: | |
| await delete_impl(m.id, skip_tx=True) | |
| self.kg.remove_memory_from_tags(m.id) | |
| await self.memory_store.db.commit() | |
| except BaseException: | |
| await self.memory_store.db.rollback() | |
| raise | |
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for m in memories: | |
| await delete_impl(m.id, skip_tx=True) | |
| await self.memory_store.db.commit() | |
| except BaseException: | |
| await self.memory_store.db.rollback() | |
| raise | |
| for m in memories: | |
| self.kg.remove_memory_from_tags(m.id) |
| elif resolution == "discard": | ||
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | ||
| async with kg_lock: | ||
| self.kg.remove_memory_from_tags(memory.id) | ||
| if callable(delete_impl): | ||
| await delete_impl(memory.id) | ||
| else: | ||
| await self.memory_store.delete(memory.id) | ||
| self.kg.save() |
There was a problem hiding this comment.
If the delete operation fails or raises an exception, the in-memory knowledge graph will have already been mutated by self.kg.remove_memory_from_tags. Perform the graph mutation only after the delete operation has successfully completed.
| elif resolution == "discard": | |
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | |
| async with kg_lock: | |
| self.kg.remove_memory_from_tags(memory.id) | |
| if callable(delete_impl): | |
| await delete_impl(memory.id) | |
| else: | |
| await self.memory_store.delete(memory.id) | |
| self.kg.save() | |
| elif resolution == "discard": | |
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | |
| async with kg_lock: | |
| if callable(delete_impl): | |
| await delete_impl(memory.id) | |
| else: | |
| await self.memory_store.delete(memory.id) | |
| self.kg.remove_memory_from_tags(memory.id) | |
| self.kg.save() |
| if callable(create_impl) and callable(delete_impl) and callable(begin): | ||
| # SQLite: shared lock + single BEGIN IMMEDIATE wrapping both the | ||
| # SQL create and SQL delete so they commit as one atomic unit | ||
| # (Issue #36 — crash-consistent across store + graph). | ||
| async with kg_lock: | ||
| await begin() | ||
| try: | ||
| new_memory = await create_impl( | ||
| data=MemoryCreate( | ||
| content=consolidated_content, | ||
| partition_id=target_partition, | ||
| importance_score=importance, | ||
| tags=tags, | ||
| metadata=memory.metadata, | ||
| source="consolidation", | ||
| ), | ||
| embedding=embedding, | ||
| skip_tx=True, | ||
| ) | ||
| self.kg.update_from_tags(tags, new_memory.id) | ||
| self.kg.remove_memory_from_tags(memory.id) | ||
| await delete_impl(memory.id, skip_tx=True) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise | ||
| self.kg.save() |
There was a problem hiding this comment.
If the database transaction fails or rolls back, the in-memory knowledge graph will have already been mutated by self.kg.update_from_tags and self.kg.remove_memory_from_tags. Perform the graph mutations only after the database transaction has successfully committed.
if callable(create_impl) and callable(delete_impl) and callable(begin):
# SQLite: shared lock + single BEGIN IMMEDIATE wrapping both the
# SQL create and SQL delete so they commit as one atomic unit
# (Issue #36 — crash-consistent across store + graph).
async with kg_lock:
await begin()
try:
new_memory = await create_impl(
data=MemoryCreate(
content=consolidated_content,
partition_id=target_partition,
importance_score=importance,
tags=tags,
metadata=memory.metadata,
source="consolidation",
),
embedding=embedding,
skip_tx=True,
)
await delete_impl(memory.id, skip_tx=True)
await self.memory_store.db.commit()
except BaseException:
await self.memory_store.db.rollback()
raise
self.kg.update_from_tags(tags, new_memory.id)
self.kg.remove_memory_from_tags(memory.id)
self.kg.save()| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for mid in to_delete: | ||
| await delete_impl(mid, skip_tx=True) | ||
| self.knowledge_graph.remove_memory_from_tags(mid) | ||
| total_deleted += 1 | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise |
There was a problem hiding this comment.
If the database transaction fails or rolls back, the in-memory knowledge graph will have already been mutated by self.knowledge_graph.remove_memory_from_tags. Perform the graph mutations and increment total_deleted only after the database transaction has successfully committed.
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for mid in to_delete: | |
| await delete_impl(mid, skip_tx=True) | |
| self.knowledge_graph.remove_memory_from_tags(mid) | |
| total_deleted += 1 | |
| await self.memory_store.db.commit() | |
| except BaseException: | |
| await self.memory_store.db.rollback() | |
| raise | |
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for mid in to_delete: | |
| await delete_impl(mid, skip_tx=True) | |
| await self.memory_store.db.commit() | |
| except BaseException: | |
| await self.memory_store.db.rollback() | |
| raise | |
| for mid in to_delete: | |
| self.knowledge_graph.remove_memory_from_tags(mid) | |
| total_deleted += 1 |
| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for mid in to_delete: | ||
| await delete_impl(mid, skip_tx=True) | ||
| kg.remove_memory_from_tags(mid) | ||
| deleted += 1 | ||
| await memory_store.db.commit() | ||
| except BaseException: | ||
| await memory_store.db.rollback() | ||
| raise |
There was a problem hiding this comment.
If the database transaction fails or rolls back, the in-memory knowledge graph will have already been mutated by kg.remove_memory_from_tags. Perform the graph mutations and increment deleted only after the database transaction has successfully committed.
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for mid in to_delete: | |
| await delete_impl(mid, skip_tx=True) | |
| kg.remove_memory_from_tags(mid) | |
| deleted += 1 | |
| await memory_store.db.commit() | |
| except BaseException: | |
| await memory_store.db.rollback() | |
| raise | |
| if callable(begin) and callable(delete_impl): | |
| await begin() | |
| try: | |
| for mid in to_delete: | |
| await delete_impl(mid, skip_tx=True) | |
| await memory_store.db.commit() | |
| except BaseException: | |
| await memory_store.db.rollback() | |
| raise | |
| for mid in to_delete: | |
| kg.remove_memory_from_tags(mid) | |
| deleted += 1 |
|
Warning Review limit reached
Next review available in: 32 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughSQLite writes, knowledge-graph mutations, consolidation, and forgetting now use shared process-level locking and explicit transactions. Composite operations coordinate memory, embedding, FTS, and graph updates. Fallback schema migration and concurrency rollback tests were expanded. ChangesTransactional coordination
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant ConsolidationAgent
participant SQLiteMemoryStore
participant KnowledgeGraph
ConsolidationAgent->>KnowledgeGraph: acquire shared lock
ConsolidationAgent->>SQLiteMemoryStore: begin transaction
ConsolidationAgent->>SQLiteMemoryStore: create memory and delete sources
SQLiteMemoryStore-->>ConsolidationAgent: commit or rollback
ConsolidationAgent->>KnowledgeGraph: update tags and save graph
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
This PR addresses Issue #36 / audit C1 by introducing a single process-level asyncio lock intended to serialize SQLite mutations and KnowledgeGraph mutations together, reducing split-lock crash windows during consolidation/forgetting.
Changes:
- Introduces a shared process-level
asyncio.Locksurfaced asStorageContext.write_lock, and wires it through SQLite stores and the KnowledgeGraph. - Refactors
SQLiteMemoryStore(andSQLitePartitionStore) to accept an injected lock and adds_*_implmethods to compose multi-step operations under one externally-held critical section. - Updates consolidation/forgetting/purge paths and expands tests with a new concurrent write+consolidate+forget stress test plus shared-lock fixtures.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_audit_storage.py | Adds concurrent write+consolidate+forget stress test targeting Issue #36 acceptance criteria. |
| tests/unit/test_audit_embedding.py | Updates fake storage/KG scaffolding to reflect new shared-lock plumbing. |
| tests/unit/test_audit_consolidation.py | Passes a shared lock into KnowledgeGraph and updates conflict-update spying to match new single-tx path. |
| tests/integration/storage/test_purge.py | Updates purge integration tests to construct KG with shared_lock. |
| tests/integration/scheduler/test_scheduler_manager.py | Updates scheduler integration test to use KG with shared_lock. |
| tests/integration/agents/test_consolidation_agent.py | Updates agent integration tests to use KG with shared_lock. |
| tests/conftest.py | Adds shared_lock fixture and injects it into stores/KG for test parity with production. |
| src/hebb/storage/sqlite_store.py | Adds injectable lock + _*_impl methods and composite single-transaction update+embedding helper. |
| src/hebb/storage/purge.py | Attempts to make purge crash-consistent by using a shared lock path when available. |
| src/hebb/storage/partition_store.py | Adds injectable lock + explicit transactions for partition mutations. |
| src/hebb/storage/factory.py | Creates and returns a shared SQLite lock via StorageContext.write_lock. |
| src/hebb/server/routers/admin.py | Batches forgetting deletes under the KG/shared lock with an explicit transaction. |
| src/hebb/server/app.py | Constructs KnowledgeGraph with ctx.write_lock for shared critical section. |
| src/hebb/scheduler/manager.py | Batches scheduled forgetting deletes under the shared lock with an explicit transaction. |
| src/hebb/graph/knowledge_graph.py | Adds optional injected lock and defaults to its own lock when absent. |
| src/hebb/api.py | Constructs KnowledgeGraph with ctx.write_lock for shared critical section. |
| src/hebb/agents/consolidation_agent.py | Reworks consolidation paths to use store impl methods and shared locking/transactions. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # SQLite backend: _write_lock is the shared process-level lock (= kg.lock). | ||
| # Acquire it once so the SQL delete + KG mutation form one critical section. | ||
| shared_lock = getattr(store, "_write_lock", None) | ||
| if shared_lock is not None and hasattr(store, "_delete_impl"): | ||
| async with shared_lock: | ||
| deleted = await store._delete_impl(memory_id) # type: ignore[attr-defined] | ||
| kg.remove_memory_from_tags(memory_id) | ||
| if save: | ||
| kg.save() | ||
| else: | ||
| # PostgreSQL pool backend: no shared lock; KG guarded by its own lock. | ||
| deleted = await store.delete(memory_id) | ||
| async with kg.lock: | ||
| kg.remove_memory_from_tags(memory_id) | ||
| if save: | ||
| kg.save() | ||
| return deleted |
| if to_delete: | ||
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | ||
| begin = getattr(self.memory_store, "_begin", None) | ||
| async with self.knowledge_graph.lock: | ||
| for mid in to_delete: | ||
| # save=False: the graph is persisted once after the loop. | ||
| await purge_memory(self.memory_store, self.knowledge_graph, mid, save=False) | ||
| total_deleted += 1 | ||
| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for mid in to_delete: | ||
| await delete_impl(mid, skip_tx=True) | ||
| self.knowledge_graph.remove_memory_from_tags(mid) | ||
| total_deleted += 1 | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise | ||
| else: | ||
| for mid in to_delete: | ||
| await self.memory_store.delete(mid) | ||
| self.knowledge_graph.remove_memory_from_tags(mid) | ||
| total_deleted += 1 | ||
| self.knowledge_graph.save() |
| if to_delete: | ||
| delete_impl = getattr(memory_store, "_delete_impl", None) | ||
| begin = getattr(memory_store, "_begin", None) | ||
| async with kg.lock: | ||
| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for mid in to_delete: | ||
| await delete_impl(mid, skip_tx=True) | ||
| kg.remove_memory_from_tags(mid) | ||
| deleted += 1 | ||
| await memory_store.db.commit() | ||
| except BaseException: | ||
| await memory_store.db.rollback() | ||
| raise | ||
| else: | ||
| for mid in to_delete: | ||
| await memory_store.delete(mid) | ||
| kg.remove_memory_from_tags(mid) | ||
| deleted += 1 | ||
| kg.save() |
| create_impl = getattr(self.memory_store, "_create_impl", None) | ||
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | ||
| begin = getattr(self.memory_store, "_begin", None) | ||
| if callable(create_impl) and callable(delete_impl) and callable(begin): | ||
| async with self.kg.lock: | ||
| await begin() | ||
| try: | ||
| new_memory = await create_impl( | ||
| data=MemoryCreate( | ||
| content=consolidated_content, | ||
| partition_id=target_partition, | ||
| importance_score=importance, | ||
| tags=tags, | ||
| metadata=memory.metadata, | ||
| source="consolidation", | ||
| ), | ||
| embedding=embedding, | ||
| skip_tx=True, | ||
| ) | ||
| self.kg.update_from_tags(tags, new_memory.id) | ||
| self.kg.remove_memory_from_tags(memory.id) | ||
| await delete_impl(memory.id, skip_tx=True) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise | ||
| self.kg.save() | ||
| else: | ||
| # Other backends (PG pool): no shared lock. | ||
| new_memory = await self.memory_store.create( | ||
| data=MemoryCreate( | ||
| content=consolidated_content, | ||
| partition_id=target_partition, | ||
| importance_score=importance, | ||
| tags=tags, | ||
| metadata=memory.metadata, | ||
| source="consolidation", | ||
| ), | ||
| embedding=embedding, | ||
| ) | ||
| self.kg.update_from_tags(tags, new_memory.id) | ||
| await purge_memory(self.memory_store, self.kg, memory.id) |
| if results: | ||
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | ||
| begin = getattr(self.memory_store, "_begin", None) | ||
| async with kg_lock: | ||
| for m in memories: | ||
| await self.memory_store.delete(m.id) | ||
| # Strip the source id from the graph (no-op unless the | ||
| # source was itself graphed, i.e. in-partition | ||
| # re-consolidation). | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| # Persist the graph immediately after the SQL deletes, under | ||
| # the same lock, to shrink the crash window between "source | ||
| # row deleted" and "graph reference removed". The caller's | ||
| # end-of-batch save() is now a redundant safety net rather | ||
| # than the only persistence point. | ||
| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for m in memories: | ||
| await delete_impl(m.id, skip_tx=True) | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise | ||
| else: | ||
| for m in memories: | ||
| await self.memory_store.delete(m.id) | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| self.kg.save() |
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | ||
| async with kg_lock: | ||
| self.kg.remove_memory_from_tags(memory.id) | ||
| if callable(delete_impl): | ||
| await delete_impl(memory.id) | ||
| else: | ||
| await self.memory_store.delete(memory.id) | ||
| self.kg.save() |
| async with kg_lock: | ||
| for m in memories: | ||
| await self.memory_store.delete(m.id) | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | ||
| begin = getattr(self.memory_store, "_begin", None) | ||
| if callable(begin) and callable(delete_impl): | ||
| await begin() | ||
| try: | ||
| for m in memories: | ||
| await delete_impl(m.id, skip_tx=True) | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise | ||
| else: | ||
| for m in memories: | ||
| await self.memory_store.delete(m.id) | ||
| self.kg.remove_memory_from_tags(m.id) | ||
| self.kg.save() |
| embedding = await self.embedder.embed(consolidated_content) | ||
| new_memory = await self.memory_store.create( | ||
| data=MemoryCreate( | ||
| content=consolidated_content, | ||
| partition_id=target_partition, | ||
| importance_score=importance, | ||
| tags=tags, | ||
| metadata=memory.metadata, | ||
| source="consolidation", | ||
| ), | ||
| embedding=embedding, | ||
| ) | ||
|
|
||
| async with kg_lock: | ||
| self.kg.update_from_tags(tags, new_memory.id) | ||
| # Strip the source id under the same lock (no-op unless the | ||
| # source was graphed via in-partition re-consolidation). | ||
| self.kg.remove_memory_from_tags(memory.id) | ||
| await self.memory_store.delete(memory.id) | ||
| # Persist the graph immediately after the SQL delete, under the | ||
| # same lock, to shrink the crash window. The end-of-batch save() | ||
| # is now a redundant safety net. | ||
| self.kg.save() | ||
| create_impl = getattr(self.memory_store, "_create_impl", None) | ||
| delete_impl = getattr(self.memory_store, "_delete_impl", None) | ||
| begin = getattr(self.memory_store, "_begin", None) | ||
| if callable(create_impl) and callable(delete_impl) and callable(begin): | ||
| # SQLite: shared lock + single BEGIN IMMEDIATE wrapping both the | ||
| # SQL create and SQL delete so they commit as one atomic unit | ||
| # (Issue #36 — crash-consistent across store + graph). | ||
| async with kg_lock: | ||
| await begin() | ||
| try: | ||
| new_memory = await create_impl( | ||
| data=MemoryCreate( | ||
| content=consolidated_content, | ||
| partition_id=target_partition, | ||
| importance_score=importance, | ||
| tags=tags, | ||
| metadata=memory.metadata, | ||
| source="consolidation", | ||
| ), | ||
| embedding=embedding, | ||
| skip_tx=True, | ||
| ) | ||
| self.kg.update_from_tags(tags, new_memory.id) | ||
| self.kg.remove_memory_from_tags(memory.id) | ||
| await delete_impl(memory.id, skip_tx=True) | ||
| await self.memory_store.db.commit() | ||
| except BaseException: | ||
| await self.memory_store.db.rollback() | ||
| raise | ||
| self.kg.save() |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/hebb/agents/consolidation_agent.py`:
- Around line 174-195: Move all in-memory knowledge-graph mutations until after
the corresponding SQL commit succeeds, keeping them under the existing shared
lock. In src/hebb/agents/consolidation_agent.py#L174-L195 and `#L742-L764`, commit
before calling update_from_tags/remove_memory_from_tags; in `#L442-L493`, collect
IDs for both delete branches, commit, then remove their tags. Apply the same
collect-then-commit-then-remove sequence in
src/hebb/scheduler/manager.py#L346-L363 and
src/hebb/server/routers/admin.py#L192-L212 before kg.save().
- Around line 208-210: Wrap the self.kg.update_from_tags(tags, new_memory.id)
call in the PG fallback with self.kg.lock, matching the locking pattern used by
_consolidate_one and _consolidate_session_chunk. Keep purge_memory outside the
graph-mutation lock unless required by existing locking conventions.
In `@tests/conftest.py`:
- Around line 54-70: Complete strict type annotations across the affected
fixtures and tests: in tests/conftest.py lines 54-70, annotate db as an
async-generator returning the connection type and type the dependent fixture
parameters; in tests/integration/agents/test_consolidation_agent.py lines
18-269, annotate injected fixtures and test methods with -> None; in
tests/integration/scheduler/test_scheduler_manager.py lines 15-16, type fixture
parameters and the SchedulerManager return; in
tests/integration/storage/test_purge.py lines 18-38, type fixture parameters and
add -> None; in tests/unit/test_audit_consolidation.py lines 255-260, type
injected fixtures; and in tests/unit/test_audit_storage.py lines 150-165, use an
AsyncIterator fixture return and type the test parameter.
In `@tests/integration/agents/test_consolidation_agent.py`:
- Around line 18-20: Add complete type annotations to the changed test methods,
including explicit types for each fixture parameter and a -> None return type;
apply this consistently to test_consolidate_memory and the other referenced
methods while preserving the existing shared-lock wiring and test behavior.
In `@tests/integration/storage/test_purge.py`:
- Around line 18-19: Add type annotations to the affected purge test functions,
including fixture parameters such as memory_store, shared_lock, and tmp_path,
and annotate each function with -> None. Apply this consistently to
test_purge_removes_sql_row_and_graph_node and the other changed purge tests
referenced by the comment, using the repository’s established fixture types.
In `@tests/unit/test_audit_storage.py`:
- Around line 163-165: Annotate the shared_store_and_kg parameter in
test_no_orphans_after_concurrent_operations with the appropriate fixture type,
matching the type used by the shared_store_and_kg fixture and the project’s
existing test annotations.
- Around line 182-214: Update api_write, consolidation_sim, and forgetting_sim
to exercise one transaction per operation rather than independently committed
SQL and graph changes: hold the lock across the complete SQL-plus-graph update,
explicitly begin the transaction, call _create_impl/_delete_impl with
skip_tx=True, and commit or roll back once. Add failure injection and assertions
that reloaded SQL and graph state are unchanged after rollback, while preserving
successful-operation verification.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e7bb5827-80fe-42a4-b1b0-808768360876
📒 Files selected for processing (17)
src/hebb/agents/consolidation_agent.pysrc/hebb/api.pysrc/hebb/graph/knowledge_graph.pysrc/hebb/scheduler/manager.pysrc/hebb/server/app.pysrc/hebb/server/routers/admin.pysrc/hebb/storage/factory.pysrc/hebb/storage/partition_store.pysrc/hebb/storage/purge.pysrc/hebb/storage/sqlite_store.pytests/conftest.pytests/integration/agents/test_consolidation_agent.pytests/integration/scheduler/test_scheduler_manager.pytests/integration/storage/test_purge.pytests/unit/test_audit_consolidation.pytests/unit/test_audit_embedding.pytests/unit/test_audit_storage.py
…n, guards, types, docs Address the 22 review comments raised on PR afx-team#52 by gemini-code-assist, Copilot, and coderabbitai (topics 1–7 of the modification plan). No behavior change beyond the safety / typing / doc gaps the bots flagged. Data integrity (Issue afx-team#36; Gemini + CodeRabbit) * consolidation_agent.consolidate_memory, _consolidate_session_chunk, _consolidate_one (incl. discard), scheduler.manager scheduled forgetting, server.routers.admin manual forgetting, storage.purge_memory: move KnowledgeGraph mutations (update_from_tags / remove_memory_from_tags) and counters to AFTER db.commit() — still under the shared lock — so a SQL rollback can no longer diverge the in-memory NetworkX graph. The narrow commit → save() crash window is acknowledged; KG→SQL orphans left by a crash there are self-healed by reconcile() (KG→SQL is self-healing; SQL→KG lost forward refs are not, so the post-commit ordering is safer). Lock discipline (Issue afx-team#36; Copilot) * storage.purge_memory: explicitly acquire kg.lock even when the store's _write_lock differs, and handle the non-reentrant same-lock case (no second async with) so SDK / test wiring cannot race the graph mutation. PG fallback (CodeRabbit) * consolidation_agent non-SQLite fast-path: guard self.kg.update_from_tags with self.kg.lock so concurrent forgetting / consolidation cannot mutate the NetworkX graph unserialized. Typing (CodeRabbit + Pylance red flags) * consolidation_agent / scheduler.manager / routers.admin / purge: replace the getattr(store, "_create_impl", None) + callable() duck-type probe with isinstance(self.memory_store, SQLiteMemoryStore). Pylance / mypy now narrow .db / ._create_impl / ._delete_impl / ._begin / ._update_content_and_embedding_impl to the concrete store type, removing the 20 "await … not awaitable" / "MemoryStore has no attribute db" red flags the probe caused — no behavior change (same SQLite vs. PG split, same lock composition). * tests/conftest + the integration/audit regression tests: complete strict type annotations (fixture params + -> None / AsyncIterator[...]) so mypy is clean across the PR's test footprint. Drop two `# type: ignore` comments mypy 2.1 now reports as unused. Tests (CodeRabbit) * tests/unit/test_audit_storage.test_no_orphans_after_concurrent_operations: rewrite to exercise one transaction per operation (BEGIN + skip_tx=True + commit/rollback once) under the shared lock — previously the simulated consolidation/forgetting committed each _impl independently and released the lock between SQL and graph updates, so the test passed even when the split-lock / partial-commit failure modes regressed. * Add test_rollback_on_injected_failure_in_consolidation: inject a failure mid-transaction and assert both SQLite and the reloaded on-disk KG remain unchanged after rollback. Documentation (CodeRabbit pre-merge check, 55% -> 100% docstring coverage on the PR-touched files) * Every public callable in the PR-touched source files (consolidation_agent, scheduler.manager, server.app, server.routers.admin, storage.factory, storage.partition_store, storage.sqlite_store) now carries a docstring with Args/Returns/Raises where the signature warrants it. PR-touched coverage 70% -> 100%; src/hebb overall 68.44% -> 72.52%. CHANGELOG: add ## [Unreleased] section summarizing the above. Verified: * ruff check src/ tests/ → All checks passed * mypy src/hebb → Success: no issues (132 source files) * pytest tests/ → 779 passed, 1 skipped Refs afx-team#36 (PR afx-team#52).
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/hebb/storage/purge.py (1)
27-49: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocstring missing
Raisessection.
purge_memorydocumentsArgsandReturnsbut omitsRaises, even though SQL errors fromstore._delete_impl/store.deletepropagate uncaught.As per coding guidelines, "All public APIs in Python MUST have docstrings with Args, Returns, and Raises sections" (
src/**/*.py).📝 Proposed addition
Returns: ``True`` if the SQL row existed and was deleted, ``False`` otherwise. Graph cleanup runs regardless of the row's existence, so a pre-existing orphan referencing ``memory_id`` is still swept. + + Raises: + Exception: Propagated from the underlying SQL delete (e.g. driver + errors) if the transaction fails. """🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/hebb/storage/purge.py` around lines 27 - 49, Add a Raises section to the purge_memory docstring documenting the SQL/database exceptions propagated from store._delete_impl or store.delete, while preserving the existing Args and Returns documentation.Source: Coding guidelines
tests/unit/test_audit_storage.py (1)
35-39: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUntyped
tmp_pathparameter (partial annotation).
tmp_pathhas no type annotation whileembedding_dim: intdoes — under mypy strict this is a partially-annotated function. Same gap recurs in theshared_store_and_kgfixture (Line 154).Also applies to: 152-155.
As per coding guidelines, "Add type hints on all public functions (mypy strict is enabled in pyproject.toml)."
🛠️ Proposed fix
-async def store(tmp_path, embedding_dim: int) -> AsyncIterator[SQLiteMemoryStore]: +async def store(tmp_path: Path, embedding_dim: int) -> AsyncIterator[SQLiteMemoryStore]:async def shared_store_and_kg( - self, tmp_path, embedding_dim: int + self, tmp_path: Path, embedding_dim: int ) -> AsyncIterator[tuple[SQLiteMemoryStore, KnowledgeGraph, asyncio.Lock]]:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_audit_storage.py` around lines 35 - 39, Add an explicit pathlib path type annotation to the tmp_path parameter of the store fixture and the corresponding shared_store_and_kg fixture, preserving their existing behavior and return annotations so both functions are fully typed under mypy strict.Source: Coding guidelines
🧹 Nitpick comments (2)
src/hebb/storage/purge.py (2)
71-80: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo test exercises the differing-lock branch.
async with shared_lock, kg.lock:only runs when the store's_write_lockdiffers fromkg.lock. Intests/integration/storage/test_purge.pyboth fixtures wire the sameshared_lockintomemory_storeandkg, soshared_lock is kg.lockis alwaysTrueand this branch is never hit.Add a purge test that constructs
SQLiteMemoryStore/KnowledgeGraphwith two distinctasyncio.Lockinstances to exercise this dual-lock path and its documented lock-ordering discipline.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/hebb/storage/purge.py` around lines 71 - 80, Add a test in the purge test suite that constructs SQLiteMemoryStore and KnowledgeGraph with distinct asyncio.Lock instances, invokes the purge flow with the differing-lock configuration, and verifies the memory is deleted and removed from the knowledge graph. Ensure the test exercises the shared_lock-then-kg.lock acquisition order documented in the dual-lock branch.
57-58: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
storage.purgestill duck-types; changelog claims it was migrated toisinstance.CHANGELOG.mdstates thegetattr(..., "_create_impl", None)+callable()probe was replaced byisinstance(self.memory_store, SQLiteMemoryStore)acrossconsolidation_agent,scheduler.manager,server.routers.admin, andstorage.purge— butpurge_memorystill usesgetattr(store, "_write_lock", None)/hasattr(store, "_delete_impl"), notisinstance.
src/hebb/storage/purge.py#L57-L58: replace withisinstance(store, SQLiteMemoryStore)(importingSQLiteMemoryStore) to match the pattern used inconsolidation_agent.pyand gain proper mypy narrowing onstore._write_lock/store._delete_impl, or leave duck-typing intentionally and correct the changelog wording.CHANGELOG.md#L42-L51: dropstorage.purgefrom the isinstance-migration bullet if the duck-typing here is intentional (e.g. becausepurge_memorymust also support non-SQLiteMemoryStorebackends exposing_write_lock), to keep the changelog accurate.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/hebb/storage/purge.py` around lines 57 - 58, Update purge_memory in src/hebb/storage/purge.py at lines 57-58 to import SQLiteMemoryStore and use isinstance(store, SQLiteMemoryStore), enabling narrowing for _write_lock and _delete_impl; keep the existing purge behavior. Update CHANGELOG.md at lines 42-51 to retain storage.purge in the isinstance-migration entry only after this change; otherwise remove it if duck-typing remains intentional.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/hebb/storage/purge.py`:
- Around line 27-49: Add a Raises section to the purge_memory docstring
documenting the SQL/database exceptions propagated from store._delete_impl or
store.delete, while preserving the existing Args and Returns documentation.
In `@tests/unit/test_audit_storage.py`:
- Around line 35-39: Add an explicit pathlib path type annotation to the
tmp_path parameter of the store fixture and the corresponding
shared_store_and_kg fixture, preserving their existing behavior and return
annotations so both functions are fully typed under mypy strict.
---
Nitpick comments:
In `@src/hebb/storage/purge.py`:
- Around line 71-80: Add a test in the purge test suite that constructs
SQLiteMemoryStore and KnowledgeGraph with distinct asyncio.Lock instances,
invokes the purge flow with the differing-lock configuration, and verifies the
memory is deleted and removed from the knowledge graph. Ensure the test
exercises the shared_lock-then-kg.lock acquisition order documented in the
dual-lock branch.
- Around line 57-58: Update purge_memory in src/hebb/storage/purge.py at lines
57-58 to import SQLiteMemoryStore and use isinstance(store, SQLiteMemoryStore),
enabling narrowing for _write_lock and _delete_impl; keep the existing purge
behavior. Update CHANGELOG.md at lines 42-51 to retain storage.purge in the
isinstance-migration entry only after this change; otherwise remove it if
duck-typing remains intentional.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 75f6ea7b-7f71-4c8a-8b34-d0d9b2c9db99
📒 Files selected for processing (15)
CHANGELOG.mdsrc/hebb/agents/consolidation_agent.pysrc/hebb/scheduler/manager.pysrc/hebb/server/app.pysrc/hebb/server/routers/admin.pysrc/hebb/storage/factory.pysrc/hebb/storage/partition_store.pysrc/hebb/storage/purge.pysrc/hebb/storage/sqlite_store.pytests/conftest.pytests/integration/agents/test_consolidation_agent.pytests/integration/scheduler/test_scheduler_manager.pytests/integration/storage/test_purge.pytests/unit/test_audit_consolidation.pytests/unit/test_audit_storage.py
🚧 Files skipped from review as they are similar to previous changes (10)
- src/hebb/server/app.py
- tests/integration/scheduler/test_scheduler_manager.py
- src/hebb/storage/factory.py
- src/hebb/server/routers/admin.py
- src/hebb/scheduler/manager.py
- tests/unit/test_audit_consolidation.py
- tests/conftest.py
- src/hebb/storage/partition_store.py
- src/hebb/agents/consolidation_agent.py
- src/hebb/storage/sqlite_store.py
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/hebb/storage/migrations.py`:
- Around line 245-249: Update the migration that adds partition_id to backfill
each existing memory_embeddings row from its corresponding
memories.partition_id, using 'default' only when no matching memory exists.
Preserve the NOT NULL constraint and ensure the downstream sqlite_store insert
contract receives the migrated partition as authoritative metadata.
- Around line 235-241: The schema check before the migration DDL must reliably
identify vec0 virtual tables. Update the sqlite_master query and detection logic
in the migration flow to use table metadata or normalized, case-insensitive
matching for a USING vec0 clause, including spacing or trailing options, and
return before ALTER TABLE or CREATE INDEX when the table is virtual.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: fe8b1da7-231d-40a9-bf8b-b829df543048
📒 Files selected for processing (1)
src/hebb/storage/migrations.py
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate file =="
wc -l src/hebb/storage/migrations.py
echo
echo "== relevant migrations.py snippets =="
sed -n '200,270p' src/hebb/storage/migrations.py
echo
sed -n '290,330p' src/hebb/storage/migrations.py
echo
echo "== references to fallback partition / vec table helpers =="
rg -n "_ensure_fallback_partition_column|_ensure_vec_table|memory_embeddings|partition_id|vec0|frozen_data" src/hebb/storage/migrations.py src/hebb -S
echo
echo "== git diff stat/name =="
git diff --stat || true
git diff -- src/hebb/storage/migrations.py | sed -n '1,220p' || trueRepository: afx-team/hebb-mind
Length of output: 47589
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== check local sqlite availability =="
python3 - <<'PY'
import sqlite3
print("sqlite3.version", sqlite3.sqlite_version)
conn = sqlite3.connect(":memory:")
cursor = conn.cursor()
for schema, expected_err_type in [
("CREATE VIRTUAL TABLE memory_embeddings USING vec0(memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding float[384])", sqlite3.OperationalError),
("CREATE TABLE memory_embeddings (memory_id TEXT PRIMARY KEY, embedding BLOB)", None),
]:
cursor.execute("DROP TABLE IF EXISTS memory_embeddings")
cursor.executescript(schema)
print("schema", schema)
try:
cursor.executescript("ALTER TABLE memory_embeddings ADD COLUMN added TEXT")
cursor.executescript("CREATE INDEX idx_memory_embeddings_added ON memory_embeddings(added)")
print("alter+index succeeded")
except Exception as err:
print(type(err).__name__, str(err))
print("expected", expected_err_type and issubclass(type(err), expected_err_type))
PY
echo
echo "== inspect sqlite_master.sql for vec0 CREATE form in this repository/code/tests =="
rg -n -i "CREATE\s+VIRTUAL\s+TABLE\s+memory_embeddings\s+USING\s+vec0|USING\s+vec0|vec0" . -S --glob '!node_modules' --glob '!venv' | head -80Repository: afx-team/hebb-mind
Length of output: 338
🌐 Web query:
SQLite CREATE VIRTUAL TABLE sql stored normalized USING module uppercase documentation
💡 Result:
In SQLite, there is no single "normalized" or "uppercase" module for virtual tables. Instead, normalization (such as case folding or removing diacritics) is typically handled by the tokenizer within the Full-Text Search (FTS5) virtual table module [1][2]. Virtual tables are created using the CREATE VIRTUAL TABLE statement, which interfaces with an external module [3][4]. The module name is specified after the USING keyword (e.g., USING fts5) [4][5]. For FTS5 virtual tables, you can control text normalization through the tokenize option [6][1]: 1. Default Normalization (unicode61): By default, the FTS5 module uses the unicode61 tokenizer [1][7]. This tokenizer automatically performs case normalization, converting characters to lowercase [1][2]. It also handles Unicode characters based on their general category [1]. 2. Customizing Tokenization: If you need specific normalization (e.g., forcing uppercase or custom diacritic handling), you can implement a custom tokenizer in C or, in some cases, use auxiliary functions to process query terms during search [1][7]. 3. Implementation: When defining the table, you specify the tokenizer configuration as part of the module arguments [5][6]: CREATE VIRTUAL TABLE my_table USING fts5(content_col, tokenize='unicode61'); The built-in tokenizers (unicode61, ascii, porter, trigram) are designed to handle common normalization tasks like case folding [1][2]. The ascii tokenizer also performs case normalization, similar to unicode61, but differs in how it handles non-ASCII characters [2]. If your goal is specifically to have data stored or indexed in uppercase, you would typically handle this by using a custom tokenizer that transforms tokens before they are indexed, or by preprocessing the data before inserting it into the virtual table [1].
Citations:
- 1: https://www.sqlite.org/fts5.html
- 2: https://audrey.feldroy.com/articles/2025-01-13-SQLite-FTS5-Tokenizers-unicode61-and-ascii
- 3: https://sqlite.org/lang_createvtab.html
- 4: https://www.sqlite.org/vtab.html
- 5: https://www2.sqlite.org/draft/fts3.html
- 6: https://www2.sqlite.org/matrix/fts5.html
- 7: https://www2.sqlite.org/draft/fts5.html
🌐 Web query:
GitHub sqlite-vec "CREATE VIRTUAL TABLE memory_embeddings USING vec0"
💡 Result:
The command CREATE VIRTUAL TABLE name USING vec0(...) is the standard syntax for defining a vector-enabled virtual table in the sqlite-vec extension [1][2]. The vec0 virtual table allows you to store vector data (float, int8, or binary) and perform similarity search (KNN) queries [1][3]. Below is a breakdown of how to use this syntax: Basic Structure To create a table for vector storage, you define it with at least one vector column specifying the data type and dimension [1][4]: CREATE VIRTUAL TABLE memory_embeddings USING vec0( embedding_column float); Advanced Features The vec0 table also supports additional column types for metadata, partition keys, and auxiliary data [5][6]: 1. Metadata Columns: Standard columns that are indexed alongside your vectors, allowing you to use them in WHERE clauses for filtered similarity searches [6]. 2. Partition Keys: Defined with the partition key keyword, these internally shard the vector index to optimize queries for large datasets [5][6]. 3. Auxiliary Columns: Defined with a + prefix (e.g., +text_data text), these store unindexed data separately from the vector index. This allows you to retrieve large blobs or text fields directly in your KNN result set without requiring a separate JOIN [5][6]. Example with multiple column types: CREATE VIRTUAL TABLE memory_embeddings USING vec0( user_id integer partition key, category text, -- Metadata column embedding float, +raw_text text -- Auxiliary column); Once created, you can insert data and perform similarity searches using the MATCH operator [1][2]: SELECT rowid, distance, raw_text FROM memory_embeddings WHERE embedding MATCH '[...]' -- vector input ORDER BY distance LIMIT 10; For detailed documentation, refer to the official sqlite-vec repository or the project's documentation site [1][6].
Citations:
- 1: https://github.com/asg017/sqlite-vec
- 2: https://github.com/asg017/sqlite-vec/blob/main/README.md
- 3: https://github.com/asg017/sqlite-vec/tree/refs/heads/main
- 4: https://alexgarcia.xyz/sqlite-vec/
- 5: https://github.com/asg017/sqlite-vec/blob/main/site/features/vec0.md
- 6: https://alexgarcia.xyz/sqlite-vec/features/vec0.html
🌐 Web query:
sqlite-vec source CREATE VIRTUAL TABLE vec0 USING keyword stored schema
💡 Result:
The vec0 virtual table is the core component of the sqlite-vec extension used for high-performance vector search [1][2]. It allows for storing and querying float, int8, and binary vectors [3][4]. Syntax and Configuration: To create a vec0 table, use the standard SQLite CREATE VIRTUAL TABLE syntax [1][5]. The basic syntax is: CREATE VIRTUAL TABLE table_name USING vec0( vector_column_name [type][dimension]); Example: CREATE VIRTUAL TABLE vec_movies USING vec0( synopsis_embedding float); Non-Vector Data Support: You can store non-vector data alongside your vectors using three specific column types [1][6]: 1. Metadata Columns: Standard columns that are indexed along with the vector columns. These allow you to include extra WHERE constraints during K-Nearest-Neighbors (KNN) queries [1][6]. 2. Partition Keys: Columns marked with the partition key keyword. These internally shard the vector index on the given key, which can improve performance for selective queries [1][6]. 3. Auxiliary Columns: Columns prefixed with a + sign (e.g., +contents text). These are stored in a separate internal table, meaning they are not indexed for vector search but can be retrieved in result sets without requiring an explicit JOIN [1][6]. Example with various column types: CREATE VIRTUAL TABLE vec_chunks USING vec0( document_id integer partition key, contents_embedding float, label text, +contents text); When querying, you perform KNN searches using the MATCH operator on the vector column [3][7]. Note that sqlite-vec is a pre-v1 project, so users should expect potential breaking changes in future updates [4].
Citations:
- 1: https://alexgarcia.xyz/sqlite-vec/features/vec0.html
- 2: https://alexgarcia.xyz/sqlite-vec/features/knn.html
- 3: https://github.com/asg017/sqlite-vec
- 4: https://github.com/asg017/sqlite-vec?tab=readme-ov-file
- 5: https://alexgarcia.xyz/sqlite-vec/
- 6: https://github.com/asg017/sqlite-vec/blob/main/site/features/vec0.md
- 7: https://github.com/asg017/sqlite-vec/tree/refs/heads/main
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== local sqlite structural behavior for virtual-table detection =="
python3 - <<'PY'
import sqlite3
import sqlite3.sqlite
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id TEXT PRIMARY KEY, embedding BLOB, UNINDEXED)")
row = conn.execute("SELECT type FROM sqlite_master WHERE name = 'mem'").fetchone()
print("virtual table sqlite_master.type:", row[0] if row else None)
try:
conn.execute("ALTER TABLE mem ADD COLUMN added TEXT")
print("ALTER TABLE virtual table succeeded unexpectedly")
except Exception as e:
print("ALTER virtual table raised:", type(e).__name__, e)
try:
conn.execute("CREATE INDEX idx_mem_added ON mem(added)")
print("CREATE INDEX virtual table succeeded unexpectedly")
except Exception as e:
print("CREATE INDEX virtual table raised:", type(e).__name__, e)
try:
conn.execute("ALTER TABLE memory_embeddings ADD COLUMN partition_id TEXT NOT NULL DEFAULT 'default'")
print("ALTER TABLE nonexistent table succeeded unexpectedly")
except Exception as e:
print("ALTER missing table raised:", type(e).__name__, e)
PY
echo
echo "== local sqlite_master.sql shape for deterministic virtual-table SQL with fts5 module =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id TEXT PRIMARY KEY, embedding BLOB UNINDEXED)")
row = conn.execute("SELECT sql FROM sqlite_master WHERE name = 'mem'").fetchone()
print("CREATE VIRTUAL TABLE SQLite returns:", repr(row[0] if row else None))
PY
echo
echo "== check sqlite_master virtual table type in migration helper context =="
sed -n '232,244p' src/hebb/storage/migrations.pyRepository: afx-team/hebb-mind
Length of output: 350
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== local sqlite structural behavior for virtual-table detection =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id TEXT PRIMARY KEY, embedding BLOB UNINDEXED)")
row = conn.execute("SELECT type FROM sqlite_master WHERE name = 'mem'").fetchone()
print("virtual table sqlite_master.type:", row[0] if row else None)
try:
conn.execute("ALTER TABLE mem ADD COLUMN added TEXT")
print("ALTER TABLE virtual table succeeded unexpectedly")
except Exception as e:
print("ALTER virtual table raised:", type(e).__name__, e)
try:
conn.execute("CREATE INDEX idx_mem_added ON mem(added)")
print("CREATE INDEX virtual table succeeded unexpectedly")
except Exception as e:
print("CREATE INDEX virtual table raised:", type(e).__name__, e)
try:
conn.execute("ALTER TABLE memory_embeddings ADD COLUMN partition_id TEXT NOT NULL DEFAULT 'default'")
print("ALTER TABLE nonexistent table succeeded unexpectedly")
except Exception as e:
print("ALTER missing table raised:", type(e).__name__, e)
PY
echo
echo "== local sqlite_master.sql shape for deterministic virtual-table SQL with fts5 module =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id TEXT PRIMARY KEY, embedding BLOB UNINDEXED)")
row = conn.execute("SELECT sql FROM sqlite_master WHERE name = 'mem'").fetchone()
print("CREATE VIRTUAL TABLE SQLite returns:", repr(row[0] if row else None))
PY
echo
echo "== check sqlite_master virtual table type in migration helper context =="
sed -n '232,244p' src/hebb/storage/migrations.pyRepository: afx-team/hebb-mind
Length of output: 366
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== local sqlite structural behavior with valid FTS5 primary key =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id tokenizer unicode61, embedding BLOB, UNINDEXED)")
row = conn.execute("SELECT type FROM sqlite_master WHERE name = 'mem'").fetchone()
print("virtual table sqlite_master.type:", row[0] if row else None)
try:
conn.execute("ALTER TABLE mem ADD COLUMN added TEXT")
print("ALTER TABLE virtual table succeeded unexpectedly")
except Exception as e:
print("ALTER virtual table raised:", type(e).__name__, str(e))
try:
conn.execute("CREATE INDEX idx_mem_added ON mem(added)")
print("CREATE INDEX virtual table succeeded unexpectedly")
except Exception as e:
print("CREATE INDEX virtual table raised:", type(e).__name__, str(e))
try:
conn.execute("ALTER TABLE memory_embeddings ADD COLUMN partition_id TEXT NOT NULL DEFAULT 'default'")
print("ALTER TABLE nonexistent table succeeded unexpectedly")
except Exception as e:
print("ALTER missing table raised:", type(e).__name__, str(e))
PY
echo
echo "== local sqlite_master.sql shape for deterministic virtual-table SQL with fts5 module =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id tokenizer unicode61, embedding BLOB, UNINDEXED)")
row = conn.execute("SELECT sql FROM sqlite_master WHERE name = 'mem'").fetchone()
print("CREATE VIRTUAL TABLE SQLite returns:", repr(row[0] if row else None))
PY
echo
echo "== sqlite_master table-info shape for regular-table detection =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE TABLE mem (memory_id TEXT PRIMARY KEY, embedding BLOB)")
row = conn.execute("SELECT type FROM sqlite_master WHERE name = 'mem'").fetchone()
print("regular table sqlite_master.type:", row[0] if row else None)
PY
echo
echo "== sqlite_master virtual module-name shape =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id tokenizer unicode61, embedding BLOB, UNINDEXED)")
row = conn.execute("SELECT sql FROM sqlite_master WHERE name = 'mem'").fetchone()
print("sql:", row[0])
print("case-sensitive using fts5 present:", "using fts5" in (row[0] or ""))
PYRepository: afx-team/hebb-mind
Length of output: 369
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== probe sqlite_master type and ALTER/CREATE INDEX behavior for FTS5 virtual tables =="
python3 - <<'PY'
import sqlite3
conn = sqlite3.connect(":memory:")
conn.execute("CREATE VIRTUAL TABLE mem USING fts5(memory_id, embedding content=mem)")
row = conn.execute("SELECT type FROM sqlite_master WHERE name = 'mem'").fetchone()
print("virtual table sqlite_master.type:", row[0] if row else None)
print("virtual table sqlite_master.sql:", conn.execute("SELECT sql FROM sqlite_master WHERE name = 'mem'").fetchone()[0])
try:
conn.execute("ALTER TABLE mem ADD COLUMN added TEXT")
print("ALTER TABLE virtual table succeeded unexpectedly")
except Exception as e:
print("ALTER virtual table raised:", type(e).__name__, str(e))
try:
conn.execute("CREATE INDEX idx_mem_added ON mem(added)")
print("CREATE INDEX virtual table succeeded unexpectedly")
except Exception as e:
print("CREATE INDEX virtual table raised:", type(e).__name__, str(e))
try:
conn.execute("ALTER TABLE memory_embeddings ADD COLUMN partition_id TEXT NOT NULL DEFAULT 'default'")
print("ALTER TABLE nonexistent table succeeded unexpectedly")
except Exception as e:
print("ALTER missing table raised:", type(e).__name__, str(e))
conn.execute("CREATE TABLE regular_mem (memory_id TEXT PRIMARY KEY, embedding BLOB)")
row = conn.execute("SELECT type FROM sqlite_master WHERE name = 'regular_mem'").fetchone()
print("regular table sqlite_master.type:", row[0] if row else None)
print("regular table sqlite_master.sql:", conn.execute("SELECT sql FROM sqlite_master WHERE name = 'regular_mem'").fetchone()[0])
PYRepository: afx-team/hebb-mind
Length of output: 381
Use table metadata or a robust vec0 match, not exact SQL substring detection.
sqlite_master.sql is preserved schema SQL and the helper still executes DDL on the virtual table when the substring misses. Rely on SQLite metadata such as sqlite_master.type = 'table' or a normalized case-insensitive USING vec0... pattern before running ALTER TABLE / CREATE INDEX.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/hebb/storage/migrations.py` around lines 235 - 241, The schema check
before the migration DDL must reliably identify vec0 virtual tables. Update the
sqlite_master query and detection logic in the migration flow to use table
metadata or normalized, case-insensitive matching for a USING vec0 clause,
including spacing or trailing options, and return before ALTER TABLE or CREATE
INDEX when the table is virtual.
|
I believe the root cause of the e2e failure is that upstream mcp 2.0.0 (released 2026-07-28) removed the v1 import paths (mcp.server.fastmcp,mcp.ClientSession, etc.), and since this project specifies mcp>=1.0.0 with no upper bound, CI resolves to 2.0 and fails with ImportError (server.py:21,tests/e2e/test_mcp_server.py:29-30). For the change below, I therefore chose to cap the mcp version in pyproject.toml. |
The macOS-only test failures (test_dim_mismatch_raises_and_writes_nothing, test_dimension_mismatch_raises_without_opt_in, test_dimension_mismatch_drops_with_opt_in) come from a pre-existing latent gap: the embedding-dim guard is skipped on the BLOB fallback path used when sqlite-vec can't load (macOS CI), because the fallback table has no inline width to read from sqlite_master. Adding partition_id to the fallback schema made the probe INSERT succeed on the 3-column fallback and return early, bypassing the mismatch check entirely. Persist the configured embedding_dim in a new schema_meta table so both the write-path guard (SQLiteMemoryStore._get_vec_dim) and the migration mismatch check (_ensure_vec_table) can read the existing width even on the fallback path. _ensure_vec_table now also creates the BLOB fallback itself when vec0 is unavailable, so direct calls (the audit-write-api tests) work without initialize_schema's except handler. Verified on both vec0-enabled and vec0-disabled (macOS CI sim) paths. Also pin mcp>=1.0.0,<2: upstream mcp 2.0.0 (2026-07-28) removed mcp.server.fastmcp and the v1 client API (mcp.ClientSession / mcp.client.stdio.stdio_client) that hebb.mcp.server and its e2e test import, breaking e2e on all platforms. Pinning unblocks CI on main and every PR with no code change; migrating to standalone fastmcp / mcp 2.0 MCPServer is a separate follow-up. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/hebb/storage/migrations.py`:
- Around line 239-290: Update initialize_schema’s vec0 fallback and
probe-exception handling to avoid unconditional drops: preserve populated
memory_embeddings unless a genuine dimension mismatch is explicitly permitted by
_ALLOW_EMBED_DROP_ENV, while still allowing first-time creation and the
empty-table partition_id migration. Check actual table existence via
sqlite_master in addition to existing_dim so pre-schema_meta fallback tables are
protected, and ensure probe failures only recreate when the schema is genuinely
incompatible. Add a regression test that initializes twice on a populated
vec0-unavailable fallback table and verifies the data remains.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: bd91618e-5bac-4f89-b5d2-278057b0e773
📒 Files selected for processing (3)
pyproject.tomlsrc/hebb/storage/migrations.pysrc/hebb/storage/sqlite_store.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/hebb/storage/sqlite_store.py
| # --- vec0 path: try the virtual table. --- | ||
| try: | ||
| if true_dim_mismatch: | ||
| # Empty-table dim change or opt-in destructive rebuild at ``dim``. | ||
| await db.execute("DROP TABLE IF EXISTS memory_embeddings") | ||
| await db.execute( | ||
| f"CREATE VIRTUAL TABLE IF NOT EXISTS memory_embeddings USING vec0(" | ||
| f"memory_id TEXT PRIMARY KEY, partition_id TEXT, embedding float[{dim}])" | ||
| ) | ||
| except Exception: | ||
| # vec0 unavailable (extension not loaded) -> BLOB fallback below. If we | ||
| # already dropped for a mismatch, the fallback recreates the table empty. | ||
| pass | ||
| else: | ||
| # Probe the partition_id column by inserting a dim-width vector. A | ||
| # pre-refactor table (no partition_id) accepts CREATE IF NOT EXISTS | ||
| # silently but rejects this insert; we then recreate cleanly. On a | ||
| # matching-dim table with the column, the probe round-trips and we're | ||
| # done. (The probe double-checks the dim on the vec0 path only — the | ||
| # BLOB fallback accepts any width, so dim enforcement there relies on | ||
| # the ``schema_meta`` mismatch check above.) | ||
| probe = np.zeros(dim, dtype=np.float32).tobytes() | ||
| try: | ||
| await db.execute( | ||
| "INSERT INTO memory_embeddings(memory_id, partition_id, embedding) " | ||
| "VALUES ('__schema_probe__', '__probe__', ?)", | ||
| (probe,), | ||
| ) | ||
| await db.execute("DELETE FROM memory_embeddings WHERE memory_id = '__schema_probe__'") | ||
| except Exception: | ||
| await db.execute("DROP TABLE IF EXISTS memory_embeddings") | ||
| await db.execute(_VEC_CREATE_SQL.format(dim=dim)) | ||
| await _write_meta_dim(db, dim) | ||
| return | ||
|
|
||
| # --- Fallback path: vec0 unavailable -> regular BLOB table at ``dim``. --- | ||
| if true_dim_mismatch and row_count > 0: | ||
| logger.warning( | ||
| "Dropping %d embeddings: dim %d -> %d (%s=1 opt-in)", | ||
| row_count, | ||
| declared, | ||
| existing_dim, | ||
| dim, | ||
| _ALLOW_EMBED_DROP_ENV, | ||
| ) | ||
| else: | ||
| elif true_dim_mismatch: | ||
| logger.warning( | ||
| "Recreating vec0 table (dim=%d, partition_id column) — any existing embeddings will be lost", | ||
| dim, | ||
| ) | ||
| await db.execute("DROP TABLE IF EXISTS memory_embeddings") | ||
| await db.execute(_VEC_CREATE_SQL.format(dim=dim)) | ||
| await db.execute(_FALLBACK_CREATE_SQL) | ||
| await _write_meta_dim(db, dim) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
Fix the destructive drop-and-recreate in the fallback and probe paths.
Two gaps let this function silently discard populated memory_embeddings data outside the guarded mismatch case.
- Lines 274-290: When vec0 stays unavailable (e.g., a macOS python.org build without loadable-extension support), the
CREATE VIRTUAL TABLEattempt always raises, and control falls through unconditionally toDROP TABLE IF EXISTS memory_embeddingsfollowed by_FALLBACK_CREATE_SQL. There is no check for "the fallback table already exists at the correct dimension." This means every call toinitialize_schemaon a vec0-unavailable install — i.e., every process restart — wipes and recreates the table empty, even whentrue_dim_mismatchisFalseandrow_count > 0. No warning is logged and the_ALLOW_EMBED_DROP_ENVopt-in is never consulted for this path, unlike the guarded case above. - Lines 260-272: the probe's
except Exceptionbranch triggers the same unconditionalDROP TABLE+ recreate for any failure during the probe insert/delete (not just a genuine missing-partition_idschema), bypassing therow_count/opt-in guard entirely for populated legacy vec0 tables.
Both paths contradict the function's own documented invariant that "First-time creation and the additive partition_id migration (empty or correctly-dimensioned tables) are unaffected."
🛡️ Proposed fix
probe = np.zeros(dim, dtype=np.float32).tobytes()
try:
await db.execute(
"INSERT INTO memory_embeddings(memory_id, partition_id, embedding) "
"VALUES ('__schema_probe__', '__probe__', ?)",
(probe,),
)
await db.execute("DELETE FROM memory_embeddings WHERE memory_id = '__schema_probe__'")
except Exception:
+ if row_count > 0 and os.getenv(_ALLOW_EMBED_DROP_ENV) != "1":
+ raise EmbeddingDimensionMismatchError(
+ f"memory_embeddings has {row_count} vectors but is missing the "
+ f"partition_id column. Dropping the table would lose every "
+ f"embedding. Run `hebb memory reembed` or set "
+ f"{_ALLOW_EMBED_DROP_ENV}=1 to drop and rebuild empty."
+ )
+ logger.warning(
+ "Recreating vec0 table to add partition_id column; %d existing embeddings will be lost",
+ row_count,
+ )
await db.execute("DROP TABLE IF EXISTS memory_embeddings")
await db.execute(_VEC_CREATE_SQL.format(dim=dim))
await _write_meta_dim(db, dim)
return
# --- Fallback path: vec0 unavailable -> regular BLOB table at ``dim``. ---
+ if existing_dim is not None and not true_dim_mismatch:
+ # Fallback table already exists at the configured width. Recreating
+ # unconditionally here would wipe every embedding on every restart in
+ # vec0-unavailable environments.
+ await _write_meta_dim(db, dim)
+ return
if true_dim_mismatch and row_count > 0:Note: existing_dim is None on a pre-existing fallback table created before schema_meta tracking was added, so this specific guard alone does not protect the first upgrade run for older installs. Verify table existence directly (e.g., via sqlite_master) in addition to existing_dim, and add a regression test that calls initialize_schema twice on a populated vec0-unavailable table to confirm no data loss.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/hebb/storage/migrations.py` around lines 239 - 290, Update
initialize_schema’s vec0 fallback and probe-exception handling to avoid
unconditional drops: preserve populated memory_embeddings unless a genuine
dimension mismatch is explicitly permitted by _ALLOW_EMBED_DROP_ENV, while still
allowing first-time creation and the empty-table partition_id migration. Check
actual table existence via sqlite_master in addition to existing_dim so
pre-schema_meta fallback tables are protected, and ensure probe failures only
recreate when the schema is genuinely incompatible. Add a regression test that
initializes twice on a populated vec0-unavailable fallback table and verifies
the data remains.
…t guard Addresses CodeRabbit review on the dim-guard fix: the vec0-unavailable fallback path unconditionally executed DROP TABLE before _FALLBACK_CREATE_SQL. On builds where CREATE VIRTUAL TABLE IF NOT EXISTS short-circuits (the BLOB table already exists), the probe returns early so the DROP is never reached; but on builds that resolve the vec0 module before the existence check, the CREATE raises and control falls through to the fallback — where an unconditional DROP would wipe a same-dim populated table. Gate the DROP on true_dim_mismatch so a same-dim (or first-time) table is never dropped on this path; _FALLBACK_CREATE_SQL is IF NOT EXISTS, so skipping the DROP is a no-op when the table exists and a clean create when it doesn't. Add test_same_dim_reinit_preserves_embeddings as an invariant guard: a populated table must survive a same-dim re-init on both the vec0 and fallback paths. Verified on vec0-enabled and vec0-disabled (macOS CI sim). Finding 2 from the review (probe-except DROP for a populated legacy vec0 table missing the partition_id column) is pre-existing — the original code took the same path, and vec0 virtual tables cannot be ALTERed to add the column. Changing it would alter the upgrade behavior for legacy installs, so it is left for a separate follow-up rather than bundled into this fix. Co-Authored-By: Claude <noreply@anthropic.com>
Route every SQL write and its KnowledgeGraph mutation through a single shared asyncio.Lock so they form one crash-consistent critical section, closing the split-lock gap where interrupted consolidation left graph orphans.
Fixes #36
Summary
Motivation
Checklist
pytest tests/ -v)repo_pages/, docstrings, or README) if user-facingCHANGELOG.mdupdated under## [Unreleased]ruff check src/passesmypy src/hebb/passes (project usesstrict = true)Notes for reviewers
Summary by CodeRabbit