From c271fbfe75df79cf1736d5d52ca0b829ad0a2fce Mon Sep 17 00:00:00 2001 From: Alberto Gimeno Date: Wed, 17 Jun 2026 14:51:49 +0200 Subject: [PATCH 1/3] java: plumb AbortSignal through ToolInvocation for cooperative cancellation Add AbortSignal to ToolInvocation so tool handlers can observe when session.abort() is called and stop in-flight work cooperatively. - New AbortSignal class in com.github.copilot.rpc with isAborted(), onAborted(Runnable), and abort() (SDK-internal) methods - ToolInvocation.getAbortSignal() returns the signal; initialized to a fresh (non-aborted) instance by default - CopilotSession tracks active signals per requestId; injects the signal into each ToolInvocation before calling the handler - CopilotSession.abort() now fires all tracked active tool signals before sending the RPC abort request, enabling cooperative cancellation - Signals are removed from the tracking map when a tool invocation completes (success or error) to avoid leaks - Unit tests for all AbortSignal behaviours added to ToolInvocationTest Fixes #1433 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../com/github/copilot/CopilotSession.java | 13 +- .../com/github/copilot/rpc/AbortSignal.java | 113 ++++++++++++++++++ .../github/copilot/rpc/ToolInvocation.java | 67 ++++++++++- .../com/github/copilot/rpc/package-info.java | 4 + .../github/copilot/ToolInvocationTest.java | 73 +++++++++++ 5 files changed, 267 insertions(+), 3 deletions(-) create mode 100644 java/src/main/java/com/github/copilot/rpc/AbortSignal.java diff --git a/java/src/main/java/com/github/copilot/CopilotSession.java b/java/src/main/java/com/github/copilot/CopilotSession.java index adfeac013..14c3b04bf 100644 --- a/java/src/main/java/com/github/copilot/CopilotSession.java +++ b/java/src/main/java/com/github/copilot/CopilotSession.java @@ -169,6 +169,7 @@ public final class CopilotSession implements AutoCloseable { private final Set> eventHandlers = ConcurrentHashMap.newKeySet(); private final Map toolHandlers = new ConcurrentHashMap<>(); private final Map commandHandlers = new ConcurrentHashMap<>(); + private final Map activeToolSignals = new ConcurrentHashMap<>(); private final AtomicReference permissionHandler = new AtomicReference<>(); private final AtomicReference userInputHandler = new AtomicReference<>(); private final AtomicReference elicitationHandler = new AtomicReference<>(); @@ -882,15 +883,19 @@ private void handleBroadcastEventAsync(SessionEvent event) { */ private void executeToolAndRespondAsync(String requestId, String toolName, String toolCallId, Object arguments, ToolDefinition tool) { + var signal = new com.github.copilot.rpc.AbortSignal(); + activeToolSignals.put(requestId, signal); Runnable task = () -> { try { JsonNode argumentsNode = arguments instanceof JsonNode jn ? jn : (arguments != null ? MAPPER.valueToTree(arguments) : null); var invocation = new com.github.copilot.rpc.ToolInvocation().setSessionId(sessionId) - .setToolCallId(toolCallId).setToolName(toolName).setArguments(argumentsNode); + .setToolCallId(toolCallId).setToolName(toolName).setArguments(argumentsNode) + .setAbortSignal(signal); tool.handler().invoke(invocation).thenAccept(result -> { + activeToolSignals.remove(requestId); try { ToolResultObject toolResult; if (result instanceof ToolResultObject tr) { @@ -905,6 +910,7 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin LOG.log(Level.WARNING, "Error sending tool result for requestId=" + requestId, e); } }).exceptionally(ex -> { + activeToolSignals.remove(requestId); try { getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId, requestId, null, ex.getMessage() != null ? ex.getMessage() : ex.toString())); @@ -914,6 +920,7 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin return null; }); } catch (Exception e) { + activeToolSignals.remove(requestId); LOG.log(Level.WARNING, "Error executing tool for requestId=" + requestId, e); try { getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId, @@ -1796,6 +1803,9 @@ public CompletableFuture> getMessages() { */ public CompletableFuture abort() { ensureNotTerminated(); + for (com.github.copilot.rpc.AbortSignal signal : activeToolSignals.values()) { + signal.abort(); + } return rpc.invoke("session.abort", Map.of("sessionId", sessionId), Void.class); } @@ -2136,6 +2146,7 @@ public void close() { eventHandlers.clear(); toolHandlers.clear(); commandHandlers.clear(); + activeToolSignals.clear(); permissionHandler.set(null); userInputHandler.set(null); elicitationHandler.set(null); diff --git a/java/src/main/java/com/github/copilot/rpc/AbortSignal.java b/java/src/main/java/com/github/copilot/rpc/AbortSignal.java new file mode 100644 index 000000000..01482328b --- /dev/null +++ b/java/src/main/java/com/github/copilot/rpc/AbortSignal.java @@ -0,0 +1,113 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.rpc; + +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * A signal that indicates whether a tool invocation has been aborted. + *

