From 44b41d6ee47241fc96aaba2e1ae8790f0d8e0587 Mon Sep 17 00:00:00 2001 From: SakethSumanBathini Date: Fri, 10 Jul 2026 22:17:05 +0530 Subject: [PATCH 1/2] feat: active circuit breaker with safe pause & resume (closes #483) --- agentwatch/circuit_breaker/__init__.py | 24 ++ agentwatch/circuit_breaker/breaker.py | 415 ++++++++++++++++++++ agentwatch/circuit_breaker/thresholds.py | 53 +++ tests/test_circuit_breaker.py | 464 +++++++++++++++++++++++ 4 files changed, 956 insertions(+) create mode 100644 agentwatch/circuit_breaker/__init__.py create mode 100644 agentwatch/circuit_breaker/breaker.py create mode 100644 agentwatch/circuit_breaker/thresholds.py create mode 100644 tests/test_circuit_breaker.py diff --git a/agentwatch/circuit_breaker/__init__.py b/agentwatch/circuit_breaker/__init__.py new file mode 100644 index 00000000..0f74c05b --- /dev/null +++ b/agentwatch/circuit_breaker/__init__.py @@ -0,0 +1,24 @@ +"""Active circuit breaker for AgentWatch. + +Provides a CLOSED -> OPEN -> PAUSED -> HALF_OPEN -> CLOSED state machine that +watches agent execution for threshold breaches (tokens, errors, hallucinations), +can safely pause and resume an agent by checkpointing its state through the +existing rollback engine, and records every transition for EU AI Act Article 12 +record-keeping through the existing governance package. +""" + +from agentwatch.circuit_breaker.breaker import ( + CircuitBreaker, + CircuitState, + TransitionRecord, + TripReason, +) +from agentwatch.circuit_breaker.thresholds import CircuitThresholds + +__all__ = [ + "CircuitBreaker", + "CircuitState", + "TripReason", + "TransitionRecord", + "CircuitThresholds", +] diff --git a/agentwatch/circuit_breaker/breaker.py b/agentwatch/circuit_breaker/breaker.py new file mode 100644 index 00000000..a5ed52e5 --- /dev/null +++ b/agentwatch/circuit_breaker/breaker.py @@ -0,0 +1,415 @@ +"""Active circuit breaker state machine. + +State model +---------- + CLOSED normal operation; events flow through and thresholds are watched + OPEN a threshold tripped; the agent should stop issuing new actions + PAUSED state has been checkpointed and the agent is safely suspended + HALF_OPEN probing recovery; a limited number of events are allowed through + CLOSED probes succeeded; counters reset and normal operation resumes + +Allowed transitions:: + + CLOSED -> OPEN (threshold breach) + OPEN -> PAUSED (pause(): checkpoint saved) + PAUSED -> HALF_OPEN (resume(): begin probing) + OPEN -> HALF_OPEN (resume() without an intervening pause) + HALF_OPEN -> CLOSED (enough consecutive successful probes) + HALF_OPEN -> OPEN (a probe failed) + any -> CLOSED (manual reset) + +The breaker is deliberately framework-agnostic: it consumes ``AgentEvent`` +objects and never calls into a specific agent runtime. Pausing delegates to the +existing :class:`~agentwatch.rollback.engine.RollbackEngine` for checkpointing, +and every transition is recorded for EU AI Act Article 12 record-keeping via the +existing :class:`~agentwatch.governance.eu_ai_act.EUAIActPackage`. +""" + +from __future__ import annotations + +import hashlib +import logging +import uuid +from dataclasses import dataclass, field +from datetime import UTC, datetime +from enum import Enum +from typing import Any + +from agentwatch.circuit_breaker.thresholds import CircuitThresholds +from agentwatch.core.schema import AgentEvent, EventType, ExecutionStatus + +logger = logging.getLogger(__name__) + + +class CircuitState(str, Enum): + CLOSED = "closed" + OPEN = "open" + PAUSED = "paused" + HALF_OPEN = "half_open" + + +class TripReason(str, Enum): + TOKEN_BUDGET = "token_budget_exceeded" # noqa: S105 # nosec B105 — enum label, not a secret + CONSECUTIVE_ERRORS = "consecutive_errors_exceeded" + TOTAL_ERRORS = "total_errors_exceeded" + HALLUCINATIONS = "hallucination_threshold_exceeded" + PROBE_FAILED = "half_open_probe_failed" + MANUAL = "manual_trip" + + +# Which allowed target states each state may move to. Used to reject invalid +# transitions loudly instead of silently corrupting the breaker's state. +_ALLOWED: dict[CircuitState, set[CircuitState]] = { + CircuitState.CLOSED: {CircuitState.OPEN, CircuitState.CLOSED}, + CircuitState.OPEN: {CircuitState.PAUSED, CircuitState.HALF_OPEN, CircuitState.CLOSED}, + CircuitState.PAUSED: {CircuitState.HALF_OPEN, CircuitState.CLOSED}, + CircuitState.HALF_OPEN: {CircuitState.CLOSED, CircuitState.OPEN}, +} + + +@dataclass +class TransitionRecord: + """A single audited state transition.""" + + when: datetime + from_state: CircuitState + to_state: CircuitState + reason: str + session_id: str + checkpoint_id: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "when": self.when.isoformat(), + "from_state": self.from_state.value, + "to_state": self.to_state.value, + "reason": self.reason, + "session_id": self.session_id, + "checkpoint_id": self.checkpoint_id, + "metadata": dict(self.metadata), + } + + +class CircuitBreaker: + """Active circuit breaker for a single agent session.""" + + def __init__( + self, + session_id: str, + thresholds: CircuitThresholds | None = None, + *, + rollback_engine: Any | None = None, + eu_ai_act_package: Any | None = None, + ) -> None: + self.session_id = session_id + self.thresholds = thresholds or CircuitThresholds() + self._rollback_engine = rollback_engine + self._eu_ai_act = eu_ai_act_package + + self._state = CircuitState.CLOSED + self._history: list[TransitionRecord] = [] + + # Window counters, reset whenever the breaker (re)enters CLOSED. + self._total_tokens = 0 + self._consecutive_errors = 0 + self._total_errors = 0 + self._hallucinations = 0 + self._events_observed = 0 + + # HALF_OPEN probe tracking. + self._probe_successes = 0 + + # The most recent checkpoint created by pause(), used by resume(). + self._last_checkpoint_id: str | None = None + self._last_step_number = 0 + + # ------------------------------------------------------------------ state + + @property + def state(self) -> CircuitState: + return self._state + + @property + def is_tripped(self) -> bool: + """True when the agent should not be issuing new actions.""" + return self._state in (CircuitState.OPEN, CircuitState.PAUSED) + + @property + def history(self) -> list[TransitionRecord]: + return list(self._history) + + def status(self) -> dict[str, Any]: + """A snapshot of the breaker suitable for a dashboard payload.""" + return { + "session_id": self.session_id, + "state": self._state.value, + "is_tripped": self.is_tripped, + "counters": { + "total_tokens": self._total_tokens, + "consecutive_errors": self._consecutive_errors, + "total_errors": self._total_errors, + "hallucinations": self._hallucinations, + "events_observed": self._events_observed, + }, + "probe_successes": self._probe_successes, + "last_checkpoint_id": self._last_checkpoint_id, + "transitions": len(self._history), + } + + # ------------------------------------------------------------- observe + + def observe(self, event: AgentEvent) -> CircuitState: + """Feed a single event to the breaker and return the resulting state. + + In CLOSED state, thresholds are evaluated and the breaker may trip to + OPEN. In HALF_OPEN state, the event is treated as a recovery probe. In + OPEN/PAUSED state, events are recorded in the counters but do not change + state (the agent is expected to have stopped; callers drive recovery via + :meth:`pause` / :meth:`resume`). + """ + if event.session_id != self.session_id: + # Ignore events for other sessions rather than corrupting counters. + return self._state + + self._events_observed += 1 + self._accumulate(event) + + if self._state is CircuitState.CLOSED: + reason = self._evaluate_thresholds() + if reason is not None: + self._transition(CircuitState.OPEN, reason.value) + elif self._state is CircuitState.HALF_OPEN: + self._evaluate_probe(event) + + return self._state + + def _accumulate(self, event: AgentEvent) -> None: + if event.token_usage is not None: + self._total_tokens += int(getattr(event.token_usage, "total_tokens", 0) or 0) + + if self._is_error(event): + self._total_errors += 1 + self._consecutive_errors += 1 + elif self._is_success(event): + self._consecutive_errors = 0 + + if self._is_hallucination(event): + self._hallucinations += 1 + + @staticmethod + def _is_error(event: AgentEvent) -> bool: + return event.status in ( + ExecutionStatus.FAILURE, + ExecutionStatus.TIMEOUT, + ) or event.event_type in ( + EventType.AGENT_ERROR, + EventType.TOOL_ERROR, + EventType.TASK_FAIL, + ) + + @staticmethod + def _is_success(event: AgentEvent) -> bool: + return event.status is ExecutionStatus.SUCCESS or event.event_type in ( + EventType.AGENT_END, + EventType.TOOL_RESULT, + EventType.TASK_COMPLETE, + ) + + def _is_hallucination(self, event: AgentEvent) -> bool: + if event.confidence is None: + return False + if event.confidence.anomaly_flags: + return True + if self.thresholds.min_confidence <= 0.0: + return False + return event.confidence.overall_score < self.thresholds.min_confidence + + def _evaluate_thresholds(self) -> TripReason | None: + t = self.thresholds + if t.max_total_tokens and self._total_tokens > t.max_total_tokens: + return TripReason.TOKEN_BUDGET + if t.max_consecutive_errors and self._consecutive_errors >= t.max_consecutive_errors: + return TripReason.CONSECUTIVE_ERRORS + if t.max_total_errors and self._total_errors >= t.max_total_errors: + return TripReason.TOTAL_ERRORS + if t.max_hallucinations and self._hallucinations >= t.max_hallucinations: + return TripReason.HALLUCINATIONS + return None + + def _evaluate_probe(self, event: AgentEvent) -> None: + """In HALF_OPEN, count successful probes; a failure re-opens.""" + if self._is_error(event): + self._probe_successes = 0 + self._transition(CircuitState.OPEN, TripReason.PROBE_FAILED.value) + return + if self._is_success(event): + self._probe_successes += 1 + if self._probe_successes >= self.thresholds.half_open_probe_successes: + self._transition(CircuitState.CLOSED, "recovery_probes_succeeded") + + # --------------------------------------------------------- pause/resume + + async def pause( + self, + *, + step_number: int | None = None, + working_dir: str | None = None, + memory_snapshot: dict[str, Any] | None = None, + ) -> str | None: + """Checkpoint current state and move to PAUSED. + + Only valid from OPEN. Delegates checkpoint creation to the rollback + engine if one was supplied; otherwise the breaker still transitions to + PAUSED but without a durable snapshot (and records that fact). + + Returns the checkpoint id, or ``None`` if no rollback engine is wired. + """ + if self._state is not CircuitState.OPEN: + raise RuntimeError(f"pause() is only valid from OPEN, not {self._state.value}") + + step = step_number if step_number is not None else self._last_step_number + checkpoint_id: str | None = None + + if self._rollback_engine is not None: + cp = await self._rollback_engine.create_checkpoint( + session_id=self.session_id, + step_number=step, + working_dir=working_dir, + memory_snapshot=memory_snapshot, + label=f"circuit-breaker-pause-{self.session_id}", + ) + checkpoint_id = cp.checkpoint_id + self._last_checkpoint_id = checkpoint_id + else: + logger.warning( + "CircuitBreaker.pause() with no rollback engine: session %s paused " + "without a durable checkpoint", + self.session_id, + ) + + self._transition( + CircuitState.PAUSED, + "paused_with_checkpoint" if checkpoint_id else "paused_without_checkpoint", + checkpoint_id=checkpoint_id, + ) + return checkpoint_id + + async def resume(self, *, restore: bool = False) -> CircuitState: + """Begin probing recovery by moving to HALF_OPEN. + + Valid from PAUSED or OPEN. When ``restore`` is True and a checkpoint was + saved by :meth:`pause`, the rollback engine restores that checkpoint + before probing begins. + """ + if self._state not in (CircuitState.PAUSED, CircuitState.OPEN): + raise RuntimeError(f"resume() is only valid from PAUSED/OPEN, not {self._state.value}") + + if restore and self._last_checkpoint_id and self._rollback_engine is not None: + await self._rollback_engine.rollback(self._last_checkpoint_id) + + self._probe_successes = 0 + self._transition( + CircuitState.HALF_OPEN, + "resume_probing", + checkpoint_id=self._last_checkpoint_id if restore else None, + ) + return self._state + + def trip(self, reason: str = TripReason.MANUAL.value) -> CircuitState: + """Manually force the breaker OPEN (e.g. operator kill-switch).""" + if self._state is CircuitState.CLOSED or self._state is CircuitState.HALF_OPEN: + self._transition(CircuitState.OPEN, reason) + return self._state + + def reset(self) -> CircuitState: + """Manually force the breaker back to CLOSED and clear the window.""" + self._transition(CircuitState.CLOSED, "manual_reset") + return self._state + + # --------------------------------------------------------- transitions + + def _transition( + self, + to_state: CircuitState, + reason: str, + *, + checkpoint_id: str | None = None, + ) -> None: + from_state = self._state + if to_state not in _ALLOWED[from_state]: + raise RuntimeError( + f"illegal circuit transition {from_state.value} -> {to_state.value}" + ) + + self._state = to_state + if to_state is CircuitState.CLOSED: + self._reset_window() + + record = TransitionRecord( + when=datetime.now(UTC), + from_state=from_state, + to_state=to_state, + reason=reason, + session_id=self.session_id, + checkpoint_id=checkpoint_id, + metadata={ + "total_tokens": self._total_tokens, + "total_errors": self._total_errors, + "hallucinations": self._hallucinations, + }, + ) + self._history.append(record) + self._record_for_eu_ai_act(record) + + logger.info( + "Circuit %s: %s -> %s (%s)", + self.session_id, + from_state.value, + to_state.value, + reason, + ) + + def _reset_window(self) -> None: + self._total_tokens = 0 + self._consecutive_errors = 0 + self._total_errors = 0 + self._hallucinations = 0 + self._events_observed = 0 + self._probe_successes = 0 + + def _record_for_eu_ai_act(self, record: TransitionRecord) -> None: + """Log the transition for EU AI Act Article 12 record-keeping. + + Best-effort: a governance package is optional, and a logging failure + must never break the breaker's control flow. + """ + if self._eu_ai_act is None: + return + try: + from agentwatch.governance.eu_ai_act import DecisionLogEntry + + payload = record.to_dict() + entry = DecisionLogEntry( + when=record.when, + decision_id=f"cb-{uuid.uuid4().hex[:12]}", + inputs_hash=self._hash(record.from_state.value + record.reason), + outputs_hash=self._hash(record.to_state.value), + confidence=1.0, + safety_checks_passed=record.to_state is not CircuitState.OPEN, + human_oversight_required=record.to_state + in (CircuitState.OPEN, CircuitState.PAUSED), + explanation=( + f"Circuit breaker transitioned {record.from_state.value} -> " + f"{record.to_state.value} for session {record.session_id} " + f"(reason: {record.reason})" + ), + ) + self._eu_ai_act.log_decision(entry) + logger.debug("Recorded circuit transition to EU AI Act log: %s", payload["reason"]) + except Exception as exc: # noqa: BLE001 - logging must not break control flow + logger.warning("Failed to record circuit transition for EU AI Act: %s", exc) + + @staticmethod + def _hash(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() diff --git a/agentwatch/circuit_breaker/thresholds.py b/agentwatch/circuit_breaker/thresholds.py new file mode 100644 index 00000000..447ad44e --- /dev/null +++ b/agentwatch/circuit_breaker/thresholds.py @@ -0,0 +1,53 @@ +"""Configurable thresholds that govern when the circuit breaker trips.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class CircuitThresholds: + """Thresholds that decide when the breaker trips from CLOSED to OPEN. + + All thresholds are evaluated over the window of events observed since the + breaker last entered the CLOSED state (a successful recovery resets the + counters). A threshold of ``0`` or ``None`` disables that particular check. + + Attributes: + max_total_tokens: Trip once cumulative tokens across observed events + exceed this value. Guards runaway generation cost. + max_consecutive_errors: Trip once this many error/failure events occur + in a row (a success resets the run). + max_total_errors: Trip once this many error/failure events occur in + total within the window, even if interspersed with successes. + max_hallucinations: Trip once this many low-confidence / anomaly-flagged + events are observed within the window. + min_confidence: Any single event whose confidence ``overall_score`` is + strictly below this value counts as a hallucination signal. Set to + ``0.0`` to disable per-event confidence checking. + half_open_probe_successes: Number of consecutive successful probe events + required in HALF_OPEN before the breaker closes again. + """ + + max_total_tokens: int = 100_000 + max_consecutive_errors: int = 3 + max_total_errors: int = 5 + max_hallucinations: int = 3 + min_confidence: float = 0.35 + half_open_probe_successes: int = 2 + + def __post_init__(self) -> None: + # Fail loudly on nonsensical configuration rather than silently + # producing a breaker that can never trip or never recover. + if self.max_total_tokens < 0: + raise ValueError("max_total_tokens must be >= 0") + if self.max_consecutive_errors < 0: + raise ValueError("max_consecutive_errors must be >= 0") + if self.max_total_errors < 0: + raise ValueError("max_total_errors must be >= 0") + if self.max_hallucinations < 0: + raise ValueError("max_hallucinations must be >= 0") + if not (0.0 <= self.min_confidence <= 1.0): + raise ValueError("min_confidence must be within [0.0, 1.0]") + if self.half_open_probe_successes < 1: + raise ValueError("half_open_probe_successes must be >= 1") diff --git a/tests/test_circuit_breaker.py b/tests/test_circuit_breaker.py new file mode 100644 index 00000000..eb1503b9 --- /dev/null +++ b/tests/test_circuit_breaker.py @@ -0,0 +1,464 @@ +"""Tests for the active circuit breaker (Issue #483).""" + +from __future__ import annotations + +import pytest + +from agentwatch.circuit_breaker import ( + CircuitBreaker, + CircuitState, + CircuitThresholds, + TripReason, +) +from agentwatch.core.schema import ( + AgentEvent, + ConfidenceData, + EventType, + ExecutionStatus, + TokenUsage, +) + +# --------------------------------------------------------------------------- helpers + + +def make_event( + *, + session_id: str = "sess-1", + event_type: EventType = EventType.TOOL_RESULT, + status: ExecutionStatus = ExecutionStatus.SUCCESS, + total_tokens: int = 0, + confidence: float | None = None, + anomaly_flags: list[str] | None = None, + step_number: int = 0, +) -> AgentEvent: + conf = None + if confidence is not None or anomaly_flags is not None: + conf = ConfidenceData( + overall_score=1.0 if confidence is None else confidence, + anomaly_flags=anomaly_flags or [], + ) + token_usage = TokenUsage(total_tokens=total_tokens) if total_tokens else None + return AgentEvent( + session_id=session_id, + agent_id="agent-1", + event_type=event_type, + status=status, + step_number=step_number, + confidence=conf, + token_usage=token_usage, + ) + + +def error_event(**kwargs) -> AgentEvent: + kwargs.setdefault("event_type", EventType.TOOL_ERROR) + kwargs.setdefault("status", ExecutionStatus.FAILURE) + return make_event(**kwargs) + + +def success_event(**kwargs) -> AgentEvent: + kwargs.setdefault("event_type", EventType.TOOL_RESULT) + kwargs.setdefault("status", ExecutionStatus.SUCCESS) + return make_event(**kwargs) + + +class FakeCheckpoint: + def __init__(self, checkpoint_id: str) -> None: + self.checkpoint_id = checkpoint_id + + +class FakeRollbackEngine: + """Records calls so tests can assert pause/resume delegate correctly.""" + + def __init__(self) -> None: + self.created: list[dict] = [] + self.rolled_back: list[str] = [] + self._counter = 0 + + async def create_checkpoint(self, **kwargs) -> FakeCheckpoint: + self._counter += 1 + self.created.append(kwargs) + return FakeCheckpoint(f"ckpt-{self._counter}") + + async def rollback(self, checkpoint_id: str) -> None: + self.rolled_back.append(checkpoint_id) + + +class FakeEUAIAct: + def __init__(self) -> None: + self.entries: list = [] + + def log_decision(self, entry) -> None: + self.entries.append(entry) + + +# --------------------------------------------------------------------------- config + + +class TestThresholdsValidation: + def test_defaults_are_valid(self): + t = CircuitThresholds() + assert t.max_total_tokens > 0 + + @pytest.mark.parametrize( + "kwargs", + [ + {"max_total_tokens": -1}, + {"max_consecutive_errors": -1}, + {"max_total_errors": -1}, + {"max_hallucinations": -1}, + {"min_confidence": 1.5}, + {"min_confidence": -0.1}, + {"half_open_probe_successes": 0}, + ], + ) + def test_invalid_config_raises(self, kwargs): + with pytest.raises(ValueError): + CircuitThresholds(**kwargs) + + +# --------------------------------------------------------------------------- tripping + + +class TestTripping: + def test_starts_closed(self): + cb = CircuitBreaker("sess-1") + assert cb.state is CircuitState.CLOSED + assert cb.is_tripped is False + + def test_trips_on_total_tokens(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_total_tokens=100)) + cb.observe(make_event(total_tokens=60)) + assert cb.state is CircuitState.CLOSED + cb.observe(make_event(total_tokens=60)) # cumulative 120 > 100 + assert cb.state is CircuitState.OPEN + assert cb.history[-1].reason == TripReason.TOKEN_BUDGET.value + + def test_trips_on_consecutive_errors(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_consecutive_errors=3, max_total_errors=0)) + cb.observe(error_event()) + cb.observe(error_event()) + assert cb.state is CircuitState.CLOSED + cb.observe(error_event()) + assert cb.state is CircuitState.OPEN + assert cb.history[-1].reason == TripReason.CONSECUTIVE_ERRORS.value + + def test_consecutive_error_run_reset_by_success(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_consecutive_errors=3, max_total_errors=0)) + cb.observe(error_event()) + cb.observe(error_event()) + cb.observe(success_event()) # resets the consecutive run + cb.observe(error_event()) + cb.observe(error_event()) + assert cb.state is CircuitState.CLOSED # never hit 3 in a row + + def test_trips_on_total_errors_even_with_successes(self): + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=3, max_consecutive_errors=0), + ) + cb.observe(error_event()) + cb.observe(success_event()) + cb.observe(error_event()) + cb.observe(success_event()) + assert cb.state is CircuitState.CLOSED + cb.observe(error_event()) # 3 total + assert cb.state is CircuitState.OPEN + assert cb.history[-1].reason == TripReason.TOTAL_ERRORS.value + + def test_trips_on_hallucination_anomaly_flags(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_hallucinations=2)) + cb.observe(make_event(anomaly_flags=["contradiction"])) + assert cb.state is CircuitState.CLOSED + cb.observe(make_event(anomaly_flags=["fabricated_citation"])) + assert cb.state is CircuitState.OPEN + assert cb.history[-1].reason == TripReason.HALLUCINATIONS.value + + def test_trips_on_low_confidence(self): + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_hallucinations=1, min_confidence=0.5), + ) + cb.observe(make_event(confidence=0.4)) # below 0.5 → hallucination + assert cb.state is CircuitState.OPEN + + def test_high_confidence_does_not_trip(self): + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_hallucinations=1, min_confidence=0.5), + ) + cb.observe(make_event(confidence=0.9)) + assert cb.state is CircuitState.CLOSED + + def test_min_confidence_zero_disables_confidence_check(self): + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_hallucinations=1, min_confidence=0.0), + ) + cb.observe(make_event(confidence=0.01)) # would trip if enabled + assert cb.state is CircuitState.CLOSED + + def test_disabled_threshold_never_trips(self): + cb = CircuitBreaker( + "sess-1", + CircuitThresholds( + max_total_tokens=0, + max_consecutive_errors=0, + max_total_errors=0, + max_hallucinations=0, + min_confidence=0.0, + ), + ) + for _ in range(50): + cb.observe(error_event(total_tokens=10_000)) + assert cb.state is CircuitState.CLOSED + + def test_events_for_other_sessions_ignored(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_total_errors=1)) + cb.observe(error_event(session_id="other-session")) + assert cb.state is CircuitState.CLOSED + assert cb.status()["counters"]["events_observed"] == 0 + + +# --------------------------------------------------------------------------- pause/resume + + +class TestPauseResume: + async def test_pause_creates_checkpoint(self): + engine = FakeRollbackEngine() + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=1), + rollback_engine=engine, + ) + cb.observe(error_event()) + assert cb.state is CircuitState.OPEN + ckpt = await cb.pause(step_number=7) + assert cb.state is CircuitState.PAUSED + assert ckpt == "ckpt-1" + assert engine.created[0]["session_id"] == "sess-1" + assert engine.created[0]["step_number"] == 7 + + async def test_pause_without_engine_still_pauses(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_total_errors=1)) + cb.observe(error_event()) + ckpt = await cb.pause() + assert cb.state is CircuitState.PAUSED + assert ckpt is None + + async def test_pause_only_valid_from_open(self): + cb = CircuitBreaker("sess-1") + with pytest.raises(RuntimeError): + await cb.pause() + + async def test_resume_from_paused_enters_half_open(self): + engine = FakeRollbackEngine() + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=1), + rollback_engine=engine, + ) + cb.observe(error_event()) + await cb.pause() + state = await cb.resume() + assert state is CircuitState.HALF_OPEN + + async def test_resume_with_restore_calls_rollback(self): + engine = FakeRollbackEngine() + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=1), + rollback_engine=engine, + ) + cb.observe(error_event()) + await cb.pause() + await cb.resume(restore=True) + assert engine.rolled_back == ["ckpt-1"] + + async def test_resume_from_open_without_pause(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_total_errors=1)) + cb.observe(error_event()) + assert cb.state is CircuitState.OPEN + state = await cb.resume() + assert state is CircuitState.HALF_OPEN + + async def test_resume_only_valid_from_paused_or_open(self): + cb = CircuitBreaker("sess-1") + with pytest.raises(RuntimeError): + await cb.resume() + + +# --------------------------------------------------------------------------- half-open recovery + + +class TestHalfOpenRecovery: + async def _open_and_probe(self, probe_successes: int = 2) -> CircuitBreaker: + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=1, half_open_probe_successes=probe_successes), + ) + cb.observe(error_event()) + await cb.resume() # OPEN -> HALF_OPEN + return cb + + async def test_probes_succeed_closes_circuit(self): + cb = await self._open_and_probe(probe_successes=2) + assert cb.state is CircuitState.HALF_OPEN + cb.observe(success_event()) + assert cb.state is CircuitState.HALF_OPEN # need 2 + cb.observe(success_event()) + assert cb.state is CircuitState.CLOSED + assert cb.history[-1].reason == "recovery_probes_succeeded" + + async def test_probe_failure_reopens(self): + cb = await self._open_and_probe(probe_successes=2) + cb.observe(success_event()) + cb.observe(error_event()) # a failed probe + assert cb.state is CircuitState.OPEN + assert cb.history[-1].reason == TripReason.PROBE_FAILED.value + + async def test_probe_success_run_resets_on_failure(self): + cb = await self._open_and_probe(probe_successes=3) + cb.observe(success_event()) + cb.observe(success_event()) + cb.observe(error_event()) # reopens, probe count resets + assert cb.state is CircuitState.OPEN + await cb.resume() + cb.observe(success_event()) + cb.observe(success_event()) + assert cb.state is CircuitState.HALF_OPEN # only 2, need 3 again + + async def test_counters_reset_after_recovery(self): + cb = await self._open_and_probe(probe_successes=1) + cb.observe(success_event()) # closes + assert cb.state is CircuitState.CLOSED + assert cb.status()["counters"]["total_errors"] == 0 + assert cb.status()["counters"]["total_tokens"] == 0 + + +# --------------------------------------------------------------------------- manual controls + + +class TestManualControls: + def test_manual_trip_from_closed(self): + cb = CircuitBreaker("sess-1") + cb.trip() + assert cb.state is CircuitState.OPEN + assert cb.history[-1].reason == TripReason.MANUAL.value + + def test_manual_reset_from_open(self): + cb = CircuitBreaker("sess-1") + cb.trip() + cb.reset() + assert cb.state is CircuitState.CLOSED + + def test_trip_is_noop_when_already_open(self): + cb = CircuitBreaker("sess-1") + cb.trip() + before = len(cb.history) + cb.trip() # no double transition + assert len(cb.history) == before + + +# --------------------------------------------------------------------------- transitions / audit + + +class TestTransitionsAndAudit: + def test_illegal_transition_raises(self): + cb = CircuitBreaker("sess-1") + # CLOSED -> PAUSED is not allowed + with pytest.raises(RuntimeError): + cb._transition(CircuitState.PAUSED, "illegal") + + def test_history_records_transitions(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_total_errors=1)) + cb.observe(error_event()) + cb.reset() + assert len(cb.history) == 2 + assert cb.history[0].to_state is CircuitState.OPEN + assert cb.history[1].to_state is CircuitState.CLOSED + + def test_transition_record_to_dict_roundtrips(self): + cb = CircuitBreaker("sess-1", CircuitThresholds(max_total_errors=1)) + cb.observe(error_event()) + d = cb.history[-1].to_dict() + assert d["from_state"] == "closed" + assert d["to_state"] == "open" + assert d["session_id"] == "sess-1" + + def test_status_snapshot_shape(self): + cb = CircuitBreaker("sess-1") + s = cb.status() + assert s["state"] == "closed" + assert "counters" in s + assert set(s["counters"]) == { + "total_tokens", + "consecutive_errors", + "total_errors", + "hallucinations", + "events_observed", + } + + def test_eu_ai_act_logging_on_transition(self): + eu = FakeEUAIAct() + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=1), + eu_ai_act_package=eu, + ) + cb.observe(error_event()) + assert len(eu.entries) == 1 + entry = eu.entries[0] + assert entry.human_oversight_required is True # OPEN requires oversight + assert "closed -> open" in entry.explanation + + def test_eu_ai_act_logging_failure_does_not_break(self): + class BrokenEU: + def log_decision(self, entry): + raise RuntimeError("boom") + + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=1), + eu_ai_act_package=BrokenEU(), + ) + # Should trip fine despite the logging sink raising. + cb.observe(error_event()) + assert cb.state is CircuitState.OPEN + + +# --------------------------------------------------------------------------- full lifecycle + + +class TestFullLifecycle: + async def test_closed_open_paused_halfopen_closed(self): + engine = FakeRollbackEngine() + eu = FakeEUAIAct() + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=1, half_open_probe_successes=1), + rollback_engine=engine, + eu_ai_act_package=eu, + ) + # CLOSED -> OPEN + cb.observe(error_event()) + assert cb.state is CircuitState.OPEN + # OPEN -> PAUSED + await cb.pause(step_number=3) + assert cb.state is CircuitState.PAUSED + # PAUSED -> HALF_OPEN (with restore) + await cb.resume(restore=True) + assert cb.state is CircuitState.HALF_OPEN + assert engine.rolled_back == ["ckpt-1"] + # HALF_OPEN -> CLOSED + cb.observe(success_event()) + assert cb.state is CircuitState.CLOSED + # Full transition trail recorded for audit + trail = [(r.from_state.value, r.to_state.value) for r in cb.history] + assert trail == [ + ("closed", "open"), + ("open", "paused"), + ("paused", "half_open"), + ("half_open", "closed"), + ] + # Every transition logged for EU AI Act Article 12 + assert len(eu.entries) == 4 From ffb621b504e6d03d437dd12099e7a58140669958 Mon Sep 17 00:00:00 2001 From: SakethSumanBathini Date: Fri, 10 Jul 2026 22:41:51 +0530 Subject: [PATCH 2/2] fix: preserve pre-reset counters in CLOSED transition audit record; track step_number in observe(); format (CodeRabbit #588) --- agentwatch/circuit_breaker/breaker.py | 20 ++++++++++-------- tests/test_circuit_breaker.py | 29 +++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 10 deletions(-) diff --git a/agentwatch/circuit_breaker/breaker.py b/agentwatch/circuit_breaker/breaker.py index a5ed52e5..4eedccf4 100644 --- a/agentwatch/circuit_breaker/breaker.py +++ b/agentwatch/circuit_breaker/breaker.py @@ -173,6 +173,7 @@ def observe(self, event: AgentEvent) -> CircuitState: return self._state self._events_observed += 1 + self._last_step_number = event.step_number self._accumulate(event) if self._state is CircuitState.CLOSED: @@ -338,9 +339,16 @@ def _transition( ) -> None: from_state = self._state if to_state not in _ALLOWED[from_state]: - raise RuntimeError( - f"illegal circuit transition {from_state.value} -> {to_state.value}" - ) + raise RuntimeError(f"illegal circuit transition {from_state.value} -> {to_state.value}") + + # Snapshot the counters *before* a CLOSED transition resets the window, so + # the audit record reflects the state at the moment of transition (e.g. a + # recovery record shows the error count that triggered the trip, not 0). + transition_metadata = { + "total_tokens": self._total_tokens, + "total_errors": self._total_errors, + "hallucinations": self._hallucinations, + } self._state = to_state if to_state is CircuitState.CLOSED: @@ -353,11 +361,7 @@ def _transition( reason=reason, session_id=self.session_id, checkpoint_id=checkpoint_id, - metadata={ - "total_tokens": self._total_tokens, - "total_errors": self._total_errors, - "hallucinations": self._hallucinations, - }, + metadata=transition_metadata, ) self._history.append(record) self._record_for_eu_ai_act(record) diff --git a/tests/test_circuit_breaker.py b/tests/test_circuit_breaker.py index eb1503b9..810131e9 100644 --- a/tests/test_circuit_breaker.py +++ b/tests/test_circuit_breaker.py @@ -134,7 +134,9 @@ def test_trips_on_total_tokens(self): assert cb.history[-1].reason == TripReason.TOKEN_BUDGET.value def test_trips_on_consecutive_errors(self): - cb = CircuitBreaker("sess-1", CircuitThresholds(max_consecutive_errors=3, max_total_errors=0)) + cb = CircuitBreaker( + "sess-1", CircuitThresholds(max_consecutive_errors=3, max_total_errors=0) + ) cb.observe(error_event()) cb.observe(error_event()) assert cb.state is CircuitState.CLOSED @@ -143,7 +145,9 @@ def test_trips_on_consecutive_errors(self): assert cb.history[-1].reason == TripReason.CONSECUTIVE_ERRORS.value def test_consecutive_error_run_reset_by_success(self): - cb = CircuitBreaker("sess-1", CircuitThresholds(max_consecutive_errors=3, max_total_errors=0)) + cb = CircuitBreaker( + "sess-1", CircuitThresholds(max_consecutive_errors=3, max_total_errors=0) + ) cb.observe(error_event()) cb.observe(error_event()) cb.observe(success_event()) # resets the consecutive run @@ -334,6 +338,27 @@ async def test_counters_reset_after_recovery(self): assert cb.status()["counters"]["total_errors"] == 0 assert cb.status()["counters"]["total_tokens"] == 0 + async def test_closed_transition_record_preserves_pre_reset_counters(self): + # The audit record for a CLOSED transition must reflect the counters at + # the moment of transition, not the post-reset zeros. This guards the + # EU AI Act Article 12 trail: a recovery record should still show the + # error count that triggered the original trip. + cb = CircuitBreaker( + "sess-1", + CircuitThresholds(max_total_errors=3, half_open_probe_successes=1), + ) + cb.observe(error_event()) + cb.observe(error_event()) + cb.observe(error_event()) # trips OPEN with 3 total errors + await cb.resume() + cb.observe(success_event()) # HALF_OPEN -> CLOSED + assert cb.state is CircuitState.CLOSED + closed_record = cb.history[-1] + assert closed_record.to_state is CircuitState.CLOSED + assert closed_record.metadata["total_errors"] == 3 # not 0 + # Live counters are still reset for the next window. + assert cb.status()["counters"]["total_errors"] == 0 + # --------------------------------------------------------------------------- manual controls