14
14
)
15
15
"""
16
16
17
+ import asyncio
17
18
from typing import Awaitable , Callable , Dict , List , Mapping , Optional , Protocol , Union
18
19
19
20
from mcp .types import Prompt , PromptMessage
38
39
39
40
class PromptProvider (Protocol ):
40
41
"""Protocol for objects that can provide prompt functionality."""
41
-
42
- async def list_prompts (self , server_name : Optional [str ] = None , agent_name : Optional [str ] = None ) -> Mapping [str , List [Prompt ]]:
42
+
43
+ async def list_prompts (
44
+ self , server_name : Optional [str ] = None , agent_name : Optional [str ] = None
45
+ ) -> Mapping [str , List [Prompt ]]:
43
46
"""List available prompts."""
44
47
...
45
-
46
- async def apply_prompt (self , prompt_name : str , arguments : Optional [Dict [str , str ]] = None , agent_name : Optional [str ] = None , ** kwargs ) -> str :
48
+
49
+ async def apply_prompt (
50
+ self ,
51
+ prompt_name : str ,
52
+ arguments : Optional [Dict [str , str ]] = None ,
53
+ agent_name : Optional [str ] = None ,
54
+ ** kwargs ,
55
+ ) -> str :
47
56
"""Apply a prompt."""
48
57
...
49
58
@@ -160,17 +169,19 @@ async def prompt_loop(
160
169
await self ._list_prompts (prompt_provider , agent )
161
170
else :
162
171
# Use the name-based selection
163
- await self ._select_prompt (
164
- prompt_provider , agent , prompt_name
165
- )
172
+ await self ._select_prompt (prompt_provider , agent , prompt_name )
166
173
continue
167
174
168
175
# Skip further processing if:
169
176
# 1. The command was handled (command_result is truthy)
170
177
# 2. The original input was a dictionary (special command like /prompt)
171
178
# 3. The command result itself is a dictionary (special command handling result)
172
179
# This fixes the issue where /prompt without arguments gets sent to the LLM
173
- if command_result or isinstance (user_input , dict ) or isinstance (command_result , dict ):
180
+ if (
181
+ command_result
182
+ or isinstance (user_input , dict )
183
+ or isinstance (command_result , dict )
184
+ ):
174
185
continue
175
186
176
187
if user_input .upper () == "STOP" :
@@ -179,11 +190,45 @@ async def prompt_loop(
179
190
continue
180
191
181
192
# Send the message to the agent
182
- result = await send_func (user_input , agent )
193
+ try :
194
+ result = await send_func (user_input , agent )
195
+ except KeyboardInterrupt :
196
+ rich_print ("\n [yellow]Request cancelled by user (Ctrl+C).[/yellow]" )
197
+ result = "" # Ensure result has a benign value for the loop
198
+ # Attempt to stop progress display safely
199
+ try :
200
+ # For rich.progress.Progress, 'progress_display.live.is_started' is a common check
201
+ if hasattr (progress_display , "live" ) and progress_display .live .is_started :
202
+ progress_display .stop ()
203
+ # Fallback for older rich or different progress setup
204
+ elif hasattr (progress_display , "is_running" ) and progress_display .is_running :
205
+ progress_display .stop ()
206
+ else : # If unsure, try stopping directly if stop() is available
207
+ if hasattr (progress_display , "stop" ):
208
+ progress_display .stop ()
209
+ except Exception :
210
+ pass # Continue anyway, don't let progress display crash the cancel
211
+ continue
212
+ except asyncio .CancelledError :
213
+ rich_print ("\n [yellow]Request task was cancelled.[/yellow]" )
214
+ result = ""
215
+ try :
216
+ if hasattr (progress_display , "live" ) and progress_display .live .is_started :
217
+ progress_display .stop ()
218
+ elif hasattr (progress_display , "is_running" ) and progress_display .is_running :
219
+ progress_display .stop ()
220
+ else :
221
+ if hasattr (progress_display , "stop" ):
222
+ progress_display .stop ()
223
+ except Exception :
224
+ pass
225
+ continue
183
226
184
227
return result
185
228
186
- async def _get_all_prompts (self , prompt_provider : PromptProvider , agent_name : Optional [str ] = None ):
229
+ async def _get_all_prompts (
230
+ self , prompt_provider : PromptProvider , agent_name : Optional [str ] = None
231
+ ):
187
232
"""
188
233
Get a list of all available prompts.
189
234
@@ -196,8 +241,10 @@ async def _get_all_prompts(self, prompt_provider: PromptProvider, agent_name: Op
196
241
"""
197
242
try :
198
243
# Call list_prompts on the provider
199
- prompt_servers = await prompt_provider .list_prompts (server_name = None , agent_name = agent_name )
200
-
244
+ prompt_servers = await prompt_provider .list_prompts (
245
+ server_name = None , agent_name = agent_name
246
+ )
247
+
201
248
all_prompts = []
202
249
203
250
# Process the returned prompt servers
@@ -326,9 +373,11 @@ async def _select_prompt(
326
373
try :
327
374
# Get all available prompts directly from the prompt provider
328
375
rich_print (f"\n [bold]Fetching prompts for agent [cyan]{ agent_name } [/cyan]...[/bold]" )
329
-
376
+
330
377
# Call list_prompts on the provider
331
- prompt_servers = await prompt_provider .list_prompts (server_name = None , agent_name = agent_name )
378
+ prompt_servers = await prompt_provider .list_prompts (
379
+ server_name = None , agent_name = agent_name
380
+ )
332
381
333
382
if not prompt_servers :
334
383
rich_print ("[yellow]No prompts available for this agent[/yellow]" )
0 commit comments