From 82aa53e09bdb5883d69f0a4c7e8529901a93937f Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 6 May 2025 21:24:29 -0400 Subject: [PATCH] Refactor context to avoid unecessary passing --- src/fastmcp/__init__.py | 3 +- src/fastmcp/prompts/prompt.py | 25 +++--- src/fastmcp/prompts/prompt_manager.py | 13 +-- src/fastmcp/resources/resource.py | 9 +- src/fastmcp/resources/resource_manager.py | 6 +- src/fastmcp/resources/template.py | 35 +++----- src/fastmcp/resources/types.py | 59 ++++-------- src/fastmcp/server/__init__.py | 1 + src/fastmcp/server/context.py | 78 ++++++++-------- src/fastmcp/server/dependencies.py | 35 ++++++++ src/fastmcp/server/http.py | 28 +++--- src/fastmcp/server/openapi.py | 21 ++--- src/fastmcp/server/proxy.py | 17 +--- src/fastmcp/tools/tool.py | 37 +++----- src/fastmcp/tools/tool_manager.py | 12 +-- tests/prompts/test_prompt_manager.py | 51 +++++------ tests/resources/test_resource_template.py | 54 ++++++----- tests/tools/test_tool_manager.py | 105 ++++++++++++---------- 18 files changed, 262 insertions(+), 327 deletions(-) create mode 100644 src/fastmcp/server/dependencies.py diff --git a/src/fastmcp/__init__.py b/src/fastmcp/__init__.py index 63b7d8e51..aaeab2340 100644 --- a/src/fastmcp/__init__.py +++ b/src/fastmcp/__init__.py @@ -2,9 +2,10 @@ from importlib.metadata import version - from fastmcp.server.server import FastMCP from fastmcp.server.context import Context +import fastmcp.server + from fastmcp.client import Client from fastmcp.utilities.types import Image from . import client, settings diff --git a/src/fastmcp/prompts/prompt.py b/src/fastmcp/prompts/prompt.py index 763c12021..dafe699ab 100644 --- a/src/fastmcp/prompts/prompt.py +++ b/src/fastmcp/prompts/prompt.py @@ -12,6 +12,7 @@ from mcp.types import Prompt as MCPPrompt from mcp.types import PromptArgument as MCPPromptArgument from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call +from fastmcp.server.dependencies import get_context from fastmcp.utilities.json_schema import prune_params from fastmcp.utilities.types import ( _convert_set_defaults, @@ -20,10 +21,7 @@ from fastmcp.utilities.types import ( ) if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - - from fastmcp.server import Context + pass CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource @@ -76,9 +74,6 @@ class Prompt(BaseModel): None, description="Arguments that can be passed to the prompt" ) fn: Callable[..., PromptResult | Awaitable[PromptResult]] - context_kwarg: str | None = Field( - None, description="Name of the kwarg that should receive context" - ) @classmethod def from_function( @@ -87,7 +82,6 @@ class Prompt(BaseModel): name: str | None = None, description: str | None = None, tags: set[str] | None = None, - context_kwarg: str | None = None, ) -> Prompt: """Create a Prompt from a function. @@ -97,7 +91,7 @@ class Prompt(BaseModel): - A dict (converted to a message) - A sequence of any of the above """ - from fastmcp import Context + from fastmcp.server.context import Context func_name = name or fn.__name__ @@ -115,8 +109,8 @@ class Prompt(BaseModel): parameters = type_adapter.json_schema() # Auto-detect context parameter if not provided - if context_kwarg is None: - context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) + + context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) if context_kwarg: parameters = prune_params(parameters, params=[context_kwarg]) @@ -141,15 +135,15 @@ class Prompt(BaseModel): arguments=arguments, fn=fn, tags=tags or set(), - context_kwarg=context_kwarg, ) async def render( self, arguments: dict[str, Any] | None = None, - context: Context[ServerSessionT, LifespanContextT] | None = None, ) -> list[PromptMessage]: """Render the prompt with arguments.""" + from fastmcp.server.context import Context + # Validate required arguments if self.arguments: required = {arg.name for arg in self.arguments if arg.required} @@ -161,8 +155,9 @@ class Prompt(BaseModel): try: # Prepare arguments with context kwargs = arguments.copy() if arguments else {} - if self.context_kwarg is not None and context is not None: - kwargs[self.context_kwarg] = context + context_kwarg = find_kwarg_by_type(self.fn, kwarg_type=Context) + if context_kwarg and context_kwarg not in kwargs: + kwargs[context_kwarg] = get_context() # Call function and check if result is a coroutine result = self.fn(**kwargs) diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index 4de06c03c..8102cd364 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -13,10 +13,7 @@ from fastmcp.settings import DuplicateBehavior from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - - from fastmcp.server import Context + pass logger = get_logger(__name__) @@ -82,19 +79,15 @@ class PromptManager: self, name: str, arguments: dict[str, Any] | None = None, - context: Context[ServerSessionT, LifespanContextT] | None = None, ) -> GetPromptResult: """Render a prompt by name with arguments.""" prompt = self.get_prompt(name) if not prompt: raise NotFoundError(f"Unknown prompt: {name}") - messages = await prompt.render(arguments, context=context) + messages = await prompt.render(arguments) - return GetPromptResult( - description=prompt.description, - messages=messages, - ) + return GetPromptResult(description=prompt.description, messages=messages) def has_prompt(self, key: str) -> bool: """Check if a prompt exists.""" diff --git a/src/fastmcp/resources/resource.py b/src/fastmcp/resources/resource.py index e4d69638d..95bb7b034 100644 --- a/src/fastmcp/resources/resource.py +++ b/src/fastmcp/resources/resource.py @@ -20,10 +20,7 @@ from pydantic import ( from fastmcp.utilities.types import _convert_set_defaults if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - - from fastmcp.server import Context + pass class Resource(BaseModel, abc.ABC): @@ -66,9 +63,7 @@ class Resource(BaseModel, abc.ABC): raise ValueError("Either name or uri must be provided") @abc.abstractmethod - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> str | bytes: + async def read(self) -> str | bytes: """Read the resource content.""" pass diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index d27f6247d..aa841cc1a 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -109,7 +109,7 @@ class ResourceManager: The added resource. If a resource with the same URI already exists, returns the existing resource. """ - resource = FunctionResource.from_function( + resource = FunctionResource( fn=fn, uri=AnyUrl(uri), name=name, @@ -219,12 +219,11 @@ class ResourceManager: return True return False - async def get_resource(self, uri: AnyUrl | str, context=None) -> Resource: + async def get_resource(self, uri: AnyUrl | str) -> Resource: """Get resource by URI, checking concrete resources first, then templates. Args: uri: The URI of the resource to get - context: Optional context object to pass to template resources Raises: NotFoundError: If no resource or template matching the URI is found. @@ -244,7 +243,6 @@ class ResourceManager: return await template.create_resource( 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 04fae970f..7bae3554e 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -5,7 +5,7 @@ from __future__ import annotations import inspect import re from collections.abc import Callable -from typing import TYPE_CHECKING, Annotated, Any +from typing import Annotated, Any from urllib.parse import unquote from mcp.types import ResourceTemplate as MCPResourceTemplate @@ -20,17 +20,12 @@ from pydantic import ( ) from fastmcp.resources.types import FunctionResource, Resource +from fastmcp.server.dependencies import get_context from fastmcp.utilities.types import ( _convert_set_defaults, find_kwarg_by_type, ) -if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - - from fastmcp.server import Context - def build_regex(template: str) -> re.Pattern: parts = re.split(r"(\{[^}]+\})", template) @@ -79,9 +74,6 @@ class ResourceTemplate(BaseModel): parameters: dict[str, Any] = Field( description="JSON schema for function parameters" ) - context_kwarg: str | None = Field( - None, description="Name of the kwarg that should receive context" - ) @field_validator("mime_type", mode="before") @classmethod @@ -100,10 +92,9 @@ class ResourceTemplate(BaseModel): description: str | None = None, mime_type: str | None = None, tags: set[str] | None = None, - context_kwarg: str | None = None, ) -> ResourceTemplate: """Create a template from a function.""" - from fastmcp import Context + from fastmcp.server.context import Context func_name = name or fn.__name__ if func_name == "": @@ -119,8 +110,8 @@ class ResourceTemplate(BaseModel): ) # Auto-detect context parameter if not provided - if context_kwarg is None: - context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) + + context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) # Validate that URI params match function params uri_params = set(re.findall(r"{(\w+)(?:\*)?}", uri_template)) @@ -170,25 +161,22 @@ class ResourceTemplate(BaseModel): fn=fn, parameters=parameters, tags=tags or set(), - context_kwarg=context_kwarg, ) def matches(self, uri: str) -> dict[str, Any] | None: """Check if URI matches template and extract parameters.""" return match_uri_template(uri, self.uri_template) - async def create_resource( - self, - uri: str, - params: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, - ) -> Resource: + async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: """Create a resource from the template with the given parameters.""" + from fastmcp.server.context import Context + try: # Add context to parameters if needed kwargs = params.copy() - if self.context_kwarg is not None and context is not None: - kwargs[self.context_kwarg] = context + context_kwarg = find_kwarg_by_type(self.fn, kwarg_type=Context) + if context_kwarg and context_kwarg not in kwargs: + kwargs[context_kwarg] = get_context() # Call function and check if result is a coroutine result = self.fn(**kwargs) @@ -202,7 +190,6 @@ class ResourceTemplate(BaseModel): mime_type=self.mime_type, fn=lambda **kwargs: result, # Capture result in closure tags=self.tags, - context_kwarg=self.context_kwarg, ) except Exception as e: raise ValueError(f"Error creating resource from template: {e}") diff --git a/src/fastmcp/resources/types.py b/src/fastmcp/resources/types.py index c6fda510a..cec2ca816 100644 --- a/src/fastmcp/resources/types.py +++ b/src/fastmcp/resources/types.py @@ -15,14 +15,12 @@ import pydantic.json import pydantic_core from pydantic import Field, ValidationInfo -import fastmcp from fastmcp.resources.resource import Resource +from fastmcp.server.dependencies import get_context +from fastmcp.utilities.types import find_kwarg_by_type if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - - from fastmcp.server import Context + pass class TextResource(Resource): @@ -30,9 +28,7 @@ class TextResource(Resource): text: str = Field(description="Text content of the resource") - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> str: + async def read(self) -> str: """Read the text content.""" return self.text @@ -42,9 +38,7 @@ class BinaryResource(Resource): data: bytes = Field(description="Binary content of the resource") - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> bytes: + async def read(self) -> bytes: """Read the binary content.""" return self.data @@ -63,40 +57,23 @@ class FunctionResource(Resource): """ fn: Callable[[], Any] - context_kwarg: str | None = Field( - 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, - ) -> str | bytes: + async def read(self) -> str | bytes: """Read the resource by calling the wrapped function.""" + from fastmcp.server.context import Context + try: kwargs = {} - if self.context_kwarg is not None: - kwargs[self.context_kwarg] = context + context_kwarg = find_kwarg_by_type(self.fn, kwarg_type=Context) + if context_kwarg is not None: + kwargs[context_kwarg] = get_context() result = self.fn(**kwargs) if inspect.iscoroutinefunction(self.fn): result = await result if isinstance(result, Resource): - return await result.read(context=context) + return await result.read() elif isinstance(result, bytes): return result elif isinstance(result, str): @@ -140,9 +117,7 @@ class FileResource(Resource): mime_type = info.data.get("mime_type", "text/plain") return not mime_type.startswith("text/") - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> str | bytes: + async def read(self) -> str | bytes: """Read the file content.""" try: if self.is_binary: @@ -160,9 +135,7 @@ class HttpResource(Resource): default="application/json", description="MIME type of the resource content" ) - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> str | bytes: + async def read(self) -> str | bytes: """Read the HTTP content.""" async with httpx.AsyncClient() as client: response = await client.get(self.url) @@ -214,9 +187,7 @@ class DirectoryResource(Resource): except Exception as e: raise ValueError(f"Error listing directory {self.path}: {e}") - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> str: # Always returns JSON string + async def read(self) -> 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/__init__.py b/src/fastmcp/server/__init__.py index e1d937863..c17dd0e4e 100644 --- a/src/fastmcp/server/__init__.py +++ b/src/fastmcp/server/__init__.py @@ -1,5 +1,6 @@ from .server import FastMCP from .context import Context +from . import dependencies __all__ = ["FastMCP", "Context"] diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 6bb45432e..31a19d0e3 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -1,11 +1,13 @@ from __future__ import annotations as _annotations -from typing import Any, Generic +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass from mcp import LoggingLevel from mcp.server.lowlevel.helper_types import ReadResourceContents -from mcp.server.session import ServerSessionT -from mcp.shared.context import LifespanContextT, RequestContext +from mcp.shared.context import RequestContext from mcp.types import ( CreateMessageResult, ImageContent, @@ -13,18 +15,29 @@ from mcp.types import ( SamplingMessage, TextContent, ) -from pydantic import BaseModel, ConfigDict from pydantic.networks import AnyUrl from starlette.requests import Request -from fastmcp.server.http import get_current_starlette_request +import fastmcp.server.dependencies from fastmcp.server.server import FastMCP from fastmcp.utilities.logging import get_logger logger = get_logger(__name__) +_current_context: ContextVar[Context | None] = ContextVar("context", default=None) -class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]): + +@contextmanager +def set_context(context: Context) -> Generator[Context, None, None]: + token = _current_context.set(context) + try: + yield context + finally: + _current_context.reset(token) + + +@dataclass +class Context: """Context object providing access to MCP capabilities. This provides a cleaner interface to MCP's RequestContext functionality. @@ -56,37 +69,30 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]): The context parameter name can be anything as long as it's annotated with Context. The context is optional - tools that don't need it can omit the parameter. + """ - _request_context: RequestContext[ServerSessionT, LifespanContextT] | None - _fastmcp: FastMCP | None + def __init__(self, fastmcp: FastMCP): + self.fastmcp = fastmcp + self._tokens: list[Token] = [] - model_config = ConfigDict(arbitrary_types_allowed=True) + def __enter__(self) -> Context: + """Enter the context manager and set this context as the current context.""" + # Always set this context and save the token + token = _current_context.set(self) + self._tokens.append(token) + return self - def __init__( - self, - *, - request_context: RequestContext[ServerSessionT, LifespanContextT] | None = None, - fastmcp: FastMCP | None = None, - **kwargs: Any, - ): - super().__init__(**kwargs) - self._request_context = request_context - self._fastmcp = fastmcp + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + """Exit the context manager and reset the most recent token.""" + if self._tokens: + token = self._tokens.pop() + _current_context.reset(token) @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[ServerSessionT, LifespanContextT]: + 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 + return self.fastmcp._mcp_server.request_context async def report_progress( self, progress: float, total: float | None = None @@ -120,10 +126,8 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]): Returns: The resource content as either text or bytes """ - assert self._fastmcp is not None, ( - "Context is not available outside of a request" - ) - return await self._fastmcp._mcp_read_resource(uri) + assert self.fastmcp is not None, "Context is not available outside of a request" + return await self.fastmcp._mcp_read_resource(uri) async def log( self, @@ -229,7 +233,5 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]): def get_http_request(self) -> Request: """Get the active starlette request.""" - request = get_current_starlette_request() - if request is None: - raise ValueError("Request is not available outside a Starlette request") - return request + + return fastmcp.server.dependencies.get_http_request() diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py new file mode 100644 index 000000000..a06560be1 --- /dev/null +++ b/src/fastmcp/server/dependencies.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, ParamSpec, TypeVar + +from starlette.requests import Request + +if TYPE_CHECKING: + from fastmcp.server.context import Context + +P = ParamSpec("P") +R = TypeVar("R") + + +# --- Context --- + + +def get_context() -> Context: + from fastmcp.server.context import _current_context + + context = _current_context.get() + if context is None: + raise RuntimeError("No active context found.") + return context + + +# --- HTTP Request --- + + +def get_http_request() -> Request: + from fastmcp.server.http import _current_http_request + + request = _current_http_request.get() + if request is None: + raise RuntimeError("No active HTTP request found.") + return request diff --git a/src/fastmcp/server/http.py b/src/fastmcp/server/http.py index 31f6d404e..b7398178d 100644 --- a/src/fastmcp/server/http.py +++ b/src/fastmcp/server/http.py @@ -1,8 +1,7 @@ from __future__ import annotations -from contextlib import ( - asynccontextmanager, -) +from collections.abc import Generator +from contextlib import contextmanager from contextvars import ContextVar from starlette.requests import Request @@ -11,27 +10,22 @@ from fastmcp.utilities.logging import get_logger logger = get_logger(__name__) - -_current_starlette_request: ContextVar[Request | None] = ContextVar( - "starlette_request", +_current_http_request: ContextVar[Request | None] = ContextVar( + "http_request", default=None, ) -@asynccontextmanager -async def starlette_request_context(request: Request): - token = _current_starlette_request.set(request) +@contextmanager +def set_http_request(request: Request) -> Generator[Request, None, None]: + token = _current_http_request.set(request) try: - yield + yield request finally: - _current_starlette_request.reset(token) + _current_http_request.reset(token) -def get_current_starlette_request() -> Request | None: - return _current_starlette_request.get() - - -class RequestMiddleware: +class RequestContextMiddleware: """ Middleware that stores each request in a ContextVar """ @@ -40,5 +34,5 @@ class RequestMiddleware: self.app = app async def __call__(self, scope, receive, send): - async with starlette_request_context(Request(scope)): + with set_http_request(Request(scope)): await self.app(scope, receive, send) diff --git a/src/fastmcp/server/openapi.py b/src/fastmcp/server/openapi.py index 82a04307b..b6316c892 100644 --- a/src/fastmcp/server/openapi.py +++ b/src/fastmcp/server/openapi.py @@ -25,9 +25,6 @@ from fastmcp.utilities.openapi import ( ) if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - from fastmcp.server import Context logger = get_logger(__name__) @@ -132,7 +129,6 @@ class OpenAPITool(Tool): description=description, parameters=parameters, fn=self._execute_request, # We'll use an instance method instead of a global function - context_kwarg="context", # Default context keyword argument tags=tags, annotations=annotations, serializer=serializer, @@ -258,12 +254,10 @@ class OpenAPITool(Tool): raise ValueError(f"Request error: {str(e)}") async def run( - self, - arguments: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, + self, arguments: dict[str, Any] ) -> list[TextContent | ImageContent | EmbeddedResource]: """Run the tool with arguments and optional context.""" - response = await self._execute_request(**arguments, context=context) + response = await self._execute_request(**arguments) return _convert_to_content(response) @@ -292,9 +286,7 @@ class OpenAPIResource(Resource): self._route = route self._timeout = timeout - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> str | bytes: + async def read(self) -> str | bytes: """Fetch the resource data by making an HTTP request.""" try: # Extract path parameters from the URI if present @@ -399,7 +391,6 @@ class OpenAPIResourceTemplate(ResourceTemplate): fn=lambda **kwargs: None, parameters=parameters, tags=tags, - context_kwarg=None, ) self._client = client self._route = route @@ -409,7 +400,7 @@ class OpenAPIResourceTemplate(ResourceTemplate): self, uri: str, params: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, + context: Context | None = None, ) -> Resource: """Create a resource with the given parameters.""" # Generate a URI for this resource instance @@ -650,7 +641,5 @@ class FastMCPOpenAPI(FastMCP): async def _mcp_call_tool(self, name: str, arguments: dict[str, Any]) -> Any: """Override the call_tool method to return the raw result without converting to content.""" - - context = self.get_context() - result = await self._tool_manager.call_tool(name, arguments, context=context) + result = await self._tool_manager.call_tool(name, arguments) return result diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 038623aab..da895a180 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -27,9 +27,6 @@ from fastmcp.tools.tool import Tool from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - from fastmcp.server import Context logger = get_logger(__name__) @@ -57,7 +54,7 @@ class ProxyTool(Tool): async def run( self, arguments: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, + context: Context | None = None, ) -> list[TextContent | ImageContent | EmbeddedResource]: # the client context manager will swallow any exceptions inside a TaskGroup # so we return the raw result and raise an exception ourselves @@ -89,9 +86,7 @@ class ProxyResource(Resource): mime_type=resource.mimeType, ) - async def read( - self, context: Context[ServerSessionT, LifespanContextT] | None = None - ) -> str | bytes: + async def read(self) -> str | bytes: if self._value is not None: return self._value @@ -127,7 +122,7 @@ class ProxyTemplate(ResourceTemplate): self, uri: str, params: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, + context: Context | None = None, ) -> ProxyResource: # dont use the provided uri, because it may not be the same as the # uri_template on the remote server. @@ -171,11 +166,7 @@ class ProxyPrompt(Prompt): fn=_proxy_passthrough, ) - async def render( - self, - arguments: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, - ) -> list[PromptMessage]: + async def render(self, arguments: dict[str, Any]) -> list[PromptMessage]: async with self._client: result = await self._client.get_prompt(self.name, arguments) return result.messages diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index 073496a90..4fb6b7eb4 100644 --- a/src/fastmcp/tools/tool.py +++ b/src/fastmcp/tools/tool.py @@ -12,6 +12,7 @@ from pydantic import BaseModel, BeforeValidator, Field import fastmcp from fastmcp.exceptions import ToolError +from fastmcp.server.dependencies import get_context from fastmcp.utilities.json_schema import prune_params from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import ( @@ -22,10 +23,7 @@ from fastmcp.utilities.types import ( ) if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - from mcp.shared.context import LifespanContextT - - from fastmcp.server import Context + pass logger = get_logger(__name__) @@ -41,9 +39,6 @@ class Tool(BaseModel): name: str = Field(description="Name of the tool") description: str = Field(description="Description of what the tool does") parameters: dict[str, Any] = Field(description="JSON schema for tool parameters") - context_kwarg: str | None = Field( - None, description="Name of the kwarg that should receive context" - ) tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field( default_factory=set, description="Tags for the tool" ) @@ -60,13 +55,12 @@ class Tool(BaseModel): fn: Callable[..., Any], name: str | None = None, description: str | None = None, - context_kwarg: str | None = None, tags: set[str] | None = None, annotations: ToolAnnotations | None = None, serializer: Callable[[Any], str] | None = None, ) -> Tool: """Create a Tool from a function.""" - from fastmcp import Context + from fastmcp.server.context import Context # Reject functions with *args or **kwargs sig = inspect.signature(fn) @@ -86,8 +80,7 @@ class Tool(BaseModel): type_adapter = get_cached_typeadapter(fn) schema = type_adapter.json_schema() - if context_kwarg is None: - context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) + context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) if context_kwarg: schema = prune_params(schema, params=[context_kwarg]) @@ -96,25 +89,23 @@ class Tool(BaseModel): name=func_name, description=func_doc, parameters=schema, - context_kwarg=context_kwarg, tags=tags or set(), annotations=annotations, serializer=serializer, ) async def run( - self, - arguments: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, + self, arguments: dict[str, Any] ) -> list[TextContent | ImageContent | EmbeddedResource]: """Run the tool with arguments.""" + from fastmcp.server.context import Context + + arguments = arguments.copy() try: - injected_args = ( - {self.context_kwarg: context} if self.context_kwarg is not None else {} - ) - - parsed_args = arguments.copy() + context_kwarg = find_kwarg_by_type(self.fn, kwarg_type=Context) + if context_kwarg and context_kwarg not in arguments: + arguments[context_kwarg] = get_context() if fastmcp.settings.settings.tool_attempt_parse_json_args: # Pre-parse data from JSON in order to handle cases like `["a", "b", "c"]` @@ -125,7 +116,7 @@ class Tool(BaseModel): # which can be pre-parsed here. signature = inspect.signature(self.fn) for param_name in self.parameters["properties"]: - arg = parsed_args.get(param_name, None) + arg = arguments.get(param_name, None) # if not in signature, we won't have annotations, so skip logic if param_name not in signature.parameters: continue @@ -140,13 +131,13 @@ class Tool(BaseModel): ): continue try: - parsed_args[param_name] = json.loads(arg) + arguments[param_name] = json.loads(arg) except json.JSONDecodeError: pass type_adapter = get_cached_typeadapter(self.fn) - result = type_adapter.validate_python(parsed_args | injected_args) + result = type_adapter.validate_python(arguments) if inspect.isawaitable(result): result = await result diff --git a/src/fastmcp/tools/tool_manager.py b/src/fastmcp/tools/tool_manager.py index e096fae32..21e38c5f2 100644 --- a/src/fastmcp/tools/tool_manager.py +++ b/src/fastmcp/tools/tool_manager.py @@ -3,7 +3,6 @@ from __future__ import annotations as _annotations from collections.abc import Callable from typing import TYPE_CHECKING, Any -from mcp.shared.context import LifespanContextT from mcp.types import EmbeddedResource, ImageContent, TextContent, ToolAnnotations from fastmcp.exceptions import NotFoundError @@ -12,9 +11,7 @@ from fastmcp.tools.tool import Tool from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: - from mcp.server.session import ServerSessionT - - from fastmcp.server import Context + pass logger = get_logger(__name__) @@ -98,14 +95,11 @@ class ToolManager: return tool async def call_tool( - self, - key: str, - arguments: dict[str, Any], - context: Context[ServerSessionT, LifespanContextT] | None = None, + self, key: str, arguments: dict[str, Any] ) -> list[TextContent | ImageContent | EmbeddedResource]: """Call a tool by name with arguments.""" tool = self.get_tool(key) if not tool: raise NotFoundError(f"Unknown tool: {key}") - return await tool.run(arguments, context=context) + return await tool.run(arguments) diff --git a/tests/prompts/test_prompt_manager.py b/tests/prompts/test_prompt_manager.py index 3a41003e7..c887fdca5 100644 --- a/tests/prompts/test_prompt_manager.py +++ b/tests/prompts/test_prompt_manager.py @@ -1,8 +1,6 @@ from typing import Annotated import pytest -from mcp.server.session import ServerSessionT -from mcp.shared.context import LifespanContextT from fastmcp import Context from fastmcp.exceptions import NotFoundError @@ -308,38 +306,30 @@ class TestContextHandling: def prompt_with_context(x: int, ctx: Context) -> str: return str(x) - prompt = Prompt.from_function(prompt_with_context) - assert prompt.context_kwarg == "ctx" + Prompt.from_function(prompt_with_context) def prompt_without_context(x: int) -> str: return str(x) - prompt = Prompt.from_function(prompt_without_context) - assert prompt.context_kwarg is None + Prompt.from_function(prompt_without_context) def test_parameterized_context_parameter_detection(self): """Test that parameterized context parameters are properly detected in Prompt.from_function().""" - def prompt_with_context( - x: int, ctx: Context[ServerSessionT, LifespanContextT] - ) -> str: + def prompt_with_context(x: int, ctx: Context) -> str: return str(x) - prompt = Prompt.from_function(prompt_with_context) - assert prompt.context_kwarg == "ctx" + Prompt.from_function(prompt_with_context) def test_parameterized_union_context_parameter_detection(self): """Test that context parameters in a union are properly detected in Prompt.from_function().""" - def prompt_with_context( - x: int, ctx: Context[ServerSessionT, LifespanContextT] | None - ) -> str: + def prompt_with_context(x: int, ctx: Context | None) -> str: return str(x) - prompt = Prompt.from_function(prompt_with_context) - assert prompt.context_kwarg == "ctx" + Prompt.from_function(prompt_with_context) async def test_context_injection(self): """Test that context is properly injected during prompt rendering.""" @@ -349,17 +339,15 @@ class TestContextHandling: return str(x) prompt = Prompt.from_function(prompt_with_context) - assert prompt.context_kwarg == "ctx" from fastmcp import FastMCP mcp = FastMCP() - ctx = mcp.get_context() + context = Context(fastmcp=mcp) + + with context: + messages = await prompt.render(arguments={"x": 42}) - messages = await prompt.render( - arguments={"x": 42}, - context=ctx, - ) assert len(messages) == 1 assert isinstance(messages[0].content, TextContent) assert messages[0].content.text == "42" @@ -371,12 +359,18 @@ class TestContextHandling: return str(x) prompt = Prompt.from_function(prompt_with_context) - assert prompt.context_kwarg == "ctx" - # Should not raise an error when context is not provided - messages = await prompt.render( - arguments={"x": 42}, - ) + # Even for optional context, we need to provide a context + from fastmcp import FastMCP + + mcp = FastMCP() + context = Context(fastmcp=mcp) + + with context: + messages = await prompt.render( + arguments={"x": 42}, + ) + assert len(messages) == 1 assert isinstance(messages[0].content, TextContent) assert messages[0].content.text == "42" @@ -388,5 +382,4 @@ class TestContextHandling: def prompt_with_context(x: int, ctx: Annotated[Context, "ctx"]) -> str: return str(x) - prompt = Prompt.from_function(prompt_with_context) - assert prompt.context_kwarg == "ctx" + Prompt.from_function(prompt_with_context) diff --git a/tests/resources/test_resource_template.py b/tests/resources/test_resource_template.py index 563876f25..00f2b01b6 100644 --- a/tests/resources/test_resource_template.py +++ b/tests/resources/test_resource_template.py @@ -2,8 +2,6 @@ import json from urllib.parse import quote import pytest -from mcp.server.session import ServerSessionT -from mcp.shared.context import LifespanContextT from pydantic import BaseModel from fastmcp import Context @@ -560,54 +558,46 @@ class TestContextHandling: def template_with_context(x: int, ctx: Context) -> str: return str(x) - template = ResourceTemplate.from_function( + ResourceTemplate.from_function( fn=template_with_context, uri_template="test://{x}", name="test", ) - assert template.context_kwarg == "ctx" def template_without_context(x: int) -> str: return str(x) - template = ResourceTemplate.from_function( + ResourceTemplate.from_function( fn=template_without_context, uri_template="test://{x}", name="test", ) - assert template.context_kwarg is None def test_parameterized_context_parameter_detection(self): """Test that parameterized context parameters are properly detected in ResourceTemplate.from_function().""" - def template_with_context( - x: int, ctx: Context[ServerSessionT, LifespanContextT] - ) -> str: + def template_with_context(x: int, ctx: Context) -> str: return str(x) - template = ResourceTemplate.from_function( + ResourceTemplate.from_function( fn=template_with_context, uri_template="test://{x}", name="test", ) - assert template.context_kwarg == "ctx" def test_parameterized_union_context_parameter_detection(self): """Test that context parameters in a union are properly detected in ResourceTemplate.from_function().""" - def template_with_context( - x: int, ctx: Context[ServerSessionT, LifespanContextT] | None - ) -> str: + def template_with_context(x: int, ctx: Context | None) -> str: return str(x) - template = ResourceTemplate.from_function( + ResourceTemplate.from_function( fn=template_with_context, uri_template="test://{x}", name="test", ) - assert template.context_kwarg == "ctx" async def test_context_injection(self): """Test that context is properly injected during resource creation.""" @@ -621,18 +611,18 @@ class TestContextHandling: uri_template="test://{x}", name="test", ) - assert template.context_kwarg == "ctx" from fastmcp import FastMCP mcp = FastMCP() - ctx = mcp.get_context() + context = Context(fastmcp=mcp) + + with context: + resource = await template.create_resource( + "test://42", + {"x": 42}, + ) - resource = await template.create_resource( - "test://42", - {"x": 42}, - context=ctx, - ) assert isinstance(resource, FunctionResource) content = await resource.read() assert content == "42" @@ -648,13 +638,19 @@ class TestContextHandling: uri_template="test://{x}", name="test", ) - assert template.context_kwarg == "ctx" - # Should not raise an error when context is not provided - resource = await template.create_resource( - "test://42", - {"x": 42}, - ) + # Even for optional context, we need to provide a context + from fastmcp import FastMCP + + mcp = FastMCP() + context = Context(fastmcp=mcp) + + with context: + resource = await template.create_resource( + "test://42", + {"x": 42}, + ) + assert isinstance(resource, FunctionResource) content = await resource.read() assert content == "42" diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index 5a4f3d122..4d2bc2711 100644 --- a/tests/tools/test_tool_manager.py +++ b/tests/tools/test_tool_manager.py @@ -5,8 +5,6 @@ from typing import Annotated, Any import pydantic_core import pytest -from mcp.server.session import ServerSessionT -from mcp.shared.context import LifespanContextT from mcp.types import ImageContent, TextContent from pydantic import BaseModel @@ -403,10 +401,20 @@ class TestCallTools: manager = ToolManager() manager.add_tool_from_fn(name_shrimp) - result = await manager.call_tool( - "name_shrimp", - {"tank": {"x": None, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}}, - ) + mcp = FastMCP() + context = Context(fastmcp=mcp) + + with context: + result = await manager.call_tool( + "name_shrimp", + { + "tank": { + "x": None, + "shrimp": [{"name": "rex"}, {"name": "gertrude"}], + } + }, + ) + assert isinstance(result, list) assert len(result) == 1 assert isinstance(result[0], TextContent) @@ -498,14 +506,12 @@ class TestContextHandling: return str(x) manager = ToolManager() - tool = manager.add_tool_from_fn(tool_with_context) - assert tool.context_kwarg == "ctx" + manager.add_tool_from_fn(tool_with_context) def tool_without_context(x: int) -> str: return str(x) - tool = manager.add_tool_from_fn(tool_without_context) - assert tool.context_kwarg is None + manager.add_tool_from_fn(tool_without_context) async def test_context_injection(self): """Test that context is properly injected during tool execution.""" @@ -515,16 +521,17 @@ class TestContextHandling: return str(x) manager = ToolManager() - tool = manager.add_tool_from_fn(tool_with_context) - assert tool.context_kwarg == "ctx" + manager.add_tool_from_fn(tool_with_context) mcp = FastMCP() - ctx = mcp.get_context() - result = await manager.call_tool("tool_with_context", {"x": 42}, context=ctx) - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], TextContent) - assert result[0].text == "42" + context = Context(fastmcp=mcp) + + with context: + result = await manager.call_tool("tool_with_context", {"x": 42}) + assert isinstance(result, list) + assert len(result) == 1 + assert isinstance(result[0], TextContent) + assert result[0].text == "42" async def test_context_injection_async(self): """Test that context is properly injected in async tools.""" @@ -534,16 +541,17 @@ class TestContextHandling: return str(x) manager = ToolManager() - tool = manager.add_tool_from_fn(async_tool) - assert tool.context_kwarg == "ctx" + manager.add_tool_from_fn(async_tool) mcp = FastMCP() - ctx = mcp.get_context() - result = await manager.call_tool("async_tool", {"x": 42}, context=ctx) - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], TextContent) - assert result[0].text == "42" + context = Context(fastmcp=mcp) + + with context: + result = await manager.call_tool("async_tool", {"x": 42}) + assert isinstance(result, list) + assert len(result) == 1 + assert isinstance(result[0], TextContent) + assert result[0].text == "42" async def test_context_optional(self): """Test that context is optional when calling tools.""" @@ -553,48 +561,45 @@ class TestContextHandling: return x manager = ToolManager() - tool = manager.add_tool_from_fn(tool_with_context) - assert tool.context_kwarg == "ctx" + manager.add_tool_from_fn(tool_with_context) # Should not raise an error when context is not provided - result = await manager.call_tool("tool_with_context", {"x": 42}) - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], TextContent) - assert result[0].text == "42" + + mcp = FastMCP() + context = Context(fastmcp=mcp) + + with context: + result = await manager.call_tool("tool_with_context", {"x": 42}) + assert isinstance(result, list) + assert len(result) == 1 + assert isinstance(result[0], TextContent) + assert result[0].text == "42" def test_parameterized_context_parameter_detection(self): """Test that context parameters are properly detected in Tool.from_function().""" - def tool_with_context( - x: int, ctx: Context[ServerSessionT, LifespanContextT] - ) -> str: + def tool_with_context(x: int, ctx: Context) -> str: return str(x) manager = ToolManager() - tool = manager.add_tool_from_fn(tool_with_context) - assert tool.context_kwarg == "ctx" + manager.add_tool_from_fn(tool_with_context) def test_annotated_context_parameter_detection(self): def tool_with_context(x: int, ctx: Annotated[Context, "ctx"]) -> str: return str(x) manager = ToolManager() - tool = manager.add_tool_from_fn(tool_with_context) - assert tool.context_kwarg == "ctx" + manager.add_tool_from_fn(tool_with_context) def test_parameterized_union_context_parameter_detection(self): """Test that context parameters are properly detected in Tool.from_function().""" - def tool_with_context( - x: int, ctx: Context[ServerSessionT, LifespanContextT] | None - ) -> str: + def tool_with_context(x: int, ctx: Context | None) -> str: return str(x) manager = ToolManager() - tool = manager.add_tool_from_fn(tool_with_context) - assert tool.context_kwarg == "ctx" + manager.add_tool_from_fn(tool_with_context) async def test_context_error_handling(self): """Test error handling when context injection fails.""" @@ -606,9 +611,13 @@ class TestContextHandling: manager.add_tool_from_fn(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) + context = Context(fastmcp=mcp) + + with context: + with pytest.raises( + ToolError, match="Error executing tool tool_with_context" + ): + await manager.call_tool("tool_with_context", {"x": 42}) class TestCustomToolNames: