aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/core/execution/step_forward.py

@@ -0,0 +1,411 @@
+from datetime import datetime, timedelta
+import functools
+import inspect
+from typing import Any
+
+from aethergraph.services.continuations.continuation import Continuation
+
+from ..graph.graph_refs import RESERVED_INJECTABLES  # {"context", "resume", "self"}
+from ..graph.task_node import NodeStatus, TaskNodeRuntime
+from ..runtime.execution_context import ExecutionContext
+from ..runtime.node_context import NodeContext
+from .retry_policy import RetryPolicy
+from .step_result import StepResult
+from .wait_types import WaitRequested
+
+
+async def maybe_await(func, *args, **kwargs):
+    if inspect.iscoroutinefunction(func):
+        return await func(*args, **kwargs)
+    return func(*args, **kwargs)
+
+
+def _normalize_result(res):
+    if res is None:
+        return {}
+    if isinstance(res, dict):
+        return res
+    if isinstance(res, tuple):
+        return {f"out{i}": v for i, v in enumerate(res)}
+    return {"result": res}
+
+
+def _waiting_status(kind: str) -> str:
+    return NodeStatus.from_kind(kind) if kind else NodeStatus.WAITING_EXTERNAL  # maps to WAITING_*
+
+
+def unwrap_callable(fn):
+    """Unwrap a callable from various wrapper types.
+        This includes:
+        - functions decorated with @tool or @waitable_tool (have __aether_impl__)
+        - functools.partial
+        - bound methods (unwrap to function)
+        Returns the innermost callable.
+
+        The function works as follows:
+        - If the callable has already been seen (to prevent infinite loops), return it as is.
+        - If the callable has an attribute __aether_impl__, unwrap it to that attribute.
+        - If the callable is a functools.partial, unwrap it to its func attribute.
+        - If the callable is a bound method, unwrap it to its __func__ attribute.
+        - If none of the above, return the callable as is.
+        This function is useful for extracting the core logic function from various
+        wrappers that may have been applied to it.
+    Args:
+        fn: The callable to unwrap.
+    """
+    seen = set()
+    while True:
+        if id(fn) in seen:
+            return fn
+        seen.add(id(fn))
+        if hasattr(fn, "__aether_impl__"):
+            fn = fn.__aether_impl__
+            continue
+        if isinstance(fn, functools.partial):
+            fn = fn.func
+            continue
+        if inspect.ismethod(fn):
+            fn = fn.__func__
+            continue
+        return fn
+
+
+def _flatten_inputs(resolved_inputs: dict[str, Any]) -> dict[str, Any]:
+    """Copy, then expand nested 'kwargs' dict into top-level keys."""
+    out = dict(resolved_inputs) if resolved_inputs else {}
+    nested = out.pop("kwargs", None)
+    if isinstance(nested, dict):
+        # only fill missing keys to let explicit top-level override nested
+        for k, v in nested.items():
+            out.setdefault(k, v)
+    return out
+
+
+def build_call_kwargs(
+    logic_fn,
+    resolved_inputs: dict[str, Any],
+    *,
+    node_ctx: NodeContext,
+    runtime_ctx: "ExecutionContext" = None,
+) -> dict[str, Any]:
+    """Build kwargs to call a logic function:
+    - flatten resolved_inputs (expand nested kwargs)
+    - inject framework args by name (node/context/logger/resume)
+    - validate required args
+    Returns a dict of kwargs to call the logic function.
+
+    NOTE: the input context is the full ExecutionContext, not a limited NodeContext. The 'context' param in the output kwargs
+    will be a NodeContext if the callee wants it. NodeContext is used when calling the logic function if it
+    accepts a 'context' parameter.
+
+    Raises TypeError if required args are missing.
+    """
+    import inspect
+
+    if runtime_ctx is None or node_ctx is None:
+        raise RuntimeError("build_call_kwargs: node_ctx and runtime_ctx are required")
+
+    sig = inspect.signature(logic_fn)
+    params = sig.parameters
+    has_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
+
+    flat = _flatten_inputs(resolved_inputs)
+    if "kwargs" in flat and isinstance(flat["kwargs"], dict):
+        flat = {**flat, **flat["kwargs"]}
+        flat.pop("kwargs", None)
+
+    # Framework injectables (authoritative)
+    inject_pool = {
+        "context": node_ctx,  # always NodeContext
+        "resume": getattr(runtime_ctx, "resume_payload", None),
+    }
+
+    merged = dict(flat)
+    for k in RESERVED_INJECTABLES:
+        if k == "self":
+            continue
+        if k in params or has_var_kw:
+            merged[k] = inject_pool.get(k)
+
+    if not has_var_kw:
+        merged = {k: v for k, v in merged.items() if k in params}
+    merged.pop("self", None)
+    merged.pop("kwargs", None)
+
+    required = [
+        name
+        for name, p in params.items()
+        if name != "self"
+        and p.default is inspect._empty
+        and p.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
+    ]
+    missing = [k for k in required if k not in merged]
+    if missing:
+        raise TypeError(
+            f"{getattr(logic_fn, '__name__', type(logic_fn).__name__)} missing required arguments: {missing}. "
+            f"Provided keys: {sorted(merged.keys())}"
+        )
+    return merged
+
+
+async def step_forward(
+    *, node: "TaskNodeRuntime", ctx: "ExecutionContext", retry_policy: "RetryPolicy"
+) -> StepResult:
+    """
+    Execute one node forward:
+      - resolve & inject kwargs (node/context/memory/logger/resume)
+      - await async logic
+      - apply should_run gate
+      - route subgraph to a dedicated handler (NotImplemented here)
+      - distinguish waits vs failures
+      - persist Continuation on wait (token, deadline/poll, channel)
+    Returns a StepResult; the runner is responsible for mutating node state.
+    """
+    lg = None
+    if getattr(ctx, "logger_factory", None) and hasattr(ctx.logger_factory, "for_node_ctx"):
+        lg = ctx.logger_factory.for_node_ctx(
+            run_id=ctx.run_id, node_id=node.node_id, graph_id=getattr(ctx, "graph_id", None)
+        )
+    attempts = getattr(node, "attempts", 0)
+
+    logic_fn = unwrap_callable(ctx.get_logic(node.logic))
+
+    # Resolve graph inputs
+    try:
+        resolved_inputs = await ctx.resolve_inputs(node)
+    except Exception as e:
+        if lg:
+            lg.exception("input resolution error")
+        return StepResult(status=NodeStatus.FAILED, error=e)
+
+    # should_run gate (unchanged) ...
+    should = True
+    if hasattr(ctx, "should_run") and callable(ctx.should_run):
+        try:
+            should = (
+                await ctx.should_run(node, resolved_inputs)
+                if inspect.iscoroutinefunction(ctx.should_run)
+                else ctx.should_run(node, resolved_inputs)
+            )
+        except Exception as e:
+            if lg:
+                lg.warning(f"should_run raised {e!r}; defaulting to run=True")
+    if not should:
+        return StepResult(
+            status=getattr(NodeStatus, "SKIPPED", "SKIPPED"), outputs={"skipped": True}
+        )
+
+    # create NodeContext once
+    node_ctx = ctx.create_node_context(node)
+
+    # Build kwargs with node_ctx as 'context' and the full ctx as 'runtime'
+    kwargs = build_call_kwargs(
+        logic_fn,
+        resolved_inputs=resolved_inputs,
+        node_ctx=node_ctx,  # <-- pass node_ctx explicitly for convenience
+        runtime_ctx=ctx,  # <-- pass runtime explicitly to resolve resume payload
+    )
+    try:
+        result = (
+            await logic_fn(**kwargs)
+            if inspect.iscoroutinefunction(logic_fn)
+            or (callable(logic_fn) and inspect.iscoroutinefunction(logic_fn.__call__))
+            else logic_fn(**kwargs)
+        )
+
+        outputs = _normalize_result(result)
+        if lg:
+            lg.info("done")
+        return StepResult(status=NodeStatus.DONE, outputs=outputs)
+
+    except WaitRequested as w:
+        # persist a Continuation and return StepResult with WAITING_*
+        if lg:
+            lg.info("wait requested: %s", getattr(w, "kind", None))
+        return await _enter_wait(
+            node=node, ctx=ctx, node_ctx=node_ctx, lg=lg, spec=w.to_dict(), attempts=attempts
+        )
+
+    except Exception as e:
+        if lg:
+            lg.exception("tool error")
+        if attempts < retry_policy.max_attempts and retry_policy.should_retry(e):
+            backoff = retry_policy.backoff(attempts)
+            if lg:
+                lg.warning(f"retry scheduled in {backoff}")
+            node.attempts = attempts + 1
+        # import traceback; traceback.print_exc()
+        return StepResult(status=NodeStatus.FAILED, error=e)
+
+
+# ---- wait path ---------------------------------------------------------------
+def _parse_deadline(deadline: Any, now_fn) -> datetime | None:
+    if not deadline:
+        return None
+    if isinstance(deadline, datetime):
+        return deadline
+    try:
+        return datetime.fromisoformat(deadline)
+    except Exception:
+        # allow "in N seconds" style if ever passed
+        try:
+            sec = int(deadline)
+            return now_fn() + timedelta(seconds=sec)
+        except Exception:
+            return None
+
+
+def normalize_wait_spec(spec: dict[str, Any], *, node_ctx: "NodeContext") -> dict[str, Any]:
+    """Normalize wait spec from WaitRequested to a canonical dict that used in channel/continuation:
+    In WaitSpec, we allow:
+        - kind: str e.g.  "approval" | "user_input" | "human" | "robot" | "external" | "time" | "event" | ...
+        - prompt: str | dict
+        - resume_schema: dict
+        - channel: str | None (it may be None)
+        - deadline: datetime | str (ISO) | int (seconds from now)
+        - poll: dict
+
+
+    In the normalized dict, we ensure:
+        - kind: str (default "external")
+        - prompt: str | dict | None
+        - resume_schema: dict | None
+        - channel: str (default from node_ctx or "console:stdin")
+        - deadline: datetime | None
+        - poll: dict | None
+
+    NOTE: in channel, we only allow kind to be "approval" or "user_input" for external interaction. Other kinds will
+    simply push a notification without expecting a user response.
+    """
+    from datetime import datetime, timezone
+
+    out = dict(spec or {})
+    out["kind"] = out.get("kind") or "external"
+    out["prompt"] = out.get("prompt")
+    out["resume_schema"] = out.get("resume_schema")
+
+    # Channel resolution via node_ctx
+    ch = out.get("channel")
+    if isinstance(ch, dict):
+        ch = None
+    if not ch:
+        ch = node_ctx.channel()._resolve_default_key() or "console:stdin"
+    out["channel"] = ch
+
+    # Deadline
+    now_fn = getattr(node_ctx, "_now", None)
+    if now_fn is None:
+
+        def now_fn():
+            return datetime.now(timezone.utc)
+
+    out["deadline"] = _parse_deadline(out.get("deadline"), now_fn)
+
+    # Poll
+    poll = out.get("poll")
+    if poll:
+        try:
+            poll["interval_sec"] = int(poll.get("interval_sec", 30))
+        except Exception:
+            poll["interval_sec"] = 30
+        out["poll"] = poll
+    return out
+
+
+async def _enter_wait(
+    *, node, ctx, node_ctx, lg, spec: dict[str, Any], attempts: int
+) -> StepResult:
+    spec = normalize_wait_spec(spec, node_ctx=node_ctx)
+
+    # 1) Reuse token if present
+    token = spec.get("token")
+    store = ctx.services.continuation_store
+
+    # Add wait spec in node state for reference -> This has not been used anywhere yet, We need save it with TaskGraph when state changes to WAITING_*
+    node.state.wait_spec = {
+        "kind": spec["kind"],  # "text" | "approval" | "files" | ...
+        "channel": spec.get("channel"),
+        "prompt": spec.get("prompt"),
+        "options": spec.get("options"),
+        "meta": spec.get("meta", {}),
+    }
+
+    cont = None
+    if token:
+        try:
+            cont = await store.get_by_token(token)
+        except Exception:
+            cont = None
+
+    if cont is None:
+        # fall back to minting (legacy path)
+        token = token or await store.mint_token(ctx.run_id, node.node_id, attempts)
+        cont = Continuation(
+            run_id=ctx.run_id,
+            node_id=node.node_id,
+            kind=spec["kind"],
+            token=token,
+            prompt=spec.get("prompt"),
+            resume_schema=spec.get("resume_schema"),
+            channel=spec["channel"],
+            deadline=spec.get("deadline"),
+            poll=spec.get("poll"),
+            next_wakeup_at=None,
+            created_at=ctx.now(),
+            attempts=attempts,
+        )
+    else:
+        # update mutable fields
+        cont.kind = spec.get("kind", cont.kind)
+        cont.prompt = spec.get("prompt", cont.prompt)
+        cont.resume_schema = spec.get("resume_schema", cont.resume_schema)
+        cont.channel = spec.get("channel", cont.channel)
+        cont.deadline = spec.get("deadline", cont.deadline)
+        cont.poll = spec.get("poll", cont.poll)
+        cont.attempts = attempts
+
+    # schedule next wakeup
+    if cont.poll and "interval_sec" in cont.poll:
+        from datetime import timedelta
+
+        cont.next_wakeup_at = ctx.now() + timedelta(seconds=int(cont.poll["interval_sec"]))
+    elif cont.deadline:
+        cont.next_wakeup_at = cont.deadline
+    else:
+        cont.next_wakeup_at = None
+
+    # persist (create or update)
+    await store.save(cont)
+
+    # 2) If inline payload was captured during setup, resume immediately
+    inline = spec.get("inline_payload")
+    if inline is not None:
+        try:
+            await ctx.resume_router.resume(cont.run_id, cont.node_id, cont.token, inline)
+            if lg:
+                lg.debug("inline resume dispatched for token=%s", cont.token)
+            # No need to notify again
+            return StepResult(
+                status=_waiting_status(cont.kind),
+                continuation=cont,
+                next_wakeup_at=cont.next_wakeup_at,
+            )
+        except Exception as e:
+            if lg:
+                lg.warning(f"inline resume failed: {e!r}; will proceed without it")
+
+    # 3) Notify only if the tool hasn't already done it
+    if not spec.get("notified", False):
+        try:
+            await ctx.channels.notify(cont)
+            if lg:
+                lg.debug("notified channel=%s", cont.channel)
+        except Exception as e:
+            if lg:
+                lg.error(f"notify failed: {e}")
+
+    return StepResult(
+        status=_waiting_status(cont.kind),
+        continuation=cont,
+        next_wakeup_at=cont.next_wakeup_at,
+    )

