diff --git a/docs/docs.json b/docs/docs.json index 7a9b88311..e833496c6 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -124,6 +124,7 @@ "servers/context", "servers/elicitation", "servers/icons", + "servers/lifespan", "servers/logging", "servers/middleware", "servers/progress", diff --git a/docs/integrations/fastapi.mdx b/docs/integrations/fastapi.mdx index bb1ec2e37..2ce28336d 100644 --- a/docs/integrations/fastapi.mdx +++ b/docs/integrations/fastapi.mdx @@ -410,41 +410,31 @@ If you need CORS on your own FastAPI routes, use the sub-app pattern: mount your ### Combining Lifespans -If your FastAPI app already has a lifespan (for database connections, startup tasks, etc.), you can't simply replace it with the MCP lifespan. Instead, you need to create a new lifespan function that manages both contexts. This ensures that both your app's initialization logic and the MCP server's session manager run properly: +If your FastAPI app already has a lifespan (for database connections, startup tasks, etc.), you can't simply replace it with the MCP lifespan. Use `combine_lifespans` to run both: ```python -from contextlib import asynccontextmanager from fastapi import FastAPI from fastmcp import FastMCP +from fastmcp.utilities.lifespan import combine_lifespans +from contextlib import asynccontextmanager # Your existing lifespan @asynccontextmanager async def app_lifespan(app: FastAPI): - # Startup print("Starting up the app...") - # Initialize database, cache, etc. yield - # Shutdown print("Shutting down the app...") # Create MCP server mcp = FastMCP("Tools") -mcp_app = mcp.http_app(path='/mcp') +mcp_app = mcp.http_app() # Combine both lifespans -@asynccontextmanager -async def combined_lifespan(app: FastAPI): - # Run both lifespans - async with app_lifespan(app): - async with mcp_app.lifespan(app): - yield - -# Use the combined lifespan -app = FastAPI(lifespan=combined_lifespan) -app.mount("/mcp", mcp_app) +app = FastAPI(lifespan=combine_lifespans(app_lifespan, mcp_app.lifespan)) +app.mount("/mcp", mcp_app) # MCP endpoint at /mcp ``` -This pattern ensures both your app's initialization logic and the MCP server's session manager are properly managed. The key is using nested `async with` statements - the inner context (MCP) will be initialized after the outer context (your app), and cleaned up before it. This maintains the correct initialization and cleanup order for all your resources. +`combine_lifespans` enters lifespans in order and exits in reverse order. ### Performance Tips diff --git a/docs/servers/lifespan.mdx b/docs/servers/lifespan.mdx new file mode 100644 index 000000000..9171ed892 --- /dev/null +++ b/docs/servers/lifespan.mdx @@ -0,0 +1,141 @@ +--- +title: Lifespans +description: Server-level setup and teardown with composable lifespans +--- + +Lifespans let you run code once when the server starts and clean up when it stops. Unlike per-session handlers, lifespans run exactly once regardless of how many clients connect. + +## Basic Usage + +Use the `@lifespan` decorator to define a lifespan: + +```python +from fastmcp import FastMCP +from fastmcp.server.lifespan import lifespan + +@lifespan +async def app_lifespan(server): + # Setup: runs once when server starts + print("Starting up...") + try: + yield {"started_at": "2024-01-01"} + finally: + # Teardown: runs when server stops + print("Shutting down...") + +mcp = FastMCP("MyServer", lifespan=app_lifespan) +``` + +The dict you yield becomes the **lifespan context**, accessible from tools. + + +Always use `try/finally` for cleanup code to ensure it runs even if the server is cancelled. + + +## Accessing Lifespan Context + +Access the lifespan context in tools via `ctx.lifespan_context`: + +```python +from fastmcp import FastMCP, Context +from fastmcp.server.lifespan import lifespan + +@lifespan +async def app_lifespan(server): + # Initialize shared state + data = {"users": ["alice", "bob"]} + yield {"data": data} + +mcp = FastMCP("MyServer", lifespan=app_lifespan) + +@mcp.tool +def list_users(ctx: Context) -> list[str]: + data = ctx.lifespan_context["data"] + return data["users"] +``` + +## Composing Lifespans + +Compose multiple lifespans with the `|` operator: + +```python +from fastmcp import FastMCP +from fastmcp.server.lifespan import lifespan + +@lifespan +async def config_lifespan(server): + config = {"debug": True, "version": "1.0"} + yield {"config": config} + +@lifespan +async def data_lifespan(server): + data = {"items": []} + yield {"data": data} + +# Compose with | +mcp = FastMCP("MyServer", lifespan=config_lifespan | data_lifespan) +``` + +Composed lifespans: +- Enter in order (left to right) +- Exit in reverse order (right to left) +- Merge their context dicts (later values overwrite earlier on conflict) + +## Backwards Compatibility + +Existing `@asynccontextmanager` lifespans still work when passed directly to FastMCP: + +```python +from contextlib import asynccontextmanager +from fastmcp import FastMCP + +@asynccontextmanager +async def legacy_lifespan(server): + yield {"key": "value"} + +mcp = FastMCP("MyServer", lifespan=legacy_lifespan) +``` + +To compose an `@asynccontextmanager` function with `@lifespan` functions, wrap it with `ContextManagerLifespan`: + +```python +from contextlib import asynccontextmanager +from fastmcp.server.lifespan import lifespan, ContextManagerLifespan + +@asynccontextmanager +async def legacy_lifespan(server): + yield {"legacy": True} + +@lifespan +async def new_lifespan(server): + yield {"new": True} + +# Wrap the legacy lifespan explicitly for composition +combined = ContextManagerLifespan(legacy_lifespan) | new_lifespan +``` + +## With FastAPI + +When mounting FastMCP into FastAPI, use `combine_lifespans` to run both your app's lifespan and the MCP server's lifespan: + +```python +from contextlib import asynccontextmanager + +from fastapi import FastAPI +from fastmcp import FastMCP +from fastmcp.utilities.lifespan import combine_lifespans + +@asynccontextmanager +async def app_lifespan(app): + print("FastAPI starting...") + yield + print("FastAPI shutting down...") + +mcp = FastMCP("Tools") +mcp_app = mcp.http_app() + +app = FastAPI(lifespan=combine_lifespans(app_lifespan, mcp_app.lifespan)) +app.mount("/mcp", mcp_app) +``` + +See the [FastAPI integration guide](/integrations/fastapi#combining-lifespans) for full details. diff --git a/docs/servers/server.mdx b/docs/servers/server.mdx index 9db7b51d4..a7c2f4878 100644 --- a/docs/servers/server.mdx +++ b/docs/servers/server.mdx @@ -60,8 +60,8 @@ The `FastMCP` constructor accepts several arguments: Authentication provider for securing HTTP-based transports. See [Authentication](/servers/auth/authentication) for configuration options - - An async context manager function for server startup and shutdown logic + + Server-level setup and teardown logic. See [Lifespans](/servers/lifespan) for composable lifespans diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 5271b523a..3789adaaa 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -248,6 +248,29 @@ class Context: except LookupError: return None + @property + def lifespan_context(self) -> dict[str, Any]: + """Access the server's lifespan context. + + Returns the context dict yielded by the server's lifespan function. + Returns an empty dict if no lifespan was configured or if the MCP + session is not yet established. + + Example: + ```python + @server.tool + def my_tool(ctx: Context) -> str: + db = ctx.lifespan_context.get("db") + if db: + return db.query("SELECT 1") + return "No database connection" + ``` + """ + rc = self.request_context + if rc is None: + return {} + return rc.lifespan_context + async def report_progress( self, progress: float, total: float | None = None, message: str | None = None ) -> None: diff --git a/src/fastmcp/server/lifespan.py b/src/fastmcp/server/lifespan.py new file mode 100644 index 000000000..03aea526a --- /dev/null +++ b/src/fastmcp/server/lifespan.py @@ -0,0 +1,198 @@ +"""Composable lifespans for FastMCP servers. + +This module provides a `@lifespan` decorator for creating composable server lifespans +that can be combined using the `|` operator. + +Example: + ```python + from fastmcp import FastMCP + from fastmcp.server.lifespan import lifespan + + @lifespan + async def db_lifespan(server): + conn = await connect_db() + yield {"db": conn} + await conn.close() + + @lifespan + async def cache_lifespan(server): + cache = await connect_cache() + yield {"cache": cache} + await cache.close() + + mcp = FastMCP("server", lifespan=db_lifespan | cache_lifespan) + ``` + +To compose with existing `@asynccontextmanager` lifespans, wrap them explicitly: + + ```python + from contextlib import asynccontextmanager + from fastmcp.server.lifespan import lifespan, ContextManagerLifespan + + @asynccontextmanager + async def legacy_lifespan(server): + yield {"legacy": True} + + @lifespan + async def new_lifespan(server): + yield {"new": True} + + # Wrap the legacy lifespan explicitly + combined = ContextManagerLifespan(legacy_lifespan) | new_lifespan + ``` +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from fastmcp.server.server import FastMCP + + +LifespanFn = Callable[["FastMCP[Any]"], AsyncIterator[dict[str, Any] | None]] +LifespanContextManagerFn = Callable[ + ["FastMCP[Any]"], AbstractAsyncContextManager[dict[str, Any] | None] +] + + +class Lifespan: + """Composable lifespan wrapper. + + Wraps an async generator function and enables composition via the `|` operator. + The wrapped function should yield a dict that becomes part of the lifespan context. + """ + + def __init__(self, fn: LifespanFn) -> None: + """Initialize a Lifespan wrapper. + + Args: + fn: An async generator function that takes a FastMCP server and yields + a dict for the lifespan context. + """ + self._fn = fn + + @asynccontextmanager + async def __call__(self, server: FastMCP[Any]) -> AsyncIterator[dict[str, Any]]: + """Execute the lifespan as an async context manager. + + Args: + server: The FastMCP server instance. + + Yields: + The lifespan context dict. + """ + async with asynccontextmanager(self._fn)(server) as result: + yield result if result is not None else {} + + def __or__(self, other: Lifespan) -> ComposedLifespan: + """Compose with another lifespan using the | operator. + + Args: + other: Another Lifespan instance. + + Returns: + A ComposedLifespan that runs both lifespans. + + Raises: + TypeError: If other is not a Lifespan instance. + """ + if not isinstance(other, Lifespan): + raise TypeError( + f"Cannot compose Lifespan with {type(other).__name__}. " + f"Use @lifespan decorator or wrap with ContextManagerLifespan()." + ) + return ComposedLifespan(self, other) + + +class ContextManagerLifespan(Lifespan): + """Lifespan wrapper for already-wrapped context manager functions. + + Use this for functions already decorated with @asynccontextmanager. + """ + + _fn: LifespanContextManagerFn # Override type for this subclass + + def __init__(self, fn: LifespanContextManagerFn) -> None: + """Initialize with a context manager factory function.""" + self._fn = fn # type: ignore[assignment] + + @asynccontextmanager + async def __call__(self, server: FastMCP[Any]) -> AsyncIterator[dict[str, Any]]: + """Execute the lifespan as an async context manager. + + Args: + server: The FastMCP server instance. + + Yields: + The lifespan context dict. + """ + # self._fn is already a context manager factory, just call it + async with self._fn(server) as result: + yield result if result is not None else {} + + +class ComposedLifespan(Lifespan): + """Two lifespans composed together. + + Enters the left lifespan first, then the right. Exits in reverse order. + Results are shallow-merged into a single dict. + """ + + def __init__(self, left: Lifespan, right: Lifespan) -> None: + """Initialize a composed lifespan. + + Args: + left: The first lifespan to enter. + right: The second lifespan to enter. + """ + # Don't call super().__init__ since we override __call__ + self._left = left + self._right = right + + @asynccontextmanager + async def __call__(self, server: FastMCP[Any]) -> AsyncIterator[dict[str, Any]]: + """Execute both lifespans, merging their results. + + Args: + server: The FastMCP server instance. + + Yields: + The merged lifespan context dict from both lifespans. + """ + async with ( + self._left(server) as left_result, + self._right(server) as right_result, + ): + yield {**left_result, **right_result} + + +def lifespan(fn: LifespanFn) -> Lifespan: + """Decorator to create a composable lifespan. + + Use this decorator on an async generator function to make it composable + with other lifespans using the `|` operator. + + Example: + ```python + @lifespan + async def my_lifespan(server): + # Setup + resource = await create_resource() + yield {"resource": resource} + # Teardown + await resource.close() + + mcp = FastMCP("server", lifespan=my_lifespan | other_lifespan) + ``` + + Args: + fn: An async generator function that takes a FastMCP server and yields + a dict for the lifespan context. + + Returns: + A composable Lifespan wrapper. + """ + return Lifespan(fn) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index d14d24452..964864c7c 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -76,6 +76,7 @@ from fastmcp.server.http import ( create_sse_app, create_streamable_http_app, ) +from fastmcp.server.lifespan import Lifespan from fastmcp.server.low_level import LowLevelServer from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.server.providers import LocalProvider, Provider @@ -203,7 +204,7 @@ class FastMCP(Generic[LifespanResultT]): auth: AuthProvider | None = None, middleware: Sequence[Middleware] | None = None, providers: Sequence[Provider] | None = None, - lifespan: LifespanCallable | None = None, + lifespan: LifespanCallable | Lifespan | None = None, mask_error_details: bool | None = None, tools: Sequence[Tool | Callable[..., Any]] | None = None, tool_transformations: Mapping[str, ToolTransformConfig] | None = None, @@ -281,7 +282,11 @@ class FastMCP(Generic[LifespanResultT]): ) self._tool_serializer: Callable[[Any], str] | None = tool_serializer - self._lifespan: LifespanCallable[LifespanResultT] = lifespan or default_lifespan + # Handle Lifespan instances (they're callable) or regular lifespan functions + if lifespan is not None: + self._lifespan: LifespanCallable[LifespanResultT] = lifespan + else: + self._lifespan = cast(LifespanCallable[LifespanResultT], default_lifespan) self._lifespan_result: LifespanResultT | None = None self._lifespan_result_set: bool = False self._started: asyncio.Event = asyncio.Event() diff --git a/src/fastmcp/utilities/lifespan.py b/src/fastmcp/utilities/lifespan.py new file mode 100644 index 000000000..09434d736 --- /dev/null +++ b/src/fastmcp/utilities/lifespan.py @@ -0,0 +1,56 @@ +"""Lifespan utilities for combining async context manager lifespans.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable +from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager +from typing import Any, TypeVar + +AppT = TypeVar("AppT") + + +def combine_lifespans( + *lifespans: Callable[[AppT], AbstractAsyncContextManager[dict[str, Any] | None]], +) -> Callable[[AppT], AbstractAsyncContextManager[dict[str, Any]]]: + """Combine multiple lifespans into a single lifespan. + + Useful when mounting FastMCP into FastAPI and you need to run + both your app's lifespan and the MCP server's lifespan. + + Works with both FastAPI-style lifespans (yield None) and FastMCP-style + lifespans (yield dict). Results are merged; later lifespans override + earlier ones on key conflicts. + + Lifespans are entered in order and exited in reverse order (LIFO). + + Example: + ```python + from fastmcp import FastMCP + from fastmcp.utilities.lifespan import combine_lifespans + from fastapi import FastAPI + + mcp = FastMCP("Tools") + mcp_app = mcp.http_app() + + app = FastAPI(lifespan=combine_lifespans(app_lifespan, mcp_app.lifespan)) + app.mount("/mcp", mcp_app) # MCP endpoint at /mcp + ``` + + Args: + *lifespans: Lifespan context manager factories to combine. + + Returns: + A combined lifespan context manager factory. + """ + + @asynccontextmanager + async def combined(app: AppT) -> AsyncIterator[dict[str, Any]]: + merged: dict[str, Any] = {} + async with AsyncExitStack() as stack: + for ls in lifespans: + result = await stack.enter_async_context(ls(app)) + if result is not None: + merged.update(result) + yield merged + + return combined diff --git a/tests/server/test_server_lifespan.py b/tests/server/test_server_lifespan.py index 68ed395e5..a6d4bb637 100644 --- a/tests/server/test_server_lifespan.py +++ b/tests/server/test_server_lifespan.py @@ -4,8 +4,12 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from typing import Any +import pytest + from fastmcp import Client, FastMCP from fastmcp.server.context import Context +from fastmcp.server.lifespan import ContextManagerLifespan, lifespan +from fastmcp.utilities.lifespan import combine_lifespans class TestServerLifespan: @@ -62,10 +66,399 @@ class TestServerLifespan: @mcp.tool def get_db_info(ctx: Context) -> str: # Access the server lifespan context - assert ctx.request_context is not None + assert ctx.request_context is not None # type narrowing for type checker lifespan_context = ctx.request_context.lifespan_context return lifespan_context.get("db_connection", "no_db") async with Client(mcp) as client: result = await client.call_tool("get_db_info", {}) assert result.data == "mock_db" + + +class TestComposableLifespans: + """Test composable lifespan functionality.""" + + async def test_lifespan_decorator_basic(self): + """Test that the @lifespan decorator works like @asynccontextmanager.""" + events: list[str] = [] + + @lifespan + async def my_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("enter") + try: + yield {"key": "value"} + finally: + events.append("exit") + + mcp = FastMCP("TestServer", lifespan=my_lifespan) + + @mcp.tool + def get_info(ctx: Context) -> str: + assert ctx.request_context is not None + lifespan_context = ctx.request_context.lifespan_context + return lifespan_context.get("key", "missing") + + assert events == [] + + async with Client(mcp) as client: + result = await client.call_tool("get_info", {}) + assert result.data == "value" + assert events == ["enter"] + + assert events == ["enter", "exit"] + + async def test_lifespan_composition_two(self): + """Test composing two lifespans with |.""" + events: list[str] = [] + + @lifespan + async def first_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("first_enter") + try: + yield {"first": "a"} + finally: + events.append("first_exit") + + @lifespan + async def second_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("second_enter") + try: + yield {"second": "b"} + finally: + events.append("second_exit") + + composed = first_lifespan | second_lifespan + mcp = FastMCP("TestServer", lifespan=composed) + + @mcp.tool + def get_both(ctx: Context) -> dict: + assert ctx.request_context is not None + return dict(ctx.request_context.lifespan_context) + + async with Client(mcp) as client: + result = await client.call_tool("get_both", {}) + # Results should be merged + assert result.data == {"first": "a", "second": "b"} + # Should enter in order + assert events == ["first_enter", "second_enter"] + + # Should exit in reverse order (LIFO) + assert events == ["first_enter", "second_enter", "second_exit", "first_exit"] + + async def test_lifespan_composition_three(self): + """Test composing three lifespans with |.""" + events: list[str] = [] + + @lifespan + async def ls_a(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("a_enter") + try: + yield {"a": 1} + finally: + events.append("a_exit") + + @lifespan + async def ls_b(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("b_enter") + try: + yield {"b": 2} + finally: + events.append("b_exit") + + @lifespan + async def ls_c(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("c_enter") + try: + yield {"c": 3} + finally: + events.append("c_exit") + + composed = ls_a | ls_b | ls_c + mcp = FastMCP("TestServer", lifespan=composed) + + @mcp.tool + def get_all(ctx: Context) -> dict: + assert ctx.request_context is not None + return dict(ctx.request_context.lifespan_context) + + async with Client(mcp) as client: + result = await client.call_tool("get_all", {}) + assert result.data == {"a": 1, "b": 2, "c": 3} + assert events == ["a_enter", "b_enter", "c_enter"] + + assert events == [ + "a_enter", + "b_enter", + "c_enter", + "c_exit", + "b_exit", + "a_exit", + ] + + async def test_lifespan_result_merge_later_wins(self): + """Test that later lifespans overwrite earlier ones on key conflict.""" + + @lifespan + async def first(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + yield {"key": "first", "only_first": "yes"} + + @lifespan + async def second(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + yield {"key": "second", "only_second": "yes"} + + composed = first | second + mcp = FastMCP("TestServer", lifespan=composed) + + @mcp.tool + def get_context(ctx: Context) -> dict: + assert ctx.request_context is not None + return dict(ctx.request_context.lifespan_context) + + async with Client(mcp) as client: + result = await client.call_tool("get_context", {}) + # "key" should be overwritten by second + assert result.data == { + "key": "second", + "only_first": "yes", + "only_second": "yes", + } + + async def test_lifespan_composition_with_context_manager_lifespan(self): + """Test composing with ContextManagerLifespan for @asynccontextmanager functions.""" + events: list[str] = [] + + @asynccontextmanager + async def legacy_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("legacy_enter") + try: + yield {"legacy": True} + finally: + events.append("legacy_exit") + + @lifespan + async def new_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + events.append("new_enter") + try: + yield {"new": True} + finally: + events.append("new_exit") + + # Wrap the @asynccontextmanager function explicitly + composed = ContextManagerLifespan(legacy_lifespan) | new_lifespan + mcp = FastMCP("TestServer", lifespan=composed) + + @mcp.tool + def get_context(ctx: Context) -> dict: + assert ctx.request_context is not None + return dict(ctx.request_context.lifespan_context) + + async with Client(mcp) as client: + result = await client.call_tool("get_context", {}) + assert result.data == {"legacy": True, "new": True} + + assert events == [ + "legacy_enter", + "new_enter", + "new_exit", + "legacy_exit", + ] + + async def test_backwards_compatibility_asynccontextmanager(self): + """Test that existing @asynccontextmanager lifespans still work.""" + + @asynccontextmanager + async def old_style_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + yield {"old_style": True} + + mcp = FastMCP("TestServer", lifespan=old_style_lifespan) + + @mcp.tool + def get_context(ctx: Context) -> dict: + assert ctx.request_context is not None + return dict(ctx.request_context.lifespan_context) + + async with Client(mcp) as client: + result = await client.call_tool("get_context", {}) + assert result.data == {"old_style": True} + + async def test_lifespan_or_requires_lifespan_instance(self): + """Test that | operator requires Lifespan instances and gives helpful error.""" + + @lifespan + async def my_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + yield {"key": "value"} + + @asynccontextmanager + async def regular_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: + yield {"regular": True} + + # Composing with non-Lifespan should raise TypeError with helpful message + with pytest.raises(TypeError) as exc_info: + my_lifespan | regular_lifespan # type: ignore[operator] + + assert "ContextManagerLifespan" in str(exc_info.value) + + +class TestCombineLifespans: + """Test combine_lifespans utility function.""" + + async def test_combine_lifespans_fastapi_style(self): + """Test combining lifespans that yield None (FastAPI-style).""" + events: list[str] = [] + + @asynccontextmanager + async def first_lifespan(app: Any) -> AsyncIterator[None]: + events.append("first_enter") + try: + yield + finally: + events.append("first_exit") + + @asynccontextmanager + async def second_lifespan(app: Any) -> AsyncIterator[None]: + events.append("second_enter") + try: + yield + finally: + events.append("second_exit") + + combined = combine_lifespans(first_lifespan, second_lifespan) + + async with combined("mock_app") as result: + assert result == {} # Empty dict when lifespans yield None + assert events == ["first_enter", "second_enter"] + + # LIFO exit order + assert events == ["first_enter", "second_enter", "second_exit", "first_exit"] + + async def test_combine_lifespans_fastmcp_style(self): + """Test combining lifespans that yield dicts (FastMCP-style).""" + events: list[str] = [] + + @asynccontextmanager + async def db_lifespan(app: Any) -> AsyncIterator[dict[str, Any]]: + events.append("db_enter") + try: + yield {"db": "connected"} + finally: + events.append("db_exit") + + @asynccontextmanager + async def cache_lifespan(app: Any) -> AsyncIterator[dict[str, Any]]: + events.append("cache_enter") + try: + yield {"cache": "ready"} + finally: + events.append("cache_exit") + + combined = combine_lifespans(db_lifespan, cache_lifespan) + + async with combined("mock_app") as result: + assert result == {"db": "connected", "cache": "ready"} + assert events == ["db_enter", "cache_enter"] + + assert events == ["db_enter", "cache_enter", "cache_exit", "db_exit"] + + async def test_combine_lifespans_mixed_styles(self): + """Test combining FastAPI-style (yield None) and FastMCP-style (yield dict).""" + events: list[str] = [] + + @asynccontextmanager + async def fastapi_lifespan(app: Any) -> AsyncIterator[None]: + events.append("fastapi_enter") + try: + yield # FastAPI-style: yield None + finally: + events.append("fastapi_exit") + + @asynccontextmanager + async def fastmcp_lifespan(app: Any) -> AsyncIterator[dict[str, Any]]: + events.append("fastmcp_enter") + try: + yield {"mcp": "initialized"} # FastMCP-style: yield dict + finally: + events.append("fastmcp_exit") + + combined = combine_lifespans(fastapi_lifespan, fastmcp_lifespan) + + async with combined("mock_app") as result: + # Only the dict from fastmcp_lifespan should be present + assert result == {"mcp": "initialized"} + assert events == ["fastapi_enter", "fastmcp_enter"] + + assert events == [ + "fastapi_enter", + "fastmcp_enter", + "fastmcp_exit", + "fastapi_exit", + ] + + async def test_combine_lifespans_result_merge_later_wins(self): + """Test that later lifespans overwrite earlier ones on key conflict.""" + + @asynccontextmanager + async def first(app: Any) -> AsyncIterator[dict[str, Any]]: + yield {"key": "first", "only_first": "yes"} + + @asynccontextmanager + async def second(app: Any) -> AsyncIterator[dict[str, Any]]: + yield {"key": "second", "only_second": "yes"} + + combined = combine_lifespans(first, second) + + async with combined("mock_app") as result: + assert result == { + "key": "second", # Overwritten by later lifespan + "only_first": "yes", + "only_second": "yes", + } + + async def test_combine_lifespans_three(self): + """Test combining three lifespans.""" + events: list[str] = [] + + @asynccontextmanager + async def ls_a(app: Any) -> AsyncIterator[dict[str, Any]]: + events.append("a_enter") + try: + yield {"a": 1} + finally: + events.append("a_exit") + + @asynccontextmanager + async def ls_b(app: Any) -> AsyncIterator[dict[str, Any]]: + events.append("b_enter") + try: + yield {"b": 2} + finally: + events.append("b_exit") + + @asynccontextmanager + async def ls_c(app: Any) -> AsyncIterator[dict[str, Any]]: + events.append("c_enter") + try: + yield {"c": 3} + finally: + events.append("c_exit") + + combined = combine_lifespans(ls_a, ls_b, ls_c) + + async with combined("mock_app") as result: + assert result == {"a": 1, "b": 2, "c": 3} + assert events == ["a_enter", "b_enter", "c_enter"] + + assert events == [ + "a_enter", + "b_enter", + "c_enter", + "c_exit", + "b_exit", + "a_exit", + ] + + async def test_combine_lifespans_empty(self): + """Test combining zero lifespans.""" + combined = combine_lifespans() + + async with combined("mock_app") as result: + assert result == {}