mcp-server-llmling

"""MCP protocol server implementation.""" from __future__ import annotations import asyncio from collections import defaultdict import contextlib from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any, Literal, Self from fastmcp import FastMCP from llmling.config.manager import ConfigManager from llmling.config.runtime import RuntimeConfig from mcp.server import NotificationOptions from pydantic import AnyUrl from mcp_server_llmling import constants from mcp_server_llmling.handlers import register_handlers from mcp_server_llmling.log import get_logger from mcp_server_llmling.server_federation import FederatedServers, ServerFederation if TYPE_CHECKING: from collections.abc import AsyncIterator, Callable, Coroutine from contextlib import AbstractAsyncContextManager import os from fastmcp.server.auth import AuthProvider from llmling.config.models import BaseResource from llmling.prompts.models import BasePrompt from llmling.tools.base import LLMCallableTool import mcp from mcp import Implementation from mcp.server.lowlevel.server import LifespanResultT logger = get_logger(__name__) TransportType = Literal["stdio", "streamable-http", "sse"] class LLMLingServer: """MCP protocol server implementation.""" def __init__( self, runtime: RuntimeConfig, *, transport: TransportType = "stdio", name: str = constants.SERVER_NAME, instructions: str | None = None, tags: set[str] | None = None, dependencies: list[str] | None = None, tool_serializer: Callable[[Any], str] | None = None, lifespan: ( Callable[ [FastMCP[LifespanResultT]], AbstractAsyncContextManager[LifespanResultT], ] | None ) = None, auth_server_provider: AuthProvider | None = None, # event_store: EventStore | None = None, transport_options: dict[str, Any] | None = None, enable_injection: bool = False, injection_port: int = 8765, zed_mode: bool = False, ) -> None: """Initialize server with runtime configuration. Args: runtime: Fully initialized runtime configuration transport: Transport type to use ("stdio", "streamable-http", or "sse") name: Server name for MCP protocol instructions: Instructions for client tags: Tags for server dependencies: Dependencies for server tool_serializer: Tool serializer function lifespan: Lifespan function for server auth_server_provider: Auth server provider transport_options: Additional options for transport enable_injection: Whether to enable config injection injection_port: Port for injection server zed_mode: Enable Zed editor compatibility mode """ self.name = name self.runtime = runtime self.zed_mode = zed_mode if zed_mode: from mcp_server_llmling.zed_wrapper import prepare_runtime_for_zed prepare_runtime_for_zed(runtime) self._subscriptions: defaultdict[str, set[mcp.ServerSession]] = defaultdict(set) self._tasks: set[asyncio.Task[Any]] = set() self.federation = ServerFederation() self.transport: TransportType = transport # Create MCP server self.fastmcp = FastMCP( name, instructions=instructions, lifespan=lifespan, auth=auth_server_provider, include_tags=tags, dependencies=dependencies, tool_serializer=tool_serializer, ) self.server = self.fastmcp._mcp_server self.server.notification_options = NotificationOptions( prompts_changed=True, resources_changed=True, tools_changed=True, ) # Setup injection if enabled self.injection_server = None if enable_injection and transport == "stdio": from mcp_server_llmling.injection import ConfigInjectionServer self.injection_server = ConfigInjectionServer(self, port=injection_port) self._setup_handlers() self._setup_events() @classmethod @asynccontextmanager async def from_config_file( cls, config_path: str | os.PathLike[str], *, transport: TransportType = "stdio", name: str = constants.SERVER_NAME, transport_options: dict[str, Any] | None = None, ) -> AsyncIterator[LLMLingServer]: """Create and run server from config file with proper context management.""" manager = ConfigManager.load(config_path) async with RuntimeConfig.from_config(manager.config) as runtime: server = cls( runtime, transport=transport, name=name, transport_options=transport_options, ) try: yield server finally: await server.shutdown() def get_client_info(self) -> Implementation: """Get information about the connected client. Returns: Implementation: Client name and version information {name: str, version: str} Raises: RuntimeError: If there is no active client connection or client info not available """ try: session = self.current_session if not session.client_params: msg = "No client initialization parameters available" raise RuntimeError(msg) # noqa: TRY301 except Exception as exc: msg = "Failed to get client information" raise RuntimeError(msg) from exc else: return session.client_params.clientInfo def _create_task(self, coro: Coroutine[None, None, Any]) -> asyncio.Task[Any]: """Create and track an asyncio task.""" task = asyncio.create_task(coro) self._tasks.add(task) task.add_done_callback(self._tasks.discard) return task def _setup_handlers(self) -> None: """Register MCP protocol handlers.""" register_handlers(self) def _setup_events(self) -> None: """Set up registry event handlers.""" # Resource events resource_registry = self.runtime._resource_registry resource_registry.events.added.connect(self._handle_resource_added) resource_registry.events.removed.connect(self._handle_resource_removed) resource_registry.events.changed.connect(self._handle_resource_modified) # Prompt events prompt_registry = self.runtime._prompt_registry prompt_registry.events.added.connect(self._handle_prompt_change) prompt_registry.events.removed.connect(self._handle_prompt_change) prompt_registry.events.changed.connect(self._handle_prompt_change) # Tool events tool_registry = self.runtime._tool_registry tool_registry.events.added.connect(self._handle_tool_change) tool_registry.events.removed.connect(self._handle_tool_change) tool_registry.events.changed.connect(self._handle_tool_change) async def start(self, *, raise_exceptions: bool = False) -> None: """Start the server.""" try: if (extra := self.runtime._config.model_extra) and ( external_servers := extra.get("external_servers") ): config = FederatedServers(external_servers=external_servers) await self.federation.connect_servers(config) # Start injection server in a separate task if enabled injection_task = None if self.injection_server: try: # Create task but don't await it directly injection_task = asyncio.create_task(self.injection_server.start()) msg = "Config injection server listening on port %d" logger.info(msg, self.injection_server.port) except Exception: logger.exception("Failed to start injection server") if raise_exceptions: raise # Run main server await self.fastmcp.run_async(transport=self.transport, show_banner=False) finally: if injection_task: injection_task.cancel() with contextlib.suppress(asyncio.CancelledError): await injection_task await self.shutdown() async def shutdown(self) -> None: """Shutdown the server.""" try: if self.injection_server: await self.injection_server.stop() # await self.server.shutdown() await self.federation.close() # Cancel all pending tasks if self._tasks: for task in self._tasks: task.cancel() await asyncio.gather(*self._tasks, return_exceptions=True) self._tasks.clear() # Shutdown runtime await self.runtime.shutdown() finally: self._tasks.clear() async def __aenter__(self) -> Self: """Async context manager entry.""" return self async def __aexit__(self, *exc: object) -> None: """Async context manager exit.""" await self.shutdown() @property def current_session(self) -> mcp.ServerSession: """Get client info from request context.""" try: return self.server.request_context.session except LookupError as exc: msg = "No active request context" raise RuntimeError(msg) from exc @property def client_info(self) -> mcp.Implementation | None: """Get current session from request context.""" session = self.current_session if not session.client_params: return None return session.client_params.clientInfo def notify_progress( self, token: str, progress: float, total: float | None = None, description: str | None = None, ) -> None: """Send progress notification to client.""" try: session = self.current_session task = session.send_progress_notification( progress_token=token, progress=progress, total=total, ) self._create_task(task) # Optionally send description as log message if description: coro = session.send_log_message(level="info", data=description) self._create_task(coro) except Exception: logger.exception("Failed to send progress notification") async def notify_resource_list_changed(self) -> None: """Notify clients about resource list changes.""" try: await self.current_session.send_resource_list_changed() except RuntimeError: logger.debug("No active session for notification") except Exception: logger.exception("Failed to send resource list change notification") async def notify_resource_change(self, uri: str) -> None: """Notify subscribers about resource changes.""" if uri in self._subscriptions: try: await self.current_session.send_resource_updated(AnyUrl(uri)) except Exception: msg = "Failed to notify subscribers about resource change: %s" logger.exception(msg, uri) async def notify_prompt_list_changed(self) -> None: """Notify clients about prompt list changes.""" try: self._create_task(self.current_session.send_prompt_list_changed()) except RuntimeError: logger.debug("No active session for notification") except Exception: logger.exception("Failed to send prompt list change notification") async def notify_tool_list_changed(self) -> None: """Notify clients about tool list changes.""" try: self._create_task(self.current_session.send_tool_list_changed()) except RuntimeError: logger.debug("No active session for notification") except Exception: logger.exception("Failed to send tool list change notification") def _handle_resource_added(self, key: str, resource: BaseResource) -> None: """Handle resource addition.""" self._create_task(self.notify_resource_list_changed()) def _handle_resource_modified(self, key: str, resource: BaseResource) -> None: """Handle resource modification.""" loader = self.runtime.get_resource_loader(resource) uri = loader.create_uri(name=key) self._create_task(self.notify_resource_change(uri)) def _handle_resource_removed(self, key: str, resource: BaseResource) -> None: """Handle resource removal.""" self._create_task(self.notify_resource_list_changed()) def _handle_prompt_change(self, key: str, prompt: BasePrompt) -> None: """Handle any prompt change.""" self._create_task(self.notify_prompt_list_changed()) def _handle_tool_change(self, key: str, tool: LLMCallableTool) -> None: """Handle any tool change.""" self._create_task(self.notify_tool_list_changed()) if __name__ == "__main__": import sys from llmling import config_resources config_path = sys.argv[1] if len(sys.argv) > 1 else config_resources.TEST_CONFIG