Merge pull request #236 from ai-agent-assembly/v0.1.0/AAASM-4472/fix/langgraph_subgraph_lineage

[AAASM-4472] 🐛 (langgraph): Fix subgraph-as-node lineage tracking gap

Author: Chisanan232

agent_assembly/adapters/langgraph/patch.py

@@ -360,6 +360,58 @@ def _wrap_subgraph_spawn_node(node_map: Any, node_name: Any, node_executor: Any,
     return True
 
 
+def _wrap_subgraph_invoke_methods_in_place(node_name: Any, subgraph: Any, process_agent_id: str | None) -> bool:
+    """Wrap a compiled subgraph's own invoke/ainvoke in place for lineage.
+
+    AAASM-4472: used when the subgraph is reached via an already-unwrapped
+    ``PregelNode.bound`` rather than arriving bare in the node map. Unlike
+    ``_wrap_subgraph_spawn_node`` (which replaces the whole node_map entry for a
+    bare subgraph), the outer ``PregelNode`` container here must stay intact —
+    its channels/triggers/mapper/writers are what the Pregel runtime uses to
+    build the executable task — so this mutates the bound Runnable's own
+    invoke/ainvoke attributes instead. The original methods are captured before
+    reassignment so the wrapper calls the real invoke/ainvoke rather than
+    recursing into itself once installed.
+    """
+    node_delegation_reason = f"langgraph_node:{node_name}"
+    wrapped_any = False
+
+    original_invoke = getattr(subgraph, "invoke", None)
+    if callable(original_invoke):
+
+        def sync_subgraph_wrapper(*args: Any, **kwargs: Any) -> Any:
+            spawn_ctx = SpawnContext(
+                parent_agent_id=process_agent_id or "",
+                depth=_current_spawn_depth(),
+                spawned_by_tool=None,
+                delegation_reason=node_delegation_reason,
+            )
+            with spawn_context_scope(spawn_ctx):
+                return original_invoke(*args, **kwargs)
+
+        subgraph.invoke = sync_subgraph_wrapper
+        wrapped_any = True
+
+    original_ainvoke = getattr(subgraph, "ainvoke", None)
+    if callable(original_ainvoke) and not getattr(subgraph, "_agent_assembly_ainvoke_spawned", False):
+
+        async def async_subgraph_wrapper(*args: Any, **kwargs: Any) -> Any:
+            spawn_ctx = SpawnContext(
+                parent_agent_id=process_agent_id or "",
+                depth=_current_spawn_depth(),
+                spawned_by_tool=None,
+                delegation_reason=node_delegation_reason,
+            )
+            with spawn_context_scope(spawn_ctx):
+                return await original_ainvoke(*args, **kwargs)
+
+        subgraph.ainvoke = async_subgraph_wrapper
+        subgraph._agent_assembly_ainvoke_spawned = True
+        wrapped_any = True
+
+    return wrapped_any
+
+
 def _wrap_node_invoke_methods(node_name: Any, node_executor: Any, callback_handler: Any) -> bool:
     """Wrap a node executor's invoke/ainvoke methods. Return True if either was wrapped."""
     wrapped_any = False
@@ -384,19 +436,28 @@ def _wrap_node_entry(
     if _is_tool_node(node_executor):
         return _wrap_tool_node_subgraphs(node_executor, process_agent_id)
 
+    # AAASM-4434/AAASM-4472: current langgraph node-map entries are PregelNode
+    # wrappers whose own .invoke/.ainvoke are never called by the Pregel runtime
+    # — it dispatches through PregelNode.bound. Unwrap *before* any other
+    # detection (including _is_compiled_subgraph) and dispatch against the
+    # unwrapped object: checking _is_compiled_subgraph() against the outer
+    # PregelNode itself always fails (only the wrapped object exposes `.nodes`),
+    # so subgraph detection must run post-unwrap to work whether a subgraph
+    # arrives bare or already PregelNode-wrapped.
+    if _is_pregel_node_wrapper(node_executor):
+        bound = node_executor.bound
+        if _is_compiled_subgraph(bound):
+            return _wrap_subgraph_invoke_methods_in_place(node_name, bound, process_agent_id)
+        return _wrap_node_invoke_methods(node_name, bound, callback_handler)
+
     if callable(node_executor):
         return _wrap_callable_node_executor(node_map, node_name, node_executor, callback_handler)
 
-    # Spawn point: node is itself a compiled subgraph — wrap for lineage.
+    # Spawn point: node is itself a compiled subgraph, not PregelNode-wrapped —
+    # wrap for lineage.
     if _is_compiled_subgraph(node_executor):
         return _wrap_subgraph_spawn_node(node_map, node_name, node_executor, process_agent_id)
 
-    # AAASM-4434: current langgraph node-map entries are PregelNode wrappers whose
-    # own .invoke/.ainvoke are never called by the Pregel runtime — it dispatches
-    # through PregelNode.bound. Wrap that inner Runnable instead, or hooks never fire.
-    if _is_pregel_node_wrapper(node_executor):
-        return _wrap_node_invoke_methods(node_name, node_executor.bound, callback_handler)
-
     return _wrap_node_invoke_methods(node_name, node_executor, callback_handler)
 
 

test/integration/test_langgraph_real_package_smoke.py

@@ -37,6 +37,7 @@
     patch as langgraph_patch_module,
 )
 from agent_assembly.adapters.langgraph.patch import LangGraphPatch  # noqa: E402
+from agent_assembly.core.spawn import _SPAWN_CTX, SpawnContext  # noqa: E402
 
 
 @pytest.fixture(autouse=True)
@@ -149,6 +150,53 @@ def test_real_stategraph_async_ainvoke_fires_node_hooks() -> None:
         patch.revert()
 
 
+@pytest.mark.integration
+def test_real_stategraph_subgraph_as_node_propagates_spawn_lineage() -> None:
+    """A compiled subgraph used as a parent-graph node must get spawn lineage.
+
+    Regresses AAASM-4472: ``_wrap_node_entry`` checked ``_is_compiled_subgraph()``
+    against the *outer* ``PregelNode`` wrapper before unwrapping it via ``.bound``.
+    That outer wrapper never exposes ``.nodes``, so the check failed and the code
+    fell through to generic hook-wrapping instead of ``_wrap_subgraph_spawn_node``.
+    Coarse hooks still fired, but ``_SPAWN_CTX`` was never set inside the
+    subgraph's own nodes, so parent-agent lineage was lost for the standard
+    multi-agent delegation pattern (``parent_graph.add_node("sub", compiled_subgraph)``).
+    """
+    captured_spawn_ctx: list[SpawnContext | None] = []
+
+    def _inner_node(state: _CounterState) -> dict[str, int]:
+        captured_spawn_ctx.append(_SPAWN_CTX.get())
+        return {"count": state["count"] + 1}
+
+    sub_graph = StateGraph(_CounterState)
+    sub_graph.add_node("inner", _inner_node)
+    sub_graph.add_edge(START, "inner")
+    sub_graph.add_edge("inner", END)
+    compiled_subgraph = sub_graph.compile()
+
+    recorder = _EventRecorder()
+    patch = LangGraphPatch(callback_handler=recorder, process_agent_id="parent-agent-42")
+    assert patch.apply() is True
+
+    try:
+        parent_graph = StateGraph(_CounterState)
+        parent_graph.add_node("sub", compiled_subgraph)
+        parent_graph.add_edge(START, "sub")
+        parent_graph.add_edge("sub", END)
+        compiled_parent = parent_graph.compile()
+
+        result = compiled_parent.invoke(_CounterState(count=0))
+
+        assert result == {"count": 1}
+        assert len(captured_spawn_ctx) == 1
+        spawn_ctx = captured_spawn_ctx[0]
+        assert spawn_ctx is not None, "expected _SPAWN_CTX to be populated inside the subgraph's own node"
+        assert spawn_ctx.parent_agent_id == "parent-agent-42"
+        assert spawn_ctx.delegation_reason == "langgraph_node:sub"
+    finally:
+        patch.revert()
+
+
 def test_revert_restores_the_real_stategraph_compile() -> None:
     """``revert()`` must restore the genuine (unpatched) ``StateGraph.compile``."""
     original_compile = StateGraph.compile