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 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 getMLTaskResponseActionListener(ActionListener listener) { + ActionListener internalListener = ActionListener + .wrap(getResponse -> { listener.onResponse(getResponse.getMlTask()); }, listener::onFailure); + ActionListener actionListener = wrapActionListener(internalListener, response -> { + MLTaskGetResponse getResponse = MLTaskGetResponse.fromActionResponse(response); + return getResponse; + }); + return actionListener; + } + private ActionListener getMlDeployModelResponseActionListener(ActionListener listener) { ActionListener actionListener = wrapActionListener(listener, response -> { MLDeployModelResponse deployModelResponse = MLDeployModelResponse.fromActionResponse(response);