Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions src/basic_memory/mcp/async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,11 +128,14 @@ async def get_cloud_control_plane_client(
yield client


# Optional factory override for dependency injection
_client_factory: Optional[Callable[[], AbstractAsyncContextManager[AsyncClient]]] = None
# Optional factory override for dependency injection.
# The factory accepts an optional workspace keyword argument so that MCP tools
# can route individual requests to a different workspace than the one set at
# connection time. See basic-memory-cloud main.py tenant_asgi_client_factory.
_client_factory: Optional[Callable[..., AbstractAsyncContextManager[AsyncClient]]] = None


def set_client_factory(factory: Callable[[], AbstractAsyncContextManager[AsyncClient]]) -> None:
def set_client_factory(factory: Callable[..., AbstractAsyncContextManager[AsyncClient]]) -> None:
"""Override the default client factory (for cloud app, testing, etc)."""
global _client_factory
_client_factory = factory
Expand Down Expand Up @@ -173,7 +176,7 @@ async def get_client(
4. Local ASGI transport by default.
"""
if _client_factory:
async with _client_factory() as client:
async with _client_factory(workspace=workspace) as client:
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve compatibility with no-arg client factories

get_client() now unconditionally calls the injected factory with workspace=..., which breaks existing factories that implement the previous zero-argument contract and are still valid according to prior versions of set_client_factory. Those integrations will now fail at runtime with an unexpected-keyword TypeError before any request handling. Please keep a compatibility path for factories that do not accept the new keyword.

Useful? React with 👍 / 👎.

yield client
return

Expand Down
10 changes: 5 additions & 5 deletions src/basic_memory/mcp/project_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -622,10 +622,10 @@ async def get_project_client(

# Step 1b: Factory injection (in-process cloud server)
# Trigger: set_client_factory() was called (e.g., by cloud MCP server)
# Why: the transport layer already resolved workspace and tenant context;
# attempting cloud workspace resolution here would call the production
# control-plane API with no valid credentials and fail with 401
# Outcome: use the factory client directly, skip workspace resolution
# Why: the factory's transport layer handles auth and tenant resolution;
# we pass workspace through so the transport can route to the correct
# workspace when the tool specifies one different from the connection default
# Outcome: factory client with optional workspace override via inner request headers
if is_factory_mode():
route_mode = "factory"
with telemetry.scope(
Expand All @@ -635,7 +635,7 @@ async def get_project_client(
workspace_id=workspace,
):
logger.debug("Using injected client factory for project routing")
async with get_client() as client:
async with get_client(workspace=workspace) as client:
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Invalidate cached project when workspace override changes

This new factory-mode call path now forwards workspace, but get_active_project() still reuses active_project cache entries keyed only by project identifier, not workspace. In one MCP session, a call for project="..." in workspace A followed by workspace B can return the cached project from A, so downstream tools use the wrong external_id under the new workspace header and may operate on the wrong project (or fail with confusing 404/authorization errors). The cache should be re-keyed or bypassed when the requested workspace differs.

Useful? React with 👍 / 👎.

active_project = await get_active_project(client, resolved_project, context)
yield client, active_project
return
Expand Down
4 changes: 2 additions & 2 deletions tests/mcp/test_async_client_modes.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ async def test_get_client_uses_injected_factory(monkeypatch):
seen = {"used": False}

@asynccontextmanager
async def factory():
async def factory(workspace=None):
seen["used"] = True
async with httpx.AsyncClient(base_url="https://example.test") as client:
yield client
Expand Down Expand Up @@ -185,7 +185,7 @@ async def test_get_client_factory_overrides_per_project_routing(config_manager):
config_manager.save_config(cfg)

@asynccontextmanager
async def factory():
async def factory(workspace=None):
async with httpx.AsyncClient(base_url="https://factory.test") as client:
yield client

Expand Down
2 changes: 1 addition & 1 deletion tests/mcp/test_project_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -709,7 +709,7 @@ async def test_factory_mode_skips_workspace_resolution(self, config_manager, mon

# Set up a factory (simulates what cloud MCP server does)
@asynccontextmanager
async def fake_factory():
async def fake_factory(workspace=None):
from httpx import ASGITransport, AsyncClient
from basic_memory.api.app import app as fastapi_app

Expand Down
Loading