fastmcp 2.14.4__py3-none-any.whl → 3.0.0b1__py3-none-any.whl

fastmcp/server/context.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import copy
-import json
 import logging
 import weakref
 from collections.abc import Callable, Generator, Mapping, Sequence
@@ -9,35 +7,25 @@ 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 anyio
+import mcp.types
 from mcp import LoggingLevel, ServerSession
-from mcp.server.lowlevel.helper_types import ReadResourceContents
 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,
     CancelledElicitation,
@@ -47,17 +35,30 @@ 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,
+    sample_impl,
+    sample_step_impl,
 )
-from fastmcp.server.sampling.run import (
-    execute_tools as run_sampling_tools,
+from fastmcp.server.server import FastMCP, StateValue
+from fastmcp.server.transforms.visibility import (
+    Visibility,
+)
+from fastmcp.server.transforms.visibility import (
+    disable_components as _disable_components,
+)
+from fastmcp.server.transforms.visibility import (
+    enable_components as _enable_components,
+)
+from fastmcp.server.transforms.visibility import (
+    get_session_transforms as _get_session_transforms,
+)
+from fastmcp.server.transforms.visibility import (
+    get_visibility_rules as _get_visibility_rules,
+)
+from fastmcp.server.transforms.visibility import (
+    reset_visibility as _reset_visibility,
 )
-from fastmcp.server.server import FastMCP
-from fastmcp.utilities.json_schema import compress_schema
 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__)
 to_client_logger: Logger = logger.getChild(suffix="to_client")
@@ -71,13 +72,27 @@ _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)
+
+TransportType = Literal["stdio", "sse", "streamable-http"]
+_current_transport: ContextVar[TransportType | None] = ContextVar(
+    "transport", default=None
+)
+
 
-_current_context: ContextVar[Context | None] = ContextVar("context", default=None)  # type: ignore[assignment]
+def set_transport(
+    transport: TransportType,
+) -> Token[TransportType | None]:
+    """Set the current transport type. Returns token for reset."""
+    return _current_transport.set(transport)
 
 
-_flush_lock = anyio.Lock()
+def reset_transport(token: Token[TransportType | None]) -> None:
+    """Reset transport to previous value."""
+    _current_transport.reset(token)
 
 
 @dataclass
@@ -141,29 +156,35 @@ class Context:
         request_id = ctx.request_id
         client_id = ctx.client_id
 
