Skip to content

refactor(memory): introduce a pluggable MemoryBackend trait#1

Draft
slvnlrt wants to merge 1 commit into
pr/ingestion-fixesfrom
pr/memory-abstraction
Draft

refactor(memory): introduce a pluggable MemoryBackend trait#1
slvnlrt wants to merge 1 commit into
pr/ingestion-fixesfrom
pr/memory-abstraction

Conversation

@slvnlrt

@slvnlrt slvnlrt commented Jun 24, 2026

Copy link
Copy Markdown
Owner

Summary

Decouple the long-term memory layer behind a single MemoryBackend trait so the storage backend is pluggable and the rest of the app is storage-agnostic. The existing SQLite + LanceDB stack becomes SqliteBackend (the production implementation); hybrid search and maintenance (decay / prune / merge) now operate over dyn MemoryBackend instead of the concrete store.

Behaviour-preserving: the default build is unchanged, no new dependencies are added, and all existing tests pass.

What changes

  • New MemoryBackend trait (src/memory/backend.rs) — a complete memory interface: structured storage (save / load / update / forget / record_access), vector + full-text search, the association graph (create / query / traverse), and maintenance ops (prune / merge).
  • SqliteBackend implements the trait over the current MemoryStore (SQLite) + EmbeddingTable (LanceDB) — the production default.
  • MemorySearch (hybrid search) and the maintenance routines are now generic over Arc<dyn MemoryBackend>. Callers (memory tools, cortex, API) go through the trait instead of touching SQLite/LanceDB directly.

Why

The memory layer was hard-wired to SQLite + LanceDB across search, maintenance, the tools and the cortex. Putting it behind a trait makes the storage backend swappable and the memory subsystem independently testable, without committing the project to any particular store.

For context: we maintain a SurrealDB-based unified memory backend (structured + vector + graph in a single embedded store) implemented against this trait in a fork. It could be submitted as a dedicated follow-up PR if there's interest — but this PR is intentionally just the seam: no alternative backend is added here.

Notes

  • Behaviour-preserving refactor — no schema change, no new dependency, default build identical.
  • cargo test --lib — all green.

Draft — stacked on spacedriveapp#604. Base is the pr/ingestion-fixes branch (the ingestion-reliability fixes). This PR shows only the abstraction commit on top. Once spacedriveapp#604 lands upstream, this will be rebased onto main and retargeted to the upstream repo.

Summary by CodeRabbit

  • Bug Fixes

    • Improved memory save, recall, delete, and maintenance reliability across the app.
    • Search and graph-related memory features now behave more consistently, especially after updates, merges, and deletions.
    • Batch memory lookups now return more complete results and handle larger sets more efficiently.
  • New Features

    • Added broader support for memory graph and association lookups, improving memory-related responses and navigation.

@coderabbitai

coderabbitai Bot commented Jun 24, 2026

Copy link
Copy Markdown

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: c8b7288e-bb88-456f-87ab-ff9c49b3a872

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • ✅ Review completed - (🔄 Check again to review again)
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch pr/memory-abstraction

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 6

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
src/memory/maintenance.rs (1)

93-105: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Avoid decaying only the first 1,000 memories per type.

This says it decays all non-identity memories, but get_by_type(..., 1000) returns only a capped page; larger stores can permanently skip lower-ranked memories. Extend the backend with pagination/streaming or a backend-side decay operation so maintenance covers the full set.

🤖 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/memory/maintenance.rs` around lines 93 - 105, The maintenance loop in
`maintain_memories` only processes a capped page because
`backend.get_by_type(mem_type, 1000)` limits each `MemoryType`, so larger stores
will miss later memories. Update the decay path to paginate or stream through
all results for each non-identity type, or add a backend-side bulk decay
operation, and keep the existing `maintenance_cancelable_op` flow while
replacing the fixed-size fetch in `get_by_type` usage.
🤖 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/api/agents.rs`:
- Around line 834-835: The backend store created in the agent setup is losing
the agent identifier, so runtime-created agents can appear as unknown in
backend-scoped context. Update the `sqlite_backend_arc` setup in
`src/api/agents.rs` to preserve the ID by chaining `with_agent_id` on the
`MemoryStore` created by `MemoryStore::new` before passing it into
`crate::memory::sqlite_backend_arc`. Keep the fix localized to the agent
creation path so `SqliteBackend::agent_id()` forwards the correct value.

In `@src/memory/backend.rs`:
- Around line 164-179: The save flow in MemoryBackend::save writes the
structured memory before storing the embedding, so a failure in embeddings.store
can leave a partially persisted record. Update save(...) to own the compensation
logic by rolling back the earlier self.store.save(memory) when the embedding
write fails, using the existing MemoryBackend and embeddings symbols to locate
the change. Keep the FTS ensure logic after a successful embedding write, and
make sure any rollback path preserves the current error return semantics.
- Around line 347-358: The helper uses two different sources for the agent
identity: `sqlite_backend_arc` receives `agent_id`, but
`SqliteBackend::agent_id()` later reads `store.agent_id()`, which can leave
backend labels as "unknown". Update `sqlite_backend_arc` and the
`SqliteBackend::new`/`agent_id()` path so the store and backend share one source
of truth, either by deriving the backend agent ID from the store or by ensuring
the store is initialized with the passed `agent_id`. Make the change in the
`sqlite_backend_arc`/`SqliteBackend` flow so downstream code calling
`backend.agent_id()` always sees the real agent ID.

