diff --git a/CHANGELOG.md b/CHANGELOG.md index 59961b4f0..788b98362 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 ea15df11c..cbc7d5fa2 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(); @@ -2693,6 +2697,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); @@ -2705,17 +2710,50 @@ 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); + if (rpc.Completion.Exception is { } exception) + { + _logger.LogDebug(exception, "JSON-RPC connection completed with an error"); + } + + var connectionTask = _connectionTask; + if (connectionTask is null + || connectionTask.Status != System.Threading.Tasks.TaskStatus.RanToCompletion + || !ReferenceEquals(connectionTask.Result.Rpc, rpc)) + { + return; + } + foreach (var session in _sessions.Values) + { + session.CancelPendingExternalTools(); + } + } + private static JsonSerializerOptions SerializerOptionsForMessageFormatter { get; } = CreateSerializerOptions(); /// diff --git a/dotnet/src/Session.cs b/dotnet/src/Session.cs index b01432a5f..8637da948 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; @@ -60,6 +61,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; @@ -226,6 +229,7 @@ internal void CloseEventChannel() /// internal void Unregister() { + CancelPendingExternalTools(); CloseEventChannel(); RemoveFromClient(); } @@ -675,6 +679,10 @@ private async Task HandleBroadcastEventAsync(SessionEvent sessionEvent) break; } + case ExternalToolCompletedEvent completedEvent: + CancelExternalTool(completedEvent.Data.RequestId); + break; + case PermissionRequestedEvent permEvent: { var data = permEvent.Data; @@ -884,8 +892,24 @@ private async Task TryCancelMcpAuthRequestAsync(string requestId) /// private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, string toolCallId, JsonElement? arguments, AIFunction tool) { + if (_externalToolLifetime.IsCancellationRequested) + { + return; + } + + using var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(_externalToolLifetime.Token); + if (!_pendingExternalTools.TryAdd(requestId, cancellationSource)) + { + return; + } + try { + if (cancellationSource.IsCancellationRequested) + { + return; + } + var invocation = new ToolInvocation { SessionId = SessionId, @@ -903,7 +927,7 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, { try { - var metadata = await Rpc.Tools.GetCurrentMetadataAsync(); + var metadata = await Rpc.Tools.GetCurrentMetadataAsync(cancellationSource.Token); invocation.AvailableTools = metadata.Tools; } catch (Exception ex) when (ex is RemoteRpcException or IOException or ObjectDisposedException or JsonException) @@ -938,7 +962,7 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, } var toolTimestamp = Stopwatch.GetTimestamp(); - var result = await tool.InvokeAsync(aiFunctionArgs); + var result = await tool.InvokeAsync(aiFunctionArgs, cancellationSource.Token); LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null, "CopilotSession.ExecuteToolAndRespondAsync tool dispatch. Elapsed={Elapsed}, SessionId={SessionId}, RequestId={RequestId}, ToolCallId={ToolCallId}, Tool={ToolName}", toolTimestamp, @@ -948,9 +972,14 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, toolName); var toolResultObject = ToolResultObject.ConvertFromInvocationResult(result, tool.JsonSerializerOptions); + if (!TryClaimExternalTool(requestId, cancellationSource)) + { + return; + } var responseRpcTimestamp = Stopwatch.GetTimestamp(); - await Rpc.Tools.HandlePendingToolCallAsync(requestId, toolResultObject, error: null); + await Rpc.Tools.HandlePendingToolCallAsync( + requestId, toolResultObject, error: null, cancellationSource.Token); LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null, "CopilotSession.ExecuteToolAndRespondAsync response sent successfully. Elapsed={Elapsed}, SessionId={SessionId}, RequestId={RequestId}, ToolCallId={ToolCallId}, Tool={ToolName}", responseRpcTimestamp, @@ -959,11 +988,33 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, toolCallId, toolName); } - catch (Exception ex) + catch (OperationCanceledException) when (cancellationSource.IsCancellationRequested) + { + // The runtime has already completed the request or the session is shutting down. + } + catch (RemoteRpcException) when (cancellationSource.IsCancellationRequested) + { + // Another client answered after this invocation completed locally. + } + catch (Exception) when (cancellationSource.IsCancellationRequested) + { + // Cancellation won the request; no response or error should escape. + } + catch (Exception ex) when (!cancellationSource.IsCancellationRequested) { + if (!TryClaimExternalTool(requestId, cancellationSource)) + { + return; + } + try { - await Rpc.Tools.HandlePendingToolCallAsync(requestId, result: null, error: ex.Message); + await Rpc.Tools.HandlePendingToolCallAsync( + requestId, result: null, error: ex.Message, cancellationSource.Token); + } + catch (OperationCanceledException) + { + // Teardown canceled the in-flight error response. } catch (IOException) { @@ -973,6 +1024,15 @@ private async Task ExecuteToolAndRespondAsync(string requestId, string toolName, { // Connection already disposed — nothing we can do } + catch (RemoteRpcException) + { + // Another client may have answered the broadcast request first. + } + } + finally + { + ((ICollection>)_pendingExternalTools) + .Remove(new(requestId, cancellationSource)); } static string GetSingleParameterName(AIFunction tool) @@ -1025,6 +1085,49 @@ static string GetSingleParameterName(AIFunction tool) } } + private bool TryClaimExternalTool(string requestId, CancellationTokenSource cancellationSource) + => ((ICollection>)_pendingExternalTools) + .Remove(new(requestId, cancellationSource)); + + private void CancelExternalTool(string requestId) + { + if (!string.IsNullOrEmpty(requestId) && _pendingExternalTools.TryRemove(requestId, out var cancellationSource)) + { + _ = Task.Run(() => + { + try + { + cancellationSource.Cancel(); + } + catch (AggregateException) + { + // Cancellation callbacks are consumer code and must not disrupt event dispatch. + } + catch (ObjectDisposedException) + { + // The invocation completed while cancellation was being delivered. + } + }); + } + } + + internal void CancelPendingExternalTools() + { + try + { + _externalToolLifetime.Cancel(); + } + catch (AggregateException) + { + // User cancellation callbacks must not prevent session teardown. + } + catch (ObjectDisposedException) + { + // Session teardown already completed. + } + _pendingExternalTools.Clear(); + } + /// /// Executes a permission handler and sends the result back via the HandlePendingPermissionRequest RPC. /// @@ -2105,6 +2208,7 @@ public async ValueTask DisposeAsync() return; } + CancelPendingExternalTools(); CloseEventChannel(); try diff --git a/dotnet/test/E2E/ExternalToolCancellationE2ETests.cs b/dotnet/test/E2E/ExternalToolCancellationE2ETests.cs new file mode 100644 index 000000000..b7f34fd03 --- /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/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index 7029d1ea0..04ef6e640 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -1216,6 +1216,252 @@ 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 ExternalToolCompleted_Does_Not_Block_Event_Dispatch_On_Cancellation_Callback() + { + 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 callbackStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseCallback = 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-blocking-callback")); + await toolStarted.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + var dispatchTask = Task.Run(() => DispatchEvent(session, new ExternalToolCompletedEvent + { + Data = new ExternalToolCompletedData { RequestId = "request-blocking-callback" } + })); + await callbackStarted.Task.WaitAsync(TimeSpan.FromSeconds(5)); + try + { + await dispatchTask.WaitAsync(TimeSpan.FromSeconds(5)); + } + finally + { + releaseCallback.TrySetResult(); + } + await toolCancelled.Task.WaitAsync(TimeSpan.FromSeconds(5)); + + async Task BlockedTool(CancellationToken cancellationToken) + { + toolStarted.TrySetResult(); + using var registration = cancellationToken.Register(() => + { + callbackStarted.TrySetResult(); + releaseCallback.Task.GetAwaiter().GetResult(); + }); + 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 ForceStopAsync_Does_Not_Start_Late_External_Tool() + { + 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 session = await client.CreateSessionAsync(new SessionConfig + { + Tools = [AIFunctionFactory.Create(Tool, "late_tool")], + OnPermissionRequest = PermissionHandler.ApproveAll + }); + + await client.ForceStopAsync(); + DispatchEvent(session, ExternalToolRequested("request-after-force-stop", "late_tool")); + + Assert.False(toolStarted.Task.IsCompleted); + + string Tool() + { + toolStarted.TrySetResult(); + return "unexpected"; + } + } + + [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, string toolName = "blocked_tool") => + new() + { + Data = new ExternalToolRequestedData + { + RequestId = requestId, + SessionId = "session-1", + ToolCallId = "tool-call-1", + ToolName = toolName + } + }; + [Fact] public async Task Generated_Session_Rpc_Throws_When_Session_Disposed() { @@ -1780,6 +2026,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 45de4f6e9..6cac724f0 100644 --- a/go/client.go +++ b/go/client.go @@ -701,8 +701,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() @@ -1012,6 +1019,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses delete(c.sessions, sessionID) } c.sessionsMux.Unlock() + s.cancelPendingExternalTools() s.stopEventProcessing() } @@ -1157,6 +1165,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses "sessionId": session.SessionID, "eventType": "mcp.oauth_required", }); err != nil { + unregisterSession(registeredSessionID, session) return nil, err } } @@ -1171,6 +1180,7 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses ManageScheduleEnabled: config.ManageScheduleEnabled, IncludedBuiltinSkills: config.IncludedBuiltinSkills, }); err != nil { + unregisterSession(registeredSessionID, session) return nil, err } @@ -1431,6 +1441,7 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string, } } c.sessionsMux.Unlock() + session.cancelPendingExternalTools() // newSession starts processEvents eagerly, before the RPC confirms the // resume; every failure path here restores the previously-registered // session (if any) but never returns this failed one to the caller, so @@ -2520,6 +2531,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 f8415be6b..52587c846 100644 --- a/go/client_test.go +++ b/go/client_test.go @@ -541,6 +541,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.detach", func(params json.RawMessage) (json.RawMessage, *jsonrpc2.Error) { + return []byte(`{"success":true}`), 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/internal/e2e/external_tool_cancellation_e2e_test.go b/go/internal/e2e/external_tool_cancellation_e2e_test.go new file mode 100644 index 000000000..1c865ca27 --- /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/go/session.go b/go/session.go index cb64bbc60..f4bb0d38c 100644 --- a/go/session.go +++ b/go/session.go @@ -66,6 +66,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 @@ -108,6 +111,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. @@ -1393,7 +1401,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) + } select { case s.eventCh <- event: @@ -1450,20 +1470,6 @@ func (s *Session) stopEventProcessing() { // 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. @@ -1546,11 +1552,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, @@ -1577,8 +1662,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{ @@ -1741,6 +1832,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() result, err := s.client.Request(context.Background(), "session.detach", sessionDetachRequest{SessionID: s.SessionID}) if err == nil { var response sessionDetachResponse diff --git a/go/session_test.go b/go/session_test.go index 1547f8237..bdf14887a 100644 --- a/go/session_test.go +++ b/go/session_test.go @@ -34,6 +34,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 TestDispatchEventReturnsAfterEventProcessingStops(t *testing.T) { session := &Session{ eventCh: make(chan SessionEvent), diff --git a/go/types.go b/go/types.go index dfedc95e7..4a45a4019 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 fbd472873..9f7d8ebcf 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java @@ -553,6 +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)); // Register handlers for server-to-client calls RpcHandlerDispatcher dispatcher = new RpcHandlerDispatcher(sessions, lifecycleManager::dispatch, executor, @@ -746,7 +747,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 @@ -1018,6 +1021,7 @@ public CompletableFuture createSession(SessionConfig config) { CopilotSession session = preRegisteredSessionHolder[0] != null ? preRegisteredSessionHolder[0] : initializeSession.apply(returnedId); + preRegisteredSessionHolder[0] = session; if (tokenRegistration != null) { session.setGitHubTokenProviderRegistration(tokenRegistration); } @@ -1049,6 +1053,9 @@ public CompletableFuture createSession(SessionConfig config) { return session; }); }).exceptionally(ex -> { + if (preRegisteredSessionHolder[0] != null) { + preRegisteredSessionHolder[0].cancelPendingExternalTools(); + } if (registeredIdHolder[0] != null) { sessions.remove(registeredIdHolder[0]); } @@ -1229,6 +1236,7 @@ public CompletableFuture resumeSession(String sessionId, ResumeS return session; }); }).exceptionally(ex -> { + session.cancelPendingExternalTools(); sessions.remove(sessionId); // Also remove the re-keyed entry if the server returned a different ID String activeId = session.getSessionId(); 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 55a83137b..c072e31ec 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; @@ -31,6 +32,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; 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; @@ -187,6 +189,8 @@ 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 boolean externalToolsClosed; private final Map commandHandlers = new ConcurrentHashMap<>(); private final Map bearerTokenProviders = new ConcurrentHashMap<>(); private final AtomicReference permissionHandler = new AtomicReference<>(); @@ -207,6 +211,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. *

