From df2d3a57c27ff8ee1d25e089198c80bade52318f Mon Sep 17 00:00:00 2001
From: David Soria Parra <167242713+dsp-ant@users.noreply.github.com>
Date: Wed, 26 Mar 2025 14:18:41 +0000
Subject: [PATCH] Fix #355: Fix type error with lifespan context (#368)

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
---
 src/mcp/server/fastmcp/server.py    |  6 ++--
 src/mcp/shared/context.py           |  2 +-
 src/mcp/shared/progress.py          |  5 +--
 tests/issues/test_355_type_error.py | 50 +++++++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 6 deletions(-)
 create mode 100644 tests/issues/test_355_type_error.py

diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
index c4b282c1..bf0ce880 100644
--- a/src/mcp/server/fastmcp/server.py
+++ b/src/mcp/server/fastmcp/server.py
@@ -652,9 +652,9 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
         Returns:
             The resource content as either text or bytes
         """
-        assert self._fastmcp is not None, (
-            "Context is not available outside of a request"
-        )
+        assert (
+            self._fastmcp is not None
+        ), "Context is not available outside of a request"
         return await self._fastmcp.read_resource(uri)
 
     async def log(
diff --git a/src/mcp/shared/context.py b/src/mcp/shared/context.py
index 63759ca4..ae85d3a1 100644
--- a/src/mcp/shared/context.py
+++ b/src/mcp/shared/context.py
@@ -7,7 +7,7 @@ from mcp.shared.session import BaseSession
 from mcp.types import RequestId, RequestParams
 
 SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
-LifespanContextT = TypeVar("LifespanContextT", default=None)
+LifespanContextT = TypeVar("LifespanContextT")
 
 
 @dataclass
diff --git a/src/mcp/shared/progress.py b/src/mcp/shared/progress.py
index db502cf0..52e0017d 100644
--- a/src/mcp/shared/progress.py
+++ b/src/mcp/shared/progress.py
@@ -5,7 +5,7 @@ from typing import Generic
 
 from pydantic import BaseModel
 
-from mcp.shared.context import RequestContext
+from mcp.shared.context import LifespanContextT, RequestContext
 from mcp.shared.session import (
     BaseSession,
     ReceiveNotificationT,
@@ -60,7 +60,8 @@ def progress(
             SendResultT,
             ReceiveRequestT,
             ReceiveNotificationT,
-        ]
+        ],
+        LifespanContextT,
     ],
     total: float | None = None,
 ) -> Generator[
diff --git a/tests/issues/test_355_type_error.py b/tests/issues/test_355_type_error.py
new file mode 100644
index 00000000..91416e5c
--- /dev/null
+++ b/tests/issues/test_355_type_error.py
@@ -0,0 +1,50 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+
+from mcp.server.fastmcp import Context, FastMCP
+
+
+class Database:  # Replace with your actual DB type
+    @classmethod
+    async def connect(cls):
+        return cls()
+
+    async def disconnect(self):
+        pass
+
+    def query(self):
+        return "Hello, World!"
+
+
+# Create a named server
+mcp = FastMCP("My App")
+
+
+@dataclass
+class AppContext:
+    db: Database
+
+
+@asynccontextmanager
+async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
+    """Manage application lifecycle with type-safe context"""
+    # Initialize on startup
+    db = await Database.connect()
+    try:
+        yield AppContext(db=db)
+    finally:
+        # Cleanup on shutdown
+        await db.disconnect()
+
+
+# Pass lifespan to server
+mcp = FastMCP("My App", lifespan=app_lifespan)
+
+
+# Access type-safe lifespan context in tools
+@mcp.tool()
+def query_db(ctx: Context) -> str:
+    """Tool that uses initialized resources"""
+    db = ctx.request_context.lifespan_context.db
+    return db.query()