diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 00757da28..9dbc9376c 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -1,6 +1,5 @@ from __future__ import annotations -import json import logging import weakref from collections.abc import Callable, Generator, Mapping, Sequence @@ -8,34 +7,24 @@ from contextlib import contextmanager from contextvars import ContextVar, Token from dataclasses import dataclass from logging import Logger -from typing import Any, Literal, cast, overload +from typing import Any, Literal, overload import mcp.types from mcp import LoggingLevel, ServerSession from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext from mcp.types import ( - CreateMessageResult, - CreateMessageResultWithTools, GetPromptResult, ModelPreferences, Root, SamplingMessage, - SamplingMessageContentBlock, - TextContent, - ToolChoice, - ToolResultContent, - ToolUseContent, ) from mcp.types import Prompt as SDKPrompt from mcp.types import Resource as SDKResource -from mcp.types import Tool as SDKTool -from pydantic import ValidationError from pydantic.networks import AnyUrl from starlette.requests import Request from typing_extensions import TypeVar -from fastmcp import settings from fastmcp.resources.resource import ResourceResult from fastmcp.server.elicitation import ( AcceptedElicitation, @@ -46,19 +35,29 @@ from fastmcp.server.elicitation import ( ) from fastmcp.server.sampling import SampleStep, SamplingResult, SamplingTool from fastmcp.server.sampling.run import ( - _parse_model_preferences, - call_sampling_handler, - determine_handler_mode, -) -from fastmcp.server.sampling.run import ( - execute_tools as run_sampling_tools, + sample_impl, + sample_step_impl, ) from fastmcp.server.server import FastMCP, StateValue -from fastmcp.server.transforms.enabled import Enabled -from fastmcp.utilities.components import FastMCPComponent -from fastmcp.utilities.json_schema import compress_schema +from fastmcp.server.transforms.enabled import ( + Enabled, +) +from fastmcp.server.transforms.enabled import ( + disable_components as _disable_components, +) +from fastmcp.server.transforms.enabled import ( + enable_components as _enable_components, +) +from fastmcp.server.transforms.enabled import ( + get_session_transforms as _get_session_transforms, +) +from fastmcp.server.transforms.enabled import ( + get_visibility_rules as _get_visibility_rules, +) +from fastmcp.server.transforms.enabled import ( + reset_components as _reset_components, +) from fastmcp.utilities.logging import _clamp_logger, get_logger -from fastmcp.utilities.types import get_cached_typeadapter from fastmcp.utilities.versions import VersionSpec logger: Logger = get_logger(name=__name__) @@ -73,8 +72,8 @@ _clamp_logger(logger=to_client_logger, max_level="DEBUG") T = TypeVar("T", default=Any) ResultT = TypeVar("ResultT", default=str) -# Simplified tool choice type - just the mode string instead of the full MCP object -ToolChoiceOption = Literal["auto", "required", "none"] +# Import ToolChoiceOption from sampling module (after other imports) +from fastmcp.server.sampling.run import ToolChoiceOption # noqa: E402 _current_context: ContextVar[Context | None] = ContextVar("context", default=None) @@ -325,26 +324,48 @@ class Context: related_request_id=self.request_id, ) + async def _paginate_list( + self, + request_factory: Callable[[str | None], Any], + call_method: Callable[[Any], Any], + extract_items: Callable[[Any], list[Any]], + ) -> list[Any]: + """Generic pagination helper for list operations. + + Args: + request_factory: Function that creates a request from a cursor + call_method: Async method to call with the request + extract_items: Function to extract items from the result + + Returns: + List of all items across all pages + """ + all_items: list[Any] = [] + cursor: str | None = None + while True: + request = request_factory(cursor) + result = await call_method(request) + all_items.extend(extract_items(result)) + if result.nextCursor is None: + break + cursor = result.nextCursor + return all_items + async def list_resources(self) -> list[SDKResource]: """List all available resources from the server. Returns: List of Resource objects available on the server """ - all_resources: list[SDKResource] = [] - cursor: str | None = None - while True: - request = mcp.types.ListResourcesRequest( - params=mcp.types.ListResourcesRequestParams(cursor=cursor) + return await self._paginate_list( + request_factory=lambda cursor: mcp.types.ListResourcesRequest( + params=mcp.types.PaginatedRequestParams(cursor=cursor) if cursor else None - ) - result = await self.fastmcp._list_resources_mcp(request) - all_resources.extend(result.resources) - if result.nextCursor is None: - break - cursor = result.nextCursor - return all_resources + ), + call_method=self.fastmcp._list_resources_mcp, + extract_items=lambda result: result.resources, + ) async def list_prompts(self) -> list[SDKPrompt]: """List all available prompts from the server. @@ -352,20 +373,15 @@ class Context: Returns: List of Prompt objects available on the server """ - all_prompts: list[SDKPrompt] = [] - cursor: str | None = None - while True: - request = mcp.types.ListPromptsRequest( - params=mcp.types.ListPromptsRequestParams(cursor=cursor) + return await self._paginate_list( + request_factory=lambda cursor: mcp.types.ListPromptsRequest( + params=mcp.types.PaginatedRequestParams(cursor=cursor) if cursor else None - ) - result = await self.fastmcp._list_prompts_mcp(request) - all_prompts.extend(result.prompts) - if result.nextCursor is None: - break - cursor = result.nextCursor - return all_prompts + ), + call_method=self.fastmcp._list_prompts_mcp, + extract_items=lambda result: result.prompts, + ) async def get_prompt( self, name: str, arguments: dict[str, Any] | None = None @@ -707,100 +723,18 @@ class Context: # Continue with tool results messages = step.history """ - # Convert messages to SamplingMessage objects - current_messages = _prepare_messages(messages) - - # Convert tools to SamplingTools - sampling_tools = _prepare_tools(tools) - sdk_tools: list[SDKTool] | None = ( - [t._to_sdk_tool() for t in sampling_tools] if sampling_tools else None + return await sample_step_impl( + self, + messages=messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=max_tokens, + model_preferences=model_preferences, + tools=tools, + tool_choice=tool_choice, + auto_execute_tools=execute_tools, + mask_error_details=mask_error_details, ) - tool_map: dict[str, SamplingTool] = ( - {t.name: t for t in sampling_tools} if sampling_tools else {} - ) - - # Determine whether to use fallback handler or client - use_fallback = determine_handler_mode(self, bool(sampling_tools)) - - # Build tool choice - effective_tool_choice: ToolChoice | None = None - if tool_choice is not None: - if tool_choice not in ("auto", "required", "none"): - raise ValueError( - f"Invalid tool_choice: {tool_choice!r}. " - "Must be 'auto', 'required', or 'none'." - ) - effective_tool_choice = ToolChoice( - mode=cast(Literal["auto", "required", "none"], tool_choice) - ) - - # Effective max_tokens - effective_max_tokens = max_tokens if max_tokens is not None else 512 - - # Make the LLM call - if use_fallback: - response = await call_sampling_handler( - self, - current_messages, - system_prompt=system_prompt, - temperature=temperature, - max_tokens=effective_max_tokens, - model_preferences=model_preferences, - sdk_tools=sdk_tools, - tool_choice=effective_tool_choice, - ) - else: - response = await self.session.create_message( - messages=current_messages, - system_prompt=system_prompt, - temperature=temperature, - max_tokens=effective_max_tokens, - model_preferences=_parse_model_preferences(model_preferences), - tools=sdk_tools, - tool_choice=effective_tool_choice, - related_request_id=self.request_id, - ) - - # Check if this is a tool use response - is_tool_use_response = ( - isinstance(response, CreateMessageResultWithTools) - and response.stopReason == "toolUse" - ) - - # Always include the assistant response in history - current_messages.append( - SamplingMessage(role="assistant", content=response.content) - ) - - # If not a tool use, return immediately - if not is_tool_use_response: - return SampleStep(response=response, history=current_messages) - - # If not executing tools, return with assistant message but no tool results - if not execute_tools: - return SampleStep(response=response, history=current_messages) - - # Execute tools and add results to history - step_tool_calls = _extract_tool_calls(response) - if step_tool_calls: - effective_mask = ( - mask_error_details - if mask_error_details is not None - else settings.mask_error_details - ) - tool_results: list[SamplingMessageContentBlock] = await run_sampling_tools( # type: ignore[assignment] - step_tool_calls, tool_map, mask_error_details=effective_mask - ) - - if tool_results: - current_messages.append( - SamplingMessage( - role="user", - content=tool_results, - ) - ) - - return SampleStep(response=response, history=current_messages) @overload async def sample( @@ -880,113 +814,17 @@ class Context: - .result: The typed result (str for text, parsed object for structured) - .history: All messages exchanged during sampling """ - # Safety limit to prevent infinite loops - max_iterations = 100 - - # Convert tools to SamplingTools - sampling_tools = _prepare_tools(tools) - - # Handle structured output with result_type - tool_choice: str | None = None - if result_type is not None and result_type is not str: - final_response_tool = _create_final_response_tool(result_type) - sampling_tools = list(sampling_tools) if sampling_tools else [] - sampling_tools.append(final_response_tool) - - # Always require tool calls when result_type is set - the LLM must - # eventually call final_response (text responses are not accepted) - tool_choice = "required" - - # Convert messages for the loop - current_messages: str | Sequence[str | SamplingMessage] = messages - - for _iteration in range(max_iterations): - step = await self.sample_step( - messages=current_messages, - system_prompt=system_prompt, - temperature=temperature, - max_tokens=max_tokens, - model_preferences=model_preferences, - tools=sampling_tools, - tool_choice=tool_choice, - mask_error_details=mask_error_details, - ) - - # Check for final_response tool call for structured output - if result_type is not None and result_type is not str and step.is_tool_use: - for tool_call in step.tool_calls: - if tool_call.name == "final_response": - # Validate and return the structured result - type_adapter = get_cached_typeadapter(result_type) - - # Unwrap if we wrapped primitives (non-object schemas) - input_data = tool_call.input - original_schema = compress_schema( - type_adapter.json_schema(), prune_titles=True - ) - if ( - original_schema.get("type") != "object" - and isinstance(input_data, dict) - and "value" in input_data - ): - input_data = input_data["value"] - - try: - validated_result = type_adapter.validate_python(input_data) - text = json.dumps( - type_adapter.dump_python(validated_result, mode="json") - ) - return SamplingResult( - text=text, - result=validated_result, - history=step.history, - ) - except ValidationError as e: - # Validation failed - add error as tool result - step.history.append( - SamplingMessage( - role="user", - content=[ - ToolResultContent( - type="tool_result", - toolUseId=tool_call.id, - content=[ - TextContent( - type="text", - text=( - f"Validation error: {e}. " - "Please try again with valid data." - ), - ) - ], - isError=True, - ) - ], - ) - ) - - # If not a tool use response, we're done - if not step.is_tool_use: - # For structured output, the LLM must use the final_response tool - if result_type is not None and result_type is not str: - raise RuntimeError( - f"Expected structured output of type {result_type.__name__}, " - "but the LLM returned a text response instead of calling " - "the final_response tool." - ) - return SamplingResult( - text=step.text, - result=cast(ResultT, step.text if step.text else ""), - history=step.history, - ) - - # Continue with the updated history - current_messages = step.history - - # After first iteration, reset tool_choice to auto - tool_choice = None - - raise RuntimeError(f"Sampling exceeded maximum iterations ({max_iterations})") + return await sample_impl( + self, + messages=messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=max_tokens, + model_preferences=model_preferences, + tools=tools, + result_type=result_type, + mask_error_details=mask_error_details, + ) @overload async def elicit( @@ -1150,71 +988,11 @@ class Context: async def _get_visibility_rules(self) -> list[dict[str, Any]]: """Load visibility rule dicts from session state.""" - return await self.get_state("_visibility_rules") or [] - - async def _save_visibility_rules( - self, - rules: list[dict[str, Any]], - *, - components: set[Literal["tool", "resource", "template", "prompt"]] - | None = None, - ) -> None: - """Save visibility rule dicts to session state and send notifications. - - Args: - rules: The visibility rules to save. - components: Optional hint about which component types are affected. - If None, sends notifications for all types (safe default). - If provided, only sends notifications for specified types. - """ - await self.set_state("_visibility_rules", rules) - - # Send notifications based on components hint - # Note: MCP has no separate template notification - templates use ResourceListChangedNotification - if components is None or "tool" in components: - await self.send_notification(mcp.types.ToolListChangedNotification()) - if components is None or "resource" in components or "template" in components: - await self.send_notification(mcp.types.ResourceListChangedNotification()) - if components is None or "prompt" in components: - await self.send_notification(mcp.types.PromptListChangedNotification()) - - def _create_enabled_transforms(self, rules: list[dict[str, Any]]) -> list[Enabled]: - """Convert rule dicts to Enabled transforms.""" - transforms = [] - for params in rules: - version = None - if params.get("version"): - version_dict = params["version"] - version = VersionSpec( - gte=version_dict.get("gte"), - lt=version_dict.get("lt"), - eq=version_dict.get("eq"), - ) - transforms.append( - Enabled( - params["enabled"], - names=set(params["names"]) if params.get("names") else None, - keys=set(params["keys"]) if params.get("keys") else None, - version=version, - tags=set(params["tags"]) if params.get("tags") else None, - components=( - set(params["components"]) if params.get("components") else None - ), - match_all=params.get("match_all", False), - ) - ) - return transforms + return await _get_visibility_rules(self) async def _get_session_transforms(self) -> list[Enabled]: """Get session-specific Enabled transforms from state store.""" - try: - # Will raise RuntimeError if no session available - _ = self.session_id - except RuntimeError: - return [] - - rules = await self._get_visibility_rules() - return self._create_enabled_transforms(rules) + return await _get_session_transforms(self) async def enable_components( self, @@ -1244,30 +1022,15 @@ class Context: components: Component types to match (e.g., {"tool", "prompt"}). match_all: If True, matches all components regardless of other criteria. """ - # Normalize empty sets to None (empty = match all) - components = components if components else None - - # Load current rules - rules = await self._get_visibility_rules() - - # Create new rule dict - rule: dict[str, Any] = { - "enabled": True, - "names": list(names) if names else None, - "keys": list(keys) if keys else None, - "version": ( - {"gte": version.gte, "lt": version.lt, "eq": version.eq} - if version - else None - ), - "tags": list(tags) if tags else None, - "components": list(components) if components else None, - "match_all": match_all, - } - - # Add and save (notifications sent by _save_visibility_rules) - rules.append(rule) - await self._save_visibility_rules(rules, components=components) + await _enable_components( + self, + names=names, + keys=keys, + version=version, + tags=tags, + components=components, + match_all=match_all, + ) async def disable_components( self, @@ -1297,30 +1060,15 @@ class Context: components: Component types to match (e.g., {"tool", "prompt"}). match_all: If True, matches all components regardless of other criteria. """ - # Normalize empty sets to None (empty = match all) - components = components if components else None - - # Load current rules - rules = await self._get_visibility_rules() - - # Create new rule dict - rule: dict[str, Any] = { - "enabled": False, - "names": list(names) if names else None, - "keys": list(keys) if keys else None, - "version": ( - {"gte": version.gte, "lt": version.lt, "eq": version.eq} - if version - else None - ), - "tags": list(tags) if tags else None, - "components": list(components) if components else None, - "match_all": match_all, - } - - # Add and save (notifications sent by _save_visibility_rules) - rules.append(rule) - await self._save_visibility_rules(rules, components=components) + await _disable_components( + self, + names=names, + keys=keys, + version=version, + tags=tags, + components=components, + match_all=match_all, + ) async def reset_components(self) -> None: """Clear all session visibility rules. @@ -1330,7 +1078,7 @@ class Context: Sends notifications to this session only: ToolListChangedNotification, ResourceListChangedNotification, and PromptListChangedNotification. """ - await self._save_visibility_rules([]) + await _reset_components(self) async def _log_to_server_and_client( @@ -1359,137 +1107,3 @@ async def _log_to_server_and_client( logger=logger_name, related_request_id=related_request_id, ) - - -def _create_final_response_tool(result_type: type) -> SamplingTool: - """Create a synthetic 'final_response' tool for structured output. - - This tool is used to capture structured responses from the LLM. - The tool's schema is derived from the result_type. - """ - type_adapter = get_cached_typeadapter(result_type) - schema = type_adapter.json_schema() - schema = compress_schema(schema, prune_titles=True) - - # Tool parameters must be object-shaped. Wrap primitives in {"value": } - if schema.get("type") != "object": - schema = { - "type": "object", - "properties": {"value": schema}, - "required": ["value"], - } - - # The fn just returns the input as-is (validation happens in the loop) - def final_response(**kwargs: Any) -> dict[str, Any]: - return kwargs - - return SamplingTool( - name="final_response", - description=( - "Call this tool to provide your final response. " - "Use this when you have completed the task and are ready to return the result." - ), - parameters=schema, - fn=final_response, - ) - - -def _extract_text_from_content( - content: SamplingMessageContentBlock | list[SamplingMessageContentBlock], -) -> str | None: - """Extract text from content block(s). - - Returns the text if content is a TextContent or list containing TextContent, - otherwise returns None. - """ - if isinstance(content, list): - for block in content: - if isinstance(block, TextContent): - return block.text - return None - elif isinstance(content, TextContent): - return content.text - return None - - -def _prepare_messages( - messages: str | Sequence[str | SamplingMessage], -) -> list[SamplingMessage]: - """Convert various message formats to a list of SamplingMessage objects.""" - if isinstance(messages, str): - return [ - SamplingMessage( - content=TextContent(text=messages, type="text"), role="user" - ) - ] - else: - return [ - SamplingMessage(content=TextContent(text=m, type="text"), role="user") - if isinstance(m, str) - else m - for m in messages - ] - - -def _prepare_tools( - tools: Sequence[SamplingTool | Callable[..., Any]] | None, -) -> list[SamplingTool] | None: - """Convert tools to SamplingTool objects.""" - if tools is None: - return None - - sampling_tools: list[SamplingTool] = [] - for t in tools: - if isinstance(t, SamplingTool): - sampling_tools.append(t) - elif callable(t): - sampling_tools.append(SamplingTool.from_function(t)) - else: - raise TypeError(f"Expected SamplingTool or callable, got {type(t)}") - - return sampling_tools if sampling_tools else None - - -def _extract_tool_calls( - response: CreateMessageResult | CreateMessageResultWithTools, -) -> list[ToolUseContent]: - """Extract tool calls from a response.""" - content = response.content - if isinstance(content, list): - return [c for c in content if isinstance(c, ToolUseContent)] - elif isinstance(content, ToolUseContent): - return [content] - return [] - - -ComponentT = TypeVar("ComponentT", bound="FastMCPComponent") - - -async def apply_session_transforms( - components: Sequence[ComponentT], -) -> Sequence[ComponentT]: - """Apply session-specific visibility transforms to components. - - This helper applies session-level enable/disable rules by marking - components with their enabled state. Session transforms override - global transforms due to mark-based semantics (later marks win). - - Args: - components: The components to apply session transforms to. - - Returns: - The components with session transforms applied. - """ - current_ctx = _current_context.get() - if current_ctx is None: - return components - - session_transforms = await current_ctx._get_session_transforms() - if not session_transforms: - return components - - # Apply each transform's marking to each component - result = list(components) - for transform in session_transforms: - result = [transform._mark_component(c) for c in result] - return result diff --git a/src/fastmcp/server/sampling/run.py b/src/fastmcp/server/sampling/run.py index 70eccd47a..f5c73afef 100644 --- a/src/fastmcp/server/sampling/run.py +++ b/src/fastmcp/server/sampling/run.py @@ -3,8 +3,10 @@ from __future__ import annotations import inspect +import json +from collections.abc import Callable, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Generic +from typing import TYPE_CHECKING, Any, Generic, Literal, cast from mcp.types import ( ClientCapabilities, @@ -14,6 +16,7 @@ from mcp.types import ( ModelPreferences, SamplingCapability, SamplingMessage, + SamplingMessageContentBlock, SamplingToolsCapability, TextContent, ToolChoice, @@ -22,18 +25,25 @@ from mcp.types import ( ) from mcp.types import CreateMessageRequestParams as SamplingParams from mcp.types import Tool as SDKTool +from pydantic import ValidationError from typing_extensions import TypeVar +from fastmcp import settings from fastmcp.exceptions import ToolError from fastmcp.server.sampling.sampling_tool import SamplingTool +from fastmcp.utilities.json_schema import compress_schema from fastmcp.utilities.logging import get_logger +from fastmcp.utilities.types import get_cached_typeadapter logger = get_logger(__name__) if TYPE_CHECKING: from fastmcp.server.context import Context -ResultT = TypeVar("ResultT", default=str) +ResultT = TypeVar("ResultT") + +# Simplified tool choice type - just the mode string instead of the full MCP object +ToolChoiceOption = Literal["auto", "required", "none"] @dataclass @@ -299,3 +309,332 @@ async def execute_tools( ) return tool_results + + +# --- Helper functions for sampling --- + + +def prepare_messages( + messages: str | Sequence[str | SamplingMessage], +) -> list[SamplingMessage]: + """Convert various message formats to a list of SamplingMessage objects.""" + if isinstance(messages, str): + return [ + SamplingMessage( + content=TextContent(text=messages, type="text"), role="user" + ) + ] + else: + return [ + SamplingMessage(content=TextContent(text=m, type="text"), role="user") + if isinstance(m, str) + else m + for m in messages + ] + + +def prepare_tools( + tools: Sequence[SamplingTool | Callable[..., Any]] | None, +) -> list[SamplingTool] | None: + """Convert tools to SamplingTool objects.""" + if tools is None: + return None + + sampling_tools: list[SamplingTool] = [] + for t in tools: + if isinstance(t, SamplingTool): + sampling_tools.append(t) + elif callable(t): + sampling_tools.append(SamplingTool.from_function(t)) + else: + raise TypeError(f"Expected SamplingTool or callable, got {type(t)}") + + return sampling_tools if sampling_tools else None + + +def extract_tool_calls( + response: CreateMessageResult | CreateMessageResultWithTools, +) -> list[ToolUseContent]: + """Extract tool calls from a response.""" + content = response.content + if isinstance(content, list): + return [c for c in content if isinstance(c, ToolUseContent)] + elif isinstance(content, ToolUseContent): + return [content] + return [] + + +def create_final_response_tool(result_type: type) -> SamplingTool: + """Create a synthetic 'final_response' tool for structured output. + + This tool is used to capture structured responses from the LLM. + The tool's schema is derived from the result_type. + """ + type_adapter = get_cached_typeadapter(result_type) + schema = type_adapter.json_schema() + schema = compress_schema(schema, prune_titles=True) + + # Tool parameters must be object-shaped. Wrap primitives in {"value": } + if schema.get("type") != "object": + schema = { + "type": "object", + "properties": {"value": schema}, + "required": ["value"], + } + + # The fn just returns the input as-is (validation happens in the loop) + def final_response(**kwargs: Any) -> dict[str, Any]: + return kwargs + + return SamplingTool( + name="final_response", + description=( + "Call this tool to provide your final response. " + "Use this when you have completed the task and are ready to return the result." + ), + parameters=schema, + fn=final_response, + ) + + +# --- Implementation functions for Context methods --- + + +async def sample_step_impl( + context: Context, + messages: str | Sequence[str | SamplingMessage], + *, + system_prompt: str | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + model_preferences: ModelPreferences | str | list[str] | None = None, + tools: Sequence[SamplingTool | Callable[..., Any]] | None = None, + tool_choice: ToolChoiceOption | str | None = None, + auto_execute_tools: bool = True, + mask_error_details: bool | None = None, +) -> SampleStep: + """Implementation of Context.sample_step(). + + Make a single LLM sampling call. This is a stateless function that makes + exactly one LLM call and optionally executes any requested tools. + """ + # Convert messages to SamplingMessage objects + current_messages = prepare_messages(messages) + + # Convert tools to SamplingTools + sampling_tools = prepare_tools(tools) + sdk_tools: list[SDKTool] | None = ( + [t._to_sdk_tool() for t in sampling_tools] if sampling_tools else None + ) + tool_map: dict[str, SamplingTool] = ( + {t.name: t for t in sampling_tools} if sampling_tools else {} + ) + + # Determine whether to use fallback handler or client + use_fallback = determine_handler_mode(context, bool(sampling_tools)) + + # Build tool choice + effective_tool_choice: ToolChoice | None = None + if tool_choice is not None: + if tool_choice not in ("auto", "required", "none"): + raise ValueError( + f"Invalid tool_choice: {tool_choice!r}. " + "Must be 'auto', 'required', or 'none'." + ) + effective_tool_choice = ToolChoice( + mode=cast(Literal["auto", "required", "none"], tool_choice) + ) + + # Effective max_tokens + effective_max_tokens = max_tokens if max_tokens is not None else 512 + + # Make the LLM call + if use_fallback: + response = await call_sampling_handler( + context, + current_messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=effective_max_tokens, + model_preferences=model_preferences, + sdk_tools=sdk_tools, + tool_choice=effective_tool_choice, + ) + else: + response = await context.session.create_message( + messages=current_messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=effective_max_tokens, + model_preferences=_parse_model_preferences(model_preferences), + tools=sdk_tools, + tool_choice=effective_tool_choice, + related_request_id=context.request_id, + ) + + # Check if this is a tool use response + is_tool_use_response = ( + isinstance(response, CreateMessageResultWithTools) + and response.stopReason == "toolUse" + ) + + # Always include the assistant response in history + current_messages.append(SamplingMessage(role="assistant", content=response.content)) + + # If not a tool use, return immediately + if not is_tool_use_response: + return SampleStep(response=response, history=current_messages) + + # If not executing tools, return with assistant message but no tool results + if not auto_execute_tools: + return SampleStep(response=response, history=current_messages) + + # Execute tools and add results to history + step_tool_calls = extract_tool_calls(response) + if step_tool_calls: + effective_mask = ( + mask_error_details + if mask_error_details is not None + else settings.mask_error_details + ) + tool_results: list[ToolResultContent] = await execute_tools( + step_tool_calls, tool_map, mask_error_details=effective_mask + ) + + if tool_results: + current_messages.append( + SamplingMessage( + role="user", + content=cast(list[SamplingMessageContentBlock], tool_results), + ) + ) + + return SampleStep(response=response, history=current_messages) + + +async def sample_impl( + context: Context, + messages: str | Sequence[str | SamplingMessage], + *, + system_prompt: str | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + model_preferences: ModelPreferences | str | list[str] | None = None, + tools: Sequence[SamplingTool | Callable[..., Any]] | None = None, + result_type: type[ResultT] | None = None, + mask_error_details: bool | None = None, +) -> SamplingResult[ResultT]: + """Implementation of Context.sample(). + + Send a sampling request to the client and await the response. This method + runs to completion automatically, executing a tool loop until the LLM + provides a final text response. + """ + # Safety limit to prevent infinite loops + max_iterations = 100 + + # Convert tools to SamplingTools + sampling_tools = prepare_tools(tools) + + # Handle structured output with result_type + tool_choice: str | None = None + if result_type is not None and result_type is not str: + final_response_tool = create_final_response_tool(result_type) + sampling_tools = list(sampling_tools) if sampling_tools else [] + sampling_tools.append(final_response_tool) + + # Always require tool calls when result_type is set - the LLM must + # eventually call final_response (text responses are not accepted) + tool_choice = "required" + + # Convert messages for the loop + current_messages: str | Sequence[str | SamplingMessage] = messages + + for _iteration in range(max_iterations): + step = await sample_step_impl( + context, + messages=current_messages, + system_prompt=system_prompt, + temperature=temperature, + max_tokens=max_tokens, + model_preferences=model_preferences, + tools=sampling_tools, + tool_choice=tool_choice, + mask_error_details=mask_error_details, + ) + + # Check for final_response tool call for structured output + if result_type is not None and result_type is not str and step.is_tool_use: + for tool_call in step.tool_calls: + if tool_call.name == "final_response": + # Validate and return the structured result + type_adapter = get_cached_typeadapter(result_type) + + # Unwrap if we wrapped primitives (non-object schemas) + input_data = tool_call.input + original_schema = compress_schema( + type_adapter.json_schema(), prune_titles=True + ) + if ( + original_schema.get("type") != "object" + and isinstance(input_data, dict) + and "value" in input_data + ): + input_data = input_data["value"] + + try: + validated_result = type_adapter.validate_python(input_data) + text = json.dumps( + type_adapter.dump_python(validated_result, mode="json") + ) + return SamplingResult( + text=text, + result=validated_result, + history=step.history, + ) + except ValidationError as e: + # Validation failed - add error as tool result + step.history.append( + SamplingMessage( + role="user", + content=[ + ToolResultContent( + type="tool_result", + toolUseId=tool_call.id, + content=[ + TextContent( + type="text", + text=( + f"Validation error: {e}. " + "Please try again with valid data." + ), + ) + ], + isError=True, + ) + ], + ) + ) + + # If not a tool use response, we're done + if not step.is_tool_use: + # For structured output, the LLM must use the final_response tool + if result_type is not None and result_type is not str: + raise RuntimeError( + f"Expected structured output of type {result_type.__name__}, " + "but the LLM returned a text response instead of calling " + "the final_response tool." + ) + return SamplingResult( + text=step.text, + result=cast(ResultT, step.text if step.text else ""), + history=step.history, + ) + + # Continue with the updated history + current_messages = step.history + + # After first iteration, reset tool_choice to auto + tool_choice = None + + raise RuntimeError(f"Sampling exceeded maximum iterations ({max_iterations})") diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index a2e500198..83eb94b3a 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -88,7 +88,7 @@ from fastmcp.server.transforms import ( ToolTransform, Transform, ) -from fastmcp.server.transforms.enabled import is_enabled +from fastmcp.server.transforms.enabled import apply_session_transforms, is_enabled from fastmcp.settings import DuplicateBehavior as DuplicateBehaviorSetting from fastmcp.settings import Settings from fastmcp.tools.function_tool import FunctionTool @@ -1027,8 +1027,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): ) # Get all tools, apply session transforms, then filter enabled - from fastmcp.server.context import apply_session_transforms - tools = list(await super().list_tools()) tools = await apply_session_transforms(tools) tools = [t for t in tools if is_enabled(t)] @@ -1098,8 +1096,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): return None # Apply session transforms to single item - from fastmcp.server.context import apply_session_transforms - tools = await apply_session_transforms([tool]) if not tools or not is_enabled(tools[0]): return None @@ -1129,8 +1125,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): ) # Get all resources, apply session transforms, then filter enabled - from fastmcp.server.context import apply_session_transforms - resources = list(await super().list_resources()) resources = await apply_session_transforms(resources) resources = [r for r in resources if is_enabled(r)] @@ -1199,8 +1193,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): return None # Apply session transforms to single item - from fastmcp.server.context import apply_session_transforms - resources = await apply_session_transforms([resource]) if not resources or not is_enabled(resources[0]): return None @@ -1232,8 +1224,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): ) # Get all templates, apply session transforms, then filter enabled - from fastmcp.server.context import apply_session_transforms - templates = list(await super().list_resource_templates()) templates = await apply_session_transforms(templates) templates = [t for t in templates if is_enabled(t)] @@ -1302,8 +1292,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): return None # Apply session transforms to single item - from fastmcp.server.context import apply_session_transforms - templates = await apply_session_transforms([template]) if not templates or not is_enabled(templates[0]): return None @@ -1331,8 +1319,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): ) # Get all prompts, apply session transforms, then filter enabled - from fastmcp.server.context import apply_session_transforms - prompts = list(await super().list_prompts()) prompts = await apply_session_transforms(prompts) prompts = [p for p in prompts if is_enabled(p)] @@ -1401,8 +1387,6 @@ class FastMCP(AggregateProvider, Generic[LifespanResultT]): return None # Apply session transforms to single item - from fastmcp.server.context import apply_session_transforms - prompts = await apply_session_transforms([prompt]) if not prompts or not is_enabled(prompts[0]): return None diff --git a/src/fastmcp/server/transforms/enabled.py b/src/fastmcp/server/transforms/enabled.py index cd9f023e6..6b855e806 100644 --- a/src/fastmcp/server/transforms/enabled.py +++ b/src/fastmcp/server/transforms/enabled.py @@ -8,7 +8,9 @@ Final filtering happens at the Provider level. from __future__ import annotations from collections.abc import Sequence -from typing import TYPE_CHECKING, Literal, TypeVar +from typing import TYPE_CHECKING, Any, Literal, TypeVar + +import mcp.types from fastmcp.resources.resource import Resource from fastmcp.resources.template import ResourceTemplate @@ -27,6 +29,7 @@ from fastmcp.utilities.versions import VersionSpec if TYPE_CHECKING: from fastmcp.prompts.prompt import Prompt + from fastmcp.server.context import Context from fastmcp.tools.tool import Tool from fastmcp.utilities.components import FastMCPComponent @@ -290,3 +293,240 @@ def is_enabled(component: FastMCPComponent) -> bool: fastmcp = meta.get(_FASTMCP_KEY, {}) internal = fastmcp.get(_INTERNAL_KEY, {}) return internal.get("enabled", True) # Default True if not set + + +# ------------------------------------------------------------------------- +# Session visibility control +# ------------------------------------------------------------------------- + +if TYPE_CHECKING: + from fastmcp.server.context import Context + + +async def get_visibility_rules(context: Context) -> list[dict[str, Any]]: + """Load visibility rule dicts from session state.""" + return await context.get_state("_visibility_rules") or [] + + +async def save_visibility_rules( + context: Context, + rules: list[dict[str, Any]], + *, + components: set[Literal["tool", "resource", "template", "prompt"]] | None = None, +) -> None: + """Save visibility rule dicts to session state and send notifications. + + Args: + context: The context to save rules for. + rules: The visibility rules to save. + components: Optional hint about which component types are affected. + If None, sends notifications for all types (safe default). + If provided, only sends notifications for specified types. + """ + await context.set_state("_visibility_rules", rules) + + # Send notifications based on components hint + # Note: MCP has no separate template notification - templates use ResourceListChangedNotification + if components is None or "tool" in components: + await context.send_notification(mcp.types.ToolListChangedNotification()) + if components is None or "resource" in components or "template" in components: + await context.send_notification(mcp.types.ResourceListChangedNotification()) + if components is None or "prompt" in components: + await context.send_notification(mcp.types.PromptListChangedNotification()) + + +def create_enabled_transforms(rules: list[dict[str, Any]]) -> list[Enabled]: + """Convert rule dicts to Enabled transforms.""" + transforms = [] + for params in rules: + version = None + if params.get("version"): + version_dict = params["version"] + version = VersionSpec( + gte=version_dict.get("gte"), + lt=version_dict.get("lt"), + eq=version_dict.get("eq"), + ) + transforms.append( + Enabled( + params["enabled"], + names=set(params["names"]) if params.get("names") else None, + keys=set(params["keys"]) if params.get("keys") else None, + version=version, + tags=set(params["tags"]) if params.get("tags") else None, + components=( + set(params["components"]) if params.get("components") else None + ), + match_all=params.get("match_all", False), + ) + ) + return transforms + + +async def get_session_transforms(context: Context) -> list[Enabled]: + """Get session-specific Enabled transforms from state store.""" + try: + # Will raise RuntimeError if no session available + _ = context.session_id + except RuntimeError: + return [] + + rules = await get_visibility_rules(context) + return create_enabled_transforms(rules) + + +async def enable_components( + context: Context, + *, + names: set[str] | None = None, + keys: set[str] | None = None, + version: VersionSpec | None = None, + tags: set[str] | None = None, + components: set[Literal["tool", "resource", "template", "prompt"]] | None = None, + match_all: bool = False, +) -> None: + """Enable components matching criteria for this session only. + + Session rules override global transforms. Rules accumulate - each call + adds a new rule to the session. Later marks override earlier ones + (Enabled transform semantics). + + Sends notifications to this session only: ToolListChangedNotification, + ResourceListChangedNotification, and PromptListChangedNotification. + + Args: + context: The context for this session. + names: Component names or URIs to match. + keys: Component keys to match (e.g., {"tool:my_tool@v1"}). + version: Component version spec to match. + tags: Tags to match (component must have at least one). + components: Component types to match (e.g., {"tool", "prompt"}). + match_all: If True, matches all components regardless of other criteria. + """ + # Normalize empty sets to None (empty = match all) + components = components if components else None + + # Load current rules + rules = await get_visibility_rules(context) + + # Create new rule dict + rule: dict[str, Any] = { + "enabled": True, + "names": list(names) if names else None, + "keys": list(keys) if keys else None, + "version": ( + {"gte": version.gte, "lt": version.lt, "eq": version.eq} + if version + else None + ), + "tags": list(tags) if tags else None, + "components": list(components) if components else None, + "match_all": match_all, + } + + # Add and save (notifications sent by save_visibility_rules) + rules.append(rule) + await save_visibility_rules(context, rules, components=components) + + +async def disable_components( + context: Context, + *, + names: set[str] | None = None, + keys: set[str] | None = None, + version: VersionSpec | None = None, + tags: set[str] | None = None, + components: set[Literal["tool", "resource", "template", "prompt"]] | None = None, + match_all: bool = False, +) -> None: + """Disable components matching criteria for this session only. + + Session rules override global transforms. Rules accumulate - each call + adds a new rule to the session. Later marks override earlier ones + (Enabled transform semantics). + + Sends notifications to this session only: ToolListChangedNotification, + ResourceListChangedNotification, and PromptListChangedNotification. + + Args: + context: The context for this session. + names: Component names or URIs to match. + keys: Component keys to match (e.g., {"tool:my_tool@v1"}). + version: Component version spec to match. + tags: Tags to match (component must have at least one). + components: Component types to match (e.g., {"tool", "prompt"}). + match_all: If True, matches all components regardless of other criteria. + """ + # Normalize empty sets to None (empty = match all) + components = components if components else None + + # Load current rules + rules = await get_visibility_rules(context) + + # Create new rule dict + rule: dict[str, Any] = { + "enabled": False, + "names": list(names) if names else None, + "keys": list(keys) if keys else None, + "version": ( + {"gte": version.gte, "lt": version.lt, "eq": version.eq} + if version + else None + ), + "tags": list(tags) if tags else None, + "components": list(components) if components else None, + "match_all": match_all, + } + + # Add and save (notifications sent by save_visibility_rules) + rules.append(rule) + await save_visibility_rules(context, rules, components=components) + + +async def reset_components(context: Context) -> None: + """Clear all session visibility rules. + + Use this to reset session visibility back to global defaults. + + Sends notifications to this session only: ToolListChangedNotification, + ResourceListChangedNotification, and PromptListChangedNotification. + + Args: + context: The context for this session. + """ + await save_visibility_rules(context, []) + + +ComponentT = TypeVar("ComponentT", bound="FastMCPComponent") + + +async def apply_session_transforms( + components: Sequence[ComponentT], +) -> Sequence[ComponentT]: + """Apply session-specific visibility transforms to components. + + This helper applies session-level enable/disable rules by marking + components with their enabled state. Session transforms override + global transforms due to mark-based semantics (later marks win). + + Args: + components: The components to apply session transforms to. + + Returns: + The components with session transforms applied. + """ + from fastmcp.server.context import _current_context + + current_ctx = _current_context.get() + if current_ctx is None: + return components + + session_transforms = await get_session_transforms(current_ctx) + if not session_transforms: + return components + + # Apply each transform's marking to each component + result = list(components) + for transform in session_transforms: + result = [transform._mark_component(c) for c in result] + return result diff --git a/tests/server/test_context.py b/tests/server/test_context.py index 89c55fdc3..373286d8a 100644 --- a/tests/server/test_context.py +++ b/tests/server/test_context.py @@ -6,10 +6,10 @@ from mcp.types import ModelPreferences from fastmcp.server.context import ( Context, - _parse_model_preferences, reset_transport, set_transport, ) +from fastmcp.server.sampling.run import _parse_model_preferences from fastmcp.server.server import FastMCP