Support kernel interrupt for %azure.execute

src/AzureClient/AzureClient.cs

@@ -7,6 +7,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Net;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Azure.Quantum;
 using Microsoft.Azure.Quantum.Client.Models;
@@ -110,7 +111,11 @@ public async Task<ExecutionResult> GetConnectionStatusAsync(IChannel channel)
             return ValidExecutionTargets.ToExecutionResult();
         }
 
-        private async Task<ExecutionResult> SubmitOrExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext, bool execute)
+        private async Task<ExecutionResult> SubmitOrExecuteJobAsync(
+            IChannel channel,
+            AzureSubmissionContext submissionContext,
+            bool execute,
+            CancellationToken cancellationToken)
         {
             if (ActiveWorkspace == null)
             {
@@ -178,38 +183,45 @@ private async Task<ExecutionResult> SubmitOrExecuteJobAsync(IChannel channel, Az
                 return AzureClientError.JobSubmissionFailed.ToExecutionResult();
             }
 
+            // If the command was not %azure.execute, simply return the job status.
             if (!execute)
             {
                 return await GetJobStatusAsync(channel, MostRecentJobId);
             }
 
+            // If the command was %azure.execute, wait for the job to complete and return the job output.
             channel.Stdout($"Waiting up to {submissionContext.ExecutionTimeout} seconds for Azure Quantum job to complete...");
 
-            using (var cts = new System.Threading.CancellationTokenSource(TimeSpan.FromSeconds(submissionContext.ExecutionTimeout)))
+            using var executionTimeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(submissionContext.ExecutionTimeout));
+            using var executionCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(executionTimeoutTokenSource.Token, cancellationToken);
             {
-                CloudJob? cloudJob = null;
-                do
+                try
                 {
-                    // TODO: Once jupyter-core supports interrupt requests (https://github.com/microsoft/jupyter-core/issues/55),
-                    //       handle Jupyter kernel interrupt here and break out of this loop 
-                    await Task.Delay(TimeSpan.FromSeconds(submissionContext.ExecutionPollingInterval));
-                    if (cts.IsCancellationRequested) break;
-                    cloudJob = await ActiveWorkspace.GetJobAsync(MostRecentJobId);
-                    channel.Stdout($"[{DateTime.Now.ToLongTimeString()}] Current job status: {cloudJob?.Status ?? "Unknown"}");
+                    CloudJob? cloudJob = null;
+                    while (cloudJob == null || cloudJob.InProgress)
+                    {
+                        executionCancellationTokenSource.Token.ThrowIfCancellationRequested();
+                        await Task.Delay(TimeSpan.FromSeconds(submissionContext.ExecutionPollingInterval), executionCancellationTokenSource.Token);
+                        cloudJob = await ActiveWorkspace.GetJobAsync(MostRecentJobId);
+                        channel.Stdout($"[{DateTime.Now.ToLongTimeString()}] Current job status: {cloudJob?.Status ?? "Unknown"}");
+                    }
+                }
+                catch (Exception e) when (e is TaskCanceledException || e is OperationCanceledException)
+                {
+                    Logger?.LogInformation($"Operation canceled while waiting for job execution to complete: {e.Message}");
                 }
-                while (cloudJob == null || cloudJob.InProgress);
             }
 
             return await GetJobResultAsync(channel, MostRecentJobId);
         }
 
         /// <inheritdoc/>
-        public async Task<ExecutionResult> SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext) =>
-            await SubmitOrExecuteJobAsync(channel, submissionContext, execute: false);
+        public async Task<ExecutionResult> SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? cancellationToken = null) =>
+            await SubmitOrExecuteJobAsync(channel, submissionContext, execute: false, cancellationToken ?? CancellationToken.None);
 
         /// <inheritdoc/>
-        public async Task<ExecutionResult> ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext) =>
-            await SubmitOrExecuteJobAsync(channel, submissionContext, execute: true);
+        public async Task<ExecutionResult> ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? cancellationToken = null) =>
+            await SubmitOrExecuteJobAsync(channel, submissionContext, execute: true, cancellationToken ?? CancellationToken.None);
 
         /// <inheritdoc/>
         public async Task<ExecutionResult> GetActiveTargetAsync(IChannel channel)

src/AzureClient/IAzureClient.cs

@@ -5,6 +5,7 @@
 
 using System.Collections.Generic;
 using System.ComponentModel;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 
@@ -44,7 +45,7 @@ public Task<ExecutionResult> ConnectAsync(IChannel channel,
         /// <returns>
         /// Details of the submitted job, or an error if submission failed.
         /// </returns>
-        public Task<ExecutionResult> SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext);
+        public Task<ExecutionResult> SubmitJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? token);
 
         /// <summary>
         /// Executes the specified Q# operation as a job to the currently active target
