feat: 10 infrastructure improvements from zeroclaw (#126)

* refactor: break up agent_loop.rs into four focused modules

Split the monolithic 2835-line agent_loop.rs into:
- agent_loop.rs (722L): Agent struct, event loop, message dispatch
- dispatcher.rs (635L): Agentic tool loop, tool execution, auth detection
- commands.rs (484L): System commands, job handlers, heartbeat, summarize
- thread_ops.rs (1059L): Thread lifecycle, approval, undo/redo, persistence

Each module gets its own impl Agent block. Agent fields changed to
pub(super) so sibling modules in the agent package can access them.
All 16 existing tests pass in their new locations.

Inspired by ZeroClaw's agent module split (agent.rs, loop_.rs,
dispatcher.rs, prompt.rs, memory_loader.rs).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add cost caps and guardrails for autonomous agent spending

Daily budget (MAX_COST_PER_DAY_CENTS) and hourly action rate
(MAX_ACTIONS_PER_HOUR) limits prevent runaway agents from burning
through API credits, especially in daemon/heartbeat modes.

- CostGuard with pre-flight check and post-call recording
- Sliding window for hourly rate, midnight-UTC daily reset
- 80% threshold warning, atomic fast-path for exceeded budget
- Wired into dispatcher loop (check before LLM call, record after)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add circuit breaker on LLM providers

Wraps LlmProvider with a Closed/Open/HalfOpen state machine that
trips after consecutive transient failures, preventing request storms
against a degraded backend. Automatically probes for recovery.

- CircuitBreakerProvider implements LlmProvider (drop-in wrapper)
- Transient error classification (server, rate-limit, network, auth infra)
- Client errors (wrong model, context overflow) don't trip the breaker
- Configurable via CIRCUIT_BREAKER_THRESHOLD and CIRCUIT_BREAKER_RECOVERY_SECS
- Composes with existing FailoverProvider (circuit breaker wraps failover)
- 12 tests covering full state machine and error classification

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add tunnel abstraction for remote access

Trait-based tunnel system with lifecycle management (start/stop/health)
for exposing the agent to the internet through external tunnel binaries.

Five providers:
- Cloudflare Tunnel (cloudflared, Zero Trust token auth)
- Tailscale (serve for tailnet, funnel for public)
- ngrok (with optional custom domain)
- Custom (arbitrary command with {host}/{port} placeholders)
- None (local-only, no external exposure)

Config via TUNNEL_PROVIDER + provider-specific env vars. Extends
existing TunnelConfig with optional managed provider alongside the
static TUNNEL_URL path. Factory, shared process management, and
37 tests covering all providers and edge cases.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add OS service management (launchd/systemd)

Adds `ironclaw service {install,start,stop,status,uninstall}` for
running the agent as a background daemon. macOS uses launchd plists
under ~/Library/LaunchAgents, Linux uses systemd user units.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add observability trait system with noop, log, and multi backends

Introduces an Observer trait for recording agent lifecycle events and
metrics, with pluggable backends. The noop backend compiles to zero
overhead, log backend uses tracing, and multi fans out to multiple
observers. Configured via OBSERVABILITY_BACKEND env var.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add in-memory LLM response cache with TTL and LRU eviction

CachedProvider wraps any LlmProvider and caches complete() responses
keyed by SHA-256(model + messages). Tool-calling requests are never
cached since they trigger side effects. Configurable via
RESPONSE_CACHE_ENABLED, RESPONSE_CACHE_TTL_SECS, and
RESPONSE_CACHE_MAX_ENTRIES env vars.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add memory hygiene with cadence-gated daily log cleanup

Adds workspace::hygiene module that automatically deletes daily log
documents older than a configurable retention period (default 30 days).
Runs on a 12-hour cadence tracked via a local state file to avoid
redundant passes. Best-effort design: failures are logged, never fatal.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add doctor diagnostics command for active health probing

Probes external dependencies (Docker, cloudflared, ngrok, tailscale),
validates NEAR AI session, checks database connectivity, and verifies
workspace directory. Complements the passive `status` command.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: add structured TOML config file support

Adds ~/.ironclaw/config.toml as a configuration layer between env vars
and database settings. Priority: env var > TOML file > DB > defaults.

- `ironclaw config init` generates a commented config.toml from current settings
- `ironclaw --config path/to/config.toml` loads a custom config file
- Settings.merge_from() only overlays non-default values from the TOML file
- `ironclaw config path` now shows TOML file status

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address codex review findings

