|
3 | 3 | import inspect |
4 | 4 | import json |
5 | 5 | import re |
| 6 | +from collections.abc import AsyncIterator |
| 7 | +from contextlib import ( |
| 8 | + AbstractAsyncContextManager, |
| 9 | + asynccontextmanager, |
| 10 | +) |
6 | 11 | from itertools import chain |
7 | | -from typing import Any, Callable, Literal, Sequence |
| 12 | +from typing import Any, Callable, Generic, Literal, Sequence |
8 | 13 |
|
9 | 14 | import anyio |
10 | 15 | import pydantic_core |
|
19 | 24 | from mcp.server.fastmcp.tools import ToolManager |
20 | 25 | from mcp.server.fastmcp.utilities.logging import configure_logging, get_logger |
21 | 26 | from mcp.server.fastmcp.utilities.types import Image |
22 | | -from mcp.server.lowlevel import Server as MCPServer |
23 | 27 | from mcp.server.lowlevel.helper_types import ReadResourceContents |
| 28 | +from mcp.server.lowlevel.server import ( |
| 29 | + LifespanResultT, |
| 30 | +) |
| 31 | +from mcp.server.lowlevel.server import ( |
| 32 | + Server as MCPServer, |
| 33 | +) |
| 34 | +from mcp.server.lowlevel.server import ( |
| 35 | + lifespan as default_lifespan, |
| 36 | +) |
24 | 37 | from mcp.server.sse import SseServerTransport |
25 | 38 | from mcp.server.stdio import stdio_server |
26 | 39 | from mcp.shared.context import RequestContext |
|
50 | 63 | logger = get_logger(__name__) |
51 | 64 |
|
52 | 65 |
|
53 | | -class Settings(BaseSettings): |
| 66 | +class Settings(BaseSettings, Generic[LifespanResultT]): |
54 | 67 | """FastMCP server settings. |
55 | 68 |
|
56 | 69 | All settings can be configured via environment variables with the prefix FASTMCP_. |
@@ -85,13 +98,36 @@ class Settings(BaseSettings): |
85 | 98 | description="List of dependencies to install in the server environment", |
86 | 99 | ) |
87 | 100 |
|
| 101 | + lifespan: ( |
| 102 | + Callable[["FastMCP"], AbstractAsyncContextManager[LifespanResultT]] | None |
| 103 | + ) = Field(None, description="Lifespan contexte manager") |
| 104 | + |
| 105 | + |
| 106 | +def lifespan_wrapper( |
| 107 | + app: "FastMCP", |
| 108 | + lifespan: Callable[["FastMCP"], AbstractAsyncContextManager[LifespanResultT]], |
| 109 | +) -> Callable[[MCPServer], AbstractAsyncContextManager[object]]: |
| 110 | + @asynccontextmanager |
| 111 | + async def wrap(s: MCPServer) -> AsyncIterator[object]: |
| 112 | + async with lifespan(app) as context: |
| 113 | + yield context |
| 114 | + |
| 115 | + return wrap |
| 116 | + |
88 | 117 |
|
89 | 118 | class FastMCP: |
90 | 119 | def __init__( |
91 | 120 | self, name: str | None = None, instructions: str | None = None, **settings: Any |
92 | 121 | ): |
93 | 122 | self.settings = Settings(**settings) |
94 | | - self._mcp_server = MCPServer(name=name or "FastMCP", instructions=instructions) |
| 123 | + |
| 124 | + self._mcp_server = MCPServer( |
| 125 | + name=name or "FastMCP", |
| 126 | + instructions=instructions, |
| 127 | + lifespan=lifespan_wrapper(self, self.settings.lifespan) |
| 128 | + if self.settings.lifespan |
| 129 | + else default_lifespan, |
| 130 | + ) |
95 | 131 | self._tool_manager = ToolManager( |
96 | 132 | warn_on_duplicate_tools=self.settings.warn_on_duplicate_tools |
97 | 133 | ) |
|
0 commit comments