OpenSpace/openspace/mcp_proxy.py

from __future__ import annotations

import asyncio
import argparse
import inspect
import json
import os
from pathlib import Path
from typing import Any

from mcp.server.fastmcp import FastMCP
from mcp.types import TextContent

from openspace.mcp_stdio import maybe_redirect_stderr_to_file
from openspace.grounding.backends.mcp.client import MCPClient
from openspace.mcp_tool_registration import (
    register_evolution_tools,
    register_main_tools,
)
from openspace.shared_mcp_runtime import ServerKind, ensure_daemon


_LOG_DIR = Path(__file__).resolve().parent.parent / "logs"
maybe_redirect_stderr_to_file(_LOG_DIR, "mcp_proxy_stderr.log")


def _proxy_mode_for(server_kind: ServerKind) -> str:
    raw = os.environ.get("OPENSPACE_MCP_PROXY_MODE", "").strip().lower()
    if raw in {"daemon", "direct"}:
        return raw
    return "daemon"


def _json_error(error: Any, **extra: Any) -> str:
    return json.dumps({"error": str(error), **extra}, ensure_ascii=False)


def _extract_text_payload(result: Any) -> str:
    text_parts: list[str] = []
    for item in getattr(result, "content", []):
        if isinstance(item, TextContent):
            text_parts.append(item.text)
            continue
        text = getattr(item, "text", None)
        if text is not None:
            text_parts.append(text)
    if not text_parts:
        raise RuntimeError("Remote MCP tool returned no text payload")
    return "\n".join(text_parts)


class _RemoteProxyBase:
    def __init__(self, server_kind: ServerKind):
        self._server_kind = server_kind

    async def _call_remote_tool_once(self, tool_name: str, args: dict[str, Any]) -> str:
        record = await ensure_daemon(self._server_kind)
        client = MCPClient(
            config={"mcpServers": {"daemon": {"url": record.url}}},
            timeout=10.0,
            sse_read_timeout=60 * 60.0,
            check_dependencies=False,
        )
        try:
            session = await client.create_session("daemon", auto_initialize=True)
            if session is None:
                raise RuntimeError("Failed to create daemon MCP session")
            result = await session.connector.call_tool(tool_name, args)
            return _extract_text_payload(result)
        finally:
            await client.close_all_sessions()

    def _call_remote_tool_blocking(self, tool_name: str, args: dict[str, Any]) -> str:
        return asyncio.run(self._call_remote_tool_once(tool_name, args))

    async def _call_remote_tool(self, tool_name: str, args: dict[str, Any]) -> str:
        for attempt in range(2):
            try:
                return await asyncio.to_thread(
                    self._call_remote_tool_blocking,
                    tool_name,
                    args,
                )
            except Exception as exc:
                if attempt == 1:
                    return _json_error(exc, status="error")
        return _json_error("Unreachable proxy retry path", status="error")


class _MainProxyImplementation(_RemoteProxyBase):
    def __init__(self):
        super().__init__("main")

    async def execute_task(
        self,
        task: str,
        workspace_dir: str | None = None,
        max_iterations: int | None = None,
        skill_dirs: list[str] | None = None,
        search_scope: str = "all",
    ) -> str:
        return await self._call_remote_tool(
            "execute_task",
            {
                "task": task,
                "workspace_dir": workspace_dir,
                "max_iterations": max_iterations,
                "skill_dirs": skill_dirs,
                "search_scope": search_scope,
            },
        )

    async def search_skills(
        self,
        query: str,
        source: str = "all",
        limit: int = 20,
        auto_import: bool = True,
    ) -> str:
        return await self._call_remote_tool(
            "search_skills",
            {
                "query": query,
                "source": source,
                "limit": limit,
                "auto_import": auto_import,
            },
        )

    async def fix_skill(
        self,
        skill_dir: str,
        direction: str,
    ) -> str:
        return await self._call_remote_tool(
            "fix_skill",
            {
                "skill_dir": skill_dir,
                "direction": direction,
            },
        )

    async def upload_skill(
        self,
        skill_dir: str,
        visibility: str = "public",
        origin: str | None = None,
        parent_skill_ids: list[str] | None = None,
        tags: list[str] | None = None,
        created_by: str | None = None,
        change_summary: str | None = None,
    ) -> str:
        return await self._call_remote_tool(
            "upload_skill",
            {
                "skill_dir": skill_dir,
                "visibility": visibility,
                "origin": origin,
                "parent_skill_ids": parent_skill_ids,
                "tags": tags,
                "created_by": created_by,
                "change_summary": change_summary,
            },
        )


class _EvolutionProxyImplementation(_RemoteProxyBase):
    def __init__(self):
        super().__init__("evolution")

    async def evolve_from_context(
        self,
        task: str,
        summary: str,
        workspace_dir: str | None = None,
        file_paths: list[str] | None = None,
        max_skills: int = 3,
        skill_dirs: list[str] | None = None,
        output_dir: str | None = None,
    ) -> str:
        return await self._call_remote_tool(
            "evolve_from_context",
            {
                "task": task,
                "summary": summary,
                "workspace_dir": workspace_dir,
                "file_paths": file_paths,
                "max_skills": max_skills,
                "skill_dirs": skill_dirs,
                "output_dir": output_dir,
            },
        )


def _build_fastmcp(server_kind: ServerKind) -> FastMCP:
    kwargs: dict[str, Any] = {}
    if "description" in inspect.signature(FastMCP.__init__).parameters:
        if server_kind == "main":
            kwargs["description"] = (
                "OpenSpace: Unite the Agents. Evolve the Mind. Rebuild the World."
            )
        else:
            kwargs["description"] = (
                "OpenSpace evolution sidecar: capture reusable skills from host-agent work."
            )
    name = "OpenSpace" if server_kind == "main" else "OpenSpace Evolution"
    return FastMCP(name, **kwargs)


def _run_proxy(server_kind: ServerKind) -> None:
    if _proxy_mode_for(server_kind) == "direct":
        if server_kind == "main":
            from openspace.mcp_server import run_mcp_server
        else:
            from openspace.evolution_mcp_server import run_mcp_server
        run_mcp_server()
        return

    parser = argparse.ArgumentParser(description="OpenSpace MCP proxy")
    parser.add_argument("--transport", choices=["stdio"], default="stdio")
    parser.parse_args()

    mcp = _build_fastmcp(server_kind)
    if server_kind == "main":
        register_main_tools(mcp, _MainProxyImplementation())
    else:
        register_evolution_tools(mcp, _EvolutionProxyImplementation())
    mcp.run(transport="stdio")


def run_main_mcp_proxy() -> None:
    _run_proxy("main")


def run_evolution_mcp_proxy() -> None:
    _run_proxy("evolution")


def main() -> None:
    parser = argparse.ArgumentParser(description="OpenSpace MCP proxy")
    parser.add_argument("--kind", choices=["main", "evolution"], required=True)
    parser.add_argument("--transport", choices=["stdio"], default="stdio")
    args = parser.parse_args()

    # Rebuild argv for the generic runner so direct fallback can reuse legacy entrypoints.
    transport = args.transport
    os.environ.setdefault("OPENSPACE_MCP_PROXY_MODE", _proxy_mode_for(args.kind))
    import sys

    sys.argv = [sys.argv[0], "--transport", transport]
    _run_proxy(args.kind)


if __name__ == "__main__":
    main()