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

[Chisanan232]

What changed

_wrap_node_entry in agent_assembly/adapters/langgraph/patch.py checked _is_compiled_subgraph() against the outer PregelNode wrapper before unwrapping it via .bound. When a compiled subgraph is used as a node (parent_graph.add_node("sub", compiled_subgraph)) — the standard multi-agent delegation pattern — LangGraph also wraps that subgraph in a PregelNode. The outer wrapper never exposes .nodes (only the wrapped object does), so the check failed and the code fell through to generic hook-wrapping instead of the dedicated _wrap_subgraph_spawn_node() path.

Effect: coarse governance hooks still fired correctly, but _SPAWN_CTX (the parent-agent lineage/spawn-tracking context) was never set inside the subgraph's own nodes — any policy/audit logic depending on "which parent agent spawned this subgraph invocation" silently lost that lineage.

This is a second, narrower instance of the same fail-open class as AAASM-4434 (which fixed the top-level-node case: PregelNode.invoke vs PregelNode.bound.invoke).

Why

Governance lineage tracking must hold for the standard subgraph-as-node delegation pattern, not just top-level callable nodes.

The fix

In _wrap_node_entry, the PregelNode-wrapper check now runs first, before the callable() and _is_compiled_subgraph() checks. It unwraps via .bound and then runs _is_compiled_subgraph() against the unwrapped object — so subgraph detection works whether the subgraph arrives bare or already PregelNode-wrapped (LangGraph's own wrapping behavior varies by call site/version).

Because the outer PregelNode container (channels/triggers/mapper/writers) must stay intact for the Pregel runtime to keep dispatching correctly, the fix does not replace the whole node-map entry (which is what _wrap_subgraph_spawn_node does for bare, non-wrapped subgraphs — safe there since the map entry is the subgraph). Instead, a new _wrap_subgraph_invoke_methods_in_place() mutates the bound subgraph's own invoke/ainvoke attributes in place, capturing the original methods first so the wrapper calls the real implementation rather than recursing into itself once installed.

Test evidence

Reproduced against the real installed langgraph==1.2.9 package per the ticket's repro steps (parent StateGraph with add_node("sub", compiled_subgraph), invoked, _SPAWN_CTX inspected from inside the subgraph's own node).

Before fix (new test failing):

test/integration/test_langgraph_real_package_smoke.py::test_real_stategraph_subgraph_as_node_propagates_spawn_lineage FAILED
AssertionError: expected _SPAWN_CTX to be populated inside the subgraph's own node
assert None is not None

After fix (same test passing):

test/integration/test_langgraph_real_package_smoke.py::test_real_stategraph_sync_invoke_fires_node_hooks PASSED
test/integration/test_langgraph_real_package_smoke.py::test_real_stategraph_async_ainvoke_fires_node_hooks PASSED
test/integration/test_langgraph_real_package_smoke.py::test_real_stategraph_subgraph_as_node_propagates_spawn_lineage PASSED
test/integration/test_langgraph_real_package_smoke.py::test_revert_restores_the_real_stategraph_compile PASSED
4 passed

The existing AAASM-4434 real-package smoke test (top-level node case) passes unchanged. Also manually verified the async (ainvoke) path for a subgraph-as-node case sets _SPAWN_CTX correctly (not covered by an existing test pattern, verified via a standalone repro script).

Full check suite results

• pytest test/: 780 passed, 16 skipped (pre-existing environment gaps: native _core module not built, optional framework deps like crewai/google.adk/haystack not installed — unrelated to this change), 0 failed.
• ruff check .: all checks passed.
• mypy agent_assembly: 4 pre-existing errors, all in agent_assembly/types.py, agent_assembly/__init__.py, agent_assembly/op_control.py, agent_assembly/core/runtime_interceptor.py — all related to the native _core module not being built / missing grpc stubs, present identically on master before this change. mypy agent_assembly/adapters/langgraph/patch.py (the changed file): 0 issues.
• pre-commit run --files agent_assembly/adapters/langgraph/patch.py test/integration/test_langgraph_real_package_smoke.py: all hooks passed (isort, autoflake, black, mypy, etc.).

Related

Closes AAASM-4472. Follow-up to AAASM-4434.

[codecov]

