gax-cli 0.4.0__py3-none-any.whl

gax/adapters/mcp_bridge.py

@@ -0,0 +1,108 @@
+"""MCP → GAX adapter: one registered command → one MCP tool call."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from gax.mcp_client import McpStdioClient, github_mcp_env
+from gax.registry import CommandManifest
+
+
+def _mcp_cfg(manifest: CommandManifest) -> dict[str, Any]:
+    return dict((manifest.raw or {}).get("mcp") or {})
+
+
+def _server_cmd(cfg: dict[str, Any]) -> list[str]:
+    cmd = cfg.get("server_command", "npx")
+    args = list(cfg.get("server_args") or ["-y", "@modelcontextprotocol/server-github"])
+    return [str(cmd), *[str(a) for a in args]]
+
+
+def _map_args(manifest: CommandManifest, args: dict[str, Any]) -> dict[str, Any]:
+    out = dict(args)
+    repo = args.get("repo")
+    if isinstance(repo, str) and "/" in repo:
+        owner, name = repo.split("/", 1)
+        out["owner"] = owner
+        out["repo"] = name
+    mapping = (manifest.raw or {}).get("mcp_arg_map") or {}
+    if mapping:
+        mapped: dict[str, Any] = {}
+        for src, dst in mapping.items():
+            if src in out:
+                mapped[str(dst)] = out[src]
+        for k, v in out.items():
+            mapped.setdefault(k, v)
+        out = mapped
+    return out
+
+
+def _normalize_list_pulls(raw: Any) -> dict[str, Any]:
+    """Map GitHub MCP list_pull_requests output to gh.pr.list envelope shape."""
+    rows: list[Any]
+    if isinstance(raw, list):
+        rows = raw
+    elif isinstance(raw, dict):
+        if "items" in raw and isinstance(raw["items"], list):
+            rows = raw["items"]
+        elif "pull_requests" in raw:
+            rows = raw["pull_requests"]
+        else:
+            rows = [raw]
+    else:
+        rows = []
+
+    items = []
+    for row in rows:
+        if not isinstance(row, dict):
+            continue
+        user = row.get("user") or row.get("author") or {}
+        login = user.get("login") if isinstance(user, dict) else str(user or "unknown")
+        items.append(
+            {
+                "number": row.get("number") or row.get("pull_number"),
+                "title": row.get("title", ""),
+                "state": str(row.get("state", "OPEN")).upper(),
+                "url": row.get("html_url") or row.get("url", ""),
+                "author": login,
+                "draft": bool(row.get("draft") or row.get("isDraft")),
+            }
+        )
+    return {"items": items}
+
+
+def _normalize(manifest: CommandManifest, raw: Any) -> dict[str, Any]:
+    if manifest.command in ("mcp.github.list_pulls", "mcp.github.list_pulls.mock"):
+        if isinstance(raw, dict) and isinstance(raw.get("result"), list):
+            raw = raw["result"]
+        return _normalize_list_pulls(raw)
+    if isinstance(raw, dict):
+        return raw
+    return {"result": raw}
+
+
+def run(
+    manifest: CommandManifest,
+    args: dict[str, Any],
+    *,
+    tenant_id: str | None = None,
+) -> dict[str, Any]:
+    cfg = _mcp_cfg(manifest)
+    tool_name = str(cfg.get("tool_name", ""))
+    if not tool_name:
+        raise RuntimeError(f"mcp.tool_name missing in manifest {manifest.command}")
+
+    env = dict(cfg.get("env") or {})
+    cmd_line = " ".join(_server_cmd(cfg))
+    needs_github = cfg.get("require_github_token")
+    if needs_github is None:
+        needs_github = "server-github" in cmd_line or "@modelcontextprotocol/server-github" in cmd_line
+    if needs_github:
+        env.update(github_mcp_env())
+
+    client = McpStdioClient(_server_cmd(cfg), env=env, timeout=float(cfg.get("timeout", 120)))
+    try:
+        raw = client.call_tool(tool_name, _map_args(manifest, args))
+        return _normalize(manifest, raw)
+    finally:
+        client.close()

gax/adapters/mock_adapter.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import Any
+
+from gax.registry import CommandManifest
+
+
+def run(manifest: CommandManifest, args: dict[str, Any]) -> dict[str, Any]:
+    if manifest.command == "demo.echo":
+        return {"echo": args.get("message", "hello")}
+    if manifest.command == "gh.pr.list":
+        repo = args.get("repo", "octocat/Hello-World")
+        return {
+            "items": [
+                {
+                    "number": 1,
+                    "title": f"[mock] Sample PR for {repo}",
+                    "state": "OPEN",
+                    "url": f"https://github.com/{repo}/pull/1",
+                    "author": "mock-user",
+                    "draft": False,
+                }
+            ],
+            "_mock": True,
+        }
+    if manifest.command == "gh.pr.view":
+        return {
+            "number": args.get("number", 1),
+            "title": "[mock] PR title",
+            "body": "Mock body",
+            "state": "OPEN",
+            "url": f"https://github.com/{args.get('repo', 'o/r')}/pull/{args.get('number', 1)}",
+            "author": "mock-user",
+            "_mock": True,
+        }
+    if manifest.command == "kubectl.get.pods":
+        ns = args.get("namespace", "default")
+        return {
+            "items": [{"name": "mock-pod-1", "namespace": ns, "status": "Running"}],
+            "_mock": True,
+        }
+    if manifest.command == "aws.s3.list":
+        return {"buckets": [{"name": "mock-bucket-a"}, {"name": "mock-bucket-b"}], "_mock": True}
+    if manifest.command == "jira.issue.get":
+        key = args.get("issue_key", "DEMO-1")
+        return {"key": key, "summary": "[mock] Issue", "status": "Open", "_mock": True}
+    if manifest.command.startswith("api."):
+        return {"_mock": True, "command": manifest.command, "args": args}
+    return {"result": "ok", "command": manifest.command, "args": args}

gax/audit.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from typing import Any
+
+from gax.paths import AUDIT_PATH, ensure_gax_home
+
+
+def log_event(
+    *,
+    audit_id: str,
+    tenant_id: str,
+    subject: str,
+    command: str,
+    args: dict[str, Any],
+    ok: bool,
+    error_kind: str | None = None,
+    duration_ms: float | None = None,
+) -> None:
+    ensure_gax_home()
+    # Never log raw secrets; args should already be sanitized by caller
+    safe_args = {k: v for k, v in args.items() if k.lower() not in ("token", "secret", "password")}
+    from gax.spiffe import attest_metadata
+
+    record = {
+        "ts": datetime.now(timezone.utc).isoformat(),
+        "audit_id": audit_id,
+        "tenant_id": tenant_id,
+        "subject": subject,
+        "command": command,
+        "args": safe_args,
+        "ok": ok,
+        "error_kind": error_kind,
+        "duration_ms": duration_ms,
+        **attest_metadata(),
+    }
+    with AUDIT_PATH.open("a") as f:
+        f.write(json.dumps(record) + "\n")
+
+    try:
+        from gax.otel_audit import export_audit_event
+
+        export_audit_event(
+            audit_id=audit_id,
+            tenant_id=tenant_id,
+            subject=subject,
+            command=command,
+            ok=ok,
+            duration_ms=duration_ms,
+            error_kind=error_kind,
+        )
+    except Exception:
+        pass

gax/caps.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import json
+import os
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+import jwt
+
+from gax.paths import CONFIG_PATH, ensure_gax_home
+from gax.side_effects import DEFAULT_CEILING
+from gax.side_effects import normalize as normalize_side_effect
+
+DEFAULT_SECRET = "gax-dev-secret-change-in-production"
+DEFAULT_TENANT = "default"
+
+
+def _load_config() -> dict[str, Any]:
+    ensure_gax_home()
+    if CONFIG_PATH.exists():
+        return json.loads(CONFIG_PATH.read_text())
+    cfg = {
+        "jwt_secret": DEFAULT_SECRET,
+        "tenant_id": DEFAULT_TENANT,
+        "subject": "dev@local",
+    }
+    CONFIG_PATH.write_text(json.dumps(cfg, indent=2))
+    return cfg
+
+
+def jwt_secret() -> str:
+    return _load_config().get("jwt_secret", DEFAULT_SECRET)
+
+
+def mint_capability(
+    *,
+    tenant_id: str | None = None,
+    subject: str | None = None,
+    commands: list[str] | None = None,
+    scopes: list[str] | None = None,
+    ttl_seconds: int = 3600,
+    budget: dict[str, int] | None = None,
+    max_side_effect: str | None = None,
+) -> str:
+    """
+    Mint a capability token.
+
+    `max_side_effect` is the danger ceiling (`read` | `write` | `destructive`).
+    It defaults to `read`, so a token minted without thinking about danger cannot
+    invoke a write or destructive command even if that command is in `commands`.
+    """
+    cfg = _load_config()
+    now = datetime.now(timezone.utc)
+    payload: dict[str, Any] = {
+        "sub": subject or cfg.get("subject", "dev@local"),
+        "tenant_id": tenant_id or cfg.get("tenant_id", DEFAULT_TENANT),
+        "iat": now,
+        "exp": now + timedelta(seconds=ttl_seconds),
+        "scopes": scopes or ["*"],
+        "commands": commands or ["*"],
+        "budget": budget or {"max_calls": 1000, "max_rows": 10_000},
+        "max_side_effect": normalize_side_effect(max_side_effect or DEFAULT_CEILING),
+    }
+    return jwt.encode(payload, jwt_secret(), algorithm="HS256")
+
+
+def decode_capability(token: str | None) -> dict[str, Any]:
+    if not token:
+        token = os.environ.get("GAX_CAP", "").strip()
+    if not token:
+        raise ValueError("missing capability token (set GAX_CAP or pass GAX-Capability header)")
+    # JWT has three dot-separated segments; macaroon is a single base64 blob
+    if token.count(".") == 2:
+        return jwt.decode(token, jwt_secret(), algorithms=["HS256"])
+    from gax.macaroons_cap import verify_macaroon
+
+    return verify_macaroon(token)
+
+
+def cap_allows_command(claims: dict[str, Any], command: str) -> bool:
+    allowed = claims.get("commands") or []
+    if "*" in allowed:
+        return True
+    return command in allowed
+
+
+def cap_allows_scopes(claims: dict[str, Any], required: list[str]) -> bool:
+    held = set(claims.get("scopes") or [])
+    if "*" in held:
+        return True
+    return set(required).issubset(held)