Pass initialization options from run to a server session

Author: dsp-ant

mcp_python/server/__init__.py

@@ -8,7 +8,7 @@
 from pydantic import AnyUrl
 
 from mcp_python.server import types
-from mcp_python.server.session import ServerSession
+from mcp_python.server.session import InitializationOptions, ServerSession
 from mcp_python.server.stdio import stdio_server as stdio_server
 from mcp_python.shared.context import RequestContext
 from mcp_python.shared.session import RequestResponder
@@ -32,6 +32,7 @@
     ReadResourceResult,
     Resource,
     ResourceReference,
+    ServerCapabilities,
     ServerResult,
     SetLevelRequest,
     SubscribeRequest,
@@ -40,7 +41,6 @@
 
 logger = logging.getLogger(__name__)
 
-
 request_ctx: contextvars.ContextVar[RequestContext] = contextvars.ContextVar(
     "request_ctx"
 )
@@ -276,13 +276,27 @@ async def handler(req: CompleteRequest):
 
         return decorator
 
+    def get_capabilities(self) -> ServerCapabilities:
+        """Convert existing handlers to a ServerCapabilities object."""
+        capabilities: dict[str, dict[str, Any] | None] = {
+            "prompts": {} if ListPromptsRequest in self.request_handlers else None,
+            "resources": {} if ListResourcesRequest in self.request_handlers else None,
+            "tools": {} if ListPromptsRequest in self.request_handlers else None,
+            "logging": {} if SetLevelRequest in self.request_handlers else None,
+        }
+
+        capabilities = {k: v for k, v in capabilities.items() if v is not None}
+
+        return ServerCapabilities(**capabilities)
+
     async def run(
         self,
         read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception],
         write_stream: MemoryObjectSendStream[JSONRPCMessage],
+        initialization_options: InitializationOptions
     ):
         with warnings.catch_warnings(record=True) as w:
-            async with ServerSession(read_stream, write_stream) as session:
+            async with ServerSession(read_stream, write_stream, initialization_options) as session:
                 async for message in session.incoming_messages:
                     logger.debug(f"Received message: {message}")