diff --git a/src/mcp/server/mcpserver/utilities/context_injection.py b/src/mcp/server/mcpserver/utilities/context_injection.py index ac7ab82d05..b2bee54608 100644 --- a/src/mcp/server/mcpserver/utilities/context_injection.py +++ b/src/mcp/server/mcpserver/utilities/context_injection.py @@ -29,6 +29,10 @@ def find_context_parameter(fn: Callable[..., Any]) -> str | None: # If we can't resolve type hints, we can't find the context parameter return None + # `get_type_hints` reports the return annotation under the "return" key, which is + # not a parameter and must not be mistaken for one. + hints.pop("return", None) + # Check each parameter's type hint for param_name, annotation in hints.items(): # Handle direct Context type diff --git a/tests/server/mcpserver/test_context_injection.py b/tests/server/mcpserver/test_context_injection.py new file mode 100644 index 0000000000..955aa467f6 --- /dev/null +++ b/tests/server/mcpserver/test_context_injection.py @@ -0,0 +1,90 @@ +"""Tests for context parameter discovery and injection.""" + +from mcp_types import TextContent, TextResourceContents + +from mcp.client.client import Client +from mcp.server.mcpserver import Context, MCPServer +from mcp.server.mcpserver.utilities.context_injection import find_context_parameter + + +def test_context_parameter_is_found_by_its_annotation() -> None: + """SDK-defined: a parameter annotated with `Context` is the one context is injected into.""" + + def fn(value: int, ctx: Context) -> str: + raise NotImplementedError + + assert find_context_parameter(fn) == "ctx" + + +def test_context_return_annotation_is_not_reported_as_a_parameter() -> None: + """SDK-defined: a `Context` return annotation is not a parameter, so nothing is injected.""" + + def fn(value: int) -> Context: + raise NotImplementedError + + assert find_context_parameter(fn) is None + + +def test_context_in_a_union_return_annotation_is_not_reported_as_a_parameter() -> None: + """SDK-defined: a union return annotation mentioning `Context` is still not a parameter.""" + + def fn(value: int) -> Context | None: + raise NotImplementedError + + assert find_context_parameter(fn) is None + + +def test_context_parameter_wins_over_a_context_return_annotation() -> None: + """SDK-defined: the real parameter is found even when the return annotation is also `Context`.""" + + def fn(ctx: Context) -> Context: + raise NotImplementedError + + assert find_context_parameter(fn) == "ctx" + + +async def test_tool_returning_context_is_called_with_only_its_own_arguments() -> None: + """A tool whose return annotation mentions `Context` runs without a spurious `return` argument.""" + server = MCPServer("test") + + @server.tool() + def maybe_context(value: int) -> Context | str: + return f"got {value}" + + async with Client(server) as client: + result = await client.call_tool("maybe_context", {"value": 7}) + + assert result.is_error is False + assert result.structured_content == {"result": "got 7"} + + +async def test_prompt_returning_context_is_called_with_only_its_own_arguments() -> None: + """A prompt whose return annotation mentions `Context` runs without a spurious `return` argument.""" + server = MCPServer("test") + + @server.prompt() + def maybe_context(value: str) -> Context | str: + return f"got {value}" + + async with Client(server) as client: + result = await client.get_prompt("maybe_context", {"value": "seven"}) + + content = result.messages[0].content + assert isinstance(content, TextContent) + assert content.text == "got seven" + + +async def test_resource_template_returning_context_is_called_with_only_its_own_arguments() -> None: + """A resource template whose return annotation mentions `Context` runs without a spurious `return`.""" + server = MCPServer("test") + + @server.resource("res://{value}") + def maybe_context(value: str) -> Context | str: + return f"got {value}" + + async with Client(server) as client: + result = await client.read_resource("res://seven") + + contents = result.contents[0] + assert isinstance(contents, TextResourceContents) + assert contents.text == "got seven"