feat: improve interactive CLI Ctrl+C handling

- Modify enhanced_prompt.py to allow session continuation when Ctrl+C
  is pressed at the input prompt.
- Modify interactive_prompt.py to catch KeyboardInterrupt and
  asyncio.CancelledError during agent tasks, allowing the session to
  continue after task cancellation.

Author: monotykamary

src/mcp_agent/core/enhanced_prompt.py

@@ -320,10 +320,12 @@ def pre_process_input(text):
         result = await session.prompt_async(HTML(prompt_text), default=default)
         return pre_process_input(result)
     except KeyboardInterrupt:
-        # Handle Ctrl+C gracefully
-        return "STOP"
+        # Handle Ctrl+C gracefully at the prompt
+        rich_print("\n[yellow]Input cancelled. Type a command or 'STOP' to exit session.[/yellow]")
+        return ""  # Return empty string to re-prompt
     except EOFError:
         # Handle Ctrl+D gracefully
+        rich_print("\n[yellow]EOF received. Type 'STOP' to exit session.[/yellow]")
         return "STOP"
     except Exception as e:
         # Log and gracefully handle other exceptions

src/mcp_agent/core/interactive_prompt.py

@@ -14,6 +14,7 @@
     )
 """
 
+import asyncio
 from typing import Awaitable, Callable, Dict, List, Mapping, Optional, Protocol, Union
 
 from mcp.types import Prompt, PromptMessage
@@ -38,12 +39,20 @@
 
 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]]:
+
+    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:
+
+    async def apply_prompt(
+        self,
+        prompt_name: str,
+        arguments: Optional[Dict[str, str]] = None,
+        agent_name: Optional[str] = None,
+        **kwargs,
+    ) -> str:
         """Apply a prompt."""
         ...
 
@@ -160,17 +169,19 @@ async def prompt_loop(
                                 await self._list_prompts(prompt_provider, agent)
                         else:
                             # Use the name-based selection
-                            await self._select_prompt(
-                                prompt_provider, agent, prompt_name
-                            )
+                            await self._select_prompt(prompt_provider, agent, prompt_name)
                         continue
 
                 # Skip further processing if:
                 # 1. The command was handled (command_result is truthy)
                 # 2. The original input was a dictionary (special command like /prompt)
                 # 3. The command result itself is a dictionary (special command handling result)
                 # This fixes the issue where /prompt without arguments gets sent to the LLM
-                if command_result or isinstance(user_input, dict) or isinstance(command_result, dict):
+                if (
+                    command_result
+                    or isinstance(user_input, dict)
+                    or isinstance(command_result, dict)
+                ):
                     continue
 
                 if user_input.upper() == "STOP":
@@ -179,11 +190,45 @@ async def prompt_loop(
                     continue
 
             # Send the message to the agent
-            result = await send_func(user_input, agent)
+            try:
+                result = await send_func(user_input, agent)
+            except KeyboardInterrupt:
+                rich_print("\n[yellow]Request cancelled by user (Ctrl+C).[/yellow]")
+                result = ""  # Ensure result has a benign value for the loop
+                # Attempt to stop progress display safely
+                try:
+                    # For rich.progress.Progress, 'progress_display.live.is_started' is a common check
+                    if hasattr(progress_display, "live") and progress_display.live.is_started:
+                        progress_display.stop()
+                    # Fallback for older rich or different progress setup
+                    elif hasattr(progress_display, "is_running") and progress_display.is_running:
+                        progress_display.stop()
+                    else:  # If unsure, try stopping directly if stop() is available
+                        if hasattr(progress_display, "stop"):
+                            progress_display.stop()
+                except Exception:
+                    pass  # Continue anyway, don't let progress display crash the cancel
+                continue
+            except asyncio.CancelledError:
+                rich_print("\n[yellow]Request task was cancelled.[/yellow]")
+                result = ""
+                try:
+                    if hasattr(progress_display, "live") and progress_display.live.is_started:
+                        progress_display.stop()
+                    elif hasattr(progress_display, "is_running") and progress_display.is_running:
+                        progress_display.stop()
+                    else:
+                        if hasattr(progress_display, "stop"):
+                            progress_display.stop()
+                except Exception:
+                    pass
+                continue
 
         return result
 
-    async def _get_all_prompts(self, prompt_provider: PromptProvider, agent_name: Optional[str] = None):
+    async def _get_all_prompts(
+        self, prompt_provider: PromptProvider, agent_name: Optional[str] = None
+    ):
         """
         Get a list of all available prompts.
 
