Incrementing reference count for body content and fixing tests

Author: jonathan-buttner

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java

@@ -360,8 +360,7 @@ protected Deque<ServerSentEvent> unifiedCompletionInferOnMockService(
         List<String> input,
         @Nullable Consumer<Response> responseConsumerCallback
     ) throws Exception {
-        var route = randomBoolean() ? "_stream" : "_unified"; // TODO remove unified route
-        var endpoint = Strings.format("_inference/%s/%s/%s", taskType, modelId, route);
+        var endpoint = Strings.format("_inference/%s/%s/_stream", taskType, modelId);
         return callAsyncUnified(endpoint, input, "user", responseConsumerCallback);
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceActionProxy.java

@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.core.inference.action.InferenceActionProxy;
 import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionAction;
+import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
 import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 
 import java.io.IOException;
@@ -77,27 +78,33 @@ protected void doExecute(Task task, InferenceActionProxy.Request request, Action
         }
     }
 
-    private void sendUnifiedCompletionRequest(InferenceActionProxy.Request request, ActionListener<InferenceAction.Response> listener)
-        throws IOException {
+    private void sendUnifiedCompletionRequest(InferenceActionProxy.Request request, ActionListener<InferenceAction.Response> listener) {
+        // format any validation exceptions from the rest -> transport path as UnifiedChatCompletionException
+        var unifiedErrorFormatListener = listener.delegateResponse((l, e) -> l.onFailure(UnifiedChatCompletionException.fromThrowable(e)));
 
-        if (request.isStreaming() == false) {
-            throw new ElasticsearchStatusException(
-                "The [chat_completion] task type only supports streaming, please try again with the _stream API",
-                RestStatus.BAD_REQUEST
-            );
-        }
+        try {
+            if (request.isStreaming() == false) {
+                unifiedErrorFormatListener.onFailure(new ElasticsearchStatusException(
+                    "The [chat_completion] task type only supports streaming, please try again with the _stream API",
+                    RestStatus.BAD_REQUEST
+                ));
+                return;
+            }
 
-        UnifiedCompletionAction.Request unifiedRequest;
-        try (var parser = XContentHelper.createParser(XContentParserConfiguration.EMPTY, request.getContent(), request.getContentType())) {
-            unifiedRequest = UnifiedCompletionAction.Request.parseRequest(
-                request.getInferenceEntityId(),
-                request.getTaskType(),
-                request.getTimeout(),
-                parser
-            );
-        }
+            UnifiedCompletionAction.Request unifiedRequest;
+            try (var parser = XContentHelper.createParser(XContentParserConfiguration.EMPTY, request.getContent(), request.getContentType())) {
+                unifiedRequest = UnifiedCompletionAction.Request.parseRequest(
+                    request.getInferenceEntityId(),
+                    request.getTaskType(),
+                    request.getTimeout(),
+                    parser
+                );
+            }
 
-        executeAsyncWithOrigin(client, INFERENCE_ORIGIN, UnifiedCompletionAction.INSTANCE, unifiedRequest, listener);
+            executeAsyncWithOrigin(client, INFERENCE_ORIGIN, UnifiedCompletionAction.INSTANCE, unifiedRequest, unifiedErrorFormatListener);
+        } catch (Exception e) {
+            unifiedErrorFormatListener.onFailure(e);
+        }
     }
 
     private void sendInferenceActionRequest(InferenceActionProxy.Request request, ActionListener<InferenceAction.Response> listener)

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

@@ -54,7 +54,7 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient
             shouldStream()
         );
 
-        return channel -> client.execute(InferenceActionProxy.INSTANCE, request, listener(channel));
+        return channel -> client.execute(InferenceActionProxy.INSTANCE, request, ActionListener.withRef(listener(channel), content));
     }
 
     protected abstract boolean shouldStream();

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

@@ -14,7 +14,6 @@
 import org.elasticsearch.rest.ServerlessScope;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
-import org.elasticsearch.xpack.core.inference.results.UnifiedChatCompletionException;
 
 import java.util.List;
 import java.util.Objects;
@@ -44,11 +43,7 @@ public List<Route> routes() {
 
     @Override
     protected ActionListener<InferenceAction.Response> listener(RestChannel channel) {
-        // TODO confirm with Pat that this will work
-        return new ServerSentEventsRestActionListener(channel, threadPool).delegateResponse((l, e) -> {
-            // format any validation exceptions from the rest -> transport path as UnifiedChatCompletionException
-            l.onFailure(UnifiedChatCompletionException.fromThrowable(e));
-        });
+        return new ServerSentEventsRestActionListener(channel, threadPool);
     }
 
     @Override