Block Unsafe AG2 Tool Calls Before They Run

The model chooses which tool to call and with what arguments. In this guide's calculator agent, add and multiply are the intended tools, yet the model can also call divide. Run the agent with no check and divide goes through: nothing stopped it. For a calculator that is harmless. Swap divide for issue_refund or delete_user and it would run just the same.

In an AG2 agent, a tool call crosses two places. register_function tells the calling agent which functions exist, and the executor agent runs the one the model picks. The wrapper goes on the function before it is registered, so the check runs on every call the executor makes.

What happens on one tool call

1. Intercept. The wrapped function receives the call. Your tool's code has not run.
2. Ask. The wrapper sends the tool name and its arguments, with the agent id, to the AgentTrust gateway. This is the pre-check.
3. Compare. The gateway checks the call against the agent's declared tools. A tool that is not declared, or an amount above its declared limit, is refused by the tool-trust hard gate.
4. Record. The gateway writes an event for the decision, allowed or refused.
5. Run or refuse. An allowed call runs, and its result goes back to the gateway for a post-check, a second event. A refused call never runs. The agent receives Blocked by AgentTrust as the tool's result.

What the wrapper must cover

The check covers only the tools you wrap. A tool you register unwrapped is not checked, so wrap every tool the agent has. Give each agent its own agent_id so the audit trail says who acted. The manifest matters too: an agent with no declared tools is not restricted, so every tool is allowed.

The same four steps, applied to a small calculator agent. You run it three ways and watch a tool call get refused. About 10 minutes. Each step has one thing to do and one thing you should see.

Before you start. Python 3.10 or later. Steps 1 to 4 need nothing else when you add --scripted to each command; without it they need an OpenAI key (OPENAI_API_KEY). Steps 5 to 7 use an AgentTrust console account.

Create the project

Make a folder with its own Python environment and install what the program needs:

mkdir agentrust-quickstart && cd agentrust-quickstart
python -m venv .venv && source .venv/bin/activate     # on Windows: .venv\Scripts\activate
pip install "agentrust-py[ag2,embedded]"

Save the program below as ag2_quickstart.py. It is the example examples/ag2_quickstart.py from the AgentTrust repository.

