beaver_project/app-instance/backend/beaver/interfaces/mcp/tools_server.py

"""Beaver local tools as real stdio MCP servers."""

from __future__ import annotations

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

import mcp.types as types
from mcp.server.lowlevel import Server
from mcp.server.lowlevel.server import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server

from beaver.engine.session import SessionManager
from beaver.memory.curated.store import MemoryStore
from beaver.services.cron_service import CronService
from beaver.skills import SkillsLoader
from beaver.skills.drafts import DraftService
from beaver.skills.specs import SkillSpecStore
from beaver.tools.base import BaseTool, ObjectBackedTool, ToolContext
from beaver.tools.builtins import (
    ClarifyTool,
    CronTool,
    DelegateTool,
    ExecuteCodeTool,
    MemoryTool,
    ProcessTool,
    SendMessageTool,
    SkillManageTool,
    SkillViewTool,
    SkillsListTool,
    SpawnTool,
    TerminalTool,
    TodoTool,
    UserFilesCopyToWorkspaceTool,
    UserFilesListTool,
    UserFilesMkdirTool,
    UserFilesPublishOutputTool,
    UserFilesReadTool,
    UserFilesWriteTool,
    WebFetchTool,
    WebSearchTool,
    WriteFileTool,
)


LOCAL_TOOL_CATEGORIES = {
    "filesystem": "Beaver Personal Agent Filesystem Tools",
    "runtime": "Beaver Local Runtime Tools",
    "memory": "Beaver Local Memory Tools",
    "skills": "Beaver Local Skills Tools",
    "coordination": "Beaver Local Coordination Tools",
    "scheduler": "Beaver Local Scheduler Tools",
    "web": "Beaver Local Web Tools",
}


def _workspace_path(value: str | None = None) -> Path:
    raw = value or os.getenv("BEAVER_WORKSPACE")
    if raw:
        return Path(raw).expanduser().resolve()
    return Path.cwd()


def _json_content(value: str) -> dict[str, Any]:
    try:
        parsed = json.loads(value)
        return parsed if isinstance(parsed, dict) else {"success": True, "result": parsed}
    except json.JSONDecodeError:
        return {"success": True, "content": value}


def _category_tools(category: str, workspace: Path) -> tuple[list[BaseTool], ToolContext]:
    skill_store = SkillSpecStore(workspace)
    skills_loader = SkillsLoader(workspace, skill_store=skill_store)
    draft_service = DraftService(skill_store)
    services = {
        "skills_loader": skills_loader,
        "draft_service": draft_service,
    }
    context = ToolContext(workspace=str(workspace), services=services)

    if category == "filesystem":
        tools: list[BaseTool] = [
            ObjectBackedTool(UserFilesListTool()),
            ObjectBackedTool(UserFilesReadTool()),
            ObjectBackedTool(UserFilesWriteTool()),
            ObjectBackedTool(UserFilesMkdirTool()),
            ObjectBackedTool(UserFilesCopyToWorkspaceTool()),
            ObjectBackedTool(UserFilesPublishOutputTool()),
        ]
    elif category == "runtime":
        tools = [
            ObjectBackedTool(TerminalTool()),
            ObjectBackedTool(ProcessTool()),
            ObjectBackedTool(ExecuteCodeTool()),
        ]
    elif category == "memory":
        session_manager = SessionManager(workspace)
        memory_store = MemoryStore(workspace / "memory" / "curated")
        memory_store.load_from_disk()
        tools = [
            ObjectBackedTool(MemoryTool(store=memory_store)),
            ObjectBackedTool(__import__("beaver.tools.builtins.session_search", fromlist=["SessionSearchTool"]).SessionSearchTool(db=session_manager)),
        ]
    elif category == "skills":
        tools = [
            ObjectBackedTool(SkillViewTool(loader=skills_loader)),
            SkillsListTool(),
            SkillManageTool(),
        ]
    elif category == "coordination":
        tools = [
            ObjectBackedTool(TodoTool()),
            ObjectBackedTool(ClarifyTool()),
            ObjectBackedTool(DelegateTool()),
            ObjectBackedTool(SpawnTool()),
            ObjectBackedTool(SendMessageTool()),
        ]
    elif category == "scheduler":
        services["cron_service"] = CronService(workspace / "cron" / "jobs.json")
        tools = [CronTool()]
    elif category == "web":
        tools = [
            ObjectBackedTool(WebFetchTool()),
            ObjectBackedTool(WebSearchTool()),
        ]
    else:
        raise ValueError(f"Unknown local tool category: {category}")
    return tools, context


def create_tools_server(*, category: str, workspace: str | None = None) -> Server:
    workspace_path = _workspace_path(workspace)
    tools, context = _category_tools(category, workspace_path)
    tool_map = {tool.spec.name: tool for tool in tools}
    server = Server(LOCAL_TOOL_CATEGORIES.get(category, f"Beaver Local {category} Tools"))

    @server.list_tools()
    async def list_tools() -> list[types.Tool]:
        return [
            types.Tool(
                name=tool.spec.name,
                description=tool.spec.description,
                inputSchema=tool.spec.input_schema,
            )
            for tool in tools
        ]

    @server.call_tool(validate_input=True)
    async def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
        tool = tool_map.get(name)
        if tool is None:
            return {"success": False, "error": f"Unknown tool: {name}"}
        result = await tool.invoke(arguments or {}, context)
        if result.raw_output is not None and isinstance(result.raw_output, dict):
            return result.raw_output
        payload = _json_content(result.content)
        if "success" not in payload:
            payload["success"] = bool(result.success)
        if result.error and "error" not in payload:
            payload["error"] = result.error
        return payload

    return server


async def _run_stdio(category: str, workspace: str | None) -> None:
    server = create_tools_server(category=category, workspace=workspace)
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name=LOCAL_TOOL_CATEGORIES.get(category, f"beaver-{category}"),
                server_version="0.1.0",
                capabilities=server.get_capabilities(notification_options=NotificationOptions(), experimental_capabilities={}),
            ),
        )


def main() -> None:
    parser = argparse.ArgumentParser(description="Run a Beaver local tool category as a stdio MCP server.")
    parser.add_argument("--category", choices=sorted(LOCAL_TOOL_CATEGORIES), required=True)
    parser.add_argument("--workspace", default=None)
    args = parser.parse_args()
    asyncio.run(_run_stdio(args.category, args.workspace))


if __name__ == "__main__":
    main()