Cancel completed host tool invocations across SDKs

CHANGELOG.md

@@ -7,6 +7,10 @@ See [GitHub Releases](https://github.com/github/copilot-sdk/releases) for the fu
 
 ## [Unreleased]
 
+### Feature: cancellation for host-owned external tools
+
+Host-owned external tool callbacks are now cancelled when their runtime request completes or their SDK session terminates. The cancellation primitive is idiomatic per SDK: .NET passes a request token to `AIFunction`, Node.js exposes `ToolInvocation.signal`, Go cancels `ToolInvocation.TraceContext`, Java cancels the returned `CompletableFuture`, Python cancels the handler task, and Rust drops the handler future. Go handlers that retain `TraceContext` for background work must derive a separate lifetime because the invocation context is cancelled when the request ends.
+
 ### Feature: declare application identity with client info
 
 Client options now accept optional client info (application name and version, integration name and version) across all six SDKs, exposed idiomatically per language (`clientInfo` in Node.js, `client_info` in Python and Rust, `ClientInfo` in Go and .NET, `setClientInfo` in Java). When set, the SDK forwards it on the `server.connect` handshake so the telemetry the runtime emits on the connection is attributed to the application and its Copilot integration instead of the runtime's own build. All fields are optional, and leaving client info unset keeps the runtime's default attribution. See [Client info](./docs/features/client-info.md).

dotnet/src/Client.cs

@@ -597,6 +597,10 @@ public async Task StopAsync()
     /// </example>
     public async Task ForceStopAsync()
     {
+        foreach (var session in _sessions.Values)
+        {
+            session.CancelPendingExternalTools();
+        }
         _sessions.Clear();
         ClearGitHubTokenProviders();
 
@@ -2693,6 +2697,7 @@ private async Task<Connection> ConnectToServerAsync(Process? cliProcess, string?
                 ClientGlobalApiRegistration.RegisterClientGlobalApiHandlers(rpc, _clientGlobalApis);
             }
             rpc.StartListening();
+            _ = CancelExternalToolsWhenConnectionClosesAsync(rpc);
             LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
                 "CopilotClient.ConnectToServerAsync transport setup complete. Elapsed={Elapsed}",
                 setupTimestamp);
@@ -2705,17 +2710,50 @@ private async Task<Connection> ConnectToServerAsync(Process? cliProcess, string?
         catch
         {
             try { rpc?.Dispose(); }
-            catch (Exception ex) { _logger.LogDebug(ex, "Failed to dispose JSON-RPC connection after startup failure"); }
+            catch (Exception ex) when (IsRecoverableConnectionCleanupFailure(ex))
+            {
+                _logger.LogDebug(ex, "Failed to dispose JSON-RPC connection after startup failure");
+            }
 
             if (networkStream is not null)
             {
                 try { await networkStream.DisposeAsync(); }
-                catch (Exception ex) { _logger.LogDebug(ex, "Failed to dispose TCP stream after startup failure"); }
+                catch (Exception ex) when (IsRecoverableConnectionCleanupFailure(ex))
+                {
+                    _logger.LogDebug(ex, "Failed to dispose TCP stream after startup failure");
+                }
             }
             throw;
         }
     }
 
+    private static bool IsRecoverableConnectionCleanupFailure(Exception exception)
+        => exception is not OutOfMemoryException
+            and not StackOverflowException
+            and not AccessViolationException
+            and not AppDomainUnloadedException;
+
+    private async Task CancelExternalToolsWhenConnectionClosesAsync(JsonRpc rpc)
+    {
+        await Task.WhenAny(rpc.Completion).ConfigureAwait(false);
+        if (rpc.Completion.Exception is { } exception)
+        {
+            _logger.LogDebug(exception, "JSON-RPC connection completed with an error");
+        }
+
+        var connectionTask = _connectionTask;
+        if (connectionTask is null
+            || connectionTask.Status != System.Threading.Tasks.TaskStatus.RanToCompletion
+            || !ReferenceEquals(connectionTask.Result.Rpc, rpc))
+        {
+            return;
+        }
+        foreach (var session in _sessions.Values)
+        {
+            session.CancelPendingExternalTools();
+        }
+    }
+
     private static JsonSerializerOptions SerializerOptionsForMessageFormatter { get; } = CreateSerializerOptions();
 
     /// <summary>

dotnet/src/Session.cs

@@ -5,6 +5,7 @@
 using GitHub.Copilot.Rpc;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.Logging;
+using System.Collections.Concurrent;
 using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
@@ -60,6 +61,8 @@ public sealed partial class CopilotSession : IAsyncDisposable
     private readonly Dictionary<string, AIFunction> _toolHandlers = [];
     private readonly Dictionary<string, Func<CommandContext, Task>> _commandHandlers = [];
     private readonly Dictionary<string, Func<ProviderTokenArgs, Task<string>>> _bearerTokenProviders = new(StringComparer.Ordinal);
+    private readonly ConcurrentDictionary<string, CancellationTokenSource> _pendingExternalTools = new(StringComparer.Ordinal);
+    private readonly CancellationTokenSource _externalToolLifetime = new();
     private readonly ILogger _logger;
     private readonly CopilotClient _parentClient;
 
@@ -226,6 +229,7 @@ internal void CloseEventChannel()
     /// </summary>
     internal void Unregister()
     {
+        CancelPendingExternalTools();
         CloseEventChannel();
         RemoveFromClient();
     }
@@ -675,6 +679,10 @@ private async Task HandleBroadcastEventAsync(SessionEvent sessionEvent)
                         break;
                     }
 
