From 00a44692e74be46d24e2c0747d9b3211595a55d7 Mon Sep 17 00:00:00 2001
From: zzstoatzz <thrast36@gmail.com>
Date: Wed, 29 Jan 2025 23:01:07 -0600
Subject: [PATCH] init

---
 .gitignore                       |  3 +++
 src/mcp/server/fastmcp/server.py | 21 +++++++++++++--------
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index f27f895..bb25f9e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -162,3 +162,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# vscode
+.vscode/
\ No newline at end of file
diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
index 06b16e0..d8d5416 100644
--- a/src/mcp/server/fastmcp/server.py
+++ b/src/mcp/server/fastmcp/server.py
@@ -12,6 +12,7 @@ import uvicorn
 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
@@ -48,6 +49,8 @@ from mcp.types import (
 
 logger = get_logger(__name__)
 
+_Function: TypeAlias = Callable[..., Any]
+
 
 class Settings(BaseSettings):
     """FastMCP server settings.
@@ -165,7 +168,7 @@ class FastMCP:
         return Context(request_context=request_context, fastmcp=self)
 
     async def call_tool(
-        self, name: str, arguments: dict
+        self, name: str, arguments: dict[str, Any]
     ) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
         """Call a tool by name with arguments."""
         context = self.get_context()
@@ -214,7 +217,7 @@ class FastMCP:
 
     def add_tool(
         self,
-        fn: Callable,
+        fn: _Function,
         name: str | None = None,
         description: str | None = None,
     ) -> None:
@@ -230,7 +233,9 @@ class FastMCP:
         """
         self._tool_manager.add_tool(fn, name=name, description=description)
 
-    def tool(self, name: str | None = None, description: str | None = None) -> Callable:
+    def tool(
+        self, name: str | None = None, description: str | None = None
+    ) -> Callable[[_Function], _Function]:
         """Decorator to register a tool.
 
         Tools can optionally request a Context object by adding a parameter with the
@@ -263,7 +268,7 @@ class FastMCP:
                 "Did you forget to call it? Use @tool() instead of @tool"
             )
 
-        def decorator(fn: Callable) -> Callable:
+        def decorator(fn: _Function) -> _Function:
             self.add_tool(fn, name=name, description=description)
             return fn
 
@@ -284,7 +289,7 @@ class FastMCP:
         name: str | None = None,
         description: str | None = None,
         mime_type: str | None = None,
-    ) -> Callable:
+    ) -> Callable[[_Function], _Function]:
         """Decorator to register a function as a resource.
 
         The function will be called when the resource is read to generate its content.
@@ -328,7 +333,7 @@ class FastMCP:
                 "Did you forget to call it? Use @resource('uri') instead of @resource"
             )
 
-        def decorator(fn: Callable) -> Callable:
+        def decorator(fn: _Function) -> _Function:
             # Check if this should be a template
             has_uri_params = "{" in uri and "}" in uri
             has_func_params = bool(inspect.signature(fn).parameters)
@@ -376,7 +381,7 @@ class FastMCP:
 
     def prompt(
         self, name: str | None = None, description: str | None = None
-    ) -> Callable:
+    ) -> Callable[[_Function], _Function]:
         """Decorator to register a prompt.
 
         Args:
@@ -417,7 +422,7 @@ class FastMCP:
                 "Did you forget to call it? Use @prompt() instead of @prompt"
             )
 
-        def decorator(func: Callable) -> Callable:
+        def decorator(func: _Function) -> _Function:
             prompt = Prompt.from_function(func, name=name, description=description)
             self.add_prompt(prompt)
             return func