diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py index c23f2523e..a430533b3 100644 --- a/src/mcp/client/session_group.py +++ b/src/mcp/client/session_group.py @@ -77,10 +77,10 @@ class ClientSessionGroup: the client and can be accessed via the session. Example Usage: - name_fn = lambda name, server_info: f"{(server_info.name)}-{name}" + name_fn = lambda name, server_info: f"{(server_info.name)}_{name}" async with ClientSessionGroup(component_name_hook=name_fn) as group: for server_params in server_params: - group.connect_to_server(server_param) + await group.connect_to_server(server_param) ... """ @@ -145,14 +145,15 @@ async def __aexit__( ) -> bool | None: """Closes session exit stacks and main exit stack upon completion.""" + # Only close the main exit stack if we created it + if self._owns_exit_stack: + await self._exit_stack.aclose() + # Concurrently close session stacks. async with anyio.create_task_group() as tg: for exit_stack in self._session_exit_stacks.values(): tg.start_soon(exit_stack.aclose) - # Only close the main exit stack if we created it - if self._owns_exit_stack: - await self._exit_stack.aclose() @property def sessions(self) -> list[mcp.ClientSession]: diff --git a/src/mcp/server/fastmcp/prompts/manager.py b/src/mcp/server/fastmcp/prompts/manager.py index 7ccbdef36..d0fb4bc08 100644 --- a/src/mcp/server/fastmcp/prompts/manager.py +++ b/src/mcp/server/fastmcp/prompts/manager.py @@ -18,6 +18,9 @@ def __init__(self, warn_on_duplicate_prompts: bool = True): def get_prompt(self, name: str) -> Prompt | None: """Get prompt by name.""" return self._prompts.get(name) + + def get_all_prompts(self) -> dict[str, Prompt]: + return self._prompts def list_prompts(self) -> list[Prompt]: """List all registered prompts.""" diff --git a/src/mcp/server/fastmcp/resources/resource_manager.py b/src/mcp/server/fastmcp/resources/resource_manager.py index d27e6ac12..1a4450dd1 100644 --- a/src/mcp/server/fastmcp/resources/resource_manager.py +++ b/src/mcp/server/fastmcp/resources/resource_manager.py @@ -20,6 +20,11 @@ def __init__(self, warn_on_duplicate_resources: bool = True): self._templates: dict[str, ResourceTemplate] = {} self.warn_on_duplicate_resources = warn_on_duplicate_resources + def get_all_resources( + self, + ) -> tuple[dict[str, Resource], dict[str, ResourceTemplate]]: + return self._resources, self._templates + def add_resource(self, resource: Resource) -> Resource: """Add a resource to the manager. diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index 3282baae6..961147e9f 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -38,6 +38,7 @@ from mcp.server.fastmcp.prompts import Prompt, PromptManager from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager from mcp.server.fastmcp.tools import Tool, ToolManager +from mcp.server.fastmcp.utilities.bundler import Bundler from mcp.server.fastmcp.utilities.logging import configure_logging, get_logger from mcp.server.fastmcp.utilities.types import Image from mcp.server.lowlevel.helper_types import ReadResourceContents @@ -587,6 +588,28 @@ def decorator( return decorator + def include_bundler(self, bundler: Bundler) -> None: + """Add bundler of resources, tools and prompts to the server.""" + bundler_tools = bundler.get_tools() + for tool_name, tool in bundler_tools.items(): + self.add_tool(tool.fn, tool_name, tool.description, tool.annotations) + + bundler_resources, bundler_templates = bundler.get_resources() + for resource in bundler_resources.values(): + self.add_resource(resource) + for template_name, template in bundler_templates.items(): + self._resource_manager.add_template( + template.fn, + template.uri_template, + template_name, + template.description, + template.mime_type, + ) + + bundler_prompts = bundler.get_prompts() + for prompt in bundler_prompts.values(): + self.add_prompt(prompt) + async def run_stdio_async(self) -> None: """Run the server using stdio transport.""" async with stdio_server() as (read_stream, write_stream): diff --git a/src/mcp/server/fastmcp/tools/tool_manager.py b/src/mcp/server/fastmcp/tools/tool_manager.py index 6ec4fd151..0ce11eb6b 100644 --- a/src/mcp/server/fastmcp/tools/tool_manager.py +++ b/src/mcp/server/fastmcp/tools/tool_manager.py @@ -37,6 +37,9 @@ def __init__( def get_tool(self, name: str) -> Tool | None: """Get tool by name.""" return self._tools.get(name) + + def get_all_tools(self) -> dict[str, Tool]: + return self._tools def list_tools(self) -> list[Tool]: """List all registered tools.""" diff --git a/src/mcp/server/fastmcp/utilities/bundler.py b/src/mcp/server/fastmcp/utilities/bundler.py new file mode 100644 index 000000000..94ed6015f --- /dev/null +++ b/src/mcp/server/fastmcp/utilities/bundler.py @@ -0,0 +1,259 @@ +import inspect +import re +from collections.abc import Callable + +from mcp.server.fastmcp.prompts import Prompt, PromptManager +from mcp.server.fastmcp.resources import ( + FunctionResource, + Resource, + ResourceManager, + ResourceTemplate, +) +from mcp.server.fastmcp.tools import Tool, ToolManager +from mcp.types import ( + AnyFunction, + ToolAnnotations, +) + + +class Bundler: + def __init__( + self, + tools: list[Tool] | None = None, + ): + self._tool_manager = ToolManager(tools=tools, warn_on_duplicate_tools=True) + self._resource_manager = ResourceManager(warn_on_duplicate_resources=True) + self._prompt_manager = PromptManager(warn_on_duplicate_prompts=True) + + def add_tool( + self, + fn: AnyFunction, + name: str | None = None, + description: str | None = None, + annotations: ToolAnnotations | None = None, + ) -> None: + """Add a tool to the bundler. + + The tool function can optionally request a Context object by adding a parameter + with the Context type annotation. See the @tool decorator for examples. + + Args: + fn: The function to register as a tool + name: Optional name for the tool (defaults to function name) + description: Optional description of what the tool does + annotations: Optional ToolAnnotations providing additional tool information + """ + self._tool_manager.add_tool( + fn, name=name, description=description, annotations=annotations + ) + + def tool( + self, + name: str | None = None, + description: str | None = None, + annotations: ToolAnnotations | None = None, + ) -> Callable[[AnyFunction], AnyFunction]: + """Decorator to register a tool. + + Tools 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 resource access. + + Args: + name: Optional name for the tool (defaults to function name) + description: Optional description of what the tool does + annotations: Optional ToolAnnotations providing additional tool information + + Example: + @server.tool() + def my_tool(x: int) -> str: + return str(x) + + @server.tool() + def tool_with_context(x: int, ctx: Context) -> str: + ctx.info(f"Processing {x}") + return str(x) + + @server.tool() + async def async_tool(x: int, context: Context) -> str: + await context.report_progress(50, 100) + return str(x) + """ + # Check if user passed function directly instead of calling decorator + if callable(name): + raise TypeError( + "The @tool decorator was used incorrectly. " + "Did you forget to call it? Use @tool() instead of @tool" + ) + + def decorator(fn: AnyFunction) -> AnyFunction: + self.add_tool( + fn, name=name, description=description, annotations=annotations + ) + return fn + + return decorator + + def get_tools(self) -> dict[str, Tool]: + return self._tool_manager.get_all_tools() + + def add_resource(self, resource: Resource) -> None: + """Add a resource to the server. + + Args: + resource: A Resource instance to add + """ + self._resource_manager.add_resource(resource) + + def resource( + self, + uri: str, + *, + name: str | None = None, + description: str | None = None, + mime_type: str | None = None, + ) -> Callable[[AnyFunction], AnyFunction]: + """Decorator to register a function as a resource. + + The function will be called when the resource is read to generate its content. + The function can return: + - str for text content + - bytes for binary content + - other types will be converted to JSON + + If the URI contains parameters (e.g. "resource://{param}") or the function + has parameters, it will be registered as a template resource. + + Args: + uri: URI for the resource (e.g. "resource://my-resource" or "resource://{param}") + name: Optional name for the resource + description: Optional description of the resource + mime_type: Optional MIME type for the resource + + Example: + @server.resource("resource://my-resource") + def get_data() -> str: + return "Hello, world!" + + @server.resource("resource://my-resource") + async get_data() -> str: + data = await fetch_data() + return f"Hello, world! {data}" + + @server.resource("resource://{city}/weather") + def get_weather(city: str) -> str: + return f"Weather for {city}" + + @server.resource("resource://{city}/weather") + async def get_weather(city: str) -> str: + data = await fetch_weather(city) + return f"Weather for {city}: {data}" + """ + # Check if user passed function directly instead of calling decorator + if callable(uri): + raise TypeError( + "The @resource decorator was used incorrectly. " + "Did you forget to call it? Use @resource('uri') instead of @resource" + ) + + def decorator(fn: AnyFunction) -> AnyFunction: + # Check if this should be a template + has_uri_params = "{" in uri and "}" in uri + has_func_params = bool(inspect.signature(fn).parameters) + + if has_uri_params or has_func_params: + # Validate that URI params match function params + uri_params = set(re.findall(r"{(\w+)}", uri)) + func_params = set(inspect.signature(fn).parameters.keys()) + + if uri_params != func_params: + raise ValueError( + f"Mismatch between URI parameters {uri_params} " + f"and function parameters {func_params}" + ) + + # Register as template + self._resource_manager.add_template( + fn=fn, + uri_template=uri, + name=name, + description=description, + mime_type=mime_type, + ) + else: + # Register as regular resource + resource = FunctionResource.from_function( + fn=fn, + uri=uri, + name=name, + description=description, + mime_type=mime_type, + ) + self.add_resource(resource) + return fn + + return decorator + + def get_resources(self) -> tuple[dict[str, Resource], dict[str, ResourceTemplate]]: + return self._resource_manager.get_all_resources() + + def add_prompt(self, prompt: Prompt) -> None: + """Add a prompt to the server. + + Args: + prompt: A Prompt instance to add + """ + self._prompt_manager.add_prompt(prompt) + + def prompt( + self, name: str | None = None, description: str | None = None + ) -> Callable[[AnyFunction], AnyFunction]: + """Decorator to register a prompt. + + Args: + name: Optional name for the prompt (defaults to function name) + description: Optional description of what the prompt does + + Example: + @server.prompt() + def analyze_table(table_name: str) -> list[Message]: + 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) + return [ + { + "role": "user", + "content": { + "type": "resource", + "resource": { + "uri": f"file://{path}", + "text": content + } + } + } + ] + """ + # Check if user passed function directly instead of calling decorator + if callable(name): + raise TypeError( + "The @prompt decorator was used incorrectly. " + "Did you forget to call it? Use @prompt() instead of @prompt" + ) + + def decorator(func: AnyFunction) -> AnyFunction: + prompt = Prompt.from_function(func, name=name, description=description) + self.add_prompt(prompt) + return func + + return decorator + + def get_prompts(self) -> dict[str, Prompt]: + return self._prompt_manager.get_all_prompts()