agentrust-py 0.0.7__tar.gz → 0.0.8__tar.gz

{agentrust_py-0.0.7 → agentrust_py-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: agentrust-py
-Version: 0.0.7
+Version: 0.0.8
 Summary: AgentTrust Edge — Full Governance Platform. Policy Engine + MAD Classifier + HITL + Kill Switch.
 Project-URL: Homepage, https://agentrust.io
 Project-URL: Documentation, https://docs.agentrust.io
@@ -25,6 +25,8 @@ Requires-Dist: pydantic<3.0.0,>=2.0.0
 Requires-Dist: pyyaml<7.0,>=6.0.3
 Provides-Extra: ag2
 Requires-Dist: ag2<1.0,>=0.9; extra == 'ag2'
+Provides-Extra: autogen
+Requires-Dist: autogen-agentchat<1.0,>=0.4; extra == 'autogen'
 Provides-Extra: autoload
 Provides-Extra: crewai
 Requires-Dist: crewai<2.0,>=1.0; extra == 'crewai'

{agentrust_py-0.0.7 → agentrust_py-0.0.8}/agentrust_sdk/__init__.py

@@ -43,6 +43,8 @@ from .adapters.autogen import (
     AgentTrustReplyHook,
     governed_function,
     governed_functions,
+    governed_tool as governed_autogen_tool,
+    governed_tools as governed_autogen_tools,
 )
 from .adapters.claude_agents import ClaudeAgentGuard
 from .adapters.openai_agents import agent_guard as openai_agent_guard, OpenAIAgentsGuard, SwarmGuard
@@ -154,4 +156,4 @@ __all__ = [
     "validate",
 ]
 
-__version__ = "0.0.7"
+__version__ = "0.0.8"

{agentrust_py-0.0.7 → agentrust_py-0.0.8}/agentrust_sdk/adapters/autogen.py

@@ -10,6 +10,10 @@ Pattern 0 — governed_function (AG2 / pyautogen 0.2 tools):
   held for approval when needed, run, post-checked and audited. A refusal
   is returned as the tool result, so the chat keeps going.
 
+Pattern 0b — governed_tool (AutoGen 0.4+ / autogen-agentchat tools):
+  Wrap each FunctionTool (or plain function) before giving it to an
+  AssistantAgent. Same checks as governed_function.
+
 Pattern 1 — AgentTrustReplyHook (v0.2 + v0.4):
   Attach to any ConversableAgent / AssistantAgent to validate every reply
   before it is sent.
@@ -42,6 +46,7 @@ from ._core import Allowed, Governor, OnBlock, Refused, check_adapter_tier
 logger = logging.getLogger(__name__)
 
 _FRAMEWORK = "ag2"
+_FRAMEWORK_V04 = "autogen"
 
 
 def governed_function(
@@ -132,6 +137,102 @@ def governed_functions(
     return [governed_function(fn, agent_id=agent_id, governor=gov) for fn in fns]
 
 
+def _governor(
+    agent_id: str, framework: str, on_block: OnBlock, review_timeout_s: float,
+    user: str, base_url: str | None, api_key: str | None,
+) -> Governor:
+    return Governor(
+        agent_id, framework=framework, on_block=on_block, review_timeout_s=review_timeout_s,
+        user=user, base_url=base_url, api_key=api_key,
+    )
+
+
+def governed_tool(
+    tool: Any,
+    *,
+    agent_id: str,
+    on_block: OnBlock = "refuse",
+    review_timeout_s: float = 300.0,
+    user: str = "autogen",
+    base_url: str | None = None,
+    api_key: str | None = None,
+    governor: Governor | None = None,
+) -> Any:
+    """Wrap an AutoGen 0.4+ tool so AgentTrust governs every call.
+
+    Accepts a ``BaseTool`` such as ``FunctionTool``, or a plain function (its
+    docstring becomes the description). The result has the same name,
+    description and argument schema, so the model sees the same tool. Pass it
+    in ``AssistantAgent(tools=[...])``. A refused call reaches the agent as the
+    tool result.
+    """
+    check_adapter_tier(Capability.AUTOGEN_ADAPTER, api_key)
+    gov = governor or _governor(
+        agent_id, _FRAMEWORK_V04, on_block, review_timeout_s, user, base_url, api_key,
+    )
+    return _governed_tool_class().wrap(_as_tool(tool), gov)
+
+
+def governed_tools(
+    tools: Sequence[Any],
+    *,
+    agent_id: str,
+    on_block: OnBlock = "refuse",
+    review_timeout_s: float = 300.0,
+    user: str = "autogen",
+    base_url: str | None = None,
+    api_key: str | None = None,
+    governor: Governor | None = None,
+) -> list[Any]:
+    """Wrap each AutoGen 0.4+ tool with :func:`governed_tool`, sharing one Governor."""
+    check_adapter_tier(Capability.AUTOGEN_ADAPTER, api_key)
+    gov = governor or _governor(
+        agent_id, _FRAMEWORK_V04, on_block, review_timeout_s, user, base_url, api_key,
+    )
+    return [governed_tool(t, agent_id=agent_id, governor=gov) for t in tools]
+
+
+def _as_tool(tool: Any) -> Any:
+    from autogen_core.tools import BaseTool, FunctionTool
+
+    if isinstance(tool, BaseTool):
+        return tool
+    return FunctionTool(tool, description=inspect.getdoc(tool) or tool.__name__)
+
+
+@functools.lru_cache(maxsize=None)
+def _governed_tool_class() -> type:
+    """Build GovernedTool on first use, so autogen-core stays optional."""
+    from autogen_core import CancellationToken
+    from autogen_core.tools import BaseTool
+
+    class GovernedTool(BaseTool):
+        """A tool whose every call goes through an AgentTrust Governor."""
+
+        def __init__(self, inner: Any, governor: Governor) -> None:
+            super().__init__(
+                inner.args_type(), inner.return_type(), inner.name, inner.description,
+                strict=getattr(inner, "_strict", False),
+            )
+            self._inner = inner
+            self._governor = governor
+
+        @classmethod
+        def wrap(cls, inner: Any, governor: Governor) -> "GovernedTool":
+            return cls(inner, governor)
+
+        async def run(self, args: Any, cancellation_token: CancellationToken) -> Any:
+            params = args.model_dump()
+            outcome = await self._governor.aguard_tool(
+                self.name, params, lambda: self._inner.run(args, cancellation_token),
+                input=_describe_call(self.name, params),
+            )
+            # The refusal text is the tool result the agent reads next.
+            return outcome.message if isinstance(outcome, Refused) else outcome.result
+
+    return GovernedTool
+
+
 def _describe_call(name: str, args: dict[str, Any]) -> str:
     """The proposed call as the pre-check's request input."""
     return f"tool call {name}({', '.join(f'{k}={v!r}' for k, v in args.items())})"
@@ -169,6 +270,39 @@ def _validate_sync(
         logger.warning("[AgentTrust] AutoGen validate failed (non-fatal): %s", exc)
 
 
+async def _avalidate(
+    agent_id: str,
+    base_url: str,
+    api_key: str | None,
+    user: str,
+    input_text: str,
+    output: dict,
+    framework: str,
+    block_on_block: bool,
+) -> None:
+    """Async counterpart of :func:`_validate_sync`, safe to await in an event loop."""
+    from agentrust_sdk.client import AsyncAgentTrustClient, GatewayUnavailableError
+    from agentrust_sdk.config import SDK_CONFIG
+    from agentrust_sdk.decorator import _check_decision, BlockedError
+
+    try:
+        async with AsyncAgentTrustClient(base_url=base_url, api_key=api_key) as client:
+            resp = await client.validate(
+                agent_id=agent_id,
+                user=user,
+                input=input_text,
+                output=output,
+                framework=framework,
+            )
+        _check_decision(resp, block_on_block, False)
+    except (BlockedError, GatewayUnavailableError):
+        raise
+    except Exception as exc:
+        if SDK_CONFIG.failure_mode == "closed":
+            raise
+        logger.warning("[AgentTrust] AutoGen validate failed (non-fatal): %s", exc)
+
+
 class AgentTrustReplyHook:
     """
     Reply hook that validates every AutoGen agent reply via AgentTrust.
@@ -259,7 +393,7 @@ class AgentTrustReplyHook:
             if hasattr(result, "chat_message") and result.chat_message:
                 content = getattr(result.chat_message, "content", "") or ""
             output = {"content": content[:4000]}
-            _validate_sync(
+            await _avalidate(
                 self._agent_id, self._base_url, self._api_key,
                 self._user, content[:500], output, self._framework, self._block_on_block,
             )

{agentrust_py-0.0.7 → agentrust_py-0.0.8}/agentrust_sdk/config.py

@@ -47,7 +47,7 @@ from pathlib import Path
 logger = logging.getLogger(__name__)
 
 _VALID_FAILURE_MODES = frozenset({"open", "closed", "queue"})
-_SDK_VERSION = "0.0.7"
+_SDK_VERSION = "0.0.8"
 _MIN_GATEWAY_VERSION = "0.0.1a1"
 
 

{agentrust_py-0.0.7 → agentrust_py-0.0.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "agentrust-py"
-version = "0.0.7"
+version = "0.0.8"
 description = "AgentTrust Edge — Full Governance Platform. Policy Engine + MAD Classifier + HITL + Kill Switch."
 readme = "README.md"
 license = { text = "Apache-2.0" }
@@ -43,6 +43,8 @@ langgraph = ["langgraph>=0.2.0,<2.0.0"]
 crewai    = ["crewai>=1.0,<2.0"]
 # AG2 0.x keeps the ConversableAgent API (module `autogen`); 1.x is a new API
 ag2       = ["ag2>=0.9,<1.0"]
+# Microsoft AutoGen 0.4+ (module autogen_agentchat)
+autogen   = ["autogen-agentchat>=0.4,<1.0"]
 openai    = ["openai>=2.41.1,<3.0.0"]
 langchain = ["langchain>=0.2.0,<1.0.0"]
 # Retry/backoff (recommended for production)

agentrust_py-0.0.8/tests/adapters/test_ag2_quickstart_example.py

@@ -0,0 +1,145 @@
+"""examples/ag2_quickstart.py: the calculator quickstart under governed_functions.
+
+The example's own agents run with a scripted model against an embedded
+gateway that has the example's manifest (add and multiply declared, divide
+not).
+"""
+from __future__ import annotations
+
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+
+pytest.importorskip("autogen")
+pytestmark = pytest.mark.framework
+
+from tests.adapters import _ag2
+from tests.adapters._ag2 import final, tool_call
+
+EXAMPLE = Path(__file__).resolve().parents[3] / "examples" / "ag2_quickstart.py"
+
+
+@pytest.fixture
+def quickstart():
+    if not EXAMPLE.exists():
+        pytest.skip("examples/ is not part of this checkout")
+    spec = importlib.util.spec_from_file_location("ag2_quickstart", EXAMPLE)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture
+def restore_model(monkeypatch):
+    """--scripted patches AG2's model call for good; undo that after the test."""
+    from autogen.oai.client import OpenAIWrapper
+
+    monkeypatch.setattr(OpenAIWrapper, "create", OpenAIWrapper.create)
+    monkeypatch.setattr(
+        OpenAIWrapper, "extract_text_or_completion_object", OpenAIWrapper.extract_text_or_completion_object,
+    )
+
+
+@pytest.fixture
+def gateway(embedded_gw, quickstart):
+    return embedded_gw(manifests={quickstart.AGENT_ID: quickstart.MANIFEST})
+
+
+def _ledger(gateway):
+    """The gateway's decisions for the example's agent, oldest first: phase, tool, decision, envelope, parent."""
+    import httpx
+
+    resp = httpx.get(f"{gateway.url}/v1/audit/executions", params={"limit": 100},
+                     headers={"Authorization": f"Bearer {gateway._token}"})
+    rows = sorted(resp.json()["items"], key=lambda r: r["id"])
+    return [[r["phase"], json.loads(r["request_json"] or "{}").get("metadata", {}).get("action", "-"),
+             r["decision"], r["envelope_id"][:8], (r["parent_envelope_id"] or "-")[:8]] for r in rows]
+
+
+def _ask(quickstart, request):
+    return quickstart.ask(request, quickstart.llm_config(scripted=True))
+
+
+def test_declared_tool_runs_and_both_checks_are_recorded(quickstart, gateway, monkeypatch):
+    model = _ag2.install(monkeypatch, tool_call("add", {"a": 3, "b": 4}), final("3 + 4 = 7. DONE"))
+
+    assert "3 + 4 = 7" in _ask(quickstart, "Add 3 and 4.")
+
+    assert any('"7"' in p or "7" in p.split("tool")[-1] for p in model.prompts[1:])
+    pre, post = _ledger(gateway)
+    assert pre[:3] == ["pre_check", "add", "approve"]
+    assert post[:3] == ["post_check", "add", "approve"]
+    assert post[4] == pre[3]  # the post-check row points at its pre-check
+
+
+def test_undeclared_tool_is_refused_before_it_runs(quickstart, gateway, monkeypatch):
+    ran = []
+
+    def divide(a: int, b: int) -> float:
+        """Divide `a` by `b`."""
+        ran.append((a, b))
+        return a / b
+
+    from agentrust_sdk.adapters.autogen import governed_function
+
+    monkeypatch.setattr(quickstart, "tools", [governed_function(divide, agent_id=quickstart.AGENT_ID)])
+    model = _ag2.install(monkeypatch, tool_call("divide", {"a": 10, "b": 2}), final("I can't divide. DONE"))
+
+    assert "I can't divide." in _ask(quickstart, "Divide 10 by 2.")  # the chat still finished
+
+    assert ran == []
+    refusal = [p for p in model.prompts if "Blocked by AgentTrust" in p]
+    assert refusal and "'divide' is not in this agent's declared_tools" in refusal[-1]
+    [pre] = _ledger(gateway)
+    assert pre[:3] == ["pre_check", "divide", "block"]
+
+
+def test_agentrust_off_lets_every_tool_run(quickstart, monkeypatch):
+    monkeypatch.setenv("AGENTRUST_ENABLED", "false")
+    model = _ag2.install(monkeypatch, tool_call("divide", {"a": 10, "b": 2}), final("10 / 2 = 5. DONE"))
+
+    assert "10 / 2 = 5" in _ask(quickstart, "Divide 10 by 2.")
+
+    assert any("5.0" in p for p in model.prompts[1:])
+    assert not any("Blocked by AgentTrust" in p for p in model.prompts)
+
+
+def test_scripted_flag_runs_the_whole_quickstart_with_no_model(
+    quickstart, gateway, restore_model, capsys, monkeypatch,
+):
+    # --console talks to the gateway the fixture started (embed_gateway keeps one per process).
+    monkeypatch.setenv("AGENTRUST_AGENT_ID", quickstart.AGENT_ID)
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+
+    assert quickstart.main(["--console", "--scripted"]) == 0
+
+    out = capsys.readouterr().out
+    assert "model scripted" in out
+    assert "Result: 3 + 4 = 7." in out
+    assert "Result: That is everything you asked for." in out
+    assert [r[:3] for r in _ledger(gateway)] == [
+        ["pre_check", "add", "approve"], ["post_check", "add", "approve"], ["pre_check", "divide", "block"],
+    ]
+
+
+def test_without_agentrust_flag_runs_with_agentrust_off(quickstart, restore_model, capsys, monkeypatch):
+    monkeypatch.setenv("AGENTRUST_ENABLED", "true")  # main() turns it off; monkeypatch restores it
+
+    assert quickstart.main(["--without-agentrust", "--scripted"]) == 0
+
+    out = capsys.readouterr().out
+    assert "AgentTrust: off" in out
+    assert "Blocked by AgentTrust" not in out
+
+
+def test_console_flag_names_the_missing_settings(quickstart, monkeypatch, capsys):
+    for name in ("AGENTRUST_GATEWAY_URL", "AGENTRUST_KEY", "AGENTRUST_AGENT_ID"):
+        monkeypatch.delenv(name, raising=False)
+
+    assert quickstart.main(["--console"]) == 2
+
+    err = capsys.readouterr().err
+    assert "AGENTRUST_GATEWAY_URL, AGENTRUST_KEY, AGENTRUST_AGENT_ID" in err
+    assert "--local" in err

agentrust_py-0.0.8/tests/adapters/test_autogen_quickstart_example.py

@@ -0,0 +1,138 @@
+"""examples/autogen_quickstart.py: the calculator quickstart under governed_tools (AutoGen 0.4+).
+
+The example's own agent runs with a scripted model against an embedded
+gateway that has the example's manifest (add and multiply declared, divide
+not).
+"""
+from __future__ import annotations
+
+import asyncio
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+
+pytest.importorskip("autogen_agentchat")
+pytest.importorskip("autogen_ext")
+pytestmark = pytest.mark.framework
+
+EXAMPLE = Path(__file__).resolve().parents[3] / "examples" / "autogen_quickstart.py"
+
+
+@pytest.fixture
+def quickstart():
+    if not EXAMPLE.exists():
+        pytest.skip("examples/ is not part of this checkout")
+    spec = importlib.util.spec_from_file_location("autogen_quickstart", EXAMPLE)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture
+def gateway(embedded_gw, quickstart):
+    return embedded_gw(manifests={quickstart.AGENT_ID: quickstart.MANIFEST})
+
+
+def _ledger(gateway):
+    """The gateway's decisions for the example's agent, oldest first: phase, tool, decision, envelope, parent."""
+    import httpx
+
+    resp = httpx.get(f"{gateway.url}/v1/audit/executions", params={"limit": 100},
+                     headers={"Authorization": f"Bearer {gateway._token}"})
+    rows = sorted(resp.json()["items"], key=lambda r: r["id"])
+    return [[r["phase"], json.loads(r["request_json"] or "{}").get("metadata", {}).get("action", "-"),
+             r["decision"], r["envelope_id"][:8], (r["parent_envelope_id"] or "-")[:8]] for r in rows]
+
+
+def _run(quickstart, request):
+    agent = quickstart.build_agent(quickstart.scripted_client())
+    return asyncio.run(agent.run(task=request))
+
+
+def _tool_results(result):
+    return [r.content for m in result.messages if m.type == "ToolCallExecutionEvent" for r in m.content]
+
+
+def test_declared_tool_runs_and_both_checks_are_recorded(quickstart, gateway):
+    result = _run(quickstart, "Add 3 and 4.")
+
+    assert "3 + 4 = 7" in result.messages[-1].content
+    assert _tool_results(result) == ["7"]
+    pre, post = _ledger(gateway)
+    assert pre[:3] == ["pre_check", "add", "approve"]
+    assert post[:3] == ["post_check", "add", "approve"]
+    assert post[4] == pre[3]  # the post-check row points at its pre-check
+
+
+def test_undeclared_tool_is_refused_before_it_runs(quickstart, gateway, monkeypatch):
+    ran = []
+    inner = quickstart.tools[2]._inner  # divide, behind the AgentTrust wrapper
+    original = inner.run
+
+    async def spy(args, token):
+        ran.append(args)
+        return await original(args, token)
+
+    monkeypatch.setattr(inner, "run", spy)
+    client = quickstart.scripted_client()
+    asyncio.run(quickstart.build_agent(client).run(task="Add 3 and 4."))  # consume the first script pair
+    ledger_before = len(_ledger(gateway))
+
+    result = asyncio.run(quickstart.build_agent(client).run(task="Divide 10 by 2."))
+
+    assert ran == []
+    assert "That is everything you asked for." in result.messages[-1].content  # the run still finished
+    [refusal] = _tool_results(result)
+    assert "Blocked by AgentTrust" in refusal and "'divide' is not in this agent's declared_tools" in refusal
+    assert _ledger(gateway)[ledger_before][:3] == ["pre_check", "divide", "block"]
+
+
+def test_agentrust_off_lets_every_tool_run(quickstart, monkeypatch):
+    monkeypatch.setenv("AGENTRUST_ENABLED", "false")
+    client = quickstart.scripted_client()
+    asyncio.run(quickstart.build_agent(client).run(task="Add 3 and 4."))
+
+    result = asyncio.run(quickstart.build_agent(client).run(task="Divide 10 by 2."))
+
+    assert _tool_results(result) == ["5.0"]
+
+
+def test_scripted_flag_runs_the_whole_quickstart_with_no_model(
+    quickstart, gateway, capsys, monkeypatch,
+):
+    # --console talks to the gateway the fixture started (embed_gateway keeps one per process).
+    monkeypatch.setenv("AGENTRUST_AGENT_ID", quickstart.AGENT_ID)
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+
+    assert quickstart.main(["--console", "--scripted"]) == 0
+
+    out = capsys.readouterr().out
+    assert "model scripted" in out
+    assert "Result: 3 + 4 = 7." in out
+    assert "Result: That is everything you asked for." in out
+    assert [r[:3] for r in _ledger(gateway)] == [
+        ["pre_check", "add", "approve"], ["post_check", "add", "approve"], ["pre_check", "divide", "block"],
+    ]
+
+
+def test_without_agentrust_flag_runs_with_agentrust_off(quickstart, capsys, monkeypatch):
+    monkeypatch.setenv("AGENTRUST_ENABLED", "true")  # main() turns it off; monkeypatch restores it
+
+    assert quickstart.main(["--without-agentrust", "--scripted"]) == 0
+
+    out = capsys.readouterr().out
+    assert "AgentTrust: off" in out
+    assert "divide returned: 5.0" in out
+
+
+def test_console_flag_names_the_missing_settings(quickstart, monkeypatch, capsys):
+    for name in ("AGENTRUST_GATEWAY_URL", "AGENTRUST_KEY", "AGENTRUST_AGENT_ID"):
+        monkeypatch.delenv(name, raising=False)
+
+    assert quickstart.main(["--console"]) == 2
+
+    err = capsys.readouterr().err
+    assert "AGENTRUST_GATEWAY_URL, AGENTRUST_KEY, AGENTRUST_AGENT_ID" in err
+    assert "--local" in err

agentrust_py-0.0.8/tests/adapters/test_autogen_v04_governed_tool.py

@@ -0,0 +1,243 @@
+"""AutoGen 0.4+ (autogen-agentchat): governed_tool and the async reply hook."""
+from __future__ import annotations
+
+import asyncio
+
+import httpx
+import pytest
+
+pytest.importorskip("autogen_agentchat")
+pytest.importorskip("autogen_ext")
+pytestmark = pytest.mark.framework
+
+from autogen_agentchat.agents import AssistantAgent
+from autogen_core import CancellationToken, FunctionCall
+from autogen_core.models import CreateResult, RequestUsage
+from autogen_core.tools import FunctionTool
+from autogen_ext.models.replay import ReplayChatCompletionClient
+
+from agentrust_sdk.adapters.autogen import AgentTrustReplyHook, governed_tool, governed_tools
+from tests.test_governor import FakeAsyncClient, _gov, _resp
+
+AGENT = "autogen-v04-agent"
+TOKEN = CancellationToken()
+
+
+def _refund_tool(calls):
+    async def refund(amount: float, reason: str = "") -> str:
+        """Refund an amount to the customer."""
+        calls.append(amount)
+        return f"refunded {amount}"
+
+    return FunctionTool(refund, description="Refund an amount to the customer.")
+
+
+def _call(name, args):
+    import json
+
+    return CreateResult(
+        finish_reason="function_calls",
+        content=[FunctionCall(id="c1", name=name, arguments=json.dumps(args))],
+        usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
+        cached=False,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Unit: a real Governor on a scripted async client
+# ---------------------------------------------------------------------------
+
+def test_wrapped_tool_keeps_name_description_and_schema():
+    inner = _refund_tool([])
+    wrapped = governed_tool(inner, agent_id="a", governor=_gov(FakeAsyncClient()))
+
+    assert wrapped.name == "refund"
+    assert wrapped.description == inner.description
+    assert wrapped.schema == inner.schema
+
+
+def test_allowed_call_is_pre_checked_runs_and_is_post_checked():
+    calls, client = [], FakeAsyncClient()
+    wrapped = governed_tool(_refund_tool(calls), agent_id="a", governor=_gov(client))
+
+    result = asyncio.run(wrapped.run_json({"amount": 5, "reason": "late"}, TOKEN))
+
+    assert result == "refunded 5.0" and calls == [5.0]
+    [pre] = client.pre_calls
+    assert (pre["action"], pre["params"]) == ("refund", {"amount": 5.0, "reason": "late"})
+    assert len(client.post_calls) == 1
+
+
+def test_blocked_call_does_not_run_and_returns_the_refusal_text():
+    calls = []
+    client = FakeAsyncClient(pre=[_resp("block", reason="over the limit")])
+    wrapped = governed_tool(_refund_tool(calls), agent_id="a", governor=_gov(client))
+
+    result = asyncio.run(wrapped.run_json({"amount": 9999}, TOKEN))
+
+    assert calls == [] and "over the limit" in result
+
+
+def test_plain_function_is_accepted_and_keeps_its_docstring_as_description():
+    calls = []
+
+    async def lookup(order_id: str) -> str:
+        """Look up an order."""
+        calls.append(order_id)
+        return f"order {order_id}"
+
+    [wrapped] = governed_tools([lookup], agent_id="a", governor=_gov(FakeAsyncClient()))
+
+    assert wrapped.description == "Look up an order."
+    assert asyncio.run(wrapped.run_json({"order_id": "42"}, TOKEN)) == "order 42"
+
+
+def test_sync_function_tool_runs_without_blocking_the_loop():
+    calls = []
+
+    def ping(host: str) -> str:
+        """Ping a host."""
+        calls.append(host)
+        return f"pong {host}"
+
+    [wrapped] = governed_tools([ping], agent_id="a", governor=_gov(FakeAsyncClient()))
+
+    assert asyncio.run(wrapped.run_json({"host": "h"}, TOKEN)) == "pong h"
+    assert calls == ["h"]
+
+
+def test_kill_switch_off_runs_plainly(monkeypatch):
+    monkeypatch.setenv("AGENTRUST_ENABLED", "false")
+    calls, client = [], FakeAsyncClient()
+    wrapped = governed_tool(_refund_tool(calls), agent_id="a", governor=_gov(client))
+
+    assert asyncio.run(wrapped.run_json({"amount": 1}, TOKEN)) == "refunded 1.0"
+    assert client.pre_calls == []
+
+
+def test_gateway_down_in_closed_mode_refuses():
+    from agentrust_sdk.client import GatewayUnavailableError
+
+    calls = []
+    client = FakeAsyncClient(pre_error=GatewayUnavailableError("down"))
+    wrapped = governed_tool(_refund_tool(calls), agent_id="a", governor=_gov(client))
+
+    assert "gateway unavailable" in asyncio.run(wrapped.run_json({"amount": 1}, TOKEN))
+    assert calls == []
+
+
+# ---------------------------------------------------------------------------
+# Real agent on an embedded gateway
+# ---------------------------------------------------------------------------
+
+def _ledger(gw):
+    rows = httpx.get(f"{gw.url}/v1/audit/executions", params={"limit": 100},
+                     headers={"Authorization": f"Bearer {gw._token}"}).json()["items"]
+    return sorted(rows, key=lambda r: r["id"])
+
+
+def _agent(tool, *replies):
+    model = ReplayChatCompletionClient(list(replies), model_info={
+        "vision": False, "function_calling": True, "json_output": False,
+        "family": "unknown", "structured_output": False,
+    })
+    return AssistantAgent("assistant", model_client=model, tools=[tool], reflect_on_tool_use=True)
+
+
+def test_agent_runs_a_declared_tool_with_pre_and_post_check(embedded_gw):
+    gw = embedded_gw(manifests={AGENT: {"declared_tools": [{"name": "refund"}]}})
+    calls = []
+    agent = _agent(governed_tool(_refund_tool(calls), agent_id=AGENT),
+                   _call("refund", {"amount": 5}), "All done.")
+
+    asyncio.run(agent.run(task="refund 5"))
+
+    assert calls == [5.0]
+    assert [(r["phase"], r["decision"]) for r in _ledger(gw)] == [
+        ("pre_check", "approve"), ("post_check", "approve"),
+    ]
+
+
+def test_agent_gets_a_refusal_for_an_undeclared_tool_and_still_finishes(embedded_gw):
+    gw = embedded_gw(manifests={AGENT: {"declared_tools": [{"name": "calculator"}]}})
+    calls = []
+    agent = _agent(governed_tool(_refund_tool(calls), agent_id=AGENT),
+                   _call("refund", {"amount": 5}), "Could not refund.")
+
+    result = asyncio.run(agent.run(task="refund 5"))
+
+    assert calls == []
+    assert "Could not refund." in result.messages[-1].content
+    first = _ledger(gw)[0]
+    assert (first["phase"], first["decision"]) == ("pre_check", "block")
+    assert any("declared_tools" in str(m.content) for m in result.messages)
+
+
+# ---------------------------------------------------------------------------
+# attach_v4: async validate, fails closed
+# ---------------------------------------------------------------------------
+
+class _AsyncClient:
+    outcome = "approve"
+    error: Exception | None = None
+    calls: list[dict] = []
+
+    def __init__(self, *a, **kw):
+        pass
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *a):
+        return False
+
+    async def validate(self, **kw):
+        type(self).calls.append(kw)
+        if type(self).error:
+            raise type(self).error
+        return _resp(type(self).outcome, reason="not allowed")
+
+
+@pytest.fixture
+def aclient(monkeypatch):
+    _AsyncClient.outcome, _AsyncClient.error, _AsyncClient.calls = "approve", None, []
+    monkeypatch.setattr("agentrust_sdk.client.AsyncAgentTrustClient", _AsyncClient)
+    return _AsyncClient
+
+
+def _plain_agent(*replies):
+    model = ReplayChatCompletionClient(list(replies))
+    return AssistantAgent("assistant", model_client=model)
+
+
+def test_attach_v4_validates_the_reply_with_the_async_client(aclient):
+    agent = _plain_agent("hello there")
+    AgentTrustReplyHook(agent_id="a").attach_v4(agent)
+
+    asyncio.run(agent.run(task="hi"))
+
+    [call] = aclient.calls
+    assert call["output"]["content"] == "hello there"
+
+
+def test_attach_v4_raises_when_the_reply_is_blocked(aclient):
+    from agentrust_sdk.decorator import BlockedError
+
+    aclient.outcome = "block"
+    agent = _plain_agent("leak it")
+    AgentTrustReplyHook(agent_id="a").attach_v4(agent)
+
+    with pytest.raises(BlockedError):
+        asyncio.run(agent.run(task="hi"))
+
+
+def test_attach_v4_fails_closed_when_the_gateway_is_down(aclient, monkeypatch):
+    from agentrust_sdk.client import GatewayUnavailableError
+
+    monkeypatch.setenv("AGENTRUST_FAILURE_MODE", "closed")
+    aclient.error = GatewayUnavailableError("down")
+    agent = _plain_agent("hello")
+    AgentTrustReplyHook(agent_id="a").attach_v4(agent)
+
+    with pytest.raises(GatewayUnavailableError):
+        asyncio.run(agent.run(task="hi"))

{agentrust_py-0.0.7 → agentrust_py-0.0.8}/tests/test_callback_url.py

@@ -152,7 +152,7 @@ class TestAnswerFields:
 # Release hygiene: one version everywhere, and the right install name
 # ---------------------------------------------------------------------------
 REPO = Path(__file__).resolve().parents[2]
-RELEASE = "0.0.7"
+RELEASE = "0.0.8"
 
 
 class TestReleaseHygiene: