fix(177): Fix prompt listing across agents (#178)

* fix(177): Fix prompt listing across agents

* fix per-agent prompt listing and selection

* bit more code tidy-up

---------

Co-authored-by: evalstate <[email protected]>

Author: wreed4

src/mcp_agent/agents/agent.py

@@ -75,27 +75,11 @@ async def prompt(self, default_prompt: str = "", agent_name: Optional[str] = Non
         async def send_wrapper(message, agent_name):
             return await self.send(message)
 
-        # Define wrapper for apply_prompt function
-        async def apply_prompt_wrapper(prompt_name, args, agent_name):
-            # Just apply the prompt directly
-            return await self.apply_prompt(prompt_name, args)
-
-        # Define wrapper for list_prompts function
-        async def list_prompts_wrapper(agent_name):
-            # Always call list_prompts on this agent regardless of agent_name
-            return await self.list_prompts()
-
-        # Define wrapper for list_resources function
-        async def list_resources_wrapper(agent_name):
-            # Always call list_resources on this agent regardless of agent_name
-            return await self.list_resources()
-
         # Start the prompt loop with just this agent
         return await prompt.prompt_loop(
             send_func=send_wrapper,
             default_agent=agent_name_str,
             available_agents=[agent_name_str],  # Only this agent
-            apply_prompt_func=apply_prompt_wrapper,
-            list_prompts_func=list_prompts_wrapper,
+            prompt_provider=self,  # Pass self as the prompt provider since we implement the protocol
             default=default_prompt,
         )

src/mcp_agent/agents/base_agent.py

@@ -456,6 +456,7 @@ async def apply_prompt(
         self,
         prompt_name: str,
         arguments: Dict[str, str] | None = None,
+        agent_name: str | None = None,
         server_name: str | None = None,
     ) -> str:
         """
@@ -468,6 +469,7 @@ async def apply_prompt(
         Args:
             prompt_name: The name of the prompt to apply
             arguments: Optional dictionary of string arguments to pass to the prompt template
+            agent_name: Optional agent name (ignored at this level, used by multi-agent apps)
             server_name: Optional name of the server to get the prompt from
 
         Returns:

src/mcp_agent/core/agent_app.py

@@ -129,6 +129,12 @@ async def list_prompts(self, server_name: str | None = None, agent_name: str | N
         Returns:
             Dictionary mapping server names to lists of available prompts
         """
+        if not agent_name:
+            results = {}
+            for agent in self._agents.values():
+                curr_prompts = await agent.list_prompts(server_name=server_name)
+                results.update(curr_prompts)
+            return results
         return await self._agent(agent_name).list_prompts(server_name=server_name)
 
     async def get_prompt(
@@ -262,7 +268,6 @@ async def send_wrapper(message, agent_name):
             send_func=send_wrapper,
             default_agent=target_name,  # Pass the agent name, not the agent object
             available_agents=list(self._agents.keys()),
-            apply_prompt_func=self.apply_prompt,
-            list_prompts_func=self.list_prompts,
+            prompt_provider=self,  # Pass self as the prompt provider
             default=default_prompt,
         )

src/mcp_agent/core/interactive_prompt.py

@@ -10,12 +10,13 @@
         send_func=agent_app.send,
         default_agent="default_agent",
         available_agents=["agent1", "agent2"],
-        apply_prompt_func=agent_app.apply_prompt
+        prompt_provider=agent_app
     )
 """
 
-from typing import Dict, List, Optional
+from typing import Awaitable, Callable, Dict, List, Mapping, Optional, Protocol, Union
 
+from mcp.types import Prompt, PromptMessage
 from rich import print as rich_print
 from rich.console import Console
 from rich.table import Table
@@ -28,8 +29,24 @@
     handle_special_commands,
 )
 from mcp_agent.mcp.mcp_aggregator import SEP  # Import SEP once at the top
+from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 from mcp_agent.progress_display import progress_display
 
+# Type alias for the send function
+SendFunc = Callable[[Union[str, PromptMessage, PromptMessageMultipart], str], Awaitable[str]]
+
+
+class PromptProvider(Protocol):
+    """Protocol for objects that can provide prompt functionality."""
+    
+    async def list_prompts(self, server_name: Optional[str] = None, agent_name: Optional[str] = None) -> Mapping[str, List[Prompt]]:
+        """List available prompts."""
+        ...
+    
+    async def apply_prompt(self, prompt_name: str, arguments: Optional[Dict[str, str]] = None, agent_name: Optional[str] = None, **kwargs) -> str:
+        """Apply a prompt."""
+        ...
+
 
 class InteractivePrompt:
     """
@@ -48,22 +65,20 @@ def __init__(self, agent_types: Optional[Dict[str, AgentType]] = None) -> None:
 
     async def prompt_loop(
         self,
-        send_func,
+        send_func: SendFunc,
         default_agent: str,
         available_agents: List[str],
-        apply_prompt_func=None,
-        list_prompts_func=None,
+        prompt_provider: Optional[PromptProvider] = None,
         default: str = "",
     ) -> str:
         """
         Start an interactive prompt session.
 
         Args:
-            send_func: Function to send messages to agents (signature: async (message, agent_name))
+            send_func: Function to send messages to agents
             default_agent: Name of the default agent to use
             available_agents: List of available agent names
-            apply_prompt_func: Optional function to apply prompts (signature: async (name, args, agent))
-            list_prompts_func: Optional function to list available prompts (signature: async (agent_name))
+            prompt_provider: Optional provider that implements list_prompts and apply_prompt
             default: Default message to use when user presses enter
 
         Returns:
@@ -110,21 +125,19 @@ async def prompt_loop(
                             rich_print(f"[red]Agent '{new_agent}' not found[/red]")
                             continue
                     # Keep the existing list_prompts handler for backward compatibility
-                    elif "list_prompts" in command_result and list_prompts_func:
-                        # Use the list_prompts_func directly
-                        await self._list_prompts(list_prompts_func, agent)
+                    elif "list_prompts" in command_result and prompt_provider:
+                        # Use the prompt_provider directly
+                        await self._list_prompts(prompt_provider, agent)
                         continue
-                    elif "select_prompt" in command_result and (
-                        list_prompts_func and apply_prompt_func
-                    ):
+                    elif "select_prompt" in command_result and prompt_provider:
                         # Handle prompt selection, using both list_prompts and apply_prompt
                         prompt_name = command_result.get("prompt_name")
                         prompt_index = command_result.get("prompt_index")
 
                         # If a specific index was provided (from /prompt <number>)
                         if prompt_index is not None:
                             # First get a list of all prompts to look up the index
-                            all_prompts = await self._get_all_prompts(list_prompts_func, agent)
+                            all_prompts = await self._get_all_prompts(prompt_provider, agent)
                             if not all_prompts:
                                 rich_print("[yellow]No prompts available[/yellow]")
                                 continue
@@ -135,8 +148,7 @@ async def prompt_loop(
                                 selected_prompt = all_prompts[prompt_index - 1]
                                 # Use the already created namespaced_name to ensure consistency
                                 await self._select_prompt(
-                                    list_prompts_func,
-                                    apply_prompt_func,
+                                    prompt_provider,
                                     agent,
                                     selected_prompt["namespaced_name"],
                                 )
@@ -145,11 +157,11 @@ async def prompt_loop(
                                     f"[red]Invalid prompt number: {prompt_index}. Valid range is 1-{len(all_prompts)}[/red]"
                                 )
                                 # Show the prompt list for convenience
-                                await self._list_prompts(list_prompts_func, agent)
+                                await self._list_prompts(prompt_provider, agent)
                         else:
                             # Use the name-based selection
                             await self._select_prompt(
-                                list_prompts_func, apply_prompt_func, agent, prompt_name
+                                prompt_provider, agent, prompt_name
                             )
                         continue
 
@@ -171,21 +183,21 @@ async def prompt_loop(
 
         return result
 
-    async def _get_all_prompts(self, list_prompts_func, agent_name):
+    async def _get_all_prompts(self, prompt_provider: PromptProvider, agent_name: Optional[str] = None):
         """
         Get a list of all available prompts.
 
         Args:
-            list_prompts_func: Function to get available prompts
-            agent_name: Name of the agent
+            prompt_provider: Provider that implements list_prompts
+            agent_name: Optional agent name (for multi-agent apps)
 
         Returns:
             List of prompt info dictionaries, sorted by server and name
         """
         try:
-            # Pass None instead of agent_name to get prompts from all servers
-            # the agent_name parameter should never be used as a server name
-            prompt_servers = await list_prompts_func(None)
+            # Call list_prompts on the provider
+            prompt_servers = await prompt_provider.list_prompts(server_name=None, agent_name=agent_name)
+            
             all_prompts = []
 
             # Process the returned prompt servers
@@ -219,14 +231,18 @@ async def _get_all_prompts(self, list_prompts_func, agent_name):
                                     }
                                 )
                             else:
+                                # Handle Prompt objects from mcp.types
+                                prompt_name = getattr(prompt, "name", str(prompt))
+                                description = getattr(prompt, "description", "No description")
+                                arguments = getattr(prompt, "arguments", [])
                                 all_prompts.append(
                                     {
                                         "server": server_name,
-                                        "name": str(prompt),
-                                        "namespaced_name": f"{server_name}{SEP}{str(prompt)}",
-                                        "description": "No description",
-                                        "arg_count": 0,
-                                        "arguments": [],
+                                        "name": prompt_name,
+                                        "namespaced_name": f"{server_name}{SEP}{prompt_name}",
+                                        "description": description,
+                                        "arg_count": len(arguments),
+                                        "arguments": arguments,
                                     }
                                 )
 
@@ -244,27 +260,22 @@ async def _get_all_prompts(self, list_prompts_func, agent_name):
             rich_print(f"[dim]{traceback.format_exc()}[/dim]")
             return []
 
-    async def _list_prompts(self, list_prompts_func, agent_name) -> None:
+    async def _list_prompts(self, prompt_provider: PromptProvider, agent_name: str) -> None:
         """
         List available prompts for an agent.
 
         Args:
-            list_prompts_func: Function to get available prompts
+            prompt_provider: Provider that implements list_prompts
             agent_name: Name of the agent
         """
-        from rich import print as rich_print
-        from rich.console import Console
-        from rich.table import Table
-
         console = Console()
 
         try:
             # Directly call the list_prompts function for this agent
             rich_print(f"\n[bold]Fetching prompts for agent [cyan]{agent_name}[/cyan]...[/bold]")
 
-            # Get all prompts using the helper function - pass None as server name
-            # to get prompts from all available servers
-            all_prompts = await self._get_all_prompts(list_prompts_func, None)
+            # Get all prompts using the helper function
+            all_prompts = await self._get_all_prompts(prompt_provider, agent_name)
 
             if all_prompts:
                 # Create a table for better display
@@ -300,28 +311,24 @@ async def _list_prompts(self, list_prompts_func, agent_name) -> None:
             rich_print(f"[dim]{traceback.format_exc()}[/dim]")
 
     async def _select_prompt(
-        self, list_prompts_func, apply_prompt_func, agent_name, requested_name=None
+        self, prompt_provider: PromptProvider, agent_name: str, requested_name: Optional[str] = None
     ) -> None:
         """
         Select and apply a prompt.
 
         Args:
-            list_prompts_func: Function to get available prompts
-            apply_prompt_func: Function to apply prompts
+            prompt_provider: Provider that implements list_prompts and apply_prompt
             agent_name: Name of the agent
             requested_name: Optional name of the prompt to apply
         """
-        # We already imported these at the top
-        from rich import print as rich_print
-
         console = Console()
 
         try:
-            # Get all available prompts directly from the list_prompts function
+            # Get all available prompts directly from the prompt provider
             rich_print(f"\n[bold]Fetching prompts for agent [cyan]{agent_name}[/cyan]...[/bold]")
-            # IMPORTANT: list_prompts_func gets MCP server prompts, not agent prompts
-            # So we pass None to get prompts from all servers, not using agent_name as server name
-            prompt_servers = await list_prompts_func(None)
+            
+            # Call list_prompts on the provider
+            prompt_servers = await prompt_provider.list_prompts(server_name=None, agent_name=agent_name)
 
             if not prompt_servers:
                 rich_print("[yellow]No prompts available for this agent[/yellow]")
@@ -542,8 +549,8 @@ async def _select_prompt(
             namespaced_name = selected_prompt["namespaced_name"]
             rich_print(f"\n[bold]Applying prompt [cyan]{namespaced_name}[/cyan]...[/bold]")
 
-            # Call apply_prompt function with the prompt name and arguments
-            await apply_prompt_func(namespaced_name, arg_values, agent_name)
+            # Call apply_prompt on the provider with the prompt name and arguments
+            await prompt_provider.apply_prompt(namespaced_name, arg_values, agent_name)
 
         except Exception as e:
             import traceback

src/mcp_agent/mcp/mcp_aggregator.py

@@ -811,12 +811,13 @@ async def get_prompt(
             messages=[],
         )
 
-    async def list_prompts(self, server_name: str | None = None) -> Mapping[str, List[Prompt]]:
+    async def list_prompts(self, server_name: str | None = None, agent_name: str | None = None) -> Mapping[str, List[Prompt]]:
         """
         List available prompts from one or all servers.
 
         :param server_name: Optional server name to list prompts from. If not provided,
                            lists prompts from all servers.
+        :param agent_name: Optional agent name (ignored at this level, used by multi-agent apps)
         :return: Dictionary mapping server names to lists of Prompt objects
         """
         if not self.initialized:

tests/integration/api/fastagent.config.yaml

@@ -34,6 +34,9 @@ mcp:
     prompts:
       command: "prompt-server"
       args: ["playback.md"]
+    prompts2:
+      command: "prompt-server"
+      args: ["prompt.txt"]
     std_io:
       command: "uv"
       args: ["run", "integration_agent.py", "--server", "--transport", "stdio"]

tests/integration/api/test_prompt_listing.py

@@ -9,34 +9,38 @@
 
 @pytest.mark.integration
 @pytest.mark.asyncio
-async def test_get_all_prompts_with_none_server(fast_agent):
+async def test_multi_agent_prompt_listing(fast_agent):
     """Test the _get_all_prompts function with None as server name."""
     fast = fast_agent
 
-    @fast.agent(name="test", servers=["prompts"])
+    @fast.agent(name="agent1", servers=["prompts"])
+    @fast.agent(name="agent2", servers=["prompts2"])
+    @fast.agent(name="agent3")
     async def agent_function():
         async with fast.run() as agent:
             # Create instance of InteractivePrompt
             prompt_ui = InteractivePrompt()
 
-            # Get the list_prompts function from the agent
-            list_prompts_func = agent.test.list_prompts
-
-            # Call _get_all_prompts directly with None as server name
-            all_prompts = await prompt_ui._get_all_prompts(list_prompts_func, None)
-
-            # Verify we got results
-            assert len(all_prompts) > 0
-
-            # Verify each prompt has the correct format
-            for prompt in all_prompts:
-                assert "server" in prompt
-                assert "name" in prompt
-                assert "namespaced_name" in prompt
-                assert prompt["server"] == "prompts"  # From our test config
-
-                # Check namespace format
-                assert prompt["namespaced_name"] == f"prompts-{prompt['name']}"
+            # Test listing prompts for each agent separately
+            # Agent1 should have prompts from "prompts" server (playback.md -> playback)
+            agent1_prompts = await prompt_ui._get_all_prompts(agent, "agent1")
+            assert len(agent1_prompts) == 1
+            assert agent1_prompts[0]["server"] == "prompts"
+            assert agent1_prompts[0]["name"] == "playback"
+            assert agent1_prompts[0]["description"] == "[USER] user1 assistant1 user2"
+            assert agent1_prompts[0]["arg_count"] == 0
+
+            # Agent2 should have prompts from "prompts2" server (prompt.txt -> prompt)
+            agent2_prompts = await prompt_ui._get_all_prompts(agent, "agent2")
+            assert len(agent2_prompts) == 1
+            assert agent2_prompts[0]["server"] == "prompts2"
+            assert agent2_prompts[0]["name"] == "prompt"
+            assert agent2_prompts[0]["description"] == "this is from the prompt file"
+            assert agent2_prompts[0]["arg_count"] == 0
+
+            # Agent3 should have no prompts (no servers configured)
+            agent3_prompts = await prompt_ui._get_all_prompts(agent, "agent3")
+            assert len(agent3_prompts) == 0
 
     await agent_function()