-        # Manage state across the request
-        ctx.set_state("key", "value")
-        value = ctx.get_state("key")
+        # Manage state across the session (persists across requests)
+        await ctx.set_state("key", "value")
+        value = await ctx.get_state("key")
 
         return str(x)
     ```
 
     State Management:
-    Context objects maintain a state dictionary that can be used to store and share
-    data across middleware and tool calls within a request. When a new context
-    is created (nested contexts), it inherits a copy of its parent's state, ensuring
-    that modifications in child contexts don't affect parent contexts.
+    Context provides session-scoped state that persists across requests within
+    the same MCP session. State is automatically keyed by session, ensuring
+    isolation between different clients.
+
+    State set during `on_initialize` middleware will persist to subsequent tool
+    calls when using the same session object (STDIO, SSE, single-server HTTP).
+    For distributed/serverless HTTP deployments where different machines handle
+    the init and tool calls, state is isolated by the mcp-session-id header.
 
     The context parameter name can be anything as long as it's annotated with Context.
     The context is optional - tools that don't need it can omit the parameter.
 
     """
 
-    def __init__(self, fastmcp: FastMCP):
+    # Default TTL for session state: 1 day in seconds
+    _STATE_TTL_SECONDS: int = 86400
+
+    def __init__(self, fastmcp: FastMCP, session: ServerSession | None = None):
         self._fastmcp: weakref.ref[FastMCP] = weakref.ref(fastmcp)
+        self._session: ServerSession | None = session  # For state ops during init
         self._tokens: list[Token] = []
-        self._notification_queue: set[str] = set()  # Dedupe notifications
-        self._state: dict[str, Any] = {}
 
     @property
     def fastmcp(self) -> FastMCP:
@@ -175,11 +196,6 @@ class Context:
 
     async def __aenter__(self) -> Context:
         """Enter the context manager and set this context as the current context."""
-        parent_context = _current_context.get(None)
-        if parent_context is not None:
-            # Inherit state from parent context
-            self._state = copy.deepcopy(parent_context._state)
-
         # Always set this context and save the token
         token = _current_context.set(self)
         self._tokens.append(token)
@@ -194,8 +210,8 @@ class Context:
         self._server_token = _current_server.set(weakref.ref(self.fastmcp))
 
         # Set docket/worker from server instance for this request's context.
-        # This ensures ContextVars work even in environments (like Lambda) where
-        # lifespan ContextVars don't propagate to request handlers.
+        # This ensures ContextVars work even in ASGI environments (Lambda, FastAPI mount)
+        # where lifespan ContextVars don't propagate to request handlers.
         server = self.fastmcp
         if server._docket is not None:
             self._docket_token = _current_docket.set(server._docket)
@@ -207,9 +223,6 @@ class Context:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
         """Exit the context manager and reset the most recent token."""
-        # Flush any remaining notifications before exiting
-        await self._flush_notifications()
-
         # Reset server/docket/worker tokens
         from fastmcp.server.dependencies import (
             _current_docket,
@@ -261,6 +274,29 @@ class Context:
         except LookupError:
             return None
 
+    @property
+    def lifespan_context(self) -> dict[str, Any]:
+        """Access the server's lifespan context.
+
+        Returns the context dict yielded by the server's lifespan function.
+        Returns an empty dict if no lifespan was configured or if the MCP
+        session is not yet established.
+
+        Example:
+        ```python
+        @server.tool
+        def my_tool(ctx: Context) -> str:
+            db = ctx.lifespan_context.get("db")
+            if db:
+                return db.query("SELECT 1")
+            return "No database connection"
+        ```
+        """
+        rc = self.request_context
+        if rc is None:
+            return {}
+        return rc.lifespan_context
+
     async def report_progress(
         self, progress: float, total: float | None = None, message: str | None = None
     ) -> None:
@@ -288,13 +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
         """
-        return await self.fastmcp._list_resources_mcp()
+        return await self._paginate_list(
+            request_factory=lambda cursor: mcp.types.ListResourcesRequest(
+                params=mcp.types.PaginatedRequestParams(cursor=cursor)
+                if cursor
+                else None
+            ),
+            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.
@@ -302,7 +373,15 @@ class Context:
         Returns:
             List of Prompt objects available on the server
         """
-        return await self.fastmcp._list_prompts_mcp()
+        return await self._paginate_list(
+            request_factory=lambda cursor: mcp.types.ListPromptsRequest(
+                params=mcp.types.PaginatedRequestParams(cursor=cursor)
+                if cursor
+                else None
+            ),
+            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
@@ -316,19 +395,28 @@ class Context:
         Returns:
             The prompt result
         """
-        return await self.fastmcp._get_prompt_mcp(name, arguments)
+        result = await self.fastmcp.render_prompt(name, arguments)
+        if isinstance(result, mcp.types.CreateTaskResult):
+            raise RuntimeError(
+                "Unexpected CreateTaskResult: Context calls should not have task metadata"
+            )
+        return result.to_mcp_prompt_result()
 
-    async def read_resource(self, uri: str | AnyUrl) -> list[ReadResourceContents]:
+    async def read_resource(self, uri: str | AnyUrl) -> ResourceResult:
         """Read a resource by URI.
 
         Args:
             uri: Resource URI to read
 
         Returns:
-            The resource content as either text or bytes
+            ResourceResult with contents
         """
-        # Context calls don't have task metadata, so always returns list
-        return await self.fastmcp._read_resource_mcp(uri)  # type: ignore[return-value]
+        result = await self.fastmcp.read_resource(str(uri))
+        if isinstance(result, mcp.types.CreateTaskResult):
+            raise RuntimeError(
+                "Unexpected CreateTaskResult: Context calls should not have task metadata"
+            )
+        return result
 
     async def log(
         self,
@@ -358,6 +446,15 @@ class Context:
             related_request_id=self.request_id,
         )
 
+    @property
+    def transport(self) -> TransportType | None:
+        """Get the current transport type.
+
+        Returns the transport type used to run this server: "stdio", "sse",
+        or "streamable-http". Returns None if called outside of a server context.
+        """
+        return _current_transport.get()
+
     @property
     def client_id(self) -> str | None:
         """Get the client ID if available."""
@@ -393,7 +490,7 @@ class Context:
             for other transports.
 
         Raises:
-            RuntimeError if MCP request context is not available.
+            RuntimeError if no session is available.
 
         Example:
             ```python
