PM-pinou
85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
"""MCP Server embedded in Django.
|
|
|
|
Uses asyncio stdio transport. Registered via the Django management command
|
|
``python manage.py start_mcp``.
|
|
|
|
Django ORM initialisation happens at the management-command level so that
|
|
the server thread has access to models when needed (e.g. during feature
|
|
extraction).
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
|
|
from mcp.server import Server
|
|
from mcp.server.stdio import stdio_server
|
|
from mcp.types import TextContent
|
|
|
|
# Ensure Django is configured when this module is imported from the
|
|
# management command. The management command calls django.setup()
|
|
# before constructing the server.
|
|
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'tianxuan.settings')
|
|
|
|
from .session_store import SessionStore
|
|
from .tool_registry import get_tools_meta, handle_call
|
|
|
|
|
|
class TLSAnalyzerMCPServer:
|
|
"""Embedded MCP server for the TLS Analyzer.
|
|
|
|
Exposes 11 tools over stdio transport. The :class:`SessionStore`
|
|
singleton holds datasets and analysis results in memory for the
|
|
lifetime of the server process.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.server = Server("tls-analyzer")
|
|
self.store = SessionStore()
|
|
self._setup_tools()
|
|
|
|
# ── tool registration ───────────────────────────────────────────────
|
|
|
|
def _setup_tools(self) -> None:
|
|
"""Register all tools with the MCP server via the SDK decorator API.
|
|
|
|
Uses :func:`get_tools_meta` for tool metadata (name, description,
|
|
inputSchema) and :func:`handle_call` for request routing.
|
|
"""
|
|
tools_meta = get_tools_meta()
|
|
_tools_map = {t.name: t for t in tools_meta}
|
|
|
|
@self.server.list_tools()
|
|
async def list_tools():
|
|
return tools_meta
|
|
|
|
@self.server.call_tool()
|
|
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
|
|
try:
|
|
result = await handle_call(name, arguments)
|
|
return [TextContent(type="text", text=json.dumps(result, default=str))]
|
|
except Exception as exc:
|
|
error_body = json.dumps({
|
|
'error': str(exc),
|
|
'tool': name,
|
|
'truncated': False,
|
|
}, default=str)
|
|
return [TextContent(type="text", text=error_body)]
|
|
|
|
# ── lifecycle ───────────────────────────────────────────────────────
|
|
|
|
async def run(self) -> None:
|
|
"""Start the MCP server with stdio transport.
|
|
|
|
Blocks until the transport is closed.
|
|
"""
|
|
async with stdio_server() as (read, write):
|
|
await self.server.run(
|
|
read,
|
|
write,
|
|
self.server.create_initialization_options(),
|
|
)
|
|
|
|
def run_sync(self) -> None:
|
|
"""Synchronous entry point for the management command."""
|
|
asyncio.run(self.run())
|