agent-safety-box 1.1.0

Battle-hardened safety wrapper for AI agents

🛡️ agent_safety_box

Battle-hardened safety wrapper for AI agents — zero external dependencies, production-grade, fully tested.

pip install agent-safety-box

---

##The Problem AI agents are powerful, but they fail unpredictably in production. They hallucinate, get stuck in loops, blow through budgets, and crash silently. As one developer put it, "we're spending 40% of our budget on garbage output." Debugging becomes guesswork. agent_safety_box is the answer .

Features

Feature	AgentWrapper	AgentWrapperV2
Budget enforcement (thread-safe)	✅	✅
Per-run cost ceiling	✅	✅
Timeout (sync & async)	✅	✅
Retry + exponential backoff + jitter	✅	✅
Circuit Breaker	✅	✅
Rate Limiter (token-bucket)	✅	✅
Output validation	✅	✅
JSONL audit logging + rotation	✅	✅
Dry-run mode	✅	✅
Tags & metadata	✅	✅
Context manager	✅	✅
Lifecycle hooks (before/after/error)	—	✅
Bulkhead (concurrency limiter)	—	✅
Adaptive timeout (auto p95-tuned)	—	✅
Rolling metrics (p50/p95/p99)	—	✅
Cost forecaster	—	✅
Plugin registry	—	✅
Budget warning callback	—	✅
Graceful shutdown	—	✅
Fallback chain (multi-agent failover)	—	✅

---

Quick Start

from agent_safety_box import AgentWrapper, CircuitBreaker, RateLimiter

cb = CircuitBreaker(failure_threshold=3, reset_timeout=60)
rl = RateLimiter(rate=5.0)

wrapper = AgentWrapper(
    agent=my_agent,
    budget=100.0,
    max_runtime=10.0,
    max_cost_per_run=5.0,
    max_retries=2,
    retry_delay=1.0,
    circuit_breaker=cb,
    rate_limiter=rl,
    output_validator=lambda o: isinstance(o, str) and len(o) > 0,
    tags={"env": "prod", "model": "gpt-4o"},
    log_path="logs/audit.jsonl",
)

with wrapper:
    result = wrapper.run("Summarise this document", cost=1.5)

---

AgentWrapperV2 — Advanced Pipeline

from agent_safety_box.wrapper_v2 import AgentWrapperV2
from agent_safety_box.middleware import HookManager, Bulkhead, AdaptiveTimeout

hooks = HookManager()
hooks.on_before(lambda ctx: print(f"→ {ctx['task']}"))
hooks.on_after(lambda ctx: print(f"✓ done"))
hooks.on_error(lambda ctx, exc: print(f"✗ {exc}"))

w = AgentWrapperV2(
    agent=my_agent,
    budget=500.0,
    hooks=hooks,
    bulkhead=Bulkhead(max_concurrent=10),
    adaptive_timeout=AdaptiveTimeout(initial=5.0),
    on_budget_warning=lambda rem, bud: print(f"⚠ Only {rem:.2f} left!"),
    budget_warning_pct=0.20,
    plugins={"logger": lambda ctx: my_logger.info(ctx)},
)

result = w.run("task", cost=1.0)
print(w.metrics)
# {
#   'total_runs': 1, 'successes': 1, 'errors': 0,
#   'spent': 1.0, 'remaining': 499.0,
#   'rolling': {'p50_s': 0.003, 'p95_s': 0.003, 'p99_s': 0.003,
#               'error_rate': 0.0, 'throughput_rps': 312.5, 'sample_size': 1},
#   'avg_cost_forecast': 1.0,
#   'bulkhead_active': 0, 'bulkhead_max': 10,
#   'adaptive_timeout_s': 5.0,
# }

---

Fallback Chain

from agent_safety_box.middleware import FallbackChain

chain = FallbackChain(
    primary=gpt4_agent,
    fallbacks=[gpt35_agent, local_llm_agent],
)
result = chain.run("task")
print(f"Used: {chain.last_used}")
print(chain.stats)

---

Async Support

from agent_safety_box import AsyncAgentWrapper

wrapper = AsyncAgentWrapper(agent=my_async_agent, budget=50.0)

async def main():
    result = await wrapper.run("Classify this text", cost=0.5)

asyncio.run(main())

---

Exceptions

All exceptions inherit from AgentSafetyBoxError:

Exception	When raised
BudgetExceededError	Total budget would be breached
MaxCostExceededError	Per-run cost ceiling exceeded
TimeoutExceededError	Agent exceeded max_runtime
RetryExhaustedError	All retry attempts failed
AgentValidationError	output_validator rejected output
CircuitOpenError	Circuit breaker is OPEN
RateLimitError	Rate limiter rejected (non-blocking mode)

---

Safety Pipeline (per run)

1. max_cost_per_run check
2. budget check  (thread/async-safe lock)
3. rate_limiter.acquire()
4. circuit_breaker.allow_request()
5. agent.run(task)  [with timeout]
6. output_validator(result)
7. actual_cost adjustment (if agent returns {"actual_cost": X})
8. budget commit
9. circuit_breaker.record_success/failure()
10. audit log write

---

Running Tests

# All 263 tests
python -m unittest discover -s agent_safety_box/tests -p "test_*.py" -v

# Single module
python -m unittest agent_safety_box.tests.test_ultimate -v

---

Architecture

agent_safety_box/
├── __init__.py          # Public API
├── wrapper.py           # AgentWrapper + AsyncAgentWrapper (v1)
├── wrapper_v2.py        # AgentWrapperV2 + AsyncAgentWrapperV2 (v2)
├── middleware.py        # HookManager, RollingMetrics, Bulkhead,
│                        # AdaptiveTimeout, CostForecaster, FallbackChain
├── safety.py            # CircuitBreaker + RateLimiter
├── exceptions.py        # All custom exceptions
├── logger.py            # AuditLogger (JSONL + rotation)
└── tests/
    ├── conftest.py      # Shared stubs and base class
    ├── test_budget.py   # Budget enforcement (31 tests)
    ├── test_execution.py# Timeout + retry (30 tests)
    ├── test_logger.py   # Audit logging (21 tests)
    ├── test_safety.py   # Circuit + rate limiter (31 tests)
    ├── test_stress.py   # Concurrency + chaos (18 tests)
    └── test_ultimate.py # Deep edge cases (132 tests)

---

License

MIT