ESQL: Fix async query inconsistent headers (elastic#135078)

Fixes elastic#135042

This PR:
- Fixes the get-result not always returning the expected headers
- Fixes the non-async query incorrectly returning the "is running" async header

Author: ivancea

docs/changelog/135078.yaml

@@ -0,0 +1,6 @@
+pr: 135078
+summary: Fix async get results with inconsistent headers
+area: ES|QL
+type: bug
+issues:
+ - 135042

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java

@@ -1411,6 +1411,10 @@ public static Map<String, Object> runEsqlSync(
             profileLogger.extractProfile(json, profileEnabled);
         }
 
+        var supportsAsyncHeadersFix = hasCapabilities(adminClient(), List.of("async_query_status_headers_fix"));
+        if (supportsAsyncHeadersFix) {
+            assertNoAsyncHeaders(response);
+        }
         assertWarnings(response, assertWarnings);
 
         return json;
@@ -1448,17 +1452,18 @@ public static Map<String, Object> runEsqlAsync(
         checkKeepOnCompletion(requestObject, json, keepOnCompletion);
         String id = (String) json.get("id");
 
-        var supportsAsyncHeaders = hasCapabilities(adminClient(), List.of("async_query_status_headers"));
+        var supportsAsyncHeaders = hasCapabilities(adminClient(), List.of("async_query_status_headers_fix"));
         var supportsSuggestedCast = hasCapabilities(adminClient(), List.of("suggested_cast"));
 
+        // Check headers on initial query call
+        if (supportsAsyncHeaders) {
+            assertAsyncHeaders(response, id, (boolean) json.get("is_running"));
+        }
+
         if (id == null) {
             // no id returned from an async call, must have completed immediately and without keep_on_completion
             assertThat(requestObject.keepOnCompletion(), either(nullValue()).or(is(false)));
             assertThat((boolean) json.get("is_running"), is(false));
-            if (supportsAsyncHeaders) {
-                assertThat(response.getHeader("X-Elasticsearch-Async-Id"), nullValue());
-                assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), is("?0"));
-            }
             if (profileLogger != null) {
                 profileLogger.extractProfile(json, profileEnabled);
             }
@@ -1485,11 +1490,6 @@ public static Map<String, Object> runEsqlAsync(
                 assertThat(json.get("pages"), nullValue());
             }
 
-            if (supportsAsyncHeaders) {
-                assertThat(response.getHeader("X-Elasticsearch-Async-Id"), is(id));
-                assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), is(isRunning ? "?1" : "?0"));
-            }
-
             // issue a second request to "async get" the results
             Request getRequest = prepareAsyncGetRequest(id);
             getRequest.setOptions(request.getOptions());
