165 lines
4.8 KiB
Python
165 lines
4.8 KiB
Python
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_util 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,
|
|
query_store=self.database,
|
|
query_allow_write=config.query.allow_write,
|
|
query_max_rows=config.query.max_rows,
|
|
)
|
|
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
|
|
|