From 351e402dddc1eb62f6f3d67e25bc3aeef7564ea7 Mon Sep 17 00:00:00 2001 From: Bryant Date: Sat, 11 Jul 2026 23:01:49 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20(langgraph):=20Add=20missing=20i?= =?UTF-8?q?dempotency=20guard=20for=20sync=20subgraph=20invoke?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AAASM-4472 added _wrap_subgraph_invoke_methods_in_place() to inject lineage tracking into a compiled subgraph's own invoke/ainvoke when it is reused as a node. The async path guards against double-wrapping via an _agent_assembly_ainvoke_spawned marker, but the sync path had no equivalent guard: reusing the same compiled subgraph as a node in two different parent graphs re-wrapped .invoke on the second compile, stacking spawn-context depth (2 instead of 1) on a single sync invoke. Add a matching _agent_assembly_invoke_spawned marker checked before wrapping .invoke, mirroring the existing async guard exactly. New regression tests: a sync test reproducing the double-wrap (mount the same compiled subgraph in two parent graphs, invoke synchronously, assert lineage depth stays 1) and a symmetric async test confirming the already-correct async path stays correct. Closes AAASM-4474 --- agent_assembly/adapters/langgraph/patch.py | 3 +- .../test_langgraph_real_package_smoke.py | 105 ++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langgraph/patch.py b/agent_assembly/adapters/langgraph/patch.py index 43e5dfe4..46153e26 100644 --- a/agent_assembly/adapters/langgraph/patch.py +++ b/agent_assembly/adapters/langgraph/patch.py @@ -377,7 +377,7 @@ def _wrap_subgraph_invoke_methods_in_place(node_name: Any, subgraph: Any, proces wrapped_any = False original_invoke = getattr(subgraph, "invoke", None) - if callable(original_invoke): + if callable(original_invoke) and not getattr(subgraph, "_agent_assembly_invoke_spawned", False): def sync_subgraph_wrapper(*args: Any, **kwargs: Any) -> Any: spawn_ctx = SpawnContext( @@ -390,6 +390,7 @@ def sync_subgraph_wrapper(*args: Any, **kwargs: Any) -> Any: return original_invoke(*args, **kwargs) subgraph.invoke = sync_subgraph_wrapper + subgraph._agent_assembly_invoke_spawned = True wrapped_any = True original_ainvoke = getattr(subgraph, "ainvoke", None) diff --git a/test/integration/test_langgraph_real_package_smoke.py b/test/integration/test_langgraph_real_package_smoke.py index 807ff179..36b7cfe3 100644 --- a/test/integration/test_langgraph_real_package_smoke.py +++ b/test/integration/test_langgraph_real_package_smoke.py @@ -197,6 +197,111 @@ def _inner_node(state: _CounterState) -> dict[str, int]: patch.revert() +@pytest.mark.integration +def test_real_stategraph_reused_subgraph_across_parents_sync_invoke_lineage_depth_stable() -> None: + """A compiled subgraph reused as a node in two different parent graphs must not + accumulate wrapping layers on ``.invoke``. + + Regresses AAASM-4474: ``_wrap_subgraph_invoke_methods_in_place`` guards the + async path against double-wrapping via ``_agent_assembly_ainvoke_spawned``, + but has no equivalent guard for the sync path. Compiling a second parent + graph that reuses the same compiled subgraph instance re-wraps ``.invoke`` + around the already-wrapped ``.invoke``, so a single synchronous invoke of the + first parent graph corrupts the lineage depth captured inside the subgraph's + own node (comes back as 2 instead of the expected 1). + """ + captured_depths: list[int | None] = [] + + def _inner_node(state: _CounterState) -> dict[str, int]: + spawn_ctx = _SPAWN_CTX.get() + captured_depths.append(spawn_ctx.depth if spawn_ctx is not None else None) + 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-a") + assert patch.apply() is True + + try: + parent_graph_a = StateGraph(_CounterState) + parent_graph_a.add_node("sub", compiled_subgraph) + parent_graph_a.add_edge(START, "sub") + parent_graph_a.add_edge("sub", END) + compiled_parent_a = parent_graph_a.compile() + + # Reuse the same compiled subgraph instance as a node in a second, + # unrelated parent graph. This is a legitimate pattern (a shared + # reusable subgraph mounted under multiple parent orchestrators) and is + # what triggers the double-wrap when compiled a second time. + parent_graph_b = StateGraph(_CounterState) + parent_graph_b.add_node("sub", compiled_subgraph) + parent_graph_b.add_edge(START, "sub") + parent_graph_b.add_edge("sub", END) + parent_graph_b.compile() + + result = compiled_parent_a.invoke(_CounterState(count=0)) + + assert result == {"count": 1} + assert captured_depths == [1], ( + "expected lineage depth 1 for a sync invoke of a subgraph reused " + "across two parent graphs, not an accumulated double-wrap depth" + ) + finally: + patch.revert() + + +@pytest.mark.integration +def test_real_stategraph_reused_subgraph_across_parents_async_ainvoke_lineage_depth_stable() -> None: + """Async counterpart of the sync test above. + + The existing ``_agent_assembly_ainvoke_spawned`` guard already prevents + double-wrapping ``.ainvoke``, so this must already pass before the sync fix + lands and must keep passing afterward as a regression guard confirming the + async path stays correct. + """ + captured_depths: list[int | None] = [] + + def _inner_node(state: _CounterState) -> dict[str, int]: + spawn_ctx = _SPAWN_CTX.get() + captured_depths.append(spawn_ctx.depth if spawn_ctx is not None else None) + 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-a") + assert patch.apply() is True + + try: + parent_graph_a = StateGraph(_CounterState) + parent_graph_a.add_node("sub", compiled_subgraph) + parent_graph_a.add_edge(START, "sub") + parent_graph_a.add_edge("sub", END) + compiled_parent_a = parent_graph_a.compile() + + parent_graph_b = StateGraph(_CounterState) + parent_graph_b.add_node("sub", compiled_subgraph) + parent_graph_b.add_edge(START, "sub") + parent_graph_b.add_edge("sub", END) + parent_graph_b.compile() + + result = asyncio.run(compiled_parent_a.ainvoke(_CounterState(count=0))) + + assert result == {"count": 1} + assert captured_depths == [1] + finally: + patch.revert() + + def test_revert_restores_the_real_stategraph_compile() -> None: """``revert()`` must restore the genuine (unpatched) ``StateGraph.compile``.""" original_compile = StateGraph.compile