refactor: Simplify async polling logic and improve task status handling

Simplified the async polling logic in `ConvertOperations` by replacing the custom scheduler with `CompletionStage` chaining. Improved task status handling by centralizing logic for success, failure, and retry behavior. Cleaned up imports and adjusted tests to maintain consistency with updated logic.

Author: edeandrea

docling-serve/docling-serve-api/src/main/java/ai/docling/serve/api/DoclingServeApi.java

@@ -1,7 +1,11 @@
 package ai.docling.serve.api;
 
+import java.time.Duration;
+
 import org.jspecify.annotations.Nullable;
 
+import ai.docling.serve.api.convert.request.ConvertDocumentRequest;
+
 /**
  * Docling Serve API interface.
  */
@@ -104,6 +108,30 @@ default B prettyPrint() {
       return prettyPrint(true);
     }
 
+    /**
+     * Sets the polling interval for async operations.
+     *
+     * <p>This configures how frequently the client will check the status of async
+     * conversion tasks when using {@link DoclingServeApi#convertSourceAsync(ConvertDocumentRequest)} (ConvertDocumentRequest)}.
+     *
+     * @param asyncPollInterval the polling interval (must not be null or negative)
+     * @return this builder instance for method chaining
+     * @throws IllegalArgumentException if asyncPollInterval is null or negative
+     */
+    B asyncPollInterval(Duration asyncPollInterval);
+
+    /**
+     * Sets the timeout for async operations.
+     *
+     * <p>This configures the maximum time to wait for an async conversion task to complete
+     * when using {@link DoclingServeApi#convertSourceAsync(ConvertDocumentRequest)} (ConvertDocumentRequest)}.
+     *
+     * @param asyncTimeout the timeout duration (must not be null or negative)
+     * @return this builder instance for method chaining
+     * @throws IllegalArgumentException if asyncTimeout is null or negative
+     */
+    B asyncTimeout(Duration asyncTimeout);
+
     /**
      * Builds and returns an instance of the specified type, representing the completed configuration
      * of the builder. The returned instance is typically an implementation of the Docling API.

docling-serve/docling-serve-client/src/main/java/ai/docling/serve/client/DoclingServeClient.java

@@ -5,14 +5,14 @@
 
 import java.net.URI;
 import java.net.http.HttpClient;
-import java.time.Duration;
 import java.net.http.HttpClient.Redirect;
 import java.net.http.HttpRequest;
 import java.net.http.HttpRequest.BodyPublisher;
 import java.net.http.HttpRequest.BodyPublishers;
 import java.net.http.HttpResponse;
 import java.net.http.HttpResponse.BodyHandlers;
 import java.nio.ByteBuffer;
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -449,16 +449,18 @@ public B prettyPrint(boolean prettyPrint) {
      * Sets the polling interval for async operations.
      *
      * <p>This configures how frequently the client will check the status of async
-     * conversion tasks when using {@link DoclingServeApi#convertSourceAsyncAndWait(ConvertDocumentRequest)}.
+     * conversion tasks when using {@link DoclingServeApi#convertSourceAsync(ConvertDocumentRequest)} (ConvertDocumentRequest)}.
      *
      * @param asyncPollInterval the polling interval (must not be null or negative)
      * @return this builder instance for method chaining
      * @throws IllegalArgumentException if asyncPollInterval is null or negative
      */
+    @Override
     public B asyncPollInterval(Duration asyncPollInterval) {
       if (asyncPollInterval == null || asyncPollInterval.isNegative() || asyncPollInterval.isZero()) {
         throw new IllegalArgumentException("asyncPollInterval must be a positive duration");
       }
+
       this.asyncPollInterval = asyncPollInterval;
       return (B) this;
     }
