Skip to content

Commit 53e74d8

Browse files
stephentoubCopilotSteveSandersonMS
authored
Cancel completed host tool invocations across SDKs (#2509)
* Cancel host tool invocations on completion Track request-scoped cancellation across each SDK and terminate pending host callbacks when the runtime completes a request or the session transport shuts down. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Fix cancellation CI validation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Address cancellation cleanup alerts Use lexical ownership for request cancellation sources and observe JSON-RPC completion faults without a generic catch. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * test: add cross-sdk external tool cancellation e2e Add new standalone E2E coverage for long-running external tools being cancelled when the session is terminated/disposed, without modifying existing E2E tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix external tool cancellation E2E validation Apply language formatters and remove an unused Vitest import from the cross-SDK E2E coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * test: fix java e2e formatting for spotless Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Narrow connection cleanup exception handling Avoid swallowing fatal runtime exceptions while retaining best-effort cleanup after connection startup failures. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Fix Java close cancellation test after detach migration Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Co-authored-by: Steve Sanderson <SteveSandersonMS@users.noreply.github.com>
1 parent 604c716 commit 53e74d8

36 files changed

Lines changed: 2507 additions & 68 deletions

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ See [GitHub Releases](https://github.com/github/copilot-sdk/releases) for the fu
77

88
## [Unreleased]
99

10+
### Feature: cancellation for host-owned external tools
11+
12+
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.
13+
1014
### Feature: declare application identity with client info
1115

1216
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‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -597,6 +597,10 @@ public async Task StopAsync()
597597
/// </example>
598598
public async Task ForceStopAsync()
599599
{
600+
foreach (var session in _sessions.Values)
601+
{
602+
session.CancelPendingExternalTools();
603+
}
600604
_sessions.Clear();
601605
ClearGitHubTokenProviders();
602606

@@ -2693,6 +2697,7 @@ private async Task<Connection> ConnectToServerAsync(Process? cliProcess, string?
26932697
ClientGlobalApiRegistration.RegisterClientGlobalApiHandlers(rpc, _clientGlobalApis);
26942698
}
26952699
rpc.StartListening();
2700+
_ = CancelExternalToolsWhenConnectionClosesAsync(rpc);
26962701
LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
26972702
"CopilotClient.ConnectToServerAsync transport setup complete. Elapsed={Elapsed}",
26982703
setupTimestamp);
@@ -2705,17 +2710,50 @@ private async Task<Connection> ConnectToServerAsync(Process? cliProcess, string?
27052710
catch
27062711
{
27072712
try { rpc?.Dispose(); }
2708-
catch (Exception ex) { _logger.LogDebug(ex, "Failed to dispose JSON-RPC connection after startup failure"); }
2713+
catch (Exception ex) when (IsRecoverableConnectionCleanupFailure(ex))
2714+
{
2715+
_logger.LogDebug(ex, "Failed to dispose JSON-RPC connection after startup failure");
2716+
}
27092717

27102718
if (networkStream is not null)
27112719
{
27122720
try { await networkStream.DisposeAsync(); }
2713-
catch (Exception ex) { _logger.LogDebug(ex, "Failed to dispose TCP stream after startup failure"); }
2721+
catch (Exception ex) when (IsRecoverableConnectionCleanupFailure(ex))
2722+
{
2723+
_logger.LogDebug(ex, "Failed to dispose TCP stream after startup failure");
2724+
}
27142725
}
27152726
throw;
27162727
}
27172728
}
27182729

2730+
private static bool IsRecoverableConnectionCleanupFailure(Exception exception)
2731+
=> exception is not OutOfMemoryException
2732+
and not StackOverflowException
2733+
and not AccessViolationException
2734+
and not AppDomainUnloadedException;
2735+
2736+
private async Task CancelExternalToolsWhenConnectionClosesAsync(JsonRpc rpc)
2737+
{
2738+
await Task.WhenAny(rpc.Completion).ConfigureAwait(false);
2739+
if (rpc.Completion.Exception is { } exception)
2740+
{
2741+
_logger.LogDebug(exception, "JSON-RPC connection completed with an error");
2742+
}
2743+
2744+
var connectionTask = _connectionTask;
2745+
if (connectionTask is null
2746+
|| connectionTask.Status != System.Threading.Tasks.TaskStatus.RanToCompletion
2747+
|| !ReferenceEquals(connectionTask.Result.Rpc, rpc))
2748+
{
2749+
return;
2750+
}
2751+
foreach (var session in _sessions.Values)
2752+
{
2753+
session.CancelPendingExternalTools();
2754+
}
2755+
}
2756+
27192757
private static JsonSerializerOptions SerializerOptionsForMessageFormatter { get; } = CreateSerializerOptions();
27202758

27212759
/// <summary>

‎dotnet/src/Session.cs‎

Lines changed: 109 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using GitHub.Copilot.Rpc;
66
using Microsoft.Extensions.AI;
77
using Microsoft.Extensions.Logging;
8+
using System.Collections.Concurrent;
89
using System.Collections.Immutable;
910
using System.Diagnostics;
1011
using System.Diagnostics.CodeAnalysis;
@@ -60,6 +61,8 @@ public sealed partial class CopilotSession : IAsyncDisposable
6061
private readonly Dictionary<string, AIFunction> _toolHandlers = [];
6162
private readonly Dictionary<string, Func<CommandContext, Task>> _commandHandlers = [];
6263
private readonly Dictionary<string, Func<ProviderTokenArgs, Task<string>>> _bearerTokenProviders = new(StringComparer.Ordinal);
64+
private readonly ConcurrentDictionary<string, CancellationTokenSource> _pendingExternalTools = new(StringComparer.Ordinal);
65+
private readonly CancellationTokenSource _externalToolLifetime = new();
6366
private readonly ILogger _logger;
6467
private readonly CopilotClient _parentClient;
6568

@@ -226,6 +229,7 @@ internal void CloseEventChannel()
226229
/// </summary>
227230
internal void Unregister()
228231
{
232+
CancelPendingExternalTools();
229233
CloseEventChannel();
230234
RemoveFromClient();
231235
}
@@ -675,6 +679,10 @@ private async Task HandleBroadcastEventAsync(SessionEvent sessionEvent)
675679
break;
676680
}
677681

682+
case ExternalToolCompletedEvent completedEvent:
683+
CancelExternalTool(completedEvent.Data.RequestId);
684+
break;
685+
678686
case PermissionRequestedEvent permEvent:
679687
{
680688
var data = permEvent.Data;
@@ -884,8 +892,24 @@ private async Task TryCancelMcpAuthRequestAsync(string requestId)
884892
/// </summary>
885893
private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, string toolCallId, JsonElement? arguments, AIFunction tool)
886894
{
895+
if (_externalToolLifetime.IsCancellationRequested)
896+
{
897+
return;
898+
}
899+
900+
using var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(_externalToolLifetime.Token);
901+
if (!_pendingExternalTools.TryAdd(requestId, cancellationSource))
902+
{
903+
return;
904+
}
905+
887906
try
888907
{
908+
if (cancellationSource.IsCancellationRequested)
909+
{
910+
return;
911+
}
912+
889913
var invocation = new ToolInvocation
890914
{
891915
SessionId = SessionId,
@@ -903,7 +927,7 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
903927
{
904928
try
905929
{
906-
var metadata = await Rpc.Tools.GetCurrentMetadataAsync();
930+
var metadata = await Rpc.Tools.GetCurrentMetadataAsync(cancellationSource.Token);
907931
invocation.AvailableTools = metadata.Tools;
908932
}
909933
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,
938962
}
939963

940964
var toolTimestamp = Stopwatch.GetTimestamp();
941-
var result = await tool.InvokeAsync(aiFunctionArgs);
965+
var result = await tool.InvokeAsync(aiFunctionArgs, cancellationSource.Token);
942966
LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
943967
"CopilotSession.ExecuteToolAndRespondAsync tool dispatch. Elapsed={Elapsed}, SessionId={SessionId}, RequestId={RequestId}, ToolCallId={ToolCallId}, Tool={ToolName}",
944968
toolTimestamp,
@@ -948,9 +972,14 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
948972
toolName);
949973

950974
var toolResultObject = ToolResultObject.ConvertFromInvocationResult(result, tool.JsonSerializerOptions);
975+
if (!TryClaimExternalTool(requestId, cancellationSource))
976+
{
977+
return;
978+
}
951979

952980
var responseRpcTimestamp = Stopwatch.GetTimestamp();
953-
await Rpc.Tools.HandlePendingToolCallAsync(requestId, toolResultObject, error: null);
981+
await Rpc.Tools.HandlePendingToolCallAsync(
982+
requestId, toolResultObject, error: null, cancellationSource.Token);
954983
LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
955984
"CopilotSession.ExecuteToolAndRespondAsync response sent successfully. Elapsed={Elapsed}, SessionId={SessionId}, RequestId={RequestId}, ToolCallId={ToolCallId}, Tool={ToolName}",
956985
responseRpcTimestamp,
@@ -959,11 +988,33 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
959988
toolCallId,
960989
toolName);
961990
}
962-
catch (Exception ex)
991+
catch (OperationCanceledException) when (cancellationSource.IsCancellationRequested)
992+
{
993+
// The runtime has already completed the request or the session is shutting down.
994+
}
995+
catch (RemoteRpcException) when (cancellationSource.IsCancellationRequested)
996+
{
997+
// Another client answered after this invocation completed locally.
998+
}
999+
catch (Exception) when (cancellationSource.IsCancellationRequested)
1000+
{
1001+
// Cancellation won the request; no response or error should escape.
1002+
}
1003+
catch (Exception ex) when (!cancellationSource.IsCancellationRequested)
9631004
{
1005+
if (!TryClaimExternalTool(requestId, cancellationSource))
1006+
{
1007+
return;
1008+
}
1009+
9641010
try
9651011
{
966-
await Rpc.Tools.HandlePendingToolCallAsync(requestId, result: null, error: ex.Message);
1012+
await Rpc.Tools.HandlePendingToolCallAsync(
1013+
requestId, result: null, error: ex.Message, cancellationSource.Token);
1014+
}
1015+
catch (OperationCanceledException)
1016+
{
1017+
// Teardown canceled the in-flight error response.
9671018
}
9681019
catch (IOException)
9691020
{
@@ -973,6 +1024,15 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName,
9731024
{
9741025
// Connection already disposed — nothing we can do
9751026
}
1027+
catch (RemoteRpcException)
1028+
{
1029+
// Another client may have answered the broadcast request first.
1030+
}
1031+
}
1032+
finally
1033+
{
1034+
((ICollection<KeyValuePair<string, CancellationTokenSource>>)_pendingExternalTools)
1035+
.Remove(new(requestId, cancellationSource));
9761036
}
9771037

9781038
static string GetSingleParameterName(AIFunction tool)
@@ -1025,6 +1085,49 @@ static string GetSingleParameterName(AIFunction tool)
10251085
}
10261086
}
10271087

