gax-cli 0.4.0__py3-none-any.whl

gax/mcp_server.py

@@ -0,0 +1,299 @@
+"""
+gax-mcp — expose GAX as an MCP server (stdio JSON-RPC).
+
+Any MCP client (Claude Code, Cursor, OpenAI Agent SDK) gets governed shell
+execution with zero GAX-specific integration:
+
+    claude mcp add gax -- gax-mcp
+
+**Why three tools and not one per command.** A naive MCP server publishes one tool
+per capability, so a 43-command registry costs ~44k tokens of schema before the
+first turn. GAX publishes exactly three — `gax_search`, `gax_doc`, `gax_invoke` —
+and keeps the registry behind them. That is the same shape as Anthropic's Tool
+Search / `defer_loading` pattern, so this composes with the platform fix rather
+than competing with it: the client's own tool-search sees a constant-size surface
+no matter how many commands are registered.
+
+**The governance boundary is unchanged.** `gax_invoke` calls the same
+`gax.executor.invoke` as the CLI and the daemon, so capability checks, scope
+checks, and policy all run *before* any adapter, and every invoke emits an
+`audit_id`. An MCP client cannot reach a command that is not in the registry, and
+cannot bypass the capability check by phrasing the request differently — the model
+only ever proposes a command name plus args.
+
+Transport is stdio JSON-RPC 2.0 written against the stdlib, mirroring
+`gax.mcp_client`, so adding an MCP server adds no dependency.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from typing import Any
+
+from gax.executor import invoke
+from gax.registry import Registry
+
+PROTOCOL_VERSION = "2024-11-05"
+SERVER_NAME = "gax"
+SERVER_VERSION = "0.4"
+
+# JSON-RPC 2.0 reserved codes.
+PARSE_ERROR = -32700
+INVALID_REQUEST = -32600
+METHOD_NOT_FOUND = -32601
+INTERNAL_ERROR = -32603
+
+TOOLS: list[dict[str, Any]] = [
+    {
+        "name": "gax_search",
+        "description": (
+            "Search registered GAX commands by keyword. Use this first to find a "
+            "command; it returns names and one-line descriptions only, not full "
+            "schemas. Only commands returned here can be invoked."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Keywords, e.g. 'pull requests' or 'pods'.",
+                },
+                "limit": {
+                    "type": "integer",
+                    "description": "Max results (default 5).",
+                },
+            },
+            "required": ["query"],
+        },
+    },
+    {
+        "name": "gax_doc",
+        "description": (
+            "Get arguments, required scopes, and side effects for one GAX command. "
+            "Call this after gax_search and before gax_invoke."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "command": {
+                    "type": "string",
+                    "description": "Command id, e.g. 'gh.pr.list'.",
+                }
+            },
+            "required": ["command"],
+        },
+    },
+    {
+        "name": "gax_invoke",
+        "description": (
+            "Invoke a registered GAX command. The request is checked against the "
+            "capability token, required scopes, and policy BEFORE any backend runs; "
+            "denials return ok=false with an error.kind. Every call returns an "
+            "audit_id. Unregistered commands cannot be invoked."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "command": {
+                    "type": "string",
+                    "description": "Command id, e.g. 'gh.pr.list'.",
+                },
+                "args": {
+                    "type": "object",
+                    "description": "Arguments per gax_doc.",
+                },
+                "surface": {
+                    "type": "string",
+                    "enum": ["model", "human", "full"],
+                    "description": "Projection; 'model' truncates for the LLM.",
+                },
+            },
+            "required": ["command"],
+        },
+    },
+]
+
+
+class GaxMcpServer:
+    """MCP protocol surface over the GAX registry and executor."""
+
+    def __init__(
+        self,
+        registry: Registry | None = None,
+        *,
+        capability: str | None = None,
+    ) -> None:
+        self._registry = registry or Registry()
+        # Read GAX_CAP lazily at call time if not injected, so a capability minted
+        # after the server starts is still picked up.
+        self._capability = capability
+
+    # -- tool implementations ---------------------------------------------
+
+    def _cap(self) -> str | None:
+        return self._capability or os.environ.get("GAX_CAP") or None
+
+    def tool_search(self, args: dict[str, Any]) -> dict[str, Any]:
+        query = str(args.get("query", ""))
+        limit = int(args.get("limit") or 5)
+        hits = self._registry.search(query, limit=limit)
+        return {
+            "query": query,
+            "results": [
+                {
+                    "command": m.command,
+                    "description": m.description,
+                    "category": m.category,
+                }
+                for m in hits
+            ],
+        }
+
+    def tool_doc(self, args: dict[str, Any]) -> dict[str, Any]:
+        command = str(args.get("command", ""))
+        doc = self._registry.doc_stub(command)
+        if not doc:
+            return {
+                "error": "not_found",
+                "message": f"unknown command: {command}",
+                "hint": "use gax_search to list registered commands",
+            }
+        return doc
+
+    def tool_invoke(self, args: dict[str, Any]) -> dict[str, Any]:
+        command = str(args.get("command", ""))
+        call_args = dict(args.get("args") or {})
+        surface = str(args.get("surface") or "model")
+        envelope, _exit_code = invoke(
+            self._registry,
+            command=command,
+            args=call_args,
+            surface=surface,
+            capability=self._cap(),
+        )
+        return envelope
+
+    def call_tool(self, name: str, args: dict[str, Any]) -> tuple[dict[str, Any], bool]:
+        """Return (payload, is_error). is_error drives MCP's isError flag."""
+        if name == "gax_search":
+            return self.tool_search(args), False
+        if name == "gax_doc":
+            result = self.tool_doc(args)
+            return result, bool(result.get("error"))
+        if name == "gax_invoke":
+            envelope = self.tool_invoke(args)
+            # A policy denial is a legitimate protocol-level result, but flagging it
+            # as isError makes clients surface it instead of treating it as data.
+            return envelope, not bool(envelope.get("ok"))
+        raise ValueError(f"unknown tool: {name}")
+
+    # -- JSON-RPC dispatch -------------------------------------------------
+
+    def handle(self, req: dict[str, Any]) -> dict[str, Any] | None:
+        """Handle one request. Returns None for notifications (no id)."""
+        method = req.get("method")
+        rid = req.get("id")
+        params = req.get("params") or {}
+
+        # Notifications carry no id and must never get a response.
+        if rid is None:
+            return None
+
+        if method == "initialize":
+            return _result(
+                rid,
+                {
+                    "protocolVersion": PROTOCOL_VERSION,
+                    "capabilities": {"tools": {"listChanged": False}},
+                    "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
+                },
+            )
+
+        if method == "tools/list":
+            return _result(rid, {"tools": TOOLS})
+
+        if method == "tools/call":
+            name = str(params.get("name", ""))
+            args = dict(params.get("arguments") or {})
+            try:
+                payload, is_error = self.call_tool(name, args)
+            except ValueError as e:
+                return _error(rid, METHOD_NOT_FOUND, str(e))
+            except Exception as e:  # adapter blew up — report, don't crash the server
+                return _result(
+                    rid,
+                    {
+                        "content": [{"type": "text", "text": json.dumps({"error": str(e)})}],
+                        "isError": True,
+                    },
+                )
+            return _result(
+                rid,
+                {
+                    "content": [
+                        {"type": "text", "text": json.dumps(payload, ensure_ascii=False)}
+                    ],
+                    "isError": is_error,
+                },
+            )
+
+        if method == "ping":
+            return _result(rid, {})
+
+        return _error(rid, METHOD_NOT_FOUND, f"unknown method: {method}")
+
+
+def _result(rid: Any, result: dict[str, Any]) -> dict[str, Any]:
+    return {"jsonrpc": "2.0", "id": rid, "result": result}
+
+
+def _error(rid: Any, code: int, message: str) -> dict[str, Any]:
+    return {"jsonrpc": "2.0", "id": rid, "error": {"code": code, "message": message}}
+
+
+def serve(
+    stdin: Any = None,
+    stdout: Any = None,
+    *,
+    registry: Registry | None = None,
+) -> None:
+    """Run the stdio loop until EOF."""
+    stdin = stdin or sys.stdin
+    stdout = stdout or sys.stdout
+    server = GaxMcpServer(registry=registry)
+
+    for line in stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            req = json.loads(line)
+        except json.JSONDecodeError:
+            _write(stdout, _error(None, PARSE_ERROR, "invalid JSON"))
+            continue
+        if not isinstance(req, dict):
+            _write(stdout, _error(None, INVALID_REQUEST, "request must be an object"))
+            continue
+        try:
+            resp = server.handle(req)
+        except Exception as e:  # never let one bad request kill the session
+            _write(stdout, _error(req.get("id"), INTERNAL_ERROR, str(e)))
+            continue
+        if resp is not None:
+            _write(stdout, resp)
+
+
+def _write(stream: Any, msg: dict[str, Any]) -> None:
+    stream.write(json.dumps(msg, ensure_ascii=False) + "\n")
+    stream.flush()
+
+
+def main() -> None:
+    """Console entry point: `gax-mcp`."""
+    serve()
+
+
+if __name__ == "__main__":
+    main()

