Skip to main content

Google ADK + AxonFlow Integration

Prerequisites: Python 3.10+, AxonFlow running (Getting Started), pip install google-adk>=2.0 axonflow>=8.0

What This Integration Gives You​

Google's Agent Development Kit (ADK) ships a first-class plugin model: BasePlugin exposes six hooks that fire around every model call and every tool call across every agent on a Runner. AxonFlow's plugin (axonflow_adk.AxonFlowPlugin) implements all six hooks and routes them through the existing axonflow Python SDK.

The practical result: register the plugin once on your Runner, and AxonFlow governs every interaction — without touching individual agents, tool definitions, or chain code. PII detection, deny short-circuits, HITL approval gates, audit logging, and tool-output redaction all light up automatically.

CapabilityHow it surfaces
Pre-LLM policy enforcementbefore_model_callback calls pre_check. Denials become an LlmResponse short-circuit the agent sees as the model output.
HITL approval gateWhen policy returns require_approval, the plugin polls get_hitl_request until approved/rejected/expired.
Tool-input policybefore_tool_callback calls check_tool_input. Denials return {"error": "[AxonFlow] <reason>"}.
Tool-output redactionafter_tool_callback calls check_tool_output. PII-redacted output replaces the original.
Audit trailafter_model_callback + on_tool_error_callback route to audit_llm_call / audit_tool_call.
MCP toolsetaxonflow_mcp_toolset() returns an McpToolset pointed at AxonFlow's MCP endpoint over Streamable HTTP.

The plugin reuses the published axonflow Python SDK (v8.0+) and inherits its auth, retry, and observability. There is no second HTTP client.

Quickstart​

Five lines of plugin registration is the entire integration:

from google.adk.runners import InMemoryRunner
from google.adk.agents import LlmAgent
from axonflow_adk import AxonFlowPlugin

agent = LlmAgent(model="gemini-2.0-flash", name="loan_desk", instruction="...")
runner = InMemoryRunner(
    agent=agent,
    app_name="loan_desk",
    plugins=[AxonFlowPlugin(
        endpoint="http://localhost:8080",
        client_id="loan-desk",
        client_secret="secret-from-axonflow",
    )],
)

That is it. Every call the runner makes through agent (and any sub-agents whose tools fire on the same Runner) goes through AxonFlow.

Reuse an existing AxonFlow client when you already have one in scope:

from axonflow import AxonFlow
from axonflow_adk import AxonFlowPlugin

axon = AxonFlow(endpoint="...", client_id="...", client_secret="...")
runner = InMemoryRunner(
    agent=agent,
    app_name="loan_desk",
    plugins=[AxonFlowPlugin.from_client(axon)],
)

Hook-to-Endpoint Reference​

The plugin implements all six BasePlugin hooks from google/adk-python. Signatures match exactly (keyword-only args, async).

ADK hookAxonFlow callReturns on denyFails open?
before_model_callbackpre_checkLlmResponse with policy-denial textYes (except require_approval — fail closed)
after_model_callbackaudit_llm_callnever blocks (audit only)Yes
before_tool_callbackcheck_tool_input{"error": "[AxonFlow] <reason>"}Yes (except require_approval — fail closed)
after_tool_callbackcheck_tool_outputredacted dict OR {"error": ...} on hard denyYes
on_tool_error_callbackaudit_tool_callnever blocks (audit only)Yes
on_user_message_callbackno-op (v1)n/an/a

The on_user_message_callback hook is intentionally a no-op in v1. Returning non-None Content there would silently replace the user's message — the wrong tool for governance. Audit of the user prompt happens at before_model_callback time via pre_check.

MCP Toolset Helper​

If you want AxonFlow's governed MCP connectors (PostgreSQL, Snowflake, GCS, …) as ADK tools, axonflow_mcp_toolset() returns a ready-to-use McpToolset configured against the AxonFlow agent's MCP endpoint over Streamable HTTP. The helper emits a canonical Authorization: Basic <base64(client_id:client_secret)> header (or Bearer <token> when bearer_token= is passed) — these are the only auth shapes AxonFlow's MCP server recognizes.

from google.adk.agents import LlmAgent
from axonflow_adk import AxonFlowPlugin, axonflow_mcp_toolset