@@ -860,6 +904,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) { @@ -931,9 +983,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 @@ -958,6 +1010,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 || externalToolsClosed || pendingExternalTools.putIfAbsent(requestId, pending) != null) { + return; + } + } Runnable task = () -> { try { JsonNode argumentsNode = arguments instanceof JsonNode jn @@ -966,9 +1024,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) { @@ -983,6 +1064,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())); @@ -990,8 +1074,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, @@ -1013,6 +1100,16 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin } } + void cancelPendingExternalTools() { + List pending; + synchronized (this) { + externalToolsClosed = true; + pending = new ArrayList<>(pendingExternalTools.values()); + pendingExternalTools.clear(); + } + pending.forEach(PendingExternalTool::cancel); + } + /** * Builds a {@link SessionUiHandlePendingElicitationParams} carrying a * {@code cancel} action, used when an elicitation handler throws or the handler @@ -2429,6 +2526,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 7eda069d2..f78ce0042 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 ba161ef99..6b8e861e8 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,14 +16,19 @@ 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; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.*; @@ -113,6 +119,49 @@ 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); + var lateStarted = new CountDownLatch(1); + var invocations = new AtomicInteger(); + session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks", Map.of(), invocation -> { + if (invocations.incrementAndGet() == 1) { + started.countDown(); + } else { + lateStarted.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()); + + 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.dispatchEvent(lateRequest); + assertFalse(lateStarted.await(100, TimeUnit.MILLISECONDS)); + } + @Test @SuppressWarnings("unchecked") void testDeleteSessionReleasesGitHubTokenProvider() throws Exception { 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 000000000..d5ae322e3 --- /dev/null +++ b/java/sdk/src/test/java/com/github/copilot/ExternalToolCancellationE2ETest.java @@ -0,0 +1,91 @@ +/*--------------------------------------------------------------------------------------------- + * 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/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java b/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java index 009f15c20..13668cd97 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 5e502459c..83913e82b 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,103 @@ 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 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 -> { + 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 306e1ee05..eb92cf0be 100644 --- a/nodejs/src/client.ts +++ b/nodejs/src/client.ts @@ -443,6 +443,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"; @@ -935,6 +936,7 @@ export class CopilotClient { } this.forceStopping = false; + this.connectionClosed = false; this.processTransportError = null; this.state = "connecting"; @@ -1070,7 +1072,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); @@ -1750,6 +1757,7 @@ export class CopilotClient { await this.updateSessionOptionsForMode(session, config); this.commitGitHubTokenProvider(returnedSessionId, gitHubTokenProviderRegistrationId); } catch (e) { + session?._markDisconnected(); if (registeredId !== undefined) { this.sessions.delete(registeredId); } @@ -2028,6 +2036,7 @@ export class CopilotClient { await this.updateSessionOptionsForMode(session, config); this.commitGitHubTokenProvider(sessionId, gitHubTokenProviderRegistrationId); } catch (e) { + session._markDisconnected(); this.sessions.delete(sessionId); if (gitHubTokenProviderRegistrationId !== undefined) { this.githubTokenProviders.delete(gitHubTokenProviderRegistrationId); @@ -3052,13 +3061,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 d68517f82..b7fc7837a 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -422,6 +422,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(); @@ -441,6 +442,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. */ @@ -820,6 +822,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(); @@ -996,6 +1002,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; @@ -1105,6 +1120,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 @@ -1113,13 +1134,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, @@ -1128,6 +1163,7 @@ export class CopilotSession { availableTools, traceparent, tracestate, + signal: controller.signal, }); let result: ToolResult; if (rawResult == null) { @@ -1139,12 +1175,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); @@ -1156,9 +1192,22 @@ export class CopilotSession { } // Connection lost or RPC error — nothing we can do } + } finally { + if (this.pendingExternalTools.get(requestId) === controller) { + this.pendingExternalTools.delete(requestId); + } + controller.abort(); } } + 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 @@ -2016,22 +2065,27 @@ export class CopilotSession { * ``` */ async disconnect(): Promise { - if (this.disconnected) { + if (this.disconnected || this.disconnecting) { return; } - let response: { success: boolean; error?: string } = { success: false }; - for (let attempt = 0; attempt < 2 && !response.success; attempt++) { - response = (await this.connection.sendRequest("session.detach", { - sessionId: this.sessionId, - })) as { success: boolean; error?: string }; - } - if (!response.success) { - throw new Error( - `Failed to disconnect session ${this.sessionId}: ${response.error || "Unknown error"}` - ); + this.disconnecting = true; + try { + let response: { success: boolean; error?: string } = { success: false }; + for (let attempt = 0; attempt < 2 && !response.success; attempt++) { + response = (await this.connection.sendRequest("session.detach", { + sessionId: this.sessionId, + })) as { success: boolean; error?: string }; + } + if (!response.success) { + throw new Error( + `Failed to disconnect session ${this.sessionId}: ${response.error || "Unknown error"}` + ); + } + this._markDisconnected(); + } catch (error) { + this.disconnecting = false; + throw error; } - this._markDisconnected(); - this.onDisconnected?.(); } /** Enables `await using session = ...` syntax for automatic cleanup. */ diff --git a/nodejs/src/types.ts b/nodejs/src/types.ts index 82c6fe0e5..0f15749f7 100644 --- a/nodejs/src/types.ts +++ b/nodejs/src/types.ts @@ -692,6 +692,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 ba96292ba..3db96ea47 100644 --- a/nodejs/test/client.test.ts +++ b/nodejs/test/client.test.ts @@ -3343,6 +3343,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"); @@ -3354,6 +3390,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/e2e/external-tool-cancellation.e2e.test.ts b/nodejs/test/e2e/external-tool-cancellation.e2e.test.ts new file mode 100644 index 000000000..4e71ac60a --- /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, 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/nodejs/test/external-tool-cancellation.test.ts b/nodejs/test/external-tool-cancellation.test.ts new file mode 100644 index 000000000..35d15bd8e --- /dev/null +++ b/nodejs/test/external-tool-cancellation.test.ts @@ -0,0 +1,228 @@ +/* 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("aborts the invocation signal after a normal tool result", async () => { + const sendRequest = vi.fn().mockResolvedValue(undefined); + const session = new CopilotSession("session-1", { sendRequest } as never); + let invocation: ToolInvocation | undefined; + (session as any).toolHandlers.set( + "completed_tool", + async (_args: unknown, context: ToolInvocation) => { + invocation = context; + return "done"; + } + ); + + (session as any)._handleBroadcastEvent({ + type: "external_tool.requested", + data: { + requestId: "request-completed", + sessionId: "session-1", + toolCallId: "tool-call-completed", + toolName: "completed_tool", + arguments: {}, + }, + }); + + await vi.waitFor(() => expect(sendRequest).toHaveBeenCalledTimes(1)); + expect(invocation?.signal?.aborted).toBe(true); +}); + +it("remains retryable when disconnect fails", async () => { + const sendRequest = vi + .fn() + .mockRejectedValueOnce(new Error("transient")) + .mockResolvedValueOnce({ success: true }); + 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 rejectDetach!: (error: Error) => void; + const sendRequest = vi.fn( + () => + new Promise((_, reject) => { + rejectDetach = 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)); + + rejectDetach(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 2df2db80a..dd531a2be 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -2192,7 +2192,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() @@ -4860,8 +4863,22 @@ def _register_github_token_provider( def _handle_connection_close(self) -> None: self._state = "disconnected" + with self._sessions_lock: + sessions = list(self._sessions.values()) 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 cancel_pending_external_tools() -> None: + for session in sessions: + session._cancel_pending_external_tools() + + try: + loop.call_soon_threadsafe(cancel_pending_external_tools) + 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 20c470302..b5823923e 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -74,6 +74,7 @@ CapabilitiesChangedData, CommandExecuteData, ElicitationRequestedData, + ExternalToolCompletedData, ExternalToolRequestedData, McpOauthRequiredData, PermissionRequest, @@ -1607,6 +1608,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 @@ -1648,6 +1650,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.""" @@ -1966,7 +1981,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) @@ -1977,11 +1992,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): @@ -2189,6 +2219,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( @@ -2207,6 +2239,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( @@ -2225,6 +2259,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( @@ -2235,6 +2271,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, @@ -3020,6 +3067,7 @@ async def disconnect(self) -> None: detail = response.get("error") or "unknown error" raise RuntimeError(f"Failed to detach session {self.session_id}: {detail}") + self._cancel_pending_external_tools() self._run_disconnect_callback() with self._event_handlers_lock: self._destroyed = True 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 000000000..aca28cead --- /dev/null +++ b/python/e2e/test_external_tool_cancellation_e2e.py @@ -0,0 +1,62 @@ +""" +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/python/test_client.py b/python/test_client.py index 47a1aa1f7..2e3868ef1 100644 --- a/python/test_client.py +++ b/python/test_client.py @@ -41,7 +41,7 @@ ModelSupports, ) from copilot.generated.rpc import AutoTier as AutoTierEnum -from copilot.session import PermissionHandler +from copilot.session import CopilotSession, PermissionHandler from copilot.session_events import ( McpOauthRequestReason, McpOauthRequiredData, @@ -198,6 +198,67 @@ 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 not session._destroyed + assert client._sessions == {"session-1": session} + + 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 == {"session-1": session} + assert client._github_token_providers == {} + class TestPermissionHandlerOptional: @pytest.mark.asyncio diff --git a/python/test_session.py b/python/test_session.py index dd2d0a72f..d58ce9409 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_detach_request(): + client = Mock() + client.request = AsyncMock(return_value={"success": True}) + 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.detach", {"sessionId": "session-1"}) diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 25a405080..48e6090ae 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 8e0ac31ea..0e64d6061 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -94,6 +94,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>, @@ -134,6 +169,7 @@ struct PendingSessionRegistration { client: Client, session_id: PendingSessionId, shutdown: CancellationToken, + external_tools_shutdown: CancellationToken, disarmed: bool, } @@ -159,11 +195,13 @@ impl PendingSessionRegistration { session_id: SessionId, token: crate::router::RegistrationToken, shutdown: CancellationToken, + external_tools_shutdown: CancellationToken, ) -> Self { Self { client, session_id: PendingSessionId::Known(session_id, token), shutdown, + external_tools_shutdown, disarmed: false, } } @@ -173,11 +211,13 @@ impl PendingSessionRegistration { client: Client, stash: Arc>>, shutdown: CancellationToken, + external_tools_shutdown: CancellationToken, ) -> Self { Self { client, session_id: PendingSessionId::Deferred(stash), shutdown, + external_tools_shutdown, disarmed: false, } } @@ -197,6 +237,7 @@ impl PendingSessionRegistration { } async fn cleanup(mut self, event_loop: JoinHandle<()>) { + self.external_tools_shutdown.cancel(); self.shutdown.cancel(); let _ = event_loop.await; if let Some(id) = self.registered_id() { @@ -219,6 +260,7 @@ impl PendingSessionRegistration { impl Drop for PendingSessionRegistration { fn drop(&mut self) { if !self.disarmed { + self.external_tools_shutdown.cancel(); self.shutdown.cancel(); if let Some(id) = self.registered_id() { if let PendingSessionId::Known(_, token) = self.session_id { @@ -270,6 +312,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 @@ -739,6 +784,7 @@ impl Session { /// [`send_and_wait`]: Self::send_and_wait pub async fn disconnect(&self) -> Result<(), Error> { self.client.detach_session(&self.id).await?; + self.external_tools_shutdown.cancel(); self.stop_event_loop().await; self.github_token_registration.lock().take(); self.client @@ -817,6 +863,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.github_token_registration.lock().take(); self.client .unregister_session_owned(&self.id, self.registration_token); @@ -1260,6 +1307,7 @@ impl Client { let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default())); let idle_waiter = Arc::new(ParkingLotMutex::new(None)); let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new())); + let external_tools_shutdown = self.inner.rpc.connection_closed_token(); // For cloud sessions (use_server_generated_id), defer session // registration to the inline callback so the read task registers @@ -1322,12 +1370,19 @@ impl Client { .expect("session registration must exist") .1 .token; - PendingSessionRegistration::new(self.clone(), sid.clone(), token, shutdown.clone()) + PendingSessionRegistration::new( + self.clone(), + sid.clone(), + token, + shutdown.clone(), + external_tools_shutdown.clone(), + ) } None => PendingSessionRegistration::deferred( self.clone(), inline_stash.clone(), shutdown.clone(), + external_tools_shutdown.clone(), ), }; @@ -1374,6 +1429,7 @@ impl Client { open_canvases.clone(), event_tx.clone(), shutdown.clone(), + external_tools_shutdown.clone(), ); tracing::debug!( elapsed_ms = setup_start.elapsed().as_millis(), @@ -1405,6 +1461,7 @@ impl Client { client: self.clone(), event_loop: ParkingLotMutex::new(Some(event_loop)), shutdown, + external_tools_shutdown, idle_waiter, capabilities, open_canvases, @@ -1572,6 +1629,7 @@ impl Client { let channels = registration.channels; let idle_waiter = Arc::new(ParkingLotMutex::new(None)); let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new())); + let external_tools_shutdown = self.inner.rpc.connection_closed_token(); let event_loop = spawn_event_loop( session_id.clone(), self.clone(), @@ -1588,12 +1646,14 @@ impl Client { open_canvases.clone(), event_tx.clone(), shutdown.clone(), + external_tools_shutdown.clone(), ); let mut registration = PendingSessionRegistration::new( self.clone(), session_id.clone(), registration_token, shutdown.clone(), + external_tools_shutdown.clone(), ); tracing::debug!( elapsed_ms = setup_start.elapsed().as_millis(), @@ -1692,6 +1752,7 @@ impl Client { client: self.clone(), event_loop: ParkingLotMutex::new(Some(event_loop)), shutdown, + external_tools_shutdown, idle_waiter, capabilities, open_canvases, @@ -1993,11 +2054,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( @@ -2030,7 +2094,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() => { @@ -2197,6 +2261,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(); @@ -2310,6 +2376,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; @@ -2453,8 +2526,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, @@ -2462,11 +2546,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 @@ -2496,13 +2591,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() @@ -2525,9 +2622,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(), @@ -2537,6 +2638,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/e2e.rs b/rust/tests/e2e.rs index eb4e75099..9d1c868fe 100644 --- a/rust/tests/e2e.rs +++ b/rust/tests/e2e.rs @@ -33,6 +33,8 @@ mod elicitation; mod error_resilience; #[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"] diff --git a/rust/tests/e2e/external_tool_cancellation.rs b/rust/tests/e2e/external_tool_cancellation.rs new file mode 100644 index 000000000..eb9b1b665 --- /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/rust/tests/session_test.rs b/rust/tests/session_test.rs index ad4c8abe4..e7bafc683 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -771,6 +771,108 @@ async fn create_session_registers_mcp_auth_interest_only_with_handler() { let _session = timeout(TIMEOUT, create_handle).await.unwrap().unwrap(); } +#[tokio::test] +async fn create_session_mcp_auth_registration_failure_cancels_external_tools() { + 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 (client, mut server_read, mut server_write) = make_client(); + let create_handle = tokio::spawn({ + let client = client.clone(); + async move { + client + .create_session( + SessionConfig::default() + .with_permission_handler(Arc::new(ApproveAllHandler)) + .with_mcp_auth_handler(Arc::new(CancelMcpAuthHandler)) + .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 create_req = read_framed(&mut server_read).await; + let session_id = requested_session_id(&create_req).to_string(); + server_respond_create(&mut server_write, &create_req, &session_id).await; + let interest_req = read_framed(&mut server_read).await; + assert_eq!(interest_req["method"], "session.eventLog.registerInterest"); + + let event = serde_json::json!({ + "jsonrpc": "2.0", + "method": "session.event", + "params": { + "sessionId": session_id, + "event": { + "id": "evt-registration-failure", + "timestamp": "2025-01-01T00:00:00Z", + "type": "external_tool.requested", + "data": { + "requestId": "request-registration-failure", + "sessionId": session_id, + "toolCallId": "tool-call-registration-failure", + "toolName": "blocked_tool", + "arguments": {}, + }, + }, + }, + }); + write_framed(&mut server_write, &serde_json::to_vec(&event).unwrap()).await; + timeout(TIMEOUT, started_rx).await.unwrap().unwrap(); + + let interest_id = interest_req["id"].as_u64().unwrap(); + let error = serde_json::json!({ + "jsonrpc": "2.0", + "id": interest_id, + "error": { "code": -32603, "message": "registration failed" }, + }); + write_framed(&mut server_write, &serde_json::to_vec(&error).unwrap()).await; + + assert!( + timeout(TIMEOUT, create_handle) + .await + .unwrap() + .unwrap() + .is_err() + ); + timeout(TIMEOUT, cancelled_rx).await.unwrap().unwrap(); +} + #[tokio::test] async fn cloud_create_session_registers_mcp_auth_interest_after_create_only_with_handler() { let cloud = || { @@ -3911,6 +4013,229 @@ 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.detach"); + assert!(!lifetime.is_cancelled()); + assert!( + timeout(Duration::from_millis(50), &mut cancelled_rx) + .await + .is_err() + ); + + server + .respond(&request, serde_json::json!({"success": true})) + .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 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 000000000..aa37004d6 --- /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 000000000..028b44e73 --- /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.