lineage helper api

src/wf_core/runtime/lineage.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterator, Mapping, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any
 
-from wf_core.run_state import ExecutionFrame, RunState, StateWrite
+from wf_core.run_state import ExecutionFrame, LineageState, RunState, StateWrite
 from wf_core.runtime.foreach_state import ForeachBarrierState, item_frame_owner
+from wf_core.runtime.ops.state import StatePatch
 from wf_core.runtime.ops.state import safe_set_nested_value
 
 
@@ -58,3 +59,84 @@ def lineage_writes_for_frame(
     if pending is None:
         return ()
     return pending.patch.writes
+
+
+def add_lineage(
+    run: RunState,
+    *,
+    scope_id: str,
+    lineage_id: str,
+    parent_id: str | None,
+) -> None:
+    """Create one lineage record inside an existing runtime scope."""
+    if scope_id not in run.scopes:
+        raise ValueError(f"unknown scope {scope_id!r}")
+    if lineage_id in run.lineages:
+        raise ValueError(f"duplicate lineage {lineage_id!r}")
+    if parent_id is not None and parent_id not in run.lineages:
+        raise ValueError(f"unknown parent lineage {parent_id!r}")
+    run.lineages[lineage_id] = LineageState(
+        id=lineage_id,
+        scope_id=scope_id,
+        parent_id=parent_id,
+    )
+
+
+def append_lineage_writes(
+    run: RunState,
+    *,
+    scope_id: str,
+    lineage_id: str,
+    writes: Sequence[StateWrite],
+) -> None:
+    """Append ordered writes to an existing lineage without committing state."""
+    lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
+    lineage.writes.extend(writes)
+
+
+def lineage_patch(
+    run: RunState,
+    *,
+    scope_id: str,
+    lineage_id: str,
+) -> StatePatch:
+    """Return a replayable patch for one lineage's pending writes."""
+    lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
+    return StatePatch(writes=list(lineage.writes))
+
+
+def lineage_state_view(
+    run: RunState,
+    *,
+    scope_id: str,
+    lineage_id: str,
+) -> dict[str, Any]:
+    """Materialize scope committed state plus ancestor/current lineage writes."""
+    scope = run.scopes.get(scope_id)
+    if scope is None:
+        raise ValueError(f"unknown scope {scope_id!r}")
+    writes: list[StateWrite] = []
+    for lineage in _lineage_chain(run, scope_id=scope_id, lineage_id=lineage_id):
+        writes.extend(lineage.writes)
+    return LineageStateView(scope.committed_state, writes).to_state_dict()
+
+
+def _lineage(run: RunState, *, scope_id: str, lineage_id: str) -> LineageState:
+    lineage = run.lineages.get(lineage_id)
+    if lineage is None:
+        raise ValueError(f"unknown lineage {lineage_id!r}")
+    if lineage.scope_id != scope_id:
+        raise ValueError(
+            f"lineage {lineage_id!r} belongs to scope {lineage.scope_id!r}, "
+            f"not {scope_id!r}"
+        )
+    return lineage
+
+
+def _lineage_chain(
+    run: RunState, *, scope_id: str, lineage_id: str
+) -> Iterator[LineageState]:
+    lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
+    if lineage.parent_id is not None:
+        yield from _lineage_chain(run, scope_id=scope_id, lineage_id=lineage.parent_id)
+    yield lineage