Reliable AI Agent Pipelines: Orchestration, Retries, Circuit Breakers, and HumanintheLoop

from enum import Enum
from dataclasses import dataclass
from typing import Optional
 
class ErrorType(Enum):
    TRANSIENT_INFRASTRUCTURE = "transient_infrastructure"
    INPUT_VALIDATION = "input_validation"
    TOOL_EXECUTION = "tool_execution"
    LLM_OUTPUT_QUALITY = "llm_output_quality"
    IRREVERSIBLE_ACTION = "irreversible_action"
    CASCADING_CONTEXT = "cascading_context"
 
class ErrorSeverity(Enum):
    LOW = "low"        # Retry transparently
    MEDIUM = "medium"  # Retry with notification
    HIGH = "high"      # Escalate to human
    CRITICAL = "critical"  # Stop pipeline immediately
 
@dataclass
class AgentError:
    error_type: ErrorType
    severity: ErrorSeverity
    message: str
    tool_name: Optional[str] = None
    retry_count: int = 0
    is_retryable: bool = True
    context: Optional[dict] = None
 
def classify_error(exception: Exception,
                    tool_name: Optional[str] = None,
                    retry_count: int = 0) -> AgentError:
    """Classify an exception into the agent error taxonomy."""
    error_msg = str(exception)
 
    # Infrastructure errors
    if isinstance(exception, TimeoutError):
        return AgentError(ErrorType.TRANSIENT_INFRASTRUCTURE, ErrorSeverity.LOW,
                         error_msg, tool_name, retry_count, is_retryable=True)
 
    if hasattr(exception, 'status_code'):
        status = exception.status_code
        if status == 429:
            return AgentError(ErrorType.TRANSIENT_INFRASTRUCTURE, ErrorSeverity.LOW,
                             f"Rate limit exceeded", tool_name, retry_count, is_retryable=True)
        if status in (500, 502, 503, 504):
            return AgentError(ErrorType.TRANSIENT_INFRASTRUCTURE, ErrorSeverity.MEDIUM,
                             f"Service unavailable: {status}", tool_name, retry_count,
                             is_retryable=retry_count < 3)
 
    # Validation errors
    if isinstance(exception, (ValueError, TypeError)) and tool_name:
        return AgentError(ErrorType.INPUT_VALIDATION, ErrorSeverity.MEDIUM,
                         error_msg, tool_name, retry_count,
                         is_retryable=retry_count < 1)  # One correction attempt only
 
    # Default: medium severity, limited retries
    return AgentError(ErrorType.TOOL_EXECUTION, ErrorSeverity.MEDIUM,
                     error_msg, tool_name, retry_count,
                     is_retryable=retry_count < 2)

import asyncio
import random
import time
from functools import wraps
from typing import TypeVar, Callable, Awaitable
 
T = TypeVar("T")
 