Codecov Report

❌ Patch coverage is 88.46154% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
agent_assembly/adapters/langgraph/patch.py	88.46%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
88.5% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[Chisanan232]

Claude Code — PR review

1. CI check

Re-ran gh pr checks 236 fresh: 19/19 passed, 0 failed. No pending/failing jobs (coverage-upload jobs referenced as "pending" earlier have since completed green). No action needed.

2. Scope check (AAASM-4472 AC)

Compared the ticket's AC against the diff:

• ✅ New test exercising the exact repro (test_real_stategraph_subgraph_as_node_propagates_spawn_lineage) — asserts _SPAWN_CTX is populated inside the subgraph's own node, with correct parent_agent_id and delegation_reason, matching the rigor requested (real hook/context assertion, not just "no exception").
• ✅ Existing AAASM-4434 smoke test (test_real_stategraph_sync_invoke_fires_node_hooks, ..._async_ainvoke_fires_node_hooks) still passes unchanged.
• ✅ Ran the full langgraph test surface locally: pytest test/ -k langgraph → 54 passed, 1 skipped (skip unrelated to this change). No regression to general hook-wrapping.

PR fully covers the stated AC.

3. Side-effect check — three-case dispatch trace (core governance logic)

Traced _wrap_node_entry for all three cases the reorder must preserve, and verified each empirically by executing the patched module against a real langgraph 1.2.9 install (not just reading the diff):