@@ -467,12 +469,13 @@ public B asyncPollInterval(Duration asyncPollInterval) {
      * Sets the timeout for async operations.
      *
      * <p>This configures the maximum time to wait for an async conversion task to complete
-     * when using {@link DoclingServeApi#convertSourceAsyncAndWait(ConvertDocumentRequest)}.
+     * when using {@link DoclingServeApi#convertSourceAsync(ConvertDocumentRequest)} (ConvertDocumentRequest)}.
      *
      * @param asyncTimeout the timeout duration (must not be null or negative)
      * @return this builder instance for method chaining
      * @throws IllegalArgumentException if asyncTimeout is null or negative
      */
+    @Override
     public B asyncTimeout(Duration asyncTimeout) {
       if (asyncTimeout == null || asyncTimeout.isNegative() || asyncTimeout.isZero()) {
         throw new IllegalArgumentException("asyncTimeout must be a positive duration");

docling-serve/docling-serve-client/src/main/java/ai/docling/serve/client/operations/ConvertOperations.java

@@ -2,19 +2,18 @@
 
 import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import ai.docling.serve.api.DoclingServeConvertApi;
+import ai.docling.serve.api.DoclingServeTaskApi;
 import ai.docling.serve.api.convert.request.ConvertDocumentRequest;
 import ai.docling.serve.api.convert.response.ConvertDocumentResponse;
 import ai.docling.serve.api.task.request.TaskResultRequest;
 import ai.docling.serve.api.task.request.TaskStatusPollRequest;
-import ai.docling.serve.api.task.response.TaskStatus;
 import ai.docling.serve.api.task.response.TaskStatusPollResponse;
 import ai.docling.serve.api.util.ValidationUtils;
 
@@ -26,22 +25,22 @@ public final class ConvertOperations implements DoclingServeConvertApi {
   private static final Logger LOG = LoggerFactory.getLogger(ConvertOperations.class);
 
   private final HttpOperations httpOperations;
-  private final TaskOperations taskOperations;
+  private final DoclingServeTaskApi taskApi;
   private final Duration asyncPollInterval;
   private final Duration asyncTimeout;
 
   /**
    * Creates a new ConvertOperations instance.
    *
    * @param httpOperations the HTTP operations handler for executing requests
-   * @param taskOperations the task operations handler for polling and retrieving results
+   * @param taskApi the task operations handler for polling and retrieving results
    * @param asyncPollInterval the interval between status polls for async operations
    * @param asyncTimeout the maximum time to wait for async operations to complete
    */
-  public ConvertOperations(HttpOperations httpOperations, TaskOperations taskOperations,
+  public ConvertOperations(HttpOperations httpOperations, DoclingServeTaskApi taskApi,
                            Duration asyncPollInterval, Duration asyncTimeout) {
     this.httpOperations = httpOperations;
-    this.taskOperations = taskOperations;
+    this.taskApi = taskApi;
     this.asyncPollInterval = asyncPollInterval;
     this.asyncTimeout = asyncTimeout;
   }
@@ -68,85 +67,69 @@ private <I> RequestContext<I, TaskStatusPollResponse> createAsyncRequestContext(
         .build();
   }
 
-  @Override
+    @Override
   public CompletableFuture<ConvertDocumentResponse> convertSourceAsync(ConvertDocumentRequest request) {
     ValidationUtils.ensureNotNull(request, "request");
 
-    CompletableFuture<ConvertDocumentResponse> resultFuture = new CompletableFuture<>();
-
-    ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
-      Thread t = new Thread(r, "docling-async-poller");
-      t.setDaemon(true);
-      return t;
-    });
-
-    // Start the async conversion
-    CompletableFuture.supplyAsync(() ->
+    // Start the async conversion and chain the polling logic
+    return CompletableFuture.supplyAsync(() ->
         this.httpOperations.executePost(createAsyncRequestContext("/v1/convert/source/async", request))
-    ).thenAccept(taskResponse -> {
+    ).thenCompose(taskResponse -> {
       var taskId = taskResponse.getTaskId();
       LOG.info("Started async conversion with task ID: {}", taskId);
 
       long startTime = System.currentTimeMillis();
-
-      // Schedule the polling task
-      Runnable pollTask = new Runnable() {
-        @Override
-        public void run() {
-          try {
-            // Check if we've timed out
-            if (System.currentTimeMillis() - startTime > asyncTimeout.toMillis()) {
-              resultFuture.completeExceptionally(
-                  new RuntimeException("Async conversion timed out after " + asyncTimeout + " for task: " + taskId));
-              scheduler.shutdown();
-              return;
-            }
-
-            // Use the existing pollTaskStatus method from TaskOperations
-            var pollRequest = TaskStatusPollRequest.builder().taskId(taskId).build();
-            var statusResponse = taskOperations.pollTaskStatus(pollRequest);
-            var status = statusResponse.getTaskStatus();
-
-            LOG.debug("Task {} status: {}", taskId, status);
-
-            if (status == TaskStatus.SUCCESS) {
-              LOG.info("Task {} completed successfully", taskId);
-              // Use the existing convertTaskResult method from TaskOperations
-              var resultRequest = TaskResultRequest.builder().taskId(taskId).build();
-              try {
-                var result = taskOperations.convertTaskResult(resultRequest);
-                resultFuture.complete(result);
-              } catch (Exception e) {
-                resultFuture.completeExceptionally(e);
-              }
-              scheduler.shutdown();
-            } else if (status == TaskStatus.FAILURE) {
-              String errorMessage = "Task failed";
-              if (statusResponse.getTaskStatusMetadata() != null) {
-                errorMessage = "Task failed: " + statusResponse.getTaskStatusMetadata();
-              }
-              resultFuture.completeExceptionally(
-                  new RuntimeException("Async conversion failed for task " + taskId + ": " + errorMessage));
-              scheduler.shutdown();
-            } else {
-              // Still in progress (PENDING or STARTED), schedule next poll
-              scheduler.schedule(this, asyncPollInterval.toMillis(), TimeUnit.MILLISECONDS);
-            }
-          } catch (Exception e) {
-            resultFuture.completeExceptionally(e);
-            scheduler.shutdown();
-          }
-        }
-      };
-
-      // Start polling immediately
-      scheduler.schedule(pollTask, 0, TimeUnit.MILLISECONDS);
-    }).exceptionally(e -> {
-      resultFuture.completeExceptionally(e);
-      scheduler.shutdown();
-      return null;
+      return pollTaskUntilComplete(taskId, startTime);
     });
+  }
 
-    return resultFuture;
+  /**
+   * Recursively polls a task until it completes, fails, or times out.
+   *
+   * @param taskId the ID of the task to poll
+   * @param startTime the timestamp when polling started (for timeout calculation)
+   * @return a CompletableFuture that completes with the conversion result
+   */
+  private CompletableFuture<ConvertDocumentResponse> pollTaskUntilComplete(String taskId, long startTime) {
+    // Check if we've timed out
+    if (System.currentTimeMillis() - startTime > asyncTimeout.toMillis()) {
+      return CompletableFuture.failedFuture(
+          new RuntimeException("Async conversion timed out after " + asyncTimeout + " for task: " + taskId));
+    }
+
+    // Poll the task status
+    var pollRequest = TaskStatusPollRequest.builder().taskId(taskId).build();
+
+    return CompletableFuture.supplyAsync(() -> taskApi.pollTaskStatus(pollRequest))
+        .thenCompose(statusResponse -> pollTaskStatus(statusResponse, startTime));
+  }
+
+  private CompletionStage<ConvertDocumentResponse> pollTaskStatus(TaskStatusPollResponse statusResponse, long startTime) {
+    var status = statusResponse.getTaskStatus();
+    var taskId = statusResponse.getTaskId();
+    LOG.debug("Task {} status: {}", taskId, status);
+
+    return switch (statusResponse.getTaskStatus()) {
+      case SUCCESS -> {
+        LOG.info("Task {} completed successfully", taskId);
+        // Retrieve the result
+        var resultRequest = TaskResultRequest.builder().taskId(taskId).build();
+        yield CompletableFuture.supplyAsync(() -> taskApi.convertTaskResult(resultRequest));
+      }
+      case FAILURE -> {
+        var errorMessage = "Task failed";
+        if (statusResponse.getTaskStatusMetadata()!=null) {
+          errorMessage = "Task failed: " + statusResponse.getTaskStatusMetadata();
+        }
+        yield CompletableFuture.failedFuture(
+            new RuntimeException("Async conversion failed for task " + taskId + ": " + errorMessage));
+      }
+      default ->
+        // Still in progress (PENDING or STARTED), schedule next poll after delay
+        CompletableFuture.supplyAsync(
+            () -> null,
+            CompletableFuture.delayedExecutor(asyncPollInterval.toMillis(), TimeUnit.MILLISECONDS)
+        ).thenCompose(v -> pollTaskUntilComplete(taskId, startTime));
+    };
   }
 }