diff --git a/docs/docs.json b/docs/docs.json index 345050c8f..08e99918c 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -163,10 +163,10 @@ "servers/transforms/namespace", "servers/transforms/tool-transformation", "servers/visibility", + "servers/transforms/code-mode", "servers/transforms/tool-search", "servers/transforms/resources-as-tools", - "servers/transforms/prompts-as-tools", - "servers/transforms/code-mode" + "servers/transforms/prompts-as-tools" ] }, { diff --git a/docs/servers/transforms/code-mode.mdx b/docs/servers/transforms/code-mode.mdx index f15e36e1d..b542de94b 100644 --- a/docs/servers/transforms/code-mode.mdx +++ b/docs/servers/transforms/code-mode.mdx @@ -1,26 +1,32 @@ --- -title: Code Mode (Experimental) +title: Code Mode sidebarTitle: Code Mode description: Let LLMs write Python to orchestrate tools in a sandbox icon: flask -tag: EXPERIMENTAL +tag: NEW --- import { VersionBadge } from '/snippets/version-badge.mdx' + +CodeMode is experimental. The core interface is stable, but the specific discovery tools and their parameters may evolve as we learn more about what works best in practice. + + Standard MCP tool usage has two scaling problems. First, every tool in the catalog is loaded into the LLM's context upfront — with hundreds of tools, that's tens of thousands of tokens spent before the LLM even reads the user's request. Second, every tool call is a round-trip: the LLM calls a tool, the result passes back through the context window, the LLM reasons about it, calls another tool, and so on. Intermediate results that only exist to feed the next step still burn tokens flowing through the model. -`CodeMode` solves both problems by replacing the tool catalog with two meta-tools — `search` for discovering tools by keyword, and `execute` for running Python scripts that chain tool calls in a sandbox. The LLM discovers what it needs, writes a script, and gets back only the final answer. One round-trip instead of ten; tool definitions loaded only when needed instead of all at once. +CodeMode solves both problems. Instead of seeing your entire tool catalog, the LLM gets meta-tools for discovering what's available and for writing and executing code that calls the tools it needs. It discovers on demand, writes a script that chains tool calls in a sandbox, and gets back only the final answer. The approach was introduced by Cloudflare in [Code Mode](https://blog.cloudflare.com/code-mode/) and explored further by Anthropic in [Code Execution with MCP](https://www.anthropic.com/engineering/code-execution-with-mcp). - -CodeMode requires a sandbox to execute LLM-generated code safely. The default sandbox uses [pydantic-monty](https://github.com/pydantic/pydantic-monty), installed via `pip install "fastmcp[code-mode]"`. You can also provide your own sandbox — see [Custom Sandbox Providers](#custom-sandbox-providers). - +## Getting Started -## Basic Usage + +CodeMode requires the `code-mode` extra for sandbox support. Install it with `pip install "fastmcp[code-mode]"`. + + +You take a normal server with normally registered tools and add a `CodeMode` transform. The transform wraps your existing tools in the code mode machinery — your tool functions don't change at all: ```python from fastmcp import FastMCP @@ -39,157 +45,225 @@ def multiply(x: int, y: int) -> int: return x * y ``` -Clients now see only two tools. The LLM discovers the real tools through `search`, then orchestrates them through `execute`: +Clients connecting to this server no longer see `add` and `multiply` directly. Instead, they see the meta-tools that CodeMode provides — tools for discovering what's available and executing code against it. The original tools are still there, but they're accessed through the CodeMode layer. -```python -# Discover tools with a keyword search -result = await client.call_tool("search", {"query": "add multiply numbers"}) -# [{"name": "add", "inputSchema": {...}, ...}, {"name": "multiply", ...}] +## Discovery -# Chain tool calls in a single round-trip -result = await client.call_tool("execute", { - "code": """ -a = await call_tool("add", {"x": 3, "y": 4}) -b = await call_tool("multiply", {"x": a["result"], "y": 2}) -return b -""" -}) -# {"result": 14} -``` +Before the LLM can write code that calls your tools, it needs to know what tools exist and how to call them. This is the **discovery** process — the LLM uses meta-tools to learn about your tool catalog, then writes code against what it finds. -## Search +The fundamental tradeoff is **tokens vs. round-trips**. Each discovery step is an LLM round-trip: the model calls a tool, waits for the response, reasons about it, then decides what to do next. More steps mean less wasted context (each step is targeted) but more latency and API calls. Fewer steps mean the LLM gets information upfront but pays for detail it might not need. -The `search` meta-tool takes a `query` string and returns matching tools ranked by relevance, including their full `inputSchema` and `outputSchema`. The LLM uses these schemas to understand how to call each tool and what to expect back. +By default, CodeMode gives the LLM three tools — `search`, `get_schema`, and `execute` — creating a three-stage discovery flow: -Search uses BM25 ranking by default, matching against tool names and descriptions. You can swap in any [search transform](/servers/transforms/tool-search): + + +First, the LLM uses the `search` meta-tool to find tools by keyword. -```python -from fastmcp.server.transforms.search import RegexSearchTransform - -mcp = FastMCP( - "Server", - transforms=[CodeMode(search_transform=RegexSearchTransform())], -) -``` - -### Search Result Format - -By default, `search` returns tool definitions in the same JSON format as `list_tools` — full `inputSchema`, `outputSchema`, and metadata. This is complete, but verbose: each tool definition carries significant structural overhead that the LLM doesn't need just to decide which tool to call next. - -`serialize_tools_for_output_markdown` is a built-in alternative that renders the same information as compact markdown. In benchmarks across typical tool catalogs it reduces search result tokens by ~65-70%. - -```python -from fastmcp import FastMCP -from fastmcp.experimental.transforms import CodeMode -from fastmcp.server.transforms.search import serialize_tools_for_output_markdown - -mcp = FastMCP( - "Server", - transforms=[CodeMode(search_result_serializer=serialize_tools_for_output_markdown)], -) -``` - -For a tool like `create_document`, the output looks like: +For example, it might do `search(query="math numbers")` and receive the following response: ``` -### create_document +- add: Add two numbers. +- multiply: Multiply two numbers. +``` -Create a new document in the workspace with the given metadata. +This lets the LLM know which tools are available and what they do, significantly reducing the surface area it needs to consider. + + + +Next, the LLM calls `get_schema` to get parameter details for the tools it found in the previous step. + +For example, it might do `get_schema(tools=["add", "multiply"])` and receive the following response: + +``` +### add + +Add two numbers. **Parameters** -- `title` (string, required) -- `content` (string, required) -- `tags` (string[], required) -- `author` (string, required) -- `published` (boolean) -- `parent_id` (string?) +- `x` (integer, required) +- `y` (integer, required) + +### multiply + +Multiply two numbers. + +**Parameters** +- `x` (integer, required) +- `y` (integer, required) ``` -You can also supply any callable that takes a sequence of tools and returns a string or a list of dicts. Async callables are supported too: +Now the LLM knows the parameters for the tools it found, and can write code that chains the tool calls. If it needed more detail, it could have called `get_schema` with `detail="full"` to get the complete JSON schema. + + + +Finally, the LLM writes and executes code that chains the tool calls in a Python sandbox. Inside the sandbox, `call_tool(name, params)` is the only function available. The LLM uses this to compose tools into a workflow and return a final result. + +For example, it might write the following code and call the `execute` tool with it: ```python -from fastmcp.server.transforms.search import SearchResultSerializer - -def my_serializer(tools) -> str: - return "\n".join(f"{t.name}: {t.description}" for t in tools) - -mcp = FastMCP( - "Server", - transforms=[CodeMode(search_result_serializer=my_serializer)], -) +a = await call_tool("add", {"x": 3, "y": 4}) +b = await call_tool("multiply", {"x": a, "y": 2}) +return b ``` -The default JSON format is unchanged when `search_result_serializer` is not set — this is purely opt-in. +The result is returned to the LLM. + + -## Execute +This three-stage flow works well for most servers — each step pulls in only the information needed for the next one, keeping context usage minimal. But CodeMode's discovery surface is fully configurable. The sections below explain each built-in discovery tool and how to combine them into different patterns. -The `execute` meta-tool takes a `code` string containing async Python. Inside the sandbox, one function is available: +## Discovery Tools -```python -await call_tool(tool_name, params) # -> dict | str +CodeMode ships with three built-in discovery tools: `Search`, `GetSchemas`, and `GetTags`. By default, only `Search` and `GetSchemas` are enabled. Each tool supports a `default_detail` parameter that sets the default verbosity level, and the LLM can override the detail level on any individual call. + +### Detail Levels + +`Search` and `GetSchemas` share the same three detail levels, so the same `detail` value produces the same output format regardless of which tool the LLM calls: + +| Level | Output | Token cost | +|---|---|---| +| `"brief"` | Tool names and one-line descriptions | Cheapest — good for scanning | +| `"detailed"` | Compact markdown with parameter names, types, and required markers | Medium — often enough to write code | +| `"full"` | Complete JSON schema | Most expensive — everything | + +`Search` defaults to `"brief"` and `GetSchemas` defaults to `"detailed"`. + +### Search + +`Search` finds tools by natural-language query using BM25 ranking. At its default `"brief"` detail, results include just tool names and descriptions — enough to decide which tools are worth inspecting further. The LLM can request `"detailed"` to get parameter schemas inline, or `"full"` for the complete JSON. + +If your tools use [tags](/servers/tools#tags), Search also accepts a `tags` parameter so the LLM can narrow results to specific categories before searching. + +### GetSchemas + +`GetSchemas` returns parameter details for specific tools by name. At its default `"detailed"` level, it renders compact markdown with parameter names, types, and required markers. At `"full"`, it returns the complete JSON schema — useful when tools have deeply nested parameters that the compact format doesn't capture. + +### GetTags + +`GetTags` lets the LLM browse tools by category using [tag](/servers/tools#tags) metadata. At brief detail, the LLM sees tag names with counts. At full detail, it sees tools listed under each tag: + +``` +- math (3 tools) +- text (2 tools) +- untagged (1 tool) ``` -The return type depends on whether the tool declares an output schema. When it does, `call_tool` returns the structured content dict exactly as the schema describes — for example, `{"result": 42}` for a tool returning `int`. When there's no output schema, `call_tool` returns the text content as a string. The LLM can tell which to expect from the `outputSchema` field in search results. +`GetTags` isn't included in the defaults — add it when browsing by category would help the LLM orient itself in a large catalog. The LLM can browse tags first, then pass specific tags into Search to narrow results. -Use `return` to produce the final output from the script. +## Discovery Patterns -### Default Arguments +The right discovery configuration depends on your server — how many tools you have and how complex their parameters are. It may be tempting to minimize round-trips by collapsing everything into fewer steps, but for the complex servers that benefit most from CodeMode, our experience is that staged discovery leads to better results. Flooding the LLM with detailed schemas for tools it doesn't end up using can hurt more than the extra round-trip costs. Each pattern below is a complete, copyable configuration. -When tools share common parameters (like workspace IDs or API keys), `default_arguments` injects them automatically: +### Three-Stage + +The default. The LLM searches for candidates, inspects schemas for the ones it wants, then writes code. Best for **large or complex tool sets** where you want to minimize context usage — the LLM only pays for schemas it actually needs. ```python -mcp = FastMCP( - "Server", - transforms=[CodeMode(default_arguments={"workspace_id": "ws-123"})], -) -``` - -Defaults are only injected when the tool actually accepts the parameter and the LLM hasn't provided it explicitly. Parameters that a tool doesn't accept are silently skipped. - -## OpenAPI Integration - -`CodeMode` pairs naturally with OpenAPI-backed providers, where a single API spec can expose hundreds of endpoints as tools: - -```python -import httpx - from fastmcp import FastMCP from fastmcp.experimental.transforms import CodeMode -from fastmcp.server.providers.openapi import OpenAPIProvider -openapi_spec = httpx.get("https://api.example.com/openapi.json").json() -api_client = httpx.AsyncClient(base_url="https://api.example.com") - -provider = OpenAPIProvider( - openapi_spec=openapi_spec, - client=api_client, -) - -mcp = FastMCP("API Code Mode", providers=[provider], transforms=[CodeMode()]) +mcp = FastMCP("Server", transforms=[CodeMode()]) ``` -## Configuration - -### Custom Tool Names - -The default `search` and `execute` names can be changed: +If your tools use [tags](/servers/tools#tags), add `GetTags` so the LLM can browse by category before searching — giving it four stages of progressive disclosure: ```python -mcp = FastMCP( - "Server", - transforms=[ - CodeMode( - search_tool_name="find_tools", - execute_tool_name="run_workflow", - execute_description="Run multi-step API workflows", - ) - ], +from fastmcp import FastMCP +from fastmcp.experimental.transforms import CodeMode +from fastmcp.experimental.transforms.code_mode import GetTags, Search, GetSchemas + +code_mode = CodeMode( + discovery_tools=[GetTags(), Search(), GetSchemas()], ) + +mcp = FastMCP("Server", transforms=[code_mode]) ``` +### Two-Stage + +Search returns parameter schemas inline, so the LLM can go straight from search to execute. Best for **smaller catalogs** where the extra tokens per search result are a reasonable price for one fewer round-trip. + +```python +from fastmcp import FastMCP +from fastmcp.experimental.transforms import CodeMode +from fastmcp.experimental.transforms.code_mode import Search, GetSchemas + +code_mode = CodeMode( + discovery_tools=[Search(default_detail="detailed"), GetSchemas()], +) + +mcp = FastMCP("Server", transforms=[code_mode]) +``` + +`GetSchemas` is still available as a fallback — the LLM can call it with `detail="full"` if it encounters a tool with complex nested parameters where the compact markdown isn't enough. + +### Single-Stage + +Skip discovery entirely and bake tool instructions into the execute tool's description. Best for **very simple servers** where the LLM already knows what tools are available — maybe there are only a few, or they're described in the system prompt. + +```python +from fastmcp import FastMCP +from fastmcp.experimental.transforms import CodeMode + +code_mode = CodeMode( + discovery_tools=[], + execute_description=( + "Available tools:\n" + "- add(x: int, y: int) -> int: Add two numbers\n" + "- multiply(x: int, y: int) -> int: Multiply two numbers\n\n" + "Write Python using `await call_tool(name, params)` and `return` the result." + ), +) + +mcp = FastMCP("Server", transforms=[code_mode]) +``` + +## Custom Discovery Tools + +Discovery tools are composable — you can mix the built-ins with your own. Each discovery tool is a callable that receives catalog access and returns a `Tool`. The catalog accessor is a function (not the catalog itself) because the catalog is request-scoped — different users may see different tools based on auth. + +Here's a minimal example: + +```python +from fastmcp.experimental.transforms import CodeMode +from fastmcp.experimental.transforms.code_mode import GetToolCatalog, GetSchemas +from fastmcp.server.context import Context +from fastmcp.tools.tool import Tool + +def list_all_tools(get_catalog: GetToolCatalog) -> Tool: + async def list_tools(ctx: Context) -> str: + """List all available tool names.""" + tools = await get_catalog(ctx) + return ", ".join(t.name for t in tools) + + return Tool.from_function(fn=list_tools, name="list_tools") + +code_mode = CodeMode(discovery_tools=[list_all_tools, GetSchemas()]) +``` + +The LLM sees the docstring of each discovery tool's inner function as its description — that's how it learns what each tool does and when to use it. Write docstrings that explain what the tool returns and when the LLM should call it. + +Discovery tools and the execute tool can also have custom names: + +```python +from fastmcp.experimental.transforms.code_mode import Search, GetSchemas + +code_mode = CodeMode( + discovery_tools=[ + Search(name="find_tools"), + GetSchemas(name="describe"), + ], + execute_tool_name="run_workflow", +) + +mcp = FastMCP("Server", transforms=[code_mode]) +``` + +## Sandbox Configuration + ### Resource Limits -The default `MontySandboxProvider` can enforce execution limits on sandboxed code — timeouts, memory caps, recursion depth, and more. Without limits, LLM-generated scripts can run indefinitely. +The default `MontySandboxProvider` can enforce execution limits — timeouts, memory caps, recursion depth, and more. Without limits, LLM-generated scripts can run indefinitely. ```python from fastmcp.experimental.transforms import CodeMode, MontySandboxProvider @@ -213,7 +287,7 @@ All keys are optional — omit any to leave that dimension uncapped: ### Custom Sandbox Providers -The default `MontySandboxProvider` uses [pydantic-monty](https://github.com/pydantic/pydantic-monty) for sandboxed execution. You can replace it with any object implementing the `SandboxProvider` protocol: +You can replace the default sandbox with any object implementing the `SandboxProvider` protocol: ```python from collections.abc import Callable @@ -232,7 +306,10 @@ class RemoteSandboxProvider: # Send code to your remote sandbox runtime ... -mcp = FastMCP("Server", transforms=[CodeMode(sandbox_provider=RemoteSandboxProvider())]) +mcp = FastMCP( + "Server", + transforms=[CodeMode(sandbox_provider=RemoteSandboxProvider())], +) ``` The `external_functions` dict contains async callables injected into the sandbox scope — `execute` uses this to provide `call_tool`. diff --git a/src/fastmcp/experimental/transforms/code_mode.py b/src/fastmcp/experimental/transforms/code_mode.py index 8aaebee52..8d204a8de 100644 --- a/src/fastmcp/experimental/transforms/code_mode.py +++ b/src/fastmcp/experimental/transforms/code_mode.py @@ -1,8 +1,8 @@ import asyncio import importlib -import logging -from collections.abc import Callable, Sequence -from typing import Annotated, Any, Protocol +import json +from collections.abc import Awaitable, Callable, Sequence +from typing import Annotated, Any, Literal, Protocol from mcp.types import TextContent from pydantic import Field @@ -12,16 +12,29 @@ from fastmcp.server.context import Context from fastmcp.server.transforms import GetToolNext from fastmcp.server.transforms.catalog import CatalogTransform from fastmcp.server.transforms.search.base import ( - BaseSearchTransform, - SearchResultSerializer, - _invoke_serializer, serialize_tools_for_output_json, serialize_tools_for_output_markdown, ) from fastmcp.tools.tool import Tool, ToolResult from fastmcp.utilities.versions import VersionSpec -logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Type aliases +# --------------------------------------------------------------------------- + +GetToolCatalog = Callable[[Context], Awaitable[Sequence[Tool]]] +"""Async callable that returns the auth-filtered tool catalog.""" + +SearchFn = Callable[[Sequence[Tool], str], Awaitable[Sequence[Tool]]] +"""Async callable that searches a tool sequence by query string.""" + +DiscoveryToolFactory = Callable[[GetToolCatalog], Tool] +"""Factory that receives catalog access and returns a synthetic Tool.""" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- def _ensure_async(fn: Callable[..., Any]) -> Callable[..., Any]: @@ -52,6 +65,11 @@ def _unwrap_tool_result(result: ToolResult) -> dict[str, Any] | str: return "\n".join(parts) +# --------------------------------------------------------------------------- +# Sandbox providers +# --------------------------------------------------------------------------- + + class SandboxProvider(Protocol): """Interface for executing LLM-generated Python code in a sandbox. @@ -122,45 +140,297 @@ class MontySandboxProvider: return await pydantic_monty.run_monty_async(monty, **run_kwargs) -class CodeMode(CatalogTransform): - """Transform that collapses all tools into `search` + `execute` meta-tools.""" +# --------------------------------------------------------------------------- +# Built-in discovery tools +# --------------------------------------------------------------------------- + + +ToolDetailLevel = Literal["brief", "detailed", "full"] +"""Detail level for discovery tool output. + +- ``"brief"``: tool names and one-line descriptions +- ``"detailed"``: compact markdown with parameter names, types, and required markers +- ``"full"``: complete JSON schema +""" + + +def _render_tools(tools: Sequence[Tool], detail: ToolDetailLevel) -> str: + """Render tools at the requested detail level. + + The same detail value produces the same output format regardless of + which discovery tool calls this, so ``detail="detailed"`` on Search + gives identical formatting to ``detail="detailed"`` on GetSchemas. + """ + if not tools: + if detail == "full": + return json.dumps([], indent=2) + return "No tools matched the query." + if detail == "full": + return json.dumps(serialize_tools_for_output_json(tools), indent=2) + if detail == "detailed": + return serialize_tools_for_output_markdown(tools) + # brief + lines: list[str] = [] + for tool in tools: + desc = f": {tool.description}" if tool.description else "" + lines.append(f"- {tool.name}{desc}") + return "\n".join(lines) + + +class Search: + """Discovery tool factory that searches the catalog by query. + + Args: + search_fn: Async callable ``(tools, query) -> matching_tools``. + Defaults to BM25 ranking. + name: Name of the synthetic tool exposed to the LLM. + default_detail: Default detail level for search results. + ``"brief"`` returns tool names and descriptions only. + ``"detailed"`` returns compact markdown with parameter schemas. + ``"full"`` returns complete JSON tool definitions. + """ + + def __init__( + self, + *, + search_fn: SearchFn | None = None, + name: str = "search", + default_detail: ToolDetailLevel | None = None, + ) -> None: + if search_fn is None: + from fastmcp.server.transforms.search.bm25 import BM25SearchTransform + + _bm25 = BM25SearchTransform() + search_fn = _bm25._search + self._search_fn = search_fn + self._name = name + self._default_detail: ToolDetailLevel = default_detail or "brief" + + def __call__(self, get_catalog: GetToolCatalog) -> Tool: + search_fn = self._search_fn + default_detail = self._default_detail + + async def search( + query: Annotated[str, "Search query to find available tools"], + tags: Annotated[ + list[str] | None, + "Filter to tools with any of these tags before searching", + ] = None, + detail: Annotated[ + ToolDetailLevel, + "'brief' for names and descriptions, 'detailed' for parameter schemas as markdown, 'full' for complete JSON schemas", + ] = default_detail, + ctx: Context = None, # type: ignore[assignment] + ) -> str: + """Search for available tools by query. + + Returns matching tools ranked by relevance. + """ + tools = await get_catalog(ctx) + if tags: + tag_set = set(tags) + has_untagged = "untagged" in tag_set + real_tags = tag_set - {"untagged"} + tools = [ + t + for t in tools + if (t.tags & real_tags) or (has_untagged and not t.tags) + ] + results = await search_fn(tools, query) + return _render_tools(results, detail) + + return Tool.from_function(fn=search, name=self._name) + + +class GetSchemas: + """Discovery tool factory that returns schemas for tools by name. + + Args: + name: Name of the synthetic tool exposed to the LLM. + default_detail: Default detail level for schema results. + ``"brief"`` returns tool names and descriptions only. + ``"detailed"`` renders compact markdown with parameter names, + types, and required markers. + ``"full"`` returns the complete JSON schema. + """ + + def __init__( + self, + *, + name: str = "get_schema", + default_detail: ToolDetailLevel | None = None, + ) -> None: + self._name = name + self._default_detail: ToolDetailLevel = default_detail or "detailed" + + def __call__(self, get_catalog: GetToolCatalog) -> Tool: + default_detail = self._default_detail + + async def get_schema( + tools: Annotated[ + list[str], + "List of tool names to get schemas for", + ], + detail: Annotated[ + ToolDetailLevel, + "'brief' for names and descriptions, 'detailed' for parameter schemas as markdown, 'full' for complete JSON schemas", + ] = default_detail, + ctx: Context = None, # type: ignore[assignment] + ) -> str: + """Get parameter schemas for specific tools. + + Use after searching to get the detail needed to call a tool. + """ + catalog = await get_catalog(ctx) + catalog_by_name = {t.name: t for t in catalog} + matched = [catalog_by_name[n] for n in tools if n in catalog_by_name] + not_found = [n for n in tools if n not in catalog_by_name] + + if not matched and not_found: + return f"Tools not found: {', '.join(not_found)}" + + if detail == "full": + data = serialize_tools_for_output_json(matched) + if not_found: + data.append({"not_found": not_found}) + return json.dumps(data, indent=2) + + result = _render_tools(matched, detail) + if not_found: + result += f"\n\nTools not found: {', '.join(not_found)}" + return result + + return Tool.from_function(fn=get_schema, name=self._name) + + +class GetTags: + """Discovery tool factory that lists tool tags from the catalog. + + Reads ``tool.tags`` from the catalog and groups tools by tag. Tools + without tags appear under ``"untagged"``. + + Args: + name: Name of the synthetic tool exposed to the LLM. + default_detail: Default detail level. + ``"brief"`` returns tag names with tool counts. + ``"full"`` lists all tools under each tag. + """ + + def __init__( + self, + *, + name: str = "tags", + default_detail: Literal["brief", "full"] | None = None, + ) -> None: + self._name = name + self._default_detail: Literal["brief", "full"] = default_detail or "brief" + + def __call__(self, get_catalog: GetToolCatalog) -> Tool: + default_detail = self._default_detail + + async def tags( + detail: Annotated[ + Literal["brief", "full"], + "Level of detail: 'brief' for tag names and counts, 'full' for tools listed under each tag", + ] = default_detail, + ctx: Context = None, # type: ignore[assignment] + ) -> str: + """List available tool tags. + + Use to browse available tools by tag before searching. + """ + catalog = await get_catalog(ctx) + by_tag: dict[str, list[Tool]] = {} + for tool in catalog: + if tool.tags: + for tag in tool.tags: + by_tag.setdefault(tag, []).append(tool) + else: + by_tag.setdefault("untagged", []).append(tool) + + if not by_tag: + return "No tools available." + + if detail == "brief": + lines = [ + f"- {tag} ({len(tools)} tool{'s' if len(tools) != 1 else ''})" + for tag, tools in sorted(by_tag.items()) + ] + return "\n".join(lines) + + blocks: list[str] = [] + for tag, tools in sorted(by_tag.items()): + lines = [f"### {tag}"] + for tool in tools: + desc = f": {tool.description}" if tool.description else "" + lines.append(f"- {tool.name}{desc}") + blocks.append("\n".join(lines)) + return "\n\n".join(blocks) + + return Tool.from_function(fn=tags, name=self._name) + + +# --------------------------------------------------------------------------- +# CodeMode +# --------------------------------------------------------------------------- + + +def _default_discovery_tools() -> list[DiscoveryToolFactory]: + return [Search(), GetSchemas()] + + +class CodeMode(CatalogTransform): + """Transform that collapses all tools into discovery + execute meta-tools. + + Discovery tools are composable via the ``discovery_tools`` parameter. + Each is a callable that receives catalog access and returns a ``Tool``. + By default, ``Search`` and ``GetSchemas`` are included for + progressive disclosure: search finds candidates, get_schema retrieves + parameter details, and execute runs code. + + The ``execute`` tool is always present and provides a sandboxed Python + environment with ``call_tool(name, params)`` in scope. + """ def __init__( self, *, - default_arguments: dict[str, Any] | None = None, sandbox_provider: SandboxProvider | None = None, - search_transform: BaseSearchTransform | None = None, - search_tool_name: str = "search", + discovery_tools: list[DiscoveryToolFactory] | None = None, execute_tool_name: str = "execute", execute_description: str | None = None, - search_result_serializer: SearchResultSerializer | None = None, ) -> None: - if search_tool_name == execute_tool_name: - raise ValueError( - "search_tool_name and execute_tool_name must be different." - ) - super().__init__() - self._default_arguments = default_arguments or {} - self.search_tool_name = search_tool_name self.execute_tool_name = execute_tool_name self.execute_description = execute_description self.sandbox_provider = sandbox_provider or MontySandboxProvider() - self._cached_search_tool: Tool | None = None - self._cached_execute_tool: Tool | None = None - self._search_result_serializer: SearchResultSerializer | None = ( - search_result_serializer + + self._discovery_factories = ( + discovery_tools + if discovery_tools is not None + else _default_discovery_tools() ) + self._built_discovery_tools: list[Tool] | None = None + self._cached_execute_tool: Tool | None = None - if search_transform is None: - from fastmcp.server.transforms.search.bm25 import BM25SearchTransform - - search_transform = BM25SearchTransform() - self._search_transform = search_transform + def _build_discovery_tools(self) -> list[Tool]: + if self._built_discovery_tools is None: + tools = [ + factory(self.get_tool_catalog) for factory in self._discovery_factories + ] + names = {t.name for t in tools} + if self.execute_tool_name in names: + raise ValueError( + f"Discovery tool name '{self.execute_tool_name}' " + f"collides with execute_tool_name." + ) + if len(names) != len(tools): + raise ValueError("Discovery tools must have unique names.") + self._built_discovery_tools = tools + return self._built_discovery_tools async def transform_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]: - return [self._get_search_tool(), self._get_execute_tool()] + return [*self._build_discovery_tools(), self._get_execute_tool()] async def get_tool( self, @@ -169,8 +439,9 @@ class CodeMode(CatalogTransform): *, version: VersionSpec | None = None, ) -> Tool | None: - if name == self.search_tool_name: - return self._get_search_tool() + for tool in self._build_discovery_tools(): + if tool.name == name: + return tool if name == self.execute_tool_name: return self._get_execute_tool() return await call_next(name, version=version) @@ -193,41 +464,11 @@ class CodeMode(CatalogTransform): return tool return None - def _get_search_tool(self) -> Tool: - if self._cached_search_tool is None: - self._cached_search_tool = self._make_search_tool() - return self._cached_search_tool - def _get_execute_tool(self) -> Tool: if self._cached_execute_tool is None: self._cached_execute_tool = self._make_execute_tool() return self._cached_execute_tool - def _make_search_tool(self) -> Tool: - transform = self - - async def search( - query: Annotated[ - str, - "Search query to find available tools", - ], - ctx: Context = None, # type: ignore[assignment] - ) -> str | list[dict[str, Any]]: - """Search for available tools by query. - - Returns matching tool definitions ranked by relevance, - in the same format as list_tools. - """ - tools = await transform.get_tool_catalog(ctx) - results = await transform._search_transform._search(tools, query) - if transform._search_result_serializer is not None: - return await _invoke_serializer( - transform._search_result_serializer, results - ) - return await transform._search_transform._render_results(results) - - return Tool.from_function(fn=search, name=self.search_tool_name) - def _make_execute_tool(self) -> Tool: transform = self @@ -243,9 +484,6 @@ class CodeMode(CatalogTransform): ctx: Context = None, # type: ignore[assignment] ) -> Any: """Execute tool calls using Python code.""" - defaults = transform._default_arguments - # Cache the tool catalog for the duration of this execute block - # so multiple call_tool() invocations don't each trigger list_tools(). cached_tools: Sequence[Tool] | None = None async def _get_cached_tools() -> Sequence[Tool]: @@ -260,26 +498,7 @@ class CodeMode(CatalogTransform): if tool is None: raise NotFoundError(f"Unknown tool: {tool_name}") - accepted_args = set(tool.parameters.get("properties", {}).keys()) - skipped = { - key - for key in defaults - if key not in params and key not in accepted_args - } - if skipped: - logger.debug( - "default_arguments keys %s not accepted by tool %r, skipping", - skipped, - tool.name, - ) - merged = { - key: value - for key, value in defaults.items() - if key not in params and key in accepted_args - } - merged.update(params) - - result = await ctx.fastmcp.call_tool(tool.name, merged) + result = await ctx.fastmcp.call_tool(tool.name, params) return _unwrap_tool_result(result) return await transform.sandbox_provider.run( @@ -296,9 +515,10 @@ class CodeMode(CatalogTransform): __all__ = [ "CodeMode", + "GetSchemas", + "GetTags", + "GetToolCatalog", "MontySandboxProvider", "SandboxProvider", - "SearchResultSerializer", - "serialize_tools_for_output_json", - "serialize_tools_for_output_markdown", + "Search", ] diff --git a/tests/experimental/transforms/test_code_mode.py b/tests/experimental/transforms/test_code_mode.py index 24d909c0c..b9b6b9154 100644 --- a/tests/experimental/transforms/test_code_mode.py +++ b/tests/experimental/transforms/test_code_mode.py @@ -8,13 +8,15 @@ from mcp.types import ImageContent, TextContent from fastmcp import FastMCP from fastmcp.exceptions import ToolError from fastmcp.experimental.transforms import CodeMode, MontySandboxProvider -from fastmcp.experimental.transforms.code_mode import _ensure_async -from fastmcp.server.transforms.search.base import ( - _schema_section, - _schema_type, - serialize_tools_for_output_markdown, +from fastmcp.experimental.transforms.code_mode import ( + GetSchemas, + GetTags, + GetToolCatalog, + Search, + _ensure_async, ) -from fastmcp.tools.tool import ToolResult +from fastmcp.server.context import Context +from fastmcp.tools.tool import Tool, ToolResult def _unwrap_result(result: ToolResult) -> Any: @@ -43,18 +45,17 @@ def _unwrap_result(result: ToolResult) -> Any: return values -def _unwrap_search_results(result: ToolResult) -> list[dict[str, Any]]: - """Extract the list of tool dicts from a search ToolResult. +def _unwrap_string_result(result: ToolResult) -> str: + """Extract a string result from a ToolResult. - The search tool returns ``list[dict]`` which gets wrapped in - ``{"result": [...]}`` by the structured-output convention. + String results are wrapped in ``{"result": "..."}`` by the + structured-output convention. """ data = _unwrap_result(result) if isinstance(data, dict) and "result" in data: return data["result"] - if isinstance(data, list): - return data - raise AssertionError(f"Unexpected search result shape: {data!r}") + assert isinstance(data, str) + return data class _UnsafeTestSandboxProvider: @@ -91,50 +92,14 @@ async def _run_tool( return await server.call_tool(name, arguments) -async def test_code_mode_transform_hides_backend_tools_and_supports_defaults() -> None: - mcp = FastMCP("CodeMode Test") - - @mcp.tool - def add(x: int, y: int, workspace_id: str) -> str: - """Add two numbers with workspace context.""" - return f"{workspace_id}:{x + y}" - - @mcp.tool - def status() -> str: - """Get current status.""" - return "ok" - - mcp.add_transform( - CodeMode( - default_arguments={"workspace_id": "ws-default"}, - sandbox_provider=_UnsafeTestSandboxProvider(), - ) - ) - - listed_tools = await mcp.list_tools(run_middleware=False) - assert {tool.name for tool in listed_tools} == {"search", "execute"} - - search_result = await _run_tool(mcp, "search", {"query": "add numbers"}) - names = [t["name"] for t in _unwrap_search_results(search_result)] - assert "add" in names - - execute_result = await _run_tool( - mcp, - "execute", - {"code": "return await call_tool('add', {'x': 2, 'y': 3})"}, - ) - assert _unwrap_result(execute_result) == {"result": "ws-default:5"} - - status_result = await _run_tool( - mcp, - "execute", - {"code": "return await call_tool('status', {})"}, - ) - assert _unwrap_result(status_result) == {"result": "ok"} +# --------------------------------------------------------------------------- +# CodeMode core tests +# --------------------------------------------------------------------------- -async def test_code_mode_transform_replaces_listed_tools() -> None: - mcp = FastMCP("CodeMode Transform") +async def test_code_mode_default_tools() -> None: + """Default CodeMode exposes search, get_schema, and execute.""" + mcp = FastMCP("CodeMode Default") @mcp.tool def ping() -> str: @@ -143,54 +108,11 @@ async def test_code_mode_transform_replaces_listed_tools() -> None: mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) listed_tools = await mcp.list_tools(run_middleware=False) - assert {tool.name for tool in listed_tools} == {"search", "execute"} + assert {tool.name for tool in listed_tools} == {"search", "get_schema", "execute"} -async def test_code_mode_tool_descriptions_are_configurable() -> None: - mcp = FastMCP("CodeMode Descriptions") - - @mcp.tool - def ping() -> str: - return "pong" - - mcp.add_transform( - CodeMode( - sandbox_provider=_UnsafeTestSandboxProvider(), - search_tool_name="search_meta", - execute_tool_name="execute_meta", - execute_description="Custom execute description", - ) - ) - - listed_tools = await mcp.list_tools(run_middleware=False) - by_name = {tool.name: tool for tool in listed_tools} - - assert by_name["execute_meta"].description == "Custom execute description" - - -async def test_code_mode_default_execute_description() -> None: - mcp = FastMCP("CodeMode Defaults") - - @mcp.tool - def ping() -> str: - return "pong" - - mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) - - listed_tools = await mcp.list_tools(run_middleware=False) - by_name = {tool.name: tool for tool in listed_tools} - - execute_description = by_name["execute"].description or "" - - assert "single block" in execute_description - assert "Use `return` to produce output." in execute_description - assert ( - "Only `call_tool(tool_name: str, params: dict) -> Any` is available in scope." - in execute_description - ) - - -async def test_code_mode_search_returns_matching_tools() -> None: +async def test_code_mode_search_returns_lightweight_results() -> None: + """Default search returns tool names and descriptions, not full schemas.""" mcp = FastMCP("CodeMode Search") @mcp.tool @@ -206,13 +128,16 @@ async def test_code_mode_search_returns_matching_tools() -> None: mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) result = await _run_tool(mcp, "search", {"query": "square number"}) - tools = _unwrap_search_results(result) - assert len(tools) > 0 - assert tools[0]["name"] == "square" + text = _unwrap_string_result(result) + assert "square" in text + assert "Compute the square" in text + # Should NOT contain full schema details + assert "inputSchema" not in text -async def test_code_mode_search_results_include_schema() -> None: - mcp = FastMCP("CodeMode Output Schema") +async def test_code_mode_get_schema_brief() -> None: + """get_schema with detail=brief returns names and descriptions only.""" + mcp = FastMCP("CodeMode Schema Brief") @mcp.tool def square(x: int) -> int: @@ -221,11 +146,569 @@ async def test_code_mode_search_results_include_schema() -> None: mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + result = await _run_tool( + mcp, "get_schema", {"tools": ["square"], "detail": "brief"} + ) + text = _unwrap_string_result(result) + assert "square" in text + assert "Compute the square" in text + # brief should NOT include parameter details + assert "**Parameters**" not in text + + +async def test_code_mode_get_schema_detailed() -> None: + """get_schema with detail=detailed returns markdown with parameter info.""" + mcp = FastMCP("CodeMode Schema Detailed") + + @mcp.tool + def square(x: int) -> int: + """Compute the square of a number.""" + return x * x + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool( + mcp, "get_schema", {"tools": ["square"], "detail": "detailed"} + ) + text = _unwrap_string_result(result) + assert "### square" in text + assert "Compute the square" in text + assert "**Parameters**" in text + assert "`x` (integer, required)" in text + + +async def test_code_mode_get_schema_full() -> None: + """get_schema with detail=full returns JSON schema.""" + mcp = FastMCP("CodeMode Schema Full") + + @mcp.tool + def square(x: int) -> int: + """Compute the square of a number.""" + return x * x + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "get_schema", {"tools": ["square"], "detail": "full"}) + text = _unwrap_string_result(result) + parsed = json.loads(text) + assert isinstance(parsed, list) + assert parsed[0]["name"] == "square" + assert "inputSchema" in parsed[0] + + +async def test_code_mode_get_schema_default_is_detailed() -> None: + """get_schema defaults to detailed (markdown with parameters).""" + mcp = FastMCP("CodeMode Schema Default") + + @mcp.tool + def square(x: int) -> int: + """Compute the square of a number.""" + return x * x + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "get_schema", {"tools": ["square"]}) + text = _unwrap_string_result(result) + assert "### square" in text + assert "**Parameters**" in text + + +async def test_code_mode_get_schema_not_found() -> None: + """get_schema reports tools that don't exist in the catalog.""" + mcp = FastMCP("CodeMode Schema NotFound") + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "get_schema", {"tools": ["nonexistent"]}) + text = _unwrap_string_result(result) + assert "not found" in text.lower() + assert "nonexistent" in text + + +async def test_code_mode_get_schema_partial_match() -> None: + """get_schema returns schemas for found tools and reports missing ones.""" + mcp = FastMCP("CodeMode Schema Partial") + + @mcp.tool + def square(x: int) -> int: + """Compute the square.""" + return x * x + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "get_schema", {"tools": ["square", "nonexistent"]}) + text = _unwrap_string_result(result) + assert "### square" in text + assert "nonexistent" in text + + +async def test_code_mode_execute_works() -> None: + """Execute tool can call backend tools through the sandbox.""" + mcp = FastMCP("CodeMode Execute") + + @mcp.tool + def add(x: int, y: int) -> int: + return x + y + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool( + mcp, "execute", {"code": "return await call_tool('add', {'x': 2, 'y': 3})"} + ) + assert _unwrap_result(result) == {"result": 5} + + +# --------------------------------------------------------------------------- +# Tool naming and configuration +# --------------------------------------------------------------------------- + + +async def test_code_mode_custom_execute_name() -> None: + mcp = FastMCP("CodeMode Custom Execute") + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform( + CodeMode( + sandbox_provider=_UnsafeTestSandboxProvider(), + execute_tool_name="run_code", + ) + ) + + listed = await mcp.list_tools(run_middleware=False) + names = {t.name for t in listed} + assert "run_code" in names + assert "execute" not in names + + +async def test_code_mode_custom_execute_description() -> None: + mcp = FastMCP("CodeMode Custom Desc") + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform( + CodeMode( + sandbox_provider=_UnsafeTestSandboxProvider(), + execute_description="Custom execute description", + ) + ) + + listed = await mcp.list_tools(run_middleware=False) + by_name = {t.name: t for t in listed} + assert by_name["execute"].description == "Custom execute description" + + +async def test_code_mode_default_execute_description() -> None: + mcp = FastMCP("CodeMode Defaults") + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + listed = await mcp.list_tools(run_middleware=False) + by_name = {t.name: t for t in listed} + desc = by_name["execute"].description or "" + + assert "single block" in desc + assert "Use `return` to produce output." in desc + assert ( + "Only `call_tool(tool_name: str, params: dict) -> Any` is available in scope." + in desc + ) + + +# --------------------------------------------------------------------------- +# Discovery tool customization +# --------------------------------------------------------------------------- + + +async def test_code_mode_no_discovery_tools() -> None: + """CodeMode with empty discovery_tools exposes only execute.""" + mcp = FastMCP("CodeMode No Discovery") + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform( + CodeMode( + discovery_tools=[], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + listed = await mcp.list_tools(run_middleware=False) + assert {t.name for t in listed} == {"execute"} + + +async def test_code_mode_custom_discovery_tool_function() -> None: + """A plain function can serve as a discovery tool factory.""" + mcp = FastMCP("CodeMode Custom Discovery") + + @mcp.tool + def square(x: int) -> int: + """Compute the square.""" + return x * x + + def list_all(get_catalog: GetToolCatalog) -> Tool: + async def list_tools( + ctx: Context = None, # type: ignore[assignment] + ) -> str: + """List all available tools.""" + tools = await get_catalog(ctx) + return ", ".join(t.name for t in tools) + + return Tool.from_function(fn=list_tools, name="list_all") + + mcp.add_transform( + CodeMode( + discovery_tools=[list_all], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + listed = await mcp.list_tools(run_middleware=False) + assert {t.name for t in listed} == {"list_all", "execute"} + + result = await _run_tool(mcp, "list_all", {}) + text = _unwrap_string_result(result) + assert "square" in text + + +async def test_code_mode_search_detailed() -> None: + """Search with detail='detailed' returns markdown with parameter info.""" + mcp = FastMCP("CodeMode Search Detailed") + + @mcp.tool + def square(x: int) -> int: + """Compute the square.""" + return x * x + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "search", {"query": "square", "detail": "detailed"}) + text = _unwrap_string_result(result) + assert "### square" in text + assert "Compute the square" in text + assert "**Parameters**" in text + assert "`x` (integer, required)" in text + + +async def test_code_mode_search_tool_full_detail() -> None: + """Search with detail='full' includes JSON schemas.""" + mcp = FastMCP("CodeMode Search Full") + + @mcp.tool + def square(x: int) -> int: + """Compute the square.""" + return x * x + + mcp.add_transform( + CodeMode( + discovery_tools=[Search(default_detail="full")], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + result = await _run_tool(mcp, "search", {"query": "square"}) - tools = _unwrap_search_results(result) - assert len(tools) > 0 - tool_dict = tools[0] - assert "inputSchema" in tool_dict + text = _unwrap_string_result(result) + parsed = json.loads(text) + assert isinstance(parsed, list) + assert parsed[0]["name"] == "square" + assert "inputSchema" in parsed[0] + + +async def test_code_mode_custom_search_tool_name() -> None: + """Search and GetSchemas support custom names.""" + mcp = FastMCP("CodeMode Custom Names") + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform( + CodeMode( + discovery_tools=[ + Search(name="find"), + GetSchemas(name="describe"), + ], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + listed = await mcp.list_tools(run_middleware=False) + assert {t.name for t in listed} == {"find", "describe", "execute"} + + +def test_code_mode_rejects_discovery_execute_name_collision() -> None: + """CodeMode raises ValueError when a discovery tool collides with execute.""" + cm = CodeMode( + discovery_tools=[Search(name="execute")], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + with pytest.raises(ValueError, match="collides"): + cm._build_discovery_tools() + + +def test_code_mode_rejects_duplicate_discovery_names() -> None: + """CodeMode raises ValueError when discovery tools have duplicate names.""" + cm = CodeMode( + discovery_tools=[Search(name="search"), Search(name="search")], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + with pytest.raises(ValueError, match="unique"): + cm._build_discovery_tools() + + +# --------------------------------------------------------------------------- +# Tags discovery tool +# --------------------------------------------------------------------------- + + +async def test_categories_brief_shows_tag_counts() -> None: + mcp = FastMCP("Tags Brief") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + return x + y + + @mcp.tool(tags={"math"}) + def multiply(x: int, y: int) -> int: + return x * y + + @mcp.tool(tags={"text"}) + def greet(name: str) -> str: + return f"Hello, {name}!" + + mcp.add_transform( + CodeMode( + discovery_tools=[GetTags()], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + result = await _run_tool(mcp, "tags", {}) + text = _unwrap_string_result(result) + assert "math (2 tools)" in text + assert "text (1 tool)" in text + + +async def test_categories_full_lists_tools_per_tag() -> None: + mcp = FastMCP("Tags Full") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + """Add two numbers.""" + return x + y + + @mcp.tool(tags={"text"}) + def greet(name: str) -> str: + """Say hello.""" + return f"Hello, {name}!" + + mcp.add_transform( + CodeMode( + discovery_tools=[GetTags(default_detail="full")], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + result = await _run_tool(mcp, "tags", {}) + text = _unwrap_string_result(result) + assert "### math" in text + assert "- add: Add two numbers." in text + assert "### text" in text + assert "- greet: Say hello." in text + + +async def test_categories_includes_untagged() -> None: + mcp = FastMCP("Tags Untagged") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + return x + y + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform( + CodeMode( + discovery_tools=[GetTags()], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + result = await _run_tool(mcp, "tags", {}) + text = _unwrap_string_result(result) + assert "math" in text + assert "untagged (1 tool)" in text + + +async def test_categories_tool_in_multiple_tags() -> None: + mcp = FastMCP("Tags Multi-tag") + + @mcp.tool(tags={"math", "core"}) + def add(x: int, y: int) -> int: + return x + y + + mcp.add_transform( + CodeMode( + discovery_tools=[GetTags(default_detail="full")], + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + result = await _run_tool(mcp, "tags", {}) + text = _unwrap_string_result(result) + assert "### core" in text + assert "### math" in text + # Tool appears under both tags + assert text.count("- add") == 2 + + +async def test_categories_detail_override_per_call() -> None: + """LLM can override default_detail on a per-call basis.""" + mcp = FastMCP("Tags Override") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + """Add numbers.""" + return x + y + + mcp.add_transform( + CodeMode( + discovery_tools=[GetTags()], # default_detail="brief" + sandbox_provider=_UnsafeTestSandboxProvider(), + ) + ) + + # Override to full + result = await _run_tool(mcp, "tags", {"detail": "full"}) + text = _unwrap_string_result(result) + assert "### math" in text + assert "- add: Add numbers." in text + + +# --------------------------------------------------------------------------- +# Search with tags filtering +# --------------------------------------------------------------------------- + + +async def test_search_with_tags_filter() -> None: + mcp = FastMCP("Search Tags") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + """Add two numbers.""" + return x + y + + @mcp.tool(tags={"text"}) + def greet(name: str) -> str: + """Say hello.""" + return f"Hello, {name}!" + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "search", {"query": "add hello", "tags": ["math"]}) + text = _unwrap_string_result(result) + assert "add" in text + assert "greet" not in text + + +async def test_search_with_tags_filter_no_matches() -> None: + mcp = FastMCP("Search Tags Empty") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + """Add two numbers.""" + return x + y + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "search", {"query": "add", "tags": ["nonexistent"]}) + text = _unwrap_string_result(result) + assert "add" not in text or "No tools" in text + + +async def test_search_without_tags_returns_all() -> None: + """Search without tags parameter searches the full catalog.""" + mcp = FastMCP("Search No Tags") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + """Add two numbers.""" + return x + y + + @mcp.tool(tags={"text"}) + def greet(name: str) -> str: + """Say hello.""" + return f"Hello, {name}!" + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "search", {"query": "add hello"}) + text = _unwrap_string_result(result) + assert "add" in text + assert "greet" in text + + +async def test_search_with_untagged_filter() -> None: + """Search with tags=["untagged"] matches tools that have no tags.""" + mcp = FastMCP("Search Untagged") + + @mcp.tool(tags={"math"}) + def add(x: int, y: int) -> int: + """Add two numbers.""" + return x + y + + @mcp.tool + def ping() -> str: + """Ping.""" + return "pong" + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool(mcp, "search", {"query": "ping add", "tags": ["untagged"]}) + text = _unwrap_string_result(result) + assert "ping" in text + assert "add" not in text + + +async def test_get_schema_full_partial_match_returns_valid_json() -> None: + """get_schema with detail=full and missing tools returns valid JSON.""" + mcp = FastMCP("Schema Full Partial") + + @mcp.tool + def square(x: int) -> int: + """Compute the square.""" + return x * x + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + result = await _run_tool( + mcp, "get_schema", {"tools": ["square", "nonexistent"], "detail": "full"} + ) + text = _unwrap_string_result(result) + parsed = json.loads(text) + assert isinstance(parsed, list) + assert parsed[0]["name"] == "square" + assert parsed[-1] == {"not_found": ["nonexistent"]} + + +# --------------------------------------------------------------------------- +# Visibility and auth +# --------------------------------------------------------------------------- async def test_code_mode_execute_respects_disabled_tool_visibility() -> None: @@ -240,9 +723,7 @@ async def test_code_mode_execute_respects_disabled_tool_visibility() -> None: with pytest.raises(ToolError, match=r"Unknown tool"): await _run_tool( - mcp, - "execute", - {"code": "return await call_tool('secret', {})"}, + mcp, "execute", {"code": "return await call_tool('secret', {})"} ) @@ -258,8 +739,8 @@ async def test_code_mode_search_respects_disabled_tool_visibility() -> None: mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) result = await _run_tool(mcp, "search", {"query": "secret"}) - tools = _unwrap_search_results(result) - assert tools == [] + text = _unwrap_string_result(result) + assert "secret" not in text or "No tools" in text async def test_code_mode_execute_respects_tool_auth() -> None: @@ -273,9 +754,7 @@ async def test_code_mode_execute_respects_tool_auth() -> None: with pytest.raises(ToolError, match=r"Unknown tool"): await _run_tool( - mcp, - "execute", - {"code": "return await call_tool('protected', {})"}, + mcp, "execute", {"code": "return await call_tool('protected', {})"} ) @@ -290,12 +769,12 @@ async def test_code_mode_search_respects_tool_auth() -> None: mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) result = await _run_tool(mcp, "search", {"query": "protected"}) - tools = _unwrap_search_results(result) - assert tools == [] + text = _unwrap_string_result(result) + assert "protected" not in text or "No tools" in text async def test_code_mode_shadows_colliding_tool_names() -> None: - """Backend tools with the same name as meta-tools are shadowed, not rejected.""" + """Backend tools with the same name as meta-tools are shadowed.""" mcp = FastMCP("CodeMode Collision") @mcp.tool @@ -310,7 +789,7 @@ async def test_code_mode_shadows_colliding_tool_names() -> None: tools = await mcp.list_tools(run_middleware=False) tool_names = {t.name for t in tools} - assert tool_names == {"search", "execute"} + assert "execute" in tool_names result = await _run_tool( mcp, "execute", {"code": 'return await call_tool("ping", {})'} @@ -318,6 +797,43 @@ async def test_code_mode_shadows_colliding_tool_names() -> None: assert _unwrap_result(result) == {"result": "pong"} +# --------------------------------------------------------------------------- +# get_tool pass-through +# --------------------------------------------------------------------------- + + +async def test_code_mode_get_tool_returns_meta_tools_and_passes_through() -> None: + """get_tool returns meta-tools by name and passes through backend tools.""" + mcp = FastMCP("CodeMode GetTool") + + @mcp.tool + def ping() -> str: + return "pong" + + mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) + + search_tool = await mcp.get_tool("search") + assert search_tool is not None + assert search_tool.name == "search" + + schema_tool = await mcp.get_tool("get_schema") + assert schema_tool is not None + assert schema_tool.name == "get_schema" + + execute_tool = await mcp.get_tool("execute") + assert execute_tool is not None + assert execute_tool.name == "execute" + + ping_tool = await mcp.get_tool("ping") + assert ping_tool is not None + assert ping_tool.name == "ping" + + +# --------------------------------------------------------------------------- +# Execute edge cases +# --------------------------------------------------------------------------- + + async def test_code_mode_execute_non_text_content_stringified() -> None: mcp = FastMCP("CodeMode NonText") @@ -328,32 +844,13 @@ async def test_code_mode_execute_non_text_content_stringified() -> None: mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) result = await _run_tool( - mcp, - "execute", - {"code": "return await call_tool('image_tool', {})"}, + mcp, "execute", {"code": "return await call_tool('image_tool', {})"} ) unwrapped = _unwrap_result(result) assert isinstance(unwrapped, str) assert "base64data" in unwrapped -async def test_monty_provider_raises_informative_error_when_missing( - monkeypatch: pytest.MonkeyPatch, -) -> None: - provider = MontySandboxProvider() - real_import_module = importlib.import_module - - def _fake_import_module(name: str, package: str | None = None): - if name == "pydantic_monty": - raise ModuleNotFoundError("No module named 'pydantic_monty'") - return real_import_module(name, package) - - monkeypatch.setattr(importlib, "import_module", _fake_import_module) - - with pytest.raises(ImportError, match=r"fastmcp\[code-mode\]"): - await provider.run("return 1") - - async def test_code_mode_execute_multi_tool_chaining() -> None: """Execute block can chain multiple call_tool() calls.""" mcp = FastMCP("CodeMode Chaining") @@ -382,63 +879,7 @@ async def test_code_mode_execute_multi_tool_chaining() -> None: assert _unwrap_result(result) == {"result": 7} -async def test_code_mode_execute_default_arguments_overridden_by_explicit() -> None: - """Explicit params in call_tool() override default_arguments.""" - mcp = FastMCP("CodeMode Override") - - @mcp.tool - def greet(name: str, greeting: str) -> str: - return f"{greeting}, {name}!" - - mcp.add_transform( - CodeMode( - default_arguments={"greeting": "Hello"}, - sandbox_provider=_UnsafeTestSandboxProvider(), - ) - ) - - result = await _run_tool( - mcp, - "execute", - {"code": "return await call_tool('greet', {'name': 'World'})"}, - ) - assert _unwrap_result(result) == {"result": "Hello, World!"} - - result = await _run_tool( - mcp, - "execute", - { - "code": "return await call_tool('greet', {'name': 'World', 'greeting': 'Hi'})" - }, - ) - assert _unwrap_result(result) == {"result": "Hi, World!"} - - -async def test_code_mode_get_tool_returns_meta_tools_and_passes_through() -> None: - """get_tool returns meta-tools by name and passes through backend tools.""" - mcp = FastMCP("CodeMode GetTool") - - @mcp.tool - def ping() -> str: - return "pong" - - mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) - - search_tool = await mcp.get_tool("search") - assert search_tool is not None - assert search_tool.name == "search" - - execute_tool = await mcp.get_tool("execute") - assert execute_tool is not None - assert execute_tool.name == "execute" - - ping_tool = await mcp.get_tool("ping") - assert ping_tool is not None - assert ping_tool.name == "ping" - - async def test_code_mode_sandbox_error_surfaces_as_tool_error() -> None: - """Runtime errors in sandbox code surface as ToolError.""" mcp = FastMCP("CodeMode Errors") @mcp.tool @@ -451,8 +892,29 @@ async def test_code_mode_sandbox_error_surfaces_as_tool_error() -> None: await _run_tool(mcp, "execute", {"code": "raise ValueError('boom')"}) +# --------------------------------------------------------------------------- +# Sandbox provider tests +# --------------------------------------------------------------------------- + + +async def test_monty_provider_raises_informative_error_when_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = MontySandboxProvider() + real_import_module = importlib.import_module + + def _fake_import_module(name: str, package: str | None = None): + if name == "pydantic_monty": + raise ModuleNotFoundError("No module named 'pydantic_monty'") + return real_import_module(name, package) + + monkeypatch.setattr(importlib, "import_module", _fake_import_module) + + with pytest.raises(ImportError, match=r"fastmcp\[code-mode\]"): + await provider.run("return 1") + + async def test_monty_provider_forwards_limits() -> None: - """MontySandboxProvider passes limits through to pydantic-monty.""" provider = MontySandboxProvider(limits={"max_duration_secs": 0.1}) with pytest.raises(Exception, match="time limit exceeded"): @@ -460,321 +922,6 @@ async def test_monty_provider_forwards_limits() -> None: async def test_monty_provider_no_limits_by_default() -> None: - """Without limits, a simple script completes normally.""" provider = MontySandboxProvider() result = await provider.run("return 1 + 2") assert result == 3 - - -def test_code_mode_rejects_identical_tool_names() -> None: - """CodeMode raises ValueError when search and execute names collide.""" - with pytest.raises(ValueError, match="must be different"): - CodeMode( - search_tool_name="tools", - execute_tool_name="tools", - sandbox_provider=_UnsafeTestSandboxProvider(), - ) - - -# --------------------------------------------------------------------------- -# _schema_type unit tests -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "schema,expected", - [ - ({"type": "string"}, "string"), - ({"type": "integer"}, "integer"), - ({"type": "boolean"}, "boolean"), - ({"type": "null"}, "null"), - ({"type": "array", "items": {"type": "string"}}, "string[]"), - ({"type": "array", "items": {"type": "integer"}}, "integer[]"), - ({"type": "array"}, "any[]"), - ({"$ref": "#/$defs/Foo"}, "object"), - ({"properties": {"x": {"type": "int"}}}, "object"), - ({}, "any"), - (None, "any"), - ("not a dict", "any"), - ], -) -def test_schema_type_basic(schema: Any, expected: str) -> None: - assert _schema_type(schema) == expected - - -@pytest.mark.parametrize( - "schema,expected", - [ - # anyOf: optional field — null stripped, "?" appended - ({"anyOf": [{"type": "string"}, {"type": "null"}]}, "string?"), - # anyOf: non-nullable union — all branches kept - ({"anyOf": [{"type": "string"}, {"type": "integer"}]}, "string | integer"), - # anyOf: multiple branches including null - ( - {"anyOf": [{"type": "string"}, {"type": "integer"}, {"type": "null"}]}, - "string | integer?", - ), - # anyOf: all-null (edge case) - ({"anyOf": [{"type": "null"}]}, "null"), - # anyOf: empty - ({"anyOf": []}, "any"), - # oneOf: treated identically to anyOf - ({"oneOf": [{"type": "string"}, {"type": "null"}]}, "string?"), - ({"oneOf": [{"type": "string"}, {"type": "integer"}]}, "string | integer"), - # allOf: always "object" (Pydantic composite / intersection type) - ({"allOf": [{"type": "object"}]}, "object"), - ({"allOf": [{"$ref": "#/$defs/Foo"}, {"$ref": "#/$defs/Bar"}]}, "object"), - ], -) -def test_schema_type_unions(schema: Any, expected: str) -> None: - assert _schema_type(schema) == expected - - -# --------------------------------------------------------------------------- -# _schema_section unit tests -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "schema,expected_lines", - [ - # None → generic fallback - (None, ["**Parameters**", "- `value` (any)"]), - # Non-dict → generic fallback - ("string", ["**Parameters**", "- `value` (any)"]), - # Scalar schema without properties → type label used - ({"type": "string"}, ["**Parameters**", "- `value` (string)"]), - # Empty properties dict → zero-argument tool - ( - {"type": "object", "properties": {}}, - ["**Parameters**", "*(no parameters)*"], - ), - ], -) -def test_schema_section_fallbacks(schema: Any, expected_lines: list[str]) -> None: - assert _schema_section(schema, "Parameters") == expected_lines - - -def test_schema_section_lists_fields_with_required_marker() -> None: - schema = { - "type": "object", - "properties": { - "name": {"type": "string"}, - "age": {"type": "integer"}, - }, - "required": ["name"], - } - lines = _schema_section(schema, "Parameters") - assert lines[0] == "**Parameters**" - assert "- `name` (string, required)" in lines - assert "- `age` (integer)" in lines - - -# --------------------------------------------------------------------------- -# serialize_tools_for_output_markdown unit tests -# --------------------------------------------------------------------------- - - -def test_serialize_tools_for_output_markdown_empty_list() -> None: - assert serialize_tools_for_output_markdown([]) == "No tools matched the query." - - -async def test_serialize_tools_for_output_markdown_basic_tool() -> None: - mcp = FastMCP("MD Basic") - - @mcp.tool - def square(x: int) -> int: - """Compute the square of a number.""" - return x * x - - tools = await mcp.list_tools() - result = serialize_tools_for_output_markdown(tools) - - assert "### square" in result - assert "Compute the square of a number." in result - assert "**Parameters**" in result - assert "`x` (integer, required)" in result - - -async def test_serialize_tools_for_output_markdown_omits_output_section_when_no_schema() -> ( - None -): - """Tools without output_schema should not render a Returns section.""" - mcp = FastMCP("MD No Output") - - @mcp.tool - def ping() -> None: - pass - - tools = await mcp.list_tools() - result = serialize_tools_for_output_markdown(tools) - - assert "**Returns**" not in result - - -async def test_serialize_tools_for_output_markdown_includes_output_section_when_schema_present() -> ( - None -): - mcp = FastMCP("MD With Output") - - @mcp.tool - def double(x: int) -> int: - return x * 2 - - tools = await mcp.list_tools() - result = serialize_tools_for_output_markdown(tools) - - assert "**Returns**" in result - - -async def test_serialize_tools_for_output_markdown_omits_description_when_absent() -> ( - None -): - mcp = FastMCP("MD No Desc") - - @mcp.tool - def ping() -> None: - pass - - tools = await mcp.list_tools() - result = serialize_tools_for_output_markdown(tools) - - # Header present, no extra blank description line injected - assert "### ping" in result - - -async def test_serialize_tools_for_output_markdown_optional_field_uses_question_mark() -> ( - None -): - mcp = FastMCP("MD Optional") - - @mcp.tool - def greet(name: str, greeting: str | None = None) -> str: - return f"{greeting or 'Hello'}, {name}!" - - tools = await mcp.list_tools() - result = serialize_tools_for_output_markdown(tools) - - assert "`greeting` (string?)" in result - - -async def test_serialize_tools_for_output_markdown_multiple_tools_separated() -> None: - mcp = FastMCP("MD Multi") - - @mcp.tool - def add(a: int, b: int) -> int: - return a + b - - @mcp.tool - def subtract(a: int, b: int) -> int: - return a - b - - tools = await mcp.list_tools() - result = serialize_tools_for_output_markdown(tools) - - assert "### add" in result - assert "### subtract" in result - # Tools separated by double newline - assert "\n\n" in result - - -# --------------------------------------------------------------------------- -# CodeMode search_result_serializer integration tests -# --------------------------------------------------------------------------- - - -def _unwrap_serializer_result(result: ToolResult) -> str: - """Extract a string result returned by a custom search serializer. - - Custom serializers returning str are wrapped in {"result": "..."} by the - output schema, so we need one extra level of unwrapping compared to - _unwrap_result. - """ - data = _unwrap_result(result) - if isinstance(data, dict) and "result" in data: - return data["result"] - assert isinstance(data, str) - return data - - -async def test_code_mode_search_supports_custom_serializer() -> None: - mcp = FastMCP("CodeMode Custom Serializer") - - @mcp.tool - def square(x: int) -> int: - return x * x - - mcp.add_transform( - CodeMode( - sandbox_provider=_UnsafeTestSandboxProvider(), - search_result_serializer=lambda tools: "\n".join(t.name for t in tools), - ) - ) - - result = await _run_tool(mcp, "search", {"query": "square"}) - text = _unwrap_serializer_result(result) - assert isinstance(text, str) - assert "square" in text - - -async def test_code_mode_search_supports_async_custom_serializer() -> None: - mcp = FastMCP("CodeMode Async Serializer") - - @mcp.tool - def square(x: int) -> int: - return x * x - - async def async_serializer(tools: Any) -> str: - return ", ".join(t.name for t in tools) - - mcp.add_transform( - CodeMode( - sandbox_provider=_UnsafeTestSandboxProvider(), - search_result_serializer=async_serializer, - ) - ) - - result = await _run_tool(mcp, "search", {"query": "square"}) - text = _unwrap_serializer_result(result) - assert isinstance(text, str) - assert "square" in text - - -async def test_code_mode_search_markdown_serializer() -> None: - mcp = FastMCP("CodeMode Markdown Serializer") - - @mcp.tool - def square(x: int) -> int: - """Compute the square of a number.""" - return x * x - - mcp.add_transform( - CodeMode( - sandbox_provider=_UnsafeTestSandboxProvider(), - search_result_serializer=serialize_tools_for_output_markdown, - ) - ) - - result = await _run_tool(mcp, "search", {"query": "square"}) - text = _unwrap_serializer_result(result) - assert isinstance(text, str) - assert "### square" in text - assert "Compute the square of a number." in text - assert "**Parameters**" in text - - -async def test_code_mode_search_default_serializer_returns_list() -> None: - """Default (no custom serializer) still returns the JSON list format.""" - mcp = FastMCP("CodeMode Default Serializer") - - @mcp.tool - def square(x: int) -> int: - """Compute the square of a number.""" - return x * x - - mcp.add_transform(CodeMode(sandbox_provider=_UnsafeTestSandboxProvider())) - - result = await _run_tool(mcp, "search", {"query": "square"}) - tools = _unwrap_search_results(result) - assert isinstance(tools, list) - assert tools[0]["name"] == "square" diff --git a/tests/experimental/transforms/test_code_mode_serialization.py b/tests/experimental/transforms/test_code_mode_serialization.py new file mode 100644 index 000000000..a589ad790 --- /dev/null +++ b/tests/experimental/transforms/test_code_mode_serialization.py @@ -0,0 +1,197 @@ +from typing import Any + +import pytest + +from fastmcp import FastMCP +from fastmcp.server.transforms.search.base import ( + _schema_section, + _schema_type, + serialize_tools_for_output_markdown, +) + +# --------------------------------------------------------------------------- +# _schema_type unit tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "schema,expected", + [ + ({"type": "string"}, "string"), + ({"type": "integer"}, "integer"), + ({"type": "boolean"}, "boolean"), + ({"type": "null"}, "null"), + ({"type": "array", "items": {"type": "string"}}, "string[]"), + ({"type": "array", "items": {"type": "integer"}}, "integer[]"), + ({"type": "array"}, "any[]"), + ({"$ref": "#/$defs/Foo"}, "object"), + ({"properties": {"x": {"type": "int"}}}, "object"), + ({}, "any"), + (None, "any"), + ("not a dict", "any"), + ], +) +def test_schema_type_basic(schema: Any, expected: str) -> None: + assert _schema_type(schema) == expected + + +@pytest.mark.parametrize( + "schema,expected", + [ + ({"anyOf": [{"type": "string"}, {"type": "null"}]}, "string?"), + ({"anyOf": [{"type": "string"}, {"type": "integer"}]}, "string | integer"), + ( + {"anyOf": [{"type": "string"}, {"type": "integer"}, {"type": "null"}]}, + "string | integer?", + ), + ({"anyOf": [{"type": "null"}]}, "null"), + ({"anyOf": []}, "any"), + ({"oneOf": [{"type": "string"}, {"type": "null"}]}, "string?"), + ({"oneOf": [{"type": "string"}, {"type": "integer"}]}, "string | integer"), + ({"allOf": [{"type": "object"}]}, "object"), + ({"allOf": [{"$ref": "#/$defs/Foo"}, {"$ref": "#/$defs/Bar"}]}, "object"), + ], +) +def test_schema_type_unions(schema: Any, expected: str) -> None: + assert _schema_type(schema) == expected + + +# --------------------------------------------------------------------------- +# _schema_section unit tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "schema,expected_lines", + [ + (None, ["**Parameters**", "- `value` (any)"]), + ("string", ["**Parameters**", "- `value` (any)"]), + ({"type": "string"}, ["**Parameters**", "- `value` (string)"]), + ( + {"type": "object", "properties": {}}, + ["**Parameters**", "*(no parameters)*"], + ), + ], +) +def test_schema_section_fallbacks(schema: Any, expected_lines: list[str]) -> None: + assert _schema_section(schema, "Parameters") == expected_lines + + +def test_schema_section_lists_fields_with_required_marker() -> None: + schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + }, + "required": ["name"], + } + lines = _schema_section(schema, "Parameters") + assert lines[0] == "**Parameters**" + assert "- `name` (string, required)" in lines + assert "- `age` (integer)" in lines + + +# --------------------------------------------------------------------------- +# serialize_tools_for_output_markdown unit tests +# --------------------------------------------------------------------------- + + +def test_serialize_tools_for_output_markdown_empty_list() -> None: + assert serialize_tools_for_output_markdown([]) == "No tools matched the query." + + +async def test_serialize_tools_for_output_markdown_basic_tool() -> None: + mcp = FastMCP("MD Basic") + + @mcp.tool + def square(x: int) -> int: + """Compute the square of a number.""" + return x * x + + tools = await mcp.list_tools() + result = serialize_tools_for_output_markdown(tools) + + assert "### square" in result + assert "Compute the square of a number." in result + assert "**Parameters**" in result + assert "`x` (integer, required)" in result + + +async def test_serialize_tools_for_output_markdown_omits_output_section_when_no_schema() -> ( + None +): + mcp = FastMCP("MD No Output") + + @mcp.tool + def ping() -> None: + pass + + tools = await mcp.list_tools() + result = serialize_tools_for_output_markdown(tools) + + assert "**Returns**" not in result + + +async def test_serialize_tools_for_output_markdown_includes_output_section_when_schema_present() -> ( + None +): + mcp = FastMCP("MD With Output") + + @mcp.tool + def double(x: int) -> int: + return x * 2 + + tools = await mcp.list_tools() + result = serialize_tools_for_output_markdown(tools) + + assert "**Returns**" in result + + +async def test_serialize_tools_for_output_markdown_omits_description_when_absent() -> ( + None +): + mcp = FastMCP("MD No Desc") + + @mcp.tool + def ping() -> None: + pass + + tools = await mcp.list_tools() + result = serialize_tools_for_output_markdown(tools) + + assert "### ping" in result + + +async def test_serialize_tools_for_output_markdown_optional_field_uses_question_mark() -> ( + None +): + mcp = FastMCP("MD Optional") + + @mcp.tool + def greet(name: str, greeting: str | None = None) -> str: + return f"{greeting or 'Hello'}, {name}!" + + tools = await mcp.list_tools() + result = serialize_tools_for_output_markdown(tools) + + assert "`greeting` (string?)" in result + + +async def test_serialize_tools_for_output_markdown_multiple_tools_separated() -> None: + mcp = FastMCP("MD Multi") + + @mcp.tool + def add(a: int, b: int) -> int: + return a + b + + @mcp.tool + def subtract(a: int, b: int) -> int: + return a - b + + tools = await mcp.list_tools() + result = serialize_tools_for_output_markdown(tools) + + assert "### add" in result + assert "### subtract" in result + assert "\n\n" in result