[ML] Retry on streaming errors (elastic#123076)

We now always retry based on the provider's configured retry logic
rather than the HTTP status code. Some providers (e.g. Cohere,
Anthropic) will return 200 status codes with error bodies, others (e.g.
OpenAI, Azure) will return non-200 status codes with non-streaming
bodies.

Notes:
- Refactored from HttpResult to StreamingHttpResult, the byte body is
  now the streaming element while the http response lives outside the
  stream.
- Refactored StreamingHttpResultPublisher so that it only pushes byte
  body into a queue.
- Tests all now have to wait for the response to be fully consumed
  before closing the service, otherwise the close method will shut down
  the mock web server and apache will throw an error.

Author: prwhelan

docs/changelog/123076.yaml

@@ -0,0 +1,5 @@
+pr: 123076
+summary: Retry on streaming errors
+area: Machine Learning
+type: bug
+issues: []

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpClient.java

@@ -27,7 +27,6 @@
 import java.io.IOException;
 import java.util.Objects;
 import java.util.concurrent.CancellationException;
-import java.util.concurrent.Flow;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static org.elasticsearch.core.Strings.format;
@@ -154,15 +153,15 @@ private void failUsingUtilityThread(Exception exception, ActionListener<?> liste
         threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(() -> listener.onFailure(exception));
     }
 
-    public void stream(HttpRequest request, HttpContext context, ActionListener<Flow.Publisher<HttpResult>> listener) throws IOException {
+    public void stream(HttpRequest request, HttpContext context, ActionListener<StreamingHttpResult> listener) throws IOException {
         // The caller must call start() first before attempting to send a request
         assert status.get() == Status.STARTED : "call start() before attempting to send a request";
 
         var streamingProcessor = new StreamingHttpResultPublisher(threadPool, settings, listener);
 
         SocketAccess.doPrivileged(() -> client.execute(request.requestProducer(), streamingProcessor, context, new FutureCallback<>() {
             @Override
-            public void completed(HttpResponse response) {
+            public void completed(Void response) {
                 streamingProcessor.close();
             }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/StreamingHttpResult.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.http;
+
+import org.apache.http.HttpResponse;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.rest.RestStatus;
+
+import java.io.ByteArrayOutputStream;
+import java.util.concurrent.Flow;
+import java.util.concurrent.atomic.AtomicReference;
+
+public record StreamingHttpResult(HttpResponse response, Flow.Publisher<byte[]> body) {
+    public boolean isSuccessfulResponse() {
+        return RestStatus.isSuccessful(response.getStatusLine().getStatusCode());
+    }
+
+    public Flow.Publisher<HttpResult> toHttpResult() {
+        return subscriber -> body().subscribe(new Flow.Subscriber<>() {
+            @Override
+            public void onSubscribe(Flow.Subscription subscription) {
+                subscriber.onSubscribe(subscription);
+            }
+
+            @Override
+            public void onNext(byte[] item) {
+                subscriber.onNext(new HttpResult(response(), item));
+            }
+
+            @Override
+            public void onError(Throwable throwable) {
+                subscriber.onError(throwable);
+            }
+
+            @Override
+            public void onComplete() {
+                subscriber.onComplete();
+            }
+        });
+    }
+
+    public void readFullResponse(ActionListener<HttpResult> fullResponse) {
+        var stream = new ByteArrayOutputStream();
+        AtomicReference<Flow.Subscription> upstream = new AtomicReference<>(null);
+        body.subscribe(new Flow.Subscriber<>() {
+            @Override
+            public void onSubscribe(Flow.Subscription subscription) {
+                upstream.set(subscription);
+                upstream.get().request(1);
+            }
+
+            @Override
+            public void onNext(byte[] item) {
+                stream.writeBytes(item);
+                upstream.get().request(1);
+            }
+
+            @Override
+            public void onError(Throwable throwable) {
+                ExceptionsHelper.maybeError(throwable).ifPresent(ExceptionsHelper::maybeDieOnAnotherThread);
+                fullResponse.onFailure(new RuntimeException("Fatal while fully consuming stream", throwable));
+            }
+
+            @Override
+            public void onComplete() {
+                fullResponse.onResponse(new HttpResult(response, stream.toByteArray()));
+            }
+        });
+    }
+}