1088+
private bool TryClaimExternalTool(string requestId, CancellationTokenSource cancellationSource)
1089+
=> ((ICollection<KeyValuePair<string, CancellationTokenSource>>)_pendingExternalTools)
1090+
.Remove(new(requestId, cancellationSource));
1091+
1092+
private void CancelExternalTool(string requestId)
1093+
{
1094+
if (!string.IsNullOrEmpty(requestId) && _pendingExternalTools.TryRemove(requestId, out var cancellationSource))
1095+
{
1096+
_ = Task.Run(() =>
1097+
{
1098+
try
1099+
{
1100+
cancellationSource.Cancel();
1101+
}
1102+
catch (AggregateException)
1103+
{
1104+
// Cancellation callbacks are consumer code and must not disrupt event dispatch.
1105+
}
1106+
catch (ObjectDisposedException)
1107+
{
1108+
// The invocation completed while cancellation was being delivered.
1109+
}
1110+
});
1111+
}
1112+
}
1113+
1114+
internal void CancelPendingExternalTools()
1115+
{
1116+
try
1117+
{
1118+
_externalToolLifetime.Cancel();
1119+
}
1120+
catch (AggregateException)
1121+
{
1122+
// User cancellation callbacks must not prevent session teardown.
1123+
}
1124+
catch (ObjectDisposedException)
1125+
{
1126+
// Session teardown already completed.
1127+
}
1128+
_pendingExternalTools.Clear();
1129+
}
1130+
10281131
/// <summary>
10291132
/// Executes a permission handler and sends the result back via the HandlePendingPermissionRequest RPC.
10301133
/// </summary>
@@ -2105,6 +2208,7 @@ public async ValueTask DisposeAsync()
21052208
return;
21062209
}
21072210