- apply_toml_overlay now returns Result and errors on explicit missing
  or invalid config paths (was log-only, violating the documented
  contract that explicit paths are fatal)
- custom tunnel url_pattern is now used to filter extracted URLs, not
  just as a gate for scanning stdout
- systemd ExecStart path is now quoted to handle spaces in paths

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review feedback

- Cache key now includes max_tokens, temperature, and stop_sequences
  so different request parameters produce distinct keys
- to_cents() uses .trunc() + parse::<u64> instead of f64 intermediary,
  avoiding precision loss for large values
- Tailscale public URL no longer includes local port (serve/funnel
  expose on standard HTTPS port 443)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat: wire up tunnel lifecycle and fix audit findings

Connect the tunnel module to the rest of the application so that
setting TUNNEL_PROVIDER actually starts a managed tunnel at boot and
stops it on shutdown. Previously create_tunnel() was never called
outside tests.

Changes:
- Expand TunnelSettings with provider credential fields (settings.rs)
- TunnelConfig::resolve() falls back to DB settings when env vars unset
- Start tunnel at boot, stop on shutdown, show URL in boot screen
- Setup wizard collects provider-specific credentials (ngrok, cloudflare,
  tailscale, custom, static URL)
- Fix public_url() returning None under lock contention (SharedUrl)
- Fix local_host parameter ignored by cloudflare/ngrok/tailscale
- Fix tailscale silent fallback to "localhost" on bad JSON
- Fix ngrok globally mutating config via add-authtoken (use env var)
- Add 10s timeout to tailscale status --json

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review comments

- Document split_whitespace limitation in CustomTunnel doc comment
- Remove unnecessary quotes from systemd ExecStart directive

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address PR review feedback (round 3)

