diff --git a/agent_assembly/adapters/haystack/patch.py b/agent_assembly/adapters/haystack/patch.py index af39089e..b1e88944 100644 --- a/agent_assembly/adapters/haystack/patch.py +++ b/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 diff --git a/agent_assembly/adapters/langchain/runtime.py b/agent_assembly/adapters/langchain/runtime.py index 6b0fa0c4..1981471e 100644 --- a/agent_assembly/adapters/langchain/runtime.py +++ b/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 diff --git a/test/unit/adapters/haystack/test_patch.py b/test/unit/adapters/haystack/test_patch.py index 7e626bfb..f9f17d7d 100644 --- a/test/unit/adapters/haystack/test_patch.py +++ b/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]] = [] diff --git a/test/unit/adapters/langchain/test_runtime.py b/test/unit/adapters/langchain/test_runtime.py index eef381aa..41e54bd0 100644 --- a/test/unit/adapters/langchain/test_runtime.py +++ b/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()