def with_exponential_backoff(
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    jitter: bool = True,
    retryable_types: tuple = (TimeoutError,),
):
    """
    Decorator for async tool functions that adds exponential backoff retry.
    Only retries on specified exception types.
    """
    def decorator(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
        @wraps(func)
        async def wrapper(*args, **kwargs) -> T:
            last_exception = None
            for attempt in range(max_retries + 1):
                try:
                    return await func(*args, **kwargs)
                except retryable_types as e:
                    last_exception = e
                    if attempt == max_retries:
                        break
 
                    delay = min(base_delay * (2 ** attempt), max_delay)
                    if jitter:
                        delay *= (0.5 + random.random() * 0.5)  # ±50% jitter
 
                    import structlog
                    log = structlog.get_logger()
                    log.warning("tool_retry",
                               func=func.__name__,
                               attempt=attempt + 1,
                               delay_seconds=round(delay, 2),
                               error=str(e))
 
                    await asyncio.sleep(delay)
 
            raise last_exception
 
        return wrapper
    return decorator
 
# Usage:
@with_exponential_backoff(max_retries=3, retryable_types=(TimeoutError, ConnectionError))
async def call_external_api(url: str, params: dict) -> dict:
    async with aiohttp.ClientSession() as session:
        async with session.get(url, params=params, timeout=aiohttp.ClientTimeout(total=10)) as resp:
            resp.raise_for_status()
            return await resp.json()

async def regenerate_with_correction(
    original_prompt: str,
    bad_output: str,
    error_description: str,
    llm_client,
    max_correction_attempts: int = 2,
) -> str:
    """
    Regenerate LLM output with an explicit correction prompt.
    Used when the first generation is invalid (wrong format, constraint violation, etc.)
    """
    correction_prompt = f"""{original_prompt}
 
---
CORRECTION REQUIRED:
 
Your previous response was:
{bad_output}
 
Problem: {error_description}
 
Please try again, addressing the specific problem described above.
"""
 
    for attempt in range(max_correction_attempts):
        response = await llm_client.complete(correction_prompt)
 
        # Validate the new response
        validation_error = validate_output(response)
        if validation_error is None:
            return response
 
        # Update correction prompt with new error
        correction_prompt = f"""{correction_prompt}
 
---
SECOND CORRECTION ATTEMPT:
 
Your second response was:
{response}
 
Still has problem: {validation_error}
 
This is your final attempt. Be very explicit about following the required format.
"""
 
    # Both attempts failed: return last output for human review
    raise ValueError(f"Failed to generate valid output after {max_correction_attempts} attempts")
 
def validate_output(output: str) -> Optional[str]:
    """
    Validate LLM output. Return error description string if invalid, None if valid.
    Override this with your specific validation logic.
    """
    # Example: validate JSON tool call format
    try:
        if output.strip().startswith("{"):
            json.loads(output)
    except json.JSONDecodeError as e:
        return f"Output is not valid JSON: {e}"
 
    return None  # Valid

class AgentLoopDetector:
    """
    Detects when an agent is stuck in a repetitive loop
    and forces a context reset.
    """
 
    def __init__(self, window_size: int = 5, similarity_threshold: float = 0.9):
        self.window_size = window_size
        self.similarity_threshold = similarity_threshold
        self.recent_actions: list[str] = []
 
    def record_action(self, action_signature: str) -> bool:
        """
        Record an action and return True if a loop is detected.
        action_signature: a hash or summary of the tool call + parameters
        """
        self.recent_actions.append(action_signature)
        if len(self.recent_actions) > self.window_size:
            self.recent_actions.pop(0)
 
        if len(self.recent_actions) < 3:
            return False
 
        # Check if the last N actions are all the same
        unique_actions = set(self.recent_actions[-3:])
        if len(unique_actions) == 1:
            return True  # Exact loop: same action 3 times in a row
 
        # Check for alternating pattern (A, B, A, B, A)
        if len(self.recent_actions) >= 5:
            pattern = self.recent_actions[-5:]
            if pattern[0] == pattern[2] == pattern[4] and pattern[1] == pattern[3]:
                return True
 
        return False
 
    def reset(self):
        self.recent_actions = []

import time
from enum import Enum
from dataclasses import dataclass, field
from threading import Lock
 
class CircuitState(Enum):
    CLOSED = "closed"         # Normal operation
    OPEN = "open"             # Blocking calls
    HALF_OPEN = "half_open"   # Testing if service recovered
 
@dataclass
class CircuitBreaker:
    """
    Circuit breaker for a named tool or external service.
    Thread-safe implementation.
    """
    name: str
    failure_threshold: int = 5        # Failures before opening
    recovery_timeout: float = 60.0    # Seconds before attempting recovery
    success_threshold: int = 2        # Successes needed to close from half-open
 
    state: CircuitState = field(default=CircuitState.CLOSED, init=False)
    failure_count: int = field(default=0, init=False)
    success_count: int = field(default=0, init=False)
    last_failure_time: Optional[float] = field(default=None, init=False)
    _lock: Lock = field(default_factory=Lock, init=False, repr=False)
 
    def call_allowed(self) -> bool:
        """Check if a call should be allowed through."""
        with self._lock:
            if self.state == CircuitState.CLOSED:
                return True
 
            if self.state == CircuitState.OPEN:
                # Check if recovery timeout has passed
                if (self.last_failure_time and
                        time.monotonic() - self.last_failure_time > self.recovery_timeout):
                    self.state = CircuitState.HALF_OPEN
                    self.success_count = 0
                    return True
                return False
 
            if self.state == CircuitState.HALF_OPEN:
                return True  # Allow test calls
 
        return False
 
    def record_success(self):
        with self._lock:
            if self.state == CircuitState.HALF_OPEN:
                self.success_count += 1
                if self.success_count >= self.success_threshold:
                    self.state = CircuitState.CLOSED
                    self.failure_count = 0
            elif self.state == CircuitState.CLOSED:
                self.failure_count = max(0, self.failure_count - 1)  # Decay failures
 
    def record_failure(self):
        with self._lock:
            self.failure_count += 1
            self.last_failure_time = time.monotonic()
 
            if (self.state == CircuitState.CLOSED and
                    self.failure_count >= self.failure_threshold):
                self.state = CircuitState.OPEN
            elif self.state == CircuitState.HALF_OPEN:
                self.state = CircuitState.OPEN  # Test failed, back to open
 
    def status(self) -> dict:
        return {
            "name": self.name,
            "state": self.state.value,
            "failure_count": self.failure_count,
            "time_until_retry": max(
                0,
                self.recovery_timeout - (time.monotonic() - (self.last_failure_time or 0))
            ) if self.state == CircuitState.OPEN else 0,
        }
 
class CircuitBreakerRegistry:
    """Central registry for all circuit breakers in the agent pipeline."""
 
    def __init__(self):
        self._breakers: dict[str, CircuitBreaker] = {}
 
    def get(self, tool_name: str,
             failure_threshold: int = 5,
             recovery_timeout: float = 60.0) -> CircuitBreaker:
        if tool_name not in self._breakers:
            self._breakers[tool_name] = CircuitBreaker(
                name=tool_name,
                failure_threshold=failure_threshold,
                recovery_timeout=recovery_timeout,
            )
        return self._breakers[tool_name]
 
    def all_statuses(self) -> list[dict]:
        return [b.status() for b in self._breakers.values()]
 
    def any_open(self) -> bool:
        return any(b.state == CircuitState.OPEN for b in self._breakers.values())
 
# Integration with tool execution:
async def execute_tool_with_circuit_breaker(
    tool_name: str,
    tool_fn: callable,
    args: dict,
    registry: CircuitBreakerRegistry,
) -> dict:
    breaker = registry.get(tool_name)
 
    if not breaker.call_allowed():
        status = breaker.status()
        raise RuntimeError(
            f"Circuit breaker OPEN for tool '{tool_name}'. "
            f"Retry in {status['time_until_retry']:.0f}s. "
            f"Do not call this tool again until the circuit recovers."
        )
 
    try:
        result = await tool_fn(**args)
        breaker.record_success()
        return result
    except Exception as e:
        breaker.record_failure()
        raise

import hashlib
import json
from datetime import datetime, timezone, timedelta
import sqlite3
 
class IdempotencyStore:
    """
    Stores idempotency keys and previous tool results.
    Prevents duplicate execution of the same tool call on retry.
    """
 
    def __init__(self, db_path: str, ttl_hours: int = 24):
        self.ttl = timedelta(hours=ttl_hours)
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS idempotency_keys (
                key         TEXT PRIMARY KEY,
                tool_name   TEXT NOT NULL,
                result      TEXT NOT NULL,  -- JSON-encoded result
                created_at  TEXT NOT NULL,
                expires_at  TEXT NOT NULL
            )
        """)
        self.conn.commit()
 
    def generate_key(self, tool_name: str, args: dict) -> str:
        """Generate a stable idempotency key from tool name + args."""
        # Sort args for consistent hashing regardless of dict order
        canonical = json.dumps({"tool": tool_name, "args": args}, sort_keys=True)
        return hashlib.sha256(canonical.encode()).hexdigest()
 
    def check(self, key: str) -> Optional[dict]:
        """Return previous result if key exists and is not expired. Else None."""
        now = datetime.now(timezone.utc).isoformat()
        row = self.conn.execute(
            "SELECT result, expires_at FROM idempotency_keys WHERE key=? AND expires_at>?",
            (key, now)
        ).fetchone()
 
        if row:
            return json.loads(row[0])
        return None
 
    def store(self, key: str, tool_name: str, result: dict):
        """Store a tool result with the given idempotency key."""
        now = datetime.now(timezone.utc)
        expires_at = (now + self.ttl).isoformat()
 
        self.conn.execute("""
            INSERT OR REPLACE INTO idempotency_keys VALUES (?,?,?,?,?)
        """, (key, tool_name, json.dumps(result), now.isoformat(), expires_at))
        self.conn.commit()
 
    def cleanup_expired(self):
        """Remove expired keys (run periodically)."""
        now = datetime.now(timezone.utc).isoformat()
        self.conn.execute("DELETE FROM idempotency_keys WHERE expires_at < ?", (now,))
        self.conn.commit()
 
def idempotent_tool(idempotency_store: IdempotencyStore):
    """
    Decorator that makes a tool function idempotent.
    Checks for a previous result before executing, stores result after.
    """
    def decorator(func: callable) -> callable:
        @wraps(func)
        async def wrapper(**kwargs) -> dict:
            tool_name = func.__name__
            key = idempotency_store.generate_key(tool_name, kwargs)
 
            # Check for previous result
            previous_result = idempotency_store.check(key)
            if previous_result is not None:
                import structlog
                log = structlog.get_logger()
                log.info("idempotency_hit",
                         tool=tool_name,
                         key=key[:8],
                         result_source="cache")
                return {**previous_result, "_idempotency_hit": True}
 
            # Execute the tool
            result = await func(**kwargs)
 
            # Store the result
            idempotency_store.store(key, tool_name, result)
            return result
 
        return wrapper
    return decorator
 
# Usage:
idempotency = IdempotencyStore("agent_idempotency.db", ttl_hours=24)
 
@idempotent_tool(idempotency)
async def send_notification(user_id: str, message: str, channel: str) -> dict:
    # This function will only execute once per unique (user_id, message, channel) combination
    # within the TTL window, regardless of how many times the agent calls it.
    result = await notification_service.send(user_id, message, channel)
    return {"sent": True, "message_id": result.id, "timestamp": result.timestamp}

from enum import Enum
from dataclasses import dataclass
from typing import Optional, Callable, Awaitable
import asyncio
 
class EscalationReason(Enum):
    IRREVERSIBLE_ACTION = "irreversible_action"
    LOW_CONFIDENCE = "low_confidence"
    REPEATED_FAILURE = "repeated_failure"
    CONTEXT_CONTRADICTION = "context_contradiction"
    POLICY_BOUNDARY = "policy_boundary"
 
class EscalationOutcome(Enum):
    APPROVED = "approved"
    REJECTED = "rejected"
    MODIFIED = "modified"    # Human approves with modifications
    TIMEOUT = "timeout"      # Human did not respond in time
 
@dataclass
class EscalationRequest:
    reason: EscalationReason
    action_description: str    # What the agent wants to do
    context_summary: str       # Why it wants to do it
    options: list[str]         # Suggested options for the human reviewer
    timeout_seconds: float = 300.0  # 5 minute default
    default_on_timeout: EscalationOutcome = EscalationOutcome.REJECTED
 
@dataclass
class EscalationResponse:
    outcome: EscalationOutcome
    instructions: Optional[str] = None    # For MODIFIED outcome
    reviewer_id: Optional[str] = None
    reviewed_at: Optional[str] = None
 
class HumanEscalationGate:
    """
    Gate that pauses agent execution pending human review.
    Integrates with whatever notification/review system you use.
    """
 
    def __init__(self,
                 notify_fn: Callable[[EscalationRequest], Awaitable[str]],
                 poll_fn: Callable[[str], Awaitable[Optional[EscalationResponse]]]):
        """
        notify_fn: send escalation to human, returns a request_id
        poll_fn: check if human has responded, returns response or None
        """
        self.notify = notify_fn
        self.poll = poll_fn
        self.pending_requests: dict[str, EscalationRequest] = {}
 
    async def request_approval(self, request: EscalationRequest) -> EscalationResponse:
        """
        Send an escalation request and wait for human response.
        Returns response or default on timeout.
        """
        # Send notification
        request_id = await self.notify(request)
        self.pending_requests[request_id] = request
 
        # Poll for response
        deadline = asyncio.get_event_loop().time() + request.timeout_seconds
        poll_interval = 5.0  # Check every 5 seconds
 
        while asyncio.get_event_loop().time() < deadline:
            response = await self.poll(request_id)
            if response is not None:
                del self.pending_requests[request_id]
                return response
            await asyncio.sleep(poll_interval)
 
        # Timeout
        del self.pending_requests.get(request_id, None)
        return EscalationResponse(
            outcome=request.default_on_timeout,
            instructions=f"Auto-{request.default_on_timeout.value} due to timeout after {request.timeout_seconds}s"
        )
 
def requires_approval(action_category: str,
                       risk_level: str = "high",
                       gate: Optional[HumanEscalationGate] = None):
    """
    Decorator for tool functions that require human approval before execution.
    """
    def decorator(func: callable) -> callable:
        @wraps(func)
        async def wrapper(**kwargs) -> dict:
            if gate is None:
                # No gate configured: allow through (development mode)
                return await func(**kwargs)
 
            # Build escalation request
            args_summary = json.dumps(kwargs, indent=2)[:500]
            request = EscalationRequest(
                reason=EscalationReason.IRREVERSIBLE_ACTION,
                action_description=f"Execute {func.__name__}({action_category})",
                context_summary=f"Tool: {func.__name__}\nRisk: {risk_level}\nParameters:\n{args_summary}",
                options=["Approve", "Reject", "Modify parameters"],
                default_on_timeout=EscalationOutcome.REJECTED,
            )
 
            response = await gate.request_approval(request)
 
            if response.outcome == EscalationOutcome.APPROVED:
                return await func(**kwargs)
            elif response.outcome == EscalationOutcome.MODIFIED:
                # Parse modified instructions and re-execute with new params
                # This requires structured modification format from the reviewer
                modified_kwargs = parse_modification_instructions(
                    response.instructions, kwargs
                )
                return await func(**modified_kwargs)
            else:  # REJECTED or TIMEOUT
                return {
                    "status": "rejected",
                    "reason": response.instructions or "Rejected by human reviewer",
                    "escalation_outcome": response.outcome.value,
                }
 
        return wrapper
    return decorator
 
# Usage:
@requires_approval(action_category="email_send", risk_level="high")
async def send_customer_email(customer_id: str, subject: str, body: str) -> dict:
    return await email_service.send(customer_id, subject, body)

class SequentialPipeline:
    def __init__(self, steps: list[callable], circuit_registry: CircuitBreakerRegistry):
        self.steps = steps
        self.circuits = circuit_registry
 
    async def run(self, initial_input: dict) -> dict:
        context = initial_input
        for step_fn in self.steps:
            if self.circuits.any_open():
                open_circuits = [s["name"] for s in self.circuits.all_statuses()
                                if s["state"] == "open"]
                raise RuntimeError(f"Pipeline halted: open circuits: {open_circuits}")
            context = await step_fn(context)
        return context

async def parallel_with_partial_failure(
    tasks: dict[str, Callable[[], Awaitable[dict]]],
    required_tasks: set[str],
) -> dict[str, dict]:
    """
    Run tasks in parallel. Required tasks must succeed; optional tasks
    may fail without halting the pipeline.
    """
    results = await asyncio.gather(
        *[task() for task in tasks.values()],
        return_exceptions=True
    )
 
    output = {}
    failed_required = []
 
    for (task_name, _), result in zip(tasks.items(), results):
        if isinstance(result, Exception):
            if task_name in required_tasks:
                failed_required.append((task_name, str(result)))
            else:
                output[task_name] = {"error": str(result), "skipped": True}
        else:
            output[task_name] = result
 
    if failed_required:
        raise RuntimeError(
            f"Required pipeline tasks failed: "
            + "; ".join(f"{name}: {err}" for name, err in failed_required)
        )
 
    return output

async def hierarchical_agent_run(goal: str,
                                  planner_agent,
                                  executor_pool: dict[str, callable],
                                  critic_agent,
                                  human_gate: HumanEscalationGate,
                                  max_plan_revisions: int = 2) -> dict:
    """Hierarchical pipeline: plan → validate → execute → critique."""
 
    # 1. Generate plan
    plan = await planner_agent.plan(goal)
 
    # 2. Validate plan before execution
    for revision in range(max_plan_revisions + 1):
        validation_errors = validate_plan(plan, executor_pool)
 
        if not validation_errors:
            break
 
        if revision == max_plan_revisions:
            # Plan never validated: escalate to human
            response = await human_gate.request_approval(EscalationRequest(
                reason=EscalationReason.REPEATED_FAILURE,
                action_description=f"Execute plan for: {goal}",
                context_summary=f"Plan failed validation after {max_plan_revisions} revisions.\nErrors: {validation_errors}",
                options=["Provide corrected plan", "Cancel goal"],
            ))
            if response.outcome != EscalationOutcome.APPROVED:
                return {"status": "cancelled", "reason": "Plan validation failed"}
 
        plan = await planner_agent.revise_plan(plan, validation_errors)
 
    # 3. Execute plan steps
    results = {}
    for step in plan.steps:
        executor = executor_pool.get(step.executor_type)
        if not executor:
            raise ValueError(f"No executor for type: {step.executor_type}")
        results[step.id] = await executor(step, results)
 
    # 4. Critique output
    critique = await critic_agent.evaluate(goal, plan, results)
    if critique.has_issues:
        results["critique"] = critique.issues
 
    return {"status": "complete", "results": results, "critique": critique}

import structlog
import uuid
from contextvars import ContextVar
 
current_trace_id: ContextVar[str] = ContextVar("trace_id", default="")
current_span_id: ContextVar[str] = ContextVar("span_id", default="")
 
def get_logger():
    return structlog.get_logger().bind(
        trace_id=current_trace_id.get(),
        span_id=current_span_id.get(),
    )
 
class AgentTracer:
    """Structured tracing for agent pipeline execution."""
 
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.trace_id = str(uuid.uuid4())[:8]
        current_trace_id.set(self.trace_id)
        self.log = get_logger().bind(session_id=session_id)
 
    def log_turn_start(self, user_message: str, context_tokens: int):
        self.log.info("turn_start",
                      message_preview=user_message[:100],
                      context_tokens=context_tokens)
 
    def log_tool_call(self, tool_name: str, args: dict) -> str:
        span_id = str(uuid.uuid4())[:8]
        current_span_id.set(span_id)
        self.log.info("tool_call",
                      tool=tool_name,
                      args_summary={k: str(v)[:50] for k, v in args.items()},
                      span_id=span_id)
        return span_id
 
    def log_tool_result(self, span_id: str, tool_name: str,
                        result: dict, duration_ms: float):
        self.log.info("tool_result",
                      tool=tool_name,
                      span_id=span_id,
                      duration_ms=round(duration_ms, 1),
                      result_preview=str(result)[:100],
                      success=True)
 
    def log_tool_error(self, span_id: str, tool_name: str,
                       error: AgentError, duration_ms: float):
        self.log.warning("tool_error",
                         tool=tool_name,
                         span_id=span_id,
                         duration_ms=round(duration_ms, 1),
                         error_type=error.error_type.value,
                         error_severity=error.severity.value,
                         message=error.message,
                         retry_count=error.retry_count)
 
    def log_escalation(self, reason: EscalationReason, outcome: EscalationOutcome):
        self.log.warning("human_escalation",
                         reason=reason.value,
                         outcome=outcome.value)
 
    def log_turn_end(self, response_preview: str, total_tool_calls: int,
                      total_duration_ms: float):
        self.log.info("turn_end",
                      response_preview=response_preview[:100],
                      total_tool_calls=total_tool_calls,
                      total_duration_ms=round(total_duration_ms, 1))

class FaultInjector:
    """
    Test utility: inject configured failures into tool calls.
    Use to validate that retry logic, circuit breakers, and escalation work.
    """
 
    def __init__(self):
        self.fault_configs: dict[str, list[dict]] = {}
 
    def configure(self, tool_name: str, faults: list[dict]):
        """
        faults: list of {
            "on_call": int,     # Which call number triggers this fault (1-indexed)
            "exception": Exception,  # Exception to raise
            "delay": float,    # Optional extra delay in seconds
        }
        """
        self.fault_configs[tool_name] = faults
        self._call_counts = {}
 
    def wrap_tool(self, tool_name: str, tool_fn: callable) -> callable:
        """Wrap a tool function with fault injection."""
        @wraps(tool_fn)
        async def wrapper(**kwargs):
            call_count = self._call_counts.get(tool_name, 0) + 1
            self._call_counts[tool_name] = call_count
 
            faults = self.fault_configs.get(tool_name, [])
            for fault in faults:
                if fault.get("on_call") == call_count:
                    if fault.get("delay"):
                        await asyncio.sleep(fault["delay"])
                    if fault.get("exception"):
                        raise fault["exception"]
 
            return await tool_fn(**kwargs)
        return wrapper
 
# Test: verify retry on transient failure
async def test_retry_on_rate_limit():
    injector = FaultInjector()
    injector.configure("search_web", [
        {"on_call": 1, "exception": RateLimitError("429 Too Many Requests")},
        {"on_call": 2, "exception": RateLimitError("429 Too Many Requests")},
        # Call 3 succeeds
    ])
 
    wrapped_search = injector.wrap_tool("search_web", real_search_web)
    result = await retry_with_backoff(wrapped_search, query="test query")
    assert result is not None, "Should succeed on third attempt"
    assert injector._call_counts["search_web"] == 3, "Should have made exactly 3 calls"
 
# Test: verify circuit breaker opens after threshold
async def test_circuit_breaker_opens():
    registry = CircuitBreakerRegistry()
    breaker = registry.get("flaky_api", failure_threshold=3)
 
    # Record failures
    for _ in range(3):
        breaker.record_failure()
 
    assert breaker.state == CircuitState.OPEN, "Circuit should be open after threshold"
 
    # Verify call is blocked
    assert not breaker.call_allowed(), "Calls should be blocked when circuit is open"
 
# Test: verify idempotency prevents double-write
async def test_idempotent_tool_deduplicates():
    store = IdempotencyStore(":memory:", ttl_hours=1)  # In-memory SQLite for testing
    call_count = 0
 
    @idempotent_tool(store)
    async def write_record(record_id: str, data: dict) -> dict:
        nonlocal call_count
        call_count += 1
        return {"written": True, "record_id": record_id}
 
    # Call twice with same parameters
    result1 = await write_record(record_id="test-123", data={"value": 42})
    result2 = await write_record(record_id="test-123", data={"value": 42})
 
    assert call_count == 1, "Tool should only execute once despite two calls"
    assert result2.get("_idempotency_hit"), "Second call should be a cache hit"