- doctor: missing libSQL DB on fresh install is Pass, not Fail
- service: quote ExecStart path for systemd space handling

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: correct cost guard doc comment (LLM calls, not LLM/tool)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-17 19:35:56 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 436dda0f2f
commit a158eee1b0
37 changed files with 7272 additions and 2177 deletions
+674
View File
@@ -0,0 +1,674 @@
//! Circuit breaker for LLM providers.
//!
//! Wraps any `LlmProvider` with a state machine that trips open after
//! consecutive transient failures, preventing request storms against a
//! degraded backend. Automatically probes for recovery via half-open state.
//!
//! ```text
//! Closed ──(failures >= threshold)──► Open
//! ▲ │
//! │ (recovery timeout)
//! │ ▼
//! └──(probe succeeds)──── HalfOpen ──(probe fails)──► Open
//! ```
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use rust_decimal::Decimal;
use tokio::sync::Mutex;
use crate::error::LlmError;
use crate::llm::provider::{
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
ToolCompletionResponse,
};
/// Configuration for the circuit breaker.
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
/// Consecutive transient failures before the circuit opens.
pub failure_threshold: u32,
/// How long the circuit stays open before allowing a probe.
pub recovery_timeout: Duration,
/// Successful probes needed in half-open to close the circuit.
pub half_open_successes_needed: u32,
}
impl Default for CircuitBreakerConfig {
fn default() -> Self {
Self {
failure_threshold: 5,
recovery_timeout: Duration::from_secs(30),
half_open_successes_needed: 2,
}
}
}
/// Circuit breaker states.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
/// Normal operation; tracking consecutive failures.
Closed,
/// Rejecting all calls; waiting for recovery timeout to elapse.
Open,
/// Allowing probe calls to test whether the backend recovered.
HalfOpen,
}
/// Internal mutable state.
struct BreakerState {
state: CircuitState,
consecutive_failures: u32,
opened_at: Option<Instant>,
half_open_successes: u32,
}
impl BreakerState {
fn new() -> Self {
Self {
state: CircuitState::Closed,
consecutive_failures: 0,
opened_at: None,
half_open_successes: 0,
}
}
}
/// Wraps an `LlmProvider` with circuit breaker protection.
///
/// Tracks consecutive transient failures. After `failure_threshold` failures
/// the circuit opens and all requests are rejected for `recovery_timeout`.
/// After that timeout a probe call is allowed through (half-open); if it
/// succeeds the circuit closes, otherwise it reopens.
pub struct CircuitBreakerProvider {
inner: Arc<dyn LlmProvider>,
state: Mutex<BreakerState>,
config: CircuitBreakerConfig,
}
impl CircuitBreakerProvider {
pub fn new(inner: Arc<dyn LlmProvider>, config: CircuitBreakerConfig) -> Self {
Self {
inner,
state: Mutex::new(BreakerState::new()),
config,
}
}
/// Current circuit state (for observability / health checks).
pub async fn circuit_state(&self) -> CircuitState {
self.state.lock().await.state
}
/// Number of consecutive failures recorded so far.
pub async fn consecutive_failures(&self) -> u32 {
self.state.lock().await.consecutive_failures
}
/// Pre-flight: is a call allowed right now?
async fn check_allowed(&self) -> Result<(), LlmError> {
let mut state = self.state.lock().await;
match state.state {
CircuitState::Closed | CircuitState::HalfOpen => Ok(()),
CircuitState::Open => {
if let Some(opened_at) = state.opened_at {
if opened_at.elapsed() >= self.config.recovery_timeout {
state.state = CircuitState::HalfOpen;
state.half_open_successes = 0;
tracing::info!(
provider = self.inner.model_name(),
"Circuit breaker: Open -> HalfOpen, allowing probe"
);
Ok(())
} else {
let remaining = self.config.recovery_timeout - opened_at.elapsed();
Err(LlmError::RequestFailed {
provider: self.inner.model_name().to_string(),
reason: format!(
"Circuit breaker open ({} consecutive failures, \
recovery in {:.0}s)",
state.consecutive_failures,
remaining.as_secs_f64()
),
})
}
} else {
// opened_at should always be Some when Open; recover gracefully
state.state = CircuitState::Closed;
Ok(())
}
}
}
}
/// Record a successful call.
async fn record_success(&self) {
let mut state = self.state.lock().await;
match state.state {
CircuitState::Closed => {
state.consecutive_failures = 0;
}
CircuitState::HalfOpen => {
state.half_open_successes += 1;
if state.half_open_successes >= self.config.half_open_successes_needed {
state.state = CircuitState::Closed;
state.consecutive_failures = 0;
state.opened_at = None;
tracing::info!(
provider = self.inner.model_name(),
"Circuit breaker: HalfOpen -> Closed (recovered)"
);
}
}
CircuitState::Open => {
// Shouldn't get here (check_allowed blocks Open), but recover
state.state = CircuitState::Closed;
state.consecutive_failures = 0;
state.opened_at = None;
}
}
}
/// Record a failed call; only transient errors count toward the threshold.
async fn record_failure(&self, err: &LlmError) {
if !is_transient(err) {
return;
}
let mut state = self.state.lock().await;
match state.state {
CircuitState::Closed => {
state.consecutive_failures += 1;
if state.consecutive_failures >= self.config.failure_threshold {
state.state = CircuitState::Open;
state.opened_at = Some(Instant::now());
tracing::warn!(
provider = self.inner.model_name(),
failures = state.consecutive_failures,
"Circuit breaker: Closed -> Open"
);
}
}
CircuitState::HalfOpen => {
state.state = CircuitState::Open;
state.opened_at = Some(Instant::now());
state.half_open_successes = 0;
tracing::warn!(
provider = self.inner.model_name(),
"Circuit breaker: HalfOpen -> Open (probe failed)"
);
}
CircuitState::Open => {}
}
}
}
/// Returns `true` for errors that indicate the provider is degraded
/// (server errors, rate limits, network failures, auth infrastructure down).
///
/// Client errors (wrong model, bad credentials, context overflow) are NOT
/// transient: they are the caller's problem, not a sign of backend trouble.
fn is_transient(err: &LlmError) -> bool {
matches!(
err,
LlmError::RequestFailed { .. }
| LlmError::RateLimited { .. }
| LlmError::InvalidResponse { .. }
| LlmError::SessionExpired { .. }
| LlmError::SessionRenewalFailed { .. }
| LlmError::Http(_)
| LlmError::Json(_)
| LlmError::Io(_)
)
}
#[async_trait]
impl LlmProvider for CircuitBreakerProvider {
fn model_name(&self) -> &str {
self.inner.model_name()
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
self.inner.cost_per_token()
}
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
self.check_allowed().await?;
match self.inner.complete(request).await {
Ok(resp) => {
self.record_success().await;
Ok(resp)
}
Err(err) => {
self.record_failure(&err).await;
Err(err)
}
}
}
async fn complete_with_tools(
&self,
request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
self.check_allowed().await?;
match self.inner.complete_with_tools(request).await {
Ok(resp) => {
self.record_success().await;
Ok(resp)
}
Err(err) => {
self.record_failure(&err).await;
Err(err)
}
}
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
self.inner.list_models().await
}
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
self.inner.model_metadata().await
}
fn active_model_name(&self) -> String {
self.inner.active_model_name()
}
fn set_model(&self, model: &str) -> Result<(), LlmError> {
self.inner.set_model(model)
}
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
self.inner.seed_response_chain(thread_id, response_id)
}
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
self.inner.get_response_chain_id(thread_id)
}
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
self.inner.calculate_cost(input_tokens, output_tokens)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, Ordering};
use crate::llm::provider::{CompletionResponse, FinishReason, ToolCompletionResponse};
/// A test stub that either always succeeds or always fails with a
/// configurable error. The `should_fail` flag can be flipped at
/// runtime for half-open recovery tests.
struct StubProvider {
name: String,
should_fail: AtomicBool,
error_kind: StubError,
}
#[derive(Clone)]
enum StubError {
Transient,
NonTransient,
}
impl StubProvider {
fn always_ok(name: &str) -> Arc<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(false),
error_kind: StubError::Transient,
})
}
fn always_fail(name: &str) -> Arc<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(true),
error_kind: StubError::Transient,
})
}
fn always_fail_non_transient(name: &str) -> Arc<Self> {
Arc::new(Self {
name: name.to_string(),
should_fail: AtomicBool::new(true),
error_kind: StubError::NonTransient,
})
}
fn set_failing(&self, fail: bool) {
self.should_fail.store(fail, Ordering::Relaxed);
}
fn make_error(&self) -> LlmError {
match self.error_kind {
StubError::Transient => LlmError::RequestFailed {
provider: self.name.clone(),
reason: "server error".to_string(),
},
StubError::NonTransient => LlmError::ContextLengthExceeded {
used: 100_000,
limit: 50_000,
},
}
}
fn ok_response() -> CompletionResponse {
CompletionResponse {
content: "ok".to_string(),
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
}
}
fn ok_tool_response() -> ToolCompletionResponse {
ToolCompletionResponse {
content: Some("ok".to_string()),
tool_calls: vec![],
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
}
}
}
#[async_trait]
impl LlmProvider for StubProvider {
fn model_name(&self) -> &str {
&self.name
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
(Decimal::ZERO, Decimal::ZERO)
}
async fn complete(
&self,
_request: CompletionRequest,
) -> Result<CompletionResponse, LlmError> {
if self.should_fail.load(Ordering::Relaxed) {
Err(self.make_error())
} else {
Ok(Self::ok_response())
}
}
async fn complete_with_tools(
&self,
_request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
if self.should_fail.load(Ordering::Relaxed) {
Err(self.make_error())
} else {
Ok(Self::ok_tool_response())
}
}
}
fn make_request() -> CompletionRequest {
CompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")])
}
fn make_tool_request() -> ToolCompletionRequest {
ToolCompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")], vec![])
}
fn fast_config(threshold: u32) -> CircuitBreakerConfig {
CircuitBreakerConfig {
failure_threshold: threshold,
recovery_timeout: Duration::from_millis(50),
half_open_successes_needed: 1,
}
}
// -- State machine tests --
#[tokio::test]
async fn closed_allows_calls_and_resets_on_success() {
let stub = StubProvider::always_ok("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(3));
let resp = cb.complete(make_request()).await;
assert!(resp.is_ok());
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn failures_accumulate_then_trip_to_open() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(3));
// First 2 failures: still closed
for i in 0..2 {
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, i + 1);
}
// 3rd failure: trips to open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn open_rejects_immediately() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(
stub,
CircuitBreakerConfig {
failure_threshold: 1,
recovery_timeout: Duration::from_secs(60),
half_open_successes_needed: 1,
},
);
// Trip the breaker
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
// Next call should fail with circuit breaker message
let err = cb.complete(make_request()).await.unwrap_err();
match err {
LlmError::RequestFailed { reason, .. } => {
assert!(
reason.contains("Circuit breaker open"),
"Expected circuit breaker message, got: {}",
reason
);
}
other => panic!("Expected RequestFailed, got: {:?}", other),
}
}
#[tokio::test]
async fn recovery_timeout_transitions_to_half_open() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(1));
// Trip to open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
// Wait for recovery timeout
tokio::time::sleep(Duration::from_millis(60)).await;
// Next call should transition to half-open (and fail, since stub fails)
let _ = cb.complete(make_request()).await;
// Failed probe sends it back to Open
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn half_open_success_closes_circuit() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(1));
// Trip to open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
// Wait for recovery, then make the stub succeed
tokio::time::sleep(Duration::from_millis(60)).await;
stub.set_failing(false);
// Probe should succeed, closing the circuit
let resp = cb.complete(make_request()).await;
assert!(resp.is_ok());
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn half_open_failure_reopens_circuit() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(1));
// Trip to open
let _ = cb.complete(make_request()).await;
// Wait for recovery timeout
tokio::time::sleep(Duration::from_millis(60)).await;
// Probe fails (stub still failing)
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn non_transient_errors_do_not_trip_breaker() {
let stub = StubProvider::always_fail_non_transient("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(1));
// ContextLengthExceeded is not transient; breaker should stay closed
for _ in 0..5 {
let _ = cb.complete(make_request()).await;
}
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn success_resets_failure_count() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub.clone(), fast_config(3));
// Accumulate 2 failures
let _ = cb.complete(make_request()).await;
let _ = cb.complete(make_request()).await;
assert_eq!(cb.consecutive_failures().await, 2);
// One success resets the counter
stub.set_failing(false);
let resp = cb.complete(make_request()).await;
assert!(resp.is_ok());
assert_eq!(cb.consecutive_failures().await, 0);
}
#[tokio::test]
async fn complete_with_tools_uses_same_breaker_logic() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(stub, fast_config(2));
let _ = cb.complete_with_tools(make_tool_request()).await;
let _ = cb.complete_with_tools(make_tool_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Open);
}
#[tokio::test]
async fn multiple_half_open_successes_needed() {
let stub = StubProvider::always_fail("test");
let cb = CircuitBreakerProvider::new(
stub.clone(),
CircuitBreakerConfig {
failure_threshold: 1,
recovery_timeout: Duration::from_millis(50),
half_open_successes_needed: 3,
},
);
// Trip to open
let _ = cb.complete(make_request()).await;
// Wait and flip to succeed
tokio::time::sleep(Duration::from_millis(60)).await;
stub.set_failing(false);
// First probe: half-open, success but not enough yet
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen);
// Second probe: still half-open
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::HalfOpen);
// Third probe: closes
let _ = cb.complete(make_request()).await;
assert_eq!(cb.circuit_state().await, CircuitState::Closed);
}
// -- Error classification tests --
#[test]
fn transient_classification() {
// Transient
assert!(is_transient(&LlmError::RequestFailed {
provider: "p".into(),
reason: "err".into(),
}));
assert!(is_transient(&LlmError::RateLimited {
provider: "p".into(),
retry_after: None,
}));
assert!(is_transient(&LlmError::InvalidResponse {
provider: "p".into(),
reason: "bad".into(),
}));
assert!(is_transient(&LlmError::SessionExpired {
provider: "p".into(),
}));
assert!(is_transient(&LlmError::SessionRenewalFailed {
provider: "p".into(),
reason: "timeout".into(),
}));
assert!(is_transient(&LlmError::Io(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset"
))));
// NOT transient
assert!(!is_transient(&LlmError::AuthFailed {
provider: "p".into(),
}));
assert!(!is_transient(&LlmError::ContextLengthExceeded {
used: 100_000,
limit: 50_000,
}));
assert!(!is_transient(&LlmError::ModelNotAvailable {
provider: "p".into(),
model: "m".into(),
}));
}
// -- Passthrough delegation tests --
#[tokio::test]
async fn passthrough_methods_delegate_to_inner() {
let stub = StubProvider::always_ok("my-model");
let cb = CircuitBreakerProvider::new(stub, fast_config(3));
assert_eq!(cb.model_name(), "my-model");
assert_eq!(cb.active_model_name(), "my-model");
assert_eq!(cb.cost_per_token(), (Decimal::ZERO, Decimal::ZERO));
assert_eq!(cb.calculate_cost(100, 50), Decimal::ZERO);
}
}
+10 -1
View File
@@ -7,16 +7,19 @@
//! - **Ollama**: Local model inference
//! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API
mod costs;
pub mod circuit_breaker;
pub mod costs;
pub mod failover;
mod nearai;
mod nearai_chat;
mod provider;
mod reasoning;
pub mod response_cache;
mod retry;
mod rig_adapter;
pub mod session;
pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider};
pub use failover::{CooldownConfig, FailoverProvider};
pub use nearai::{ModelInfo, NearAiProvider};
pub use nearai_chat::NearAiChatProvider;
@@ -28,6 +31,7 @@ pub use reasoning::{
ActionPlan, Reasoning, ReasoningContext, RespondOutput, RespondResult, TokenUsage,
ToolSelection,
};
pub use response_cache::{CachedProvider, ResponseCacheConfig};
pub use rig_adapter::RigAdapter;
pub use session::{SessionConfig, SessionManager, create_session_manager};
@@ -235,6 +239,11 @@ mod tests {
api_key: None,
fallback_model: None,
max_retries: 3,
circuit_breaker_threshold: None,
circuit_breaker_recovery_secs: 30,
response_cache_enabled: false,
response_cache_ttl_secs: 3600,
response_cache_max_entries: 1000,
failover_cooldown_secs: 300,
failover_cooldown_threshold: 3,
}
+526
View File
@@ -0,0 +1,526 @@
//! In-memory LLM response cache with TTL and LRU eviction.
//!
//! Wraps any [`LlmProvider`] and caches [`complete()`] responses keyed
//! by a SHA-256 hash of the messages and model name. Tool-calling
//! requests are never cached since they can trigger side effects.
//!
//! ```text
//! ┌──────────────────────────────────────────────────┐
//! │ CachedProvider │
//! │ complete() ──► cache lookup ──► hit? return │
//! │ miss? call inner │
//! │ store response │
//! │ │
//! │ complete_with_tools() ──► always call inner │
//! └──────────────────────────────────────────────────┘
//! ```
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use rust_decimal::Decimal;
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;
use crate::error::LlmError;
use crate::llm::provider::{
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
ToolCompletionResponse,
};
/// Configuration for the response cache.
#[derive(Debug, Clone)]
pub struct ResponseCacheConfig {
/// Time-to-live for cache entries.
pub ttl: Duration,
/// Maximum number of cached entries before LRU eviction.
pub max_entries: usize,
}
impl Default for ResponseCacheConfig {
fn default() -> Self {
Self {
ttl: Duration::from_secs(3600), // 1 hour
max_entries: 1000,
}
}
}
struct CacheEntry {
response: CompletionResponse,
created_at: Instant,
last_accessed: Instant,
hit_count: u64,
}
/// LLM provider wrapper that caches `complete()` responses.
///
/// Tool completion requests are always forwarded without caching since
/// tool calls can have side effects that should not be replayed.
pub struct CachedProvider {
inner: Arc<dyn LlmProvider>,
cache: Mutex<HashMap<String, CacheEntry>>,
config: ResponseCacheConfig,
}
impl CachedProvider {
/// Wrap an existing provider with response caching.
pub fn new(inner: Arc<dyn LlmProvider>, config: ResponseCacheConfig) -> Self {
Self {
inner,
cache: Mutex::new(HashMap::new()),
config,
}
}
/// Number of entries currently in the cache.
pub async fn len(&self) -> usize {
self.cache.lock().await.len()
}
/// Whether the cache is empty.
pub async fn is_empty(&self) -> bool {
self.cache.lock().await.is_empty()
}
/// Total cache hits across all entries.
pub async fn total_hits(&self) -> u64 {
self.cache.lock().await.values().map(|e| e.hit_count).sum()
}
/// Clear all cached entries.
pub async fn clear(&self) {
self.cache.lock().await.clear();
}
}
/// Build a deterministic cache key from a completion request.
///
/// Hashes the model name, messages, and response-affecting parameters
/// (max_tokens, temperature, stop_sequences) via SHA-256. Two requests
/// with identical content and parameters produce the same key.
fn cache_key(model: &str, request: &CompletionRequest) -> String {
let mut hasher = Sha256::new();
hasher.update(model.as_bytes());
hasher.update(b"|");
// Messages are Serialize, so we can deterministically hash them.
// serde_json produces stable output for the same input structure.
if let Ok(json) = serde_json::to_string(&request.messages) {
hasher.update(json.as_bytes());
}
// Include response-affecting parameters so different temperatures,
// max_tokens, or stop sequences produce distinct cache keys.
hasher.update(b"|");
if let Some(max_tokens) = request.max_tokens {
hasher.update(max_tokens.to_le_bytes());
}
hasher.update(b"|");
if let Some(temp) = request.temperature {
hasher.update(temp.to_le_bytes());
}
hasher.update(b"|");
if let Some(ref stops) = request.stop_sequences {
for s in stops {
hasher.update(s.as_bytes());
hasher.update(b"\x00");
}
}
format!("{:x}", hasher.finalize())
}
#[async_trait]
impl LlmProvider for CachedProvider {
fn model_name(&self) -> &str {
self.inner.model_name()
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
self.inner.cost_per_token()
}
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
let key = cache_key(self.inner.model_name(), &request);
let now = Instant::now();
// Check cache
{
let mut guard = self.cache.lock().await;
if let Some(entry) = guard.get_mut(&key) {
if now.duration_since(entry.created_at) < self.config.ttl {
entry.last_accessed = now;
entry.hit_count += 1;
tracing::debug!(hits = entry.hit_count, "response cache hit");
return Ok(entry.response.clone());
}
// Expired, remove it
guard.remove(&key);
}
}
// Cache miss, call the real provider
let response = self.inner.complete(request).await?;
// Store in cache
{
let mut guard = self.cache.lock().await;
// Evict expired entries
guard.retain(|_, entry| now.duration_since(entry.created_at) < self.config.ttl);
// LRU eviction if over capacity
while guard.len() >= self.config.max_entries {
let oldest_key = guard
.iter()
.min_by_key(|(_, entry)| entry.last_accessed)
.map(|(k, _)| k.clone());
if let Some(k) = oldest_key {
guard.remove(&k);
} else {
break;
}
}
guard.insert(
key,
CacheEntry {
response: response.clone(),
created_at: now,
last_accessed: now,
hit_count: 0,
},
);
}
Ok(response)
}
async fn complete_with_tools(
&self,
request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
// Never cache tool calls; they can trigger side effects.
self.inner.complete_with_tools(request).await
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
self.inner.list_models().await
}
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
self.inner.model_metadata().await
}
fn active_model_name(&self) -> String {
self.inner.active_model_name()
}
fn set_model(&self, model: &str) -> Result<(), LlmError> {
self.inner.set_model(model)
}
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
self.inner.seed_response_chain(thread_id, response_id);
}
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
self.inner.get_response_chain_id(thread_id)
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use crate::llm::provider::{ChatMessage, FinishReason};
use crate::llm::response_cache::*;
/// Controllable stub provider for testing cache behavior.
struct StubProvider {
call_count: AtomicU32,
should_fail: AtomicBool,
}
impl StubProvider {
fn new() -> Self {
Self {
call_count: AtomicU32::new(0),
should_fail: AtomicBool::new(false),
}
}
fn calls(&self) -> u32 {
self.call_count.load(Ordering::Relaxed)
}
}
#[async_trait]
impl LlmProvider for StubProvider {
fn model_name(&self) -> &str {
"stub-model"
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
(Decimal::ZERO, Decimal::ZERO)
}
async fn complete(
&self,
_request: CompletionRequest,
) -> Result<CompletionResponse, LlmError> {
self.call_count.fetch_add(1, Ordering::Relaxed);
if self.should_fail.load(Ordering::Relaxed) {
return Err(LlmError::RequestFailed {
provider: "stub".into(),
reason: "forced failure".into(),
});
}
Ok(CompletionResponse {
content: "cached response".into(),
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
})
}
async fn complete_with_tools(
&self,
_request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
self.call_count.fetch_add(1, Ordering::Relaxed);
Ok(ToolCompletionResponse {
content: Some("tool response".into()),
tool_calls: vec![],
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
response_id: None,
})
}
}
fn simple_request() -> CompletionRequest {
CompletionRequest {
messages: vec![ChatMessage::user("hello")],
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: Default::default(),
}
}
fn different_request() -> CompletionRequest {
CompletionRequest {
messages: vec![ChatMessage::user("goodbye")],
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: Default::default(),
}
}
#[test]
fn cache_key_is_deterministic() {
let req = simple_request();
let k1 = cache_key("model-a", &req);
let k2 = cache_key("model-a", &req);
assert_eq!(k1, k2);
assert_eq!(k1.len(), 64); // SHA-256 hex
}
#[test]
fn cache_key_varies_by_model() {
let req = simple_request();
let k1 = cache_key("model-a", &req);
let k2 = cache_key("model-b", &req);
assert_ne!(k1, k2);
}
#[test]
fn cache_key_varies_by_messages() {
let k1 = cache_key("model-a", &simple_request());
let k2 = cache_key("model-a", &different_request());
assert_ne!(k1, k2);
}
#[test]
fn cache_key_varies_by_temperature() {
let mut req_a = simple_request();
req_a.temperature = Some(0.0);
let mut req_b = simple_request();
req_b.temperature = Some(1.0);
assert_ne!(cache_key("m", &req_a), cache_key("m", &req_b));
}
#[test]
fn cache_key_varies_by_max_tokens() {
let mut req_a = simple_request();
req_a.max_tokens = Some(100);
let mut req_b = simple_request();
req_b.max_tokens = Some(500);
assert_ne!(cache_key("m", &req_a), cache_key("m", &req_b));
}
#[tokio::test]
async fn cache_hit_avoids_provider_call() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_secs(60),
max_entries: 100,
},
);
// First call: cache miss
let r1 = cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 1);
assert_eq!(r1.content, "cached response");
// Second call: cache hit
let r2 = cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 1); // still 1
assert_eq!(r2.content, "cached response");
assert_eq!(cached.total_hits().await, 1);
}
#[tokio::test]
async fn different_messages_get_different_entries() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
cached.complete(simple_request()).await.unwrap();
cached.complete(different_request()).await.unwrap();
assert_eq!(stub.calls(), 2);
assert_eq!(cached.len().await, 2);
}
#[tokio::test]
async fn expired_entries_are_evicted() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_millis(1),
max_entries: 100,
},
);
cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 1);
// Wait for TTL to expire
tokio::time::sleep(Duration::from_millis(10)).await;
// Should be a cache miss now
cached.complete(simple_request()).await.unwrap();
assert_eq!(stub.calls(), 2);
}
#[tokio::test]
async fn lru_eviction_removes_oldest() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_secs(60),
max_entries: 2,
},
);
// Fill cache with 2 entries
cached.complete(simple_request()).await.unwrap();
cached.complete(different_request()).await.unwrap();
assert_eq!(cached.len().await, 2);
// Add a third: should evict the oldest
let third = CompletionRequest {
messages: vec![ChatMessage::user("third")],
max_tokens: None,
temperature: None,
stop_sequences: None,
metadata: Default::default(),
};
cached.complete(third).await.unwrap();
assert_eq!(cached.len().await, 2);
assert_eq!(stub.calls(), 3);
}
#[tokio::test]
async fn tool_calls_are_never_cached() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
let req = ToolCompletionRequest {
messages: vec![ChatMessage::user("use tool")],
tools: vec![],
max_tokens: None,
temperature: None,
tool_choice: None,
metadata: Default::default(),
};
cached.complete_with_tools(req.clone()).await.unwrap();
cached.complete_with_tools(req).await.unwrap();
// Both should have called through
assert_eq!(stub.calls(), 2);
assert!(cached.is_empty().await);
}
#[tokio::test]
async fn provider_errors_are_not_cached() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(
stub.clone(),
ResponseCacheConfig {
ttl: Duration::from_secs(60),
max_entries: 100,
},
);
stub.should_fail.store(true, Ordering::Relaxed);
let result = cached.complete(simple_request()).await;
assert!(result.is_err());
assert!(cached.is_empty().await);
// After fixing the provider, should succeed and cache
stub.should_fail.store(false, Ordering::Relaxed);
cached.complete(simple_request()).await.unwrap();
assert_eq!(cached.len().await, 1);
}
#[tokio::test]
async fn clear_empties_cache() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
cached.complete(simple_request()).await.unwrap();
assert_eq!(cached.len().await, 1);
cached.clear().await;
assert!(cached.is_empty().await);
}
#[test]
fn default_config_is_reasonable() {
let cfg = ResponseCacheConfig::default();
assert_eq!(cfg.ttl, Duration::from_secs(3600));
assert_eq!(cfg.max_entries, 1000);
}
#[tokio::test]
async fn delegates_model_name() {
let stub = Arc::new(StubProvider::new());
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
assert_eq!(cached.model_name(), "stub-model");
}
}
+3 -3
View File
@@ -432,9 +432,9 @@ impl SessionManager {
.get_setting(&user_id, "nearai.session_token")
.await
.map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
reason: format!("DB query failed: {}", e),
})? {
provider: "nearai".to_string(),
reason: format!("DB query failed: {}", e),
})? {
value
} else {
tracing::warn!(