2211+
CancelPendingExternalTools();
21082212
CloseEventChannel();
21092213

21102214
try

‎dotnet/test/E2E/ExternalToolCancellationE2ETests.cs‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
*--------------------------------------------------------------------------------------------*/
4+
5+
using Microsoft.Extensions.AI;
6+
using System.ComponentModel;
7+
using Xunit;
8+
using Xunit.Abstractions;
9+
10+
namespace GitHub.Copilot.Test.E2E;
11+
12+
public class ExternalToolCancellationE2ETests(E2ETestFixture fixture, ITestOutputHelper output)
13+
: E2ETestBase(fixture, "external_tool_cancellation", output)
14+
{
15+
[Fact]
16+
public async Task Should_Cancel_Tool_Handler_When_Session_Disposes()
17+
{
18+
var toolStarted = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
19+
var toolCancelled = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
20+
var releaseTool = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
21+
22+
var session = await CreateSessionAsync(new SessionConfig
23+
{
24+
Tools = [AIFunctionFactory.Create(SlowTool, "slow_analysis")],
25+
OnPermissionRequest = PermissionHandler.ApproveAll,
26+
});
27+
28+
_ = session.SendAsync(new MessageOptions
29+
{
30+
Prompt = "Use slow_analysis with value 'test_abort'. Wait for the result.",
31+
});
32+
33+
var startedValue = await toolStarted.Task.WaitAsync(TimeSpan.FromSeconds(60));
34+
Assert.Equal("test_abort", startedValue);
35+
36+
await session.DisposeAsync();
37+
await toolCancelled.Task.WaitAsync(TimeSpan.FromSeconds(60));
38+
39+
releaseTool.TrySetResult("RELEASED");
40+
41+
[Description("A slow analysis tool that blocks until released")]
42+
async Task<string> SlowTool([Description("Value to analyze")] string value, CancellationToken cancellationToken)
43+
{
44+
toolStarted.TrySetResult(value);
45+
try
46+
{
47+
var completed = await Task.WhenAny(releaseTool.Task, Task.Delay(Timeout.Infinite, cancellationToken));
48+
if (completed == releaseTool.Task)
49+
{
50+
return await releaseTool.Task;
51+
}
52+
53+
throw new OperationCanceledException(cancellationToken);
54+
}
55+
catch (OperationCanceledException)
56+
{
57+
toolCancelled.TrySetResult(true);
58+
throw;
59+
}
60+
}
61+
}
62+
}

0 commit comments

Comments
 (0)