diff --git a/src/fastmcp/prompts/prompt.py b/src/fastmcp/prompts/prompt.py index c6bcefc62..0cf37403d 100644 --- a/src/fastmcp/prompts/prompt.py +++ b/src/fastmcp/prompts/prompt.py @@ -1,9 +1,11 @@ """Base classes for FastMCP prompts.""" +from __future__ import annotations as _annotations + import inspect import json from collections.abc import Awaitable, Callable, Sequence -from typing import Annotated, Any, Literal +from typing import TYPE_CHECKING, Annotated, Any, Literal import pydantic_core from mcp.types import EmbeddedResource, ImageContent, TextContent @@ -13,6 +15,12 @@ from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_ca 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 + CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource @@ -72,6 +80,9 @@ 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( @@ -80,7 +91,8 @@ class Prompt(BaseModel): name: str | None = None, description: str | None = None, tags: set[str] | None = None, - ) -> "Prompt": + context_kwarg: str | None = None, + ) -> Prompt: """Create a Prompt from a function. The function can return: @@ -89,11 +101,24 @@ class Prompt(BaseModel): - A dict (converted to a message) - A sequence of any of the above """ + from fastmcp import Context + func_name = name or fn.__name__ if func_name == "": raise ValueError("You must provide a name for lambda functions") + # Auto-detect context parameter if not provided + if context_kwarg is None: + if inspect.ismethod(fn) and hasattr(fn, "__func__"): + sig = inspect.signature(fn.__func__) + else: + sig = inspect.signature(fn) + for param_name, param in sig.parameters.items(): + if param.annotation is Context: + context_kwarg = param_name + break + # Get schema from TypeAdapter - will fail if function isn't properly typed parameters = TypeAdapter(fn).json_schema() @@ -101,6 +126,10 @@ class Prompt(BaseModel): arguments: list[PromptArgument] = [] if "properties" in parameters: for param_name, param in parameters["properties"].items(): + # Skip context parameter + if param_name == context_kwarg: + continue + required = param_name in parameters.get("required", []) arguments.append( PromptArgument( @@ -119,9 +148,14 @@ 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) -> list[Message]: + async def render( + self, + arguments: dict[str, Any] | None = None, + context: Context[ServerSessionT, LifespanContextT] | None = None, + ) -> list[Message]: """Render the prompt with arguments.""" # Validate required arguments if self.arguments: @@ -132,8 +166,13 @@ class Prompt(BaseModel): raise ValueError(f"Missing required arguments: {missing}") 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 + # Call function and check if result is a coroutine - result = self.fn(**(arguments or {})) + result = self.fn(**kwargs) if inspect.iscoroutine(result): result = await result diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index a48ec91b9..db16d4398 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -1,13 +1,21 @@ """Prompt management functionality.""" +from __future__ import annotations as _annotations + from collections.abc import Awaitable, Callable -from typing import Any +from typing import TYPE_CHECKING, Any from fastmcp.exceptions import NotFoundError from fastmcp.prompts.prompt import Message, Prompt, PromptResult 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 + logger = get_logger(__name__) @@ -69,14 +77,17 @@ class PromptManager: return prompt async def render_prompt( - self, name: str, arguments: dict[str, Any] | None = None + self, + name: str, + arguments: dict[str, Any] | None = None, + context: Context[ServerSessionT, LifespanContextT] | None = None, ) -> list[Message]: """Render a prompt by name with arguments.""" prompt = self.get_prompt(name) if not prompt: raise NotFoundError(f"Unknown prompt: {name}") - return await prompt.render(arguments) + return await prompt.render(arguments, context=context) 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 b79381bbd..e4d69638d 100644 --- a/src/fastmcp/resources/resource.py +++ b/src/fastmcp/resources/resource.py @@ -1,7 +1,9 @@ """Base classes and interfaces for FastMCP resources.""" +from __future__ import annotations + import abc -from typing import Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any from mcp.types import Resource as MCPResource from pydantic import ( @@ -17,6 +19,12 @@ 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 + class Resource(BaseModel, abc.ABC): """Base class for all resources.""" @@ -58,7 +66,9 @@ class Resource(BaseModel, abc.ABC): raise ValueError("Either name or uri must be provided") @abc.abstractmethod - async def read(self) -> str | bytes: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str | bytes: """Read the resource content.""" pass diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index 613cd728b..1cf42d839 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -212,9 +212,13 @@ class ResourceManager: return True return False - async def get_resource(self, uri: AnyUrl | str) -> Resource: + async def get_resource(self, uri: AnyUrl | str, context=None) -> 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. """ @@ -230,7 +234,9 @@ class ResourceManager: # Try to match against the storage key (which might be a custom key) if params := match_uri_template(uri_str, storage_key): try: - return await template.create_resource(uri_str, params) + return await template.create_resource( + uri_str, 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 0d804f3da..41516dd91 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 Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any from urllib.parse import unquote from mcp.types import ResourceTemplate as MCPResourceTemplate @@ -22,6 +22,12 @@ from pydantic import ( from fastmcp.resources.types import FunctionResource, Resource 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 + def build_regex(template: str) -> re.Pattern: parts = re.split(r"(\{[^}]+\})", template) @@ -70,6 +76,9 @@ 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 @@ -88,18 +97,34 @@ 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 + func_name = name or fn.__name__ if func_name == "": raise ValueError("You must provide a name for lambda functions") + # Auto-detect context parameter if not provided + if context_kwarg is None: + if inspect.ismethod(fn) and hasattr(fn, "__func__"): + sig = inspect.signature(fn.__func__) + else: + sig = inspect.signature(fn) + for param_name, param in sig.parameters.items(): + if param.annotation is Context: + context_kwarg = param_name + break + # Validate that URI params match function params uri_params = set(re.findall(r"{(\w+)(?:\*)?}", uri_template)) if not uri_params: raise ValueError("URI template must contain at least one parameter") func_params = set(inspect.signature(fn).parameters.keys()) + if context_kwarg: + func_params.discard(context_kwarg) # get the parameters that are required required_params = { @@ -107,6 +132,8 @@ class ResourceTemplate(BaseModel): for p in func_params if inspect.signature(fn).parameters[p].default is inspect.Parameter.empty } + if context_kwarg and context_kwarg in required_params: + required_params.discard(context_kwarg) if not required_params.issubset(uri_params): raise ValueError( @@ -132,17 +159,28 @@ 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]) -> Resource: + async def create_resource( + self, + uri: str, + params: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, + ) -> Resource: """Create a resource from the template with the given parameters.""" 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 + # Call function and check if result is a coroutine - result = self.fn(**params) + result = self.fn(**kwargs) if inspect.iscoroutine(result): result = await result @@ -153,6 +191,7 @@ class ResourceTemplate(BaseModel): mime_type=self.mime_type, fn=lambda: 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 30d168bee..e86733e8a 100644 --- a/src/fastmcp/resources/types.py +++ b/src/fastmcp/resources/types.py @@ -1,10 +1,12 @@ """Concrete resource implementations.""" +from __future__ import annotations + import inspect import json from collections.abc import Callable from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import anyio import anyio.to_thread @@ -15,13 +17,21 @@ from pydantic import Field, ValidationInfo from fastmcp.resources.resource import Resource +if TYPE_CHECKING: + from mcp.server.session import ServerSessionT + from mcp.shared.context import LifespanContextT + + from fastmcp.server import Context + class TextResource(Resource): """A resource that reads from a string.""" text: str = Field(description="Text content of the resource") - async def read(self) -> str: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str: """Read the text content.""" return self.text @@ -31,7 +41,9 @@ class BinaryResource(Resource): data: bytes = Field(description="Binary content of the resource") - async def read(self) -> bytes: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> bytes: """Read the binary content.""" return self.data @@ -50,13 +62,23 @@ class FunctionResource(Resource): """ fn: Callable[[], Any] + context_kwarg: str | None = Field( + default=None, description="Name of the kwarg that should receive context" + ) - async def read(self) -> str | bytes: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str | bytes: """Read the resource by calling the wrapped function.""" try: - result = ( - await self.fn() if inspect.iscoroutinefunction(self.fn) else self.fn() - ) + kwargs = {} + if self.context_kwarg is not None: + kwargs[self.context_kwarg] = context + + result = self.fn(**kwargs) + if inspect.iscoroutinefunction(self.fn): + result = await result + if isinstance(result, Resource): return await result.read() if isinstance(result, bytes): diff --git a/src/fastmcp/server/openapi.py b/src/fastmcp/server/openapi.py index 0911cae8a..ae2776af8 100644 --- a/src/fastmcp/server/openapi.py +++ b/src/fastmcp/server/openapi.py @@ -1,11 +1,13 @@ """FastMCP server implementation for OpenAPI integration.""" +from __future__ import annotations + import enum import json import re from dataclasses import dataclass from re import Pattern -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal import httpx from mcp.types import TextContent @@ -22,6 +24,12 @@ from fastmcp.utilities.openapi import ( format_description_with_responses, ) +if TYPE_CHECKING: + from mcp.server.session import ServerSessionT + from mcp.shared.context import LifespanContextT + + from fastmcp.server import Context + logger = get_logger(__name__) HttpMethod = Literal["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"] @@ -347,11 +355,17 @@ class OpenAPIResourceTemplate(ResourceTemplate): fn=lambda **kwargs: None, parameters=parameters, tags=tags, + context_kwarg=None, ) self._client = client self._route = route - async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: + async def create_resource( + self, + uri: str, + params: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, + ) -> Resource: """Create a resource with the given parameters.""" # Generate a URI for this resource instance uri_parts = [] diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index b3d9d8dce..e7cd35d68 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -1,4 +1,6 @@ -from typing import Any, cast +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, cast from urllib.parse import quote import mcp.types @@ -25,6 +27,12 @@ from fastmcp.tools.tool import Tool from fastmcp.utilities.func_metadata import func_metadata 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__) @@ -33,12 +41,12 @@ def _proxy_passthrough(): class ProxyTool(Tool): - def __init__(self, client: "Client", **kwargs): + def __init__(self, client: Client, **kwargs): super().__init__(**kwargs) self._client = client @classmethod - async def from_client(cls, client: "Client", tool: mcp.types.Tool) -> "ProxyTool": + async def from_client(cls, client: Client, tool: mcp.types.Tool) -> ProxyTool: return cls( client=client, name=tool.name, @@ -50,7 +58,9 @@ class ProxyTool(Tool): ) async def run( - self, arguments: dict[str, Any], context: Context | None = None + self, + arguments: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, ) -> Any: # the client context manager will swallow any exceptions inside a TaskGroup # so we return the raw result and raise an exception ourselves @@ -64,17 +74,15 @@ class ProxyTool(Tool): class ProxyResource(Resource): - def __init__( - self, client: "Client", *, _value: str | bytes | None = None, **kwargs - ): + def __init__(self, client: Client, *, _value: str | bytes | None = None, **kwargs): super().__init__(**kwargs) self._client = client self._value = _value @classmethod async def from_client( - cls, client: "Client", resource: mcp.types.Resource - ) -> "ProxyResource": + cls, client: Client, resource: mcp.types.Resource + ) -> ProxyResource: return cls( client=client, uri=resource.uri, @@ -83,7 +91,9 @@ class ProxyResource(Resource): mime_type=resource.mimeType, ) - async def read(self) -> str | bytes: + async def read( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> str | bytes: if self._value is not None: return self._value @@ -98,14 +108,14 @@ class ProxyResource(Resource): class ProxyTemplate(ResourceTemplate): - def __init__(self, client: "Client", **kwargs): + def __init__(self, client: Client, **kwargs): super().__init__(**kwargs) self._client = client @classmethod async def from_client( - cls, client: "Client", template: mcp.types.ResourceTemplate - ) -> "ProxyTemplate": + cls, client: Client, template: mcp.types.ResourceTemplate + ) -> ProxyTemplate: return cls( client=client, uri_template=template.uriTemplate, @@ -115,7 +125,12 @@ class ProxyTemplate(ResourceTemplate): parameters={}, ) - async def create_resource(self, uri: str, params: dict[str, Any]) -> ProxyResource: + async def create_resource( + self, + uri: str, + params: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, + ) -> ProxyResource: # dont use the provided uri, because it may not be the same as the # uri_template on the remote server. # quote params to ensure they are valid for the uri_template @@ -144,14 +159,12 @@ class ProxyTemplate(ResourceTemplate): class ProxyPrompt(Prompt): - def __init__(self, client: "Client", **kwargs): + def __init__(self, client: Client, **kwargs): super().__init__(**kwargs) self._client = client @classmethod - async def from_client( - cls, client: "Client", prompt: mcp.types.Prompt - ) -> "ProxyPrompt": + async def from_client(cls, client: Client, prompt: mcp.types.Prompt) -> ProxyPrompt: return cls( client=client, name=prompt.name, @@ -160,14 +173,18 @@ class ProxyPrompt(Prompt): fn=_proxy_passthrough, ) - async def render(self, arguments: dict[str, Any]) -> list[Message]: + async def render( + self, + arguments: dict[str, Any], + context: Context[ServerSessionT, LifespanContextT] | None = None, + ) -> list[Message]: async with self._client: result = await self._client.get_prompt(self.name, arguments) return [Message(role=m.role, content=m.content) for m in result] class FastMCPProxy(FastMCP): - def __init__(self, client: "Client", **kwargs): + def __init__(self, client: Client, **kwargs): super().__init__(**kwargs) self.client = client diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 6f2167410..1ca3b116a 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -398,7 +398,8 @@ class FastMCP(Generic[LifespanResultT]): server. """ if self._resource_manager.has_resource(uri): - resource = await self._resource_manager.get_resource(uri) + context = self.get_context() + resource = await self._resource_manager.get_resource(uri, context=context) try: content = await resource.read() return [ @@ -424,7 +425,10 @@ class FastMCP(Generic[LifespanResultT]): """ if self._prompt_manager.has_prompt(name): - messages = await self._prompt_manager.render_prompt(name, arguments) + context = self.get_context() + messages = await self._prompt_manager.render_prompt( + name, arguments, context=context + ) return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages)) else: for server in self._mounted_servers.values(): @@ -562,6 +566,10 @@ class FastMCP(Generic[LifespanResultT]): - bytes for binary content - other types will be converted to JSON + Resources can optionally request a Context object by adding a parameter with the + Context type annotation. The context provides access to MCP capabilities like + logging, progress reporting, and session information. + If the URI contains parameters (e.g. "resource://{param}") or the function has parameters, it will be registered as a template resource. @@ -586,6 +594,11 @@ class FastMCP(Generic[LifespanResultT]): def get_weather(city: str) -> str: return f"Weather for {city}" + @server.resource("resource://{city}/weather") + def get_weather_with_context(city: str, ctx: Context) -> str: + ctx.info(f"Fetching weather for {city}") + return f"Weather for {city}" + @server.resource("resource://{city}/weather") async def get_weather(city: str) -> str: data = await fetch_weather(city) @@ -639,6 +652,10 @@ class FastMCP(Generic[LifespanResultT]): ) -> Callable[[AnyFunction], AnyFunction]: """Decorator to register a prompt. + Prompts can optionally request a Context object by adding a parameter with the + Context type annotation. The context provides access to MCP capabilities like + logging, progress reporting, and session information. + Args: name: Optional name for the prompt (defaults to function name) description: Optional description of what the prompt does @@ -655,6 +672,17 @@ class FastMCP(Generic[LifespanResultT]): } ] + @server.prompt() + def analyze_with_context(table_name: str, ctx: Context) -> list[Message]: + ctx.info(f"Analyzing table {table_name}") + schema = read_table_schema(table_name) + return [ + { + "role": "user", + "content": f"Analyze this schema:\n{schema}" + } + ] + @server.prompt() async def analyze_file(path: str) -> list[Message]: content = await read_file(path)