fix: simplify context injection tests and improve server implementation

- Use find_context_parameter utility in server.py for cleaner code
- Simplify test assertions to check for context injection consistently
- Remove conditional checks and try-catch blocks in tests
- Improve readability by removing unnecessary complexity

Co-Authored-By: David S <[email protected]>

Author: dsp-ant

src/mcp/server/fastmcp/server.py

@@ -30,6 +30,7 @@
 from mcp.server.fastmcp.prompts import Prompt, PromptManager
 from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager
 from mcp.server.fastmcp.tools import Tool, ToolManager
+from mcp.server.fastmcp.utilities.context_injection import find_context_parameter
 from mcp.server.fastmcp.utilities.logging import configure_logging, get_logger
 from mcp.server.lowlevel.helper_types import ReadResourceContents
 from mcp.server.lowlevel.server import LifespanResultT
@@ -511,31 +512,19 @@ async def get_weather(city: str) -> str:
 
         def decorator(fn: AnyFunction) -> AnyFunction:
             # Check if this should be a template
+            sig = inspect.signature(fn)
             has_uri_params = "{" in uri and "}" in uri
-            has_func_params = bool(inspect.signature(fn).parameters)
+            has_func_params = bool(sig.parameters)
 
             if has_uri_params or has_func_params:
                 # Check for Context parameter to exclude from validation
-                from typing import get_origin
-
-                sig = inspect.signature(fn)
-                context_param = None
-                for param_name, param in sig.parameters.items():
-                    if param.annotation is not inspect.Parameter.empty:
-                        # Check if it's a Context type
-                        if get_origin(param.annotation) is None:
-                            # Try to check if it's a Context without importing
-                            # (to avoid circular imports)
-                            annotation_str = str(param.annotation)
-                            if "Context" in annotation_str:
-                                context_param = param_name
-                                break
+                context_param = find_context_parameter(fn)
 
                 # Validate that URI params match function params (excluding context)
                 uri_params = set(re.findall(r"{(\w+)}", uri))
-                func_params = set(sig.parameters.keys())
-                if context_param:
-                    func_params.discard(context_param)
+                # We need to remove the context_param from the resource function if
+                # there is any.
+                func_params = {p for p in sig.parameters.keys() if p != context_param}
 
                 if uri_params != func_params:
                     raise ValueError(
@@ -1001,8 +990,7 @@ async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -
             if not prompt:
                 raise ValueError(f"Unknown prompt: {name}")
 
-            context = self.get_context()
-            messages = await prompt.render(arguments, context=context)
+            messages = await prompt.render(arguments, context=self.get_context())
 
             return GetPromptResult(
                 description=prompt.description,

tests/server/fastmcp/test_server.py

@@ -958,14 +958,7 @@ async def test_resource_with_context(self):
         def resource_with_context(name: str, ctx: Context[ServerSession, None]) -> str:
             """Resource that receives context."""
             assert ctx is not None
-            # Context should be provided even if request_id might not be accessible
-            # in all contexts
-            try:
-                request_id = ctx.request_id
-                return f"Resource {name} - request_id: {request_id}"
-            except (AttributeError, ValueError):
-                # Context was injected but request context not available
-                return f"Resource {name} - context injected"
+            return f"Resource {name} - context injected"
 
         # Verify template has context_kwarg set
         templates = mcp._resource_manager.list_templates()
@@ -981,8 +974,7 @@ def resource_with_context(name: str, ctx: Context[ServerSession, None]) -> str:
             content = result.contents[0]
             assert isinstance(content, TextResourceContents)
             # Should have either request_id or indication that context was injected
-            assert "Resource test" in content.text
-            assert "request_id:" in content.text or "context injected" in content.text
+            assert "Resource test - context injected" == content.text
 
     @pytest.mark.anyio
     async def test_resource_without_context(self):
@@ -998,8 +990,7 @@ def resource_no_context(name: str) -> str:
         templates = mcp._resource_manager.list_templates()
         assert len(templates) == 1
         template = templates[0]
-        if hasattr(template, "context_kwarg"):
-            assert template.context_kwarg is None
+        assert template.context_kwarg is None
 
         # Test via client
         async with client_session(mcp._mcp_server) as client:
@@ -1024,8 +1015,7 @@ def resource_custom_ctx(id: str, my_ctx: Context[ServerSession, None]) -> str:
         templates = mcp._resource_manager.list_templates()
         assert len(templates) == 1
         template = templates[0]
-        if hasattr(template, "context_kwarg"):
-            assert template.context_kwarg == "my_ctx"
+        assert template.context_kwarg == "my_ctx"
 
         # Test via client
         async with client_session(mcp._mcp_server) as client:
@@ -1044,42 +1034,21 @@ async def test_prompt_with_context(self):
         def prompt_with_context(text: str, ctx: Context[ServerSession, None]) -> str:
             """Prompt that expects context."""
             assert ctx is not None
-            try:
-                request_id = ctx.request_id
-                return f"Prompt '{text}' with context: {request_id}"
-            except (AttributeError, ValueError):
-                return f"Prompt '{text}' - context injected"
+            return f"Prompt '{text}' - context injected"
 
         # Check if prompt has context parameter detection
         prompts = mcp._prompt_manager.list_prompts()
         assert len(prompts) == 1
-        prompt = prompts[0]
-
-        # Check if context_kwarg attribute exists (for future implementation)
-        has_context_kwarg = hasattr(prompt, "context_kwarg")
 
         # Test via client
         async with client_session(mcp._mcp_server) as client:
-            try:
-                # Try calling without passing ctx explicitly
-                result = await client.get_prompt("prompt_with_ctx", {"text": "test"})
-                # If this succeeds, check if context was injected
-                assert len(result.messages) == 1
-                message = result.messages[0]
-                content = message.content
-                assert isinstance(content, TextContent)
-                if "context injected" in content.text or "with context:" in content.text:
-                    # Context injection is working for prompts
-                    assert has_context_kwarg, "Prompt should have context_kwarg attribute"
-                else:
-                    # Context was not injected
-                    pytest.skip("Prompt context injection not yet implemented")
-            except Exception as e:
-                if "Missing required arguments" in str(e) and "ctx" in str(e):
-                    # Context injection not working - expected for now
-                    pytest.skip("Prompt context injection not yet implemented")
-                else:
-                    raise
+            # Try calling without passing ctx explicitly
+            result = await client.get_prompt("prompt_with_ctx", {"text": "test"})
+            # If this succeeds, check if context was injected
+            assert len(result.messages) == 1
+            content = result.messages[0].content
+            assert isinstance(content, TextContent)
+            assert "Prompt 'test' - context injected" in content.text
 
     @pytest.mark.anyio
     async def test_prompt_without_context(self):