Safety Guardrails & Token Budgeting for ADK Agents

Production agent harnesses require robust operational protections beyond system prompts and tool definitions. To run AI agents safely at enterprise scale, developers must implement layered guardrails to block prompt injections, scrub sensitive personally identifiable information (PII), enforce token budgets, and prevent runaway execution loops.

This guide explores the three core security and reliability pillars provided by Google ADK callbacks: input/output guardrails, token budgeting, and automatic circuit breakers.

Input & Output Guardrails

Guardrails execute as interceptor callbacks within the agent lifecycle — checking user inputs before model processing (input guardrails) and sanitizing model outputs before user delivery (output guardrails).

Input Guardrails (before_agent_callback)

Input guardrails inspect incoming user prompts and reject malicious payloads before consuming LLM tokens:

async def input_safety_check(callback_context: CallbackContext) -> None:
    """Intercepts incoming prompts and checks for jailbreak or injection attacks."""
    user_message = callback_context.user_message

    if contains_injection_pattern(user_message):
        # Immediately halt execution and return safe fallback response
        callback_context.stop_execution = True
        callback_context.response = "I cannot process this request due to safety policies."

Essential input checks include:

• Prompt Injection Detection: Scanning for phrases like "ignore all previous instructions".
• Payload Size Capping: Rejecting abnormally large inputs designed to overflow context windows.
• Identity & Authorization: Verifying user session permissions before routing to sensitive domains.

Output Guardrails (after_agent_callback)

Output guardrails sanitize generated text to prevent data leakage or policy violations:

import re

async def output_pii_redactor(callback_context: CallbackContext) -> None:
    """Redacts PII patterns (Credit Cards, SSNs) from agent outputs."""
    response = callback_context.agent_response

    # Mask Credit Card Numbers (16 digits with optional dashes/spaces)
    response = re.sub(r'\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b', '****-****-****-****', response)

    # Mask Social Security Numbers (SSN: XXX-XX-XXXX)
    response = re.sub(r'\b\d{3}-\d{2}-\d{4}\b', '***-**-****', response)

    callback_context.agent_response = response

Token Budgeting & Circuit Breakers

Uncontrolled agent reasoning loops can rapidly degrade application performance and incur massive API bills. The agent harness must enforce strict operational boundaries.

Session Turn Budget Limits

Track and restrict maximum turns per user session turn using session state:

async def enforce_token_budget(callback_context: CallbackContext) -> bool:
    """Limits the maximum number of interaction turns per session."""
    state = callback_context.state
    turn_count = state.get("turn_count", 0) + 1
    state["turn_count"] = turn_count

    if turn_count > 10:
        print("[CIRCUIT BREAKER] Exceeded maximum turn limit (10). Halting agent.")
        return False  # Cancels further agent execution
    return True

Step Count Caps

Limit maximum consecutive tool execution iterations to stop agents stuck in recursive loops:

MAX_CONSECUTIVE_TOOL_STEPS = 10

Automated Circuit Breakers

Trip and temporarily suspend agent execution when tool failure rates exceed safety thresholds:

TOOL_FAILURE_THRESHOLD = 0.20  # 20% failure rate threshold
ERROR_WINDOW_SECONDS = 300     # 5-minute evaluation window

async def circuit_breaker_monitor(callback_context: CallbackContext) -> bool:
    """Monitors rolling failure rate and trips circuit breaker if tool errors spike."""
    recent_failures = get_recent_tool_failures(window_seconds=ERROR_WINDOW_SECONDS)
    total_calls = get_recent_tool_calls(window_seconds=ERROR_WINDOW_SECONDS)

    if total_calls > 0 and (recent_failures / total_calls) > TOOL_FAILURE_THRESHOLD:
        logger.critical("Circuit breaker TRIPPED: High tool failure rate detected.")
        return False
    return True

Wiring Up Complete Defense in Depth

Integrate all safety layers directly into the ADK Agent definition:

agent = Agent(
    name="safe_enterprise_agent",
    model="gemini-2.5-flash",
    instruction="You are an enterprise support assistant.",
    tools=[...],
    before_agent_callback=input_safety_check,
    after_agent_callback=output_pii_redactor,
    before_tool_callback=circuit_breaker_monitor
)

Defense-in-Depth Security Matrix

Layer	Interception Point	Primary Defense Function
Input Guardrail	before_agent_callback	Blocks prompt injections and malicious payloads
Token Budget	Session state counter	Prevents unexpected cost spikes
Circuit Breaker	before_tool_callback	Prevents cascading infrastructure failures
Output Guardrail	after_agent_callback	Redacts PII and sensitive data leakage
Step Cap	max_iterations setting	Prevents infinite tool-execution loops

Key Takeaway

System security and stability cannot rely on single-point assertions. By combining ADK callbacks across input validation, PII redaction, token budgets, and circuit breakers, you create a resilient, production-ready AI agent harness.