ripperdoc 0.1.0__py3-none-any.whl

ripperdoc/utils/mcp.py

@@ -0,0 +1,427 @@
+"""MCP configuration loader, connection manager, and prompt helpers."""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import json
+from contextlib import AsyncExitStack
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from ripperdoc import __version__
+from ripperdoc.utils.log import get_logger
+from ripperdoc.utils.message_compaction import estimate_tokens_from_text
+
+try:
+    import mcp.types as mcp_types
+    from mcp.client.session import ClientSession
+    from mcp.client.sse import sse_client
+    from mcp.client.stdio import StdioServerParameters, stdio_client
+    from mcp.client.streamable_http import streamablehttp_client
+
+    MCP_AVAILABLE = True
+except Exception:  # pragma: no cover - handled gracefully at runtime
+    MCP_AVAILABLE = False
+    ClientSession = object  # type: ignore
+    mcp_types = None  # type: ignore
+
+
+logger = get_logger()
+
+
+@dataclass
+class McpToolInfo:
+    name: str
+    description: str = ""
+    input_schema: Optional[Dict[str, Any]] = None
+    annotations: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class McpResourceInfo:
+    uri: str
+    name: Optional[str] = None
+    description: str = ""
+    mime_type: Optional[str] = None
+    size: Optional[int] = None
+    text: Optional[str] = None
+
+
+@dataclass
+class McpServerInfo:
+    name: str
+    type: str = "stdio"
+    url: Optional[str] = None
+    description: str = ""
+    command: Optional[str] = None
+    args: List[str] = field(default_factory=list)
+    env: Dict[str, str] = field(default_factory=dict)
+    headers: Dict[str, str] = field(default_factory=dict)
+    tools: List[McpToolInfo] = field(default_factory=list)
+    resources: List[McpResourceInfo] = field(default_factory=list)
+    status: str = "configured"
+    error: Optional[str] = None
+    instructions: Optional[str] = None
+    server_version: Optional[str] = None
+    capabilities: Dict[str, Any] = field(default_factory=dict)
+
+
+def _load_json_file(path: Path) -> Dict[str, Any]:
+    if not path.exists():
+        return {}
+    try:
+        return json.loads(path.read_text())
+    except (OSError, json.JSONDecodeError) as exc:
+        logger.error(f"Failed to load JSON from {path}: {exc}")
+        return {}
+
+
+def _ensure_str_dict(raw: object) -> Dict[str, str]:
+    if not isinstance(raw, dict):
+        return {}
+    result: Dict[str, str] = {}
+    for key, value in raw.items():
+        try:
+            result[str(key)] = str(value)
+        except Exception:
+            continue
+    return result
+
+
+def _parse_server(name: str, raw: Dict[str, Any]) -> McpServerInfo:
+    server_type = str(raw.get("type") or raw.get("transport") or "").strip().lower()
+    command = raw.get("command")
+    args = raw.get("args") if isinstance(raw.get("args"), list) else []
+    url = str(raw.get("url") or raw.get("uri") or "").strip() or None
+
+    if not server_type:
+        if url:
+            server_type = "sse"
+        elif command:
+            server_type = "stdio"
+        else:
+            server_type = "stdio"
+
+    description = str(raw.get("description") or "")
+    env = _ensure_str_dict(raw.get("env"))
+    headers = _ensure_str_dict(raw.get("headers"))
+    instructions = raw.get("instructions")
+
+    return McpServerInfo(
+        name=name,
+        type=server_type,
+        url=url,
+        description=description,
+        command=str(command) if isinstance(command, str) else None,
+        args=[str(a) for a in args] if args else [],
+        env=env,
+        headers=headers,
+        instructions=str(instructions) if isinstance(instructions, str) else None,
+    )
+
+
+def _parse_servers(data: Dict[str, Any]) -> Dict[str, McpServerInfo]:
+    servers: Dict[str, McpServerInfo] = {}
+    for key in ("servers", "mcpServers"):
+        raw_servers = data.get(key)
+        if not isinstance(raw_servers, dict):
+            continue
+        for name, raw in raw_servers.items():
+            if not isinstance(raw, dict):
+                continue
+            server_name = str(name).strip()
+            if not server_name:
+                continue
+            servers[server_name] = _parse_server(server_name, raw)
+    return servers
+
+
+def _load_server_configs(project_path: Optional[Path]) -> Dict[str, McpServerInfo]:
+    project_path = project_path or Path.cwd()
+    candidates = [
+        Path.home() / ".ripperdoc" / "mcp.json",
+        Path.home() / ".mcp.json",
+        project_path / ".ripperdoc" / "mcp.json",
+        project_path / ".mcp.json",
+    ]
+
+    merged: Dict[str, McpServerInfo] = {}
+    for path in candidates:
+        data = _load_json_file(path)
+        merged.update(_parse_servers(data))
+    return merged
+
+
+class McpRuntime:
+    """Manages live MCP connections for the current event loop."""
+
+    def __init__(self, project_path: Path):
+        self.project_path = project_path
+        self._exit_stack = AsyncExitStack()
+        self.sessions: Dict[str, ClientSession] = {}
+        self.servers: List[McpServerInfo] = []
+        self._closed = False
+
+    async def connect(self, configs: Dict[str, McpServerInfo]) -> List[McpServerInfo]:
+        await self._exit_stack.__aenter__()
+        if not MCP_AVAILABLE:
+            for config in configs.values():
+                self.servers.append(
+                    replace(
+                        config,
+                        status="unavailable",
+                        error="MCP Python SDK not installed; install `mcp[cli]` with Python 3.10+.",
+                    )
+                )
+            return self.servers
+
+        for config in configs.values():
+            self.servers.append(await self._connect_server(config))
+        return self.servers
+
+    async def _list_roots_callback(self, *_: Any, **__: Any):
+        if not mcp_types:
+            return None
+        return mcp_types.ListRootsResult(
+            roots=[mcp_types.Root(uri=Path(self.project_path).resolve().as_uri())]
+        )
+
+    async def _connect_server(self, config: McpServerInfo) -> McpServerInfo:
+        info = replace(config, tools=[], resources=[])
+        if not MCP_AVAILABLE or not mcp_types:
+            info.status = "unavailable"
+            info.error = "MCP Python SDK not installed."
+            return info
+
+        try:
+            read_stream = None
+            write_stream = None
+
+            if config.type in ("sse", "sse-ide"):
+                if not config.url:
+                    raise ValueError("SSE MCP server requires a 'url'.")
+                read_stream, write_stream = await self._exit_stack.enter_async_context(
+                    sse_client(config.url, headers=config.headers or None)
+                )
+            elif config.type in ("http", "streamable-http"):
+                if not config.url:
+                    raise ValueError("HTTP MCP server requires a 'url'.")
+                read_stream, write_stream, _ = await self._exit_stack.enter_async_context(
+                    streamablehttp_client(
+                        url=config.url,
+                        headers=config.headers or None,
+                        terminate_on_close=True,
+                    )
+                )
+            else:
+                if not config.command:
+                    raise ValueError("Stdio MCP server requires a 'command'.")
+                stdio_params = StdioServerParameters(
+                    command=config.command,
+                    args=config.args,
+                    env=config.env or None,
+                    cwd=self.project_path,
+                )
+                read_stream, write_stream = await self._exit_stack.enter_async_context(
+                    stdio_client(stdio_params)
+                )
+
+            session = await self._exit_stack.enter_async_context(
+                ClientSession(
+                    read_stream,
+                    write_stream,
+                    list_roots_callback=self._list_roots_callback,
+                    client_info=mcp_types.Implementation(name="ripperdoc", version=__version__),
+                )
+            )
+
+            init_result = await session.initialize()
+            capabilities = session.get_server_capabilities()
+
+            info.status = "connected"
+            info.instructions = init_result.instructions or info.instructions
+            info.server_version = getattr(init_result.serverInfo, "version", None)
+            info.capabilities = (
+                capabilities.model_dump() if hasattr(capabilities, "model_dump") else {}
+            )
+            self.sessions[config.name] = session
+
+            tools_result = await session.list_tools()
+            info.tools = [
+                McpToolInfo(
+                    name=tool.name,
+                    description=tool.description or "",
+                    input_schema=tool.inputSchema,
+                    annotations=(tool.annotations.model_dump() if tool.annotations else {}),
+                )
+                for tool in tools_result.tools
+            ]
+
+            if capabilities and getattr(capabilities, "resources", None):
+                resources_result = await session.list_resources()
+                info.resources = [
+                    McpResourceInfo(
+                        uri=str(resource.uri),
+                        name=resource.name,
+                        description=resource.description or "",
+                        mime_type=resource.mimeType,
+                        size=resource.size,
+                    )
+                    for resource in resources_result.resources
+                ]
+
+        except Exception as exc:  # pragma: no cover - network/process errors
+            logger.error(f"Failed to connect to MCP server '{config.name}': {exc}")
+            info.status = "failed"
+            info.error = str(exc)
+
+        return info
+
+    async def aclose(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            await self._exit_stack.aclose()
+        finally:
+            self.sessions.clear()
+            self.servers.clear()
+
+
+_runtime_var: contextvars.ContextVar[Optional[McpRuntime]] = contextvars.ContextVar(
+    "ripperdoc_mcp_runtime", default=None
+)
+
+
+def _get_runtime() -> Optional[McpRuntime]:
+    return _runtime_var.get()
+
+
+def get_existing_mcp_runtime() -> Optional[McpRuntime]:
+    """Return the current MCP runtime if it has already been initialized."""
+    return _get_runtime()
+
+
+async def ensure_mcp_runtime(project_path: Optional[Path] = None) -> McpRuntime:
+    runtime = _get_runtime()
+    project_path = project_path or Path.cwd()
+    if runtime and not runtime._closed and runtime.project_path == project_path:
+        return runtime
+
+    if runtime:
+        await runtime.aclose()
+
+    runtime = McpRuntime(project_path)
+    configs = _load_server_configs(project_path)
+    await runtime.connect(configs)
+    _runtime_var.set(runtime)
+    return runtime
+
+
+async def shutdown_mcp_runtime() -> None:
+    runtime = _get_runtime()
+    if not runtime:
+        return
+    await runtime.aclose()
+    _runtime_var.set(None)
+
+
+async def load_mcp_servers_async(project_path: Optional[Path] = None) -> List[McpServerInfo]:
+    runtime = await ensure_mcp_runtime(project_path)
+    return list(runtime.servers)
+
+
+def _config_only_servers(project_path: Optional[Path]) -> List[McpServerInfo]:
+    return list(_load_server_configs(project_path).values())
+
+
+def load_mcp_servers(project_path: Optional[Path] = None) -> List[McpServerInfo]:
+    """Synchronous wrapper primarily for legacy call sites."""
+    try:
+        loop = asyncio.get_running_loop()
+        if loop.is_running():
+            runtime = _get_runtime()
+            if runtime and runtime.servers:
+                return list(runtime.servers)
+            return _config_only_servers(project_path)
+    except RuntimeError:
+        pass
+
+    return asyncio.run(load_mcp_servers_async(project_path))
+
+
+def find_mcp_resource(
+    servers: List[McpServerInfo],
+    server_name: str,
+    uri: str,
+) -> Optional[McpResourceInfo]:
+    server = next((s for s in servers if s.name == server_name), None)
+    if not server:
+        return None
+    return next((r for r in server.resources if r.uri == uri), None)
+
+
+def _summarize_tools(server: McpServerInfo) -> str:
+    if not server.tools:
+        return "no tools"
+    names = [tool.name for tool in server.tools[:6]]
+    suffix = ", ".join(names)
+    if len(server.tools) > 6:
+        suffix += f", and {len(server.tools) - 6} more"
+    return suffix
+
+
+def format_mcp_instructions(servers: List[McpServerInfo]) -> str:
+    """Build a concise MCP instruction block for the system prompt."""
+    if not servers:
+        return ""
+
+    lines: List[str] = [
+        "MCP servers are available. Call tools via CallMcpTool by specifying server, tool, and arguments.",
+        "Use ListMcpServers to inspect servers and ListMcpResources/ReadMcpResource when a server exposes resources.",
+    ]
+
+    for server in servers:
+        status = server.status or "unknown"
+        prefix = f"- {server.name} [{status}]"
+        if server.url:
+            prefix += f" {server.url}"
+        lines.append(prefix)
+
+        if status == "connected":
+            if server.instructions:
+                trimmed = server.instructions.strip()
+                if len(trimmed) > 260:
+                    trimmed = trimmed[:257] + "..."
+                lines.append(f"  Instructions: {trimmed}")
+            tool_summary = _summarize_tools(server)
+            lines.append(f"  Tools: {tool_summary}")
+            if server.resources:
+                lines.append(f"  Resources: {len(server.resources)} available")
+        elif server.error:
+            lines.append(f"  Error: {server.error}")
+
+    return "\n".join(lines)
+
+
+def estimate_mcp_tokens(servers: List[McpServerInfo]) -> int:
+    """Estimate token usage for MCP instructions."""
+    mcp_text = format_mcp_instructions(servers)
+    return estimate_tokens_from_text(mcp_text)
+
+
+__all__ = [
+    "McpServerInfo",
+    "McpToolInfo",
+    "McpResourceInfo",
+    "get_existing_mcp_runtime",
+    "load_mcp_servers",
+    "load_mcp_servers_async",
+    "ensure_mcp_runtime",
+    "shutdown_mcp_runtime",
+    "find_mcp_resource",
+    "format_mcp_instructions",
+    "estimate_mcp_tokens",
+]

ripperdoc/utils/memory.py

@@ -0,0 +1,239 @@
+"""Helpers for loading AGENTS.md memory files."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Set
+
+MEMORY_FILE_NAME = "AGENTS.md"
+LOCAL_MEMORY_FILE_NAME = "AGENTS.local.md"
+
+MEMORY_INSTRUCTIONS = (
+    "Codebase and user instructions are shown below. Be sure to adhere to these "
+    "instructions. IMPORTANT: These instructions OVERRIDE any default behavior "
+    "and you MUST follow them exactly as written."
+)
+
+MAX_CONTENT_LENGTH = 40_000
+MAX_INCLUDE_DEPTH = 5
+
+_CODE_FENCE_RE = re.compile(r"```.*?```", flags=re.DOTALL)
+_INLINE_CODE_RE = re.compile(r"`[^`]*`")
+_MENTION_RE = re.compile(r"(?:^|\s)@((?:[^\s\\]|\\ )+)")
+_PUNCT_START_RE = re.compile(r"^[#%^&*()]+")
+_VALID_START_RE = re.compile(r"^[A-Za-z0-9._-]")
+
+
+@dataclass
+class MemoryFile:
+    """Representation of a loaded memory file."""
+
+    path: str
+    type: str
+    content: str
+    parent: Optional[str] = None
+    is_nested: bool = False
+
+
+def _is_path_under_directory(path: Path, directory: Path) -> bool:
+    """Return True if path is inside directory (after resolving)."""
+    try:
+        path.resolve().relative_to(directory.resolve())
+        return True
+    except Exception:
+        return False
+
+
+def _resolve_relative_path(raw_path: str, base_path: Path) -> Path:
+    """Resolve a mention (./foo, ~/bar, /abs, or relative) against a base file."""
+    normalized = raw_path.replace("\\ ", " ")
+    if normalized.startswith("~/"):
+        return (Path.home() / normalized[2:]).resolve()
+    candidate = Path(normalized)
+    if not candidate.is_absolute():
+        return (base_path.parent / candidate).resolve()
+    return candidate.resolve()
+
+
+def _read_file_with_type(file_path: Path, file_type: str) -> Optional[MemoryFile]:
+    """Read a file if it exists, returning a MemoryFile entry."""
+    try:
+        if not file_path.exists() or not file_path.is_file():
+            return None
+        content = file_path.read_text(encoding="utf-8", errors="ignore")
+        return MemoryFile(path=str(file_path), type=file_type, content=content)
+    except PermissionError:
+        return None
+    except OSError:
+        return None
+
+
+def _extract_relative_paths_from_markdown(markdown_content: str, base_path: Path) -> List[Path]:
+    """Extract @-mentions that look like file paths from markdown content."""
+    if not markdown_content:
+        return []
+
+    cleaned = _CODE_FENCE_RE.sub("", markdown_content)
+    cleaned = _INLINE_CODE_RE.sub("", cleaned)
+
+    relative_paths: Set[Path] = set()
+    for match in _MENTION_RE.finditer(cleaned):
+        mention = (match.group(1) or "").replace("\\ ", " ").strip()
+        if not mention or mention.startswith("@"):
+            continue
+
+        if not (
+            mention.startswith("./")
+            or mention.startswith("~/")
+            or (mention.startswith("/") and mention != "/")
+            or (not _PUNCT_START_RE.match(mention) and _VALID_START_RE.match(mention))
+        ):
+            continue
+
+        resolved = _resolve_relative_path(mention, base_path)
+        relative_paths.add(resolved)
+
+    return list(relative_paths)
+
+
+def _collect_files(
+    file_path: Path,
+    file_type: str,
+    visited: Set[str],
+    allow_outside_cwd: bool,
+    depth: int = 0,
+    parent_path: Optional[Path] = None,
+) -> List[MemoryFile]:
+    """Collect a memory file and any nested references."""
+    if depth >= MAX_INCLUDE_DEPTH:
+        return []
+
+    resolved_path = file_path.expanduser()
+    try:
+        resolved_path = resolved_path.resolve()
+    except Exception:
+        pass
+
+    resolved_key = str(resolved_path)
+    if resolved_key in visited:
+        return []
+
+    current_file = _read_file_with_type(resolved_path, file_type)
+    if not current_file or not current_file.content.strip():
+        return []
+
+    if parent_path is not None:
+        current_file.parent = str(parent_path)
+        current_file.is_nested = depth > 0
+
+    visited.add(resolved_key)
+
+    collected: List[MemoryFile] = [current_file]
+    relative_paths = _extract_relative_paths_from_markdown(current_file.content, resolved_path)
+    for nested_path in relative_paths:
+        if not allow_outside_cwd and not _is_path_under_directory(nested_path, Path.cwd()):
+            continue
+        collected.extend(
+            _collect_files(
+                nested_path,
+                file_type,
+                visited,
+                allow_outside_cwd,
+                depth + 1,
+                resolved_path,
+            )
+        )
+
+    return collected
+
+
+def collect_all_memory_files(force_include_external: bool = False) -> List[MemoryFile]:
+    """Collect all AGENTS memory files reachable from the working directory."""
+    visited: Set[str] = set()
+    files: List[MemoryFile] = []
+
+    # Global/user-level memories live in home and ~/.ripperdoc.
+    user_memory_paths = [
+        Path.home() / ".ripperdoc" / MEMORY_FILE_NAME,
+        Path.home() / MEMORY_FILE_NAME,
+    ]
+    for user_memory_path in user_memory_paths:
+        files.extend(
+            _collect_files(
+                user_memory_path,
+                "User",
+                visited,
+                allow_outside_cwd=True,
+            )
+        )
+
+    # Project memories from the current working directory up to the filesystem root.
+    ancestor_dirs: List[Path] = []
+    current_dir = Path.cwd()
+    while True:
+        ancestor_dirs.append(current_dir)
+        if current_dir.parent == current_dir:
+            break
+        current_dir = current_dir.parent
+
+    for directory in reversed(ancestor_dirs):
+        files.extend(
+            _collect_files(
+                directory / MEMORY_FILE_NAME,
+                "Project",
+                visited,
+                allow_outside_cwd=force_include_external,
+            )
+        )
+        files.extend(
+            _collect_files(
+                directory / LOCAL_MEMORY_FILE_NAME,
+                "Local",
+                visited,
+                allow_outside_cwd=force_include_external,
+            )
+        )
+
+    return files
+
+
+def get_oversized_memory_files() -> List[MemoryFile]:
+    """Return memory files that exceed the recommended length."""
+    return [file for file in collect_all_memory_files() if len(file.content) > MAX_CONTENT_LENGTH]
+
+
+def build_memory_instructions() -> str:
+    """Build the instruction block to append to the system prompt."""
+    memory_files = collect_all_memory_files()
+    snippets: List[str] = []
+    for memory_file in memory_files:
+        if not memory_file.content:
+            continue
+        type_description = (
+            " (project instructions, checked into the codebase)"
+            if memory_file.type == "Project"
+            else (
+                " (user's private project instructions, not checked in)"
+                if memory_file.type == "Local"
+                else " (user's private global instructions)"
+            )
+        )
+        snippets.append(
+            f"Contents of {memory_file.path}{type_description}:\n\n{memory_file.content}"
+        )
+
+    if not snippets:
+        return ""
+
+    return f"{MEMORY_INSTRUCTIONS}\n\n" + "\n\n".join(snippets)
+
+
+__all__ = [
+    "MemoryFile",
+    "collect_all_memory_files",
+    "build_memory_instructions",
+    "get_oversized_memory_files",
+    "MAX_CONTENT_LENGTH",
+]