From 6efb1ebeefb8a7930f6767262457d0d5b568ac75 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Fri, 10 Nov 2023 16:49:20 -0800
Subject: [PATCH] Fixing class casting exception for getTask (#1616) (#1618)

Signed-off-by: Joshua Palis <jpalis@amazon.com>
(cherry picked from commit 8f690fc84c0f93fec67a70b74ee01f3b1361033d)

Co-authored-by: Joshua Palis <jpalis@amazon.com>
---
 .../ml/client/MachineLearningNodeClient.java       | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java b/client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java
index 4c83b5060b..1c66a94403 100644
--- a/client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java
+++ b/client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java
@@ -208,9 +208,7 @@ public void execute(FunctionName name, Input input, ActionListener<MLExecuteTask
     public void getTask(String taskId, ActionListener<MLTask> listener) {
         MLTaskGetRequest mlTaskGetRequest = MLTaskGetRequest.builder().taskId(taskId).build();
 
-        client.execute(MLTaskGetAction.INSTANCE, mlTaskGetRequest, ActionListener.wrap(response -> {
-            listener.onResponse(MLTaskGetResponse.fromActionResponse(response).getMlTask());
-        }, listener::onFailure));
+        client.execute(MLTaskGetAction.INSTANCE, mlTaskGetRequest, getMLTaskResponseActionListener(listener));
     }
 
     @Override
@@ -250,6 +248,16 @@ public void createConnector(MLCreateConnectorInput mlCreateConnectorInput, Actio
         client.execute(MLCreateConnectorAction.INSTANCE, createConnectorRequest, getMlCreateConnectorResponseActionListener(listener));
     }
 
+    private ActionListener<MLTaskGetResponse> getMLTaskResponseActionListener(ActionListener<MLTask> listener) {
+        ActionListener<MLTaskGetResponse> internalListener = ActionListener
+            .wrap(getResponse -> { listener.onResponse(getResponse.getMlTask()); }, listener::onFailure);
+        ActionListener<MLTaskGetResponse> actionListener = wrapActionListener(internalListener, response -> {
+            MLTaskGetResponse getResponse = MLTaskGetResponse.fromActionResponse(response);
+            return getResponse;
+        });
+        return actionListener;
+    }
+
     private ActionListener<MLDeployModelResponse> getMlDeployModelResponseActionListener(ActionListener<MLDeployModelResponse> listener) {
         ActionListener<MLDeployModelResponse> actionListener = wrapActionListener(listener, response -> {
             MLDeployModelResponse deployModelResponse = MLDeployModelResponse.fromActionResponse(response);