From 4f1ce04cc005c6d4364726d17c50e36f649f2474 Mon Sep 17 00:00:00 2001
From: Nidhi Nandwani <nidhiii@google.com>
Date: Mon, 26 Oct 2026 08:04:05 +0000
Subject: [PATCH] fix(storage): handle write_handle-only responses while
 awaiting flush ack in BidiUploadState

Allow AppendableUploadState.onResponse to update writeHandle and return early when an intermediate BidiWriteObjectResponse carries only write_handle without persisted_size or resource, preventing IllegalStateException during flush/awaitAck.

[Generated-by: AI]
---
 .../google/cloud/storage/BidiUploadState.java |  6 ++
 .../google/cloud/storage/BidiUploadTest.java  | 73 +++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/BidiUploadState.java b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/BidiUploadState.java
index d964339ca0f6..353e9072496a 100644
--- a/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/BidiUploadState.java
+++ b/java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/BidiUploadState.java
@@ -982,6 +982,12 @@ private AppendableUploadState(
 
         if (response.hasWriteHandle()) {
           this.writeHandle = response.getWriteHandle();
+          if (!response.hasPersistedSize()
+              && !response.hasResource()
+              && state != State.TAKEOVER
+              && stateToReturnToAfterRetry != State.TAKEOVER) {
+            return null;
+          }
         }
 
         boolean incremental = !response.hasResource();
diff --git a/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/BidiUploadTest.java b/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/BidiUploadTest.java
index 2774981ecd62..948f91928436 100644
--- a/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/BidiUploadTest.java
+++ b/java-storage/google-cloud-storage/src/test/java/com/google/cloud/storage/BidiUploadTest.java
@@ -777,6 +777,79 @@ public void onResponse_writeHandleUpdated() {
       assertThat(state.writeHandle).isEqualTo(handle);
     }
 
+    @Test
+    public void onResponse_writeHandleOnly() {
+      BidiWriteHandle handle =
+          BidiWriteHandle.newBuilder().setHandle(ByteString.copyFromUtf8("new-handle")).build();
+      AppendableUploadState state = getAppendable();
+      BidiUploadTestUtils.sinkToList(state);
+      runExpectNoException(
+          state, BidiWriteObjectResponse.newBuilder().setWriteHandle(handle).build());
+      assertThat(state.writeHandle).isEqualTo(handle);
+      assertThat(state.confirmedBytes).isEqualTo(0);
+    }
+
+    @Test
+    public void awaitAck_flushWithIntermediateWriteHandleAndSmallerPersistedSize()
+        throws InterruptedException, ExecutionException, TimeoutException {
+      SettableApiFuture<BidiWriteObjectResponse> resultFuture = SettableApiFuture.create();
+      AppendableUploadState state =
+          appendableNew(
+              appendRequestNew,
+              GrpcCallContext::createDefault,
+              50_000,
+              resultFuture,
+              Crc32cValue.zero());
+      state.enqueueFirstMessageAndGetGrpcCallContext();
+      BidiUploadTestUtils.sinkToList(state);
+      assertThat(state.onResponse(resourceWithSize(0))).isNull();
+
+      assertThat(state.offer(createSegment(4))).isTrue();
+      assertThat(
+              state.offer(
+                  BidiWriteObjectRequest.newBuilder()
+                      .setWriteOffset(4)
+                      .setFlush(true)
+                      .setStateLookup(true)
+                      .build()))
+          .isTrue();
+      BidiUploadTestUtils.sinkToList(state);
+
+      BidiWriteHandle handle =
+          BidiWriteHandle.newBuilder().setHandle(ByteString.copyFromUtf8("flush-handle")).build();
+      ExecutorService exec = Executors.newSingleThreadExecutor();
+      try {
+        Future<Integer> f =
+            exec.submit(
+                () -> {
+                  try {
+                    Thread.sleep(10);
+                    assertThat(
+                            state.onResponse(
+                                BidiWriteObjectResponse.newBuilder()
+                                    .setWriteHandle(handle)
+                                    .build()))
+                        .isNull();
+                    Thread.sleep(10);
+                    assertThat(state.onResponse(incremental(2))).isNull();
+                    Thread.sleep(10);
+                    assertThat(state.onResponse(incremental(4))).isNull();
+                    return 3;
+                  } catch (Throwable t) {
+                    resultFuture.setException(t);
+                    throw new RuntimeException(t);
+                  }
+                });
+
+        state.awaitAck(4);
+        assertThat(f.get(3, TimeUnit.SECONDS)).isEqualTo(3);
+        assertThat(state.writeHandle).isEqualTo(handle);
+        assertThat(state.confirmedBytes).isEqualTo(4);
+      } finally {
+        exec.shutdownNow();
+      }
+    }
+
     @Test
     public void onResponse_notFinalizing_incremental_ackEq() {
       AppendableUploadState state = getAppendable();