Add streaming cancellation support

docs/content/services/chat-service.md

@@ -289,6 +289,41 @@ chatService.chatStreaming(
 
 > **Threading note:** Callbacks run on the callback executor (virtual threads on Java 21+, configurable via the `CallbackExecutorProvider` SPI). Within a single request every callback is delivered sequentially, in the order the events were emitted, and never concurrently with another one: all the `onPartialToolCall` fragments of a tool call reach the handler before its `onCompleteToolCall`, tool calls arrive in index order, and `onCompleteResponse` comes last. Consecutive callbacks may run on different threads, but each one returns before the next one starts and what it wrote is visible to the next one, so your handler can accumulate into unsynchronized fields. The flip side is that a callback that blocks delays the ones after it in the same request.
 
+### Cancelling a Stream
+
+Cancel the returned future to stop a stream early, for example when the user navigates away or your own deadline expires:
+
+```java
+CompletableFuture<ChatResponse> future = chatService.chatStreaming(
+    List.of(UserMessage.text("Tell me a very long story")),
+    System.out::print
+);
+
+future.cancel(true);
+```
+
+Cancellation aborts the response body subscription and closes the connection, so the model stops streaming. After `cancel(...)` returns, no further callback reaches the handler, not even `onError`, because stopping the stream is a decision of the caller rather than a failure. A callback that is already running is allowed to finish.
+
+The future ends in the cancelled state, so a later `get()` or `join()` throws a `CancellationException`. Cancelling twice is a no-op, and so is cancelling a stream that has already completed. The `mayInterruptIfRunning` flag is ignored, therefore `cancel(false)` behaves exactly like `cancel(true)`.
+
+Calling `cancel(...)` from inside a callback is safe, which is the usual way to stop as soon as a condition is met:
+
+```java
+var future = new AtomicReference<CompletableFuture<ChatResponse>>();
+var chunks = new AtomicInteger();
+
+future.set(chatService.chatStreaming(messages, new ChatHandler() {
+    @Override
+    public void onPartialResponse(String text, PartialChatResponse partial) {
+        System.out.print(text);
+        if (chunks.incrementAndGet() == 10)
+            future.get().cancel(true);
+    }
+}));
+```
+
+> **Note:** The request has already been sent when you cancel, so the model may have generated tokens that you never receive. Cancellation stops the delivery of the response, it does not undo the work already done on the server.
+
 ---
 
 ## Tool Calling

docs/content/services/deployment-service.md

@@ -121,6 +121,18 @@ CompletableFuture<ChatResponse> future = deploymentService.chatStreaming(chatReq
 future.join(); // wait for completion
 ```
 
+### Cancelling a Stream
+
+Both `chatStreaming` and `generateStreaming` return a future that can be cancelled to stop the stream early:
+
+```java
+future.cancel(true);
+```
+
+Cancellation aborts the response body subscription and closes the connection, so the model stops streaming. After `cancel(...)` returns, no further callback reaches the handler, not even `onError`, because stopping the stream is a decision of the caller rather than a failure. A callback that is already running is allowed to finish.
+
+The future ends in the cancelled state, so a later `get()` or `join()` throws a `CancellationException`. Cancelling twice is a no-op, and so is cancelling a stream that has already completed. Calling `cancel(...)` from inside a callback is safe. See [Chat](../chat-service#cancelling-a-stream) for the full contract.
+
 ---
 
 ## Time Series Forecasting

docs/content/services/model-gateway/chat.md

@@ -192,6 +192,23 @@ service.chatStreaming(
 );
 ```
 
+### Cancelling a Stream
+
+Cancel the returned future to stop a stream early:
+
+```java
+CompletableFuture<ChatResponse> future = service.chatStreaming(
+    "Tell me a very long story",
+    System.out::print
+);
+
+future.cancel(true);
+```
+
+Cancellation aborts the response body subscription and closes the connection, so the model stops streaming. After `cancel(...)` returns, no further callback reaches the handler, not even `onError`, because stopping the stream is a decision of the caller rather than a failure. A callback that is already running is allowed to finish.
+
+The future ends in the cancelled state, so a later `get()` or `join()` throws a `CancellationException`. Cancelling twice is a no-op, and so is cancelling a stream that has already completed. Calling `cancel(...)` from inside a callback is safe. See [Chat](../chat-service#cancelling-a-stream) for the full contract.
+
 ---
 
 ## Model Gateway Parameters

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatProvider.java

@@ -32,6 +32,9 @@ public interface ChatProvider<R extends BaseChatRequest, C extends ChatResponse>
      * <p>
      * This method initiates an asynchronous chat operation where partial responses are delivered incrementally through the provided
      * {@link ChatHandler}.
+     * <p>
+     * Calling {@code cancel(...)} on the returned future stops the stream: the response body subscription is aborted and no further callback is
+     * dispatched to the handler.
      *
      * @param chatRequest the chat request object
      * @param handler a {@link ChatHandler} implementation that receives partial responses, the complete response, and error notifications

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatService.java

@@ -134,6 +134,9 @@ public TextChatResponse chat(ChatRequest chatRequest) {
 
     /**
      * Sends a streaming chat request.
+     * <p>
+     * Calling {@code cancel(...)} on the returned future stops the stream: the response body subscription is aborted and no further callback is
+     * dispatched to the handler.
      *
      * @param chatRequest the {@link ChatRequest}
      * @param handler a {@link ChatHandler} implementation that receives partial responses, the complete response, and error notifications

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/DefaultRestClient.java

@@ -88,17 +88,30 @@ public CompletableFuture<ChatResponse> chatStreaming(
             );
 
         var subscriber = chatSubscriber.asFlowSubscriber(response, !handler.failOnFirstError());
-        asyncHttpClient.send(httpRequest.build(), responseInfo -> logResponses
+        var httpFuture = asyncHttpClient.send(httpRequest.build(), responseInfo -> logResponses
             ? BodySubscribers.fromLineSubscriber(new SseEventLogger(subscriber, responseInfo.statusCode(), responseInfo.headers()))
-            : BodySubscribers.fromLineSubscriber(subscriber))
+            : BodySubscribers.fromLineSubscriber(subscriber));
+
+        httpFuture
             .thenAccept(r -> {})
             .exceptionally(t -> {
+                if (chatSubscriber.isCancelled())
+                    return null;
+
                 Throwable cause = nonNull(t.getCause()) ? t.getCause() : t;
                 if (chatSubscriber.markErrorReported())
                     handler.onError(cause);
                 response.completeExceptionally(cause);
                 return null;
             });
+
+        response.whenComplete((r, t) -> {
+            if (response.isCancelled()) {
+                chatSubscriber.cancelStream();
+                httpFuture.cancel(true);
+            }
+        });
+
         return response;
     }
 

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecorator.java

@@ -8,6 +8,7 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import com.ibm.watsonx.ai.chat.BaseChatRequest;
 import com.ibm.watsonx.ai.chat.ChatHandler;
@@ -69,6 +70,11 @@ public class ChatHandlerDecorator<R extends BaseChatRequest> implements ChatHand
     private final AtomicReference<CompletableFuture<Void>> callbackChain =
         new AtomicReference<>(CompletableFuture.completedFuture(null));
 
+    /**
+     * Set once the streaming request has been cancelled, after which no callback is delivered to the delegate.
+     */
+    private final AtomicBoolean cancelled = new AtomicBoolean(false);
+
     /**
      * Constructs a new {@code ChatHandlerDecorator}.
      *
@@ -121,6 +127,27 @@ public boolean failOnFirstError() {
         return delegate.failOnFirstError();
     }
 
+    /**
+     * Stops the delivery of every callback that has not started yet.
+     * <p>
+     * A callback already running is allowed to finish, but nothing further is handed to the delegate, including {@link ChatHandler#onError}. Safe to
+     * call from any thread and idempotent.
+     *
+     * @return {@code true} if this call cancelled the delivery, {@code false} if it was already cancelled
+     */
+    public boolean cancel() {
+        return cancelled.compareAndSet(false, true);
+    }
+
+    /**
+     * Returns whether callback delivery has been cancelled.
+     *
+     * @return {@code true} if {@link #cancel()} has been called
+     */
+    public boolean isCancelled() {
+        return cancelled.get();
+    }
+
     /**
      * Waits for every callback scheduled so far to be delivered.
      * <p>
@@ -160,7 +187,8 @@ private void scheduleCallback(Runnable callback) {
 
         previous.thenRunAsync(() -> {
             try {
-                callback.run();
+                if (!cancelled.get())
+                    callback.run();
             } catch (RuntimeException | Error e) {
                 safeOnError(e);
             } finally {
@@ -173,6 +201,9 @@ private void scheduleCallback(Runnable callback) {
      * Reports an error to the delegate, ignoring any failure of the error callback itself.
      */
     private void safeOnError(Throwable error) {
+        if (cancelled.get())
+            return;
+
         try {
             delegate.onError(error);
         } catch (RuntimeException | Error ignored) {

modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/streaming/DefaultChatSubscriber.java

@@ -10,6 +10,7 @@
 import java.util.concurrent.Flow;
 import java.util.concurrent.Flow.Subscription;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import com.ibm.watsonx.ai.chat.ChatHandler;
 import com.ibm.watsonx.ai.chat.ChatResponse;
 import com.ibm.watsonx.ai.chat.SseEventProcessor;
@@ -32,18 +33,54 @@ public class DefaultChatSubscriber extends ChatSubscriber {
      */
     private final AtomicBoolean errorReported = new AtomicBoolean(false);
 
+    /**
+     * The same handler as {@link #handler}, typed to expose callback cancellation and {@code awaitCallbacks}.
+     */
+    private final ChatHandlerDecorator<?> decorator;
+
+    /**
+     * The subscription of the body stream, set once {@code onSubscribe} fires and read by {@link #cancelStream()}.
+     */
+    private final AtomicReference<Flow.Subscription> subscriptionRef = new AtomicReference<>();
+
     /**
      * Creates a new DefaultChatSubscriber.
      *
      * @param processor the stream processor that parses SSE chunks
      * @param handler the decorated handler that executes user callbacks
      */
-    public DefaultChatSubscriber(SseEventProcessor processor, ChatHandlerDecorator handler) {
+    public DefaultChatSubscriber(SseEventProcessor processor, ChatHandlerDecorator<?> handler) {
         super(processor, handler);
+        this.decorator = handler;
+    }
+
+    /**
+     * Stops the stream: no further callback is delivered to the handler and the body subscription is cancelled.
+     * <p>
+     * Safe to call from any thread, before the subscription is established, and more than once.
+     */
+    public void cancelStream() {
+        decorator.cancel();
+
+        var subscription = subscriptionRef.get();
+        if (nonNull(subscription))
+            subscription.cancel();
+    }
+
+    /**
+     * Returns whether the stream has been cancelled.
+     *
+     * @return {@code true} if {@link #cancelStream()} has been called
+     */
+    public boolean isCancelled() {
+        return decorator.isCancelled();
     }
 
     @Override
     public CompletableFuture<ChatResponse> onComplete() {
+        if (isCancelled())
+            return CompletableFuture.completedFuture(null);
+
         return awaitCallbacks()
             .thenCompose(completeToolCalls -> {
                 var response = processor.buildResponse();
@@ -84,11 +121,24 @@ public Flow.Subscriber<String> asFlowSubscriber(
             @Override
             public void onSubscribe(Subscription subscription) {
                 this.subscription = subscription;
+                subscriptionRef.set(subscription);
+
+                if (isCancelled()) {
+                    subscription.cancel();
+                    return;
+                }
+
                 this.subscription.request(1);
             }
 
             @Override
             public void onNext(String partialMessage) {
+
+                if (isCancelled()) {
+                    subscription.cancel();
+                    return;
+                }
+
                 try {
 
                     DefaultChatSubscriber.this.onNext(partialMessage);
@@ -107,7 +157,7 @@ public void onNext(String partialMessage) {
                     }
 
                 } finally {
-                    if (continueProcessing)
+                    if (continueProcessing && !isCancelled())
                         subscription.request(1);
                     else {
                         subscription.cancel();
@@ -117,6 +167,10 @@ public void onNext(String partialMessage) {
 
             @Override
             public void onError(Throwable throwable) {
+
+                if (isCancelled())
+                    return;
+
                 Throwable t = nonNull(throwable.getCause()) ? throwable.getCause() : throwable;
                 if (markErrorReported())
                     DefaultChatSubscriber.this.onError(t);
@@ -150,11 +204,7 @@ public void onComplete() {
      * @return a CompletableFuture that resolves to a list of all processed {@link CompletedToolCall} objects
      */
     private CompletableFuture<List<CompletedToolCall>> awaitCallbacks() {
-        // This cast is safe because the constructor enforces ChatHandlerDecorator type
-        if (!(handler instanceof ChatHandlerDecorator handlerDecorator))
-            throw new IllegalStateException("Handler must be a ChatHandlerDecorator");
-
-        return handlerDecorator.awaitCallbacks();
+        return decorator.awaitCallbacks();
     }
 
     /**