mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 15:40:18 +00:00
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:
co-authored by
Claude Opus 4.6
parent
436dda0f2f
commit
a158eee1b0
@@ -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
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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
@@ -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!(
|
||||
|
||||
Reference in New Issue
Block a user