From d76b2d9cfa79e00b5ad9dd91f46a50e26716aaa5 Mon Sep 17 00:00:00 2001 From: Matt McCarthy Date: Tue, 28 Jul 2026 16:59:12 -0500 Subject: [PATCH 1/2] fix: use topological order and scoped context for agent graph traversal --- .../src/ldai/agent_graph/__init__.py | 260 ++++++++++-------- .../sdk/server-ai/tests/test_agent_graph.py | 256 +++++++++++++++-- 2 files changed, 390 insertions(+), 126 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py index 5980952a..4a0296ae 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -121,28 +121,29 @@ def get_parent_nodes(self, node_key: str) -> List[AgentGraphNode]: nodes.append(node) return nodes - def _collect_nodes( - self, - node: AgentGraphNode, - node_depths: Dict[str, int], - nodes_by_depth: Dict[int, List[AgentGraphNode]], - visited: Set[str], - max_depth: int, - ) -> None: - """Collect all reachable nodes from the given node and group them by depth.""" - node_key = node.get_key() - if node_key in visited: - return - visited.add(node_key) - - # Use max_depth for nodes not in node_depths to ensure they execute last - node_depth = node_depths.get(node_key, max_depth) - if node_depth not in nodes_by_depth: - nodes_by_depth[node_depth] = [] - nodes_by_depth[node_depth].append(node) - - for child in self.get_child_nodes(node_key): - self._collect_nodes(child, node_depths, nodes_by_depth, visited, max_depth) + def _reachable_and_discovery( + self, root_key: str + ) -> tuple[Set[str], List[str]]: + """Return reachable node keys and BFS discovery order from root.""" + reachable: Set[str] = set() + order: List[str] = [] + queue: List[str] = [root_key] + reachable.add(root_key) + order.append(root_key) + while queue: + key = queue.pop(0) + node = self.get_node(key) + if node is None: + continue + for edge in node.get_edges(): + if ( + self.get_node(edge.target_config) is not None + and edge.target_config not in reachable + ): + reachable.add(edge.target_config) + order.append(edge.target_config) + queue.append(edge.target_config) + return reachable, order def terminal_nodes(self) -> List[AgentGraphNode]: """Get the terminal nodes of the graph, meaning any nodes without children.""" @@ -161,8 +162,18 @@ def traverse( fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], execution_context: Optional[Dict[str, Any]] = None, ) -> Any: - """Traverse from the root down to terminal nodes, visiting nodes in order of depth. - Nodes with the longest paths from the root (deepest nodes) will always be visited last.""" + """Visit each reachable node in topological (predecessors-first) order. + + The root is visited first. A node is visited only after all of its + reachable predecessors. When multiple nodes are ready, discovery order + (BFS from root following declared edge order) breaks ties. Cycles are + broken at the lowest remaining in-degree, also tie-broken by discovery + order. + + Each node's ``fn`` receives a fresh context containing the caller-provided + initial context plus the returned results of exactly that node's reachable + predecessors — not unrelated parallel-branch nodes. + """ if execution_context is None: execution_context = {} @@ -170,104 +181,135 @@ def traverse( if root_node is None: return - node_depths: Dict[str, int] = {root_node.get_key(): 0} - current_level: List[AgentGraphNode] = [root_node] - depth = 0 - max_depth_limit = 10 # Infinite loop protection limit - max_depth_encountered = 0 - seen_nodes: Set[str] = {root_node.get_key()} - - while current_level: - next_level: List[AgentGraphNode] = [] - depth += 1 - - for node in current_level: - node_key = node.get_key() - for child in self.get_child_nodes(node_key): - child_key = child.get_key() - if depth <= max_depth_limit: - # Defer this child to the next level if it's at a longer path - if child_key not in node_depths or depth > node_depths[child_key]: - node_depths[child_key] = depth - max_depth_encountered = max(max_depth_encountered, depth) - # Add to next level if not already visited (prevents cycles) - if child_key not in seen_nodes: - seen_nodes.add(child_key) - next_level.append(child) - else: - max_depth_encountered = max(max_depth_encountered, depth) - if child_key not in seen_nodes: - # Push this to the next level to be visited - seen_nodes.add(child_key) - next_level.append(child) - - current_level = next_level - - # Use max_depth_limit + 1 to ensure they execute after all recorded nodes - max_depth = max(max_depth_limit + 1, max_depth_encountered + 1) - - # Group all nodes by depth - nodes_by_depth: Dict[int, List[AgentGraphNode]] = {} - # New visited for children nodes - visited: Set[str] = set() + reachable, order = self._reachable_and_discovery(root_node.get_key()) - self._collect_nodes(root_node, node_depths, nodes_by_depth, visited, max_depth) - # Execute the lambda at this level for the nodes at this depth - for depth_level in sorted(nodes_by_depth.keys()): - for node in nodes_by_depth[depth_level]: - execution_context[node.get_key()] = fn(node, execution_context) + indeg = {k: 0 for k in reachable} + for k in reachable: + node = self.get_node(k) + if node is None: + continue + for edge in node.get_edges(): + if edge.target_config in reachable: + indeg[edge.target_config] += 1 + indeg[root_node.get_key()] = 0 # force root - return execution_context[self._agent_graph.root_config_key] + visited: Set[str] = set() + results: Dict[str, Any] = {} + ancestors: Dict[str, Set[str]] = {} + + def scoped(deps: Set[str]) -> Dict[str, Any]: + c = dict(execution_context) + for dep_key in deps: + c[dep_key] = results[dep_key] + return c + + while len(visited) < len(reachable): + nxt = next( + (k for k in order if k not in visited and indeg[k] <= 0), None + ) + if nxt is None: # cycle break + nxt = min( + (k for k in order if k not in visited), key=lambda k: indeg[k] + ) + visited.add(nxt) + anc: Set[str] = set() + for parent in self.get_parent_nodes(nxt): + pk = parent.get_key() + if pk not in visited: + continue + anc.add(pk) + anc |= ancestors.get(pk, set()) + ancestors[nxt] = anc + node = self.get_node(nxt) + assert node is not None + results[nxt] = fn(node, scoped(anc)) + for edge in node.get_edges(): + if edge.target_config in reachable: + indeg[edge.target_config] -= 1 + return results.get(self._agent_graph.root_config_key) def reverse_traverse( self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], execution_context: Optional[Dict[str, Any]] = None, ) -> Any: - """Traverse from terminal nodes up to the root, visiting nodes level by level. - The root node will always be visited last, even if multiple paths converge at it.""" + """Visit each reachable node in reverse topological (descendants-first) order. + + The root is visited last. A node is visited only after all of its + reachable descendants. When multiple nodes are ready, discovery order + (BFS from root following declared edge order) breaks ties. Cycles among + non-root nodes are broken at the lowest remaining out-degree, also + tie-broken by discovery order. Graphs with no terminals (pure cycles) + still visit every reachable node, root last. + + Each node's ``fn`` receives a fresh context containing the caller-provided + initial context plus the returned results of exactly that node's reachable + descendants — not unrelated parallel-branch nodes. + """ if execution_context is None: execution_context = {} - terminal_nodes = self.terminal_nodes() - if not terminal_nodes: + root_node = self.root() + if root_node is None: return - visited: Set[str] = set() - current_level: List[AgentGraphNode] = terminal_nodes root_key = self._agent_graph.root_config_key - root_node_seen = False - - while current_level: - next_level: List[AgentGraphNode] = [] - - for node in current_level: - node_key = node.get_key() - if node_key in visited: - continue - - visited.add(node_key) - # Skip the root node if we reach a terminus, it will be visited last - if node_key == root_key: - root_node_seen = True - continue - - execution_context[node_key] = fn(node, execution_context) - - for parent in self.get_parent_nodes(node_key): - parent_key = parent.get_key() - if parent_key not in visited: - next_level.append(parent) - - current_level = next_level + reachable, order = self._reachable_and_discovery(root_key) + + outdeg: Dict[str, int] = {} + for k in reachable: + node = self.get_node(k) + assert node is not None + outdeg[k] = sum( + 1 for e in node.get_edges() if e.target_config in reachable + ) - # If we saw the root node, append it at the end as it'll always be the last node in a - # reverse traversal (this should always happen, non-contiguous graphs are invalid) - if root_node_seen: - root_node = self.root() - if root_node is not None: - execution_context[root_node.get_key()] = fn( - root_node, execution_context + visited: Set[str] = set() + results: Dict[str, Any] = {} + descendants: Dict[str, Set[str]] = {} + + def scoped(deps: Set[str]) -> Dict[str, Any]: + c = dict(execution_context) + for dep_key in deps: + c[dep_key] = results[dep_key] + return c + + def non_root_remaining() -> bool: + return any(k != root_key and k not in visited for k in reachable) + + while non_root_remaining(): + nxt = next( + ( + k + for k in order + if k != root_key and k not in visited and outdeg[k] <= 0 + ), + None, + ) + if nxt is None: # cycle break + nxt = min( + (k for k in order if k != root_key and k not in visited), + key=lambda k: outdeg[k], ) - - return execution_context[self._agent_graph.root_config_key] + visited.add(nxt) + desc: Set[str] = set() + node = self.get_node(nxt) + assert node is not None + for edge in node.get_edges(): + ck = edge.target_config + if ck not in reachable or ck not in visited: + continue + desc.add(ck) + desc |= descendants.get(ck, set()) + descendants[nxt] = desc + results[nxt] = fn(node, scoped(desc)) + for parent in self.get_parent_nodes(nxt): + pk = parent.get_key() + if pk != root_key and pk in reachable: + outdeg[pk] -= 1 + + # root last, once; root depends on every reachable non-root node + root_deps = {k for k in reachable if k != root_key} + visited.add(root_key) + results[root_key] = fn(root_node, scoped(root_deps)) + return results.get(root_key) diff --git a/packages/sdk/server-ai/tests/test_agent_graph.py b/packages/sdk/server-ai/tests/test_agent_graph.py index 0f2a7ee6..e007c2d0 100644 --- a/packages/sdk/server-ai/tests/test_agent_graph.py +++ b/packages/sdk/server-ai/tests/test_agent_graph.py @@ -353,21 +353,14 @@ def test_agent_graph_traverse(ldai_client: LDAIClient): "test-agent-graph-depth-3", Context.create("user-key") ) - context = {} order = [] def handle_traverse(node, context): - # Asserting that returned values are included in the context - for previousKey in order: - assert previousKey in context - assert context[previousKey] == previousKey + "-test" order.append(node.get_key()) return node.get_key() + "-test" - graph.traverse(handle_traverse, context) - # Asserting that we traverse in the expected order - # This config specifically has nodes connecting from depth 2->3 and root->3 to ensure the root node is visited first - # and minimal-agent is visited last + graph.traverse(handle_traverse, {}) + # Nested diamond: root first, convergence node (minimal-agent) last assert order == [ "customer-support-agent", "personalized-agent", @@ -381,20 +374,14 @@ def test_agent_graph_reverse_traverse(ldai_client: LDAIClient): "test-agent-graph-depth-3", Context.create("user-key") ) - context = {} order = [] def handle_reverse_traverse(node, context): - # Asserting that returned values are included in the context - for previousKey in order: - assert previousKey in context - assert context[previousKey] == previousKey + "-test" order.append(node.get_key()) return node.get_key() + "-test" - graph.reverse_traverse(handle_reverse_traverse, context) - # Asserting that we traverse in the expected order - # This config specifically has nodes connecting from depth 2->3 and root->3 to ensure the root node is visited last + graph.reverse_traverse(handle_reverse_traverse, {}) + # Nested diamond: terminals first, root last assert order == [ "minimal-agent", "multi-context-agent", @@ -445,3 +432,238 @@ def handle_traverse(node, context): return None graph.traverse(handle_traverse, context) + + +# --------------------------------------------------------------------------- +# AIGRAPH parity fixtures (G1–G6 / G2b) — topological order + scoped context +# --------------------------------------------------------------------------- + +_TRAVERSAL_VECTORS = [ + { + "id": "G1", + "root": "a", + "nodes": ["a", "b", "c"], + "edges": [["a", "b"], ["b", "c"]], + "traverse": ["a", "b", "c"], + "reverse_traverse": ["c", "b", "a"], + "traverse_context": {"a": [], "b": ["a"], "c": ["a", "b"]}, + "reverse_traverse_context": {"a": ["b", "c"], "b": ["c"], "c": []}, + }, + { + "id": "G2", + "root": "a", + "nodes": ["a", "b", "c", "d", "e"], + "edges": [["a", "b"], ["a", "c"], ["c", "d"], ["d", "e"], ["b", "e"]], + "traverse": ["a", "b", "c", "d", "e"], + "reverse_traverse": ["e", "b", "d", "c", "a"], + "traverse_context": { + "a": [], + "b": ["a"], + "c": ["a"], + "d": ["a", "c"], + "e": ["a", "b", "c", "d"], + }, + "reverse_traverse_context": { + "a": ["b", "c", "d", "e"], + "b": ["e"], + "c": ["d", "e"], + "d": ["e"], + "e": [], + }, + }, + { + "id": "G2b", + "root": "a", + "nodes": ["a", "b", "c", "d", "e"], + "edges": [["a", "c"], ["a", "b"], ["c", "d"], ["d", "e"], ["b", "e"]], + "traverse": ["a", "c", "b", "d", "e"], + "reverse_traverse": ["e", "b", "d", "c", "a"], + "traverse_context": { + "a": [], + "b": ["a"], + "c": ["a"], + "d": ["a", "c"], + "e": ["a", "b", "c", "d"], + }, + "reverse_traverse_context": { + "a": ["b", "c", "d", "e"], + "b": ["e"], + "c": ["d", "e"], + "d": ["e"], + "e": [], + }, + }, + { + "id": "G3", + "root": "a", + "nodes": ["a", "b", "c", "d"], + "edges": [["a", "b"], ["a", "c"], ["b", "d"], ["c", "d"]], + "traverse": ["a", "b", "c", "d"], + "reverse_traverse": ["d", "b", "c", "a"], + "traverse_context": { + "a": [], + "b": ["a"], + "c": ["a"], + "d": ["a", "b", "c"], + }, + "reverse_traverse_context": { + "a": ["b", "c", "d"], + "b": ["d"], + "c": ["d"], + "d": [], + }, + }, + { + "id": "G4", + "root": "a", + "nodes": ["a", "n", "m", "t"], + "edges": [["a", "n"], ["n", "m"], ["n", "t"], ["m", "t"]], + "traverse": ["a", "n", "m", "t"], + "reverse_traverse": ["t", "m", "n", "a"], + "traverse_context": { + "a": [], + "n": ["a"], + "m": ["a", "n"], + "t": ["a", "m", "n"], + }, + "reverse_traverse_context": { + "a": ["m", "n", "t"], + "n": ["m", "t"], + "m": ["t"], + "t": [], + }, + }, + { + "id": "G5", + "root": "a", + "nodes": ["a", "b", "c", "d"], + "edges": [["a", "b"], ["a", "c"], ["b", "d"]], + "traverse": ["a", "b", "c", "d"], + "reverse_traverse": ["c", "d", "b", "a"], + "traverse_context": { + "a": [], + "b": ["a"], + "c": ["a"], + "d": ["a", "b"], + }, + "reverse_traverse_context": { + "a": ["b", "c", "d"], + "b": ["d"], + "c": [], + "d": [], + }, + }, + { + "id": "G6", + "root": "a", + "nodes": ["a", "b", "c"], + "edges": [["a", "b"], ["b", "c"], ["c", "b"]], + "traverse": ["a", "b", "c"], + "reverse_traverse": ["b", "c", "a"], + "traverse_context": {"a": [], "b": ["a"], "c": ["a", "b"]}, + "reverse_traverse_context": {"a": ["b", "c"], "b": [], "c": ["b"]}, + }, +] + + +def _make_graph(root: str, nodes: list, edges: list) -> AgentGraphDefinition: + edge_objs = [ + Edge( + key=f"{src}-{tgt}", + source_config=src, + target_config=tgt, + handoff={}, + ) + for src, tgt in edges + ] + graph_config = AIAgentGraphConfig( + key="parity-graph", + root_config_key=root, + edges=edge_objs, + ) + graph_nodes = { + key: AIAgentConfig( + key=key, + enabled=True, + create_tracker=MagicMock(), + evaluator=Evaluator.noop(), + ) + for key in nodes + } + built = AgentGraphDefinition.build_nodes(graph_config, graph_nodes) + return AgentGraphDefinition( + agent_graph=graph_config, + nodes=built, + context=Context.create("parity-user"), + enabled=True, + create_tracker=MagicMock(), + ) + + +@pytest.mark.parametrize( + "vector", _TRAVERSAL_VECTORS, ids=[v["id"] for v in _TRAVERSAL_VECTORS] +) +def test_traversal_parity_order_and_context(vector): + graph = _make_graph(vector["root"], vector["nodes"], vector["edges"]) + node_keys = set(vector["nodes"]) + initial = {"seed": "shared"} + + # Forward traverse + fwd_order = [] + fwd_contexts = {} + + def fwd_fn(node, ctx): + key = node.get_key() + fwd_order.append(key) + fwd_contexts[key] = dict(ctx) + assert ctx.get("seed") == "shared" + return f"{key}-result" + + graph.traverse(fwd_fn, dict(initial)) + assert fwd_order == vector["traverse"] + for key, expected_deps in vector["traverse_context"].items(): + observed = {k for k in fwd_contexts[key] if k in node_keys} + assert observed == set(expected_deps), ( + f"{vector['id']} forward {key}: expected {set(expected_deps)}, " + f"got {observed}" + ) + + # Reverse traverse + rev_order = [] + rev_contexts = {} + + def rev_fn(node, ctx): + key = node.get_key() + rev_order.append(key) + rev_contexts[key] = dict(ctx) + assert ctx.get("seed") == "shared" + return f"{key}-result" + + graph.reverse_traverse(rev_fn, dict(initial)) + assert rev_order == vector["reverse_traverse"] + for key, expected_deps in vector["reverse_traverse_context"].items(): + observed = {k for k in rev_contexts[key] if k in node_keys} + assert observed == set(expected_deps), ( + f"{vector['id']} reverse {key}: expected {set(expected_deps)}, " + f"got {observed}" + ) + + +@pytest.mark.parametrize( + "vector", _TRAVERSAL_VECTORS, ids=[v["id"] for v in _TRAVERSAL_VECTORS] +) +def test_traversal_parity_determinism(vector): + graph = _make_graph(vector["root"], vector["nodes"], vector["edges"]) + + def collect(method): + order = [] + + def fn(node, _ctx): + order.append(node.get_key()) + return node.get_key() + + method(fn, {}) + return order + + assert collect(graph.traverse) == collect(graph.traverse) + assert collect(graph.reverse_traverse) == collect(graph.reverse_traverse) From 3f801da138a8c106832407f9275e72c55e5a3df7 Mon Sep 17 00:00:00 2001 From: Matt McCarthy Date: Wed, 29 Jul 2026 09:43:01 -0500 Subject: [PATCH 2/2] fix: use == 0 for agent graph ready-check parity --- packages/sdk/server-ai/src/ldai/agent_graph/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py index 4a0296ae..967e7689 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -205,7 +205,7 @@ def scoped(deps: Set[str]) -> Dict[str, Any]: while len(visited) < len(reachable): nxt = next( - (k for k in order if k not in visited and indeg[k] <= 0), None + (k for k in order if k not in visited and indeg[k] == 0), None ) if nxt is None: # cycle break nxt = min( @@ -282,7 +282,7 @@ def non_root_remaining() -> bool: ( k for k in order - if k != root_key and k not in visited and outdeg[k] <= 0 + if k != root_key and k not in visited and outdeg[k] == 0 ), None, )