lda-wf/src/wf_mcp/admin_surface/tools.py

from __future__ import annotations

from typing import Annotated, Any

from pydantic import Field

from wf_api.models import (
    AdminEventPayload,
    ConnectionPayload,
    ConnectionStatusPayload,
    InspectSourceResult,
    ListSourcesResult,
)
from wf_mcp.broker.service import WfMcpService

from .handlers.broker import BrokerAdminHandlers


def register_service_admin_tools(
    server: Any,
    service: WfMcpService,
    *,
    namespace: str | None = "wf.admin",
    legacy_names: bool = False,
    include_connection_tools: bool = True,
) -> None:
    """Register service-backed admin/control tools on an MCP server.

    The public server uses dotted `wf.admin.*` names. The retired broker server
    constructor still asks for bare compatibility names, so the namespace stays
    configurable while the implementation remains single-sourced.
    """
    handlers = BrokerAdminHandlers(service)

    legacy_name_map = {
        "read_resource": "read_broker_resource",
        "render_prompt": "render_broker_prompt",
        "invoke_method": "invoke_broker_method",
        "get_events": "get_broker_events",
    }

    def name(local_name: str) -> str:
        visible_name = (
            legacy_name_map.get(local_name, local_name) if legacy_names else local_name
        )
        return visible_name if namespace is None else f"{namespace}.{visible_name}"

    if include_connection_tools:

        @server.tool(
            name=name("list_connections"),
            title="List Connections",
            description="List configured MCP connections known to this server.",
        )
        async def list_connections() -> list[ConnectionPayload]:
            return await handlers.list_connections()

        @server.tool(
            name=name("get_connection_statuses"),
            title="Get Connection Statuses",
            description="Show configured MCP connection status and catalog counts.",
        )
        async def get_connection_statuses() -> list[ConnectionStatusPayload]:
            return await handlers.get_connection_statuses()

    @server.tool(
        name=name("refresh_connection_catalog"),
        title="Refresh Connection Catalog",
        description="Refresh one connection catalog snapshot from its upstream MCP server.",
    )
    async def refresh_connection_catalog(connection_id: str) -> dict[str, Any]:
        return await handlers.refresh_connection_catalog(connection_id)

    @server.tool(
        name=name("get_catalog"),
        title="Get Catalog",
        description="Return the current upstream MCP capability catalog.",
    )
    async def get_catalog() -> dict[str, Any]:
        return handlers.get_catalog()

    @server.tool(
        name=name("get_planner_catalog"),
        title="Get Planner Catalog",
        description="Return the planner catalog including local workflow sources.",
    )
    async def get_planner_catalog() -> dict[str, Any]:
        return handlers.get_planner_catalog()

    @server.tool(
        name=name("list_sources"),
        title="List Sources",
        description="List compact configured capability source summaries.",
    )
    async def list_sources(
        cursor: Annotated[
            str | None,
            Field(
                description=(
                    "Opaque pagination cursor returned by a previous list_sources "
                    "call. Omit for the first page."
                )
            ),
        ] = None,
        limit: Annotated[
            int,
            Field(
                ge=1,
                le=100,
                description=(
                    "Maximum source summaries to return. Use inspect_source for "
                    "one full source inventory."
                ),
            ),
        ] = 50,
    ) -> ListSourcesResult:
        return await handlers.list_sources(cursor=cursor, limit=limit)

    @server.tool(
        name=name("inspect_source"),
        title="Inspect Source",
        description="Return the full inventory for one configured capability source.",
    )
    async def inspect_source(
        source_id: Annotated[
            str,
            Field(
                description=(
                    "Exact source id from list_sources, such as wf.std, wf.docs, "
                    "or an enabled connection id like demo.personal."
                )
            ),
        ],
    ) -> InspectSourceResult:
        return await handlers.inspect_source(source_id)

    @server.tool(
        name=name("read_resource"),
        title="Read Resource",
        description=(
            "Read a local docs or broker-catalog resource by qualified name, "
            "for example wf.docs.workflow_capabilities."
        ),
    )
    async def read_resource(qualified_name: str) -> dict[str, Any]:
        return await handlers.read_broker_resource(qualified_name)

    @server.tool(
        name=name("render_prompt"),
        title="Render Prompt",
        description=(
            "Render a local docs or broker-catalog prompt by qualified name, "
            "for example wf.docs.workflow_authoring_guide."
        ),
    )
    async def render_prompt(
        qualified_name: str,
        arguments: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        return await handlers.render_broker_prompt(
            qualified_name,
            arguments=arguments,
        )

    @server.tool(
        name=name("invoke_method"),
        title="Invoke Method",
        description="Invoke a raw MCP method on one configured connection.",
    )
    async def invoke_method(
        connection_id: str,
        method: str,
        params: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        return await handlers.invoke_broker_method(
            connection_id,
            method,
            params=params,
        )

    @server.tool(
        name=name("get_events"),
        title="Get Events",
        description="Return locally recorded broker/platform events.",
    )
    async def get_events() -> list[AdminEventPayload]:
        return await handlers.get_broker_events()