gax/oauth.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import json
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import httpx
+import yaml
+
+from gax.paths import CONFIG_DIR, GAX_HOME, ensure_gax_home
+
+PROVIDERS_PATH = CONFIG_DIR / "oauth_providers.yaml"
+
+
+@dataclass
+class OAuthProvider:
+    name: str
+    device_authorization_url: str
+    token_url: str
+    client_id_env: str
+    scopes: list[str]
+    audience: str | None = None
+
+    @property
+    def client_id(self) -> str:
+        val = os.environ.get(self.client_id_env, "").strip()
+        if not val:
+            raise ValueError(
+                f"Set {self.client_id_env} for OAuth client id (register an OAuth app for provider {self.name})"
+            )
+        return val
+
+
+def load_providers() -> dict[str, OAuthProvider]:
+    path = PROVIDERS_PATH
+    if not path.exists():
+        return {}
+    raw = yaml.safe_load(path.read_text()) or {}
+    out: dict[str, OAuthProvider] = {}
+    for name, cfg in (raw.get("providers") or {}).items():
+        out[name] = OAuthProvider(
+            name=name,
+            device_authorization_url=str(cfg["device_authorization_url"]),
+            token_url=str(cfg["token_url"]),
+            client_id_env=str(cfg.get("client_id_env", f"GAX_{name.upper()}_CLIENT_ID")),
+            scopes=list(cfg.get("scopes") or []),
+            audience=cfg.get("audience"),
+        )
+    return out
+
+
+def token_path(tenant: str, provider: str) -> Path:
+    ensure_gax_home()
+    d = GAX_HOME / "tokens" / tenant
+    d.mkdir(parents=True, exist_ok=True)
+    return d / f"{provider}.json"
+
+
+def save_tokens(tenant: str, provider: str, data: dict[str, Any]) -> Path:
+    from gax.secrets import save_secret
+
+    path = token_path(tenant, provider)
+    save_secret(tenant, provider, data, file_path=path)
+    return path
+
+
+def load_tokens(tenant: str, provider: str) -> dict[str, Any] | None:
+    from gax.secrets import load_secret
+
+    return load_secret(tenant, provider, file_path=token_path(tenant, provider))
+
+
+def device_flow_login(
+    provider: OAuthProvider,
+    *,
+    tenant: str,
+    open_browser: bool = True,
+    poll_interval: float | None = None,
+) -> dict[str, Any]:
+    """RFC 8628 device authorization grant."""
+    body: dict[str, Any] = {
+        "client_id": provider.client_id,
+        "scope": " ".join(provider.scopes),
+    }
+    headers = {"Accept": "application/json"}
+    if provider.name == "github":
+        headers["Content-Type"] = "application/json"
+        r = httpx.post(provider.device_authorization_url, json=body, headers=headers, timeout=30.0)
+    else:
+        r = httpx.post(provider.device_authorization_url, data=body, headers=headers, timeout=30.0)
+    r.raise_for_status()
+    device = r.json()
+
+    verification_uri = device.get("verification_uri_complete") or device.get("verification_uri")
+    user_code = device.get("user_code", "")
+    device_code = device["device_code"]
+    interval = poll_interval or float(device.get("interval", 5))
+    expires_in = int(device.get("expires_in", 900))
+
+    click_echo = __import__("click").echo
+    click_echo(f"\nVisit: {verification_uri}")
+    if user_code:
+        click_echo(f"Code: {user_code}\n")
+    if open_browser and verification_uri:
+        try:
+            import webbrowser
+
+            webbrowser.open(verification_uri)
+        except Exception:
+            pass
+
+    deadline = time.time() + expires_in
+    token_body: dict[str, Any] = {
+        "client_id": provider.client_id,
+        "device_code": device_code,
+        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+    }
+
+    while time.time() < deadline:
+        time.sleep(interval)
+        if provider.name == "github":
+            tr = httpx.post(
+                provider.token_url,
+                json=token_body,
+                headers={"Accept": "application/json", "Content-Type": "application/json"},
+                timeout=30.0,
+            )
+        else:
+            tr = httpx.post(
+                provider.token_url,
+                data=token_body,
+                headers={"Accept": "application/json"},
+                timeout=30.0,
+            )
+        if tr.status_code == 200:
+            tokens = tr.json()
+            if "access_token" in tokens:
+                tokens["obtained_at"] = time.time()
+                tokens["provider"] = provider.name
+                tokens["scopes"] = provider.scopes
+                save_tokens(tenant, provider.name, tokens)
+                return tokens
+        err = tr.json() if tr.headers.get("content-type", "").startswith("application/json") else {}
+        error = err.get("error", "")
+        if error in ("authorization_pending", "slow_down"):
+            if error == "slow_down":
+                interval = min(interval + 2, 30)
+            continue
+        if error == "access_denied":
+            raise RuntimeError("User denied authorization")
+        if error == "expired_token":
+            raise RuntimeError("Device code expired; run login again")
+        if tr.status_code >= 400 and error not in ("authorization_pending",):
+            raise RuntimeError(f"Token poll failed: {tr.text}")
+
+    raise RuntimeError("Device flow timed out")
+
+
+def mint_cap_from_oauth(
+    tenant: str,
+    provider: str,
+    *,
+    commands: list[str] | None = None,
+    ttl_seconds: int = 3600,
+) -> str:
+    """After OAuth login, mint GAX capability JWT scoped to tenant + provider scopes."""
+    from gax.caps import mint_capability
+
+    tokens = load_tokens(tenant, provider)
+    if not tokens:
+        raise ValueError(f"No tokens for tenant={tenant} provider={provider}; run gax auth login first")
+    scopes = tokens.get("scopes") or tokens.get("scope") or []
+    if isinstance(scopes, str):
+        scopes = [s for s in scopes.replace(",", " ").split() if s]
+    return mint_capability(
+        tenant_id=tenant,
+        subject=tokens.get("subject") or f"oauth:{provider}",
+        commands=commands or ["*"],
+        scopes=list(scopes) if scopes else ["*"],
+        ttl_seconds=ttl_seconds,
+    )