@@ -53,7 +54,7 @@ public Task<ExecutionResult> ConnectAsync(IChannel channel,
         /// <returns>
         /// The result of the executed job, or an error if execution failed.
         /// </returns>
-        public Task<ExecutionResult> ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext);
+        public Task<ExecutionResult> ExecuteJobAsync(IChannel channel, AzureSubmissionContext submissionContext, CancellationToken? token);
 
         /// <summary>
         /// Sets the specified target for job submission.

src/AzureClient/Magic/AzureClientMagicBase.cs

@@ -4,6 +4,7 @@
 #nullable enable
 
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -35,11 +36,15 @@ public AzureClientMagicBase(IAzureClient azureClient, string keyword, Documentat
 
         /// <inheritdoc/>
         public override ExecutionResult Run(string input, IChannel channel) =>
-            RunAsync(input, channel).GetAwaiter().GetResult();
+            RunCancellable(input, channel, CancellationToken.None);
+
+        /// <inheritdoc/>
+        public override ExecutionResult RunCancellable(string input, IChannel channel, CancellationToken cancellationToken) =>
+            RunAsync(input, channel, cancellationToken).GetAwaiter().GetResult();
 
         /// <summary>
         ///     Executes the magic command functionality for the given input.
         /// </summary>
-        public abstract Task<ExecutionResult> RunAsync(string input, IChannel channel);
+        public abstract Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken);
     }
 }

src/AzureClient/Magic/ConnectMagic.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -90,7 +91,7 @@ credentials when connecting to Azure.
         ///     Connects to an Azure workspace given a subscription ID, resource group name,
         ///     workspace name, and connection string as a JSON-encoded object.
         /// </summary>
-        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel)
+        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken)
         {
             var inputParameters = ParseInputParameters(input);
 

src/AzureClient/Magic/ExecuteMagic.cs

@@ -5,6 +5,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -56,9 +57,9 @@ The Azure Quantum workspace must previously have been initialized
         ///     name that is present in the current Q# Jupyter workspace, and
         ///     waits for the job to complete before returning.
         /// </summary>
-        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel)
+        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken)
         {
-            return await AzureClient.ExecuteJobAsync(channel, AzureSubmissionContext.Parse(input));
+            return await AzureClient.ExecuteJobAsync(channel, AzureSubmissionContext.Parse(input), cancellationToken);
         }
     }
 }

src/AzureClient/Magic/JobsMagic.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -52,7 +53,7 @@ The Azure Quantum workspace must previously have been initialized
         /// <summary>
         ///     Lists all jobs in the active workspace.
         /// </summary>
-        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel) =>
+        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken) =>
             await AzureClient.GetJobListAsync(channel);
     }
 }

src/AzureClient/Magic/OutputMagic.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -68,7 +69,7 @@ The Azure Quantum workspace must previously have been initialized
         ///     Displays the output of a given completed job ID, if provided,
         ///     or all jobs submitted in the current session.
         /// </summary>
-        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel)
+        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken)
         {
             var inputParameters = ParseInputParameters(input, firstParameterInferredName: ParameterNameJobId);
             string jobId = inputParameters.DecodeParameter<string>(ParameterNameJobId);

src/AzureClient/Magic/StatusMagic.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -66,7 +67,7 @@ The Azure Quantum workspace must previously have been initialized
         ///     Displays the status corresponding to a given job ID, if provided,
         ///     or the most recently-submitted job in the current session.
         /// </summary>
-        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel)
+        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken)
         {
             var inputParameters = ParseInputParameters(input, firstParameterInferredName: ParameterNameJobId);
             string jobId = inputParameters.DecodeParameter<string>(ParameterNameJobId);

src/AzureClient/Magic/SubmitMagic.cs

@@ -4,6 +4,7 @@
 #nullable enable
 
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -52,9 +53,9 @@ The Azure Quantum workspace must previously have been initialized
         ///     Submits a new job to an Azure Quantum workspace given a Q# operation
         ///     name that is present in the current Q# Jupyter workspace.
         /// </summary>
-        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel)
+        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken)
         {
-            return await AzureClient.SubmitJobAsync(channel, AzureSubmissionContext.Parse(input));
+            return await AzureClient.SubmitJobAsync(channel, AzureSubmissionContext.Parse(input), cancellationToken);
         }
     }
 }

src/AzureClient/Magic/TargetMagic.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Jupyter;
@@ -63,7 +64,7 @@ available in the workspace.
         /// <summary>
         ///     Sets or views the target for job submission to the current Azure Quantum workspace.
         /// </summary>
