From f819a09f4c4a8cfcad9c73ee0ac7ee8cbc867a38 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:41:33 +0800 Subject: [PATCH 01/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20empty=20l?= =?UTF-8?q?angchain=20adapter=20package=20scaffold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/__init__.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 agent_assembly/adapters/langchain/__init__.py diff --git a/agent_assembly/adapters/langchain/__init__.py b/agent_assembly/adapters/langchain/__init__.py new file mode 100644 index 0000000..3db47b0 --- /dev/null +++ b/agent_assembly/adapters/langchain/__init__.py @@ -0,0 +1 @@ +"""LangChain adapter package.""" From fa35cf1124d8c5c0eb0a4bba6686f38724398972 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:41:47 +0800 Subject: [PATCH 02/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20empty=20c?= =?UTF-8?q?allback=20handler=20module=20scaffold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/callback_handler.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 agent_assembly/adapters/langchain/callback_handler.py diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py new file mode 100644 index 0000000..52a7145 --- /dev/null +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -0,0 +1 @@ +"""LangChain callback handler module.""" From 4d8b9770eac9169546bf483000272135e9fb0675 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:41:59 +0800 Subject: [PATCH 03/37] =?UTF-8?q?=E2=9C=A8=20(errors):=20Add=20ToolExecuti?= =?UTF-8?q?onBlockedError=20for=20governance=20deny?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/exceptions/__init__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/agent_assembly/exceptions/__init__.py b/agent_assembly/exceptions/__init__.py index b973d5d..a3ab7f4 100644 --- a/agent_assembly/exceptions/__init__.py +++ b/agent_assembly/exceptions/__init__.py @@ -9,6 +9,7 @@ "GatewayError", "ConfigurationError", "AdapterValidationError", + "ToolExecutionBlockedError", ] @@ -40,3 +41,8 @@ class ConfigurationError(AssemblyError): class AdapterValidationError(AssemblyError): """Exception raised when an adapter contract is invalid.""" pass + + +class ToolExecutionBlockedError(AssemblyError): + """Exception raised when a tool run is blocked by governance.""" + pass From 0730e0da8f062d2a07a6cef514e7e31530cb5277 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:42:16 +0800 Subject: [PATCH 04/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20AssemblyC?= =?UTF-8?q?allbackHandler=20class=20skeleton?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 52a7145..ab86235 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -1 +1,63 @@ """LangChain callback handler module.""" + +from __future__ import annotations + +from typing import Any +from uuid import UUID + +try: # pragma: no cover - import availability depends on installed extras. + from langchain_core.callbacks import BaseCallbackHandler +except ImportError: # pragma: no cover - fallback keeps runtime import-safe. + class BaseCallbackHandler: # type: ignore[no-redef] + """Fallback base type when langchain-core is not installed.""" + + pass + + +class AssemblyCallbackHandler(BaseCallbackHandler): + """Callback handler that delegates runtime events to governance interception.""" + + def __init__(self, interceptor: Any) -> None: + self._interceptor = interceptor + + def on_tool_start( + self, + serialized: dict[str, Any], + input_str: str, + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + del serialized, input_str, run_id, kwargs + return None + + def on_tool_end( + self, + output: Any, + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + del output, run_id, kwargs + return None + + def on_llm_start( + self, + serialized: dict[str, Any], + prompts: list[str], + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + del serialized, prompts, run_id, kwargs + return None + + def on_llm_end( + self, + response: Any, + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + del response, run_id, kwargs + return None From 601962cd771d074f74bae160d5e9d5b6abf3bca4 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:42:32 +0800 Subject: [PATCH 05/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20governanc?= =?UTF-8?q?e=20decision=20normalization=20helper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index ab86235..b3c9c83 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from typing import Any, Literal, Mapping from uuid import UUID try: # pragma: no cover - import availability depends on installed extras. @@ -20,6 +20,30 @@ class AssemblyCallbackHandler(BaseCallbackHandler): def __init__(self, interceptor: Any) -> None: self._interceptor = interceptor + def _normalize_decision( + self, + decision: object, + ) -> tuple[Literal["allow", "deny", "pending"], str | None]: + if isinstance(decision, str): + normalized = decision.strip().lower() + if normalized in {"allow", "deny", "pending"}: + return normalized, None + return "allow", None + + if isinstance(decision, Mapping): + raw_status = str(decision.get("status", "allow")).strip().lower() + status: Literal["allow", "deny", "pending"] + if raw_status in {"allow", "deny", "pending"}: + status = raw_status + else: + status = "allow" + + reason_value = decision.get("reason") + reason = str(reason_value) if reason_value is not None else None + return status, reason + + return "allow", None + def on_tool_start( self, serialized: dict[str, Any], From da77a3cb3fa5c47319fe9c5c814e6e38d0023301 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:42:45 +0800 Subject: [PATCH 06/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Implement=20on?= =?UTF-8?q?=5Ftool=5Fstart=20deny=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index b3c9c83..3b46f5b 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -5,6 +5,8 @@ from typing import Any, Literal, Mapping from uuid import UUID +from agent_assembly.exceptions import ToolExecutionBlockedError + try: # pragma: no cover - import availability depends on installed extras. from langchain_core.callbacks import BaseCallbackHandler except ImportError: # pragma: no cover - fallback keeps runtime import-safe. @@ -52,7 +54,20 @@ def on_tool_start( run_id: UUID, **kwargs: Any, ) -> None: - del serialized, input_str, run_id, kwargs + method = getattr(self._interceptor, "check_tool_start", None) + if not callable(method): + return None + + decision = method( + serialized=serialized, + input_str=input_str, + run_id=run_id, + **kwargs, + ) + status, reason = self._normalize_decision(decision) + if status == "deny": + raise ToolExecutionBlockedError(reason or "Tool execution blocked by governance.") + return None def on_tool_end( From 38e9ec403c6579ec9a561a335e8efb0c8146feb1 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:42:57 +0800 Subject: [PATCH 07/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Implement=20on?= =?UTF-8?q?=5Ftool=5Fstart=20pending=20approval=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 3b46f5b..430e620 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -67,9 +67,40 @@ def on_tool_start( status, reason = self._normalize_decision(decision) if status == "deny": raise ToolExecutionBlockedError(reason or "Tool execution blocked by governance.") + if status == "pending": + approval = self._resolve_pending_approval( + serialized=serialized, + input_str=input_str, + run_id=run_id, + **kwargs, + ) + approval_status, approval_reason = self._normalize_decision(approval) + if approval_status != "allow": + raise ToolExecutionBlockedError( + approval_reason or reason or "Tool execution was not approved by governance." + ) return None + def _resolve_pending_approval( + self, + *, + serialized: dict[str, Any], + input_str: str, + run_id: UUID, + **kwargs: Any, + ) -> object: + wait_method = getattr(self._interceptor, "wait_for_tool_approval", None) + if not callable(wait_method): + return "deny" + + return wait_method( + serialized=serialized, + input_str=input_str, + run_id=run_id, + **kwargs, + ) + def on_tool_end( self, output: Any, From 9f99d0465f41760511686a3b29835b90e316a0c3 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:43:06 +0800 Subject: [PATCH 08/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Implement=20on?= =?UTF-8?q?=5Ftool=5Fend=20interception=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/callback_handler.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 430e620..d0880bd 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -108,7 +108,15 @@ def on_tool_end( run_id: UUID, **kwargs: Any, ) -> None: - del output, run_id, kwargs + method = getattr(self._interceptor, "on_tool_end", None) + if not callable(method): + return None + + method( + output=output, + run_id=run_id, + **kwargs, + ) return None def on_llm_start( From c361bbdd2c2e377b868f33bc8cb2b46885ee0604 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:43:17 +0800 Subject: [PATCH 09/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Implement=20on?= =?UTF-8?q?=5Fllm=5Fstart=20scan-only=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/callback_handler.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index d0880bd..30e31f2 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -127,7 +127,16 @@ def on_llm_start( run_id: UUID, **kwargs: Any, ) -> None: - del serialized, prompts, run_id, kwargs + method = getattr(self._interceptor, "on_llm_start_scan", None) + if not callable(method): + return None + + method( + serialized=serialized, + prompts=prompts, + run_id=run_id, + **kwargs, + ) return None def on_llm_end( From 20562cd1a2efc4a9623ab06deabf1f99b6de262e Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:43:29 +0800 Subject: [PATCH 10/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Implement=20on?= =?UTF-8?q?=5Fllm=5Fend=20interception=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/callback_handler.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 30e31f2..4793b36 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -146,5 +146,13 @@ def on_llm_end( run_id: UUID, **kwargs: Any, ) -> None: - del response, run_id, kwargs + method = getattr(self._interceptor, "on_llm_end", None) + if not callable(method): + return None + + method( + response=response, + run_id=run_id, + **kwargs, + ) return None From 7c9e2013306e3c2c1af9765da797c7a459fdf4c5 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:43:48 +0800 Subject: [PATCH 11/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20async=20o?= =?UTF-8?q?n=5Ftool=5Fstart=20interception?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 4793b36..a3d7792 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect from typing import Any, Literal, Mapping from uuid import UUID @@ -101,6 +102,47 @@ def _resolve_pending_approval( **kwargs, ) + async def aon_tool_start( + self, + serialized: dict[str, Any], + input_str: str, + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + method = getattr(self._interceptor, "check_tool_start", None) + if not callable(method): + return None + + decision = method( + serialized=serialized, + input_str=input_str, + run_id=run_id, + **kwargs, + ) + if inspect.isawaitable(decision): + decision = await decision + + status, reason = self._normalize_decision(decision) + if status == "deny": + raise ToolExecutionBlockedError(reason or "Tool execution blocked by governance.") + if status == "pending": + approval = self._resolve_pending_approval( + serialized=serialized, + input_str=input_str, + run_id=run_id, + **kwargs, + ) + if inspect.isawaitable(approval): + approval = await approval + approval_status, approval_reason = self._normalize_decision(approval) + if approval_status != "allow": + raise ToolExecutionBlockedError( + approval_reason or reason or "Tool execution was not approved by governance." + ) + + return None + def on_tool_end( self, output: Any, From 262a47aec610d3c76474368d7ec8a1497d0df55d Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:43:57 +0800 Subject: [PATCH 12/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20async=20o?= =?UTF-8?q?n=5Ftool=5Fend=20interception?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index a3d7792..aa19e9e 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -161,6 +161,26 @@ def on_tool_end( ) return None + async def aon_tool_end( + self, + output: Any, + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + method = getattr(self._interceptor, "on_tool_end", None) + if not callable(method): + return None + + result = method( + output=output, + run_id=run_id, + **kwargs, + ) + if inspect.isawaitable(result): + await result + return None + def on_llm_start( self, serialized: dict[str, Any], From 903eeb2c90de89d7eb814399ae1d9cd1297bb82f Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:44:06 +0800 Subject: [PATCH 13/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20async=20o?= =?UTF-8?q?n=5Fllm=5Fstart=20interception?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index aa19e9e..35980cb 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -201,6 +201,28 @@ def on_llm_start( ) return None + async def aon_llm_start( + self, + serialized: dict[str, Any], + prompts: list[str], + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + method = getattr(self._interceptor, "on_llm_start_scan", None) + if not callable(method): + return None + + result = method( + serialized=serialized, + prompts=prompts, + run_id=run_id, + **kwargs, + ) + if inspect.isawaitable(result): + await result + return None + def on_llm_end( self, response: Any, From 729ea886e964b53831db181dd95639de33110604 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:44:17 +0800 Subject: [PATCH 14/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20async=20o?= =?UTF-8?q?n=5Fllm=5Fend=20interception?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 35980cb..c388b12 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -240,3 +240,23 @@ def on_llm_end( **kwargs, ) return None + + async def aon_llm_end( + self, + response: Any, + *, + run_id: UUID, + **kwargs: Any, + ) -> None: + method = getattr(self._interceptor, "on_llm_end", None) + if not callable(method): + return None + + result = method( + response=response, + run_id=run_id, + **kwargs, + ) + if inspect.isawaitable(result): + await result + return None From df1e2e25b28e5877d0d50a946828d7bae0a9b23f Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:44:39 +0800 Subject: [PATCH 15/37] =?UTF-8?q?=E2=9C=A8=20(runtime):=20Auto-inject=20ca?= =?UTF-8?q?llback=20handler=20in=20init=5Fassembly?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/runtime.py | 23 ++++++++++++++++++++ agent_assembly/core/assembly.py | 5 ++++- 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 agent_assembly/adapters/langchain/runtime.py diff --git a/agent_assembly/adapters/langchain/runtime.py b/agent_assembly/adapters/langchain/runtime.py new file mode 100644 index 0000000..e9209f6 --- /dev/null +++ b/agent_assembly/adapters/langchain/runtime.py @@ -0,0 +1,23 @@ +"""LangChain runtime wiring helpers.""" + +from __future__ import annotations + +from typing import Any + +from agent_assembly.adapters.langchain.callback_handler import AssemblyCallbackHandler + +_ACTIVE_CALLBACK_HANDLER: AssemblyCallbackHandler | None = None + + +def auto_inject_callback_handler(interceptor: Any) -> AssemblyCallbackHandler: + """Create and register the active callback handler instance.""" + global _ACTIVE_CALLBACK_HANDLER + + handler = AssemblyCallbackHandler(interceptor) + _ACTIVE_CALLBACK_HANDLER = handler + return handler + + +def get_active_callback_handler() -> AssemblyCallbackHandler | None: + """Return the current callback handler instance when one is registered.""" + return _ACTIVE_CALLBACK_HANDLER diff --git a/agent_assembly/core/assembly.py b/agent_assembly/core/assembly.py index 88afbe9..be49337 100644 --- a/agent_assembly/core/assembly.py +++ b/agent_assembly/core/assembly.py @@ -4,6 +4,7 @@ from typing import Optional +from agent_assembly.adapters.langchain.runtime import auto_inject_callback_handler from agent_assembly.client.gateway import GatewayClient from agent_assembly.exceptions import ConfigurationError @@ -32,8 +33,10 @@ def init_assembly( if not agent_id: raise ConfigurationError("agent_id is required") - return GatewayClient( + client = GatewayClient( gateway_url=gateway_url, agent_id=agent_id, api_key=api_key, ) + auto_inject_callback_handler(interceptor=object()) + return client From 92966e21cf77dcce965e232f4faf0d767a0f6c7b Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:44:50 +0800 Subject: [PATCH 16/37] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(runtime):=20Make=20?= =?UTF-8?q?callback=20injection=20idempotent?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/runtime.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/agent_assembly/adapters/langchain/runtime.py b/agent_assembly/adapters/langchain/runtime.py index e9209f6..c1fadef 100644 --- a/agent_assembly/adapters/langchain/runtime.py +++ b/agent_assembly/adapters/langchain/runtime.py @@ -2,20 +2,26 @@ from __future__ import annotations +from threading import Lock from typing import Any from agent_assembly.adapters.langchain.callback_handler import AssemblyCallbackHandler _ACTIVE_CALLBACK_HANDLER: AssemblyCallbackHandler | None = None +_RUNTIME_LOCK = Lock() def auto_inject_callback_handler(interceptor: Any) -> AssemblyCallbackHandler: """Create and register the active callback handler instance.""" global _ACTIVE_CALLBACK_HANDLER - handler = AssemblyCallbackHandler(interceptor) - _ACTIVE_CALLBACK_HANDLER = handler - return handler + with _RUNTIME_LOCK: + if _ACTIVE_CALLBACK_HANDLER is not None: + return _ACTIVE_CALLBACK_HANDLER + + handler = AssemblyCallbackHandler(interceptor) + _ACTIVE_CALLBACK_HANDLER = handler + return handler def get_active_callback_handler() -> AssemblyCallbackHandler | None: From 67f89f1094bc6bfe00edc396d8170ab24ca3f4a3 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:44:59 +0800 Subject: [PATCH 17/37] =?UTF-8?q?=E2=9C=A8=20(langgraph):=20Add=20StateGra?= =?UTF-8?q?ph=20compile=20patch=20scaffold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/langgraph_patch.py | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 agent_assembly/adapters/langchain/langgraph_patch.py diff --git a/agent_assembly/adapters/langchain/langgraph_patch.py b/agent_assembly/adapters/langchain/langgraph_patch.py new file mode 100644 index 0000000..c934cc3 --- /dev/null +++ b/agent_assembly/adapters/langchain/langgraph_patch.py @@ -0,0 +1,11 @@ +"""LangGraph compile-time patching for governance interception.""" + +from __future__ import annotations + +from typing import Any + + +def patch_stategraph_compile(callback_handler: Any) -> bool: + """Patch `StateGraph.compile()` to attach runtime governance hooks.""" + del callback_handler + return False From 8c007daccecef75b5eaff9a082c92686c0a75b11 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:45:15 +0800 Subject: [PATCH 18/37] =?UTF-8?q?=E2=9C=A8=20(langgraph):=20Add=20pre-node?= =?UTF-8?q?=20governance=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/langgraph_patch.py | 48 ++++++++++++++++++- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/agent_assembly/adapters/langchain/langgraph_patch.py b/agent_assembly/adapters/langchain/langgraph_patch.py index c934cc3..e981a32 100644 --- a/agent_assembly/adapters/langchain/langgraph_patch.py +++ b/agent_assembly/adapters/langchain/langgraph_patch.py @@ -2,10 +2,54 @@ from __future__ import annotations +import importlib +import inspect from typing import Any +_PATCHED_FLAG = "_agent_assembly_compile_patched" +_ORIGINAL_COMPILE = "_agent_assembly_original_compile" + + +def _invoke_pre_node_hook(callback_handler: Any, state: Any) -> None: + method = getattr(callback_handler, "on_graph_node_start", None) + if not callable(method): + return None + + result = method(node_name="graph.invoke", state=state) + if inspect.isawaitable(result): + return None + + return None + def patch_stategraph_compile(callback_handler: Any) -> bool: """Patch `StateGraph.compile()` to attach runtime governance hooks.""" - del callback_handler - return False + 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) + invoke = getattr(compiled_graph, "invoke", None) + if callable(invoke): + def wrapped_invoke(state: Any, *invoke_args: Any, **invoke_kwargs: Any) -> Any: + _invoke_pre_node_hook(callback_handler, state) + return invoke(state, *invoke_args, **invoke_kwargs) + + 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 0cb1ae9c9a032f4563be7fa746b129003bedc8ab Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:45:27 +0800 Subject: [PATCH 19/37] =?UTF-8?q?=E2=9C=A8=20(langgraph):=20Add=20post-nod?= =?UTF-8?q?e=20governance=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/langgraph_patch.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langchain/langgraph_patch.py b/agent_assembly/adapters/langchain/langgraph_patch.py index e981a32..04e9531 100644 --- a/agent_assembly/adapters/langchain/langgraph_patch.py +++ b/agent_assembly/adapters/langchain/langgraph_patch.py @@ -22,6 +22,18 @@ def _invoke_pre_node_hook(callback_handler: Any, state: Any) -> None: return None +def _invoke_post_node_hook(callback_handler: Any, 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="graph.invoke", state=state, result=result) + if inspect.isawaitable(callback_result): + return None + + return None + + def patch_stategraph_compile(callback_handler: Any) -> bool: """Patch `StateGraph.compile()` to attach runtime governance hooks.""" try: @@ -44,7 +56,9 @@ def patched_compile(self: Any, *args: Any, **kwargs: Any) -> Any: if callable(invoke): def wrapped_invoke(state: Any, *invoke_args: Any, **invoke_kwargs: Any) -> Any: _invoke_pre_node_hook(callback_handler, state) - return invoke(state, *invoke_args, **invoke_kwargs) + result = invoke(state, *invoke_args, **invoke_kwargs) + _invoke_post_node_hook(callback_handler, state, result) + return result setattr(compiled_graph, "invoke", wrapped_invoke) return compiled_graph From fc34ba963d226a63dd8f653c94aeda2264dcfd7b Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:45:49 +0800 Subject: [PATCH 20/37] =?UTF-8?q?=E2=9C=A8=20(adapters):=20Add=20graph=20n?= =?UTF-8?q?ode=20interception=20methods?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index c388b12..2be480c 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -260,3 +260,17 @@ async def aon_llm_end( if inspect.isawaitable(result): await result return None + + def on_graph_node_start(self, node_name: str, state: Any) -> None: + method = getattr(self._interceptor, "on_graph_node_start", None) + if not callable(method): + return None + method(node_name=node_name, state=state) + return None + + def on_graph_node_end(self, node_name: str, state: Any, result: Any) -> 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) + return None From 97b4743db55d44d63874a6dfdab1533053eae5e2 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:45:57 +0800 Subject: [PATCH 21/37] =?UTF-8?q?=E2=9C=A8=20(runtime):=20Auto-patch=20Lan?= =?UTF-8?q?gGraph=20compile=20during=20injection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/runtime.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/agent_assembly/adapters/langchain/runtime.py b/agent_assembly/adapters/langchain/runtime.py index c1fadef..15830f1 100644 --- a/agent_assembly/adapters/langchain/runtime.py +++ b/agent_assembly/adapters/langchain/runtime.py @@ -6,6 +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 _ACTIVE_CALLBACK_HANDLER: AssemblyCallbackHandler | None = None _RUNTIME_LOCK = Lock() @@ -17,10 +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) return _ACTIVE_CALLBACK_HANDLER handler = AssemblyCallbackHandler(interceptor) _ACTIVE_CALLBACK_HANDLER = handler + patch_stategraph_compile(handler) return handler From 9e21f08498223b2da104853521c1f562ad6c9029 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:46:06 +0800 Subject: [PATCH 22/37] =?UTF-8?q?=E2=9C=A8=20(runtime):=20Add=20runtime=20?= =?UTF-8?q?reset=20helper=20for=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/runtime.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/agent_assembly/adapters/langchain/runtime.py b/agent_assembly/adapters/langchain/runtime.py index 15830f1..81c2e1a 100644 --- a/agent_assembly/adapters/langchain/runtime.py +++ b/agent_assembly/adapters/langchain/runtime.py @@ -30,3 +30,10 @@ def auto_inject_callback_handler(interceptor: Any) -> AssemblyCallbackHandler: def get_active_callback_handler() -> AssemblyCallbackHandler | None: """Return the current callback handler instance when one is registered.""" return _ACTIVE_CALLBACK_HANDLER + + +def _reset_runtime_state_for_tests() -> None: + global _ACTIVE_CALLBACK_HANDLER + + with _RUNTIME_LOCK: + _ACTIVE_CALLBACK_HANDLER = None From 1297f435befa942e7a45a38ecab671a4f50d6c62 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:46:15 +0800 Subject: [PATCH 23/37] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(adapters):=20Export?= =?UTF-8?q?=20LangChain=20adapter=20runtime=20symbols?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/__init__.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/agent_assembly/adapters/langchain/__init__.py b/agent_assembly/adapters/langchain/__init__.py index 3db47b0..b2a365f 100644 --- a/agent_assembly/adapters/langchain/__init__.py +++ b/agent_assembly/adapters/langchain/__init__.py @@ -1 +1,15 @@ """LangChain adapter package.""" + +from agent_assembly.adapters.langchain.callback_handler import AssemblyCallbackHandler +from agent_assembly.adapters.langchain.langgraph_patch import patch_stategraph_compile +from agent_assembly.adapters.langchain.runtime import ( + auto_inject_callback_handler, + get_active_callback_handler, +) + +__all__ = [ + "AssemblyCallbackHandler", + "patch_stategraph_compile", + "auto_inject_callback_handler", + "get_active_callback_handler", +] From 9356821c954a3e6a65bbcb14f1976783f1f9d472 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:46:47 +0800 Subject: [PATCH 24/37] =?UTF-8?q?=E2=9C=85=20(tests):=20Add=20unit=20tests?= =?UTF-8?q?=20for=20callback=20sync=20methods?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langchain/test_callback_handler_sync.py | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 test/unit/adapters/langchain/test_callback_handler_sync.py diff --git a/test/unit/adapters/langchain/test_callback_handler_sync.py b/test/unit/adapters/langchain/test_callback_handler_sync.py new file mode 100644 index 0000000..f8bfd79 --- /dev/null +++ b/test/unit/adapters/langchain/test_callback_handler_sync.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from uuid import uuid4 + +import pytest + +from agent_assembly.adapters.langchain import AssemblyCallbackHandler +from agent_assembly.exceptions import ToolExecutionBlockedError + + +class SyncInterceptor: + def __init__(self) -> None: + self.tool_end_calls = 0 + self.llm_scan_calls = 0 + self.llm_end_calls = 0 + self.pending_wait_calls = 0 + self.last_prompts: list[str] | None = None + + def check_tool_start(self, **kwargs: object) -> object: + return kwargs.get("decision", {"status": "allow"}) + + def wait_for_tool_approval(self, **kwargs: object) -> object: + self.pending_wait_calls += 1 + return kwargs.get("approval_decision", {"status": "allow"}) + + def on_tool_end(self, **kwargs: object) -> None: + self.tool_end_calls += 1 + + def on_llm_start_scan(self, **kwargs: object) -> None: + self.llm_scan_calls += 1 + prompts = kwargs.get("prompts") + if isinstance(prompts, list): + self.last_prompts = prompts + + def on_llm_end(self, **kwargs: object) -> None: + self.llm_end_calls += 1 + + +def test_on_tool_start_raises_when_governance_denies() -> None: + handler = AssemblyCallbackHandler(SyncInterceptor()) + + with pytest.raises(ToolExecutionBlockedError): + handler.on_tool_start( + serialized={"name": "web_search"}, + input_str="query", + run_id=uuid4(), + decision={"status": "deny", "reason": "blocked"}, + ) + + +def test_on_tool_start_waits_for_pending_approval() -> None: + interceptor = SyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + handler.on_tool_start( + serialized={"name": "calendar_write"}, + input_str="create event", + run_id=uuid4(), + decision={"status": "pending"}, + approval_decision={"status": "allow"}, + ) + + assert interceptor.pending_wait_calls == 1 + + +def test_on_tool_start_blocks_when_pending_never_approved() -> None: + handler = AssemblyCallbackHandler(SyncInterceptor()) + + with pytest.raises(ToolExecutionBlockedError): + handler.on_tool_start( + serialized={"name": "calendar_write"}, + input_str="create event", + run_id=uuid4(), + decision={"status": "pending"}, + approval_decision={"status": "deny"}, + ) + + +def test_on_tool_end_delegates_to_interceptor() -> None: + interceptor = SyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + handler.on_tool_end(output={"ok": True}, run_id=uuid4()) + + assert interceptor.tool_end_calls == 1 + + +def test_on_llm_start_scans_without_mutating_prompts() -> None: + interceptor = SyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + prompts = ["hello", "world"] + + handler.on_llm_start( + serialized={"name": "gpt"}, + prompts=prompts, + run_id=uuid4(), + ) + + assert interceptor.llm_scan_calls == 1 + assert interceptor.last_prompts is prompts + assert prompts == ["hello", "world"] + + +def test_on_llm_end_delegates_to_interceptor() -> None: + interceptor = SyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + handler.on_llm_end(response={"text": "done"}, run_id=uuid4()) + + assert interceptor.llm_end_calls == 1 From e14272fc90856afbafdf8db0b045081bad432c75 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:47:07 +0800 Subject: [PATCH 25/37] =?UTF-8?q?=E2=9C=85=20(tests):=20Add=20unit=20tests?= =?UTF-8?q?=20for=20callback=20async=20methods?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langchain/test_callback_handler_async.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) create mode 100644 test/unit/adapters/langchain/test_callback_handler_async.py diff --git a/test/unit/adapters/langchain/test_callback_handler_async.py b/test/unit/adapters/langchain/test_callback_handler_async.py new file mode 100644 index 0000000..89ffeef --- /dev/null +++ b/test/unit/adapters/langchain/test_callback_handler_async.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from uuid import uuid4 + +import pytest + +from agent_assembly.adapters.langchain import AssemblyCallbackHandler +from agent_assembly.exceptions import ToolExecutionBlockedError + + +class AsyncInterceptor: + def __init__(self) -> None: + self.tool_end_calls = 0 + self.llm_scan_calls = 0 + self.llm_end_calls = 0 + self.pending_wait_calls = 0 + + async def check_tool_start(self, **kwargs: object) -> object: + return kwargs.get("decision", {"status": "allow"}) + + async def wait_for_tool_approval(self, **kwargs: object) -> object: + self.pending_wait_calls += 1 + return kwargs.get("approval_decision", {"status": "allow"}) + + async def on_tool_end(self, **kwargs: object) -> None: + self.tool_end_calls += 1 + + async def on_llm_start_scan(self, **kwargs: object) -> None: + self.llm_scan_calls += 1 + + async def on_llm_end(self, **kwargs: object) -> None: + self.llm_end_calls += 1 + + +@pytest.mark.asyncio +async def test_aon_tool_start_raises_when_governance_denies() -> None: + handler = AssemblyCallbackHandler(AsyncInterceptor()) + + with pytest.raises(ToolExecutionBlockedError): + await handler.aon_tool_start( + serialized={"name": "web_search"}, + input_str="query", + run_id=uuid4(), + decision={"status": "deny", "reason": "blocked"}, + ) + + +@pytest.mark.asyncio +async def test_aon_tool_start_waits_for_pending_approval() -> None: + interceptor = AsyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + await handler.aon_tool_start( + serialized={"name": "calendar_write"}, + input_str="create event", + run_id=uuid4(), + decision={"status": "pending"}, + approval_decision={"status": "allow"}, + ) + + assert interceptor.pending_wait_calls == 1 + + +@pytest.mark.asyncio +async def test_aon_tool_end_delegates_to_interceptor() -> None: + interceptor = AsyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + await handler.aon_tool_end(output={"ok": True}, run_id=uuid4()) + + assert interceptor.tool_end_calls == 1 + + +@pytest.mark.asyncio +async def test_aon_llm_start_delegates_to_interceptor() -> None: + interceptor = AsyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + await handler.aon_llm_start( + serialized={"name": "gpt"}, + prompts=["hello", "world"], + run_id=uuid4(), + ) + + assert interceptor.llm_scan_calls == 1 + + +@pytest.mark.asyncio +async def test_aon_llm_end_delegates_to_interceptor() -> None: + interceptor = AsyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + await handler.aon_llm_end(response={"text": "done"}, run_id=uuid4()) + + assert interceptor.llm_end_calls == 1 From aaa1c7716a43f0b9df884a35dc1fcffb28f22bcb Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:47:24 +0800 Subject: [PATCH 26/37] =?UTF-8?q?=E2=9C=85=20(tests):=20Add=20unit=20tests?= =?UTF-8?q?=20for=20auto-injection=20and=20idempotency?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/unit/adapters/langchain/test_runtime.py | 42 ++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 test/unit/adapters/langchain/test_runtime.py diff --git a/test/unit/adapters/langchain/test_runtime.py b/test/unit/adapters/langchain/test_runtime.py new file mode 100644 index 0000000..51707a9 --- /dev/null +++ b/test/unit/adapters/langchain/test_runtime.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from agent_assembly import init_assembly +from agent_assembly.adapters.langchain.runtime import ( + _reset_runtime_state_for_tests, + auto_inject_callback_handler, + get_active_callback_handler, +) + + +def test_auto_inject_callback_handler_is_idempotent() -> None: + _reset_runtime_state_for_tests() + + first = auto_inject_callback_handler(interceptor=object()) + second = auto_inject_callback_handler(interceptor=object()) + + assert first is second + assert get_active_callback_handler() is first + + +def test_init_assembly_auto_injects_callback_handler() -> None: + _reset_runtime_state_for_tests() + + client = init_assembly(gateway_url="http://localhost:8080", agent_id="test-agent") + try: + assert get_active_callback_handler() is not None + finally: + client.close() + + +def test_init_assembly_reuses_existing_callback_handler() -> None: + _reset_runtime_state_for_tests() + + first_client = init_assembly(gateway_url="http://localhost:8080", agent_id="test-agent-a") + second_client = init_assembly(gateway_url="http://localhost:8080", agent_id="test-agent-b") + try: + first_handler = get_active_callback_handler() + assert first_handler is not None + assert get_active_callback_handler() is first_handler + finally: + first_client.close() + second_client.close() From 4c0b52e9968511aab62a62c19415980bdce3ae2b Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:47:43 +0800 Subject: [PATCH 27/37] =?UTF-8?q?=E2=9C=85=20(tests):=20Add=20integration?= =?UTF-8?q?=20test=20for=20LangGraph=20interception=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...test_langgraph_interception_integration.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 test/integration/test_langgraph_interception_integration.py diff --git a/test/integration/test_langgraph_interception_integration.py b/test/integration/test_langgraph_interception_integration.py new file mode 100644 index 0000000..239dcc3 --- /dev/null +++ b/test/integration/test_langgraph_interception_integration.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from agent_assembly.adapters.langchain import AssemblyCallbackHandler, patch_stategraph_compile + + +class GraphInterceptor: + def __init__(self) -> None: + self.events: list[str] = [] + + def on_graph_node_start(self, **kwargs: object) -> None: + self.events.append(f"start:{kwargs.get('node_name')}") + + def on_graph_node_end(self, **kwargs: object) -> None: + self.events.append(f"end:{kwargs.get('node_name')}") + + +@pytest.mark.integration +def test_langgraph_compile_patch_invokes_pre_post_hooks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + interceptor = GraphInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + class FakeCompiledGraph: + def invoke(self, state: dict[str, object]) -> dict[str, object]: + return {"ok": True, "input": state} + + class FakeStateGraph: + def compile(self) -> FakeCompiledGraph: + return FakeCompiledGraph() + + fake_module = SimpleNamespace(StateGraph=FakeStateGraph) + + def fake_import_module(module_name: str) -> object: + if module_name == "langgraph.graph.state": + return fake_module + raise ImportError(module_name) + + monkeypatch.setattr( + "agent_assembly.adapters.langchain.langgraph_patch.importlib.import_module", + fake_import_module, + ) + + patched = patch_stategraph_compile(handler) + assert patched is True + + compiled = FakeStateGraph().compile() + result = compiled.invoke({"step": "run"}) + + assert result == {"ok": True, "input": {"step": "run"}} + assert interceptor.events == ["start:graph.invoke", "end:graph.invoke"] From 475c91bd9fafaa90b98a723315bd38d7c2331e3b Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:48:08 +0800 Subject: [PATCH 28/37] =?UTF-8?q?=F0=9F=93=9D=20(docs):=20Document=20LangC?= =?UTF-8?q?hain=20and=20LangGraph=20interception=20behavior?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/contents/document/api-references/index.mdx | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/contents/document/api-references/index.mdx b/docs/contents/document/api-references/index.mdx index 23a4d8b..8ed37d7 100644 --- a/docs/contents/document/api-references/index.mdx +++ b/docs/contents/document/api-references/index.mdx @@ -60,5 +60,15 @@ from agent_assembly.exceptions import ( PolicyError, GatewayError, ConfigurationError, + ToolExecutionBlockedError, ) ``` + +## LangChain runtime interception + +`init_assembly(...)` auto-injects a LangChain callback handler when the SDK runtime starts. + +- Tool start checks support governance decisions: `allow`, `deny`, and `pending`. +- `deny` (or unresolved `pending`) raises `ToolExecutionBlockedError`. +- LLM start interception is scan-only and does not mutate prompt content. +- LangGraph `StateGraph.compile()` is patched to add pre/post invocation governance hooks. From 957f115a2756e2285c5b8eaf648c80e8a1cfdc62 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:49:24 +0800 Subject: [PATCH 29/37] =?UTF-8?q?=F0=9F=94=A7=20(typing):=20Make=20BaseCal?= =?UTF-8?q?lbackHandler=20import=20mypy-safe=20without=20extras?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 2be480c..cf2adf0 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -2,22 +2,31 @@ from __future__ import annotations +import importlib import inspect -from typing import Any, Literal, Mapping +from typing import Any, Literal, Mapping, cast from uuid import UUID from agent_assembly.exceptions import ToolExecutionBlockedError + +class _FallbackBaseCallbackHandler: + """Fallback base type when langchain-core is not installed.""" + + pass + + +_CallbackHandlerBase: type[object] = _FallbackBaseCallbackHandler try: # pragma: no cover - import availability depends on installed extras. - from langchain_core.callbacks import BaseCallbackHandler + callbacks_module = importlib.import_module("langchain_core.callbacks") + maybe_base = getattr(callbacks_module, "BaseCallbackHandler", _FallbackBaseCallbackHandler) + if isinstance(maybe_base, type): + _CallbackHandlerBase = cast(type[object], maybe_base) except ImportError: # pragma: no cover - fallback keeps runtime import-safe. - class BaseCallbackHandler: # type: ignore[no-redef] - """Fallback base type when langchain-core is not installed.""" - - pass + pass -class AssemblyCallbackHandler(BaseCallbackHandler): +class AssemblyCallbackHandler(_CallbackHandlerBase): """Callback handler that delegates runtime events to governance interception.""" def __init__(self, interceptor: Any) -> None: From 166f185a6f6cb0a29709c2e62003e7bc897bf4c8 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:49:38 +0800 Subject: [PATCH 30/37] =?UTF-8?q?=F0=9F=94=A7=20(typing):=20Tighten=20deci?= =?UTF-8?q?sion=20status=20literals=20for=20mypy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/callback_handler.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index cf2adf0..5ac91d9 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -38,15 +38,22 @@ def _normalize_decision( ) -> tuple[Literal["allow", "deny", "pending"], str | None]: if isinstance(decision, str): normalized = decision.strip().lower() - if normalized in {"allow", "deny", "pending"}: - return normalized, None + if normalized == "allow": + return "allow", None + if normalized == "deny": + return "deny", None + if normalized == "pending": + return "pending", None return "allow", None if isinstance(decision, Mapping): raw_status = str(decision.get("status", "allow")).strip().lower() - status: Literal["allow", "deny", "pending"] - if raw_status in {"allow", "deny", "pending"}: - status = raw_status + if raw_status == "allow": + status: Literal["allow", "deny", "pending"] = "allow" + elif raw_status == "deny": + status = "deny" + elif raw_status == "pending": + status = "pending" else: status = "allow" From 1eed3b9a65afb68617830109e2ec459ec0a88850 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:50:03 +0800 Subject: [PATCH 31/37] =?UTF-8?q?=F0=9F=94=A7=20(typing):=20Silence=20dyna?= =?UTF-8?q?mic=20callback=20base=20mypy=20base-class=20checks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/adapters/langchain/callback_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent_assembly/adapters/langchain/callback_handler.py b/agent_assembly/adapters/langchain/callback_handler.py index 5ac91d9..751feec 100644 --- a/agent_assembly/adapters/langchain/callback_handler.py +++ b/agent_assembly/adapters/langchain/callback_handler.py @@ -26,7 +26,7 @@ class _FallbackBaseCallbackHandler: pass -class AssemblyCallbackHandler(_CallbackHandlerBase): +class AssemblyCallbackHandler(_CallbackHandlerBase): # type: ignore[valid-type,misc] """Callback handler that delegates runtime events to governance interception.""" def __init__(self, interceptor: Any) -> None: From ae6b118d3b6af40082cf755f87938bf56d15ddf0 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:50:22 +0800 Subject: [PATCH 32/37] =?UTF-8?q?=F0=9F=94=A7=20(tests):=20Register=20inte?= =?UTF-8?q?gration=20pytest=20marker?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pytest.ini | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pytest.ini b/pytest.ini index dfef326..e293d55 100644 --- a/pytest.ini +++ b/pytest.ini @@ -12,3 +12,5 @@ log_cli = 1 log_cli_level = INFO log_cli_format = %(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s) log_cli_date_format=%Y-%m-%d %H:%M:%S +markers = + integration: marks tests as integration tests From a2ce8674ca7a19b84e5c388583fa0fc39ff10a26 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 14:50:42 +0800 Subject: [PATCH 33/37] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(exports):=20Re-expo?= =?UTF-8?q?rt=20ToolExecutionBlockedError=20from=20package=20root?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent_assembly/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agent_assembly/__init__.py b/agent_assembly/__init__.py index eaf5edd..071f5cb 100644 --- a/agent_assembly/__init__.py +++ b/agent_assembly/__init__.py @@ -9,6 +9,7 @@ ConfigurationError, GatewayError, PolicyError, + ToolExecutionBlockedError, ) __version__ = "0.0.0" @@ -24,4 +25,5 @@ "GatewayError", "ConfigurationError", "AdapterValidationError", + "ToolExecutionBlockedError", ] From 5fb8dd706c9264d4f04c7befb89a1823ce6ed1ed Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 15:03:11 +0800 Subject: [PATCH 34/37] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20(langgraph):=20Wrap?= =?UTF-8?q?=20compiled=20node=20executors=20with=20per-node=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../adapters/langchain/langgraph_patch.py | 157 ++++++++++++++++-- 1 file changed, 144 insertions(+), 13 deletions(-) diff --git a/agent_assembly/adapters/langchain/langgraph_patch.py b/agent_assembly/adapters/langchain/langgraph_patch.py index 04e9531..80da5d8 100644 --- a/agent_assembly/adapters/langchain/langgraph_patch.py +++ b/agent_assembly/adapters/langchain/langgraph_patch.py @@ -8,32 +8,140 @@ _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 _invoke_pre_node_hook(callback_handler: Any, state: Any) -> None: +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="graph.invoke", state=state) + result = method(node_name=node_name, state=state) if inspect.isawaitable(result): return None return None -def _invoke_post_node_hook(callback_handler: Any, state: Any, result: Any) -> 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="graph.invoke", state=state, result=result) + 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: @@ -52,15 +160,38 @@ def patch_stategraph_compile(callback_handler: Any) -> bool: def patched_compile(self: Any, *args: Any, **kwargs: Any) -> Any: compiled_graph = original_compile(self, *args, **kwargs) - invoke = getattr(compiled_graph, "invoke", None) - if callable(invoke): - def wrapped_invoke(state: Any, *invoke_args: Any, **invoke_kwargs: Any) -> Any: - _invoke_pre_node_hook(callback_handler, state) - result = invoke(state, *invoke_args, **invoke_kwargs) - _invoke_post_node_hook(callback_handler, state, result) - return result - - setattr(compiled_graph, "invoke", wrapped_invoke) + 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) From ed7cfa5a52c9bd0fcd50d0aefe388d77cd516c70 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 15:03:47 +0800 Subject: [PATCH 35/37] =?UTF-8?q?=E2=9C=85=20(tests):=20Cover=20multi-node?= =?UTF-8?q?=20LangGraph=20tool-block=20integration=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...test_langgraph_interception_integration.py | 49 ++++++++++++++++--- 1 file changed, 43 insertions(+), 6 deletions(-) diff --git a/test/integration/test_langgraph_interception_integration.py b/test/integration/test_langgraph_interception_integration.py index 239dcc3..606af97 100644 --- a/test/integration/test_langgraph_interception_integration.py +++ b/test/integration/test_langgraph_interception_integration.py @@ -1,16 +1,24 @@ from __future__ import annotations from types import SimpleNamespace +from uuid import uuid4 import pytest from agent_assembly.adapters.langchain import AssemblyCallbackHandler, patch_stategraph_compile +from agent_assembly.exceptions import ToolExecutionBlockedError class GraphInterceptor: def __init__(self) -> None: self.events: list[str] = [] + def check_tool_start(self, **kwargs: object) -> dict[str, str]: + serialized = kwargs.get("serialized") + if isinstance(serialized, dict) and serialized.get("name") == "blocked_tool": + return {"status": "deny", "reason": "blocked by policy"} + return {"status": "allow"} + def on_graph_node_start(self, **kwargs: object) -> None: self.events.append(f"start:{kwargs.get('node_name')}") @@ -19,15 +27,40 @@ def on_graph_node_end(self, **kwargs: object) -> None: @pytest.mark.integration -def test_langgraph_compile_patch_invokes_pre_post_hooks( +def test_langgraph_compile_patch_wraps_multi_node_graph_and_blocks_denied_tool( monkeypatch: pytest.MonkeyPatch, ) -> None: interceptor = GraphInterceptor() handler = AssemblyCallbackHandler(interceptor) class FakeCompiledGraph: + def __init__(self) -> None: + self.nodes = { + "node_a": self._node_a, + "node_b": self._node_b, + } + + def _node_a(self, state: dict[str, object]) -> dict[str, object]: + handler.on_tool_start( + serialized={"name": "safe_tool"}, + input_str="{}", + run_id=uuid4(), + ) + return {**state, "node_a": "ok"} + + def _node_b(self, state: dict[str, object]) -> dict[str, object]: + handler.on_tool_start( + serialized={"name": "blocked_tool"}, + input_str="{}", + run_id=uuid4(), + ) + return {**state, "node_b": "should-not-complete"} + def invoke(self, state: dict[str, object]) -> dict[str, object]: - return {"ok": True, "input": state} + current_state = state + for node_name in ("node_a", "node_b"): + current_state = self.nodes[node_name](current_state) + return current_state class FakeStateGraph: def compile(self) -> FakeCompiledGraph: @@ -49,7 +82,11 @@ def fake_import_module(module_name: str) -> object: assert patched is True compiled = FakeStateGraph().compile() - result = compiled.invoke({"step": "run"}) - - assert result == {"ok": True, "input": {"step": "run"}} - assert interceptor.events == ["start:graph.invoke", "end:graph.invoke"] + with pytest.raises(ToolExecutionBlockedError): + compiled.invoke({"step": "run"}) + + assert interceptor.events == [ + "start:node_a", + "end:node_a", + "start:node_b", + ] From 4fbd1530ec0ea9f232f5c9ea94a05ba3ab902a0f Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 15:04:07 +0800 Subject: [PATCH 36/37] =?UTF-8?q?=E2=9C=85=20(tests):=20Add=20explicit=20t?= =?UTF-8?q?ool-allow=20unit=20coverage=20for=20callback=20handler?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langchain/test_callback_handler_sync.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/test/unit/adapters/langchain/test_callback_handler_sync.py b/test/unit/adapters/langchain/test_callback_handler_sync.py index f8bfd79..d42a782 100644 --- a/test/unit/adapters/langchain/test_callback_handler_sync.py +++ b/test/unit/adapters/langchain/test_callback_handler_sync.py @@ -48,6 +48,20 @@ def test_on_tool_start_raises_when_governance_denies() -> None: ) +def test_on_tool_start_allows_when_governance_allows() -> None: + interceptor = SyncInterceptor() + handler = AssemblyCallbackHandler(interceptor) + + handler.on_tool_start( + serialized={"name": "web_search"}, + input_str="query", + run_id=uuid4(), + decision={"status": "allow"}, + ) + + assert interceptor.pending_wait_calls == 0 + + def test_on_tool_start_waits_for_pending_approval() -> None: interceptor = SyncInterceptor() handler = AssemblyCallbackHandler(interceptor) From 318663bb7631a35dfebf01d9bd9977d72d410aa1 Mon Sep 17 00:00:00 2001 From: Chisanan232 Date: Mon, 27 Apr 2026 22:27:36 +0800 Subject: [PATCH 37/37] =?UTF-8?q?=E2=9C=85=20(tests):=20Add=20branch=20cov?= =?UTF-8?q?erage=20for=20LangGraph=20patch=20fallback=20paths?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langchain/test_langgraph_patch.py | 237 ++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 test/unit/adapters/langchain/test_langgraph_patch.py diff --git a/test/unit/adapters/langchain/test_langgraph_patch.py b/test/unit/adapters/langchain/test_langgraph_patch.py new file mode 100644 index 0000000..2c735e6 --- /dev/null +++ b/test/unit/adapters/langchain/test_langgraph_patch.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from agent_assembly.adapters.langchain import langgraph_patch + + +class _TestAwaitable: + def __await__(self): # type: ignore[no-untyped-def] + if False: + yield None + return None + + +class GraphEventRecorder: + def __init__(self, await_callbacks: bool = False) -> None: + self.await_callbacks = await_callbacks + self.events: list[tuple[str, str, object]] = [] + + def on_graph_node_start(self, *, node_name: str, state: object) -> object: + self.events.append(("start", node_name, state)) + if self.await_callbacks: + return _TestAwaitable() + return None + + def on_graph_node_end(self, *, node_name: str, state: object, result: object) -> object: + self.events.append(("end", node_name, result)) + if self.await_callbacks: + return _TestAwaitable() + return None + + +def test_extract_state_prefers_args_and_falls_back_to_kwargs() -> None: + assert langgraph_patch._extract_state(({"from": "args"},), {"state": {"from": "kwargs"}}) == { + "from": "args" + } + assert langgraph_patch._extract_state((), {"state": {"from": "kwargs"}}) == { + "from": "kwargs" + } + + +def test_invoke_hooks_handle_missing_methods_and_awaitables() -> None: + # Missing hook methods should no-op. + langgraph_patch._invoke_pre_node_hook(object(), "n1", {"state": 1}) + langgraph_patch._invoke_post_node_hook(object(), "n1", {"state": 1}, {"result": 1}) + + recorder = GraphEventRecorder(await_callbacks=True) + langgraph_patch._invoke_pre_node_hook(recorder, "n2", {"state": 2}) + langgraph_patch._invoke_post_node_hook(recorder, "n2", {"state": 2}, {"result": 2}) + + assert recorder.events == [ + ("start", "n2", {"state": 2}), + ("end", "n2", {"result": 2}), + ] + + +@pytest.mark.asyncio +async def test_wrap_node_callable_handles_already_wrapped_and_async_results() -> None: + recorder = GraphEventRecorder() + + async def async_node(state: dict[str, object]) -> dict[str, object]: + return {"ok": state} + + wrapped_async = langgraph_patch._wrap_node_callable("async_node", async_node, recorder) + async_result = wrapped_async({"v": 1}) + assert async_result is not None + assert await async_result == {"ok": {"v": 1}} + + def already_wrapped(state: dict[str, object]) -> dict[str, object]: + return {"state": state} + + setattr(already_wrapped, langgraph_patch._NODE_WRAPPED_FLAG, True) + assert langgraph_patch._wrap_node_callable("wrapped", already_wrapped, recorder) is already_wrapped + + assert recorder.events == [ + ("start", "async_node", {"v": 1}), + ("end", "async_node", {"ok": {"v": 1}}), + ] + + +def test_wrap_node_map_covers_non_mapping_and_assignment_failure() -> None: + assert langgraph_patch._wrap_node_map(object(), GraphEventRecorder()) is False + + class FailingNodeMap: + def __init__(self) -> None: + self._items = {"node": lambda state: state} + + def items(self) -> Any: + return self._items.items() + + def __setitem__(self, key: object, value: object) -> None: + del key, value + raise RuntimeError("cannot assign") + + assert langgraph_patch._wrap_node_map(FailingNodeMap(), GraphEventRecorder()) is False + + +@pytest.mark.asyncio +async def test_wrap_node_map_wraps_invoke_and_ainvoke_members() -> None: + recorder = GraphEventRecorder() + + class InvokeNode: + def invoke(self, state: dict[str, object]) -> dict[str, object]: + return {"invoke": state} + + async def ainvoke(self, state: dict[str, object]) -> dict[str, object]: + return {"ainvoke": state} + + node = InvokeNode() + node_map: dict[str, object] = {"node": node} + assert langgraph_patch._wrap_node_map(node_map, recorder) is True + + invoke_result = node.invoke({"v": 1}) + ainvoke_result = await node.ainvoke({"v": 2}) + + assert invoke_result == {"invoke": {"v": 1}} + assert ainvoke_result == {"ainvoke": {"v": 2}} + assert recorder.events == [ + ("start", "node", {"v": 1}), + ("end", "node", {"invoke": {"v": 1}}), + ("start", "node", {"v": 2}), + ("end", "node", {"ainvoke": {"v": 2}}), + ] + + +def test_wrap_compiled_graph_nodes_supports_pregel_fallback() -> None: + recorder = GraphEventRecorder() + + class FakeCompiledGraph: + def __init__(self) -> None: + self.nodes = None + self._nodes = None + self.pregel = None + self._pregel = SimpleNamespace( + nodes={"pregel_node": lambda state: {"pregel": state}}, + _nodes=None, + ) + + compiled_graph = FakeCompiledGraph() + assert langgraph_patch._wrap_compiled_graph_nodes(compiled_graph, recorder) is True + + result = compiled_graph._pregel.nodes["pregel_node"]({"k": 1}) + assert result == {"pregel": {"k": 1}} + assert recorder.events == [ + ("start", "pregel_node", {"k": 1}), + ("end", "pregel_node", {"pregel": {"k": 1}}), + ] + + +def test_patch_stategraph_compile_handles_import_and_stategraph_edge_cases( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "agent_assembly.adapters.langchain.langgraph_patch.importlib.import_module", + lambda name: (_ for _ in ()).throw(ImportError(name)), + ) + assert langgraph_patch.patch_stategraph_compile(GraphEventRecorder()) is False + + monkeypatch.setattr( + "agent_assembly.adapters.langchain.langgraph_patch.importlib.import_module", + lambda name: SimpleNamespace(), + ) + assert langgraph_patch.patch_stategraph_compile(GraphEventRecorder()) is False + + class AlreadyPatchedStateGraph: + compile = lambda self: object() + + setattr(AlreadyPatchedStateGraph, langgraph_patch._PATCHED_FLAG, True) + monkeypatch.setattr( + "agent_assembly.adapters.langchain.langgraph_patch.importlib.import_module", + lambda name: SimpleNamespace(StateGraph=AlreadyPatchedStateGraph), + ) + assert langgraph_patch.patch_stategraph_compile(GraphEventRecorder()) is True + + +def test_patch_stategraph_compile_fallback_wraps_sync_invoke( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder = GraphEventRecorder() + + class FallbackCompiledGraph: + def invoke(self, *, state: dict[str, object]) -> dict[str, object]: + return {"invoke": state} + + class StateGraphWithFallbackInvoke: + def compile(self) -> FallbackCompiledGraph: + return FallbackCompiledGraph() + + fake_module = SimpleNamespace(StateGraph=StateGraphWithFallbackInvoke) + monkeypatch.setattr( + "agent_assembly.adapters.langchain.langgraph_patch.importlib.import_module", + lambda name: fake_module, + ) + + assert langgraph_patch.patch_stategraph_compile(recorder) is True + compiled_graph = StateGraphWithFallbackInvoke().compile() + result = compiled_graph.invoke(state={"n": 1}) + + assert result == {"invoke": {"n": 1}} + assert recorder.events == [ + ("start", "graph.invoke", {"n": 1}), + ("end", "graph.invoke", {"invoke": {"n": 1}}), + ] + + +@pytest.mark.asyncio +async def test_patch_stategraph_compile_fallback_wraps_async_invoke( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder = GraphEventRecorder() + + class AsyncFallbackCompiledGraph: + async def invoke(self, state: dict[str, object]) -> dict[str, object]: + return {"invoke": state} + + class StateGraphWithAsyncFallbackInvoke: + def compile(self) -> AsyncFallbackCompiledGraph: + return AsyncFallbackCompiledGraph() + + fake_module = SimpleNamespace(StateGraph=StateGraphWithAsyncFallbackInvoke) + monkeypatch.setattr( + "agent_assembly.adapters.langchain.langgraph_patch.importlib.import_module", + lambda name: fake_module, + ) + + assert langgraph_patch.patch_stategraph_compile(recorder) is True + compiled_graph = StateGraphWithAsyncFallbackInvoke().compile() + result = await compiled_graph.invoke({"n": 2}) + + assert result == {"invoke": {"n": 2}} + assert recorder.events == [ + ("start", "graph.invoke", {"n": 2}), + ("end", "graph.invoke", {"invoke": {"n": 2}}), + ]