Jeremiah Lowin commited on
Commit
80723cc
·
1 Parent(s): 816878f

Update prompt manager

Browse files
src/fastmcp/prompts/prompt_manager.py CHANGED
@@ -2,7 +2,7 @@ from __future__ import annotations as _annotations
2
 
3
  import warnings
4
  from collections.abc import Awaitable, Callable
5
- from typing import TYPE_CHECKING, Any
6
 
7
  from mcp import GetPromptResult
8
 
@@ -13,7 +13,7 @@ from fastmcp.settings import DuplicateBehavior
13
  from fastmcp.utilities.logging import get_logger
14
 
15
  if TYPE_CHECKING:
16
- pass
17
 
18
  logger = get_logger(__name__)
19
 
@@ -27,6 +27,7 @@ class PromptManager:
27
  mask_error_details: bool | None = None,
28
  ):
29
  self._prompts: dict[str, Prompt] = {}
 
30
  self.mask_error_details = mask_error_details or settings.mask_error_details
31
 
32
  # Default to "warn" if None is provided
@@ -41,15 +42,80 @@ class PromptManager:
41
 
42
  self.duplicate_behavior = duplicate_behavior
43
 
44
- def get_prompt(self, key: str) -> Prompt:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
45
  """Get prompt by key."""
46
- if key in self._prompts:
47
- return self._prompts[key]
 
48
  raise NotFoundError(f"Unknown prompt: {key}")
49
 
50
- def get_prompts(self) -> dict[str, Prompt]:
51
- """Get all registered prompts, indexed by registered key."""
52
- return self._prompts
 
 
 
 
 
 
 
 
 
53
 
54
  def add_prompt_from_fn(
55
  self,
@@ -96,30 +162,44 @@ class PromptManager:
96
  name: str,
97
  arguments: dict[str, Any] | None = None,
98
  ) -> GetPromptResult:
99
- """Render a prompt by name with arguments."""
100
- prompt = self.get_prompt(name)
101
- if not prompt:
102
- raise NotFoundError(f"Unknown prompt: {name}")
103
-
104
- try:
105
- messages = await prompt.render(arguments)
106
- return GetPromptResult(description=prompt.description, messages=messages)
107
-
108
- # Pass through PromptErrors as-is
109
- except PromptError as e:
110
- logger.exception(f"Error rendering prompt {name!r}: {e}")
111
- raise e
112
-
113
- # Handle other exceptions
114
- except Exception as e:
115
- logger.exception(f"Error rendering prompt {name!r}: {e}")
116
- if self.mask_error_details:
117
- # Mask internal details
118
- raise PromptError(f"Error rendering prompt {name!r}")
119
- else:
120
- # Include original error details
121
- raise PromptError(f"Error rendering prompt {name!r}: {e}")
122
-
123
- def has_prompt(self, key: str) -> bool:
124
- """Check if a prompt exists."""
125
- return key in self._prompts
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
  import warnings
4
  from collections.abc import Awaitable, Callable
5
+ from typing import TYPE_CHECKING, Any, Literal
6
 
7
  from mcp import GetPromptResult
8
 
 
13
  from fastmcp.utilities.logging import get_logger
14
 
15
  if TYPE_CHECKING:
16
+ from fastmcp.server.server import MountedServer
17
 
18
  logger = get_logger(__name__)
19
 
 
27
  mask_error_details: bool | None = None,
28
  ):
29
  self._prompts: dict[str, Prompt] = {}
30
+ self._mounted_sources: list[MountedServer] = []
31
  self.mask_error_details = mask_error_details or settings.mask_error_details
32
 
33
  # Default to "warn" if None is provided
 
42
 
43
  self.duplicate_behavior = duplicate_behavior
44
 