+                case ExternalToolCompletedEvent completedEvent:
+                    CancelExternalTool(completedEvent.Data.RequestId);
+                    break;
+
                 case PermissionRequestedEvent permEvent:
                     {
                         var data = permEvent.Data;
@@ -884,8 +892,24 @@ private async Task TryCancelMcpAuthRequestAsync(string requestId)
     /// </summary>
     private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, string toolCallId, JsonElement? arguments, AIFunction tool)
     {
+        if (_externalToolLifetime.IsCancellationRequested)
+        {
+            return;
+        }
+
+        using var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(_externalToolLifetime.Token);
+        if (!_pendingExternalTools.TryAdd(requestId, cancellationSource))
+        {
+            return;
+        }
+
         try
         {
+            if (cancellationSource.IsCancellationRequested)
+            {
+                return;
+            }
+
             var invocation = new ToolInvocation
             {
                 SessionId = SessionId,
@@ -903,7 +927,7 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
             {
                 try
                 {
-                    var metadata = await Rpc.Tools.GetCurrentMetadataAsync();
+                    var metadata = await Rpc.Tools.GetCurrentMetadataAsync(cancellationSource.Token);
                     invocation.AvailableTools = metadata.Tools;
                 }
                 catch (Exception ex) when (ex is RemoteRpcException or IOException or ObjectDisposedException or JsonException)
@@ -938,7 +962,7 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
             }
 
             var toolTimestamp = Stopwatch.GetTimestamp();
-            var result = await tool.InvokeAsync(aiFunctionArgs);
+            var result = await tool.InvokeAsync(aiFunctionArgs, cancellationSource.Token);
             LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
                 "CopilotSession.ExecuteToolAndRespondAsync tool dispatch. Elapsed={Elapsed}, SessionId={SessionId}, RequestId={RequestId}, ToolCallId={ToolCallId}, Tool={ToolName}",
                 toolTimestamp,
@@ -948,9 +972,14 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
                 toolName);
 
             var toolResultObject = ToolResultObject.ConvertFromInvocationResult(result, tool.JsonSerializerOptions);
+            if (!TryClaimExternalTool(requestId, cancellationSource))
+            {
+                return;
+            }
 
             var responseRpcTimestamp = Stopwatch.GetTimestamp();
-            await Rpc.Tools.HandlePendingToolCallAsync(requestId, toolResultObject, error: null);
+            await Rpc.Tools.HandlePendingToolCallAsync(
+                requestId, toolResultObject, error: null, cancellationSource.Token);
             LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
                 "CopilotSession.ExecuteToolAndRespondAsync response sent successfully. Elapsed={Elapsed}, SessionId={SessionId}, RequestId={RequestId}, ToolCallId={ToolCallId}, Tool={ToolName}",
                 responseRpcTimestamp,
@@ -959,11 +988,33 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
                 toolCallId,
                 toolName);
         }
