feat: add should_complete to CompletionConfig

[ayushiahjolia]

Issue #, if available: #519

Description of changes:
Adds a should_complete option to CompletionConfig that lets users write custom logic for when a map or parallel batch should stop early.

Today you can only use fixed thresholds (min_successful=3, tolerated_failure_count=2). With this change, you can pass any function that looks at current progress and decides "stop now" or "keep going":

config = CompletionConfig(
    should_complete=lambda status: status.success_count >= 2
)

The function receives a CompletionStatus snapshot with counts and per-item statuses, so you can write rules like "stop when branch A succeeds OR both B and C succeed".

Key design decisions -

• When should_complete is set, threshold fields are ignored (predicate takes full precedence)
• The predicate only runs during live execution, never during replay (replay uses the checkpointed decision)
• The predicate must be deterministic and side-effect-free
• A new CUSTOM_COMPLETION completion reason is reported when the predicate stops the batch

Testing -

• Unit tests, integration tests and examples

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

items snapshot never reflects STARTED — is_started/None distinction is unreliable. needs_snapshot_rebuild is set only on the COMPLETED (321) and FAILED (327) cases. It is not set when a branch is first submitted (submit() flips it PENDING → RUNNING, i.e. None → STARTED) nor on the SUSPENDED / SUSPENDED_UNTIL cases (a status change to a started/suspended state). As a result the snapshot passed to the predicate keeps reporting started/suspended branches as status=None until some other branch reaches a terminal state and forces a rebuild.

This directly contradicts the advertised contract:

• CompletionItemStatus.is_started — "True when this branch is actively running or suspended."
• _build_items_snapshot docstring — "Distinguishes None (not yet scheduled) from STARTED (running or suspended) so predicates can reason about scheduling state."
• CompletionStatus.items docstring encourages index-based reasoning over per-branch state.

Failure scenario: max_concurrency=2, 4 branches, predicate sum(i.is_started for i in status.items) >= 2. After both branches are submitted (and even after one suspends), the loop re-enters is_complete, but the snapshot still shows every branch as None, so is_started is False for all of them and the predicate never fires on scheduling state — the batch only completes via the all-items path. Success/failure-count predicates are unaffected (those transitions do rebuild), but any predicate inspecting scheduling state gets stale/incorrect data.

