Local runner does not complete a due wait until the invocation returns

[yaythomas]

Summary

In DurableFunctionTestRunner, a wait that becomes due while the handler is
still running is not completed until that invocation returns. A parallel or map
branch that waits therefore cannot resume in the same invocation, even though the
sibling branches keep it alive. Against the deployed service the same workload
resumes inline in a single invocation.

The same applies to a step retry whose next-attempt time elapses mid-invocation.

Impact

Any local test where one branch waits while another keeps working sees a
different invocation count and a different timeline than the deployed service.
Tests that assert on invocation counts, or that assert a branch made progress
while a sibling was still running, either fail locally or pass for the wrong
reason. The divergence is silent: nothing errors, the execution just takes an
extra invocation.

Reproduction

_SIBLING_WORK_SECONDS = 3.0
_WAIT_SECONDS = 1

@durable_step
def do_work(step_context: StepContext, seconds: float) -> str:
    time.sleep(seconds)
    return "work-done"

@durable_execution
def parallel_wait_and_work(event, context: DurableContext) -> list[str]:
    def branch_a(ctx: DurableContext) -> str:
        print("A: before wait")
        ctx.wait(Duration.from_seconds(_WAIT_SECONDS))
        print("A: after wait")
        return "a-done"

    def branch_b(ctx: DurableContext) -> str:
        print("B: start work")
        ctx.step(do_work(_SIBLING_WORK_SECONDS))
        print("B: finished work")
        return "b-done"

    return [item.result for item in context.parallel([branch_a, branch_b]).all]

with DurableFunctionTestRunner(handler=parallel_wait_and_work, skip_time=False) as runner:
    result = runner.run(input="in")

skip_time=False is required; with the skip clock the wait is fast-forwarded and
the problem is hidden.

Observed

Two invocations. Branch A re-enters on schedule, finds its wait still STARTED,
and re-suspends. It only gets past the wait in a second invocation, after B has
finished and the first invocation has returned.

   0.00s  INVOCATION START / A: before wait / B: start work
   2.21s  A: before wait          <- resumed, wait still STARTED, re-suspends
   4.31s  A: before wait
   6.41s  A: before wait
   8.51s  A: before wait
  10.00s  B: finished work
  10.21s  INVOCATION START (#2)
  10.21s  A: after wait

(Timeline above from a 2 second wait and a 10 second sibling, to make the repeated
re-entry visible.)

Polling the execution store during the run confirms the wait operation stays
STARTED past its scheduled end time and flips to SUCCEEDED only after the
first invocation returns.

Expected

One invocation. The wait completes when it is due and the branch continues
alongside its still-running sibling. This is what the deployed service does. The
same workload on Lambda (wait(5), sibling busy 20s), with the invocation id
printed on every line:

PROBE inv=35078b21 t=   0.00s INVOCATION START
PROBE inv=35078b21 t=   0.00s A: before wait(5)
PROBE inv=35078b21 t=   0.00s B: start 20s of work
PROBE inv=35078b21 t=   5.34s A: before wait(5)
PROBE inv=35078b21 t=   5.34s A: after wait
PROBE inv=35078b21 t=  20.00s B: finished work
PROBE inv=35078b21 t=  20.32s parallel returned reason=CompletionReason.ALL_COMPLETED

One invocation id throughout. The branch resumes at 5.34s, while the sibling
still has 15 seconds of work left.

Cause

Executor._schedule_earliest_pending arms the single wake-up timer that later
completes due waits and step retries. It is called from
Executor._checkpoint_execution (reached only over HTTP, from
web/handlers.py), from fire_due re-arming itself, and from
_finish_invocation after an invocation returns.

DurableFunctionTestRunner does not use the HTTP path. Its service client is
InMemoryServiceClient, whose checkpoint() submits a CheckpointTask that
calls CheckpointProcessor.process_checkpoint directly, and that flow never arms
the timer. So no timer is armed at any point during an in-process invocation, and
the first arm for a wait started in that invocation happens only at
_finish_invocation.

WebRunner is unaffected, since its checkpoints arrive over HTTP.

Suggested fix

Arm the earliest-pending wake-up after an in-process checkpoint commits, the same
way _checkpoint_execution already does on the HTTP path.

Not covered by this issue

• When a re-entering branch finds its wait still pending, the SDK re-suspends for
  the full configured duration rather than the remaining time
  (aws_durable_execution_sdk_python/operation/wait.py). Against the service the
  first re-check already sees the wait completed, so this is not observable there,
  but it is why the timeline above shows re-entry every 2 seconds rather than a
  tight retry.
• The local runner does not simulate invocation timeouts. Separate gap.

Environment

• aws-durable-execution-sdk-python-testing at main (fdc952b)
• Python 3.14.6, Linux
• Deployed comparison run in us-west-2, python3.13 runtime