fix(mcp): use lightweight graph hydration lookup (#806)

Signed-off-by: phernandez <paul@basicmachines.co>

Author: phernandez

src/basic_memory/api/v2/utils.py

@@ -18,7 +18,7 @@
 
 
 class EntityBatchLookup(Protocol):
-    async def find_by_ids(self, ids: List[int]) -> Sequence[Any]: ...
+    async def find_by_ids_for_hydration(self, ids: List[int]) -> Sequence[Any]: ...
 
 
 class EntityServiceBatchLookup(Protocol):
@@ -76,7 +76,7 @@ async def to_graph_context(
                     if item.to_id:
                         entity_ids_needed.add(item.to_id)
 
-        # Batch fetch all entities at once - get both title and external_id
+        # Batch fetch just the entity fields needed to shape the response.
         entity_title_lookup: dict[int, str] = {}
         entity_external_id_lookup: dict[int, str] = {}
         if entity_ids_needed:
@@ -87,7 +87,9 @@ async def to_graph_context(
                 phase="lookup_entities",
                 result_count=len(entity_ids_needed),
             ):
-                entities = await entity_repository.find_by_ids(list(entity_ids_needed))
+                entities = await entity_repository.find_by_ids_for_hydration(
+                    list(entity_ids_needed)
+                )
             for e in entities:
                 entity_title_lookup[e.id] = e.title
                 entity_external_id_lookup[e.id] = e.external_id

src/basic_memory/repository/entity_repository.py

@@ -8,7 +8,7 @@
 from sqlalchemy import select, func
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
-from sqlalchemy.orm import selectinload
+from sqlalchemy.orm import load_only, selectinload
 from sqlalchemy.orm.interfaces import LoaderOption
 from sqlalchemy.engine import Row
 
@@ -178,6 +178,24 @@ async def get_all_permalinks(self) -> List[str]:
         result = await self.execute_query(query, use_query_options=False)
         return list(result.scalars().all())
 
+    async def find_by_ids_for_hydration(self, ids: List[int]) -> Sequence[Entity]:
+        """Fetch minimal entity fields needed for context hydration.
+
+        Context hydration only needs an entity's primary key, title, and external
+        UUID. Keeping this separate from find_by_ids avoids the relationship eager
+        loads that are useful for full entity reads but expensive for response shaping.
+        """
+        if not ids:
+            return []
+
+        query = (
+            self.select()
+            .where(Entity.id.in_(ids))
+            .options(load_only(Entity.id, Entity.title, Entity.external_id))
+        )
+        result = await self.execute_query(query, use_query_options=False)
+        return list(result.scalars().all())
+
     async def get_permalink_to_file_path_map(self) -> dict[str, str]:
         """Get a mapping of permalink -> file_path for all entities.
 

tests/api/v2/test_memory_hydration.py

@@ -55,6 +55,25 @@ async def find_by_ids(self, ids: list[int]):
         self.calls.append(ids)
         return [self.entities_by_id[i] for i in ids if i in self.entities_by_id]
 
+    async def find_by_ids_for_hydration(self, ids: list[int]):
+        self.calls.append(ids)
+        return [self.entities_by_id[i] for i in ids if i in self.entities_by_id]
+
+
+class LightweightOnlyEntityRepository:
+    """Raises if graph hydration uses the eager-loading repository method."""
+
+    def __init__(self, entities_by_id: dict[int, SimpleNamespace]):
+        self.entities_by_id = entities_by_id
+        self.hydration_calls: list[list[int]] = []
+
+    async def find_by_ids(self, ids: list[int]):
+        raise AssertionError("graph hydration must use the lightweight hydration lookup")
+
+    async def find_by_ids_for_hydration(self, ids: list[int]):
+        self.hydration_calls.append(ids)
+        return [self.entities_by_id[i] for i in ids if i in self.entities_by_id]
+
 
 # --- Single batch fetch (N+1 elimination) ---
 
@@ -198,3 +217,63 @@ async def test_to_graph_context_empty_results_skip_entity_lookup():
 
     assert repo.calls == []
     assert list(graph.results) == []
+
+
+@pytest.mark.asyncio
+async def test_to_graph_context_uses_lightweight_hydration_lookup():
+    """Hydration should not load observations/relations when only entity fields are needed."""
+    repo = LightweightOnlyEntityRepository(
+        {
+            1: _make_entity(1, "Root", "ext-root"),
+            2: _make_entity(2, "Child", "ext-child"),
+        }
+    )
+    now = datetime.now(timezone.utc)
+
+    context = ServiceContextResult(
+        results=[
+            ContextResultItem(
+                primary_result=_make_row(
+                    type="entity",
+                    id=1,
+                    root_id=1,
+                    title="Root",
+                    permalink="notes/root",
+                    file_path="notes/root.md",
+                    created_at=now,
+                ),
+                observations=[],
+                related_results=[
+                    _make_row(
+                        type="relation",
+                        id=20,
+                        root_id=1,
+                        title="links_to: Child",
+                        permalink="notes/root",
+                        file_path="notes/root.md",
+                        relation_type="links_to",
+                        from_id=1,
+                        to_id=2,
+                        depth=1,
+                        created_at=now,
+                    )
+                ],
+            )
+        ],
+        metadata=ContextMetadata(
+            types=[SearchItemType.ENTITY, SearchItemType.RELATION],
+            depth=1,
+            primary_count=1,
+            related_count=1,
+            total_relations=1,
+        ),
+    )
+
+    graph = await to_graph_context(context, entity_repository=repo)
+
+    assert len(repo.hydration_calls) == 1
+    assert set(repo.hydration_calls[0]) == {1, 2}
+    relation = graph.results[0].related_results[0]
+    assert isinstance(relation, RelationSummary)
+    assert relation.from_entity_external_id == "ext-root"
+    assert relation.to_entity_external_id == "ext-child"

tests/repository/test_entity_repository.py

@@ -1050,6 +1050,25 @@ async def test_get_all_permalinks(entity_repository: EntityRepository, session_m
         assert isinstance(permalink, str)
 
 
+@pytest.mark.asyncio
+async def test_find_by_ids_for_hydration_skips_eager_load_options(
+    entity_repository: EntityRepository, sample_entity: Entity, monkeypatch: pytest.MonkeyPatch
+):
+    """Context hydration should bypass relationship loader options."""
+
+    def fail_get_load_options():
+        raise AssertionError("hydration lookup must not eager load entity relationships")
+
+    monkeypatch.setattr(entity_repository, "get_load_options", fail_get_load_options)
+
+    found = await entity_repository.find_by_ids_for_hydration([sample_entity.id])
+
+    assert len(found) == 1
+    assert found[0].id == sample_entity.id
+    assert found[0].title == sample_entity.title
+    assert found[0].external_id == sample_entity.external_id
+
+
 @pytest.mark.asyncio
 async def test_get_permalink_to_file_path_map(entity_repository: EntityRepository, session_maker):
     """Test getting permalink -> file_path mapping for bulk operations."""