feat(gax): add recoverable error query and buffer realignment loop

Implements Category 2 (recoverable) error recovery for chunk uploads by
querying server upload status, realigning the buffer to the committed offset,
topping up from the payload stream, and re-transmitting.

- In ChunkAttemptCallable, asynchronously query upload status via
  queryStatusCallable, handle server finalized sessions, validate committed
  offset invariants (raising UploadProtocolViolationException on violations),
  and realign RewindableStreamBuffer to the server committed byte offset.
- Wrap QueryStatusCallable in ResumableUploadCallableImpl with a RetryingCallable
  using UploadResultRetryAlgorithm(UploadCommand.QUERY) and ExponentialRetryAlgorithm.
  Resolves queryStatusCallable eagerly at construction time in
  ResumableUploadCallableImpl, with CallableTest stubbing queryStatusCallable().
- Propagate query future cancellation immediately to the in-flight future.

Author: whowes

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/resumable/ChunkUploadResponse.java

@@ -69,7 +69,7 @@ public static <ResponseT> Builder<ResponseT> newBuilder() {
 
   public static <ResponseT> ChunkUploadResponse<ResponseT> create(
       boolean isComplete, @Nullable ResponseT response) {
-    return create(isComplete, response, null);
+    return create(isComplete, response, isComplete ? "final" : "active");
   }
 
   public static <ResponseT> ChunkUploadResponse<ResponseT> create(

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ChunkAttemptCallable.java

@@ -0,0 +1,298 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.rpc;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.resumable.ChunkUploadRequest;
+import com.google.api.gax.resumable.ChunkUploadResponse;
+import com.google.api.gax.resumable.QueryStatusRequest;
+import com.google.api.gax.resumable.QueryStatusResponse;
+import com.google.api.gax.retrying.RetryingFuture;
+import com.google.common.util.concurrent.MoreExecutors;
+import java.time.Duration;
+import java.util.concurrent.Callable;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+/**
+ * A {@link Callable} representing an attempt to transmit a single chunk in a resumable upload
+ * session. Used with {@link com.google.api.gax.retrying.ScheduledRetryingExecutor}.
+ *
+ * <p>Execution follows the standard attempt template with pre-attempt recovery handling. When the
+ * previous attempt failed with a Category 2 (recoverable) error or missing status header, {@code
+ * prepareAttempt} queries session status, realigns the buffer window, tops up from the stream, and
+ * dispatches the chunk upload request. The callable never blocks on {@code .get()}; results and
+ * cancellations propagate asynchronously.
+ *
+ * @param <ResponseT> the type of the final response message once the upload completes
+ */
+@NullMarked
+class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<ResponseT>> {
+
+  private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
+      uploadChunkCallable;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      queryStatusCallable;
+  private final RewindableStreamBuffer buffer;
+  private final String uploadUrl;
+  private final ApiCallContext originalCallContext;
+
+  private volatile ChunkUploadRequest currentRequest;
+  private volatile UploadCommand currentCommand;
+
+  private volatile @Nullable RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture;
+  private volatile @Nullable ApiFuture<?> inFlightFuture;
+  private volatile @Nullable Throwable lastFailure;
+  private volatile @Nullable ChunkUploadResponse<ResponseT> lastResponse;
+
+  ChunkAttemptCallable(
+      UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
+      RewindableStreamBuffer buffer,
+      String uploadUrl,
+      ChunkUploadRequest request,
+      ApiCallContext callContext,
+      UploadCommand command) {
+    this.uploadChunkCallable =
+        checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
+    this.queryStatusCallable =
+        checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
+    this.buffer = checkNotNull(buffer, "buffer must not be null");
+    this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
+    this.currentRequest = checkNotNull(request, "request must not be null");
+    this.originalCallContext = checkNotNull(callContext, "callContext must not be null");
+    this.currentCommand = checkNotNull(command, "command must not be null");
+  }
+
+  void setRetryingFuture(RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture) {
+    this.retryingFuture = checkNotNull(retryingFuture, "retryingFuture must not be null");
+  }
+
+  private boolean needsRecovery() {
+    if (lastFailure != null) {
+      UploadErrorClassifier.Category category =
+          UploadErrorClassifier.classify(lastFailure, currentCommand);
+      return category == UploadErrorClassifier.Category.RECOVERABLE;
+    }
+    if (lastResponse != null && lastResponse.getUploadStatus() == null) {
+      UploadErrorClassifier.Category category =
+          UploadErrorClassifier.classifyMissingStatusHeader(currentCommand);
+      return category == UploadErrorClassifier.Category.RECOVERABLE;
+    }
+    return false;
+  }
+
+  private void failAttempt(
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture, Throwable t) {
+    lastFailure = t;
+    lastResponse = null;
+    attemptFuture.setException(t);
+  }
+
+  /**
+   * Pre-attempt recovery step invoked before transmitting an attempt when the previous attempt
+   * encountered a Category 2 (recoverable) error or missing status header.
+   */
+  private void prepareAttempt(
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
+      ApiCallContext attemptContext,
+      RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    QueryStatusRequest queryRequest = QueryStatusRequest.create(uploadUrl);
+    ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
+        queryStatusCallable.futureCall(queryRequest, attemptContext);
+    if (queryFuture == null) {
+      failAttempt(
+          attemptFuture, new IllegalStateException("queryStatusCallable returned a null future"));
+      return;
+    }
+    this.inFlightFuture = queryFuture;
+
+    ApiFutures.addCallback(
+        queryFuture,
+        new ApiFutureCallback<QueryStatusResponse<ResponseT>>() {
+          @Override
+          public void onSuccess(QueryStatusResponse<ResponseT> queryResponse) {
+            handleQuerySuccess(queryResponse, attemptFuture, attemptContext, currentRetryingFuture);
+          }
+
+          @Override
+          public void onFailure(Throwable t) {
+            failAttempt(attemptFuture, t);
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  private void handleQuerySuccess(
+      QueryStatusResponse<ResponseT> queryResponse,
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
+      ApiCallContext attemptContext,
+      RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    if (currentRetryingFuture.isDone()) {
+      return;
+    }
+
+    if (queryResponse.getUploadStatus() == null) {
+      failAttempt(
+          attemptFuture,
+          UploadErrors.protocolViolation(
+              "Query status response missing X-Goog-Upload-Status header for upload URL: "
+                  + uploadUrl));
+      return;
+    }
+
+    // Server already finalized the upload.
+    if (queryResponse.isComplete()) {
+      ChunkUploadResponse<ResponseT> response =
+          ChunkUploadResponse.create(
+              true, queryResponse.getResponse(), queryResponse.getUploadStatus());
+      lastFailure = null;
+      lastResponse = response;
+      attemptFuture.set(response);
+      return;
+    }
+
+    // Incomplete query response with null committed offset violates the protocol invariant.
+    Long committedOffset = queryResponse.getCommittedOffset();
+    if (committedOffset == null) {
+      failAttempt(
+          attemptFuture,
+          UploadErrors.protocolViolation(
+              "Incomplete query status response did not include a committed offset for upload URL: "
+                  + uploadUrl));
+      return;
+    }
+
+    // Normal path: realign buffer to committedOffset, compact and top up.
+    try {
+      buffer.realignTo(committedOffset);
+    } catch (Throwable e) {
+      failAttempt(attemptFuture, e);
+      return;
+    }
+
+    // Determine the upload command for the realigned buffer.
+    // Preserve upload,finalize for a trailing partial after realignment.
+    UploadCommand realignedCommand;
+    if (buffer.isFinal()) {
+      realignedCommand = buffer.isEmpty() ? UploadCommand.FINALIZE : UploadCommand.UPLOAD_FINALIZE;
+    } else {
+      realignedCommand = UploadCommand.UPLOAD;
+    }
+
+    ChunkUploadRequest realignedRequest =
+        ChunkUploadRequest.newBuilder()
+            .setUploadUrl(uploadUrl)
+            .setPayload(buffer.getBuffer())
+            .setPayloadLength(buffer.getPayloadLength())
+            .setOffset(buffer.getBufferBaseOffset())
+            .setFinal(buffer.isFinal())
+            .build();
+
+    this.currentRequest = realignedRequest;
+    this.currentCommand = realignedCommand;
+
+    dispatchChunkUpload(attemptFuture, attemptContext, currentRetryingFuture);
+  }
+
+  private void dispatchChunkUpload(
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
+      ApiCallContext attemptContext,
+      RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
+        uploadChunkCallable.futureCall(currentRequest, attemptContext);
+    this.inFlightFuture = chunkFuture;
+
+    ApiFutures.addCallback(
+        chunkFuture,
+        new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
+          @Override
+          public void onSuccess(ChunkUploadResponse<ResponseT> response) {
+            lastFailure = null;
+            lastResponse = response;
+            attemptFuture.set(response);
+          }
+
+          @Override
+          public void onFailure(Throwable t) {
+            failAttempt(attemptFuture, t);
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  @Override
+  public @Nullable ChunkUploadResponse<ResponseT> call() {
+    RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture =
+        checkNotNull(retryingFuture, "retryingFuture must be set before call()");
+    ApiCallContext attemptContext = originalCallContext;
+
+    Duration rpcTimeout = currentRetryingFuture.getAttemptSettings().getRpcTimeoutDuration();
+    if (!rpcTimeout.isZero() && attemptContext.getTimeoutDuration() == null) {
+      attemptContext = attemptContext.withTimeoutDuration(rpcTimeout);
+    }
+
+    SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture = SettableApiFuture.create();
+    currentRetryingFuture.setAttemptFuture(attemptFuture);
+
+    if (currentRetryingFuture.isDone()) {
+      return null;
+    }
+
+    currentRetryingFuture.addListener(
+        () -> {
+          if (currentRetryingFuture.isCancelled()) {
+            ApiFuture<?> inFlight = inFlightFuture;
+            if (inFlight != null) {
+              inFlight.cancel(true);
+            }
+            attemptFuture.cancel(true);
+          }
+        },
+        MoreExecutors.directExecutor());
+
+    try {
+      if (needsRecovery()) {
+        prepareAttempt(attemptFuture, attemptContext, currentRetryingFuture);
+      } else {
+        dispatchChunkUpload(attemptFuture, attemptContext, currentRetryingFuture);
+      }
+    } catch (Throwable t) {
+      failAttempt(attemptFuture, t);
+    }
+
+    return null;
+  }
+}

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java

@@ -35,9 +35,16 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.api.gax.resumable.QueryStatusRequest;
+import com.google.api.gax.resumable.QueryStatusResponse;
 import com.google.api.gax.resumable.ResumableUploadClient;
 import com.google.api.gax.resumable.ResumableUploadSession;
+import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
+import com.google.api.gax.retrying.RetryAlgorithm;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import java.io.InputStream;
+import java.time.Duration;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -54,9 +61,22 @@
 public class ResumableUploadCallableImpl<RequestT, ResponseT>
     extends ResumableUploadCallable<RequestT, ResponseT> {
 
+  static final RetrySettings DEFAULT_QUERY_RETRY_SETTINGS =
+      RetrySettings.newBuilder()
+          .setInitialRetryDelayDuration(Duration.ofMillis(100))
+          .setRetryDelayMultiplier(1.3)
+          .setMaxRetryDelayDuration(Duration.ofMinutes(1))
+          .setInitialRpcTimeoutDuration(Duration.ofSeconds(30))
+          .setRpcTimeoutMultiplier(1.0)
+          .setMaxRpcTimeoutDuration(Duration.ofSeconds(30))
+          .setTotalTimeoutDuration(Duration.ofMinutes(5))
+          .build();
+
   private final ResumableUploadClient<RequestT, ResponseT> client;
   private final ResumableUploadCallSettings defaultCallSettings;
   private final ClientContext clientContext;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      retryingQueryCallable;
 
   public ResumableUploadCallableImpl(
       ResumableUploadClient<RequestT, ResponseT> client,
@@ -66,6 +86,17 @@ public ResumableUploadCallableImpl(
     this.defaultCallSettings =
         checkNotNull(defaultCallSettings, "defaultCallSettings must not be null");
     this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
+
+    RetryAlgorithm<QueryStatusResponse<ResponseT>> queryRetryAlgorithm =
+        new RetryAlgorithm<>(
+            new UploadResultRetryAlgorithm<>(UploadCommand.QUERY),
+            new ExponentialRetryAlgorithm(DEFAULT_QUERY_RETRY_SETTINGS, clientContext.getClock()));
+
+    this.retryingQueryCallable =
+        new RetryingCallable<>(
+            clientContext.getDefaultCallContext(),
+            checkNotNull(client.queryStatusCallable(), "queryStatusCallable must not be null"),
+            new ScheduledRetryingExecutor<>(queryRetryAlgorithm, clientContext.getExecutor()));
   }
 
   @Override
@@ -89,6 +120,7 @@ public ResumableUploadFuture<ResponseT> futureCall(
     return ResumableUploadFutureImpl.create(
         startFuture,
         client.uploadChunkCallable(),
+        retryingQueryCallable,
         payload,
         effectiveSettings,
         clientContext.getDefaultCallContext(),