• (a) Plain callable/Runnable node, PregelNode-wrapped (AAASM-4434's case): _is_pregel_node_wrapper → True → unwrap .bound → _is_compiled_subgraph(bound) → False (RunnableCallable has .invoke but no .nodes) → falls to _wrap_node_invoke_methods(bound, ...). Confirmed unchanged/correct — verified live, hooks fire as before.
• (b) Bare compiled subgraph as node, not PregelNode-wrapped (pre-existing working path): _is_pregel_node_wrapper → False (CompiledStateGraph has no .bound) → callable(node_executor) → False (confirmed empirically — CompiledStateGraph is not directly callable despite defining __call__ at the class level) → _is_compiled_subgraph → True → _wrap_subgraph_spawn_node. Confirmed unchanged/correct. Note: empirically, current pinned langgraph (1.2.9) actually PregelNode-wraps every node including subgraph nodes, so this bare path is effectively dead code against the pinned version today — but it's preserved correctly for other langgraph versions per the module's documented duck-typing/graceful-degradation design, and the ordering relative to case (a)/(c) is untouched by this PR.
• (c) PregelNode-wrapped compiled subgraph (this PR's fix): _is_pregel_node_wrapper → True → unwrap .bound → _is_compiled_subgraph(bound) → True (bound is CompiledStateGraph, has .nodes+.invoke) → _wrap_subgraph_invoke_methods_in_place. Verified live: _SPAWN_CTX now correctly populates inside the subgraph's own node with depth=1, parent_agent_id and delegation_reason all correct. This is the actual bug fix, and it works.

None of the three cases fall through to the wrong branch. The reorder is safe.

⚠️ Found a real issue while verifying the "capture originals to avoid recursion" claim in _wrap_subgraph_invoke_methods_in_place:

The docstring's claim that capturing original_invoke/original_ainvoke before reassignment prevents the wrapper from recursing into itself is correct — that part is not a bug, and there's no infinite-recursion risk within a single wrap call.

However, I found an asymmetry: the ainvoke path guards against re-wrapping with not getattr(subgraph, "_agent_assembly_ainvoke_spawned", False), but the sync invoke path has no equivalent guard. I reproduced this live: if the same compiled subgraph instance is used as a node in two different parent graphs (e.g. a shared/reusable subgraph mounted into multiple orchestrators — a plausible pattern, not contrived), _wrap_subgraph_invoke_methods_in_place runs twice against the same object and double-wraps .invoke:

Single parent, subgraph used once:         captured depth = 1  (correct)
Same subgraph reused across 2 parents:     captured depth = 2  (sync invoke — WRONG, double-wrapped)
Same subgraph reused across 2 parents:     captured depth = 1  (async ainvoke — correct, guard works)

This isn't infinite recursion (bounded, no stack overflow), but it's silent N-fold instrumentation stacking that corrupts the lineage depth and duplicates spawn-context construction on every reuse — exactly the kind of governance/audit correctness issue this module exists to prevent. Every other wrap path in this file has an idempotency guard (_NODE_WRAPPED_FLAG for callables, _agent_assembly_ainvoke_spawned for async subgraph paths); this is the one path that's missing one.

Suggested fix: add a matching guard, e.g. _agent_assembly_invoke_spawned, checked/set the same way _agent_assembly_ainvoke_spawned already is, before wrapping subgraph.invoke.

This doesn't invalidate the fix itself (the new test's scenario — single subgraph instance, single parent — is unaffected and passes correctly), but it's a latent correctness gap in the new code that should be closed, since a shared-subgraph-across-orchestrators pattern is realistic for this SDK's use case.

4. Front-end check

N/A — pure Python governance library code, no UI surface.

Verdict

CI is green and the AC is fully met. The core three-case dispatch reorder is correct and was verified empirically against a real langgraph 1.2.9 install, not just read from the diff. Recommend a small follow-up fix before merge: add an idempotency guard to the sync invoke wrap in _wrap_subgraph_invoke_methods_in_place (mirroring the existing ainvoke guard) to prevent depth-corrupting double-wrap when a compiled subgraph is reused as a node across multiple parent graphs. Everything else is solid.

[Chisanan232]

Review — Claude Code (AAASM-4472)

Verdict: Approve. The fix cleanly addresses the ticket's root cause and matches the suggested fix direction exactly.

CI

All real checks green — CI (unit + integration), CodeQL, docs-build, pip-audit, LangChain contract test all SUCCESS. Only codecov/SonarCloud are informational. mergeStateStatus: BLOCKED is just the pending required approval, not a failure.

Scope — does it fix the ticket?

Yes. _wrap_node_entry now runs the _is_pregel_node_wrapper check first, unwraps via .bound, then applies _is_compiled_subgraph() to the unwrapped object — precisely the ticket's suggested direction. Because the outer PregelNode container (channels/triggers/mapper/writers) must stay intact for the Pregel runtime, the fix mutates the bound subgraph's own invoke/ainvoke in place (_wrap_subgraph_invoke_methods_in_place) rather than replacing the node-map entry — the right call for the wrapped case. Parent→child lineage (_SPAWN_CTX with parent_agent_id / delegation_reason) now propagates into the subgraph's own nodes.

Side-effects — normal / other paths

No regression to the non-subgraph path. A plain top-level node is still PregelNode-wrapped with bound = a Runnable (not a compiled subgraph), so it falls to _wrap_node_invoke_methods(bound) — identical to AAASM-4434's behavior. Reordering the pregel check ahead of the callable() check is safe: PregelNodes aren't callable (otherwise AAASM-4434's post-callable branch would have been dead, yet it was verified working). The bare (non-wrapped) compiled-subgraph path is preserved unchanged.

Regression test

Yes, and it covers both directions the AC asked for: the new test_real_stategraph_subgraph_as_node_propagates_spawn_lineage exercises the exact repro against real langgraph and asserts _SPAWN_CTX IS populated (parent_agent_id == "parent-agent-42", delegation_reason == "langgraph_node:sub") inside the subgraph's own node — not merely "no exception." The existing AAASM-4434 top-level smoke tests remain and pass, proving the normal case is unchanged.

Validation

Ran locally in the worktree: pytest test/integration/test_langgraph_real_package_smoke.py → 4 passed (new subgraph test + 3 existing AAASM-4434 tests). Matches CI's green integration job.

Minor, non-blocking observations

1. Idempotency asymmetry: ainvoke wrapping is guarded by _agent_assembly_ainvoke_spawned, but the sync invoke wrapping has no such guard. A double-apply() on the same subgraph object would re-wrap sync invoke, nesting spawn-context scopes. Low practical risk (apply is normally once-per-process), but the two paths could be made symmetric.
2. Revert coverage: the in-place invoke/ainvoke mutations aren't restored by revert() (which restores StateGraph.compile). This is the same pre-existing pattern as _wrap_node_invoke_methods, so not a new regression — just noting a wrapped subgraph instance retains the wrapper after revert.

Neither blocks merge.