| """ |
| Persistent CLI wrapper for interacting with a running Workspace MCP server. |
| |
| Reuses the project's existing FileTreeStore and FernetEncryptionWrapper to |
| cache OAuth tokens on disk so that ``fastmcp list/call`` does not re-trigger |
| the full browser-based OAuth flow on every invocation. |
| |
| Usage:: |
| |
| uv run workspace-cli list |
| uv run workspace-cli call search_gmail_messages query="is:unread" max_results=5 |
| """ |
|
|
| import argparse |
| import asyncio |
| import json |
| import logging |
| import os |
| import stat |
| import sys |
| from typing import Any |
|
|
| from cryptography.fernet import Fernet |
| from fastmcp import Client |
| from fastmcp.client.auth import OAuth |
| from key_value.aio.wrappers.encryption import FernetEncryptionWrapper |
|
|
| from core.storage import make_sanitized_file_store |
|
|
| logger = logging.getLogger(__name__) |
|
|
| DEFAULT_URL = "http://localhost:8000/mcp" |
| CLI_HOME = os.path.expanduser("~/.workspace-mcp") |
| TOKEN_DIR = os.path.join(CLI_HOME, "cli-tokens") |
| KEY_PATH = os.path.join(CLI_HOME, ".cli-encryption-key") |
|
|
|
|
| def _get_token_storage() -> FernetEncryptionWrapper: |
| """Return an encrypted, disk-backed token store. |
| |
| On first run the directory tree and a random Fernet key are created. |
| The key file is restricted to owner-only access (0o600). |
| """ |
| os.makedirs(TOKEN_DIR, exist_ok=True) |
|
|
| try: |
| fd = os.open(KEY_PATH, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) |
| except FileExistsError: |
| with open(KEY_PATH, "rb") as fh: |
| key = fh.read() |
| else: |
| key = Fernet.generate_key() |
| try: |
| os.write(fd, key) |
| finally: |
| os.close(fd) |
| os.chmod(KEY_PATH, stat.S_IRUSR | stat.S_IWUSR) |
|
|
| return FernetEncryptionWrapper( |
| key_value=make_sanitized_file_store(TOKEN_DIR), |
| fernet=Fernet(key), |
| ) |
|
|
|
|
| def _build_oauth() -> OAuth: |
| """Build an OAuth helper with persistent encrypted token storage.""" |
| storage = _get_token_storage() |
| return OAuth(token_storage=storage) |
|
|
|
|
| async def _list_tools(url: str) -> None: |
| """Connect, authenticate once, and print available tools.""" |
| try: |
| auth = _build_oauth() |
| except Exception as e: |
| print( |
| f"Error: failed to initialize OAuth ({type(e).__name__})", file=sys.stderr |
| ) |
| sys.exit(1) |
| try: |
| async with Client(url, auth=auth) as client: |
| tools = await client.list_tools() |
| except Exception as e: |
| print(f"Error: failed to list tools ({type(e).__name__})", file=sys.stderr) |
| sys.exit(1) |
| for tool in tools: |
| desc = (tool.description or "").split("\n")[0] |
| print(f" {tool.name:40s} {desc}") |
| print(f"\n{len(tools)} tools available") |
|
|
|
|
| async def _call_tool(url: str, tool_name: str, raw_args: list[str]) -> None: |
| """Connect, authenticate once, call a single tool, and print the result.""" |
| kwargs: dict[str, Any] = {} |
| for arg in raw_args: |
| if "=" not in arg: |
| print(f"Error: argument '{arg}' must be in key=value form", file=sys.stderr) |
| sys.exit(1) |
| k, v = arg.split("=", 1) |
| try: |
| kwargs[k] = json.loads(v) |
| except json.JSONDecodeError: |
| kwargs[k] = v |
|
|
| async with Client(url, auth=_build_oauth()) as client: |
| result = await client.call_tool(tool_name, kwargs) |
| for block in result.content: |
| if hasattr(block, "text"): |
| try: |
| parsed = json.loads(block.text) |
| print(json.dumps(parsed, indent=2)) |
| except (json.JSONDecodeError, TypeError): |
| print(block.text) |
| else: |
| print(block) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser( |
| prog="workspace-cli", |
| description="CLI for Workspace MCP with persistent OAuth token caching", |
| ) |
| parser.add_argument( |
| "--url", |
| default=os.getenv("WORKSPACE_MCP_URL", DEFAULT_URL), |
| help=f"MCP server URL (default: {DEFAULT_URL})", |
| ) |
|
|
| sub = parser.add_subparsers(dest="command") |
|
|
| sub.add_parser("list", help="List available tools") |
|
|
| call_parser = sub.add_parser("call", help="Call a tool") |
| call_parser.add_argument("tool", help="Tool name") |
| call_parser.add_argument("args", nargs="*", help="key=value arguments") |
|
|
| args = parser.parse_args() |
|
|
| if not args.command: |
| parser.print_help() |
| sys.exit(1) |
|
|
| if args.command == "list": |
| asyncio.run(_list_tools(args.url)) |
| elif args.command == "call": |
| asyncio.run(_call_tool(args.url, args.tool, args.args)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