45
+ def mount(self, server: MountedServer) -> None:
46
+ """Adds a mounted server as a source for prompts."""
47
+ self._mounted_sources.append(server)
48
+
49
+ async def _load_prompts(
50
+ self, *, mode: Literal["inventory", "protocol"]
51
+ ) -> dict[str, Prompt]:
52
+ """
53
+ The single, consolidated recursive method for fetching prompts. The 'mode'
54
+ parameter determines the communication path.
55
+
56
+ - mode="inventory": Manager-to-manager path for complete, unfiltered inventory
57
+ - mode="protocol": Server-to-server path for filtered MCP requests
58
+ """
59
+ all_prompts: dict[str, Prompt] = {}
60
+
61
+ for mounted in self._mounted_sources:
62
+ try:
63
+ if mode == "protocol":
64
+ # PATH 2: Use the server-to-server filtered path
65
+ child_results = await mounted.server._list_prompts()
66
+ else: # mode == "inventory"
67
+ # PATH 1: Use the manager-to-manager unfiltered path
68
+ child_results = await mounted.server._prompt_manager.get_prompts()
69
+
70
+ # The combination logic is the same for both paths
71
+ child_dict = (
72
+ {p.key: p for p in child_results}
73
+ if isinstance(child_results, list)
74
+ else child_results
75
+ )
76
+ if mounted.prefix:
77
+ for prompt in child_dict.values():
78
+ prefixed_prompt = prompt.with_key(
79
+ f"{mounted.prefix}_{prompt.key}"
80
+ )
81
+ all_prompts[prefixed_prompt.key] = prefixed_prompt
82
+ else:
83
+ all_prompts.update(child_dict)
84
+ except Exception as e:
85
+ # Skip failed mounts silently, matches existing behavior
86
+ logger.warning(
87
+ f"Failed to get prompts from mounted server '{mounted.prefix}': {e}"
88
+ )
89
+ continue
90
+
91
+ # Finally, add local prompts, which always take precedence
92
+ all_prompts.update(self._prompts)
93
+ return all_prompts
94
+
95
+ async def has_prompt(self, key: str) -> bool:
96
+ """Check if a prompt exists."""
97
+ prompts = await self.get_prompts()
98
+ return key in prompts
99
+
100
+ async def get_prompt(self, key: str) -> Prompt:
101
  """Get prompt by key."""
102
+ prompts = await self.get_prompts()
103
+ if key in prompts:
104
+ return prompts[key]
105
  raise NotFoundError(f"Unknown prompt: {key}")
106
 
107
+ async def get_prompts(self) -> dict[str, Prompt]:
108
+ """
109
+ Gets the complete, unfiltered inventory of all prompts.
110
+ """
111
+ return await self._load_prompts(mode="inventory")
112
+
113
+ async def list_prompts(self) -> list[Prompt]:
114
+ """
115
+ Lists all prompts, applying protocol filtering.
116
+ """
117
+ prompts_dict = await self._load_prompts(mode="protocol")
118
+ return list(prompts_dict.values())
119
 
120
  def add_prompt_from_fn(
121
  self,
 
162
  name: str,
163
  arguments: dict[str, Any] | None = None,
164
  ) -> GetPromptResult:
165
+ """
166
+ Internal API for servers: Finds and renders a prompt, respecting the
167
+ filtered protocol path.
168
+ """
169
+ # 1. Check local prompts first. The server will have already applied its filter.
170
+ if name in self._prompts:
171
+ prompt = await self.get_prompt(name)
172
+ if not prompt:
173
+ raise NotFoundError(f"Unknown prompt: {name}")
174
+
175
+ try:
176
+ messages = await prompt.render(arguments)
177
+ return GetPromptResult(
178
+ description=prompt.description, messages=messages
179
+ )
180
+
181
+ # Pass through PromptErrors as-is
182
+ except PromptError as e:
183
+ logger.exception(f"Error rendering prompt {name!r}: {e}")
184
+ raise e
185
+
186
+ # Handle other exceptions
187
+ except Exception as e:
188
+ logger.exception(f"Error rendering prompt {name!r}: {e}")
189
+ if self.mask_error_details:
190
+ # Mask internal details
191
+ raise PromptError(f"Error rendering prompt {name!r}") from e
192
+ else:
193
+ # Include original error details
194
+ raise PromptError(f"Error rendering prompt {name!r}: {e}") from e
195
+
196
+ # 2. Check mounted servers using the filtered protocol path.
197
+ for mounted in reversed(self._mounted_sources):
198
+ if mounted.prefix and name.startswith(f"{mounted.prefix}_"):
199
+ name_on_child = name.removeprefix(f"{mounted.prefix}_")
200
+ try:
201
+ return await mounted.server._get_prompt(name_on_child, arguments)
202
+ except NotFoundError:
203
+ continue
204
+
205
+ raise NotFoundError(f"Unknown prompt: {name}")
src/fastmcp/server/server.py CHANGED
@@ -341,8 +341,7 @@ class FastMCP(Generic[LifespanResultT]):
341
 