agent = LlmAgent(
    model="gemini-2.0-flash",
    name="postgres_governed",
    instruction="Answer questions about the production DB.",
    tools=[axonflow_mcp_toolset(
        endpoint="http://localhost:8080",
        client_id="my-app",
        client_secret="secret",
    )],
)

MCP tools surfaced this way are governed twice: once by AxonFlow's MCP connector layer (e.g. SQL injection detection, row-cap enforcement) and again by the plugin's before_tool_callback / after_tool_callback. Two independent gates, by design — the connector layer enforces connector-specific policy, the plugin layer enforces ADK-tool-shaped policy.

Reference: Streamable HTTP transport on the ADK side.

HITL Approval Flow — 4-step​

When AxonFlow policy evaluates to require_approval, the plugin runs the full 4-step HITL flow by default (enable_hitl_polling=True):

before_model_callback / before_tool_callback
    │
    ├─ STEP 1 — gate (pre_check / check_tool_input)
    │            returns blocked, BlockReason == "require_approval"
    │
    ├─ STEP 2 — create_hitl_request POST /api/v1/hitl/queue
    │            returns approval_id (uuid)
    │            (the gate-minted context_id / decision_id is for audit
    │             linkage only — the canonical handle is what create returns)
    │
    ├─ STEP 3 — poll get_hitl_request(approval_id)
    │            every approval_poll_interval_seconds (default 2s)
    │            uses a LOCAL counter; polling errors do NOT trip the
    │            shared circuit breaker
    │
    └─ STEP 4 — resume or deny:
        ├─ status="approved"            → return None (continue)
        ├─ status="rejected" | "expired" → return deny short-circuit
        ├─ N consecutive poll failures   → return deny short-circuit
        └─ time > approval_max_wait_seconds → return deny short-circuit

Defaults: 2-second poll interval, 300-second total ceiling. Both tunable via AxonFlowPluginConfig.

The platform's pre_check and check_tool_input gates set BlockReason="require_approval" (matched exactly against the wire-stable sentinel string — substring matching previously false-positived on reasons containing the word "approval") and mint a correlation context_id / decision_id. They do not create the HITL queue row at those sites — that is the plugin's responsibility via step 2. This is intentional platform architecture: the decision lives in pre_check, the row creation is an explicit POST so callers can carry application-specific context (triggered_policy_id, severity, compliance_framework).

The polling path (step 3) is the only fail-closed path in the plugin. Approvals are safety-critical; defaulting to "allow" on an AxonFlow outage during an approval gate would defeat the gate.

Approving / rejecting out-of-band​

On entry to step 3 the plugin emits a single INFO log:

axonflow hitl AWAITING APPROVAL: request_id=<uuid>; approve via
POST /api/v1/hitl/queue/<uuid>/{approve|reject} (poll_interval=2.0s, max_wait=300s)

Default-configured loggers will surface it (INFO, not DEBUG). The request_id is the canonical handle for both approve and reject — pipe it into your operator UI / Slack bot / pager.

The reviewer (UI, Slack bot, internal portal) then posts the decision via:

curl -X POST $AXONFLOW_ENDPOINT/api/v1/hitl/queue/<approval_id>/approve \
     -H 'Content-Type: application/json' \
     -d '{"reviewer_id":"compliance","reviewer_email":"[email protected]"}'

(or …/<approval_id>/reject for the reject path.) The plugin's polling loop sees the status change on its next iteration and resumes the agent's call (or short-circuits with a deny).

Opting out — deny-fast mode​

Set enable_hitl_polling=False on the config to short-circuit require_approval immediately without enqueuing a row. The host app then drives its own approval workflow (e.g. ticket system, Slack, internal portal). Use this when AxonFlow's HITL queue is not your system of record for approvals.

For a working example see examples/loan_disbursement_agent.py.

Failure Semantics​

A buggy or unreachable AxonFlow must not break the agent. The plugin defends with:

  • Per-hook timeout — default 5 s, configurable via call_timeout_seconds. AxonFlow calls that exceed this are abandoned.
  • Half-open circuit breaker — default opens after 5 consecutive failures, recovers after 30 s. While open, every hook skips the AxonFlow call entirely (no timeout overhead, no log spam).
  • Fail open by default — every hook except _await_hitl_decision returns None on error/timeout/open-circuit, letting the model or tool call proceed unguarded.

