Skip to content

Commit ee33683

Browse files
committed
feat: add workspace selection and routing for MCP tools
Signed-off-by: phernandez <paul@basicmachines.co>
1 parent 8c05a9e commit ee33683

28 files changed

Lines changed: 829 additions & 58 deletions

src/basic_memory/cli/commands/__init__.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,16 @@
11
"""CLI commands for basic-memory."""
22

33
from . import status, db, doctor, import_memory_json, mcp, import_claude_conversations
4-
from . import import_claude_projects, import_chatgpt, tool, project, format, schema, watch
4+
from . import (
5+
import_claude_projects,
6+
import_chatgpt,
7+
tool,
8+
project,
9+
format,
10+
schema,
11+
watch,
12+
workspace,
13+
)
514

615
__all__ = [
716
"status",
@@ -17,4 +26,5 @@
1726
"format",
1827
"schema",
1928
"watch",
29+
"workspace",
2030
]

src/basic_memory/cli/commands/tool.py

Lines changed: 88 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -99,14 +99,26 @@ def _parse_opening_frontmatter(content: str) -> tuple[str, dict[str, Any] | None
9999

100100

101101
async def _write_note_json(
102-
title: str, content: str, folder: str, project_name: Optional[str], tags: Optional[List[str]]
102+
title: str,
103+
content: str,
104+
folder: str,
105+
project_name: Optional[str],
106+
workspace: Optional[str],
107+
tags: Optional[List[str]],
103108
) -> dict:
104109
"""Write a note and return structured JSON metadata."""
105110
# Use the MCP tool to create/update the entity (handles create-or-update logic)
106-
await mcp_write_note.fn(title, content, folder, project_name, tags)
111+
await mcp_write_note.fn(
112+
title=title,
113+
content=content,
114+
directory=folder,
115+
project=project_name,
116+
workspace=workspace,
117+
tags=tags,
118+
)
107119

108120
# Resolve the entity to get metadata back
109-
async with get_client(project_name=project_name) as client:
121+
async with get_client(project_name=project_name, workspace=workspace) as client:
110122
active_project = await get_active_project(client, project_name)
111123
knowledge_client = KnowledgeClient(client, active_project.external_id)
112124

@@ -125,10 +137,14 @@ async def _write_note_json(
125137

126138

127139
async def _read_note_json(
128-
identifier: str, project_name: Optional[str], page: int, page_size: int
140+
identifier: str,
141+
project_name: Optional[str],
142+
workspace: Optional[str],
143+
page: int,
144+
page_size: int,
129145
) -> dict:
130146
"""Read a note and return structured JSON with content and metadata."""
131-
async with get_client(project_name=project_name) as client:
147+
async with get_client(project_name=project_name, workspace=workspace) as client:
132148
active_project = await get_active_project(client, project_name)
133149
knowledge_client = KnowledgeClient(client, active_project.external_id)
134150
resource_client = ResourceClient(client, active_project.external_id)
@@ -146,7 +162,10 @@ async def _read_note_json(
146162
from basic_memory.mcp.tools.search import search_notes as mcp_search_tool
147163

148164
title_results = await mcp_search_tool.fn(
149-
query=identifier, search_type="title", project=project_name
165+
query=identifier,
166+
search_type="title",
167+
project=project_name,
168+
workspace=workspace,
150169
)
151170
if title_results and hasattr(title_results, "results") and title_results.results:
152171
result = title_results.results[0]
@@ -172,12 +191,13 @@ async def _edit_note_json(
172191
operation: str,
173192
content: str,
174193
project_name: Optional[str],
194+
workspace: Optional[str],
175195
section: Optional[str],
176196
find_text: Optional[str],
177197
expected_replacements: int,
178198
) -> dict:
179199
"""Edit a note and return structured JSON metadata."""
180-
async with get_client(project_name=project_name) as client:
200+
async with get_client(project_name=project_name, workspace=workspace) as client:
181201
active_project = await get_active_project(client, project_name)
182202
knowledge_client = KnowledgeClient(client, active_project.external_id)
183203

@@ -227,11 +247,12 @@ async def _recent_activity_json(
227247
depth: Optional[int],
228248
timeframe: Optional[TimeFrame],
229249
project_name: Optional[str] = None,
250+
workspace: Optional[str] = None,
230251
page: int = 1,
231252
page_size: int = 50,
232253
) -> list:
233254
"""Get recent activity and return structured JSON list."""
234-
async with get_client(project_name=project_name) as client:
255+
async with get_client(project_name=project_name, workspace=workspace) as client:
235256
# Build query params matching the MCP tool's logic
236257
params: dict = {"page": page, "page_size": page_size, "max_related": 10}
237258
if depth:
@@ -275,6 +296,10 @@ def write_note(
275296
help="The project to write to. If not provided, the default project will be used."
276297
),
277298
] = None,
299+
workspace: Annotated[
300+
Optional[str],
301+
typer.Option(help="Cloud workspace tenant ID or unique name to route this request."),
302+
] = None,
278303
content: Annotated[
279304
Optional[str],
280305
typer.Option(
@@ -362,12 +387,19 @@ def write_note(
362387
with force_routing(local=local, cloud=cloud):
363388
if format == "json":
364389
result = run_with_cleanup(
365-
_write_note_json(title, content, folder, project_name, tags)
390+
_write_note_json(title, content, folder, project_name, workspace, tags)
366391
)
367392
print(json.dumps(result, indent=2, ensure_ascii=True, default=str))
368393
else:
369394
note = run_with_cleanup(
370-
mcp_write_note.fn(title, content, folder, project_name, tags)
395+
mcp_write_note.fn(
396+
title=title,
397+
content=content,
398+
directory=folder,
399+
project=project_name,
400+
workspace=workspace,
401+
tags=tags,
402+
)
371403
)
372404
rprint(note)
373405
except ValueError as e:
@@ -389,6 +421,10 @@ def read_note(
389421
help="The project to use for the note. If not provided, the default project will be used."
390422
),
391423
] = None,
424+
workspace: Annotated[
425+
Optional[str],
426+
typer.Option(help="Cloud workspace tenant ID or unique name to route this request."),
427+
] = None,
392428
page: int = 1,
393429
page_size: int = 10,
394430
format: str = typer.Option("text", "--format", help="Output format: text or json"),
@@ -429,15 +465,23 @@ def read_note(
429465
with force_routing(local=local, cloud=cloud):
430466
if format == "json":
431467
result = run_with_cleanup(
432-
_read_note_json(identifier, project_name, page, page_size)
468+
_read_note_json(identifier, project_name, workspace, page, page_size)
433469
)
434470
stripped_content, parsed_frontmatter = _parse_opening_frontmatter(result["content"])
435471
result["frontmatter"] = parsed_frontmatter
436472
if strip_frontmatter:
437473
result["content"] = stripped_content
438474
print(json.dumps(result, indent=2, ensure_ascii=True, default=str))
439475
else:
440-
note = run_with_cleanup(mcp_read_note.fn(identifier, project_name, page, page_size))
476+
note = run_with_cleanup(
477+
mcp_read_note.fn(
478+
identifier=identifier,
479+
project=project_name,
480+
workspace=workspace,
481+
page=page,
482+
page_size=page_size,
483+
)
484+
)
441485
if strip_frontmatter:
442486
note, _ = _parse_opening_frontmatter(note)
443487
rprint(note)
@@ -462,6 +506,10 @@ def edit_note(
462506
help="The project to edit. If not provided, the default project will be used."
463507
),
464508
] = None,
509+
workspace: Annotated[
510+
Optional[str],
511+
typer.Option(help="Cloud workspace tenant ID or unique name to route this request."),
512+
] = None,
465513
find_text: Annotated[
466514
Optional[str], typer.Option("--find-text", help="Text to find for find_replace operation")
467515
] = None,
@@ -509,6 +557,7 @@ def edit_note(
509557
operation=operation,
510558
content=content,
511559
project_name=project_name,
560+
workspace=workspace,
512561
section=section,
513562
find_text=find_text,
514563
expected_replacements=expected_replacements,
@@ -522,6 +571,7 @@ def edit_note(
522571
operation=operation,
523572
content=content,
524573
project=project_name,
574+
workspace=workspace,
525575
section=section,
526576
find_text=find_text,
527577
expected_replacements=expected_replacements,
@@ -547,6 +597,10 @@ def build_context(
547597
Optional[str],
548598
typer.Option(help="The project to use. If not provided, the default project will be used."),
549599
] = None,
600+
workspace: Annotated[
601+
Optional[str],
602+
typer.Option(help="Cloud workspace tenant ID or unique name to route this request."),
603+
] = None,
550604
depth: Optional[int] = 1,
551605
timeframe: Optional[TimeFrame] = "7d",
552606
page: int = 1,
@@ -582,6 +636,7 @@ def build_context(
582636
result = run_with_cleanup(
583637
mcp_build_context.fn(
584638
project=project_name,
639+
workspace=workspace,
585640
url=url,
586641
depth=depth,
587642
timeframe=timeframe,
@@ -609,6 +664,10 @@ def recent_activity(
609664
Optional[str],
610665
typer.Option(help="The project to use. If not provided, the default project will be used."),
611666
] = None,
667+
workspace: Annotated[
668+
Optional[str],
669+
typer.Option(help="Cloud workspace tenant ID or unique name to route this request."),
670+
] = None,
612671
depth: Optional[int] = 1,
613672
timeframe: Optional[TimeFrame] = "7d",
614673
page: int = typer.Option(1, "--page", help="Page number for pagination (JSON format)"),
@@ -642,7 +701,15 @@ def recent_activity(
642701
with force_routing(local=local, cloud=cloud):
643702
if format == "json":
644703
result = run_with_cleanup(
645-
_recent_activity_json(type, depth, timeframe, project_name, page, page_size)
704+
_recent_activity_json(
705+
type=type,
706+
depth=depth,
707+
timeframe=timeframe,
708+
project_name=project_name,
709+
workspace=workspace,
710+
page=page,
711+
page_size=page_size,
712+
)
646713
)
647714
print(json.dumps(result, indent=2, ensure_ascii=True, default=str))
648715
else:
@@ -652,6 +719,7 @@ def recent_activity(
652719
depth=depth,
653720
timeframe=timeframe,
654721
project=project_name,
722+
workspace=workspace,
655723
)
656724
)
657725
# The tool returns a formatted string directly
@@ -682,6 +750,10 @@ def search_notes(
682750
help="The project to use for the note. If not provided, the default project will be used."
683751
),
684752
] = None,
753+
workspace: Annotated[
754+
Optional[str],
755+
typer.Option(help="Cloud workspace tenant ID or unique name to route this request."),
756+
] = None,
685757
after_date: Annotated[
686758
Optional[str],
687759
typer.Option("--after_date", help="Search results after date, eg. '2d', '1 week'"),
@@ -793,8 +865,9 @@ def search_notes(
793865
with force_routing(local=local, cloud=cloud):
794866
results = run_with_cleanup(
795867
mcp_search.fn(
796-
query or "",
797-
project_name,
868+
query=query or "",
869+
project=project_name,
870+
workspace=workspace,
798871
search_type=search_type,
799872
page=page,
800873
after_date=after_date,
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
"""Workspace commands for Basic Memory cloud workspaces."""
2+
3+
import typer
4+
from rich.console import Console
5+
from rich.table import Table
6+
7+
from basic_memory.cli.app import app
8+
from basic_memory.cli.commands.command_utils import run_with_cleanup
9+
from basic_memory.mcp.project_context import get_available_workspaces
10+
11+
console = Console()
12+
13+
workspace_app = typer.Typer(help="Manage cloud workspaces")
14+
app.add_typer(workspace_app, name="workspace")
15+
16+
17+
@workspace_app.command("list")
18+
def list_workspaces() -> None:
19+
"""List cloud workspaces available to the current OAuth session."""
20+
21+
async def _list():
22+
return await get_available_workspaces()
23+
24+
try:
25+
workspaces = run_with_cleanup(_list())
26+
except RuntimeError as exc:
27+
console.print(f"[red]Error: {exc}[/red]")
28+
raise typer.Exit(1)
29+
except Exception as exc: # pragma: no cover
30+
console.print(f"[red]Error listing workspaces: {exc}[/red]")
31+
raise typer.Exit(1)
32+
33+
if not workspaces:
34+
console.print("[yellow]No accessible workspaces found.[/yellow]")
35+
return
36+
37+
table = Table(title="Available Workspaces")
38+
table.add_column("Name", style="cyan")
39+
table.add_column("Type", style="blue")
40+
table.add_column("Role", style="green")
41+
table.add_column("Tenant ID", style="yellow")
42+
43+
for workspace in workspaces:
44+
table.add_row(
45+
workspace.name,
46+
workspace.workspace_type,
47+
workspace.role,
48+
workspace.tenant_id,
49+
)
50+
51+
console.print(table)
52+
53+
54+
@app.command("workspaces")
55+
def workspaces_alias() -> None:
56+
"""Alias for `bm workspace list`."""
57+
list_workspaces()

src/basic_memory/cli/main.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ def _version_only_invocation(argv: list[str]) -> bool:
2828
schema,
2929
status,
3030
tool,
31+
workspace,
3132
)
3233

3334
warnings.filterwarnings("ignore") # pragma: no cover

0 commit comments

Comments
 (0)