From ca060014bbfeaf262bb7d16bd1f12302c5aafb81 Mon Sep 17 00:00:00 2001
From: zzstoatzz <thrast36@gmail.com>
Date: Mon, 3 Feb 2025 14:27:07 -0600
Subject: [PATCH] centralize type

---
 src/mcp/server/fastmcp/server.py | 18 ++++++++----------
 src/mcp/types.py                 |  3 ++-
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
index d8d5416..aa7c79b 100644
--- a/src/mcp/server/fastmcp/server.py
+++ b/src/mcp/server/fastmcp/server.py
@@ -12,7 +12,6 @@
 from pydantic import BaseModel, Field
 from pydantic.networks import AnyUrl
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from typing_extensions import TypeAlias
 
 from mcp.server.fastmcp.exceptions import ResourceError
 from mcp.server.fastmcp.prompts import Prompt, PromptManager
@@ -26,6 +25,7 @@
 from mcp.server.stdio import stdio_server
 from mcp.shared.context import RequestContext
 from mcp.types import (
+    AnyFunction,
     EmbeddedResource,
     GetPromptResult,
     ImageContent,
@@ -49,8 +49,6 @@
 
 logger = get_logger(__name__)
 
-_Function: TypeAlias = Callable[..., Any]
-
 
 class Settings(BaseSettings):
     """FastMCP server settings.
@@ -217,7 +215,7 @@ async def read_resource(self, uri: AnyUrl | str) -> ReadResourceContents:
 
     def add_tool(
         self,
-        fn: _Function,
+        fn: AnyFunction,
         name: str | None = None,
         description: str | None = None,
     ) -> None:
@@ -235,7 +233,7 @@ def add_tool(
 
     def tool(
         self, name: str | None = None, description: str | None = None
-    ) -> Callable[[_Function], _Function]:
+    ) -> Callable[[AnyFunction], AnyFunction]:
         """Decorator to register a tool.
 
         Tools can optionally request a Context object by adding a parameter with the
@@ -268,7 +266,7 @@ async def async_tool(x: int, context: Context) -> str:
                 "Did you forget to call it? Use @tool() instead of @tool"
             )
 
-        def decorator(fn: _Function) -> _Function:
+        def decorator(fn: AnyFunction) -> AnyFunction:
             self.add_tool(fn, name=name, description=description)
             return fn
 
@@ -289,7 +287,7 @@ def resource(
         name: str | None = None,
         description: str | None = None,
         mime_type: str | None = None,
-    ) -> Callable[[_Function], _Function]:
+    ) -> Callable[[AnyFunction], AnyFunction]:
         """Decorator to register a function as a resource.
 
         The function will be called when the resource is read to generate its content.
@@ -333,7 +331,7 @@ async def get_weather(city: str) -> str:
                 "Did you forget to call it? Use @resource('uri') instead of @resource"
             )
 
-        def decorator(fn: _Function) -> _Function:
+        def decorator(fn: AnyFunction) -> AnyFunction:
             # Check if this should be a template
             has_uri_params = "{" in uri and "}" in uri
             has_func_params = bool(inspect.signature(fn).parameters)
@@ -381,7 +379,7 @@ def add_prompt(self, prompt: Prompt) -> None:
 
     def prompt(
         self, name: str | None = None, description: str | None = None
-    ) -> Callable[[_Function], _Function]:
+    ) -> Callable[[AnyFunction], AnyFunction]:
         """Decorator to register a prompt.
 
         Args:
@@ -422,7 +420,7 @@ async def analyze_file(path: str) -> list[Message]:
                 "Did you forget to call it? Use @prompt() instead of @prompt"
             )
 
-        def decorator(func: _Function) -> _Function:
+        def decorator(func: AnyFunction) -> AnyFunction:
             prompt = Prompt.from_function(func, name=name, description=description)
             self.add_prompt(prompt)
             return func
diff --git a/src/mcp/types.py b/src/mcp/types.py
index a4fa254..d1157aa 100644
--- a/src/mcp/types.py
+++ b/src/mcp/types.py
@@ -1,4 +1,4 @@
-from typing import Annotated, Any, Generic, Literal, TypeVar
+from typing import Annotated, Any, Callable, Generic, Literal, TypeAlias, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel
 from pydantic.networks import AnyUrl
@@ -27,6 +27,7 @@
 Cursor = str
 Role = Literal["user", "assistant"]
 RequestId = str | int
+AnyFunction: TypeAlias = Callable[..., Any]
 
 
 class RequestParams(BaseModel):