[ML] Handle Errors and pre-streaming exceptions (#115868)

If we fail to establish a connection to bedrock, the error is returned in the client's CompletableFuture. We will forward it to the listener via the stream processor. Any Errors are thrown on another thread.

Author: prwhelan

docs/changelog/115868.yaml

@@ -0,0 +1,5 @@
+pr: 115868
+summary: Forward bedrock connection errors to user
+area: Machine Learning
+type: bug
+issues: []

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockInferenceClient.java

@@ -23,6 +23,7 @@
 import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelResponse;
 
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.SpecialPermission;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.xcontent.ChunkedToXContent;
@@ -93,11 +94,15 @@ public Flow.Publisher<? extends ChunkedToXContent> converseStream(ConverseStream
         internalClient.converseStream(
             request,
             ConverseStreamResponseHandler.builder().subscriber(() -> FlowAdapters.toSubscriber(awsResponseProcessor)).build()
-        );
+        ).exceptionally(e -> {
+            awsResponseProcessor.onError(e);
+            return null; // Void
+        });
         return awsResponseProcessor;
     }
 
     private void onFailure(ActionListener<?> listener, Throwable t, String method) {
+        ExceptionsHelper.maybeDieOnAnotherThread(t);
         var unwrappedException = t;
         if (t instanceof CompletionException || t instanceof ExecutionException) {
             unwrappedException = t.getCause() != null ? t.getCause() : t;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/amazonbedrock/AmazonBedrockStreamingChatProcessor.java

@@ -12,6 +12,7 @@
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseStreamResponseHandler;
 
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.logging.LogManager;
@@ -89,14 +90,15 @@ private void sendDownstreamOnAnotherThread(ContentBlockDeltaEvent event) {
 
     @Override
     public void onError(Throwable amazonBedrockRuntimeException) {
+        ExceptionsHelper.maybeDieOnAnotherThread(amazonBedrockRuntimeException);
         error.set(
             new ElasticsearchException(
                 Strings.format("AmazonBedrock StreamingChatProcessor failure: [%s]", amazonBedrockRuntimeException.getMessage()),
                 amazonBedrockRuntimeException
             )
         );
         if (isDone.compareAndSet(false, true) && checkAndResetDemand() && onErrorCalled.compareAndSet(false, true)) {
-            downstream.onError(error.get());
+            runOnUtilityThreadPool(() -> downstream.onError(amazonBedrockRuntimeException));
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/ServerSentEventsRestActionListener.java

@@ -223,7 +223,7 @@ public void onNext(ChunkedToXContent item) {
         @Override
         public void onError(Throwable throwable) {
             if (isLastPart.compareAndSet(false, true)) {
-                logger.error("A failure occurred in ElasticSearch while streaming the response.", throwable);
+                logger.warn("A failure occurred in ElasticSearch while streaming the response.", throwable);
                 nextBodyPartListener().onResponse(new ServerSentEventResponseBodyPart(ServerSentEvents.ERROR, errorChunk(throwable)));
             }
         }