combine json chunks

Signed-off-by: Jiaping Zeng <jpz@amazon.com>

Author: jiapingzeng

plugin/src/main/java/org/opensearch/ml/rest/RestMLExecuteStreamAction.java

@@ -17,6 +17,7 @@
 import static org.opensearch.ml.utils.RestActionUtils.isAsync;
 import static org.opensearch.ml.utils.TenantAwareHelper.getTenantID;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.ByteBuffer;
@@ -48,7 +49,6 @@
 import org.opensearch.ml.common.MLModel;
 import org.opensearch.ml.common.agent.MLAgent;
 import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
-import org.opensearch.ml.common.exception.MLException;
 import org.opensearch.ml.common.input.Input;
 import org.opensearch.ml.common.input.MLInput;
 import org.opensearch.ml.common.input.execute.agent.AgentMLInput;
@@ -158,10 +158,12 @@ public RestChannelConsumer prepareRequest(RestRequest request, NodeClient client
                 );
             channel.prepareResponse(RestStatus.OK, headers);
 
-            Flux.from(channel).ofType(HttpChunk.class).concatMap(chunk -> {
-                final CompletableFuture<HttpChunk> future = new CompletableFuture<>();
+            Flux.from(channel).ofType(HttpChunk.class).collectList().flatMap(chunks -> {
                 try {
-                    MLExecuteTaskRequest mlExecuteTaskRequest = getRequest(agentId, request, chunk.content());
+                    BytesReference completeContent = combineChunks(chunks);
+                    MLExecuteTaskRequest mlExecuteTaskRequest = getRequest(agentId, request, completeContent);
+
+                    final CompletableFuture<HttpChunk> future = new CompletableFuture<>();
                     StreamTransportResponseHandler<MLTaskResponse> handler = new StreamTransportResponseHandler<MLTaskResponse>() {
                         @Override
                         public void handleStreamResponse(StreamTransportResponse<MLTaskResponse> streamResponse) {
@@ -214,11 +216,10 @@ public MLTaskResponse read(StreamInput in) throws IOException {
                             handler
                         );
 
+                    return Mono.fromCompletionStage(future);
                 } catch (IOException e) {
-                    throw new MLException("Got an exception in flux.", e);
+                    return Mono.error(new OpenSearchStatusException("Failed to parse request", RestStatus.BAD_REQUEST, e));
                 }
-
-                return Mono.fromCompletionStage(future);
             }).doOnNext(channel::sendChunk).onErrorComplete(ex -> {
                 // Error handling
                 try {
@@ -402,6 +403,19 @@ private String extractTensorResult(MLTaskResponse response, String tensorName) {
         return Map.of();
     }
 
+    private BytesReference combineChunks(List<HttpChunk> chunks) {
+        try {
+            ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+            for (HttpChunk chunk : chunks) {
+                chunk.content().writeTo(buffer);
+            }
+            return BytesReference.fromByteBuffer(ByteBuffer.wrap(buffer.toByteArray()));
+        } catch (IOException e) {
+            log.error("Failed to combine chunks", e);
+            throw new OpenSearchStatusException("Failed to combine request chunks", RestStatus.INTERNAL_SERVER_ERROR, e);
+        }
+    }
+
     private HttpChunk createHttpChunk(String sseData, boolean isLast) {
         BytesReference bytesRef = BytesReference.fromByteBuffer(ByteBuffer.wrap(sseData.getBytes()));
         return new HttpChunk() {