From 571126c1a99ae73136efb6139d9e21f42f762ebe Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Thu, 3 Sep 2026 16:05:26 -0400 Subject: [PATCH 1/8] 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> --- CHANGELOG.md | 4 + dotnet/src/Client.cs | 21 ++ dotnet/src/Session.cs | 100 +++++++- .../test/Unit/ClientSessionLifetimeTests.cs | 174 ++++++++++++++ go/client.go | 16 ++ go/client_test.go | 61 +++++ go/session.go | 126 ++++++++-- go/session_test.go | 36 +++ go/types.go | 4 +- .../com/github/copilot/CopilotClient.java | 4 + .../com/github/copilot/CopilotSession.java | 107 ++++++++- .../com/github/copilot/JsonRpcClient.java | 35 +++ .../com/github/copilot/CopilotClientTest.java | 35 +++ .../com/github/copilot/JsonRpcClientTest.java | 31 +++ .../copilot/SessionEventHandlingTest.java | 106 ++++++++- nodejs/src/client.ts | 30 ++- nodejs/src/session.ts | 72 +++++- nodejs/src/types.ts | 2 + nodejs/test/client.test.ts | 37 +++ .../test/external-tool-cancellation.test.ts | 201 ++++++++++++++++ python/copilot/client.py | 18 ++ python/copilot/session.py | 53 ++++- python/test_client.py | 62 ++++- python/test_session.py | 59 +++++ rust/src/jsonrpc.rs | 10 + rust/src/session.rs | 107 ++++++++- rust/tests/session_test.rs | 221 ++++++++++++++++++ 27 files changed, 1675 insertions(+), 57 deletions(-) create mode 100644 nodejs/test/external-tool-cancellation.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index f7e753beda..a35962b69b 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index 92650ee60c..8d7fee8595 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -597,6 +597,10 @@ public async Task StopAsync() /// public async Task ForceStopAsync() { + foreach (var session in _sessions.Values) + { + session.CancelPendingExternalTools(); + } _sessions.Clear(); ClearGitHubTokenProviders(); @@ -2680,6 +2684,7 @@ private async Task 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); @@ -2703,6 +2708,22 @@ private async Task ConnectToServerAsync(Process? cliProcess, string? } } + private async Task CancelExternalToolsWhenConnectionClosesAsync(JsonRpc rpc) + { + await rpc.Completion.ConfigureAwait(false); + 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(); /// diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index 5995abaaff..77af34a24c 100644 --- a/dotnet/src/Session.cs +++ b/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; @@ -59,6 +60,8 @@ public sealed partial class CopilotSession : IAsyncDisposable private readonly Dictionary _toolHandlers = []; private readonly Dictionary> _commandHandlers = []; private readonly Dictionary>> _bearerTokenProviders = new(StringComparer.Ordinal); + private readonly ConcurrentDictionary _pendingExternalTools = new(StringComparer.Ordinal); + private readonly CancellationTokenSource _externalToolLifetime = new(); private readonly ILogger _logger; private readonly CopilotClient _parentClient; @@ -649,6 +652,10 @@ private async Task HandleBroadcastEventAsync(SessionEvent sessionEvent) break; } + case ExternalToolCompletedEvent completedEvent: + CancelExternalTool(completedEvent.Data.RequestId); + break; + case PermissionRequestedEvent permEvent: { var data = permEvent.Data; @@ -858,6 +865,12 @@ private async Task TryCancelMcpAuthRequestAsync(string requestId) /// private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, string toolCallId, JsonElement? arguments, AIFunction tool) { + using var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(_externalToolLifetime.Token); + if (!_pendingExternalTools.TryAdd(requestId, cancellationSource)) + { + return; + } + try { var invocation = new ToolInvocation @@ -877,7 +890,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) @@ -905,7 +918,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, @@ -915,9 +928,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, @@ -926,11 +944,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) { @@ -940,7 +980,56 @@ 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>)_pendingExternalTools) + .Remove(new(requestId, cancellationSource)); + } + } + + private bool TryClaimExternalTool(string requestId, CancellationTokenSource cancellationSource) + => ((ICollection>)_pendingExternalTools) + .Remove(new(requestId, cancellationSource)); + + private void CancelExternalTool(string requestId) + { + if (!string.IsNullOrEmpty(requestId) && _pendingExternalTools.TryRemove(requestId, out var cancellationSource)) + { + 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(); } /// @@ -1933,6 +2022,7 @@ public async ValueTask DisposeAsync() return; } + CancelPendingExternalTools(); _eventChannel.Writer.TryComplete(); try diff --git a/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index a59257b3da..fd221d6bf9 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -12,6 +12,7 @@ using System.Text; using System.Text.Json; using GitHub.Copilot.Rpc; +using Microsoft.Extensions.AI; using Xunit; namespace GitHub.Copilot.Test.Unit; @@ -950,6 +951,174 @@ public async Task McpAuth_Handler_Exception_Cancels_Pending_Request() Assert.Equal("cancelled", request.Params.GetProperty("result").GetProperty("kind").GetString()); } + [Fact] + public async Task ExternalToolCompleted_Cancels_Blocked_Tool_When_Cancellation_Callback_Throws() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + var toolStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var toolCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + await using var session = await client.CreateSessionAsync(new SessionConfig + { + Tools = [AIFunctionFactory.Create(BlockedTool, "blocked_tool")], + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + DispatchEvent(session, ExternalToolRequested("request-1")); + await toolStarted.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + DispatchEvent(session, new ExternalToolCompletedEvent + { + Data = new ExternalToolCompletedData { RequestId = "request-1" } + }); + + await toolCancelled.Task.WaitAsync(TimeSpan.FromSeconds(5)); + await Task.Delay(50); + Assert.DoesNotContain(server.Requests, + request => request.Method == "session.tools.handlePendingToolCall" + && request.Params.GetProperty("requestId").GetString() == "request-1"); + + async Task BlockedTool(CancellationToken cancellationToken) + { + toolStarted.TrySetResult(); + using var registration = cancellationToken.Register( + () => throw new InvalidOperationException("cancellation callback failed")); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + return "unreachable"; + } + catch (OperationCanceledException) + { + toolCancelled.TrySetResult(); + throw; + } + } + } + + [Fact] + public async Task ForceStopAsync_Cancels_Blocked_Tool_When_Cancellation_Callback_Throws() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + var toolStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var toolCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var session = await client.CreateSessionAsync(new SessionConfig + { + Tools = [AIFunctionFactory.Create(BlockedTool, "blocked_tool")], + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + DispatchEvent(session, ExternalToolRequested("request-force-stop")); + await toolStarted.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + await client.ForceStopAsync(); + + await toolCancelled.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + async Task BlockedTool(CancellationToken cancellationToken) + { + toolStarted.TrySetResult(); + using var registration = cancellationToken.Register( + () => throw new InvalidOperationException("cancellation callback failed")); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + return "unreachable"; + } + catch (OperationCanceledException) + { + toolCancelled.TrySetResult(); + throw; + } + } + } + + [Fact] + public async Task ConnectionClose_Cancels_Blocked_Tool_Delegate() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + var toolStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var toolCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var session = await client.CreateSessionAsync(new SessionConfig + { + Tools = [AIFunctionFactory.Create(BlockedTool, "blocked_tool")], + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + DispatchEvent(session, ExternalToolRequested("request-connection-close")); + await toolStarted.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + server.CloseConnection(); + + await toolCancelled.Task.WaitAsync(TimeSpan.FromSeconds(5)); + await client.ForceStopAsync(); + + async Task BlockedTool(CancellationToken cancellationToken) + { + toolStarted.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + return "unreachable"; + } + catch (OperationCanceledException) + { + toolCancelled.TrySetResult(); + throw; + } + } + } + + [Fact] + public async Task DisposeAsync_Cancels_Blocked_Tool_Delegate() + { + await using var server = await FakeCopilotServer.StartAsync(); + await using var client = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri(server.Url) }); + var toolStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var toolCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var session = await client.CreateSessionAsync(new SessionConfig + { + Tools = [AIFunctionFactory.Create(BlockedTool, "blocked_tool")], + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + DispatchEvent(session, ExternalToolRequested("request-2")); + await toolStarted.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + await session.DisposeAsync(); + + await toolCancelled.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + async Task BlockedTool(CancellationToken cancellationToken) + { + toolStarted.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + return "unreachable"; + } + catch (OperationCanceledException) + { + toolCancelled.TrySetResult(); + throw; + } + } + } + + private static ExternalToolRequestedEvent ExternalToolRequested(string requestId) => + new() + { + Data = new ExternalToolRequestedData + { + RequestId = requestId, + SessionId = "session-1", + ToolCallId = "tool-call-1", + ToolName = "blocked_tool" + } + }; + [Fact] public async Task Generated_Session_Rpc_Throws_When_Session_Disposed() { @@ -1514,6 +1683,11 @@ public void FailSessionCreate() _failSessionCreate = true; } + public void CloseConnection() + { + _stream?.Dispose(); + } + public async Task SendRequestAsync(string method, Dictionary parameters) { var stream = _stream ?? throw new InvalidOperationException("Client is not connected."); diff --git a/go/client.go b/go/client.go index 735e8ae414..ec4eeaf5ea 100644 --- a/go/client.go +++ b/go/client.go @@ -684,8 +684,15 @@ func (c *Client) ForceStop() { // Clear sessions immediately without trying to destroy them c.sessionsMux.Lock() + sessions := make([]*Session, 0, len(c.sessions)) + for _, session := range c.sessions { + sessions = append(sessions, session) + } c.sessions = make(map[string]*Session) c.sessionsMux.Unlock() + for _, session := range sessions { + session.cancelPendingExternalTools() + } c.clearGitHubTokenProviders() c.startStopMux.Lock() @@ -2489,6 +2496,15 @@ func (c *Client) clearGitHubTokenProviders() { func (c *Client) handleConnectionClose() { c.clearGitHubTokenProviders() + c.sessionsMux.Lock() + sessions := make([]*Session, 0, len(c.sessions)) + for _, session := range c.sessions { + sessions = append(sessions, session) + } + c.sessionsMux.Unlock() + for _, session := range sessions { + session.cancelPendingExternalTools() + } // Avoid deadlocking with Stop/ForceStop, which hold startStopMux while // waiting for the JSON-RPC read loop to finish. go func() { diff --git a/go/client_test.go b/go/client_test.go index e74c234637..cfda162205 100644 --- a/go/client_test.go +++ b/go/client_test.go @@ -514,6 +514,67 @@ func TestClient_ForceStopAndExternalStopDoNotRequestRuntimeShutdown(t *testing.T externalServer.Stop() } +func TestClient_ForceStopCancelsPendingExternalTools(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + session := &Session{ + pendingExternalTools: map[string]*pendingExternalTool{ + "request-1": {ctx: ctx, cancel: cancel}, + }, + } + + client := &Client{sessions: map[string]*Session{"session-1": session}} + + client.ForceStop() + + if len(client.sessions) != 0 { + t.Fatal("ForceStop did not clear sessions") + } + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("ForceStop did not cancel the pending external tool") + } +} + +func TestClient_ConnectionCloseCancelsPendingExternalTools(t *testing.T) { + rpcClient, server, _ := newRuntimeShutdownRpcPair(t) + server.SetRequestHandler("session.destroy", func(params json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + return []byte(`{}`), nil + }) + ctx, cancel := context.WithCancel(context.Background()) + session := newSession("session-1", rpcClient, "", false) + session.pendingExternalTools = map[string]*pendingExternalTool{ + "request-1": {ctx: ctx, cancel: cancel}, + } + client := &Client{ + client: rpcClient, + RPC: rpc.NewServerRPC(rpcClient), + sessions: map[string]*Session{"session-1": session}, + isExternalServer: true, + } + + client.handleConnectionClose() + + if len(client.sessions) != 1 { + t.Fatal("connection close removed sessions before Stop could clean them up") + } + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("connection close did not cancel the pending external tool") + } + + if err := client.Stop(); err != nil { + t.Fatalf("Stop failed after connection close: %v", err) + } + session.toolHandlersM.RLock() + defer session.toolHandlersM.RUnlock() + if session.toolHandlers != nil { + t.Fatal("Stop did not clean up the retained session") + } + server.Stop() +} + func newRuntimeShutdownRpcPair(t *testing.T) (*jsonrpc2.Client, *jsonrpc2.Client, chan struct{}) { t.Helper() diff --git a/go/session.go b/go/session.go index 5d45d19ef2..5e4c7788f7 100644 --- a/go/session.go +++ b/go/session.go @@ -65,6 +65,9 @@ type Session struct { handlerMutex sync.RWMutex toolHandlers map[string]ToolHandler toolHandlersM sync.RWMutex + pendingExternalTools map[string]*pendingExternalTool + pendingExternalToolsM sync.Mutex + externalToolsClosed bool permissionHandler PermissionHandlerFunc permissionMux sync.RWMutex managedSettings bool @@ -105,6 +108,11 @@ type Session struct { RPC *rpc.SessionRPC } +type pendingExternalTool struct { + ctx context.Context + cancel context.CancelFunc +} + // WorkspacePath returns the path to the session workspace directory when infinite // sessions are enabled. Contains checkpoints/, plan.md, and files/ subdirectories. // Returns empty string if infinite sessions are disabled. @@ -1389,7 +1397,19 @@ func fromRPCElicitationRequestedSchema(schema *rpc.ElicitationRequestedSchema) * // serial, FIFO dispatch without blocking the read loop. func (s *Session) dispatchEvent(event SessionEvent) { s.updateOpenCanvasesFromEvent(event) - go s.handleBroadcastEvent(event) + + broadcastHandled := false + switch data := event.Data.(type) { + case *ExternalToolRequestedData: + s.startExternalTool(data) + broadcastHandled = true + case *ExternalToolCompletedData: + s.cancelExternalTool(data.RequestID) + broadcastHandled = true + } + if !broadcastHandled { + go s.handleBroadcastEvent(event) + } // Send to the event channel in a closure with a recover guard. // Disconnect closes eventCh, and in Go sending on a closed channel @@ -1436,20 +1456,6 @@ func (s *Session) processEvents() { // cause RPC deadlocks. func (s *Session) handleBroadcastEvent(event SessionEvent) { switch d := event.Data.(type) { - case *ExternalToolRequestedData: - handler, ok := s.getToolHandler(d.ToolName) - if !ok { - return - } - var tp, ts string - if d.Traceparent != nil { - tp = *d.Traceparent - } - if d.Tracestate != nil { - ts = *d.Tracestate - } - s.executeToolAndRespond(d.RequestID, d.ToolName, d.ToolCallID, d.Arguments, handler, tp, ts) - case *PermissionRequestedData: if d.ResolvedByHook != nil && *d.ResolvedByHook { return // Already resolved by a permissionRequest hook; no client action needed. @@ -1532,11 +1538,90 @@ func (s *Session) handleBroadcastEvent(event SessionEvent) { } } +func (s *Session) startExternalTool(data *ExternalToolRequestedData) { + handler, ok := s.getToolHandler(data.ToolName) + if !ok { + return + } + + var traceparent, tracestate string + if data.Traceparent != nil { + traceparent = *data.Traceparent + } + if data.Tracestate != nil { + tracestate = *data.Tracestate + } + traceCtx := contextWithTraceParent(context.Background(), traceparent, tracestate) + ctx, cancel := context.WithCancel(traceCtx) + pending := &pendingExternalTool{ctx: ctx, cancel: cancel} + + s.pendingExternalToolsM.Lock() + if s.externalToolsClosed { + s.pendingExternalToolsM.Unlock() + cancel() + return + } + if s.pendingExternalTools == nil { + s.pendingExternalTools = make(map[string]*pendingExternalTool) + } + if _, exists := s.pendingExternalTools[data.RequestID]; exists { + s.pendingExternalToolsM.Unlock() + cancel() + return + } + s.pendingExternalTools[data.RequestID] = pending + s.pendingExternalToolsM.Unlock() + + go s.executeToolAndRespond(data.RequestID, data.ToolName, data.ToolCallID, data.Arguments, handler, pending) +} + +func (s *Session) cancelExternalTool(requestID string) { + s.pendingExternalToolsM.Lock() + pending := s.pendingExternalTools[requestID] + delete(s.pendingExternalTools, requestID) + s.pendingExternalToolsM.Unlock() + if pending != nil { + pending.cancel() + } +} + +func (s *Session) cancelPendingExternalTools() { + s.pendingExternalToolsM.Lock() + s.externalToolsClosed = true + pendingTools := s.pendingExternalTools + s.pendingExternalTools = nil + s.pendingExternalToolsM.Unlock() + for _, pending := range pendingTools { + pending.cancel() + } +} + +func (s *Session) claimExternalTool(requestID string, pending *pendingExternalTool) bool { + s.pendingExternalToolsM.Lock() + defer s.pendingExternalToolsM.Unlock() + if s.pendingExternalTools[requestID] != pending { + return false + } + delete(s.pendingExternalTools, requestID) + return true +} + // executeToolAndRespond executes a tool handler and sends the result back via RPC. -func (s *Session) executeToolAndRespond(requestID, toolName, toolCallID string, arguments any, handler ToolHandler, traceparent, tracestate string) { - ctx := contextWithTraceParent(context.Background(), traceparent, tracestate) +func (s *Session) executeToolAndRespond(requestID, toolName, toolCallID string, arguments any, handler ToolHandler, pending *pendingExternalTool) { + ctx := pending.ctx + defer func() { + s.pendingExternalToolsM.Lock() + if s.pendingExternalTools[requestID] == pending { + delete(s.pendingExternalTools, requestID) + } + s.pendingExternalToolsM.Unlock() + pending.cancel() + }() defer func() { if r := recover(); r != nil { + if !s.claimExternalTool(requestID, pending) { + return + } errMsg := fmt.Sprintf("tool panic: %v", r) s.RPC.Tools.HandlePendingToolCall(ctx, &rpc.HandlePendingToolCallRequest{ RequestID: requestID, @@ -1563,8 +1648,14 @@ func (s *Session) executeToolAndRespond(requestID, toolName, toolCallID string, invocation.AvailableTools = metadata.Tools } } + if ctx.Err() != nil { + return + } result, err := handler(invocation) + if !s.claimExternalTool(requestID, pending) { + return + } if err != nil { errMsg := err.Error() s.RPC.Tools.HandlePendingToolCall(ctx, &rpc.HandlePendingToolCallRequest{ @@ -1727,6 +1818,7 @@ func (s *Session) GetEvents(ctx context.Context) ([]SessionEvent, error) { // log.Printf("Failed to disconnect session: %v", err) // } func (s *Session) Disconnect() error { + s.cancelPendingExternalTools() _, err := s.client.Request(context.Background(), "session.destroy", sessionDestroyRequest{SessionID: s.SessionID}) s.closeOnce.Do(func() { close(s.eventCh) }) diff --git a/go/session_test.go b/go/session_test.go index 74e212c418..fabfccbc42 100644 --- a/go/session_test.go +++ b/go/session_test.go @@ -33,6 +33,42 @@ func newTestEvent() SessionEvent { return SessionEvent{Data: &SessionIdleData{}} } +func TestExternalToolCompletedCancelsBlockedHandler(t *testing.T) { + session, cleanup := newTestSession() + defer cleanup() + + started := make(chan struct{}) + cancelled := make(chan struct{}) + session.registerTools([]Tool{{ + Name: "blocked_tool", + Handler: func(invocation ToolInvocation) (ToolResult, error) { + close(started) + <-invocation.TraceContext.Done() + close(cancelled) + return ToolResult{}, invocation.TraceContext.Err() + }, + }}) + + session.dispatchEvent(SessionEvent{Data: &ExternalToolRequestedData{ + RequestID: "request-1", + SessionID: "session-1", + ToolCallID: "tool-call-1", + ToolName: "blocked_tool", + }}) + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("tool handler did not start") + } + + session.dispatchEvent(SessionEvent{Data: &ExternalToolCompletedData{RequestID: "request-1"}}) + select { + case <-cancelled: + case <-time.After(time.Second): + t.Fatal("tool handler was not cancelled") + } +} + func ptr[T any](value T) *T { return &value } diff --git a/go/types.go b/go/types.go index 772202a0bc..680f260756 100644 --- a/go/types.go +++ b/go/types.go @@ -1743,7 +1743,9 @@ type ToolInvocation struct { // TraceContext carries the W3C Trace Context propagated from the CLI's // execute_tool span. Pass this to OpenTelemetry-aware code so that // child spans created inside the handler are parented to the CLI span. - // When no trace context is available this will be context.Background(). + // It is cancelled when the external tool request completes or the session + // disconnects, so background work must derive its own lifetime if it should + // outlive the invocation. TraceContext context.Context } diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java index c986e829a4..882dfaed8a 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java @@ -550,6 +550,8 @@ private Connection startCoreBody() { JsonRpcClient connectedRpc = rpc; Connection connection = new Connection(connectedRpc, process, new ServerRpc(connectedRpc::invoke), inProcessTransport == null ? null : inProcessTransport.host()); + connectedRpc.setCloseHandler( + () -> sessions.values().forEach(CopilotSession::cancelPendingExternalTools)); // Register handlers for server-to-client calls RpcHandlerDispatcher dispatcher = new RpcHandlerDispatcher(sessions, lifecycleManager::dispatch, executor, @@ -742,7 +744,9 @@ public CompletableFuture stop() { */ public CompletableFuture forceStop() { disposed = true; + var activeSessions = new ArrayList<>(sessions.values()); sessions.clear(); + activeSessions.forEach(CopilotSession::cancelPendingExternalTools); gitHubTokenProviders.clear(); // Dispatch the blocking shutdownOwnedExecutor() on a dedicated thread: // cleanupConnection() is chained off async work running on the owned diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java index f3a35967d3..fb37bf99b5 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java @@ -20,6 +20,7 @@ import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.logging.Level; @@ -30,6 +31,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.github.copilot.generated.AssistantMessageEvent; +import com.github.copilot.generated.ExternalToolCompletedEvent; import com.github.copilot.generated.rpc.SessionCommandsHandlePendingCommandParams; import com.github.copilot.generated.rpc.SessionLogParams; import com.github.copilot.generated.rpc.SessionLogLevel; @@ -184,6 +186,7 @@ public final class CopilotSession implements AutoCloseable { private volatile SessionRpc sessionRpc; private final Set> eventHandlers = ConcurrentHashMap.newKeySet(); private final Map toolHandlers = new ConcurrentHashMap<>(); + private final Map pendingExternalTools = new ConcurrentHashMap<>(); private final Map commandHandlers = new ConcurrentHashMap<>(); private final Map bearerTokenProviders = new ConcurrentHashMap<>(); private final AtomicReference permissionHandler = new AtomicReference<>(); @@ -204,6 +207,46 @@ public final class CopilotSession implements AutoCloseable { /** Tracks whether this session instance has been terminated via close(). */ private volatile boolean isTerminated = false; + private static final class PendingExternalTool { + private static final int WAITING = 0; + private static final int STARTED = 1; + private static final int CANCELLED = 2; + + private final AtomicInteger state = new AtomicInteger(WAITING); + private final AtomicReference> future = new AtomicReference<>(); + + T join(CompletableFuture operation) { + future.set(operation); + if (state.get() == CANCELLED) { + operation.cancel(true); + } + try { + return operation.join(); + } finally { + future.compareAndSet(operation, null); + } + } + + boolean tryStart() { + return state.compareAndSet(WAITING, STARTED); + } + + void attach(CompletableFuture toolFuture) { + future.set(toolFuture); + if (state.get() == CANCELLED) { + toolFuture.cancel(true); + } + } + + void cancel() { + state.set(CANCELLED); + CompletableFuture activeFuture = future.get(); + if (activeFuture != null) { + activeFuture.cancel(true); + } + } + } + /** * Creates a new session with the given ID and RPC client. *

