From bacd3a86b4b27baf4b858f51653591a9d4f26899 Mon Sep 17 00:00:00 2001 From: Uchebuzz Date: Thu, 25 Jun 2026 20:34:33 +0100 Subject: [PATCH 1/7] feat(sdk/python): add Anthropic SDK instrumentation Patches anthropic.Anthropic and anthropic.AsyncAnthropic so every messages.create call is captured as an Adrian PairedEvent and emitted through the hook registry, with the same lifecycle guarantees as the existing LangChain/LangGraph integration. - sdk/python/adrian/anthropic_handler.py: core patch module - _flatten_content / _flatten_anthropic_messages: normalise Anthropic message params (strings, TypedDict blocks, SDK objects) to ChatMessage - _extract_anthropic_tool_calls: pull ToolUseBlock records from response - _extract_anthropic_usage: map input/output tokens to TokenUsage - build_anthropic_llm_pair: assemble PairedEvent from request + response - _emit_pair / _schedule_emit: async and sync emission paths - patch_anthropic: idempotent monkey-patch; reads hooks/config at call time so shutdown + re-init cycles are handled correctly - anthropic_invocation / anthropic_invocation_sync: context managers to group multi-turn calls under a single invocation_id - sdk/python/adrian/__init__.py: wire auto-instrumentation - patch_anthropic() public function (manual opt-in path) - _auto_instrument_anthropic() called by init() when auto_instrument=True - sdk/python/pyproject.toml: add anthropic optional-dep group and to dev - sdk/python/tests/test_anthropic_handler.py: 61 unit tests covering all helper functions, PairedEvent assembly, emission, patching, and context managers; all passing locally - examples/python/anthropic_quickstart.py: minimal two-turn async example All new source files carry the LICENSE_HEADER.txt SPDX header per CONTRIBUTING.md. --- examples/python/anthropic_quickstart.py | 84 +++ sdk/python/adrian/__init__.py | 34 + sdk/python/adrian/anthropic_handler.py | 550 ++++++++++++++++ sdk/python/pyproject.toml | 4 + sdk/python/tests/test_anthropic_handler.py | 728 +++++++++++++++++++++ 5 files changed, 1400 insertions(+) create mode 100644 examples/python/anthropic_quickstart.py create mode 100644 sdk/python/adrian/anthropic_handler.py create mode 100644 sdk/python/tests/test_anthropic_handler.py diff --git a/examples/python/anthropic_quickstart.py b/examples/python/anthropic_quickstart.py new file mode 100644 index 0000000..7af2590 --- /dev/null +++ b/examples/python/anthropic_quickstart.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 SecureAgentics +# +# Licensed under the Apache Licence, Version 2.0 (the "Licence"). +# You may not use this file except in compliance with the Licence. +# A copy of the Licence is included at LICENSE in the repository root. +"""Minimal quickstart: monitor Anthropic API calls with Adrian. + +Run:: + + export ANTHROPIC_API_KEY="sk-ant-..." + export ADRIAN_API_KEY="..." # optional -- omit to collect locally only + python examples/python/anthropic_quickstart.py +""" + +from __future__ import annotations + +import asyncio +import os + +import anthropic +import adrian + +# ------------------------------------------------------------------ +# 1. Initialise Adrian. This auto-instruments Anthropic by default. +# ------------------------------------------------------------------ +adrian.init( + api_key=os.environ.get("ADRIAN_API_KEY", ""), + session_id="anthropic-quickstart-session", +) + +# ------------------------------------------------------------------ +# 2. Create an Anthropic client as normal. +# ------------------------------------------------------------------ +client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"]) + + +async def main() -> None: + print("Sending first request...") + + # ------------------------------------------------------------------ + # 3. Wrap related calls in an invocation context so Adrian groups them. + # ------------------------------------------------------------------ + async with adrian.anthropic_invocation(): + response = await client.messages.create( + model="claude-haiku-4-5-20251001", + max_tokens=256, + system="You are a concise assistant.", + messages=[{"role": "user", "content": "What is 2 + 2? Answer in one sentence."}], + ) + + text = next( + (block.text for block in response.content if hasattr(block, "text")), + "", + ) + print(f"Model says: {text}") + + # A second call in the same invocation -- same invocation_id in Adrian. + follow_up = await client.messages.create( + model="claude-haiku-4-5-20251001", + max_tokens=256, + system="You are a concise assistant.", + messages=[ + {"role": "user", "content": "What is 2 + 2? Answer in one sentence."}, + {"role": "assistant", "content": text}, + {"role": "user", "content": "Now multiply that result by 10."}, + ], + ) + + follow_text = next( + (block.text for block in follow_up.content if hasattr(block, "text")), + "", + ) + print(f"Follow-up: {follow_text}") + + # ------------------------------------------------------------------ + # 4. Always shut down Adrian cleanly to flush any pending events. + # ------------------------------------------------------------------ + await adrian.shutdown() + print("Done. Check your Adrian dashboard for the captured events.") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/sdk/python/adrian/__init__.py b/sdk/python/adrian/__init__.py index 2cb0d23..983d5d4 100644 --- a/sdk/python/adrian/__init__.py +++ b/sdk/python/adrian/__init__.py @@ -30,6 +30,7 @@ from pathlib import Path from typing import Any +from adrian.anthropic_handler import anthropic_invocation, anthropic_invocation_sync from adrian.config import ( AdrianConfig, OnAuditCallback, @@ -106,6 +107,9 @@ "__version__", "mcp_servers", "redact_text", + "patch_anthropic", + "anthropic_invocation", + "anthropic_invocation_sync", ] logger = logging.getLogger("adrian") @@ -346,6 +350,7 @@ def init( if auto_instrument: _auto_instrument_langchain() + _auto_instrument_anthropic() # MCP server tracking is independent of LangChain auto-instrumentation, # it observes a different library (langchain-mcp-adapters) and is the @@ -379,6 +384,26 @@ def shutdown() -> None: set_config(None) +def patch_anthropic() -> None: + """Apply Anthropic SDK instrumentation. + + Monkey-patches ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic`` so + that every ``messages.create`` call is captured as an Adrian ``PairedEvent``. + Called automatically by :func:`init` when ``auto_instrument=True``. + + Call explicitly only when ``auto_instrument=False``:: + + adrian.init(api_key="...", auto_instrument=False) + adrian.patch_anthropic() + """ + from adrian.anthropic_handler import patch_anthropic as _patch + + _patch( + hooks_getter=lambda: _hooks, + config_getter=lambda: get_config() if is_initialized() else None, + ) + + def get_handler() -> AdrianCallbackHandler | None: """Return the SDK's callback handler, or ``None`` if uninitialised. @@ -501,6 +526,15 @@ def patch_langchain() -> None: # ------------------------------------------------------------------ +def _auto_instrument_anthropic() -> None: + """Apply Anthropic SDK monkey-patches if the package is installed.""" + try: + patch_anthropic() + logger.debug("Anthropic auto-instrumentation applied") + except Exception: + logger.exception("Anthropic auto-instrumentation failed") + + def _auto_instrument_langchain() -> None: """Apply LangChain / LangGraph monkey-patches if the libraries are present.""" try: diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py new file mode 100644 index 0000000..f8daa51 --- /dev/null +++ b/sdk/python/adrian/anthropic_handler.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 SecureAgentics +# +# Licensed under the Apache Licence, Version 2.0 (the "Licence"). +# You may not use this file except in compliance with the Licence. +# A copy of the Licence is included at LICENSE in the repository root. +"""Anthropic SDK instrumentation for Adrian. + +Patches ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic`` so that every +``messages.create`` call is captured as an Adrian ``PairedEvent`` and emitted +through the hook registry. The patch is idempotent; calling +:func:`patch_anthropic` again after a shutdown / re-init only updates the +internal getters, it does not re-wrap the already-patched method. + +Usage without auto-instrumentation:: + + import anthropic + import adrian + + adrian.init(api_key="...", auto_instrument=False) + adrian.patch_anthropic() + + client = anthropic.AsyncAnthropic() + response = await client.messages.create(model="...", ...) + +To group multi-turn calls under a single invocation ID:: + + async with adrian.anthropic_invocation(): + r1 = await client.messages.create(...) + r2 = await client.messages.create(...) # same invocation_id as r1 +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Callable +from contextlib import asynccontextmanager, contextmanager +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any +from uuid import uuid4 + +from adrian.config import AdrianConfig +from adrian.context import get_invocation_id, set_invocation_id +from adrian.format.types import AgentContext, LlmPairData, PairedEvent +from adrian.hooks import HookRegistry +from adrian.types import ChatMessage, EventData, TokenUsage, ToolCallRecord + +if TYPE_CHECKING: + from contextvars import Token + +logger = logging.getLogger("adrian.anthropic") + +# Set once by patch_anthropic(); read at call time so shutdown + re-init works. +_hooks_getter: Callable[[], HookRegistry | None] | None = None +_config_getter: Callable[[], AdrianConfig | None] | None = None + + +# ------------------------------------------------------------------ +# Message format conversion +# ------------------------------------------------------------------ + + +def _flatten_content(content: str | list[Any]) -> str: + """Flatten Anthropic message content to a plain string. + + Anthropic messages carry either a plain string or a list of content + blocks (``TextBlockParam``, ``ToolUseBlockParam``, ``ToolResultBlockParam``, + and so on). Both forms are normalised to a plain string for + ``ChatMessage.content``. + + Args: + content: Anthropic message content -- a string or a block list. + + Returns: + Plain string representation. + """ + if isinstance(content, str): + return content + + if not isinstance(content, list): + return str(content) + + parts: list[str] = [] + + for block in content: + if hasattr(block, "type"): + # SDK typed objects (TextBlock, ToolUseBlock, ToolResultBlock, …) + btype = block.type + + if btype == "text": + parts.append(getattr(block, "text", "")) + elif btype == "tool_use": + name = getattr(block, "name", "unknown") + args = getattr(block, "input", {}) + parts.append(f"[tool_use: {name} args={args}]") + elif btype == "tool_result": + inner = getattr(block, "content", "") + parts.append(_flatten_content(inner)) + elif isinstance(block, dict): + btype = block.get("type", "") + + if btype == "text": + parts.append(str(block.get("text", ""))) + elif btype == "tool_use": + name = block.get("name", "unknown") + args = block.get("input", {}) + parts.append(f"[tool_use: {name} args={args}]") + elif btype == "tool_result": + inner = block.get("content", "") + parts.append(_flatten_content(inner)) + + return "\n".join(p for p in parts if p) + + +def _flatten_anthropic_messages( + messages: list[dict[str, Any]], + system: str | list[Any] | None, +) -> list[ChatMessage]: + """Convert Anthropic message params to a flat ``ChatMessage`` list. + + Prepends the system prompt (if any) as a ``"system"`` role entry, + then converts each user / assistant turn in order. + + Args: + messages: Anthropic ``messages`` parameter -- a list of dicts with + ``role`` and ``content`` keys. + system: Anthropic ``system`` parameter -- a string, a block list, or + ``None``. + + Returns: + Flat list of ``ChatMessage`` dicts compatible with the Adrian format. + """ + result: list[ChatMessage] = [] + + if system is not None: + result.append(ChatMessage(role="system", content=_flatten_content(system))) + + for msg in messages: + role = str(msg.get("role", "unknown")) + content = _flatten_content(msg.get("content", "")) + result.append(ChatMessage(role=role, content=content)) + + return result + + +def _extract_anthropic_tool_calls(content: list[Any]) -> list[ToolCallRecord]: + """Extract tool call records from an Anthropic response content list. + + Scans for ``ToolUseBlock`` SDK objects or ``tool_use`` dicts and converts + each to a ``ToolCallRecord``. + + Args: + content: ``Message.content`` from the Anthropic response. + + Returns: + List of ``ToolCallRecord`` dicts, empty when no tool calls are present. + """ + records: list[ToolCallRecord] = [] + + for block in content: + if hasattr(block, "type") and block.type == "tool_use": + args = getattr(block, "input", {}) + + if not isinstance(args, dict): + try: + args = dict(args) + except (TypeError, ValueError): + args = {} + + records.append( + ToolCallRecord( + id=str(getattr(block, "id", "")), + name=str(getattr(block, "name", "unknown")), + args=args, + ) + ) + elif isinstance(block, dict) and block.get("type") == "tool_use": + args = block.get("input", {}) + + if not isinstance(args, dict): + args = {} + + records.append( + ToolCallRecord( + id=str(block.get("id", "")), + name=str(block.get("name", "unknown")), + args=args, + ) + ) + + return records + + +def _extract_anthropic_usage(response: Any) -> TokenUsage | None: + """Extract token usage from an Anthropic ``Message`` response object. + + Args: + response: ``anthropic.types.Message`` or any object with a ``usage`` + attribute carrying ``input_tokens`` and ``output_tokens``. + + Returns: + ``TokenUsage`` TypedDict, or ``None`` if usage data is absent. + """ + usage = getattr(response, "usage", None) + + if usage is None: + return None + + input_tokens: int = getattr(usage, "input_tokens", 0) or 0 + output_tokens: int = getattr(usage, "output_tokens", 0) or 0 + + return TokenUsage( + prompt_tokens=input_tokens, + completion_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ) + + +def _extract_response_text(content: list[Any]) -> str: + """Extract plain text output from an Anthropic response content list. + + Args: + content: ``Message.content`` from the Anthropic response. + + Returns: + Concatenated text from all ``TextBlock`` entries, joined by newlines. + """ + parts: list[str] = [] + + for block in content: + if hasattr(block, "type") and block.type == "text": + parts.append(getattr(block, "text", "")) + elif isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + + return "\n".join(p for p in parts if p) + + +def _derive_agent_id(messages: list[ChatMessage]) -> str: + """Derive a stable agent identity from the system prompt. + + Without LangGraph checkpoint metadata, the system prompt is the primary + signal for agent identity. Returns ``"default"`` when no system message + is present. + + Args: + messages: Flattened message list, may contain a ``"system"`` entry. + + Returns: + Agent identifier string (at most 64 characters). + """ + for msg in messages: + if msg.get("role") == "system": + content = msg["content"].strip() + + if content: + return content[:64].replace("\n", " ") + + return "default" + + +# ------------------------------------------------------------------ +# PairedEvent assembly +# ------------------------------------------------------------------ + + +def build_anthropic_llm_pair( + *, + flat_messages: list[ChatMessage], + response: Any, + model: str, + session_id: str, + invocation_id: str, + run_id: str, +) -> PairedEvent: + """Assemble a ``PairedEvent`` from an Anthropic ``messages.create`` call. + + Args: + flat_messages: Converted input messages (includes system prompt at index 0 + when present). + response: Raw ``anthropic.types.Message`` response object. + model: Model identifier from the request parameters. + session_id: Adrian session identifier. + invocation_id: Invocation correlation ID. + run_id: Per-call unique identifier generated by the patch. + + Returns: + Assembled ``PairedEvent`` with ``pair_type="llm"``. + """ + system_prompt = "" + user_instruction = "" + + for msg in flat_messages: + if msg.get("role") == "system" and not system_prompt: + system_prompt = msg["content"] + + for msg in reversed(flat_messages): + if msg.get("role") == "user": + user_instruction = msg["content"] + break + + content: list[Any] = getattr(response, "content", []) + output_text = _extract_response_text(content) + tool_calls = _extract_anthropic_tool_calls(content) + usage = _extract_anthropic_usage(response) + + # Prefer the model identifier echoed by the server; fall back to the request param. + response_model: str = getattr(response, "model", "") or model + + return PairedEvent( + event_id=str(uuid4()), + invocation_id=invocation_id, + session_id=session_id, + run_id=run_id, + timestamp=datetime.now(UTC).isoformat(), + pair_type="llm", + agent=AgentContext( + agent_id=_derive_agent_id(flat_messages), + system_prompt=system_prompt, + user_instruction=user_instruction, + ), + parent=None, + data=LlmPairData( + model=response_model, + messages=flat_messages, + output=output_text, + tool_calls=tool_calls, + usage=usage, + ), + ) + + +# ------------------------------------------------------------------ +# Emission helpers +# ------------------------------------------------------------------ + + +async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: + """Assemble and emit a ``PairedEvent`` for a completed ``messages.create`` call. + + Reads hooks and config at call time so the correct state is used even if + :func:`~adrian.shutdown` and :func:`~adrian.init` have been called since + the patch was applied. + + Args: + response: Anthropic ``Message`` response object. + kwargs: Original ``messages.create`` keyword arguments. + """ + if _hooks_getter is None or _config_getter is None: + return + + hooks = _hooks_getter() + config = _config_getter() + + if hooks is None or config is None: + return + + try: + session_id = config.session_id + messages_param: list[dict[str, Any]] = list(kwargs.get("messages") or []) + system_param: str | list[Any] | None = kwargs.get("system") + model_param: str = str(kwargs.get("model", "unknown")) + + flat_messages = _flatten_anthropic_messages(messages_param, system_param) + invocation_id = get_invocation_id() or "no_invocation" + run_id = str(uuid4()) + + pair = build_anthropic_llm_pair( + flat_messages=flat_messages, + response=response, + model=model_param, + session_id=session_id, + invocation_id=invocation_id, + run_id=run_id, + ) + + await hooks.emit(pair) + + if config.on_event is not None: + from typing import cast + + result = config.on_event( + pair.pair_type, + cast(EventData, pair.data), + pair.run_id, + None, + pair.event_id, + ) + + if asyncio.iscoroutine(result): + await result + + except Exception: + logger.exception("Failed to emit Anthropic paired event") + + +def _schedule_emit(response: Any, kwargs: dict[str, Any]) -> None: + """Schedule event emission from a synchronous call site. + + When inside a running event loop, schedules a fire-and-forget task so + the sync caller is not blocked. When no loop is running, blocks until + emission completes so the event is not silently dropped. + + Args: + response: Anthropic ``Message`` response object. + kwargs: Original ``messages.create`` keyword arguments. + """ + coro = _emit_pair(response, kwargs) + + try: + loop = asyncio.get_running_loop() + loop.create_task(coro) + except RuntimeError: + try: + asyncio.run(coro) + except Exception: + logger.exception("Failed to emit Anthropic event (sync path)") + + +# ------------------------------------------------------------------ +# SDK patching +# ------------------------------------------------------------------ + + +def patch_anthropic( + hooks_getter: Callable[[], HookRegistry | None], + config_getter: Callable[[], AdrianConfig | None], +) -> None: + """Monkey-patch ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic``. + + Wraps ``messages.create`` on both the sync and async Anthropic resource + classes so every API call is captured as an Adrian ``PairedEvent``. + + The patch is idempotent: subsequent calls update the internal getters but + do not re-wrap the already-patched method. If the ``anthropic`` package is + not installed the call is a silent no-op. + + This function is called automatically by :func:`~adrian.init` when + ``auto_instrument=True`` (the default). + + Args: + hooks_getter: Zero-arg callable returning the current ``HookRegistry``, + or ``None`` when the SDK is not initialised. + config_getter: Zero-arg callable returning the current ``AdrianConfig``, + or ``None`` when the SDK is not initialised. + """ + global _hooks_getter, _config_getter # noqa: PLW0603 + + _hooks_getter = hooks_getter + _config_getter = config_getter + + try: + import anthropic + except ImportError: + logger.debug("anthropic package not installed; skipping Anthropic patching") + return + + # ---- sync Messages.create ---- + try: + sync_cls = anthropic.resources.Messages + + if not getattr(sync_cls, "_adrian_patched", False): + _original_sync = sync_cls.create + + def _patched_sync_create( + self: Any, *args: Any, **kwargs: Any # noqa: ANN401 + ) -> Any: # noqa: ANN401 + response = _original_sync(self, *args, **kwargs) + _schedule_emit(response, kwargs) + return response + + sync_cls.create = _patched_sync_create # type: ignore[method-assign] + sync_cls._adrian_patched = True # type: ignore[attr-defined] + logger.debug("Patched anthropic.resources.Messages.create") + except AttributeError: + logger.warning( + "Could not patch anthropic.resources.Messages; " + "the SDK structure may have changed" + ) + + # ---- async AsyncMessages.create ---- + try: + async_cls = anthropic.resources.AsyncMessages + + if not getattr(async_cls, "_adrian_patched", False): + _original_async = async_cls.create + + async def _patched_async_create( + self: Any, *args: Any, **kwargs: Any # noqa: ANN401 + ) -> Any: # noqa: ANN401 + response = await _original_async(self, *args, **kwargs) + await _emit_pair(response, kwargs) + return response + + async_cls.create = _patched_async_create # type: ignore[method-assign] + async_cls._adrian_patched = True # type: ignore[attr-defined] + logger.debug("Patched anthropic.resources.AsyncMessages.create") + except AttributeError: + logger.warning( + "Could not patch anthropic.resources.AsyncMessages; " + "the SDK structure may have changed" + ) + + +# ------------------------------------------------------------------ +# Invocation context managers +# ------------------------------------------------------------------ + + +@asynccontextmanager +async def anthropic_invocation(): # type: ignore[return] + """Group async Anthropic API calls under a single invocation ID. + + Sets the ``invocation_id`` context variable so all ``messages.create`` + calls within the block share the same ID, enabling multi-turn agent + conversations to be correlated in the Adrian dashboard. + + Usage:: + + async with adrian.anthropic_invocation(): + r1 = await client.messages.create(...) + r2 = await client.messages.create(...) # same invocation_id as r1 + """ + token: Token[str | None] = set_invocation_id(str(uuid4())) + + try: + yield + finally: + token.var.reset(token) + + +@contextmanager +def anthropic_invocation_sync(): # type: ignore[return] + """Group synchronous Anthropic API calls under a single invocation ID. + + The sync counterpart to :func:`anthropic_invocation`. + + Usage:: + + with adrian.anthropic_invocation_sync(): + r1 = client.messages.create(...) + r2 = client.messages.create(...) # same invocation_id as r1 + """ + token: Token[str | None] = set_invocation_id(str(uuid4())) + + try: + yield + finally: + token.var.reset(token) diff --git a/sdk/python/pyproject.toml b/sdk/python/pyproject.toml index 8ee52ba..82e40c1 100644 --- a/sdk/python/pyproject.toml +++ b/sdk/python/pyproject.toml @@ -41,6 +41,9 @@ dependencies = [ ] [project.optional-dependencies] +anthropic = [ + "anthropic>=0.40.0", +] dev = [ "pytest>=8.0.0", "pytest-asyncio>=0.24.0", @@ -51,6 +54,7 @@ dev = [ "langgraph==1.1.2", "langgraph-prebuilt==1.0.8", "langchain-mcp-adapters>=0.2.2", + "anthropic>=0.40.0", ] [project.urls] diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py new file mode 100644 index 0000000..483c077 --- /dev/null +++ b/sdk/python/tests/test_anthropic_handler.py @@ -0,0 +1,728 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 SecureAgentics +# +# Licensed under the Apache Licence, Version 2.0 (the "Licence"). +# You may not use this file except in compliance with the Licence. +# A copy of the Licence is included at LICENSE in the repository root. +"""Tests for the Anthropic SDK instrumentation (adrian.anthropic_handler).""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +import adrian.anthropic_handler as _ah +from adrian.anthropic_handler import ( + _derive_agent_id, + _emit_pair, + _extract_anthropic_tool_calls, + _extract_anthropic_usage, + _extract_response_text, + _flatten_anthropic_messages, + _flatten_content, + anthropic_invocation, + anthropic_invocation_sync, + build_anthropic_llm_pair, + patch_anthropic, +) +from adrian.config import AdrianConfig +from adrian.context import get_invocation_id, set_invocation_id +from adrian.format.types import LlmPairData, PairedEvent +from adrian.hooks import HookRegistry +from adrian.types import ChatMessage + + +# ------------------------------------------------------------------ +# Shared helpers +# ------------------------------------------------------------------ + + +class _Collector: + """Minimal EventHandler that accumulates paired events.""" + + def __init__(self) -> None: + self.events: list[PairedEvent] = [] + + async def on_paired_event(self, event: PairedEvent) -> None: + self.events.append(event) + + async def close(self) -> None: + return None + + +def _make_text_response( + *, + model: str = "claude-opus-4-6", + text: str = "Hello!", + input_tokens: int = 10, + output_tokens: int = 5, +) -> MagicMock: + """Build a minimal mock Anthropic Message response.""" + text_block = MagicMock() + text_block.type = "text" + text_block.text = text + + usage = MagicMock() + usage.input_tokens = input_tokens + usage.output_tokens = output_tokens + + response = MagicMock() + response.model = model + response.content = [text_block] + response.usage = usage + + return response + + +def _wired_hooks(config: AdrianConfig) -> tuple[HookRegistry, _Collector]: + """Return a HookRegistry + Collector pair and wire them into the handler.""" + collector = _Collector() + hooks = HookRegistry() + hooks.register(collector) + _ah._hooks_getter = lambda: hooks + _ah._config_getter = lambda: config + return hooks, collector + + +# ------------------------------------------------------------------ +# _flatten_content +# ------------------------------------------------------------------ + + +class TestFlattenContent: + def test_plain_string_passthrough(self) -> None: + assert _flatten_content("hello world") == "hello world" + + def test_non_list_non_str_coerced(self) -> None: + assert _flatten_content(42) == "42" # type: ignore[arg-type] + + def test_text_block_dict(self) -> None: + result = _flatten_content([{"type": "text", "text": "hi"}]) + assert result == "hi" + + def test_tool_use_dict(self) -> None: + blocks = [{"type": "tool_use", "name": "search", "input": {"q": "test"}}] + result = _flatten_content(blocks) + assert "tool_use: search" in result + + def test_tool_result_string_dict(self) -> None: + blocks = [{"type": "tool_result", "content": "42"}] + assert _flatten_content(blocks) == "42" + + def test_sdk_text_object(self) -> None: + block = MagicMock() + block.type = "text" + block.text = "SDK text" + assert _flatten_content([block]) == "SDK text" + + def test_sdk_tool_use_object(self) -> None: + block = MagicMock() + block.type = "tool_use" + block.name = "my_tool" + block.input = {"x": 1} + result = _flatten_content([block]) + assert "tool_use: my_tool" in result + + def test_sdk_tool_result_delegates_recursively(self) -> None: + inner = MagicMock() + inner.type = "text" + inner.text = "inner text" + outer = MagicMock() + outer.type = "tool_result" + outer.content = [inner] + result = _flatten_content([outer]) + assert "inner text" in result + + def test_mixed_blocks_joined_by_newline(self) -> None: + blocks = [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "second"}, + ] + result = _flatten_content(blocks) + assert result == "first\nsecond" + + def test_empty_list(self) -> None: + assert _flatten_content([]) == "" + + +# ------------------------------------------------------------------ +# _flatten_anthropic_messages +# ------------------------------------------------------------------ + + +class TestFlattenAnthropicMessages: + def test_no_system(self) -> None: + msgs = [{"role": "user", "content": "hi"}] + result = _flatten_anthropic_messages(msgs, None) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"] == "hi" + + def test_system_prepended_as_first_entry(self) -> None: + msgs = [{"role": "user", "content": "hello"}] + result = _flatten_anthropic_messages(msgs, "You are helpful.") + assert result[0]["role"] == "system" + assert result[0]["content"] == "You are helpful." + assert result[1]["role"] == "user" + + def test_system_as_block_list(self) -> None: + system = [{"type": "text", "text": "block system"}] + result = _flatten_anthropic_messages([], system) + assert result[0]["role"] == "system" + assert result[0]["content"] == "block system" + + def test_assistant_role_preserved(self) -> None: + msgs = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + result = _flatten_anthropic_messages(msgs, None) + assert result[-1]["role"] == "assistant" + + def test_multi_turn_order(self) -> None: + msgs = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"}, + {"role": "user", "content": "third"}, + ] + result = _flatten_anthropic_messages(msgs, "sys") + assert len(result) == 4 + assert result[0]["role"] == "system" + assert result[1]["content"] == "first" + assert result[3]["content"] == "third" + + def test_empty_messages_with_system(self) -> None: + result = _flatten_anthropic_messages([], "only system") + assert len(result) == 1 + assert result[0]["role"] == "system" + + +# ------------------------------------------------------------------ +# _extract_anthropic_tool_calls +# ------------------------------------------------------------------ + + +class TestExtractAnthropicToolCalls: + def test_empty_content(self) -> None: + assert _extract_anthropic_tool_calls([]) == [] + + def test_text_block_ignored(self) -> None: + block = MagicMock() + block.type = "text" + block.text = "hello" + assert _extract_anthropic_tool_calls([block]) == [] + + def test_sdk_tool_use_object(self) -> None: + block = MagicMock() + block.type = "tool_use" + block.id = "call_abc" + block.name = "get_weather" + block.input = {"city": "London"} + result = _extract_anthropic_tool_calls([block]) + assert len(result) == 1 + assert result[0]["id"] == "call_abc" + assert result[0]["name"] == "get_weather" + assert result[0]["args"] == {"city": "London"} + + def test_dict_tool_use(self) -> None: + block = {"type": "tool_use", "id": "c1", "name": "search", "input": {"q": "x"}} + result = _extract_anthropic_tool_calls([block]) + assert len(result) == 1 + assert result[0]["name"] == "search" + assert result[0]["id"] == "c1" + + def test_multiple_tool_calls(self) -> None: + def _make(name: str, id_: str) -> MagicMock: + b = MagicMock() + b.type = "tool_use" + b.id = id_ + b.name = name + b.input = {} + return b + + result = _extract_anthropic_tool_calls([_make("tool_a", "c1"), _make("tool_b", "c2")]) + assert len(result) == 2 + assert {r["name"] for r in result} == {"tool_a", "tool_b"} + + def test_non_dict_input_coerced(self) -> None: + block = MagicMock() + block.type = "tool_use" + block.id = "c1" + block.name = "t" + block.input = [("key", "val")] + result = _extract_anthropic_tool_calls([block]) + assert isinstance(result[0]["args"], dict) + + +# ------------------------------------------------------------------ +# _extract_anthropic_usage +# ------------------------------------------------------------------ + + +class TestExtractAnthropicUsage: + def test_none_when_usage_attribute_missing(self) -> None: + assert _extract_anthropic_usage(object()) is None + + def test_none_when_usage_is_none(self) -> None: + response = MagicMock() + response.usage = None + assert _extract_anthropic_usage(response) is None + + def test_extracts_tokens_correctly(self) -> None: + usage = MagicMock() + usage.input_tokens = 150 + usage.output_tokens = 30 + response = MagicMock() + response.usage = usage + result = _extract_anthropic_usage(response) + assert result is not None + assert result["prompt_tokens"] == 150 + assert result["completion_tokens"] == 30 + assert result["total_tokens"] == 180 + + def test_zero_tokens_handled(self) -> None: + usage = MagicMock() + usage.input_tokens = 0 + usage.output_tokens = 0 + response = MagicMock() + response.usage = usage + result = _extract_anthropic_usage(response) + assert result is not None + assert result["total_tokens"] == 0 + + +# ------------------------------------------------------------------ +# _extract_response_text +# ------------------------------------------------------------------ + + +class TestExtractResponseText: + def test_single_text_block(self) -> None: + block = MagicMock() + block.type = "text" + block.text = "The answer." + assert _extract_response_text([block]) == "The answer." + + def test_multiple_text_blocks_joined(self) -> None: + def _tb(text: str) -> MagicMock: + b = MagicMock() + b.type = "text" + b.text = text + return b + + result = _extract_response_text([_tb("line1"), _tb("line2")]) + assert result == "line1\nline2" + + def test_non_text_blocks_skipped(self) -> None: + tool = MagicMock() + tool.type = "tool_use" + text = MagicMock() + text.type = "text" + text.text = "answer" + assert _extract_response_text([tool, text]) == "answer" + + def test_empty_content(self) -> None: + assert _extract_response_text([]) == "" + + def test_dict_text_block(self) -> None: + assert _extract_response_text([{"type": "text", "text": "dict text"}]) == "dict text" + + +# ------------------------------------------------------------------ +# _derive_agent_id +# ------------------------------------------------------------------ + + +class TestDeriveAgentId: + def test_default_when_no_system(self) -> None: + msgs: list[ChatMessage] = [ChatMessage(role="user", content="hi")] + assert _derive_agent_id(msgs) == "default" + + def test_uses_system_prompt(self) -> None: + msgs: list[ChatMessage] = [ + ChatMessage(role="system", content="You are a code assistant."), + ChatMessage(role="user", content="Help me."), + ] + assert _derive_agent_id(msgs) == "You are a code assistant." + + def test_truncates_at_64_chars(self) -> None: + msgs: list[ChatMessage] = [ChatMessage(role="system", content="x" * 100)] + assert len(_derive_agent_id(msgs)) == 64 + + def test_newlines_replaced_with_spaces(self) -> None: + msgs: list[ChatMessage] = [ChatMessage(role="system", content="line1\nline2")] + result = _derive_agent_id(msgs) + assert "\n" not in result + + def test_empty_system_prompt_falls_back(self) -> None: + msgs: list[ChatMessage] = [ + ChatMessage(role="system", content=" "), + ChatMessage(role="user", content="hi"), + ] + assert _derive_agent_id(msgs) == "default" + + +# ------------------------------------------------------------------ +# build_anthropic_llm_pair +# ------------------------------------------------------------------ + + +class TestBuildAnthropicLlmPair: + def test_pair_type_is_llm(self) -> None: + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="hi")], + response=_make_text_response(), + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert pair.pair_type == "llm" + + def test_ids_propagated(self) -> None: + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="hi")], + response=_make_text_response(), + model="m", + session_id="sess-abc", + invocation_id="inv-xyz", + run_id="run-1", + ) + assert pair.session_id == "sess-abc" + assert pair.invocation_id == "inv-xyz" + assert pair.run_id == "run-1" + + def test_model_from_response_preferred_over_request(self) -> None: + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="hi")], + response=_make_text_response(model="claude-haiku-4-5"), + model="request-model", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert isinstance(pair.data, LlmPairData) + assert pair.data.model == "claude-haiku-4-5" + + def test_fallback_to_request_model_when_response_empty(self) -> None: + resp = _make_text_response() + resp.model = "" + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="hi")], + response=resp, + model="fallback-model", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert isinstance(pair.data, LlmPairData) + assert pair.data.model == "fallback-model" + + def test_system_prompt_extracted(self) -> None: + flat_msgs: list[ChatMessage] = [ + ChatMessage(role="system", content="You are a triage agent."), + ChatMessage(role="user", content="Help."), + ] + pair = build_anthropic_llm_pair( + flat_messages=flat_msgs, + response=_make_text_response(), + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert pair.agent.system_prompt == "You are a triage agent." + + def test_last_user_message_is_user_instruction(self) -> None: + flat_msgs: list[ChatMessage] = [ + ChatMessage(role="system", content="sys"), + ChatMessage(role="user", content="first question"), + ChatMessage(role="assistant", content="answer"), + ChatMessage(role="user", content="follow-up"), + ] + pair = build_anthropic_llm_pair( + flat_messages=flat_msgs, + response=_make_text_response(), + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert pair.agent.user_instruction == "follow-up" + + def test_output_text_captured(self) -> None: + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="hi")], + response=_make_text_response(text="The answer is 42."), + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert isinstance(pair.data, LlmPairData) + assert pair.data.output == "The answer is 42." + + def test_token_usage_populated(self) -> None: + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="hi")], + response=_make_text_response(input_tokens=200, output_tokens=50), + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert isinstance(pair.data, LlmPairData) + assert pair.data.usage is not None + assert pair.data.usage["prompt_tokens"] == 200 + assert pair.data.usage["completion_tokens"] == 50 + assert pair.data.usage["total_tokens"] == 250 + + def test_tool_calls_in_data(self) -> None: + tool_block = MagicMock() + tool_block.type = "tool_use" + tool_block.id = "c1" + tool_block.name = "search" + tool_block.input = {"query": "test"} + + usage = MagicMock() + usage.input_tokens = 10 + usage.output_tokens = 5 + + resp = MagicMock() + resp.model = "claude-opus-4-6" + resp.content = [tool_block] + resp.usage = usage + + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="find it")], + response=resp, + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert isinstance(pair.data, LlmPairData) + assert len(pair.data.tool_calls) == 1 + assert pair.data.tool_calls[0]["name"] == "search" + + def test_event_id_is_unique_per_call(self) -> None: + kwargs: Any = dict( + flat_messages=[ChatMessage(role="user", content="hi")], + response=_make_text_response(), + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert build_anthropic_llm_pair(**kwargs).event_id != build_anthropic_llm_pair(**kwargs).event_id + + def test_parent_is_none(self) -> None: + pair = build_anthropic_llm_pair( + flat_messages=[ChatMessage(role="user", content="hi")], + response=_make_text_response(), + model="m", + session_id="s", + invocation_id="i", + run_id="r", + ) + assert pair.parent is None + + +# ------------------------------------------------------------------ +# _emit_pair +# ------------------------------------------------------------------ + + +class TestEmitPair: + async def test_emits_event_to_hooks(self) -> None: + config = AdrianConfig(session_id="sess-emit") + _, collector = _wired_hooks(config) + + await _emit_pair( + _make_text_response(text="Reply"), + { + "model": "claude-opus-4-6", + "messages": [{"role": "user", "content": "Question"}], + "system": "You are helpful.", + }, + ) + + assert len(collector.events) == 1 + event = collector.events[0] + assert event.pair_type == "llm" + assert event.session_id == "sess-emit" + assert isinstance(event.data, LlmPairData) + assert event.data.output == "Reply" + assert event.agent.system_prompt == "You are helpful." + + async def test_skips_silently_when_hooks_none(self) -> None: + _ah._hooks_getter = lambda: None + _ah._config_getter = lambda: None + await _emit_pair(_make_text_response(), {"model": "m", "messages": []}) + + async def test_skips_silently_when_getters_not_set(self) -> None: + _ah._hooks_getter = None + _ah._config_getter = None + await _emit_pair(_make_text_response(), {"model": "m", "messages": []}) + + async def test_fires_on_event_callback(self) -> None: + fired: list[str] = [] + + def on_event( + event_type: str, + data: Any, + run_id: str, + parent_run_id: str | None, + event_id: str | None, + ) -> None: + fired.append(event_type) + + config = AdrianConfig(session_id="s", on_event=on_event) + _wired_hooks(config) + + await _emit_pair( + _make_text_response(), + {"model": "m", "messages": [{"role": "user", "content": "q"}]}, + ) + + assert fired == ["llm"] + + async def test_uses_invocation_id_from_context(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + token = set_invocation_id("fixed-inv-id") + + try: + await _emit_pair( + _make_text_response(), + {"model": "m", "messages": [{"role": "user", "content": "q"}]}, + ) + finally: + token.var.reset(token) + + assert collector.events[0].invocation_id == "fixed-inv-id" + + async def test_defaults_to_no_invocation_outside_context(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + await _emit_pair( + _make_text_response(), + {"model": "m", "messages": [{"role": "user", "content": "q"}]}, + ) + + assert collector.events[0].invocation_id == "no_invocation" + + async def test_token_usage_in_emitted_event(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + await _emit_pair( + _make_text_response(input_tokens=100, output_tokens=40), + {"model": "m", "messages": [{"role": "user", "content": "q"}]}, + ) + + data = collector.events[0].data + assert isinstance(data, LlmPairData) + assert data.usage is not None + assert data.usage["total_tokens"] == 140 + + +# ------------------------------------------------------------------ +# patch_anthropic +# ------------------------------------------------------------------ + + +class TestPatchAnthropicGetters: + def test_getters_updated_on_each_call(self) -> None: + hooks_a: list[HookRegistry] = [HookRegistry()] + config_a: list[AdrianConfig] = [AdrianConfig()] + + patch_anthropic(hooks_getter=lambda: hooks_a[0], config_getter=lambda: config_a[0]) + + assert _ah._hooks_getter is not None + assert _ah._config_getter is not None + assert _ah._hooks_getter() is hooks_a[0] + assert _ah._config_getter() is config_a[0] + + hooks_b = HookRegistry() + config_b = AdrianConfig() + + patch_anthropic(hooks_getter=lambda: hooks_b, config_getter=lambda: config_b) + + assert _ah._hooks_getter() is hooks_b + assert _ah._config_getter() is config_b + + def test_no_op_when_anthropic_not_installed(self, monkeypatch: pytest.MonkeyPatch) -> None: + import sys + + saved = sys.modules.pop("anthropic", None) + monkeypatch.setitem(sys.modules, "anthropic", None) # type: ignore[arg-type] + + try: + patch_anthropic(hooks_getter=lambda: None, config_getter=lambda: None) + finally: + if saved is not None: + sys.modules["anthropic"] = saved + else: + sys.modules.pop("anthropic", None) + + +# ------------------------------------------------------------------ +# anthropic_invocation / anthropic_invocation_sync +# ------------------------------------------------------------------ + + +class TestAnthropicInvocationContext: + async def test_async_sets_invocation_id(self) -> None: + assert get_invocation_id() is None + + async with anthropic_invocation(): + inv_id = get_invocation_id() + assert inv_id is not None + assert len(inv_id) > 0 + + assert get_invocation_id() is None + + async def test_async_id_is_uuid_format(self) -> None: + import re + + uuid_re = re.compile( + r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$" + ) + + async with anthropic_invocation(): + assert uuid_re.match(get_invocation_id() or "") is not None + + async def test_async_resets_on_exit(self) -> None: + outer_token = set_invocation_id("outer") + + async with anthropic_invocation(): + inner_id = get_invocation_id() + assert inner_id != "outer" + + assert get_invocation_id() == "outer" + outer_token.var.reset(outer_token) + + def test_sync_sets_invocation_id(self) -> None: + assert get_invocation_id() is None + + with anthropic_invocation_sync(): + inv_id = get_invocation_id() + assert inv_id is not None + + assert get_invocation_id() is None + + async def test_two_consecutive_invocations_have_different_ids(self) -> None: + ids: list[str] = [] + + async with anthropic_invocation(): + ids.append(get_invocation_id() or "") + + async with anthropic_invocation(): + ids.append(get_invocation_id() or "") + + assert ids[0] != ids[1] From 904991ebd33797a5fa3b168c35efe9e4f38b6c7b Mon Sep 17 00:00:00 2001 From: Uchebuzz Date: Fri, 26 Jun 2026 14:12:15 +0100 Subject: [PATCH 2/7] style: apply ruff format to anthropic_handler and tests --- sdk/python/adrian/anthropic_handler.py | 8 ++++++-- sdk/python/tests/test_anthropic_handler.py | 22 +++++++++++++++++----- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py index f8daa51..f7c5201 100644 --- a/sdk/python/adrian/anthropic_handler.py +++ b/sdk/python/adrian/anthropic_handler.py @@ -464,7 +464,9 @@ def patch_anthropic( _original_sync = sync_cls.create def _patched_sync_create( - self: Any, *args: Any, **kwargs: Any # noqa: ANN401 + self: Any, + *args: Any, + **kwargs: Any, # noqa: ANN401 ) -> Any: # noqa: ANN401 response = _original_sync(self, *args, **kwargs) _schedule_emit(response, kwargs) @@ -487,7 +489,9 @@ def _patched_sync_create( _original_async = async_cls.create async def _patched_async_create( - self: Any, *args: Any, **kwargs: Any # noqa: ANN401 + self: Any, + *args: Any, + **kwargs: Any, # noqa: ANN401 ) -> Any: # noqa: ANN401 response = await _original_async(self, *args, **kwargs) await _emit_pair(response, kwargs) diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py index 483c077..2dabf49 100644 --- a/sdk/python/tests/test_anthropic_handler.py +++ b/sdk/python/tests/test_anthropic_handler.py @@ -242,7 +242,9 @@ def _make(name: str, id_: str) -> MagicMock: b.input = {} return b - result = _extract_anthropic_tool_calls([_make("tool_a", "c1"), _make("tool_b", "c2")]) + result = _extract_anthropic_tool_calls( + [_make("tool_a", "c1"), _make("tool_b", "c2")] + ) assert len(result) == 2 assert {r["name"] for r in result} == {"tool_a", "tool_b"} @@ -327,7 +329,10 @@ def test_empty_content(self) -> None: assert _extract_response_text([]) == "" def test_dict_text_block(self) -> None: - assert _extract_response_text([{"type": "text", "text": "dict text"}]) == "dict text" + assert ( + _extract_response_text([{"type": "text", "text": "dict text"}]) + == "dict text" + ) # ------------------------------------------------------------------ @@ -516,7 +521,10 @@ def test_event_id_is_unique_per_call(self) -> None: invocation_id="i", run_id="r", ) - assert build_anthropic_llm_pair(**kwargs).event_id != build_anthropic_llm_pair(**kwargs).event_id + assert ( + build_anthropic_llm_pair(**kwargs).event_id + != build_anthropic_llm_pair(**kwargs).event_id + ) def test_parent_is_none(self) -> None: pair = build_anthropic_llm_pair( @@ -641,7 +649,9 @@ def test_getters_updated_on_each_call(self) -> None: hooks_a: list[HookRegistry] = [HookRegistry()] config_a: list[AdrianConfig] = [AdrianConfig()] - patch_anthropic(hooks_getter=lambda: hooks_a[0], config_getter=lambda: config_a[0]) + patch_anthropic( + hooks_getter=lambda: hooks_a[0], config_getter=lambda: config_a[0] + ) assert _ah._hooks_getter is not None assert _ah._config_getter is not None @@ -656,7 +666,9 @@ def test_getters_updated_on_each_call(self) -> None: assert _ah._hooks_getter() is hooks_b assert _ah._config_getter() is config_b - def test_no_op_when_anthropic_not_installed(self, monkeypatch: pytest.MonkeyPatch) -> None: + def test_no_op_when_anthropic_not_installed( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: import sys saved = sys.modules.pop("anthropic", None) From 83be809d2241ae83cc45d2bdae59611e388d7066 Mon Sep 17 00:00:00 2001 From: Uchebuzz Date: Fri, 26 Jun 2026 14:54:24 +0100 Subject: [PATCH 3/7] fix(lint): resolve ruff import order and basedpyright strict-mode errors --- sdk/python/adrian/anthropic_handler.py | 13 +++++++++---- sdk/python/tests/test_anthropic_handler.py | 4 +--- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py index f7c5201..f068f0f 100644 --- a/sdk/python/adrian/anthropic_handler.py +++ b/sdk/python/adrian/anthropic_handler.py @@ -30,6 +30,11 @@ r2 = await client.messages.create(...) # same invocation_id as r1 """ +# pyright: reportUnknownVariableType=false +# pyright: reportUnknownMemberType=false +# pyright: reportUnknownArgumentType=false +# pyright: reportUnknownLambdaType=false + from __future__ import annotations import asyncio @@ -61,7 +66,7 @@ # ------------------------------------------------------------------ -def _flatten_content(content: str | list[Any]) -> str: +def _flatten_content(content: Any) -> str: # noqa: ANN401 """Flatten Anthropic message content to a plain string. Anthropic messages carry either a plain string or a list of content @@ -451,14 +456,14 @@ def patch_anthropic( _config_getter = config_getter try: - import anthropic + from anthropic.resources.messages import AsyncMessages, Messages except ImportError: logger.debug("anthropic package not installed; skipping Anthropic patching") return # ---- sync Messages.create ---- try: - sync_cls = anthropic.resources.Messages + sync_cls = Messages if not getattr(sync_cls, "_adrian_patched", False): _original_sync = sync_cls.create @@ -483,7 +488,7 @@ def _patched_sync_create( # ---- async AsyncMessages.create ---- try: - async_cls = anthropic.resources.AsyncMessages + async_cls = AsyncMessages if not getattr(async_cls, "_adrian_patched", False): _original_async = async_cls.create diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py index 2dabf49..e2dc1ba 100644 --- a/sdk/python/tests/test_anthropic_handler.py +++ b/sdk/python/tests/test_anthropic_handler.py @@ -11,9 +11,8 @@ from typing import Any from unittest.mock import MagicMock -import pytest - import adrian.anthropic_handler as _ah +import pytest from adrian.anthropic_handler import ( _derive_agent_id, _emit_pair, @@ -33,7 +32,6 @@ from adrian.hooks import HookRegistry from adrian.types import ChatMessage - # ------------------------------------------------------------------ # Shared helpers # ------------------------------------------------------------------ From 5212849cd52ad2396963a9d70151623039b812e1 Mon Sep 17 00:00:00 2001 From: uchebuzz-coder Date: Wed, 22 Jul 2026 10:32:34 +0100 Subject: [PATCH 4/7] feat(sdk/anthropic): gate tool calls on classifier verdict (BLOCK/HITL) The Anthropic handler was audit-only: it emitted a PairedEvent after messages.create() but never acted on the verdict. This adds the verdict gate requested in review. The Anthropic Messages API has no discrete tool-execution step to intercept (unlike LangChain's BaseTool.invoke or the TS OpenAI captureTool), so the gate runs at the only choke point the SDK owns: the return of messages.create(). After emitting (which registers the verdict future), the async path awaits each tool_use block's verdict and, when it halts, rewrites that block into a "[BLOCKED by security policy]" text block -- consistent with the LangChain gate, so the model sees its action was stopped. Fail-closed rules mirror langchain_handler: missing LoginAck within 5s blocks; MODE_BLOCK verdict timeout blocks; MODE_HITL waits indefinitely. Scope: async (AsyncMessages) only; sync Messages.create stays audit-only pending maintainer confirmation on parity. _should_halt is a self-contained copy so the Anthropic integration does not depend on the LangChain module. Adds 18 tests: ALERT passthrough, BLOCK halt/allow, verdict-timeout fail-closed, partial block, dict/object block shapes, HITL reject/approve/holds-until-human, and an end-to-end emit->gate seam test. --- sdk/python/adrian/__init__.py | 1 + sdk/python/adrian/anthropic_handler.py | 215 +++++++++++- sdk/python/tests/test_anthropic_handler.py | 389 ++++++++++++++++++++- sdk/python/uv.lock | 122 ++++++- 4 files changed, 721 insertions(+), 6 deletions(-) diff --git a/sdk/python/adrian/__init__.py b/sdk/python/adrian/__init__.py index 983d5d4..09d4b6f 100644 --- a/sdk/python/adrian/__init__.py +++ b/sdk/python/adrian/__init__.py @@ -401,6 +401,7 @@ def patch_anthropic() -> None: _patch( hooks_getter=lambda: _hooks, config_getter=lambda: get_config() if is_initialized() else None, + ws_getter=lambda: _ws_client, ) diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py index f068f0f..83c6465 100644 --- a/sdk/python/adrian/anthropic_handler.py +++ b/sdk/python/adrian/anthropic_handler.py @@ -42,6 +42,7 @@ from collections.abc import Callable from contextlib import asynccontextmanager, contextmanager from datetime import UTC, datetime +from types import SimpleNamespace from typing import TYPE_CHECKING, Any from uuid import uuid4 @@ -49,16 +50,25 @@ from adrian.context import get_invocation_id, set_invocation_id from adrian.format.types import AgentContext, LlmPairData, PairedEvent from adrian.hooks import HookRegistry +from adrian.proto import event_pb2 as pb from adrian.types import ChatMessage, EventData, TokenUsage, ToolCallRecord if TYPE_CHECKING: from contextvars import Token + from adrian.ws import WebSocketClient + logger = logging.getLogger("adrian.anthropic") +# Content substituted for a tool_use block whose verdict says halt. A text +# block (rather than a raised error) keeps the response usable and lets the +# model see that its action was stopped, mirroring the LangChain gate. +_BLOCKED_CONTENT = "[BLOCKED by security policy]" + # Set once by patch_anthropic(); read at call time so shutdown + re-init works. _hooks_getter: Callable[[], HookRegistry | None] | None = None _config_getter: Callable[[], AdrianConfig | None] | None = None +_ws_getter: Callable[[], WebSocketClient | None] | None = None # ------------------------------------------------------------------ @@ -423,6 +433,195 @@ def _schedule_emit(response: Any, kwargs: dict[str, Any]) -> None: logger.exception("Failed to emit Anthropic event (sync path)") +# ------------------------------------------------------------------ +# Verdict gate (MODE_BLOCK / MODE_HITL) +# ------------------------------------------------------------------ + + +def _should_halt(verdict: pb.Verdict) -> bool: + """Decide whether a verdict should halt a tool call. + + Kept local to this module so the Anthropic integration does not depend on + the LangChain handler; the logic mirrors ``langchain_handler._should_halt``. + HITL resolutions override per-MAD policy when present. + + Args: + verdict: The classifier verdict for the producing LLM event. + + Returns: + ``True`` when the tool call must be blocked. + """ + if verdict.HasField("hitl"): + return not verdict.hitl.continue_execution + + mad_prefix = verdict.mad_code[:2] + return { + "M0": verdict.policy.policy_m0, + "M2": verdict.policy.policy_m2, + "M3": verdict.policy.policy_m3, + "M4": verdict.policy.policy_m4, + }.get(mad_prefix, False) + + +def _blocked_text_block(original_block: Any) -> Any: # noqa: ANN401 + """Build a text block used to replace a blocked ``tool_use`` block. + + Matches the shape of ``original_block`` so the rewritten response stays + consistent: a dict block is replaced with a dict, an SDK block with a real + ``anthropic.types.TextBlock`` when the package is importable, falling back + to a lightweight attribute shim otherwise (e.g. under test). + + Args: + original_block: The ``tool_use`` block being replaced. + + Returns: + A ``"text"`` block carrying :data:`_BLOCKED_CONTENT`. + """ + if isinstance(original_block, dict): + return {"type": "text", "text": _BLOCKED_CONTENT} + + try: + from anthropic.types import TextBlock + + return TextBlock(type="text", text=_BLOCKED_CONTENT) + except Exception: # noqa: BLE001 - anthropic missing or signature drift + return SimpleNamespace(type="text", text=_BLOCKED_CONTENT) + + +def _rewrite_blocked_response(response: Any, blocked_ids: set[str]) -> Any: # noqa: ANN401 + """Replace blocked ``tool_use`` blocks in a response with ``[BLOCKED]`` text. + + Rebuilds ``response.content`` in place, swapping every ``tool_use`` block + whose id is in ``blocked_ids`` for a text block. When no ``tool_use`` block + survives, ``stop_reason`` is downgraded from ``"tool_use"`` to ``"end_turn"`` + so the caller's agent loop terminates cleanly instead of expecting a tool + result. + + Args: + response: Anthropic ``Message`` response object. + blocked_ids: ``tool_use`` block ids to replace. + + Returns: + The same ``response`` object, mutated. + """ + content: list[Any] = getattr(response, "content", []) + new_content: list[Any] = [] + tool_use_remains = False + + for block in content: + if isinstance(block, dict): + btype = block.get("type") + bid = str(block.get("id", "")) + else: + btype = getattr(block, "type", None) + bid = str(getattr(block, "id", "")) + + if btype == "tool_use" and bid in blocked_ids: + new_content.append(_blocked_text_block(block)) + else: + new_content.append(block) + if btype == "tool_use": + tool_use_remains = True + + response.content = new_content + + if not tool_use_remains and getattr(response, "stop_reason", None) == "tool_use": + response.stop_reason = "end_turn" + + return response + + +async def _gate_response(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: ANN401 + """Hold a response on the classifier verdict before returning it. + + Runs after :func:`_emit_pair` (which registers the pending verdict future + for the producing LLM event). In ``MODE_ALERT`` / unset state the response + passes through unchanged. In ``MODE_BLOCK`` / ``MODE_HITL`` each + ``tool_use`` block waits for its verdict; blocked blocks are rewritten to a + ``[BLOCKED]`` text block. + + Fail-closed rules match the LangChain gate: a missing ``LoginAck`` within 5s + blocks all tool calls, and in ``MODE_BLOCK`` a verdict timeout blocks the + tool call (absence of a verdict is a policy violation). ``MODE_HITL`` waits + indefinitely. + + Args: + response: Anthropic ``Message`` response object. + kwargs: Original ``messages.create`` keyword arguments (unused; kept for + symmetry with :func:`_emit_pair`). + + Returns: + The response, rewritten in place when any tool call was blocked. + """ + if _ws_getter is None: + return response + + ws = _ws_getter() + + if ws is None: + return response + + content: Any = getattr(response, "content", None) + + if not isinstance(content, list): + return response + + tool_ids = [ + rec["id"] for rec in _extract_anthropic_tool_calls(content) if rec["id"] + ] + + if not tool_ids: + return response + + # Refuse to run without a verified policy: block if LoginAck is late. + if not ws._login_ack_received.is_set(): # pyright: ignore[reportPrivateUsage] + try: + await asyncio.wait_for( + ws._login_ack_received.wait(), # pyright: ignore[reportPrivateUsage] + timeout=5.0, + ) + except TimeoutError: + logger.warning( + "Anthropic gate: LoginAck not received within 5s; " + "blocking all tool calls (refusing to run without policy)" + ) + return _rewrite_blocked_response(response, set(tool_ids)) + + # ALERT / unset: observe only, no gating. + if not ws.policy_active(): + return response + + config = _config_getter() if _config_getter is not None else None + timeout = ws.block_timeout(config.block_timeout if config else 30.0) + + blocked: set[str] = set() + + for tc_id in tool_ids: + verdict = await ws.wait_for_tool_call_verdict(tc_id, timeout) + + if verdict is None: + # Fail-closed in block mode: no verdict = block. + logger.warning( + "Anthropic gate: verdict timeout for tool_call_id=%s; " + "blocking (fail-closed in MODE_BLOCK)", + tc_id, + ) + blocked.add(tc_id) + elif _should_halt(verdict): + logger.warning( + "Anthropic gate: halting tool_call_id=%s event_id=%s mad_code=%s", + tc_id, + verdict.event_id, + verdict.mad_code, + ) + blocked.add(tc_id) + + if blocked: + return _rewrite_blocked_response(response, blocked) + + return response + + # ------------------------------------------------------------------ # SDK patching # ------------------------------------------------------------------ @@ -431,11 +630,14 @@ def _schedule_emit(response: Any, kwargs: dict[str, Any]) -> None: def patch_anthropic( hooks_getter: Callable[[], HookRegistry | None], config_getter: Callable[[], AdrianConfig | None], + ws_getter: Callable[[], WebSocketClient | None] | None = None, ) -> None: """Monkey-patch ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic``. Wraps ``messages.create`` on both the sync and async Anthropic resource - classes so every API call is captured as an Adrian ``PairedEvent``. + classes so every API call is captured as an Adrian ``PairedEvent``. On the + async path, calls are additionally gated on the classifier verdict under + ``MODE_BLOCK`` / ``MODE_HITL`` (see :func:`_gate_response`). The patch is idempotent: subsequent calls update the internal getters but do not re-wrap the already-patched method. If the ``anthropic`` package is @@ -449,11 +651,15 @@ def patch_anthropic( or ``None`` when the SDK is not initialised. config_getter: Zero-arg callable returning the current ``AdrianConfig``, or ``None`` when the SDK is not initialised. + ws_getter: Zero-arg callable returning the current ``WebSocketClient``, + or ``None``. Required for verdict gating; when absent the handler + stays audit-only (ALERT-mode behaviour). """ - global _hooks_getter, _config_getter # noqa: PLW0603 + global _hooks_getter, _config_getter, _ws_getter # noqa: PLW0603 _hooks_getter = hooks_getter _config_getter = config_getter + _ws_getter = ws_getter try: from anthropic.resources.messages import AsyncMessages, Messages @@ -499,8 +705,11 @@ async def _patched_async_create( **kwargs: Any, # noqa: ANN401 ) -> Any: # noqa: ANN401 response = await _original_async(self, *args, **kwargs) + # Emit first so the verdict future is registered, then gate: + # under BLOCK/HITL this holds the response until the verdict + # arrives and rewrites blocked tool calls to a [BLOCKED] block. await _emit_pair(response, kwargs) - return response + return await _gate_response(response, kwargs) async_cls.create = _patched_async_create # type: ignore[method-assign] async_cls._adrian_patched = True # type: ignore[attr-defined] diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py index e2dc1ba..f54fe96 100644 --- a/sdk/python/tests/test_anthropic_handler.py +++ b/sdk/python/tests/test_anthropic_handler.py @@ -8,12 +8,15 @@ from __future__ import annotations -from typing import Any -from unittest.mock import MagicMock +import asyncio +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock import adrian.anthropic_handler as _ah import pytest from adrian.anthropic_handler import ( + _BLOCKED_CONTENT, + _blocked_text_block, _derive_agent_id, _emit_pair, _extract_anthropic_tool_calls, @@ -21,6 +24,9 @@ _extract_response_text, _flatten_anthropic_messages, _flatten_content, + _gate_response, + _rewrite_blocked_response, + _should_halt, anthropic_invocation, anthropic_invocation_sync, build_anthropic_llm_pair, @@ -30,7 +36,9 @@ from adrian.context import get_invocation_id, set_invocation_id from adrian.format.types import LlmPairData, PairedEvent from adrian.hooks import HookRegistry +from adrian.proto import event_pb2 as pb from adrian.types import ChatMessage +from adrian.ws import WebSocketClient # ------------------------------------------------------------------ # Shared helpers @@ -736,3 +744,380 @@ async def test_two_consecutive_invocations_have_different_ids(self) -> None: ids.append(get_invocation_id() or "") assert ids[0] != ids[1] + + +# ------------------------------------------------------------------ +# Verdict gate (MODE_BLOCK / MODE_HITL) +# ------------------------------------------------------------------ + + +def _apply_mode( + ws: WebSocketClient, + mode: int, + *, + policy_m0: bool = False, + policy_m2: bool = False, + policy_m3: bool = False, + policy_m4: bool = False, +) -> pb.PolicySnapshot: + """Drive the ws mode/policy state as if a LoginAck had arrived.""" + policy = pb.PolicySnapshot( + mode=cast("pb.Mode", mode), + policy_m0=policy_m0, + policy_m2=policy_m2, + policy_m3=policy_m3, + policy_m4=policy_m4, + ) + ws._mode = mode # pyright: ignore[reportPrivateUsage] + ws._policy = policy # pyright: ignore[reportPrivateUsage] + ws._login_ack_received.set() # pyright: ignore[reportPrivateUsage] + return policy + + +def _make_tool_response( + *, + tool_id: str = "tc-1", + tool_name: str = "run_shell", + stop_reason: str = "tool_use", + as_dict: bool = False, + with_text: bool = False, +) -> MagicMock: + """Build a mock Anthropic Message carrying a single ``tool_use`` block.""" + if as_dict: + tool_block: Any = { + "type": "tool_use", + "id": tool_id, + "name": tool_name, + "input": {"cmd": "ls"}, + } + else: + tool_block = MagicMock() + tool_block.type = "tool_use" + tool_block.id = tool_id + tool_block.name = tool_name + tool_block.input = {"cmd": "ls"} + + content: list[Any] = [] + if with_text: + text_block = MagicMock() + text_block.type = "text" + text_block.text = "Let me run that." + content.append(text_block) + content.append(tool_block) + + response = MagicMock() + response.model = "claude-opus-4-6" + response.content = content + response.stop_reason = stop_reason + usage = MagicMock() + usage.input_tokens = 1 + usage.output_tokens = 1 + response.usage = usage + return response + + +def _wire_gate(ws: WebSocketClient | None, config: AdrianConfig | None) -> None: + """Wire the ws / config getters the gate reads at call time.""" + _ah._ws_getter = (lambda: ws) if ws is not None else (lambda: None) + _ah._config_getter = (lambda: config) if config is not None else (lambda: None) + + +def _block_types(response: Any) -> list[str]: # noqa: ANN401 + """Return the ``type`` of each content block, dict- or object-shaped.""" + out: list[str] = [] + for block in response.content: + out.append(block["type"] if isinstance(block, dict) else block.type) + return out + + +class TestVerdictGate: + @pytest.fixture(autouse=True) # pyright: ignore[reportUntypedFunctionDecorator] + def _reset_getters(self) -> Any: # noqa: ANN401 + """Reset module getters so gate wiring never leaks between tests.""" + yield + _ah._ws_getter = None + _ah._config_getter = None + _ah._hooks_getter = None + + async def test_alert_mode_passes_through(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_ALERT) + _wire_gate(ws, AdrianConfig(session_id="s")) + + response = _make_tool_response() + result = await _gate_response(response, {}) + + # ALERT observes only: tool_use survives untouched. + assert _block_types(result) == ["tool_use"] + + async def test_no_ws_passes_through(self) -> None: + _wire_gate(None, AdrianConfig(session_id="s")) + response = _make_tool_response() + result = await _gate_response(response, {}) + assert _block_types(result) == ["tool_use"] + + async def test_no_tool_calls_passes_through(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + _wire_gate(ws, AdrianConfig(session_id="s")) + + response = _make_text_response(text="Just text.") + result = await _gate_response(response, {}) + assert _block_types(result) == ["text"] + + async def test_block_halt_rewrites_tool_use(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + _wire_gate(ws, AdrianConfig(session_id="s")) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + fut.set_result(pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy)) + + response = _make_tool_response() + result = await _gate_response(response, {}) + + assert _block_types(result) == ["text"] + assert result.content[0].text == _BLOCKED_CONTENT + # No tool_use survives -> stop_reason downgraded. + assert result.stop_reason == "end_turn" + + async def test_block_allow_passes_through(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + # M4 policy inactive -> verdict does not halt. + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=False) + _wire_gate(ws, AdrianConfig(session_id="s")) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + fut.set_result(pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy)) + + response = _make_tool_response() + result = await _gate_response(response, {}) + + assert _block_types(result) == ["tool_use"] + assert result.stop_reason == "tool_use" + + async def test_block_verdict_timeout_fails_closed(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + # Tiny timeout so the fail-closed path resolves fast. + _wire_gate(ws, AdrianConfig(session_id="s", block_timeout=0.05)) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + ws.register_pending("llm-evt") # never resolved -> times out + + response = _make_tool_response() + result = await _gate_response(response, {}) + + assert _block_types(result) == ["text"] + assert result.content[0].text == _BLOCKED_CONTENT + + async def test_partial_block_keeps_other_tool_use(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + _wire_gate(ws, AdrianConfig(session_id="s")) + + # Two tool calls sharing one producing LLM event; only tc-1 blocked + # (tc-2 has no verdict mapping -> ... also None -> blocked). To isolate + # "one blocked, one allowed" we map both and resolve distinct verdicts. + ws._tool_call_id_to_event_id["tc-1"] = "evt-a" # pyright: ignore[reportPrivateUsage] + ws._tool_call_id_to_event_id["tc-2"] = "evt-b" # pyright: ignore[reportPrivateUsage] + halt = ws.register_pending("evt-a") + halt.set_result(pb.Verdict(event_id="evt-a", mad_code="M4_a", policy=policy)) + allow_policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m4=False) + ok = ws.register_pending("evt-b") + ok.set_result( + pb.Verdict(event_id="evt-b", mad_code="M4_a", policy=allow_policy) + ) + + blocked_tool = MagicMock() + blocked_tool.type = "tool_use" + blocked_tool.id = "tc-1" + blocked_tool.name = "danger" + blocked_tool.input = {} + ok_tool = MagicMock() + ok_tool.type = "tool_use" + ok_tool.id = "tc-2" + ok_tool.name = "safe" + ok_tool.input = {} + response = MagicMock() + response.content = [blocked_tool, ok_tool] + response.stop_reason = "tool_use" + + result = await _gate_response(response, {}) + + assert _block_types(result) == ["text", "tool_use"] + assert result.content[0].text == _BLOCKED_CONTENT + assert result.content[1].id == "tc-2" + # A tool_use still remains -> stop_reason preserved. + assert result.stop_reason == "tool_use" + + async def test_dict_shaped_block_rewritten_as_dict(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + _wire_gate(ws, AdrianConfig(session_id="s")) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + fut.set_result(pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy)) + + response = _make_tool_response(as_dict=True) + result = await _gate_response(response, {}) + + assert result.content[0] == {"type": "text", "text": _BLOCKED_CONTENT} + + async def test_hitl_reject_blocks(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_HITL, policy_m4=True) + _wire_gate(ws, AdrianConfig(session_id="s")) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + verdict = pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy) + verdict.hitl.continue_execution = False + fut.set_result(verdict) + + result = await _gate_response(_make_tool_response(), {}) + assert result.content[0].text == _BLOCKED_CONTENT + + async def test_hitl_approve_passes_through(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_HITL, policy_m4=True) + _wire_gate(ws, AdrianConfig(session_id="s")) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + verdict = pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy) + verdict.hitl.continue_execution = True + fut.set_result(verdict) + + result = await _gate_response(_make_tool_response(), {}) + assert _block_types(result) == ["tool_use"] + + async def test_hitl_holds_until_human_then_blocks(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_HITL, policy_m4=True) + _wire_gate(ws, AdrianConfig(session_id="s")) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + + gate = asyncio.ensure_future(_gate_response(_make_tool_response(), {})) + # Gate must still be waiting: no verdict resolved yet. + await asyncio.sleep(0.05) + assert not gate.done() + + verdict = pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy) + verdict.hitl.continue_execution = False + fut.set_result(verdict) + + result = await gate + assert result.content[0].text == _BLOCKED_CONTENT + + async def test_emit_then_gate_end_to_end(self) -> None: + """The real seam: _emit_pair populates the verdict map the gate reads. + + Rather than hand-populating ``_tool_call_id_to_event_id`` and + pre-registering the future (as the unit tests do), this drives the + actual emission path -- ``_emit_pair`` -> ``hooks.emit`` -> + ``ws.on_paired_event`` -- with only the network send stubbed, then + verifies the gate finds the verdict and rewrites the blocked call. + """ + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._send_frame = AsyncMock() # pyright: ignore[reportPrivateUsage] - no network + + hooks = HookRegistry() + hooks.register(ws) + _ah._hooks_getter = lambda: hooks + _wire_gate(ws, AdrianConfig(session_id="s")) + + response = _make_tool_response(tool_id="tc-1") + kwargs: dict[str, Any] = { + "model": "claude-opus-4-6", + "messages": [{"role": "user", "content": "run it"}], + } + + # Emission maps tc-1 -> event_id and pre-registers the wait future. + await _emit_pair(response, kwargs) + event_id = ws._tool_call_id_to_event_id["tc-1"] # pyright: ignore[reportPrivateUsage] + ws._pending_verdicts[event_id].set_result( # pyright: ignore[reportPrivateUsage] + pb.Verdict(event_id=event_id, mad_code="M4_a", policy=policy) + ) + + result = await _gate_response(response, kwargs) + + assert result.content[0].text == _BLOCKED_CONTENT + assert result.stop_reason == "end_turn" + ws._send_frame.assert_awaited() # pyright: ignore[reportPrivateUsage] + + async def test_patch_anthropic_stores_ws_getter(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + patch_anthropic( + hooks_getter=lambda: None, + config_getter=lambda: None, + ws_getter=lambda: ws, + ) + assert _ah._ws_getter is not None + assert _ah._ws_getter() is ws + + +class TestShouldHalt: + def test_hitl_reject_halts(self) -> None: + v = pb.Verdict(event_id="e", mad_code="M4_a") + v.hitl.continue_execution = False + assert _should_halt(v) is True + + def test_hitl_approve_does_not_halt(self) -> None: + v = pb.Verdict(event_id="e", mad_code="M4_a") + v.hitl.continue_execution = True + assert _should_halt(v) is False + + def test_policy_flag_on_halts(self) -> None: + policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m3=True) + v = pb.Verdict(event_id="e", mad_code="M3_x", policy=policy) + assert _should_halt(v) is True + + def test_policy_flag_off_does_not_halt(self) -> None: + policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m3=False) + v = pb.Verdict(event_id="e", mad_code="M3_x", policy=policy) + assert _should_halt(v) is False + + def test_unknown_mad_prefix_does_not_halt(self) -> None: + policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m4=True) + v = pb.Verdict(event_id="e", mad_code="M9_z", policy=policy) + assert _should_halt(v) is False + + +class TestBlockedTextBlock: + def test_dict_block_returns_dict(self) -> None: + result = _blocked_text_block({"type": "tool_use", "id": "x"}) + assert result == {"type": "text", "text": _BLOCKED_CONTENT} + + def test_object_block_returns_text_shaped_object(self) -> None: + original = MagicMock() + original.type = "tool_use" + result = _blocked_text_block(original) + assert result.type == "text" + assert result.text == _BLOCKED_CONTENT + + +class TestRewriteBlockedResponse: + def test_downgrades_stop_reason_when_no_tool_use_left(self) -> None: + block = {"type": "tool_use", "id": "tc-1", "name": "x", "input": {}} + response = MagicMock() + response.content = [block] + response.stop_reason = "tool_use" + _rewrite_blocked_response(response, {"tc-1"}) + assert response.stop_reason == "end_turn" + + def test_preserves_unblocked_blocks(self) -> None: + keep = {"type": "text", "text": "hi"} + drop = {"type": "tool_use", "id": "tc-1", "name": "x", "input": {}} + response = MagicMock() + response.content = [keep, drop] + response.stop_reason = "tool_use" + _rewrite_blocked_response(response, {"tc-1"}) + assert response.content[0] == keep + assert response.content[1] == {"type": "text", "text": _BLOCKED_CONTENT} diff --git a/sdk/python/uv.lock b/sdk/python/uv.lock index 1940089..0d39736 100644 --- a/sdk/python/uv.lock +++ b/sdk/python/uv.lock @@ -13,7 +13,11 @@ dependencies = [ ] [package.optional-dependencies] +anthropic = [ + { name = "anthropic" }, +] dev = [ + { name = "anthropic" }, { name = "basedpyright" }, { name = "langchain-mcp-adapters" }, { name = "langgraph" }, @@ -27,6 +31,8 @@ dev = [ [package.metadata] requires-dist = [ + { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.40.0" }, + { name = "anthropic", marker = "extra == 'dev'", specifier = ">=0.40.0" }, { name = "basedpyright", marker = "extra == 'dev'", specifier = "==1.39.3" }, { name = "langchain-core", specifier = ">=1.2.19,<2.0" }, { name = "langchain-mcp-adapters", marker = "extra == 'dev'", specifier = ">=0.2.2" }, @@ -40,7 +46,7 @@ requires-dist = [ { name = "ruff", marker = "extra == 'dev'", specifier = "==0.15.12" }, { name = "websockets", specifier = ">=16.0" }, ] -provides-extras = ["dev"] +provides-extras = ["anthropic", "dev"] [[package]] name = "annotated-types" @@ -51,6 +57,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "anthropic" +version = "0.117.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "distro" }, + { name = "docstring-parser" }, + { name = "httpx" }, + { name = "jiter" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/41/0d/8f71d535edb0d438f023bd825fb65f67c14fa88a2bd6b75f292a58a63de4/anthropic-0.117.0.tar.gz", hash = "sha256:98107f2b76439641e0ae2a1754087534b8f178dbab99d6eb1bc4b7bc8c744496", size = 989933, upload-time = "2026-07-16T19:36:13.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/4c/917d21d6619a4475cdafc6d13a69fdb3b901ddac57e76caca5a25c117b6d/anthropic-0.117.0-py3-none-any.whl", hash = "sha256:451a0a6905f11dff7663d13e4ee5dbf909eb8942b1d049803c7b937a13ac47ec", size = 998327, upload-time = "2026-07-16T19:36:11.225Z" }, +] + [[package]] name = "anyio" version = "4.13.0" @@ -400,6 +425,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/08/9c41fb51ab5b43eb21674aff13df270e8ba6c4b29c8624e328dc7a9482af/distlib-0.4.3-py2.py3-none-any.whl", hash = "sha256:4b0ce306c966eb73bc3a7b6abad017c556dadd92c44701562cd528ac7fde4d5b", size = 470628, upload-time = "2026-06-12T08:04:50.506Z" }, ] +[[package]] +name = "distro" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fc/f8/98eea607f65de6527f8a2e8885fc8015d3e6f5775df186e443e0964a11c3/distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed", size = 60722, upload-time = "2023-12-24T09:54:32.31Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" }, +] + +[[package]] +name = "docstring-parser" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/4d/f332313098c1de1b2d2ff91cf2674415cc7cddab2ca1b01ae29774bd5fdf/docstring_parser-0.18.0.tar.gz", hash = "sha256:292510982205c12b1248696f44959db3cdd1740237a968ea1e2e7a900eeb2015", size = 29341, upload-time = "2026-04-14T04:09:19.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, +] + [[package]] name = "filelock" version = "3.29.4" @@ -482,6 +525,74 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "jiter" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1d/1f/10936e16d8860c70698a1aa939a46aa0224813b782bce4e000e637da0b2d/jiter-0.16.0.tar.gz", hash = "sha256:7b24c3492c5f4f84a37946ad9cf504910cf6a782d6a4e0689b6673c5894b4a1c", size = 176431, upload-time = "2026-06-29T13:05:13.657Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/2b/52ace16ed031354f0539749a49e4bf33797d82bea5137910835fa4b09793/jiter-0.16.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:67c3bc1760f8c99d805dcab4e644027142a53b1d5d861f18780ebdbd5d40b72a", size = 306943, upload-time = "2026-06-29T13:03:14.035Z" }, + { url = "https://files.pythonhosted.org/packages/94/2e/34957c2c1b661c252ba9bcc60ae0bddc27e0f7202c6073326a13c5390eec/jiter-0.16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5af7780e4a26bd7d0d989592bf9ef12ebf806b74ab709223ecca37c749872ea9", size = 307779, upload-time = "2026-06-29T13:03:15.418Z" }, + { url = "https://files.pythonhosted.org/packages/88/6c/59bd309cab4460c54cf1079f3eb7fe7af6a4c895c5c957a53378693bad2b/jiter-0.16.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d5bf78d0e05e45cfdd66558893938d59afe3d1b1a824a202039b20e607d25a72", size = 335826, upload-time = "2026-06-29T13:03:17.11Z" }, + { url = "https://files.pythonhosted.org/packages/3b/8c/f5ef7b65f0df47afa16596969defb281ebb86e96df346d62be6fd853d620/jiter-0.16.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f4444a83f946605990c98f625cdd3d2725bfb818158760c5748c653170a20e0e", size = 362573, upload-time = "2026-06-29T13:03:18.781Z" }, + { url = "https://files.pythonhosted.org/packages/2b/0b/ace4354da061ee38844a0c27dc2c21eecd27aea119e8da324bea987522d0/jiter-0.16.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3a23f0e4f957e1be65752d2dfac9a5a06b1917af8dc85deb639c3b9d02e31290", size = 457979, upload-time = "2026-06-29T13:03:20.293Z" }, + { url = "https://files.pythonhosted.org/packages/55/40/c0253d3772eb9dcd8e6606ee9b2d53ec8e5b814589c47f140aa585f21eaa/jiter-0.16.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c22a488f7b9218e245a0025a9ba6b100e2e54700831cf4cf16833a27fba3ad01", size = 372302, upload-time = "2026-06-29T13:03:21.739Z" }, + { url = "https://files.pythonhosted.org/packages/a8/d2/4839422241aa12860ce597b20068727094ba0bc480723c74924ca5bad483/jiter-0.16.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:46add52f4ad47a08bfb1219f3e673da972191489a33016edefdb5ea55bfa8c48", size = 343805, upload-time = "2026-06-29T13:03:23.384Z" }, + { url = "https://files.pythonhosted.org/packages/e2/59/e196888a05befdda7dbe299b722d56f2f6eec65402bc34c0a3306d595feb/jiter-0.16.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:9c8a956fd72c2cf1e730d01ea080341f13aa0a97a4a33b51abebe725b7ae9ca9", size = 351107, upload-time = "2026-06-29T13:03:24.815Z" }, + { url = "https://files.pythonhosted.org/packages/ec/74/4cd9e0fca65232136400354b630fbfcd2de634e22ccbb96567725981b548/jiter-0.16.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:561926e0573ffe4a32498420a76d64b16c513e1ab413b9d28158a8764ac701e5", size = 388441, upload-time = "2026-06-29T13:03:26.266Z" }, + { url = "https://files.pythonhosted.org/packages/d9/8c/554691e48bc711299c0a293dd8a6179e24b2d66a54dc295421fcf64569c0/jiter-0.16.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:44d019fa8cdaf89bf29c71b39e3712143fdd0ac76725c6ef954f9957a5ea8730", size = 516354, upload-time = "2026-06-29T13:03:28.02Z" }, + { url = "https://files.pythonhosted.org/packages/a4/cb/01e9d69dc2cc6759d4f91e230b34489c4fdb2518992650633f9e20bece89/jiter-0.16.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:0df91907609837f33341b8e6fe73b95991fdaa57caf1a0fbd343dffe826f386f", size = 547880, upload-time = "2026-06-29T13:03:29.534Z" }, + { url = "https://files.pythonhosted.org/packages/79/70/2953195f1c6ad00f49fa67e13df7e60acb3dd4f387101bc15abccddd905e/jiter-0.16.0-cp312-cp312-win32.whl", hash = "sha256:51d7b836acb0108d7c77df1742332cac2a1fa04a74d6dacec46e7091f0e91274", size = 203473, upload-time = "2026-06-29T13:03:31.025Z" }, + { url = "https://files.pythonhosted.org/packages/2d/05/2909a8b10699a4d560f8c502b6b2c5f3991b682b1922c1eedda242b225bd/jiter-0.16.0-cp312-cp312-win_amd64.whl", hash = "sha256:1878349266f8ee36ecb1375cc5ba2f115f35fd9f0a1a4119e725e379126647f7", size = 196905, upload-time = "2026-06-29T13:03:32.472Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a9/6b82bb1c8d7790d602489b967b982a909e5d092875a6c2ade96444c8dfc5/jiter-0.16.0-cp312-cp312-win_arm64.whl", hash = "sha256:2ed5738ae4af18271a51a528b8811b0cbfa4a1858de9d83359e4169855d6a331", size = 190618, upload-time = "2026-06-29T13:03:34.672Z" }, + { url = "https://files.pythonhosted.org/packages/91/c0/555fc60473d30d66894ba825e63615e3be7524fac23858356afa7a38906c/jiter-0.16.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:41977aa5654023948c2dae2a81cbf9c43343954bef1cd59a154dd15a4d84c195", size = 306203, upload-time = "2026-06-29T13:03:36.243Z" }, + { url = "https://files.pythonhosted.org/packages/d0/2b/c3eaf16f5d7c9bad66ea32f40a95bd169b29a91217fcc7f081375157e99c/jiter-0.16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d28bb3c26762358dadf3e5bf0bccd29ae987d65e6988d2e6f49829c76b003c09", size = 306489, upload-time = "2026-06-29T13:03:37.846Z" }, + { url = "https://files.pythonhosted.org/packages/96/3f/02fdfc6705cad96127d883af5c34e4867f554f29ec7705ec1a46156400a9/jiter-0.16.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0542a7189c26920778658fc8fcf2af8bae05bae9924577f71804acef37996536", size = 335453, upload-time = "2026-06-29T13:03:39.221Z" }, + { url = "https://files.pythonhosted.org/packages/b2/a6/e4bda5920d4b0d7c5dfb7174ce4a6b2e4d3e11c9162c452ef0eab4cdbdbd/jiter-0.16.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8fb8de1e23a0cb2a7f53c335049c7b72b6db41aa6227cdcc0972a1de5cb39450", size = 361625, upload-time = "2026-06-29T13:03:40.597Z" }, + { url = "https://files.pythonhosted.org/packages/b7/97/4e6b59b2c6e55cbb3e183595f81ad65dcfb21c915fee5e19e335df21bc55/jiter-0.16.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b72d0b2990ca754a9102779ac98d8597b7cb31678958562214a007f909eab78e", size = 456958, upload-time = "2026-06-29T13:03:42.074Z" }, + { url = "https://files.pythonhosted.org/packages/15/e0/97e9557686d2f94f4b93786eccb7eed28e9228ad132ea8237f44727314a7/jiter-0.16.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d5f91b1c27fc22a57993d5a5cb8a627cb8ed4b10502716fac1ffbfe1d19d84e8", size = 372017, upload-time = "2026-06-29T13:03:43.658Z" }, + { url = "https://files.pythonhosted.org/packages/0f/94/db768b6938e0df35c86beeba3dfbbb025c9ee5c19e1aa271f2396e50864d/jiter-0.16.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c682bea068a90b764577bdb78a60a4c1d1606daf9cd4c893832a37c7cc9d9026", size = 343320, upload-time = "2026-06-29T13:03:45.226Z" }, + { url = "https://files.pythonhosted.org/packages/c1/d6/5a59d938244a30735fe62d9433fd325f9021ea29d89780ea4596ea93bc89/jiter-0.16.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:8d031aabecc4f1b6276adfb42e3aabb77c89d468bf616600e8d3a11328929053", size = 350520, upload-time = "2026-06-29T13:03:46.671Z" }, + { url = "https://files.pythonhosted.org/packages/67/f8/c4a857f49c9af125f6bbcac7e3eee7f7978ed89682833062e2dbf62576b1/jiter-0.16.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:eab2cd170150e70153de16896a1774e3a1dca80154c56b54d7a812c479a7165e", size = 387550, upload-time = "2026-06-29T13:03:48.361Z" }, + { url = "https://files.pythonhosted.org/packages/8b/d6/5fbc2f7d6b67b754caa61a993a2e626e815dec47ffc2f9e35f01adfebec7/jiter-0.16.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:6edb63a46e65a82c26800a868e49b2cac30dd5a4218b88d74bc2c848c8ad60bb", size = 515424, upload-time = "2026-06-29T13:03:49.881Z" }, + { url = "https://files.pythonhosted.org/packages/ed/54/284f0164b64a5fed915fea6ba7e9ba9b3d8d37c67d59cf2e3bb99d45cdfe/jiter-0.16.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:659039cc50b5addcc35fcc87ae2c1833b7c0a8e5326ef631a75e4478447bcf84", size = 546981, upload-time = "2026-06-29T13:03:51.363Z" }, + { url = "https://files.pythonhosted.org/packages/13/c5/2a467585a576594384e1d2c43e1224deaafc085f24e243529cf98beef8e1/jiter-0.16.0-cp313-cp313-win32.whl", hash = "sha256:c9c53be232c2e206ef9cdbad81a48bfa74c3d3f08bcf8124630a8a748aad993e", size = 202853, upload-time = "2026-06-29T13:03:53.015Z" }, + { url = "https://files.pythonhosted.org/packages/88/6a/de61d04b9eec69c71719968d2f716532a3bc121170c44a39e14979c6be81/jiter-0.16.0-cp313-cp313-win_amd64.whl", hash = "sha256:baad945ed47f163ad833314f8e3288c396118934f94e7bbb9e243ce4b341a4fd", size = 196160, upload-time = "2026-06-29T13:03:54.447Z" }, + { url = "https://files.pythonhosted.org/packages/19/4b/b390ed59bafb3f31d008d1218578f10327714484b334439947f7e5b11e7f/jiter-0.16.0-cp313-cp313-win_arm64.whl", hash = "sha256:3c1fd2dbe1b0af19e987f03fe66c5f5bd105a2229c1aff4ab14890b24f41d21a", size = 189862, upload-time = "2026-06-29T13:03:55.754Z" }, + { url = "https://files.pythonhosted.org/packages/a7/89/bc4f1b57d5da938fd344a466396541e586d161320d70bffd929aaafcd8f4/jiter-0.16.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:b2c61484666ad42726029af0c00ef4541f0f3b5cdc550221f56c2343208018ee", size = 308239, upload-time = "2026-06-29T13:03:57.205Z" }, + { url = "https://files.pythonhosted.org/packages/65/7a/c415453e5213001bf3b411ff65dec3d303b0e76a4a2cfea9768cd4960994/jiter-0.16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:63efadc657488f45db1c676d81e704cac2abf3fdb892def1faea61db053127e2", size = 308928, upload-time = "2026-06-29T13:03:58.643Z" }, + { url = "https://files.pythonhosted.org/packages/11/fc/1f4fb7ebf9a724c7741994f4aae18fba1e2f3133df14521a79194952c34a/jiter-0.16.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cf0d73f50e7b6935677854f6e8e31d499ca7064dd24734f703e060f5b237d883", size = 336998, upload-time = "2026-06-29T13:04:00.071Z" }, + { url = "https://files.pythonhosted.org/packages/a0/8d/72cadaac05ccfa7cc3a0a2232862e6c72443ca40cf300ba8b57f9f18b69b/jiter-0.16.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bf3ea07d9bc8e7d03a9fbc051295462e6dbc295b894fd72457c3136e3e43d898", size = 362112, upload-time = "2026-06-29T13:04:01.52Z" }, + { url = "https://files.pythonhosted.org/packages/58/4a/c4b0d5f651fda90a24ffce9f8d56cde462a2e09d31ae3de3c68cef34c04e/jiter-0.16.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:26798522707abb47d767db536e4148ceac1b14446bf028ee85e579a2e043cfe5", size = 459807, upload-time = "2026-06-29T13:04:03.214Z" }, + { url = "https://files.pythonhosted.org/packages/80/58/ef77879ea9aa56b50824edc5a445e226422c7a8d211f3fd2a56bcb9493cf/jiter-0.16.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:bc837c1b9631be10abfe0191537fe8009838204cec7e44827401ace390ddb567", size = 373181, upload-time = "2026-06-29T13:04:04.629Z" }, + { url = "https://files.pythonhosted.org/packages/49/2e/ffbc3f254e4d8a66da3062c624a7df4b7c2b2cf9e1fe43cf394b3e104041/jiter-0.16.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:49060fd70737fad59d33ba9dcc0d83247dc9e77187de26053a19c16c9f32bd69", size = 344927, upload-time = "2026-06-29T13:04:06.067Z" }, + { url = "https://files.pythonhosted.org/packages/9a/f6/0be5dc6d64a89f80aa8fec984f94dedb2973e251edcae55841d60786d578/jiter-0.16.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:adbb8edeadd431bc4477879d5d371ece7cb1334486584e0f252656dd7ffada29", size = 352754, upload-time = "2026-06-29T13:04:07.477Z" }, + { url = "https://files.pythonhosted.org/packages/da/6e/7d31243b3b91cd261dd19e9d3557fc3251a80883d3d8049c86174e7ab7af/jiter-0.16.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:31aaee5b80f672c1dc21272bcfb9cbdcfc1ea04ff50f00ed5af500b80c44fa93", size = 390553, upload-time = "2026-06-29T13:04:08.92Z" }, + { url = "https://files.pythonhosted.org/packages/25/33/51ae371fde3c88897520f62b4d5f8b27ad7103e2bb10812ff52195609853/jiter-0.16.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:6722bcef4ffc86c835574b1b2fac6b33b9fb4a889c781e67950e891591f3c55a", size = 516900, upload-time = "2026-06-29T13:04:10.407Z" }, + { url = "https://files.pythonhosted.org/packages/a0/45/6449b3d123ea439ba79507c657288f461d55049e7bcbdc2cf8eb8210f491/jiter-0.16.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:5ab4f50ff971b611d656554ea10b75f80097392c827bc32923c6eeb6386c8b00", size = 548754, upload-time = "2026-06-29T13:04:12.046Z" }, + { url = "https://files.pythonhosted.org/packages/9b/e7/fd2fb11ae3e2649333da3aa170d04d7b3000bbdc3b270f6513382fdf4e04/jiter-0.16.0-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:710cc51d4ebdcd3c1f70b232c1db1ea1344a075770422bbd4bede5708335acbe", size = 122381, upload-time = "2026-06-29T13:04:13.413Z" }, + { url = "https://files.pythonhosted.org/packages/26/80/f0b147a62c315a164ed2168908286ca302310824c218d3aae52b06c0c9a9/jiter-0.16.0-cp314-cp314-win32.whl", hash = "sha256:57b37fc887a32d44798e4d8ebfa7c9683ff3da1d5bf38f08d1bb3573ccb39106", size = 204578, upload-time = "2026-06-29T13:04:14.813Z" }, + { url = "https://files.pythonhosted.org/packages/5e/e6/4758a14304b4523a6f5adb2419340086aa3593bd4327c2b25b5948a90548/jiter-0.16.0-cp314-cp314-win_amd64.whl", hash = "sha256:cbd18dd5e2df96b580487b5745adf57ef64ad89ba2d9662fc3c19386acce7db8", size = 198154, upload-time = "2026-06-29T13:04:16.272Z" }, + { url = "https://files.pythonhosted.org/packages/26/be/41fa54a2e7ea41d6c99f1dc5b1f0fd4cb474680304b5d268dd518e81da3a/jiter-0.16.0-cp314-cp314-win_arm64.whl", hash = "sha256:a32d2027a9fa67f109ff245a3252ece3ccc32cc56703e1deab6cc846a59e0585", size = 191458, upload-time = "2026-06-29T13:04:17.707Z" }, + { url = "https://files.pythonhosted.org/packages/81/6b/59127338b86d9fe4d99418f5a15118bea778103ee0fe9d9dd7e0af174e95/jiter-0.16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2577196f4474ef3fc4779a088a23b0897bbf86f9ea3679c372d45b8383b43207", size = 316739, upload-time = "2026-06-29T13:04:19.663Z" }, + { url = "https://files.pythonhosted.org/packages/2d/95/49461034d5388196d3dabf98748935f017b7785d8f3f5349f834bcc4ed0d/jiter-0.16.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:616e89e008a93c01104161c75b4988e58716b01d62307ebfe161e52a56d2a818", size = 340911, upload-time = "2026-06-29T13:04:21.257Z" }, + { url = "https://files.pythonhosted.org/packages/cd/97/a4369f2fb82cb3dda13b98622f31249b2e014b223fe64ee534413ad72294/jiter-0.16.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0e2e9efbe042210df657bade597f66d6d75723e3d8f45a12ea6d8167ff8bbce3", size = 361747, upload-time = "2026-06-29T13:04:22.677Z" }, + { url = "https://files.pythonhosted.org/packages/28/51/49b6ed456261646e1906016a6760367a28aacd3c24805e4e5fe64116c1db/jiter-0.16.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3f4d9e473a5ce7d27fef8b848df4dc16e283893d3f53b4a585e72c9595f3c284", size = 460225, upload-time = "2026-06-29T13:04:24.441Z" }, + { url = "https://files.pythonhosted.org/packages/33/b5/5689aff4f66c5b60be63106e591dbfcba2190df97d2c9c7cf052361ddb98/jiter-0.16.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8d30a4a1c87713060c8d1cc59a7b6c8fb6b8ef0a6900368014c76c87922a2929", size = 373169, upload-time = "2026-06-29T13:04:25.884Z" }, + { url = "https://files.pythonhosted.org/packages/a2/96/3ae1b85ee0d6d6cab254fb7f8da018272b932bbf2d69b07e98aa2a96c746/jiter-0.16.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bae96332410f866e5900d809298b1ed82735932986c672495f9701daacd80620", size = 350332, upload-time = "2026-06-29T13:04:27.302Z" }, + { url = "https://files.pythonhosted.org/packages/15/32/c99d7bafd78986556c95bf60ce84c6cc98786eac56066c12d7f828bb6747/jiter-0.16.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:da3d7ec75dc83bb18bca888b5edfae0656a26849056c59e05a7728badd17e7af", size = 353377, upload-time = "2026-06-29T13:04:28.731Z" }, + { url = "https://files.pythonhosted.org/packages/0e/4b/f99a8e571287c3dec766bcc18528bbe8e8fb5365522ab5e6d64c93e87066/jiter-0.16.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:ee6162b77d49a9939229df666dfa8af3e656b6701b54c4c84966d740e189264e", size = 387746, upload-time = "2026-06-29T13:04:30.319Z" }, + { url = "https://files.pythonhosted.org/packages/75/69/c78a5b3f71040e34eb5917df26fb7ae9a2174cad1ccbf277512507c53a6e/jiter-0.16.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:63ffdbdae7d4499f4cda14eadc12ddcabef0fc0c081191bdc2247489cb698077", size = 517292, upload-time = "2026-06-29T13:04:31.709Z" }, + { url = "https://files.pythonhosted.org/packages/c2/f7/095b38eda4c70d03651c403f29a5590f16d12ddc5d544aac9f9cddf72277/jiter-0.16.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:a111256a7193bea0759267b10385e5870949c239ed7b6ddbaaf57573edb38734", size = 549259, upload-time = "2026-06-29T13:04:33.721Z" }, + { url = "https://files.pythonhosted.org/packages/2e/c5/6a0207d90e5f656d95af98ebd0934f382d37674416f215aeda2ff8063e51/jiter-0.16.0-cp314-cp314t-win32.whl", hash = "sha256:de5ba8763e56b793561f43bed197c9ea55776daa5e9a6b91eed68a909bc9cdbf", size = 206523, upload-time = "2026-06-29T13:04:35.068Z" }, + { url = "https://files.pythonhosted.org/packages/a5/31/c757d5f30a8980fd945ce7b98be10be9e4ff59c7c42f5fd86804c2e87db8/jiter-0.16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b8a3f9a6008048fe9def7bf465180564a6e458047d2ce499149cfbe73c3ae9db", size = 200366, upload-time = "2026-06-29T13:04:36.61Z" }, + { url = "https://files.pythonhosted.org/packages/7c/a2/d88de6d313d734a544a7901353ad5db67cb38dcfcd91713b7979dafc345d/jiter-0.16.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0fa25b09b13075c46f5bc174f2690525a925a4fc2f7c82969a2bbabff22386ce", size = 190516, upload-time = "2026-06-29T13:04:38.004Z" }, + { url = "https://files.pythonhosted.org/packages/98/ab/664fd8c4be028b2bedd3d2ff08769c4ede23d0dbc87a77c62384a0515b5d/jiter-0.16.0-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:f17d61a28b4b3e0e3e2ba98490c70501403b4d196f78732439160e7fd3678127", size = 303106, upload-time = "2026-06-29T13:05:07.118Z" }, + { url = "https://files.pythonhosted.org/packages/1a/07/421f1d5b65493a76e16027b848aba6a7d28073ae75944fa4289cc914d39f/jiter-0.16.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:96e38eea538c8ddf853a35727c7be0741c76c13f04148ac5c116222f50ece3b3", size = 304658, upload-time = "2026-06-29T13:05:08.708Z" }, + { url = "https://files.pythonhosted.org/packages/0a/db/bba1155f01a01c3c37a89425d571da751bbedf5c54247b831a04cb971798/jiter-0.16.0-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d284fb8d94d5855d60c44fefcab4bf966f1da6fada73992b01f6f0c9bc0c6702", size = 339719, upload-time = "2026-06-29T13:05:10.41Z" }, + { url = "https://files.pythonhosted.org/packages/78/f7/18a1afcd64f35314b68c1f23afcd9994d0bc13e65cc77517afff4e83986d/jiter-0.16.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64d613743df53199b1aa256a7d328340da6d7078aac7705a7db9d7a791e9cfd2", size = 343885, upload-time = "2026-06-29T13:05:12.087Z" }, +] + [[package]] name = "jsonpatch" version = "1.33" @@ -1258,6 +1369,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/98/6beb4b351e472e5f4c4613f7c35a5290b8be2497e183825310c4c3a3984b/ruff-0.15.12-py3-none-win_arm64.whl", hash = "sha256:a538f7a82d061cee7be55542aca1d86d1393d55d81d4fcc314370f4340930d4f", size = 11120821, upload-time = "2026-04-24T18:16:57.979Z" }, ] +[[package]] +name = "sniffio" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372, upload-time = "2024-02-25T23:20:04.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, +] + [[package]] name = "sse-starlette" version = "3.4.1" From 13f17eab55e54066d801695d6a1cbb9e53d99099 Mon Sep 17 00:00:00 2001 From: uchebuzz-coder Date: Thu, 23 Jul 2026 08:17:01 +0100 Subject: [PATCH 5/7] feat(sdk/anthropic): sync gating, verdict callbacks, shared should_halt Addresses review feedback on the Anthropic integration to reach full parity with the LangChain handler. 1. Modularise should_halt: promoted from langchain_handler into ws.py as a public helper next to policy_active / block_timeout (mirrors the TS core shouldHalt). Both handlers now import it; the Anthropic integration no longer needs a local copy or a dependency on the LangChain module. 2. Fire on_verdict / on_block / on_audit for Anthropic. These callbacks only fire for events registered in the callback handler's event map, and the Anthropic handler previously emitted straight to the hook registry, bypassing it -- so the callbacks never reached the developer. Emission now delegates to the callback handler when available (registering the event and firing on_event), with the direct-hooks path kept as a fallback. 3. Gate the sync Messages.create path. Verdict futures live on the WS event loop, so sync emission + gating are bridged onto it via run_coroutine_threadsafe and the caller blocks for the verdict, matching the LangChain _sync_gate. Degrades to audit-only when gating isn't possible (no WS client, ALERT/pre-login, or a call from an event-loop thread). Tests: should_halt cases moved to test_ws.py; added notification-callback parity tests (block/audit tiers) and sync-gate tests (bridge fail-closed rewrite, event-loop-thread passthrough, _should_gate_sync matrix). --- sdk/python/adrian/__init__.py | 1 + sdk/python/adrian/anthropic_handler.py | 143 ++++++++++---- sdk/python/adrian/langchain_handler.py | 23 +-- sdk/python/adrian/ws.py | 21 +++ sdk/python/tests/test_anthropic_handler.py | 207 ++++++++++++++++++--- sdk/python/tests/test_ws.py | 37 ++++ 6 files changed, 346 insertions(+), 86 deletions(-) diff --git a/sdk/python/adrian/__init__.py b/sdk/python/adrian/__init__.py index 09d4b6f..55f4ce8 100644 --- a/sdk/python/adrian/__init__.py +++ b/sdk/python/adrian/__init__.py @@ -402,6 +402,7 @@ def patch_anthropic() -> None: hooks_getter=lambda: _hooks, config_getter=lambda: get_config() if is_initialized() else None, ws_getter=lambda: _ws_client, + handler_getter=lambda: _handler, ) diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py index 83c6465..43767d3 100644 --- a/sdk/python/adrian/anthropic_handler.py +++ b/sdk/python/adrian/anthropic_handler.py @@ -50,12 +50,13 @@ from adrian.context import get_invocation_id, set_invocation_id from adrian.format.types import AgentContext, LlmPairData, PairedEvent from adrian.hooks import HookRegistry -from adrian.proto import event_pb2 as pb from adrian.types import ChatMessage, EventData, TokenUsage, ToolCallRecord +from adrian.ws import should_halt if TYPE_CHECKING: from contextvars import Token + from adrian.handler import AdrianCallbackHandler from adrian.ws import WebSocketClient logger = logging.getLogger("adrian.anthropic") @@ -69,6 +70,7 @@ _hooks_getter: Callable[[], HookRegistry | None] | None = None _config_getter: Callable[[], AdrianConfig | None] | None = None _ws_getter: Callable[[], WebSocketClient | None] | None = None +_handler_getter: Callable[[], AdrianCallbackHandler | None] | None = None # ------------------------------------------------------------------ @@ -354,9 +356,16 @@ def build_anthropic_llm_pair( async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: """Assemble and emit a ``PairedEvent`` for a completed ``messages.create`` call. - Reads hooks and config at call time so the correct state is used even if - :func:`~adrian.shutdown` and :func:`~adrian.init` have been called since - the patch was applied. + Reads hooks / config / handler at call time so the correct state is used + even if :func:`~adrian.shutdown` and :func:`~adrian.init` have been called + since the patch was applied. + + When the callback handler is available, emission is delegated to it so the + event is registered in the handler's event map -- this is what lets the + developer ``on_verdict`` / ``on_block`` / ``on_audit`` callbacks fire for + Anthropic calls (they are keyed on that map), matching the LangChain + integration. Without a handler (e.g. ``auto_instrument=False`` before + ``init``) it falls back to emitting straight through the hook registry. Args: response: Anthropic ``Message`` response object. @@ -371,6 +380,8 @@ async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: if hooks is None or config is None: return + handler = _handler_getter() if _handler_getter is not None else None + try: session_id = config.session_id messages_param: list[dict[str, Any]] = list(kwargs.get("messages") or []) @@ -390,6 +401,13 @@ async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: run_id=run_id, ) + if handler is not None: + # Emits through the same hooks, registers the event for verdict + # enrichment, and fires on_event -- so the notification callbacks + # reach the developer for Anthropic events too. + await handler._emit_pair(pair) # pyright: ignore[reportPrivateUsage] + return + await hooks.emit(pair) if config.on_event is not None: @@ -438,31 +456,6 @@ def _schedule_emit(response: Any, kwargs: dict[str, Any]) -> None: # ------------------------------------------------------------------ -def _should_halt(verdict: pb.Verdict) -> bool: - """Decide whether a verdict should halt a tool call. - - Kept local to this module so the Anthropic integration does not depend on - the LangChain handler; the logic mirrors ``langchain_handler._should_halt``. - HITL resolutions override per-MAD policy when present. - - Args: - verdict: The classifier verdict for the producing LLM event. - - Returns: - ``True`` when the tool call must be blocked. - """ - if verdict.HasField("hitl"): - return not verdict.hitl.continue_execution - - mad_prefix = verdict.mad_code[:2] - return { - "M0": verdict.policy.policy_m0, - "M2": verdict.policy.policy_m2, - "M3": verdict.policy.policy_m3, - "M4": verdict.policy.policy_m4, - }.get(mad_prefix, False) - - def _blocked_text_block(original_block: Any) -> Any: # noqa: ANN401 """Build a text block used to replace a blocked ``tool_use`` block. @@ -531,7 +524,7 @@ def _rewrite_blocked_response(response: Any, blocked_ids: set[str]) -> Any: # n return response -async def _gate_response(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: ANN401 +async def _gate_response(response: Any, _kwargs: dict[str, Any]) -> Any: # noqa: ANN401 """Hold a response on the classifier verdict before returning it. Runs after :func:`_emit_pair` (which registers the pending verdict future @@ -547,7 +540,7 @@ async def _gate_response(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: Args: response: Anthropic ``Message`` response object. - kwargs: Original ``messages.create`` keyword arguments (unused; kept for + _kwargs: Original ``messages.create`` keyword arguments (unused; kept for symmetry with :func:`_emit_pair`). Returns: @@ -607,7 +600,7 @@ async def _gate_response(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: tc_id, ) blocked.add(tc_id) - elif _should_halt(verdict): + elif should_halt(verdict): logger.warning( "Anthropic gate: halting tool_call_id=%s event_id=%s mad_code=%s", tc_id, @@ -622,6 +615,72 @@ async def _gate_response(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: return response +def _emit_and_gate_sync(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: ANN401 + """Emit and (under BLOCK/HITL) gate a response from a synchronous call site. + + The verdict futures live on the WebSocket client's event loop, so emission + and gating must run there together: emitting on a different loop would + register the wait future where the verdict frame never resolves it. When a + gate-capable context exists -- a WS client with a running loop on another + thread, and this thread is not itself an event loop -- both steps are + bridged onto that loop via ``run_coroutine_threadsafe`` and this caller + blocks until the (possibly rewritten) response comes back. This mirrors the + LangChain ``_sync_gate`` bridge. + + In every other case (no WS client, ALERT / pre-login state, or a call made + from an event-loop thread that must not be blocked) it degrades to + audit-only emission via :func:`_schedule_emit` and returns the response + unchanged. + + Args: + response: Anthropic ``Message`` response object. + kwargs: Original ``messages.create`` keyword arguments. + + Returns: + The response, rewritten in place when a tool call was blocked. + """ + ws = _ws_getter() if _ws_getter is not None else None + + if ws is not None and _should_gate_sync(ws): + main_loop = getattr(ws, "_loop", None) + + if main_loop is not None and main_loop.is_running(): + + async def _emit_then_gate() -> Any: # noqa: ANN401 + await _emit_pair(response, kwargs) + return await _gate_response(response, kwargs) + + try: + future = asyncio.run_coroutine_threadsafe(_emit_then_gate(), main_loop) + return future.result() + except Exception: + logger.exception( + "Anthropic sync gate bridge failed; emitting audit-only" + ) + + _schedule_emit(response, kwargs) + return response + + +def _should_gate_sync(ws: WebSocketClient) -> bool: + """Whether the sync path should gate rather than emit audit-only. + + Gates only when a policy may be active (BLOCK / HITL, or pre-login where + the mode is not yet known and the async gate would fail-closed) and this + thread is not itself running an event loop -- blocking the event-loop + thread would deadlock, so those callers are left to emit and pass through. + """ + if not ws.policy_active() and ws._login_ack_received.is_set(): # pyright: ignore[reportPrivateUsage] + return False + + try: + asyncio.get_running_loop() + except RuntimeError: + return True # no loop on this thread: worker or pure-sync caller + else: + return False # on an event-loop thread: must not block it + + # ------------------------------------------------------------------ # SDK patching # ------------------------------------------------------------------ @@ -631,13 +690,15 @@ def patch_anthropic( hooks_getter: Callable[[], HookRegistry | None], config_getter: Callable[[], AdrianConfig | None], ws_getter: Callable[[], WebSocketClient | None] | None = None, + handler_getter: Callable[[], AdrianCallbackHandler | None] | None = None, ) -> None: """Monkey-patch ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic``. Wraps ``messages.create`` on both the sync and async Anthropic resource - classes so every API call is captured as an Adrian ``PairedEvent``. On the - async path, calls are additionally gated on the classifier verdict under - ``MODE_BLOCK`` / ``MODE_HITL`` (see :func:`_gate_response`). + classes so every API call is captured as an Adrian ``PairedEvent`` and, + under ``MODE_BLOCK`` / ``MODE_HITL``, gated on the classifier verdict (see + :func:`_gate_response`). Both the sync and async paths gate; the sync path + bridges onto the WebSocket loop (see :func:`_emit_and_gate_sync`). The patch is idempotent: subsequent calls update the internal getters but do not re-wrap the already-patched method. If the ``anthropic`` package is @@ -654,12 +715,17 @@ def patch_anthropic( ws_getter: Zero-arg callable returning the current ``WebSocketClient``, or ``None``. Required for verdict gating; when absent the handler stays audit-only (ALERT-mode behaviour). + handler_getter: Zero-arg callable returning the current + ``AdrianCallbackHandler``, or ``None``. Enables the developer + ``on_verdict`` / ``on_block`` / ``on_audit`` callbacks for Anthropic + events by registering them in the handler's event map. """ - global _hooks_getter, _config_getter, _ws_getter # noqa: PLW0603 + global _hooks_getter, _config_getter, _ws_getter, _handler_getter # noqa: PLW0603 _hooks_getter = hooks_getter _config_getter = config_getter _ws_getter = ws_getter + _handler_getter = handler_getter try: from anthropic.resources.messages import AsyncMessages, Messages @@ -680,8 +746,9 @@ def _patched_sync_create( **kwargs: Any, # noqa: ANN401 ) -> Any: # noqa: ANN401 response = _original_sync(self, *args, **kwargs) - _schedule_emit(response, kwargs) - return response + # Emit + gate together on the WS loop (BLOCK/HITL); degrades to + # audit-only emission when gating isn't possible. + return _emit_and_gate_sync(response, kwargs) sync_cls.create = _patched_sync_create # type: ignore[method-assign] sync_cls._adrian_patched = True # type: ignore[attr-defined] diff --git a/sdk/python/adrian/langchain_handler.py b/sdk/python/adrian/langchain_handler.py index f406791..4b3c977 100644 --- a/sdk/python/adrian/langchain_handler.py +++ b/sdk/python/adrian/langchain_handler.py @@ -40,7 +40,7 @@ from langchain_core.runnables.config import ensure_config from adrian.context import get_invocation_id, set_invocation_id -from adrian.proto import event_pb2 as pb +from adrian.ws import should_halt if TYPE_CHECKING: from adrian.config import AdrianConfig @@ -437,23 +437,6 @@ def _extract_tool_calls( # pyright: ignore[reportUnusedFunction] return [] -def _should_halt(verdict: pb.Verdict) -> bool: - """Decide whether a verdict should halt tool execution. - - HITL resolutions override per-MAD policy when present. - """ - if verdict.HasField("hitl"): - return not verdict.hitl.continue_execution - - mad_prefix = verdict.mad_code[:2] - return { - "M0": verdict.policy.policy_m0, - "M2": verdict.policy.policy_m2, - "M3": verdict.policy.policy_m3, - "M4": verdict.policy.policy_m4, - }.get(mad_prefix, False) - - def _patch_tool_node() -> None: """Patch ToolNode for callback injection + async verdict gate. @@ -592,7 +575,7 @@ async def _async_gate(tool_call_id: str) -> bool: ) return True - if _should_halt(verdict): + if should_halt(verdict): logger.warning( "halting tool execution for event_id=%s mad_code=%s", verdict.event_id, @@ -787,7 +770,7 @@ async def patched_aperform( return AgentStep( action=agent_action, observation=_BLOCKED_CONTENT ) - if _should_halt(verdict): + if should_halt(verdict): logger.warning( "halting tool execution for event_id=%s mad_code=%s", verdict.event_id, diff --git a/sdk/python/adrian/ws.py b/sdk/python/adrian/ws.py index 9b26d70..3708b8e 100644 --- a/sdk/python/adrian/ws.py +++ b/sdk/python/adrian/ws.py @@ -176,6 +176,27 @@ def _paired_event_to_proto(event: PairedEvent) -> pb.PairedEvent: return proto +def should_halt(verdict: pb.Verdict) -> bool: + """Decide whether a verdict should halt tool execution. + + Shared halt decision consumed by every integration handler (LangChain, + Anthropic, …) so the block/HITL policy is evaluated identically across + providers. Mirrors the TypeScript core ``shouldHalt`` helper. + + HITL resolutions override per-MAD policy when present. + """ + if verdict.HasField("hitl"): + return not verdict.hitl.continue_execution + + mad_prefix = verdict.mad_code[:2] + return { + "M0": verdict.policy.policy_m0, + "M2": verdict.policy.policy_m2, + "M3": verdict.policy.policy_m3, + "M4": verdict.policy.policy_m4, + }.get(mad_prefix, False) + + class WebSocketClient: """Streams ``PairedEvent`` instances to the worker core API. diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py index f54fe96..b65c240 100644 --- a/sdk/python/tests/test_anthropic_handler.py +++ b/sdk/python/tests/test_anthropic_handler.py @@ -18,6 +18,7 @@ _BLOCKED_CONTENT, _blocked_text_block, _derive_agent_id, + _emit_and_gate_sync, _emit_pair, _extract_anthropic_tool_calls, _extract_anthropic_usage, @@ -26,7 +27,6 @@ _flatten_content, _gate_response, _rewrite_blocked_response, - _should_halt, anthropic_invocation, anthropic_invocation_sync, build_anthropic_llm_pair, @@ -838,6 +838,7 @@ def _reset_getters(self) -> Any: # noqa: ANN401 _ah._ws_getter = None _ah._config_getter = None _ah._hooks_getter = None + _ah._handler_getter = None async def test_alert_mode_passes_through(self) -> None: ws = WebSocketClient("ws://x", "s", api_key="k") @@ -1063,33 +1064,6 @@ async def test_patch_anthropic_stores_ws_getter(self) -> None: assert _ah._ws_getter() is ws -class TestShouldHalt: - def test_hitl_reject_halts(self) -> None: - v = pb.Verdict(event_id="e", mad_code="M4_a") - v.hitl.continue_execution = False - assert _should_halt(v) is True - - def test_hitl_approve_does_not_halt(self) -> None: - v = pb.Verdict(event_id="e", mad_code="M4_a") - v.hitl.continue_execution = True - assert _should_halt(v) is False - - def test_policy_flag_on_halts(self) -> None: - policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m3=True) - v = pb.Verdict(event_id="e", mad_code="M3_x", policy=policy) - assert _should_halt(v) is True - - def test_policy_flag_off_does_not_halt(self) -> None: - policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m3=False) - v = pb.Verdict(event_id="e", mad_code="M3_x", policy=policy) - assert _should_halt(v) is False - - def test_unknown_mad_prefix_does_not_halt(self) -> None: - policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m4=True) - v = pb.Verdict(event_id="e", mad_code="M9_z", policy=policy) - assert _should_halt(v) is False - - class TestBlockedTextBlock: def test_dict_block_returns_dict(self) -> None: result = _blocked_text_block({"type": "tool_use", "id": "x"}) @@ -1121,3 +1095,180 @@ def test_preserves_unblocked_blocks(self) -> None: _rewrite_blocked_response(response, {"tc-1"}) assert response.content[0] == keep assert response.content[1] == {"type": "text", "text": _BLOCKED_CONTENT} + + +# ------------------------------------------------------------------ +# Developer notification callbacks (on_verdict / on_block / on_audit) +# ------------------------------------------------------------------ + + +def _make_handler(config: AdrianConfig) -> Any: # noqa: ANN401 + """Build a real AdrianCallbackHandler and wire the Anthropic getters to it.""" + from adrian.context import AgentContextTracker + from adrian.handler import AdrianCallbackHandler + from adrian.pairing import EventPairBuffer + + hooks = HookRegistry() + handler = AdrianCallbackHandler( + pair_buffer=EventPairBuffer(), + context_tracker=AgentContextTracker(), + hooks=hooks, + config=config, + ) + _ah._hooks_getter = lambda: hooks + _ah._config_getter = lambda: config + _ah._handler_getter = lambda: handler + return handler + + +class TestNotificationCallbacks: + """Anthropic events must reach on_verdict/on_block/on_audit (parity).""" + + @pytest.fixture(autouse=True) # pyright: ignore[reportUntypedFunctionDecorator] + def _reset_getters(self) -> Any: # noqa: ANN401 + yield + _ah._ws_getter = None + _ah._config_getter = None + _ah._hooks_getter = None + _ah._handler_getter = None + + async def test_emit_registers_event_in_handler_map(self) -> None: + # Without this registration the verdict callbacks can never fire. + config = AdrianConfig(session_id="s") + handler = _make_handler(config) + + await _emit_pair( + _make_text_response(text="hi"), + {"model": "m", "messages": [{"role": "user", "content": "q"}]}, + ) + + assert len(handler._event_map) == 1 # pyright: ignore[reportPrivateUsage] + + async def test_block_tier_verdict_fires_verdict_and_block(self) -> None: + fired: dict[str, int] = {"verdict": 0, "block": 0, "audit": 0} + config = AdrianConfig( + session_id="s", + on_verdict=lambda _ctx: fired.__setitem__("verdict", fired["verdict"] + 1), + on_block=lambda _ctx: fired.__setitem__("block", fired["block"] + 1), + on_audit=lambda _ctx: fired.__setitem__("audit", fired["audit"] + 1), + ) + handler = _make_handler(config) + + await _emit_pair( + _make_text_response(), + {"model": "m", "messages": [{"role": "user", "content": "q"}]}, + ) + event_id = next(iter(handler._event_map)) # pyright: ignore[reportPrivateUsage] + + policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m4=True) + await handler.handle_verdict( + pb.Verdict( + event_id=event_id, session_id="s", mad_code="M4_a", policy=policy + ) + ) + + assert fired == {"verdict": 1, "block": 1, "audit": 0} + + async def test_audit_tier_verdict_fires_verdict_and_audit(self) -> None: + fired: dict[str, int] = {"verdict": 0, "block": 0, "audit": 0} + config = AdrianConfig( + session_id="s", + on_verdict=lambda _ctx: fired.__setitem__("verdict", fired["verdict"] + 1), + on_block=lambda _ctx: fired.__setitem__("block", fired["block"] + 1), + on_audit=lambda _ctx: fired.__setitem__("audit", fired["audit"] + 1), + ) + handler = _make_handler(config) + + await _emit_pair( + _make_text_response(), + {"model": "m", "messages": [{"role": "user", "content": "q"}]}, + ) + event_id = next(iter(handler._event_map)) # pyright: ignore[reportPrivateUsage] + + policy = pb.PolicySnapshot(mode=pb.MODE_ALERT, policy_m2=True) + await handler.handle_verdict( + pb.Verdict(event_id=event_id, session_id="s", mad_code="M2", policy=policy) + ) + + assert fired == {"verdict": 1, "block": 0, "audit": 1} + + +# ------------------------------------------------------------------ +# Sync-path gating (_emit_and_gate_sync / _should_gate_sync) +# ------------------------------------------------------------------ + + +class TestSyncGate: + @pytest.fixture(autouse=True) # pyright: ignore[reportUntypedFunctionDecorator] + def _reset_getters(self) -> Any: # noqa: ANN401 + yield + _ah._ws_getter = None + _ah._config_getter = None + _ah._hooks_getter = None + _ah._handler_getter = None + + def test_should_gate_when_policy_active_and_no_loop(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + assert _ah._should_gate_sync(ws) is True + + def test_should_not_gate_in_alert(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_ALERT) + assert _ah._should_gate_sync(ws) is False + + async def test_should_not_gate_on_event_loop_thread(self) -> None: + # Called from within the running test loop: must not block it. + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + assert _ah._should_gate_sync(ws) is False + + def test_sync_bridge_gates_and_rewrites(self) -> None: + """A sync caller blocks on the WS loop and gets a rewritten response.""" + import threading + + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + try: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._loop = loop # pyright: ignore[reportPrivateUsage] - WS loop + ws._send_frame = AsyncMock() # pyright: ignore[reportPrivateUsage] + + hooks = HookRegistry() + hooks.register(ws) + _ah._hooks_getter = lambda: hooks + _ah._ws_getter = lambda: ws + _ah._config_getter = lambda: AdrianConfig(session_id="s", block_timeout=0.1) + + response = _make_tool_response(tool_id="tc-sync") + result = _emit_and_gate_sync( + response, + {"model": "m", "messages": [{"role": "user", "content": "go"}]}, + ) + + # No verdict resolves within block_timeout -> fail-closed rewrite. + assert result.content[0].text == _BLOCKED_CONTENT + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=2) + loop.close() + + async def test_sync_helper_audit_only_on_event_loop_thread(self) -> None: + """On an event-loop thread the sync path emits but does not block/gate.""" + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + hooks = HookRegistry() + _ah._hooks_getter = lambda: hooks + _ah._ws_getter = lambda: ws + _ah._config_getter = lambda: AdrianConfig(session_id="s") + + response = _make_tool_response(tool_id="tc-x") + result = _emit_and_gate_sync( + response, + {"model": "m", "messages": [{"role": "user", "content": "go"}]}, + ) + + # Passed through untouched (no blocking gate on the loop thread). + assert _block_types(result) == ["tool_use"] diff --git a/sdk/python/tests/test_ws.py b/sdk/python/tests/test_ws.py index c28762e..3b49238 100644 --- a/sdk/python/tests/test_ws.py +++ b/sdk/python/tests/test_ws.py @@ -22,6 +22,7 @@ WebSocketClient, _derive_provider, _paired_event_to_proto, + should_halt, ) @@ -364,3 +365,39 @@ async def test_empty_tool_call_id_is_skipped(self) -> None: ) assert "" not in client._tool_call_id_to_event_id + + +class TestShouldHalt: + """The shared halt decision consumed by every integration handler.""" + + def test_hitl_reject_halts(self) -> None: + v = pb.Verdict(event_id="e", mad_code="M4_a") + v.hitl.continue_execution = False + assert should_halt(v) is True + + def test_hitl_approve_does_not_halt(self) -> None: + v = pb.Verdict(event_id="e", mad_code="M4_a") + v.hitl.continue_execution = True + assert should_halt(v) is False + + def test_hitl_overrides_policy(self) -> None: + # HITL approval wins even when the per-MAD policy flag is set. + policy = pb.PolicySnapshot(mode=pb.MODE_HITL, policy_m4=True) + v = pb.Verdict(event_id="e", mad_code="M4_a", policy=policy) + v.hitl.continue_execution = True + assert should_halt(v) is False + + def test_policy_flag_on_halts(self) -> None: + policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m3=True) + v = pb.Verdict(event_id="e", mad_code="M3_x", policy=policy) + assert should_halt(v) is True + + def test_policy_flag_off_does_not_halt(self) -> None: + policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m3=False) + v = pb.Verdict(event_id="e", mad_code="M3_x", policy=policy) + assert should_halt(v) is False + + def test_unknown_mad_prefix_does_not_halt(self) -> None: + policy = pb.PolicySnapshot(mode=pb.MODE_BLOCK, policy_m4=True) + v = pb.Verdict(event_id="e", mad_code="M9_z", policy=policy) + assert should_halt(v) is False From a1eae31511ffbbf92f776baededc001429cc20cb Mon Sep 17 00:00:00 2001 From: uchebuzz-coder Date: Fri, 24 Jul 2026 09:20:07 +0100 Subject: [PATCH 6/7] fix(sdk/anthropic): sync gate must fail closed on error / no WS loop _emit_and_gate_sync had two fail-open holes once gating was engaged: a bridge/gate exception was logged and fell through to returning the original (unblocked) response, and the loop-less path degraded to audit-only instead of gating. Under BLOCK, both let a halted tool call through. Restructured so that once _should_gate_sync engages, every exit fails closed: the no-WS-loop path now runs asyncio.run(_emit_then_gate()) (gate fail-closes via verdict timeout, matching LangChain's asyncio.run fallback), and any bridge/gate exception rewrites the response's tool calls to [BLOCKED]. The audit-only pass-through now applies only when gating is not engaged (no WS client, inactive policy, or an event-loop thread). Matches langchain_handler._sync_gate (except Exception: return True). Tests: bridge exception under BLOCK -> [BLOCKED]; no running WS loop -> fail-closed via timeout. --- sdk/python/adrian/anthropic_handler.py | 75 ++++++++++++++-------- sdk/python/tests/test_anthropic_handler.py | 60 +++++++++++++++++ 2 files changed, 108 insertions(+), 27 deletions(-) diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py index 43767d3..c8b6341 100644 --- a/sdk/python/adrian/anthropic_handler.py +++ b/sdk/python/adrian/anthropic_handler.py @@ -621,45 +621,66 @@ def _emit_and_gate_sync(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: The verdict futures live on the WebSocket client's event loop, so emission and gating must run there together: emitting on a different loop would register the wait future where the verdict frame never resolves it. When a - gate-capable context exists -- a WS client with a running loop on another - thread, and this thread is not itself an event loop -- both steps are - bridged onto that loop via ``run_coroutine_threadsafe`` and this caller - blocks until the (possibly rewritten) response comes back. This mirrors the - LangChain ``_sync_gate`` bridge. - - In every other case (no WS client, ALERT / pre-login state, or a call made - from an event-loop thread that must not be blocked) it degrades to - audit-only emission via :func:`_schedule_emit` and returns the response - unchanged. + WS loop is running on another thread both steps are bridged onto it via + ``run_coroutine_threadsafe``; otherwise they run to completion here via + ``asyncio.run``. Either way the caller blocks until the (possibly + rewritten) response comes back. This mirrors the LangChain ``_sync_gate``. + + Once gating is engaged (see :func:`_should_gate_sync`) every exit fails + **closed**: a bridge/gate error blocks the response's tool calls rather than + letting them through, matching ``langchain_handler._sync_gate``. Only when + gating is not engaged at all -- no WS client, ALERT / post-login inactive + policy, or a call from an event-loop thread that must not be blocked -- does + it degrade to audit-only emission and return the response unchanged. Args: response: Anthropic ``Message`` response object. kwargs: Original ``messages.create`` keyword arguments. Returns: - The response, rewritten in place when a tool call was blocked. + The response, rewritten in place when a tool call was (or must be, + on error) blocked. """ ws = _ws_getter() if _ws_getter is not None else None - if ws is not None and _should_gate_sync(ws): - main_loop = getattr(ws, "_loop", None) + # Not gating: no backend, inactive policy, or an event-loop thread we must + # not block -- emit for audit and pass the response through unchanged. + if ws is None or not _should_gate_sync(ws): + _schedule_emit(response, kwargs) + return response - if main_loop is not None and main_loop.is_running(): + # Gating is engaged: every path below must fail closed. + async def _emit_then_gate() -> Any: # noqa: ANN401 + await _emit_pair(response, kwargs) + return await _gate_response(response, kwargs) - async def _emit_then_gate() -> Any: # noqa: ANN401 - await _emit_pair(response, kwargs) - return await _gate_response(response, kwargs) - - try: - future = asyncio.run_coroutine_threadsafe(_emit_then_gate(), main_loop) - return future.result() - except Exception: - logger.exception( - "Anthropic sync gate bridge failed; emitting audit-only" - ) + main_loop = getattr(ws, "_loop", None) - _schedule_emit(response, kwargs) - return response + try: + if main_loop is not None and main_loop.is_running(): + # WS loop runs on another thread: bridge onto it and block here. + return asyncio.run_coroutine_threadsafe( + _emit_then_gate(), main_loop + ).result() + + # No WS loop on another thread: run to completion on a temporary loop. + # With no live connection no verdict arrives, so the gate fail-closes + # via timeout -- the same outcome as LangChain's asyncio.run fallback. + return asyncio.run(_emit_then_gate()) + except Exception: + # Fail closed: a bridge/gate failure must not let a halted tool call + # through. Block every tool call in the response. + logger.exception( + "Anthropic sync gate failed; failing closed (blocking tool calls)" + ) + blocked = { + rec["id"] + for rec in _extract_anthropic_tool_calls( + getattr(response, "content", []) or [] + ) + if rec["id"] + } + return _rewrite_blocked_response(response, blocked) def _should_gate_sync(ws: WebSocketClient) -> bool: diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py index b65c240..ae9cf20 100644 --- a/sdk/python/tests/test_anthropic_handler.py +++ b/sdk/python/tests/test_anthropic_handler.py @@ -1272,3 +1272,63 @@ async def test_sync_helper_audit_only_on_event_loop_thread(self) -> None: # Passed through untouched (no blocking gate on the loop thread). assert _block_types(result) == ["tool_use"] + + def test_sync_bridge_exception_fails_closed( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A bridge/gate error must fail closed, not let the tool call through.""" + import threading + + async def _boom(*_args: Any, **_kwargs: Any) -> Any: # noqa: ANN401 + raise RuntimeError("gate exploded") + + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + try: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._loop = loop # pyright: ignore[reportPrivateUsage] - WS loop + ws._send_frame = AsyncMock() # pyright: ignore[reportPrivateUsage] + + hooks = HookRegistry() + hooks.register(ws) + _ah._hooks_getter = lambda: hooks + _ah._ws_getter = lambda: ws + _ah._config_getter = lambda: AdrianConfig(session_id="s") + # Gate raises after emit -> the bridge re-raises on .result(). + monkeypatch.setattr(_ah, "_gate_response", _boom) + + response = _make_tool_response(tool_id="tc-err") + result = _emit_and_gate_sync( + response, + {"model": "m", "messages": [{"role": "user", "content": "go"}]}, + ) + + assert result.content[0].text == _BLOCKED_CONTENT + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=2) + loop.close() + + def test_sync_no_ws_loop_fails_closed_via_timeout(self) -> None: + """No running WS loop -> the asyncio.run path fail-closes via timeout.""" + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._loop = None # pyright: ignore[reportPrivateUsage] - no WS loop + ws._send_frame = AsyncMock() # pyright: ignore[reportPrivateUsage] + + hooks = HookRegistry() + hooks.register(ws) + _ah._hooks_getter = lambda: hooks + _ah._ws_getter = lambda: ws + _ah._config_getter = lambda: AdrianConfig(session_id="s", block_timeout=0.1) + + response = _make_tool_response(tool_id="tc-noloop") + result = _emit_and_gate_sync( + response, + {"model": "m", "messages": [{"role": "user", "content": "go"}]}, + ) + + # No verdict can arrive with no live loop -> fail-closed rewrite. + assert result.content[0].text == _BLOCKED_CONTENT From 157a196bb1a1fc5894b8fce8937c63b1f3967a8e Mon Sep 17 00:00:00 2001 From: Uchebuzz Date: Wed, 29 Jul 2026 10:50:53 +0100 Subject: [PATCH 7/7] feat(sdk/anthropic): instrument and gate messages.stream() Addresses two review findings on the Anthropic integration. invocation_id was silently no_invocation for any call not wrapped in anthropic_invocation(). Anthropic exposes no graph-like unit of work to scope an invocation to (unlike Pregel.ainvoke), so the sentinel stays -- _resolve_invocation_id() now mirrors the LangChain handler and logs at INFO, naming the wrapper that fixes it. Nothing is minted, so invocation semantics are unchanged. client.messages.stream() was not instrumented at all: no event, no gate, so streaming bypassed BLOCK/HITL entirely. stream() returns a context manager rather than a response, so the manager is wrapped and __enter__ hands back the real MessageStream with get_final_message / until_done rerouted through the existing _emit_pair + _gate_response. text_stream, response, request_id and close are untouched. get_final_text() delegates to get_final_message(), so it is covered too. Notes: - get_final_message() calls self.until_done() internally, an instance attribute lookup that reaches the override, so emission is once-guarded. Safe because _rewrite_blocked_response mutates in place and both paths hold the same snapshot object. - The invocation id is sampled at stream() call time, not at consumption time: the caller may read the final message after leaving the anthropic_invocation() block. Hence the optional invocation_id keyword on _emit_pair / _emit_and_gate_sync. - A text_stream-only consumer still leaves an audit trail via an audit-only emit on the manager __exit__. Raw event iteration stays audited-but-not-gated; a consumer acting directly on content_block_stop events bypasses the gate. Documented in the module docstring, to be closed in a follow-up. --- examples/python/anthropic_quickstart.py | 3 +- examples/python/anthropic_streaming.py | 74 +++ sdk/python/adrian/anthropic_handler.py | 360 +++++++++++++- sdk/python/tests/test_anthropic_handler.py | 536 +++++++++++++++++++++ 4 files changed, 951 insertions(+), 22 deletions(-) create mode 100644 examples/python/anthropic_streaming.py diff --git a/examples/python/anthropic_quickstart.py b/examples/python/anthropic_quickstart.py index 7af2590..cd3ec7b 100644 --- a/examples/python/anthropic_quickstart.py +++ b/examples/python/anthropic_quickstart.py @@ -9,7 +9,6 @@ Run:: export ANTHROPIC_API_KEY="sk-ant-..." - export ADRIAN_API_KEY="..." # optional -- omit to collect locally only python examples/python/anthropic_quickstart.py """ @@ -76,7 +75,7 @@ async def main() -> None: # ------------------------------------------------------------------ # 4. Always shut down Adrian cleanly to flush any pending events. # ------------------------------------------------------------------ - await adrian.shutdown() + adrian.shutdown() print("Done. Check your Adrian dashboard for the captured events.") diff --git a/examples/python/anthropic_streaming.py b/examples/python/anthropic_streaming.py new file mode 100644 index 0000000..5c5bdc1 --- /dev/null +++ b/examples/python/anthropic_streaming.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 SecureAgentics +# +# Licensed under the Apache Licence, Version 2.0 (the "Licence"). +# You may not use this file except in compliance with the Licence. +# A copy of the Licence is included at LICENSE in the repository root. +"""Streaming quickstart: monitor Anthropic streamed calls with Adrian. + +The streaming counterpart to ``anthropic_quickstart.py``. Text deltas arrive as +usual; the Adrian event is emitted -- and, under BLOCK / HITL, the verdict gate +runs -- when the final message is requested. + +Run:: + + export ANTHROPIC_API_KEY="sk-ant-..." + python examples/python/anthropic_streaming.py +""" + +from __future__ import annotations + +import asyncio +import os + +import anthropic +import adrian + +# ------------------------------------------------------------------ +# 1. Initialise Adrian. This auto-instruments Anthropic by default. +# ------------------------------------------------------------------ +adrian.init( + api_key=os.environ.get("ADRIAN_API_KEY", ""), + session_id="anthropic-streaming-session", +) + +client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"]) + + +async def main() -> None: + print("Streaming response...\n") + + # ------------------------------------------------------------------ + # 2. Wrap the call so the event carries a real invocation_id. Without + # this it is emitted as "no_invocation" -- a raw Anthropic call has no + # unit of work for Adrian to scope an invocation to. + # ------------------------------------------------------------------ + async with adrian.anthropic_invocation(): + async with client.messages.stream( + model="claude-haiku-4-5-20251001", + max_tokens=256, + system="You are a concise assistant.", + messages=[{"role": "user", "content": "Count to five, one word per line."}], + ) as stream: + # Text deltas stream through untouched. + async for text in stream.text_stream: + print(text, end="", flush=True) + + # ------------------------------------------------------------------ + # 3. The Adrian event is emitted here. Under BLOCK / HITL this also + # holds for the classifier verdict and rewrites any halted + # tool_use block to "[BLOCKED by security policy]". + # ------------------------------------------------------------------ + message = await stream.get_final_message() + + print(f"\n\nStop reason: {message.stop_reason}") + + # ------------------------------------------------------------------ + # 4. Always shut down Adrian cleanly to flush any pending events. + # ------------------------------------------------------------------ + adrian.shutdown() + print("Done. Check your Adrian dashboard for the captured event.") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py index c8b6341..a405a3f 100644 --- a/sdk/python/adrian/anthropic_handler.py +++ b/sdk/python/adrian/anthropic_handler.py @@ -7,10 +7,10 @@ """Anthropic SDK instrumentation for Adrian. Patches ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic`` so that every -``messages.create`` call is captured as an Adrian ``PairedEvent`` and emitted -through the hook registry. The patch is idempotent; calling -:func:`patch_anthropic` again after a shutdown / re-init only updates the -internal getters, it does not re-wrap the already-patched method. +``messages.create`` and ``messages.stream`` call is captured as an Adrian +``PairedEvent`` and emitted through the hook registry. The patch is idempotent; +calling :func:`patch_anthropic` again after a shutdown / re-init only updates the +internal getters, it does not re-wrap the already-patched methods. Usage without auto-instrumentation:: @@ -28,6 +28,36 @@ async with adrian.anthropic_invocation(): r1 = await client.messages.create(...) r2 = await client.messages.create(...) # same invocation_id as r1 + +Without that wrapper each call is emitted with ``invocation_id="no_invocation"`` +and an INFO log line, because the Anthropic SDK -- unlike LangGraph, where +``Pregel.ainvoke`` bounds a unit of work -- exposes only point requests with no +boundary to scope an invocation to. + +Streaming +--------- + +``client.messages.stream(...)`` returns a context manager rather than a +response, so the returned manager is wrapped and the stream's terminal methods +are rerouted through the same emit + gate path as ``create``:: + + with client.messages.stream(model="...", ...) as stream: + for text in stream.text_stream: + print(text, end="") + + message = stream.get_final_message() # emitted and gated here + +Under ``MODE_BLOCK`` / ``MODE_HITL`` a halted ``tool_use`` block in the final +message is rewritten to a ``[BLOCKED]`` text block exactly as in the +non-streaming path. Known limits: + +* Gating happens at ``get_final_message()`` / ``until_done()`` (and so also at + ``get_final_text()``, which delegates to the former). A consumer that acts on + raw ``content_block_stop`` events instead is emitted for audit (on + context-manager exit) but not gated -- raw event iteration is planned for a + follow-up change. +* A ``stream()`` call used outside a ``with`` block is not instrumented at all; + ``__enter__`` is the seam. """ # pyright: reportUnknownVariableType=false @@ -353,7 +383,45 @@ def build_anthropic_llm_pair( # ------------------------------------------------------------------ -async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: +def _resolve_invocation_id(captured: str | None = None) -> str: + """Resolve the invocation ID for an emitted event. + + Unlike LangGraph -- where ``Pregel.ainvoke`` is a natural unit of work the + patch can scope an invocation to -- the Anthropic SDK exposes only point + requests, so an unwrapped call genuinely has no invocation to belong to. + Mirrors ``AdrianCallbackHandler._resolve_invocation_id``: log and fall back + to the sentinel rather than inventing an ID. + + Args: + captured: Invocation ID sampled earlier, used by the streaming path + where the context variable may be out of scope by the time the + stream is consumed. ``None`` falls back to the current context. + + Returns: + The resolved invocation ID, or ``"no_invocation"``. + """ + invocation_id = captured or get_invocation_id() + + if invocation_id is None: + # The event is still emitted; it just cannot be correlated with the + # other calls in the same logical task. INFO not WARN: nothing is + # dropped. + logger.info( + "Anthropic call made outside of an invocation context; event will " + "be emitted with invocation_id=no_invocation. Wrap related calls " + "in adrian.anthropic_invocation() to correlate them." + ) + return "no_invocation" + + return invocation_id + + +async def _emit_pair( + response: Any, # noqa: ANN401 + kwargs: dict[str, Any], + *, + invocation_id: str | None = None, +) -> None: """Assemble and emit a ``PairedEvent`` for a completed ``messages.create`` call. Reads hooks / config / handler at call time so the correct state is used @@ -370,6 +438,8 @@ async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: Args: response: Anthropic ``Message`` response object. kwargs: Original ``messages.create`` keyword arguments. + invocation_id: Invocation ID sampled at call time; see + :func:`_resolve_invocation_id`. ``None`` reads the current context. """ if _hooks_getter is None or _config_getter is None: return @@ -389,7 +459,7 @@ async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: model_param: str = str(kwargs.get("model", "unknown")) flat_messages = _flatten_anthropic_messages(messages_param, system_param) - invocation_id = get_invocation_id() or "no_invocation" + resolved_invocation_id = _resolve_invocation_id(invocation_id) run_id = str(uuid4()) pair = build_anthropic_llm_pair( @@ -397,7 +467,7 @@ async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: response=response, model=model_param, session_id=session_id, - invocation_id=invocation_id, + invocation_id=resolved_invocation_id, run_id=run_id, ) @@ -428,7 +498,12 @@ async def _emit_pair(response: Any, kwargs: dict[str, Any]) -> None: logger.exception("Failed to emit Anthropic paired event") -def _schedule_emit(response: Any, kwargs: dict[str, Any]) -> None: +def _schedule_emit( + response: Any, # noqa: ANN401 + kwargs: dict[str, Any], + *, + invocation_id: str | None = None, +) -> None: """Schedule event emission from a synchronous call site. When inside a running event loop, schedules a fire-and-forget task so @@ -438,8 +513,9 @@ def _schedule_emit(response: Any, kwargs: dict[str, Any]) -> None: Args: response: Anthropic ``Message`` response object. kwargs: Original ``messages.create`` keyword arguments. + invocation_id: Invocation ID sampled at call time, or ``None``. """ - coro = _emit_pair(response, kwargs) + coro = _emit_pair(response, kwargs, invocation_id=invocation_id) try: loop = asyncio.get_running_loop() @@ -615,7 +691,12 @@ async def _gate_response(response: Any, _kwargs: dict[str, Any]) -> Any: # noqa return response -def _emit_and_gate_sync(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: ANN401 +def _emit_and_gate_sync( + response: Any, # noqa: ANN401 + kwargs: dict[str, Any], + *, + invocation_id: str | None = None, +) -> Any: # noqa: ANN401 """Emit and (under BLOCK/HITL) gate a response from a synchronous call site. The verdict futures live on the WebSocket client's event loop, so emission @@ -636,6 +717,7 @@ def _emit_and_gate_sync(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: Args: response: Anthropic ``Message`` response object. kwargs: Original ``messages.create`` keyword arguments. + invocation_id: Invocation ID sampled at call time, or ``None``. Returns: The response, rewritten in place when a tool call was (or must be, @@ -646,12 +728,12 @@ def _emit_and_gate_sync(response: Any, kwargs: dict[str, Any]) -> Any: # noqa: # Not gating: no backend, inactive policy, or an event-loop thread we must # not block -- emit for audit and pass the response through unchanged. if ws is None or not _should_gate_sync(ws): - _schedule_emit(response, kwargs) + _schedule_emit(response, kwargs, invocation_id=invocation_id) return response # Gating is engaged: every path below must fail closed. async def _emit_then_gate() -> Any: # noqa: ANN401 - await _emit_pair(response, kwargs) + await _emit_pair(response, kwargs, invocation_id=invocation_id) return await _gate_response(response, kwargs) main_loop = getattr(ws, "_loop", None) @@ -702,6 +784,210 @@ def _should_gate_sync(ws: WebSocketClient) -> bool: return False # on an event-loop thread: must not block it +# ------------------------------------------------------------------ +# Streaming (messages.stream) +# ------------------------------------------------------------------ + + +def _safe_snapshot(stream: Any) -> Any: # noqa: ANN401 + """Read a stream's accumulated message snapshot without raising. + + ``MessageStream.current_message_snapshot`` asserts the snapshot exists, so + it raises for a stream that was never consumed. + + Args: + stream: ``MessageStream`` or ``AsyncMessageStream``. + + Returns: + The accumulated ``Message``, or ``None`` when unavailable. + """ + try: + return stream.current_message_snapshot + except Exception: # noqa: BLE001 - AssertionError when nothing accumulated + return None + + +class _StreamGateState: + """Runs emit + gate exactly once for one streamed message. + + ``MessageStream.get_final_message`` calls ``self.until_done()`` internally, + and both are instrumented, so the run is once-guarded. The guard is safe + because :func:`_rewrite_blocked_response` mutates the message in place and + both paths hold the same snapshot object: a second call returns a message + that has already been gated and rewritten. + """ + + def __init__(self, kwargs: dict[str, Any], invocation_id: str | None) -> None: + self._kwargs = kwargs + self._invocation_id = invocation_id + self._done = False + + @property + def done(self) -> bool: + """Whether emit + gate has already run for this stream.""" + return self._done + + def run_sync(self, message: Any) -> Any: # noqa: ANN401 + """Emit and gate ``message`` from a synchronous consumer.""" + if self._done or message is None: + return message + + self._done = True + + return _emit_and_gate_sync( + message, self._kwargs, invocation_id=self._invocation_id + ) + + async def run_async(self, message: Any) -> Any: # noqa: ANN401 + """Emit and gate ``message`` from an asynchronous consumer.""" + if self._done or message is None: + return message + + self._done = True + + await _emit_pair(message, self._kwargs, invocation_id=self._invocation_id) + + return await _gate_response(message, self._kwargs) + + def emit_audit_only_sync(self, message: Any) -> None: # noqa: ANN401 + """Emit without gating, for a stream consumed without a final message.""" + if self._done or message is None: + return + + self._done = True + _schedule_emit(message, self._kwargs, invocation_id=self._invocation_id) + + async def emit_audit_only_async(self, message: Any) -> None: # noqa: ANN401 + """Async counterpart to :meth:`emit_audit_only_sync`.""" + if self._done or message is None: + return + + self._done = True + await _emit_pair(message, self._kwargs, invocation_id=self._invocation_id) + + +def _instrument_sync_stream(stream: Any, state: _StreamGateState) -> None: # noqa: ANN401 + """Route a ``MessageStream``'s terminal methods through emit + gate. + + Rebinds ``get_final_message`` / ``until_done`` as instance attributes so the + real SDK object is handed back to the caller untouched otherwise -- + ``text_stream``, ``response``, ``request_id`` and ``close`` keep working. + + Args: + stream: The ``MessageStream`` returned by the manager's ``__enter__``. + state: Once-guarded emit + gate runner for this stream. + """ + original_final = stream.get_final_message + original_until = stream.until_done + + def until_done() -> None: + original_until() + state.run_sync(_safe_snapshot(stream)) + + def get_final_message() -> Any: # noqa: ANN401 + # original_final() calls until_done() above, which already ran the + # gate; the once-guard makes this a pass-through of the gated message. + return state.run_sync(original_final()) + + stream.until_done = until_done + stream.get_final_message = get_final_message + + +def _instrument_async_stream(stream: Any, state: _StreamGateState) -> None: # noqa: ANN401 + """Async counterpart to :func:`_instrument_sync_stream`.""" + original_final = stream.get_final_message + original_until = stream.until_done + + async def until_done() -> None: + await original_until() + await state.run_async(_safe_snapshot(stream)) + + async def get_final_message() -> Any: # noqa: ANN401 + return await state.run_async(await original_final()) + + stream.until_done = until_done + stream.get_final_message = get_final_message + + +class _GatedMessageStreamManager: + """Wraps ``MessageStreamManager`` so the streamed message is emitted + gated. + + ``client.messages.stream(...)`` returns a context manager rather than a + response, so the instrumentation seam is ``__enter__``: the real + ``MessageStream`` is handed back with its terminal methods rerouted (see + :func:`_instrument_sync_stream`). + + ``__exit__`` is a safety net -- a consumer that only reads ``text_stream`` + never calls a terminal method, and would otherwise leave no audit trail. + That emission is deliberately audit-only: by ``__exit__`` the caller has + already seen every block, so gating there would change nothing. + """ + + def __init__( + self, + inner: Any, # noqa: ANN401 + kwargs: dict[str, Any], + invocation_id: str | None, + ) -> None: + self._inner = inner + self._stream: Any = None + self._state = _StreamGateState(kwargs, invocation_id) + + def __enter__(self) -> Any: # noqa: ANN401 + stream = self._inner.__enter__() + self._stream = stream + + try: + _instrument_sync_stream(stream, self._state) + except Exception: + logger.exception("Failed to instrument Anthropic message stream") + + return stream + + def __exit__(self, *exc_info: Any) -> Any: # noqa: ANN401 + try: + if not self._state.done and self._stream is not None: + self._state.emit_audit_only_sync(_safe_snapshot(self._stream)) + except Exception: + logger.exception("Failed to emit Anthropic stream event on exit") + + return self._inner.__exit__(*exc_info) + + +class _GatedAsyncMessageStreamManager: + """Async counterpart to :class:`_GatedMessageStreamManager`.""" + + def __init__( + self, + inner: Any, # noqa: ANN401 + kwargs: dict[str, Any], + invocation_id: str | None, + ) -> None: + self._inner = inner + self._stream: Any = None + self._state = _StreamGateState(kwargs, invocation_id) + + async def __aenter__(self) -> Any: # noqa: ANN401 + stream = await self._inner.__aenter__() + self._stream = stream + + try: + _instrument_async_stream(stream, self._state) + except Exception: + logger.exception("Failed to instrument Anthropic message stream") + + return stream + + async def __aexit__(self, *exc_info: Any) -> Any: # noqa: ANN401 + try: + if not self._state.done and self._stream is not None: + await self._state.emit_audit_only_async(_safe_snapshot(self._stream)) + except Exception: + logger.exception("Failed to emit Anthropic stream event on exit") + + return await self._inner.__aexit__(*exc_info) + + # ------------------------------------------------------------------ # SDK patching # ------------------------------------------------------------------ @@ -715,14 +1001,16 @@ def patch_anthropic( ) -> None: """Monkey-patch ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic``. - Wraps ``messages.create`` on both the sync and async Anthropic resource - classes so every API call is captured as an Adrian ``PairedEvent`` and, - under ``MODE_BLOCK`` / ``MODE_HITL``, gated on the classifier verdict (see - :func:`_gate_response`). Both the sync and async paths gate; the sync path - bridges onto the WebSocket loop (see :func:`_emit_and_gate_sync`). + Wraps ``messages.create`` and ``messages.stream`` on both the sync and async + Anthropic resource classes so every API call is captured as an Adrian + ``PairedEvent`` and, under ``MODE_BLOCK`` / ``MODE_HITL``, gated on the + classifier verdict (see :func:`_gate_response`). Both the sync and async + paths gate; the sync path bridges onto the WebSocket loop (see + :func:`_emit_and_gate_sync`). For ``stream`` the gate runs when the caller + asks for the final message (see :class:`_GatedMessageStreamManager`). The patch is idempotent: subsequent calls update the internal getters but - do not re-wrap the already-patched method. If the ``anthropic`` package is + do not re-wrap the already-patched methods. If the ``anthropic`` package is not installed the call is a silent no-op. This function is called automatically by :func:`~adrian.init` when @@ -771,9 +1059,25 @@ def _patched_sync_create( # audit-only emission when gating isn't possible. return _emit_and_gate_sync(response, kwargs) + _original_sync_stream = sync_cls.stream + + def _patched_sync_stream( + self: Any, + *args: Any, + **kwargs: Any, # noqa: ANN401 + ) -> Any: # noqa: ANN401 + # Sampled now, not at consumption time: the caller may read the + # final message after leaving the anthropic_invocation() block. + captured = get_invocation_id() + + return _GatedMessageStreamManager( + _original_sync_stream(self, *args, **kwargs), kwargs, captured + ) + sync_cls.create = _patched_sync_create # type: ignore[method-assign] + sync_cls.stream = _patched_sync_stream # type: ignore[method-assign] sync_cls._adrian_patched = True # type: ignore[attr-defined] - logger.debug("Patched anthropic.resources.Messages.create") + logger.debug("Patched anthropic.resources.Messages.create / stream") except AttributeError: logger.warning( "Could not patch anthropic.resources.Messages; " @@ -799,9 +1103,25 @@ async def _patched_async_create( await _emit_pair(response, kwargs) return await _gate_response(response, kwargs) + _original_async_stream = async_cls.stream + + def _patched_async_stream( + self: Any, + *args: Any, + **kwargs: Any, # noqa: ANN401 + ) -> Any: # noqa: ANN401 + # Not async: stream() returns an async context manager without + # being awaited itself. + captured = get_invocation_id() + + return _GatedAsyncMessageStreamManager( + _original_async_stream(self, *args, **kwargs), kwargs, captured + ) + async_cls.create = _patched_async_create # type: ignore[method-assign] + async_cls.stream = _patched_async_stream # type: ignore[method-assign] async_cls._adrian_patched = True # type: ignore[attr-defined] - logger.debug("Patched anthropic.resources.AsyncMessages.create") + logger.debug("Patched anthropic.resources.AsyncMessages.create / stream") except AttributeError: logger.warning( "Could not patch anthropic.resources.AsyncMessages; " diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py index ae9cf20..eb926de 100644 --- a/sdk/python/tests/test_anthropic_handler.py +++ b/sdk/python/tests/test_anthropic_handler.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import logging from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -26,7 +27,11 @@ _flatten_anthropic_messages, _flatten_content, _gate_response, + _GatedAsyncMessageStreamManager, + _GatedMessageStreamManager, + _resolve_invocation_id, _rewrite_blocked_response, + _safe_snapshot, anthropic_invocation, anthropic_invocation_sync, build_anthropic_llm_pair, @@ -1332,3 +1337,534 @@ def test_sync_no_ws_loop_fails_closed_via_timeout(self) -> None: # No verdict can arrive with no live loop -> fail-closed rewrite. assert result.content[0].text == _BLOCKED_CONTENT + + +# ------------------------------------------------------------------ +# Invocation ID resolution +# ------------------------------------------------------------------ + + +_KWARGS: dict[str, Any] = { + "model": "m", + "messages": [{"role": "user", "content": "q"}], +} + + +class TestResolveInvocationId: + def test_returns_context_id(self) -> None: + token = set_invocation_id("ctx-id") + + try: + assert _resolve_invocation_id() == "ctx-id" + finally: + token.var.reset(token) + + def test_captured_wins_over_context(self) -> None: + """The streaming path samples the ID early; that sample takes priority.""" + token = set_invocation_id("ctx-id") + + try: + assert _resolve_invocation_id("captured-id") == "captured-id" + finally: + token.var.reset(token) + + def test_falls_back_to_sentinel(self) -> None: + assert _resolve_invocation_id() == "no_invocation" + + def test_logs_info_when_uncorrelated( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """A developer with no invocation context is told why, at INFO.""" + with caplog.at_level(logging.INFO, logger="adrian.anthropic"): + _resolve_invocation_id() + + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.INFO + assert "anthropic_invocation()" in caplog.records[0].getMessage() + + def test_logs_nothing_when_correlated( + self, caplog: pytest.LogCaptureFixture + ) -> None: + token = set_invocation_id("ctx-id") + + try: + with caplog.at_level(logging.INFO, logger="adrian.anthropic"): + _resolve_invocation_id() + finally: + token.var.reset(token) + + assert caplog.records == [] + + async def test_emit_pair_honours_captured_id(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + try: + await _emit_pair( + _make_text_response(), _KWARGS, invocation_id="captured-id" + ) + finally: + _ah._hooks_getter = None + _ah._config_getter = None + + assert collector.events[0].invocation_id == "captured-id" + + +# ------------------------------------------------------------------ +# Streaming (messages.stream) +# ------------------------------------------------------------------ + + +class _FakeMessageStream: + """Stand-in for ``anthropic.lib.streaming.MessageStream``. + + Reproduces the one behaviour the instrumentation depends on: + ``get_final_message()`` calls ``self.until_done()``, so an instance-level + override of ``until_done`` is reached from inside it. + """ + + def __init__(self, message: Any) -> None: # noqa: ANN401 + self._message = message + self.until_done_calls = 0 + self.closed = False + + @property + def current_message_snapshot(self) -> Any: # noqa: ANN401 + return self._message + + def until_done(self) -> None: + self.until_done_calls += 1 + + def get_final_message(self) -> Any: # noqa: ANN401 + self.until_done() + return self._message + + def close(self) -> None: + self.closed = True + + +class _FakeStreamManager: + """Stand-in for ``MessageStreamManager``.""" + + def __init__(self, stream: Any) -> None: # noqa: ANN401 + self._stream = stream + self.exited = False + + def __enter__(self) -> Any: # noqa: ANN401 + return self._stream + + def __exit__(self, *_exc: Any) -> None: + self.exited = True + + +class _FakeAsyncMessageStream: + """Async counterpart to :class:`_FakeMessageStream`.""" + + def __init__(self, message: Any) -> None: # noqa: ANN401 + self._message = message + self.until_done_calls = 0 + self.closed = False + + @property + def current_message_snapshot(self) -> Any: # noqa: ANN401 + return self._message + + async def until_done(self) -> None: + self.until_done_calls += 1 + + async def get_final_message(self) -> Any: # noqa: ANN401 + await self.until_done() + return self._message + + async def close(self) -> None: + self.closed = True + + +class _FakeAsyncStreamManager: + """Stand-in for ``AsyncMessageStreamManager``.""" + + def __init__(self, stream: Any) -> None: # noqa: ANN401 + self._stream = stream + self.exited = False + + async def __aenter__(self) -> Any: # noqa: ANN401 + return self._stream + + async def __aexit__(self, *_exc: Any) -> None: + self.exited = True + + +def _gated_async( + message: Any, # noqa: ANN401 + invocation_id: str | None = None, +) -> tuple[_GatedAsyncMessageStreamManager, _FakeAsyncStreamManager]: + """Build a gated async manager around a fake stream carrying ``message``.""" + inner = _FakeAsyncStreamManager(_FakeAsyncMessageStream(message)) + return ( + _GatedAsyncMessageStreamManager(inner, dict(_KWARGS), invocation_id), + inner, + ) + + +class TestSafeSnapshot: + def test_returns_snapshot(self) -> None: + message = _make_text_response() + assert _safe_snapshot(_FakeMessageStream(message)) is message + + def test_swallows_assertion_from_unconsumed_stream(self) -> None: + """The real property asserts when nothing has accumulated yet.""" + + class _Unconsumed: + @property + def current_message_snapshot(self) -> Any: # noqa: ANN401 + raise AssertionError + + assert _safe_snapshot(_Unconsumed()) is None + + +class TestAsyncStreamGate: + @pytest.fixture(autouse=True) # pyright: ignore[reportUntypedFunctionDecorator] + def _reset_getters(self) -> Any: # noqa: ANN401 + yield + _ah._ws_getter = None + _ah._config_getter = None + _ah._hooks_getter = None + _ah._handler_getter = None + + async def test_get_final_message_emits_once(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + manager, _ = _gated_async(_make_text_response()) + + async with manager as stream: + await stream.get_final_message() + + # get_final_message() drives until_done() internally; the once-guard + # must keep that from emitting twice. + assert len(collector.events) == 1 + + async def test_until_done_then_get_final_message_emits_once(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + manager, _ = _gated_async(_make_text_response()) + + async with manager as stream: + await stream.until_done() + await stream.get_final_message() + + assert len(collector.events) == 1 + + async def test_exit_emits_audit_only_when_terminal_never_called(self) -> None: + """A text_stream-only consumer still leaves an audit trail.""" + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + manager, inner = _gated_async(_make_text_response()) + + async with manager: + pass + + assert len(collector.events) == 1 + assert inner.exited is True + + async def test_captured_invocation_id_survives_context_exit(self) -> None: + """The ID is sampled at stream() time, not when the caller reads it.""" + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + + async with anthropic_invocation(): + captured = get_invocation_id() + manager, _ = _gated_async(_make_text_response(), captured) + ctx = await manager.__aenter__() + + # Invocation context is gone by the time the final message is read. + assert get_invocation_id() is None + await ctx.get_final_message() + await manager.__aexit__(None, None, None) + + assert collector.events[0].invocation_id == captured + + async def test_block_halt_rewrites_tool_use(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + config = AdrianConfig(session_id="s") + _wired_hooks(config) + _wire_gate(ws, config) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + fut.set_result(pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy)) + + manager, _ = _gated_async(_make_tool_response()) + + async with manager as stream: + message = await stream.get_final_message() + + assert _block_types(message) == ["text"] + assert message.content[0].text == _BLOCKED_CONTENT + assert message.stop_reason == "end_turn" + + async def test_block_allow_passes_through(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=False) + config = AdrianConfig(session_id="s") + _wired_hooks(config) + _wire_gate(ws, config) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + fut.set_result(pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy)) + + manager, _ = _gated_async(_make_tool_response()) + + async with manager as stream: + message = await stream.get_final_message() + + assert _block_types(message) == ["tool_use"] + assert message.stop_reason == "tool_use" + + async def test_verdict_timeout_fails_closed(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + config = AdrianConfig(session_id="s", block_timeout=0.05) + _wired_hooks(config) + _wire_gate(ws, config) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + ws.register_pending("llm-evt") # never resolved -> times out + + manager, _ = _gated_async(_make_tool_response()) + + async with manager as stream: + message = await stream.get_final_message() + + assert message.content[0].text == _BLOCKED_CONTENT + + async def test_hitl_reject_blocks(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_HITL) + config = AdrianConfig(session_id="s") + _wired_hooks(config) + _wire_gate(ws, config) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + verdict = pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy) + verdict.hitl.continue_execution = False + fut.set_result(verdict) + + manager, _ = _gated_async(_make_tool_response()) + + async with manager as stream: + message = await stream.get_final_message() + + assert message.content[0].text == _BLOCKED_CONTENT + + async def test_hitl_approve_passes_through(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + policy = _apply_mode(ws, pb.MODE_HITL) + config = AdrianConfig(session_id="s") + _wired_hooks(config) + _wire_gate(ws, config) + + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" # pyright: ignore[reportPrivateUsage] + fut = ws.register_pending("llm-evt") + verdict = pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy) + verdict.hitl.continue_execution = True + fut.set_result(verdict) + + manager, _ = _gated_async(_make_tool_response()) + + async with manager as stream: + message = await stream.get_final_message() + + assert _block_types(message) == ["tool_use"] + + async def test_instrumentation_failure_does_not_break_the_stream(self) -> None: + """A stream shape the patch cannot instrument is still usable.""" + + class _Odd: + """No get_final_message / until_done to rebind.""" + + inner = _FakeAsyncStreamManager(_Odd()) + manager = _GatedAsyncMessageStreamManager(inner, dict(_KWARGS), None) + + async with manager as stream: + assert isinstance(stream, _Odd) + + assert inner.exited is True + + +class TestSyncStreamGate: + """Sync stream tests run without a running loop so emission is blocking.""" + + @pytest.fixture(autouse=True) # pyright: ignore[reportUntypedFunctionDecorator] + def _reset_getters(self) -> Any: # noqa: ANN401 + yield + _ah._ws_getter = None + _ah._config_getter = None + _ah._hooks_getter = None + _ah._handler_getter = None + + def test_get_final_message_emits_once(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + _wire_gate(None, config) + + inner = _FakeStreamManager(_FakeMessageStream(_make_text_response())) + manager = _GatedMessageStreamManager(inner, dict(_KWARGS), None) + + with manager as stream: + stream.get_final_message() + + assert len(collector.events) == 1 + + def test_exit_emits_audit_only_when_terminal_never_called(self) -> None: + config = AdrianConfig(session_id="s") + _, collector = _wired_hooks(config) + _wire_gate(None, config) + + inner = _FakeStreamManager(_FakeMessageStream(_make_text_response())) + manager = _GatedMessageStreamManager(inner, dict(_KWARGS), None) + + with manager: + pass + + assert len(collector.events) == 1 + assert inner.exited is True + + def test_no_ws_loop_fails_closed_via_timeout(self) -> None: + ws = WebSocketClient("ws://x", "s", api_key="k") + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._loop = None # pyright: ignore[reportPrivateUsage] - no WS loop + ws._send_frame = AsyncMock() # pyright: ignore[reportPrivateUsage] + + hooks = HookRegistry() + hooks.register(ws) + _ah._hooks_getter = lambda: hooks + config = AdrianConfig(session_id="s", block_timeout=0.05) + _wire_gate(ws, config) + + inner = _FakeStreamManager(_FakeMessageStream(_make_tool_response())) + manager = _GatedMessageStreamManager(inner, dict(_KWARGS), None) + + with manager as stream: + message = stream.get_final_message() + + # No verdict can arrive with no live loop -> fail-closed rewrite. + assert message.content[0].text == _BLOCKED_CONTENT + + def test_exception_inside_with_block_propagates(self) -> None: + config = AdrianConfig(session_id="s") + _wired_hooks(config) + _wire_gate(None, config) + + inner = _FakeStreamManager(_FakeMessageStream(_make_text_response())) + manager = _GatedMessageStreamManager(inner, dict(_KWARGS), None) + + with pytest.raises(RuntimeError, match="boom"), manager: + raise RuntimeError("boom") + + assert inner.exited is True + + +class TestPatchAnthropicStream: + @pytest.fixture(autouse=True) # pyright: ignore[reportUntypedFunctionDecorator] + def _restore_sdk(self) -> Any: # noqa: ANN401 + """Save / restore the real SDK methods this test rewraps.""" + from anthropic.resources.messages import AsyncMessages, Messages + + saved = [ + (Messages, "create", Messages.create), + (Messages, "stream", Messages.stream), + (AsyncMessages, "create", AsyncMessages.create), + (AsyncMessages, "stream", AsyncMessages.stream), + ] + flags = [ + (cls, getattr(cls, "_adrian_patched", None)) + for cls in (Messages, AsyncMessages) + ] + + for cls, _flag in flags: + if hasattr(cls, "_adrian_patched"): + delattr(cls, "_adrian_patched") + + yield + + for cls, name, original in saved: + setattr(cls, name, original) + + for cls, flag in flags: + if hasattr(cls, "_adrian_patched"): + delattr(cls, "_adrian_patched") + if flag is not None: + cls._adrian_patched = flag # type: ignore[attr-defined] + + _ah._hooks_getter = None + _ah._config_getter = None + _ah._ws_getter = None + _ah._handler_getter = None + + def test_stream_is_wrapped_on_both_classes(self) -> None: + from anthropic.resources.messages import AsyncMessages, Messages + + before = (Messages.stream, AsyncMessages.stream) + patch_anthropic(lambda: HookRegistry(), lambda: AdrianConfig(session_id="s")) + + assert Messages.stream is not before[0] + assert AsyncMessages.stream is not before[1] + + def test_stream_patch_is_idempotent(self) -> None: + from anthropic.resources.messages import AsyncMessages, Messages + + hooks_getter: Any = lambda: HookRegistry() # noqa: E731 + config_getter: Any = lambda: AdrianConfig(session_id="s") # noqa: E731 + + patch_anthropic(hooks_getter, config_getter) + after_first = (Messages.stream, AsyncMessages.stream) + + patch_anthropic(hooks_getter, config_getter) + + assert Messages.stream is after_first[0] + assert AsyncMessages.stream is after_first[1] + + def test_patched_stream_wraps_the_original_result(self) -> None: + """The wrapper wraps whatever the original ``stream()`` returned.""" + from anthropic.resources.messages import Messages + + sentinel = _FakeStreamManager(_FakeMessageStream(_make_text_response())) + seen: list[dict[str, Any]] = [] + + def _fake_stream(_self: Any, **kwargs: Any) -> Any: # noqa: ANN401 + seen.append(kwargs) + return sentinel + + # Stand in for the network call: patch_anthropic captures whatever + # Messages.stream is at patch time, so the wrapper calls this. + Messages.stream = _fake_stream # type: ignore[method-assign, assignment] + + patch_anthropic(lambda: HookRegistry(), lambda: AdrianConfig(session_id="s")) + + manager = Messages.stream(MagicMock(), **_KWARGS) + + assert isinstance(manager, _GatedMessageStreamManager) + assert seen == [_KWARGS] + + def test_patched_stream_captures_invocation_id_eagerly(self) -> None: + from anthropic.resources.messages import Messages + + Messages.stream = lambda _self, **_kw: _FakeStreamManager( # type: ignore[method-assign, assignment] + _FakeMessageStream(_make_text_response()) + ) + + patch_anthropic(lambda: HookRegistry(), lambda: AdrianConfig(session_id="s")) + + with anthropic_invocation_sync(): + expected = get_invocation_id() + manager = Messages.stream(MagicMock(), **_KWARGS) + + # Sampled inside the block, so it survives the context exit. + assert get_invocation_id() is None + assert manager._state._invocation_id == expected # pyright: ignore[reportPrivateUsage]