feat(mcp): add rebuild_embeddings tool (issue #39)

app/protocols/memory_protocol.py

@@ -183,6 +183,65 @@ async def get_memories_for_reembedding(self, limit: int, offset: int) -> list[Me
         """Fetch memories in batches for re-embedding (all users, ordered by id)"""
         ...
 
+    async def count_memories_for_targeted_rebuild(
+            self,
+            user_id: UUID,
+            memory_ids: list[int] | None = None,
+            project_id: int | None = None,
+    ) -> int:
+        """Count memories eligible for a *user-scoped, targeted* embedding rebuild.
+
+        Used by `rebuild_embeddings` to avoid the destructive global
+        `re_embed_all` path. Filters are user-scoped; user_id is mandatory and
+        scopes ownership. Matching memories are rebuilt unconditionally.
+
+        Args:
+            user_id: Authenticated user id; ownership filter, never optional.
+            memory_ids: Restrict to explicit ids (already user-scoped server-side).
+            project_id: Restrict to a single project owned by `user_id`.
+
+        Returns:
+            Number of non-obsolete memories matching the filters and owned by
+            `user_id`.
+        """
+        ...
+
+    async def get_memories_for_targeted_rebuild(
+            self,
+            user_id: UUID,
+            limit: int,
+            after_id: int | None = None,
+            memory_ids: list[int] | None = None,
+            project_id: int | None = None,
+    ) -> list[Memory]:
+        """Page through memories selected by `count_memories_for_targeted_rebuild`.
+
+        Same semantics and same filters as the count variant. Uses keyset
+        pagination by memory id so callers can process deterministically.
+        """
+        ...
+
+    async def upsert_targeted_embeddings(
+            self,
+            user_id: UUID,
+            updates: list[tuple[int, list[float]]],
+    ) -> list[int]:
+        """Idempotently write a batch of new embeddings for `user_id`.
+
+        Unlike `bulk_update_embeddings`, this method:
+        - is user-scoped (rejects ids not owned by `user_id`),
+        - never resets vector storage,
+        - is idempotent for repeated calls on the same memory id.
+
+        Args:
+            user_id: Authenticated user id; ownership filter, never optional.
+            updates: list of (memory_id, embedding) tuples.
+
+        Returns:
+            list of memory ids whose embedding was actually written.
+        """
+        ...
+
     async def reset_embedding_storage(self) -> None:
         """Prepare vector storage for new dimensions.
         SQLite: DROP + CREATE vec_memories with new dimensions

app/repositories/postgres/memory_repository.py

@@ -964,6 +964,115 @@ async def bulk_update_embeddings(self, updates: list[tuple[int, list[float]]]) -
                     .values(embedding=embedding),
                 )
 
+    # ---- Targeted (user-scoped) rebuild support, never resets vec storage ----
+
+    def _build_targeted_rebuild_filter(
+        self,
+        user_id: UUID,
+        memory_ids: list[int] | None,
+        project_id: int | None,
+    ):
+        conditions = [
+            MemoryTable.user_id == user_id,
+            MemoryTable.is_obsolete.is_(False),
+        ]
+        if memory_ids:
+            conditions.append(MemoryTable.id.in_(memory_ids))
+        stmt = select(MemoryTable).where(*conditions)
+        if project_id is not None:
+            stmt = stmt.join(MemoryTable.projects).where(
+                ProjectsTable.id == project_id,
+                ProjectsTable.user_id == user_id,
+            )
+        return stmt
+
+    async def count_memories_for_targeted_rebuild(
+        self,
+        user_id: UUID,
+        memory_ids: list[int] | None = None,
+        project_id: int | None = None,
+    ) -> int:
+        from sqlalchemy import func
+        async with self.db_adapter.system_session() as session:
+            base = self._build_targeted_rebuild_filter(
+                user_id=user_id,
+                memory_ids=memory_ids,
+                project_id=project_id,
+            )
+            count_stmt = select(func.count()).select_from(base.subquery())
+            return (await session.execute(count_stmt)).scalar() or 0
+
+    async def get_memories_for_targeted_rebuild(
+        self,
+        user_id: UUID,
+        limit: int,
+        after_id: int | None = None,
+        memory_ids: list[int] | None = None,
+        project_id: int | None = None,
+    ) -> list[Memory]:
+        async with self.db_adapter.system_session() as session:
+            stmt = self._build_targeted_rebuild_filter(
+                user_id=user_id,
+                memory_ids=memory_ids,
+                project_id=project_id,
+            )
+            if after_id is not None:
+                stmt = stmt.where(MemoryTable.id > after_id)
+            stmt = (
+                stmt
+                .options(
+                    selectinload(MemoryTable.projects),
+                    selectinload(MemoryTable.linked_memories),
+                    selectinload(MemoryTable.linking_memories),
+                    selectinload(MemoryTable.code_artifacts),
+                    selectinload(MemoryTable.documents),
+                    selectinload(MemoryTable.files),
+                )
+                .order_by(MemoryTable.id.asc())
+                .limit(limit)
+            )
+            result = await session.execute(stmt)
+            memories_orm = result.scalars().unique().all()
+            return [Memory.model_validate(m) for m in memories_orm]
+
+    async def upsert_targeted_embeddings(
+        self,
+        user_id: UUID,
+        updates: list[tuple[int, list[float]]],
+    ) -> list[int]:
+        """Idempotently set embeddings for memories owned by `user_id`.
+
+        Postgres stores embeddings inline on `memories.embedding`, so an UPDATE
+        is naturally idempotent. We still verify ownership server-side instead
+        of trusting the input list.
+        """
+        if not updates:
+            return []
+        ids = [mid for mid, _ in updates]
+        async with self.db_adapter.system_session() as session:
+            owned_rows = await session.execute(
+                select(MemoryTable.id).where(
+                    MemoryTable.user_id == user_id,
+                    MemoryTable.is_obsolete.is_(False),
+                    MemoryTable.id.in_(ids),
+                ),
+            )
+            owned_ids = {row[0] for row in owned_rows}
+            written: list[int] = []
+            for memory_id, embedding in updates:
+                if memory_id not in owned_ids:
+                    continue
+                await session.execute(
+                    update(MemoryTable)
+                    .where(
+                        MemoryTable.id == memory_id,
+                        MemoryTable.user_id == user_id,
+                    )
+                    .values(embedding=embedding),
+                )
+                written.append(memory_id)
+            return written
+
     async def validate_embedding_count(self) -> bool:
         """Check embedding count matches non-obsolete memory count"""
         from sqlalchemy import func

app/repositories/sqlite/memory_repository.py

@@ -983,6 +983,127 @@ async def bulk_update_embeddings(self, updates: list[tuple[int, list[float]]]) -
                     {"memory_id": str(memory_id), "embedding": embedding_bytes},
                 )
 
+    # ---- Targeted (user-scoped) rebuild support, never resets vec storage ----
+
+    def _build_targeted_rebuild_filter(
+        self,
+        user_id: UUID,
+        memory_ids: list[int] | None,
+        project_id: int | None,
+    ):
+        """Shared SQLAlchemy filter for targeted rebuild count/select queries."""
+        conditions = [
+            MemoryTable.user_id == str(user_id),
+            MemoryTable.is_obsolete.is_(False),
+        ]
+        if memory_ids:
+            conditions.append(MemoryTable.id.in_(memory_ids))
+        stmt = select(MemoryTable).where(*conditions)
+        if project_id is not None:
+            stmt = stmt.join(MemoryTable.projects).where(
+                ProjectsTable.id == project_id,
+                ProjectsTable.user_id == str(user_id),
+            )
+        return stmt
+
+    async def count_memories_for_targeted_rebuild(
+        self,
+        user_id: UUID,
+        memory_ids: list[int] | None = None,
+        project_id: int | None = None,
+    ) -> int:
+        from sqlalchemy import func
+        async with self.db_adapter.system_session() as session:
+            base = self._build_targeted_rebuild_filter(
+                user_id=user_id,
+                memory_ids=memory_ids,
+                project_id=project_id,
+            )
+            count_stmt = select(func.count()).select_from(base.subquery())
+            return (await session.execute(count_stmt)).scalar() or 0
+
+    async def get_memories_for_targeted_rebuild(
+        self,
+        user_id: UUID,
+        limit: int,
+        after_id: int | None = None,
+        memory_ids: list[int] | None = None,
+        project_id: int | None = None,
+    ) -> list[Memory]:
+        async with self.db_adapter.system_session() as session:
+            stmt = self._build_targeted_rebuild_filter(
+                user_id=user_id,
+                memory_ids=memory_ids,
+                project_id=project_id,
+            )
+            if after_id is not None:
+                stmt = stmt.where(MemoryTable.id > after_id)
+            stmt = (
+                stmt
+                .options(
+                    selectinload(MemoryTable.projects),
+                    selectinload(MemoryTable.linked_memories),
+                    selectinload(MemoryTable.linking_memories),
+                    selectinload(MemoryTable.code_artifacts),
+                    selectinload(MemoryTable.documents),
+                    selectinload(MemoryTable.files),
+                )
+                .order_by(MemoryTable.id.asc())
+                .limit(limit)
+            )
+            result = await session.execute(stmt)
+            memories_orm = result.scalars().unique().all()
+            return [Memory.model_validate(m) for m in memories_orm]
+
+    async def upsert_targeted_embeddings(
+        self,
+        user_id: UUID,
+        updates: list[tuple[int, list[float]]],
+    ) -> list[int]:
+        """Idempotently write embeddings for memories owned by `user_id`.
+
+        Implementation note: SQLite-vec's `vec_memories` virtual table is
+        keyed by `memory_id` and rejects duplicate inserts. We delete the
+        existing row (if any) before inserting the new embedding. We also
+        verify ownership of every memory id against the canonical `memories`
+        table so a malicious caller cannot poison another user's vector.
+        """
+        if not updates:
+            return []
+        ids = [mid for mid, _ in updates]
+        async with self.db_adapter.system_session() as session:
+            owned_rows = await session.execute(
+                select(MemoryTable.id).where(
+                    MemoryTable.user_id == str(user_id),
+                    MemoryTable.is_obsolete.is_(False),
+                    MemoryTable.id.in_(ids),
+                ),
+            )
+            owned_ids = {row[0] for row in owned_rows}
+            written: list[int] = []
+            for memory_id, embedding in updates:
+                if memory_id not in owned_ids:
+                    continue
+                embedding_bytes = sqlite_vec.serialize_float32(embedding)
+                await session.execute(
+                    text(
+                        "DELETE FROM vec_memories WHERE memory_id = :memory_id "
+                        "AND EXISTS ("
+                        "SELECT 1 FROM memories "
+                        "WHERE memories.id = CAST(:memory_id AS INTEGER) "
+                        "AND memories.user_id = :user_id "
+                        "AND memories.is_obsolete = 0"
+                        ")",
+                    ),
+                    {"memory_id": str(memory_id), "user_id": str(user_id)},
+                )
+                await session.execute(
+                    text("INSERT INTO vec_memories (memory_id, embedding) VALUES (:memory_id, :embedding)"),
+                    {"memory_id": str(memory_id), "embedding": embedding_bytes},
+                )
+                written.append(memory_id)
+            return written
+
     async def validate_embedding_count(self) -> bool:
         """Check embedding count matches non-obsolete memory count"""
         async with self.db_adapter.system_session() as session: