[AAASM-6445] 🐛 (python-sdk): Govern LangChain tool calls from init_assembly()

README.md

@@ -177,10 +177,18 @@ with init_assembly(
 
 What this does:
 
-1. `init_assembly()` registers the agent with the gateway and auto-loads the LangChain adapter — each tool call the wrapped agent makes from now on goes through the policy gate.
+1. `init_assembly()` registers the agent with the gateway and auto-loads the LangChain adapter. The adapter registers one process-wide callback handler through `langchain_core`'s public `register_configure_hook`, so each tool call a LangChain run makes from now on — in any thread or async task — goes through the policy gate. Under `enforce` a deny raises `ToolExecutionBlockedError` before the tool body runs, and a runtime or handler failure blocks the tool too. You do not pass `callbacks` yourself.
 2. The `FakeListLLM` replays canned responses so the example runs **offline** with no real LLM.
 3. The `with` block tears down the gateway connection and unwinds adapter hooks on exit.
 
+> **Bounds of the LangChain gate.** It is a LangChain *callback*, so it is exactly as strong as that path:
+>
+> - A tool body called directly (`tool.func(...)`, `tool._run(...)`) never goes through LangChain's callback manager and is **not** checked.
+> - It covers `langchain_core` runs only. Other frameworks are governed by their own adapters, not by this hook.
+> - It governs tool calls only. LLM-side hooks (`on_llm_start` / `on_llm_end`, prompt scanning) are **not** enabled by `init_assembly()`; they run only for an `AssemblyCallbackHandler` you pass in `callbacks` yourself.
+> - It stops governing after `shutdown()` / the end of the `with` block, and it is **never installed** if the LangChain adapter fails to register or `langchain-core` is not importable. A failed adapter registration only writes a `framework adapter 'langchain' failed to register governance hooks` warning to stderr and `init_assembly()` still returns, even under `enforce`; the tool then runs ungoverned. `ctx.registered` reports the *gateway* registration, not this. Check that the warning is absent and that `any(a.get_framework_name() == "langchain" for a in ctx.adapters)` is true.
+> - A handler you pass yourself is neither checked nor ended twice (one policy check and one `on_tool_end` per tool run) when it wraps the active interceptor, and only after it allowed the run. A handler built around a different interceptor is checked in addition, never instead.
+
 ## Public API
 
 - `init_assembly(gateway_url, api_key, agent_id=None, mode="auto", *, control_plane_url=None, allow_insecure=False) -> AssemblyContext`

agent_assembly/adapters/langchain/callback_handler.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import importlib
+from collections import OrderedDict
 from collections.abc import Mapping
+from threading import Lock
 from typing import Any, Literal, NoReturn, cast
 from uuid import UUID
 
@@ -18,6 +20,46 @@
 
 _KNOWN_STATUSES: frozenset[str] = frozenset({"allow", "deny", "pending"})
 
+# run_id -> interceptor that already allowed that tool run, so the process-wide
+# configure-hook proxy (configure_hook.py) can skip a run an explicitly passed handler
+# already checked *with the active interceptor*, instead of checking it twice. Keyed by
+# interceptor so a handler built around a weaker or different one never suppresses the
+# active governance. Bounded: entries are consumed by the proxy and the oldest are
+# dropped if no proxy ever runs.
+_CHECKED_RUNS: OrderedDict[UUID, object] = OrderedDict()
+# Runs the proxy skipped at start; it must skip their on_tool_end too, or the wrapped
+# interceptor would see two ends (duplicate audit / budget side effects) for one run.
+_END_SKIPPED_RUNS: OrderedDict[UUID, None] = OrderedDict()
+_CHECKED_RUNS_LOCK = Lock()
+_CHECKED_RUNS_MAX = 4096
+# Returned on a registry miss. A unique object, so a miss can never compare equal to an
+# interceptor (``None`` would, for a handler built around ``interceptor=None``).
+CHECKED_RUN_MISS: object = object()
+
+
+def _bounded_put(registry: OrderedDict[UUID, Any], run_id: UUID, value: Any) -> None:
+    registry[run_id] = value
+    while len(registry) > _CHECKED_RUNS_MAX:
+        registry.popitem(last=False)
+
+
+def consume_checked_run(run_id: UUID) -> object:
+    """Pop the interceptor that already allowed ``run_id``, or ``CHECKED_RUN_MISS``."""
+    with _CHECKED_RUNS_LOCK:
+        return _CHECKED_RUNS.pop(run_id, CHECKED_RUN_MISS)
+
+
+def mark_end_skipped(run_id: UUID) -> None:
+    """Record that the proxy skipped ``run_id``'s start, so it skips the matching end."""
+    with _CHECKED_RUNS_LOCK:
+        _bounded_put(_END_SKIPPED_RUNS, run_id, None)
+
+
+def consume_end_skipped(run_id: UUID) -> bool:
+    """Whether ``run_id``'s start was skipped; clears the mark."""
+    with _CHECKED_RUNS_LOCK:
+        return _END_SKIPPED_RUNS.pop(run_id, CHECKED_RUN_MISS) is not CHECKED_RUN_MISS
+
 
 class _FallbackBaseCallbackHandler:
     """Fallback base type when langchain-core is not installed."""
@@ -202,6 +244,8 @@ def on_tool_start(
                     approval_reason or reason or "Tool execution was not approved by governance.",
                 )
 
+        with _CHECKED_RUNS_LOCK:
+            _bounded_put(_CHECKED_RUNS, run_id, self._interceptor)
         return None
 
     @staticmethod

agent_assembly/adapters/langchain/configure_hook.py

@@ -0,0 +1,102 @@
+"""Process-wide LangChain callback registration through a public extension point.
+
+``langchain_core`` lets a library contribute a callback handler to every run's
+configuration with ``register_configure_hook(ContextVar, inheritable=True)``. This is
+how LangChain's own tracers are attached, and it is documented API, not patching.
+
+Two properties of the registration matter for governance:
+
+* The ``ContextVar`` carries a **default** (a process-wide proxy handler) instead of
+  being ``set()`` by ``init_assembly``. A value ``set()`` in one context is invisible to
+  a plain ``threading.Thread`` (threads start with an empty context), so a tool run
+  there would silently skip the gate. A default is seen from every thread and task.
+* The proxy skips a run only when an explicitly passed ``AssemblyCallbackHandler`` already
+  allowed it with the *active* interceptor, so a tool is not checked twice. A handler built
+  around a weaker or different interceptor never suppresses the active governance
+  (registering with ``handle_class`` would have: LangChain then skips the proxy for any
+  ``AssemblyCallbackHandler`` instance).
+
+The hook cannot be unregistered, so the proxy resolves the active handler on every
+call and does nothing once ``init_assembly``'s context is shut down.
+"""
+
+from __future__ import annotations
+
+import importlib
+from contextvars import ContextVar
+from threading import Lock
+from typing import Any
+
+from agent_assembly.adapters.langchain.callback_handler import (
+    AssemblyCallbackHandler,
+    consume_checked_run,
+    consume_end_skipped,
+    mark_end_skipped,
+)
+from agent_assembly.adapters.langchain.runtime import get_active_callback_handler
+
+
+class ProcessGovernanceHandler(AssemblyCallbackHandler):
+    """Forwards LangChain tool callbacks to the currently active governance handler.
+
+    ``raise_error`` stays True (inherited): LangChain swallows a callback's exception
+    otherwise, so any failure here, including a runtime/gateway error raised by the
+    interceptor, aborts the tool call instead of letting it run ungoverned.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(interceptor=None)
+
+    def on_tool_start(self, serialized: dict[str, Any], input_str: str, **kwargs: Any) -> None:
+        active = get_active_callback_handler()
+        if active is None:
+            return
+        run_id = kwargs.get("run_id")
+        if run_id is not None and consume_checked_run(run_id) is active._interceptor:
+            mark_end_skipped(run_id)
+            return
+        active.on_tool_start(serialized, input_str, **kwargs)
+
+    def on_tool_end(self, output: Any, **kwargs: Any) -> None:
+        run_id = kwargs.get("run_id")
+        if run_id is not None and consume_end_skipped(run_id):
+            return
+        active = get_active_callback_handler()
+        if active is not None:
+            active.on_tool_end(output, **kwargs)
+
+    def on_tool_error(self, _error: BaseException, **kwargs: Any) -> None:
+        # Only clears the per-run mark so a failed tool does not leak it.
+        run_id = kwargs.get("run_id")
+        if run_id is not None:
+            consume_end_skipped(run_id)
+
+
+_REGISTER_LOCK = Lock()
+_registered_var: ContextVar[Any] | None = None
+
+
+def ensure_configure_hook_registered() -> bool:
+    """Register the process-wide governance handler with LangChain once.
+
+    Returns ``True`` when the hook is in place and ``False`` when ``langchain_core``
+    is not installed (nothing to govern). Safe to call on every ``init_assembly``.
+    """
+    global _registered_var
+
+    with _REGISTER_LOCK:
+        if _registered_var is not None:
+            return True
+        try:
+            context = importlib.import_module("langchain_core.tracers.context")
+        except ImportError:
+            return False
+        # B039: the default is deliberately a shared mutable handler; a per-context
+        # ``set()`` would not reach plain threads (see module docstring).
+        var: ContextVar[Any] = ContextVar(
+            "agent_assembly_langchain_governance",
+            default=ProcessGovernanceHandler(),  # noqa: B039
+        )
+        context.register_configure_hook(var, True)
+        _registered_var = var
+        return True

agent_assembly/adapters/langchain/patch.py

@@ -6,6 +6,7 @@
 from typing import Any
 
 from agent_assembly.adapters.langchain.callback_handler import AssemblyCallbackHandler
+from agent_assembly.adapters.langchain.configure_hook import ensure_configure_hook_registered
 from agent_assembly.adapters.langchain.runtime import (
     _reset_runtime_state_for_tests,
     auto_inject_callback_handler,
@@ -25,6 +26,7 @@ def apply(self) -> bool:
             self.interceptor,
             process_agent_id=self.process_agent_id,
         )
+        ensure_configure_hook_registered()
         return True
 
     def revert(self) -> None:

docs/index.md

@@ -24,8 +24,9 @@ You keep writing agents the way you always have. The SDK is the seatbelt you cli
 
 The SDK does not replace your agent framework — it **intercepts** it. When you call
 `init_assembly()`, the SDK detects which supported framework is importable in your process
-and monkey-patches that framework's tool-call entry points so each call passes through a
-policy gate first. Today it ships adapters for **LangChain, LangGraph, CrewAI, OpenAI Agents,
+and hooks that framework's tool-call entry points so each call passes through a
+policy gate first (LangChain through its public `register_configure_hook` callback extension point,
+the other adapters by patching the framework). The LangChain gate is a callback, so a tool body called directly (`tool.func(...)`) is not checked, LLM-side hooks are not enabled by `init_assembly()`, and if the adapter fails to register (a stderr warning, init still succeeds) LangChain runs ungoverned — see the README Quick Start for the full list of bounds. Today it ships adapters for **LangChain, LangGraph, CrewAI, OpenAI Agents,
 Pydantic AI, Google ADK, and MCP servers**.
 
 ```mermaid