Skip to content

Commit b37ad4f

Browse files
daviwildaviwil
committed
Add improved prompt cancellation handling
This change improves the current handling of PSHostUserInterface prompt cancellation so that cancelled prompts are cleaned up effectively and don't block further command execution.
1 parent 8589a61 commit b37ad4f

File tree

10 files changed

+212
-103
lines changed

10 files changed

+212
-103
lines changed

‎src/PowerShellEditorServices.Protocol/MessageProtocol/MessageDispatcher.cs

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -234,12 +234,14 @@ protected async Task DispatchMessage(
234234
Message messageToDispatch,
235235
MessageWriter messageWriter)
236236
{
237+
Task handlerToAwait = null;
238+
237239
if (messageToDispatch.MessageType == MessageType.Request)
238240
{
239241
Func<Message, MessageWriter, Task> requestHandler = null;
240242
if (this.requestHandlers.TryGetValue(messageToDispatch.Method, out requestHandler))
241243
{
242-
await requestHandler(messageToDispatch, messageWriter);
244+
handlerToAwait = requestHandler(messageToDispatch, messageWriter);
243245
}
244246
else
245247
{
@@ -258,7 +260,7 @@ protected async Task DispatchMessage(
258260
Func<Message, MessageWriter, Task> eventHandler = null;
259261
if (this.eventHandlers.TryGetValue(messageToDispatch.Method, out eventHandler))
260262
{
261-
await eventHandler(messageToDispatch, messageWriter);
263+
handlerToAwait = eventHandler(messageToDispatch, messageWriter);
262264
}
263265
else
264266
{
@@ -269,6 +271,28 @@ protected async Task DispatchMessage(
269271
{
270272
// TODO: Return message not supported
271273
}
274+
275+
if (handlerToAwait != null)
276+
{
277+
try
278+
{
279+
await handlerToAwait;
280+
}
281+
catch (TaskCanceledException)
282+
{
283+
// Some tasks may be cancelled due to legitimate
284+
// timeouts so don't let those exceptions go higher.
285+
}
286+
catch (AggregateException e)
287+
{
288+
if (!(e.InnerExceptions[0] is TaskCanceledException))
289+
{
290+
// Cancelled tasks aren't a problem, so rethrow
291+
// anything that isn't a TaskCanceledException
292+
throw e;
293+
}
294+
}
295+
}
272296
}
273297

274298
private void OnListenTaskCompleted(Task listenTask)

‎src/PowerShellEditorServices/Console/ChoicePromptHandler.cs

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
//
55

66
using System;
7+
using System.Threading;
78
using System.Threading.Tasks;
89

910
namespace Microsoft.PowerShell.EditorServices.Console
@@ -32,7 +33,7 @@ public enum PromptStyle
3233
/// that present the user a set of options from which a selection
3334
/// should be made.
3435
/// </summary>
35-
public abstract class ChoicePromptHandler : IPromptHandler
36+
public abstract class ChoicePromptHandler : PromptHandler
3637
{
3738
#region Private Fields
3839

@@ -82,6 +83,9 @@ public abstract class ChoicePromptHandler : IPromptHandler
8283
/// <param name="defaultChoice">
8384
/// The default choice to highlight for the user.
8485
/// </param>
86+
/// <param name="cancellationToken">
87+
/// A CancellationToken that can be used to cancel the prompt.
88+
/// </param>
8589
/// <returns>
8690
/// A Task instance that can be monitored for completion to get
8791
/// the user's choice.
@@ -90,7 +94,8 @@ public Task<int> PromptForChoice(
9094
string promptCaption,
9195
string promptMessage,
9296
ChoiceDetails[] choices,
93-
int defaultChoice)
97+
int defaultChoice,
98+
CancellationToken cancellationToken)
9499
{
95100
// TODO: Guard against multiple calls
96101

@@ -100,6 +105,9 @@ public Task<int> PromptForChoice(
100105
this.DefaultChoice = defaultChoice;
101106
this.promptTask = new TaskCompletionSource<int>();
102107

108+
// Cancel the TaskCompletionSource if the caller cancels the task
109+
cancellationToken.Register(this.CancelPrompt, true);
110+
103111
// Show the prompt to the user
104112
this.ShowPrompt(PromptStyle.Full);
105113

@@ -114,7 +122,7 @@ public Task<int> PromptForChoice(
114122
/// True if the prompt is complete, false if the prompt is
115123
/// still waiting for a valid response.
116124
/// </returns>
117-
public virtual bool HandleResponse(string responseString)
125+
public override bool HandleResponse(string responseString)
118126
{
119127
int choiceIndex = -1;
120128

@@ -145,7 +153,7 @@ public virtual bool HandleResponse(string responseString)
145153
/// <summary>
146154
/// Called when the active prompt should be cancelled.
147155
/// </summary>
148-
public void CancelPrompt()
156+
protected override void OnPromptCancelled()
149157
{
150158
// Cancel the prompt task
151159
this.promptTask.TrySetCanceled();

‎src/PowerShellEditorServices/Console/ConsoleService.cs

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
using Microsoft.PowerShell.EditorServices.Utility;
77
using System;
88
using System.Collections.Generic;
9-
using System.Threading.Tasks;
109

1110
namespace Microsoft.PowerShell.EditorServices.Console
1211
{
@@ -20,7 +19,7 @@ public class ConsoleService : IConsoleHost
2019

2120
private PowerShellContext powerShellContext;
2221

23-
private IPromptHandler activePromptHandler;
22+
private PromptHandler activePromptHandler;
2423
private Stack<IPromptHandlerContext> promptHandlerContextStack =
2524
new Stack<IPromptHandlerContext>();
2625

@@ -193,7 +192,7 @@ void IConsoleHost.UpdateProgress(long sourceId, ProgressDetails progressDetails)
193192

194193
void IConsoleHost.ExitSession(int exitCode)
195194
{
196-
throw new NotImplementedException();
195+
//throw new NotImplementedException();
197196
}
198197

199198
ChoicePromptHandler IConsoleHost.GetChoicePromptHandler()
@@ -210,7 +209,7 @@ InputPromptHandler IConsoleHost.GetInputPromptHandler()
210209

211210
private TPromptHandler GetPromptHandler<TPromptHandler>(
212211
Func<IPromptHandlerContext, TPromptHandler> factoryInvoker)
213-
where TPromptHandler : IPromptHandler
212+
where TPromptHandler : PromptHandler
214213
{
215214
if (this.activePromptHandler != null)
216215
{
@@ -225,11 +224,23 @@ private TPromptHandler GetPromptHandler<TPromptHandler>(
225224

226225
TPromptHandler promptHandler = factoryInvoker(promptHandlerContext);
227226
this.activePromptHandler = promptHandler;
227+
this.activePromptHandler.PromptCancelled += activePromptHandler_PromptCancelled;
228228

229229
return promptHandler;
230230
}
231231

232232
#endregion
233+
234+
#region Event Handlers
235+
236+
private void activePromptHandler_PromptCancelled(object sender, EventArgs e)
237+
{
238+
// Clean up the existing prompt
239+
this.activePromptHandler.PromptCancelled -= activePromptHandler_PromptCancelled;
240+
this.activePromptHandler = null;
241+
}
242+
243+
#endregion
233244
}
234245
}
235246

‎src/PowerShellEditorServices/Console/IPromptHandler.cs

Lines changed: 0 additions & 29 deletions
This file was deleted.

‎src/PowerShellEditorServices/Console/InputPromptHandler.cs

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
using System.Globalization;
1010
using System.Linq;
1111
using System.Management.Automation;
12+
using System.Threading;
1213
using System.Threading.Tasks;
1314

1415
namespace Microsoft.PowerShell.EditorServices.Console
@@ -18,7 +19,7 @@ namespace Microsoft.PowerShell.EditorServices.Console
1819
/// that present the user a set of fields for which values
1920
/// should be entered.
2021
/// </summary>
21-
public abstract class InputPromptHandler : IPromptHandler
22+
public abstract class InputPromptHandler : PromptHandler
2223
{
2324
#region Private Fields
2425

@@ -51,13 +52,15 @@ public abstract class InputPromptHandler : IPromptHandler
5152
/// A Task instance that can be monitored for completion to get
5253
/// the user's input.
5354
/// </returns>
54-
public Task<string> PromptForInput()
55+
public Task<string> PromptForInput(
56+
CancellationToken cancellationToken)
5557
{
5658
Task<Dictionary<string, object>> innerTask =
5759
this.PromptForInput(
5860
null,
5961
null,
60-
new FieldDetails[] { new FieldDetails("", "", typeof(string), false, "") });
62+
new FieldDetails[] { new FieldDetails("", "", typeof(string), false, "") },
63+
cancellationToken);
6164

6265
return
6366
innerTask.ContinueWith<string>(
@@ -67,6 +70,10 @@ public Task<string> PromptForInput()
6770
{
6871
throw task.Exception;
6972
}
73+
else if (task.IsCanceled)
74+
{
75+
throw new TaskCanceledException(task);
76+
}
7077

7178
// Return the value of the sole field
7279
return (string)task.Result[""];
@@ -86,17 +93,24 @@ public Task<string> PromptForInput()
8693
/// An array of FieldDetails items to be displayed which prompt the
8794
/// user for input of a specific type.
8895
/// </param>
96+
/// <param name="cancellationToken">
97+
/// A CancellationToken that can be used to cancel the prompt.
98+
/// </param>
8999
/// <returns>
90100
/// A Task instance that can be monitored for completion to get
91101
/// the user's input.
92102
/// </returns>
93103
public Task<Dictionary<string, object>> PromptForInput(
94104
string promptCaption,
95105
string promptMessage,
96-
FieldDetails[] fields)
106+
FieldDetails[] fields,
107+
CancellationToken cancellationToken)
97108
{
98109
this.promptTask = new TaskCompletionSource<Dictionary<string, object>>();
99110

111+
// Cancel the TaskCompletionSource if the caller cancels the task
112+
cancellationToken.Register(this.CancelPrompt, true);
113+
100114
this.Fields = fields;
101115

102116
this.ShowPromptMessage(promptCaption, promptMessage);
@@ -113,7 +127,7 @@ public Task<Dictionary<string, object>> PromptForInput(
113127
/// True if the prompt is complete, false if the prompt is
114128
/// still waiting for a valid response.
115129
/// </returns>
116-
public bool HandleResponse(string responseString)
130+
public override bool HandleResponse(string responseString)
117131
{
118132
if (this.currentField == null)
119133
{
@@ -188,7 +202,7 @@ public bool HandleResponse(string responseString)
188202
/// <summary>
189203
/// Called when the active prompt should be cancelled.
190204
/// </summary>
191-
public void CancelPrompt()
205+
protected override void OnPromptCancelled()
192206
{
193207
// Cancel the prompt task
194208
this.promptTask.TrySetCanceled();

‎src/PowerShellEditorServices/Console/PromptHandler.cs

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
//
2+
// Copyright (c) Microsoft. All rights reserved.
3+
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
4+
//
5+
6+
using System;
7+
8+
namespace Microsoft.PowerShell.EditorServices.Console
9+
{
10+
/// <summary>
11+
/// Defines an abstract base class for prompt handler implementations.
12+
/// </summary>
13+
public abstract class PromptHandler
14+
{
15+
/// <summary>
16+
/// Implements behavior to handle the user's response.
17+
/// </summary>
18+
/// <param name="responseString">The string representing the user's response.</param>
19+
/// <returns>
20+
/// True if the prompt is complete, false if the prompt is
21+
/// still waiting for a valid response.
22+
/// </returns>
23+
public abstract bool HandleResponse(string responseString);
24+
25+
/// <summary>
26+
/// Called when the active prompt should be cancelled.
27+
/// </summary>
28+
public void CancelPrompt()
29+
{
30+
// Allow the implementation to clean itself up
31+
this.OnPromptCancelled();
32+
33+
if (this.PromptCancelled != null)
34+
{
35+
this.PromptCancelled(this, new EventArgs());
36+
}
37+
}
38+
39+
/// <summary>
40+
/// An event that gets raised if the prompt is cancelled, either
41+
/// by the user or due to a timeout.
42+
/// </summary>
43+
public event EventHandler PromptCancelled;
44+
45+
/// <summary>
46+
/// Implementation classes may override this method to perform
47+
/// cleanup when the CancelPrompt method gets called.
48+
/// </summary>
49+
protected virtual void OnPromptCancelled()
50+
{
51+
}
52+
}
53+
}
54+

‎src/PowerShellEditorServices/PowerShellEditorServices.csproj

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,8 @@
5757
<Compile Include="Console\ConsoleService.cs" />
5858
<Compile Include="Console\FieldDetails.cs" />
5959
<Compile Include="Console\InputPromptHandler.cs" />
60-
<Compile Include="Console\IPromptHandler.cs" />
6160
<Compile Include="Console\IPromptHandlerContext.cs" />
61+
<Compile Include="Console\PromptHandler.cs" />
6262
<Compile Include="Debugging\BreakpointDetails.cs" />
6363
<Compile Include="Debugging\DebugService.cs" />
6464
<Compile Include="Debugging\StackFrameDetails.cs" />

0 commit comments

Comments
 (0)