beaver-gateway/src/beaver_gateway/core/gateway_tools.py

"""In-process MCP server with the gateway's own tools (§3.1, §3.2).

One server per live session so every tool knows which conversation is
calling; ``alwaysLoad`` keeps the tools out of tool search. Which names a
session gets comes from ``ClaudeAgent.gateway_tools``.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, cast

from claude_agent_sdk import create_sdk_mcp_server, tool

if TYPE_CHECKING:
    from collections.abc import Iterable

    from claude_agent_sdk import McpSdkServerConfig, SdkMcpTool

    from beaver_gateway.core.conversations import Conversations

__all__ = ["SERVER_NAME", "TOOL_NAMES", "build_tool_server"]

_log = logging.getLogger("beaver_gateway.core.gateway_tools")

SERVER_NAME = "gateway"
TOOL_NAMES = ("read_conversation", "spawn", "say", "schedule", "inject")


def build_tool_server(
    conversations: Conversations, *, conversation_key: str, names: Iterable[str]
) -> McpSdkServerConfig | None:
    wanted = set(names)
    unknown = wanted - set(TOOL_NAMES)
    if unknown:
        msg = f"unknown gateway tools: {sorted(unknown)}"
        raise ValueError(msg)
    tools = [t for t in _tools(conversations, conversation_key) if t.name in wanted]
    if not tools:
        return None
    server = create_sdk_mcp_server(SERVER_NAME, tools=tools)
    return cast("McpSdkServerConfig", {**server, "alwaysLoad": True})


def _tools(conversations: Conversations, key: str) -> list[SdkMcpTool[Any]]:
    async def current() -> Any:
        conv = await conversations.get(key)
        if conv is None:
            msg = f"conversation {key} not found"
            raise LookupError(msg)
        return conv

    @tool(
        "read_conversation",
        "Read another conversation (a branch, the master, a deep chat) as plain "
        "text. `window` limits it to the last N user turns.",
        {
            "type": "object",
            "properties": {
                "id": {"type": "string", "description": "conversation id"},
                "window": {"type": "integer", "minimum": 1},
            },
            "required": ["id"],
        },
    )
    async def read_conversation(args: dict[str, Any]) -> dict[str, Any]:
        conv = await conversations.get(str(args["id"]))
        if conv is None:
            return _error(f"conversation {args['id']} not found")
        text = await conversations.read(conv, window=args.get("window"))
        return _text(text or "(empty)")

    @tool(
        "spawn",
        "Open a new conversation of the given kind (branch = your own thread, "
        "deep = a long research chat, job = a headless task). `seed` is how it "
        "starts: clean (nothing), morning (handout), copy (copy of this "
        "conversation, last `window` turns), brief (your `text`). Returns the id.",
        {
            "type": "object",
            "properties": {
                "kind": {"type": "string", "enum": ["branch", "deep", "job"]},
                "seed": {
                    "type": "string",
                    "enum": ["clean", "morning", "copy", "brief"],
                    "default": "clean",
                },
                "text": {"type": "string", "description": "brief for seed=brief"},
                "title": {"type": "string"},
                "window": {"type": "integer", "minimum": 1},
            },
            "required": ["kind"],
        },
    )
    async def spawn(args: dict[str, Any]) -> dict[str, Any]:
        parent = await current()
        child = await conversations.spawn(
            kind=str(args["kind"]),
            agent=parent.agent_name,
            seed=str(args.get("seed") or "clean"),
            parent=parent,
            text=args.get("text"),
            title=args.get("title"),
            window=args.get("window"),
            origin="mcp",
        )
        return _text(f"spawned {child.kind} {child.external_id}")

    @tool(
        "say",
        "Say something to the human in the frontend this conversation is bound "
        "to. The only way an inject-started turn can speak; silence is simply "
        "not calling it.",
        {"text": str},
    )
    async def say(args: dict[str, Any]) -> dict[str, Any]:
        conv = await current()
        await conversations.say(conv, str(args["text"]))
        return _text("ok")

    @tool(
        "schedule",
        "Promise yourself an inject later: `at` is `+15m`, `+2h`, `+1d` or an "
        "ISO datetime; `text` arrives in this conversation at that time.",
        {"at": str, "text": str},
    )
    async def schedule(args: dict[str, Any]) -> dict[str, Any]:
        conv = await current()
        row = await conversations.schedule(conv, str(args["at"]), str(args["text"]))
        return _text(f"scheduled #{row.id} at {row.execute_at.isoformat()}")

    @tool(
        "inject",
        "Put a system-origin message into another conversation's queue.",
        {
            "type": "object",
            "properties": {
                "conversation": {"type": "string"},
                "text": {"type": "string"},
                "urgency": {
                    "type": "string",
                    "enum": ["normal", "urgent"],
                    "default": "normal",
                },
            },
            "required": ["conversation", "text"],
        },
    )
    async def inject(args: dict[str, Any]) -> dict[str, Any]:
        target = await conversations.get(str(args["conversation"]))
        if target is None:
            return _error(f"conversation {args['conversation']} not found")
        item = await conversations.inject(
            target,
            str(args["text"]),
            urgency=cast("Any", args.get("urgency") or "normal"),
            origin="агент",
        )
        return _text(f"queued #{item.id}")

    return [read_conversation, spawn, say, schedule, inject]


def _text(text: str) -> dict[str, Any]:
    return {"content": [{"type": "text", "text": text}]}


def _error(text: str) -> dict[str, Any]:
    return {"content": [{"type": "text", "text": text}], "is_error": True}