From ef432ade0f240dc273d492a44be5c75e6abf5456 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 25 Apr 2025 20:57:49 -0400 Subject: [PATCH] Add context injection to all objects --- src/fastmcp/resources/resource_manager.py | 17 +- src/fastmcp/resources/template.py | 2 +- src/fastmcp/resources/types.py | 32 ++- src/fastmcp/server/openapi.py | 4 +- src/fastmcp/server/server.py | 4 +- tests/server/test_server_interactions.py | 316 +++++++++++++--------- 6 files changed, 227 insertions(+), 148 deletions(-) diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 1cf42d839..d27f6247d 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -61,9 +61,16 @@ class ResourceManager: The added resource or template. If a resource or template with the same URI already exists, returns the existing resource or template. """ + from fastmcp.server.context import Context + # Check if this should be a template has_uri_params = "{" in uri and "}" in uri - has_func_params = bool(inspect.signature(fn).parameters) + # check if the function has any parameters (other than injected context) + has_func_params = any( + p + for p in inspect.signature(fn).parameters.values() + if p.annotation is not Context + ) if has_uri_params or has_func_params: return self.add_template_from_fn( @@ -102,12 +109,12 @@ class ResourceManager: The added resource. If a resource with the same URI already exists, returns the existing resource. """ - resource = FunctionResource( + resource = FunctionResource.from_function( + fn=fn, uri=AnyUrl(uri), name=name, description=description, mime_type=mime_type or "text/plain", - fn=fn, tags=tags or set(), ) return self.add_resource(resource) @@ -235,7 +242,9 @@ class ResourceManager: if params := match_uri_template(uri_str, storage_key): try: return await template.create_resource( - uri_str, params, context=context + uri_str, + params=params, + context=context, ) except Exception as e: raise ValueError(f"Error creating resource from template: {e}") diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index 41516dd91..d0ec24fbd 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -189,7 +189,7 @@ class ResourceTemplate(BaseModel): name=self.name, description=self.description, mime_type=self.mime_type, - fn=lambda: result, # Capture result in closure + fn=lambda **kwargs: result, # Capture result in closure tags=self.tags, context_kwarg=self.context_kwarg, ) diff --git a/src/fastmcp/resources/types.py b/src/fastmcp/resources/types.py index e86733e8a..f7384a85a 100644 --- a/src/fastmcp/resources/types.py +++ b/src/fastmcp/resources/types.py @@ -15,6 +15,7 @@ import pydantic.json import pydantic_core from pydantic import Field, ValidationInfo +import fastmcp from fastmcp.resources.resource import Resource if TYPE_CHECKING: @@ -66,8 +67,23 @@ class FunctionResource(Resource): default=None, description="Name of the kwarg that should receive context" ) + @classmethod + def from_function( + cls, fn: Callable[[], Any], context_kwarg: str | None = None, **kwargs + ) -> FunctionResource: + if context_kwarg is None: + parameters = inspect.signature(fn).parameters + context_param = next( + (p for p in parameters.values() if p.annotation is fastmcp.Context), + None, + ) + if context_param is not None: + context_kwarg = context_param.name + return cls(fn=fn, context_kwarg=context_kwarg, **kwargs) + async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None + self, + context: Context[ServerSessionT, LifespanContextT] | None = None, ) -> str | bytes: """Read the resource by calling the wrapped function.""" try: @@ -80,7 +96,7 @@ class FunctionResource(Resource): result = await result if isinstance(result, Resource): - return await result.read() + return await result.read(context=context) if isinstance(result, bytes): return result if isinstance(result, str): @@ -127,7 +143,9 @@ class FileResource(Resource): mime_type = info.data.get("mime_type", "text/plain") return not mime_type.startswith("text/") - async def read(self) -> str | bytes: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str | bytes: """Read the file content.""" try: if self.is_binary: @@ -145,7 +163,9 @@ class HttpResource(Resource): default="application/json", description="MIME type of the resource content" ) - async def read(self) -> str | bytes: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str | bytes: """Read the HTTP content.""" async with httpx.AsyncClient() as client: response = await client.get(self.url) @@ -197,7 +217,9 @@ class DirectoryResource(Resource): except Exception as e: raise ValueError(f"Error listing directory {self.path}: {e}") - async def read(self) -> str: # Always returns JSON string + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str: # Always returns JSON string """Read the directory listing.""" try: files = await anyio.to_thread.run_sync(self.list_files) diff --git a/src/fastmcp/server/openapi.py b/src/fastmcp/server/openapi.py index ae2776af8..b3396b368 100644 --- a/src/fastmcp/server/openapi.py +++ b/src/fastmcp/server/openapi.py @@ -265,7 +265,9 @@ class OpenAPIResource(Resource): self._client = client self._route = route - async def read(self) -> str | bytes: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str | bytes: """Fetch the resource data by making an HTTP request.""" try: # Extract path parameters from the URI if present diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 1ca3b116a..2983932a1 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -401,7 +401,7 @@ class FastMCP(Generic[LifespanResultT]): context = self.get_context() resource = await self._resource_manager.get_resource(uri, context=context) try: - content = await resource.read() + content = await resource.read(context=context) return [ ReadResourceContents(content=content, mime_type=resource.mime_type) ] @@ -427,7 +427,7 @@ class FastMCP(Generic[LifespanResultT]): if self._prompt_manager.has_prompt(name): context = self.get_context() messages = await self._prompt_manager.render_prompt( - name, arguments, context=context + name, arguments=arguments or {}, context=context ) return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages)) else: diff --git a/tests/server/test_server_interactions.py b/tests/server/test_server_interactions.py index 506f1a65a..2653ea1a9 100644 --- a/tests/server/test_server_interactions.py +++ b/tests/server/test_server_interactions.py @@ -682,7 +682,147 @@ class TestToolParameters: assert result[0].text == "0:16:40" -class TestResources: +class TestToolContextInjection: + """Test context injection in tools.""" + + async def test_context_detection(self): + """Test that context parameters are properly detected.""" + mcp = FastMCP() + + def tool_with_context(x: int, ctx: Context) -> str: + return f"Request {ctx.request_id}: {x}" + + mcp.add_tool(tool_with_context) + async with Client(mcp) as client: + tools = await client.list_tools() + assert len(tools) == 1 + assert tools[0].name == "tool_with_context" + + async def test_context_injection(self): + """Test that context is properly injected into tool calls.""" + mcp = FastMCP() + + @mcp.tool() + def tool_with_context(x: int, ctx: Context) -> str: + assert isinstance(ctx, Context) + assert ctx.request_id is not None + return ctx.request_id + + async with Client(mcp) as client: + result = await client.call_tool("tool_with_context", {"x": 42}) + assert len(result) == 1 + content = result[0] + assert isinstance(content, TextContent) + assert content.text == "1" + + async def test_async_context(self): + """Test that context works in async functions.""" + mcp = FastMCP() + + async def async_tool(x: int, ctx: Context) -> str: + assert ctx.request_id is not None + return f"Async request {ctx.request_id}: {x}" + + mcp.add_tool(async_tool) + async with Client(mcp) as client: + result = await client.call_tool("async_tool", {"x": 42}) + assert len(result) == 1 + content = result[0] + assert isinstance(content, TextContent) + assert "Async request" in content.text + assert "42" in content.text + + async def test_context_logging(self): + from unittest.mock import patch + + import mcp.server.session + + """Test that context logging methods work.""" + mcp = FastMCP() + + async def logging_tool(msg: str, ctx: Context) -> str: + await ctx.debug("Debug message") + await ctx.info("Info message") + await ctx.warning("Warning message") + await ctx.error("Error message") + return f"Logged messages for {msg}" + + mcp.add_tool(logging_tool) + + with patch("mcp.server.session.ServerSession.send_log_message") as mock_log: + async with Client(mcp) as client: + result = await client.call_tool("logging_tool", {"msg": "test"}) + assert len(result) == 1 + content = result[0] + assert isinstance(content, TextContent) + assert "Logged messages for test" in content.text + + assert mock_log.call_count == 4 + mock_log.assert_any_call( + level="debug", data="Debug message", logger=None + ) + mock_log.assert_any_call(level="info", data="Info message", logger=None) + mock_log.assert_any_call( + level="warning", data="Warning message", logger=None + ) + mock_log.assert_any_call( + level="error", data="Error message", logger=None + ) + + async def test_optional_context(self): + """Test that context is optional.""" + mcp = FastMCP() + + def no_context(x: int) -> int: + return x * 2 + + mcp.add_tool(no_context) + async with Client(mcp) as client: + result = await client.call_tool("no_context", {"x": 21}) + assert len(result) == 1 + content = result[0] + assert isinstance(content, TextContent) + assert content.text == "42" + + async def test_context_resource_access(self): + """Test that context can access resources.""" + mcp = FastMCP() + + @mcp.resource("test://data") + def test_resource() -> str: + return "resource data" + + @mcp.tool() + async def tool_with_resource(ctx: Context) -> str: + r_iter = await ctx.read_resource("test://data") + r_list = list(r_iter) + assert len(r_list) == 1 + r = r_list[0] + return f"Read resource: {r.content} with mime type {r.mime_type}" + + async with Client(mcp) as client: + result = await client.call_tool("tool_with_resource", {}) + assert len(result) == 1 + content = result[0] + assert isinstance(content, TextContent) + assert "Read resource: resource data" in content.text + + async def test_tool_decorator_with_tags(self): + """Test that the tool decorator properly sets tags.""" + mcp = FastMCP() + + @mcp.tool(tags={"example", "test-tag"}) + def sample_tool(x: int) -> int: + return x * 2 + + # Verify the tool exists + async with Client(mcp) as client: + tools = await client.list_tools() + assert len(tools) == 1 + # Note: MCPTool from the client API doesn't expose tags + + +class TestResource: async def test_text_resource(self): mcp = FastMCP() @@ -756,6 +896,21 @@ class TestResources: assert result[0].blob == base64.b64encode(b"Binary file data").decode() +class TestResourceContext: + async def test_resource_with_context_annotation_gets_context(self): + mcp = FastMCP() + + @mcp.resource("resource://test") + def resource_with_context(ctx: Context) -> str: + assert isinstance(ctx, Context) + return ctx.request_id + + async with Client(mcp) as client: + result = await client.read_resource(AnyUrl("resource://test")) + assert isinstance(result[0], TextResourceContents) + assert result[0].text == "1" + + class TestResourceTemplates: async def test_resource_with_params_not_in_uri(self): """Test that a resource with function parameters raises an error if the URI @@ -1026,144 +1181,19 @@ class TestResourceTemplates: assert result[0].text == "Template resource 1: a/b" -class TestContextInjection: - """Test context injection in tools.""" - - async def test_context_detection(self): - """Test that context parameters are properly detected.""" +class TestResourceTemplateContext: + async def test_resource_template_context(self): mcp = FastMCP() - def tool_with_context(x: int, ctx: Context) -> str: - return f"Request {ctx.request_id}: {x}" - - mcp.add_tool(tool_with_context) - async with Client(mcp) as client: - tools = await client.list_tools() - assert len(tools) == 1 - assert tools[0].name == "tool_with_context" - - async def test_context_injection(self): - """Test that context is properly injected into tool calls.""" - mcp = FastMCP() - - def tool_with_context(x: int, ctx: Context) -> str: - assert ctx.request_id is not None - return f"Request {ctx.request_id}: {x}" - - mcp.add_tool(tool_with_context) - async with Client(mcp) as client: - result = await client.call_tool("tool_with_context", {"x": 42}) - assert len(result) == 1 - content = result[0] - assert isinstance(content, TextContent) - assert "Request" in content.text - assert "42" in content.text - - async def test_async_context(self): - """Test that context works in async functions.""" - mcp = FastMCP() - - async def async_tool(x: int, ctx: Context) -> str: - assert ctx.request_id is not None - return f"Async request {ctx.request_id}: {x}" - - mcp.add_tool(async_tool) - async with Client(mcp) as client: - result = await client.call_tool("async_tool", {"x": 42}) - assert len(result) == 1 - content = result[0] - assert isinstance(content, TextContent) - assert "Async request" in content.text - assert "42" in content.text - - async def test_context_logging(self): - from unittest.mock import patch - - import mcp.server.session - - """Test that context logging methods work.""" - mcp = FastMCP() - - async def logging_tool(msg: str, ctx: Context) -> str: - await ctx.debug("Debug message") - await ctx.info("Info message") - await ctx.warning("Warning message") - await ctx.error("Error message") - return f"Logged messages for {msg}" - - mcp.add_tool(logging_tool) - - with patch("mcp.server.session.ServerSession.send_log_message") as mock_log: - async with Client(mcp) as client: - result = await client.call_tool("logging_tool", {"msg": "test"}) - assert len(result) == 1 - content = result[0] - assert isinstance(content, TextContent) - assert "Logged messages for test" in content.text - - assert mock_log.call_count == 4 - mock_log.assert_any_call( - level="debug", data="Debug message", logger=None - ) - mock_log.assert_any_call(level="info", data="Info message", logger=None) - mock_log.assert_any_call( - level="warning", data="Warning message", logger=None - ) - mock_log.assert_any_call( - level="error", data="Error message", logger=None - ) - - async def test_optional_context(self): - """Test that context is optional.""" - mcp = FastMCP() - - def no_context(x: int) -> int: - return x * 2 - - mcp.add_tool(no_context) - async with Client(mcp) as client: - result = await client.call_tool("no_context", {"x": 21}) - assert len(result) == 1 - content = result[0] - assert isinstance(content, TextContent) - assert content.text == "42" - - async def test_context_resource_access(self): - """Test that context can access resources.""" - mcp = FastMCP() - - @mcp.resource("test://data") - def test_resource() -> str: - return "resource data" - - @mcp.tool() - async def tool_with_resource(ctx: Context) -> str: - r_iter = await ctx.read_resource("test://data") - r_list = list(r_iter) - assert len(r_list) == 1 - r = r_list[0] - return f"Read resource: {r.content} with mime type {r.mime_type}" + @mcp.resource("resource://{param}") + def resource_template(param: str, ctx: Context) -> str: + assert isinstance(ctx, Context) + return f"Resource template: {param} {ctx.request_id}" async with Client(mcp) as client: - result = await client.call_tool("tool_with_resource", {}) - assert len(result) == 1 - content = result[0] - assert isinstance(content, TextContent) - assert "Read resource: resource data" in content.text - - async def test_tool_decorator_with_tags(self): - """Test that the tool decorator properly sets tags.""" - mcp = FastMCP() - - @mcp.tool(tags={"example", "test-tag"}) - def sample_tool(x: int) -> int: - return x * 2 - - # Verify the tool exists - async with Client(mcp) as client: - tools = await client.list_tools() - assert len(tools) == 1 - # Note: MCPTool from the client API doesn't expose tags + result = await client.read_resource(AnyUrl("resource://test")) + assert isinstance(result[0], TextResourceContents) + assert result[0].text == "Resource template: test 1" class TestPrompts: @@ -1350,3 +1380,19 @@ class TestPrompts: assert len(prompts_dict) == 1 prompt = prompts_dict["sample_prompt"] assert prompt.tags == {"example", "test-tag"} + + +class TestPromptContext: + async def test_prompt_context(self): + mcp = FastMCP() + + @mcp.prompt() + def prompt_fn(name: str, ctx: Context) -> str: + assert isinstance(ctx, Context) + return f"Hello, {name}! {ctx.request_id}" + + async with Client(mcp) as client: + result = await client.get_prompt("prompt_fn", {"name": "World"}) + assert len(result) == 1 + message = result[0] + assert message.role == "user"