diff --git a/docs/content/services/chat-service.md b/docs/content/services/chat-service.md index 14d79c27..93f23796 100644 --- a/docs/content/services/chat-service.md +++ b/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 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>(); +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 diff --git a/docs/content/services/deployment-service.md b/docs/content/services/deployment-service.md index d073c413..4932191d 100644 --- a/docs/content/services/deployment-service.md +++ b/docs/content/services/deployment-service.md @@ -121,6 +121,18 @@ CompletableFuture 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 diff --git a/docs/content/services/model-gateway/chat.md b/docs/content/services/model-gateway/chat.md index 39937125..bfa23850 100644 --- a/docs/content/services/model-gateway/chat.md +++ b/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 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 diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatProvider.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatProvider.java index d84854fd..21669b53 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatProvider.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatProvider.java @@ -32,6 +32,9 @@ public interface ChatProvider *

* This method initiates an asynchronous chat operation where partial responses are delivered incrementally through the provided * {@link ChatHandler}. + *

+ * 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 diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatService.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatService.java index b9050190..5926d99d 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/ChatService.java +++ b/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. + *

+ * 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 diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/DefaultRestClient.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/DefaultRestClient.java index 8606a28b..d32e0105 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/DefaultRestClient.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/DefaultRestClient.java @@ -88,17 +88,30 @@ public CompletableFuture 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; } diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecorator.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecorator.java index c15ca32f..4ba92ee9 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecorator.java +++ b/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 implements ChatHand private final AtomicReference> 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. + *

+ * 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. *

