From 8e64caf7845f4f5fbc2bc15b6a59e5da188b8709 Mon Sep 17 00:00:00 2001 From: phernandez Date: Sun, 1 Mar 2026 12:07:24 -0600 Subject: [PATCH] fix: list_workspaces bypasses factory pattern on cloud MCP server (#636) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add set_workspace_provider() injection point so the cloud MCP server can list workspaces by querying its own database directly, instead of making an HTTP round-trip to the control-plane API with credentials it doesn't have. ๐Ÿ”ง Mirrors the existing set_client_factory() pattern in async_client.py ๐Ÿงช Adds 3 tests for provider injection, fallback, and context caching ๐Ÿฉน Updates build_context test assertions for v0.18 backward compat fields Co-Authored-By: Claude Opus 4.6 Signed-off-by: phernandez --- src/basic_memory/mcp/project_context.py | 27 ++++++- tests/mcp/test_tool_build_context.py | 28 +++---- tests/mcp/test_tool_workspace_management.py | 87 +++++++++++++++++++++ 3 files changed, 126 insertions(+), 16 deletions(-) diff --git a/src/basic_memory/mcp/project_context.py b/src/basic_memory/mcp/project_context.py index 0ccc17c8..02f56134 100644 --- a/src/basic_memory/mcp/project_context.py +++ b/src/basic_memory/mcp/project_context.py @@ -9,7 +9,7 @@ compatibility with existing MCP tools. """ from contextlib import asynccontextmanager -from typing import AsyncIterator, Optional, List, Tuple +from typing import AsyncIterator, Awaitable, Callable, Optional, List, Tuple from httpx import AsyncClient from httpx._types import ( @@ -27,6 +27,18 @@ from basic_memory.schemas.v2 import ProjectResolveResponse from basic_memory.schemas.memory import memory_url_path from basic_memory.utils import generate_permalink, normalize_project_reference +# --- Workspace provider injection --- +# Mirrors the set_client_factory() pattern in async_client.py. +# The cloud MCP server sets a provider that queries its own database directly, +# avoiding the control-plane HTTP round-trip that requires local credentials. +_workspace_provider: Optional[Callable[[], Awaitable[list[WorkspaceInfo]]]] = None + + +def set_workspace_provider(provider: Callable[[], Awaitable[list[WorkspaceInfo]]]) -> None: + """Override workspace discovery (for cloud app, testing, etc).""" + global _workspace_provider + _workspace_provider = provider + async def resolve_project_parameter( project: Optional[str] = None, @@ -103,6 +115,19 @@ async def get_available_workspaces(context: Optional[Context] = None) -> list[Wo if isinstance(cached_raw, list): return [WorkspaceInfo.model_validate(item) for item in cached_raw] + # Trigger: workspace provider was injected (e.g., by cloud MCP server) + # Why: the cloud server IS the cloud โ€” it can query its own database + # directly instead of making an HTTP round-trip that requires local credentials + # Outcome: use provider result, cache in context, skip control-plane client + if _workspace_provider is not None: + workspaces = await _workspace_provider() + if context: + await context.set_state( + "available_workspaces", + [ws.model_dump() for ws in workspaces], + ) + return workspaces + from basic_memory.mcp.async_client import get_cloud_control_plane_client from basic_memory.mcp.tools.utils import call_get diff --git a/tests/mcp/test_tool_build_context.py b/tests/mcp/test_tool_build_context.py index 8ea297a8..64d0fbec 100644 --- a/tests/mcp/test_tool_build_context.py +++ b/tests/mcp/test_tool_build_context.py @@ -9,7 +9,7 @@ from basic_memory.mcp.tools import build_context @pytest.mark.asyncio async def test_get_basic_discussion_context(client, test_graph, test_project): - """Test getting basic discussion context returns JSON dict with excluded fields removed.""" + """Test getting basic discussion context returns JSON dict with expected fields.""" result = await build_context(project=test_project.name, url="memory://test/root") assert isinstance(result, dict) @@ -19,29 +19,27 @@ async def test_get_basic_discussion_context(client, test_graph, test_project): assert primary["permalink"] == f"{test_project.name}/test/root" assert len(result["results"][0]["related_results"]) > 0 - # Verify metadata โ€” excluded fields should be absent + # Verify metadata fields meta = result["metadata"] assert meta["uri"] == f"{test_project.name}/test/root" assert meta["depth"] == 1 # default depth assert meta["timeframe"] is not None assert meta["primary_count"] == 1 - assert "generated_at" not in meta - assert "total_results" not in meta + # COMPAT(v0.18): generated_at and total_results restored for old clients + assert "generated_at" in meta + assert "total_results" in meta - # Entity: entity_id excluded, created_at kept (needed for related results) - assert "entity_id" not in primary + # Entity fields present + assert "entity_id" in primary assert "created_at" in primary - # Verify observation-level fields: internal IDs excluded, file_path/created_at kept + # Verify observation-level fields if result["results"][0]["observations"]: obs = result["results"][0]["observations"][0] - assert "observation_id" not in obs - assert "entity_id" not in obs - assert "title" not in obs - # file_path and created_at kept (needed when observation is primary_result) + assert "observation_id" in obs + assert "entity_id" in obs assert "file_path" in obs assert "created_at" in obs - # Other kept fields assert "permalink" in obs assert "category" in obs assert "content" in obs @@ -53,14 +51,14 @@ async def test_get_basic_discussion_context(client, test_graph, test_project): assert "title" in related assert "file_path" in related assert "created_at" in related - assert "entity_id" not in related # excluded + assert "entity_id" in related elif item_type == "relation": assert "relation_type" in related assert "title" in related assert "file_path" in related assert "created_at" in related - assert "relation_id" not in related # excluded - assert "entity_id" not in related # excluded + assert "relation_id" in related + assert "entity_id" in related @pytest.mark.asyncio diff --git a/tests/mcp/test_tool_workspace_management.py b/tests/mcp/test_tool_workspace_management.py index ff186b1c..0cf3f8c8 100644 --- a/tests/mcp/test_tool_workspace_management.py +++ b/tests/mcp/test_tool_workspace_management.py @@ -2,6 +2,7 @@ import pytest +from basic_memory.mcp.project_context import get_available_workspaces, set_workspace_provider from basic_memory.mcp.tools.workspaces import list_workspaces from basic_memory.schemas.cloud import WorkspaceInfo @@ -105,3 +106,89 @@ async def test_list_workspaces_uses_context_cache_path(monkeypatch): assert "# Available Workspaces (1)" in first assert "# Available Workspaces (1)" in second assert call_count["fetches"] == 1 + + +# --- Workspace provider injection tests --- + + +@pytest.fixture +def _reset_workspace_provider(monkeypatch): + """Ensure _workspace_provider is reset after each test.""" + import basic_memory.mcp.project_context as _mod + + monkeypatch.setattr(_mod, "_workspace_provider", None) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_reset_workspace_provider") +async def test_get_available_workspaces_uses_provider_when_set(): + """When a workspace provider is injected, it is called instead of the control-plane client.""" + expected = [ + WorkspaceInfo( + tenant_id="aaaa-bbbb", + workspace_type="personal", + name="Injected", + role="owner", + ), + ] + + async def fake_provider() -> list[WorkspaceInfo]: + return expected + + set_workspace_provider(fake_provider) + + result = await get_available_workspaces() + assert len(result) == 1 + assert result[0].tenant_id == "aaaa-bbbb" + assert result[0].name == "Injected" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_reset_workspace_provider") +async def test_get_available_workspaces_falls_back_without_provider(monkeypatch): + """Without a provider, get_available_workspaces uses the control-plane client (existing path).""" + called = {"control_plane": False} + + async def fake_control_plane_path(context=None): + called["control_plane"] = True + return [] + + # Patch the entire function to avoid needing real credentials + monkeypatch.setattr( + "basic_memory.mcp.tools.workspaces.get_available_workspaces", + fake_control_plane_path, + ) + + result = await list_workspaces() + assert called["control_plane"] + assert "# No Workspaces Available" in result + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_reset_workspace_provider") +async def test_get_available_workspaces_provider_caches_in_context(): + """Provider results are cached in the MCP context for subsequent calls.""" + call_count = {"provider": 0} + workspace = WorkspaceInfo( + tenant_id="cccc-dddd", + workspace_type="organization", + name="Cached Provider", + role="editor", + ) + + async def counting_provider() -> list[WorkspaceInfo]: + call_count["provider"] += 1 + return [workspace] + + set_workspace_provider(counting_provider) + context = _ContextState() + + # First call: provider is invoked, result cached + first = await get_available_workspaces(context=context) + assert len(first) == 1 + assert call_count["provider"] == 1 + + # Second call: served from context cache, provider not called again + second = await get_available_workspaces(context=context) + assert len(second) == 1 + assert call_count["provider"] == 1