feat(graph): parallel branches (proposal 0011)

[chris-colinsky]

Summary

• Implements spec proposal 0011 (parallel-branches dispatch primitive) in openarmature.graph. PR-5 of the five-PR batch following PR-1 (feat(llm): structured output (proposal 0016) #42), PR-2 (feat(llm): image content blocks (proposal 0015) #44), PR-3 (feat(prompts): prompt-management core (proposal 0017) #45), PR-4 (feat(checkpoint): state migration for checkpoints (proposal 0014) #46). With this PR, the batch is complete.
• New GraphBuilder.add_parallel_branches_node(name, *, branches, error_policy, errors_field, middleware) surface dispatches M heterogeneous compiled subgraphs concurrently per pipeline-utilities §11. BranchSpec (frozen dataclass with subgraph + inputs/outputs projection + branch middleware) and ParallelBranchesNode types exported from openarmature.graph.
• Branch insertion order determines fan-in merge order regardless of inner-node completion timing (§11.8). Single-branch field contributions flow through the parent's reducer as plain values; multi-branch contributions to one field use a _MultiContribution sentinel that _merge_partial folds through the reducer in branch insertion order.
• Two error policies: "fail_fast" raises ParallelBranchesBranchFailed (a NodeException subtype) with branch_name, original cause as __cause__, and recoverable_state carrying the parent's pre-dispatch snapshot — no buffered branch contributions are applied (§11.5 buffer-and-apply). "collect" runs every branch to completion, records per-branch failures in an optional errors_field ({branch_name, category, message, cause_type}), and continues.
• Two new error categories: ParallelBranchesNoBranches (compile time, empty branches map) and ParallelBranchesBranchFailed (runtime fail_fast).
• NodeEvent.branch_name: str | None populated on events from nodes inside a parallel-branches branch (graph-engine §6, proposal 0011). Independent of fan_out_index — both may be present simultaneously when a branch contains a fan-out (or a fan-out instance contains a parallel-branches node).
• openarmature.branch_name OTel span attribute parallel to the existing openarmature.node.fan_out_index. Both coexist on inner nodes of a fan-out-inside-a-branch composition.
• Retry-MW attempt-index ContextVar propagation (graph-engine §6 clarified in spec v0.16.1). Retry middleware sets _attempt_index_var via the existing ContextVar token pattern around each await next_(state) call; the engine reads current_attempt_index() when emitting events and when recording the checkpoint-save attempt_index, instead of writing its own local counter. Innermost-wins precedence falls out of Python's ContextVar token-stack semantics. Branch and instance-middleware retry now correctly surface their counter on inner-node events, matching the v0.16.1 clarification (and fixture 036's alpha_inner_attempt_indices_seen: [0, 1] invariant).

What's new

Surface	Change
BranchSpec, ParallelBranchesNode	New types exported from openarmature.graph. BranchSpec is a frozen dataclass; ParallelBranchesNode holds the branches dict, error policy, errors_field, and dispatcher-level middleware.
GraphBuilder.add_parallel_branches_node	Registration surface. Validates empty branches map → ParallelBranchesNoBranches; empty branch names → ValueError; inputs/outputs reference declared fields via the existing mapping_references_undeclared_field path.
ParallelBranchesNoBranches, ParallelBranchesBranchFailed	Two new canonical categories per spec §11.9. The second carries branch_name as a structured field, preserves the original cause via __cause__, and inherits transient classification from the wrapped exception.
NodeEvent.branch_name: str | None	New field on the event shape per graph-engine §6 v0.11.0. Populated only inside a parallel-branches branch.
openarmature.branch_name OTel attribute	Set on every inner-node span when event.branch_name is not None. NOT set on the parallel-branches node's own span.
_attempt_index_var propagation	RetryMiddleware.__call__ now sets/resets the ContextVar around each next_(); engine reads current_attempt_index() for event emission and checkpoint save attribution. Innermost-wins via Python's ContextVar token-stack.
descend_into_parallel_branch on _InvocationContext	Atomic-restart descent: drops checkpointer + resume state per §10.7.
Spec pin	v0.16.0 → v0.16.1.