-        catch (Exception ex)
+        catch (OperationCanceledException) when (cancellationSource.IsCancellationRequested)
+        {
+            // The runtime has already completed the request or the session is shutting down.
+        }
+        catch (RemoteRpcException) when (cancellationSource.IsCancellationRequested)
+        {
+            // Another client answered after this invocation completed locally.
+        }
+        catch (Exception) when (cancellationSource.IsCancellationRequested)
+        {
+            // Cancellation won the request; no response or error should escape.
+        }
+        catch (Exception ex) when (!cancellationSource.IsCancellationRequested)
         {
+            if (!TryClaimExternalTool(requestId, cancellationSource))
+            {
+                return;
+            }
+
             try
             {
-                await Rpc.Tools.HandlePendingToolCallAsync(requestId, result: null, error: ex.Message);
+                await Rpc.Tools.HandlePendingToolCallAsync(
+                    requestId, result: null, error: ex.Message, cancellationSource.Token);
+            }
+            catch (OperationCanceledException)
+            {
+                // Teardown canceled the in-flight error response.
             }
             catch (IOException)
             {
@@ -973,6 +1024,15 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
             {
                 // Connection already disposed — nothing we can do
             }
+            catch (RemoteRpcException)
+            {
+                // Another client may have answered the broadcast request first.
+            }
+        }
+        finally
+        {
+            ((ICollection<KeyValuePair<string, CancellationTokenSource>>)_pendingExternalTools)
+                .Remove(new(requestId, cancellationSource));
         }
 
         static string GetSingleParameterName(AIFunction tool)
@@ -1025,6 +1085,49 @@ static string GetSingleParameterName(AIFunction tool)
         }
     }
 
+    private bool TryClaimExternalTool(string requestId, CancellationTokenSource cancellationSource)
+        => ((ICollection<KeyValuePair<string, CancellationTokenSource>>)_pendingExternalTools)
+            .Remove(new(requestId, cancellationSource));
+
+    private void CancelExternalTool(string requestId)
+    {
+        if (!string.IsNullOrEmpty(requestId) && _pendingExternalTools.TryRemove(requestId, out var cancellationSource))
+        {
+            _ = Task.Run(() =>
+            {
+                try
+                {
+                    cancellationSource.Cancel();
+                }
+                catch (AggregateException)
+                {
+                    // Cancellation callbacks are consumer code and must not disrupt event dispatch.
+                }
+                catch (ObjectDisposedException)
+                {
+                    // The invocation completed while cancellation was being delivered.
+                }
+            });
+        }
+    }
+
+    internal void CancelPendingExternalTools()
+    {
+        try
+        {
+            _externalToolLifetime.Cancel();
+        }
+        catch (AggregateException)
+        {
+            // User cancellation callbacks must not prevent session teardown.
+        }
+        catch (ObjectDisposedException)
+        {
+            // Session teardown already completed.
+        }
+        _pendingExternalTools.Clear();
+    }
+
     /// <summary>
     /// Executes a permission handler and sends the result back via the HandlePendingPermissionRequest RPC.
     /// </summary>
@@ -2105,6 +2208,7 @@ public async ValueTask DisposeAsync()
             return;
         }
 
+        CancelPendingExternalTools();
         CloseEventChannel();
 
         try

dotnet/test/E2E/ExternalToolCancellationE2ETests.cs

@@ -0,0 +1,62 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+using Microsoft.Extensions.AI;
+using System.ComponentModel;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace GitHub.Copilot.Test.E2E;
+
+public class ExternalToolCancellationE2ETests(E2ETestFixture fixture, ITestOutputHelper output)
+    : E2ETestBase(fixture, "external_tool_cancellation", output)
+{
+    [Fact]
+    public async Task Should_Cancel_Tool_Handler_When_Session_Disposes()
+    {
+        var toolStarted = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var toolCancelled = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        var releaseTool = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var session = await CreateSessionAsync(new SessionConfig
+        {
+            Tools = [AIFunctionFactory.Create(SlowTool, "slow_analysis")],
+            OnPermissionRequest = PermissionHandler.ApproveAll,
+        });
+
+        _ = session.SendAsync(new MessageOptions
+        {
+            Prompt = "Use slow_analysis with value 'test_abort'. Wait for the result.",
+        });
+
+        var startedValue = await toolStarted.Task.WaitAsync(TimeSpan.FromSeconds(60));
+        Assert.Equal("test_abort", startedValue);
+
+        await session.DisposeAsync();
+        await toolCancelled.Task.WaitAsync(TimeSpan.FromSeconds(60));
+
+        releaseTool.TrySetResult("RELEASED");
+
+        [Description("A slow analysis tool that blocks until released")]
+        async Task<string> SlowTool([Description("Value to analyze")] string value, CancellationToken cancellationToken)
+        {
+            toolStarted.TrySetResult(value);
+            try
+            {
+                var completed = await Task.WhenAny(releaseTool.Task, Task.Delay(Timeout.Infinite, cancellationToken));
+                if (completed == releaseTool.Task)
+                {
+                    return await releaseTool.Task;
+                }
+
+                throw new OperationCanceledException(cancellationToken);
+            }
+            catch (OperationCanceledException)
+            {
+                toolCancelled.TrySetResult(true);
+                throw;
+            }
+        }
+    }
+}