[AAASM-51] ✨ (langgraph): Add node execution hooks with metadata

.github/workflows/rw_run_all_test_and_record.yaml

@@ -164,6 +164,6 @@ jobs:
       - name: SonarCloud Scan
         uses: SonarSource/sonarqube-scan-action@v7.1.0
         env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GITHUB_TOKEN: ${{ github.token }}
           SONAR_TOKEN: ${{ secrets.sonar_token }}
           SONAR_HOST_URL: https://sonarcloud.io

agent_assembly/adapters/langchain/__init__.py

@@ -1,6 +1,7 @@
 """LangChain adapter package."""
 
 from agent_assembly.adapters.langchain.callback_handler import AssemblyCallbackHandler
+from agent_assembly.adapters.langgraph import LangGraphPatch
 from agent_assembly.adapters.langchain.langgraph_patch import patch_stategraph_compile
 from agent_assembly.adapters.langchain.runtime import (
     auto_inject_callback_handler,
@@ -9,6 +10,7 @@
 
 __all__ = [
     "AssemblyCallbackHandler",
+    "LangGraphPatch",
     "patch_stategraph_compile",
     "auto_inject_callback_handler",
     "get_active_callback_handler",

agent_assembly/adapters/langchain/callback_handler.py

@@ -277,16 +277,46 @@ async def aon_llm_end(
             await result
         return None
 
-    def on_graph_node_start(self, node_name: str, state: Any) -> None:
+    def on_graph_node_start(
+        self,
+        node_name: str,
+        state: Any,
+        *,
+        agent_id: str | None = None,
+        state_keys: list[str] | None = None,
+        config: Any = None,
+    ) -> None:
         method = getattr(self._interceptor, "on_graph_node_start", None)
         if not callable(method):
             return None
-        method(node_name=node_name, state=state)
+        method(
+            node_name=node_name,
+            agent_id=agent_id,
+            state=state,
+            state_keys=state_keys,
+            config=config,
+        )
         return None
 
-    def on_graph_node_end(self, node_name: str, state: Any, result: Any) -> None:
+    def on_graph_node_end(
+        self,
+        node_name: str,
+        state: Any,
+        result: Any,
+        *,
+        agent_id: str | None = None,
+        state_delta: dict[str, Any] | None = None,
+        config: Any = None,
+    ) -> None:
         method = getattr(self._interceptor, "on_graph_node_end", None)
         if not callable(method):
             return None
-        method(node_name=node_name, state=state, result=result)
+        method(
+            node_name=node_name,
+            agent_id=agent_id,
+            state=state,
+            result=result,
+            state_delta=state_delta,
+            config=config,
+        )
         return None

agent_assembly/adapters/langchain/langgraph_patch.py

@@ -1,200 +1,33 @@
-"""LangGraph compile-time patching for governance interception."""
-
-from __future__ import annotations
+"""Backward-compatible shim for LangGraph patch utilities."""
 
 import importlib
-import inspect
-from typing import Any
-
-_PATCHED_FLAG = "_agent_assembly_compile_patched"
-_ORIGINAL_COMPILE = "_agent_assembly_original_compile"
-_NODE_WRAPPED_FLAG = "_agent_assembly_node_wrapped"
-_INVOKE_WRAPPED_FLAG = "_agent_assembly_invoke_wrapped"
-
-
-def _extract_state(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
-    if args:
-        return args[0]
-    return kwargs.get("state")
-
-
-def _invoke_pre_node_hook(callback_handler: Any, node_name: str, state: Any) -> None:
-    method = getattr(callback_handler, "on_graph_node_start", None)
-    if not callable(method):
-        return None
-
-    result = method(node_name=node_name, state=state)
-    if inspect.isawaitable(result):
-        return None
-
-    return None
-
-
-def _invoke_post_node_hook(callback_handler: Any, node_name: str, state: Any, result: Any) -> None:
-    method = getattr(callback_handler, "on_graph_node_end", None)
-    if not callable(method):
-        return None
-
-    callback_result = method(node_name=node_name, state=state, result=result)
-    if inspect.isawaitable(callback_result):
-        return None
-
-    return None
-
-
-def _wrap_node_callable(node_name: str, node_callable: Any, callback_handler: Any) -> Any:
-    if getattr(node_callable, _NODE_WRAPPED_FLAG, False):
-        return node_callable
-
-    def wrapped_node(*node_args: Any, **node_kwargs: Any) -> Any:
-        state = _extract_state(node_args, node_kwargs)
-        _invoke_pre_node_hook(callback_handler, node_name=node_name, state=state)
-
-        node_result = node_callable(*node_args, **node_kwargs)
-        if inspect.isawaitable(node_result):
-            async def awaited_node_result() -> Any:
-                resolved_result = await node_result
-                _invoke_post_node_hook(
-                    callback_handler,
-                    node_name=node_name,
-                    state=state,
-                    result=resolved_result,
-                )
-                return resolved_result
-
-            return awaited_node_result()
-
-        _invoke_post_node_hook(
-            callback_handler,
-            node_name=node_name,
-            state=state,
-            result=node_result,
-        )
-        return node_result
-
-    setattr(wrapped_node, _NODE_WRAPPED_FLAG, True)
-    return wrapped_node
-
-
-def _wrap_node_map(node_map: Any, callback_handler: Any) -> bool:
-    items_method = getattr(node_map, "items", None)
-    if not callable(items_method):
-        return False
-
-    wrapped_any = False
-    for node_name, node_executor in list(items_method()):
-        if callable(node_executor):
-            wrapped_executor = _wrap_node_callable(str(node_name), node_executor, callback_handler)
-            if wrapped_executor is node_executor:
-                continue
-            try:
-                node_map[node_name] = wrapped_executor
-            except Exception:
-                continue
-            wrapped_any = True
-            continue
-
-        invoke = getattr(node_executor, "invoke", None)
-        if callable(invoke):
-            setattr(
-                node_executor,
-                "invoke",
-                _wrap_node_callable(str(node_name), invoke, callback_handler),
-            )
-            wrapped_any = True
-
-        ainvoke = getattr(node_executor, "ainvoke", None)
-        if callable(ainvoke):
-            setattr(
-                node_executor,
-                "ainvoke",
-                _wrap_node_callable(str(node_name), ainvoke, callback_handler),
-            )
-            wrapped_any = True
-
-    return wrapped_any
-
-
-def _wrap_compiled_graph_nodes(compiled_graph: Any, callback_handler: Any) -> bool:
-    candidate_maps = [
-        getattr(compiled_graph, "nodes", None),
-        getattr(compiled_graph, "_nodes", None),
-    ]
-
-    pregel = getattr(compiled_graph, "pregel", None)
-    if pregel is None:
-        pregel = getattr(compiled_graph, "_pregel", None)
-    if pregel is not None:
-        candidate_maps.extend(
-            [
-                getattr(pregel, "nodes", None),
-                getattr(pregel, "_nodes", None),
-            ]
-        )
-
-    wrapped_any = False
-    for node_map in candidate_maps:
-        if node_map is None:
-            continue
-        if _wrap_node_map(node_map, callback_handler):
-            wrapped_any = True
-
-    return wrapped_any
-
-
-def patch_stategraph_compile(callback_handler: Any) -> bool:
-    """Patch `StateGraph.compile()` to attach runtime governance hooks."""
-    try:
-        module = importlib.import_module("langgraph.graph.state")
-    except ImportError:
-        return False
-
-    state_graph_cls = getattr(module, "StateGraph", None)
-    if state_graph_cls is None:
-        return False
-
-    if getattr(state_graph_cls, _PATCHED_FLAG, False):
-        return True
-
-    original_compile = state_graph_cls.compile
-
-    def patched_compile(self: Any, *args: Any, **kwargs: Any) -> Any:
-        compiled_graph = original_compile(self, *args, **kwargs)
-        nodes_wrapped = _wrap_compiled_graph_nodes(compiled_graph, callback_handler)
-        if not nodes_wrapped:
-            invoke = getattr(compiled_graph, "invoke", None)
-            if callable(invoke) and not getattr(invoke, _INVOKE_WRAPPED_FLAG, False):
-                def wrapped_invoke(*invoke_args: Any, **invoke_kwargs: Any) -> Any:
-                    state = _extract_state(invoke_args, invoke_kwargs)
-                    _invoke_pre_node_hook(callback_handler, node_name="graph.invoke", state=state)
-
-                    invoke_result = invoke(*invoke_args, **invoke_kwargs)
-                    if inspect.isawaitable(invoke_result):
-                        async def awaited_invoke_result() -> Any:
-                            resolved_result = await invoke_result
-                            _invoke_post_node_hook(
-                                callback_handler,
-                                node_name="graph.invoke",
-                                state=state,
-                                result=resolved_result,
-                            )
-                            return resolved_result
-
-                        return awaited_invoke_result()
-
-                    _invoke_post_node_hook(
-                        callback_handler,
-                        node_name="graph.invoke",
-                        state=state,
-                        result=invoke_result,
-                    )
-                    return invoke_result
-
-                setattr(wrapped_invoke, _INVOKE_WRAPPED_FLAG, True)
-                setattr(compiled_graph, "invoke", wrapped_invoke)
-        return compiled_graph
 
-    setattr(state_graph_cls, _ORIGINAL_COMPILE, original_compile)
-    setattr(state_graph_cls, "compile", patched_compile)
-    setattr(state_graph_cls, _PATCHED_FLAG, True)
-    return True
+from agent_assembly.adapters.langgraph import patch as _impl
+
+LangGraphPatch = _impl.LangGraphPatch
+patch_stategraph_compile = _impl.patch_stategraph_compile
+
+_PATCHED_FLAG = _impl._PATCHED_FLAG
+_ORIGINAL_COMPILE = _impl._ORIGINAL_COMPILE
+_NODE_WRAPPED_FLAG = _impl._NODE_WRAPPED_FLAG
+_INVOKE_WRAPPED_FLAG = _impl._INVOKE_WRAPPED_FLAG
+
+_extract_state = _impl._extract_state
+_extract_config = _impl._extract_config
+_extract_agent_id = _impl._extract_agent_id
+_summarize_state_keys = _impl._summarize_state_keys
+_compute_state_delta = _impl._compute_state_delta
+_invoke_pre_node_hook = _impl._invoke_pre_node_hook
+_invoke_post_node_hook = _impl._invoke_post_node_hook
+_wrap_node_callable = _impl._wrap_node_callable
+_make_sync_node_wrapper = _impl._make_sync_node_wrapper
+_make_async_node_wrapper = _impl._make_async_node_wrapper
+_make_assembly_node_wrapper = _impl._make_assembly_node_wrapper
+_wrap_node_map = _impl._wrap_node_map
+_wrap_compiled_graph_nodes = _impl._wrap_compiled_graph_nodes
+_discover_compiled_graph_node_maps = _impl._discover_compiled_graph_node_maps
+_apply_stategraph_compile_patch = _impl._apply_stategraph_compile_patch
+_wrap_graph_invoke_fallback = _impl._wrap_graph_invoke_fallback
+_record_node_enter = _impl._record_node_enter
+_record_node_exit = _impl._record_node_exit
+_load_stategraph_class = _impl._load_stategraph_class

agent_assembly/adapters/langchain/runtime.py

@@ -6,7 +6,7 @@
 from typing import Any
 
 from agent_assembly.adapters.langchain.callback_handler import AssemblyCallbackHandler
-from agent_assembly.adapters.langchain.langgraph_patch import patch_stategraph_compile
+from agent_assembly.adapters.langgraph import LangGraphPatch
 
 _ACTIVE_CALLBACK_HANDLER: AssemblyCallbackHandler | None = None
 _RUNTIME_LOCK = Lock()
@@ -18,12 +18,12 @@ def auto_inject_callback_handler(interceptor: Any) -> AssemblyCallbackHandler:
 
     with _RUNTIME_LOCK:
         if _ACTIVE_CALLBACK_HANDLER is not None:
-            patch_stategraph_compile(_ACTIVE_CALLBACK_HANDLER)
+            LangGraphPatch(_ACTIVE_CALLBACK_HANDLER).apply()
             return _ACTIVE_CALLBACK_HANDLER
 
         handler = AssemblyCallbackHandler(interceptor)
         _ACTIVE_CALLBACK_HANDLER = handler
-        patch_stategraph_compile(handler)
+        LangGraphPatch(handler).apply()
         return handler
 
 

agent_assembly/adapters/langgraph/__init__.py

@@ -0,0 +1,5 @@
+"""LangGraph adapter package."""
+
+from agent_assembly.adapters.langgraph.patch import LangGraphPatch
+
+__all__ = ["LangGraphPatch"]