aethergraph/core/execution/step_result.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+from aethergraph.contracts.services.artifacts import Artifact
+from aethergraph.services.continuations.continuation import Continuation
+
+
+@dataclass
+class StepResult:
+    status: str  # NodeStatus
+    outputs: dict[str, Any] | None = None  # outputs if completed
+    artifacts: list[Artifact] = field(default_factory=list)
+    error: str | None = None  # error message if failed
+    continuation: Continuation | None = None  # continuation if waiting
+    next_wakeup_at: datetime | None = None  # ISO timestamp for next wakeup (for time-based waits)

aethergraph/core/execution/wait_types.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+_WAIT_KEY = "__wait__"
+
+
+@dataclass
+class WaitSpec:
+    kind: str = "external"  # "human" | "ask_text" | "external" | "time" | "event" | ... This is more generic than channel wait kinds
+    prompt: dict[str, Any] | str | None = None  # for human/robot
+    resume_schema: dict[str, Any] | None = None  # for human/robot validation
+    channel: str | None = None  # for external/event
+    deadline: datetime | str | None = None  # ISO timestamp or datetime
+    poll: dict[str, Any] | None = (
+        None  # {"interval_sec": 30, "endpoint": "...", "extract": "$.path"}
+    )
+
+    # resume handles
+    token: str | None = None  # internal opaque continuation id (do NOT expose to untrusted clients)
+    resume_key: str | None = None  # short alias safe to surface in UI/buttons
+    notified: bool = False  # internal flag: whether continuation notification has been sent out
+    inline_payload: dict[str, Any] | None = (
+        None  # internal: optional inline payload returned from notification step
+    )
+
+    # Optional grab-bag for extensions; avoids new fields churn later
+    meta: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        # Only include non-None fields to preserve backward compatibility with consumers
+        d = {
+            "kind": self.kind,
+            "prompt": self.prompt,
+            "resume_schema": self.resume_schema,
+            "channel": self.channel,
+            "deadline": self.deadline,
+            "poll": self.poll,
+            "token": self.token,
+            "resume_key": self.resume_key,
+            "notified": self.notified,
+            "inline_payload": self.inline_payload,
+            "meta": self.meta or None,
+        }
+        return {k: v for k, v in d.items() if v is not None}
+
+    def sanitized_for_transport(self) -> dict[str, Any]:
+        """
+        Strip sensitive fields for UI/adapters/webhooks.
+        Prefer exposing `resume_key` (short alias) over raw `token`.
+        """
+        d = self.to_dict()
+        d.pop("token", None)
+        return d
+
+
+def wait_sentinel(spec: WaitSpec | dict[str, Any]) -> dict[str, Any]:
+    """Return the canonical sentinel the executor understands as 'please wait'."""
+    return {_WAIT_KEY: spec if isinstance(spec, dict) else spec.__dict__}
+
+
+class WaitRequested(RuntimeError):
+    """Exception to raise from a tool to indicate it wants to wait."""
+
+    def __init__(self, spec: dict[str, Any]):
+        self.spec = spec
+        super().__init__(f"Wait requested: {spec}")
+
+    def to_dict(self):
+        return self.spec if isinstance(self.spec, dict) else self.spec.to_dict()

