Support Ctrl+C handling

src/System.CommandLine.Tests/InvocationPipelineTests.cs

@@ -6,6 +6,7 @@
 using FluentAssertions;
 using System.Linq;
 using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
 
@@ -76,9 +77,8 @@ public void When_command_handler_throws_then_InvokeAsync_does_not_handle_the_exc
             Func<Task> invoke = async () => await parser.InvokeAsync("the-command", _console);
 
             invoke.Should()
-                  .Throw<TargetInvocationException>()
+                  .Throw<Exception>()
                   .Which
-                  .InnerException
                   .Message
                   .Should()
                   .Be("oops!");
@@ -139,5 +139,63 @@ public async Task Invocation_can_be_short_circuited_by_middleware_by_not_calling
             middlewareWasCalled.Should().BeTrue();
             handlerWasCalled.Should().BeFalse();
         }
+
+        [Fact]
+        public async Task Invocation_can_be_cancelled_by_middleware_before_cancellation_enabled()
+        {
+            var command = new Command("the-command");
+            command.Handler = CommandHandler.Create((InvocationContext context) =>
+            {
+                CancellationToken ct = context.EnableCancellation();
+                // The middleware canceled the request
+                Assert.True(ct.IsCancellationRequested);
+            });
+
+            var parser = new CommandLineBuilder()
+                         .UseMiddleware(async (context, next) =>
+                         {
+                             bool cancellationEnabled = context.Cancel();
+                             // Cancellation is not yet enabled
+                             Assert.False(cancellationEnabled);
+                             await next(context);
+                         })
+                         .AddCommand(command)
+                         .Build();
+
+            await parser.InvokeAsync("the-command", new TestConsole());
+        }
+
+        [Fact]
+        public async Task Invocation_can_be_cancelled_by_middleware_after_cancellation_enabled()
+        {
+            const int timeout = 5000;
+
+            var command = new Command("the-command");
+            command.Handler = CommandHandler.Create(async (InvocationContext context) =>
+            {
+                CancellationToken ct = context.EnableCancellation();
+                // The middleware hasn't cancelled our request yet.
+                Assert.False(ct.IsCancellationRequested);
+
+                await Task.Delay(timeout, ct);
+            });
+
+            var parser = new CommandLineBuilder()
+                         .UseMiddleware(async (context, next) =>
+                         {
+                             Task task = next(context);
+
+                             bool cancellationEnabled = context.Cancel();
+                             // Cancellation is enabled by next
+                             Assert.True(cancellationEnabled);
+
+                             await task;
+                         })
+                         .AddCommand(command)
+                         .Build();
+
+            var invocation = parser.InvokeAsync("the-command", new TestConsole());
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(() => invocation);
+        }
     }
 }

src/System.CommandLine.Tests/UseExceptionHandlerTests.cs

@@ -3,6 +3,7 @@
 
 using System.CommandLine.Builder;
 using System.CommandLine.Invocation;
+using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
 using Xunit;
@@ -73,6 +74,21 @@ public async Task UseExceptionHandler_catches_command_handler_exceptions_and_wri
             _console.Error.ToString().Should().Contain("System.Exception: oops!");
         }
 
