agent-safety-box 1.1.0
Battle-hardened safety wrapper for AI agents
License: MIT
Requires Python: >=3.10
Keywords: ai, agent, safety, wrapper, budget, circuit-breaker

# 🛡️ 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

```python
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

```python
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

```python
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

```python
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

```bash
# 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