342
  async def get_tools(self) -> dict[str, Tool]:
343
  """Get all registered tools, indexed by registered key."""
344
- tools = await self._list_tools(apply_middleware=False)
345
- return {tool.key: tool for tool in tools}
346
 
347
  async def get_tool(self, key: str) -> Tool:
348
  tools = await self.get_tools()
@@ -376,9 +375,7 @@ class FastMCP(Generic[LifespanResultT]):
376
  """
377
  List all available prompts.
378
  """
379
-
380
- prompts = await self._list_prompts(apply_middleware=False)
381
- return {prompt.key: prompt for prompt in prompts}
382
 
383
  async def get_prompt(self, key: str) -> Prompt:
384
  prompts = await self.get_prompts()
@@ -651,10 +648,10 @@ class FastMCP(Generic[LifespanResultT]):
651
  logger.debug("Handler called: list_prompts")
652
 
653
  with fastmcp.server.context.Context(fastmcp=self):
654
- prompts = await self._middleware_list_prompts()
655
  return [prompt.to_mcp_prompt(name=prompt.key) for prompt in prompts]
656
 
657
- async def _middleware_list_prompts(self) -> list[Prompt]:
658
  """
659
  List all available prompts, in the format expected by the low-level MCP
660
  server.
@@ -662,9 +659,9 @@ class FastMCP(Generic[LifespanResultT]):
662
  """
663
 
664
  async def _handler(
665
- context: MiddlewareContext[dict[str, Any]],
666
  ) -> list[Prompt]:
667
- prompts = await self._list_prompts()
668
 
669
  mcp_prompts: list[Prompt] = []
670
  for prompt in prompts:
@@ -676,7 +673,7 @@ class FastMCP(Generic[LifespanResultT]):
676
  with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx:
677
  # Create the middleware context.
678
  mw_context = MiddlewareContext(
679
- message={}, # List prompts doesn't have parameters
680
  source="client",
681
  type="request",
682
  method="prompts/list",
@@ -686,43 +683,6 @@ class FastMCP(Generic[LifespanResultT]):
686
  # Apply the middleware chain.
687
  return await self._apply_middleware(mw_context, _handler)
688
 
689
- async def _list_prompts(self, apply_middleware: bool = True) -> list[Prompt]:
690
- """
691
- List all available prompts.
692
- """
693
-
694
- if (prompts := self._cache.get("prompts")) is self._cache.NOT_FOUND:
695
- prompts: dict[str, Prompt] = {}
696
-
697
- # iterate such that new mounts overwrite older ones
698
- for mounted_server in self._mounted_servers:
699
- try:
700
- if apply_middleware:
701
- server_prompts = (
702
- await mounted_server.server._middleware_list_prompts()
703
- )
704
- else:
705
- server_prompts = await mounted_server.server._list_prompts()
706
- # Apply prefix to each prompt key if prefix exists
707
- if mounted_server.prefix:
708
- for prompt in server_prompts:
709
- prompt = prompt.with_key(
710
- f"{mounted_server.prefix}_{prompt.key}"
711
- )
712
- prompts[prompt.key] = prompt
713
- else:
714
- prompts.update(
715
- {prompt.key: prompt for prompt in server_prompts}
716
- )
717
- except Exception as e:
718
- logger.warning(
719
- f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}"
720
- )
721
- continue
722
- prompts.update(await self._prompt_manager.get_prompts())
723
- self._cache.set("prompts", prompts)
724
- return list(prompts.values())
725
-
726
  async def _mcp_call_tool(
727
  self, key: str, arguments: dict[str, Any]
728
  ) -> list[MCPContent]:
@@ -828,7 +788,7 @@ class FastMCP(Generic[LifespanResultT]):
828
  Read a resource by URI, in the format expected by the low-level MCP
829
  server.