+ * An {@code AbortSignal} is passed to tool handlers via + * {@link ToolInvocation#getAbortSignal()} and is triggered when + * {@link com.github.copilot.CopilotSession#abort()} is called while the tool is + * executing. Tool handlers can use this to implement cooperative cancellation, + * allowing them to stop long-running work gracefully when the session is aborted. + * + *

Example Usage

+ * + * 
{@code
+ * ToolHandler handler = invocation -> {
+ * 	AbortSignal signal = invocation.getAbortSignal();
+ * 	return CompletableFuture.supplyAsync(() -> {
+ * 		while (!signal.isAborted()) {
+ * 			// do incremental work here
+ * 		}
+ * 		throw new CancellationException("Tool aborted");
+ * 	});
+ * };
+ * }
+ * + *

Callback Registration

+ * + * 
{@code
+ * ToolHandler handler = invocation -> {
+ * 	AbortSignal signal = invocation.getAbortSignal();
+ * 	signal.onAborted(() -> System.out.println("Aborting tool!"));
+ * 	// ... perform work ...
+ * 	return CompletableFuture.completedFuture("done");
+ * };
+ * }
+ * + * @see ToolInvocation#getAbortSignal() + * @see com.github.copilot.CopilotSession#abort() + * @since 1.6.0 + */ +public final class AbortSignal { + + private final AtomicBoolean aborted = new AtomicBoolean(false); + private final List listeners = new CopyOnWriteArrayList<>(); + + /** + * Returns whether this signal has been aborted. + * + * @return {@code true} if {@link com.github.copilot.CopilotSession#abort()} was + * called while this tool invocation was in progress; {@code false} + * otherwise + */ + public boolean isAborted() { + return aborted.get(); + } + + /** + * Registers a callback to be invoked when this signal is aborted. + *

+ * If the signal is already aborted at the time of registration, the callback is + * invoked immediately on the calling thread. + *

+ * Exceptions thrown by the callback are silently ignored. + * + * @param listener + * the callback to invoke on abort + * @throws NullPointerException + * if listener is null + */ + public void onAborted(Runnable listener) { + Objects.requireNonNull(listener, "listener must not be null"); + listeners.add(listener); + if (aborted.get()) { + try { + listener.run(); + } catch (Exception ignored) { + // Exceptions from listeners are silently ignored + } + } + } + + /** + * Triggers this abort signal, notifying all registered listeners. + *

+ * Note: This method is intended for internal SDK use only. + * It is called by the SDK when + * {@link com.github.copilot.CopilotSession#abort()} is invoked while this tool + * invocation is in progress. + *

+ * Calling this method more than once has no effect — the signal fires exactly + * once. + */ + public void abort() { + if (aborted.compareAndSet(false, true)) { + for (Runnable listener : listeners) { + try { + listener.run(); + } catch (Exception ignored) { + // Exceptions from listeners are silently ignored + } + } + } + } +} diff --git a/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java b/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java index dddfdd06f..14f40ff29 100644 --- a/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java +++ b/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java @@ -6,6 +6,7 @@ import java.util.Map; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonSetter; import com.fasterxml.jackson.core.type.TypeReference; @@ -16,11 +17,28 @@ * Represents a tool invocation request from the AI assistant. *

* When the assistant invokes a tool, this object contains the context including - * the session ID, tool call ID, tool name, and arguments parsed from the - * assistant's request. + * the session ID, tool call ID, tool name, arguments parsed from the + * assistant's request, and an {@link AbortSignal} that is triggered when + * {@link com.github.copilot.CopilotSession#abort()} is called while the tool is + * executing. + * + *

Cooperative Cancellation

+ * + * 
{@code
+ * ToolHandler handler = invocation -> {
+ * 	AbortSignal signal = invocation.getAbortSignal();
+ * 	return CompletableFuture.supplyAsync(() -> {
+ * 		while (!signal.isAborted()) {
+ * 			// do incremental work here
+ * 		}
+ * 		throw new CancellationException("Tool aborted");
+ * 	});
+ * };
+ * }
* * @see ToolHandler * @see ToolDefinition + * @see AbortSignal * @since 1.0.0 */ @JsonInclude(JsonInclude.Include.NON_NULL) @@ -34,6 +52,7 @@ public final class ToolInvocation { private String toolCallId; private String toolName; private JsonNode argumentsNode; + private AbortSignal abortSignal = new AbortSignal(); /** * Gets the session ID where the tool was invoked. @@ -168,4 +187,48 @@ public ToolInvocation setArguments(JsonNode arguments) { this.argumentsNode = arguments; return this; } + + /** + * Returns the abort signal for this tool invocation. + *

+ * The signal is triggered when + * {@link com.github.copilot.CopilotSession#abort()} is called while this tool + * is executing. Use it to implement cooperative cancellation in your tool + * handler. + * + * 

{@code
+     * ToolHandler handler = invocation -> {
+     * 	AbortSignal signal = invocation.getAbortSignal();
+     * 	return CompletableFuture.supplyAsync(() -> {
+     * 		while (!signal.isAborted()) {
+     * 			// do incremental work here
+     * 		}
+     * 		throw new CancellationException("Tool aborted");
+     * 	});
+     * };
+     * }
+ * + * @return the abort signal; never {@code null} + * @see AbortSignal + * @since 1.6.0 + */ + @JsonIgnore + public AbortSignal getAbortSignal() { + return abortSignal; + } + + /** + * Sets the abort signal for this tool invocation. + *

+ * Note: This method is intended for internal SDK use only. + * Users do not need to call this method directly. + * + * @param abortSignal + * the abort signal to associate with this invocation + * @return this invocation for method chaining + */ + public ToolInvocation setAbortSignal(AbortSignal abortSignal) { + this.abortSignal = abortSignal; + return this; + } } diff --git a/java/src/main/java/com/github/copilot/rpc/package-info.java b/java/src/main/java/com/github/copilot/rpc/package-info.java index edc7dedcf..2339ec864 100644 --- a/java/src/main/java/com/github/copilot/rpc/package-info.java +++ b/java/src/main/java/com/github/copilot/rpc/package-info.java @@ -39,6 +39,10 @@ * tool that can be invoked by the assistant. *

  • {@link com.github.copilot.rpc.ToolInvocation} - Represents a tool * invocation request from the assistant.
    • + *
  • {@link com.github.copilot.rpc.AbortSignal} - Cancellation signal passed + * to tool handlers via {@link com.github.copilot.rpc.ToolInvocation#getAbortSignal()}, + * triggered when {@link com.github.copilot.CopilotSession#abort()} is + * called.
    • *
  • {@link com.github.copilot.rpc.Attachment} - File attachment for * messages.
    • * diff --git a/java/src/test/java/com/github/copilot/ToolInvocationTest.java b/java/src/test/java/com/github/copilot/ToolInvocationTest.java index 2bc9edb1b..a4ffec430 100644 --- a/java/src/test/java/com/github/copilot/ToolInvocationTest.java +++ b/java/src/test/java/com/github/copilot/ToolInvocationTest.java @@ -6,10 +6,13 @@ import static org.junit.jupiter.api.Assertions.*; +import java.util.concurrent.atomic.AtomicBoolean; + import org.junit.jupiter.api.Test; import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.github.copilot.rpc.AbortSignal; import com.github.copilot.rpc.ToolInvocation; /** @@ -127,6 +130,76 @@ void testGetArgumentsAsThrowsOnInvalidType() { assertTrue(exception.getMessage().contains("StrictType")); } + /** + * Test that getAbortSignal returns a non-null signal by default. + */ + @Test + void testGetAbortSignalReturnedByDefault() { + ToolInvocation invocation = new ToolInvocation().setSessionId("s1").setToolCallId("c1") + .setToolName("my_tool"); + assertNotNull(invocation.getAbortSignal(), "getAbortSignal should not return null"); + assertFalse(invocation.getAbortSignal().isAborted(), "signal should not be aborted by default"); + } + + /** + * Test that isAborted returns true after the signal is aborted. + */ + @Test + void testAbortSignalIsAbortedAfterAbort() { + AbortSignal signal = new AbortSignal(); + assertFalse(signal.isAborted()); + signal.abort(); + assertTrue(signal.isAborted()); + } + + /** + * Test that onAborted callback is invoked when signal is aborted. + */ + @Test + void testAbortSignalOnAbortedCallbackInvoked() { + AbortSignal signal = new AbortSignal(); + var called = new AtomicBoolean(false); + signal.onAborted(() -> called.set(true)); + assertFalse(called.get()); + signal.abort(); + assertTrue(called.get()); + } + + /** + * Test that onAborted callback is invoked immediately if signal is already + * aborted. + */ + @Test + void testAbortSignalOnAbortedCallbackInvokedImmediatelyIfAlreadyAborted() { + AbortSignal signal = new AbortSignal(); + signal.abort(); + var called = new AtomicBoolean(false); + signal.onAborted(() -> called.set(true)); + assertTrue(called.get(), "callback should be invoked immediately when signal is already aborted"); + } + + /** + * Test that abort() is idempotent — callbacks fire only once. + */ + @Test + void testAbortSignalAbortIsIdempotent() { + AbortSignal signal = new AbortSignal(); + var count = new java.util.concurrent.atomic.AtomicInteger(0); + signal.onAborted(count::incrementAndGet); + signal.abort(); + signal.abort(); + assertEquals(1, count.get(), "callback should be invoked exactly once even if abort() called twice"); + } + + /** + * Test that onAborted throws NullPointerException for null listener. + */ + @Test + void testAbortSignalOnAbortedRejectsNullListener() { + AbortSignal signal = new AbortSignal(); + assertThrows(NullPointerException.class, () -> signal.onAborted(null)); + } + /** * Record for testing type-safe argument deserialization. */ From d1030f8108159aec998a437198d2a762fd97b8ed Mon Sep 17 00:00:00 2001 From: Alberto Gimeno Date: Wed, 17 Jun 2026 15:08:31 +0200 Subject: [PATCH 2/3] java: add cancelToolCall and README cancellation docs - CopilotSession.cancelToolCall(toolCallId): fires the AbortSignal for only the named in-flight handler without aborting the agentic loop or other handlers; returns true if found and cancelled, false otherwise - activeToolSignals map is now keyed by toolCallId (falls back to requestId when toolCallId is null) so cancelToolCall can look up directly by the ID exposed on ToolInvocation - CancelToolCallTest: 3 unit tests covering targeted cancel, unknown id returns false, and signal removal from the tracking map - java/README.md: new 'Tool Handler Cancellation' section documenting both abort() (all handlers) and cancelToolCall() (single handler) with isAborted()/onAborted() handler examples Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- java/README.md | 43 +++++++++ .../com/github/copilot/CopilotSession.java | 49 +++++++++- .../com/github/copilot/rpc/AbortSignal.java | 10 +- .../com/github/copilot/rpc/package-info.java | 6 +- .../github/copilot/CancelToolCallTest.java | 91 +++++++++++++++++++ .../github/copilot/ToolInvocationTest.java | 3 +- 6 files changed, 188 insertions(+), 14 deletions(-) create mode 100644 java/src/test/java/com/github/copilot/CancelToolCallTest.java diff --git a/java/README.md b/java/README.md index 7c7209bc8..12aaa3553 100644 --- a/java/README.md +++ b/java/README.md @@ -242,6 +242,49 @@ public class Consumer { The gate also applies to individual methods annotated with `@CopilotExperimental` on otherwise stable types. When a type-level annotation is present, all member accesses through that type are considered experimental. `@AllowCopilotExperimental` mirrors the same declaration-level boundary: annotating a class opts in that class and its enclosed declarations, while annotating a method or constructor opts in just that executable signature. +## Tool Handler Cancellation + +Tool handlers can observe when the session is aborted and stop in-flight work cooperatively using the `AbortSignal` available on every `ToolInvocation`. + +### Cancelling all in-flight handlers: `session.abort()` + +`session.abort()` stops the agentic loop **and** immediately fires the `AbortSignal` for every tool handler currently running in the session. + +```java +// Handler using isAborted() polling +var tool = ToolDefinition.create("long_task", "A long-running task", + Map.of("type", "object", "properties", Map.of()), + invocation -> CompletableFuture.supplyAsync(() -> { + AbortSignal signal = invocation.getAbortSignal(); + while (!signal.isAborted()) { + // perform incremental work + } + return "cancelled"; + }) +); + +// Handler using onAborted() callback +var tool2 = ToolDefinition.create("http_task", "Task with cleanup", + Map.of("type", "object", "properties", Map.of()), + invocation -> { + var future = new CompletableFuture(); + invocation.getAbortSignal().onAborted(() -> future.complete("aborted")); + // start async work that completes future normally ... + return future; + } +); +``` + +### Cancelling a single handler: `session.cancelToolCall(toolCallId)` + +`cancelToolCall(toolCallId)` fires the `AbortSignal` for **only** the named in-flight invocation, without aborting the agentic loop or affecting other running handlers. Returns `true` if the call was found and cancelled, `false` if not found (already completed or unknown id). + +```java +// Cancel a specific in-flight tool call by its ID +String toolCallId = invocation.getToolCallId(); // captured from ToolInvocation +boolean wasCancelled = session.cancelToolCall(toolCallId); +``` + ## Projects Using This SDK | Project | Description | diff --git a/java/src/main/java/com/github/copilot/CopilotSession.java b/java/src/main/java/com/github/copilot/CopilotSession.java index 14c3b04bf..4c4b2bbc0 100644 --- a/java/src/main/java/com/github/copilot/CopilotSession.java +++ b/java/src/main/java/com/github/copilot/CopilotSession.java @@ -884,7 +884,8 @@ private void handleBroadcastEventAsync(SessionEvent event) { private void executeToolAndRespondAsync(String requestId, String toolName, String toolCallId, Object arguments, ToolDefinition tool) { var signal = new com.github.copilot.rpc.AbortSignal(); - activeToolSignals.put(requestId, signal); + String signalKey = toolCallId != null ? toolCallId : requestId; + activeToolSignals.put(signalKey, signal); Runnable task = () -> { try { JsonNode argumentsNode = arguments instanceof JsonNode jn @@ -895,7 +896,7 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin .setAbortSignal(signal); tool.handler().invoke(invocation).thenAccept(result -> { - activeToolSignals.remove(requestId); + activeToolSignals.remove(signalKey); try { ToolResultObject toolResult; if (result instanceof ToolResultObject tr) { @@ -910,7 +911,7 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin LOG.log(Level.WARNING, "Error sending tool result for requestId=" + requestId, e); } }).exceptionally(ex -> { - activeToolSignals.remove(requestId); + activeToolSignals.remove(signalKey); try { getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId, requestId, null, ex.getMessage() != null ? ex.getMessage() : ex.toString())); @@ -920,7 +921,7 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin return null; }); } catch (Exception e) { - activeToolSignals.remove(requestId); + activeToolSignals.remove(signalKey); LOG.log(Level.WARNING, "Error executing tool for requestId=" + requestId, e); try { getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId, @@ -1809,6 +1810,46 @@ public CompletableFuture abort() { return rpc.invoke("session.abort", Map.of("sessionId", sessionId), Void.class); } + /** + * Cancels a single in-flight tool handler by its tool call ID. + *

    + * Unlike {@link #abort()}, this method fires the + * {@link com.github.copilot.rpc.AbortSignal} for only the specified tool + * invocation and does not abort the agentic loop or affect any other in-flight + * handlers. + *

    + * The signal is fired and the entry is removed from the tracking map + * immediately. The handler is responsible for observing + * {@link com.github.copilot.rpc.AbortSignal#isAborted()} or registering an + * {@link com.github.copilot.rpc.AbortSignal#onAborted(Runnable)} callback. + * + * 

    {@code
+     * // Cancel a specific tool invocation
+     * boolean cancelled = session.cancelToolCall(toolCallId);
+     * if (!cancelled) {
+     * 	// tool call was already complete or id was not found
+     * }
+     * }
    + * + * @param toolCallId + * the tool call ID to cancel, as provided by + * {@link com.github.copilot.rpc.ToolInvocation#getToolCallId()} + * @return {@code true} if an in-flight handler was found and its signal was + * fired; {@code false} if no in-flight handler matched the given ID + * @throws IllegalStateException + * if this session has been terminated + * @since 1.6.0 + */ + public boolean cancelToolCall(String toolCallId) { + ensureNotTerminated(); + com.github.copilot.rpc.AbortSignal signal = activeToolSignals.remove(toolCallId); + if (signal != null) { + signal.abort(); + return true; + } + return false; + } + /** * Changes the model for this session with an optional reasoning effort level. *

    diff --git a/java/src/main/java/com/github/copilot/rpc/AbortSignal.java b/java/src/main/java/com/github/copilot/rpc/AbortSignal.java index 01482328b..914741f8a 100644 --- a/java/src/main/java/com/github/copilot/rpc/AbortSignal.java +++ b/java/src/main/java/com/github/copilot/rpc/AbortSignal.java @@ -16,7 +16,8 @@ * {@link ToolInvocation#getAbortSignal()} and is triggered when * {@link com.github.copilot.CopilotSession#abort()} is called while the tool is * executing. Tool handlers can use this to implement cooperative cancellation, - * allowing them to stop long-running work gracefully when the session is aborted. + * allowing them to stop long-running work gracefully when the session is + * aborted. * *

    Example Usage

    * @@ -91,10 +92,9 @@ public void onAborted(Runnable listener) { /** * Triggers this abort signal, notifying all registered listeners. *

    - * Note: This method is intended for internal SDK use only. - * It is called by the SDK when - * {@link com.github.copilot.CopilotSession#abort()} is invoked while this tool - * invocation is in progress. + * Note: This method is intended for internal SDK use only. It + * is called by the SDK when {@link com.github.copilot.CopilotSession#abort()} + * is invoked while this tool invocation is in progress. *

    * Calling this method more than once has no effect — the signal fires exactly * once. diff --git a/java/src/main/java/com/github/copilot/rpc/package-info.java b/java/src/main/java/com/github/copilot/rpc/package-info.java index 2339ec864..eda97abe4 100644 --- a/java/src/main/java/com/github/copilot/rpc/package-info.java +++ b/java/src/main/java/com/github/copilot/rpc/package-info.java @@ -40,9 +40,9 @@ *

  • {@link com.github.copilot.rpc.ToolInvocation} - Represents a tool * invocation request from the assistant.
    • *
  • {@link com.github.copilot.rpc.AbortSignal} - Cancellation signal passed - * to tool handlers via {@link com.github.copilot.rpc.ToolInvocation#getAbortSignal()}, - * triggered when {@link com.github.copilot.CopilotSession#abort()} is - * called.
    • + * to tool handlers via + * {@link com.github.copilot.rpc.ToolInvocation#getAbortSignal()}, triggered + * when {@link com.github.copilot.CopilotSession#abort()} is called. *
  • {@link com.github.copilot.rpc.Attachment} - File attachment for * messages.
    • * diff --git a/java/src/test/java/com/github/copilot/CancelToolCallTest.java b/java/src/test/java/com/github/copilot/CancelToolCallTest.java new file mode 100644 index 000000000..9f76d7a39 --- /dev/null +++ b/java/src/test/java/com/github/copilot/CancelToolCallTest.java @@ -0,0 +1,91 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import static org.junit.jupiter.api.Assertions.*; + +import java.lang.reflect.Field; +import java.util.Map; + +import org.junit.jupiter.api.Test; + +import com.github.copilot.rpc.AbortSignal; + +/** + * Unit tests for {@link CopilotSession#cancelToolCall(String)}. + *

    + * Uses reflection to inject {@link AbortSignal} instances directly into the + * session's active-tool-signal tracking map, allowing the cancellation logic to + * be verified in isolation without requiring the full E2E test harness. + */ +class CancelToolCallTest { + + /** + * Injects two signals into a session, calls cancelToolCall for one, and + * verifies that only the targeted signal is aborted while the other remains + * unaffected. + */ + @Test + void cancelToolCallFiresOnlyTargetedSignal() throws Exception { + var session = new CopilotSession("sess-cancel-test", null); + + AbortSignal signalA = new AbortSignal(); + AbortSignal signalB = new AbortSignal(); + + Map map = getActiveToolSignals(session); + map.put("call-A", signalA); + map.put("call-B", signalB); + + boolean result = session.cancelToolCall("call-A"); + + assertTrue(result, "cancelToolCall should return true for a known toolCallId"); + assertTrue(signalA.isAborted(), "signal A should be aborted after cancelToolCall(call-A)"); + assertFalse(signalB.isAborted(), "signal B must NOT be aborted — only the targeted signal fires"); + } + + /** + * Verifies that cancelToolCall returns false for an unknown tool call ID, + * without affecting any in-flight signals. + */ + @Test + void cancelToolCallReturnsFalseForUnknownId() throws Exception { + var session = new CopilotSession("sess-cancel-unknown", null); + + AbortSignal signal = new AbortSignal(); + Map map = getActiveToolSignals(session); + map.put("call-exists", signal); + + boolean result = session.cancelToolCall("call-does-not-exist"); + + assertFalse(result, "cancelToolCall should return false for an unknown toolCallId"); + assertFalse(signal.isAborted(), "existing signal must not be affected"); + } + + /** + * Verifies that a cancelled signal is removed from the tracking map so it + * cannot be double-fired. + */ + @Test + void cancelToolCallRemovesSignalFromMap() throws Exception { + var session = new CopilotSession("sess-cancel-cleanup", null); + + AbortSignal signal = new AbortSignal(); + Map map = getActiveToolSignals(session); + map.put("call-X", signal); + + session.cancelToolCall("call-X"); + + assertFalse(map.containsKey("call-X"), "signal should be removed from the map after cancellation"); + // second call must return false since the entry is gone + assertFalse(session.cancelToolCall("call-X"), "second cancelToolCall for same id should return false"); + } + + @SuppressWarnings("unchecked") + private static Map getActiveToolSignals(CopilotSession session) throws Exception { + Field f = CopilotSession.class.getDeclaredField("activeToolSignals"); + f.setAccessible(true); + return (Map) f.get(session); + } +} diff --git a/java/src/test/java/com/github/copilot/ToolInvocationTest.java b/java/src/test/java/com/github/copilot/ToolInvocationTest.java index a4ffec430..b7f386648 100644 --- a/java/src/test/java/com/github/copilot/ToolInvocationTest.java +++ b/java/src/test/java/com/github/copilot/ToolInvocationTest.java @@ -135,8 +135,7 @@ void testGetArgumentsAsThrowsOnInvalidType() { */ @Test void testGetAbortSignalReturnedByDefault() { - ToolInvocation invocation = new ToolInvocation().setSessionId("s1").setToolCallId("c1") - .setToolName("my_tool"); + ToolInvocation invocation = new ToolInvocation().setSessionId("s1").setToolCallId("c1").setToolName("my_tool"); assertNotNull(invocation.getAbortSignal(), "getAbortSignal should not return null"); assertFalse(invocation.getAbortSignal().isAborted(), "signal should not be aborted by default"); } From b57c3322b10743fce4c075af5c168fa2e5ae6734 Mon Sep 17 00:00:00 2001 From: Alberto Gimeno Date: Wed, 17 Jun 2026 16:19:47 +0200 Subject: [PATCH 3/3] java: address Copilot PR review feedback on AbortSignal - Fix at-most-once callback delivery: wrap each onAborted() listener in an AtomicBoolean-guarded Runnable so it cannot fire twice even when abort() races with onAborted() registration - Catch Throwable (not just Exception) in listener invocations to align with the 'silently ignored' Javadoc contract - Add @JsonIgnore to setAbortSignal() so it is excluded from Jackson serialization/deserialization, matching the getter annotation - setAbortSignal(null) now silently preserves the existing signal for backwards compatibility, rather than throwing NullPointerException - Add tests: at-most-once callback guarantee; null setAbortSignal is ignored and leaves the existing signal in place Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../com/github/copilot/rpc/AbortSignal.java | 34 ++++++++++++------- .../github/copilot/rpc/ToolInvocation.java | 12 +++++-- .../github/copilot/ToolInvocationTest.java | 27 +++++++++++++++ 3 files changed, 58 insertions(+), 15 deletions(-) diff --git a/java/src/main/java/com/github/copilot/rpc/AbortSignal.java b/java/src/main/java/com/github/copilot/rpc/AbortSignal.java index 914741f8a..2c7626814 100644 --- a/java/src/main/java/com/github/copilot/rpc/AbortSignal.java +++ b/java/src/main/java/com/github/copilot/rpc/AbortSignal.java @@ -4,7 +4,6 @@ package com.github.copilot.rpc; -import java.util.List; import java.util.Objects; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.atomic.AtomicBoolean; @@ -51,7 +50,7 @@ public final class AbortSignal { private final AtomicBoolean aborted = new AtomicBoolean(false); - private final List listeners = new CopyOnWriteArrayList<>(); + private final CopyOnWriteArrayList listeners = new CopyOnWriteArrayList<>(); /** * Returns whether this signal has been aborted. @@ -70,7 +69,9 @@ public boolean isAborted() { * If the signal is already aborted at the time of registration, the callback is * invoked immediately on the calling thread. *

    - * Exceptions thrown by the callback are silently ignored. + * The callback is guaranteed to be invoked at most once, regardless of + * concurrent calls to {@link #abort()} and {@code onAborted}. Any + * {@link Throwable} thrown by the callback is silently ignored. * * @param listener * the callback to invoke on abort @@ -79,13 +80,22 @@ public boolean isAborted() { */ public void onAborted(Runnable listener) { Objects.requireNonNull(listener, "listener must not be null"); - listeners.add(listener); - if (aborted.get()) { - try { - listener.run(); - } catch (Exception ignored) { - // Exceptions from listeners are silently ignored + // Wrap in an AtomicBoolean-guarded runnable so the callback fires at most once + // even if abort() races with this method between listeners.add() and the + // aborted.get() check below. + AtomicBoolean fired = new AtomicBoolean(false); + Runnable once = () -> { + if (fired.compareAndSet(false, true)) { + try { + listener.run(); + } catch (Throwable ignored) { + // Throwables from listeners are silently ignored + } } + }; + listeners.add(once); + if (aborted.get()) { + once.run(); } } @@ -97,15 +107,15 @@ public void onAborted(Runnable listener) { * is invoked while this tool invocation is in progress. *

    * Calling this method more than once has no effect — the signal fires exactly - * once. + * once. Any {@link Throwable} thrown by a listener is silently ignored. */ public void abort() { if (aborted.compareAndSet(false, true)) { for (Runnable listener : listeners) { try { listener.run(); - } catch (Exception ignored) { - // Exceptions from listeners are silently ignored + } catch (Throwable ignored) { + // Throwables from listeners are silently ignored } } } diff --git a/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java b/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java index 14f40ff29..3ef38206e 100644 --- a/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java +++ b/java/src/main/java/com/github/copilot/rpc/ToolInvocation.java @@ -221,14 +221,20 @@ public AbortSignal getAbortSignal() { * Sets the abort signal for this tool invocation. *

    * Note: This method is intended for internal SDK use only. - * Users do not need to call this method directly. + * Users do not need to call this method directly. Passing {@code null} is + * accepted for backwards compatibility and leaves the existing signal + * unchanged. * * @param abortSignal - * the abort signal to associate with this invocation + * the abort signal to associate with this invocation, or + * {@code null} to leave the existing signal unchanged * @return this invocation for method chaining */ + @JsonIgnore public ToolInvocation setAbortSignal(AbortSignal abortSignal) { - this.abortSignal = abortSignal; + if (abortSignal != null) { + this.abortSignal = abortSignal; + } return this; } } diff --git a/java/src/test/java/com/github/copilot/ToolInvocationTest.java b/java/src/test/java/com/github/copilot/ToolInvocationTest.java index b7f386648..1d3d14d9f 100644 --- a/java/src/test/java/com/github/copilot/ToolInvocationTest.java +++ b/java/src/test/java/com/github/copilot/ToolInvocationTest.java @@ -199,6 +199,33 @@ void testAbortSignalOnAbortedRejectsNullListener() { assertThrows(NullPointerException.class, () -> signal.onAborted(null)); } + /** + * Test that a callback registered via onAborted fires at most once even when + * abort() races with onAborted registration (at-most-once delivery guarantee). + */ + @Test + void testAbortSignalCallbackFiresAtMostOnce() { + AbortSignal signal = new AbortSignal(); + // Pre-abort the signal so onAborted() will fire immediately on registration + signal.abort(); + var count = new java.util.concurrent.atomic.AtomicInteger(0); + // Registering after abort fires immediately — but only once + signal.onAborted(count::incrementAndGet); + assertEquals(1, count.get(), "callback should fire exactly once when registered after abort"); + } + + /** + * Test that setAbortSignal(null) is accepted for backwards compatibility and + * leaves the existing signal unchanged. + */ + @Test + void testSetAbortSignalNullIsIgnored() { + ToolInvocation invocation = new ToolInvocation(); + AbortSignal original = invocation.getAbortSignal(); + invocation.setAbortSignal(null); // must not throw + assertSame(original, invocation.getAbortSignal(), "existing signal should be preserved when null is passed"); + } + /** * Record for testing type-safe argument deserialization. */