feat: add smart routing provider for cost-optimized model selection (#281)

* feat: add smart routing provider for cost-optimized model selection Route simple tasks (greetings, status checks, short questions) to a cheap model (e.g. Haiku) and complex tasks (code generation, analysis) to the primary model, reducing agent costs without sacrificing quality. Activates automatically when NEARAI_CHEAP_MODEL is set. Cascade mode retries uncertain cheap-model responses with the primary model. Co-Authored-By: Claude Opus 4.6 <[email protected]> * style: apply cargo fmt formatting Co-Authored-By: Claude Opus 4.6 <[email protected]> * refactor: extract provider chain into shared build_provider_chain() Consolidate the duplicated LLM provider chain construction from main.rs and app.rs into a single build_provider_chain() function in llm/mod.rs. This fixes the inconsistency where app.rs was missing retry wrapping that main.rs had, and ensures both paths apply identical decorators: retry → smart routing → failover → circuit breaker → cache. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address PR review — uncertainty detection and clippy lint - Remove false-positive short response (<20 chars) uncertainty check that would escalate "Yes.", "42" etc. Now only empty responses and explicit uncertainty phrases trigger cascade escalation. - Add #[allow(clippy::type_complexity)] to build_provider_chain() to fix CI clippy -D warnings failure. Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Claude Opus 4.6 <[email protected]>
2026-08-25 14:53:34 +00:00 · 2026-02-21 02:37:45 +00:00
parent 98ee648fcb
commit c038c7705b
7 changed files with 858 additions and 187 deletions
@@ -289,94 +289,14 @@ impl AppBuilder {

    /// Phase 3: Initialize LLM provider chain.
    ///
-    /// Creates the primary provider, then wraps with failover, circuit
-    /// breaker, and response cache as configured.
+    /// Delegates to `build_provider_chain` which applies all decorators
+    /// (retry, smart routing, failover, circuit breaker, response cache).
    #[allow(clippy::type_complexity)]
    pub fn init_llm(
        &self,
    ) -> Result<(Arc<dyn LlmProvider>, Option<Arc<dyn LlmProvider>>), anyhow::Error> {
-        use crate::llm::{
-            CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig,
-            FailoverProvider, ResponseCacheConfig, create_cheap_llm_provider, create_llm_provider,
-            create_llm_provider_with_config,
-        };
-
-        let llm = create_llm_provider(&self.config.llm, self.session.clone())?;
-        tracing::info!("LLM provider initialized: {}", llm.model_name());
-
-        // Wrap in failover if a fallback model is configured
-        let llm: Arc<dyn LlmProvider> = if let Some(fallback_model) =
-            self.config.llm.nearai.fallback_model.as_ref()
-        {
-            if fallback_model == &self.config.llm.nearai.model {
-                tracing::warn!(
-                    "fallback_model is the same as primary model, failover may not be effective"
-                );
-            }
-            let mut fallback_config = self.config.llm.nearai.clone();
-            fallback_config.model = fallback_model.clone();
-            let fallback = create_llm_provider_with_config(&fallback_config, self.session.clone())?;
-            tracing::info!(
-                primary = %llm.model_name(),
-                fallback = %fallback.model_name(),
-                "LLM failover enabled"
-            );
-            let cooldown_config = CooldownConfig {
-                cooldown_duration: std::time::Duration::from_secs(
-                    self.config.llm.nearai.failover_cooldown_secs,
-                ),
-                failure_threshold: self.config.llm.nearai.failover_cooldown_threshold,
-            };
-            Arc::new(FailoverProvider::with_cooldown(
-                vec![llm, fallback],
-                cooldown_config,
-            )?)
-        } else {
-            llm
-        };
-
-        // Wrap in circuit breaker if configured
-        let llm: Arc<dyn LlmProvider> =
-            if let Some(threshold) = self.config.llm.nearai.circuit_breaker_threshold {
-                let cb_config = CircuitBreakerConfig {
-                    failure_threshold: threshold,
-                    recovery_timeout: std::time::Duration::from_secs(
-                        self.config.llm.nearai.circuit_breaker_recovery_secs,
-                    ),
-                    ..CircuitBreakerConfig::default()
-                };
-                tracing::info!(
-                    threshold,
-                    recovery_secs = self.config.llm.nearai.circuit_breaker_recovery_secs,
-                    "LLM circuit breaker enabled"
-                );
-                Arc::new(CircuitBreakerProvider::new(llm, cb_config))
-            } else {
-                llm
-            };
-
-        // Wrap in response cache if configured
-        let llm: Arc<dyn LlmProvider> = if self.config.llm.nearai.response_cache_enabled {
-            let rc_config = ResponseCacheConfig {
-                ttl: std::time::Duration::from_secs(self.config.llm.nearai.response_cache_ttl_secs),
-                max_entries: self.config.llm.nearai.response_cache_max_entries,
-            };
-            tracing::info!(
-                ttl_secs = self.config.llm.nearai.response_cache_ttl_secs,
-                max_entries = self.config.llm.nearai.response_cache_max_entries,
-                "LLM response cache enabled"
-            );
-            Arc::new(CachedProvider::new(llm, rc_config))
-        } else {
-            llm
-        };
-
-        // Cheap LLM for lightweight tasks
-        let cheap_llm = create_cheap_llm_provider(&self.config.llm, self.session.clone())?;
-        if let Some(ref cheap) = cheap_llm {
-            tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name());
-        }
-
+        let (llm, cheap_llm) =
+            crate::llm::build_provider_chain(&self.config.llm, self.session.clone())?;
        Ok((llm, cheap_llm))
    }

@@ -167,6 +167,10 @@ pub struct NearAiConfig {
    /// Number of consecutive retryable failures before a provider enters
    /// cooldown (default: 3).
    pub failover_cooldown_threshold: u32,
+    /// Enable cascade mode for smart routing: when a moderate-complexity task
+    /// gets an uncertain response from the cheap model, re-send to primary.
+    /// Default: true.
+    pub smart_routing_cascade: bool,
 }

 impl LlmConfig {
@@ -232,6 +236,7 @@ impl LlmConfig {
            response_cache_max_entries: parse_optional_env("RESPONSE_CACHE_MAX_ENTRIES", 1000)?,
            failover_cooldown_secs: parse_optional_env("LLM_FAILOVER_COOLDOWN_SECS", 300)?,
            failover_cooldown_threshold: parse_optional_env("LLM_FAILOVER_THRESHOLD", 3)?,
+            smart_routing_cascade: parse_optional_env("SMART_ROUTING_CASCADE", true)?,
        };

        // Resolve provider-specific configs based on backend
@@ -17,6 +17,7 @@ pub mod response_cache;
 pub mod retry;
 mod rig_adapter;
 pub mod session;
+pub mod smart_routing;

 pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider};
 pub use failover::{CooldownConfig, FailoverProvider};
@@ -33,6 +34,7 @@ pub use response_cache::{CachedProvider, ResponseCacheConfig};
 pub use retry::{RetryConfig, RetryProvider};
 pub use rig_adapter::RigAdapter;
 pub use session::{SessionConfig, SessionManager, create_session_manager};
+pub use smart_routing::{SmartRoutingConfig, SmartRoutingProvider, TaskComplexity};

 use std::sync::Arc;

@@ -273,6 +275,147 @@ pub fn create_cheap_llm_provider(
    )?)))
 }

