Jeremiah Lowin commited on
Commit
60e7e71
·
1 Parent(s): 7ca2f08

Improve methods on server and registration

Browse files
Files changed (2) hide show
  1. src/fastmcp/server.py +47 -43
  2. tests/test_server.py +25 -13
src/fastmcp/server.py CHANGED
@@ -67,50 +67,54 @@ class FastMCP:
67
 
68
  def _setup_handlers(self) -> None:
69
  """Set up core MCP protocol handlers."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
 
71
- @self._mcp_server.list_tools()
72
- async def handle_list_tools() -> list[Tool]:
73
- tools = self._tool_manager.list_tools()
74
- return [
75
- Tool(
76
- name=info.name,
77
- description=info.description,
78
- inputSchema=info.parameters,
79
- )
80
- for info in tools
81
- ]
82
-
83
- @self._mcp_server.call_tool()
84
- async def handle_call_tool(
85
- name: str, arguments: dict
86
- ) -> Sequence[Union[TextContent, ImageContent, EmbeddedResource]]:
87
- result = await self._tool_manager.call_tool(name, arguments)
88
- return [self._convert_to_content(result)]
89
-
90
- @self._mcp_server.list_resources()
91
- async def handle_list_resources() -> list[MCPResource]:
92
- resources = self._resource_manager.list_resources()
93
- return [
94
- MCPResource(
95
- uri=resource.uri,
96
- name=resource.name,
97
- description=resource.description,
98
- mimeType=resource.mime_type,
99
- )
100
- for resource in resources
101
- ]
102
-
103
- @self._mcp_server.read_resource()
104
- async def handle_read_resource(uri: str) -> Union[str, bytes]:
105
- resource = self._resource_manager.get_resource(uri)
106
- if not resource:
107
- raise ResourceError(f"Unknown resource: {uri}")
108
 
109
- try:
110
- return await resource.read()
111
- except Exception as e:
112
- logger.error(f"Error reading resource {uri}: {e}")
113
- raise ResourceError(str(e))
114
 
115
  def _convert_to_content(
116
  self, value: Any
@@ -143,6 +147,7 @@ class FastMCP:
143
  self, name: Optional[str] = None, description: Optional[str] = None
144
  ) -> Callable:
145
  """Decorator to register a tool."""
 
146
 
147
  def decorator(func: Callable) -> Callable:
148
  self.add_tool(func, name=name, description=description)
@@ -308,7 +313,6 @@ class FastMCP:
308
  app: "FastMCP",
309
  ) -> None:
310
  """Run the server using SSE transport."""
311
- from mcp.server.sse import SseServerTransport
312
  from starlette.applications import Starlette
313
  from starlette.routing import Route
314
  import uvicorn
 
67
 
68
  def _setup_handlers(self) -> None:
69
  """Set up core MCP protocol handlers."""
70
+ self._mcp_server.list_tools()(self.list_tools)
71
+ self._mcp_server.call_tool()(self.call_tool)
72
+ self._mcp_server.list_resources()(self.list_resources)
73
+ self._mcp_server.read_resource()(self.read_resource)
74
+
75
+ async def list_tools(self) -> list[Tool]:
76
+ """List all available tools."""
77
+ tools = self._tool_manager.list_tools()
78
+ return [
79
+ Tool(
80
+ name=info.name,
81
+ description=info.description,
82
+ inputSchema=info.parameters,
83
+ )
84
+ for info in tools
85
+ ]
86
+
87
+ async def call_tool(
88
+ self, name: str, arguments: dict
89
+ ) -> Sequence[Union[TextContent, ImageContent, EmbeddedResource]]:
90
+ """Call a tool by name with arguments."""
91
+ result = await self._tool_manager.call_tool(name, arguments)
92
+ return [self._convert_to_content(result)]
93
+
94
+ async def list_resources(self) -> list[MCPResource]:
95
+ """List all available resources."""
96
+ resources = self._resource_manager.list_resources()
97
+ return [
98
+ MCPResource(
99
+ uri=resource.uri,
100
+ name=resource.name,
101
+ description=resource.description,
102
+ mimeType=resource.mime_type,
103
+ )
104
+ for resource in resources
105
+ ]
106
 
107
+ async def read_resource(self, uri: str) -> Union[str, bytes]:
108
+ """Read a resource by URI."""
109
+ resource = self._resource_manager.get_resource(uri)
110
+ if not resource:
111
+ raise ResourceError(f"Unknown resource: {uri}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
112
 
113
+ try:
114
+ return await resource.read()
115
+ except Exception as e:
116
+ logger.error(f"Error reading resource {uri}: {e}")
117
+ raise ResourceError(str(e))
118
 
119
  def _convert_to_content(
120
  self, value: Any
 
147
  self, name: Optional[str] = None, description: Optional[str] = None
148
  ) -> Callable:
149
  """Decorator to register a tool."""
150
+ breakpoint()
151
 
152
  def decorator(func: Callable) -> Callable:
153
  self.add_tool(func, name=name, description=description)
 
313
  app: "FastMCP",
314
  ) -> None:
315
  """Run the server using SSE transport."""
 
316
  from starlette.applications import Starlette
317
  from starlette.routing import Route
318
  import uvicorn
tests/test_server.py CHANGED
@@ -1,13 +1,25 @@
1
  from mcp.shared.memory import (
2
  create_connected_server_and_client_session as client_session,
3
  )
4
- from fastmcp.server import FastMCP
5
 
6
 
7
  class TestServer:
8
  async def test_create_server(self):
9
- server = FastMCP()
10
- assert server.name == "FastMCP"
 
 
 
 
 
 
 
 
 
 
 
 
11
 
12
 
13
  def tool_fn(x: int, y: int) -> int:
@@ -16,22 +28,22 @@ def tool_fn(x: int, y: int) -> int:
16
 
17
  class TestServerTools:
18
  async def test_add_tool(self):
19
- server = FastMCP()
20
- server.add_tool(tool_fn)
21
- server.add_tool(tool_fn)
22
- assert len(server._tool_manager.list_tools()) == 1
23
 
24
  async def test_list_tools(self):
25
- server = FastMCP()
26
- server.add_tool(tool_fn)
27
- async with client_session(server._mcp_server) as client:
28
  tools = await client.list_tools()
29
  assert len(tools.tools) == 1
30
 
31
  async def test_call_tool(self):
32
- server = FastMCP()
33
- server.add_tool(tool_fn)
34
- async with client_session(server._mcp_server) as client:
35
  result = await client.call_tool("my_tool", {"arg1": "value"})
36
  assert "error" not in result
37
  assert len(result.content) > 0
 
1
  from mcp.shared.memory import (
2
  create_connected_server_and_client_session as client_session,
3
  )
4
+ from fastmcp import FastMCP
5
 
6
 
7
  class TestServer:
8
  async def test_create_server(self):
9
+ mcp = FastMCP()
10
+ assert mcp.name == "FastMCP"
11
+
12
+ async def test_add_tool_decorator(self):
13
+ mcp = FastMCP()
14
+
15
+ @mcp.tool
16
+ def add(x: int, y: int) -> int:
17
+ return x + y
18
+
19
+ async with client_session(mcp._mcp_server) as client:
20
+ tools = await client.list_tools()
21
+ assert len(tools.tools) == 1
22
+ assert tools.tools[0].name == "add"
23
 
24
 
25
  def tool_fn(x: int, y: int) -> int:
 
28
 
29
  class TestServerTools:
30
  async def test_add_tool(self):
31
+ mcp = FastMCP()
32
+ mcp.add_tool(tool_fn)
33
+ mcp.add_tool(tool_fn)
34
+ assert len(mcp._tool_manager.list_tools()) == 1
35
 
36
  async def test_list_tools(self):
37
+ mcp = FastMCP()
38
+ mcp.add_tool(tool_fn)
39
+ async with client_session(mcp._mcp_server) as client:
40
  tools = await client.list_tools()
41
  assert len(tools.tools) == 1
42
 
43
  async def test_call_tool(self):
44
+ mcp = FastMCP()
45
+ mcp.add_tool(tool_fn)
46
+ async with client_session(mcp._mcp_server) as client:
47
  result = await client.call_tool("my_tool", {"arg1": "value"})
48
  assert "error" not in result
49
  assert len(result.content) > 0