Jeremiah Lowin commited on
Commit
a277f25
·
1 Parent(s): 2a88273

Refactor server to wrap mcp calls

Browse files
src/fastmcp/exceptions.py CHANGED
@@ -33,3 +33,7 @@ class ClientError(Exception):
33
 
34
  class NotFoundError(Exception):
35
  """Object not found."""
 
 
 
 
 
33
 
34
  class NotFoundError(Exception):
35
  """Object not found."""
36
+
37
+
38
+ class DisabledError(Exception):
39
+ """Object is disabled."""
src/fastmcp/server/proxy.py CHANGED
@@ -241,20 +241,20 @@ class FastMCPProxy(FastMCP):
241
  prompts[prompt_proxy.name] = prompt_proxy
242
  return prompts
243
 
244
- async def _mcp_call_tool(
245
  self, key: str, arguments: dict[str, Any]
246
  ) -> list[TextContent | ImageContent | EmbeddedResource]:
247
  try:
248
- result = await super()._mcp_call_tool(key, arguments)
249
  return result
250
  except NotFoundError:
251
  async with self.client:
252
  result = await self.client.call_tool(key, arguments)
253
  return result
254
 
255
- async def _mcp_read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]:
256
  try:
257
- result = await super()._mcp_read_resource(uri)
258
  return result
259
  except NotFoundError:
260
  async with self.client:
@@ -270,11 +270,11 @@ class FastMCPProxy(FastMCP):
270
  ReadResourceContents(content=content, mime_type=resource[0].mimeType)
271
  ]
272
 
273
- async def _mcp_get_prompt(
274
  self, name: str, arguments: dict[str, Any] | None = None
275
  ) -> GetPromptResult:
276
  try:
277
- result = await super()._mcp_get_prompt(name, arguments)
278
  return result
279
  except NotFoundError:
280
  async with self.client:
 
241
  prompts[prompt_proxy.name] = prompt_proxy
242
  return prompts
243
 
244
+ async def _call_tool(
245
  self, key: str, arguments: dict[str, Any]
246
  ) -> list[TextContent | ImageContent | EmbeddedResource]:
247
  try:
248
+ result = await super()._call_tool(key, arguments)
249
  return result
250
  except NotFoundError:
251
  async with self.client:
252
  result = await self.client.call_tool(key, arguments)
253
  return result
254
 
255
+ async def _read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]:
256
  try:
257
+ result = await super()._read_resource(uri)
258
  return result
259
  except NotFoundError:
260
  async with self.client:
 
270
  ReadResourceContents(content=content, mime_type=resource[0].mimeType)
271
  ]
272
 
273
+ async def _get_prompt(
274
  self, name: str, arguments: dict[str, Any] | None = None
275
  ) -> GetPromptResult:
276
  try:
277
+ result = await super()._get_prompt(name, arguments)
278
  return result
279
  except NotFoundError:
280
  async with self.client:
src/fastmcp/server/server.py CHANGED
@@ -44,7 +44,7 @@ from starlette.routing import BaseRoute, Route
44
  import fastmcp
45
  import fastmcp.server
46
  import fastmcp.settings
47
- from fastmcp.exceptions import NotFoundError
48
  from fastmcp.prompts import Prompt, PromptManager
49
  from fastmcp.prompts.prompt import FunctionPrompt
50
  from fastmcp.resources import Resource, ResourceManager
@@ -291,6 +291,12 @@ class FastMCP(Generic[LifespanResultT]):
291
  self._cache.set("resources", resources)
292
  return resources
293
 
 
 
 
 
 
 
294
  async def get_resource_templates(self) -> dict[str, ResourceTemplate]:
295
  """Get all registered resource templates, indexed by registered key."""