-        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel)
+        public override async Task<ExecutionResult> RunAsync(string input, IChannel channel, CancellationToken cancellationToken)
         {
             var inputParameters = ParseInputParameters(input, firstParameterInferredName: ParameterNameTargetId);
             if (inputParameters.ContainsKey(ParameterNameTargetId))

src/Jupyter/Jupyter.csproj

@@ -35,7 +35,7 @@
   </ItemGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Jupyter.Core" Version="1.3.60623" />
+    <PackageReference Include="Microsoft.Jupyter.Core" Version="1.4.76835" />
     <PackageReference Include="System.Reactive" Version="4.3.2" />
   </ItemGroup>
 

src/Jupyter/Magic/AbstractMagic.cs

@@ -6,18 +6,18 @@
 using System.IO;
 using System.Linq;
 using System.Text.RegularExpressions;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Jupyter.Core;
 using Microsoft.Quantum.IQSharp.Common;
 using Microsoft.Quantum.QsCompiler.Serialization;
-using Newtonsoft.Json.Linq;
 
 namespace Microsoft.Quantum.IQSharp.Jupyter
 {
     /// <summary>
     ///     Abstract base class for IQ# magic symbols.
     /// </summary>
-    public abstract class AbstractMagic : MagicSymbol
+    public abstract class AbstractMagic : CancellableMagicSymbol
     {
         /// <summary>
         ///     Constructs a new magic symbol given its name and documentation.
@@ -28,7 +28,7 @@ public AbstractMagic(string keyword, Documentation docs)
             this.Documentation = docs;
 
             this.Kind = SymbolKind.Magic;
-            this.Execute = SafeExecute(this.Run);
+            this.ExecuteCancellable = this.SafeExecute(this.RunCancellable);
         }
 
         /// <summary>
@@ -38,31 +38,32 @@ public AbstractMagic(string keyword, Documentation docs)
         ///     returned execution function displays the given exceptions to its
         ///     display channel.
         /// </summary>
-        public Func<string, IChannel, Task<ExecutionResult>> SafeExecute(Func<string, IChannel, ExecutionResult> magic) =>
-            async (input, channel) =>
-            {
-                channel = channel.WithNewLines();
-
-                try
-                {
-                    return magic(input, channel);
-                }
-                catch (InvalidWorkspaceException ws)
+        public Func<string, IChannel, CancellationToken, Task<ExecutionResult>> SafeExecute(
+            Func<string, IChannel, CancellationToken, ExecutionResult> magic) =>
+                async (input, channel, cancellationToken) =>
                 {
-                    foreach (var m in ws.Errors) channel.Stderr(m);
-                    return ExecuteStatus.Error.ToExecutionResult();
-                }
-                catch (AggregateException agg)
-                {
-                    foreach (var e in agg.InnerExceptions) channel.Stderr(e?.Message);
-                    return ExecuteStatus.Error.ToExecutionResult();
-                }
-                catch (Exception e)
-                {
-                    channel.Stderr(e.Message);
-                    return ExecuteStatus.Error.ToExecutionResult();
-                }
-            };
+                    channel = channel.WithNewLines();
+
+                    try
+                    {
+                        return magic(input, channel, cancellationToken);
+                    }
+                    catch (InvalidWorkspaceException ws)
+                    {
+                        foreach (var m in ws.Errors) channel.Stderr(m);
+                        return ExecuteStatus.Error.ToExecutionResult();
+                    }
+                    catch (AggregateException agg)
+                    {
+                        foreach (var e in agg.InnerExceptions) channel.Stderr(e?.Message);
+                        return ExecuteStatus.Error.ToExecutionResult();
+                    }
+                    catch (Exception e)
+                    {
+                        channel.Stderr(e.Message);
+                        return ExecuteStatus.Error.ToExecutionResult();
+                    }
+                };
 
         /// <summary>
         ///     Parses the input to a magic command, interpreting the input as
@@ -146,5 +147,17 @@ public static Dictionary<string, string> ParseInputParameters(string input, stri
         ///     A method to be run when the magic command is executed.
         /// </summary>
         public abstract ExecutionResult Run(string input, IChannel channel);
+
+        /// <summary>
+        ///     A method to be run when the magic command is executed, including a cancellation
+        ///     token to use for requesting cancellation.
+        /// </summary>
+        /// <remarks>
+        ///     The default implementation in <see cref="AbstractMagic"/> ignores the cancellation token.
+        ///     Derived classes should override this method and monitor the cancellation token if they
+        ///     wish to support cancellation.
+        /// </remarks>
+        public virtual ExecutionResult RunCancellable(string input, IChannel channel, CancellationToken cancellationToken) =>
+            Run(input, channel);
     }
 }

src/Kernel/Kernel.csproj

@@ -22,7 +22,6 @@
   </ItemGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Jupyter.Core" Version="1.3.60623" />
     <PackageReference Include="System.Reactive" Version="4.3.2" />
   </ItemGroup>