Release gate

PR-5 is the final piece of the five-PR batch (0016 → 0015 → 0017 → 0014 → 0011). With this PR merged, the gate is cleared and the consolidated v0.5.x release can tag the whole batch against spec v0.16.1. The CHANGELOG [Unreleased] notes flip to reflect this.

Commits

Atop commit 945b4c8 (the engine surface that landed before the v0.16.1 spec clarification reopened the attempt-index question):

1. feat(graph): propagate retry attempt_index
2. test(conformance): drive 0011 parallel-branches
3. test(unit): parallel_branches + docs + CHANGELOG

Notable implementation details

• _MultiContribution sentinel pattern. A frozen dataclass holding tuple[Any, ...] of per-branch contributions to one parent field. _merge_partial in compiled.py recognizes the sentinel and folds each value through the parent's reducer in branch insertion order. Single-branch contributions stay as plain values — the sentinel only fires when two or more branches write the same field. Lazy import in _merge_partial breaks the textual cycle (parallel_branches.py has a TYPE_CHECKING back-ref to compiled.py).
• asyncio.create_task(coro, context=ctx.copy()) per branch. Per-task ContextVar isolation: each branch's _set_branch_name stays branch-local. Outer state (the wrapping retry's _attempt_index_var) flows in because it was set BEFORE the copy. _set_branch_name is set inside the spawned task body, not in the dispatcher loop, because each task's context is fixed at spawn time via copy_context().
• Q5 cancellation drain (per coord-thread 02 + impl summary 05). _fail_fast waits for FIRST_EXCEPTION, cancels pending tasks, then drains with gather(*, return_exceptions=True). A second task may race past the cancellation point with its own exception, absorbed into the gather return list. The raise is committed to the FIRST failure observed; residual non-CancelledError exceptions are discarded silently with a DEBUG log line for post-mortem analysis.
• Engine no longer writes _attempt_index_var directly — a subtle but important shift from pre-v0.16.1 behavior. Retry middleware is now the canonical writer; the engine is pure reader via current_attempt_index(). The line-957 comment in compiled.py explicitly cites the v0.16.1 clarification so this isn't accidentally undone in a future refactor.
• deferred_info[0] captures current_attempt_index() at the moment of the successful merge (not from a stale local counter) for checkpoint-save attribution. All three save sites in compiled.py source from deferred_info[0][0] when available, falling back to 0 for the short-circuit case (middleware short-circuited without invoking next).
• descend_into_parallel_branch atomic-restart contract. Mirrors fan-out's per-instance descent: drops checkpointer + resume state on the spawned branch's invocation context. Per spec §10.7, per-branch progress is not individually persisted in v1 — a checkpoint resume landing on a parallel-branches node re-dispatches all branches from scratch.

Conformance fixtures (032–038 + graph-engine/021)

All 8 fixtures pass:

Fixture	Asserts
032 basic	Three heterogeneous branches dispatch concurrently; projected outputs merge into disjoint parent fields
033 fail-fast	Second branch raises; dispatcher raises parallel_branches_branch_failed with branch_name=beta and cause_message; recoverable_state carries no branch contributions (including the first branch's successful work, per §11.4 buffer-and-apply); third branch's slow inner node never completes
034 collect	Middle branch raises; alpha + gamma contributions land; beta's outputs do not fire (beta_result stays at default); branch_errors records the beta failure
035 different-state-schemas	Three branches with deliberately distinct state schemas; each contributes via its own outputs to differently-typed parent fields
036 with-branch-middleware-retry	Branch alpha wraps its subgraph with retry MW; inner node fails once (transient), retry MW re-invokes, inner node succeeds. alpha_inner_attempt_indices_seen: [0, 1] — the canonical v0.16.1 demonstration of transitive-retry attempt_index propagation. Branches beta and gamma show [0] (no retry)
037 determinism	Three branches with deliberately randomized inner-node completion timing (sleep_ms 50/5/25). Two branches write the same merge-reducer parent field. Insertion order alpha→beta→gamma determines fan-in order regardless of completion order (which is beta→gamma→alpha at runtime). branch_started_event_order: [alpha, beta, gamma]
038 compose-with-fan-out	One branch contains a fan-out node; inner-node events inside that fan-out carry BOTH branch_name AND fan_out_index. The plain branch's inner events carry branch_name but no fan_out_index
graph-engine/021 observer-branch-name	Outermost-graph nodes' events carry no branch_name; inner nodes of branch alpha carry branch_name=alpha; inner nodes of branch beta carry branch_name=beta

Harness extensions

• ParallelBranchSpec + ParallelBranchesSpec typed-directive models on NodeSpec. sleep_ms node-companion modifier (used by 033's slow third branch and 037's randomized completion timing). recoverable_state expected key on the PipelineUtilitiesExpected discriminator.
• adapter._add_parallel_branches_node translates the parallel_branches: block to builder.add_parallel_branches_node, resolving branch subgraph references against the case's subgraphs: block. The registered node is swapped for a _TracingParallelBranchesNode so the conformance trace records the dispatcher as one engine step.
• adapter._wrap_with_sleep wraps a node body in an asyncio.sleep(ms/1000) for the sleep_ms companion modifier.
• test_pipeline_utilities driver: lifts the fixture-number gate from 23 to 38; loads top-level plural subgraphs: blocks (alongside the legacy singular subgraph: / subgraph_with_idx:); translates per-branch middleware lists; wires graph-attached observers per run; asserts parallel-branches observer_event_invariants (branch_started_event_order, alpha_inner_attempt_indices_seen, fan-out-inside-branch invariants, plain-branch invariants); routes fail_fast assertions to surface branch_name + cause_message + recoverable_state alongside category; subset-matches errors_field records (the spec mandates branch_name + category; the engine's record carries two extras).
• test_conformance driver: drops the 021 skip and adds an invariants checker for outermost_events_have_no_branch_name and inner_events_carry_correct_branch_name.
• Deferred set update: checkpoint-resume fixtures (024-031) move into test_pipeline_utilities's deferred set since their cases-shape with first_run_expected_error / resume: blocks is driven by test_checkpoint.py.

Unit tests

13 tests under tests/unit/test_parallel_branches.py:

• Compile-time validation (5 tests): empty branches → ParallelBranchesNoBranches; empty branch_name → ValueError; inputs/outputs/errors_field projection-field-existence checks raise MappingReferencesUndeclaredField on the correct side.
• Runtime happy path: three heterogeneous branches merge into disjoint parent fields.
• fail_fast: ParallelBranchesBranchFailed surfaces branch_name; __cause__ chain walks back to the original RuntimeError; recoverable_state holds the pre-dispatch snapshot with no buffered contributions applied (alpha's contribution must NOT land even though its branch may have completed before cancellation).
• collect: per-branch failures land in errors_field with branch_name + category; failed branches' outputs do not fire (beta_result stays at default).
• Determinism: when two branches write the same merge-reducer parent field, insertion order determines fan-in order regardless of inner-node completion timing.
• Single-branch field: writes flow through the parent's reducer as a plain value (the _MultiContribution sentinel only fires for multi-branch fields).
• Cancellation drain: a second branch racing past cancellation has its exception absorbed silently; the dispatcher commits to the FIRST failure observed.

Test plan

• uv run pytest (excluding tests/test_examples_smoke.py pre-existing failures from missing openai in CI venv) — 657 pass, 58 skipped, 0 failed.
• uv run pyright — clean.
• uv run ruff check + uv run ruff format — clean.
• All 8 parallel-branches conformance fixtures pass (7 pipeline-utilities + 1 graph-engine).
• 13 unit tests under tests/unit/test_parallel_branches.py pass.
• Existing fixtures unchanged: 645+ tests still passing across graph-engine, pipeline-utilities, llm-provider, observability, and state-migration capabilities.
• Manual: walk docs/concepts/parallel-branches.md via mkdocs serve to confirm rendering.

Pre-1.0 SemVer

Additive: new types, new errors, new builder method, new NodeEvent field, new OTel attribute. Existing free-form callers (no parallel-branches nodes registered) see no behavior change.

The retry-MW attempt_index propagation is a behavior CHANGE for users with retry middleware wrapping a subgraph (branch middleware, fan-out instance_middleware): events from inner nodes of that subgraph now carry the wrapping retry's attempt counter rather than starting at 0 each re-invocation. Per-node retry behavior is unchanged. Spec v0.16.1 frames this as a clarification of the ambiguous pre-v0.16.0 wording, so the v0.16.1 pin is the framing pre-1.0 MINOR. CHANGELOG flags the change under Added.

[Copilot]

Pull request overview

Implements spec proposal 0011 (parallel branches) — the final PR in a five-PR batch — and bumps the pinned spec to v0.16.1. Adds a new graph primitive that dispatches M heterogeneous compiled subgraphs concurrently, with deterministic insertion-order fan-in, two error policies (fail_fast / collect), branch_name propagation on events and OTel spans, and refactors retry-middleware's attempt_index to flow through a ContextVar so transitive (branch/instance) retry wrapping correctly surfaces the wrapping counter on inner-node events.

Changes:

• New ParallelBranchesNode / BranchSpec types and GraphBuilder.add_parallel_branches_node builder surface, with ParallelBranchesNoBranches (compile) and ParallelBranchesBranchFailed (runtime) error categories; conformance harness + 13 unit tests + concept docs added.
• NodeEvent.branch_name field and openarmature.branch_name OTel attribute, populated on inner-node events/spans inside a branch (independent of fan_out_index).
• Retry middleware now writes _attempt_index_var via ContextVar tokens; engine reads it via current_attempt_index() for both event emission and checkpoint-save attribution (replacing its own local counter at the _step_function_node / _step_fan_out_node / new _step_parallel_branches_node sites).

Reviewed changes

Copilot reviewed 25 out of 25 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
src/openarmature/graph/parallel_branches.py	New BranchSpec, ParallelBranchesNode, _MultiContribution; fail_fast/collect dispatch with cancellation drain.
src/openarmature/graph/builder.py	add_parallel_branches_node with projection-field + branch-name + empty-branches validation.
src/openarmature/graph/compiled.py	Adds _step_parallel_branches_node; refactors function/fan-out attempt-index to ContextVar reads; _merge_partial handles _MultiContribution; branch_name stamped on events.
src/openarmature/graph/errors.py	New ParallelBranchesNoBranches, ParallelBranchesBranchFailed categories.
src/openarmature/graph/events.py	NodeEvent.branch_name field + docstring updates.
src/openarmature/graph/observer.py	descend_into_parallel_branch atomic-restart descent.
src/openarmature/graph/middleware/retry.py	Sets/resets _attempt_index_var around next_(state).
src/openarmature/graph/__init__.py	Re-exports new types/errors.
src/openarmature/observability/correlation.py	New _branch_name_var ContextVar + current_branch_name() API.
src/openarmature/observability/otel/observer.py	Emits openarmature.branch_name span attribute.
src/openarmature/__init__.py, pyproject.toml	Spec pin bumped to 0.16.1.
tests/unit/test_parallel_branches.py	13 unit tests covering compile-time validation, fail_fast, collect, determinism, cancellation drain.
tests/conformance/adapter.py	parallel_branches: directive translation, _TracingParallelBranchesNode, sleep_ms wrapper.
tests/conformance/harness/directives.py, expectations.py	New ParallelBranchSpec / ParallelBranchesSpec models, sleep_ms + recoverable_state keys.
tests/conformance/test_pipeline_utilities.py	Drives 032-038, parallel-branches invariants, branch-middleware translation, errors_field subset matching.
tests/conformance/test_conformance.py	Drives graph-engine 021 observer-branch-name invariants.
tests/conformance/test_fixture_parsing.py, tests/test_smoke.py	Removes deferred-skip entries; spec-version assertion update.
docs/concepts/parallel-branches.md, observability.md, index.md, CHANGELOG.md	Docs + changelog for the new feature and the spec-pin bump.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.