"""
ag2_quickstart.py — an AG2 agent whose tool calls AgentTrust checks.

The agent is a calculator: one AG2 (AutoGen 0.2 API) agent with add, multiply
and divide tools, asked to do two sums.

AgentTrust is one line (marked "AgentTrust" below): the tools are wrapped with
governed_functions before they are registered. Every call is checked before it
runs, and its result after. A tool the agent may not use is refused before its
code runs; the agent is told why as the tool's result and the chat still
finishes. The agent may use add and multiply, not divide.

RUN IT (in the folder that holds this file)
    pip install "agentrust-py[ag2,embedded]"
    export OPENAI_API_KEY=...                             # or use --scripted, which needs no key

    python ag2_quickstart.py --without-agentrust          # AgentTrust off: every tool runs
    python ag2_quickstart.py --local                      # a gateway starts inside this program; no account
    python ag2_quickstart.py --console                    # your AgentTrust console keeps every decision
    python ag2_quickstart.py --scripted                   # replay a fixed conversation instead of calling a model

  --console needs three settings, from the agent you create in the console:
      export AGENTRUST_GATEWAY_URL=https://console.agent-trust.tech
      export AGENTRUST_KEY=at_live_...
      export AGENTRUST_AGENT_ID=...

ANOTHER MODEL
    AG2_QUICKSTART_MODEL=gpt-4o (export OPENAI_API_KEY). Any model AG2 supports
    through llm_config works; the model must support tool calling.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import tempfile
from types import SimpleNamespace

from autogen import ConversableAgent, register_function

from agentrust_sdk.adapters.autogen import governed_functions

AGENT_ID = os.environ.get("AGENTRUST_AGENT_ID", "calculator-agent")
MODEL = os.environ.get("AG2_QUICKSTART_MODEL", "gpt-4o-mini")

# The tools this agent may call. With --local they are set here; with --console
# you declare the same list on the agent in the console. divide is left out on purpose.
MANIFEST = {"declared_tools": [{"name": "add"}, {"name": "multiply"}]}


def multiply(a: int, b: int) -> int:
    """Multiply `a` and `b`."""
    return a * b


def add(a: int, b: int) -> int:
    """Add `a` and `b`."""
    return a + b


def divide(a: int, b: int) -> float:
    """Divide `a` by `b`."""
    return a / b


# AgentTrust: every call to these tools is checked before it runs and after it
# returns. Register the result with the agent, not the plain functions.
tools = governed_functions([add, multiply, divide], agent_id=AGENT_ID)

REQUESTS = ["Add 3 and 4.", "Divide 10 by 2."]


def build_agents(llm_config) -> tuple[ConversableAgent, ConversableAgent]:
    """The calculator that asks for tools and the executor that runs them."""
    calculator = ConversableAgent(
        "calculator",
        system_message=(
            "You do the arithmetic you are asked for, using your tools, one tool call at a time. "
            "If a tool is refused, say the operation could not be done. "
            "Reply with the answer and the word DONE."
        ),
        llm_config=llm_config,
        human_input_mode="NEVER",
    )
    executor = ConversableAgent(
        "executor",
        llm_config=False,
        human_input_mode="NEVER",
        is_termination_msg=lambda m: "DONE" in (m.get("content") or ""),
    )
    for tool in tools:
        register_function(
            tool, caller=calculator, executor=executor, name=tool.__name__, description=tool.__doc__,
        )
    return calculator, executor


def llm_config(scripted: bool = False) -> dict:
    # AG2 builds its OpenAI client when the agent is created, so even the scripted
    # run needs some key to be set; it is never sent anywhere.
    key = "scripted" if scripted else os.environ.get("OPENAI_API_KEY", "")
    return {"config_list": [{"model": MODEL, "api_key": key}]}


def ask(request: str, config) -> str:
    calculator, executor = build_agents(config)
    result = executor.initiate_chat(calculator, message=request, max_turns=6, silent=True)
    return (result.summary or "").replace("DONE", "").strip()


# --scripted: replay a fixed conversation, so the quickstart runs with no model and no key.
# A replay does not read tool results, so its last line is the same with AgentTrust on or off;
# what differs is what each tool returned.
# Each reply is either a tool call (name, arguments) or the final text.
SCRIPT = [
    ("add", {"a": 3, "b": 4}), "3 + 4 = 7. DONE",
    ("divide", {"a": 10, "b": 2}), "That is everything you asked for. DONE",
]


def use_scripted_model(replies=SCRIPT) -> None:
    from autogen.oai.client import OpenAIWrapper

    queue = list(replies)

    def create(self, **kwargs):
        reply = queue.pop(0)
        if isinstance(reply, str):
            message = {"role": "assistant", "content": reply}
            calls = None
        else:
            name, args = reply
            calls = [{"id": "call-1", "type": "function",
                      "function": {"name": name, "arguments": json.dumps(args)}}]
            message = {"role": "assistant", "content": None, "tool_calls": calls}
        msg = SimpleNamespace(content=message["content"], tool_calls=calls, function_call=None,
                              role="assistant", model_dump=lambda: message)
        return SimpleNamespace(
            choices=[SimpleNamespace(message=msg, finish_reason="stop")], model="scripted", cost=0, id="r",
            usage=SimpleNamespace(prompt_tokens=0, completion_tokens=0, total_tokens=0),
        )

    OpenAIWrapper.create = create
    OpenAIWrapper.extract_text_or_completion_object = classmethod(
        lambda cls, r: [r.choices[0].message.model_dump()] if r.choices[0].message.tool_calls
        else [r.choices[0].message.content]
    )


def ask_team(where: str, scripted: bool) -> bool:
    """Ask the agent to add, then to divide. False if the model can't be reached."""
    print(f"AgentTrust: {where} · agent {AGENT_ID} · model {'scripted' if scripted else MODEL}")
    if scripted:
        use_scripted_model()
    try:
        for request in REQUESTS:
            print(f"\n=== {request} ===")
            print(f"\nResult: {ask(request, llm_config(scripted))}")
    except Exception as exc:  # noqa: BLE001 - any provider error ends the run
        print(f"\nThe model call failed ({MODEL}): {exc}\n"
              "Check OPENAI_API_KEY is set, or run with --scripted.", file=sys.stderr)
        return False
    return True


# --without-agentrust: every tool runs. For comparison.
def run_without_agentrust(scripted: bool) -> int:
    os.environ["AGENTRUST_ENABLED"] = "false"   # the SDK's off switch: wrapped tools pass every call through
    return 0 if ask_team("off (every tool runs)", scripted) else 1


# --local: a gateway runs inside this program; no account needed.
def run_local(scripted: bool) -> int:
    from agentrust_sdk.embedded import embed_gateway

    # SQLite in a temp dir, gone after the run. embed_gateway sets
    # AGENTRUST_GATEWAY_URL and AGENTRUST_KEY, which governed_functions reads.
    gateway = embed_gateway(db_path=os.path.join(tempfile.mkdtemp(), "gateway.db"))
    gateway.register_manifest(AGENT_ID, MANIFEST)
    try:
        ok = ask_team(f"local gateway {gateway.url} (no account, nothing kept)", scripted)
    finally:
        gateway.stop()
    return 0 if ok else 1