In `@src/memory/search.rs`:
- Around line 286-306: The association traversal in `search.rs` is relying on
the incidental order returned by `get_associations_for`, which makes downstream
graph/RRF ranking unstable. Make the ordering explicit before the `by_node`
grouping in the `frontier` processing path, either by sorting `all_edges`
deterministically or by tightening the backend/query contract in
`get_associations_for` so it always returns a defined order. Apply the same fix
wherever this traversal logic is duplicated, including the referenced mirrored
block, so `search` and its ranking stay stable across backends.

In `@src/memory/store.rs`:
- Around line 497-548: The batched IN queries in load_many and the association
lookup path bind the entire ID slice at once, which can exceed SQLite’s
parameter limit during large BFS traversals. Update these methods to split the
incoming ids into bounded chunks, run the existing query logic per chunk, and
merge the results before returning so large graph walks remain reliable without
hitting too many SQL variables.

In `@src/tools/memory_save.rs`:
- Around line 256-258: The FTS retry logic is missing from the `set_embedding`
save flow, so memories written through `MemorySave::save` with `embedding: None`
can stay out of FTS after a startup indexing failure. Update the backend-level
embedding persistence path (the `SqliteBackend::set_embedding`/related save
flow) to reuse the same retry behavior currently implemented in
`SqliteBackend::save(..., Some(_))`, and ensure the later embedding update also
triggers the FTS-index retry when needed.

---

Outside diff comments:
In `@src/memory/maintenance.rs`:
- Around line 93-105: The maintenance loop in `maintain_memories` only processes
a capped page because `backend.get_by_type(mem_type, 1000)` limits each
`MemoryType`, so larger stores will miss later memories. Update the decay path
to paginate or stream through all results for each non-identity type, or add a
backend-side bulk decay operation, and keep the existing
`maintenance_cancelable_op` flow while replacing the fixed-size fetch in
`get_by_type` usage.
🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 95095290-98fe-4665-aca1-6660730707e0

📥 Commits

Reviewing files that changed from the base of the PR and between e4a2eec and aa675e9.

📒 Files selected for processing (15)
  • src/agent/cortex.rs
  • src/agent/maintenance.rs
  • src/api/agents.rs
  • src/api/memories.rs
  • src/conversation/context.rs
  • src/main.rs
  • src/memory.rs
  • src/memory/backend.rs
  • src/memory/lance.rs
  • src/memory/maintenance.rs
  • src/memory/search.rs
  • src/memory/store.rs
  • src/tools/memory_delete.rs
  • src/tools/memory_recall.rs
  • src/tools/memory_save.rs

Comment thread src/api/agents.rs Outdated
Comment thread src/memory/backend.rs
Comment thread src/memory/backend.rs
Comment thread src/memory/search.rs Outdated
Comment thread src/memory/store.rs Outdated
Comment thread src/tools/memory_save.rs
Put the SQLite + LanceDB memory stack behind a single MemoryBackend trait so the memory layer is pluggable and testable. SqliteBackend is the production implementation; hybrid search and maintenance (decay/prune/merge) now operate over dyn MemoryBackend instead of the concrete store. Behaviour-preserving: the default build is unchanged and all existing tests pass.
@slvnlrt slvnlrt force-pushed the pr/memory-abstraction branch from aa675e9 to c757231 Compare June 24, 2026 22:02
@slvnlrt

slvnlrt commented Jun 24, 2026

Copy link
Copy Markdown
Owner Author

Thanks for the reviews. Addressed in the latest push (folded into the single refactor commit):

  • Agent ID on the backend store (agents.rs) — the API agent-creation path now builds MemoryStore::with_agent_id, so backend.agent_id() returns the real id instead of unknown. This also resolves the related sqlite_backend_arc agent-id note: the store is the single source of truth and both call sites now populate it.
  • Deterministic association ordering (search.rs) — traverse_graph now sorts incident edges by descending weight (tie-broken by the stable association id) before grouping, so traversal and ranking no longer depend on the backend's incidental row order or the chunking below. The golden test was updated to the resulting deterministic order.
  • Bounded IN (...) queries (store.rs) — get_associations_for, get_associations_between and load_many now chunk their id lists to stay under SQLite's bind-parameter limit. get_associations_between queries source_id IN (chunk) and filters target_id against the full set in Rust, so cross-chunk pairs are not missed.
  • FTS index on the set_embedding path (backend.rs) — set_embedding now calls ensure_fts_index, so the save(None) + set_embedding path used by the memory_save tool still triggers index creation.

On the save(..., Some(...)) half-write note: leaving the behaviour as-is intentionally. This is a behaviour-preserving refactor, and the structured store and the vector store have no shared transaction (cross-store writes are best-effort by design); the sole production caller (memory_save) already compensates for a failed embedding write. I added a doc comment on the trait method stating that save is not atomic across the two stores and that callers must compensate.

slvnlrt added a commit that referenced this pull request Jun 25, 2026
Cross-backend fixes raised by the upstream PR reviews, applied to the integration branch: ingest delete-before-clear-progress, merged_memory_content size cap on the empty-winner path, ingest delete path-component guard, agent_id on the SQLite backend store, deterministic (weight-descending) edge ordering in graph traversal, chunked IN-list queries to stay under SQLite's bind limit, and FTS-index ensure on set_embedding. Feature-off: 891 tests green. Feature-on: clippy clean.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant