diff --git a/agent_assembly/adapters/langgraph/patch.py b/agent_assembly/adapters/langgraph/patch.py index f0b1ed28..43e5dfe4 100644 --- a/agent_assembly/adapters/langgraph/patch.py +++ b/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) diff --git a/test/integration/test_langgraph_real_package_smoke.py b/test/integration/test_langgraph_real_package_smoke.py index 49fe787b..807ff179 100644 --- a/test/integration/test_langgraph_real_package_smoke.py +++ b/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