If your compliance posture requires fail-closed governance, override the defaults: set a very long breaker_recovery_seconds and add a watchdog that alerts on breaker-open transitions. The plugin logs every guarded failure at WARNING level so the runtime can be wired into your alerting.

Tuning​

All knobs live on AxonFlowPluginConfig:

from axonflow_adk import AxonFlowPlugin, AxonFlowPluginConfig

plugin = AxonFlowPlugin(
    endpoint="http://localhost:8080",
    client_id="loan-desk",
    client_secret="...",
    config=AxonFlowPluginConfig(
        call_timeout_seconds=5.0,                # per-hook deadline
        approval_poll_interval_seconds=2.0,      # HITL polling cadence
        approval_max_wait_seconds=300.0,         # HITL timeout ceiling
        breaker_failure_threshold=5,             # consecutive failures before opening
        breaker_recovery_seconds=30.0,           # how long the breaker stays open
        request_type="loan-desk-chat",           # audit label for pre_check
        tool_connector_type="loan-desk-tool",    # audit label for tool I/O
        tenant_id="bank-1",                      # multi-tenant scoping
        extra_context={"product": "loan_desk"},  # static fields on every audit
    ),
)

user_token is resolved per call in this order:

  1. callback_context.state["axonflow_user_token"]
  2. config.default_user_token (default "anonymous")

The plugin does not fall back to ADK's callback_context.user_id. In enterprise mode the platform expects a JWT signed with the tenant key — silently passing a raw identifier would 401 every call, and the plugin's fail-open default would then silently disable governance.

For enterprise mode, set state["axonflow_user_token"] to a JWT BEFORE each runner.run_async(...) call:

session = runner.session_service.create_session(
    app_name="loan_desk", user_id="cust-001", session_id="sess-A",
)
session.state["axonflow_user_token"] = generate_axonflow_jwt(user_id="cust-001")

For community mode, leave the state key unset — default_user_token="anonymous" is what the platform expects.

Audit endpoints are enterprise-only​

audit_llm_call (called by after_model_callback) and audit_tool_call (called by on_tool_error_callback) are enterprise-tier features. On a community-mode platform they return 401 and the plugin fails open — the audit hooks effectively become silent no-ops. This matches the rest of the AxonFlow SDK surface.

Known ADK Gotchas​

The plugin documents — and where applicable, regression-tests — three upstream ADK behaviors that affect governance:

IssueBehaviorPlugin stance
google/adk-python#2809 AgentTool plugin isolationAgentTool constructs an isolated inner Runner that does not inherit the parent Runner's plugins. Sub-agent tool calls are not governed by the parent plugin.Documented. Test test_agent_tool_plugin_isolation_gotcha_is_documented pins the behavior so future ADK fixes (or regressions) surface in CI.
google/adk-python#4464 InMemoryRunner plugin lifecyclePlugin lifecycle is bound to the Runner, not the session. Re-creating the Runner re-creates the AxonFlow client.Documented. For long-running services keep one Runner per app; do not re-create per request.
google/adk-python#4509 before_model_callback cornerA non-None LlmResponse return can corner-case downstream tool dispatch in older ADK builds.Documented. Mitigation: pin google-adk>=2.0 (or newer).

If you need governance on sub-agents wrapped by AgentTool, register AxonFlowPlugin on the inner Runner as well, or use RemoteA2aAgent for cross-process delegation.

What's NOT in v1​

For transparency, the v1 scope deliberately excludes:

  • User-message rewriting. on_user_message_callback returns None. Returning Content would silently replace the user's message; we treat that as a deliberate ADR-pinned future decision.
  • Output policy on model responses. after_model_callback audits but does not block or rewrite. Output policy on model output is owned by the platform layer (see the Media Governance and Policies overview docs).
  • Zero-dependency rewrite. The plugin calls the existing axonflow Python SDK rather than re-implementing the REST surface. Auth, retry, and observability are inherited.
  • Cross-language parity. ADK Java/Go/TypeScript/Kotlin plugins are a separate workstream tracked on the same epic.

Source + Issues​

Last updated on