mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
fix(mcp): address PR review feedback (P2 + multi-project search)
Two follow-up fixes on the same branch: 1. **Purge sqlite-vec embeddings during project delete** (Codex P2) sqlite-vec stores vectors in a vec0 virtual table keyed by chunk rowid with no cascade. The previous purge removed search_vector_chunks but left the embeddings behind; `_run_vector_query` then keeps returning stale vectors that crowd live results. ProjectRepository.delete now deletes embeddings first (using the same rowid-IN-chunks pattern as SQLiteSearchRepository.delete_project_vector_rows), then the chunk rows. Both deletes are skipped if the underlying table is absent on a given install. New test test_remove_project_purges_vector_embeddings covers the happy path and skips cleanly when the embeddings table isn't initialized. 2. **Fix search_all_projects=True on local installs** `_search_all_projects` recurses into search_notes with both project= and project_id= set. project_id (external UUID) routes through the cloud v2 API path, which 401s on local installs because there's no JWT to present — so the inner calls silently failed and the merged result list stayed empty. The fan-out now mirrors get_project_client's cloud_available composite (factory mode OR explicit --cloud OR has_cloud_credentials). When that composite is false we forward project= only and take the name-routed local-ASGI path. Cloud disambiguation still works because the project name in project_ref is already the workspace/project qualified_name. The existing cloud-style fan-out tests now go through a cloud_routing fixture that pins the three signals; a new local_routing test confirms project_id is dropped when no cloud route is available. Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
@@ -10,8 +10,13 @@ from loguru import logger
|
||||
from fastmcp import Context
|
||||
from pydantic import AliasChoices, BeforeValidator, Field
|
||||
|
||||
from basic_memory.config import ConfigManager
|
||||
from basic_memory.config import ConfigManager, has_cloud_credentials
|
||||
from basic_memory.utils import build_canonical_permalink, coerce_dict, coerce_list
|
||||
from basic_memory.mcp.async_client import (
|
||||
_explicit_routing,
|
||||
_force_local_mode,
|
||||
is_factory_mode,
|
||||
)
|
||||
from basic_memory.mcp.container import get_container
|
||||
from basic_memory.mcp.project_context import (
|
||||
detect_project_from_identifier_prefix,
|
||||
@@ -479,12 +484,30 @@ async def _search_all_projects(
|
||||
total = 0
|
||||
any_project_has_more = False
|
||||
|
||||
# Trigger: caller asked for an account-wide search.
|
||||
# Why: project_id (external UUID) routes through the cloud v2 API path,
|
||||
# which 401s on local installs because there's no JWT to present.
|
||||
# Project names route through the local-ASGI path and work for both
|
||||
# backends — cloud disambiguates names via the workspace/project
|
||||
# qualified_name already baked into project_ref["project"].
|
||||
# Outcome: forward project_id only when the same signals get_project_client
|
||||
# uses to pick a cloud route are present. Mirrors the cloud_available
|
||||
# composite in project_context.get_project_client (single source of
|
||||
# truth for "can we route to cloud?").
|
||||
config = ConfigManager().config
|
||||
use_cloud_routing = (
|
||||
is_factory_mode()
|
||||
or (_explicit_routing() and not _force_local_mode())
|
||||
or has_cloud_credentials(config)
|
||||
)
|
||||
|
||||
for project_ref in project_refs:
|
||||
recursive_project_id = project_ref["project_id"] if use_cloud_routing else None
|
||||
try:
|
||||
results = await search_notes(
|
||||
query=query,
|
||||
project=project_ref["project"],
|
||||
project_id=project_ref["project_id"],
|
||||
project_id=recursive_project_id,
|
||||
page=1,
|
||||
page_size=per_project_page_size,
|
||||
search_type=search_type,
|
||||
|
||||
@@ -133,10 +133,15 @@ class ProjectRepository(Repository[Project]):
|
||||
inherits the previous tenant's content. search_vector_chunks is a real
|
||||
table on both backends but only carries the FK on Postgres.
|
||||
|
||||
search_index is created at runtime by SearchRepository.init_search_index
|
||||
and search_vector_chunks only appears once semantic search initializes,
|
||||
so each table may be absent on minimal test DBs. Inspect first and
|
||||
skip whichever table isn't there.
|
||||
sqlite-vec stores embeddings in a separate vec0 virtual table keyed by
|
||||
chunk rowid with no cascade, so embeddings must be purged before the
|
||||
chunk rows or `_run_vector_query` will keep returning stale vectors
|
||||
that crowd out live results.
|
||||
|
||||
Each derived table — search_index, search_vector_chunks,
|
||||
search_vector_embeddings — is created lazily on first use, so any of
|
||||
them may be absent on minimal test DBs or installs without semantic
|
||||
search. Inspect the connection once and skip whichever is missing.
|
||||
"""
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
try:
|
||||
@@ -150,12 +155,29 @@ class ProjectRepository(Repository[Project]):
|
||||
existing_tables = await session.run_sync(
|
||||
lambda sync_session: set(sa_inspect(sync_session.connection()).get_table_names())
|
||||
)
|
||||
for table in ("search_index", "search_vector_chunks"):
|
||||
if table in existing_tables:
|
||||
|
||||
if "search_index" in existing_tables:
|
||||
await session.execute(
|
||||
text("DELETE FROM search_index WHERE project_id = :project_id"),
|
||||
{"project_id": entity_id},
|
||||
)
|
||||
|
||||
if "search_vector_chunks" in existing_tables:
|
||||
if "search_vector_embeddings" in existing_tables:
|
||||
# sqlite-vec has no CASCADE — drop embeddings first while the
|
||||
# chunk rows that name them still exist.
|
||||
await session.execute(
|
||||
text(f"DELETE FROM {table} WHERE project_id = :project_id"),
|
||||
text(
|
||||
"DELETE FROM search_vector_embeddings WHERE rowid IN ("
|
||||
"SELECT id FROM search_vector_chunks "
|
||||
"WHERE project_id = :project_id)"
|
||||
),
|
||||
{"project_id": entity_id},
|
||||
)
|
||||
await session.execute(
|
||||
text("DELETE FROM search_vector_chunks WHERE project_id = :project_id"),
|
||||
{"project_id": entity_id},
|
||||
)
|
||||
|
||||
await session.delete(project)
|
||||
return True
|
||||
|
||||
@@ -8,8 +8,38 @@ import pytest
|
||||
from basic_memory.schemas.search import SearchItemType, SearchResponse, SearchResult
|
||||
|
||||
|
||||
def _stub_routing_mode(monkeypatch, *, cloud: bool) -> None:
|
||||
"""Pin the three cloud-route signals search.py reads.
|
||||
|
||||
`_search_all_projects` only forwards project_id (external UUID) when a
|
||||
cloud route is available. The composite mirrors get_project_client:
|
||||
factory mode OR explicit --cloud OR has_cloud_credentials. Tests stub
|
||||
all three so a dev box with OAuth tokens on disk can't bleed into the
|
||||
local-mode case.
|
||||
"""
|
||||
search_mod = importlib.import_module("basic_memory.mcp.tools.search")
|
||||
monkeypatch.setattr(search_mod, "is_factory_mode", lambda: False)
|
||||
monkeypatch.setattr(search_mod, "_explicit_routing", lambda: cloud)
|
||||
monkeypatch.setattr(search_mod, "_force_local_mode", lambda: False)
|
||||
monkeypatch.setattr(search_mod, "has_cloud_credentials", lambda config: cloud)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cloud_routing(monkeypatch):
|
||||
"""Force the cloud-routing path for multi-project search tests."""
|
||||
_stub_routing_mode(monkeypatch, cloud=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_routing(monkeypatch):
|
||||
"""Force the local-routing path for multi-project search tests."""
|
||||
_stub_routing_mode(monkeypatch, cloud=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_notes_search_all_projects_qualifies_result_permalinks(monkeypatch):
|
||||
async def test_search_notes_search_all_projects_qualifies_result_permalinks(
|
||||
monkeypatch, cloud_routing
|
||||
):
|
||||
"""Multi-project search belongs to search_notes and keeps result ids routable."""
|
||||
clients_mod = importlib.import_module("basic_memory.mcp.clients")
|
||||
search_mod = importlib.import_module("basic_memory.mcp.tools.search")
|
||||
@@ -169,7 +199,7 @@ async def test_search_notes_search_all_projects_with_no_refs_returns_empty_all_p
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_notes_search_all_projects_continues_after_project_failure(
|
||||
monkeypatch,
|
||||
monkeypatch, cloud_routing
|
||||
):
|
||||
"""One failing project should not discard successful all-project search results."""
|
||||
clients_mod = importlib.import_module("basic_memory.mcp.clients")
|
||||
@@ -254,3 +284,85 @@ async def test_search_notes_search_all_projects_continues_after_project_failure(
|
||||
assert result["total"] == 1
|
||||
assert any("team-paul/main" in warning for warning in warnings)
|
||||
assert any("team index unavailable" in warning for warning in warnings)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_notes_search_all_projects_local_omits_project_id(
|
||||
monkeypatch, local_routing
|
||||
):
|
||||
"""Without a cloud route, fan-out must address each project by name only.
|
||||
|
||||
project_id (external UUID) routes through the cloud v2 API path, which
|
||||
returns 401 on local installs because there's no JWT to present. Local
|
||||
fan-out has to fall back to the name-routed path so each per-project
|
||||
search actually returns results instead of silently failing.
|
||||
"""
|
||||
clients_mod = importlib.import_module("basic_memory.mcp.clients")
|
||||
search_mod = importlib.import_module("basic_memory.mcp.tools.search")
|
||||
|
||||
project_refs = [
|
||||
{
|
||||
"project": "alpha",
|
||||
"project_id": "11111111-1111-1111-1111-111111111111",
|
||||
},
|
||||
{
|
||||
"project": "beta",
|
||||
"project_id": "22222222-2222-2222-2222-222222222222",
|
||||
},
|
||||
]
|
||||
searched_projects: list[tuple[str | None, str | None]] = []
|
||||
|
||||
async def fake_load_search_project_refs(context=None):
|
||||
return project_refs
|
||||
|
||||
class StubProject:
|
||||
def __init__(self, name: str | None, external_id: str | None):
|
||||
self.name = name or "main"
|
||||
self.external_id = external_id or "local-main"
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_get_project_client(project=None, context=None, project_id=None):
|
||||
searched_projects.append((project, project_id))
|
||||
yield object(), StubProject(project, project_id)
|
||||
|
||||
async def fake_resolve_project_and_path(client, identifier, project=None, context=None):
|
||||
return StubProject(project, None), identifier, False
|
||||
|
||||
class MockSearchClient:
|
||||
def __init__(self, client, project_id):
|
||||
self.project_id = project_id
|
||||
|
||||
async def search(self, payload, page, page_size):
|
||||
return SearchResponse(
|
||||
results=[
|
||||
SearchResult(
|
||||
title=f"Note in {self.project_id or 'local'}",
|
||||
permalink="notes/example",
|
||||
content="",
|
||||
type=SearchItemType.ENTITY,
|
||||
score=0.5,
|
||||
file_path="/notes/example.md",
|
||||
)
|
||||
],
|
||||
current_page=page,
|
||||
page_size=page_size,
|
||||
total=1,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(search_mod, "_load_search_project_refs", fake_load_search_project_refs)
|
||||
monkeypatch.setattr(search_mod, "get_project_client", fake_get_project_client)
|
||||
monkeypatch.setattr(search_mod, "resolve_project_and_path", fake_resolve_project_and_path)
|
||||
monkeypatch.setattr(clients_mod, "SearchClient", MockSearchClient)
|
||||
|
||||
result = await search_mod.search_notes(
|
||||
query="anything",
|
||||
search_all_projects=True,
|
||||
output_format="json",
|
||||
)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert searched_projects == [("alpha", None), ("beta", None)], (
|
||||
"Local fan-out must omit project_id so the recursive search_notes calls "
|
||||
"take the name-routed path."
|
||||
)
|
||||
assert result["total"] == 2
|
||||
|
||||
@@ -140,6 +140,16 @@ async def test_remove_project_with_related_entities(project_service: ProjectServ
|
||||
await project_service.repository.delete(project.id)
|
||||
|
||||
|
||||
async def _table_exists(session_maker, table: str) -> bool:
|
||||
"""Return True if the named table is present on the current connection."""
|
||||
from sqlalchemy import inspect as sa_inspect
|
||||
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
return await session.run_sync(
|
||||
lambda sync_session: table in sa_inspect(sync_session.connection()).get_table_names()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_project_purges_search_rows(project_service: ProjectService):
|
||||
"""Project deletion must sweep the derived search tables.
|
||||
@@ -242,3 +252,89 @@ async def test_remove_project_purges_search_rows(project_service: ProjectService
|
||||
f"search_vector_chunks still has {post_chunks} rows for deleted "
|
||||
f"project_id={project_id}."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_project_purges_vector_embeddings(project_service: ProjectService):
|
||||
"""Project deletion must also drop sqlite-vec embeddings keyed by chunk rowid.
|
||||
|
||||
sqlite-vec stores vectors in a vec0 virtual table that has no cascade
|
||||
behavior. If embeddings linger after the chunks they reference are gone,
|
||||
`_run_vector_query` pulls them as top-k candidates and crowds out live
|
||||
results. The test only runs when the embeddings table is present, which
|
||||
matches the install path that exercises semantic search.
|
||||
"""
|
||||
test_project_name = f"test-vec-cleanup-{os.urandom(4).hex()}"
|
||||
session_maker = project_service.repository.session_maker
|
||||
|
||||
# The embeddings table only exists once semantic search has initialized.
|
||||
# Skipping when it's absent keeps this test honest on minimal CI DBs.
|
||||
if not await _table_exists(session_maker, "search_vector_embeddings"):
|
||||
pytest.skip("search_vector_embeddings is not present on this connection")
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
test_project_path = str(Path(temp_dir) / "test-vec-cleanup")
|
||||
os.makedirs(test_project_path, exist_ok=True)
|
||||
|
||||
await project_service.add_project(test_project_name, test_project_path)
|
||||
project = await project_service.get_project(test_project_name)
|
||||
assert project is not None
|
||||
project_id = project.id
|
||||
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO search_vector_chunks "
|
||||
"(id, entity_id, project_id, chunk_key, chunk_text, source_hash, "
|
||||
" entity_fingerprint, embedding_model) "
|
||||
"VALUES (:id, :entity_id, :project_id, :chunk_key, :chunk_text, "
|
||||
" :source_hash, :entity_fingerprint, :embedding_model)"
|
||||
),
|
||||
{
|
||||
"id": 999_201,
|
||||
"entity_id": 999_201,
|
||||
"project_id": project_id,
|
||||
"chunk_key": "vec-canary",
|
||||
"chunk_text": "vec canary",
|
||||
"source_hash": "abc",
|
||||
"entity_fingerprint": "",
|
||||
"embedding_model": "",
|
||||
},
|
||||
)
|
||||
# vec0 requires a vector matching the configured dimensions, but the
|
||||
# delete path filters by rowid; a non-existing dimension would block
|
||||
# this seed step. Skip the insert if the embeddings DDL hasn't run.
|
||||
try:
|
||||
await session.execute(
|
||||
text(
|
||||
"INSERT INTO search_vector_embeddings (rowid, embedding) "
|
||||
"VALUES (:rowid, :embedding)"
|
||||
),
|
||||
{"rowid": 999_201, "embedding": "[" + ",".join(["0.0"] * 384) + "]"},
|
||||
)
|
||||
except Exception:
|
||||
pytest.skip("search_vector_embeddings rejected the synthetic seed row")
|
||||
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
pre = (
|
||||
await session.execute(
|
||||
text("SELECT COUNT(*) FROM search_vector_embeddings WHERE rowid = :rowid"),
|
||||
{"rowid": 999_201},
|
||||
)
|
||||
).scalar_one()
|
||||
assert pre >= 1, "seed embedding should exist before removal"
|
||||
|
||||
await project_service.remove_project(test_project_name)
|
||||
|
||||
async with db.scoped_session(session_maker) as session:
|
||||
post = (
|
||||
await session.execute(
|
||||
text("SELECT COUNT(*) FROM search_vector_embeddings WHERE rowid = :rowid"),
|
||||
{"rowid": 999_201},
|
||||
)
|
||||
).scalar_one()
|
||||
|
||||
assert post == 0, (
|
||||
f"search_vector_embeddings still has {post} rows for rowid 999_201 "
|
||||
"— project deletion did not sweep the sqlite-vec embeddings table."
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user