# --console: the SDK sends every check to your console, which keeps the decisions.
def run_with_console(scripted: bool) -> int:
    missing = [v for v in ("AGENTRUST_GATEWAY_URL", "AGENTRUST_KEY", "AGENTRUST_AGENT_ID") if not os.environ.get(v)]
    if missing:
        print(f"Set {', '.join(missing)} (from the agent you created in the console), or use --local.",
              file=sys.stderr)
        return 2
    console = os.environ["AGENTRUST_GATEWAY_URL"]
    ok = ask_team(f"console {console}", scripted)
    print(f"\nOpen Events in the console ({console}) and filter by agent ID {AGENT_ID}.")
    return 0 if ok else 1


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="The AG2 calculator agent, governed by AgentTrust.")
    where = parser.add_mutually_exclusive_group()
    where.add_argument("--local", action="store_true", help="run a gateway inside this program, no account (the default)")
    where.add_argument("--console", action="store_true", help="send every check to your AgentTrust console")
    where.add_argument("--without-agentrust", action="store_true", help="for comparison: AgentTrust off, every tool runs")
    parser.add_argument("--scripted", action="store_true", help="replay a fixed conversation: no model, no API key")
    args = parser.parse_args(argv)
    if args.without_agentrust:
        return run_without_agentrust(args.scripted)
    return run_with_console(args.scripted) if args.console else run_local(args.scripted)


if __name__ == "__main__":
    sys.exit(main())

Run it without AgentTrust

First, see what the agent does on its own:

python ag2_quickstart.py --without-agentrust --scripted

The agent is asked to add, then to divide. The scripted agent replays a fixed conversation and does not read tool results, so its last line is the same with AgentTrust on or off. What differs is what each tool returned.

AgentTrust: off (every tool runs) · agent calculator-agent · model scripted

=== Add 3 and 4. ===
>>>>>>>> EXECUTED FUNCTION add...
Output:
7
Result: 3 + 4 = 7.

=== Divide 10 by 2. ===
>>>>>>>> EXECUTED FUNCTION divide...
Output:
5.0
Result: That is everything you asked for.

The agent used divide and nothing stopped it. For a calculator that is harmless. Swap divide for issue_refund or delete_user and it would run just the same.

Where governed_functions sits in the AG2 file

This is the only change AgentTrust needs, and it is already in your file, so there is nothing to edit. Find the governed_functions line in the program, near the top. It sits right after the tools are defined and before the agent is built. On your own agent it looks like this:

+ from agentrust_sdk.adapters.autogen import governed_functions
- tools = [add, multiply, divide]
+ tools = governed_functions([add, multiply, divide], agent_id=AGENT_ID)

Names, descriptions and argument types stay the same, so the agent sees what it saw before. A tool you register unwrapped is not checked. Every tool call now goes through AgentTrust first. Allowed calls run and their results are checked. Refused calls never run.

Run it locally

Now run the same agent with AgentTrust on:

python ag2_quickstart.py --local --scripted

The program starts a small AgentTrust gateway inside itself and allows only the tools in MANIFEST: add and multiply. divide is not on the list. No account needed.

AgentTrust: local gateway http://127.0.0.1:8765 (no account, nothing kept) · agent calculator-agent · model scripted

=== Add 3 and 4. ===
>>>>>>>> EXECUTED FUNCTION add...
Output:
7
Result: 3 + 4 = 7.

=== Divide 10 by 2. ===
>>>>>>>> EXECUTED FUNCTION divide...
Output:
Blocked by AgentTrust: Tool-trust hard gate: action 'divide' is
not in this agent's declared_tools (envelope
ca333621-ea2b-4602-983d-8a8e942a1e7b)
Result: That is everything you asked for.

add ran. divide was refused before its code ran, and the agent was told why, so the run still finished. Nothing is kept after the program exits.

Create the agent in your console

To keep decisions and manage the agent from the console, register it there. In your AgentTrust console:

1. Open Agents, select New agent, name it ag2-quickstart, and select Create.
2. Copy the API Key and the Agent ID before you select Done. The key starts with at_live_ and is shown only once.
3. Select the agent in the list. In the panel that opens, under Manifest & fingerprint, select Declare. Under Declared tools, select + Add tool twice, type add and multiply, and select Save manifest.

Step 3 is the same list MANIFEST holds in the program, now kept in the console.

Run it against the console

Give the program the three values from step 5, then run it with --console:

export AGENTRUST_GATEWAY_URL=https://console.agent-trust.tech   # your console's address
export AGENTRUST_KEY=at_live_...                                # the API key
export AGENTRUST_AGENT_ID=ag2-quickstart                     # the Agent ID

python ag2_quickstart.py --console --scripted

The first line names the console, and the run should match step 4: add returns 7 and divide is refused with the same Blocked by AgentTrust text. The last line tells you where to look. Drop --scripted to use a real model.

See the decisions in the console

Open Events in the console and type your Agent ID into Filter by Agent ID. You will see one row per decision: add approved before and after it ran, and divide blocked. Select a row to see the scores behind the decision.

You are done. Your agent's tool calls are now checked, refused when they should not run, and recorded.

Wrap the functions first. A function you register unwrapped is not checked.

Hold risky calls for a person. Set the agent's mode to Human Review. Rejected or timed-out calls do not run. See docs/framework-adapters/autogen.md.

Check every message too. Add AgentTrustReplyHook(agent_id=...).attach(agent). It checks each message an agent sends. See docs/framework-adapters/autogen.md.

Put @harness on plain Python tools in Stop Risky AI Agent Tool Calls in Python with @harness, or build a full governed bot in Guide: Stop Your AI Support Bot from Issuing a Bad Refund.