milo-cli 0.1.0__py3-none-any.whl

milo/mcp.py

@@ -0,0 +1,299 @@
+"""MCP server — expose CLI commands as tools via JSON-RPC on stdin/stdout."""
+
+from __future__ import annotations
+
+import json
+import sys
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from milo.commands import CLI, CommandDef, LazyCommandDef
+
+_MCP_VERSION = "2025-11-25"
+_SERVER_NAME = "milo"
+_SERVER_VERSION = "0.1.0"
+
+
+def run_mcp_server(cli: CLI) -> None:
+    """Run MCP JSON-RPC server on stdin/stdout.
+
+    Implements the MCP protocol (initialize, tools/list, tools/call,
+    resources/list, resources/read, prompts/list, prompts/get).
+    """
+    tools = _list_tools(cli)
+    tool_names = [t["name"] for t in tools]
+
+    _stderr(f"MCP server ready — {cli.name}")
+    _stderr(f"  Protocol:  {_MCP_VERSION}")
+    _stderr(f"  Tools:     {len(tools)} ({', '.join(tool_names)})")
+    _stderr(f"  Resources: {len(cli._resources)}")
+    _stderr(f"  Prompts:   {len(cli._prompts)}")
+    _stderr("  Transport: stdin/stdout (JSON-RPC, one request per line)")
+    _stderr("")
+    _stderr("Send requests as JSON, for example:")
+    _stderr('  {"jsonrpc":"2.0","id":1,"method":"initialize"}')
+    _stderr('  {"jsonrpc":"2.0","id":2,"method":"tools/list"}')
+    if tool_names:
+        example_tool = tool_names[0]
+        example_schema = tools[0].get("inputSchema", {})
+        example_args = dict.fromkeys(example_schema.get("properties", {}), "...")
+        args_json = json.dumps(example_args) if example_args else "{}"
+        _stderr(
+            f'  {{"jsonrpc":"2.0","id":3,"method":"tools/call",'
+            f'"params":{{"name":"{example_tool}","arguments":{args_json}}}}}'
+        )
+    _stderr("")
+    _stderr("Or pipe from a file:")
+    _stderr(f"  cat requests.jsonl | {cli.name} --mcp")
+    _stderr("")
+    _stderr("Press Ctrl+C to stop.")
+    _stderr("")
+
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            request = json.loads(line)
+        except json.JSONDecodeError:
+            _write_error(None, -32700, "Parse error")
+            continue
+
+        req_id = request.get("id")
+        method = request.get("method", "")
+
+        try:
+            result = _handle_method(cli, method, request.get("params", {}))
+            if result is not None:
+                _write_result(req_id, result)
+        except Exception as e:
+            _write_error(req_id, -32603, str(e))
+
+
+def _handle_method(cli: CLI, method: str, params: dict[str, Any]) -> dict[str, Any] | None:
+    """Dispatch an MCP method."""
+    match method:
+        case "initialize":
+            return {
+                "protocolVersion": _MCP_VERSION,
+                "capabilities": {"tools": {}, "resources": {}, "prompts": {}},
+                "serverInfo": {
+                    "name": cli.name or _SERVER_NAME,
+                    "version": cli.version or _SERVER_VERSION,
+                    "title": cli.description,
+                },
+                "instructions": cli.description,
+            }
+        case "notifications/initialized":
+            # Client confirms initialization — no response required
+            return None
+        case "tools/list":
+            return {"tools": _list_tools(cli)}
+        case "tools/call":
+            return _call_tool(cli, params)
+        case "resources/list":
+            return {"resources": _list_resources(cli)}
+        case "resources/read":
+            return _read_resource(cli, params)
+        case "prompts/list":
+            return {"prompts": _list_prompts(cli)}
+        case "prompts/get":
+            return _get_prompt(cli, params)
+        case _:
+            raise ValueError(f"Unknown method: {method!r}")
+
+
+def _list_tools(cli: CLI) -> list[dict[str, Any]]:
+    """Generate MCP tools/list response from all commands including groups.
+
+    Group commands use dot-notation names: ``site.build``, ``site.config.show``.
+    Includes outputSchema when return type annotations are available.
+    """
+    tools = []
+    for dotted_name, cmd in cli.walk_commands():
+        if cmd.hidden:
+            continue
+        tool: dict[str, Any] = {
+            "name": dotted_name,
+            "description": cmd.description,
+            "inputSchema": cmd.schema,
+        }
+
+        # title: human-readable display name from docstring or description
+        title = _tool_title(cmd)
+        if title:
+            tool["title"] = title
+
+        # outputSchema: generated from handler return type annotation
+        output_schema = _output_schema(cmd)
+        if output_schema:
+            tool["outputSchema"] = output_schema
+
+        tools.append(tool)
+    return tools
+
+
+def _tool_title(cmd: CommandDef | LazyCommandDef) -> str:
+    """Derive a human-readable title for a tool.
+
+    Uses the handler's docstring first line if available and different
+    from the description. Falls back to a title-cased command name.
+    """
+    from milo.commands import LazyCommandDef
+
+    # For lazy commands, use name-derived title to avoid triggering imports
+    if isinstance(cmd, LazyCommandDef) and cmd._resolved is None:
+        return cmd.name.replace("-", " ").replace("_", " ").title()
+
+    doc = getattr(cmd.handler, "__doc__", None)
+    if doc:
+        first_line = doc.strip().split("\n")[0].strip().rstrip(".")
+        if first_line and first_line != cmd.description:
+            return first_line
+
+    return cmd.name.replace("-", " ").replace("_", " ").title()
+
+
+def _output_schema(cmd: CommandDef | LazyCommandDef) -> dict[str, Any] | None:
+    """Generate outputSchema from handler return type annotation.
+
+    Returns None for lazy commands that haven't been resolved (avoids imports).
+    """
+    from milo.commands import LazyCommandDef
+    from milo.schema import return_to_schema
+
+    # Don't trigger imports on lazy commands just for output schema
+    if isinstance(cmd, LazyCommandDef) and cmd._resolved is None:
+        return None
+
+    return return_to_schema(cmd.handler)
+
+
+def _call_tool(cli: CLI, params: dict[str, Any]) -> dict[str, Any]:
+    """Handle tools/call by dispatching to the command handler."""
+    tool_name = params.get("name", "")
+    arguments = params.get("arguments", {})
+
+    try:
+        result = cli.call(tool_name, **arguments)
+    except Exception as e:
+        return {
+            "content": [{"type": "text", "text": f"Error: {e}"}],
+            "isError": True,
+        }
+
+    # Convert result to MCP content with structuredContent for non-string results
+    text = result if isinstance(result, str) else json.dumps(result, indent=2, default=str)
+
+    response: dict[str, Any] = {
+        "content": [{"type": "text", "text": text}],
+    }
+
+    # Include structuredContent for structured data (dict, list, number, bool)
+    if not isinstance(result, str) and result is not None:
+        response["structuredContent"] = result
+
+    return response
+
+
+def _list_resources(cli: CLI) -> list[dict[str, Any]]:
+    """Generate MCP resources/list response from registered resources."""
+    resources = []
+    for _uri, res in cli.walk_resources():
+        resources.append(
+            {
+                "uri": res.uri,
+                "name": res.name,
+                "description": res.description,
+                "mimeType": res.mime_type,
+            }
+        )
+    return resources
+
+
+def _read_resource(cli: CLI, params: dict[str, Any]) -> dict[str, Any]:
+    """Handle resources/read by calling the resource handler."""
+    uri = params.get("uri", "")
+
+    res = cli._resources.get(uri)
+    if not res:
+        return {"contents": []}
+
+    try:
+        result = res.handler()
+    except Exception as e:
+        return {"contents": [{"uri": uri, "text": f"Error: {e}", "mimeType": "text/plain"}]}
+
+    text = result if isinstance(result, str) else json.dumps(result, indent=2, default=str)
+
+    return {"contents": [{"uri": uri, "text": text, "mimeType": res.mime_type}]}
+
+
+def _list_prompts(cli: CLI) -> list[dict[str, Any]]:
+    """Generate MCP prompts/list response from registered prompts."""
+    prompts = []
+    for _name, p in cli.walk_prompts():
+        prompt_info: dict[str, Any] = {
+            "name": p.name,
+            "description": p.description,
+        }
+        if p.arguments:
+            prompt_info["arguments"] = list(p.arguments)
+        prompts.append(prompt_info)
+    return prompts
+
+
+def _get_prompt(cli: CLI, params: dict[str, Any]) -> dict[str, Any]:
+    """Handle prompts/get by calling the prompt handler."""
+    name = params.get("name", "")
+    arguments = params.get("arguments", {})
+
+    p = cli._prompts.get(name)
+    if not p:
+        return {"messages": []}
+
+    try:
+        result = p.handler(**arguments)
+    except Exception as e:
+        return {"messages": [{"role": "user", "content": {"type": "text", "text": f"Error: {e}"}}]}
+
+    # If handler returns list of dicts, treat as messages
+    if isinstance(result, list):
+        return {"messages": result}
+
+    # If handler returns a string, wrap as single user message
+    if isinstance(result, str):
+        return {"messages": [{"role": "user", "content": {"type": "text", "text": result}}]}
+
+    return {
+        "messages": [
+            {
+                "role": "user",
+                "content": {"type": "text", "text": json.dumps(result, default=str)},
+            }
+        ]
+    }
+
+
+def _write_result(req_id: Any, result: dict[str, Any]) -> None:
+    """Write a JSON-RPC success response."""
+    response = {"jsonrpc": "2.0", "id": req_id, "result": result}
+    sys.stdout.write(json.dumps(response) + "\n")
+    sys.stdout.flush()
+
+
+def _write_error(req_id: Any, code: int, message: str) -> None:
+    """Write a JSON-RPC error response."""
+    response = {
+        "jsonrpc": "2.0",
+        "id": req_id,
+        "error": {"code": code, "message": message},
+    }
+    sys.stdout.write(json.dumps(response) + "\n")
+    sys.stdout.flush()
+
+
+def _stderr(message: str) -> None:
+    """Write an informational line to stderr."""
+    sys.stderr.write(message + "\n")
+    sys.stderr.flush()

