fix: send an oversized first operation in a batch

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/state.py

@@ -154,6 +154,12 @@ def is_full(self) -> bool:
         return self.effective_count >= self._config.max_batch_operations
 
     def fits(self, size: int) -> bool:
+        # An empty batch accepts an operation of any size. An operation above
+        # the limit can only ever be sent alone, so refusing it here would
+        # return it to its queue on every pass, and its sync caller would block
+        # forever. The service, not the batcher, enforces the per-operation limit.
+        if self.is_empty():
+            return True
         return self.total_size + size <= self._config.max_batch_size_bytes
 
     def add(self, op: QueuedOperation, size: int = 0) -> None:
@@ -1487,8 +1493,7 @@ def _wait_for_first_operation(self, batch: _BatchAccumulator) -> bool:
             if isinstance(item, _Signal):
                 self._consume_refresh_wake()
                 continue
-            # The first operation is never refused for size. It could not be
-            # sent otherwise.
+            # The first operation is never refused for size; see fits().
             batch.add(item, self._calculate_operation_size(item))
             return True
         return False

packages/aws-durable-execution-sdk-python/tests/e2e/large_step_result_int_test.py

@@ -0,0 +1,36 @@
+"""A step result above the checkpoint batch size limit completes (issue #871).
+
+The step result is 800 KB. That is above ``max_batch_size_bytes`` (750 KB) and
+below the service's 1 MB per-operation limit. With the default
+``AT_LEAST_ONCE_PER_RETRY`` semantics the START checkpoint is async, so the
+batcher still holds it when the sync SUCCEED arrives, and SUCCEED takes the
+overflow path. Before the fix the SUCCEED operation was never sent and the
+handler thread blocked forever, so ``run()`` raised ``TimeoutError``.
+"""
+
+from __future__ import annotations
+
+import json
+
+from aws_durable_execution_sdk_python import DurableContext, durable_execution
+from aws_durable_execution_sdk_python.execution import InvocationStatus
+from aws_durable_execution_sdk_python_testing.runner import (
+    DurableFunctionTestResult,
+    DurableFunctionTestRunner,
+)
+
+RESULT_SIZE_BYTES = 800 * 1024
+
+
+@durable_execution
+def _large_step_handler(_event: dict[str, str], context: DurableContext) -> str:
+    return context.step(lambda _sc: "y" * RESULT_SIZE_BYTES, name="big_step")
+
+
+def test_step_result_above_batch_size_limit_completes() -> None:
+    with DurableFunctionTestRunner(handler=_large_step_handler) as runner:
+        result: DurableFunctionTestResult = runner.run(input="{}", execution_timeout=30)
+
+    assert result.status is InvocationStatus.SUCCEEDED
+    assert result.result is not None
+    assert len(json.loads(result.result)) == RESULT_SIZE_BYTES

packages/aws-durable-execution-sdk-python/tests/state_test.py

@@ -56,6 +56,7 @@
 from aws_durable_execution_sdk_python.state import (
     _BLOCKED_CALLER_WAIT_SECONDS,
     _STOP_SIGNAL_POLL_SECONDS,
+    _BatchAccumulator,
     _Signal,
     CheckpointBatcherConfig,
     CheckpointedResult,
@@ -1367,6 +1368,96 @@ def test_collect_checkpoint_batch_uses_overflow_queue():
     assert batch[1].operation_update.operation_id == "overflow_2"
 
 
+def _oversized_queued_operation(
+    operation_id: str, max_batch_size_bytes: int
+) -> QueuedOperation:
+    """A sync operation whose serialized size alone exceeds the batch limit."""
+    return QueuedOperation(
+        OperationUpdate(
+            operation_id=operation_id,
+            operation_type=OperationType.STEP,
+            action=OperationAction.SUCCEED,
+            payload="y" * (2 * max_batch_size_bytes),
+        ),
+        CompletionEvent(),
+    )
+
+
+def test_collect_checkpoint_batch_sends_oversized_operation_from_overflow() -> None:
+    """An operation above max_batch_size_bytes that reached the overflow queue is
+    still sent, alone, in the next batch (issue #871).
+
+    This is the default AT_LEAST_ONCE_PER_RETRY step shape: a small START is
+    already in the batch when the large SUCCEED arrives, so SUCCEED does not fit
+    and moves to the overflow queue. The next batch must send it rather than
+    return it to the queue forever.
+    """
+    mock_lambda_client = Mock(spec=LambdaClient)
+    config = CheckpointBatcherConfig(max_batch_size_bytes=500)
+    state = ExecutionState(
+        durable_execution_arn="test_arn",
+        initial_checkpoint_token="token123",  # noqa: S106
+        operations={},
+        service_client=mock_lambda_client,
+        plugin_executor=PluginExecutor(plugins=None),
+        batcher_config=config,
+    )
+
+    small_op = QueuedOperation(
+        OperationUpdate(
+            operation_id="big_step",
+            operation_type=OperationType.STEP,
+            action=OperationAction.START,
+        ),
+        None,
+    )
+    large_op = _oversized_queued_operation("big_step", config.max_batch_size_bytes)
+    assert state._calculate_operation_size(large_op) > config.max_batch_size_bytes
+    state._checkpoint_queue.put(small_op)
+    state._checkpoint_queue.put(large_op)
+
+    first_batch = state._collect_checkpoint_batch()
+    assert first_batch == [small_op]
+    assert state._overflow_queue.qsize() == 1
+
+    # The handler thread is blocked on large_op, so no further operation will
+    # arrive. Stopping the batcher makes a regression return [] instead of
+    # blocking the test forever in _wait_for_first_operation.
+    state._checkpointing_stopped.set()
+    second_batch = state._collect_checkpoint_batch()
+    assert second_batch == [large_op]
+    assert state._overflow_queue.empty()
+
+
+def test_drain_overflow_takes_oversized_operation_when_batch_is_empty() -> None:
+    """_drain_overflow adds an operation above the size limit to an empty batch,
+    and refuses it once the batch holds another operation.
+    """
+    mock_lambda_client = Mock(spec=LambdaClient)
+    config = CheckpointBatcherConfig(max_batch_size_bytes=500)
+    state = ExecutionState(
+        durable_execution_arn="test_arn",
+        initial_checkpoint_token="token123",  # noqa: S106
+        operations={},
+        service_client=mock_lambda_client,
+        plugin_executor=PluginExecutor(plugins=None),
+        batcher_config=config,
+    )
+    first_large = _oversized_queued_operation("large_1", config.max_batch_size_bytes)
+    second_large = _oversized_queued_operation("large_2", config.max_batch_size_bytes)
+    state._overflow_queue.put(first_large)
+    state._overflow_queue.put(second_large)
+
+    batch = _BatchAccumulator(config)
+    state._drain_overflow(batch)
+
+    # The first operation is taken despite its size. The second does not fit
+    # beside it and stays queued for the next batch.
+    assert batch.operations == [first_large]
+    assert state._overflow_queue.qsize() == 1
+    assert state._overflow_queue.get_nowait() is second_large
+
+
 def test_collect_checkpoint_batch_handles_empty_checkpoint():
     """Test batch collection with empty checkpoints (None operation_update)."""
     mock_lambda_client = Mock(spec=LambdaClient)