830
  """
831
- if self._resource_manager.has_resource(uri):
832
  resource = await self._resource_manager.get_resource(uri)
833
  if not self._should_enable_component(resource):
834
  raise DisabledError(f"Resource {str(uri)!r} is disabled")
@@ -840,32 +800,7 @@ class FastMCP(Generic[LifespanResultT]):
840
  )
841
  ]
842
  else:
843
- # iterate such that new mounts take precedence over older ones
844
- for mounted_server in reversed(self._mounted_servers):
845
- resource_uri = uri
846
- try:
847
- if mounted_server.prefix:
848
- # If server has a prefix, check if URI matches and strip prefix
849
- if has_resource_prefix(
850
- str(resource_uri),
851
- mounted_server.prefix,
852
- self.resource_prefix_format,
853
- ):
854
- resource_uri = remove_resource_prefix(
855
- str(resource_uri),
856
- mounted_server.prefix,
857
- self.resource_prefix_format,
858
- )
859
- else:
860
- continue
861
- return await mounted_server.server._middleware_read_resource(
862
- resource_uri
863
- )
864
- except NotFoundError:
865
- # Resource not found on this server, try the next one
866
- continue
867
- else:
868
- raise NotFoundError(f"Unknown resource: {uri}")
869
 
870
  async def _mcp_get_prompt(
871
  self, name: str, arguments: dict[str, Any] | None = None
@@ -879,7 +814,7 @@ class FastMCP(Generic[LifespanResultT]):
879
 
880
  with fastmcp.server.context.Context(fastmcp=self):
881
  try:
882
- return await self._middleware_get_prompt(name, arguments)
883
  except DisabledError:
884
  # convert to NotFoundError to avoid leaking prompt presence
885
  raise NotFoundError(f"Unknown prompt: {name}")
@@ -887,21 +822,22 @@ class FastMCP(Generic[LifespanResultT]):
887
  # standardize NotFound message
888
  raise NotFoundError(f"Unknown prompt: {name}")
889
 
890
- async def _middleware_get_prompt(
891
- self,
892
- name: str,
893
- arguments: dict[str, Any] | None = None,
894
  ) -> GetPromptResult:
895
  """
896
- Get a prompt with middleware.
897
  """
898
 
899
  async def _handler(
900
  context: MiddlewareContext[mcp.types.GetPromptRequestParams],
901
  ) -> GetPromptResult:
902
- return await self._get_prompt(
903
- name=context.message.name,
904
- arguments=context.message.arguments,
 
 
 
905
  )
906
 
907
  mw_context = MiddlewareContext(
@@ -913,49 +849,6 @@ class FastMCP(Generic[LifespanResultT]):
913
  )
914
  return await self._apply_middleware(mw_context, _handler)
915
 
916
- async def _get_prompt(
917
- self, name: str, arguments: dict[str, Any] | None = None
918
- ) -> GetPromptResult:
919
- """Handle MCP 'getPrompt' requests.
920
-
921
- Args:
922
- name: The name of the prompt to render
923
- arguments: Arguments to pass to the prompt
924
-
925
- Returns:
926
- GetPromptResult containing the rendered prompt messages
927
- """
928
- logger.debug("Get prompt: %s with %s", name, arguments)
929
-
930
- # Get prompt, checking first from our prompts, then from the mounted servers
931
- if self._prompt_manager.has_prompt(name):
932
- prompt = self._prompt_manager.get_prompt(name)
933
- if not self._should_enable_component(prompt):
934
- raise DisabledError(f"Prompt {name!r} is disabled")
935
- return await self._prompt_manager.render_prompt(name, arguments)
936
-
937
- # Check mounted servers to see if they have the prompt
938
- # iterate such that new mounts take precedence over older ones
939
- for mounted_server in reversed(self._mounted_servers):
940
- prompt_name = name
941
- try:
942
- if mounted_server.prefix:
943
- # If server has a prefix, check if name matches and strip prefix
944
- if prompt_name.startswith(f"{mounted_server.prefix}_"):
945
- prompt_name = prompt_name.removeprefix(
946
- f"{mounted_server.prefix}_"
947
- )
948
- else:
949
- continue
950
- return await mounted_server.server._middleware_get_prompt(
951
- prompt_name, arguments
952
- )
953
- except NotFoundError:
954
- # Prompt not found on this server, try the next one
955
- continue
956
-
957
- raise NotFoundError(f"Unknown prompt: {name}")
958
-
959
  def add_tool(self, tool: Tool) -> None:
960
  """Add a tool to the server.