@@ -196,8 +241,10 @@ async def _get_all_prompts(self, prompt_provider: PromptProvider, agent_name: Op
         """
         try:
             # Call list_prompts on the provider
-            prompt_servers = await prompt_provider.list_prompts(server_name=None, agent_name=agent_name)
-            
+            prompt_servers = await prompt_provider.list_prompts(
+                server_name=None, agent_name=agent_name
+            )
+
             all_prompts = []
 
             # Process the returned prompt servers
@@ -326,9 +373,11 @@ async def _select_prompt(
         try:
             # Get all available prompts directly from the prompt provider
             rich_print(f"\n[bold]Fetching prompts for agent [cyan]{agent_name}[/cyan]...[/bold]")
-            
+
             # Call list_prompts on the provider
-            prompt_servers = await prompt_provider.list_prompts(server_name=None, agent_name=agent_name)
+            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]")

src/mcp_agent/llm/augmented_llm_slow.py

@@ -0,0 +1,42 @@
+import asyncio
+from typing import Any, List, Optional, Union
+
+from mcp_agent.llm.augmented_llm import (
+    MessageParamT,
+    RequestParams,
+)
+from mcp_agent.llm.augmented_llm_passthrough import PassthroughLLM
+from mcp_agent.llm.provider_types import Provider
+from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
+
+
+class SlowLLM(PassthroughLLM):
+    """
+    A specialized LLM implementation that sleeps for 3 seconds before responding like PassthroughLLM.
+
+    This is useful for testing scenarios where you want to simulate slow responses
+    or for debugging timing-related issues in parallel workflows.
+    """
+
+    def __init__(
+        self, provider=Provider.FAST_AGENT, name: str = "Slow", **kwargs: dict[str, Any]
+    ) -> None:
+        super().__init__(name=name, provider=provider, **kwargs)
+
+    async def generate_str(
+        self,
+        message: Union[str, MessageParamT, List[MessageParamT]],
+        request_params: Optional[RequestParams] = None,
+    ) -> str:
+        """Sleep for 3 seconds then return the input message as a string."""
+        await asyncio.sleep(3)
+        return await super().generate_str(message, request_params)
+
+    async def _apply_prompt_provider_specific(
+        self,
+        multipart_messages: List["PromptMessageMultipart"],
+        request_params: RequestParams | None = None,
+    ) -> PromptMessageMultipart:
+        """Sleep for 3 seconds then apply prompt like PassthroughLLM."""
+        await asyncio.sleep(3)
+        return await super()._apply_prompt_provider_specific(multipart_messages, request_params)

src/mcp_agent/llm/model_factory.py

@@ -8,6 +8,7 @@
 from mcp_agent.core.request_params import RequestParams
 from mcp_agent.llm.augmented_llm_passthrough import PassthroughLLM
 from mcp_agent.llm.augmented_llm_playback import PlaybackLLM
+from mcp_agent.llm.augmented_llm_slow import SlowLLM
 from mcp_agent.llm.provider_types import Provider
 from mcp_agent.llm.providers.augmented_llm_anthropic import AnthropicAugmentedLLM
 from mcp_agent.llm.providers.augmented_llm_azure import AzureOpenAIAugmentedLLM
@@ -29,6 +30,7 @@
     Type[OpenAIAugmentedLLM],
     Type[PassthroughLLM],
     Type[PlaybackLLM],
+    Type[SlowLLM],
     Type[DeepSeekAugmentedLLM],
     Type[OpenRouterAugmentedLLM],
     Type[TensorZeroAugmentedLLM],
@@ -73,6 +75,7 @@ class ModelFactory:
     DEFAULT_PROVIDERS = {
         "passthrough": Provider.FAST_AGENT,
         "playback": Provider.FAST_AGENT,
+        "slow": Provider.FAST_AGENT,
         "gpt-4o": Provider.OPENAI,
         "gpt-4o-mini": Provider.OPENAI,
         "gpt-4.1": Provider.OPENAI,
@@ -139,6 +142,7 @@ class ModelFactory:
     # This overrides the provider-based class selection
     MODEL_SPECIFIC_CLASSES: Dict[str, LLMClass] = {
         "playback": PlaybackLLM,
+        "slow": SlowLLM,
     }
 
     @classmethod

tests/integration/sampling/fastagent.config.yaml

@@ -23,7 +23,11 @@ mcp:
       args: ["run", "sampling_test_server.py"]
       sampling:
         model: "passthrough"
-    
+    slow_sampling:
+      command: "uv"
+      args: ["run", "sampling_test_server.py"]
+      sampling:
+        model: "slow"
     sampling_test_no_config:
       command: "uv"
       args: ["run", "sampling_test_server.py"]

tests/integration/sampling/live.py

@@ -7,13 +7,16 @@
 
 
 # Define the agent
-@fast.agent(servers=["sampling_test"])
+@fast.agent(servers=["sampling_test", "slow_sampling"])
 async def main():
     # use the --model command line switch or agent arguments to change model
     async with fast.run() as agent:
         result = await agent.send('***CALL_TOOL sampling_test-sample {"to_sample": "123foo"}')
         print(f"RESULT: {result}")
 
+        result = await agent.send('***CALL_TOOL slow_sampling-sample_parallel')
+        print(f"RESULT: {result}")
+
 
 if __name__ == "__main__":
     asyncio.run(main())

tests/integration/sampling/sampling_test_server.py

@@ -61,6 +61,47 @@ async def sample_many(ctx: Context) -> CallToolResult:
     return CallToolResult(content=[TextContent(type="text", text=str(result))])
 
 
+@mcp.tool()
+async def sample_parallel(ctx: Context, count: int = 5) -> CallToolResult:
+    """Tool that makes multiple concurrent sampling requests to test parallel processing"""
+    try:
+        logger.info(f"Making {count} concurrent sampling requests")
+
+        # Create multiple concurrent sampling requests
+        import asyncio
+
+        async def _send_sampling(request: int):
+            return await ctx.session.create_message(
+                max_tokens=100,
+                messages=[SamplingMessage(
+                    role="user",
+                    content=TextContent(type="text", text=f"Parallel request {request+1}")
+                )],
+            )
+
+
+        tasks = []
+        for i in range(count):
+            task = _send_sampling(i)
+            tasks.append(task)
+
+        # Execute all requests concurrently
+        results = await asyncio.gather(*[_send_sampling(i) for i in range(count)])
+
+        # Combine results
+        response_texts = [result.content.text for result in results]
+        combined_response = f"Completed {len(results)} parallel requests: " + ", ".join(response_texts[:3])
+        if len(response_texts) > 3:
+            combined_response += f"... and {len(response_texts) - 3} more"
+
+        logger.info(f"Parallel sampling completed: {combined_response}")
+        return CallToolResult(content=[TextContent(type="text", text=combined_response)])
+
+    except Exception as e:
+        logger.error(f"Error in sample_parallel tool: {e}", exc_info=True)
+        return CallToolResult(isError=True, content=[TextContent(type="text", text=f"Error: {str(e)}")])
+
+
 if __name__ == "__main__":
     logger.info("Starting sampling test server...")
     mcp.run()