Partial orchestration workitem completion support (Merge after next dts dp release)

samples/ScheduleWebApp/Orchestrations/CacheClearingOrchestratorV2.cs

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Microsoft.DurableTask;
+using ScheduleWebApp.Activities;
+
+namespace ScheduleWebApp.Orchestrations;
+
+public class CacheClearingOrchestratorV2 : TaskOrchestrator<string, string>
+{
+    public override async Task<string> RunAsync(TaskOrchestrationContext context, string scheduleId)
+    {
+        ILogger logger = context.CreateReplaySafeLogger(nameof(CacheClearingOrchestratorV2));
+        try
+        {
+            logger.LogInformation("Starting CacheClearingOrchestration for schedule ID: {ScheduleId}", scheduleId);
+
+            TaskOptions options = new TaskOptions(tags: new Dictionary<string, string>
+            {
+                { "scheduleId", scheduleId }
+            });
+
+            // Schedule all activities first to ensure deterministic ordering
+            Task<string>[] tasks = Enumerable.Range(0, 100)
+                .Select(i => context.CallActivityAsync<string>(nameof(CacheClearingActivity), new string('A', 4 * 1024), options))
+                .ToArray();
+
+            await Task.WhenAll(tasks);
+
+            logger.LogInformation("CacheClearingOrchestration completed for schedule ID: {ScheduleId}", scheduleId);
+
+            return "ok";
+        }
+        catch (Exception ex)
+        {
+            logger.LogError(ex, "Error in CacheClearingOrchestration for schedule ID: {ScheduleId}", scheduleId);
+            throw;
+        }
+    }
+}

src/Grpc/orchestrator_service.proto