296
  if (
@@ -311,6 +317,12 @@ class FastMCP(Generic[LifespanResultT]):
311
  self._cache.set("resource_templates", templates)
312
  return templates
313
 
 
 
 
 
 
 
314
  async def get_prompts(self) -> dict[str, Prompt]:
315
  """
316
  List all available prompts.
@@ -330,6 +342,12 @@ class FastMCP(Generic[LifespanResultT]):
330
  self._cache.set("prompts", prompts)
331
  return prompts
332
 
 
 
 
 
 
 
333
  def custom_route(
334
  self,
335
  path: str,
@@ -381,7 +399,9 @@ class FastMCP(Generic[LifespanResultT]):
381
 
382
  """
383
  tools = await self.get_tools()
384
- return [tool.to_mcp_tool(name=key) for key, tool in tools.items()]
 
 
385
 
386
  async def _mcp_list_resources(self) -> list[MCPResource]:
387
  """
@@ -391,7 +411,9 @@ class FastMCP(Generic[LifespanResultT]):
391
  """
392
  resources = await self.get_resources()
393
  return [
394
- resource.to_mcp_resource(uri=key) for key, resource in resources.items()
 
 
395
  ]
396
 
397
  async def _mcp_list_resource_templates(self) -> list[MCPResourceTemplate]:
@@ -404,6 +426,7 @@ class FastMCP(Generic[LifespanResultT]):
404
  return [
405
  template.to_mcp_template(uriTemplate=key)
406
  for key, template in templates.items()
 
407
  ]
408
 
409
  async def _mcp_list_prompts(self) -> list[MCPPrompt]:
@@ -413,12 +436,19 @@ class FastMCP(Generic[LifespanResultT]):
413
 
414
  """
415
  prompts = await self.get_prompts()
416
- return [prompt.to_mcp_prompt(name=key) for key, prompt in prompts.items()]
 
 
 
 
417
 
418
  async def _mcp_call_tool(
419
  self, key: str, arguments: dict[str, Any]
420
  ) -> list[TextContent | ImageContent | EmbeddedResource]:
421
- """Handle MCP 'callTool' requests.
 
 
 
422
 
423
  Args:
424
  key: The name of the tool to call
@@ -431,43 +461,109 @@ class FastMCP(Generic[LifespanResultT]):
431
 
432
  # Create and use context for the entire call
433
  with fastmcp.server.context.Context(fastmcp=self):
434
- # Get tool, checking first from our tools, then from the mounted servers
435
- if self._tool_manager.has_tool(key):
436
- return await self._tool_manager.call_tool(key, arguments)
 
 
 
 
 
 
 
 
 
 
 
 
437
 
438
- # Check mounted servers to see if they have the tool
439
- for server in self._mounted_servers.values():
440
- if server.match_tool(key):
441
- tool_key = server.strip_tool_prefix(key)
442
- return await server.server._mcp_call_tool(tool_key, arguments)
443
 
444
- raise NotFoundError(f"Unknown tool: {key}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
445
 
446
  async def _mcp_read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
447
  """
448
  Read a resource by URI, in the format expected by the low-level MCP
449
  server.
450
  """
451
- with fastmcp.server.context.Context(fastmcp=self):
452
- if self._resource_manager.has_resource(uri):
453
- resource = await self._resource_manager.get_resource(uri)
454
- content = await self._resource_manager.read_resource(uri)
455
- return [
456
- ReadResourceContents(
457
- content=content,
458
- mime_type=resource.mime_type,
459
- )
460
- ]
 
 
 
 
 
 
461
  else:
462
- for server in self._mounted_servers.values():
463
- if server.match_resource(str(uri)):
464
- new_uri = server.strip_resource_prefix(str(uri))
465
- return await server.server._mcp_read_resource(new_uri)
466
- else:
467
- raise NotFoundError(f"Unknown resource: {uri}")
468
 
469
  async def _mcp_get_prompt(
470
  self, name: str, arguments: dict[str, Any] | None = None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
471
  ) -> GetPromptResult:
472
  """Handle MCP 'getPrompt' requests.
473
 
@@ -480,19 +576,20 @@ class FastMCP(Generic[LifespanResultT]):
480
  """
481
  logger.debug("Get prompt: %s with %s", name, arguments)
482
 
483
- # Create and use context for the entire call
484
- with fastmcp.server.context.Context(fastmcp=self):
485
- # Get prompt, checking first from our prompts, then from the mounted servers
486
- if self._prompt_manager.has_prompt(name):
487
- return await self._prompt_manager.render_prompt(name, arguments)
 
488
 
489
- # Check mounted servers to see if they have the prompt
490
- for server in self._mounted_servers.values():
491
- if server.match_prompt(name):
492
- prompt_name = server.strip_prompt_prefix(name)
493
- return await server.server._mcp_get_prompt(prompt_name, arguments)
494
 
495
- raise NotFoundError(f"Unknown prompt: {name}")
496
 
497
  def add_tool(self, tool: Tool) -> None:
498
  """Add a tool to the server.
 
44
  import fastmcp
45
  import fastmcp.server
46
  import fastmcp.settings
47
+ from fastmcp.exceptions import DisabledError, NotFoundError
48
  from fastmcp.prompts import Prompt, PromptManager
49
  from fastmcp.prompts.prompt import FunctionPrompt
50
  from fastmcp.resources import Resource, ResourceManager
 
291
  self._cache.set("resources", resources)
292
  return resources
293
 
294
+ async def get_resource(self, key: str) -> Resource:
295
+ resources = await self.get_resources()
296
+ if key not in resources:
297
+ raise NotFoundError(f"Unknown resource: {key}")
298
+ return resources[key]
299
+
300
  async def get_resource_templates(self) -> dict[str, ResourceTemplate]:
301
  """Get all registered resource templates, indexed by registered key."""
302
  if (
 
317
  self._cache.set("resource_templates", templates)
318
  return templates
319
 
320
+ async def get_resource_template(self, key: str) -> ResourceTemplate:
321
+ templates = await self.get_resource_templates()
322
+ if key not in templates:
323
+ raise NotFoundError(f"Unknown resource template: {key}")
324
+ return templates[key]
325
+
326
  async def get_prompts(self) -> dict[str, Prompt]:
327
  """
328
  List all available prompts.
 
342
  self._cache.set("prompts", prompts)
343
  return prompts
344
 
345
+ async def get_prompt(self, key: str) -> Prompt:
346
+ prompts = await self.get_prompts()
347
+ if key not in prompts:
348
+ raise NotFoundError(f"Unknown prompt: {key}")
349
+ return prompts[key]
350
+
351
  def custom_route(
352
  self,
353
  path: str,
 
399
 
400
  """
401
  tools = await self.get_tools()
402
+ return [
403
+ tool.to_mcp_tool(name=key) for key, tool in tools.items() if tool.enabled
404
+ ]
405
 
406
  async def _mcp_list_resources(self) -> list[MCPResource]:
407
  """
 
411
  """
412
  resources = await self.get_resources()
413
  return [
414
+ resource.to_mcp_resource(uri=key)
415
+ for key, resource in resources.items()
416
+ if resource.enabled
417
  ]
418
 
419
  async def _mcp_list_resource_templates(self) -> list[MCPResourceTemplate]:
 
426
  return [
427
  template.to_mcp_template(uriTemplate=key)
428
  for key, template in templates.items()
429
+ if template.enabled
430
  ]
431
 
432
  async def _mcp_list_prompts(self) -> list[MCPPrompt]:
 
436
 
437
  """
438
  prompts = await self.get_prompts()
439
+ return [
440
+ prompt.to_mcp_prompt(name=key)
441
+ for key, prompt in prompts.items()
442
+ if prompt.enabled
443
+ ]
444
 
445
  async def _mcp_call_tool(
446
  self, key: str, arguments: dict[str, Any]
447
  ) -> list[TextContent | ImageContent | EmbeddedResource]:
448
+ """
449
+ Handle MCP 'callTool' requests.
450
+
451
+ Delegates to _call_tool, which should be overridden by FastMCP subclasses.
452
 
453
  Args:
454
  key: The name of the tool to call
 
461
 
462
  # Create and use context for the entire call
463
  with fastmcp.server.context.Context(fastmcp=self):
464
+ try:
465
+ return await self._call_tool(key, arguments)
466
+ except DisabledError:
467
+ # convert to NotFoundError to avoid leaking tool presence
468
+ raise NotFoundError(f"Unknown tool: {key}")
469
+ except NotFoundError:
470
+ # standardize NotFound message
471
+ raise NotFoundError(f"Unknown tool: {key}")
472
+
473
+ async def _call_tool(
474
+ self, key: str, arguments: dict[str, Any]
475
+ ) -> list[TextContent | ImageContent | EmbeddedResource]:
476
+ """
477
+ Call a tool with raw MCP arguments. FastMCP subclasses should override
478
+ this method, not _mcp_call_tool.
479
 
480
+ Args:
481
+ key: The name of the tool to call arguments: Arguments to pass to
482
+ the tool
 
 
483
 
484
+ Returns:
485
+ List of MCP Content objects containing the tool results
486
+ """
487
+
488
+ # Get tool, checking first from our tools, then from the mounted servers
489
+ if self._tool_manager.has_tool(key):
490
+ tool = self._tool_manager.get_tool(key)
491
+ if not tool.enabled:
492
+ raise DisabledError(f"Tool {key!r} is disabled")
493
+ return await self._tool_manager.call_tool(key, arguments)
494
+
495
+ # Check mounted servers to see if they have the tool
496
+ for server in self._mounted_servers.values():
497
+ if server.match_tool(key):
498
+ tool_key = server.strip_tool_prefix(key)
499
+ return await server.server._call_tool(tool_key, arguments)
500
+
501
+ raise NotFoundError(f"Unknown tool: {key!r}")
502
 
503
  async def _mcp_read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]:
504
+ """
505
+ Handle MCP 'readResource' requests.
506
+
507
+ Delegates to _read_resource, which should be overridden by FastMCP subclasses.
508
+ """
509
+ logger.debug("Read resource: %s", uri)
510
+
511
+ with fastmcp.server.context.Context(fastmcp=self):
512
+ try:
513
+ return await self._read_resource(uri)
514
+ except DisabledError:
515
+ # convert to NotFoundError to avoid leaking resource presence
516
+ raise NotFoundError(f"Unknown resource: {str(uri)!r}")
517
+ except NotFoundError:
518
+ # standardize NotFound message
519
+ raise NotFoundError(f"Unknown resource: {str(uri)!r}")
520
+
521
+ async def _read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]:
522
  """
523
  Read a resource by URI, in the format expected by the low-level MCP
524
  server.
525
  """
526
+ if self._resource_manager.has_resource(uri):
527
+ resource = await self._resource_manager.get_resource(uri)
528
+ if not resource.enabled:
529
+ raise DisabledError(f"Resource {str(uri)!r} is disabled")
530
+ content = await self._resource_manager.read_resource(uri)
531
+ return [
532
+ ReadResourceContents(
533
+ content=content,
534
+ mime_type=resource.mime_type,
535
+ )
536
+ ]
537
+ else:
538
+ for server in self._mounted_servers.values():
539
+ if server.match_resource(str(uri)):
540
+ new_uri = server.strip_resource_prefix(str(uri))
541
+ return await server.server._mcp_read_resource(new_uri)
542
  else:
543
+ raise NotFoundError(f"Unknown resource: {uri}")
 
 
 
 
 
544
 
545
  async def _mcp_get_prompt(
546
  self, name: str, arguments: dict[str, Any] | None = None
547
+ ) -> GetPromptResult:
548
+ """
549
+ Handle MCP 'getPrompt' requests.
550
+
551
+ Delegates to _get_prompt, which should be overridden by FastMCP subclasses.
552
+ """
553
+ logger.debug("Get prompt: %s with %s", name, arguments)
554
+
555
+ with fastmcp.server.context.Context(fastmcp=self):
556
+ try:
557
+ return await self._get_prompt(name, arguments)
558
+ except DisabledError:
559
+ # convert to NotFoundError to avoid leaking prompt presence
560
+ raise NotFoundError(f"Unknown prompt: {name}")
561
+ except NotFoundError:
562
+ # standardize NotFound message
563
+ raise NotFoundError(f"Unknown prompt: {name}")
564
+
565
+ async def _get_prompt(
566
+ self, name: str, arguments: dict[str, Any] | None = None
567
  ) -> GetPromptResult:
568
  """Handle MCP 'getPrompt' requests.
569
 
 
576
  """
577
  logger.debug("Get prompt: %s with %s", name, arguments)
578
 
579
+ # Get prompt, checking first from our prompts, then from the mounted servers
580
+ if self._prompt_manager.has_prompt(name):
581
+ prompt = self._prompt_manager.get_prompt(name)
582
+ if not prompt.enabled:
583
+ raise DisabledError(f"Prompt {name!r} is disabled")
584
+ return await self._prompt_manager.render_prompt(name, arguments)
585
 
586
+ # Check mounted servers to see if they have the prompt
587
+ for server in self._mounted_servers.values():
588
+ if server.match_prompt(name):
589
+ prompt_name = server.strip_prompt_prefix(name)
590
+ return await server.server._mcp_get_prompt(prompt_name, arguments)
591
 
592
+ raise NotFoundError(f"Unknown prompt: {name}")
593
 
594
  def add_tool(self, tool: Tool) -> None:
595
  """Add a tool to the server.