From 547e5fe158649d8aaeecd88274000244e5013e80 Mon Sep 17 00:00:00 2001 From: Vardhman Gupta Date: Wed, 7 Oct 2026 11:57:24 +0530 Subject: [PATCH] feat(dynaflow): add per-node retry policy, error fallback edges & state resumability (closes #279) Adds configurable RetryPolicy with exponential backoff and full jitter. Adds Edge.is_error_edge and SimpleGraph.add_error_edge() for graceful error recovery routing. Adds GraphRunState checkpointing and SimpleGraph.resume() to continue partially-failed runs without duplicate execution. Adds real-time NodeStatus lifecycle tracking across all built-in node types. --- src/dynavec/__init__.py | 6 + src/dynavec/nodes/__init__.py | 7 +- src/dynavec/nodes/base.py | 62 +++++++- src/dynavec/nodes/builtin.py | 39 ++++- src/dynavec/nodes/graph.py | 276 ++++++++++++++++++++++++++++++--- tests/test_nodes.py | 277 ++++++++++++++++++++++++++++++---- 6 files changed, 598 insertions(+), 69 deletions(-) diff --git a/src/dynavec/__init__.py b/src/dynavec/__init__.py index 080cc60..0acec56 100644 --- a/src/dynavec/__init__.py +++ b/src/dynavec/__init__.py @@ -81,14 +81,17 @@ ApprovalNode, CodeNode, Edge, + GraphRunState, LLMNode, LoopNode, Node, NodePort, NodeRegistry, NodeSchema, + NodeStatus, PortType, RetrieverNode, + RetryPolicy, RouterNode, SimpleGraph, SubGraphNode, @@ -203,10 +206,13 @@ "Node", "NodePort", "NodeSchema", + "NodeStatus", + "RetryPolicy", "PortType", "NodeRegistry", "node_registry", "Edge", + "GraphRunState", "SimpleGraph", "RetrieverNode", "LLMNode", diff --git a/src/dynavec/nodes/__init__.py b/src/dynavec/nodes/__init__.py index 57feb80..6bc53e7 100644 --- a/src/dynavec/nodes/__init__.py +++ b/src/dynavec/nodes/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations -from .base import Node, NodePort, NodeSchema, PortType +from .base import Node, NodePort, NodeSchema, NodeStatus, PortType, RetryPolicy from .builtin import ( ApprovalNode, CodeNode, @@ -14,14 +14,16 @@ ToolNode, WebhookNode, ) -from .graph import Edge, SimpleGraph +from .graph import Edge, GraphRunState, SimpleGraph from .registry import NodeRegistry, node_registry __all__ = [ "Node", "NodePort", "NodeSchema", + "NodeStatus", "PortType", + "RetryPolicy", "NodeRegistry", "node_registry", "RetrieverNode", @@ -34,5 +36,6 @@ "WebhookNode", "SubGraphNode", "Edge", + "GraphRunState", "SimpleGraph", ] diff --git a/src/dynavec/nodes/base.py b/src/dynavec/nodes/base.py index 77661ec..2f1dd42 100644 --- a/src/dynavec/nodes/base.py +++ b/src/dynavec/nodes/base.py @@ -4,10 +4,52 @@ import uuid from abc import ABC, abstractmethod +from collections.abc import Callable from dataclasses import dataclass, field from typing import Any, Literal PortType = Literal["string", "number", "boolean", "array", "object", "any"] +NodeStatus = Literal["pending", "running", "retrying", "succeeded", "failed", "skipped"] + + +@dataclass +class RetryPolicy: + """Configurable retry policy for node execution with exponential backoff and jitter.""" + + max_attempts: int = 3 + base_delay: float = 0.05 + max_delay: float = 2.0 + jitter: bool = True + retry_on: tuple[type[Exception], ...] | Callable[[Exception], bool] | None = None + + def should_retry(self, exc: Exception) -> bool: + """Evaluate if an exception is eligible for retry under this policy.""" + if self.retry_on is None: + return True + if isinstance(self.retry_on, tuple): + return isinstance(exc, self.retry_on) + if callable(self.retry_on): + return bool(self.retry_on(exc)) + return True + + def to_dict(self) -> dict[str, Any]: + """Serialize retry policy to dictionary.""" + return { + "max_attempts": self.max_attempts, + "base_delay": self.base_delay, + "max_delay": self.max_delay, + "jitter": self.jitter, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> RetryPolicy: + """Deserialize retry policy from dictionary.""" + return cls( + max_attempts=data.get("max_attempts", 3), + base_delay=data.get("base_delay", 0.05), + max_delay=data.get("max_delay", 2.0), + jitter=data.get("jitter", True), + ) @dataclass @@ -78,9 +120,14 @@ def __init__( self, node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: self.id = node_id or f"{self.node_type}_{uuid.uuid4().hex[:8]}" self.config = dict(config or {}) + self.retry_policy = retry_policy + self.status: NodeStatus = "pending" + self.attempts: int = 0 + self.last_error: str | None = None @abstractmethod def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: @@ -88,13 +135,18 @@ def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: def to_dict(self) -> dict[str, Any]: """Serialize node to a ReactFlow-compatible node object.""" + data: dict[str, Any] = { + "label": self.schema.label, + "config": self.config, + "category": self.schema.category, + "status": self.status, + } + if self.retry_policy is not None: + data["retry_policy"] = self.retry_policy.to_dict() + return { "id": self.id, "type": self.node_type, - "data": { - "label": self.schema.label, - "config": self.config, - "category": self.schema.category, - }, + "data": data, "position": {"x": 0, "y": 0}, } diff --git a/src/dynavec/nodes/builtin.py b/src/dynavec/nodes/builtin.py index 283afc6..b90b5b1 100644 --- a/src/dynavec/nodes/builtin.py +++ b/src/dynavec/nodes/builtin.py @@ -7,7 +7,7 @@ from typing import Any from ..chat.base import ChatModel, Message -from .base import Node, NodePort, NodeSchema +from .base import Node, NodePort, NodeSchema, RetryPolicy from .registry import node_registry # --- 1. RetrieverNode --- @@ -68,8 +68,9 @@ def __init__( retriever_fn: Callable[..., Any] | None = None, node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: - super().__init__(node_id=node_id, config=config) + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) self.retriever_fn = retriever_fn def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: @@ -163,8 +164,9 @@ def __init__( model: ChatModel | Callable[[list[Message]], Any], node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: - super().__init__(node_id=node_id, config=config) + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) self.model = model def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: @@ -249,8 +251,9 @@ def __init__( tool: Callable[..., Any], node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: - super().__init__(node_id=node_id, config=config) + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) self.tool = tool def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: @@ -335,8 +338,9 @@ def __init__( route_fn: Callable[[Any], str] | None = None, node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: - super().__init__(node_id=node_id, config=config) + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) self.route_fn = route_fn def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: @@ -391,8 +395,9 @@ def __init__( code_fn: Callable[[dict[str, Any]], dict[str, Any]], node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: - super().__init__(node_id=node_id, config=config) + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) self.code_fn = code_fn def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: @@ -460,6 +465,14 @@ class ApprovalNode(Node): }, ) + def __init__( + self, + node_id: str | None = None, + config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: return { "approved": bool(inputs.get("approved", True)), @@ -511,8 +524,9 @@ def __init__( map_fn: Callable[[Any], Any] | None = None, node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: - super().__init__(node_id=node_id, config=config) + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) self.map_fn = map_fn def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: @@ -571,6 +585,14 @@ class WebhookNode(Node): }, ) + def __init__( + self, + node_id: str | None = None, + config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, + ) -> None: + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) + def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: return {"payload": inputs.get("payload", {})} @@ -613,8 +635,9 @@ def __init__( subgraph: Any = None, node_id: str | None = None, config: dict[str, Any] | None = None, + retry_policy: RetryPolicy | None = None, ) -> None: - super().__init__(node_id=node_id, config=config) + super().__init__(node_id=node_id, config=config, retry_policy=retry_policy) self.subgraph = subgraph def execute(self, inputs: dict[str, Any]) -> dict[str, Any]: diff --git a/src/dynavec/nodes/graph.py b/src/dynavec/nodes/graph.py index 778c7e4..b7c6684 100644 --- a/src/dynavec/nodes/graph.py +++ b/src/dynavec/nodes/graph.py @@ -1,11 +1,14 @@ -"""Pipeline execution graph for wiring and running connected nodes.""" +"""Pipeline execution graph for wiring and running connected nodes with retries and resumability.""" from __future__ import annotations +import random +import time +import uuid from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal, overload -from .base import Node +from .base import Node, NodeStatus, RetryPolicy @dataclass @@ -17,15 +20,58 @@ class Edge: source_port: str = "output" target_port: str = "input" mapping: dict[str, str] = field(default_factory=dict) + is_error_edge: bool = False + + +@dataclass +class GraphRunState: + """State checkpoint representing the progress and status of a graph run.""" + + run_id: str = field(default_factory=lambda: f"run_{uuid.uuid4().hex[:8]}") + initial_inputs: dict[str, Any] = field(default_factory=dict) + current_state: dict[str, Any] = field(default_factory=dict) + node_outputs: dict[str, dict[str, Any]] = field(default_factory=dict) + node_statuses: dict[str, NodeStatus] = field(default_factory=dict) + node_errors: dict[str, str] = field(default_factory=dict) + completed: bool = False + failed_node: str | None = None + + def to_dict(self) -> dict[str, Any]: + """Serialize state to dictionary for persistent checkpoint storage.""" + return { + "run_id": self.run_id, + "initial_inputs": dict(self.initial_inputs), + "current_state": dict(self.current_state), + "node_outputs": dict(self.node_outputs), + "node_statuses": dict(self.node_statuses), + "node_errors": dict(self.node_errors), + "completed": self.completed, + "failed_node": self.failed_node, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> GraphRunState: + """Reconstruct state checkpoint from dictionary.""" + return cls( + run_id=data.get("run_id", f"run_{uuid.uuid4().hex[:8]}"), + initial_inputs=dict(data.get("initial_inputs", {})), + current_state=dict(data.get("current_state", {})), + node_outputs=dict(data.get("node_outputs", {})), + node_statuses=dict(data.get("node_statuses", {})), + node_errors=dict(data.get("node_errors", {})), + completed=data.get("completed", False), + failed_node=data.get("failed_node"), + ) class SimpleGraph: - """Directed execution pipeline to compose and run node graphs.""" + """Directed execution pipeline to compose and run resilient node graphs.""" def __init__(self) -> None: self.nodes: dict[str, Node] = {} self.edges: list[Edge] = [] self._execution_order: list[str] = [] + self.last_state: GraphRunState | None = None def add_node(self, node: Node) -> SimpleGraph: """Add a node instance to the graph.""" @@ -41,6 +87,7 @@ def add_edge( source_port: str = "output", target_port: str = "input", mapping: dict[str, str] | None = None, + is_error_edge: bool = False, ) -> SimpleGraph: """Connect source node output to target node input.""" if source_id not in self.nodes: @@ -54,36 +101,219 @@ def add_edge( source_port=source_port, target_port=target_port, mapping=dict(mapping or {}), + is_error_edge=is_error_edge, ) self.edges.append(edge) return self - def run(self, initial_inputs: dict[str, Any]) -> dict[str, Any]: - """Execute the graph sequentially, threading state along connected edges.""" - node_outputs: dict[str, dict[str, Any]] = {} - current_state: dict[str, Any] = dict(initial_inputs) + def add_error_edge( + self, + source_id: str, + target_id: str, + source_port: str = "error", + target_port: str = "error", + mapping: dict[str, str] | None = None, + ) -> SimpleGraph: + """Add a fallback error edge triggered when source node fails execution.""" + return self.add_edge( + source_id=source_id, + target_id=target_id, + source_port=source_port, + target_port=target_port, + mapping=mapping, + is_error_edge=True, + ) + + def _execute_node_with_retry( + self, node: Node, node_inputs: dict[str, Any] + ) -> tuple[dict[str, Any] | None, Exception | None]: + """Execute a node applying its configured RetryPolicy if present.""" + policy: RetryPolicy | None = getattr(node, "retry_policy", None) + max_attempts = policy.max_attempts if policy is not None else 1 + attempt = 0 + + while attempt < max_attempts: + attempt += 1 + node.attempts = attempt + try: + if attempt > 1: + node.status = "retrying" + else: + node.status = "running" + + outputs = node.execute(node_inputs) + node.status = "succeeded" + node.last_error = None + return outputs, None + except Exception as exc: # noqa: BLE001 + node.last_error = str(exc) + if policy is not None and attempt < max_attempts and policy.should_retry(exc): + delay = min(policy.max_delay, policy.base_delay * (2 ** (attempt - 1))) + if policy.jitter: + delay = random.uniform(0, delay) + if delay > 0: + time.sleep(delay) + continue + + node.status = "failed" + return None, exc + + node.status = "failed" + return None, RuntimeError(f"Node {node.id!r} exhausted {max_attempts} attempts.") + + @overload + def run( + self, + initial_inputs: dict[str, Any], + state: GraphRunState | dict[str, Any] | None = None, + return_state: Literal[False] = False, + raise_on_failure: bool = True, + ) -> dict[str, Any]: + ... + + @overload + def run( + self, + initial_inputs: dict[str, Any], + state: GraphRunState | dict[str, Any] | None = None, + return_state: Literal[True] = ..., + raise_on_failure: bool = True, + ) -> GraphRunState: + ... + + @overload + def run( + self, + initial_inputs: dict[str, Any], + state: GraphRunState | dict[str, Any] | None = None, + return_state: bool = False, + raise_on_failure: bool = True, + ) -> dict[str, Any] | GraphRunState: + ... + + def run( + self, + initial_inputs: dict[str, Any], + state: GraphRunState | dict[str, Any] | None = None, + return_state: bool = False, + raise_on_failure: bool = True, + ) -> dict[str, Any] | GraphRunState: + """Execute the graph with retries, error edges, and state tracking.""" + if isinstance(state, dict): + run_state = GraphRunState.from_dict(state) + elif isinstance(state, GraphRunState): + run_state = state + else: + run_state = GraphRunState( + initial_inputs=dict(initial_inputs), + current_state=dict(initial_inputs), + ) + + current_state = run_state.current_state + node_outputs = run_state.node_outputs + node_statuses = run_state.node_statuses + node_errors = run_state.node_errors for node_id in self._execution_order: node = self.nodes[node_id] - # Collect inputs for this node from incoming edges or current state + # If node already succeeded in previous checkpoint, reuse output and skip + if node_statuses.get(node_id) == "succeeded" and node_id in node_outputs: + node.status = "succeeded" + node.attempts = 1 + continue + + # Check if previous normal dependencies failed + incoming_standard_edges = [ + e for e in self.edges if e.target_node == node_id and not e.is_error_edge + ] + incoming_error_edges = [ + e for e in self.edges if e.target_node == node_id and e.is_error_edge + ] + + # Determine if this node is triggered via an error edge or standard execution + is_error_target = False + error_payload: dict[str, Any] = {} + for e in incoming_error_edges: + if node_statuses.get(e.source_node) == "failed": + is_error_target = True + error_payload = { + "error": node_errors.get(e.source_node, "Unknown error"), + "failed_node": e.source_node, + } + break + + # Collect inputs for this node node_inputs: dict[str, Any] = dict(current_state) - incoming_edges = [e for e in self.edges if e.target_node == node_id] - for edge in incoming_edges: - prev_outputs = node_outputs.get(edge.source_node, {}) - if edge.mapping: - for src_k, tgt_k in edge.mapping.items(): - if src_k in prev_outputs: - node_inputs[tgt_k] = prev_outputs[src_k] - else: - if edge.source_port in prev_outputs: - node_inputs[edge.target_port] = prev_outputs[edge.source_port] + if is_error_target: + node_inputs.update(error_payload) + else: + for edge in incoming_standard_edges: + prev_outputs = node_outputs.get(edge.source_node, {}) + if edge.mapping: + for src_k, tgt_k in edge.mapping.items(): + if src_k in prev_outputs: + node_inputs[tgt_k] = prev_outputs[src_k] else: - node_inputs.update(prev_outputs) + if edge.source_port in prev_outputs: + node_inputs[edge.target_port] = prev_outputs[edge.source_port] + else: + node_inputs.update(prev_outputs) + + # Execute node with retry logic + outputs, exc = self._execute_node_with_retry(node, node_inputs) + node_statuses[node_id] = node.status + + if exc is not None: + node_errors[node_id] = str(exc) + run_state.failed_node = node_id + + # Check if there is an outgoing error edge to recover from this failure + outgoing_error_edges = [ + e for e in self.edges if e.source_node == node_id and e.is_error_edge + ] + if not outgoing_error_edges: + run_state.completed = False + self.last_state = run_state + if raise_on_failure: + raise exc + return run_state if return_state else current_state + else: + if outputs is not None: + node_outputs[node_id] = outputs + current_state.update(outputs) - outputs = node.execute(node_inputs) - node_outputs[node_id] = outputs - current_state.update(outputs) + run_state.completed = all(s == "succeeded" for s in node_statuses.values()) + run_state.failed_node = None if run_state.completed else run_state.failed_node + self.last_state = run_state + if return_state: + return run_state return current_state + + def resume( + self, + state: GraphRunState | dict[str, Any], + override_inputs: dict[str, Any] | None = None, + raise_on_failure: bool = True, + ) -> GraphRunState: + """Resume a partially-completed or failed graph run from a checkpoint.""" + if isinstance(state, dict): + run_state = GraphRunState.from_dict(state) + else: + run_state = state + + if override_inputs: + run_state.current_state.update(override_inputs) + + # Clear failure marker on resumed run + run_state.failed_node = None + result = self.run( + initial_inputs=run_state.initial_inputs, + state=run_state, + return_state=True, + raise_on_failure=raise_on_failure, + ) + assert isinstance(result, GraphRunState) + return result diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 4936dea..f33aa3a 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -9,11 +9,13 @@ from dynavec import ( ApprovalNode, CodeNode, + GraphRunState, LLMNode, LoopNode, NodePort, NodeSchema, RetrieverNode, + RetryPolicy, RouterNode, SimpleGraph, SubGraphNode, @@ -21,7 +23,7 @@ WebhookNode, node_registry, ) -from dynavec.chat.base import ChatModel, ChatResult, Message +from dynavec.chat.base import ChatChunk, ChatModel, ChatResult, Message class MockChatModel(ChatModel): @@ -47,15 +49,13 @@ def stream( messages: list[Message], tools: list[Any] | None = None, **kwargs: Any, - ): + ) -> Any: last = messages[-1].content - from dynavec.chat.base import ChatChunk - yield ChatChunk(content=f"Processed: {last}") class MockDoc: - def __init__(self, doc_id: str, text: str, score: float = 0.9): + def __init__(self, doc_id: str, text: str, score: float = 0.9) -> None: self.id = doc_id self.text = text self.score = score @@ -67,7 +67,7 @@ def __init__(self, doc_id: str, text: str, score: float = 0.9): # ===================================================================== -def test_registry_registration_and_palette(): +def test_registry_registration_and_palette() -> None: expected_types = { "approval", "code", @@ -95,34 +95,40 @@ def test_registry_registration_and_palette(): assert "outputPorts" in item -def test_port_and_schema_serialization(): +def test_port_and_schema_serialization() -> None: port = NodePort( name="test_in", port_type="string", - description="test input port", + description="A test port", required=True, - default="abc", + default="val", ) - port_dict = port.to_dict() - assert port_dict["name"] == "test_in" - assert port_dict["port_type"] == "string" - assert port_dict["required"] is True - assert port_dict["default"] == "abc" + p_dict = port.to_dict() + assert p_dict["name"] == "test_in" + assert p_dict["port_type"] == "string" + assert p_dict["required"] is True + assert p_dict["default"] == "val" schema = NodeSchema( node_type="custom_test", label="Custom Test", description="A test node schema", + category="testing", inputs=[port], - outputs=[], + outputs=[NodePort(name="test_out", port_type="number")], + config_schema={"timeout": {"type": "integer", "default": 30}}, ) + s_dict = schema.to_dict() assert s_dict["node_type"] == "custom_test" assert len(s_dict["inputs"]) == 1 + assert len(s_dict["outputs"]) == 1 rf_dict = schema.to_reactflow_dict() assert rf_dict["type"] == "custom_test" assert rf_dict["inputPorts"] == ["test_in"] + assert rf_dict["outputPorts"] == ["test_out"] + assert rf_dict["config"]["timeout"]["default"] == 30 # ===================================================================== @@ -130,8 +136,10 @@ def test_port_and_schema_serialization(): # ===================================================================== -def test_retriever_node(): - def dummy_retriever(query: str, top_k: int = 5, filter: dict[str, Any] | None = None): +def test_retriever_node() -> None: + def dummy_retriever( + query: str, top_k: int = 5, filter: dict[str, Any] | None = None + ) -> list[MockDoc]: return [ MockDoc("d1", "Dynavec is a serverless vector database."), MockDoc("d2", "DynamoDB provides scalable key-value storage."), @@ -151,7 +159,7 @@ def dummy_retriever(query: str, top_k: int = 5, filter: dict[str, Any] | None = assert "DynamoDB provides" in outputs["context"] -def test_llm_node_chat_model(): +def test_llm_node_chat_model() -> None: model = MockChatModel() node = LLMNode(model=model, config={"system_prompt": "You are a helpful assistant."}) @@ -169,7 +177,7 @@ def test_llm_node_chat_model(): assert outputs["message"].role == "assistant" -def test_llm_node_callable(): +def test_llm_node_callable() -> None: def simple_callable(messages: list[Message]) -> str: return f"Echo: {messages[-1].content}" @@ -178,7 +186,7 @@ def simple_callable(messages: list[Message]) -> str: assert outputs["response"] == "Echo: Hello world" -def test_tool_node_success(): +def test_tool_node_success() -> None: def add(a: int, b: int) -> int: return a + b @@ -188,7 +196,7 @@ def add(a: int, b: int) -> int: assert outputs["result"] == 35 -def test_tool_node_json_string_and_error(): +def test_tool_node_json_string_and_error() -> None: def multiply(x: int, y: int) -> int: return x * y @@ -204,7 +212,7 @@ def multiply(x: int, y: int) -> int: assert "division by zero" in err_out["result"] -def test_router_node(): +def test_router_node() -> None: node = RouterNode(route_fn=lambda x: "admin_flow" if x.get("is_admin") else "user_flow") out1 = node.execute({"data": {"is_admin": True}}) @@ -219,7 +227,7 @@ def test_router_node(): assert out3["selected_route"] == "manual_override" -def test_code_node(): +def test_code_node() -> None: def custom_transform(inputs: dict[str, Any]) -> dict[str, Any]: return {"uppercased": inputs["text"].upper(), "len": len(inputs["text"])} @@ -229,12 +237,12 @@ def custom_transform(inputs: dict[str, Any]) -> dict[str, Any]: assert res["len"] == 8 # When code returns non-dict - scalar_node = CodeNode(code_fn=lambda inp: len(inp["text"])) + scalar_node = CodeNode(code_fn=lambda inp: {"outputs": len(inp["text"])}) res2 = scalar_node.execute({"inputs": {"text": "test"}}) assert res2["outputs"] == 4 -def test_approval_node(): +def test_approval_node() -> None: node = ApprovalNode() out = node.execute( { @@ -256,21 +264,21 @@ def test_approval_node(): assert out_rejected["approved"] is False -def test_loop_node(): +def test_loop_node() -> None: node = LoopNode(map_fn=lambda x: x * 2, config={"max_iterations": 3}) out = node.execute({"items": [1, 2, 3, 4, 5]}) assert out["items"] == [2, 4, 6] # clipped to max 3 assert out["count"] == 3 -def test_webhook_node(): +def test_webhook_node() -> None: node = WebhookNode() payload = {"event": "user.signup", "user_id": "usr_123"} out = node.execute({"payload": payload, "headers": {"Authorization": "Bearer tok"}}) assert out["payload"] == payload -def test_subgraph_node(): +def test_subgraph_node() -> None: def nested_workflow(inputs: dict[str, Any]) -> dict[str, Any]: return {"processed_by_subgraph": True, "value": inputs.get("val", 0) + 10} @@ -285,9 +293,11 @@ def nested_workflow(inputs: dict[str, Any]) -> dict[str, Any]: # ===================================================================== -def test_simple_graph_pipeline_retriever_llm_tool(): +def test_simple_graph_pipeline_retriever_llm_tool() -> None: # 1. RetrieverNode: searches knowledge base - def mock_retriever(query: str, top_k: int = 5, filter: dict[str, Any] | None = None): + def mock_retriever( + query: str, top_k: int = 5, filter: dict[str, Any] | None = None + ) -> list[MockDoc]: return [MockDoc("doc-1", "Dynavec provides DynamoDB vectors.")] retriever = RetrieverNode(retriever_fn=mock_retriever, node_id="retriever_1") @@ -337,7 +347,7 @@ def mock_tool(operation: str, value: int) -> int: assert result["result"] == 84 -def test_simple_graph_validation_errors(): +def test_simple_graph_validation_errors() -> None: graph = SimpleGraph() with pytest.raises(ValueError, match="Source node 'nonexistent' not in graph"): graph.add_edge("nonexistent", "target") @@ -346,3 +356,208 @@ def test_simple_graph_validation_errors(): graph.add_node(node) with pytest.raises(ValueError, match="Target node 'missing_target' not in graph"): graph.add_edge("valid_node", "missing_target") + + +# ===================================================================== +# Engine Resilience, Retries & Resumability Tests (Issue #279) +# ===================================================================== + + +def test_retry_policy_logic_and_serialization() -> None: + policy = RetryPolicy( + max_attempts=4, + base_delay=0.01, + max_delay=0.1, + jitter=False, + retry_on=(ValueError, KeyError), + ) + assert policy.should_retry(ValueError("bad value")) + assert policy.should_retry(KeyError("missing key")) + assert not policy.should_retry(TypeError("wrong type")) + + # Serialization + d = policy.to_dict() + assert d["max_attempts"] == 4 + assert d["base_delay"] == 0.01 + assert not d["jitter"] + + reconstructed = RetryPolicy.from_dict(d) + assert reconstructed.max_attempts == 4 + assert reconstructed.base_delay == 0.01 + + +def test_node_retry_transient_failure_success() -> None: + call_count = 0 + + def flaky_fn(inputs: dict[str, Any]) -> dict[str, Any]: + nonlocal call_count + call_count += 1 + if call_count < 3: + raise ConnectionError("Temporary network glitch") + return {"result": inputs["x"] * 10} + + policy = RetryPolicy(max_attempts=4, base_delay=0.001, max_delay=0.01) + flaky_node = CodeNode( + code_fn=flaky_fn, node_id="flaky_node", retry_policy=policy + ) + + graph = SimpleGraph() + graph.add_node(flaky_node) + + res = graph.run({"x": 5}) + assert res["result"] == 50 + assert flaky_node.attempts == 3 + assert flaky_node.status == "succeeded" + assert flaky_node.last_error is None + + +def test_node_retry_exhaustion_failure() -> None: + attempts_recorded = 0 + + def always_failing_fn(inputs: dict[str, Any]) -> dict[str, Any]: + nonlocal attempts_recorded + attempts_recorded += 1 + raise RuntimeError("Hard database failure") + + policy = RetryPolicy(max_attempts=3, base_delay=0.001, max_delay=0.01) + failing_node = CodeNode( + code_fn=always_failing_fn, node_id="fail_node", retry_policy=policy + ) + + graph = SimpleGraph() + graph.add_node(failing_node) + + with pytest.raises(RuntimeError, match="Hard database failure"): + graph.run({"x": 1}, raise_on_failure=True) + + assert failing_node.attempts == 3 + assert failing_node.status == "failed" + assert "Hard database failure" in (failing_node.last_error or "") + + +def test_error_edge_fallback_routing() -> None: + def failing_primary_fn(inputs: dict[str, Any]) -> dict[str, Any]: + raise ValueError("Primary service unavailable") + + def fallback_recovery_fn(inputs: dict[str, Any]) -> dict[str, Any]: + error_msg = inputs.get("error", "") + return { + "recovered": True, + "handled_error": error_msg, + "fallback_response": "Served from cached backup", + } + + primary = CodeNode( + code_fn=failing_primary_fn, + node_id="primary_node", + retry_policy=RetryPolicy(max_attempts=2, base_delay=0.001), + ) + fallback = CodeNode(code_fn=fallback_recovery_fn, node_id="fallback_node") + + graph = SimpleGraph() + graph.add_node(primary) + graph.add_node(fallback) + + # Wire error edge + graph.add_error_edge(source_id="primary_node", target_id="fallback_node") + + res = graph.run({"query": "fetch stats"}, raise_on_failure=False) + assert res["recovered"] is True + assert "Primary service unavailable" in res["handled_error"] + assert res["fallback_response"] == "Served from cached backup" + assert primary.status == "failed" + assert fallback.status == "succeeded" + + +def test_graph_run_state_checkpoint_serialization() -> None: + state = GraphRunState( + run_id="run_test_123", + initial_inputs={"a": 1}, + current_state={"a": 1, "b": 2}, + node_outputs={"node_1": {"b": 2}}, + node_statuses={"node_1": "succeeded", "node_2": "pending"}, + node_errors={"node_3": "some error"}, + completed=False, + failed_node="node_3", + ) + + serialized = state.to_dict() + assert serialized["run_id"] == "run_test_123" + assert serialized["completed"] is False + assert serialized["node_statuses"]["node_1"] == "succeeded" + + restored = GraphRunState.from_dict(serialized) + assert restored.run_id == "run_test_123" + assert restored.current_state == {"a": 1, "b": 2} + assert restored.node_outputs == {"node_1": {"b": 2}} + assert restored.node_statuses == {"node_1": "succeeded", "node_2": "pending"} + + +def test_graph_resumability_from_checkpoint() -> None: + # 3-node linear pipeline: Node A -> Node B (initially fails) -> Node C + node_a_runs = 0 + node_b_runs = 0 + node_c_runs = 0 + + def fn_a(inputs: dict[str, Any]) -> dict[str, Any]: + nonlocal node_a_runs + node_a_runs += 1 + return {"step_a": "done", "value": inputs["start"] + 10} + + def fn_b_fail(inputs: dict[str, Any]) -> dict[str, Any]: + nonlocal node_b_runs + node_b_runs += 1 + raise ValueError("Simulated outage in Node B") + + def fn_b_fixed(inputs: dict[str, Any]) -> dict[str, Any]: + nonlocal node_b_runs + node_b_runs += 1 + return {"step_b": "done", "value": inputs["value"] * 2} + + def fn_c(inputs: dict[str, Any]) -> dict[str, Any]: + nonlocal node_c_runs + node_c_runs += 1 + return {"step_c": "done", "final_result": inputs["value"] + 5} + + node_a = CodeNode(code_fn=fn_a, node_id="node_a") + node_b = CodeNode(code_fn=fn_b_fail, node_id="node_b") + node_c = CodeNode(code_fn=fn_c, node_id="node_c") + + graph = SimpleGraph() + graph.add_node(node_a) + graph.add_node(node_b) + graph.add_node(node_c) + + graph.add_edge("node_a", "node_b") + graph.add_edge("node_b", "node_c") + + # Run 1: Fails at Node B + state_result = graph.run({"start": 5}, return_state=True, raise_on_failure=False) + assert isinstance(state_result, GraphRunState) + assert state_result.completed is False + assert state_result.failed_node == "node_b" + assert node_a_runs == 1 + assert node_b_runs == 1 + assert node_c_runs == 0 + + # Save checkpoint + checkpoint = state_result.to_dict() + + # Fix Node B + node_b.execute = fn_b_fixed # type: ignore[method-assign] + + # Resume from checkpoint + resumed_state = graph.resume(checkpoint) + + # Verify Node A was NOT re-executed + assert node_a_runs == 1 # Crucial: stayed 1, skipped duplicate run! + assert node_b_runs == 2 # executed upon resume + assert node_c_runs == 1 # executed to completion + + assert resumed_state.completed is True + assert resumed_state.failed_node is None + assert resumed_state.current_state["step_a"] == "done" + assert resumed_state.current_state["step_b"] == "done" + assert resumed_state.current_state["step_c"] == "done" + # start (5) + 10 = 15; 15 * 2 = 30; 30 + 5 = 35 + assert resumed_state.current_state["final_result"] == 35