coco-framework 0.1.0__py3-none-any.whl

coco_graph/types.py

@@ -0,0 +1,1079 @@
+from __future__ import annotations
+
+import sys
+from collections import deque
+from collections.abc import Callable, Hashable, Sequence
+from dataclasses import asdict, dataclass
+from datetime import timedelta
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Generic,
+    Literal,
+    NamedTuple,
+    final,
+    overload,
+)
+from warnings import warn
+
+from coco_core.messages import AnyMessage
+from coco_core.runnables import Runnable, RunnableConfig
+from coco_graph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
+from pydantic import TypeAdapter
+from typing_extensions import (
+    NotRequired,
+    TypeAliasType,
+    TypedDict,
+    TypeVar,
+    Unpack,
+    deprecated,
+)
+from xxhash import xxh3_128_hexdigest
+
+from coco_graph._internal._cache import default_cache_key
+from coco_graph._internal._constants import INTERRUPT as _INTERRUPT_KEY
+from coco_graph._internal._fields import get_cached_annotated_keys, get_update_as_tuples
+from coco_graph._internal._retry import default_retry_on
+from coco_graph._internal._typing import MISSING, DeprecatedKwargs
+from coco_graph.warnings import CocoGraphDeprecatedSinceV10, CocoGraphDeprecatedSinceV11
+
+# Local TypeVars for generic stream TypedDicts.
+# We use separate TypeVars here (rather than importing from coco_graph.typing)
+# because the typing module TypeVars have defaults that cause type checker issues
+# when used in standalone type aliases.
+StateT = TypeVar("StateT")
+OutputT = TypeVar("OutputT")
+ResponseT = TypeVar("ResponseT", default=Any)
+
+if TYPE_CHECKING:
+    from coco_graph.pregel.protocol import PregelProtocol
+
+
+try:
+    from coco_core.messages.tool import ToolOutputMixin
+except ImportError:
+
+    class ToolOutputMixin:  # type: ignore[no-redef]
+        pass
+
+
+__all__ = (
+    "All",
+    "Checkpointer",
+    "StreamMode",
+    "StreamWriter",
+    "StreamPart",
+    "ValuesStreamPart",
+    "UpdatesStreamPart",
+    "MessagesStreamPart",
+    "CustomStreamPart",
+    "CheckpointStreamPart",
+    "TasksStreamPart",
+    "DebugStreamPart",
+    "TaskPayload",
+    "TaskResultPayload",
+    "CheckpointTask",
+    "CheckpointPayload",
+    "DebugPayload",
+    "RetryPolicy",
+    "TimeoutPolicy",
+    "CachePolicy",
+    "TracePolicy",
+    "omit_payload",
+    "Interrupt",
+    "StateUpdate",
+    "PregelTask",
+    "PregelExecutableTask",
+    "StateSnapshot",
+    "Send",
+    "Command",
+    "Durability",
+    "interrupt",
+    "Overwrite",
+    "GraphOutput",
+    "ensure_valid_checkpointer",
+)
+
+Durability = Literal["sync", "async", "exit"]
+"""Durability mode for the graph execution.
+
+- `'sync'`: Changes are persisted synchronously before the next step starts.
+- `'async'`: Changes are persisted asynchronously while the next step executes.
+- `'exit'`: Changes are persisted only when the graph exits.
+"""
+
+All = Literal["*"]
+"""Special value to indicate that graph should interrupt on all nodes."""
+
+Checkpointer = None | bool | BaseCheckpointSaver
+"""Type of the checkpointer to use for a subgraph.
+
+- `True` enables persistent checkpointing for this subgraph.
+- `False` disables checkpointing, even if the parent graph has a checkpointer.
+- `None` inherits checkpointer from the parent graph.
+"""
+
+
+def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer:
+    if checkpointer not in (None, True, False) and not isinstance(
+        checkpointer, BaseCheckpointSaver
+    ):
+        raise TypeError(
+            "Invalid checkpointer provided. Expected an instance of "
+            "`BaseCheckpointSaver`, `True`, `False`, or `None`. "
+            f"Received {type(checkpointer).__name__!s}. "
+            "Pass a proper saver (e.g., InMemorySaver, AsyncPostgresSaver)."
+        )
+    return checkpointer
+
+
+StreamMode = Literal[
+    "values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"
+]
+"""How the stream method should emit outputs.
+
+- `"values"`: Emit all values in the state after each step, including interrupts.
+    When used with functional API, values are emitted once at the end of the workflow.
+- `"updates"`: Emit only the node or task names and updates returned by the nodes or tasks after each step.
+    If multiple updates are made in the same step (e.g. multiple nodes are run) then those updates are emitted separately.
+- `"custom"`: Emit custom data using from inside nodes or tasks using `StreamWriter`.
+- `"messages"`: Emit LLM messages token-by-token together with metadata for any LLM invocations inside nodes or tasks.
+- `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by `get_state()`.
+- `"tasks"`: Emit events when tasks start and finish, including their results and errors.
+- `"debug"`: Emit `"checkpoints"` and `"tasks"` events for debugging purposes.
+"""
+
+StreamWriter = Callable[[Any], None]
+"""`Callable` that accepts a single argument and writes it to the output stream.
+Always injected into nodes if requested as a keyword argument, but it's a no-op
+when not using `stream_mode="custom"`."""
+
+
+class TaskPayload(TypedDict):
+    """Payload for a task start event."""
+
+    id: str
+    """Unique identifier for this task."""
+    name: str
+    """Name of the node being executed."""
+    input: Any
+    """Input data passed to the task."""
+    triggers: list[str]
+    """List of triggers that caused this task to be executed (e.g. channel writes)."""
+    metadata: NotRequired[dict[str, Any]]
+    """Framework-resolved metadata associated with the task.
+
+    Generic dict carrier following the messages-stream pattern. Populated by
+    `map_debug_tasks` from `task.config["metadata"]` when non-empty, so the
+    same keys `stream_mode="messages"` consumers see (e.g. `lc_agent_name`,
+    `coco_graph_node`, `coco_graph_step`) are available to stream transformers.
+
+    Consumers should ignore unrecognized keys.
+    """
+
+
+class TaskResultPayload(TypedDict):
+    """Payload for a task result event."""
+
+    id: str
+    """Unique identifier for this task."""
+    name: str
+    """Name of the node that was executed."""
+    error: str | None
+    """Error message if the task failed, otherwise `None`."""
+    interrupts: list[dict]
+    """List of interrupts that occurred during task execution."""
+    result: dict[str, Any]
+    """Mapping of channel names to the values written by this task."""
+
+
+class CheckpointTask(TypedDict):
+    """A task entry within a `CheckpointPayload`.
+
+    The keys present depend on the task's state:
+
+    - **Error:** `id`, `name`, `error`, `state`
+    - **Has result:** `id`, `name`, `result`, `interrupts`, `state`
+    - **Pending:** `id`, `name`, `interrupts`, `state`
+    """
+
+    id: str
+    """Unique identifier for this task."""
+    name: str
+    """Name of the node being executed."""
+    error: NotRequired[str]
+    """Error message, present only if the task failed."""
+    result: NotRequired[Any]
+    """Result of the task, present only if the task completed successfully."""
+    interrupts: NotRequired[list[dict]]
+    """List of interrupts, present when the task has been interrupted or completed."""
+    state: StateSnapshot | RunnableConfig | None
+    """Snapshot of the subgraph state, or a `RunnableConfig` pointing to it. `None` if not a subgraph."""
+
+
+class CheckpointPayload(TypedDict, Generic[StateT]):
+    """Payload for a checkpoint event."""
+
+    config: RunnableConfig | None
+    """Configuration for this checkpoint, including the `thread_id` and `checkpoint_id`."""
+    metadata: CheckpointMetadata
+    """Metadata associated with this checkpoint (e.g. step number, source, writes)."""
+    values: StateT
+    """Current state values at the time of this checkpoint."""
+    next: list[str]
+    """Names of the nodes scheduled to execute next."""
+    parent_config: RunnableConfig | None
+    """Configuration of the parent checkpoint, or `None` if this is the first checkpoint."""
+    tasks: list[CheckpointTask]
+    """List of tasks associated with this checkpoint."""
+
+
+class _DebugCheckpointPayload(TypedDict, Generic[StateT]):
+    step: int
+    """The step number in the graph execution."""
+    timestamp: str
+    """ISO 8601 timestamp of when this event occurred."""
+    type: Literal["checkpoint"]
+    """Event type discriminator, always `"checkpoint"`."""
+    payload: CheckpointPayload[StateT]
+    """The checkpoint payload."""
+
+
+class _DebugTaskPayload(TypedDict):
+    step: int
+    """The step number in the graph execution."""
+    timestamp: str
+    """ISO 8601 timestamp of when this event occurred."""
+    type: Literal["task"]
+    """Event type discriminator, always `"task"`."""
+    payload: TaskPayload
+    """The task start payload."""
+
+
+class _DebugTaskResultPayload(TypedDict):
+    step: int
+    """The step number in the graph execution."""
+    timestamp: str
+    """ISO 8601 timestamp of when this event occurred."""
+    type: Literal["task_result"]
+    """Event type discriminator, always `"task_result"`."""
+    payload: TaskResultPayload
+    """The task result payload."""
+
+
+DebugPayload = TypeAliasType(
+    "DebugPayload",
+    _DebugCheckpointPayload[StateT] | _DebugTaskPayload | _DebugTaskResultPayload,
+    type_params=(StateT,),
+)
+"""Wrapper payload for debug events. Discriminate on `type`."""
+
+
+class ValuesStreamPart(TypedDict, Generic[OutputT]):
+    """Stream part emitted for `stream_mode="values"`.
+
+    `data` contains the full state after each step, as returned by `read_channels()`.
+    """
+
+    type: Literal["values"]
+    ns: tuple[str, ...]
+    data: OutputT
+    interrupts: tuple[Interrupt, ...]
+
+
+class UpdatesStreamPart(TypedDict):
+    """Stream part emitted for `stream_mode="updates"`.
+
+    `data` maps node names to their outputs. May also contain
+    `__interrupt__` (tuple of `Interrupt` dicts) and `__metadata__` keys.
+    """
+
+    type: Literal["updates"]
+    ns: tuple[str, ...]
+    data: dict[str, Any]
+
+
+class MessagesStreamPart(TypedDict):
+    """Stream part emitted for `stream_mode="messages"`.
+
+    `data` is a 2-tuple of `(message, metadata)` where `message` is a
+    `BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is a dict containing
+    keys like `coco_graph_step`, `coco_graph_node`, `coco_graph_triggers`, etc.
+    """
+
+    type: Literal["messages"]
+    ns: tuple[str, ...]
+    data: tuple[AnyMessage, dict[str, Any]]
+
+
+class CustomStreamPart(TypedDict):
+    """Stream part emitted for `stream_mode="custom"`.
+
+    `data` is whatever value was passed to `StreamWriter` inside a node.
+    """
+
+    type: Literal["custom"]
+    ns: tuple[str, ...]
+    data: Any
+
+
+class CheckpointStreamPart(TypedDict, Generic[StateT]):
+    """Stream part emitted for `stream_mode="checkpoints"`."""
+
+    type: Literal["checkpoints"]
+    ns: tuple[str, ...]
+    data: CheckpointPayload[StateT]
+
+
+class TasksStreamPart(TypedDict):
+    """Stream part emitted for `stream_mode="tasks"`.
+
+    For task start events, `data` is a `TaskPayload` with `id`, `name`,
+    `input`, and `triggers` keys.
+
+    For task result events, `data` is a `TaskResultPayload` with `id`,
+    `name`, `error`, `interrupts`, and `result` keys.
+    """
+
+    type: Literal["tasks"]
+    ns: tuple[str, ...]
+    data: TaskPayload | TaskResultPayload
+
+
+class DebugStreamPart(TypedDict, Generic[StateT]):
+    """Stream part emitted for `stream_mode="debug"`."""
+
+    type: Literal["debug"]
+    ns: tuple[str, ...]
+    data: DebugPayload[StateT]
+
+
+StreamPart = TypeAliasType(
+    "StreamPart",
+    ValuesStreamPart[OutputT]
+    | UpdatesStreamPart
+    | MessagesStreamPart
+    | CustomStreamPart
+    | CheckpointStreamPart[StateT]
+    | TasksStreamPart
+    | DebugStreamPart[StateT],
+    type_params=(StateT, OutputT),
+)
+"""A discriminated union of all v2 stream part types.
+
+Use `part["type"]` to narrow the type:
+
+```python
+async for part in graph.astream(input, version="v2"):
+    if part["type"] == "values":
+        part["data"]  # OutputT — full state (pydantic/dataclass/dict)
+    elif part["type"] == "messages":
+        part["data"]  # tuple[BaseMessage, dict] — (message, metadata)
+    elif part["type"] == "custom":
+        part["data"]  # Any — user-defined
+```
+"""
+
+
+@dataclass(frozen=True)
+class GraphOutput(Generic[OutputT]):
+    """Typed container returned by `invoke()` / `ainvoke()` with `version="v2"`.
+
+    Attributes:
+        value: The final output of the graph (dict, Pydantic model, dataclass, etc.).
+        interrupts: Any interrupts that occurred during execution.
+    """
+
+    value: OutputT
+    interrupts: tuple[Interrupt, ...] = ()
+
+    def __getitem__(self, key: str) -> Any:
+        """Backward compat: `result['__interrupt__']` and dict-key access."""
+        warn(
+            "Accessing GraphOutput via `result[key]` is deprecated. "
+            "Use `result.value` to access the output value directly, "
+            "or `result.interrupts` for interrupts.",
+            CocoGraphDeprecatedSinceV11,
+            stacklevel=2,
+        )
+        if key == _INTERRUPT_KEY:
+            return self.interrupts
+        if isinstance(self.value, dict):
+            return self.value[key]
+        try:
+            return getattr(self.value, key)
+        except AttributeError:
+            raise KeyError(key)
+
+    def __contains__(self, key: object) -> bool:
+        warn(
+            "Accessing GraphOutput via `key in result` is deprecated. "
+            "Use `result.value` to access the output value directly, "
+            "or `result.interrupts` for interrupts.",
+            CocoGraphDeprecatedSinceV11,
+            stacklevel=2,
+        )
+        if key == _INTERRUPT_KEY:
+            return bool(self.interrupts)
+        if isinstance(self.value, dict):
+            return key in self.value
+        return isinstance(key, str) and hasattr(self.value, key)
+
+
+_DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True}
+
+
+class RetryPolicy(NamedTuple):
+    """Configuration for retrying nodes.
+
+    !!! version-added "Added in version 0.2.24"
+    """
+
+    initial_interval: float = 0.5
+    """Amount of time that must elapse before the first retry occurs. In seconds."""
+    backoff_factor: float = 2.0
+    """Multiplier by which the interval increases after each retry."""
+    max_interval: float = 128.0
+    """Maximum amount of time that may elapse between retries. In seconds."""
+    max_attempts: int = 3
+    """Maximum number of attempts to make before giving up, including the first."""
+    jitter: bool = True
+    """Whether to add random jitter to the interval between retries."""
+    retry_on: (
+        type[Exception] | Sequence[type[Exception]] | Callable[[Exception], bool]
+    ) = default_retry_on
+    """List of exception classes that should trigger a retry, or a callable that returns `True` for exceptions that should trigger a retry."""
+
+
+def _coerce_timeout_seconds(
+    value: float | timedelta | None, *, field: str
+) -> float | None:
+    if value is None:
+        return None
+    seconds = value.total_seconds() if isinstance(value, timedelta) else float(value)
+    if seconds <= 0:
+        raise ValueError(f"{field} must be greater than 0")
+    return seconds
+
+
+@dataclass(**_DC_KWARGS)
+class TimeoutPolicy:
+    """Configuration for timing out node attempts.
+
+    !!! note "Cooperative cancellation"
+
+        Timeouts rely on asyncio cancellation. If your node uses synchronous
+        time.sleep() or other CPU-bound work that blocks the GIL, the timeout will not
+        be fired until after the event loop has been released.
+
+    !!! note "Inline callback dispatch"
+
+        Under `refresh_on="auto"`, an internal handler refreshes the timeout on any
+        callback event that occurs in the execution of the node or its nested descendants.
+    """
+
+    run_timeout: float | timedelta | None = None
+    """Hard wall-clock cap (in seconds) for a single node attempt.
+
+    This timeout is never refreshed by progress signals or `runtime.heartbeat()`.
+    """
+
+    idle_timeout: float | timedelta | None = None
+    """Maximum time (in seconds) a single node attempt may go without observable progress."""
+
+    refresh_on: Literal["auto", "heartbeat"] = "auto"
+    """Which signals refresh `idle_timeout`.
+
+    `"auto"` refreshes on standard graph progress signals and explicit heartbeats.
+    `"heartbeat"` refreshes only on explicit `runtime.heartbeat()` calls.
+    """
+
+    @classmethod
+    def coerce(
+        cls, value: float | timedelta | TimeoutPolicy | None
+    ) -> TimeoutPolicy | None:
+        """Normalize a timeout value to positive-second policy fields."""
+        if value is None:
+            return None
+        if isinstance(value, TimeoutPolicy):
+            # Fast path: a policy already produced by coerce() has float
+            # timeouts and a validated refresh_on, so we can return it as-is.
+            # `frozen=True` makes this safe to share.
+            rt, it = value.run_timeout, value.idle_timeout
+            if (
+                value.refresh_on in ("auto", "heartbeat")
+                and (rt is None or (type(rt) is float and rt > 0))
+                and (it is None or (type(it) is float and it > 0))
+                and (rt is not None or it is not None)
+            ):
+                return value
+        else:
+            value = cls(run_timeout=value)
+        if value.refresh_on not in ("auto", "heartbeat"):
+            raise ValueError("refresh_on must be 'auto' or 'heartbeat'")
+        run_timeout = _coerce_timeout_seconds(value.run_timeout, field="run_timeout")
+        idle_timeout = _coerce_timeout_seconds(value.idle_timeout, field="idle_timeout")
+        if run_timeout is None and idle_timeout is None:
+            return None
+        return cls(
+            run_timeout=run_timeout,
+            idle_timeout=idle_timeout,
+            refresh_on=value.refresh_on,
+        )
+
+
+KeyFuncT = TypeVar("KeyFuncT", bound=Callable[..., str | bytes])
+
+
+@dataclass(**_DC_KWARGS)
+class CachePolicy(Generic[KeyFuncT]):
+    """Configuration for caching nodes."""
+
+    key_func: KeyFuncT = default_cache_key  # type: ignore[assignment]
+    """Function to generate a cache key from the node's input.
+    Defaults to hashing the input with pickle."""
+
+    ttl: int | None = None
+    """Time to live for the cache entry in seconds. If `None`, the entry never expires."""
+
+
+@dataclass(**_DC_KWARGS)
+class TracePolicy:
+    """Configuration for how a node's run is traced.
+
+    Scope: this only transforms what the node's *own* run records. Child runs created
+    by a traced `bound` runnable and the root graph run are not affected. Plain
+    function nodes are traced with `trace=False`, so they have no such child runs.
+
+    Not intended to redact secrets. To redact inputs/outputs across all runs
+    (children included), use the LangSmith client's
+    `hide_inputs`/`hide_outputs`/`anonymizer` instead.
+
+    Each processor receives the node's raw input/output value (not a normalized
+    kwargs dict) and returns the value to record.
+    """
+
+    process_inputs: Callable[[Any], Any] | None = None
+    """Optional callable to transform the node's input before it is recorded on the
+    node's trace run. Can be used to omit or summarize large payloads
+    (e.g. message history). Not intended to affect the value passed to the node; avoid
+    mutating arguments in place."""
+
+    process_outputs: Callable[[Any], Any] | None = None
+    """Optional callable to transform the node's output before it is recorded on the
+    node's trace run. Can be used to omit or summarize large payloads
+    (e.g. message history). Not intended to affect the value returned by the node; avoid
+    mutating arguments in place."""
+
+
+def omit_payload(_value: Any) -> dict[str, Any]:
+    """`TracePolicy` helper that records an empty payload, dropping the value entirely.
+
+    Use as `process_inputs` and/or `process_outputs` on a `TracePolicy` to keep a node's
+    span and its timing while omitting its inputs/outputs from the trace.
+    """
+    return {}
+
+
+_DEFAULT_INTERRUPT_ID = "placeholder-id"
+
+
+@final
+@dataclass(init=False, slots=True)
+class Interrupt(Generic[ResponseT]):
+    """Information about an interrupt that occurred in a node.
+
+    !!! version-added "Added in version 0.2.24"
+
+    !!! version-changed "Changed in version v0.4.0"
+        * `interrupt_id` was introduced as a property
+
+    !!! version-changed "Changed in version v0.6.0"
+
+        The following attributes have been removed:
+
+        * `ns`
+        * `when`
+        * `resumable`
+        * `interrupt_id`, deprecated in favor of `id`
+    """
+
+    value: Any
+    """The value associated with the interrupt."""
+
+    id: str
+    """The ID of the interrupt. Can be used to resume the interrupt directly."""
+
+    response_schema: type[ResponseT] | dict[str, Any] | None = None
+    """Schema for the value expected when resuming this interrupt, if the graph provided one.
+
+    A surfaced interrupt carries JSON Schema (a `dict`); `type[ResponseT]` records the
+    Python type at construction so `Interrupt[Decision]` is meaningful to type checkers."""
+
+    def __init__(
+        self,
+        value: Any,
+        id: str = _DEFAULT_INTERRUPT_ID,
+        *,
+        response_schema: type[ResponseT] | dict[str, Any] | None = None,
+        **deprecated_kwargs: Unpack[DeprecatedKwargs],
+    ) -> None:
+        self.value = value
+        self.response_schema = response_schema
+
+        if (
+            (ns := deprecated_kwargs.get("ns", MISSING)) is not MISSING
+            and (id == _DEFAULT_INTERRUPT_ID)
+            and (isinstance(ns, Sequence))
+        ):
+            self.id = xxh3_128_hexdigest("|".join(ns).encode())
+        else:
+            self.id = id
+
+    @classmethod
+    def from_ns(
+        cls,
+        value: Any,
+        ns: str,
+        *,
+        response_schema: type[ResponseT] | dict[str, Any] | None = None,
+    ) -> Interrupt[ResponseT]:
+        return cls(
+            value=value,
+            id=xxh3_128_hexdigest(ns.encode()),
+            response_schema=response_schema,
+        )
+
+    @property
+    @deprecated("`interrupt_id` is deprecated. Use `id` instead.", category=None)
+    def interrupt_id(self) -> str:
+        warn(
+            "`interrupt_id` is deprecated. Use `id` instead.",
+            CocoGraphDeprecatedSinceV10,
+            stacklevel=2,
+        )
+        return self.id
+
+
+class StateUpdate(NamedTuple):
+    values: dict[str, Any] | None
+    as_node: str | None = None
+    task_id: str | None = None
+
+
+class PregelTask(NamedTuple):
+    """A Pregel task."""
+
+    id: str
+    name: str
+    path: tuple[str | int | tuple, ...]
+    error: Exception | None = None
+    interrupts: tuple[Interrupt, ...] = ()
+    state: None | RunnableConfig | StateSnapshot = None
+    result: Any | None = None
+
+
+if sys.version_info > (3, 11):
+    _T_DC_KWARGS = {"weakref_slot": True, "slots": True, "frozen": True}
+else:
+    _T_DC_KWARGS = {"frozen": True}
+
+
+class CacheKey(NamedTuple):
+    """Cache key for a task."""
+
+    ns: tuple[str, ...]
+    """Namespace for the cache entry."""
+    key: str
+    """Key for the cache entry."""
+    ttl: int | None
+    """Time to live for the cache entry in seconds."""
+
+
+@dataclass(**_T_DC_KWARGS)
+class PregelExecutableTask:
+    name: str
+    input: Any
+    proc: Runnable
+    writes: deque[tuple[str, Any]]
+    config: RunnableConfig
+    triggers: Sequence[str]
+    retry_policy: Sequence[RetryPolicy]
+    cache_key: CacheKey | None
+    id: str
+    path: tuple[str | int | tuple, ...]
+    writers: Sequence[Runnable] = ()
+    subgraphs: Sequence[PregelProtocol] = ()
+    timeout: TimeoutPolicy | None = None
+
+
+class StateSnapshot(NamedTuple):
+    """Snapshot of the state of the graph at the beginning of a step."""
+
+    values: dict[str, Any] | Any
+    """Current values of channels."""
+    next: tuple[str, ...]
+    """The name of the node to execute in each task for this step."""
+    config: RunnableConfig
+    """Config used to fetch this snapshot."""
+    metadata: CheckpointMetadata | None
+    """Metadata associated with this snapshot."""
+    created_at: str | None
+    """Timestamp of snapshot creation."""
+    parent_config: RunnableConfig | None
+    """Config used to fetch the parent snapshot, if any."""
+    tasks: tuple[PregelTask, ...]
+    """Tasks to execute in this step. If already attempted, may contain an error."""
+    interrupts: tuple[Interrupt, ...]
+    """Interrupts that occurred in this step that are pending resolution."""
+
+
+class Send:
+    """A message or packet to send to a specific node in the graph.
+
+    The `Send` class is used within a `StateGraph`'s conditional edges to
+    dynamically invoke a node with a custom state at the next step.
+
+    Importantly, the sent state can differ from the core graph's state,
+    allowing for flexible and dynamic workflow management.
+
+    One such example is a "map-reduce" workflow where your graph invokes
+    the same node multiple times in parallel with different states,
+    before aggregating the results back into the main graph's state.
+
+    Attributes:
+        node (str): The name of the target node to send the message to.
+        arg (Any): The state or message to send to the target node.
+        timeout (TimeoutPolicy | None): Optional timeout policy for this specific
+            pushed task. If omitted, the target node's timeout policy is used.
+
+    !!! example
+
+        ```python
+        from typing import Annotated
+        from coco_graph.types import Send
+        from coco_graph.graph import END, START
+        from coco_graph.graph import StateGraph
+        import operator
+
+        class OverallState(TypedDict):
+            subjects: list[str]
+            jokes: Annotated[list[str], operator.add]
+
+        def continue_to_jokes(state: OverallState):
+            return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
+
+        builder = StateGraph(OverallState)
+        builder.add_node("generate_joke", lambda state: {"jokes": [f"Joke about {state['subject']}"]})
+        builder.add_conditional_edges(START, continue_to_jokes)
+        builder.add_edge("generate_joke", END)
+        graph = builder.compile()
+
+        # Invoking with two subjects results in a generated joke for each
+        graph.invoke({"subjects": ["cats", "dogs"]})
+        # {'subjects': ['cats', 'dogs'], 'jokes': ['Joke about cats', 'Joke about dogs']}
+        ```
+    """
+
+    __slots__ = ("node", "arg", "timeout")
+
+    node: str
+    arg: Any
+    timeout: TimeoutPolicy | None
+
+    def __init__(
+        self,
+        /,
+        node: str,
+        arg: Any,
+        *,
+        timeout: float | timedelta | TimeoutPolicy | None = None,
+    ) -> None:
+        """
+        Initialize a new instance of the `Send` class.
+
+        Args:
+            node: The name of the target node to send the message to.
+            arg: The state or message to send to the target node.
+            timeout: Optional timeout policy for this specific pushed task. A
+                number or `timedelta` is treated as a hard `run_timeout`.
+        """
+        self.node = node
+        self.arg = arg
+        self.timeout = TimeoutPolicy.coerce(timeout)
+
+    def __hash__(self) -> int:
+        return hash((self.node, self.arg, self.timeout))
+
+    def __repr__(self) -> str:
+        if self.timeout is None:
+            return f"Send(node={self.node!r}, arg={self.arg!r})"
+        return f"Send(node={self.node!r}, arg={self.arg!r}, timeout={self.timeout!r})"
+
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, Send)
+            and self.node == value.node
+            and self.arg == value.arg
+            and self.timeout == value.timeout
+        )
+
+
+N = TypeVar("N", bound=Hashable)
+
+
+@dataclass(**_DC_KWARGS)
+class Command(Generic[N], ToolOutputMixin):
+    """One or more commands to update the graph's state and send messages to nodes.
+
+    Args:
+        graph: Graph to send the command to. Supported values are:
+
+            - `None`: the current graph
+            - `Command.PARENT`: closest parent graph
+        update: Update to apply to the graph's state.
+        resume: Value to resume execution with. To be used together with [`interrupt()`][coco_graph.types.interrupt].
+            Can be one of the following:
+
+            - Mapping of interrupt ids to resume values
+            - A single value with which to resume the next interrupt
+        goto: Can be one of the following:
+
+            - Name of the node to navigate to next (any node that belongs to the specified `graph`)
+            - Sequence of node names to navigate to next
+            - `Send` object (to execute a node with the input provided)
+            - Sequence of `Send` objects
+    """
+
+    graph: str | None = None
+    update: Any | None = None
+    resume: dict[str, Any] | Any | None = None
+    goto: Send | Sequence[Send | N] | N = ()
+
+    def __repr__(self) -> str:
+        # get all non-None values
+        contents = ", ".join(
+            f"{key}={value!r}" for key, value in asdict(self).items() if value
+        )
+        return f"Command({contents})"
+
+    def _update_as_tuples(self) -> Sequence[tuple[str, Any]]:
+        if isinstance(self.update, dict):
+            return list(self.update.items())
+        elif isinstance(self.update, (list, tuple)) and all(
+            isinstance(t, tuple) and len(t) == 2 and isinstance(t[0], str)
+            for t in self.update
+        ):
+            return self.update
+        elif keys := get_cached_annotated_keys(type(self.update)):
+            return get_update_as_tuples(self.update, keys)
+        elif self.update is not None:
+            return [("__root__", self.update)]
+        else:
+            return []
+
+    PARENT: ClassVar[Literal["__parent__"]] = "__parent__"
+
+
+@overload
+def interrupt(value: Any, *, response_schema: type[ResponseT]) -> ResponseT: ...
+
+
+@overload
+def interrupt(value: Any, *, response_schema: dict[str, Any] | None = None) -> Any: ...
+
+
+def interrupt(
+    value: Any, *, response_schema: dict[str, Any] | type | None = None
+) -> Any:
+    """Interrupt the graph with a resumable exception from within a node.
+
+    The `interrupt` function enables human-in-the-loop workflows by pausing graph
+    execution and surfacing a value to the client. This value can communicate context
+    or request input required to resume execution.
+
+    In a given node, the first invocation of this function raises a `GraphInterrupt`
+    exception, halting execution. The provided `value` is included with the exception
+    and sent to the client executing the graph.
+
+    A client resuming the graph must use the [`Command`][coco_graph.types.Command]
+    primitive to specify a value for the interrupt and continue execution.
+    The graph resumes from the start of the node, **re-executing** all logic.
+
+    If a node contains multiple `interrupt` calls, CocoGraph matches resume values
+    to interrupts based on their order in the node. This list of resume values
+    is scoped to the specific task executing the node and is not shared across tasks.
+
+    To use an `interrupt`, you must enable a checkpointer, as the feature relies
+    on persisting the graph state.
+
+    !!! example
+
+        ```python
+        import uuid
+        from typing import Optional
+        from typing_extensions import TypedDict
+
+        from coco_graph.checkpoint.memory import InMemorySaver
+        from coco_graph.constants import START
+        from coco_graph.graph import StateGraph
+        from coco_graph.types import interrupt, Command
+
+
+        class State(TypedDict):
+            \"\"\"The graph state.\"\"\"
+
+            foo: str
+            human_value: Optional[str]
+            \"\"\"Human value will be updated using an interrupt.\"\"\"
+
+
+        def node(state: State):
+            answer = interrupt(
+                # This value will be sent to the client
+                # as part of the interrupt information.
+                \"what is your age?\"
+            )
+            print(f\"> Received an input from the interrupt: {answer}\")
+            return {\"human_value\": answer}
+
+
+        builder = StateGraph(State)
+        builder.add_node(\"node\", node)
+        builder.add_edge(START, \"node\")
+
+        # A checkpointer must be enabled for interrupts to work!
+        checkpointer = InMemorySaver()
+        graph = builder.compile(checkpointer=checkpointer)
+
+        config = {
+            \"configurable\": {
+                \"thread_id\": uuid.uuid4(),
+            }
+        }
+
+        for chunk in graph.stream({\"foo\": \"abc\"}, config):
+            print(chunk)
+
+        # > {'__interrupt__': (Interrupt(value='what is your age?', id='45fda8478b2ef754419799e10992af06', response_schema=None),)}
+
+        command = Command(resume=\"some input from a human!!!\")
+
+        for chunk in graph.stream(Command(resume=\"some input from a human!!!\"), config):
+            print(chunk)
+
+        # > Received an input from the interrupt: some input from a human!!!
+        # > {'node': {'human_value': 'some input from a human!!!'}}
+        ```
+
+    Args:
+        value: The value to surface to the client when the graph is interrupted.
+        response_schema: Optional schema for the value expected on resume, surfaced
+            to clients so they can render a typed input form. Accepts a JSON Schema
+            `dict` (used as-is, resume values are not validated), or a Pydantic model
+            class, `TypedDict`, or dataclass, which are converted to JSON Schema for
+            clients and used to validate the resume value; the validated object is
+            what `interrupt` returns.
+
+    Returns:
+        Any: On subsequent invocations within the same node (same task to be precise), returns the value provided during the first invocation,
+            validated against `response_schema` when one that supports validation was given.
+
+    Raises:
+        GraphInterrupt: On the first invocation within the node, halts execution and surfaces the provided value to the client.
+        pydantic.ValidationError: When a resume value does not match a Pydantic model, `TypedDict`, or dataclass `response_schema`.
+    """
+    from coco_graph._internal._constants import (
+        CONFIG_KEY_CHECKPOINT_NS,
+        CONFIG_KEY_SCRATCHPAD,
+        CONFIG_KEY_SEND,
+        RESUME,
+    )
+    from coco_graph.config import get_config
+    from coco_graph.errors import GraphInterrupt
+
+    conf = get_config()["configurable"]
+    adapter = (
+        None
+        if response_schema is None or isinstance(response_schema, dict)
+        else TypeAdapter(response_schema)
+    )
+    # track interrupt index
+    scratchpad = conf[CONFIG_KEY_SCRATCHPAD]
+    idx = scratchpad.interrupt_counter()
+    # find previous resume values
+    if scratchpad.resume:
+        if idx < len(scratchpad.resume):
+            v = scratchpad.resume[idx]
+            validated = adapter.validate_python(v) if adapter else v
+            conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume[: idx + 1])])
+            return validated
+    # find current resume value
+    v = scratchpad.get_null_resume(True)
+    if v is not None:
+        assert len(scratchpad.resume) == idx, (scratchpad.resume, idx)
+        validated = adapter.validate_python(v) if adapter else v
+        scratchpad.resume.append(v)
+        conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)])
+        return validated
+    # no resume value found
+    raise GraphInterrupt(
+        (
+            Interrupt.from_ns(
+                value=value,
+                ns=conf[CONFIG_KEY_CHECKPOINT_NS],
+                response_schema=adapter.json_schema() if adapter else response_schema,
+            ),
+        )
+    )
+
+
+@dataclass(slots=True)
+class Overwrite:
+    """Bypass a reducer and write the wrapped value directly to a `BinaryOperatorAggregate` channel.
+
+    Receiving multiple `Overwrite` values for the same channel in a single super-step
+    will raise an `InvalidUpdateError`.
+
+    !!! example
+
+        ```python
+        from typing import Annotated
+        import operator
+        from coco_graph.graph import StateGraph
+        from coco_graph.types import Overwrite
+
+        class State(TypedDict):
+            messages: Annotated[list, operator.add]
+
+        def node_a(state: TypedDict):
+            # Normal update: uses the reducer (operator.add)
+            return {"messages": ["a"]}
+
+        def node_b(state: State):
+            # Overwrite: bypasses the reducer and replaces the entire value
+            return {"messages": Overwrite(value=["b"])}
+
+        builder = StateGraph(State)
+        builder.add_node("node_a", node_a)
+        builder.add_node("node_b", node_b)
+        builder.set_entry_point("node_a")
+        builder.add_edge("node_a", "node_b")
+        graph = builder.compile()
+
+        # Without Overwrite in node_b, messages would be ["START", "a", "b"]
+        # With Overwrite, messages is just ["b"]
+        result = graph.invoke({"messages": ["START"]})
+        assert result == {"messages": ["b"]}
+        ```
+    """
+
+    value: Any
+    """The value to write directly to the channel, bypassing any reducer."""
+
+    type: Literal["__overwrite__"] = "__overwrite__"
+    """Discriminator field. Lets the channel reducer recognise an `Overwrite`
+    even after its dataclass form is JSON-serialised and the typed instance
+    is lost (e.g. an `orjson`-encoded state update routed through the
+    CocoGraph API server)."""