aethergraph/core/graph/graph_builder.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from contextlib import contextmanager
+from contextvars import ContextVar
+import itertools
+from typing import Any
+import uuid
+
+from .graph_refs import GRAPH_INPUTS_NODE_ID, RESERVED_INJECTABLES
+from .graph_spec import TaskGraphSpec
+from .node_spec import TaskNodeSpec
+from .task_graph import TaskGraph
+
+_GRAPH_CTX: ContextVar[GraphBuilder | None] = ContextVar("_GRAPH_CTX", default=None)  # Async-safe
+
+
+def current_builder() -> GraphBuilder | None:
+    return _GRAPH_CTX.get()
+
+
+class GraphBuilder:
+    _auto_counter = itertools.count(1)
+
+    def __init__(self, *, name: str = "default_graph"):
+        self.spec = TaskGraphSpec(graph_id=name, nodes={}, meta={})
+        self.graph = TaskGraph(spec=self.spec)
+        self.graph.ensure_inputs_node()
+
+        self._auto_counter_by_logic = {}  # logic_name -> counter
+
+        # index for quick lookup
+        self._alias_index: dict[str, str] = {}  # alias -> node_id
+        self._logic_index: dict[str, list[str]] = {}  # logic -> [node_id, ...]
+        self._label_index: dict[str, list[str]] = {}  # label -> {node_id, ...}
+
+    def add_node(self, node_spec: TaskNodeSpec) -> str:
+        if node_spec.node_id in self.spec.nodes:
+            raise ValueError(
+                f"Node ID '{node_spec.node_id}' already exists in graph '{self.spec.graph_id}'"
+            )
+        self.spec.nodes[node_spec.node_id] = node_spec
+        return self
+
+    def add_tool_node(
+        self,
+        *,
+        node_id: str,
+        logic: str,
+        inputs: dict,
+        expected_input_keys: Iterable[str] | None = None,
+        expected_output_keys: Iterable[str] | None = None,
+        after: Iterable[str] | None = None,
+        inject: list[str] | None = None,
+        tool_name: str | None = None,
+        tool_version: str | None = None,
+    ) -> GraphBuilder:
+        """Add a tool node to the graph."""
+
+        if node_id in self.spec.nodes:
+            raise ValueError(f"Node with id {node_id} already exists in the graph.")
+
+        # Initialize injection and pure input mappings. Injection is for reserved keywords that should be passed from the context.
+        deps = set(after or [])
+        inject = inject or []
+        pure_inputs = {}
+        has_arg = False
+
+        for k, v in list(inputs.items()):
+            if k in RESERVED_INJECTABLES:
+                inject.append(k)
+            else:
+                pure_inputs[k] = v
+
+        # infer dependencies from input Refs
+        def _walk_refs(x):
+            # Recursively walk input bindings to find Ref dependencies
+            nonlocal has_arg
+            if isinstance(x, dict):
+                if x.get("_type") == "ref" and "from" in x:
+                    yield x["from"]
+                elif x.get("_type") == "arg":
+                    has_arg = True
+                else:
+                    for v in x.values():
+                        yield from _walk_refs(v)
+            elif isinstance(x, list | tuple):
+                for v in x:
+                    yield from _walk_refs(v)
+
+        deps = set(_walk_refs(pure_inputs))
+        if has_arg:
+            deps.add(GRAPH_INPUTS_NODE_ID)  # ensure inputs node is a dependency
+        if after:
+            for a in after:
+                deps.add(a.node_id if hasattr(a, "node_id") else a)
+
+        node = TaskNodeSpec(
+            node_id=node_id,
+            type="tool",
+            logic=logic,
+            inputs=inputs,
+            dependencies=list(deps),
+            expected_input_keys=expected_input_keys,
+            expected_output_keys=expected_output_keys,
+            metadata={},
+            tool_name=tool_name or logic or "unknown_tool",
+            tool_version=tool_version,  # could be set to a version string if available
+        )
+        return self.add_node(node)
+
+    def ensure_inputs_node(self):
+        """Ensure the special inputs node exists in the graph."""
+        if GRAPH_INPUTS_NODE_ID not in self.spec.nodes:
+            self.spec.nodes[GRAPH_INPUTS_NODE_ID] = TaskNodeSpec(
+                node_id=GRAPH_INPUTS_NODE_ID,
+                type="inputs",
+                logic=None,
+                inputs={},
+                dependencies=[],
+                expected_input_keys=[],
+                expected_output_keys=[],
+                metadata={"synthetic": True},
+            )
+        return self
+
+    def freeze(self) -> TaskGraphSpec:
+        """Frozen dataclass / validate topo order"""
+        return self.spec
+
+    def expose(self, name: str, value: Any):
+        self.graph.expose(name, value)
+
+    # ---- ids and utils ----
+    def next_id(self, logic_name: str | None = None) -> str:
+        """Generate a unique node ID."""
+        base = (logic_name or "node").rstrip("_")
+        return f"{base}_{next(self._auto_counter)}_{uuid.uuid4().hex[:6]}"
+
+    def _next_readable_id(self, logic_name: str | None = None) -> str:
+        """Generate a more human-readable node ID, but may not be unique."""
+        n = self._auto_counter_by_logic.get(logic_name, 0) + 1
+        self._auto_counter_by_logic[logic_name] = n
+        return f"{logic_name}_{n}"  # deterministic and readable
+
+    def to_graph(self) -> TaskGraph:
+        self.graph.spec.metadata["graph_io"] = self.graph.io_signature()
+        return self.graph
+
+    def register_alias(self, alias: str, node_id: str):
+        if alias in self._alias_index and self._alias_index[alias] != node_id:
+            raise ValueError(
+                f"Alias '{alias}' already registered for node '{self._alias_index[alias]}', cannot re-register for '{node_id}'"
+            )
+        self._alias_index[alias] = node_id
+
+    def register_logic_name(self, logic_name: str, node_id: str):
+        self._logic_index.setdefault(logic_name, []).append(node_id)
+
+    def register_labels(self, labels: Iterable[str], node_id: str):
+        for label in labels or []:
+            self._label_index.setdefault(label, set()).add(node_id)
+
+    # ergonomic accessors
+    def find_by_alias(self, alias: str) -> str | None:
+        return self._alias_index.get(alias)
+
+    def find_by_logic(self, logic_prefix: str) -> list[str]:
+        exact = self._logic_index.get(logic_prefix, [])
+        if exact:
+            return list(exact)
+        # fuzzy match: logic_name contained in key
+        out = []
+        for k, v in self._logic_index.items():
+            if k.startswith(logic_prefix):
+                out.extend(v)
+        return out
+
+    def find_by_label(self, label: str) -> list[str]:
+        return sorted(self._label_index.get(label, set()))
+
+
+@contextmanager
+def graph(*, name: str = "default_graph"):
+    """Context manager that yields a GraphBuilder to build a TaskGraph."""
+    builder = GraphBuilder(name=name)
+    token = _GRAPH_CTX.set(builder)
+    try:
+        yield builder.graph
+    finally:
+        builder.graph.__post_init__()  # reify runtime nodes
+        _GRAPH_CTX.reset(token)