java: plumb AbortSignal through ToolInvocation for cooperative cancellation

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<Object>();
+        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                                |

java/src/main/java/com/github/copilot/CopilotSession.java

@@ -169,6 +169,7 @@ public final class CopilotSession implements AutoCloseable {
     private final Set<Consumer<SessionEvent>> eventHandlers = ConcurrentHashMap.newKeySet();
     private final Map<String, ToolDefinition> toolHandlers = new ConcurrentHashMap<>();
     private final Map<String, CommandHandler> commandHandlers = new ConcurrentHashMap<>();
+    private final Map<String, com.github.copilot.rpc.AbortSignal> activeToolSignals = new ConcurrentHashMap<>();
     private final AtomicReference<PermissionHandler> permissionHandler = new AtomicReference<>();
     private final AtomicReference<UserInputHandler> userInputHandler = new AtomicReference<>();
     private final AtomicReference<ElicitationHandler> elicitationHandler = new AtomicReference<>();
@@ -882,15 +883,20 @@ 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();
+        String signalKey = toolCallId != null ? toolCallId : requestId;
+        activeToolSignals.put(signalKey, 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(signalKey);
                     try {
                         ToolResultObject toolResult;
                         if (result instanceof ToolResultObject tr) {
@@ -905,6 +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(signalKey);
                     try {
                         getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId,
                                 requestId, null, ex.getMessage() != null ? ex.getMessage() : ex.toString()));
@@ -914,6 +921,7 @@ private void executeToolAndRespondAsync(String requestId, String toolName, Strin
                     return null;
                 });
             } catch (Exception e) {
+                activeToolSignals.remove(signalKey);
                 LOG.log(Level.WARNING, "Error executing tool for requestId=" + requestId, e);
                 try {
                     getRpc().tools.handlePendingToolCall(new SessionToolsHandlePendingToolCallParams(sessionId,
@@ -1796,9 +1804,52 @@ public CompletableFuture<List<SessionEvent>> getMessages() {
      */
     public CompletableFuture<Void> abort() {
         ensureNotTerminated();
+        for (com.github.copilot.rpc.AbortSignal signal : activeToolSignals.values()) {
+            signal.abort();
+        }
         return rpc.invoke("session.abort", Map.of("sessionId", sessionId), Void.class);
     }
 
+    /**
+     * Cancels a single in-flight tool handler by its tool call ID.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     *
+     * <pre>{@code
+     * // Cancel a specific tool invocation
+     * boolean cancelled = session.cancelToolCall(toolCallId);
+     * if (!cancelled) {
+     * 	// tool call was already complete or id was not found
+     * }
+     * }</pre>
+     *
+     * @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.
      * <p>
@@ -2136,6 +2187,7 @@ public void close() {
         eventHandlers.clear();
         toolHandlers.clear();
         commandHandlers.clear();
+        activeToolSignals.clear();
         permissionHandler.set(null);
         userInputHandler.set(null);
         elicitationHandler.set(null);

java/src/main/java/com/github/copilot/rpc/AbortSignal.java

@@ -0,0 +1,123 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+package com.github.copilot.rpc;
+
+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.
+ * <p>
+ * 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.
+ *
+ * <h2>Example Usage</h2>
+ *
+ * <pre>{@code
+ * ToolHandler handler = invocation -> {
+ * 	AbortSignal signal = invocation.getAbortSignal();
+ * 	return CompletableFuture.supplyAsync(() -> {
+ * 		while (!signal.isAborted()) {
+ * 			// do incremental work here
+ * 		}
+ * 		throw new CancellationException("Tool aborted");
+ * 	});
+ * };
+ * }</pre>
+ *
+ * <h2>Callback Registration</h2>
+ *
+ * <pre>{@code
+ * ToolHandler handler = invocation -> {
+ * 	AbortSignal signal = invocation.getAbortSignal();
+ * 	signal.onAborted(() -> System.out.println("Aborting tool!"));
+ * 	// ... perform work ...
+ * 	return CompletableFuture.completedFuture("done");
+ * };
+ * }</pre>
+ *
+ * @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 CopyOnWriteArrayList<Runnable> 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.
+     * <p>
+     * If the signal is already aborted at the time of registration, the callback is
+     * invoked immediately on the calling thread.
+     * <p>
+     * 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
+     * @throws NullPointerException
+     *             if listener is null
+     */
+    public void onAborted(Runnable listener) {
+        Objects.requireNonNull(listener, "listener must not be null");
+        // 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();
+        }
+    }
+
+    /**
+     * Triggers this abort signal, notifying all registered listeners.
+     * <p>
+     * <strong>Note:</strong> 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.
+     * <p>
+     * Calling this method more than once has no effect — the signal fires exactly
+     * 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 (Throwable ignored) {
+                    // Throwables from listeners are silently ignored
+                }
+            }
+        }
+    }
+}

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.
  * <p>
  * 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.
+ *
+ * <h2>Cooperative Cancellation</h2>
+ *
+ * <pre>{@code
+ * ToolHandler handler = invocation -> {
+ * 	AbortSignal signal = invocation.getAbortSignal();
+ * 	return CompletableFuture.supplyAsync(() -> {
+ * 		while (!signal.isAborted()) {
+ * 			// do incremental work here
+ * 		}
+ * 		throw new CancellationException("Tool aborted");
+ * 	});
+ * };
+ * }</pre>
  *
  * @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,54 @@ public ToolInvocation setArguments(JsonNode arguments) {
         this.argumentsNode = arguments;
         return this;
     }
+
+    /**
+     * Returns the abort signal for this tool invocation.
+     * <p>
+     * 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.
+     *
+     * <pre>{@code
+     * ToolHandler handler = invocation -> {
+     * 	AbortSignal signal = invocation.getAbortSignal();
+     * 	return CompletableFuture.supplyAsync(() -> {
+     * 		while (!signal.isAborted()) {
+     * 			// do incremental work here
+     * 		}
+     * 		throw new CancellationException("Tool aborted");
+     * 	});
+     * };
+     * }</pre>
+     *
+     * @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.
+     * <p>
+     * <strong>Note:</strong> This method is intended for internal SDK use only.
+     * 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, or
+     *            {@code null} to leave the existing signal unchanged
+     * @return this invocation for method chaining
+     */
+    @JsonIgnore
+    public ToolInvocation setAbortSignal(AbortSignal abortSignal) {
+        if (abortSignal != null) {
+            this.abortSignal = abortSignal;
+        }
+        return this;
+    }
 }

java/src/main/java/com/github/copilot/rpc/package-info.java

@@ -39,6 +39,10 @@
  * tool that can be invoked by the assistant.</li>
  * <li>{@link com.github.copilot.rpc.ToolInvocation} - Represents a tool
  * invocation request from the assistant.</li>
+ * <li>{@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.</li>
  * <li>{@link com.github.copilot.rpc.Attachment} - File attachment for
  * messages.</li>
  * </ul>