fix connector executor: restore context

Signed-off-by: Yaliang Wu <ylwu@amazon.com>

Author: ylwu-amzn

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java

@@ -26,6 +26,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.util.TokenBucket;
+import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.connector.AwsConnector;
 import org.opensearch.ml.common.connector.Connector;
@@ -123,14 +124,15 @@ public void invokeRemoteService(
                 default:
                     throw new IllegalArgumentException("unsupported http method");
             }
+            ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().newStoredContext(true);
             AsyncExecuteRequest executeRequest = AsyncExecuteRequest
                 .builder()
                 .request(signRequest(request))
                 .requestContentPublisher(new SimpleHttpContentPublisher(request))
                 .responseHandler(
                     new MLSdkAsyncHttpResponseHandler(
                         executionContext,
-                        actionListener,
+                        ActionListener.runBefore(actionListener, storedContext::restore),
                         parameters,
                         connector,
                         scriptService,

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutor.java

@@ -25,6 +25,7 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.util.TokenBucket;
+import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.connector.Connector;
 import org.opensearch.ml.common.connector.HttpConnector;
@@ -122,14 +123,15 @@ public void invokeRemoteService(
                 default:
                     throw new IllegalArgumentException("unsupported http method");
             }
+            ThreadContext.StoredContext storedContext = client.threadPool().getThreadContext().newStoredContext(true);
             AsyncExecuteRequest executeRequest = AsyncExecuteRequest
                 .builder()
                 .request(request)
                 .requestContentPublisher(new SimpleHttpContentPublisher(request))
                 .responseHandler(
                     new MLSdkAsyncHttpResponseHandler(
                         executionContext,
-                        actionListener,
+                        ActionListener.runBefore(actionListener, storedContext::restore),
                         parameters,
                         connector,
                         scriptService,