+/// Build the full LLM provider chain with all configured wrappers.
+///
+/// Applies decorators in this order:
+/// 1. Raw provider (from config)
+/// 2. RetryProvider (per-provider retry with exponential backoff)
+/// 3. SmartRoutingProvider (cheap/primary split when cheap model is configured)
+/// 4. FailoverProvider (fallback model when primary fails)
+/// 5. CircuitBreakerProvider (fast-fail when backend is degraded)
+/// 6. CachedProvider (in-memory response cache)
+///
+/// Also returns a separate cheap LLM provider for heartbeat/evaluation (not
+/// part of the chain — it's a standalone provider for explicitly cheap tasks).
+///
+/// This is the single source of truth for provider chain construction,
+/// called by both `main.rs` and `app.rs`.
+#[allow(clippy::type_complexity)]
+pub fn build_provider_chain(
+    config: &LlmConfig,
+    session: Arc<SessionManager>,
+) -> Result<(Arc<dyn LlmProvider>, Option<Arc<dyn LlmProvider>>), LlmError> {
+    let llm = create_llm_provider(config, session.clone())?;
+    tracing::info!("LLM provider initialized: {}", llm.model_name());
+
+    // 1. Retry
+    let retry_config = RetryConfig {
+        max_retries: config.nearai.max_retries,
+    };
+    let llm: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
+        tracing::info!(
+            max_retries = retry_config.max_retries,
+            "LLM retry wrapper enabled"
+        );
+        Arc::new(RetryProvider::new(llm, retry_config.clone()))
+    } else {
+        llm
+    };
+
+    // 2. Smart routing (cheap/primary split)
+    let llm: Arc<dyn LlmProvider> = if let Some(ref cheap_model) = config.nearai.cheap_model {
+        let mut cheap_config = config.nearai.clone();
+        cheap_config.model = cheap_model.clone();
+        let cheap = create_llm_provider_with_config(&cheap_config, session.clone())?;
+        let cheap: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
+            Arc::new(RetryProvider::new(cheap, retry_config.clone()))
+        } else {
+            cheap
+        };
+        tracing::info!(
+            primary = %llm.model_name(),
+            cheap = %cheap.model_name(),
+            "Smart routing enabled"
+        );
+        Arc::new(SmartRoutingProvider::new(
+            llm,
+            cheap,
+            SmartRoutingConfig {
+                cascade_enabled: config.nearai.smart_routing_cascade,
+                ..SmartRoutingConfig::default()
+            },
+        ))
+    } else {
+        llm
+    };
+
+    // 3. Failover
+    let llm: Arc<dyn LlmProvider> = if let Some(ref fallback_model) = config.nearai.fallback_model {
+        if fallback_model == &config.nearai.model {
+            tracing::warn!(
+                "fallback_model is the same as primary model, failover may not be effective"
+            );
+        }
+        let mut fallback_config = config.nearai.clone();
+        fallback_config.model = fallback_model.clone();
+        let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?;
+        tracing::info!(
+            primary = %llm.model_name(),
+            fallback = %fallback.model_name(),
+            "LLM failover enabled"
+        );
+        let fallback: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
+            Arc::new(RetryProvider::new(fallback, retry_config.clone()))
+        } else {
+            fallback
+        };
+        let cooldown_config = CooldownConfig {
+            cooldown_duration: std::time::Duration::from_secs(config.nearai.failover_cooldown_secs),
+            failure_threshold: config.nearai.failover_cooldown_threshold,
+        };
+        Arc::new(FailoverProvider::with_cooldown(
+            vec![llm, fallback],
+            cooldown_config,
+        )?)
+    } else {
+        llm
+    };
+
+    // 4. Circuit breaker
+    let llm: Arc<dyn LlmProvider> = if let Some(threshold) = config.nearai.circuit_breaker_threshold
+    {
+        let cb_config = CircuitBreakerConfig {
+            failure_threshold: threshold,
+            recovery_timeout: std::time::Duration::from_secs(
+                config.nearai.circuit_breaker_recovery_secs,
+            ),
+            ..CircuitBreakerConfig::default()
+        };
+        tracing::info!(
+            threshold,
+            recovery_secs = config.nearai.circuit_breaker_recovery_secs,
+            "LLM circuit breaker enabled"
+        );
+        Arc::new(CircuitBreakerProvider::new(llm, cb_config))
+    } else {
+        llm
+    };
+
+    // 5. Response cache
+    let llm: Arc<dyn LlmProvider> = if config.nearai.response_cache_enabled {
+        let rc_config = ResponseCacheConfig {
+            ttl: std::time::Duration::from_secs(config.nearai.response_cache_ttl_secs),
+            max_entries: config.nearai.response_cache_max_entries,
+        };
+        tracing::info!(
+            ttl_secs = config.nearai.response_cache_ttl_secs,
+            max_entries = config.nearai.response_cache_max_entries,
+            "LLM response cache enabled"
+        );
+        Arc::new(CachedProvider::new(llm, rc_config))
+    } else {
+        llm
+    };
+
+    // Standalone cheap LLM for heartbeat/evaluation (not part of the chain)
+    let cheap_llm = create_cheap_llm_provider(config, session)?;
+    if let Some(ref cheap) = cheap_llm {
+        tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name());
+    }
+
+    Ok((llm, cheap_llm))
+}
+
 #[cfg(test)]
 mod tests {
    use super::*;
@@ -296,6 +439,7 @@ mod tests {
            response_cache_max_entries: 1000,
            failover_cooldown_secs: 300,
            failover_cooldown_threshold: 3,
+            smart_routing_cascade: true,
        }
    }