Fix: also set needs_snapshot_rebuild = True inside submit() and on the SUSPENDED / SUSPENDED_UNTIL cases (any transition that changes a branch's snapshot status), or rebuild unconditionally before each is_complete evaluation when a predicate is active. Add a test that keys the predicate off is_started to lock this in.

[github-actions]

Low severity — reason() re-invokes the predicate with an empty items snapshot on the ORPHANED partial-completion path.

Every other call site that can reach the predicate now threads the per-branch snapshot: is_complete(... items_snapshot), the fallback in _create_result (self.policy.reason(succeeded, failed, fallback_items)), and from_items. This branch is the exception — it calls self.policy.reason(succeeded, failed) with the default items=().

When should_complete is set, the coordinator loop can only leave via is_complete returning True or an ORPHANED break (should_continue is always True while a predicate is active). On the is_complete/all-done exits, custom_predicate_fired is either True (→ short-circuits to CUSTOM_COMPLETION) or the batch is fully complete (reason() returns ALL_COMPLETED before touching the predicate). But on the ORPHANED break with succeeded + failed < total, custom_predicate_fired is False, so reason() runs the predicate branch and calls should_complete(_build_status(..., items=())).

Impact: a quorum-style predicate that indexes status.items (e.g. status.items[0].is_succeeded, as in parallel_with_should_complete.py) receives an empty tuple here. A predicate that doesn't guard for empty items raises IndexError, which escapes execute() on the coordinator thread and can turn otherwise-discarded orphan handling into a crash. The ORPHANED result is discarded upstream, so the reason value itself is immaterial — only the spurious predicate call matters.

Concrete fix — pass the snapshot already in scope, matching the other call sites:

Suggested change

	completion_reason = self.policy.reason(succeeded, failed)
	completion_reason = self.policy.reason(succeeded, failed, items_snapshot)

[github-actions]

Public docstring omits the "not evaluated until a terminal event" contract (low).

The coordinator only calls should_complete after at least one branch reaches a terminal state (is_complete(..., has_terminal_event) in models.py; _build_items_snapshot/has_terminal_event gating in executor.py). This is documented in the internal code comments but not in this user-facing CompletionConfig docstring, which only states the predicate "may be called multiple times per invocation."

Impact: a predicate written against scheduling/clean-state — e.g. lambda s: s.failure_count == 0 or a rule keyed on is_started items — behaves surprisingly. failure_count == 0 fires on the first successful branch (completing the whole batch at 1 item), and a purely is_started-based rule never fires at all (the parent just suspends when all in-flight branches suspend). Users can't reason about this from the public docs.

Fix: add a sentence to the should_complete docstring, e.g. "The predicate is not evaluated until at least one branch reaches a terminal (succeeded/failed) state; predicates that key only on scheduling state (is_started) or on the absence of failures will therefore first fire on the earliest terminal event."

[github-actions]

Codex AI review

• [P1] Restore checkpointed branch state before evaluating the predicate — packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py:264

  On resume, branches restart as PENDING/RUNNING; the first replayed terminal event enables should_complete before other previously checkpointed branches are restored. Event ordering can therefore change the completion reason and report a checkpointed success as STARTED, omitting its result. Initialize branch state from checkpoints or defer predicate evaluation until prior terminal branches are restored, and add a cross-invocation suspend/callback test with reversed replay event order.

Reviewed commit a3f32b8eb48b5dbdd51479f4fce122b3438927e3. Workflow run

[github-actions]

Claude AI review

I reviewed the changes introduced by PR #605 (feat: add should_complete to CompletionConfig), tracing the coordinator loop, the CompletionPolicy changes, both replay paths, serialization, and the public-API surface across the core, testing, OTel, and examples packages.

Assessment

The core design is sound and replay determinism is well-preserved:

• Replay determinism is intact. The predicate is never re-invoked on either real replay path: the normal path deserializes the checkpointed BatchResult directly, and the large-payload path uses executor.replay() driven by the checkpointed CompletionRecord (envelope_summary_generator always writes completionReason). test_replay_round_trip_custom_completion_without_reinvoking_predicate confirms this.
• custom_predicate_fired correctly distinguishes predicate-driven exit from all-completed exit; the succeeded + failed < total guard is right because the >= total branch short-circuits before the predicate in both is_complete and reason, so reason() isn't re-evaluated after the loop.
• The has_terminal_event gate + needs_snapshot_rebuild bookkeeping are consistent: the snapshot is set on every submit/terminal/suspend event, so it's never stale, and the predicate never fires before at least one branch runs.
• Additive-safe changes. CUSTOM_COMPLETION is additive; no exhaustive match on CompletionReason exists in the testing/OTel packages. Moving BatchItemStatus to config.py and re-exporting from concurrency.models keeps existing importers working. No circular import introduced (config.py uses only TYPE_CHECKING for Callable).

No confirmed high- or medium-severity correctness, determinism, concurrency, serialization, or API-compatibility defect survived verification.

Findings

Low — docstrings overpromise scheduling-state predicates vs. the terminal-event gate
packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/config.py (CompletionItemStatus.is_started docstring: "True when this branch is actively running or suspended") and concurrency/executor.py::_build_items_snapshot ("so predicates can reason about scheduling state") advertise that predicates can act on STARTED/PENDING state. But CompletionPolicy.is_complete only evaluates the predicate after has_terminal_event is set, and SUSPENDED/SUSPENDED_UNTIL events do not set it — only COMPLETED/FAILED do. A predicate keyed purely on started/suspended state with no branch ever reaching a terminal state is silently never evaluated (the batch suspends or runs to completion instead). This is intentional (the internal no-hang-guard comment documents it), but the public-facing docstrings don't warn about it. Fix: add a note to the should_complete / CompletionStatus / is_started docs that the predicate is only evaluated after at least one branch reaches a terminal (succeeded/failed) state.

Residual test risk

• No operation-level suspend/resume replay test for should_complete. test_suspend_resume_mid_batch_with_should_complete only exercises in-process timed resume within a single execute() call, and the replay test uses the replay_children path. There is no test that suspends the parent map/parallel operation across invocations and resumes to confirm the predicate re-evaluates deterministically and the checkpointed reason is stable — the feature hinges on this determinism.
• Integration-test timing bounds. test_map_with_should_complete asserts success_count <= 4, which is correct for max_concurrency=2 (at most one raced sibling), but the accompanying comment ("up to max_concurrency extra items") overstates the bound and could mislead future edits toward a flaky assertion. These @pytest.mark.example tests also depend on live scheduling timing.

Reviewed commit a3f32b8eb48b5dbdd51479f4fce122b3438927e3. Workflow run

[yaythomas]

JS uses shouldComplete: (status) => CompletionDecision with continueBatch() / completeBatch(outcome) factories and CompletionOutcome.SUCCEEDED | FAILED. The FAILED outcome marks the whole batch failed (throwIfError() throws) even with zero item failures for the quorum-can't-be-met case. A bool can't express that.

Is there a reason not to mirror the JS shape with a frozen CompletionDecision dataclass + factories?

[yaythomas]

Single CUSTOM_COMPLETION vs JS has CUSTOM_COMPLETION_SUCCEEDED / CUSTOM_COMPLETION_FAILED. This string is serialized into checkpoints, so it is used. Follows on from the CompletionDecision comment on config.py.

[yaythomas]

The has_terminal_event gate means an always-true predicate runs exactly 1 item, but in JS the same predicate stops the batch before any item runs (explicit hang guard, empty result). So the same customer predicate results in different behavior across SDKs. Is there a reason to diverge here? The CI Codex bot is independently also flagging a related resumed-replay partial-snapshot concern.

[yaythomas]

Python checks all-completed before the predicate → ALL_COMPLETED; JS checks the predicate first → CUSTOM_COMPLETION_* when both are true at once. Is there a reason not to match the JS ordering, since customers may branch on this field?

[yaythomas]

JS is {index, name?, status?}; this PR drops name (parallel branches can be named) and adds result/error_message, which JS deliberately doesn't expose. Is there a reason for the richer shape? Also the two snapshot builders disagree on error format (str(branch.error) vs item.error.message).

[yaythomas]

JS makes the combo a compile-time error; Python's analog is a ValidationError in __post_init__, not silently ignoring thresholds. Is there a reason not to raise? (This would also make the _validate_for_total skip unnecessary.)

[yaythomas]

hasattr(item.error, "message") can never be False — error is typed ErrorObject | None and ErrorObject is a dataclass where the message attribute always exists. The item.error is not None check is the only guard needed. Please see CONTRIBUTING → Data Structures & Typing: rely on exact type declarations rather than duck typing.

[yaythomas]

Inline from aws_durable_execution_sdk_python.config import BatchItemStatus inside the test body — please move to the top of the module with the other imports, per CONTRIBUTING → General style → Google Python Style Guide §3.13 ("imports are always put at the top of the file").

[yaythomas]

For posterity: replay normally uses the checkpointed completion reason and never re-invokes the predicate, but this fallback and BatchResult.from_items (for checkpoints that carry no recorded reason) re-infer the reason and therefore call the predicate again, against reconstructed items. JS does the same in its legacy-summary fallback, so no action on the re-invocation itself — but note these need updating if the decision-API change is made and we align with JS.

[yaythomas]

P1] Restore checkpointed branch state before evaluating the predicate — packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py:264

