File size: 4,605 Bytes
bf06a0d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """
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()
|