milo/middleware.py

@@ -0,0 +1,67 @@
+"""Middleware stack for MCP and CLI call interception."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+# Type aliases
+MiddlewareFn = Callable[..., Any]
+NextFn = Callable[["MCPCall"], Any]
+
+
+@dataclass(frozen=True, slots=True)
+class MCPCall:
+    """Represents an interceptable MCP or CLI call."""
+
+    method: str  # "tools/call", "resources/read", etc.
+    name: str  # tool/resource/prompt name
+    arguments: dict[str, Any] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class MiddlewareStack:
+    """Ordered middleware pipeline for intercepting calls.
+
+    Middleware signature: ``def mw(ctx, call: MCPCall, next_fn: NextFn) -> Any``
+
+    Usage::
+
+        stack = MiddlewareStack()
+
+        @stack.use
+        def log_calls(ctx, call, next_fn):
+            print(f"Calling {call.name}")
+            result = next_fn(call)
+            print(f"Done {call.name}")
+            return result
+
+        result = stack.execute(ctx, call, handler)
+    """
+
+    def __init__(self) -> None:
+        self._middlewares: list[MiddlewareFn] = []
+
+    def use(self, fn: MiddlewareFn) -> MiddlewareFn:
+        """Register a middleware function. Can be used as a decorator."""
+        self._middlewares.append(fn)
+        return fn
+
+    def execute(self, ctx: Any, call: MCPCall, handler: Callable[..., Any]) -> Any:
+        """Execute the middleware chain, ending with the handler."""
+        if not self._middlewares:
+            return handler(call)
+
+        def build_next(index: int) -> NextFn:
+            if index >= len(self._middlewares):
+                return handler
+            mw = self._middlewares[index]
+
+            def next_fn(c: MCPCall) -> Any:
+                return mw(ctx, c, build_next(index + 1))
+
+            return next_fn
+
+        first = self._middlewares[0]
+        return first(ctx, call, build_next(1))

milo/observability.py

@@ -0,0 +1,111 @@
+"""Observability for MCP requests — logging, stats, correlation IDs."""
+
+from __future__ import annotations
+
+import contextvars
+import threading
+import time
+import uuid
+from collections import deque
+from dataclasses import dataclass
+from typing import Any
+
+# Correlation ID for request tracing
+correlation_id: contextvars.ContextVar[str] = contextvars.ContextVar(
+    "milo_correlation_id", default=""
+)
+
+
+def new_correlation_id() -> str:
+    """Generate and set a new correlation ID. Returns the ID."""
+    cid = uuid.uuid4().hex[:12]
+    correlation_id.set(cid)
+    return cid
+
+
+@dataclass(frozen=True, slots=True)
+class RequestLog:
+    """A single MCP request log entry."""
+
+    correlation_id: str
+    method: str
+    name: str
+    latency_ms: float
+    error: str = ""
+    cli_name: str = ""
+    timestamp: float = 0.0
+
+
+class RequestLogger:
+    """Thread-safe ring buffer for MCP request logs with stats computation."""
+
+    def __init__(self, max_size: int = 1000) -> None:
+        self._buffer: deque[RequestLog] = deque(maxlen=max_size)
+        self._lock = threading.Lock()
+        self._total: int = 0
+        self._errors: int = 0
+
+    def record(self, log: RequestLog) -> None:
+        """Record a request log entry."""
+        with self._lock:
+            self._buffer.append(log)
+            self._total += 1
+            if log.error:
+                self._errors += 1
+
+    def recent(self, n: int = 20) -> list[RequestLog]:
+        """Return the N most recent log entries."""
+        with self._lock:
+            items = list(self._buffer)
+        return items[-n:]
+
+    def stats(self) -> dict[str, Any]:
+        """Compute aggregate statistics."""
+        with self._lock:
+            items = list(self._buffer)
+            total = self._total
+            errors = self._errors
+
+        if not items:
+            return {
+                "total": total,
+                "errors": errors,
+                "avg_latency_ms": 0.0,
+                "p99_latency_ms": 0.0,
+            }
+
+        latencies = sorted(log.latency_ms for log in items)
+        avg = sum(latencies) / len(latencies)
+        p99_idx = max(0, int(len(latencies) * 0.99) - 1)
+
+        return {
+            "total": total,
+            "errors": errors,
+            "avg_latency_ms": round(avg, 2),
+            "p99_latency_ms": round(latencies[p99_idx], 2),
+        }
+
+
+def log_request(
+    logger: RequestLogger,
+    method: str,
+    name: str,
+    start_time: float,
+    *,
+    error: str = "",
+    cli_name: str = "",
+) -> RequestLog:
+    """Create and record a request log entry. Returns the log."""
+    elapsed = (time.monotonic() - start_time) * 1000
+    cid = correlation_id.get("")
+    entry = RequestLog(
+        correlation_id=cid,
+        method=method,
+        name=name,
+        latency_ms=round(elapsed, 2),
+        error=error,
+        cli_name=cli_name,
+        timestamp=time.time(),
+    )
+    logger.record(entry)
+    return entry

milo/output.py

@@ -0,0 +1,106 @@
+"""Structured output formatting for CLI commands."""
+
+from __future__ import annotations
+
+import json
+import sys
+from typing import Any
+
+
+def format_output(data: Any, fmt: str = "plain", template: str = "") -> str:
+    """Format command output based on requested format.
+
+    Formats:
+        plain  — human-readable (default)
+        json   — JSON output
+        table  — tabular output (for lists of dicts)
+        template — render through a kida template
+    """
+    match fmt:
+        case "json":
+            return _format_json(data)
+        case "table":
+            return _format_table(data)
+        case "template":
+            return _format_template(data, template)
+        case _:
+            return _format_plain(data)
+
+
+def _format_json(data: Any) -> str:
+    """JSON output."""
+    return json.dumps(data, indent=2, default=str)
+
+
+def _format_plain(data: Any) -> str:
+    """Human-readable plain output."""
+    if isinstance(data, list):
+        lines = []
+        for item in data:
+            if isinstance(item, dict):
+                parts = [f"{v}" for v in item.values()]
+                lines.append("  ".join(parts))
+            else:
+                lines.append(str(item))
+        return "\n".join(lines)
+    if isinstance(data, dict):
+        lines = []
+        max_key = max((len(str(k)) for k in data), default=0)
+        for k, v in data.items():
+            lines.append(f"  {k!s:<{max_key}}  {v}")
+        return "\n".join(lines)
+    return str(data)
+
+
+def _format_table(data: Any) -> str:
+    """Table output using kida's table filter if available."""
+    if not isinstance(data, list):
+        return _format_plain(data)
+
+    if not data:
+        return "(empty)"
+
+    # Extract headers from first dict item
+    if isinstance(data[0], dict):
+        headers = list(data[0].keys())
+        rows = [[str(item.get(h, "")) for h in headers] for item in data]
+    else:
+        headers = None
+        rows = [[str(item)] for item in data]
+
+    try:
+        from kida import Environment
+
+        env = Environment(autoescape="terminal")
+        tmpl = env.from_string("{{ data | table(headers=headers) }}", name="table_fmt")
+        return tmpl.render(data=rows, headers=headers or [])
+    except ImportError:
+        # Fallback: simple column alignment
+        all_rows = [headers, *rows] if headers else rows
+        widths = [max(len(str(cell)) for cell in col) for col in zip(*all_rows, strict=False)]
+        lines = []
+        for row in all_rows:
+            lines.append(
+                "  ".join(str(cell).ljust(w) for cell, w in zip(row, widths, strict=False))
+            )
+            if row is headers:
+                lines.append("  ".join("-" * w for w in widths))
+        return "\n".join(lines)
+
+
+def _format_template(data: Any, template_name: str) -> str:
+    """Render through a kida template."""
+    if not template_name:
+        return _format_plain(data)
+    from milo.templates import get_env
+
+    env = get_env()
+    tmpl = env.get_template(template_name)
+    return tmpl.render(state=data)
+
+
+def write_output(data: Any, fmt: str = "plain", template: str = "") -> None:
+    """Format and write command output to stdout."""
+    output = format_output(data, fmt=fmt, template=template)
+    sys.stdout.write(output + "\n")
+    sys.stdout.flush()