The Codex review raised a legit issue here? When a batch that suspended midway resumes, every branch restarts as PENDING and the coordinator only rediscovers checkpointed outcomes as replayed worker events arrive, in arbitrary thread order. The predicate gate opens on the first terminal event, so should_complete can evaluate a snapshot where an already checkpointed success still shows as STARTED. That means replay arrival order can change the decision, the completion reason, and which results the batch reports.

The fix Codex suggests is to initialize branch state from checkpoints, but this is troublesome. Marking branches COMPLETED up front on the coordinator thread means materializing results outside the worker and child context path, which is more invasive and touches the threading assumptions of the coordinator.

So maybe instead: count the terminal checkpoints before the loop, and defer predicate evaluation until that many terminal events have arrived again. Something like replacing the has_terminal_event bool with terminal_events_seen >= restored_target, where restored_target is computed up front via get_checkpoint_result per executable (cheap, checkpoints are already in memory). The replayed branches still flow through the normal worker path so results materialize correctly, and the first predicate evaluation is guaranteed to see everything the previous invocation knew. On a fresh run restored_target is 0, which also happens to restore the JS zero progress semantics from my other comment, so one rule closes both.

No deadlock risk. The gate only suppresses early completion, scheduling continues, and replayed branches complete near instantly, so the gate always opens and ALL_COMPLETED stays reachable.

Plus the regression test Codex asks for: suspend and resume with two terminal checkpointed branches forced to report again in reverse order, asserting the reason and result set are order independent.