[AAASM-4831] 🐛 (python): Harden langchain handler singleton + haystack positional-arg passthrough

agent_assembly/adapters/haystack/patch.py

@@ -232,8 +232,12 @@ def _apply_tool_invoke_patch(tool_cls: type[Any], callback_handler: Any) -> None
     enforce = _interceptor_enforces(callback_handler)
 
     @wraps(original_invoke)
-    def patched_invoke(self: Any, **kwargs: Any) -> Any:
+    def patched_invoke(self: Any, *args: Any, **kwargs: Any) -> Any:
         tool_name = str(getattr(self, "name", self.__class__.__name__))
+        # Haystack's ToolInvoker calls ``invoke(**final_args)`` (keyword-only), but
+        # a direct ``Tool.invoke(x)`` positional call must be governed and forwarded
+        # too rather than raising TypeError; positional args are opaque to the
+        # governance check, so only keyword args populate the inspected tool_args.
         tool_args = dict(kwargs)
         decision = _invoke_tool_check(
             callback_handler,
@@ -258,7 +262,7 @@ def patched_invoke(self: Any, **kwargs: Any) -> Any:
                 return _format_approval_rejected_message(reason)
             return _format_blocked_message(reason)
 
-        result = original_invoke(self, **kwargs)
+        result = original_invoke(self, *args, **kwargs)
         _record_tool_result(callback_handler, tool_name=tool_name, result=result)
         return result
 

agent_assembly/adapters/langchain/runtime.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from threading import Lock
 from typing import Any
 
@@ -17,15 +18,34 @@ def auto_inject_callback_handler(
     *,
     process_agent_id: str | None = None,
 ) -> AssemblyCallbackHandler:
-    """Create and register the active callback handler instance."""
+    """Create and register the active callback handler instance.
+
+    The active handler is a module-level singleton. Re-injecting with the *same*
+    interceptor returns the existing handler unchanged (the matching-params
+    re-init path, which ``init_assembly`` guards today). If a *different*
+    interceptor is injected while one is already active, the stale handler is
+    replaced rather than silently kept: today ``init_assembly`` blocks a
+    mismatched re-init before reaching here, so this is defensive hardening for
+    future refactors that might inject a differently-postured interceptor
+    without a prior ``shutdown()`` clearing the singleton (AAASM-4831).
+    """
     global _ACTIVE_CALLBACK_HANDLER
 
     with _RUNTIME_LOCK:
         if process_agent_id is not None:
             set_process_agent_id(process_agent_id)
 
-        if _ACTIVE_CALLBACK_HANDLER is not None:
-            return _ACTIVE_CALLBACK_HANDLER
+        active = _ACTIVE_CALLBACK_HANDLER
+        if active is not None:
+            if active._interceptor is interceptor:
+                return active
+            warnings.warn(
+                "auto_inject_callback_handler called with a different interceptor "
+                "while a callback handler is already active; replacing the stale "
+                "handler. Call the assembly context's shutdown() before re-initializing.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
 
         handler = AssemblyCallbackHandler(interceptor)
         _ACTIVE_CALLBACK_HANDLER = handler

test/unit/adapters/haystack/test_patch.py

@@ -142,6 +142,75 @@ def on_tool_end(self, *, output: object, **kwargs: object) -> None:
     assert observed == [result]
 
 
+def test_positional_arg_is_governed_and_forwarded(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A positional ``Tool.invoke(x)`` must be governed and forwarded, not TypeError."""
+    received: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
+
+    class FakeTool:
+        name = "positional_tool"
+
+        def invoke(self, *args: Any, **kwargs: Any) -> dict[str, object]:
+            received.append((args, kwargs))
+            return {"args": args, "kwargs": kwargs}
+
+    monkeypatch.setattr(
+        haystack_patch.importlib,
+        "import_module",
+        lambda name: (
+            SimpleNamespace(Tool=FakeTool) if name == "haystack.tools" else (_ for _ in ()).throw(ImportError(name))
+        ),
+    )
+
+    checked: list[str] = []
+
+    class RecordingInterceptor:
+        def check_tool_start(self, **kwargs: object) -> dict[str, str]:
+            checked.append(str(kwargs.get("tool_name")))
+            return {"status": "allow"}
+
+    patcher = haystack_patch.HaystackPatch(RecordingInterceptor())
+    assert patcher.apply() is True
+
+    result = FakeTool().invoke("hello")
+
+    assert result == {"args": ("hello",), "kwargs": {}}
+    assert received == [(("hello",), {})]
+    assert checked == ["positional_tool"]  # governance fired on the positional call
+
+
+def test_positional_arg_denied_blocks_function(monkeypatch: pytest.MonkeyPatch) -> None:
+    ran: list[bool] = []
+
+    class FakeTool:
+        name = "positional_danger"
+
+        def invoke(self, *args: Any, **kwargs: Any) -> object:
+            ran.append(True)
+            return {"args": args, "kwargs": kwargs}
+
+    monkeypatch.setattr(
+        haystack_patch.importlib,
+        "import_module",
+        lambda name: (
+            SimpleNamespace(Tool=FakeTool) if name == "haystack.tools" else (_ for _ in ()).throw(ImportError(name))
+        ),
+    )
+
+    class BlockInterceptor:
+        def check_tool_start(self, **kwargs: object) -> dict[str, str]:
+            del kwargs
+            return {"status": "deny", "reason": "no positional tools"}
+
+    patcher = haystack_patch.HaystackPatch(BlockInterceptor())
+    assert patcher.apply() is True
+
+    result = FakeTool().invoke("payload")
+
+    assert isinstance(result, str)
+    assert "[BLOCKED by governance policy]" in result
+    assert ran == []  # the real tool function must NOT have executed
+
+
 def test_pending_tool_waits_and_allows_when_approved(monkeypatch: pytest.MonkeyPatch) -> None:
     FakeTool = _install_fake_haystack_module(monkeypatch)
     wait_calls: list[dict[str, object]] = []

test/unit/adapters/langchain/test_runtime.py

@@ -23,13 +23,28 @@ def test_auto_inject_callback_handler_is_idempotent() -> None:
     _reset_runtime_state_for_tests()
     _reset_assembly_state()
 
-    first = auto_inject_callback_handler(interceptor=object())
-    second = auto_inject_callback_handler(interceptor=object())
+    interceptor = object()
+    first = auto_inject_callback_handler(interceptor=interceptor)
+    second = auto_inject_callback_handler(interceptor=interceptor)
 
     assert first is second
     assert get_active_callback_handler() is first
 
 
+def test_reinject_different_interceptor_warns_and_replaces() -> None:
+    _reset_runtime_state_for_tests()
+    _reset_assembly_state()
+
+    first = auto_inject_callback_handler(interceptor=object())
+
+    with pytest.warns(RuntimeWarning):
+        second = auto_inject_callback_handler(interceptor=object())
+
+    # The stale handler must not be silently kept when the config differs.
+    assert second is not first
+    assert get_active_callback_handler() is second
+
+
 def test_init_assembly_auto_injects_callback_handler(monkeypatch: pytest.MonkeyPatch) -> None:
     _reset_runtime_state_for_tests()
     _reset_assembly_state()