mini-db-mcp/src/mini_db_mcp/server.py

from . import __version__

import json
import logging
from typing import Any, Protocol
from mcp.server import Server
from mcp.types import (
    CallToolResult,
    GetPromptResult,
    ListPromptsResult,
    ListToolsResult,
    PaginatedRequestParams,
    TextContent,
    Resource,
    Tool,
    Prompt
)

from mini_db_mcp.config import Config
from mini_db_mcp.db import Database
from mini_db_mcp.domain.domain_catalog import DOMAIN_CATALOG
from mini_db_mcp.domain.domain_handler_factory import create_domain_registry
from mini_db_mcp.utils.pagination import paginate_list
from mini_db_mcp.prompts import PromptsManager
from mini_db_mcp.tools import ToolsManager


NAME = "mini-db-mcp"


logger = logging.getLogger(__name__)


class ResourcesProtocol(Protocol):
    async def list_resources(self) -> list[Resource]: ...

    async def read_resource(self, uri: str) -> str: ...


class ToolsProtocol(Protocol):
    async def list_tools(self) -> list[Tool]: ...

    async def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]: ...


class PromptsProtocol(Protocol):
    async def list_prompts(self) -> list[Prompt]: ...

    async def get_prompt(
        self, name: str, arguments: dict[str, Any]
    ) -> GetPromptResult: ...


class MiniServer():

    def __init__(self, config: Config) -> None:
        self.config = config
        self.prompts_manager = PromptsManager()
        self.database = Database(config)
        self.handler_registry = create_domain_registry(
            metadata_store=self.database,
        )
        self.tools_manager = ToolsManager(
            catalog=DOMAIN_CATALOG,
            registry=self.handler_registry,
        )
        self.server = _create_mcp_server(
            prompts_manager=self.prompts_manager,
            tools_manager=self.tools_manager,
        )

    async def start_stdio(self) -> None:
        from mcp.server.stdio import stdio_server

        try:
            async with stdio_server() as streams:
                read_stream, write_stream = streams
                await self.server.run(
                    read_stream,
                    write_stream,
                    self.server.create_initialization_options(),
                )
        finally:
            await self.database.close()

    async def start_http(self) -> None:
        """Serve MCP over stateless Streamable HTTP at ``/mcp``."""
        import uvicorn

        app = self.server.streamable_http_app(
            streamable_http_path="/mcp",
            json_response=True,
            stateless_http=True,
            host=self.config.server.host,
        )
        server = uvicorn.Server(
            uvicorn.Config(
                app,
                host=self.config.server.host,
                port=self.config.server.port,
                log_level=self.config.server.log_level.lower(),
            )
        )
        try:
            await server.serve()
        finally:
            await self.database.close()


def _create_mcp_server(
        *,
        prompts_manager: PromptsProtocol | None = None,
        tools_manager: ToolsProtocol | None = None,
) -> Server:
    
    async def on_list_prompts(ctx, params):
        return ListPromptsResult(prompts=await prompts_manager.list_prompts())

    async def on_get_prompt(ctx, params):
        return await prompts_manager.get_prompt(params.name, dict(params.arguments or {}))

    async def list_tools(
        ctx,
        params: PaginatedRequestParams | None,
    ) -> ListToolsResult:
        tools = await tools_manager.list_tools()
        page_items, next_cursor = paginate_list(
            tools,
            cursor=params.cursor if params else None,
            page_size=10,
            identifier=lambda tool: tool.name,
        )
        logger.info("Returning %d of %d tools", len(page_items), len(tools))
        return ListToolsResult(
            tools=page_items,
            next_cursor=next_cursor,
        )

    async def call_tool(ctx, params):
        # ToolsManager returns either a manifest (discovery) or a child result.
        result = await tools_manager.call_tool(
            params.name,
            dict(params.arguments or {}),
        )
        # Publish the same payload as structured data and backward-compatible text.
        return CallToolResult(
            content=[TextContent(type="text", text=json.dumps(result, ensure_ascii=False))],
            structuredContent=result,
            isError=result.get("mode") == "error",
        )

    server = Server(
        name=NAME,
        version=__version__,
        on_list_prompts=on_list_prompts,
        on_get_prompt=on_get_prompt,
        on_list_tools=list_tools,
        on_call_tool=call_tool,
    )
    return server