Add context tests and better runtime loads

This commit is contained in:
Jeremiah Lowin 2024-11-30 16:13:07 -05:00
commit 33105dc49b
3 changed files with 99 additions and 5 deletions

View file

@ -119,12 +119,16 @@ class FastMCP:
for info in tools
]
def get_context(self) -> Optional["Context"]:
def get_context(self) -> "Context":
"""
Returns a Context object. Note that the context will only be valid
during a request; outside a request, most methods will error.
"""
try:
request_context = self._mcp_server.request_context
return Context(request_context=request_context, fastmcp=self)
except LookupError:
return None
request_context = None
return Context(request_context=request_context, fastmcp=self)
async def call_tool(
self, name: str, arguments: dict
@ -457,7 +461,11 @@ class Context(BaseModel):
_fastmcp: FastMCP
def __init__(
self, *, request_context: RequestContext, fastmcp: FastMCP, **kwargs: Any
self,
*,
request_context: RequestContext = None,
fastmcp: FastMCP = None,
**kwargs: Any,
):
super().__init__(**kwargs)
self._request_context = request_context
@ -466,11 +474,15 @@ class Context(BaseModel):
@property
def fastmcp(self) -> FastMCP:
"""Access to the FastMCP server."""
if self._fastmcp is None:
raise ValueError("Context is not available outside of a request")
return self._fastmcp
@property
def request_context(self) -> RequestContext:
"""Access to the underlying request context."""
if self._request_context is None:
raise ValueError("Context is not available outside of a request")
return self._request_context
async def report_progress(

View file

@ -71,7 +71,7 @@ class Tool(BaseModel):
"""Run the tool with arguments."""
try:
# Inject context if needed
if self.context_kwarg and context:
if self.context_kwarg:
arguments[self.context_kwarg] = context
# Call function with proper async handling

View file

@ -1,6 +1,7 @@
import logging
import pytest
from pydantic import BaseModel
from typing import Optional
from fastmcp.exceptions import ToolError
from fastmcp.tools import ToolManager
@ -153,3 +154,84 @@ class TestCallTools:
manager = ToolManager()
with pytest.raises(ToolError):
await manager.call_tool("unknown", {"a": 1})
class TestContextHandling:
"""Test context handling in the tool manager."""
def test_context_parameter_detection(self):
"""Test that context parameters are properly detected in Tool.from_function()."""
from fastmcp import Context
def tool_with_context(x: int, ctx: Context) -> str:
return str(x)
manager = ToolManager()
tool = manager.add_tool(tool_with_context)
assert tool.context_kwarg == "ctx"
def tool_without_context(x: int) -> str:
return str(x)
tool = manager.add_tool(tool_without_context)
assert tool.context_kwarg is None
async def test_context_injection(self):
"""Test that context is properly injected during tool execution."""
from fastmcp import Context, FastMCP
def tool_with_context(x: int, ctx: Context) -> str:
assert isinstance(ctx, Context)
return str(x)
manager = ToolManager()
tool = manager.add_tool(tool_with_context)
mcp = FastMCP()
ctx = mcp.get_context()
result = await manager.call_tool("tool_with_context", {"x": 42}, context=ctx)
assert result == "42"
async def test_context_injection_async(self):
"""Test that context is properly injected in async tools."""
from fastmcp import Context, FastMCP
async def async_tool(x: int, ctx: Context) -> str:
assert isinstance(ctx, Context)
return str(x)
manager = ToolManager()
tool = manager.add_tool(async_tool)
mcp = FastMCP()
ctx = mcp.get_context()
result = await manager.call_tool("async_tool", {"x": 42}, context=ctx)
assert result == "42"
async def test_context_optional(self):
"""Test that context is optional when calling tools."""
from fastmcp import Context
def tool_with_context(x: int, ctx: Optional[Context] = None) -> str:
return str(x)
manager = ToolManager()
tool = 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})
assert result == "42"
async def test_context_error_handling(self):
"""Test error handling when context injection fails."""
from fastmcp import Context, FastMCP
def tool_with_context(x: int, ctx: Context) -> str:
raise ValueError("Test error")
manager = ToolManager()
tool = manager.add_tool(tool_with_context)
mcp = FastMCP()
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)