@@ -404,32 +501,37 @@ class Context:
                 return f"Data stored for session {session_id}"
             ```
         """
+        from uuid import uuid4
+
+        # Get session from request context or _session (for on_initialize)
         request_ctx = self.request_context
-        if request_ctx is None:
+        if request_ctx is not None:
+            session = request_ctx.session
+        elif self._session is not None:
+            session = self._session
+        else:
             raise RuntimeError(
-                "session_id is not available because the MCP session has not been established yet. "
-                "Check `context.request_context` for None before accessing this attribute."
+                "session_id is not available because no session exists. "
+                "This typically means you're outside a request context."
             )
-        session = request_ctx.session
 
-        # Try to get the session ID from the session attributes
-        session_id = getattr(session, "_fastmcp_id", None)
+        # Check for cached session ID
+        session_id = getattr(session, "_fastmcp_state_prefix", None)
         if session_id is not None:
             return session_id
 
-        # Try to get the session ID from the http request headers
-        request = request_ctx.request
-        if request:
-            session_id = request.headers.get("mcp-session-id")
+        # For HTTP, try to get from header
+        if request_ctx is not None:
+            request = request_ctx.request
+            if request:
+                session_id = request.headers.get("mcp-session-id")
 
-        # Generate a session ID if it doesn't exist.
+        # For STDIO/SSE/in-memory, generate a UUID
         if session_id is None:
-            from uuid import uuid4
-
             session_id = str(uuid4())
 
-        # Save the session id to the session attributes
-        session._fastmcp_id = session_id  # type: ignore[attr-defined]
+        # Cache on session for consistency
+        session._fastmcp_state_prefix = session_id  # type: ignore[attr-defined]
         return session_id
 
     @property
@@ -515,17 +617,15 @@ class Context:
         result = await self.session.list_roots()
         return result.roots
 
-    async def send_tool_list_changed(self) -> None:
-        """Send a tool list changed notification to the client."""
-        await self.session.send_tool_list_changed()
-
-    async def send_resource_list_changed(self) -> None:
-        """Send a resource list changed notification to the client."""
-        await self.session.send_resource_list_changed()
+    async def send_notification(
+        self, notification: mcp.types.ServerNotificationType
+    ) -> None:
+        """Send a notification to the client immediately.
 
-    async def send_prompt_list_changed(self) -> None:
-        """Send a prompt list changed notification to the client."""
-        await self.session.send_prompt_list_changed()
+        Args:
+            notification: An MCP notification instance (e.g., ToolListChangedNotification())
+        """
+        await self.session.send_notification(mcp.types.ServerNotification(notification))
 
     async def close_sse_stream(self) -> None:
         """Close the current response stream to trigger client reconnection.
@@ -623,101 +723,19 @@ 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
-        )
-        tool_map: dict[str, SamplingTool] = (
-            {t.name: t for t in sampling_tools} if sampling_tools else {}
+        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,
         )
 
-        # 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 = await run_sampling_tools(
-                step_tool_calls, tool_map, mask_error_details=effective_mask
-            )
-
-            if tool_results:
-                current_messages.append(
-                    SamplingMessage(
-                        role="user",
-                        content=tool_results,  # type: ignore[arg-type]
-                    )
-                )
-
-        return SampleStep(response=response, history=current_messages)
-
     @overload
     async def sample(
         self,
@@ -796,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,
-                                        )
-                                    ],  # type: ignore[arg-type]
-                                )
-                            )
-
-            # 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(
@@ -936,10 +858,50 @@ class Context:
     """When response_type is a list of strings, the accepted elicitation will
     contain the selected string response"""
 
+    @overload
     async def elicit(
         self,
         message: str,
-        response_type: type[T] | list[str] | dict[str, dict[str, str]] | None = None,
+        response_type: dict[str, dict[str, str]],
+    ) -> AcceptedElicitation[str] | DeclinedElicitation | CancelledElicitation: ...
+
+    """When response_type is a dict mapping keys to title dicts, the accepted
+    elicitation will contain the selected key"""
+
+    @overload
+    async def elicit(
+        self,
+        message: str,
+        response_type: list[list[str]],
+    ) -> (
+        AcceptedElicitation[list[str]] | DeclinedElicitation | CancelledElicitation
+    ): ...
+
+    """When response_type is a list containing a list of strings (multi-select),
+    the accepted elicitation will contain a list of selected strings"""
+
+    @overload
+    async def elicit(
+        self,
+        message: str,
+        response_type: list[dict[str, dict[str, str]]],
+    ) -> (
+        AcceptedElicitation[list[str]] | DeclinedElicitation | CancelledElicitation
+    ): ...
+
+    """When response_type is a list containing a dict mapping keys to title dicts
+    (multi-select with titles), the accepted elicitation will contain a list of
+    selected keys"""
+
+    async def elicit(
+        self,
+        message: str,
+        response_type: type[T]
+        | list[str]
+        | dict[str, dict[str, str]]
+        | list[list[str]]
+        | list[dict[str, dict[str, str]]]
+        | None = None,
     ) -> (
         AcceptedElicitation[T]
         | AcceptedElicitation[dict[str, Any]]
@@ -988,43 +950,135 @@ class Context:
         else:
             raise ValueError(f"Unexpected elicitation action: {result.action}")
 
-    def set_state(self, key: str, value: Any) -> None:
-        """Set a value in the context state."""
-        self._state[key] = value
-
-    def get_state(self, key: str) -> Any:
-        """Get a value from the context state. Returns None if the key is not found."""
-        return self._state.get(key)
-
-    def _queue_tool_list_changed(self) -> None:
-        """Queue a tool list changed notification."""
-        self._notification_queue.add("notifications/tools/list_changed")
-
-    def _queue_resource_list_changed(self) -> None:
-        """Queue a resource list changed notification."""
-        self._notification_queue.add("notifications/resources/list_changed")
-
-    def _queue_prompt_list_changed(self) -> None:
-        """Queue a prompt list changed notification."""
-        self._notification_queue.add("notifications/prompts/list_changed")
-
-    async def _flush_notifications(self) -> None:
-        """Send all queued notifications."""
-        async with _flush_lock:
-            if not self._notification_queue:
-                return
-
-            try:
-                if "notifications/tools/list_changed" in self._notification_queue:
-                    await self.session.send_tool_list_changed()
-                if "notifications/resources/list_changed" in self._notification_queue:
-                    await self.session.send_resource_list_changed()
-                if "notifications/prompts/list_changed" in self._notification_queue:
-                    await self.session.send_prompt_list_changed()
-                self._notification_queue.clear()
-            except Exception:
-                # Don't let notification failures break the request
-                pass
+    def _make_state_key(self, key: str) -> str:
+        """Create session-prefixed key for state storage."""
+        return f"{self.session_id}:{key}"
+
+    async def set_state(self, key: str, value: Any) -> None:
+        """Set a value in the session-scoped state store.
+
+        Values persist across requests within the same MCP session.
+        The key is automatically prefixed with the session identifier.
+        State expires after 1 day to prevent unbounded memory growth.
+        """
+        prefixed_key = self._make_state_key(key)
+        await self.fastmcp._state_store.put(
+            key=prefixed_key,
+            value=StateValue(value=value),
+            ttl=self._STATE_TTL_SECONDS,
+        )
+
+    async def get_state(self, key: str) -> Any:
+        """Get a value from the session-scoped state store.
+
+        Returns None if the key is not found.
+        """
+        prefixed_key = self._make_state_key(key)
+        result = await self.fastmcp._state_store.get(key=prefixed_key)
+        return result.value if result is not None else None
+
+    async def delete_state(self, key: str) -> None:
+        """Delete a value from the session-scoped state store."""
+        prefixed_key = self._make_state_key(key)
+        await self.fastmcp._state_store.delete(key=prefixed_key)
+
+    # -------------------------------------------------------------------------
+    # Session visibility control
+    # -------------------------------------------------------------------------
+
+    async def _get_visibility_rules(self) -> list[dict[str, Any]]:
+        """Load visibility rule dicts from session state."""
+        return await _get_visibility_rules(self)
+
+    async def _get_session_transforms(self) -> list[Visibility]:
+        """Get session-specific Visibility transforms from state store."""
+        return await _get_session_transforms(self)
+
+    async def enable_components(
+        self,
+        *,
+        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
+        (Visibility transform semantics).
+
+        Sends notifications to this session only: ToolListChangedNotification,
+        ResourceListChangedNotification, and PromptListChangedNotification.
+
+        Args:
+            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.
+        """
+        await _enable_components(
+            self,
+            names=names,
+            keys=keys,
+            version=version,
+            tags=tags,
+            components=components,
+            match_all=match_all,
+        )
+
+    async def disable_components(
+        self,
+        *,
+        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
+        (Visibility transform semantics).
+
+        Sends notifications to this session only: ToolListChangedNotification,
+        ResourceListChangedNotification, and PromptListChangedNotification.
+
+        Args:
+            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.
+        """
+        await _disable_components(
+            self,
+            names=names,
+            keys=keys,
+            version=version,
+            tags=tags,
+            components=components,
+            match_all=match_all,
+        )
+
+    async def reset_visibility(self) -> 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.
+        """
+        await _reset_visibility(self)
 
 
 async def _log_to_server_and_client(
@@ -1053,104 +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": <schema>}
-    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 []