@@ -766,6 +766,7 @@ mod tests {
            response_cache_max_entries: 1000,
            failover_cooldown_secs: 300,
            failover_cooldown_threshold: 3,
+            smart_routing_cascade: true,
        }
    }

@@ -0,0 +1,700 @@
+//! Smart routing provider that routes requests to cheap or primary models based on task complexity.
+//!
+//! Inspired by RelayPlane's cost-reduction approach: simple tasks (status checks, greetings,
+//! short questions) go to a cheap model (e.g. Haiku), while complex tasks (code generation,
+//! analysis, multi-step reasoning) go to the primary model (e.g. Sonnet/Opus).
+//!
+//! This is a decorator that wraps two `LlmProvider`s and implements `LlmProvider` itself,
+//! following the same pattern as `RetryProvider`, `CachedProvider`, and `CircuitBreakerProvider`.
+
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use async_trait::async_trait;
+use rust_decimal::Decimal;
+
+use crate::error::LlmError;
+use crate::llm::provider::{
+    CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, Role, ToolCompletionRequest,
+    ToolCompletionResponse,
+};
+
+/// Classification of a request's complexity, determining which model handles it.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum TaskComplexity {
+    /// Short, simple queries -> cheap model
+    Simple,
+    /// Ambiguous complexity -> cheap model first, cascade to primary if uncertain
+    Moderate,
+    /// Code generation, analysis, multi-step reasoning -> primary model
+    Complex,
+}
+
+/// Configuration for the smart routing provider.
+#[derive(Debug, Clone)]
+pub struct SmartRoutingConfig {
+    /// Enable cascade mode: retry with primary if cheap model response seems uncertain.
+    pub cascade_enabled: bool,
+    /// Message length threshold below which a message may be classified as Simple (default: 200).
+    pub simple_max_chars: usize,
+    /// Message length threshold above which a message is classified as Complex (default: 1000).
+    pub complex_min_chars: usize,
+}
+
+impl Default for SmartRoutingConfig {
+    fn default() -> Self {
+        Self {
+            cascade_enabled: true,
+            simple_max_chars: 200,
+            complex_min_chars: 1000,
+        }
+    }
+}
+
+/// Atomic counters for routing observability.
+struct SmartRoutingStats {
+    total_requests: AtomicU64,
+    cheap_requests: AtomicU64,
+    primary_requests: AtomicU64,
+    cascade_escalations: AtomicU64,
+}
+
+impl SmartRoutingStats {
+    fn new() -> Self {
+        Self {
+            total_requests: AtomicU64::new(0),
+            cheap_requests: AtomicU64::new(0),
+            primary_requests: AtomicU64::new(0),
+            cascade_escalations: AtomicU64::new(0),
+        }
+    }
+}
+
+/// Snapshot of routing statistics for external consumption.
+#[derive(Debug, Clone)]
+pub struct SmartRoutingSnapshot {
+    pub total_requests: u64,
+    pub cheap_requests: u64,
+    pub primary_requests: u64,
+    pub cascade_escalations: u64,
+}
+
+/// Smart routing provider that classifies task complexity and routes to the appropriate model.
+///
+/// - `complete()` — classifies and routes to cheap or primary model
+/// - `complete_with_tools()` — always routes to primary (tool use requires reliable structured output)
+pub struct SmartRoutingProvider {
+    primary: Arc<dyn LlmProvider>,
+    cheap: Arc<dyn LlmProvider>,
+    config: SmartRoutingConfig,
+    stats: SmartRoutingStats,
+}
+
+impl SmartRoutingProvider {
+    /// Create a new smart routing provider wrapping a primary and cheap provider.
+    pub fn new(
+        primary: Arc<dyn LlmProvider>,
+        cheap: Arc<dyn LlmProvider>,
+        config: SmartRoutingConfig,
+    ) -> Self {
+        Self {
+            primary,
+            cheap,
+            config,
+            stats: SmartRoutingStats::new(),
+        }
+    }
+
+    /// Get a snapshot of routing statistics.
+    pub fn stats(&self) -> SmartRoutingSnapshot {
+        SmartRoutingSnapshot {
+            total_requests: self.stats.total_requests.load(Ordering::Relaxed),
+            cheap_requests: self.stats.cheap_requests.load(Ordering::Relaxed),
+            primary_requests: self.stats.primary_requests.load(Ordering::Relaxed),
+            cascade_escalations: self.stats.cascade_escalations.load(Ordering::Relaxed),
+        }
+    }
+
+    /// Classify the complexity of a request based on its last user message.
+    fn classify(&self, request: &CompletionRequest) -> TaskComplexity {
+        let last_user_msg = request
+            .messages
+            .iter()
+            .rev()
+            .find(|m| m.role == Role::User)
+            .map(|m| m.content.as_str())
+            .unwrap_or("");
+
+        classify_message(last_user_msg, &self.config)
+    }
+
+    /// Check if a response from the cheap model shows uncertainty, warranting escalation.
+    fn response_is_uncertain(response: &CompletionResponse) -> bool {
+        let content = response.content.trim();
+
+        // Empty response is always uncertain
+        if content.is_empty() {
+            return true;
+        }
+
+        let lower = content.to_lowercase();
+
+        // Uncertainty signals
+        let uncertainty_patterns = [
+            "i'm not sure",
+            "i am not sure",
+            "i don't know",
+            "i do not know",
+            "i'm unable to",
+            "i am unable to",
+            "i cannot",
+            "i can't",
+            "beyond my capabilities",
+            "beyond my ability",
+            "i'm not able to",
+            "i am not able to",
+            "i don't have enough",
+            "i do not have enough",
+            "i need more context",
+            "i need more information",
+            "could you clarify",
+            "could you provide more",
+            "i'm not confident",
+            "i am not confident",
+        ];
+
+        uncertainty_patterns.iter().any(|p| lower.contains(p))
+    }
+}
+
+/// Classify a message's complexity based on content patterns and length.
+///
+/// Exposed as a free function for testability.
+fn classify_message(msg: &str, config: &SmartRoutingConfig) -> TaskComplexity {
+    let trimmed = msg.trim();
+    let len = trimmed.len();
+
+    // Empty or very short -> Simple
+    if len == 0 {
+        return TaskComplexity::Simple;
+    }
+
+    // Check for code blocks (triple backticks) -> Complex
+    if trimmed.contains("```") {
+        return TaskComplexity::Complex;
+    }
+
+    let lower = trimmed.to_lowercase();
+
+    // Complex keywords/patterns -> Complex regardless of length
+    const COMPLEX_KEYWORDS: &[&str] = &[
+        "implement",
+        "refactor",
+        "analyze",
+        "debug",
+        "create a",
+        "build a",
+        "design",
+        "fix the",
+        "fix this",
+        "write a",
+        "write the",
+        "explain how",
+        "explain why",
+        "explain the",
+        "compare",
+        "optimize",
+        "review",
+        "rewrite",
+        "migrate",
+        "architect",
+        "integrate",
+    ];
+
+    if COMPLEX_KEYWORDS.iter().any(|k| lower.contains(k)) {
+        return TaskComplexity::Complex;
+    }
+
+    // Long messages -> Complex
+    if len >= config.complex_min_chars {
+        return TaskComplexity::Complex;
+    }
+
+    // Simple keywords/patterns for short messages
+    const SIMPLE_KEYWORDS: &[&str] = &[
+        "list",
+        "show",
+        "what is",
+        "what's",
+        "status",
+        "help",
+        "yes",
+        "no",
+        "ok",
+        "thanks",
+        "thank you",
+        "hello",
+        "hi",
+        "hey",
+        "ping",
+        "version",
+        "how many",
+        "when",
+        "where is",
+        "who",
+    ];
+
+    if len <= config.simple_max_chars && SIMPLE_KEYWORDS.iter().any(|k| lower.contains(k)) {
+        return TaskComplexity::Simple;
+    }
+
+    // Short confirmations / single words -> Simple
+    if len <= 10 {
+        return TaskComplexity::Simple;
+    }
+
+    // Everything else -> Moderate
+    TaskComplexity::Moderate
+}
+
+#[async_trait]
+impl LlmProvider for SmartRoutingProvider {
+    fn model_name(&self) -> &str {
+        self.primary.model_name()
+    }
+
+    fn cost_per_token(&self) -> (Decimal, Decimal) {
+        self.primary.cost_per_token()
+    }
+
+    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
+        self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
+
+        let complexity = self.classify(&request);
+
+        match complexity {
+            TaskComplexity::Simple => {
+                tracing::debug!(
+                    model = %self.cheap.model_name(),
+                    "Smart routing: Simple task -> cheap model"
+                );
+                self.stats.cheap_requests.fetch_add(1, Ordering::Relaxed);
+                self.cheap.complete(request).await
+            }
+            TaskComplexity::Complex => {
+                tracing::debug!(
+                    model = %self.primary.model_name(),
+                    "Smart routing: Complex task -> primary model"
+                );
+                self.stats.primary_requests.fetch_add(1, Ordering::Relaxed);
+                self.primary.complete(request).await
+            }
+            TaskComplexity::Moderate => {
+                if self.config.cascade_enabled {
+                    tracing::debug!(
+                        model = %self.cheap.model_name(),
+                        "Smart routing: Moderate task -> cheap model (cascade enabled)"
+                    );
+                    self.stats.cheap_requests.fetch_add(1, Ordering::Relaxed);
+
+                    let response = self.cheap.complete(request.clone()).await?;
+
+                    if Self::response_is_uncertain(&response) {
+                        tracing::info!(
+                            cheap_model = %self.cheap.model_name(),
+                            primary_model = %self.primary.model_name(),
+                            "Smart routing: Escalating to primary (cheap model response uncertain)"
+                        );
+                        self.stats
+                            .cascade_escalations
+                            .fetch_add(1, Ordering::Relaxed);
+                        self.stats.primary_requests.fetch_add(1, Ordering::Relaxed);
+                        self.primary.complete(request).await
+                    } else {
+                        Ok(response)
+                    }
+                } else {
+                    // Without cascade, moderate tasks go to cheap model
+                    tracing::debug!(
+                        model = %self.cheap.model_name(),
+                        "Smart routing: Moderate task -> cheap model (cascade disabled)"
+                    );
+                    self.stats.cheap_requests.fetch_add(1, Ordering::Relaxed);
+                    self.cheap.complete(request).await
+                }
+            }
+        }
+    }
+
+    /// Tool use always goes to the primary model for reliable structured output.
+    async fn complete_with_tools(
+        &self,
+        request: ToolCompletionRequest,
+    ) -> Result<ToolCompletionResponse, LlmError> {
+        self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
+        self.stats.primary_requests.fetch_add(1, Ordering::Relaxed);
+        tracing::debug!(
+            model = %self.primary.model_name(),
+            "Smart routing: Tool use -> primary model (always)"
+        );
+        self.primary.complete_with_tools(request).await
+    }
+
+    async fn list_models(&self) -> Result<Vec<String>, LlmError> {
+        self.primary.list_models().await
+    }
+
+    async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
+        self.primary.model_metadata().await
+    }
+
+    fn active_model_name(&self) -> String {
+        self.primary.active_model_name()
+    }
+
+    fn set_model(&self, model: &str) -> Result<(), LlmError> {
+        self.primary.set_model(model)
+    }
+
+    fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
+        self.primary.calculate_cost(input_tokens, output_tokens)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::llm::ChatMessage;
+    use crate::testing::StubLlm;
+
+    fn default_config() -> SmartRoutingConfig {
+        SmartRoutingConfig::default()
+    }
+
+    // -- Classification tests --
+
+    #[test]
+    fn classify_empty_message_as_simple() {
+        assert_eq!(
+            classify_message("", &default_config()),
+            TaskComplexity::Simple
+        );
+    }
+
+    #[test]
+    fn classify_greeting_as_simple() {
+        assert_eq!(
+            classify_message("hello", &default_config()),
+            TaskComplexity::Simple
+        );
+        assert_eq!(
+            classify_message("hi there", &default_config()),
+            TaskComplexity::Simple
+        );
+    }
+
+    #[test]
+    fn classify_short_question_with_simple_keyword() {
+        assert_eq!(
+            classify_message("what is the status?", &default_config()),
+            TaskComplexity::Simple
+        );
+        assert_eq!(
+            classify_message("show me the list", &default_config()),
+            TaskComplexity::Simple
+        );
+        assert_eq!(
+            classify_message("help", &default_config()),
+            TaskComplexity::Simple
+        );
+    }
+
+    #[test]
+    fn classify_yes_no_as_simple() {
+        assert_eq!(
+            classify_message("yes", &default_config()),
+            TaskComplexity::Simple
+        );
+        assert_eq!(
+            classify_message("no", &default_config()),
+            TaskComplexity::Simple
+        );
+        assert_eq!(
+            classify_message("ok", &default_config()),
+            TaskComplexity::Simple
+        );
+    }
+
+    #[test]
+    fn classify_code_generation_as_complex() {
+        assert_eq!(
+            classify_message("implement a binary search function", &default_config()),
+            TaskComplexity::Complex
+        );
+        assert_eq!(
+            classify_message("refactor the auth module", &default_config()),
+            TaskComplexity::Complex
+        );
+        assert_eq!(
+            classify_message("debug this error", &default_config()),
+            TaskComplexity::Complex
+        );
+    }
+
+    #[test]
+    fn classify_code_blocks_as_complex() {
+        let msg = "What does this do?\n```rust\nfn main() {}\n```";
+        assert_eq!(
+            classify_message(msg, &default_config()),
+            TaskComplexity::Complex
+        );
+    }
+
+    #[test]
+    fn classify_long_message_as_complex() {
+        let long_msg = "a ".repeat(600); // 1200 chars
+        assert_eq!(
+            classify_message(&long_msg, &default_config()),
+            TaskComplexity::Complex
+        );
+    }
+
+    #[test]
+    fn classify_medium_message_without_keywords_as_moderate() {
+        // > 10 chars, < 1000 chars, no simple or complex keywords
+        let msg = "Tell me about the weather patterns in the Pacific Ocean during summer months";
+        assert_eq!(
+            classify_message(msg, &default_config()),
+            TaskComplexity::Moderate
+        );
+    }
+
+    #[test]
+    fn classify_very_short_unknown_as_simple() {
+        // <= 10 chars, no keywords
+        assert_eq!(
+            classify_message("foo", &default_config()),
+            TaskComplexity::Simple
+        );
+    }
+
+    // -- Uncertainty detection tests --
+
+    #[test]
+    fn detects_uncertain_short_response() {
+        let response = CompletionResponse {
+            content: "I'm not sure.".to_string(),
+            input_tokens: 10,
+            output_tokens: 5,
+            finish_reason: crate::llm::FinishReason::Stop,
+        };
+        assert!(SmartRoutingProvider::response_is_uncertain(&response));
+    }
+
+    #[test]
+    fn detects_empty_response_as_uncertain() {
+        let response = CompletionResponse {
+            content: "".to_string(),
+            input_tokens: 10,
+            output_tokens: 0,
+            finish_reason: crate::llm::FinishReason::Stop,
+        };
+        assert!(SmartRoutingProvider::response_is_uncertain(&response));
+    }
+
+    #[test]
+    fn short_confident_response_is_not_uncertain() {
+        let response = CompletionResponse {
+            content: "Yes.".to_string(),
+            input_tokens: 10,
+            output_tokens: 1,
+            finish_reason: crate::llm::FinishReason::Stop,
+        };
+        assert!(!SmartRoutingProvider::response_is_uncertain(&response));
+    }
+
+    #[test]
+    fn confident_response_is_not_uncertain() {
+        let response = CompletionResponse {
+            content: "The answer is 42. This is a well-known constant from the Hitchhiker's Guide."
+                .to_string(),
+            input_tokens: 10,
+            output_tokens: 20,
+            finish_reason: crate::llm::FinishReason::Stop,
+        };
+        assert!(!SmartRoutingProvider::response_is_uncertain(&response));
+    }
+
+    // -- Routing tests --
+
+    fn make_request(content: &str) -> CompletionRequest {
+        CompletionRequest::new(vec![ChatMessage::user(content)])
+    }
+
+    fn make_tool_request() -> ToolCompletionRequest {
+        ToolCompletionRequest::new(vec![ChatMessage::user("implement a search")], vec![])
+    }
+
+    #[tokio::test]
+    async fn simple_task_routes_to_cheap() {
+        let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
+        let cheap = Arc::new(StubLlm::new("cheap-response").with_model_name("cheap"));
+
+        let router = SmartRoutingProvider::new(
+            primary.clone(),
+            cheap.clone(),
+            SmartRoutingConfig {
+                cascade_enabled: false,
+                ..default_config()
+            },
+        );
+
+        let resp = router.complete(make_request("hello")).await.unwrap();
+        assert_eq!(resp.content, "cheap-response");
+        assert_eq!(cheap.calls(), 1);
+        assert_eq!(primary.calls(), 0);
+    }
+
+    #[tokio::test]
+    async fn complex_task_routes_to_primary() {
+        let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
+        let cheap = Arc::new(StubLlm::new("cheap-response").with_model_name("cheap"));
+
+        let router = SmartRoutingProvider::new(primary.clone(), cheap.clone(), default_config());
+
+        let resp = router
+            .complete(make_request("implement a binary search"))
+            .await
+            .unwrap();
+        assert_eq!(resp.content, "primary-response");
+        assert_eq!(primary.calls(), 1);
+        assert_eq!(cheap.calls(), 0);
+    }
+
+    #[tokio::test]
+    async fn tool_use_always_routes_to_primary() {
+        let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
+        let cheap = Arc::new(StubLlm::new("cheap-response").with_model_name("cheap"));
+
+        let router = SmartRoutingProvider::new(primary.clone(), cheap.clone(), default_config());
+
+        let resp = router
+            .complete_with_tools(make_tool_request())
+            .await
+            .unwrap();
+        assert_eq!(resp.content, Some("primary-response".to_string()));
+        assert_eq!(primary.calls(), 1);
+        assert_eq!(cheap.calls(), 0);
+    }
+
+    #[tokio::test]
+    async fn stats_increment_correctly() {
+        let primary = Arc::new(StubLlm::new("primary").with_model_name("primary"));
+        let cheap = Arc::new(StubLlm::new("cheap").with_model_name("cheap"));
+
+        let router = SmartRoutingProvider::new(
+            primary,
+            cheap,
+            SmartRoutingConfig {
+                cascade_enabled: false,
+                ..default_config()
+            },
+        );
+
+        // Simple -> cheap
+        router.complete(make_request("hello")).await.unwrap();
+        // Complex -> primary
+        router
+            .complete(make_request("implement a search"))
+            .await
+            .unwrap();
+        // Tool use -> primary
+        router
+            .complete_with_tools(make_tool_request())
+            .await
+            .unwrap();
+
+        let stats = router.stats();
+        assert_eq!(stats.total_requests, 3);
+        assert_eq!(stats.cheap_requests, 1);
+        assert_eq!(stats.primary_requests, 2);
+        assert_eq!(stats.cascade_escalations, 0);
+    }
+
+    #[tokio::test]
+    async fn cascade_escalates_on_uncertain_response() {
+        // Cheap model returns an uncertain response
+        let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
+        let cheap = Arc::new(StubLlm::new("I'm not sure about that.").with_model_name("cheap"));
+
+        let router = SmartRoutingProvider::new(
+            primary.clone(),
+            cheap.clone(),
+            SmartRoutingConfig {
+                cascade_enabled: true,
+                ..default_config()
+            },
+        );
+
+        // A moderate task (no simple/complex keywords, medium length)
+        let resp = router
+            .complete(make_request(
+                "Tell me about the weather patterns in the Pacific Ocean during summer months",
+            ))
+            .await
+            .unwrap();
+
+        // Should have escalated to primary
+        assert_eq!(resp.content, "primary-response");
+        assert_eq!(cheap.calls(), 1);
+        assert_eq!(primary.calls(), 1);
+
+        let stats = router.stats();
+        assert_eq!(stats.cascade_escalations, 1);
+    }
+
+    #[tokio::test]
+    async fn cascade_does_not_escalate_on_confident_response() {
+        let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
+        let cheap = Arc::new(
+            StubLlm::new(
+                "The Pacific Ocean weather patterns during summer are characterized by trade winds.",
+            )
+            .with_model_name("cheap"),
+        );
+
+        let router = SmartRoutingProvider::new(
+            primary.clone(),
+            cheap.clone(),
+            SmartRoutingConfig {
+                cascade_enabled: true,
+                ..default_config()
+            },
+        );
+
+        let resp = router
+            .complete(make_request(
+                "Tell me about the weather patterns in the Pacific Ocean during summer months",
+            ))
+            .await
+            .unwrap();
+
+        // Should NOT have escalated
+        assert!(resp.content.contains("Pacific Ocean"));
+        assert_eq!(cheap.calls(), 1);
+        assert_eq!(primary.calls(), 0);
+
+        let stats = router.stats();
+        assert_eq!(stats.cascade_escalations, 0);
+    }
+
+    #[tokio::test]
+    async fn model_name_returns_primary() {
+        let primary = Arc::new(StubLlm::new("ok").with_model_name("sonnet"));
+        let cheap = Arc::new(StubLlm::new("ok").with_model_name("haiku"));
+
+        let router = SmartRoutingProvider::new(primary, cheap, default_config());
+        assert_eq!(router.model_name(), "sonnet");
+        assert_eq!(router.active_model_name(), "sonnet");
+    }
+}
@@ -24,12 +24,7 @@ use ironclaw::{
    context::ContextManager,
    extensions::ExtensionManager,
    hooks::{HookRegistry, bootstrap_hooks},
-    llm::{
-        CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig,
-        FailoverProvider, LlmProvider, ResponseCacheConfig, RetryConfig, RetryProvider,
-        SessionConfig, create_cheap_llm_provider, create_llm_provider,
-        create_llm_provider_with_config, create_session_manager,
-    },
+    llm::{SessionConfig, build_provider_chain, create_session_manager},
    orchestrator::{
        ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
        api::OrchestratorState,
@@ -626,103 +621,8 @@ async fn main() -> anyhow::Result<()> {
            None
        };

-    // Initialize LLM provider (clone session so we can reuse it for embeddings)
-    let llm = create_llm_provider(&config.llm, session.clone())?;
-    tracing::info!("LLM provider initialized: {}", llm.model_name());
-
-    // Wrap each provider with RetryProvider for automatic retries on transient errors.
-    // RetryProvider sits inside FailoverProvider so each provider in the failover chain
-    // gets its own retry attempts before the failover moves to the next provider.
-    let retry_config = RetryConfig {
-        max_retries: config.llm.nearai.max_retries,
-    };
-    let llm: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
-        tracing::info!(
-            max_retries = retry_config.max_retries,
-            "LLM retry wrapper enabled"
-        );
-        Arc::new(RetryProvider::new(llm, retry_config.clone()))
-    } else {
-        llm
-    };
-
-    // Wrap in failover if a fallback model is configured
-    let llm: Arc<dyn LlmProvider> =
-        if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() {
-            if fallback_model == &config.llm.nearai.model {
-                tracing::warn!(
-                    "fallback_model is the same as primary model, failover may not be effective"
-                );
-            }
-            let mut fallback_config = config.llm.nearai.clone();
-            fallback_config.model = fallback_model.clone();
-            let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?;
-            tracing::info!(
-                primary = %llm.model_name(),
-                fallback = %fallback.model_name(),
-                "LLM failover enabled"
-            );
-            // Wrap fallback with retry too
-            let fallback: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
-                Arc::new(RetryProvider::new(fallback, retry_config.clone()))
-            } else {
-                fallback
-            };
-            let cooldown_config = CooldownConfig {
-                cooldown_duration: std::time::Duration::from_secs(
-                    config.llm.nearai.failover_cooldown_secs,
-                ),
-                failure_threshold: config.llm.nearai.failover_cooldown_threshold,
-            };
-            Arc::new(FailoverProvider::with_cooldown(
-                vec![llm, fallback],
-                cooldown_config,
-            )?)
-        } else {
-            llm
-        };
-
-    // Wrap in circuit breaker if configured
-    let llm: Arc<dyn LlmProvider> =
-        if let Some(threshold) = config.llm.nearai.circuit_breaker_threshold {
-            let cb_config = CircuitBreakerConfig {
-                failure_threshold: threshold,
-                recovery_timeout: std::time::Duration::from_secs(
-                    config.llm.nearai.circuit_breaker_recovery_secs,
-                ),
-                ..CircuitBreakerConfig::default()
-            };
-            tracing::info!(
-                threshold,
-                recovery_secs = config.llm.nearai.circuit_breaker_recovery_secs,
-                "LLM circuit breaker enabled"
-            );
-            Arc::new(CircuitBreakerProvider::new(llm, cb_config))
-        } else {
-            llm
-        };
-
-    // Wrap in response cache if configured
-    let llm: Arc<dyn LlmProvider> = if config.llm.nearai.response_cache_enabled {
-        let rc_config = ResponseCacheConfig {
-            ttl: std::time::Duration::from_secs(config.llm.nearai.response_cache_ttl_secs),
-            max_entries: config.llm.nearai.response_cache_max_entries,
-        };
-        tracing::info!(
-            ttl_secs = config.llm.nearai.response_cache_ttl_secs,
-            max_entries = config.llm.nearai.response_cache_max_entries,
-            "LLM response cache enabled"
-        );
-        Arc::new(CachedProvider::new(llm, rc_config))
-    } else {
-        llm
-    };
-
-    // Initialize cheap LLM provider for lightweight tasks (heartbeat, evaluation)
-    let cheap_llm = create_cheap_llm_provider(&config.llm, session.clone())?;
-    if let Some(ref cheap) = cheap_llm {
-        tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name());
-    }
+    // Build the full LLM provider chain (retry → smart routing → failover → circuit breaker → cache)
+    let (llm, cheap_llm) = build_provider_chain(&config.llm, session.clone())?;

    // Initialize safety layer
    let safety = Arc::new(SafetyLayer::new(&config.safety));
@@ -1176,6 +1176,7 @@ impl SetupWizard {
                response_cache_max_entries: 1000,
                failover_cooldown_secs: 300,
                failover_cooldown_threshold: 3,
+                smart_routing_cascade: true,
            },
            openai: None,
            anthropic: None,