fix: list_workspaces bypasses factory pattern on cloud MCP server (#636)

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 <noreply@anthropic.com>
Signed-off-by: phernandez <paul@basicmachines.co>

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
 

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

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