refactor: remove request_ctx ContextVar, thread Context explicitly (#2203)

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
This commit is contained in:
Max Isbey 2026-03-04 13:23:02 +00:00 • committed by GitHub
commit cc22bf5464
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 484 additions and 413 deletions

View file

@ -3,8 +3,7 @@ from pydantic import FileUrl
from mcp import Client
from mcp.client.session import ClientSession
from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver.server import Context
from mcp.server.mcpserver import Context, MCPServer
from mcp.shared._context import RequestContext
from mcp.types import ListRootsResult, Root, TextContent

View file

@ -3,7 +3,7 @@ from typing import Any, Literal
import pytest
from mcp import Client, types
from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver import Context, MCPServer
from mcp.shared.session import RequestResponder
from mcp.types import (
LoggingMessageNotificationParams,
@ -33,14 +33,10 @@ async def test_logging_callback():
# Create a function that can send a log notification
@server.tool("test_tool_with_log")
async def test_tool_with_log(
message: str, level: Literal["debug", "info", "warning", "error"], logger: str
message: str, level: Literal["debug", "info", "warning", "error"], logger: str, ctx: Context
) -> bool:
"""Send a log notification to the client."""
await server.get_context().log(
level=level,
message=message,
logger_name=logger,
)
await ctx.log(level=level, message=message, logger_name=logger)
return True
@server.tool("test_tool_with_log_extra")
@ -50,9 +46,10 @@ async def test_logging_callback():
logger: str,
extra_string: str,
extra_dict: dict[str, Any],
ctx: Context,
) -> bool:
"""Send a log notification to the client with extra fields."""
await server.get_context().log(
await ctx.log(
level=level,
message=message,
logger_name=logger,

View file

@ -2,7 +2,7 @@ import pytest
from mcp import Client
from mcp.client.session import ClientSession
from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver import Context, MCPServer
from mcp.shared._context import RequestContext
from mcp.types import (
CreateMessageRequestParams,
@ -32,8 +32,8 @@ async def test_sampling_callback():
return callback_return
@server.tool("test_sampling")
async def test_sampling_tool(message: str):
value = await server.get_context().session.create_message(
async def test_sampling_tool(message: str, ctx: Context) -> bool:
value = await ctx.session.create_message(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text=message))],
max_tokens=100,
)
@ -77,9 +77,9 @@ async def test_create_message_backwards_compat_single_content():
return callback_return
@server.tool("test_backwards_compat")
async def test_tool(message: str):
async def test_tool(message: str, ctx: Context) -> bool:
# Call create_message WITHOUT tools
result = await server.get_context().session.create_message(
result = await ctx.session.create_message(
messages=[SamplingMessage(role="user", content=TextContent(type="text", text=message))],
max_tokens=100,
)

View file

@ -2,6 +2,7 @@ from typing import Any
import pytest
from mcp.server.mcpserver import Context
from mcp.server.mcpserver.prompts.base import AssistantMessage, Message, Prompt, UserMessage
from mcp.types import EmbeddedResource, TextContent, TextResourceContents
@ -13,7 +14,9 @@ class TestRenderPrompt:
return "Hello, world!"
prompt = Prompt.from_function(fn)
assert await prompt.render() == [UserMessage(content=TextContent(type="text", text="Hello, world!"))]
assert await prompt.render(None, Context()) == [
UserMessage(content=TextContent(type="text", text="Hello, world!"))
]
@pytest.mark.anyio
async def test_async_fn(self):
@ -21,7 +24,9 @@ class TestRenderPrompt:
return "Hello, world!"
prompt = Prompt.from_function(fn)
assert await prompt.render() == [UserMessage(content=TextContent(type="text", text="Hello, world!"))]
assert await prompt.render(None, Context()) == [
UserMessage(content=TextContent(type="text", text="Hello, world!"))
]
@pytest.mark.anyio
async def test_fn_with_args(self):
@ -29,7 +34,7 @@ class TestRenderPrompt:
return f"Hello, {name}! You're {age} years old."
prompt = Prompt.from_function(fn)
assert await prompt.render(arguments={"name": "World"}) == [
assert await prompt.render({"name": "World"}, Context()) == [
UserMessage(content=TextContent(type="text", text="Hello, World! You're 30 years old."))
]
@ -40,7 +45,7 @@ class TestRenderPrompt:
prompt = Prompt.from_function(fn)
with pytest.raises(ValueError):
await prompt.render(arguments={"age": 40})
await prompt.render({"age": 40}, Context())
@pytest.mark.anyio
async def test_fn_returns_message(self):
@ -48,7 +53,9 @@ class TestRenderPrompt:
return UserMessage(content="Hello, world!")
prompt = Prompt.from_function(fn)
assert await prompt.render() == [UserMessage(content=TextContent(type="text", text="Hello, world!"))]
assert await prompt.render(None, Context()) == [
UserMessage(content=TextContent(type="text", text="Hello, world!"))
]
@pytest.mark.anyio
async def test_fn_returns_assistant_message(self):
@ -56,7 +63,9 @@ class TestRenderPrompt:
return AssistantMessage(content=TextContent(type="text", text="Hello, world!"))
prompt = Prompt.from_function(fn)
assert await prompt.render() == [AssistantMessage(content=TextContent(type="text", text="Hello, world!"))]
assert await prompt.render(None, Context()) == [
AssistantMessage(content=TextContent(type="text", text="Hello, world!"))
]
@pytest.mark.anyio
async def test_fn_returns_multiple_messages(self):
@ -70,7 +79,7 @@ class TestRenderPrompt:
return expected
prompt = Prompt.from_function(fn)
assert await prompt.render() == expected
assert await prompt.render(None, Context()) == expected
@pytest.mark.anyio
async def test_fn_returns_list_of_strings(self):
@ -83,7 +92,7 @@ class TestRenderPrompt:
return expected
prompt = Prompt.from_function(fn)
assert await prompt.render() == [UserMessage(t) for t in expected]
assert await prompt.render(None, Context()) == [UserMessage(t) for t in expected]
@pytest.mark.anyio
async def test_fn_returns_resource_content(self):
@ -102,7 +111,7 @@ class TestRenderPrompt:
)
prompt = Prompt.from_function(fn)
assert await prompt.render() == [
assert await prompt.render(None, Context()) == [
UserMessage(
content=EmbeddedResource(
type="resource",
@ -136,7 +145,7 @@ class TestRenderPrompt:
]
prompt = Prompt.from_function(fn)
assert await prompt.render() == [
assert await prompt.render(None, Context()) == [
UserMessage(content=TextContent(type="text", text="Please analyze this file:")),
UserMessage(
content=EmbeddedResource(
@ -169,7 +178,7 @@ class TestRenderPrompt:
}
prompt = Prompt.from_function(fn)
assert await prompt.render() == [
assert await prompt.render(None, Context()) == [
UserMessage(
content=EmbeddedResource(
type="resource",

View file

@ -1,5 +1,6 @@
import pytest
from mcp.server.mcpserver import Context
from mcp.server.mcpserver.prompts.base import Prompt, UserMessage
from mcp.server.mcpserver.prompts.manager import PromptManager
from mcp.types import TextContent
@ -72,7 +73,7 @@ class TestPromptManager:
manager = PromptManager()
prompt = Prompt.from_function(fn)
manager.add_prompt(prompt)
messages = await manager.render_prompt("fn")
messages = await manager.render_prompt("fn", None, Context())
assert messages == [UserMessage(content=TextContent(type="text", text="Hello, world!"))]
@pytest.mark.anyio
@ -85,7 +86,7 @@ class TestPromptManager:
manager = PromptManager()
prompt = Prompt.from_function(fn)
manager.add_prompt(prompt)
messages = await manager.render_prompt("fn", arguments={"name": "World"})
messages = await manager.render_prompt("fn", {"name": "World"}, Context())
assert messages == [UserMessage(content=TextContent(type="text", text="Hello, World!"))]
@pytest.mark.anyio
@ -93,7 +94,7 @@ class TestPromptManager:
"""Test rendering a non-existent prompt."""
manager = PromptManager()
with pytest.raises(ValueError, match="Unknown prompt: unknown"):
await manager.render_prompt("unknown")
await manager.render_prompt("unknown", None, Context())
@pytest.mark.anyio
async def test_render_prompt_with_missing_args(self):
@ -106,4 +107,4 @@ class TestPromptManager:
prompt = Prompt.from_function(fn)
manager.add_prompt(prompt)
with pytest.raises(ValueError, match="Missing required arguments"):
await manager.render_prompt("fn")
await manager.render_prompt("fn", None, Context())

View file

@ -4,6 +4,7 @@ from tempfile import NamedTemporaryFile
import pytest
from pydantic import AnyUrl
from mcp.server.mcpserver import Context
from mcp.server.mcpserver.resources import FileResource, FunctionResource, ResourceManager, ResourceTemplate
@ -86,7 +87,7 @@ class TestResourceManager:
path=temp_file,
)
manager.add_resource(resource)
retrieved = await manager.get_resource(resource.uri)
retrieved = await manager.get_resource(resource.uri, Context())
assert retrieved == resource
@pytest.mark.anyio
@ -104,7 +105,7 @@ class TestResourceManager:
)
manager._templates[template.uri_template] = template
resource = await manager.get_resource(AnyUrl("greet://world"))
resource = await manager.get_resource(AnyUrl("greet://world"), Context())
assert isinstance(resource, FunctionResource)
content = await resource.read()
assert content == "Hello, world!"
@ -114,7 +115,7 @@ class TestResourceManager:
"""Test getting a non-existent resource."""
manager = ResourceManager()
with pytest.raises(ValueError, match="Unknown resource"):
await manager.get_resource(AnyUrl("unknown://test"))
await manager.get_resource(AnyUrl("unknown://test"), Context())
def test_list_resources(self, temp_file: Path):
"""Test listing all resources."""

View file

@ -4,7 +4,7 @@ from typing import Any
import pytest
from pydantic import BaseModel
from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver import Context, MCPServer
from mcp.server.mcpserver.resources import FunctionResource, ResourceTemplate
from mcp.types import Annotations
@ -64,6 +64,7 @@ class TestResourceTemplate:
resource = await template.create_resource(
"test://foo/123",
{"key": "foo", "value": 123},
Context(),
)
assert isinstance(resource, FunctionResource)
@ -86,7 +87,7 @@ class TestResourceTemplate:
)
with pytest.raises(ValueError, match="Error creating resource from template"):
await template.create_resource("fail://test", {"x": "test"})
await template.create_resource("fail://test", {"x": "test"}, Context())
@pytest.mark.anyio
async def test_async_text_resource(self):
@ -104,6 +105,7 @@ class TestResourceTemplate:
resource = await template.create_resource(
"greet://world",
{"name": "world"},
Context(),
)
assert isinstance(resource, FunctionResource)
@ -126,6 +128,7 @@ class TestResourceTemplate:
resource = await template.create_resource(
"bytes://test",
{"value": "test"},
Context(),
)
assert isinstance(resource, FunctionResource)
@ -152,6 +155,7 @@ class TestResourceTemplate:
resource = await template.create_resource(
"test://foo/123",
{"key": "foo", "value": 123},
Context(),
)
assert isinstance(resource, FunctionResource)
@ -183,6 +187,7 @@ class TestResourceTemplate:
resource = await template.create_resource(
"test://hello",
{"value": "hello"},
Context(),
)
assert isinstance(resource, FunctionResource)
@ -249,7 +254,7 @@ class TestResourceTemplateAnnotations:
)
# Create a resource from the template
resource = await template.create_resource("resource://items/123", {"item_id": "123"})
resource = await template.create_resource("resource://items/123", {"item_id": "123"}, Context())
# The resource should inherit the template's annotations
assert resource.annotations is not None
@ -298,7 +303,7 @@ class TestResourceTemplateMetadata:
)
# Create a resource from the template
resource = await template.create_resource("resource://items/123", {"item_id": "123"})
resource = await template.create_resource("resource://items/123", {"item_id": "123"}, Context())
# The resource should inherit the template's metadata
assert resource.meta is not None

View file

@ -891,7 +891,7 @@ class TestServerResourceTemplates:
assert len(await mcp.list_resources()) == 0
# When accessed, should create a concrete resource
resource = await mcp._resource_manager.get_resource("resource://test/data")
resource = await mcp._resource_manager.get_resource("resource://test/data", Context())
assert isinstance(resource, FunctionResource)
result = await resource.read()
assert result == "Data for test"
@ -1231,6 +1231,19 @@ class TestContextInjection:
class TestServerPrompts:
"""Test prompt functionality in MCPServer server."""
async def test_get_prompt_direct_call_without_context(self):
"""Test calling mcp.get_prompt() directly without passing context."""
mcp = MCPServer()
@mcp.prompt()
def fn() -> str:
return "Hello, world!"
result = await mcp.get_prompt("fn")
content = result.messages[0].content
assert isinstance(content, TextContent)
assert content.text == "Hello, world!"
async def test_prompt_decorator(self):
"""Test that the prompt decorator registers prompts correctly."""
mcp = MCPServer()
@ -1243,7 +1256,7 @@ class TestServerPrompts:
assert len(prompts) == 1
assert prompts[0].name == "fn"
# Don't compare functions directly since validate_call wraps them
content = await prompts[0].render()
content = await prompts[0].render(None, Context())
assert isinstance(content[0].content, TextContent)
assert content[0].content.text == "Hello, world!"
@ -1258,7 +1271,7 @@ class TestServerPrompts:
prompts = mcp._prompt_manager.list_prompts()
assert len(prompts) == 1
assert prompts[0].name == "custom_name"
content = await prompts[0].render()
content = await prompts[0].render(None, Context())
assert isinstance(content[0].content, TextContent)
assert content[0].content.text == "Hello, world!"
@ -1273,7 +1286,7 @@ class TestServerPrompts:
prompts = mcp._prompt_manager.list_prompts()
assert len(prompts) == 1
assert prompts[0].description == "A custom description"
content = await prompts[0].render()
content = await prompts[0].render(None, Context())
assert isinstance(content[0].content, TextContent)
assert content[0].content.text == "Hello, world!"

View file

@ -188,7 +188,7 @@ class TestCallTools:
manager = ToolManager()
manager.add_tool(sum)
result = await manager.call_tool("sum", {"a": 1, "b": 2})
result = await manager.call_tool("sum", {"a": 1, "b": 2}, Context())
assert result == 3
@pytest.mark.anyio
@ -199,7 +199,7 @@ class TestCallTools:
manager = ToolManager()
manager.add_tool(double)
result = await manager.call_tool("double", {"n": 5})
result = await manager.call_tool("double", {"n": 5}, Context())
assert result == 10
@pytest.mark.anyio
@ -213,7 +213,7 @@ class TestCallTools:
manager = ToolManager()
tool = manager.add_tool(MyTool())
result = await tool.run({"x": 5})
result = await tool.run({"x": 5}, Context())
assert result == 10
@pytest.mark.anyio
@ -227,7 +227,7 @@ class TestCallTools:
manager = ToolManager()
tool = manager.add_tool(MyAsyncTool())
result = await tool.run({"x": 5})
result = await tool.run({"x": 5}, Context())
assert result == 10
@pytest.mark.anyio
@ -238,7 +238,7 @@ class TestCallTools:
manager = ToolManager()
manager.add_tool(sum)
result = await manager.call_tool("sum", {"a": 1})
result = await manager.call_tool("sum", {"a": 1}, Context())
assert result == 2
@pytest.mark.anyio
@ -250,13 +250,13 @@ class TestCallTools:
manager = ToolManager()
manager.add_tool(sum)
with pytest.raises(ToolError):
await manager.call_tool("sum", {"a": 1})
await manager.call_tool("sum", {"a": 1}, Context())
@pytest.mark.anyio
async def test_call_unknown_tool(self):
manager = ToolManager()
with pytest.raises(ToolError):
await manager.call_tool("unknown", {"a": 1})
await manager.call_tool("unknown", {"a": 1}, Context())
@pytest.mark.anyio
async def test_call_tool_with_list_int_input(self):
@ -266,9 +266,9 @@ class TestCallTools:
manager = ToolManager()
manager.add_tool(sum_vals)
# Try both with plain list and with JSON list
result = await manager.call_tool("sum_vals", {"vals": "[1, 2, 3]"})
result = await manager.call_tool("sum_vals", {"vals": "[1, 2, 3]"}, Context())
assert result == 6
result = await manager.call_tool("sum_vals", {"vals": [1, 2, 3]})
result = await manager.call_tool("sum_vals", {"vals": [1, 2, 3]}, Context())
assert result == 6
@pytest.mark.anyio
@ -279,13 +279,13 @@ class TestCallTools:
manager = ToolManager()
manager.add_tool(concat_strs)
# Try both with plain python object and with JSON list
result = await manager.call_tool("concat_strs", {"vals": ["a", "b", "c"]})
result = await manager.call_tool("concat_strs", {"vals": ["a", "b", "c"]}, Context())
assert result == "abc"
result = await manager.call_tool("concat_strs", {"vals": '["a", "b", "c"]'})
result = await manager.call_tool("concat_strs", {"vals": '["a", "b", "c"]'}, Context())
assert result == "abc"
result = await manager.call_tool("concat_strs", {"vals": "a"})
result = await manager.call_tool("concat_strs", {"vals": "a"}, Context())
assert result == "a"
result = await manager.call_tool("concat_strs", {"vals": '"a"'})
result = await manager.call_tool("concat_strs", {"vals": '"a"'}, Context())
assert result == '"a"'
@pytest.mark.anyio
@ -297,7 +297,7 @@ class TestCallTools:
shrimp: list[Shrimp]
x: None
def name_shrimp(tank: MyShrimpTank, ctx: Context[ServerSessionT, None]) -> list[str]:
def name_shrimp(tank: MyShrimpTank) -> list[str]:
return [x.name for x in tank.shrimp]
manager = ToolManager()
@ -305,11 +305,13 @@ class TestCallTools:
result = await manager.call_tool(
"name_shrimp",
{"tank": {"x": None, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}},
Context(),
)
assert result == ["rex", "gertrude"]
result = await manager.call_tool(
"name_shrimp",
{"tank": '{"x": null, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}'},
Context(),
)
assert result == ["rex", "gertrude"]
@ -364,9 +366,7 @@ class TestContextHandling:
manager = ToolManager()
manager.add_tool(tool_with_context)
mcp = MCPServer()
ctx = mcp.get_context()
result = await manager.call_tool("tool_with_context", {"x": 42}, context=ctx)
result = await manager.call_tool("tool_with_context", {"x": 42}, context=Context())
assert result == "42"
@pytest.mark.anyio
@ -380,22 +380,7 @@ class TestContextHandling:
manager = ToolManager()
manager.add_tool(async_tool)
mcp = MCPServer()
ctx = mcp.get_context()
result = await manager.call_tool("async_tool", {"x": 42}, context=ctx)
assert result == "42"
@pytest.mark.anyio
async def test_context_optional(self):
"""Test that context is optional when calling tools."""
def tool_with_context(x: int, ctx: Context[ServerSessionT, None] | None = None) -> str:
return str(x)
manager = ToolManager()
manager.add_tool(tool_with_context)
# Should not raise an error when context is not provided
result = await manager.call_tool("tool_with_context", {"x": 42})
result = await manager.call_tool("async_tool", {"x": 42}, context=Context())
assert result == "42"
@pytest.mark.anyio
@ -408,10 +393,8 @@ class TestContextHandling:
manager = ToolManager()
manager.add_tool(tool_with_context)
mcp = MCPServer()
ctx = mcp.get_context()
with pytest.raises(ToolError, match="Error executing tool tool_with_context"):
await manager.call_tool("tool_with_context", {"x": 42}, context=ctx)
await manager.call_tool("tool_with_context", {"x": 42}, context=Context())
class TestToolAnnotations:
@ -471,7 +454,7 @@ class TestStructuredOutput:
manager = ToolManager()
manager.add_tool(get_user)
result = await manager.call_tool("get_user", {"user_id": 1}, convert_result=True)
result = await manager.call_tool("get_user", {"user_id": 1}, Context(), convert_result=True)
# don't test unstructured output here, just the structured conversion
assert len(result) == 2 and result[1] == {"name": "John", "age": 30}
@ -485,9 +468,9 @@ class TestStructuredOutput:
manager = ToolManager()
manager.add_tool(double_number)
result = await manager.call_tool("double_number", {"n": 5})
result = await manager.call_tool("double_number", {"n": 5}, Context())
assert result == 10
result = await manager.call_tool("double_number", {"n": 5}, convert_result=True)
result = await manager.call_tool("double_number", {"n": 5}, Context(), convert_result=True)
assert isinstance(result[0][0], TextContent) and result[1] == {"result": 10}
@pytest.mark.anyio
@ -506,7 +489,7 @@ class TestStructuredOutput:
manager = ToolManager()
manager.add_tool(get_user_dict)
result = await manager.call_tool("get_user_dict", {"user_id": 1})
result = await manager.call_tool("get_user_dict", {"user_id": 1}, Context())
assert result == expected_output
@pytest.mark.anyio
@ -526,7 +509,7 @@ class TestStructuredOutput:
manager = ToolManager()
manager.add_tool(get_person)
result = await manager.call_tool("get_person", {}, convert_result=True)
result = await manager.call_tool("get_person", {}, Context(), convert_result=True)
# don't test unstructured output here, just the structured conversion
assert len(result) == 2 and result[1] == expected_output
@ -543,9 +526,9 @@ class TestStructuredOutput:
manager = ToolManager()
manager.add_tool(get_numbers)
result = await manager.call_tool("get_numbers", {})
result = await manager.call_tool("get_numbers", {}, Context())
assert result == expected_list
result = await manager.call_tool("get_numbers", {}, convert_result=True)
result = await manager.call_tool("get_numbers", {}, Context(), convert_result=True)
assert isinstance(result[0][0], TextContent) and result[1] == expected_output
@pytest.mark.anyio
@ -558,7 +541,7 @@ class TestStructuredOutput:
manager = ToolManager()
manager.add_tool(get_dict, structured_output=False)
result = await manager.call_tool("get_dict", {})
result = await manager.call_tool("get_dict", {}, Context())
assert isinstance(result, dict)
assert result == {"key": "value"}
@ -601,12 +584,12 @@ class TestStructuredOutput:
assert "properties" not in tool.output_schema # dict[str, Any] has no constraints
# Test raw result
result = await manager.call_tool("get_config", {})
result = await manager.call_tool("get_config", {}, Context())
expected = {"debug": True, "port": 8080, "features": ["auth", "logging"]}
assert result == expected
# Test converted result
result = await manager.call_tool("get_config", {})
result = await manager.call_tool("get_config", {}, Context())
assert result == expected
@pytest.mark.anyio
@ -626,12 +609,12 @@ class TestStructuredOutput:
assert tool.output_schema["additionalProperties"]["type"] == "integer"
# Test raw result
result = await manager.call_tool("get_scores", {})
result = await manager.call_tool("get_scores", {}, Context())
expected = {"alice": 100, "bob": 85, "charlie": 92}
assert result == expected
# Test converted result
result = await manager.call_tool("get_scores", {})
result = await manager.call_tool("get_scores", {}, Context())
assert result == expected
@ -885,7 +868,7 @@ class TestRemoveTools:
manager.add_tool(greet)
# Verify tool works before removal
result = await manager.call_tool("greet", {"name": "World"})
result = await manager.call_tool("greet", {"name": "World"}, Context())
assert result == "Hello, World!"
# Remove the tool
@ -893,7 +876,7 @@ class TestRemoveTools:
# Verify calling removed tool raises error
with pytest.raises(ToolError, match="Unknown tool: greet"):
await manager.call_tool("greet", {"name": "World"})
await manager.call_tool("greet", {"name": "World"}, Context())
def test_remove_tool_case_sensitive(self):
"""Test that tool removal is case-sensitive."""

View file

View file

@ -0,0 +1,10 @@
from mcp.server.mcpserver import Context
from mcp.server.mcpserver.tools.base import Tool
def test_context_detected_in_union_annotation():
def my_tool(x: int, ctx: Context | None) -> str:
raise NotImplementedError
tool = Tool.from_function(my_tool)
assert tool.context_kwarg == "ctx"