more fixes

Author: bachuv

azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/BlobPayloadStore.java

@@ -9,9 +9,7 @@
 import com.azure.storage.blob.BlobServiceClientBuilder;
 import com.azure.storage.blob.models.BlobDownloadResponse;
 import com.azure.storage.blob.models.BlobHttpHeaders;
-import com.azure.storage.blob.models.BlobRange;
 import com.azure.storage.blob.models.BlobStorageException;
-import com.azure.storage.blob.options.BlobDownloadToFileOptions;
 import com.azure.storage.common.policy.RequestRetryOptions;
 import com.azure.storage.common.policy.RetryPolicyType;
 
@@ -20,7 +18,6 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
-import java.time.Duration;
 import java.util.UUID;
 import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;

azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/LargePayloadInterceptor.java

@@ -8,6 +8,7 @@
 
 import io.grpc.*;
 
+import java.nio.charset.StandardCharsets;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -767,23 +768,12 @@ private static String getStackTraceString(Throwable t) {
     }
 
     /**
-     * Computes the UTF-8 encoded byte length of a string without allocating a byte array.
+     * Computes the UTF-8 encoded byte length of a string.
+     * <p>
+     * Uses Java's canonical UTF-8 encoding behavior so malformed surrogate
+     * sequences are measured exactly the same way as payload serialization.
      */
     private static int utf8ByteLength(String s) {
-        int count = 0;
-        for (int i = 0; i < s.length(); i++) {
-            char ch = s.charAt(i);
-            if (ch <= 0x7F) {
-                count++;
-            } else if (ch <= 0x7FF) {
-                count += 2;
-            } else if (Character.isHighSurrogate(ch)) {
-                count += 4;
-                i++; // skip low surrogate
-            } else {
-                count += 3;
-            }
-        }
-        return count;
+        return s.getBytes(StandardCharsets.UTF_8).length;
     }
 }

client/src/main/java/com/microsoft/durabletask/DurableTaskGrpcWorker.java

@@ -649,6 +649,7 @@ private void completeOrchestratorTaskWithChunking(OrchestratorResponse response)
         int chunkIndex = 0;
         boolean isPartial = true;
         boolean isChunkedMode = false;
+        boolean hasTraceContext = response.hasOrchestrationTraceContext();
 
         while (isPartial) {
             OrchestratorResponse.Builder chunk = OrchestratorResponse.newBuilder()
@@ -657,15 +658,23 @@ private void completeOrchestratorTaskWithChunking(OrchestratorResponse response)
                 .setCompletionToken(response.getCompletionToken())
                 .setRequiresHistory(response.getRequiresHistory());
 
-            int chunkPayloadSize = 0;
             while (actionsCompleted < allActions.size()) {
-                int actionSize = allActions.get(actionsCompleted).getSerializedSize();
+                OrchestratorAction nextAction = allActions.get(actionsCompleted);
+
+                chunk.addActions(nextAction);
+
+                int estimatedSize = estimateChunkSerializedSize(
+                    chunk,
+                    chunkIndex,
+                    hasTraceContext,
+                    response.getOrchestrationTraceContext());
+
                 // Always accept the first action in an empty chunk to avoid infinite loops
-                if (chunkPayloadSize + actionSize > maxChunkBytes && chunkPayloadSize > 0) {
+                if (estimatedSize > maxChunkBytes && chunk.getActionsCount() > 1) {
+                    chunk.removeActions(chunk.getActionsCount() - 1);
                     break;
                 }
-                chunk.addActions(allActions.get(actionsCompleted));
-                chunkPayloadSize += actionSize;
+
                 actionsCompleted++;
             }
 
@@ -696,6 +705,32 @@ private void completeOrchestratorTaskWithChunking(OrchestratorResponse response)
         }
     }
 
+    /**
+     * Estimates the serialized size of a candidate chunk including envelope overhead fields
+     * that may be added later in the chunking flow.
+     */
+    private static int estimateChunkSerializedSize(
+            OrchestratorResponse.Builder chunk,
+            int chunkIndex,
+            boolean hasTraceContext,
+            OrchestrationTraceContext traceContext) {
+        OrchestratorResponse.Builder estimate = chunk.clone();
+
+        // Include potential overhead fields to avoid under-estimating chunk size.
+        estimate.setIsPartial(true);
+        estimate.setChunkIndex(Int32Value.of(chunkIndex));
+
+        if (chunkIndex == 0) {
+            if (hasTraceContext) {
+                estimate.setOrchestrationTraceContext(traceContext);
+            }
+        } else {
+            estimate.setNumEventsProcessed(Int32Value.of(0));
+        }
+
+        return estimate.build().getSerializedSize();
+    }
+
     static WorkItemFilters toProtoWorkItemFilters(WorkItemFilter filter) {
         WorkItemFilters.Builder builder = WorkItemFilters.newBuilder();
         for (WorkItemFilter.OrchestrationFilter orch : filter.getOrchestrations()) {