@@ -356,6 +356,9 @@ message OrchestratorResponse {
 
     // Whether or not a history is required to complete the original OrchestratorRequest and none was provided.
     bool requiresHistory = 7;
+
+    // True if this is a partial (chunked) completion. The backend must keep the work item open until the final chunk (isPartial=false).
+    bool isPartial = 8;
 }
 
 message CreateInstanceRequest {

src/Grpc/versions.txt

@@ -1,2 +1,2 @@
-# The following files were downloaded from branch main at 2025-11-14 16:36:47 UTC
+# The following files were downloaded from branch main at 2025-11-25 20:53:40 UTC
 https://raw.githubusercontent.com/microsoft/durabletask-protobuf/9f762f1301b91e3e7c736b9c5a29c2e09f2a850e/protos/orchestrator_service.proto

src/InProcessTestHost/Sidecar/Grpc/TaskHubGrpcServer.cs

@@ -5,6 +5,7 @@
 using System.Diagnostics;
 using System.Globalization;
 using DurableTask.Core;
+using DurableTask.Core.Command;
 using DurableTask.Core.History;
 using DurableTask.Core.Query;
 using Google.Protobuf.WellKnownTypes;
@@ -30,6 +31,16 @@ public class TaskHubGrpcServer : P.TaskHubSidecarService.TaskHubSidecarServiceBa
 
     readonly ConcurrentDictionary<string, TaskCompletionSource<GrpcOrchestratorExecutionResult>> pendingOrchestratorTasks = new(StringComparer.OrdinalIgnoreCase);
     readonly ConcurrentDictionary<string, TaskCompletionSource<ActivityExecutionResult>> pendingActivityTasks = new(StringComparer.OrdinalIgnoreCase);
+    readonly ConcurrentDictionary<string, PartialOrchestratorChunk> partialOrchestratorChunks = new(StringComparer.OrdinalIgnoreCase);
+
+    /// <summary>
+    /// Helper class to accumulate partial orchestrator chunks.
+    /// </summary>
+    sealed class PartialOrchestratorChunk
+    {
+        public TaskCompletionSource<GrpcOrchestratorExecutionResult> TaskCompletionSource { get; set; } = null!;
+        public List<OrchestratorAction> AccumulatedActions { get; } = new();
+    }
 
     readonly ILogger log;
     readonly IOrchestrationService service;
@@ -191,7 +202,7 @@ async Task WaitForWorkItemClientConnection()
     /// <returns>A create instance response.</returns>
     public override async Task<P.CreateInstanceResponse> StartInstance(P.CreateInstanceRequest request, ServerCallContext context)
     {
-        var instance = new OrchestrationInstance
+        OrchestrationInstance instance = new OrchestrationInstance
         {
             InstanceId = request.InstanceId ?? Guid.NewGuid().ToString("N"),
             ExecutionId = Guid.NewGuid().ToString(),
@@ -471,20 +482,20 @@ static P.GetInstanceResponse CreateGetInstanceResponse(OrchestrationState state,
         {
             // Get the original orchestration state
             IList<OrchestrationState> states = await this.client.GetOrchestrationStateAsync(request.InstanceId, false);
-            
+
             if (states == null || states.Count == 0)
             {
                 throw new RpcException(new Status(StatusCode.NotFound, $"An orchestration with the instanceId {request.InstanceId} was not found."));
             }
 
             OrchestrationState state = states[0];
-            
+
             // Check if the state is null
             if (state == null)
             {
                 throw new RpcException(new Status(StatusCode.NotFound, $"An orchestration with the instanceId {request.InstanceId} was not found."));
             }
-            
+
             string newInstanceId = request.RestartWithNewInstanceId ? Guid.NewGuid().ToString("N") : request.InstanceId;
 
             // Create a new orchestration instance
@@ -535,6 +546,51 @@ static P.GetInstanceResponse CreateGetInstanceResponse(OrchestrationState state,
     /// <returns>Returns an empty ack back to the remote SDK that we've received the completion.</returns>
     public override Task<P.CompleteTaskResponse> CompleteOrchestratorTask(P.OrchestratorResponse request, ServerCallContext context)
     {
+        if (request.IsPartial)
+        {
+            // This is a partial chunk - accumulate actions but don't complete yet
+            PartialOrchestratorChunk partialChunk = this.partialOrchestratorChunks.GetOrAdd(
+                request.InstanceId,
+                _ =>
+                {
+                    // First chunk - get the TCS and initialize the partial chunk
+                    if (!this.pendingOrchestratorTasks.TryGetValue(request.InstanceId, out TaskCompletionSource<GrpcOrchestratorExecutionResult>? tcs))
+                    {
+                        throw new RpcException(new Status(StatusCode.NotFound, $"Orchestration with instance ID '{request.InstanceId}' not found"));
+                    }
+
+                    return new PartialOrchestratorChunk
+                    {
+                        TaskCompletionSource = tcs,
+                    };
+                });
+
+            // Accumulate actions from this chunk
+            partialChunk.AccumulatedActions.AddRange(request.Actions.Select(ProtobufUtils.ToOrchestratorAction));
+
+            return EmptyCompleteTaskResponse;
+        }
+
+        // This is the final chunk (or a single non-chunked response)
+        if (this.partialOrchestratorChunks.TryRemove(request.InstanceId, out PartialOrchestratorChunk? existingPartialChunk))
+        {
+            // We've been accumulating chunks - combine with final chunk
+            existingPartialChunk.AccumulatedActions.AddRange(request.Actions.Select(ProtobufUtils.ToOrchestratorAction));
+
+            GrpcOrchestratorExecutionResult res = new()
+            {
+                Actions = existingPartialChunk.AccumulatedActions,
+                CustomStatus = request.CustomStatus, // Use custom status from final chunk
+            };
+
+            // Remove the TCS from pending tasks and complete it
+            this.pendingOrchestratorTasks.TryRemove(request.InstanceId, out _);
+            existingPartialChunk.TaskCompletionSource.TrySetResult(res);
+
+            return EmptyCompleteTaskResponse;
+        }
+
+        // Single non-chunked response (no partial chunks)
         if (!this.pendingOrchestratorTasks.TryRemove(
             request.InstanceId,
             out TaskCompletionSource<GrpcOrchestratorExecutionResult>? tcs))
@@ -646,6 +702,7 @@ public override async Task GetWorkItems(P.GetWorkItemsRequest request, IServerSt
         }
     }
 
+    /// <inheritdoc/>
     public override async Task StreamInstanceHistory(P.StreamInstanceHistoryRequest request, IServerStreamWriter<P.HistoryChunk> responseStream, ServerCallContext context)
     {
         if (this.streamingPastEvents.TryGetValue(request.InstanceId, out List<P.HistoryEvent>? pastEvents))
@@ -684,7 +741,7 @@ async Task<GrpcOrchestratorExecutionResult> ITaskExecutor.ExecuteOrchestrator(
         IEnumerable<HistoryEvent> pastEvents,
         IEnumerable<HistoryEvent> newEvents)
     {
-        var executionStartedEvent = pastEvents.OfType<ExecutionStartedEvent>().FirstOrDefault();
+        ExecutionStartedEvent? executionStartedEvent = pastEvents.OfType<ExecutionStartedEvent>().FirstOrDefault();
 
         P.OrchestrationTraceContext? orchestrationTraceContext = executionStartedEvent?.ParentTraceContext?.SpanId is not null
             ? new P.OrchestrationTraceContext
@@ -701,7 +758,7 @@ async Task<GrpcOrchestratorExecutionResult> ITaskExecutor.ExecuteOrchestrator(
 
         try
         {
-            var orkRequest = new P.OrchestratorRequest
+            P.OrchestratorRequest orkRequest = new P.OrchestratorRequest
             {
                 InstanceId = instance.InstanceId,
                 ExecutionId = instance.ExecutionId,
@@ -829,6 +886,7 @@ TaskCompletionSource<GrpcOrchestratorExecutionResult> CreateTaskCompletionSource
     void RemoveOrchestratorTaskCompletionSource(string instanceId)
     {
         this.pendingOrchestratorTasks.TryRemove(instanceId, out _);
+        this.partialOrchestratorChunks.TryRemove(instanceId, out _);
     }
 
     TaskCompletionSource<ActivityExecutionResult> CreateTaskCompletionSourceForActivity(string instanceId, int taskId)

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -2,11 +2,13 @@
 // Licensed under the MIT License.
 
 using System.Diagnostics;
+using System.Linq;
 using System.Text;
 using DurableTask.Core;
 using DurableTask.Core.Entities;
 using DurableTask.Core.Entities.OperationFormat;
 using DurableTask.Core.History;
+using Google.Protobuf;
 using Microsoft.DurableTask.Abstractions;
 using Microsoft.DurableTask.Entities;
 using Microsoft.DurableTask.Tracing;
@@ -755,7 +757,11 @@ await this.client.AbandonTaskOrchestratorWorkItemAsync(
                 response.Actions.Count,
                 GetActionsListForLogging(response.Actions));
 
-            await this.client.CompleteOrchestratorTaskAsync(response, cancellationToken: cancellationToken);
+            // Auto-chunk the response if it exceeds the maximum size
+            await this.CompleteOrchestratorTaskWithChunkingAsync(
+                response,
+                this.worker.grpcOptions.MaxCompleteOrchestrationWorkItemSizePerChunk,
+                cancellationToken);
         }
 
         async Task OnRunActivityAsync(P.ActivityRequest request, string completionToken, CancellationToken cancellation)
@@ -914,5 +920,101 @@ async Task OnRunEntityBatchAsync(
 
             await this.client.CompleteEntityTaskAsync(response, cancellationToken: cancellation);
         }
+
+        /// <summary>
+        /// Completes an orchestration task with automatic chunking if the response exceeds the maximum size.
+        /// </summary>
+        /// <param name="response">The orchestrator response to send.</param>
+        /// <param name="maxChunkBytes">The maximum size in bytes for each chunk.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        async Task CompleteOrchestratorTaskWithChunkingAsync(
+            P.OrchestratorResponse response,
+            int maxChunkBytes,
+            CancellationToken cancellationToken)
+        {
+            // Validate that no single action exceeds the maximum chunk size
+            static void ValidateActionsSize(IEnumerable<P.OrchestratorAction> actions, int maxChunkBytes)
+            {
+                foreach (P.OrchestratorAction action in actions)
+                {
+                    int actionSize = action.CalculateSize();
+                    if (actionSize > maxChunkBytes)
+                    {
+                        throw new InvalidOperationException(
+                            $"A single orchestrator action of type {action.OrchestratorActionTypeCase} with id {action.Id} " +
+                            $"exceeds the {maxChunkBytes / 1024.0 / 1024.0:F2}MB limit: {actionSize / 1024.0 / 1024.0:F2}MB");
+                    }
+                }
+            }
+
+            ValidateActionsSize(response.Actions, maxChunkBytes);
+
+            // Helper to add an action to the current chunk if it fits
+            static bool TryAddAction(
+                Google.Protobuf.Collections.RepeatedField<P.OrchestratorAction> dest,
+                P.OrchestratorAction action,
+                ref int currentSize,
+                int maxChunkBytes)
+            {
+                int actionSize = action.CalculateSize();
+                if (currentSize + actionSize > maxChunkBytes)
+                {
+                    return false;
+                }
+
+                dest.Add(action);
+                currentSize += actionSize;
+                return true;
+            }
+
+            // Check if the entire response fits in one chunk
+            int totalSize = response.CalculateSize();
+            if (totalSize <= maxChunkBytes)
+            {
+                // Response fits in one chunk, send it directly (isPartial defaults to false)
+                await this.client.CompleteOrchestratorTaskAsync(response, cancellationToken: cancellationToken);
+                return;
+            }
+
+            // Response is too large, split into multiple chunks
+            int actionsCompletedSoFar = 0;
+            bool isFirstChunk = true;
+            List<P.OrchestratorAction> allActions = response.Actions.ToList();
+
+            while (actionsCompletedSoFar < allActions.Count)
+            {
+                P.OrchestratorResponse chunkedResponse = new()
+                {
+                    InstanceId = response.InstanceId,
+                    CustomStatus = response.CustomStatus,
+                    CompletionToken = response.CompletionToken,
+                    OrchestrationTraceContext = isFirstChunk ? response.OrchestrationTraceContext : null, // Only include trace context in first chunk
+                    RequiresHistory = response.RequiresHistory,
+                };
+
+                int chunkSize = chunkedResponse.CalculateSize();
+
+                // Fill the chunk with actions until we reach the size limit
+                while (actionsCompletedSoFar < allActions.Count &&
+                       TryAddAction(chunkedResponse.Actions, allActions[actionsCompletedSoFar], ref chunkSize, maxChunkBytes))
+                {
+                    actionsCompletedSoFar++;
+                }
+
+                // Determine if this is a partial chunk (more actions remaining)
+                bool isPartial = actionsCompletedSoFar < allActions.Count;
+                chunkedResponse.IsPartial = isPartial;
+
+                if (isFirstChunk)
+                {
+                    isFirstChunk = false;
+                } else {
+                    chunkedResponse.NumEventsProcessed = -1;
+                }
+
+                // Send the chunk
+                await this.client.CompleteOrchestratorTaskAsync(chunkedResponse, cancellationToken: cancellationToken);
+            }
+        }
     }
 }

src/Worker/Grpc/GrpcDurableTaskWorkerOptions.cs

@@ -31,6 +31,16 @@ public sealed class GrpcDurableTaskWorkerOptions : DurableTaskWorkerOptions
     /// </summary>
     public HashSet<P.WorkerCapability> Capabilities { get; } = new() { P.WorkerCapability.HistoryStreaming };
 
+    /// <summary>
+    /// Gets or sets the maximum size of all actions in a complete orchestration work item chunk.
+    /// The default value is 3.9MB. We leave some headroom to account for request size overhead.
+    /// </summary>
+    /// <remarks>
+    /// This value is used to limit the size of the complete orchestration work item chunk request.
+    /// If the response exceeds this limit, it will be automatically split into multiple chunks.
+    /// </remarks>
+    public int MaxCompleteOrchestrationWorkItemSizePerChunk { get; set; } = 4089446; // 3.9MB
+
     /// <summary>
     /// Gets the internal protocol options. These are used to control backend-dependent features.
     /// </summary>