feat: add custom durable operation SPI

[zhongkechen]

Summary

• add a stable ExtensionContext and opaque one-shot operation reservation API for third-party durable operation libraries
• support deterministic caller-defined local operation IDs, custom operation subtypes, replay metadata validation, and stateful extension steps
• make _operation the canonical implementation layer for all SDK operations, including step(), wait(), invoke(), create_callback(), and run_in_child_context()
• route built-in primitive and composite operations through the same custom operation SPI used by third-party extensions
• keep checkpoint state machines and the stateful extension STEP executor internal to _primitive, with the SPI delegating to those executors
• preserve top-level public APIs, helper signatures, durable operation IDs, names, subtypes, serialization, and replay behavior
• keep legacy _primitive helper imports as lazy compatibility delegates and legacy _extension.<module> imports as aliases to _operation
• accept historical _extension.wait_for_condition.WaitForConditionError checkpoint metadata during replay
• preserve failed STEP history for exhausted wait_for_callback submitter retries

This follows the extension-authoring direction of aws/aws-durable-execution-sdk-java#607 without exposing the SDK checkpoint machinery.

Validation

• hatch run test:sdk -q (2,223 passed)
• hatch run test:examples -q (89 passed)
• hatch run test:typecheck
• hatch fmt --check
• hatch run docs:build
• hatch build -t wheel

Fixes #279

[zhongkechen]

Addressed in 0705af5:

• Replay transitions now happen when a reserved operation selects a primitive, using that primitive’s user-code semantics. Child contexts inherit the pre-transition replay state, and terminal explicit local IDs leave replay when selected.
• Stateful extension failures now restore and raise SDK control-error categories immediately after checkpointing, so first execution and replay expose the same exception type and catch behavior.
• Added regression tests for inert reservations, terminal explicit-ID selection, child-context replay inheritance, and exact first-run/replay control-error type equality.

Verified locally: 2382 passed, formatting/Ruff clean, mypy clean, and strict docs build successful.

[zhongkechen]

Addressed the latest Codex review in 1814408 (rebased onto current main):

• Each ExtensionOperation now snapshots replay state when reserved and passes that snapshot into child-context creation. A regression test pre-reserves two historical children, lets the first flip the mutable parent state, and verifies both children still observe replay mode.
• Public wait() now validates and converts its duration before reserving an operation ID. A regression test verifies wait(0) never calls the reservation allocator.
• Strict nonblank-name validation now applies only to the public extension-author reserve() API. SDK-owned helpers use a compatibility reservation path that normalizes "" to None, preserves whitespace names, and is used by all primitive and composite operations.

Validation: 2,223 SDK tests and 89 example tests pass; mypy, Ruff/formatting, strict docs, and wheel build are clean.

[zhongkechen]

Addressed the latest Codex findings in 389918f:\n\n- Preserved replay state when a virtual child follows a completed parent operation by carrying an in-memory replay-frontier snapshot into the virtual scope. The marker is cleared at the next known operation or virtual boundary, so prior parent-frontier, cancellation, lazy-sibling, FLAT, and restart fixes remain intact.\n- Added LambdaContext | None to the public ExtensionContext.lambda_context property and direct runtime type-hint coverage.\n- Added an end-to-end regression for completed step -> virtual child -> nested wait replay.\n\nValidation: 2,252 SDK tests, 89 example tests, mypy, Ruff lint/format, strict docs, and wheel build all pass. The prior Claude job failed during checkout before generating a review; watching the fresh review cycle now.

[zhongkechen]

AI review status on current head 7a74cfd:

• Codex completed and reported no actionable findings.
• Claude was rerun four times after the initial attempt. Every substantive attempt exhausted the trusted base workflow fixed 20-turn budget (error_max_turns, 21 turns used), returned no structured review, and posted no buffered findings. One earlier attempt also failed checkout before review generation.
• The turn limit is defined on main; because this is a pull_request_target workflow, this PR cannot change the workflow used to review itself.

No additional review comment needs a code change. The product fixes remain validated at 2,252 SDK tests, 89 examples, mypy, Ruff, strict docs, and wheel build.