+        [Fact]
+        public async Task UseExceptionHandler_doesnt_write_operationcancelledexception_details_when_context_is_cancelled()
+        {
+            var parser = new CommandLineBuilder()
+                         .UseMiddleware((InvocationContext context) => context.Cancel())
+                         .AddCommand("the-command", "",
+                                     cmd => cmd.OnExecute((CancellationToken ct) => ct.ThrowIfCancellationRequested()))
+                         .UseExceptionHandler()
+                         .Build();
+
+            await parser.InvokeAsync("the-command", _console);
+
+            _console.Error.ToString().Should().Equals("");
+        }
+
         [Fact]
         public async Task Declaration_of_UseExceptionHandler_can_come_before_other_middleware()
         {

src/System.CommandLine/Builder/CommandLineBuilder.cs

@@ -55,7 +55,8 @@ internal void AddMiddleware(
 
         internal static class MiddlewareOrder
         {
-            public const int ExceptionHandler = int.MinValue;
+            public const int ProcessExit = int.MinValue;
+            public const int ExceptionHandler = ProcessExit + 100;
             public const int Configuration = ExceptionHandler + 100;
             public const int Preprocessing = Configuration + 100;
             public const int AfterPreprocessing = Preprocessing + 100;

src/System.CommandLine/Invocation/CommandHandler.cs

@@ -20,6 +20,10 @@ public static ICommandHandler Create<T>(
             Action<T> action) => 
             new MethodBindingCommandHandler(action);
 
+        public static ICommandHandler Create<T>(
+            Func<T, Task> action) =>
+            new MethodBindingCommandHandler(action);
+
         public static ICommandHandler Create<T1, T2>(
             Action<T1, T2> action) => 
             new MethodBindingCommandHandler(action);

src/System.CommandLine/Invocation/InvocationContext.cs

@@ -1,11 +1,16 @@
 // Copyright (c) .NET Foundation and contributors. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
+using System.Threading;
+
 namespace System.CommandLine.Invocation
 {
     public sealed class InvocationContext : IDisposable
     {
         private readonly IDisposable _onDispose;
+        private CancellationTokenSource _cts;
+        private bool _isCancellationEnabled;
+        private readonly object _gate = new object();
 
         public InvocationContext(
             ParseResult parseResult,
@@ -36,6 +41,42 @@ public InvocationContext(
 
         public IInvocationResult InvocationResult { get; set; }
 
+        // Indicates the invocation can be cancelled.
+        // This returns a CancellationToken that will be set when the invocation
+        // is cancelled. The method may return a CancellationToken that is already
+        // cancelled.
+        public CancellationToken EnableCancellation()
+        {
+            lock (_gate)
+            {
+                _isCancellationEnabled = true;
+                if (_cts == null)
+                {
+                    _cts = new CancellationTokenSource();
+                }
+                return _cts.Token;
+            }
+        }
+
+        // The return value indicates if the Invocation has cancellation enabled.
+        // When Cancel returns false, the Middleware may decide to forcefully
+        // end the process, for example, by calling Environment.Exit.
+        public bool Cancel()
+        {
+            lock (_gate)
+            {
+                if (_cts == null)
+                {
+                    _cts = new CancellationTokenSource();
+                }
+                _cts.Cancel();
+                return _isCancellationEnabled;
+            }
+        }
+
+        public bool IsCancelled =>
+            _cts?.Token.IsCancellationRequested == true;
+
         public void Dispose()
         {
             _onDispose?.Dispose();

src/System.CommandLine/Invocation/InvocationExtensions.cs

@@ -7,13 +7,51 @@
 using System.Linq;
 using System.Reflection;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using static System.Environment;
 
 namespace System.CommandLine.Invocation
 {
     public static class InvocationExtensions
     {
+        public static CommandLineBuilder UseConsoleLifetime(this CommandLineBuilder builder)
+        {
+            builder.AddMiddleware(async (context, next) =>
+            {
+                ConsoleCancelEventHandler consoleHandler = (_, args) =>
+                {
+                    if (context.Cancel())
+                    {
+                        args.Cancel = true;
+                    }
+                };
+                var blockProcessExit = new ManualResetEventSlim(initialState: false);
+                EventHandler processExitHandler = (_1, _2) =>
+                {
+                    if (context.Cancel())
+                    {
+                        blockProcessExit.Wait();
+                        Environment.ExitCode = context.ResultCode;
+                    }
+                };
+                try
+                {
+                    Console.CancelKeyPress += consoleHandler;
+                    AppDomain.CurrentDomain.ProcessExit += processExitHandler;
+                    await next(context);
+                }
+                finally
+                {
+                    Console.CancelKeyPress -= consoleHandler;
+                    AppDomain.CurrentDomain.ProcessExit -= processExitHandler;
+                    blockProcessExit.Set();
+                }
+            }, CommandLineBuilder.MiddlewareOrder.ProcessExit);
+
+            return builder;
+        }
+
         public static CommandLineBuilder UseMiddleware(
             this CommandLineBuilder builder,
             InvocationMiddleware middleware)
@@ -75,11 +113,15 @@ public static CommandLineBuilder UseExceptionHandler(
                 }
                 catch (Exception exception)
                 {
-                    context.Console.ResetColor();
-                    context.Console.ForegroundColor = ConsoleColor.Red;
-                    context.Console.Error.Write("Unhandled exception: ");
-                    context.Console.Error.WriteLine(exception.ToString());
-                    context.Console.ResetColor();
+                    if (!(context.IsCancelled &&
+                          exception is OperationCanceledException))
+                    {
+                        context.Console.ResetColor();
+                        context.Console.ForegroundColor = ConsoleColor.Red;
+                        context.Console.Error.Write("Unhandled exception: ");
+                        context.Console.Error.WriteLine(exception.ToString());
+                        context.Console.ResetColor();
+                    }
                     context.ResultCode = 1;
                 }
             }, order: CommandLineBuilder.MiddlewareOrder.ExceptionHandler);

src/System.CommandLine/Invocation/MethodBinderBase.cs

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.CommandLine.Invocation
@@ -76,6 +77,11 @@ public static object[] BindMethodArguments(
                 {
                     arguments.Add(context.Console);
                 }
+                else if (parameterInfo.ParameterType == typeof(CancellationToken))
+                {
+                    CancellationToken ct = context.EnableCancellation();
+                    arguments.Add(ct);
+                }
                 else
                 {
                     var argument = context.ParseResult

src/System.CommandLine/Invocation/MethodBindingCommandHandler.cs

@@ -27,7 +27,18 @@ public MethodBindingCommandHandler(MethodBinderBase methodBinder)
 
         public Task<int> InvokeAsync(InvocationContext context)
         {
-            return  _methodBinder.InvokeAsync(context);
+            try
+            {
+                return  _methodBinder.InvokeAsync(context);
+            }
+            catch (TargetInvocationException te)
+            {
+                if (te.InnerException != null)
+                {
+                    throw te.InnerException;
+                }
+                throw;
+            }
         }
     }
 }