@@ -1499,6 +1499,11 @@ public static Map<String, Object> runEsqlAsync(
 
         var result = entityToMap(entity, requestObject.contentType());
 
+        // Check headers on get call
+        if (supportsAsyncHeaders) {
+            assertAsyncHeaders(response, id, (boolean) result.get("is_running"));
+        }
+
         // assert initial contents, if any, are the same as async get contents
         if (initialColumns != null) {
             if (supportsSuggestedCast == false) {
@@ -2005,6 +2010,16 @@ private static void createIndex(String indexName, boolean lookupMode, String map
         assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
     }
 
+    private static void assertAsyncHeaders(Response response, @Nullable String asyncId, boolean isRunning) {
+        assertThat(response.getHeader("X-Elasticsearch-Async-Id"), asyncId == null ? nullValue() : equalTo(asyncId));
+        assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), isRunning ? is("?1") : is("?0"));
+    }
+
+    private static void assertNoAsyncHeaders(Response response) {
+        assertThat(response.getHeader("X-Elasticsearch-Async-Id"), nullValue());
+        assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), nullValue());
+    }
+
     public static RequestObjectBuilder requestObjectBuilder() throws IOException {
         return new RequestObjectBuilder();
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -669,6 +669,11 @@ public enum Cap {
          */
         ASYNC_QUERY_STATUS_HEADERS,
 
+        /**
+         * Fix async headers not being sent on "get" requests
+         */
+        ASYNC_QUERY_STATUS_HEADERS_FIX,
+
         /**
          * Consider the upper bound when computing the interval in BUCKET auto mode.
          */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryResponse.java

@@ -111,7 +111,7 @@ public static Writeable.Reader<EsqlQueryResponse> reader(BlockFactory blockFacto
     }
 
     static EsqlQueryResponse deserialize(BlockStreamInput in) throws IOException {
-        String asyncExecutionId = asyncExecutionId = in.readOptionalString();
+        String asyncExecutionId = in.readOptionalString();
         boolean isRunning = in.readBoolean();
         boolean isAsync = in.readBoolean();
         List<ColumnInfoImpl> columns = in.readCollectionAsList(ColumnInfoImpl::new);
@@ -208,7 +208,7 @@ public boolean isRunning() {
     }
 
     public boolean isAsync() {
-        return isRunning;
+        return isAsync;
     }
 
     public boolean isPartial() {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlAsyncGetResultsAction.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
 import org.elasticsearch.xpack.core.async.GetAsyncResultRequest;
 import org.elasticsearch.xpack.esql.VerificationException;
 import org.elasticsearch.xpack.esql.action.EsqlAsyncGetResultAction;
@@ -35,6 +36,7 @@
 public class TransportEsqlAsyncGetResultsAction extends AbstractTransportQlAsyncGetResultsAction<EsqlQueryResponse, EsqlQueryTask> {
 
     private final BlockFactory blockFactory;
+    private final ThreadPool threadPool;
 
     @Inject
     public TransportEsqlAsyncGetResultsAction(
@@ -43,9 +45,9 @@ public TransportEsqlAsyncGetResultsAction(
         ClusterService clusterService,
         NamedWriteableRegistry registry,
         Client client,
-        ThreadPool threadPool,
         BigArrays bigArrays,
-        BlockFactoryProvider blockFactoryProvider
+        BlockFactoryProvider blockFactoryProvider,
+        ThreadPool threadPool
     ) {
         super(
             EsqlAsyncGetResultAction.NAME,
@@ -59,11 +61,12 @@ public TransportEsqlAsyncGetResultsAction(
             EsqlQueryTask.class
         );
         this.blockFactory = blockFactoryProvider.blockFactory();
+        this.threadPool = threadPool;
     }
 
     @Override
     protected void doExecute(Task task, GetAsyncResultRequest request, ActionListener<EsqlQueryResponse> listener) {
-        super.doExecute(task, request, unwrapListener(listener));
+        super.doExecute(task, request, unwrapListener(request.getId(), listener));
     }
 
     @Override
@@ -75,14 +78,21 @@ public Writeable.Reader<EsqlQueryResponse> responseReader() {
     static final String VERIFY_EX_NAME = ElasticsearchException.getExceptionName(new VerificationException(""));
 
     /**
-     * Unwraps the exception in the case of failure. This keeps the exception types
-     * the same as the sync API, namely ParsingException and VerificationException.
+     * Adds async headers, and unwraps the exception in the case of failure.
+     * <p>
+     *     This keeps the exception types the same as the sync API, namely ParsingException and VerificationException.
+     * </p>
      */
-    static <R> ActionListener<R> unwrapListener(ActionListener<R> listener) {
+    ActionListener<EsqlQueryResponse> unwrapListener(String asyncExecutionId, ActionListener<EsqlQueryResponse> listener) {
         return new ActionListener<>() {
             @Override
-            public void onResponse(R o) {
-                listener.onResponse(o);
+            public void onResponse(EsqlQueryResponse response) {
+                boolean isRunning = response.isRunning();
+                threadPool.getThreadContext()
+                    .addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_IS_RUNNING_HEADER, isRunning ? "?1" : "?0");
+                threadPool.getThreadContext().addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_ID_HEADER, asyncExecutionId);
+
+                listener.onResponse(response);
             }
 
             @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java

@@ -282,7 +282,20 @@ private void innerExecute(Task task, EsqlQueryRequest request, ActionListener<Es
             ActionListener.wrap(result -> {
                 recordCCSTelemetry(task, executionInfo, request, null);
                 planExecutor.metrics().recordTook(executionInfo.overallTook().millis());
-                listener.onResponse(toResponse(task, request, configuration, result));
+                var response = toResponse(task, request, configuration, result);
+                assert response.isAsync() == request.async() : "The response must be async if the request was async";
+
+                if (response.isAsync()) {
+                    if (response.asyncExecutionId().isPresent()) {
+                        String asyncExecutionId = response.asyncExecutionId().get();
+                        threadPool.getThreadContext().addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_ID_HEADER, asyncExecutionId);
+                    }
+                    boolean isRunning = response.isRunning();
+                    threadPool.getThreadContext()
+                        .addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_IS_RUNNING_HEADER, isRunning ? "?1" : "?0");
+                }
+
+                listener.onResponse(response);
             }, ex -> {
                 recordCCSTelemetry(task, executionInfo, request, ex);
                 listener.onFailure(ex);
@@ -380,10 +393,8 @@ private EsqlQueryResponse toResponse(Task task, EsqlQueryRequest request, Config
         EsqlQueryResponse.Profile profile = configuration.profile()
             ? new EsqlQueryResponse.Profile(result.completionInfo().driverProfiles(), result.completionInfo().planProfiles())
             : null;
-        threadPool.getThreadContext().addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_IS_RUNNING_HEADER, "?0");
         if (task instanceof EsqlQueryTask asyncTask && request.keepOnCompletion()) {
             String asyncExecutionId = asyncTask.getExecutionId().getEncoded();
-            threadPool.getThreadContext().addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_ID_HEADER, asyncExecutionId);
             return new EsqlQueryResponse(
                 columns,
                 result.pages(),