coco-framework 0.1.0__py3-none-any.whl

coco_graph/pregel/_retry.py

@@ -0,0 +1,854 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+import sys
+import threading
+import time
+import weakref
+from collections.abc import Awaitable, Callable, Sequence
+from contextlib import suppress
+from dataclasses import dataclass, replace
+from datetime import datetime, timedelta, timezone
+from typing import Any, Literal, NamedTuple
+
+from coco_core.callbacks import BaseCallbackHandler
+from coco_core.runnables import RunnableConfig
+
+from coco_graph._internal._config import (
+    merge_configs,
+    patch_configurable,
+    recast_checkpoint_ns,
+)
+from coco_graph._internal._constants import (
+    CONF,
+    CONFIG_KEY_CALL,
+    CONFIG_KEY_CHECKPOINT_ID,
+    CONFIG_KEY_CHECKPOINT_NS,
+    CONFIG_KEY_RESUMING,
+    CONFIG_KEY_RUNTIME,
+    CONFIG_KEY_SEND,
+    CONFIG_KEY_STREAM,
+    CONFIG_KEY_TASK_ID,
+    CONFIG_KEY_THREAD_ID,
+    CONFIG_KEY_TIMED_ATTEMPT_OBSERVER,
+    NS_SEP,
+)
+from coco_graph._internal._runnable import create_task_in_config_context
+from coco_graph._internal._timeout import sync_timeout_unsupported
+from coco_graph.errors import (
+    GraphBubbleUp,
+    NodeCancelledError,
+    NodeTimeoutError,
+    ParentCommand,
+)
+from coco_graph.pregel.protocol import StreamProtocol
+from coco_graph.runtime import ExecutionInfo, Runtime
+from coco_graph.types import Command, PregelExecutableTask, RetryPolicy, TimeoutPolicy
+
+logger = logging.getLogger(__name__)
+SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
+# `asyncio.Task.cancelling()` was added in Python 3.11. It reports the number of
+# pending cancel requests on the task: ``0`` means no external code asked us to
+# cancel — so a ``CancelledError`` observed here was raised by the task body
+# itself (the user's node) rather than by pregel cancelling sibling tasks.
+SUPPORTS_TASK_CANCELLING = sys.version_info >= (3, 11)
+
+
+def _timeout_secs(value: float | timedelta) -> float:
+    return value.total_seconds() if isinstance(value, timedelta) else value
+
+
+@dataclass(frozen=True, slots=True)
+class _ResolvedTimeout:
+    run_timeout_secs: float | None
+    idle_timeout_secs: float | None
+    refresh_on: Literal["auto", "heartbeat"] | None
+
+
+def _resolve_timeout(timeout: TimeoutPolicy) -> _ResolvedTimeout:
+    idle_timeout_secs = (
+        _timeout_secs(timeout.idle_timeout)
+        if timeout.idle_timeout is not None
+        else None
+    )
+    return _ResolvedTimeout(
+        run_timeout_secs=(
+            _timeout_secs(timeout.run_timeout)
+            if timeout.run_timeout is not None
+            else None
+        ),
+        idle_timeout_secs=idle_timeout_secs,
+        refresh_on=timeout.refresh_on if idle_timeout_secs is not None else None,
+    )
+
+
+class _AttemptContext(NamedTuple):
+    """Immutable per-attempt metadata shared across start/progress/finish events.
+
+    Built once at attempt start and referenced (not copied) by every emitted
+    `_AttemptEvent`, so per-event allocation is just the small event wrapper.
+
+    Intentionally underscore-prefixed: this and `_AttemptEvent` are part of an
+    internal observer contract consumed by trustcoco-graph-server. Do not move to
+    `coco_graph.types` — server imports them by this path.
+    """
+
+    task_id: str
+    task_name: str
+    attempt: int
+    run_id: str | None
+    thread_id: str | None
+    checkpoint_ns: str | None
+    started_at: datetime
+    run_timeout_secs: float | None
+    idle_timeout_secs: float | None
+    refresh_on: Literal["auto", "heartbeat"] | None
+
+
+@dataclass(frozen=True, slots=True)
+class _AttemptEvent:
+    """One lifecycle event for a timed attempt.
+
+    Holds a reference to the shared `_AttemptContext` and the event-specific
+    fields. The observer must treat this and `context` as read-only — they
+    are reused across all events for the same attempt.
+    """
+
+    context: _AttemptContext
+    event: Literal["start", "progress", "finish"]
+    progress_at: datetime | None = None
+    finished_at: datetime | None = None
+    status: Literal["success", "error"] | None = None
+    error_type: str | None = None
+    error_message: str | None = None
+
+
+class _TimedAttemptScope:
+    """Guarded-config window for timed attempts.
+
+    The wrapped config marks writes, stream events, runtime stream writer calls,
+    child task scheduling, and any Coco callback event emitted under the
+    node's run as observable progress when `refresh_on="auto"`.
+    `runtime.heartbeat()` exposes a manual progress signal for work that doesn't
+    otherwise emit any of these, and is the only progress signal when
+    `refresh_on="heartbeat"`.
+    Guarded writes are serialized with `close()` so cancelled background tasks
+    cannot persist writes past the timeout boundary. Stream/custom output is
+    best-effort: it is dropped after close is observed, but callbacks run outside
+    the lock because they may contain arbitrary user/runtime code.
+    """
+
+    __slots__ = (
+        "__weakref__",
+        "_active",
+        "_last_progress",
+        "_last_progress_emit",
+        "_lock",
+        "_on_progress",
+        "_progress_min_interval",
+        "_refresh_on",
+    )
+
+    def __init__(
+        self,
+        on_progress: Callable[[], None] | None = None,
+        progress_min_interval: float = 0.0,
+        refresh_on: Literal["auto", "heartbeat"] | None = None,
+    ) -> None:
+        self._active = True
+        self._last_progress = time.monotonic()
+        self._lock = threading.Lock()
+        self._on_progress = on_progress
+        self._progress_min_interval = progress_min_interval
+        self._refresh_on = refresh_on
+        # `-inf` so the first touch always passes the rate-limit gate.
+        self._last_progress_emit: float = float("-inf")
+
+    def wrap_config(self, config: RunnableConfig) -> RunnableConfig:
+        configurable = config.get(CONF, {})
+        patch: dict[str, Any] = {}
+        if (send := configurable.get(CONFIG_KEY_SEND)) is not None:
+            patch[CONFIG_KEY_SEND] = self._guard_send(send)
+        if (stream := configurable.get(CONFIG_KEY_STREAM)) is not None:
+            patch[CONFIG_KEY_STREAM] = self._guard_stream(stream)
+        if (call := configurable.get(CONFIG_KEY_CALL)) is not None:
+            patch[CONFIG_KEY_CALL] = self._guard_call(call)
+        if isinstance(runtime := configurable.get(CONFIG_KEY_RUNTIME), Runtime):
+            if self._refresh_on is not None:
+                patch[CONFIG_KEY_RUNTIME] = runtime.override(
+                    stream_writer=self._guard_stream_writer(runtime.stream_writer),
+                    heartbeat=self.touch,
+                )
+            else:
+                patch[CONFIG_KEY_RUNTIME] = runtime.override(
+                    stream_writer=self._guard_stream_writer(runtime.stream_writer)
+                )
+        new_config = patch_configurable(config, patch) if patch else config
+        if self._refresh_on == "auto":
+            return merge_configs(
+                new_config, {"callbacks": [_IdleProgressCallbackHandler(self)]}
+            )
+        return new_config
+
+    def touch(self) -> None:
+        # Avoid locking this hot progress path. We accept a small race window in
+        # timestamp ordering because idle_timeout is expected to be coarse compared
+        # with scheduler/thread timing.
+        now = time.monotonic()
+        self._last_progress = now
+        if self._on_progress is None:
+            return
+        # Best-effort rate limit: a benign race may emit a duplicate progress
+        # event under heavy concurrency, which observers must already tolerate
+        # (callbacks fire from arbitrary threads).
+        if now - self._last_progress_emit < self._progress_min_interval:
+            return
+        self._last_progress_emit = now
+        self._on_progress()
+
+    def close(self) -> None:
+        with self._lock:
+            self._active = False
+
+    async def wait_for_idle_timeout(self, idle_timeout_s: float) -> None:
+        while True:
+            with self._lock:
+                if not self._active:
+                    return
+                remaining = self._last_progress + idle_timeout_s - time.monotonic()
+            if remaining <= 0:
+                raise asyncio.TimeoutError
+            await asyncio.sleep(remaining)
+
+    def _guard_send(
+        self, send: Callable[[Sequence[tuple[str, Any]]], None]
+    ) -> Callable[[Sequence[tuple[str, Any]]], None]:
+        def guarded_send(writes: Sequence[tuple[str, Any]]) -> None:
+            with self._lock:
+                if self._active:
+                    if writes and self._refresh_on == "auto":
+                        self._last_progress = time.monotonic()
+                    send(writes)
+
+        return guarded_send
+
+    def _guard_stream(self, stream: StreamProtocol) -> StreamProtocol:
+        # No lock: stream callbacks fire from the event loop only, so the
+        # active-check + write happen atomically between awaits.
+        def guarded_stream(chunk: tuple[tuple[str, ...], str, Any]) -> None:
+            if not self._active:
+                return
+            if self._refresh_on == "auto":
+                self._last_progress = time.monotonic()
+            stream(chunk)
+
+        return StreamProtocol(guarded_stream, stream.modes)
+
+    def _guard_call(self, call: Callable[..., Any]) -> Callable[..., Any]:
+        # No lock: child-task scheduling happens from the event loop only.
+        def guarded_call(*args: Any, **kwargs: Any) -> Any:
+            if not self._active:
+                raise asyncio.CancelledError
+            if self._refresh_on == "auto":
+                self._last_progress = time.monotonic()
+            return call(*args, **kwargs)
+
+        return guarded_call
+
+    def _guard_stream_writer(
+        self, stream_writer: Callable[[Any], None]
+    ) -> Callable[[Any], None]:
+        def guarded_stream_writer(chunk: Any) -> None:
+            with self._lock:
+                if not self._active:
+                    return
+                if self._refresh_on == "auto":
+                    self._last_progress = time.monotonic()
+            stream_writer(chunk)
+
+        return guarded_stream_writer
+
+
+class _IdleProgressCallbackHandler(BaseCallbackHandler):
+    """Resets the idle timeout clock on any Coco callback event.
+
+    Inherits via `config["callbacks"]`, so it sees only events emitted by
+    runs descended from the node's attempt — sibling nodes do not bleed
+    through. Holds the scope by weakref so a child manager that outlives
+    the attempt cannot keep the scope alive.
+    """
+
+    # Run inline so progress is recorded in callback emission order;
+    # thread-pool dispatch would introduce extra reordering.
+    run_inline = True
+
+    def __init__(self, scope: _TimedAttemptScope) -> None:
+        self._scope_ref = weakref.ref(scope)
+
+    def _touch(self, *args: Any, **kwargs: Any) -> None:
+        if (scope := self._scope_ref()) is not None:
+            scope.touch()
+
+    on_llm_start = _touch
+    on_chat_model_start = _touch
+    on_llm_new_token = _touch
+    on_llm_end = _touch
+    on_llm_error = _touch
+    on_chain_start = _touch
+    on_chain_end = _touch
+    on_chain_error = _touch
+    on_tool_start = _touch
+    on_tool_end = _touch
+    on_tool_error = _touch
+    on_retriever_start = _touch
+    on_retriever_end = _touch
+    on_retriever_error = _touch
+    on_agent_action = _touch
+    on_agent_finish = _touch
+    on_text = _touch
+    on_retry = _touch
+    on_custom_event = _touch
+
+
+def _is_user_raised_cancelled() -> bool:
+    """Return True if the in-flight ``CancelledError`` came from the task body.
+
+    Pregel cancels sibling tasks via ``task.cancel()`` when a peer fails, which
+    increments ``asyncio.Task.cancelling()`` on the target before the cancel
+    actually fires. A user node that calls ``raise asyncio.CancelledError()``
+    from inside its own body raises while ``cancelling() == 0``, which is the
+    signal we use to convert the exception into a regular
+    :class:`NodeCancelledError`.
+
+    Returns ``False`` when we can't tell (``cancelling()`` unavailable, or no
+    current task — neither should happen in practice from ``arun_with_retry``)
+    so framework-initiated cancellation continues to propagate unchanged.
+    """
+    if not SUPPORTS_TASK_CANCELLING:
+        return False
+    current = asyncio.current_task()
+    if current is None:
+        return False
+    return current.cancelling() == 0
+
+
+def _drain_cancelled(task: asyncio.Task[Any]) -> None:
+    # Mark the abandoned task's exception as retrieved so asyncio doesn't log it.
+    with suppress(asyncio.CancelledError):
+        task.exception()
+
+
+def _start_timed_attempt(
+    task: PregelExecutableTask, config: RunnableConfig, timeout: _ResolvedTimeout
+) -> _AttemptContext | None:
+    configurable = config.get(CONF, {})
+    callback = configurable.get(CONFIG_KEY_TIMED_ATTEMPT_OBSERVER)
+    if callback is None:
+        return None
+    runtime = configurable.get(CONFIG_KEY_RUNTIME)
+    execution_info = runtime.execution_info if isinstance(runtime, Runtime) else None
+    context = _AttemptContext(
+        task_id=task.id,
+        task_name=task.name,
+        attempt=execution_info.node_attempt if execution_info is not None else 1,
+        run_id=execution_info.run_id if execution_info is not None else None,
+        thread_id=execution_info.thread_id if execution_info is not None else None,
+        checkpoint_ns=(
+            execution_info.checkpoint_ns if execution_info is not None else None
+        ),
+        started_at=datetime.now(timezone.utc),
+        run_timeout_secs=timeout.run_timeout_secs,
+        idle_timeout_secs=timeout.idle_timeout_secs,
+        refresh_on=timeout.refresh_on,
+    )
+    _dispatch_observer(callback, _AttemptEvent(context=context, event="start"))
+    return context
+
+
+def _finish_timed_attempt(
+    config: RunnableConfig,
+    context: _AttemptContext | None,
+    error: BaseException | None = None,
+) -> None:
+    if context is None:
+        return
+    callback = config.get(CONF, {}).get(CONFIG_KEY_TIMED_ATTEMPT_OBSERVER)
+    if callback is None:
+        return
+    _dispatch_observer(
+        callback,
+        _AttemptEvent(
+            context=context,
+            event="finish",
+            finished_at=datetime.now(timezone.utc),
+            status="error" if error is not None else "success",
+            error_type=type(error).__name__ if error is not None else None,
+            error_message=str(error) if error is not None else None,
+        ),
+    )
+
+
+def _emit_progress(
+    callback: Callable[[_AttemptEvent], None],
+    context: _AttemptContext,
+) -> None:
+    _dispatch_observer(
+        callback,
+        _AttemptEvent(
+            context=context,
+            event="progress",
+            progress_at=datetime.now(timezone.utc),
+        ),
+    )
+
+
+def _dispatch_observer(
+    callback: Callable[[_AttemptEvent], None],
+    event: _AttemptEvent,
+) -> None:
+    try:
+        callback(event)
+    except Exception:
+        logger.warning("Timed attempt observer failed", exc_info=True)
+
+
+async def _run_timeout_watchdog(run_timeout_s: float) -> None:
+    await asyncio.sleep(run_timeout_s)
+    raise asyncio.TimeoutError
+
+
+async def _arun_with_timeout(
+    task: PregelExecutableTask,
+    config: RunnableConfig,
+    timeout: _ResolvedTimeout,
+    attempt_ctx: _AttemptContext | None,
+    *,
+    stream: bool,
+) -> Any:
+    run_timeout_s = timeout.run_timeout_secs
+    idle_timeout_s = timeout.idle_timeout_secs
+    on_progress: Callable[[], None] | None = None
+    if attempt_ctx is not None:
+        callback = config.get(CONF, {}).get(CONFIG_KEY_TIMED_ATTEMPT_OBSERVER)
+        if callback is not None and idle_timeout_s is not None:
+            on_progress = lambda: _emit_progress(callback, attempt_ctx)  # noqa: E731
+    scope = _TimedAttemptScope(
+        on_progress=on_progress,
+        # Cap progress emission at ~4 events per idle window so token-rate
+        # callbacks don't flood the observer.
+        progress_min_interval=idle_timeout_s / 4 if idle_timeout_s is not None else 0.0,
+        refresh_on=timeout.refresh_on,
+    )
+    scoped_config = scope.wrap_config(config)
+    start = time.monotonic()
+    if stream:
+        # Yielded chunks count as progress only under `refresh_on="auto"`.
+        # `refresh_on="heartbeat"` is the strict mode where only explicit
+        # `runtime.heartbeat()` calls reset the idle clock.
+        async def run() -> Any:
+            async for _ in task.proc.astream(task.input, scoped_config):
+                if timeout.refresh_on == "auto":
+                    scope.touch()
+
+    else:
+
+        async def run() -> Any:
+            return await task.proc.ainvoke(task.input, scoped_config)
+
+    bg = create_task_in_config_context(run, scoped_config)
+    watchdogs: dict[asyncio.Task[None], Literal["idle", "run"]] = {}
+    if idle_timeout_s is not None:
+        watchdogs[asyncio.create_task(scope.wait_for_idle_timeout(idle_timeout_s))] = (
+            "idle"
+        )
+    if run_timeout_s is not None:
+        watchdogs[asyncio.create_task(_run_timeout_watchdog(run_timeout_s))] = "run"
+    try:
+        done, _ = await asyncio.wait(
+            {bg, *watchdogs}, return_when=asyncio.FIRST_COMPLETED
+        )
+        if bg in done:
+            # Task completed in time.
+            for watchdog in watchdogs:
+                watchdog.cancel()
+            # FIRST_COMPLETED can return both; a watchdog may have
+            # already raised TimeoutError before we cancelled it.
+            for watchdog in watchdogs:
+                with suppress(asyncio.CancelledError, asyncio.TimeoutError):
+                    await watchdog
+            return await bg
+        # bg was not in `done`, so every member of `done` is one of our
+        # watchdogs. Only a watchdog's TimeoutError converts to
+        # NodeTimeoutError; any TimeoutError raised by the proc itself
+        # propagates unchanged.
+        for watchdog in done:
+            kind = watchdogs[watchdog]
+            try:
+                await watchdog
+            except asyncio.TimeoutError as exc:
+                elapsed = time.monotonic() - start
+                scope.close()
+                task.writes.clear()
+                bg.cancel()
+                bg.add_done_callback(_drain_cancelled)
+                raise NodeTimeoutError(
+                    task.name,
+                    elapsed,
+                    kind=kind,
+                    idle_timeout=idle_timeout_s,
+                    run_timeout=run_timeout_s,
+                ) from exc
+            raise RuntimeError(
+                f"{kind} timeout watchdog completed without raising TimeoutError"
+            )
+        raise RuntimeError("timeout wait completed without task or watchdog")
+    except asyncio.CancelledError:
+        scope.close()
+        bg.cancel()
+        for watchdog in watchdogs:
+            watchdog.cancel()
+        bg.add_done_callback(_drain_cancelled)
+        raise
+    finally:
+        scope.close()
+        for watchdog in watchdogs:
+            watchdog.cancel()
+
+
+def _ensure_execution_info(
+    runtime: Runtime, config: RunnableConfig, task: PregelExecutableTask
+) -> Runtime:
+    """Ensure runtime has execution_info, creating one from config if needed.
+
+    In the distributed runtime (CocoGraph Platform), tasks are prepared by the
+    server and deserialized in the executor, bypassing the OSS _algo.py code
+    that normally creates ExecutionInfo. This function fills in execution_info
+    from the task config when it's missing.
+    """
+    if runtime.execution_info is not None:
+        return runtime
+    configurable = config.get(CONF, {})
+    return runtime.override(
+        execution_info=ExecutionInfo(
+            checkpoint_id=configurable.get(CONFIG_KEY_CHECKPOINT_ID) or "",
+            checkpoint_ns=configurable.get(CONFIG_KEY_CHECKPOINT_NS) or "",
+            task_id=configurable.get(CONFIG_KEY_TASK_ID) or task.id,
+            thread_id=configurable.get(CONFIG_KEY_THREAD_ID),
+            run_id=str(rid) if (rid := config.get("run_id")) else None,
+        ),
+    )
+
+
+def _checkpoint_ns_for_parent_command(ns: str) -> str:
+    """Return the checkpoint namespace for the parent graph.
+
+    The checkpoint namespace is a `|`-separated path. Each segment is usually
+    of the form `name:task_id` (e.g. `parent_first:<uuid>|node:<uuid>`), but the
+    runtime may also insert a purely-numeric segment (e.g. `|1`) to disambiguate
+    concurrent tasks (e.g. `parent_first:<uuid>|1|node:<uuid>`).
+
+    Numeric segments are not real path levels, so we drop them before computing
+    the parent namespace.
+    """
+
+    parts = ns.split(NS_SEP)
+
+    # Drop any trailing numeric selectors for the current frame (e.g. `...|node:<id>|1`).
+    while parts and parts[-1].isdigit():
+        parts.pop()
+
+    # Drop the current frame segment itself (e.g. the `node:<id>`).
+    if parts:
+        parts.pop()
+
+    # Drop any trailing numeric selectors for the parent frame (e.g. `...|1|node:<id>`).
+    while parts and parts[-1].isdigit():
+        parts.pop()
+
+    return NS_SEP.join(parts)
+
+
+def run_with_retry(
+    task: PregelExecutableTask,
+    retry_policy: Sequence[RetryPolicy] | None,
+    configurable: dict[str, Any] | None = None,
+) -> None:
+    """Run a task with retries."""
+    retry_policy = task.retry_policy or retry_policy
+    if task.timeout is not None:
+        # `validate_timeout_supported` catches sync nodes at compile time;
+        # this is a runtime safety net for paths that may bypass that validation.
+        raise sync_timeout_unsupported(task.name)
+    attempts = 0
+    node_first_attempt_time = time.time()
+    config = task.config
+    if configurable is not None:
+        config = patch_configurable(config, configurable)
+    runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
+    if isinstance(runtime, Runtime):
+        runtime = _ensure_execution_info(runtime, config, task)
+        config = patch_configurable(
+            config,
+            {
+                CONFIG_KEY_RUNTIME: runtime.patch_execution_info(
+                    node_first_attempt_time=node_first_attempt_time,
+                )
+            },
+        )
+    while True:
+        runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
+        if isinstance(runtime, Runtime):
+            config = patch_configurable(
+                config,
+                {
+                    CONFIG_KEY_RUNTIME: runtime.patch_execution_info(
+                        # node_attempt is execution count (1-indexed): 1 on first run,
+                        # then 2, 3, ... on subsequent retries.
+                        node_attempt=attempts + 1,
+                    )
+                },
+            )
+        try:
+            # clear any writes from previous attempts
+            task.writes.clear()
+            # run the task
+            return task.proc.invoke(task.input, config)
+        except ParentCommand as exc:
+            ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
+            cmd = exc.args[0]
+            # strip task_ids from namespace for comparison (ns format: "node1|node2:task_id")
+            if cmd.graph in (ns, recast_checkpoint_ns(ns), task.name):
+                # this command is for the current graph, handle it
+                for w in task.writers:
+                    w.invoke(cmd, config)
+                break
+            elif cmd.graph == Command.PARENT:
+                # this command is for the parent graph, assign it to the parent.
+                exc.args = (replace(cmd, graph=_checkpoint_ns_for_parent_command(ns)),)
+            # bubble up
+            raise
+        except GraphBubbleUp:
+            # if interrupted, end
+            raise
+        except asyncio.CancelledError as exc:
+            # A sync node has no asyncio context, so any ``CancelledError`` that
+            # reaches here was raised by the node body itself. Surface it as a
+            # regular exception so the pregel runner panics the run instead of
+            # treating the task as a silent tear-down (LSD-1507).
+            raise NodeCancelledError(task.name) from exc
+        except Exception as exc:
+            if SUPPORTS_EXC_NOTES:
+                exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
+            if not retry_policy:
+                raise
+
+            # Check which retry policy applies to this exception
+            matching_policy = None
+            for policy in retry_policy:
+                if _should_retry_on(policy, exc):
+                    matching_policy = policy
+                    break
+
+            if not matching_policy:
+                raise
+
+            # attempts tracks failed tries only; it increments after a failure.
+            attempts += 1
+            # check if we should give up
+            if attempts >= matching_policy.max_attempts:
+                raise
+            # sleep before retrying
+            interval = matching_policy.initial_interval
+            # Apply backoff factor based on attempt count
+            interval = min(
+                matching_policy.max_interval,
+                interval * (matching_policy.backoff_factor ** (attempts - 1)),
+            )
+
+            # Apply jitter if configured
+            sleep_time = (
+                interval + random.uniform(0, 1) if matching_policy.jitter else interval
+            )
+            time.sleep(sleep_time)
+
+            # log the retry
+            logger.info(
+                f"Retrying task {task.name} after {sleep_time:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}",
+                exc_info=exc,
+            )
+            # signal subgraphs to resume (if available)
+            config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
+
+
+async def arun_with_retry(
+    task: PregelExecutableTask,
+    retry_policy: Sequence[RetryPolicy] | None,
+    stream: bool = False,
+    match_cached_writes: Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
+    | None = None,
+    configurable: dict[str, Any] | None = None,
+) -> None:
+    """Run a task asynchronously with retries."""
+    retry_policy = task.retry_policy or retry_policy
+    resolved_timeout = (
+        _resolve_timeout(task.timeout) if task.timeout is not None else None
+    )
+    attempts = 0
+    node_first_attempt_time = time.time()
+    config = task.config
+    if configurable is not None:
+        config = patch_configurable(config, configurable)
+    runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
+    if isinstance(runtime, Runtime):
+        runtime = _ensure_execution_info(runtime, config, task)
+        config = patch_configurable(
+            config,
+            {
+                CONFIG_KEY_RUNTIME: runtime.patch_execution_info(
+                    node_first_attempt_time=node_first_attempt_time,
+                )
+            },
+        )
+    if match_cached_writes is not None and task.cache_key is not None:
+        for t in await match_cached_writes():
+            if t is task:
+                # if the task is already cached, return
+                return
+    while True:
+        runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
+        if isinstance(runtime, Runtime):
+            config = patch_configurable(
+                config,
+                {
+                    CONFIG_KEY_RUNTIME: runtime.patch_execution_info(
+                        # node_attempt is execution count (1-indexed): 1 on first run,
+                        # then 2, 3, ... on subsequent retries.
+                        node_attempt=attempts + 1,
+                    )
+                },
+            )
+        attempt_ctx = (
+            _start_timed_attempt(task, config, resolved_timeout)
+            if resolved_timeout is not None
+            else None
+        )
+        try:
+            task.writes.clear()
+            if resolved_timeout is None:
+                if stream:
+                    async for _ in task.proc.astream(task.input, config):
+                        pass
+                    break
+                return await task.proc.ainvoke(task.input, config)
+            result = await _arun_with_timeout(
+                task, config, resolved_timeout, attempt_ctx, stream=stream
+            )
+            _finish_timed_attempt(config, attempt_ctx)
+            if stream:
+                # if successful, end
+                break
+            return result
+        except ParentCommand as exc:
+            ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
+            cmd = exc.args[0]
+            # strip task_ids from namespace for comparison (ns format: "node1|node2:task_id")
+            if cmd.graph in (ns, recast_checkpoint_ns(ns), task.name):
+                try:
+                    # this command is for the current graph, handle it
+                    for w in task.writers:
+                        w.invoke(cmd, config)
+                except Exception as writer_exc:
+                    _finish_timed_attempt(config, attempt_ctx, writer_exc)
+                    raise
+                _finish_timed_attempt(config, attempt_ctx)
+                break
+            elif cmd.graph == Command.PARENT:
+                # this command is for the parent graph, assign it to the parent.
+                exc.args = (replace(cmd, graph=_checkpoint_ns_for_parent_command(ns)),)
+            _finish_timed_attempt(config, attempt_ctx)
+            # bubble up the exception to the parent graph
+            raise
+        except GraphBubbleUp:
+            # if interrupted, end
+            _finish_timed_attempt(config, attempt_ctx)
+            raise
+        except asyncio.CancelledError as exc:
+            # ``CancelledError`` reaches us in two very different shapes:
+            #   1. Pregel cancelled this task because a sibling failed
+            #      (``asyncio.Task.cancelling() >= 1``). The framework already
+            #      knows the run is failing and we must let cancellation
+            #      propagate so the watchdog/cleanup code in the runner sees a
+            #      cancelled future.
+            #   2. The node body itself raised ``asyncio.CancelledError`` (
+            #      ``cancelling() == 0``). The runner would otherwise treat
+            #      this as silent tear-down and the run would report
+            #      ``success`` even though the node failed (LSD-1507). Convert
+            #      it into :class:`NodeCancelledError` so it follows the same
+            #      path as any other node failure.
+            if _is_user_raised_cancelled():
+                _finish_timed_attempt(config, attempt_ctx, exc)
+                raise NodeCancelledError(task.name) from exc
+            _finish_timed_attempt(config, attempt_ctx, exc)
+            raise
+        except Exception as exc:
+            _finish_timed_attempt(config, attempt_ctx, exc)
+            if SUPPORTS_EXC_NOTES:
+                exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
+            if not retry_policy:
+                raise
+
+            # Check which retry policy applies to this exception
+            matching_policy = None
+            for policy in retry_policy:
+                if _should_retry_on(policy, exc):
+                    matching_policy = policy
+                    break
+
+            if not matching_policy:
+                raise
+
+            # attempts tracks failed tries only; it increments after a failure.
+            # The next execution's node_attempt is derived as attempts + 1.
+            attempts += 1
+            # check if we should give up
+            if attempts >= matching_policy.max_attempts:
+                raise
+            # sleep before retrying
+            interval = matching_policy.initial_interval
+            # Apply backoff factor based on attempt count
+            interval = min(
+                matching_policy.max_interval,
+                interval * (matching_policy.backoff_factor ** (attempts - 1)),
+            )
+
+            # Apply jitter if configured
+            sleep_time = (
+                interval + random.uniform(0, 1) if matching_policy.jitter else interval
+            )
+            await asyncio.sleep(sleep_time)
+
+            # log the retry
+            logger.info(
+                f"Retrying task {task.name} after {sleep_time:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}",
+                exc_info=exc,
+            )
+            # signal subgraphs to resume (if available)
+            config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
+
+
+def _should_retry_on(retry_policy: RetryPolicy, exc: Exception) -> bool:
+    """Check if the given exception should be retried based on the retry policy."""
+    if isinstance(retry_policy.retry_on, Sequence):
+        return isinstance(exc, tuple(retry_policy.retry_on))
+    elif isinstance(retry_policy.retry_on, type) and issubclass(
+        retry_policy.retry_on, Exception
+    ):
+        return isinstance(exc, retry_policy.retry_on)
+    elif callable(retry_policy.retry_on):
+        return retry_policy.retry_on(exc)  # type: ignore[call-arg]
+    else:
+        raise TypeError(
+            "retry_on must be an Exception class, a list or tuple of Exception classes, or a callable"
+        )