@@ -857,6 +900,14 @@ private void handleBroadcastEventAsync(SessionEvent event) { } executeToolAndRespondAsync(data.requestId(), data.toolName(), data.toolCallId(), data.arguments(), tool); + } else if (event instanceof ExternalToolCompletedEvent completedEvent) { + var data = completedEvent.getData(); + if (data != null && data.requestId() != null) { + PendingExternalTool pending = pendingExternalTools.remove(data.requestId()); + if (pending != null) { + pending.cancel(); + } + } } else if (event instanceof PermissionRequestedEvent permEvent) { var data = permEvent.getData(); if (data == null || data.requestId() == null || data.permissionRequest() == null) { @@ -928,9 +979,9 @@ private void handleBroadcastEventAsync(SessionEvent event) { * built-in tool-search tool, so an override can filter the live catalog without * issuing its own RPC. The snapshot is fetched only for that tool to avoid a * round-trip on every ordinary tool call; a failed fetch leaves the snapshot - * {@code null} rather than failing the tool. Shared by both server-to-client - * tool dispatch paths ({@link RpcHandlerDispatcher} and - * {@link #executeToolAndRespondAsync}). + * {@code null} rather than failing the tool. Used by the direct RPC dispatch + * path; event-dispatched tools perform the same lookup with request-scoped + * cancellation. * * @param toolName * the name of the tool being invoked @@ -955,6 +1006,12 @@ void populateToolSearchMetadata(String toolName, com.github.copilot.rpc.ToolInvo */ private void executeToolAndRespondAsync(String requestId, String toolName, String toolCallId, Object arguments, ToolDefinition tool) { + var pending = new PendingExternalTool(); + synchronized (this) { + if (isTerminated || pendingExternalTools.putIfAbsent(requestId, pending) != null) { + return; + } + } Runnable task = () -> { try { JsonNode argumentsNode = arguments instanceof JsonNode jn @@ -963,9 +1020,32 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin var invocation = new com.github.copilot.rpc.ToolInvocation().setSessionId(sessionId) .setToolCallId(toolCallId).setToolName(toolName).setArguments(argumentsNode); - populateToolSearchMetadata(toolName, invocation); + if (TOOL_SEARCH_TOOL_NAME.equals(toolName)) { + try { + var metadata = pending.join(getRpc().tools.getCurrentMetadata()); + if (metadata != null) { + invocation.setAvailableTools(metadata.tools()); + } + } catch (RuntimeException e) { + if (pendingExternalTools.get(requestId) != pending) { + return; + } + LOG.log(Level.FINE, "Failed to fetch tool metadata for tool search", e); + } + if (pendingExternalTools.get(requestId) != pending) { + return; + } + } - tool.handler().invoke(invocation).thenAccept(result -> { + if (!pending.tryStart()) { + return; + } + CompletableFuture toolFuture = tool.handler().invoke(invocation); + pending.attach(toolFuture); + toolFuture.thenAccept(result -> { + if (!pendingExternalTools.remove(requestId, pending)) { + return; + } try { ToolResultObject toolResult; if (result instanceof ToolResultObject tr) { @@ -980,6 +1060,9 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin LOG.log(Level.WARNING, "Error sending tool result for requestId=" + requestId, e); } }).exceptionally(ex -> { + if (!pendingExternalTools.remove(requestId, pending)) { + return null; + } try { getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId, requestId, null, ex.getMessage() != null ? ex.getMessage() : ex.toString())); @@ -987,8 +1070,11 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin LOG.log(Level.WARNING, "Error sending tool error for requestId=" + requestId, e); } return null; - }); + }).whenComplete((result, error) -> pendingExternalTools.remove(requestId, pending)); } catch (Exception e) { + if (!pendingExternalTools.remove(requestId, pending)) { + return; + } LOG.log(Level.WARNING, "Error executing tool for requestId=" + requestId, e); try { getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId, @@ -1010,6 +1096,14 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin } } + void cancelPendingExternalTools() { + pendingExternalTools.forEach((requestId, pending) -> { + if (pendingExternalTools.remove(requestId, pending)) { + pending.cancel(); + } + }); + } + /** * Builds a {@link SessionUiHandlePendingElicitationParams} carrying a * {@code cancel} action, used when an elicitation handler throws or the handler @@ -2315,6 +2409,7 @@ public void close() { isTerminated = true; } + cancelPendingExternalTools(); timeoutScheduler.shutdownNow(); releaseGitHubTokenProviderRegistration(); diff --git a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java index 7eda069d23..f78ce00425 100644 --- a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java +++ b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java @@ -53,7 +53,10 @@ class JsonRpcClient implements AutoCloseable { private final Map> pendingRequests = new ConcurrentHashMap<>(); private final Map> notificationHandlers = new ConcurrentHashMap<>(); private final ExecutorService readerExecutor; + private final Object closeHandlerLock = new Object(); private volatile boolean running = true; + private boolean closeNotified; + private Runnable closeHandler; private JsonRpcClient(InputStream inputStream, OutputStream outputStream, Socket socket, Process process) { this(inputStream, outputStream, socket, process, false); @@ -322,10 +325,41 @@ private void startReader() { if (running) { LOG.log(Level.SEVERE, "Error in JSON-RPC reader", e); } + } finally { + notifyClose(); } }); } + void setCloseHandler(Runnable handler) { + boolean runNow; + synchronized (closeHandlerLock) { + closeHandler = handler; + runNow = closeNotified; + } + if (runNow) { + handler.run(); + } + } + + private void notifyClose() { + Runnable handler; + synchronized (closeHandlerLock) { + if (closeNotified) { + return; + } + closeNotified = true; + handler = closeHandler; + } + if (handler != null) { + try { + handler.run(); + } catch (RuntimeException e) { + LOG.log(Level.WARNING, "Error handling JSON-RPC connection close", e); + } + } + } + private void handleMessage(String content) { try { JsonNode node = MAPPER.readTree(content); @@ -390,6 +424,7 @@ else if (node.has("method")) { public void close() { running = false; readerExecutor.shutdownNow(); + notifyClose(); // Cancel all pending requests pendingRequests.forEach((id, future) -> future.completeExceptionally(new IOException("Client closed"))); diff --git a/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java b/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java index 4cfd7f4c48..c9ce0825fb 100644 --- a/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java +++ b/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java @@ -7,6 +7,7 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import com.github.copilot.generated.ExternalToolRequestedEvent; import com.github.copilot.rpc.CopilotClientOptions; import com.github.copilot.rpc.DeleteSessionResponse; import com.github.copilot.rpc.GitHubTokenProviderResult; @@ -15,12 +16,16 @@ import com.github.copilot.rpc.SessionConfig; import com.github.copilot.rpc.SessionLifecycleEvent; import com.github.copilot.rpc.SessionLifecycleEventTypes; +import com.github.copilot.rpc.ToolDefinition; import java.lang.reflect.Field; import java.util.ArrayList; +import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; @@ -113,6 +118,36 @@ void testForceStopAndExternalStopDoNotRequestRuntimeShutdown() throws Exception verify(externalRpc, never()).invoke(eq("runtime.shutdown"), any(), eq(Void.class)); } + @Test + @SuppressWarnings("unchecked") + void testForceStopCancelsPendingExternalTools() throws Exception { + var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false)); + var rpc = mock(JsonRpcClient.class); + setConnectionFuture(client, rpc, null); + var session = new CopilotSession("force-stop-session", rpc); + var toolFuture = new CompletableFuture(); + var started = new CountDownLatch(1); + session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks", Map.of(), invocation -> { + started.countDown(); + return toolFuture; + }))); + Field sessionsField = CopilotClient.class.getDeclaredField("sessions"); + sessionsField.setAccessible(true); + var sessions = (Map) sessionsField.get(client); + sessions.put(session.getSessionId(), session); + + var requested = new ExternalToolRequestedEvent(); + requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-force-stop", + session.getSessionId(), "tool-call-force-stop", "blocked_tool", null, Map.of(), null, null, null)); + session.dispatchEvent(requested); + assertTrue(started.await(1, TimeUnit.SECONDS)); + + client.forceStop().get(); + + assertThrows(CancellationException.class, () -> toolFuture.get(1, TimeUnit.SECONDS)); + assertTrue(sessions.isEmpty()); + } + @Test @SuppressWarnings("unchecked") void testDeleteSessionReleasesGitHubTokenProvider() throws Exception { diff --git a/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java b/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java index 009f15c200..13668cd97f 100644 --- a/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java +++ b/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java @@ -16,6 +16,7 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; @@ -184,6 +185,36 @@ void testGetProcessNullForSocket() throws Exception { } } + @Test + void testCloseHandlerRunsOnceOnRemoteAndExplicitClose() throws Exception { + try (var pair = createSocketPair()) { + var closeCount = new AtomicInteger(); + var closed = new CompletableFuture(); + pair.client.setCloseHandler(() -> { + closeCount.incrementAndGet(); + closed.complete(null); + }); + + pair.serverSide.close(); + closed.get(5, TimeUnit.SECONDS); + pair.client.close(); + + assertEquals(1, closeCount.get()); + } + } + + @Test + void testCloseHandlerRunsWhenRegisteredAfterClose() throws Exception { + try (var pair = createSocketPair()) { + pair.client.close(); + var closed = new CompletableFuture(); + + pair.client.setCloseHandler(() -> closed.complete(null)); + + closed.get(5, TimeUnit.SECONDS); + } + } + // ---- invoke() edge cases ---- @Test diff --git a/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java b/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java index b75e710720..cb217b023f 100644 --- a/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java +++ b/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java @@ -10,9 +10,12 @@ import java.io.Closeable; import java.lang.reflect.Method; +import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CancellationException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -26,11 +29,15 @@ import com.github.copilot.generated.SessionEvent; import com.github.copilot.generated.AssistantMessageEvent; +import com.github.copilot.generated.ExternalToolCompletedEvent; +import com.github.copilot.generated.ExternalToolRequestedEvent; import com.github.copilot.generated.SessionIdleEvent; import com.github.copilot.generated.SessionMode; import com.github.copilot.generated.SessionStartEvent; +import com.github.copilot.generated.rpc.SessionToolsGetCurrentMetadataResult; import com.github.copilot.rpc.MessageOptions; import com.github.copilot.rpc.SendMessageResponse; +import com.github.copilot.rpc.ToolDefinition; /** * Unit tests for session event handling API. @@ -50,10 +57,14 @@ void setup() throws Exception { } private CopilotSession createTestSession() throws Exception { + return createTestSession(null); + } + + private CopilotSession createTestSession(JsonRpcClient rpc) throws Exception { // Use the package-private constructor via reflection for testing var constructor = CopilotSession.class.getDeclaredConstructor(String.class, JsonRpcClient.class, String.class); constructor.setAccessible(true); - return constructor.newInstance("test-session-id", null, null); + return constructor.newInstance("test-session-id", rpc, null); } @Test @@ -73,6 +84,99 @@ void testGenericEventHandler() { assertInstanceOf(SessionIdleEvent.class, receivedEvents.get(2)); } + @Test + void testExternalToolCompletedCancelsBlockedHandler() throws Exception { + var toolFuture = new CompletableFuture(); + var started = new CountDownLatch(1); + session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks", Map.of(), invocation -> { + started.countDown(); + return toolFuture; + }))); + + var requested = new ExternalToolRequestedEvent(); + requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-1", + "test-session-id", "tool-call-1", "blocked_tool", null, Map.of(), null, null, null)); + dispatchEvent(requested); + assertTrue(started.await(1, TimeUnit.SECONDS)); + + var completed = new ExternalToolCompletedEvent(); + completed.setData(new ExternalToolCompletedEvent.ExternalToolCompletedEventData("request-1")); + dispatchEvent(completed); + + assertThrows(CancellationException.class, () -> toolFuture.get(1, TimeUnit.SECONDS)); + } + + @Test + void testCloseCancelsBlockedExternalTool() throws Exception { + var toolFuture = new CompletableFuture(); + var started = new CountDownLatch(1); + session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks", Map.of(), invocation -> { + started.countDown(); + return toolFuture; + }))); + + var requested = new ExternalToolRequestedEvent(); + requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-close", + "test-session-id", "tool-call-close", "blocked_tool", null, Map.of(), null, null, null)); + dispatchEvent(requested); + assertTrue(started.await(1, TimeUnit.SECONDS)); + + session.close(); + + assertThrows(CancellationException.class, () -> toolFuture.get(1, TimeUnit.SECONDS)); + } + + @Test + void testExternalToolCompletedDoesNotBlockOnSynchronousHandler() throws Exception { + var started = new CountDownLatch(1); + var release = new CountDownLatch(1); + session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks synchronously", Map.of(), + invocation -> { + started.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return CompletableFuture.completedFuture("done"); + }))); + + var requested = new ExternalToolRequestedEvent(); + requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-sync", + "test-session-id", "tool-call-sync", "blocked_tool", null, Map.of(), null, null, null)); + dispatchEvent(requested); + assertTrue(started.await(1, TimeUnit.SECONDS)); + + var completed = new ExternalToolCompletedEvent(); + completed.setData(new ExternalToolCompletedEvent.ExternalToolCompletedEventData("request-sync")); + try { + assertTimeoutPreemptively(Duration.ofSeconds(1), () -> dispatchEvent(completed)); + } finally { + release.countDown(); + } + } + + @Test + void testToolSearchRunsWhenMetadataResultIsNull() throws Exception { + var rpc = mock(JsonRpcClient.class); + CompletableFuture metadata = CompletableFuture.completedFuture(null); + when(rpc.invoke(eq("session.tools.getCurrentMetadata"), any(), + eq(SessionToolsGetCurrentMetadataResult.class))).thenReturn(metadata); + session = createTestSession(rpc); + var invoked = new CountDownLatch(1); + session.registerTools(List.of(ToolDefinition.create("tool_search_tool", "Searches", Map.of(), invocation -> { + invoked.countDown(); + return new CompletableFuture<>(); + }))); + + var requested = new ExternalToolRequestedEvent(); + requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-search", + "test-session-id", "tool-call-search", "tool_search_tool", null, Map.of(), null, null, null)); + dispatchEvent(requested); + + assertTrue(invoked.await(1, TimeUnit.SECONDS)); + } + @Test void testTypedEventHandler() { var receivedMessages = new ArrayList(); diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts index 83b0a5f483..0e55e7369e 100644 --- a/nodejs/src/client.ts +++ b/nodejs/src/client.ts @@ -518,6 +518,7 @@ export class CopilotClient { private ffiHost: FfiRuntimeHost | null = null; private connection: MessageConnection | null = null; private messageWriter: TeardownResilientStreamMessageWriter | null = null; + private connectionClosed: boolean = false; private socket: Socket | null = null; private runtimePort: number | null = null; private actualHost: string = "localhost"; @@ -996,6 +997,7 @@ export class CopilotClient { } this.forceStopping = false; + this.connectionClosed = false; this.processTransportError = null; this.state = "connecting"; @@ -1131,7 +1133,12 @@ export class CopilotClient { // Ask SDK-owned runtimes to flush and clean up before we tear down // their transport/process. External runtimes may be shared, so only // close our connection to them. - if (this.connection && (this.cliProcess || this.ffiHost) && !this.isExternalServer) { + if ( + this.connection && + !this.connectionClosed && + (this.cliProcess || this.ffiHost) && + !this.isExternalServer + ) { const runtimeShutdownStart = Date.now(); const shutdownPromise = this.rpc.runtime.shutdown(); void shutdownPromise.catch(() => undefined); @@ -3101,13 +3108,24 @@ export class CopilotClient { } ); - this.connection.onClose(() => { + const connection = this.connection; + const markDisconnected = () => { + if (this.connection !== connection) { + return; + } + this.connectionClosed = true; this.state = "disconnected"; + for (const session of this.sessions.values()) { + session._markDisconnected(); + } + this.sessions.clear(); this.githubTokenProviders.clear(); - }); - - this.connection.onError((_error) => { - this.state = "disconnected"; + }; + this.connection.onClose(markDisconnected); + this.connection.onError(() => { + if (this.connection === connection) { + this.state = "disconnected"; + } }); } diff --git a/nodejs/src/session.ts b/nodejs/src/session.ts index eb4c6b7561..18521b8f56 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -420,6 +420,7 @@ export class CopilotSession { private typedEventHandlers: Map void>> = new Map(); private toolHandlers: Map = new Map(); + private pendingExternalTools: Map = new Map(); private canvases: Map = new Map(); private bearerTokenProviders: Map = new Map(); private commandHandlers: Map = new Map(); @@ -439,6 +440,7 @@ export class CopilotSession { private _capabilities: SessionCapabilities = {}; private openCanvasInstances: OpenCanvasInstance[] = []; private disconnected = false; + private disconnecting = false; private onDisconnected?: () => void; /** @internal Client session API handlers, populated by CopilotClient during create/resume. */ @@ -818,6 +820,10 @@ export class CopilotSession { return; } this.disconnected = true; + for (const controller of this.pendingExternalTools.values()) { + controller.abort(); + } + this.pendingExternalTools.clear(); this._runOnDisconnected(); this.eventHandlers.clear(); this.typedEventHandlers.clear(); @@ -994,6 +1000,15 @@ export class CopilotSession { tracestate ); } + } else if (event.type === "external_tool.completed") { + const { requestId } = event.data as { requestId?: string }; + if (requestId) { + const controller = this.pendingExternalTools.get(requestId); + if (controller) { + this.pendingExternalTools.delete(requestId); + controller.abort(); + } + } } else if (event.type === "permission.requested") { const { requestId, permissionRequest, resolvedByHook } = event.data as { requestId: string; @@ -1103,6 +1118,12 @@ export class CopilotSession { traceparent?: string, tracestate?: string ): Promise { + const controller = new AbortController(); + if (this.disconnected || this.pendingExternalTools.has(requestId)) { + return; + } + this.pendingExternalTools.set(requestId, controller); + try { // The built-in tool-search tool receives a snapshot of the session's // currently initialized tools so an override can filter the live @@ -1111,13 +1132,27 @@ export class CopilotSession { // leaves the snapshot undefined rather than failing the tool. let availableTools: CurrentToolMetadata[] | undefined; if (toolName === TOOL_SEARCH_TOOL_NAME) { + if (controller.signal.aborted) { + return; + } + const aborted = new Promise((resolve) => { + controller.signal.addEventListener("abort", () => resolve(undefined), { + once: true, + }); + }); try { - const metadata = await this.rpc.tools.getCurrentMetadata(); - availableTools = metadata.tools ?? undefined; + const metadata = await Promise.race([ + this.rpc.tools.getCurrentMetadata(), + aborted, + ]); + availableTools = metadata?.tools ?? undefined; } catch { availableTools = undefined; } } + if (controller.signal.aborted) { + return; + } const rawResult = await handler(args, { sessionId: this.sessionId, toolCallId, @@ -1126,6 +1161,7 @@ export class CopilotSession { availableTools, traceparent, tracestate, + signal: controller.signal, }); let result: ToolResult; if (rawResult == null) { @@ -1137,12 +1173,12 @@ export class CopilotSession { } else { result = JSON.stringify(rawResult); } - if (this.disconnected) { + if (!this._claimExternalTool(requestId, controller)) { return; } await this.rpc.tools.handlePendingToolCall({ requestId, result }); } catch (error) { - if (this.disconnected) { + if (!this._claimExternalTool(requestId, controller)) { return; } const message = error instanceof Error ? error.message : String(error); @@ -1154,9 +1190,21 @@ export class CopilotSession { } // Connection lost or RPC error — nothing we can do } + } finally { + if (this.pendingExternalTools.get(requestId) === controller) { + this.pendingExternalTools.delete(requestId); + } } } + private _claimExternalTool(requestId: string, controller: AbortController): boolean { + if (this.disconnected || this.pendingExternalTools.get(requestId) !== controller) { + return false; + } + this.pendingExternalTools.delete(requestId); + return true; + } + /** * Executes a permission handler and sends the result back via RPC. * @internal @@ -2014,13 +2062,19 @@ export class CopilotSession { * ``` */ async disconnect(): Promise { - if (this.disconnected) { + if (this.disconnected || this.disconnecting) { return; } - await this.connection.sendRequest("session.destroy", { - sessionId: this.sessionId, - }); - this._markDisconnected(); + this.disconnecting = true; + try { + await this.connection.sendRequest("session.destroy", { + sessionId: this.sessionId, + }); + this._markDisconnected(); + } catch (error) { + this.disconnecting = false; + throw error; + } } /** Enables `await using session = ...` syntax for automatic cleanup. */ diff --git a/nodejs/src/types.ts b/nodejs/src/types.ts index 128ced3d68..fc7b5370e7 100644 --- a/nodejs/src/types.ts +++ b/nodejs/src/types.ts @@ -687,6 +687,8 @@ export interface ToolInvocation { traceparent?: string; /** W3C Trace Context tracestate from the CLI's execute_tool span. */ tracestate?: string; + /** Aborted when the runtime completes this request or the session disconnects. */ + signal?: AbortSignal; } export type ToolHandler = ( diff --git a/nodejs/test/client.test.ts b/nodejs/test/client.test.ts index bd1cf9812f..f1ec027cc2 100644 --- a/nodejs/test/client.test.ts +++ b/nodejs/test/client.test.ts @@ -3201,6 +3201,42 @@ describe("CopilotClient", () => { const client = new CopilotClient(); await client.start(); onTestFinished(() => stopClient(client)); + let invocationSignal: AbortSignal | undefined; + let toolStarted!: () => void; + const started = new Promise((resolve) => { + toolStarted = resolve; + }); + const session = await client.createSession({ + onPermissionRequest: approveAll, + tools: [ + { + name: "blocked_tool", + description: "blocks until cancelled", + handler: async (_args, invocation) => { + invocationSignal = invocation.signal; + toolStarted(); + await new Promise((_, reject) => + invocation.signal?.addEventListener( + "abort", + () => reject(invocation.signal?.reason), + { once: true } + ) + ); + }, + }, + ], + }); + (session as any)._handleBroadcastEvent({ + type: "external_tool.requested", + data: { + requestId: "request-connection-close", + sessionId: session.sessionId, + toolCallId: "tool-call-connection-close", + toolName: "blocked_tool", + arguments: {}, + }, + }); + await started; expect((client as any).state).toBe("connected"); @@ -3212,6 +3248,7 @@ describe("CopilotClient", () => { // Wait for the connection.onClose handler to fire await vi.waitFor(() => { expect((client as any).state).toBe("disconnected"); + expect(invocationSignal?.aborted).toBe(true); }); } ); diff --git a/nodejs/test/external-tool-cancellation.test.ts b/nodejs/test/external-tool-cancellation.test.ts new file mode 100644 index 0000000000..11c81c9f9d --- /dev/null +++ b/nodejs/test/external-tool-cancellation.test.ts @@ -0,0 +1,201 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +import { expect, it, vi } from "vitest"; +import { CopilotSession } from "../src/session.js"; +import type { ToolInvocation } from "../src/types.js"; + +it("cancels a blocked external tool when completion arrives", async () => { + const session = new CopilotSession("session-1", {} as never); + let invocation: ToolInvocation | undefined; + let started!: () => void; + const toolStarted = new Promise((resolve) => { + started = resolve; + }); + + (session as any).toolHandlers.set( + "blocked_tool", + async (_args: unknown, context: ToolInvocation) => { + invocation = context; + started(); + await new Promise((_, reject) => + context.signal?.addEventListener("abort", () => reject(context.signal?.reason), { + once: true, + }) + ); + } + ); + + (session as any)._handleBroadcastEvent({ + type: "external_tool.requested", + data: { + requestId: "request-1", + sessionId: "session-1", + toolCallId: "tool-call-1", + toolName: "blocked_tool", + arguments: {}, + }, + }); + await toolStarted; + + (session as any)._handleBroadcastEvent({ + type: "external_tool.completed", + data: { requestId: "request-1" }, + }); + + expect(invocation?.signal?.aborted).toBe(true); +}); + +it("does not respond when a cancelled handler returns a late result", async () => { + const sendRequest = vi.fn().mockResolvedValue(undefined); + const session = new CopilotSession("session-1", { sendRequest } as never); + let started!: () => void; + const toolStarted = new Promise((resolve) => { + started = resolve; + }); + + (session as any).toolHandlers.set( + "late_tool", + async (_args: unknown, context: ToolInvocation) => { + started(); + await new Promise((resolve) => + context.signal?.addEventListener("abort", () => resolve(), { once: true }) + ); + return "late result"; + } + ); + + (session as any)._handleBroadcastEvent({ + type: "external_tool.requested", + data: { + requestId: "request-late", + sessionId: "session-1", + toolCallId: "tool-call-late", + toolName: "late_tool", + arguments: {}, + }, + }); + await toolStarted; + (session as any)._handleBroadcastEvent({ + type: "external_tool.completed", + data: { requestId: "request-late" }, + }); + await vi.waitFor(() => expect((session as any).pendingExternalTools.size).toBe(0)); + + expect(sendRequest).not.toHaveBeenCalled(); +}); + +it("remains retryable when disconnect fails", async () => { + const sendRequest = vi + .fn() + .mockRejectedValueOnce(new Error("transient")) + .mockResolvedValueOnce(undefined); + const session = new CopilotSession("session-1", { sendRequest } as never); + const controller = new AbortController(); + (session as any).pendingExternalTools.set("request-1", controller); + + await expect(session.disconnect()).rejects.toThrow("transient"); + expect(controller.signal.aborted).toBe(false); + expect((session as any).pendingExternalTools.get("request-1")).toBe(controller); + await session.disconnect(); + + expect(sendRequest).toHaveBeenCalledTimes(2); + expect(controller.signal.aborted).toBe(true); +}); + +it("accepts tool requests while a failing disconnect is pending", async () => { + let rejectDestroy!: (error: Error) => void; + const sendRequest = vi.fn( + () => + new Promise((_, reject) => { + rejectDestroy = reject; + }) + ); + const session = new CopilotSession("session-1", { sendRequest } as never); + const handler = vi.fn( + (_args: unknown, context: ToolInvocation) => + new Promise((_, reject) => + context.signal?.addEventListener("abort", () => reject(context.signal?.reason), { + once: true, + }) + ) + ); + (session as any).toolHandlers.set("blocked_tool", handler); + + const disconnect = session.disconnect(); + (session as any)._handleBroadcastEvent({ + type: "external_tool.requested", + data: { + requestId: "request-during-disconnect", + sessionId: "session-1", + toolCallId: "tool-call-during-disconnect", + toolName: "blocked_tool", + arguments: {}, + }, + }); + await vi.waitFor(() => expect(handler).toHaveBeenCalledTimes(1)); + + rejectDestroy(new Error("transient")); + await expect(disconnect).rejects.toThrow("transient"); + (session as any)._handleBroadcastEvent({ + type: "external_tool.completed", + data: { requestId: "request-during-disconnect" }, + }); +}); + +it("cancels tool-search metadata preflight before invoking the handler", async () => { + const sendRequest = vi.fn(() => new Promise(() => {})); + const session = new CopilotSession("session-1", { sendRequest } as never); + const handler = vi.fn(); + (session as any).toolHandlers.set("tool_search_tool", handler); + + (session as any)._handleBroadcastEvent({ + type: "external_tool.requested", + data: { + requestId: "request-search", + sessionId: "session-1", + toolCallId: "tool-call-search", + toolName: "tool_search_tool", + arguments: {}, + }, + }); + await vi.waitFor(() => expect(sendRequest).toHaveBeenCalledTimes(1)); + + (session as any)._handleBroadcastEvent({ + type: "external_tool.completed", + data: { requestId: "request-search" }, + }); + await vi.waitFor(() => expect((session as any).pendingExternalTools.size).toBe(0)); + + expect(handler).not.toHaveBeenCalled(); +}); + +it("invokes duplicate request IDs only once", async () => { + const session = new CopilotSession("session-1", {} as never); + const handler = vi.fn( + (_args: unknown, context: ToolInvocation) => + new Promise((_, reject) => + context.signal?.addEventListener("abort", () => reject(context.signal?.reason), { + once: true, + }) + ) + ); + (session as any).toolHandlers.set("blocked_tool", handler); + const requested = { + type: "external_tool.requested", + data: { + requestId: "request-duplicate", + sessionId: "session-1", + toolCallId: "tool-call-duplicate", + toolName: "blocked_tool", + arguments: {}, + }, + }; + + (session as any)._handleBroadcastEvent(requested); + (session as any)._handleBroadcastEvent(requested); + await vi.waitFor(() => expect(handler).toHaveBeenCalledTimes(1)); + + (session as any)._handleBroadcastEvent({ + type: "external_tool.completed", + data: { requestId: "request-duplicate" }, + }); +}); diff --git a/python/copilot/client.py b/python/copilot/client.py index 929194c9b7..1a4d235914 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -2182,7 +2182,10 @@ async def force_stop(self) -> None: """ # Clear sessions immediately without trying to destroy them with self._sessions_lock: + sessions = list(self._sessions.values()) self._sessions.clear() + for session in sessions: + session._mark_disconnected() with self._github_token_providers_lock: self._github_token_providers.clear() @@ -4845,8 +4848,23 @@ def _register_github_token_provider( def _handle_connection_close(self) -> None: self._state = "disconnected" + with self._sessions_lock: + sessions = list(self._sessions.values()) + self._sessions.clear() with self._github_token_providers_lock: self._github_token_providers.clear() + client = self._client + loop = client._loop if client is not None else None + if loop is not None and not loop.is_closed(): + + def mark_disconnected_sessions() -> None: + for session in sessions: + session._mark_disconnected() + + try: + loop.call_soon_threadsafe(mark_disconnected_sessions) + except RuntimeError: + logger.debug("Event loop closed while handling connection loss") def _assign_github_token_provider(self, registration_id: str | None, session_id: str) -> None: if registration_id is None: diff --git a/python/copilot/session.py b/python/copilot/session.py index 3c6d3d54a4..54be0685f5 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -69,6 +69,7 @@ CapabilitiesChangedData, CommandExecuteData, ElicitationRequestedData, + ExternalToolCompletedData, ExternalToolRequestedData, McpOauthRequiredData, PermissionRequest, @@ -1574,6 +1575,7 @@ def __init__( self._event_handlers_lock = threading.Lock() self._tool_handlers: dict[str, ToolHandler] = {} self._tool_handlers_lock = threading.Lock() + self._pending_external_tools: dict[str, asyncio.Task[None]] = {} self._permission_handler: _PermissionHandlerFn | None = None self._permission_handler_lock = threading.Lock() self._mcp_auth_handler: McpAuthHandler | None = None @@ -1614,6 +1616,19 @@ def _run_disconnect_callback(self) -> None: if callback is not None: callback() + def _cancel_pending_external_tools(self) -> None: + pending_external_tools = list(self._pending_external_tools.values()) + self._pending_external_tools.clear() + current_task = asyncio.current_task() + for task in pending_external_tools: + if task is not current_task: + task.cancel() + + def _mark_disconnected(self) -> None: + self._destroyed = True + self._cancel_pending_external_tools() + self._run_disconnect_callback() + @property def rpc(self) -> SessionRpc: """Typed session-scoped RPC methods.""" @@ -1932,7 +1947,7 @@ def _handle_broadcast_event(self, event: SessionEvent) -> None: case ExternalToolRequestedData() as data: request_id = data.request_id tool_name = data.tool_name - if not request_id or not tool_name: + if self._destroyed or not request_id or not tool_name: return handler = self._get_tool_handler(tool_name) @@ -1943,11 +1958,26 @@ def _handle_broadcast_event(self, event: SessionEvent) -> None: arguments = data.arguments tp = getattr(data, "traceparent", None) ts = getattr(data, "tracestate", None) - asyncio.ensure_future( + task = asyncio.create_task( self._execute_tool_and_respond( request_id, tool_name, tool_call_id, arguments, handler, tp, ts ) ) + if request_id in self._pending_external_tools: + task.cancel() + return + self._pending_external_tools[request_id] = task + task.add_done_callback( + lambda completed, rid=request_id: self._remove_pending_external_tool( + rid, completed + ) + ) + + case ExternalToolCompletedData() as data: + if data.request_id: + task = self._pending_external_tools.pop(data.request_id, None) + if task is not None: + task.cancel() case PermissionRequestedData() as data: if logger.isEnabledFor(logging.DEBUG): @@ -2155,6 +2185,8 @@ async def _execute_tool_and_respond( # standard "Failed to execute..." message. Deliberate user-returned # failures send the full structured result to preserve metadata. if tool_result._from_exception: + if not self._claim_external_tool(request_id): + return rpc_start = time.perf_counter() await self.rpc.tools.handle_pending_tool_call( HandlePendingToolCallRequest( @@ -2173,6 +2205,8 @@ async def _execute_tool_and_respond( tool_name=tool_name, ) else: + if not self._claim_external_tool(request_id): + return rpc_start = time.perf_counter() await self.rpc.tools.handle_pending_tool_call( HandlePendingToolCallRequest( @@ -2191,6 +2225,8 @@ async def _execute_tool_and_respond( tool_name=tool_name, ) except Exception as exc: + if not self._claim_external_tool(request_id): + return try: await self.rpc.tools.handle_pending_tool_call( HandlePendingToolCallRequest( @@ -2201,6 +2237,17 @@ async def _execute_tool_and_respond( except (JsonRpcError, ProcessExitedError, OSError): pass # Connection lost or RPC error — nothing we can do + def _remove_pending_external_tool(self, request_id: str, completed: asyncio.Task[None]) -> None: + if self._pending_external_tools.get(request_id) is completed: + self._pending_external_tools.pop(request_id, None) + + def _claim_external_tool(self, request_id: str) -> bool: + current = asyncio.current_task() + if self._destroyed or self._pending_external_tools.get(request_id) is not current: + return False + self._pending_external_tools.pop(request_id, None) + return True + async def _execute_permission_and_respond( self, request_id: str, @@ -2983,6 +3030,8 @@ async def disconnect(self) -> None: return self._destroyed = True + self._cancel_pending_external_tools() + try: await self._client.request("session.destroy", {"sessionId": self.session_id}) finally: diff --git a/python/test_client.py b/python/test_client.py index bf8de113ae..f4edda8962 100644 --- a/python/test_client.py +++ b/python/test_client.py @@ -38,7 +38,7 @@ ModelLimits, ModelSupports, ) -from copilot.session import PermissionHandler +from copilot.session import CopilotSession, PermissionHandler from copilot.session_events import ( McpOauthRequestReason, McpOauthRequiredData, @@ -195,6 +195,66 @@ async def test_force_stop_external_server_clears_process_references(self): assert client._process is None assert client._cli_process is None + @pytest.mark.asyncio + async def test_force_stop_cancels_pending_external_tools(self): + client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) + session = CopilotSession("session-1", Mock()) + cancelled = asyncio.Event() + + async def blocked(): + try: + await asyncio.Future() + finally: + cancelled.set() + + task = asyncio.create_task(blocked()) + session._pending_external_tools["request-1"] = task + client._sessions["session-1"] = session + await asyncio.sleep(0) + + await client.force_stop() + + await asyncio.wait_for(cancelled.wait(), timeout=1) + assert session._destroyed + + @pytest.mark.asyncio + async def test_connection_close_cancels_pending_external_tools(self): + client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) + session = CopilotSession("session-1", Mock()) + cancelled = asyncio.Event() + + async def blocked(): + try: + await asyncio.Future() + finally: + cancelled.set() + + task = asyncio.create_task(blocked()) + session._pending_external_tools["request-1"] = task + client._sessions["session-1"] = session + await asyncio.sleep(0) + + client._client = Mock(_loop=asyncio.get_running_loop()) + await asyncio.to_thread(client._handle_connection_close) + + await asyncio.wait_for(cancelled.wait(), timeout=1) + assert session._destroyed + + def test_connection_close_tolerates_event_loop_close_race(self): + client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) + session = CopilotSession("session-1", Mock()) + loop = Mock() + loop.is_closed.return_value = False + loop.call_soon_threadsafe.side_effect = RuntimeError("Event loop is closed") + client._client = Mock(_loop=loop) + client._sessions["session-1"] = session + client._github_token_providers["registration-1"] = Mock() + + client._handle_connection_close() + + assert client._sessions == {} + assert client._github_token_providers == {} + class TestPermissionHandlerOptional: @pytest.mark.asyncio diff --git a/python/test_session.py b/python/test_session.py index dd2d0a72f6..638f739bb0 100644 --- a/python/test_session.py +++ b/python/test_session.py @@ -10,11 +10,14 @@ from copilot.session import CopilotSession from copilot.session_events import ( AssistantMessageData, + ExternalToolCompletedData, + ExternalToolRequestedData, SessionEvent, SessionEventType, SessionIdleData, SessionMode, ) +from copilot.tools import Tool, ToolResult def _event(data, event_type: SessionEventType) -> SessionEvent: @@ -67,3 +70,59 @@ async def test_send_and_wait_skips_autopilot_continuation_idle(): assert result is not None assert isinstance(result.data, AssistantMessageData) assert result.data.content == "final" + + +@pytest.mark.asyncio +async def test_external_tool_completed_cancels_blocked_handler(): + client = Mock() + client.request = AsyncMock() + session = CopilotSession("session-1", client) + started = asyncio.Event() + cancelled = asyncio.Event() + + async def blocked_tool(_invocation): + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancelled.set() + return ToolResult(text_result_for_llm="late result") + + session._register_tools([Tool("blocked_tool", "Blocks", blocked_tool)]) + session._dispatch_event( + _event( + ExternalToolRequestedData( + request_id="request-1", + session_id="session-1", + tool_call_id="tool-call-1", + tool_name="blocked_tool", + ), + SessionEventType.EXTERNAL_TOOL_REQUESTED, + ) + ) + await asyncio.wait_for(started.wait(), timeout=1) + + session._dispatch_event( + _event( + ExternalToolCompletedData(request_id="request-1"), + SessionEventType.EXTERNAL_TOOL_COMPLETED, + ) + ) + + await asyncio.wait_for(cancelled.wait(), timeout=1) + await asyncio.sleep(0) + client.request.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_disconnect_from_tool_task_does_not_cancel_destroy_request(): + client = Mock() + client.request = AsyncMock() + session = CopilotSession("session-1", client) + current_task = asyncio.current_task() + assert current_task is not None + session._pending_external_tools["request-1"] = current_task + + await session.disconnect() + + client.request.assert_awaited_once_with("session.destroy", {"sessionId": "session-1"}) diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 25a405080b..48e6090aed 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -9,6 +9,7 @@ use serde_json::Value; use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader}; use tokio::sync::{broadcast, mpsc, oneshot}; use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; use tracing::{Instrument, debug, error, warn}; use crate::{Error, ErrorKind, ProtocolErrorKind}; @@ -266,6 +267,7 @@ pub struct JsonRpcClient { pending_requests: Arc>>, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, + connection_closed: CancellationToken, read_task: Mutex>>, write_task: Mutex>>, } @@ -294,6 +296,7 @@ impl JsonRpcClient { pending_requests: Arc::new(RwLock::new(HashMap::new())), notification_tx, request_tx, + connection_closed: CancellationToken::new(), read_task: Mutex::new(None), write_task: Mutex::new(Some(write_task)), }; @@ -301,6 +304,7 @@ impl JsonRpcClient { let pending_requests = client.pending_requests.clone(); let notification_tx_clone = client.notification_tx.clone(); let request_tx_clone = client.request_tx.clone(); + let connection_closed = client.connection_closed.clone(); let reader_span = tracing::error_span!("jsonrpc_read_loop"); let read_task = tokio::spawn( @@ -312,6 +316,7 @@ impl JsonRpcClient { request_tx_clone, ) .await; + connection_closed.cancel(); } .instrument(reader_span), ); @@ -321,6 +326,7 @@ impl JsonRpcClient { } pub(crate) fn force_close(&self) { + self.connection_closed.cancel(); if let Some(task) = self.read_task.lock().take() { task.abort(); } @@ -330,6 +336,10 @@ impl JsonRpcClient { self.pending_requests.write().clear(); } + pub(crate) fn connection_closed_token(&self) -> CancellationToken { + self.connection_closed.child_token() + } + /// Writer-actor task. Owns the `AsyncWrite`, drains the command queue, /// and writes each frame atomically (header + body + flush) before /// signaling the ack. diff --git a/rust/src/session.rs b/rust/src/session.rs index 65c2aff7c4..0c170e6849 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -66,6 +66,41 @@ pub(crate) struct SessionHandlers { pub tools: Arc>>, } +type PendingExternalTools = Arc>>>; + +struct PendingExternalToolGuard { + request_id: RequestId, + token: Arc, + pending: PendingExternalTools, +} + +impl Drop for PendingExternalToolGuard { + fn drop(&mut self) { + let mut pending = self.pending.lock(); + if pending + .get(&self.request_id) + .is_some_and(|token| Arc::ptr_eq(token, &self.token)) + { + pending.remove(&self.request_id); + } + } +} + +impl PendingExternalToolGuard { + fn claim(&self) -> bool { + let mut pending = self.pending.lock(); + if pending + .get(&self.request_id) + .is_some_and(|token| Arc::ptr_eq(token, &self.token)) + { + pending.remove(&self.request_id); + true + } else { + false + } + } +} + fn has_managed_settings( enable_managed_settings: Option, managed_settings: Option<&crate::types::ManagedSettings>, @@ -177,6 +212,9 @@ pub struct Session { /// via [`Session::cancellation_token`] to bind their own work to /// the session lifetime. shutdown: CancellationToken, + /// Cancels only host-owned external tool callbacks. Disconnect signals this + /// before the destroy RPC without stopping unrelated event delivery. + external_tools_shutdown: CancellationToken, /// Only populated while a `send_and_wait` call is in flight. /// /// Sync `parking_lot::Mutex` because the lock is never held across an @@ -584,6 +622,7 @@ impl Session { Some(serde_json::json!({ "sessionId": self.id })), ) .await?; + self.external_tools_shutdown.cancel(); self.stop_event_loop().await; self.client.unregister_session(&self.id); self.github_token_registration.lock().take(); @@ -663,6 +702,7 @@ impl Drop for Session { // tokio runtime when it next polls; we intentionally don't await // it here because Drop is sync. self.shutdown.cancel(); + self.external_tools_shutdown.cancel(); self.client.unregister_session(&self.id); self.github_token_registration.lock().take(); } @@ -969,6 +1009,7 @@ impl Client { let idle_waiter = Arc::new(ParkingLotMutex::new(None)); let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new())); let shutdown = CancellationToken::new(); + let external_tools_shutdown = self.inner.rpc.connection_closed_token(); let (event_tx, _) = tokio::sync::broadcast::channel(512); // For cloud sessions (use_server_generated_id), defer session @@ -1071,6 +1112,7 @@ impl Client { open_canvases.clone(), event_tx.clone(), shutdown.clone(), + external_tools_shutdown.clone(), ); tracing::debug!( elapsed_ms = setup_start.elapsed().as_millis(), @@ -1098,6 +1140,7 @@ impl Client { client: self.clone(), event_loop: ParkingLotMutex::new(Some(event_loop)), shutdown, + external_tools_shutdown, idle_waiter, capabilities, open_canvases, @@ -1258,6 +1301,7 @@ impl Client { let idle_waiter = Arc::new(ParkingLotMutex::new(None)); let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new())); let shutdown = CancellationToken::new(); + let external_tools_shutdown = self.inner.rpc.connection_closed_token(); let (event_tx, _) = tokio::sync::broadcast::channel(512); let event_loop = spawn_event_loop( session_id.clone(), @@ -1275,6 +1319,7 @@ impl Client { open_canvases.clone(), event_tx.clone(), shutdown.clone(), + external_tools_shutdown.clone(), ); let mut registration = PendingSessionRegistration::new(self.clone(), session_id.clone(), shutdown.clone()); @@ -1373,6 +1418,7 @@ impl Client { client: self.clone(), event_loop: ParkingLotMutex::new(Some(event_loop)), shutdown, + external_tools_shutdown, idle_waiter, capabilities, open_canvases, @@ -1530,11 +1576,14 @@ fn spawn_event_loop( open_canvases: Arc>>, event_tx: tokio::sync::broadcast::Sender, shutdown: CancellationToken, + external_tools_shutdown: CancellationToken, ) -> JoinHandle<()> { let crate::router::SessionChannels { mut notifications, mut requests, } = channels; + let pending_external_tools: PendingExternalTools = + Arc::new(ParkingLotMutex::new(HashMap::new())); let span = tracing::error_span!("session_event_loop", session_id = %session_id); tokio::spawn( @@ -1567,7 +1616,7 @@ fn spawn_event_loop( _ = shutdown.cancelled() => break, Some(notification) = notifications.recv() => { handle_notification( - &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, &shutdown, + &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, &shutdown, &external_tools_shutdown, &pending_external_tools, ).await; } Some(request) = requests.recv() => { @@ -1720,6 +1769,8 @@ async fn handle_notification( open_canvases: &Arc>>, event_tx: &tokio::sync::broadcast::Sender, shutdown: &CancellationToken, + external_tools_shutdown: &CancellationToken, + pending_external_tools: &PendingExternalTools, ) { let dispatch_start = Instant::now(); let event = notification.event.clone(); @@ -1833,6 +1884,13 @@ async fn handle_notification( // Notification-based permission/tool/elicitation requests require a // separate RPC callback. Spawn concurrently since the CLI doesn't block. match event_type { + SessionEventType::ExternalToolCompleted => { + if let Some(request_id) = extract_request_id(¬ification.event.data) + && let Some(token) = pending_external_tools.lock().remove(&request_id) + { + token.cancel(); + } + } SessionEventType::PermissionRequested => { let Some(request_id) = extract_request_id(¬ification.event.data) else { return; @@ -1976,8 +2034,19 @@ async fn handle_notification( let Some(tool_handler) = tool_handler else { return; }; + let cancellation = Arc::new(external_tools_shutdown.child_token()); + { + let mut pending = pending_external_tools.lock(); + if external_tools_shutdown.is_cancelled() || pending.contains_key(&request_id) { + return; + } + pending.insert(request_id.clone(), cancellation.clone()); + } let client = client.clone(); let sid = session_id.clone(); + let pending_external_tools = pending_external_tools.clone(); + let guard_request_id = request_id.clone(); + let guard_cancellation = cancellation.clone(); let span = tracing::error_span!( "external_tool_handler", session_id = %sid, @@ -1985,11 +2054,22 @@ async fn handle_notification( ); tokio::spawn( async move { + let guard = PendingExternalToolGuard { + request_id: guard_request_id, + token: guard_cancellation, + pending: pending_external_tools, + }; + if cancellation.is_cancelled() { + return; + } // `tool_name.is_empty()` would have produced a `None` // lookup in `handlers.tools` and short-circuited at the // outer guard above, so only the tool_call_id check is // reachable here. if data.tool_call_id.is_empty() { + if !guard.claim() { + return; + } let error_msg = "Missing toolCallId"; let rpc_start = Instant::now(); let _ = client @@ -2019,13 +2099,15 @@ async fn handle_notification( // call; a failed fetch leaves the snapshot `None` rather than // failing the tool. let available_tools = if tool_name == TOOL_SEARCH_TOOL_NAME { - match client - .call( + let metadata_result = tokio::select! { + biased; + _ = cancellation.cancelled() => return, + result = client.call( rpc_methods::SESSION_TOOLS_GETCURRENTMETADATA, Some(serde_json::json!({ "sessionId": sid })), - ) - .await - { + ) => result, + }; + match metadata_result { Ok(value) => { serde_json::from_value::(value) .ok() @@ -2048,9 +2130,13 @@ async fn handle_notification( tracestate: data.tracestate, }; let handler_start = Instant::now(); - let tool_result = match tool_handler.call(invocation).await { - Ok(r) => r, - Err(e) => tool_failure_result(e.to_string()), + let tool_result = tokio::select! { + biased; + _ = cancellation.cancelled() => return, + result = tool_handler.call(invocation) => match result { + Ok(r) => r, + Err(e) => tool_failure_result(e.to_string()), + }, }; tracing::debug!( elapsed_ms = handler_start.elapsed().as_millis(), @@ -2060,6 +2146,9 @@ async fn handle_notification( tool_name = %tool_name, "ToolHandler::call dispatch" ); + if !guard.claim() { + return; + } let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null); let rpc_start = Instant::now(); let _ = client diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index 9f86777ffc..c33d1dd126 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -3857,6 +3857,227 @@ async fn external_tool_requested_dispatches_to_handler_and_responds() { assert_eq!(rpc_call["params"]["result"], "all tests passed"); } +#[tokio::test] +async fn external_tool_completed_cancels_blocked_handler() { + struct DropProbe(Option>); + + impl Drop for DropProbe { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + + struct BlockingTool { + started: parking_lot::Mutex>>, + cancelled: parking_lot::Mutex>>, + } + + #[async_trait] + impl tool::ToolHandler for BlockingTool { + async fn call( + &self, + _invocation: ToolInvocation, + ) -> Result { + if let Some(sender) = self.started.lock().take() { + let _ = sender.send(()); + } + let _probe = DropProbe(self.cancelled.lock().take()); + std::future::pending().await + } + } + + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (cancelled_tx, cancelled_rx) = tokio::sync::oneshot::channel(); + let (_session, mut server) = create_session_pair_with_config(|cfg| { + cfg.with_tools(vec![ + Tool::new("blocked_tool") + .with_description("Blocks") + .with_parameters(serde_json::json!({"type":"object"})) + .with_handler(Arc::new(BlockingTool { + started: parking_lot::Mutex::new(Some(started_tx)), + cancelled: parking_lot::Mutex::new(Some(cancelled_tx)), + })), + ]) + }) + .await; + + server + .send_event( + "external_tool.requested", + serde_json::json!({ + "requestId": "request-cancel-1", + "sessionId": server.session_id, + "toolCallId": "tool-call-cancel-1", + "toolName": "blocked_tool", + "arguments": {}, + }), + ) + .await; + timeout(TIMEOUT, started_rx).await.unwrap().unwrap(); + + server + .send_event( + "external_tool.completed", + serde_json::json!({ "requestId": "request-cancel-1" }), + ) + .await; + + timeout(TIMEOUT, cancelled_rx).await.unwrap().unwrap(); +} + +#[tokio::test] +async fn connection_close_cancels_blocked_external_tool() { + struct DropProbe(Option>); + + impl Drop for DropProbe { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + + struct BlockingTool { + started: parking_lot::Mutex>>, + cancelled: parking_lot::Mutex>>, + } + + #[async_trait] + impl tool::ToolHandler for BlockingTool { + async fn call( + &self, + _invocation: ToolInvocation, + ) -> Result { + if let Some(sender) = self.started.lock().take() { + let _ = sender.send(()); + } + let _probe = DropProbe(self.cancelled.lock().take()); + std::future::pending().await + } + } + + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (cancelled_tx, cancelled_rx) = tokio::sync::oneshot::channel(); + let (_session, mut server) = create_session_pair_with_config(|cfg| { + cfg.with_tools(vec![ + Tool::new("blocked_tool") + .with_description("Blocks") + .with_parameters(serde_json::json!({"type":"object"})) + .with_handler(Arc::new(BlockingTool { + started: parking_lot::Mutex::new(Some(started_tx)), + cancelled: parking_lot::Mutex::new(Some(cancelled_tx)), + })), + ]) + }) + .await; + + server + .send_event( + "external_tool.requested", + serde_json::json!({ + "requestId": "request-connection-close", + "sessionId": server.session_id, + "toolCallId": "tool-call-connection-close", + "toolName": "blocked_tool", + "arguments": {}, + }), + ) + .await; + timeout(TIMEOUT, started_rx).await.unwrap().unwrap(); + + drop(server); + + timeout(TIMEOUT, cancelled_rx).await.unwrap().unwrap(); +} + +#[tokio::test] +async fn disconnect_cancels_external_tools_before_stopping_session() { + struct DropProbe(Option>); + + impl Drop for DropProbe { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + + struct BlockingTool { + started: parking_lot::Mutex>>, + cancelled: parking_lot::Mutex>>, + } + + #[async_trait] + impl tool::ToolHandler for BlockingTool { + async fn call( + &self, + _invocation: ToolInvocation, + ) -> Result { + if let Some(sender) = self.started.lock().take() { + let _ = sender.send(()); + } + let _probe = DropProbe(self.cancelled.lock().take()); + std::future::pending().await + } + } + + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (cancelled_tx, mut cancelled_rx) = tokio::sync::oneshot::channel(); + let (session, mut server) = create_session_pair_with_config(|cfg| { + cfg.with_tools(vec![ + Tool::new("blocked_tool") + .with_description("Blocks") + .with_parameters(serde_json::json!({"type":"object"})) + .with_handler(Arc::new(BlockingTool { + started: parking_lot::Mutex::new(Some(started_tx)), + cancelled: parking_lot::Mutex::new(Some(cancelled_tx)), + })), + ]) + }) + .await; + let session = Arc::new(session); + let lifetime = session.cancellation_token(); + + server + .send_event( + "external_tool.requested", + serde_json::json!({ + "requestId": "request-disconnect", + "sessionId": server.session_id, + "toolCallId": "tool-call-disconnect", + "toolName": "blocked_tool", + "arguments": {}, + }), + ) + .await; + timeout(TIMEOUT, started_rx).await.unwrap().unwrap(); + + let disconnect = tokio::spawn({ + let session = session.clone(); + async move { session.disconnect().await } + }); + + let request = timeout(TIMEOUT, server.read_request()).await.unwrap(); + assert_eq!(request["method"], "session.destroy"); + assert!(!lifetime.is_cancelled()); + assert!( + timeout(Duration::from_millis(50), &mut cancelled_rx) + .await + .is_err() + ); + + server.respond(&request, serde_json::json!({})).await; + timeout(TIMEOUT, cancelled_rx).await.unwrap().unwrap(); + timeout(TIMEOUT, disconnect) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(lifetime.is_cancelled()); +} + #[tokio::test] async fn external_tool_broadcast_for_unknown_tool_is_not_responded_to() { // Phase H multi-client safety: a handler that doesn't claim the From 3acb3990d1a37c02cd053271bf09618a59549b3a Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Thu, 3 Sep 2026 17:22:58 -0400 Subject: [PATCH 2/8] Fix cancellation CI validation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../com/github/copilot/CopilotClient.java | 3 +-- .../com/github/copilot/CopilotSession.java | 3 +-- .../com/github/copilot/CopilotClientTest.java | 3 ++- .../copilot/SessionEventHandlingTest.java | 8 ++++---- python/test_session.py | 2 +- rust/src/session.rs | 20 +++++++++---------- 6 files changed, 19 insertions(+), 20 deletions(-) diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java index 1069982b2d..5b33f48aaa 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java @@ -553,8 +553,7 @@ private Connection startCoreBody() { JsonRpcClient connectedRpc = rpc; Connection connection = new Connection(connectedRpc, process, new ServerRpc(connectedRpc::invoke), inProcessTransport == null ? null : inProcessTransport.host()); - connectedRpc.setCloseHandler( - () -> sessions.values().forEach(CopilotSession::cancelPendingExternalTools)); + connectedRpc.setCloseHandler(() -> sessions.values().forEach(CopilotSession::cancelPendingExternalTools)); // Register handlers for server-to-client calls RpcHandlerDispatcher dispatcher = new RpcHandlerDispatcher(sessions, lifecycleManager::dispatch, executor, diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java index 601237880f..68e7d4a132 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotSession.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotSession.java @@ -1009,8 +1009,7 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin ToolDefinition tool) { var pending = new PendingExternalTool(); synchronized (this) { - if (isTerminated || externalToolsClosed - || pendingExternalTools.putIfAbsent(requestId, pending) != null) { + if (isTerminated || externalToolsClosed || pendingExternalTools.putIfAbsent(requestId, pending) != null) { return; } } diff --git a/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java b/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java index 5894693781..bd633232cb 100644 --- a/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java +++ b/java/sdk/src/test/java/com/github/copilot/CopilotClientTest.java @@ -156,7 +156,8 @@ void testForceStopCancelsPendingExternalTools() throws Exception { var lateRequest = new ExternalToolRequestedEvent(); lateRequest.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-after-force-stop", - session.getSessionId(), "tool-call-after-force-stop", "blocked_tool", null, Map.of(), null, null, null)); + session.getSessionId(), "tool-call-after-force-stop", "blocked_tool", null, Map.of(), null, null, + null)); session.dispatchEvent(lateRequest); assertFalse(lateStarted.await(100, TimeUnit.MILLISECONDS)); } diff --git a/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java b/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java index cb217b023f..b22a05963f 100644 --- a/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java +++ b/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java @@ -94,8 +94,8 @@ void testExternalToolCompletedCancelsBlockedHandler() throws Exception { }))); var requested = new ExternalToolRequestedEvent(); - requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-1", - "test-session-id", "tool-call-1", "blocked_tool", null, Map.of(), null, null, null)); + requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-1", "test-session-id", + "tool-call-1", "blocked_tool", null, Map.of(), null, null, null)); dispatchEvent(requested); assertTrue(started.await(1, TimeUnit.SECONDS)); @@ -130,8 +130,8 @@ void testCloseCancelsBlockedExternalTool() throws Exception { void testExternalToolCompletedDoesNotBlockOnSynchronousHandler() throws Exception { var started = new CountDownLatch(1); var release = new CountDownLatch(1); - session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks synchronously", Map.of(), - invocation -> { + session.registerTools( + List.of(ToolDefinition.create("blocked_tool", "Blocks synchronously", Map.of(), invocation -> { started.countDown(); try { release.await(); diff --git a/python/test_session.py b/python/test_session.py index 638f739bb0..beef9c587d 100644 --- a/python/test_session.py +++ b/python/test_session.py @@ -86,7 +86,7 @@ async def blocked_tool(_invocation): await asyncio.Future() except asyncio.CancelledError: cancelled.set() - return ToolResult(text_result_for_llm="late result") + return ToolResult(text_result_for_llm="late result") session._register_tools([Tool("blocked_tool", "Blocks", blocked_tool)]) session._dispatch_event( diff --git a/rust/src/session.rs b/rust/src/session.rs index 18f058b45e..b87ebe8411 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1145,11 +1145,11 @@ impl Client { "Client::create_session local setup complete" ); *capabilities.write() = create_result.capabilities.unwrap_or_default(); - if has_mcp_auth_handler { - if let Err(error) = register_mcp_auth_interest(self, &session_id).await { - registration.cleanup(event_loop).await; - return Err(error); - } + if has_mcp_auth_handler + && let Err(error) = register_mcp_auth_interest(self, &session_id).await + { + registration.cleanup(event_loop).await; + return Err(error); } tracing::debug!( @@ -1395,11 +1395,11 @@ impl Client { }) .into()); } - if has_mcp_auth_handler { - if let Err(error) = register_mcp_auth_interest(self, &session_id).await { - registration.cleanup(event_loop).await; - return Err(error); - } + if has_mcp_auth_handler + && let Err(error) = register_mcp_auth_interest(self, &session_id).await + { + registration.cleanup(event_loop).await; + return Err(error); } // Reload skills after resume (best-effort). From 9045ec3fb9da4fcb56f761393cfdb7ed9f448d04 Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Thu, 3 Sep 2026 18:20:48 -0400 Subject: [PATCH 3/8] 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> --- dotnet/src/Client.cs | 10 ++++------ dotnet/src/Session.cs | 4 +--- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index 09c76816f2..84794384ea 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -2716,14 +2716,12 @@ private async Task ConnectToServerAsync(Process? cliProcess, string? private async Task CancelExternalToolsWhenConnectionClosesAsync(JsonRpc rpc) { - try + await Task.WhenAny(rpc.Completion).ConfigureAwait(false); + if (rpc.Completion.Exception is { } exception) { - await rpc.Completion.ConfigureAwait(false); - } - catch (Exception ex) - { - _logger.LogDebug(ex, "JSON-RPC connection completed with an error"); + _logger.LogDebug(exception, "JSON-RPC connection completed with an error"); } + var connectionTask = _connectionTask; if (connectionTask is null || connectionTask.Status != System.Threading.Tasks.TaskStatus.RanToCompletion diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index 4a96840a81..b62340609c 100644 --- a/dotnet/src/Session.cs +++ b/dotnet/src/Session.cs @@ -896,10 +896,9 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, return; } - var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(_externalToolLifetime.Token); + using var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(_externalToolLifetime.Token); if (!_pendingExternalTools.TryAdd(requestId, cancellationSource)) { - cancellationSource.Dispose(); return; } @@ -1033,7 +1032,6 @@ await Rpc.Tools.HandlePendingToolCallAsync( { ((ICollection>)_pendingExternalTools) .Remove(new(requestId, cancellationSource)); - cancellationSource.Dispose(); } static string GetSingleParameterName(AIFunction tool) From b6ac7f4f03137cb0ee7a7cfb47ed72b48c809409 Mon Sep 17 00:00:00 2001 From: Steve Sanderson Date: Fri, 4 Sep 2026 08:48:42 +0000 Subject: [PATCH 4/8] 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> --- .../E2E/ExternalToolCancellationE2ETests.cs | 62 +++++++++ .../external_tool_cancellation_e2e_test.go | 77 +++++++++++ .../ExternalToolCancellationE2ETest.java | 95 ++++++++++++++ .../external-tool-cancellation.e2e.test.ts | 87 +++++++++++++ .../test_external_tool_cancellation_e2e.py | 64 +++++++++ rust/tests/e2e.rs | 2 + rust/tests/e2e/external_tool_cancellation.rs | 123 ++++++++++++++++++ ...tool_handler_when_session_disconnects.yaml | 15 +++ ...el_tool_handler_when_session_disposes.yaml | 18 +++ 9 files changed, 543 insertions(+) create mode 100644 dotnet/test/E2E/ExternalToolCancellationE2ETests.cs create mode 100644 go/internal/e2e/external_tool_cancellation_e2e_test.go create mode 100644 java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java create mode 100644 nodejs/test/e2e/external-tool-cancellation.e2e.test.ts create mode 100644 python/e2e/test_external_tool_cancellation_e2e.py create mode 100644 rust/tests/e2e/external_tool_cancellation.rs create mode 100644 test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disconnects.yaml create mode 100644 test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disposes.yaml diff --git a/dotnet/test/E2E/ExternalToolCancellationE2ETests.cs b/dotnet/test/E2E/ExternalToolCancellationE2ETests.cs new file mode 100644 index 0000000000..b7f34fd037 --- /dev/null +++ b/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(TaskCreationOptions.RunContinuationsAsynchronously); + var toolCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseTool = new TaskCompletionSource(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 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; + } + } + } +} diff --git a/go/internal/e2e/external_tool_cancellation_e2e_test.go b/go/internal/e2e/external_tool_cancellation_e2e_test.go new file mode 100644 index 0000000000..1c865ca278 --- /dev/null +++ b/go/internal/e2e/external_tool_cancellation_e2e_test.go @@ -0,0 +1,77 @@ +package e2e + +import ( + "testing" + "time" + + copilot "github.com/github/copilot-sdk/go" + "github.com/github/copilot-sdk/go/internal/e2e/testharness" +) + +func TestExternalToolCancellationE2E(t *testing.T) { + ctx := testharness.NewTestContext(t) + client := ctx.NewClient() + t.Cleanup(func() { client.ForceStop() }) + + t.Run("should_cancel_tool_handler_when_session_disconnects", func(t *testing.T) { + ctx.ConfigureForTest(t) + + type ValueParams struct { + Value string `json:"value" jsonschema:"Value to analyze"` + } + toolStarted := make(chan struct{}, 1) + toolCancelled := make(chan struct{}, 1) + releaseTool := make(chan string, 1) + + slowTool := copilot.DefineTool("slow_analysis", "A slow analysis tool that blocks until released", + func(_ ValueParams, inv copilot.ToolInvocation) (string, error) { + select { + case toolStarted <- struct{}{}: + default: + } + select { + case value := <-releaseTool: + return value, nil + case <-inv.TraceContext.Done(): + select { + case toolCancelled <- struct{}{}: + default: + } + return "", inv.TraceContext.Err() + } + }) + slowTool.SkipPermission = true + + session, err := client.CreateSession(t.Context(), &copilot.SessionConfig{ + OnPermissionRequest: copilot.PermissionHandler.ApproveAll, + Tools: []copilot.Tool{slowTool}, + }) + if err != nil { + t.Fatalf("Failed to create session: %v", err) + } + t.Cleanup(func() { _ = session.Disconnect() }) + + go func() { + _, _ = session.Send(t.Context(), copilot.MessageOptions{ + Prompt: "Use slow_analysis with value 'test_abort'. Wait for the result.", + }) + }() + + select { + case <-toolStarted: + case <-time.After(60 * time.Second): + t.Fatal("Timed out waiting for tool handler to start") + } + + if err := session.Disconnect(); err != nil { + t.Fatalf("Disconnect failed: %v", err) + } + + select { + case <-toolCancelled: + case <-time.After(60 * time.Second): + t.Fatal("Timed out waiting for tool handler cancellation") + } + + }) +} diff --git a/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java b/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java new file mode 100644 index 0000000000..54973cf0ce --- /dev/null +++ b/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java @@ -0,0 +1,95 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import com.github.copilot.rpc.MessageOptions; +import com.github.copilot.rpc.PermissionHandler; +import com.github.copilot.rpc.SessionConfig; +import com.github.copilot.rpc.ToolDefinition; + +public class ExternalToolCancellationE2ETest { + + private static E2ETestContext ctx; + + @BeforeAll + static void setup() throws Exception { + ctx = E2ETestContext.create(); + } + + @AfterAll + static void teardown() throws Exception { + if (ctx != null) { + ctx.close(); + } + } + + @Test + void shouldCancelToolHandlerWhenSessionDisconnects() throws Exception { + ctx.configureForTest("external_tool_cancellation", "should_cancel_tool_handler_when_session_disconnects"); + + var pendingTool = new AtomicReference>(); + ToolDefinition slowTool = ToolDefinition.create( + "slow_analysis", + "A slow analysis tool that blocks until released", + slowAnalysisSchema(), + invocation -> { + CompletableFuture pending = new CompletableFuture<>(); + pendingTool.set(pending); + return pending; + }); + + try (CopilotClient client = ctx.createClient()) { + CopilotSession session = client + .createSession(new SessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL) + .setTools(List.of(slowTool))) + .get(60, TimeUnit.SECONDS); + try { + session.send( + new MessageOptions().setPrompt("Use slow_analysis with value 'test_abort'. Wait for the result.")) + .get(60, TimeUnit.SECONDS); + + waitFor(() -> pendingTool.get() != null, 60_000); + session.close(); + waitFor(() -> pendingTool.get() != null && pendingTool.get().isCancelled(), 60_000); + } finally { + if (session != null) { + session.close(); + } + } + } + } + + private static Map slowAnalysisSchema() { + Map props = new HashMap<>(); + props.put("value", Map.of("type", "string", "description", "Value to analyze")); + Map schema = new HashMap<>(); + schema.put("type", "object"); + schema.put("properties", props); + schema.put("required", List.of("value")); + return schema; + } + + private static void waitFor(BooleanSupplier predicate, long timeoutMillis) throws InterruptedException { + long deadline = System.currentTimeMillis() + timeoutMillis; + while (!predicate.getAsBoolean()) { + if (System.currentTimeMillis() > deadline) { + throw new AssertionError("waitFor timed out"); + } + Thread.sleep(50); + } + } +} diff --git a/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts b/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts new file mode 100644 index 0000000000..d86775af53 --- /dev/null +++ b/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts @@ -0,0 +1,87 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { approveAll, defineTool } from "../../src/index.js"; +import { createSdkTestContext } from "./harness/sdkTestContext.js"; + +describe("External tool cancellation", async () => { + const { copilotClient: client } = await createSdkTestContext(); + + async function withTimeout(promise: Promise, ms: number, label: string): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error(`Timeout: ${label}`)), ms); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } + } + + it("should cancel tool handler when session disconnects", { timeout: 120_000 }, async () => { + let toolStartedResolve!: () => void; + const toolStarted = new Promise((resolve) => { + toolStartedResolve = resolve; + }); + let toolCancelledResolve!: () => void; + const toolCancelled = new Promise((resolve) => { + toolCancelledResolve = resolve; + }); + let releaseToolResolve!: () => void; + const releaseTool = new Promise((resolve) => { + releaseToolResolve = resolve; + }); + + const session = await client.createSession({ + onPermissionRequest: approveAll, + tools: [ + defineTool("slow_analysis", { + description: "A slow analysis tool that blocks until released", + parameters: z.object({ + value: z.string().describe("Value to analyze"), + }), + handler: async (_args, invocation) => { + toolStartedResolve(); + await Promise.race([ + releaseTool, + new Promise((_, reject) => + setImmediate(() => { + const onAbort = () => { + toolCancelledResolve(); + reject(new Error("aborted")); + }; + if (invocation.signal?.aborted) { + onAbort(); + return; + } + invocation.signal?.addEventListener("abort", onAbort, { + once: true, + }); + }) + ), + ]); + return "RELEASED"; + }, + }), + ], + }); + + try { + void session.send({ + prompt: "Use slow_analysis with value 'test_abort'. Wait for the result.", + }); + + await withTimeout(toolStarted, 60_000, "slow_analysis start"); + await session.disconnect(); + await withTimeout(toolCancelled, 60_000, "slow_analysis cancellation"); + } finally { + releaseToolResolve(); + } + }); +}); diff --git a/python/e2e/test_external_tool_cancellation_e2e.py b/python/e2e/test_external_tool_cancellation_e2e.py new file mode 100644 index 0000000000..d447be1fe6 --- /dev/null +++ b/python/e2e/test_external_tool_cancellation_e2e.py @@ -0,0 +1,64 @@ +""" +E2E tests for external tool cancellation. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from copilot.session import PermissionHandler +from copilot.tools import Tool, ToolInvocation, ToolResult + +from .testharness import E2ETestContext + +pytestmark = pytest.mark.asyncio(loop_scope="module") + + +class TestExternalToolCancellation: + async def test_should_cancel_tool_handler_when_session_disconnects( + self, ctx: E2ETestContext + ): + tool_started = asyncio.Event() + tool_cancelled = asyncio.Event() + release_tool: asyncio.Future = asyncio.get_event_loop().create_future() + + async def slow_tool_handler(invocation: ToolInvocation) -> ToolResult: + _ = (invocation.arguments or {}).get("value", "") + tool_started.set() + try: + result = await asyncio.wait_for(release_tool, timeout=120.0) + return ToolResult(text_result_for_llm=str(result)) + except asyncio.CancelledError: + tool_cancelled.set() + raise + + session = await ctx.client.create_session( + on_permission_request=PermissionHandler.approve_all, + tools=[ + Tool( + name="slow_analysis", + description="A slow analysis tool that blocks until released", + parameters={ + "type": "object", + "properties": { + "value": {"type": "string", "description": "Value to analyze"} + }, + "required": ["value"], + }, + handler=slow_tool_handler, + ) + ], + ) + + try: + asyncio.ensure_future( + session.send("Use slow_analysis with value 'test_abort'. Wait for the result.") + ) + await asyncio.wait_for(tool_started.wait(), timeout=60.0) + await session.disconnect() + await asyncio.wait_for(tool_cancelled.wait(), timeout=60.0) + finally: + if not release_tool.done(): + release_tool.set_result("RELEASED") diff --git a/rust/tests/e2e.rs b/rust/tests/e2e.rs index 03723dfb1b..f597ffa17e 100644 --- a/rust/tests/e2e.rs +++ b/rust/tests/e2e.rs @@ -29,6 +29,8 @@ mod copilot_request_handler; mod elicitation; #[path = "e2e/error_resilience.rs"] mod error_resilience; +#[path = "e2e/external_tool_cancellation.rs"] +mod external_tool_cancellation; #[path = "e2e/event_fidelity.rs"] mod event_fidelity; #[path = "e2e/github_telemetry.rs"] diff --git a/rust/tests/e2e/external_tool_cancellation.rs b/rust/tests/e2e/external_tool_cancellation.rs new file mode 100644 index 0000000000..eb9b1b6650 --- /dev/null +++ b/rust/tests/e2e/external_tool_cancellation.rs @@ -0,0 +1,123 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use github_copilot_sdk::handler::ApproveAllHandler; +use github_copilot_sdk::tool::ToolHandler; +use github_copilot_sdk::{Error, SessionConfig, Tool, ToolInvocation, ToolResult}; +use serde_json::json; +use tokio::sync::{Mutex, mpsc, oneshot}; +use tokio::time::{Duration, timeout}; + +use super::support::DEFAULT_TEST_TOKEN; + +#[tokio::test] +async fn should_cancel_tool_handler_when_session_disconnects() { + super::support::with_dedicated_e2e_context( + "external_tool_cancellation", + "should_cancel_tool_handler_when_session_disconnects", + |ctx| { + Box::pin(async move { + ctx.set_default_copilot_user(); + let client = ctx.start_client().await; + let (started_tx, mut started_rx) = mpsc::unbounded_channel(); + let (release_tx, release_rx) = oneshot::channel(); + let (cancelled_tx, cancelled_rx) = oneshot::channel(); + let tool = Arc::new(CancelAwareSlowTool { + started_tx, + release_rx: Mutex::new(Some(release_rx)), + cancelled_tx: Mutex::new(Some(cancelled_tx)), + }); + + let session = client + .create_session( + SessionConfig::default() + .with_github_token(DEFAULT_TEST_TOKEN) + .with_permission_handler(Arc::new(ApproveAllHandler)) + .with_tools(vec![ + Tool::new("slow_analysis") + .with_description( + "A slow analysis tool that blocks until released", + ) + .with_parameters(json!({ + "type": "object", + "properties": { + "value": { + "type": "string", + "description": "Value to analyze" + } + }, + "required": ["value"] + })) + .with_handler(tool), + ]), + ) + .await + .expect("create session"); + + session + .send("Use slow_analysis with value 'test_abort'. Wait for the result.") + .await + .expect("send tool turn"); + + let started_value = timeout(Duration::from_secs(60), started_rx.recv()) + .await + .expect("tool start wait timed out") + .expect("tool start channel closed"); + assert_eq!(started_value, "test_abort"); + + session.disconnect().await.expect("disconnect session"); + timeout(Duration::from_secs(60), cancelled_rx) + .await + .expect("tool cancellation wait timed out") + .expect("tool cancellation sender dropped"); + + let _ = release_tx.send("RELEASED".to_string()); + client.stop().await.expect("stop client"); + }) + }, + ) + .await; +} + +struct CancelAwareSlowTool { + started_tx: mpsc::UnboundedSender, + release_rx: Mutex>>, + cancelled_tx: Mutex>>, +} + +struct CancelSignalGuard { + cancelled_tx: Option>, +} + +impl Drop for CancelSignalGuard { + fn drop(&mut self) { + if let Some(sender) = self.cancelled_tx.take() { + let _ = sender.send(()); + } + } +} + +#[async_trait] +impl ToolHandler for CancelAwareSlowTool { + async fn call(&self, invocation: ToolInvocation) -> Result { + let value = invocation + .arguments + .get("value") + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(); + let _ = self.started_tx.send(value); + + let cancelled_tx = self.cancelled_tx.lock().await.take(); + let _guard = CancelSignalGuard { cancelled_tx }; + + let release_rx = self + .release_rx + .lock() + .await + .take() + .expect("slow tool called once"); + let released = release_rx.await.unwrap_or_else(|_| "released".to_string()); + Ok(ToolResult::Text(released)) + } +} diff --git a/test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disconnects.yaml b/test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disconnects.yaml new file mode 100644 index 0000000000..aa37004d64 --- /dev/null +++ b/test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disconnects.yaml @@ -0,0 +1,15 @@ +models: + - claude-sonnet-5 +conversations: + - messages: + - role: system + content: ${system} + - role: user + content: Use slow_analysis with value 'test_abort'. Wait for the result. + - role: assistant + tool_calls: + - id: toolcall_0 + type: function + function: + name: slow_analysis + arguments: '{"value":"test_abort"}' diff --git a/test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disposes.yaml b/test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disposes.yaml new file mode 100644 index 0000000000..028b44e73f --- /dev/null +++ b/test/snapshots/external_tool_cancellation/should_cancel_tool_handler_when_session_disposes.yaml @@ -0,0 +1,18 @@ +models: + - claude-sonnet-5 +conversations: + - messages: + - role: system + content: ${system} + - role: user + content: Use slow_analysis with value 'test_abort'. Wait for the result. + - role: assistant + tool_calls: + - id: toolcall_0 + type: function + function: + name: slow_analysis + arguments: '{"value":"test_abort"}' + - role: tool + tool_call_id: toolcall_0 + content: The execution of this tool, or a previous tool was interrupted. From 245cb0c0220ac7d57d4082004ed42ee7abc4e0dc Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 4 Sep 2026 04:57:51 -0400 Subject: [PATCH 5/8] 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> --- nodejs/test/e2e/external-tool-cancellation.e2e.test.ts | 2 +- python/e2e/test_external_tool_cancellation_e2e.py | 4 +--- rust/tests/e2e.rs | 4 ++-- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts b/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts index d86775af53..4e71ac60a2 100644 --- a/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts +++ b/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts @@ -2,7 +2,7 @@ * Copyright (c) Microsoft Corporation. All rights reserved. *--------------------------------------------------------------------------------------------*/ -import { describe, expect, it } from "vitest"; +import { describe, it } from "vitest"; import { z } from "zod"; import { approveAll, defineTool } from "../../src/index.js"; import { createSdkTestContext } from "./harness/sdkTestContext.js"; diff --git a/python/e2e/test_external_tool_cancellation_e2e.py b/python/e2e/test_external_tool_cancellation_e2e.py index d447be1fe6..aca28ceadf 100644 --- a/python/e2e/test_external_tool_cancellation_e2e.py +++ b/python/e2e/test_external_tool_cancellation_e2e.py @@ -17,9 +17,7 @@ class TestExternalToolCancellation: - async def test_should_cancel_tool_handler_when_session_disconnects( - self, ctx: E2ETestContext - ): + async def test_should_cancel_tool_handler_when_session_disconnects(self, ctx: E2ETestContext): tool_started = asyncio.Event() tool_cancelled = asyncio.Event() release_tool: asyncio.Future = asyncio.get_event_loop().create_future() diff --git a/rust/tests/e2e.rs b/rust/tests/e2e.rs index f597ffa17e..1129d9a871 100644 --- a/rust/tests/e2e.rs +++ b/rust/tests/e2e.rs @@ -29,10 +29,10 @@ mod copilot_request_handler; mod elicitation; #[path = "e2e/error_resilience.rs"] mod error_resilience; -#[path = "e2e/external_tool_cancellation.rs"] -mod external_tool_cancellation; #[path = "e2e/event_fidelity.rs"] mod event_fidelity; +#[path = "e2e/external_tool_cancellation.rs"] +mod external_tool_cancellation; #[path = "e2e/github_telemetry.rs"] mod github_telemetry; #[path = "e2e/hooks.rs"] From ca73ddb3509cef54a918fe93be99f79a8a5081c7 Mon Sep 17 00:00:00 2001 From: Steve Sanderson Date: Fri, 4 Sep 2026 09:16:14 +0000 Subject: [PATCH 6/8] test: fix java e2e formatting for spotless Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../copilot/ExternalToolCancellationE2ETest.java | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java b/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java index 54973cf0ce..d5ae322e32 100644 --- a/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java +++ b/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java @@ -42,24 +42,20 @@ void shouldCancelToolHandlerWhenSessionDisconnects() throws Exception { ctx.configureForTest("external_tool_cancellation", "should_cancel_tool_handler_when_session_disconnects"); var pendingTool = new AtomicReference>(); - ToolDefinition slowTool = ToolDefinition.create( - "slow_analysis", - "A slow analysis tool that blocks until released", - slowAnalysisSchema(), - invocation -> { + ToolDefinition slowTool = ToolDefinition.create("slow_analysis", + "A slow analysis tool that blocks until released", slowAnalysisSchema(), invocation -> { CompletableFuture pending = new CompletableFuture<>(); pendingTool.set(pending); return pending; }); try (CopilotClient client = ctx.createClient()) { - CopilotSession session = client - .createSession(new SessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL) - .setTools(List.of(slowTool))) + CopilotSession session = client.createSession(new SessionConfig() + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL).setTools(List.of(slowTool))) .get(60, TimeUnit.SECONDS); try { - session.send( - new MessageOptions().setPrompt("Use slow_analysis with value 'test_abort'. Wait for the result.")) + session.send(new MessageOptions() + .setPrompt("Use slow_analysis with value 'test_abort'. Wait for the result.")) .get(60, TimeUnit.SECONDS); waitFor(() -> pendingTool.get() != null, 60_000); From 98db8c20b3c4ecfa82f982ca4567db6bda9d48f7 Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 4 Sep 2026 05:55:13 -0400 Subject: [PATCH 7/8] 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> --- dotnet/src/Client.cs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index 84794384ea..2e98ce507c 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -2703,17 +2703,29 @@ private async Task 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); From ee560e4b5444700d3f1df0611221c7984ac80ecb Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 4 Sep 2026 08:17:16 -0400 Subject: [PATCH 8/8] Fix Java close cancellation test after detach migration Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../java/com/github/copilot/SessionEventHandlingTest.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java b/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java index 8575ec2f7a..83913e82b1 100644 --- a/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java +++ b/java/sdk/src/test/java/com/github/copilot/SessionEventHandlingTest.java @@ -108,6 +108,10 @@ void testExternalToolCompletedCancelsBlockedHandler() throws Exception { @Test void testCloseCancelsBlockedExternalTool() throws Exception { + var rpc = mock(JsonRpcClient.class); + when(rpc.invoke(eq("session.detach"), any(), eq(CopilotSession.SessionDetachResponse.class))) + .thenReturn(CompletableFuture.completedFuture(new CopilotSession.SessionDetachResponse(true, null))); + session = createTestSession(rpc); var toolFuture = new CompletableFuture(); var started = new CountDownLatch(1); session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks", Map.of(), invocation -> {