@@ -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) { diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/streaming/DefaultChatSubscriber.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/streaming/DefaultChatSubscriber.java index 73ee243e..cecf97e8 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/chat/streaming/DefaultChatSubscriber.java +++ b/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 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. + *

+ * 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 onComplete() { + if (isCancelled()) + return CompletableFuture.completedFuture(null); + return awaitCallbacks() .thenCompose(completeToolCalls -> { var response = processor.buildResponse(); @@ -84,11 +121,24 @@ public Flow.Subscriber 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> 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(); } /** diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DefaultRestClient.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DefaultRestClient.java index 285e37fe..44dc4cf8 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DefaultRestClient.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DefaultRestClient.java @@ -19,6 +19,8 @@ import java.util.concurrent.CompletableFuture; 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.ChatClientContext; import com.ibm.watsonx.ai.chat.ChatHandler; import com.ibm.watsonx.ai.chat.ChatResponse; @@ -134,12 +136,31 @@ public CompletableFuture generateStreaming( if (nonNull(transactionId)) httpRequest.header(TRANSACTION_ID_HEADER, transactionId); - var subscriber = textGenerationSubscriber(handler); - return asyncHttpClient.send(httpRequest.build(), responseInfo -> logResponses + var response = new CompletableFuture(); + var subscriber = new CancellableTextGenerationSubscriber(handler); + var httpFuture = asyncHttpClient.send(httpRequest.build(), responseInfo -> logResponses ? BodySubscribers.fromLineSubscriber(new SseEventLogger(subscriber, responseInfo.statusCode(), responseInfo.headers())) - : BodySubscribers.fromLineSubscriber(subscriber)) - .thenAccept(r -> {}) - .exceptionally(t -> TextGenerationSubscriber.handleError(t, handler)); + : BodySubscribers.fromLineSubscriber(subscriber)); + + httpFuture + .thenAccept(r -> response.complete(null)) + .exceptionally(t -> { + if (subscriber.isCancelled()) + return null; + + TextGenerationSubscriber.handleError(t, handler); + response.completeExceptionally(nonNull(t.getCause()) ? t.getCause() : t); + return null; + }); + + response.whenComplete((r, t) -> { + if (response.isCancelled()) { + subscriber.cancelStream(); + httpFuture.cancel(true); + } + }); + + return response; } @Override @@ -197,17 +218,30 @@ public CompletableFuture 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; } @@ -237,53 +271,94 @@ public ForecastResponse forecast(String transactionId, String deploymentId, Dura } /** - * Creates a subscriber that listens to raw SSE messages from the chat stream, and delegates processing to a {@link TextGenerationSubscriber}. - * - * @param handler the handler that receives processed chat events - * @return a {@link Flow.Subscriber} suitable for consumption by the HTTP client + * A subscriber of raw SSE messages that delegates processing to a {@link TextGenerationSubscriber} and can be stopped through + * {@link #cancelStream()}. */ - private Flow.Subscriber textGenerationSubscriber(TextGenerationHandler handler) { + private static final class CancellableTextGenerationSubscriber implements Flow.Subscriber { + + private final TextGenerationHandler handler; + private final TextGenerationSubscriber chatSubscriber; + private final AtomicBoolean cancelled = new AtomicBoolean(false); + private final AtomicReference subscriptionRef = new AtomicReference<>(); + private Flow.Subscription subscription; + private volatile boolean success = true; + + CancellableTextGenerationSubscriber(TextGenerationHandler handler) { + this.handler = handler; + this.chatSubscriber = TextGenerationSubscriber.createSubscriber(handler); + } - return new Flow.Subscriber() { - private Flow.Subscription subscription; - private volatile boolean success = true; - private volatile TextGenerationSubscriber chatSubscriber = TextGenerationSubscriber.createSubscriber(handler); + /** + * Stops the stream: no further callback is delivered to the handler and the body subscription is cancelled. + */ + void cancelStream() { + cancelled.set(true); - @Override - public void onSubscribe(Subscription subscription) { - this.subscription = subscription; - this.subscription.request(1); - } + var subscription = subscriptionRef.get(); + if (nonNull(subscription)) + subscription.cancel(); + } - @Override - public void onNext(String partialMessage) { - try { + /** + * Returns whether the stream has been cancelled. + */ + boolean isCancelled() { + return cancelled.get(); + } - chatSubscriber.onNext(partialMessage); + @Override + public void onSubscribe(Subscription subscription) { + this.subscription = subscription; + subscriptionRef.set(subscription); - } catch (RuntimeException e) { + if (isCancelled()) { + subscription.cancel(); + return; + } - onError(e); - success = !handler.failOnFirstError(); + this.subscription.request(1); + } - } finally { - if (success) - subscription.request(1); - else - subscription.cancel(); - } - } + @Override + public void onNext(String partialMessage) { - @Override - public void onError(Throwable throwable) { - chatSubscriber.onError(throwable); + if (isCancelled()) { + subscription.cancel(); + return; } - @Override - public void onComplete() { - chatSubscriber.onComplete(); + try { + + chatSubscriber.onNext(partialMessage); + + } catch (RuntimeException e) { + + onError(e); + success = !handler.failOnFirstError(); + + } finally { + if (success && !isCancelled()) + subscription.request(1); + else + subscription.cancel(); } - }; + } + + @Override + public void onError(Throwable throwable) { + if (isCancelled()) + return; + + chatSubscriber.onError(throwable); + } + + @Override + public void onComplete() { + if (isCancelled()) + return; + + chatSubscriber.onComplete(); + } } /** diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DeploymentService.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DeploymentService.java index 267a7d11..cf820edd 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DeploymentService.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/deployment/DeploymentService.java @@ -220,6 +220,9 @@ public void onPartialResponse(String partialResponse, PartialChatResponse partia /** * Sends a streaming chat request to a deployment. + *

+ * 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 DeploymentChatRequest} * @param handler a {@link ChatHandler} implementation that receives partial responses, the complete response, and error notifications diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/DefaultRestClient.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/DefaultRestClient.java index 9dcadc47..e8745b90 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/DefaultRestClient.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/DefaultRestClient.java @@ -95,17 +95,30 @@ public CompletableFuture 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; } diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/ModelGatewayChatService.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/ModelGatewayChatService.java index 20addb1d..5eae6207 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/ModelGatewayChatService.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/gateway/chat/ModelGatewayChatService.java @@ -119,6 +119,9 @@ public ModelGatewayChatResponse chat(ModelGatewayChatRequest chatRequest) { /** * Sends a streaming chat request to the Model Gateway. + *

+ * 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 ModelGatewayChatRequest} * @param handler a {@link ChatHandler} implementation that receives partial responses, the complete response, and error notifications diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/DefaultRestClient.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/DefaultRestClient.java index 2a31ff72..c7c017bc 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/DefaultRestClient.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/DefaultRestClient.java @@ -20,6 +20,8 @@ import java.util.concurrent.CompletableFuture; 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.core.SseEventLogger; import com.ibm.watsonx.ai.core.factory.HttpClientFactory; import com.ibm.watsonx.ai.core.http.AsyncHttpClient; @@ -81,62 +83,122 @@ public CompletableFuture generateStreaming(String transactionId, TextReque if (nonNull(transactionId)) httpRequest.header(TRANSACTION_ID_HEADER, transactionId); - var subscriber = subscriber(handler); - return asyncHttpClient.send(httpRequest.build(), responseInfo -> logResponses + var response = new CompletableFuture(); + var subscriber = new CancellableSubscriber(handler); + var httpFuture = asyncHttpClient.send(httpRequest.build(), responseInfo -> logResponses ? BodySubscribers.fromLineSubscriber(new SseEventLogger(subscriber, responseInfo.statusCode(), responseInfo.headers())) - : BodySubscribers.fromLineSubscriber(subscriber)) - .thenAccept(r -> {}) - .exceptionally(t -> handleError(t, handler)); + : BodySubscribers.fromLineSubscriber(subscriber)); + + httpFuture + .thenAccept(r -> response.complete(null)) + .exceptionally(t -> { + if (subscriber.isCancelled()) + return null; + + handleError(t, handler); + response.completeExceptionally(nonNull(t.getCause()) ? t.getCause() : t); + return null; + }); + + response.whenComplete((r, t) -> { + if (response.isCancelled()) { + subscriber.cancelStream(); + httpFuture.cancel(true); + } + }); + + return response; } /** - * Creates a subscriber that listens to raw SSE messages from the chat stream, and delegates processing to a {@link TextGenerationSubscriber}. - * - * @param handler the handler that receives processed chat events - * @return a {@link Flow.Subscriber} suitable for consumption by the HTTP client + * A subscriber of raw SSE messages that delegates processing to a {@link TextGenerationSubscriber} and can be stopped through + * {@link #cancelStream()}. */ - private Flow.Subscriber subscriber(TextGenerationHandler handler) { + private static final class CancellableSubscriber implements Flow.Subscriber { + + private final TextGenerationHandler handler; + private final TextGenerationSubscriber chatSubscriber; + private final AtomicBoolean cancelled = new AtomicBoolean(false); + private final AtomicReference subscriptionRef = new AtomicReference<>(); + private Flow.Subscription subscription; + private volatile boolean success = true; + + CancellableSubscriber(TextGenerationHandler handler) { + this.handler = handler; + this.chatSubscriber = createSubscriber(handler); + } - return new Flow.Subscriber() { - private Flow.Subscription subscription; - private volatile boolean success = true; - private volatile TextGenerationSubscriber chatSubscriber = createSubscriber(handler); + /** + * Stops the stream: no further callback is delivered to the handler and the body subscription is cancelled. + */ + void cancelStream() { + cancelled.set(true); - @Override - public void onSubscribe(Subscription subscription) { - this.subscription = subscription; - this.subscription.request(1); - } + var subscription = subscriptionRef.get(); + if (nonNull(subscription)) + subscription.cancel(); + } + + /** + * Returns whether the stream has been cancelled. + */ + boolean isCancelled() { + return cancelled.get(); + } - @Override - public void onNext(String partialMessage) { - try { + @Override + public void onSubscribe(Subscription subscription) { + this.subscription = subscription; + subscriptionRef.set(subscription); - chatSubscriber.onNext(partialMessage); + if (isCancelled()) { + subscription.cancel(); + return; + } - } catch (RuntimeException e) { + this.subscription.request(1); + } - onError(e); - success = !handler.failOnFirstError(); + @Override + public void onNext(String partialMessage) { - } finally { - if (success) - subscription.request(1); - else - subscription.cancel(); - } + if (isCancelled()) { + subscription.cancel(); + return; } - @Override - public void onError(Throwable throwable) { - chatSubscriber.onError(throwable); - } + try { + + chatSubscriber.onNext(partialMessage); - @Override - public void onComplete() { - chatSubscriber.onComplete(); + } catch (RuntimeException e) { + + onError(e); + success = !handler.failOnFirstError(); + + } finally { + if (success && !isCancelled()) + subscription.request(1); + else + subscription.cancel(); } - }; + } + + @Override + public void onError(Throwable throwable) { + if (isCancelled()) + return; + + chatSubscriber.onError(throwable); + } + + @Override + public void onComplete() { + if (isCancelled()) + return; + + chatSubscriber.onComplete(); + } } /** diff --git a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/TextGenerationProvider.java b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/TextGenerationProvider.java index 319fc768..8b12e302 100644 --- a/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/TextGenerationProvider.java +++ b/modules/watsonx-ai/src/main/java/com/ibm/watsonx/ai/textgeneration/TextGenerationProvider.java @@ -28,6 +28,9 @@ public interface TextGenerationProvider { *

* This method initiates an asynchronous text generation operation where partial responses are delivered incrementally through the provided * {@link TextGenerationHandler}. + *

+ * 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 request the {@link TextGenerationRequest} containing input, moderation, parameters, and optional deployment ID * @param handler the handler that will receive streamed generation events diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/chat/ChatStreamingCancellationTest.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/chat/ChatStreamingCancellationTest.java new file mode 100644 index 00000000..4fd1173a --- /dev/null +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/chat/ChatStreamingCancellationTest.java @@ -0,0 +1,510 @@ +/* + * Copyright 2025 IBM Corporation + * SPDX-License-Identifier: Apache-2.0 + */ +package com.ibm.watsonx.ai.chat; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static com.github.tomakehurst.wiremock.client.WireMock.urlMatching; +import static java.util.concurrent.CompletableFuture.completedFuture; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; +import java.net.URI; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.Flow; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.parallel.Isolated; +import org.mockito.junit.jupiter.MockitoExtension; +import com.ibm.watsonx.ai.AbstractWatsonxTest; +import com.ibm.watsonx.ai.chat.decorator.ChatHandlerDecorator; +import com.ibm.watsonx.ai.chat.model.ChatMessage; +import com.ibm.watsonx.ai.chat.model.CompletedToolCall; +import com.ibm.watsonx.ai.chat.model.PartialChatResponse; +import com.ibm.watsonx.ai.chat.model.PartialToolCall; +import com.ibm.watsonx.ai.chat.model.UserMessage; +import com.ibm.watsonx.ai.chat.streaming.DefaultChatSubscriber; + +@ExtendWith(MockitoExtension.class) +@Isolated("Asserts the absence of callbacks within timing windows; must run without concurrent CPU contention.") +public class ChatStreamingCancellationTest extends AbstractWatsonxTest { + + /** + * Counts every callback and cancels the future once a given number of partial responses has been delivered. + */ + static final class Recorder implements ChatHandler { + + final CompletableFuture> future = new CompletableFuture<>(); + final AtomicInteger partialResponses = new AtomicInteger(); + final AtomicInteger partialToolCalls = new AtomicInteger(); + final AtomicInteger completeToolCalls = new AtomicInteger(); + final AtomicInteger completeResponses = new AtomicInteger(); + final AtomicInteger errors = new AtomicInteger(); + final CountDownLatch cancelled = new CountDownLatch(1); + + private final int cancelAfterPartialResponses; + private final int cancelAfterPartialToolCalls; + private final boolean mayInterruptIfRunning; + private final boolean failOnFirstError; + + Recorder(int cancelAfterPartialResponses) { + this(cancelAfterPartialResponses, 0, true, false); + } + + Recorder(int cancelAfterPartialResponses, int cancelAfterPartialToolCalls, boolean mayInterruptIfRunning, boolean failOnFirstError) { + this.cancelAfterPartialResponses = cancelAfterPartialResponses; + this.cancelAfterPartialToolCalls = cancelAfterPartialToolCalls; + this.mayInterruptIfRunning = mayInterruptIfRunning; + this.failOnFirstError = failOnFirstError; + } + + @Override + public void onPartialResponse(String partialResponse, PartialChatResponse partialChatResponse) { + if (partialResponses.incrementAndGet() == cancelAfterPartialResponses) + cancelNow(); + } + + @Override + public void onPartialToolCall(PartialToolCall partialToolCall) { + if (partialToolCalls.incrementAndGet() == cancelAfterPartialToolCalls) + cancelNow(); + } + + @Override + public void onCompleteToolCall(CompletedToolCall completeToolCall) { + completeToolCalls.incrementAndGet(); + } + + @Override + public void onCompleteResponse(ChatResponse completeResponse) { + completeResponses.incrementAndGet(); + } + + @Override + public void onError(Throwable error) { + errors.incrementAndGet(); + } + + @Override + public boolean failOnFirstError() { + return failOnFirstError; + } + + /** + * Cancels from inside the callback, which is the pattern a caller uses to stop as soon as a condition is met. + */ + private void cancelNow() { + future.join().cancel(mayInterruptIfRunning); + cancelled.countDown(); + } + } + + /** + * Records the interactions of the body subscription so that the cancellation mechanism itself can be asserted. + */ + static final class RecordingSubscription implements Flow.Subscription { + + final AtomicInteger requests = new AtomicInteger(); + final AtomicInteger cancellations = new AtomicInteger(); + + @Override + public void request(long n) { + requests.incrementAndGet(); + } + + @Override + public void cancel() { + cancellations.incrementAndGet(); + } + } + + static String contentEvent(int id, String content) { + return """ + id: %d + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[{"index":0,"finish_reason":null,"delta":{"content":"%s"}}],"created":1749736055,"created_at":"2025-06-12T13:47:35.542Z"} + + """ + .formatted(id, content); + } + + /** + * Builds an SSE body carrying the given number of content deltas, followed by the stop and usage events. + */ + static String contentBody(int deltas) { + var body = new StringBuilder(); + + for (int i = 1; i <= deltas; i++) + body.append(contentEvent(i, "chunk" + i)); + + body.append( + """ + id: %d + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[{"index":0,"finish_reason":"stop","delta":{"content":""}}],"created":1749736055,"created_at":"2025-06-12T13:47:35.563Z"} + + id: %d + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[],"created":1749736055,"created_at":"2025-06-12T13:47:35.564Z","usage":{"completion_tokens":3,"prompt_tokens":38,"total_tokens":41}} + + """ + .formatted(deltas + 1, deltas + 2)); + + return body.toString(); + } + + /** + * Builds an SSE body that streams the fragments of a single tool call, so that a cancellation can land while it is being assembled. + */ + static String toolCallBody(int fragments) { + var body = new StringBuilder(); + + body.append( + """ + id: 1 + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[{"index":0,"finish_reason":null,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"sum","arguments":""}}]}}],"created":1749764735,"created_at":"2025-06-12T21:45:35.348Z"} + + """); + + for (int i = 1; i <= fragments; i++) + body.append( + """ + id: %d + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[{"index":0,"finish_reason":null,"delta":{"tool_calls":[{"index":0,"function":{"name":"","arguments":"a"}}]}}],"created":1749764735,"created_at":"2025-06-12T21:45:35.357Z"} + + """ + .formatted(i + 1)); + + body.append( + """ + id: %d + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[{"index":0,"finish_reason":"tool_calls","delta":{"content":""}}],"created":1749764735,"created_at":"2025-06-12T21:45:35.555Z"} + + """ + .formatted(fragments + 2)); + + return body.toString(); + } + + void stubChatStream(String body, int chunks, int totalMillis) { + wireMock.stubFor(post(urlMatching("/ml/v1/text/chat_stream.*")) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withChunkedDribbleDelay(chunks, totalMillis) + .withBody(body))); + } + + ChatService chatService(boolean logResponses) { + return ChatService.builder() + .authenticator(mockAuthenticator) + .modelId("m") + .projectId("project-id") + .baseUrl(URI.create("http://localhost:%s".formatted(wireMock.getPort()))) + .version(API_VERSION) + .logResponses(logResponses) + .build(); + } + + List messages() { + return List.of(UserMessage.text("Tell me a long story")); + } + + /** + * Starts a stream through the public service API and returns the future the caller would cancel. + */ + CompletableFuture start(Recorder recorder, boolean logResponses) { + when(mockAuthenticator.tokenAsync()).thenReturn(completedFuture("my-token")); + var future = chatService(logResponses).chatStreaming(ChatRequest.builder().messages(messages()).build(), recorder); + recorder.future.complete(future); + return future; + } + + // Test 1, 5 and 11: cancellation through the service API, triggered from inside a callback. + @Test + void should_stop_delivering_callbacks_after_cancel() throws Exception { + + stubChatStream(contentBody(12), 40, 4000); + + var recorder = new Recorder(2); + var future = start(recorder, false); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + // Well inside the remaining dribble window: further chunks would still be arriving if the stream had not been stopped. + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } + + // Test 2: cancelling before the first chunk reaches the subscriber. + @Test + void should_deliver_no_callback_when_cancelled_before_the_first_chunk() throws Exception { + + wireMock.stubFor(post(urlMatching("/ml/v1/text/chat_stream.*")) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withFixedDelay(800) + .withBody(contentBody(12)))); + + var recorder = new Recorder(0); + var future = start(recorder, false); + + future.cancel(true); + Thread.sleep(2000); + + assertEquals(0, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } + + // Test 3: cancelling after normal completion changes nothing. + @Test + void should_be_a_no_op_when_cancelled_after_completion() throws Exception { + + stubChatStream(contentBody(3), 4, 100); + + var recorder = new Recorder(0); + var future = start(recorder, false); + + var response = future.get(10, TimeUnit.SECONDS); + + assertFalse(future.cancel(true)); + assertFalse(future.isCancelled()); + assertTrue(future.isDone() && !future.isCompletedExceptionally()); + assertNotNull(response); + assertEquals("chunk1chunk2chunk3", response.toAssistantMessage().content()); + assertEquals(1, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + } + + // Test 4: two threads cancelling concurrently. + @Test + void should_take_effect_once_when_two_threads_cancel() throws Exception { + + stubChatStream(contentBody(12), 40, 4000); + + var recorder = new Recorder(0); + var future = start(recorder, false); + + var start = new CountDownLatch(1); + var outcomes = new AtomicInteger(); + var executor = Executors.newFixedThreadPool(2); + + try { + var first = CompletableFuture.runAsync(() -> { + awaitQuietly(start); + if (future.cancel(true)) + outcomes.incrementAndGet(); + }, executor); + + var second = CompletableFuture.runAsync(() -> { + awaitQuietly(start); + if (future.cancel(false)) + outcomes.incrementAndGet(); + }, executor); + + start.countDown(); + CompletableFuture.allOf(first, second).get(10, TimeUnit.SECONDS); + } finally { + executor.shutdownNow(); + } + + // cancel reports the state, not the transition: once the future is cancelled every caller gets true. + assertEquals(2, outcomes.get()); + assertTrue(future.isCancelled()); + + // Let a callback that was already running finish, then check the stream stays stopped for the rest of the window. + Thread.sleep(500); + var delivered = recorder.partialResponses.get(); + Thread.sleep(1000); + + assertEquals(delivered, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + } + + // Test 6: cancelling while a tool call is being assembled. + @Test + void should_not_deliver_the_complete_tool_call_when_cancelled_while_assembling() throws Exception { + + stubChatStream(toolCallBody(12), 40, 4000); + + var recorder = new Recorder(0, 2, true, false); + var future = start(recorder, false); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial tool call was never delivered"); + + Thread.sleep(1000); + + assertEquals(2, recorder.partialToolCalls.get()); + assertEquals(0, recorder.completeToolCalls.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } + + // Test 7: failOnFirstError does not change the cancellation behaviour. + @Test + void should_cancel_the_same_way_when_fail_on_first_error_is_enabled() throws Exception { + + stubChatStream(contentBody(12), 40, 4000); + + var recorder = new Recorder(2, 0, true, true); + var future = start(recorder, false); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } + + // Test 8: the SseEventLogger path behaves the same. Log output is not part of the contract, only handler behaviour. + @Test + void should_cancel_the_same_way_when_log_responses_is_enabled() throws Exception { + + stubChatStream(contentBody(12), 40, 4000); + + var recorder = new Recorder(2); + var future = start(recorder, true); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } + + // Test 12: cancel(false) is equivalent to cancel(true), because CompletableFuture ignores the interrupt flag. + @Test + void should_treat_cancel_false_like_cancel_true() throws Exception { + + stubChatStream(contentBody(12), 40, 4000); + + var recorder = new Recorder(2, 0, false, false); + var future = start(recorder, false); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } + + // Test 9: the body subscription is really cancelled, and no further element is requested. + @Test + void should_cancel_the_body_subscription() { + + var recorder = new Recorder(0); + var decorator = new ChatHandlerDecorator(recorder, null, null); + var subscriber = new DefaultChatSubscriber(new SseEventProcessor(null, null, TextChatResponse::builder), decorator); + var response = cancellableResponse(subscriber); + var flowSubscriber = subscriber.asFlowSubscriber(response, true); + var subscription = new RecordingSubscription(); + + flowSubscriber.onSubscribe(subscription); + flowSubscriber.onNext(dataLine("chunk1")); + + assertEquals(2, subscription.requests.get()); + assertEquals(0, subscription.cancellations.get()); + + response.cancel(true); + + assertTrue(subscription.cancellations.get() > 0, "the subscription was not cancelled"); + assertTrue(subscriber.isCancelled()); + + var requestsAtCancel = subscription.requests.get(); + flowSubscriber.onNext(dataLine("chunk2")); + + assertEquals(requestsAtCancel, subscription.requests.get()); + } + + // Test 10: the signals the JDK is still allowed to deliver after cancel are dropped. + @Test + void should_drop_the_signals_delivered_after_cancel() { + + var recorder = new Recorder(0); + var decorator = new ChatHandlerDecorator(recorder, null, null); + var subscriber = new DefaultChatSubscriber(new SseEventProcessor(null, null, TextChatResponse::builder), decorator); + var response = cancellableResponse(subscriber); + var flowSubscriber = subscriber.asFlowSubscriber(response, true); + + flowSubscriber.onSubscribe(new RecordingSubscription()); + flowSubscriber.onNext(dataLine("chunk1")); + + decorator.awaitCallbacks().join(); + assertEquals(1, recorder.partialResponses.get()); + + response.cancel(true); + + flowSubscriber.onNext(dataLine("chunk2")); + flowSubscriber.onNext(dataLine("chunk3")); + flowSubscriber.onError(new RuntimeException("boom")); + flowSubscriber.onComplete(); + + decorator.awaitCallbacks().join(); + + assertEquals(1, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(response.isCancelled()); + } + + /** + * Wires the cancellation hook the routes install, so that the subscriber can be driven by hand. + */ + static CompletableFuture cancellableResponse(DefaultChatSubscriber subscriber) { + var response = new CompletableFuture(); + + response.whenComplete((r, t) -> { + if (response.isCancelled()) + subscriber.cancelStream(); + }); + + return response; + } + + static String dataLine(String content) { + return "data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"model_id\":\"m\",\"model\":\"m\"," + + "\"choices\":[{\"index\":0,\"finish_reason\":null,\"delta\":{\"content\":\"%s\"}}],\"created\":1749736055," + + "\"created_at\":\"2025-06-12T13:47:35.542Z\"}".formatted(content); + } + + static void awaitQuietly(CountDownLatch latch) { + try { + latch.await(10, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } +} diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecoratorTest.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecoratorTest.java index 63bb1610..b7bf830b 100644 --- a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecoratorTest.java +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/chat/decorator/ChatHandlerDecoratorTest.java @@ -6,9 +6,12 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.ArrayList; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import com.ibm.watsonx.ai.chat.ChatHandler; import com.ibm.watsonx.ai.chat.ChatRequest; @@ -257,4 +260,65 @@ void should_return_tool_calls_in_receive_order() { var toolCalls = decorator.awaitCallbacks().join(); assertEquals(List.of(0, 1), toolCalls.stream().map(toolCall -> toolCall.toolCall().index()).toList()); } + + @Test + void should_take_effect_only_once_when_cancelled_more_than_once() { + + var decorator = new ChatHandlerDecorator(new Recorder(), null, null); + + assertTrue(decorator.cancel()); + assertFalse(decorator.cancel()); + assertTrue(decorator.isCancelled()); + } + + @Test + void should_deliver_no_callback_after_cancel() { + + var recorder = new Recorder(); + var decorator = new ChatHandlerDecorator(recorder, null, null); + + decorator.onPartialResponse("Hello", null); + decorator.awaitCallbacks().join(); + decorator.cancel(); + + emitTwoToolCalls(decorator); + + // The callback chain must keep advancing while cancelled, otherwise awaitCallbacks would never complete. + assertDoesNotThrow(() -> decorator.awaitCallbacks().join()); + + assertEquals(List.of("partialResponse"), recorder.events); + } + + @Test + void should_not_report_the_error_of_a_callback_that_throws_after_cancel() { + + var recorder = new Recorder(); + var decorator = new AtomicReference>(); + + decorator.set(new ChatHandlerDecorator(new ChatHandler() { + + @Override + public void onPartialResponse(String partialResponse, PartialChatResponse partialChatResponse) { + recorder.onPartialResponse(partialResponse, partialChatResponse); + decorator.get().cancel(); + throw new IllegalStateException("thrown by the handler"); + } + + @Override + public void onCompleteResponse(ChatResponse completeResponse) { + recorder.onCompleteResponse(completeResponse); + } + + @Override + public void onError(Throwable error) { + recorder.onError(error); + } + }, null, null)); + + decorator.get().onPartialResponse("Hello", null); + decorator.get().onCompleteResponse(null); + decorator.get().awaitCallbacks().join(); + + assertEquals(List.of("partialResponse"), recorder.events); + } } diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/deployment/DeploymentStreamingCancellationTest.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/deployment/DeploymentStreamingCancellationTest.java new file mode 100644 index 00000000..dc7e7849 --- /dev/null +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/deployment/DeploymentStreamingCancellationTest.java @@ -0,0 +1,263 @@ +/* + * Copyright 2025 IBM Corporation + * SPDX-License-Identifier: Apache-2.0 + */ +package com.ibm.watsonx.ai.deployment; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static java.util.concurrent.CompletableFuture.completedFuture; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; +import java.net.URI; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.parallel.Isolated; +import org.mockito.junit.jupiter.MockitoExtension; +import com.ibm.watsonx.ai.AbstractWatsonxTest; +import com.ibm.watsonx.ai.chat.ChatHandler; +import com.ibm.watsonx.ai.chat.ChatResponse; +import com.ibm.watsonx.ai.chat.model.PartialChatResponse; +import com.ibm.watsonx.ai.chat.model.UserMessage; +import com.ibm.watsonx.ai.textgeneration.TextGenerationHandler; +import com.ibm.watsonx.ai.textgeneration.TextGenerationRequest; +import com.ibm.watsonx.ai.textgeneration.TextGenerationResponse; + +@ExtendWith(MockitoExtension.class) +@Isolated("Asserts the absence of callbacks within timing windows; must run without concurrent CPU contention.") +public class DeploymentStreamingCancellationTest extends AbstractWatsonxTest { + + static final String DEPLOYMENT_ID = "deployment-id"; + + /** + * Counts every callback and cancels the future once a given number of partial responses has been delivered. + */ + static final class ChatRecorder implements ChatHandler { + + final CompletableFuture> future = new CompletableFuture<>(); + final AtomicInteger partialResponses = new AtomicInteger(); + final AtomicInteger completeResponses = new AtomicInteger(); + final AtomicInteger errors = new AtomicInteger(); + final CountDownLatch cancelled = new CountDownLatch(1); + + private final int cancelAfterPartialResponses; + + ChatRecorder(int cancelAfterPartialResponses) { + this.cancelAfterPartialResponses = cancelAfterPartialResponses; + } + + @Override + public void onPartialResponse(String partialResponse, PartialChatResponse partialChatResponse) { + if (partialResponses.incrementAndGet() == cancelAfterPartialResponses) { + future.join().cancel(true); + cancelled.countDown(); + } + } + + @Override + public void onCompleteResponse(ChatResponse completeResponse) { + completeResponses.incrementAndGet(); + } + + @Override + public void onError(Throwable error) { + errors.incrementAndGet(); + } + } + + /** + * Counts every callback and cancels the future once a given number of partial responses has been delivered. + */ + static final class TextGenerationRecorder implements TextGenerationHandler { + + final CompletableFuture> future = new CompletableFuture<>(); + final AtomicInteger partialResponses = new AtomicInteger(); + final AtomicInteger completeResponses = new AtomicInteger(); + final AtomicInteger errors = new AtomicInteger(); + final CountDownLatch cancelled = new CountDownLatch(1); + + private final int cancelAfterPartialResponses; + + TextGenerationRecorder(int cancelAfterPartialResponses) { + this.cancelAfterPartialResponses = cancelAfterPartialResponses; + } + + @Override + public void onPartialResponse(String partialResponse) { + if (partialResponses.incrementAndGet() == cancelAfterPartialResponses) { + future.join().cancel(true); + cancelled.countDown(); + } + } + + @Override + public void onCompleteResponse(TextGenerationResponse completeResponse) { + completeResponses.incrementAndGet(); + } + + @Override + public void onError(Throwable error) { + errors.incrementAndGet(); + } + } + + static String chatBody(int deltas) { + var body = new StringBuilder(); + + for (int i = 1; i <= deltas; i++) + body.append( + """ + id: %d + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[{"index":0,"finish_reason":null,"delta":{"content":"chunk%d"}}],"created":1749736055,"created_at":"2025-06-12T13:47:35.542Z"} + + """ + .formatted(i, i)); + + body.append( + """ + id: %d + event: message + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model_id":"m","model":"m","choices":[{"index":0,"finish_reason":"stop","delta":{"content":""}}],"created":1749736055,"created_at":"2025-06-12T13:47:35.563Z"} + + """ + .formatted(deltas + 1)); + + return body.toString(); + } + + static String generationBody(int deltas) { + var body = new StringBuilder(); + + for (int i = 1; i <= deltas; i++) + body.append( + """ + id: %d + event: message + data: {"model_id":"m","created_at":"2025-06-24T15:30:13.552Z","results":[{"generated_text":"chunk%d","generated_token_count":%d,"input_token_count":0,"stop_reason":"not_finished"}]} + + """ + .formatted(i, i, i)); + + body.append( + """ + id: %d + event: message + data: {"model_id":"m","created_at":"2025-06-24T15:30:13.588Z","results":[{"generated_text":".","generated_token_count":%d,"input_token_count":0,"stop_reason":"eos_token"}]} + + """ + .formatted(deltas + 1, deltas + 1)); + + return body.toString(); + } + + DeploymentService service() { + return DeploymentService.builder() + .authenticator(mockAuthenticator) + .baseUrl(URI.create("http://localhost:%s".formatted(wireMock.getPort()))) + .version(API_VERSION) + .build(); + } + + @Test + void should_stop_delivering_chat_callbacks_after_cancel() throws Exception { + + wireMock.stubFor(post("/ml/v1/deployments/%s/text/chat_stream?version=%s".formatted(DEPLOYMENT_ID, API_VERSION)) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withChunkedDribbleDelay(40, 4000) + .withBody(chatBody(12)))); + + when(mockAuthenticator.tokenAsync()).thenReturn(completedFuture("my-token")); + + var recorder = new ChatRecorder(2); + var request = DeploymentChatRequest.builder() + .deploymentId(DEPLOYMENT_ID) + .messages(UserMessage.text("Tell me a long story")) + .build(); + + var future = service().chatStreaming(request, recorder); + recorder.future.complete(future); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + // Well inside the remaining dribble window: further chunks would still be arriving if the stream had not been stopped. + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } + + @Test + void should_deliver_no_chat_callback_when_cancelled_before_the_first_chunk() throws Exception { + + wireMock.stubFor(post("/ml/v1/deployments/%s/text/chat_stream?version=%s".formatted(DEPLOYMENT_ID, API_VERSION)) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withFixedDelay(800) + .withBody(chatBody(12)))); + + when(mockAuthenticator.tokenAsync()).thenReturn(completedFuture("my-token")); + + var recorder = new ChatRecorder(0); + var request = DeploymentChatRequest.builder() + .deploymentId(DEPLOYMENT_ID) + .messages(UserMessage.text("Tell me a long story")) + .build(); + + var future = service().chatStreaming(request, recorder); + recorder.future.complete(future); + + future.cancel(true); + Thread.sleep(2000); + + assertEquals(0, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } + + @Test + void should_stop_delivering_text_generation_callbacks_after_cancel() throws Exception { + + wireMock.stubFor(post("/ml/v1/deployments/%s/text/generation_stream?version=%s".formatted(DEPLOYMENT_ID, API_VERSION)) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withChunkedDribbleDelay(40, 4000) + .withBody(generationBody(12)))); + + when(mockAuthenticator.tokenAsync()).thenReturn(completedFuture("my-token")); + + var recorder = new TextGenerationRecorder(2); + var request = TextGenerationRequest.builder() + .deploymentId(DEPLOYMENT_ID) + .input("Tell me a long story") + .build(); + + var future = service().generateStreaming(request, recorder); + recorder.future.complete(future); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } +} diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/gateway/chat/ModelGatewayChatStreamingCancellationTest.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/gateway/chat/ModelGatewayChatStreamingCancellationTest.java new file mode 100644 index 00000000..0e30b8d7 --- /dev/null +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/gateway/chat/ModelGatewayChatStreamingCancellationTest.java @@ -0,0 +1,155 @@ +/* + * Copyright 2025 IBM Corporation + * SPDX-License-Identifier: Apache-2.0 + */ +package com.ibm.watsonx.ai.gateway.chat; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static java.util.concurrent.CompletableFuture.completedFuture; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; +import java.net.URI; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.parallel.Isolated; +import org.mockito.junit.jupiter.MockitoExtension; +import com.ibm.watsonx.ai.AbstractWatsonxTest; +import com.ibm.watsonx.ai.chat.ChatHandler; +import com.ibm.watsonx.ai.chat.ChatResponse; +import com.ibm.watsonx.ai.chat.model.ChatMessage; +import com.ibm.watsonx.ai.chat.model.PartialChatResponse; +import com.ibm.watsonx.ai.chat.model.UserMessage; + +@ExtendWith(MockitoExtension.class) +@Isolated("Asserts the absence of callbacks within timing windows; must run without concurrent CPU contention.") +public class ModelGatewayChatStreamingCancellationTest extends AbstractWatsonxTest { + + /** + * Counts every callback and cancels the future once a given number of partial responses has been delivered. + */ + static final class Recorder implements ChatHandler { + + final CompletableFuture> future = new CompletableFuture<>(); + final AtomicInteger partialResponses = new AtomicInteger(); + final AtomicInteger completeResponses = new AtomicInteger(); + final AtomicInteger errors = new AtomicInteger(); + final CountDownLatch cancelled = new CountDownLatch(1); + + private final int cancelAfterPartialResponses; + + Recorder(int cancelAfterPartialResponses) { + this.cancelAfterPartialResponses = cancelAfterPartialResponses; + } + + @Override + public void onPartialResponse(String partialResponse, PartialChatResponse partialChatResponse) { + if (partialResponses.incrementAndGet() == cancelAfterPartialResponses) { + future.join().cancel(true); + cancelled.countDown(); + } + } + + @Override + public void onCompleteResponse(ChatResponse completeResponse) { + completeResponses.incrementAndGet(); + } + + @Override + public void onError(Throwable error) { + errors.incrementAndGet(); + } + } + + static String body(int deltas) { + var body = new StringBuilder(); + + for (int i = 1; i <= deltas; i++) + body.append( + """ + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"chunk%d"},"finish_reason":"","logprobs":null}],"created":1749736055,"model":"gpt-4o","usage":null,"cached":false} + + """ + .formatted(i)); + + body.append( + """ + data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop","logprobs":null}],"created":1749736055,"model":"gpt-4o","usage":{"prompt_tokens":38,"completion_tokens":3,"total_tokens":41},"cached":false} + + data: [DONE] + """); + + return body.toString(); + } + + ModelGatewayChatService service() { + return ModelGatewayChatService.builder() + .authenticator(mockAuthenticator) + .modelId("gpt-4o") + .baseUrl(URI.create("http://localhost:%s".formatted(wireMock.getPort()))) + .version(API_VERSION) + .build(); + } + + CompletableFuture start(Recorder recorder) { + when(mockAuthenticator.tokenAsync()).thenReturn(completedFuture("my-token")); + var future = service().chatStreaming(List.of(UserMessage.text("Tell me a long story")), recorder); + recorder.future.complete(future); + return future; + } + + @Test + void should_stop_delivering_callbacks_after_cancel() throws Exception { + + wireMock.stubFor(post("/ml/gateway/v1/chat/completions?version=%s".formatted(API_VERSION)) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withChunkedDribbleDelay(40, 4000) + .withBody(body(12)))); + + var recorder = new Recorder(2); + var future = start(recorder); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + // Well inside the remaining dribble window: further chunks would still be arriving if the stream had not been stopped. + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } + + @Test + void should_deliver_no_callback_when_cancelled_before_the_first_chunk() throws Exception { + + wireMock.stubFor(post("/ml/gateway/v1/chat/completions?version=%s".formatted(API_VERSION)) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withFixedDelay(800) + .withBody(body(12)))); + + var recorder = new Recorder(0); + var future = start(recorder); + + future.cancel(true); + Thread.sleep(2000); + + assertEquals(0, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } +} diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ChatServiceIT.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ChatServiceIT.java index 8f7916ef..3f6bb446 100644 --- a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ChatServiceIT.java +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ChatServiceIT.java @@ -20,6 +20,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; @@ -27,6 +28,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.IntStream; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; @@ -1743,5 +1745,92 @@ public void onError(Throwable error) { assertTrue(piiFlaggedInChunk.get(), "expected at least one streaming chunk with a PII moderation match"); assertTrue(hapFlaggedInChunk.get(), "expected at least one streaming chunk with a HAP moderation match"); } + + @Test + void should_stop_the_stream_when_the_returned_future_is_cancelled() throws Exception { + + var chatService = ChatService.builder() + .baseUrl(URL) + .projectId(PROJECT_ID) + .modelId("mistralai/mistral-small-3-1-24b-instruct-2503") + .authenticator(authentication) + .logRequests(true) + .logResponses(true) + .build(); + + var partialResponses = new AtomicInteger(); + var terminalCallbacks = new AtomicInteger(); + var firstPartialResponse = new CountDownLatch(1); + + var future = chatService.chatStreaming(createLongChatRequest(), new ChatHandler() { + + @Override + public void onPartialResponse(String partialResponse, PartialChatResponse partialChatResponse) { + partialResponses.incrementAndGet(); + firstPartialResponse.countDown(); + } + + @Override + public void onCompleteResponse(ChatResponse completeResponse) { + terminalCallbacks.incrementAndGet(); + } + + @Override + public void onError(Throwable error) { + terminalCallbacks.incrementAndGet(); + } + }); + + assertTrue(firstPartialResponse.await(60, TimeUnit.SECONDS), "the model never started streaming"); + assertTrue(future.cancel(true)); + + Thread.sleep(500); + var deliveredAtCancel = partialResponses.get(); + Thread.sleep(3000); + + assertEquals(deliveredAtCancel, partialResponses.get()); + assertEquals(0, terminalCallbacks.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } + + @Test + void should_keep_the_service_usable_after_a_cancelled_stream() throws Exception { + + var chatService = ChatService.builder() + .baseUrl(URL) + .projectId(PROJECT_ID) + .modelId("mistralai/mistral-small-3-1-24b-instruct-2503") + .authenticator(authentication) + .logRequests(true) + .logResponses(true) + .build(); + + var firstPartialResponse = new CountDownLatch(1); + var cancelled = chatService.chatStreaming(createLongChatRequest(), + (partialResponse, partialChatResponse) -> firstPartialResponse.countDown()); + + assertTrue(firstPartialResponse.await(60, TimeUnit.SECONDS), "the model never started streaming"); + assertTrue(cancelled.cancel(true)); + + var chatResponse = chatService.chatStreaming("Hello!", (partialResponse, partialChatResponse) -> {}) + .get(60, TimeUnit.SECONDS); + + assertNotNull(chatResponse.toAssistantMessage().content()); + assertFalse(chatResponse.toAssistantMessage().content().isBlank()); + } + + private ChatRequest createLongChatRequest() { + + var parameters = ChatParameters.builder() + .temperature(0.0) + .maxCompletionTokens(1000) + .build(); + + return ChatRequest.builder() + .messages(UserMessage.text("Count from 1 to 300, one number per line, without any other text.")) + .parameters(parameters) + .build(); + } } } diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/DeploymentServiceIT.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/DeploymentServiceIT.java index 0e82a067..2e2fc166 100644 --- a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/DeploymentServiceIT.java +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/DeploymentServiceIT.java @@ -18,11 +18,13 @@ import java.util.HashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.IntStream; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; @@ -1160,5 +1162,65 @@ private DeploymentChatRequest createChatRequest() { .parameters(parameters) .build(); } + + @Test + void should_stop_the_stream_when_the_returned_future_is_cancelled() throws Exception { + + var chatService = DeploymentService.builder() + .baseUrl(URL) + .authenticator(authentication) + .logRequests(true) + .logResponses(true) + .build(); + + var partialResponses = new AtomicInteger(); + var terminalCallbacks = new AtomicInteger(); + var firstPartialResponse = new CountDownLatch(1); + + var future = chatService.chatStreaming(createLongChatRequest(), new ChatHandler() { + + @Override + public void onPartialResponse(String partialResponse, PartialChatResponse partialChatResponse) { + partialResponses.incrementAndGet(); + firstPartialResponse.countDown(); + } + + @Override + public void onCompleteResponse(ChatResponse completeResponse) { + terminalCallbacks.incrementAndGet(); + } + + @Override + public void onError(Throwable error) { + terminalCallbacks.incrementAndGet(); + } + }); + + assertTrue(firstPartialResponse.await(60, TimeUnit.SECONDS), "the model never started streaming"); + assertTrue(future.cancel(true)); + + Thread.sleep(500); + var deliveredAtCancel = partialResponses.get(); + Thread.sleep(3000); + + assertEquals(deliveredAtCancel, partialResponses.get()); + assertEquals(0, terminalCallbacks.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } + + private DeploymentChatRequest createLongChatRequest() { + + var parameters = ChatParameters.builder() + .temperature(0.0) + .maxCompletionTokens(400) + .build(); + + return DeploymentChatRequest.builder() + .messages(UserMessage.text("Count from 1 to 300, one number per line, without any other text.")) + .deploymentId(DEPLOYMENT_ID) + .parameters(parameters) + .build(); + } } } diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ModelGatewayChatServiceIT.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ModelGatewayChatServiceIT.java index c681e6d0..43c77635 100644 --- a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ModelGatewayChatServiceIT.java +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/it/ModelGatewayChatServiceIT.java @@ -20,6 +20,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; @@ -1413,5 +1414,57 @@ public void onCompleteToolCall(CompletedToolCall completeToolCall) { assertEquals(3, assistantMessage.toolCalls().size()); assertTrue(countries.contains("Germany") && countries.contains("Italy") && countries.contains("Japan")); } + + @Test + void should_stop_the_stream_when_the_returned_future_is_cancelled() throws Exception { + + var modelGatewayChatService = ModelGatewayChatService.builder() + .baseUrl(URL) + .modelId(CHAT_MODEL_CLAUDE) + .authenticator(authentication) + .logRequests(true) + .logResponses(true) + .build(); + + var chatRequest = ModelGatewayChatRequest.builder() + .messages(UserMessage.text("Count from 1 to 300, one number per line, without any other text.")) + .parameters(ModelGatewayChatParameters.builder().maxTokens(1000).build()) + .build(); + + var partialResponses = new AtomicInteger(); + var terminalCallbacks = new AtomicInteger(); + var firstPartialResponse = new CountDownLatch(1); + + var future = modelGatewayChatService.chatStreaming(chatRequest, new ChatHandler() { + + @Override + public void onPartialResponse(String partialResponse, PartialChatResponse partialChatResponse) { + partialResponses.incrementAndGet(); + firstPartialResponse.countDown(); + } + + @Override + public void onCompleteResponse(ChatResponse completeResponse) { + terminalCallbacks.incrementAndGet(); + } + + @Override + public void onError(Throwable error) { + terminalCallbacks.incrementAndGet(); + } + }); + + assertTrue(firstPartialResponse.await(60, TimeUnit.SECONDS), "the model never started streaming"); + assertTrue(future.cancel(true)); + + Thread.sleep(500); + var deliveredAtCancel = partialResponses.get(); + Thread.sleep(3000); + + assertEquals(deliveredAtCancel, partialResponses.get()); + assertEquals(0, terminalCallbacks.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } } } diff --git a/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textgeneration/TextGenerationStreamingCancellationTest.java b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textgeneration/TextGenerationStreamingCancellationTest.java new file mode 100644 index 00000000..0fa1049f --- /dev/null +++ b/modules/watsonx-ai/src/test/java/com/ibm/watsonx/ai/textgeneration/TextGenerationStreamingCancellationTest.java @@ -0,0 +1,212 @@ +/* + * Copyright 2025 IBM Corporation + * SPDX-License-Identifier: Apache-2.0 + */ +package com.ibm.watsonx.ai.textgeneration; + +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.post; +import static java.util.concurrent.CompletableFuture.completedFuture; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; +import java.net.URI; +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 org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.parallel.Isolated; +import org.mockito.junit.jupiter.MockitoExtension; +import com.ibm.watsonx.ai.AbstractWatsonxTest; + +@ExtendWith(MockitoExtension.class) +@Isolated("Asserts the absence of callbacks within timing windows; must run without concurrent CPU contention.") +public class TextGenerationStreamingCancellationTest extends AbstractWatsonxTest { + + /** + * Counts every callback and cancels the future once a given number of partial responses has been delivered. + */ + static final class Recorder implements TextGenerationHandler { + + final CompletableFuture> future = new CompletableFuture<>(); + final AtomicInteger partialResponses = new AtomicInteger(); + final AtomicInteger completeResponses = new AtomicInteger(); + final AtomicInteger errors = new AtomicInteger(); + final CountDownLatch cancelled = new CountDownLatch(1); + final CountDownLatch terminated = new CountDownLatch(1); + + private final int cancelAfterPartialResponses; + + Recorder(int cancelAfterPartialResponses) { + this.cancelAfterPartialResponses = cancelAfterPartialResponses; + } + + @Override + public void onPartialResponse(String partialResponse) { + if (partialResponses.incrementAndGet() == cancelAfterPartialResponses) { + future.join().cancel(true); + cancelled.countDown(); + } + } + + @Override + public void onCompleteResponse(TextGenerationResponse completeResponse) { + completeResponses.incrementAndGet(); + terminated.countDown(); + } + + @Override + public void onError(Throwable error) { + errors.incrementAndGet(); + terminated.countDown(); + } + } + + static String body(int deltas) { + var body = new StringBuilder(); + + for (int i = 1; i <= deltas; i++) + body.append( + """ + id: %d + event: message + data: {"model_id":"m","created_at":"2025-06-24T15:30:13.552Z","results":[{"generated_text":"chunk%d","generated_token_count":%d,"input_token_count":0,"stop_reason":"not_finished"}]} + + """ + .formatted(i, i, i)); + + body.append( + """ + id: %d + event: message + data: {"model_id":"m","created_at":"2025-06-24T15:30:13.588Z","results":[{"generated_text":".","generated_token_count":%d,"input_token_count":0,"stop_reason":"eos_token"}]} + + """ + .formatted(deltas + 1, deltas + 1)); + + return body.toString(); + } + + void stubGenerationStream(String body, int chunks, int totalMillis) { + wireMock.stubFor(post("/ml/v1/text/generation_stream?version=%s".formatted(API_VERSION)) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withChunkedDribbleDelay(chunks, totalMillis) + .withBody(body))); + } + + TextGenerationService service() { + return TextGenerationService.builder() + .authenticator(mockAuthenticator) + .modelId("m") + .projectId("project-id") + .baseUrl(URI.create("http://localhost:%s".formatted(wireMock.getPort()))) + .version(API_VERSION) + .build(); + } + + CompletableFuture start(Recorder recorder) { + when(mockAuthenticator.tokenAsync()).thenReturn(completedFuture("my-token")); + var future = service().generateStreaming("Tell me a long story", recorder); + recorder.future.complete(future); + return future; + } + + // Test 15 and 18: cancellation through the service API, triggered from inside a callback. + @Test + void should_stop_delivering_callbacks_after_cancel() throws Exception { + + stubGenerationStream(body(12), 40, 4000); + + var recorder = new Recorder(2); + var future = start(recorder); + + assertTrue(recorder.cancelled.await(10, TimeUnit.SECONDS), "the second partial response was never delivered"); + + // Well inside the remaining dribble window: further chunks would still be arriving if the stream had not been stopped. + Thread.sleep(1000); + + assertEquals(2, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + assertThrows(CancellationException.class, future::join); + } + + // Test 16: cancelling before the first chunk reaches the subscriber. + @Test + void should_deliver_no_callback_when_cancelled_before_the_first_chunk() throws Exception { + + wireMock.stubFor(post("/ml/v1/text/generation_stream?version=%s".formatted(API_VERSION)) + .willReturn(aResponse() + .withStatus(200) + .withHeader("Content-Type", "text/event-stream") + .withFixedDelay(800) + .withBody(body(12)))); + + var recorder = new Recorder(0); + var future = start(recorder); + + future.cancel(true); + Thread.sleep(2000); + + assertEquals(0, recorder.partialResponses.get()); + assertEquals(0, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + assertTrue(future.isCancelled()); + } + + // Test 17: the future the caller now receives still tracks the outcome of a stream nobody cancels. + @Test + void should_complete_normally_at_the_end_of_the_stream() throws Exception { + + stubGenerationStream(body(3), 4, 100); + + var recorder = new Recorder(0); + var future = start(recorder); + + future.get(10, TimeUnit.SECONDS); + + assertTrue(recorder.terminated.await(10, TimeUnit.SECONDS), "the stream never terminated"); + assertTrue(future.isDone() && !future.isCompletedExceptionally()); + assertFalse(future.isCancelled()); + assertEquals(4, recorder.partialResponses.get()); + assertEquals(1, recorder.completeResponses.get()); + assertEquals(0, recorder.errors.get()); + } + + // Test 17: a failing stream is reported both to the handler and through the returned future. + @Test + void should_complete_exceptionally_when_the_stream_fails() throws Exception { + + wireMock.stubFor(post("/ml/v1/text/generation_stream?version=%s".formatted(API_VERSION)) + .willReturn(aResponse() + .withStatus(400) + .withHeader("Content-Type", "application/json") + .withBody(""" + { + "errors": [ { "code": "invalid_request_entity", "message": "input is required" } ], + "trace": "trace-id", + "status_code": 400 + } + """))); + + var recorder = new Recorder(0); + var future = start(recorder); + + assertThrows(ExecutionException.class, () -> future.get(10, TimeUnit.SECONDS)); + + assertTrue(recorder.terminated.await(10, TimeUnit.SECONDS), "the failure was never reported to the handler"); + assertTrue(future.isCompletedExceptionally()); + assertFalse(future.isCancelled()); + assertEquals(1, recorder.errors.get()); + assertEquals(0, recorder.completeResponses.get()); + } +}