4
4
import pytest
5
5
from pydantic import BaseModel
6
6
7
+ from mcp .server .fastmcp import Context , FastMCP
7
8
from mcp .server .fastmcp .exceptions import ToolError
8
9
from mcp .server .fastmcp .tools import ToolManager
10
+ from mcp .server .session import ServerSessionT
11
+ from mcp .shared .context import LifespanContextT
9
12
10
13
11
14
class TestAddTools :
@@ -194,8 +197,6 @@ def concat_strs(vals: list[str] | str) -> str:
194
197
195
198
@pytest .mark .anyio
196
199
async def test_call_tool_with_complex_model (self ):
197
- from mcp .server .fastmcp import Context
198
-
199
200
class MyShrimpTank (BaseModel ):
200
201
class Shrimp (BaseModel ):
201
202
name : str
@@ -223,8 +224,6 @@ def name_shrimp(tank: MyShrimpTank, ctx: Context) -> list[str]:
223
224
class TestToolSchema :
224
225
@pytest .mark .anyio
225
226
async def test_context_arg_excluded_from_schema (self ):
226
- from mcp .server .fastmcp import Context
227
-
228
227
def something (a : int , ctx : Context ) -> int :
229
228
return a
230
229
@@ -241,7 +240,6 @@ class TestContextHandling:
241
240
def test_context_parameter_detection (self ):
242
241
"""Test that context parameters are properly detected in
243
242
Tool.from_function()."""
244
- from mcp .server .fastmcp import Context
245
243
246
244
def tool_with_context (x : int , ctx : Context ) -> str :
247
245
return str (x )
@@ -256,10 +254,17 @@ def tool_without_context(x: int) -> str:
256
254
tool = manager .add_tool (tool_without_context )
257
255
assert tool .context_kwarg is None
258
256
257
+ def tool_with_parametrized_context (
258
+ x : int , ctx : Context [ServerSessionT , LifespanContextT ]
259
+ ) -> str :
260
+ return str (x )
261
+
262
+ tool = manager .add_tool (tool_with_parametrized_context )
263
+ assert tool .context_kwarg == "ctx"
264
+
259
265
@pytest .mark .anyio
260
266
async def test_context_injection (self ):
261
267
"""Test that context is properly injected during tool execution."""
262
- from mcp .server .fastmcp import Context , FastMCP
263
268
264
269
def tool_with_context (x : int , ctx : Context ) -> str :
265
270
assert isinstance (ctx , Context )
@@ -276,7 +281,6 @@ def tool_with_context(x: int, ctx: Context) -> str:
276
281
@pytest .mark .anyio
277
282
async def test_context_injection_async (self ):
278
283
"""Test that context is properly injected in async tools."""
279
- from mcp .server .fastmcp import Context , FastMCP
280
284
281
285
async def async_tool (x : int , ctx : Context ) -> str :
282
286
assert isinstance (ctx , Context )
@@ -293,7 +297,6 @@ async def async_tool(x: int, ctx: Context) -> str:
293
297
@pytest .mark .anyio
294
298
async def test_context_optional (self ):
295
299
"""Test that context is optional when calling tools."""
296
- from mcp .server .fastmcp import Context
297
300
298
301
def tool_with_context (x : int , ctx : Context | None = None ) -> str :
299
302
return str (x )
@@ -307,7 +310,6 @@ def tool_with_context(x: int, ctx: Context | None = None) -> str:
307
310
@pytest .mark .anyio
308
311
async def test_context_error_handling (self ):
309
312
"""Test error handling when context injection fails."""
310
- from mcp .server .fastmcp import Context , FastMCP
311
313
312
314
def tool_with_context (x : int , ctx : Context ) -> str :
313
315
raise ValueError ("Test error" )
0 commit comments