From 3aa682082ad58006132da690e1c292508b486502 Mon Sep 17 00:00:00 2001
From: Jacob Alber <jaalber@microsoft.com>
Date: Thu, 23 Oct 2025 16:01:56 -0400
Subject: [PATCH] feat: Add Cancellation API on StreamingRun (#1675)

---
 .../Execution/AsyncRunHandle.cs                     | 13 ++++++++++---
 .../Microsoft.Agents.AI.Workflows/StreamingRun.cs   | 11 +++++++----
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs
index de1a1dfc7f..fb3d391251 100644
--- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/AsyncRunHandle.cs
@@ -156,15 +156,22 @@ internal sealed class AsyncRunHandle : ICheckpointingHandle, IAsyncDisposable
         this._eventStream.SignalInput();
     }
 
+    public async ValueTask CancelRunAsync()
+    {
+        this._endRunSource.Cancel();
+
+        await this._eventStream.StopAsync().ConfigureAwait(false);
+    }
+
     public async ValueTask DisposeAsync()
     {
         if (Interlocked.Exchange(ref this._isDisposed, 1) == 0)
         {
-            this._endRunSource.Cancel();
+            // Cancel the run if it is still running
+            await this.CancelRunAsync().ConfigureAwait(false);
 
-            await this._eventStream.StopAsync().ConfigureAwait(false);
+            // These actually release and clean up resources
             await this._stepRunner.RequestEndRunAsync().ConfigureAwait(false);
-
             this._endRunSource.Dispose();
 
             await this._eventStream.DisposeAsync().ConfigureAwait(false);
diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs
index ad6727fc54..d84ee8bf85 100644
--- a/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Workflows/StreamingRun.cs
@@ -79,11 +79,14 @@ public sealed class StreamingRun : IAsyncDisposable
         CancellationToken cancellationToken = default)
         => this._runHandle.TakeEventStreamAsync(blockOnPendingRequest, cancellationToken);
 
+    /// <summary>
+    /// Attempt to cancel the streaming run.
+    /// </summary>
+    /// <returns>A <see cref="ValueTask"/> that represents the asynchronous send operation.</returns>
+    public ValueTask CancelRunAsync() => this._runHandle.CancelRunAsync();
+
     /// <inheritdoc/>
-    public ValueTask DisposeAsync()
-    {
-        return this._runHandle.DisposeAsync();
-    }
+    public ValueTask DisposeAsync() => this._runHandle.DisposeAsync();
 }
 
 /// <summary>