[zhongkechen]

Addressed the completed-replay reservation finding in d304b29.

ParallelExecutor.replay_completed() now obtains branch operation IDs without materializing _BranchOperationReservations, so absent branches do not invoke item_namer or allocate reservations. Existing branches that require ReplayChildren still materialize normally when executed.

Added an early-completion regression with one historical branch and one unstarted item whose namer raises if touched.

Validation: 2,253 SDK tests, 89 examples, mypy, Ruff, strict docs, and wheel build pass.

[github-actions]

Codex AI review

No actionable findings. Residual risk: custom subtype and stateful retry behavior is tested through the local runner, not a cloud-backed durable execution.

Reviewed commit d304b29c91de65a390934dd8fbf49418b31c7bfc. Workflow run

[github-actions]

Claude AI review

Review Summary

I read the full PR diff (74 files, ~4.6k/0.7k +/-), AGENTS.md, and the base-revision source for the most safety-critical touched files (_core/context.py, _primitive/base.py, _primitive/step.py, _primitive/child.py, _runner/local/processors/{base,step}.py), since the sandboxed shell could not execute the test suite or run git diff directly (all Bash invocations failed with a bwrap permission error unrelated to the commands themselves).

This is a large, well-scoped refactor: it introduces a public async_durable_execution.extension SPI (ExtensionContext/ExtensionOperation/ExtensionStepResult), moves all built-in operations (step, wait, invoke, create_callback, run_in_child_context, and the composite map/parallel/flow/with_retry/wait_for_condition/wait_for_callback) onto that same reservation API, and keeps _primitive as the internal checkpoint state-machine layer with _extension/_primitive legacy import paths preserved as aliases. The replay bookkeeping (OperationIdGenerator reservation tracking, _replay_aware, the new "replay frontier" flag for virtual/lazy children) is intricate but is backed by a large, targeted new test file (tests/extension/custom_operation_test.py) plus updated tests in context_test.py, parallel_test.py, map_test.py, and the local-runner processor tests, covering launch-order changes, local-ID reservations, cancellation, virtual/flat/bounded-parallel branches, and legacy-checkpoint compatibility.

No confirmed correctness bug was found. Two low-severity, likely-intentional behavior changes are worth calling out:

• Low — async_durable_execution/_runner/model.py:~1528-1533 (events_to_operations): an unrecognized sub_type string from a backend/cloud event no longer raises InvalidParameterValueException; it is now silently kept as a raw string to support extension-defined subtypes end-to-end. This is intentional and tested (test_events_to_operations_preserves_custom_sub_type), but it removes a fail-fast guard against malformed/typo'd subtype values coming from the (simulated) backend for built-in operations too. Consider validating that the string is either a known OperationSubType or otherwise plausible (e.g., only accept unknown strings when the corresponding operation_type is one the SDK doesn't own) rather than accepting any string unconditionally.
• Low — async_durable_execution/extension.py (_normalize_sdk_operation_name, ~line 86): SDK-owned helpers (step(), wait(), invoke(), etc.) now silently convert an explicit name="" to None, whereas previously "" was stored verbatim. This is deliberate (to mirror how the real backend omits blank names) and is covered by test_sdk_reservations_preserve_legacy_blank_name_behavior, but it is an observable behavior change for any caller/checkpoint inspection code that compares operation names against the literal empty string.

Residual risk: the replay-frontier / reservation-consumption logic in _core/context.py (OperationIdGenerator._register_reservation/_consume_reservation, _replay_aware, _virtual_child_replay_snapshot) and its use in extension.py's lazy child-context/parallel-branch reservation paths is the highest-complexity and highest-impact part of this change for replay determinism. I could not execute hatch run test:sdk/test:examples in this environment to independently verify the PR's claimed "2,223 passed" result, so CI results for this PR should be treated as the authoritative confirmation before merge, with particular attention to interleaved sequential + local-ID reservations across suspend/resume boundaries in real (non-local-runner) Lambda invocations.

Reviewed commit d304b29c91de65a390934dd8fbf49418b31c7bfc. Workflow run