machinaos 0.0.76 → 0.0.78

package/server/services/cli_agent/mcp_server.py

@@ -0,0 +1,482 @@
+"""MachinaOs MCP server (VSCode pattern, no custom IPC).
+
+Hosts a `FastMCP` ASGI sub-app at ``/mcp/ide`` that spawned CLI sessions
+auto-discover via the lockfile written by ``lockfile.py``. Each session
+gets a per-batch bearer token; the middleware validates it and binds the
+matching ``BatchContext`` into a contextvar so tool implementations can
+scope to the calling session's workspace_dir / connected_skill_names /
+allowed_credentials without explicit plumbing.
+
+Tools (5 in v1, mirroring Claude Code's progressive-disclosure pattern):
+  - ``getWorkspaceFiles`` — glob/read inside the session's worktree
+  - ``listSkills`` — metadata for skills connected to the parent agent
+  - ``getSkill`` — full skill markdown + scripts + references
+  - ``getCredential`` — gated by per-batch allowlist
+  - ``broadcastLog`` — write to MachinaOs Terminal tab
+
+The server module exposes:
+  - :func:`get_mcp_app` — Starlette/ASGI sub-app for ``app.mount(...)``
+  - :func:`register_batch` / :func:`unregister_batch` — `AICliService`
+    calls these around `run_batch()` to register/expire tokens
+  - :class:`BatchContext` — scoping data attached to each token
+
+Tools deferred to v2: ``getDiagnostics``, ``executeCode``.
+"""
+
+from __future__ import annotations
+
+import contextvars
+import logging
+import secrets
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Set
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Mount
+
+from core.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# BatchContext + token registry
+# ---------------------------------------------------------------------------
+
+@dataclass
+class BatchContext:
+    """Scoping data attached to one batch's bearer token.
+
+    Populated at ``AICliService.run_batch()`` entry; deregistered in the
+    ``finally`` block. Tools dereference the calling batch via the
+    bearer token in `Authorization` header.
+    """
+    workflow_id: str
+    node_id: str
+    workspace_dir: Path
+    connected_skill_names: Set[str] = field(default_factory=set)
+    allowed_credentials: Set[str] = field(default_factory=set)
+    # Connected ``input-tools`` nodes (entries from
+    # ``services.plugin.edge_walker.collect_agent_connections``). Each
+    # dict: ``{node_id, node_type, label, parameters, ...}``. Drives the
+    # ``listMachinaOsTools`` / ``callMachinaOsTool`` MCP tools so the
+    # CLI agent sees the same tool surface the AI Agent does.
+    connected_tools: List[Dict[str, Any]] = field(default_factory=list)
+    # Optional broadcaster for `broadcastLog`. Lazily resolved from the
+    # global container if None.
+    broadcaster: Optional[Any] = None
+
+
+# Token -> BatchContext registry. Lives in-memory only; tokens never
+# touch disk or the credentials.db.
+_active_tokens: Dict[str, BatchContext] = {}
+
+
+def issue_token() -> str:
+    """Mint a new bearer token (32 bytes hex)."""
+    return secrets.token_hex(32)
+
+
+def register_batch(token: str, ctx: BatchContext) -> None:
+    """Register a batch's auth token + expose its connected workflow
+    tools on the FastMCP server. Idempotent on identical context."""
+    if token in _active_tokens:
+        existing = _active_tokens[token]
+        if existing is not ctx:
+            raise ValueError("Token collision in MCP server batch registry")
+        return
+    _active_tokens[token] = ctx
+    tool_names = [t.get("node_type") for t in ctx.connected_tools]
+    logger.info(
+        "[CC-Agent MCP register_batch] node=%s wf=%s token=%s... "
+        "skills=%d tools=%d creds=%d %s",
+        ctx.node_id, ctx.workflow_id, token[:8],
+        len(ctx.connected_skill_names), len(ctx.connected_tools),
+        len(ctx.allowed_credentials),
+        f"tools={tool_names}" if tool_names else "(no tools wired)",
+    )
+    # Per-batch workflow-tool exposure lives in workflow_tools.py.
+    from services.cli_agent.workflow_tools import expose_workflow_tools
+    expose_workflow_tools(ctx.connected_tools)
+
+
+def unregister_batch(token: str) -> None:
+    """Drop a batch's token + un-expose its tools when the refcount
+    hits zero. Safe to call twice."""
+    ctx = _active_tokens.pop(token, None)
+    if ctx is None:
+        return
+    logger.debug("[CC-Agent MCP] unregistered batch token=%s...", token[:8])
+    from services.cli_agent.workflow_tools import unexpose_workflow_tools
+    unexpose_workflow_tools(ctx.connected_tools)
+
+
+def lookup_batch(token: str) -> Optional[BatchContext]:
+    return _active_tokens.get(token)
+
+
+def active_batch_count() -> int:
+    return len(_active_tokens)
+
+
+# ---------------------------------------------------------------------------
+# ContextVar — thread/task-local handle to the current batch
+# ---------------------------------------------------------------------------
+
+_current_batch: contextvars.ContextVar[Optional[BatchContext]] = (
+    contextvars.ContextVar("machina_current_batch", default=None)
+)
+
+
+def _require_batch() -> BatchContext:
+    ctx = _current_batch.get()
+    if ctx is None:
+        raise RuntimeError(
+            "MCP tool called without an active batch context. "
+            "This indicates the auth middleware was bypassed."
+        )
+    return ctx
+
+
+# Per-batch workflow-tool exposure happens in `_expose_workflow_tools`
+# (above). Each connected node lands as its own
+# `mcp__machinaos__<node_type>` entry; FastMCP infers the inputSchema
+# from the typed `params` annotation. The legacy generic wrapper has
+# been removed.
+
+
+# ---------------------------------------------------------------------------
+# Auth middleware
+# ---------------------------------------------------------------------------
+
+class _BearerAuthMiddleware(BaseHTTPMiddleware):
+    """Validate `Authorization: Bearer <token>` against the registry.
+
+    On success: bind the matching `BatchContext` into the contextvar.
+    On failure: 401.
+
+    The MCP spec uses MCP-Protocol-Version + Authorization headers; we
+    care only about the Bearer here. Health checks under `/healthz`
+    bypass auth (they're for dev sanity).
+    """
+
+    async def dispatch(  # type: ignore[override]
+        self,
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        path = request.url.path
+        method = request.method
+        if path.endswith("/healthz"):
+            return await call_next(request)
+
+        auth = request.headers.get("authorization") or request.headers.get(
+            "Authorization"
+        )
+        token: Optional[str] = None
+        if auth and auth.lower().startswith("bearer "):
+            token = auth[7:].strip() or None
+
+        if not token:
+            ua = request.headers.get("user-agent", "")
+            logger.warning(
+                "[CC-Agent MCP auth] %s %s -> 401 (no Bearer token; "
+                "auth_header=%r ua=%r) — claude CLI either didn't read the "
+                "lockfile or is hitting the wrong URL.",
+                method, path, auth, ua,
+            )
+            return JSONResponse(
+                {"error": "missing or malformed Authorization header"},
+                status_code=401,
+            )
+
+        ctx = lookup_batch(token)
+        if ctx is None:
+            logger.warning(
+                "[CC-Agent MCP auth] %s %s -> 401 (token=%s... not in "
+                "active registry — batch may have ended)",
+                method, path, token[:8],
+            )
+            return JSONResponse(
+                {"error": "invalid or expired token"},
+                status_code=401,
+            )
+
+        logger.info(
+            "[CC-Agent MCP auth] %s %s -> OK (node=%s wf=%s token=%s...)",
+            method, path, ctx.node_id, ctx.workflow_id, token[:8],
+        )
+
+        reset_token = _current_batch.set(ctx)
+        try:
+            return await call_next(request)
+        finally:
+            _current_batch.reset(reset_token)
+
+
+# ---------------------------------------------------------------------------
+# Tool registration helper — defers FastMCP import so module import is cheap
+# ---------------------------------------------------------------------------
+
+def _build_tools(mcp: Any) -> None:  # FastMCP type
+    """Register the 5 v1 tools on a `FastMCP` instance."""
+
+    @mcp.tool(
+        name="getWorkspaceFiles",
+        description=(
+            "List or read files inside the calling session's per-task git "
+            "worktree. Use `read=False` for metadata-only listings; "
+            "`read=True` to fetch file contents (capped at 1MB per file)."
+        ),
+    )
+    def get_workspace_files(
+        path: str = ".",
+        pattern: str = "*",
+        read: bool = False,
+        max_bytes: int = 1_000_000,
+    ) -> Dict[str, Any]:
+        ctx = _require_batch()
+        try:
+            base = ctx.workspace_dir.resolve()
+            target = (base / path).resolve()
+            # Path-traversal guard
+            try:
+                target.relative_to(base)
+            except ValueError:
+                return {
+                    "error": "path escapes workspace_dir",
+                    "path": str(path),
+                }
+
+            if not target.exists():
+                return {"files": [], "path": str(path)}
+
+            entries: List[Dict[str, Any]] = []
+            if target.is_file():
+                files_iter = [target]
+            else:
+                files_iter = sorted(target.rglob(pattern))
+
+            for p in files_iter:
+                if not p.is_file():
+                    continue
+                try:
+                    rel = str(p.relative_to(base))
+                    info: Dict[str, Any] = {
+                        "path": rel,
+                        "size": p.stat().st_size,
+                        "mtime": p.stat().st_mtime,
+                    }
+                    if read and p.stat().st_size <= max_bytes:
+                        try:
+                            info["content"] = p.read_text(
+                                encoding="utf-8", errors="replace"
+                            )
+                        except OSError as exc:
+                            info["read_error"] = str(exc)
+                    entries.append(info)
+                except OSError:
+                    continue
+
+                if len(entries) >= 1000:
+                    break
+
+            return {"files": entries, "path": str(path)}
+        except Exception as exc:  # pragma: no cover — defensive
+            logger.exception("[CC-Agent MCP] getWorkspaceFiles failed")
+            return {"error": str(exc), "path": str(path)}
+
+    @mcp.tool(
+        name="listSkills",
+        description=(
+            "List skills connected to the parent agent node. Returns "
+            "metadata only (~100 tokens per skill). Call `getSkill(name)` "
+            "to fetch the full instructions."
+        ),
+    )
+    def list_skills() -> Dict[str, Any]:
+        ctx = _require_batch()
+        try:
+            from services.skill_loader import get_skill_loader
+            loader = get_skill_loader()
+            registry = loader.scan_skills()
+            results = []
+            for name in sorted(ctx.connected_skill_names):
+                meta = registry.get(name)
+                if meta is None:
+                    continue
+                results.append({
+                    "name": meta.name,
+                    "description": meta.description,
+                    "allowed_tools": list(meta.allowed_tools),
+                    "category": (
+                        meta.metadata.get("category")
+                        if isinstance(meta.metadata, dict) else None
+                    ),
+                })
+            return {"skills": results}
+        except Exception as exc:  # pragma: no cover
+            logger.exception("[CC-Agent MCP] listSkills failed")
+            return {"error": str(exc), "skills": []}
+
+    @mcp.tool(
+        name="getSkill",
+        description=(
+            "Fetch full content for one skill: instructions (markdown), "
+            "scripts (executable code samples), and references (extra "
+            "docs). The skill must be connected to the parent agent node."
+        ),
+    )
+    def get_skill(name: str) -> Dict[str, Any]:
+        ctx = _require_batch()
+        if name not in ctx.connected_skill_names:
+            return {
+                "error": f"skill {name!r} is not connected to this agent node",
+                "name": name,
+            }
+        try:
+            from services.skill_loader import get_skill_loader
+            loader = get_skill_loader()
+            skill = loader.load_skill(name)
+            if skill is None:
+                return {"error": f"skill {name!r} not found", "name": name}
+            return {
+                "name": skill.metadata.name,
+                "description": skill.metadata.description,
+                "instructions": skill.instructions,
+                "allowed_tools": list(skill.metadata.allowed_tools),
+                "metadata": dict(skill.metadata.metadata) if skill.metadata.metadata else {},
+                "scripts": dict(skill.scripts),
+                "references": dict(skill.references),
+                # `assets` (binary) excluded by default — too big for MCP responses.
+            }
+        except Exception as exc:  # pragma: no cover
+            logger.exception("[CC-Agent MCP] getSkill failed for %r", name)
+            return {"error": str(exc), "name": name}
+
+    @mcp.tool(
+        name="getCredential",
+        description=(
+            "Fetch a credential by provider name. Only credentials in the "
+            "batch's allowlist are returned; everything else returns 403. "
+            "Use sparingly — prefer the CLI's own auth where possible."
+        ),
+    )
+    async def get_credential(name: str) -> Dict[str, Any]:
+        ctx = _require_batch()
+        if name not in ctx.allowed_credentials:
+            return {
+                "error": f"credential {name!r} not in allowlist for this batch",
+                "name": name,
+                "status": 403,
+            }
+        try:
+            from core.container import container
+            auth = container.auth_service()
+            value = await auth.get_api_key(name)
+            if not value:
+                return {
+                    "error": f"credential {name!r} not configured",
+                    "name": name,
+                    "status": 404,
+                }
+            return {"name": name, "value": value}
+        except Exception as exc:  # pragma: no cover
+            logger.exception("[CC-Agent MCP] getCredential failed for %r", name)
+            return {"error": str(exc), "name": name, "status": 500}
+
+    # Per-batch workflow tools are exposed dynamically on
+    # `register_batch` via `_expose_workflow_tools` — each connected
+    # node lands as its own `mcp__machinaos__<node_type>` entry. No
+    # generic `listMachinaOsTools` / `callMachinaOsTool` wrapper.
+
+    @mcp.tool(
+        name="broadcastLog",
+        description=(
+            "Write a log line to the MachinaOs Terminal tab. Use to "
+            "surface intermediate progress that would otherwise be lost "
+            "between CLI sessions."
+        ),
+    )
+    async def broadcast_log(
+        message: str,
+        level: str = "info",
+        source: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        ctx = _require_batch()
+        if level not in ("debug", "info", "warning", "error"):
+            level = "info"
+        try:
+            broadcaster = ctx.broadcaster
+            if broadcaster is None:
+                from services.status_broadcaster import get_status_broadcaster
+                broadcaster = get_status_broadcaster()
+            payload = {
+                "source": source or f"mcp:{ctx.node_id}",
+                "level": level,
+                "message": message[:5000],
+            }
+            await broadcaster.broadcast_terminal_log(payload)
+            return {"success": True}
+        except Exception as exc:  # pragma: no cover
+            logger.exception("[CC-Agent MCP] broadcastLog failed")
+            return {"error": str(exc)}
+
+
+# ---------------------------------------------------------------------------
+# ASGI sub-app factory (mounted in main.py lifespan)
+# ---------------------------------------------------------------------------
+
+_app_singleton: Optional[Any] = None  # Starlette app
+_mcp_singleton: Optional[Any] = None  # FastMCP instance — used by per-batch
+                                       # dynamic tool (un)registration.
+
+
+def get_mcp_app() -> Any:
+    """Return the Starlette/ASGI app to mount under `/mcp/ide`.
+
+    Idempotent — multiple calls return the same instance, so the FastAPI
+    lifespan can wire it without worrying about duplicate registration.
+    """
+    global _app_singleton, _mcp_singleton
+    if _app_singleton is not None:
+        return _app_singleton
+
+    from mcp.server.fastmcp import FastMCP
+
+    mcp = FastMCP(
+        name="machinaos-cli-agent",
+        instructions=(
+            "MachinaOs IDE MCP server. Exposes workspace files, connected "
+            "skills, scoped credentials, and a Terminal-tab log channel "
+            "to the calling CLI session."
+        ),
+        # FastMCP's HTTP path defaults are fine; we mount the whole app
+        # under /mcp/ide externally.
+    )
+    _build_tools(mcp)
+    _mcp_singleton = mcp
+
+    asgi_app = mcp.streamable_http_app()
+    asgi_app.add_middleware(_BearerAuthMiddleware)
+
+    _app_singleton = asgi_app
+    return _app_singleton
+
+
+# ---------------------------------------------------------------------------
+# Test/diagnostic helpers (used by tests + verification step #11/#12)
+# ---------------------------------------------------------------------------
+
+def _reset_for_tests() -> None:  # pragma: no cover
+    """Wipe the token registry + per-batch workflow-tool refcounts.
+    ONLY use in tests."""
+    global _app_singleton, _mcp_singleton
+    _active_tokens.clear()
+    _app_singleton = None
+    _mcp_singleton = None
+    from services.cli_agent.workflow_tools import _reset_for_tests as _wt_reset
+    _wt_reset()

package/server/services/cli_agent/protocol.py

@@ -0,0 +1,173 @@
+"""Protocol for AI CLI providers (Claude Code, Codex, Gemini).
+
+Mirrors `services/llm/protocol.py` shape: a structurally-typed Protocol
+with a small, explicit surface. Each concrete provider lives under
+`providers/<vendor>.py`.
+
+The framework spawns N parallel CLI sessions (one `AICliSession` each,
+backed by `BaseProcessSupervisor`) over a list of tasks. Per-CLI
+differences (argv, JSON event schema, auth handling, feature support)
+are isolated to the provider; the session/pool/service layer is generic.
+
+Auth: native CLI handles its own tokens. We only trigger the login flow
+(`login_argv()`) and detect logged-in state (`auth_status_argv()` +
+`detect_auth_error()`). No credential wrapping, no `~/.claude-machina/`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+
+
+# ---------------------------------------------------------------------------
+# Shared data types
+# ---------------------------------------------------------------------------
+
+@dataclass
+class CanonicalUsage:
+    """Vendor-normalised token counts.
+
+    Pattern from Hermes `agent/usage_pricing.py:CanonicalUsage` — every
+    provider's usage shape (Anthropic vs Codex vs OpenAI handle cache
+    differently) maps into this so the existing `services/pricing.py`
+    can compute USD without per-vendor branches.
+    """
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read: int = 0
+    cache_write: int = 0
+    reasoning_tokens: int = 0
+    request_count: int = 0
+
+
+@dataclass
+class SessionResult:
+    """Per-task result returned by `AICliSession`.
+
+    Shared keys are the schema; vendor extras live in `provider_data`.
+    """
+    task_id: str
+    session_id: Optional[str] = None
+    provider: str = ""
+    prompt: str = ""
+    branch: Optional[str] = None
+    worktree_path: Optional[str] = None
+    response: str = ""
+    cost_usd: Optional[float] = None
+    duration_ms: Optional[int] = None
+    num_turns: Optional[int] = None
+    tool_calls: int = 0
+    canonical_usage: CanonicalUsage = field(default_factory=CanonicalUsage)
+    provider_data: Dict[str, Any] = field(default_factory=dict)
+    success: bool = False
+    error: Optional[str] = None
+
+
+@dataclass
+class BatchResult:
+    """Aggregated result returned by `AICliService.run_batch()`."""
+    tasks: List[SessionResult] = field(default_factory=list)
+    n_tasks: int = 0
+    n_succeeded: int = 0
+    n_failed: int = 0
+    total_cost_usd: Optional[float] = None
+    wall_clock_ms: int = 0
+    budget_remaining_usd: Optional[float] = None
+    provider: str = ""
+    timestamp: str = ""
+
+
+# ---------------------------------------------------------------------------
+# Provider protocol (structural typing)
+# ---------------------------------------------------------------------------
+
+@runtime_checkable
+class AICliProvider(Protocol):
+    """Structural Protocol for an AI CLI provider.
+
+    Concrete classes:
+      - `providers.anthropic_claude.AnthropicClaudeProvider`
+      - `providers.openai_codex.OpenAICodexProvider`
+      - `providers.google_gemini.GoogleGeminiProvider` (v2 stub)
+    """
+
+    name: str                          # "claude" | "codex" | "gemini"
+    package_name: str                  # npm package
+    binary_name: str                   # "claude" | "codex" | "gemini"
+    ide_lock_env_var: Optional[str]    # CLAUDE_IDE_LOCK | GEMINI_IDE_LOCK | None
+    ide_lockfile_dir: Optional[Path]   # ~/.claude/ide | <tmpdir>/gemini/ide
+
+    # ---- spawn surface ---------------------------------------------------
+
+    def binary_path(self) -> Path: ...
+        # Resolve the CLI binary. Resolution chain (Composio pattern):
+        #   1) shutil.which(<binary_name>)
+        #   2) `npx --yes <package_name>` shim path
+        # Raises FileNotFoundError if neither is available.
+
+    def headless_argv(self, task: Any, *, defaults: Dict[str, Any]) -> List[str]: ...
+        # Build the full argv (binary + flags) for a headless run of
+        # `task` (a `<Provider>TaskSpec` Pydantic model). `defaults`
+        # comes from `ai_cli_providers.json` for this provider.
+
+    # ---- native auth (no token wrapping) --------------------------------
+
+    def login_argv(self) -> List[str]: ...
+        # CLI's own login command, e.g. ["claude", "login"]. Spawned
+        # interactively from the Credentials Modal. CLI stores its own
+        # credentials in `~/.claude/`, `~/.codex/`, `~/.gemini/`.
+
+    def auth_status_argv(self) -> Optional[List[str]]: ...
+        # No-op invocation to verify auth, e.g. ["claude", "--print", "-p", "ok"].
+        # Returns None if no cheap probe exists; in that case the framework
+        # infers from the first session's stderr.
+
+    def detect_auth_error(self, stderr: str, exit_code: int) -> bool: ...
+        # Match "not logged in" patterns:
+        #   Claude: "Please run 'claude login'"
+        #   Codex:  HTTP 401 / "OPENAI_API_KEY not set"
+        #   Gemini: exit_code == 1 (FatalAuthenticationError)
+
+    # ---- streaming output parsing ---------------------------------------
+
+    def parse_event(self, line: str) -> Optional[Dict[str, Any]]: ...
+        # Parse a single NDJSON line from stdout. Return None for
+        # un-parseable garbage.
+
+    def is_final_event(self, event: Dict[str, Any]) -> bool: ...
+        # True if this event marks end-of-task. For Claude: `type=="result"`.
+        # For Gemini: `type=="result"`. For Codex: `type=="complete"` or
+        # heuristic fallback.
+
+    def event_to_session_result(
+        self,
+        events: List[Dict[str, Any]],
+        stderr: str,
+        exit_code: int,
+    ) -> Dict[str, Any]: ...
+        # Reconstruct a partial dict of `SessionResult` fields from the
+        # event stream. Returns:
+        #   {
+        #     "session_id": ..., "response": ..., "cost_usd": ...,
+        #     "duration_ms": ..., "num_turns": ..., "tool_calls": ...,
+        #     "success": bool, "error": Optional[str],
+        #     "canonical_usage": CanonicalUsage,
+        #     "provider_data": {<vendor-specific>},
+        #   }
+        # provider_data carries vendor-only metadata (Anthropic
+        # reasoning_details, Codex call_id, Gemini extra_content) without
+        # bloating the shared schema. Pattern from
+        # Hermes agent/transports/types.py NormalizedResponse.
+
+    def canonical_usage(self, events: List[Dict[str, Any]]) -> CanonicalUsage: ...
+        # Normalise vendor token-counting into the shared `CanonicalUsage`
+        # shape. Pattern from Hermes agent/usage_pricing.py.
+
+    # ---- feature gating --------------------------------------------------
+
+    def supports(self, feature: str) -> bool: ...
+        # Feature flags consulted by the session/service layer.
+        # Recognised: "max_budget", "max_turns", "session_id", "resume",
+        # "mcp_runtime", "json_cost", "ide_lockfile", "sandbox".

package/server/services/cli_agent/providers/__init__.py

@@ -0,0 +1,9 @@
+"""Concrete `AICliProvider` implementations.
+
+- `anthropic_claude.AnthropicClaudeProvider` — full v1 surface
+- `openai_codex.OpenAICodexProvider` — sandbox-first, no session
+- `google_gemini.GoogleGeminiProvider` — v2 stub raising NotImplementedError
+
+Imports are intentionally not eager — `factory.create_cli_provider()`
+lazy-imports each one.
+"""