961
 
 
341
 
342
  async def get_tools(self) -> dict[str, Tool]:
343
  """Get all registered tools, indexed by registered key."""
344
+ return await self._tool_manager.get_tools()
 
345
 
346
  async def get_tool(self, key: str) -> Tool:
347
  tools = await self.get_tools()
 
375
  """
376
  List all available prompts.
377
  """
378
+ return await self._prompt_manager.get_prompts()
 
 
379
 
380
  async def get_prompt(self, key: str) -> Prompt:
381
  prompts = await self.get_prompts()
 
648
  logger.debug("Handler called: list_prompts")
649
 
650
  with fastmcp.server.context.Context(fastmcp=self):
651
+ prompts = await self._list_prompts()
652
  return [prompt.to_mcp_prompt(name=prompt.key) for prompt in prompts]
653
 
654
+ async def _list_prompts(self) -> list[Prompt]:
655
  """
656
  List all available prompts, in the format expected by the low-level MCP
657
  server.
 
659
  """
660
 
661
  async def _handler(
662
+ context: MiddlewareContext[mcp.types.ListPromptsRequest],
663
  ) -> list[Prompt]:
664
+ prompts = await self._prompt_manager.list_prompts()
665
 
666
  mcp_prompts: list[Prompt] = []
667
  for prompt in prompts:
 
673
  with fastmcp.server.context.Context(fastmcp=self) as fastmcp_ctx:
674
  # Create the middleware context.
675
  mw_context = MiddlewareContext(
676
+ message=mcp.types.ListPromptsRequest(method="prompts/list"),
677
  source="client",
678
  type="request",
679
  method="prompts/list",
 
683
  # Apply the middleware chain.
684
  return await self._apply_middleware(mw_context, _handler)
685
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
686
  async def _mcp_call_tool(
687
  self, key: str, arguments: dict[str, Any]
688
  ) -> list[MCPContent]:
 
788
  Read a resource by URI, in the format expected by the low-level MCP
789
  server.
790
  """
791
+ if await self._resource_manager.has_resource(uri):
792
  resource = await self._resource_manager.get_resource(uri)
793
  if not self._should_enable_component(resource):
794
  raise DisabledError(f"Resource {str(uri)!r} is disabled")
 
800
  )
801
  ]
802
  else:
803
+ raise NotFoundError(f"Unknown resource: {uri}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
804
 
805
  async def _mcp_get_prompt(
806
  self, name: str, arguments: dict[str, Any] | None = None
 
814
 
815
  with fastmcp.server.context.Context(fastmcp=self):
816
  try:
817
+ return await self._get_prompt(name, arguments)
818
  except DisabledError:
819
  # convert to NotFoundError to avoid leaking prompt presence
820
  raise NotFoundError(f"Unknown prompt: {name}")
 
822
  # standardize NotFound message
823
  raise NotFoundError(f"Unknown prompt: {name}")
824
 
825
+ async def _get_prompt(
826
+ self, name: str, arguments: dict[str, Any] | None = None
 
 
827
  ) -> GetPromptResult:
828
  """
829
+ Applies this server's middleware and delegates the filtered call to the manager.
830
  """
831
 
832
  async def _handler(
833
  context: MiddlewareContext[mcp.types.GetPromptRequestParams],
834
  ) -> GetPromptResult:
835
+ prompt = await self._prompt_manager.get_prompt(context.message.name)
836
+ if not self._should_enable_component(prompt):
837
+ raise NotFoundError(f"Unknown prompt: {context.message.name!r}")
838
+
839
+ return await self._prompt_manager.render_prompt(
840
+ name=context.message.name, arguments=context.message.arguments
841
  )
842
 
843
  mw_context = MiddlewareContext(
 
849
  )
850
  return await self._apply_middleware(mw_context, _handler)
851
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
852
  def add_tool(self, tool: Tool) -> None:
853
  """Add a tool to the server.
854