ares_llm/

pool.rs

//! LLM Client Connection Pooling (DIR-44)
//!
//! This module provides connection pooling for LLM clients, enabling connection
//! reuse across requests to reduce latency and resource consumption.
//!
//! # Architecture
//!
//! The pool maintains a set of pre-initialized `LLMClient` instances per provider
//! configuration. Clients are checked out, used, and returned to the pool.
//!
//! # Features
//!
//! - Configurable maximum pool size per provider
//! - Connection health checking with configurable TTL
//! - Automatic stale connection cleanup
//! - Graceful shutdown with connection draining
//! - Fair distribution via round-robin or least-connections
//!
//! # Example
//!
//! ```rust,ignore
//! use ares::llm::pool::{ClientPool, PoolConfig};
//! use ares::llm::Provider;
//!
//! let config = PoolConfig::default();
//! let pool = ClientPool::new(config);
//!
//! // Register a provider
//! pool.register_provider("openai", provider).await?;
//!
//! // Get a pooled client
//! let guard = pool.get("openai").await?;
//! let response = guard.client().generate("Hello!").await?;
//! // Client is automatically returned to pool when guard is dropped
//! ```

use crate::client::{LLMClient, Provider};
use ares_types::types::{AppError, Result};
use parking_lot::{Mutex, RwLock};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};

/// Pool-specific errors for borrow/return operations (R42).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum PoolError {
    #[serde(rename = "pool_exhausted")]
    PoolExhausted { max: usize },
    #[serde(rename = "timeout")]
    Timeout { timeout_ms: u64 },
    #[serde(rename = "invalid_client")]
    InvalidClient { reason: String },
}

impl std::fmt::Display for PoolError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::PoolExhausted { max } => write!(f, "pool exhausted (max={max})"),
            Self::Timeout { timeout_ms } => write!(f, "pool acquire timeout after {timeout_ms}ms"),
            Self::InvalidClient { reason } => write!(f, "invalid pooled client: {reason}"),
        }
    }
}

impl std::error::Error for PoolError {}

impl From<PoolError> for AppError {
    fn from(err: PoolError) -> Self {
        AppError::LLM(err.to_string())
    }
}

mod duration_secs {
    use serde::{Deserialize, Deserializer, Serializer};
    use std::time::Duration;

    pub fn serialize<S>(value: &Duration, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_u64(value.as_secs())
    }

    pub fn deserialize<'de, D>(deserializer: D) -> std::result::Result<Duration, D::Error>
    where
        D: Deserializer<'de>,
    {
        Ok(Duration::from_secs(u64::deserialize(deserializer)?))
    }
}

/// Configuration for the client pool
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct PoolConfig {
    /// Maximum number of clients per provider (default: 10)
    pub max_connections_per_provider: usize,

    /// Minimum number of idle clients to maintain per provider (default: 2)
    pub min_idle_connections: usize,

    /// Maximum time a client can be idle before being considered stale (default: 5 minutes)
    #[serde(with = "duration_secs")]
    pub idle_timeout: Duration,

    /// Maximum lifetime of a client before forced refresh (default: 30 minutes)
    #[serde(with = "duration_secs")]
    pub max_lifetime: Duration,

    /// How often to run health checks on idle connections (default: 60 seconds)
    #[serde(with = "duration_secs")]
    pub health_check_interval: Duration,

    /// Timeout for acquiring a client from the pool (default: 30 seconds)
    #[serde(with = "duration_secs")]
    pub acquire_timeout: Duration,

    /// Whether to enable connection health checking (default: true)
    pub enable_health_check: bool,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            max_connections_per_provider: 10,
            min_idle_connections: 2,
            idle_timeout: Duration::from_secs(300), // 5 minutes
            max_lifetime: Duration::from_secs(1800), // 30 minutes
            health_check_interval: Duration::from_secs(60),
            acquire_timeout: Duration::from_secs(30),
            enable_health_check: true,
        }
    }
}

impl PoolConfig {
    /// Create a new pool config with custom max connections
    pub fn with_max_connections(mut self, max: usize) -> Self {
        self.max_connections_per_provider = max;
        self
    }

    /// Create a new pool config with custom idle timeout
    pub fn with_idle_timeout(mut self, timeout: Duration) -> Self {
        self.idle_timeout = timeout;
        self
    }

    /// Create a new pool config with custom max lifetime
    pub fn with_max_lifetime(mut self, lifetime: Duration) -> Self {
        self.max_lifetime = lifetime;
        self
    }

    /// Disable health checking (useful for testing)
    pub fn without_health_check(mut self) -> Self {
        self.enable_health_check = false;
        self
    }
}


impl std::fmt::Display for PoolConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "PoolConfig(max={}, min_idle={}, health_check={})",
            self.max_connections_per_provider,
            self.min_idle_connections,
            self.enable_health_check
        )
    }
}

/// Metadata for a pooled client
#[derive(Debug)]
struct PooledClientMeta {
    /// When this client was created
    created_at: Instant,
    /// When this client was last used
    last_used: Instant,
    /// Number of times this client has been used
    #[allow(dead_code)] // Used for metrics/debugging
    use_count: AtomicU64,
}

impl PooledClientMeta {
    fn new() -> Self {
        let now = Instant::now();
        Self {
            created_at: now,
            last_used: now,
            use_count: AtomicU64::new(0),
        }
    }

    fn mark_used(&mut self) {
        self.last_used = Instant::now();
        self.use_count.fetch_add(1, Ordering::Relaxed);
    }

    fn is_stale(&self, config: &PoolConfig) -> bool {
        let now = Instant::now();
        let idle_duration = now.duration_since(self.last_used);
        let lifetime = now.duration_since(self.created_at);

        idle_duration > config.idle_timeout || lifetime > config.max_lifetime
    }
}


#[derive(Debug)]
enum BorrowFromIdle {
    Found(PooledClient),
    Exhausted,
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ReturnDisposition {
    Returned,
    Dropped,
}

fn borrow_client(idle: &mut Vec<PooledClient>, config: &PoolConfig) -> BorrowFromIdle {
    let mut found_idx = None;
    for (idx, pooled) in idle.iter().enumerate() {
        if health_check(&pooled.meta, config) {
            found_idx = Some(idx);
            break;
        }
    }
    if let Some(idx) = found_idx {
        return BorrowFromIdle::Found(idle.swap_remove(idx));
    }
    idle.retain(|c| health_check(&c.meta, config));
    BorrowFromIdle::Exhausted
}

fn return_client(
    idle: &mut Vec<PooledClient>,
    client: Box<dyn LLMClient>,
    config: &PoolConfig,
) -> ReturnDisposition {
    if idle.len() < config.max_connections_per_provider {
        idle.push(PooledClient {
            client,
            meta: PooledClientMeta::new(),
        });
        ReturnDisposition::Returned
    } else {
        ReturnDisposition::Dropped
    }
}

fn health_check(meta: &PooledClientMeta, config: &PoolConfig) -> bool {
    if !config.enable_health_check {
        return true;
    }
    !meta.is_stale(config)
}

fn pool_stats(
    available: usize,
    in_use: usize,
    total_created: u64,
    max_size: usize,
    borrow_count: u64,
    error_count: u64,
) -> ProviderPoolStats {
    ProviderPoolStats {
        available,
        in_use,
        total: available.saturating_add(in_use),
        total_created,
        max_size,
        borrow_count,
        error_count,
    }
}

fn validate_pooled_client(client: &dyn LLMClient) -> std::result::Result<(), PoolError> {
    if client.model_name().trim().is_empty() {
        return Err(PoolError::InvalidClient {
            reason: "empty model name".to_string(),
        });
    }
    Ok(())
}

/// A pooled LLM client with its metadata
struct PooledClient {
    client: Box<dyn LLMClient>,
    meta: PooledClientMeta,
}

impl std::fmt::Debug for PooledClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PooledClient")
            .field("meta", &self.meta)
            .finish()
    }
}

/// Pool of clients for a single provider
#[derive(Debug)]
struct ProviderPool {
    /// The provider configuration for creating new clients
    provider: Provider,
    /// Pool of available clients
    clients: Mutex<Vec<PooledClient>>,
    /// Semaphore to limit concurrent connections
    semaphore: Arc<Semaphore>,
    /// Number of clients currently in use
    in_use_count: AtomicUsize,
    /// Total number of clients created (for stats)
    total_created: AtomicU64,
    borrow_count: AtomicU64,
    error_count: AtomicU64,
    /// Configuration reference
    config: PoolConfig,
}

impl ProviderPool {
    fn new(provider: Provider, config: PoolConfig) -> Self {
        let semaphore = Arc::new(Semaphore::new(config.max_connections_per_provider));
        Self {
            provider,
            clients: Mutex::new(Vec::with_capacity(config.max_connections_per_provider)),
            semaphore,
            in_use_count: AtomicUsize::new(0),
            total_created: AtomicU64::new(0),
            borrow_count: AtomicU64::new(0),
            error_count: AtomicU64::new(0),
            config,
        }
    }

    /// Get an available client from the pool, or create a new one
    async fn acquire(&self) -> std::result::Result<(Box<dyn LLMClient>, OwnedSemaphorePermit), PoolError> {
        let permit = match tokio::time::timeout(
            self.config.acquire_timeout,
            self.semaphore.clone().acquire_owned(),
        )
        .await
        {
            Ok(Ok(permit)) => permit,
            Ok(Err(_)) => {
                self.error_count.fetch_add(1, Ordering::Relaxed);
                return Err(PoolError::PoolExhausted {
                    max: self.config.max_connections_per_provider,
                });
            }
            Err(_) => {
                self.error_count.fetch_add(1, Ordering::Relaxed);
                return Err(PoolError::Timeout {
                    timeout_ms: self.config.acquire_timeout.as_millis() as u64,
                });
            }
        };

        self.borrow_count.fetch_add(1, Ordering::Relaxed);

        let client = {
            let borrowed = {
                let mut clients = self.clients.lock();
                match borrow_client(&mut clients, &self.config) {
                    BorrowFromIdle::Found(mut pooled) => {
                        validate_pooled_client(pooled.client.as_ref())?;
                        pooled.meta.mark_used();
                        Ok(pooled.client)
                    }
                    BorrowFromIdle::Exhausted => Err(()),
                }
            };
            match borrowed {
                Ok(client) => client,
                Err(()) => {
                    self.total_created.fetch_add(1, Ordering::Relaxed);
                    let created = self.provider.create_client().await.map_err(|e| {
                        self.error_count.fetch_add(1, Ordering::Relaxed);
                        PoolError::InvalidClient {
                            reason: e.to_string(),
                        }
                    })?;
                    validate_pooled_client(created.as_ref())?;
                    created
                }
            }
        };

        self.in_use_count.fetch_add(1, Ordering::Relaxed);
        Ok((client, permit))
    }

    async fn try_acquire(&self) -> std::result::Result<(Box<dyn LLMClient>, OwnedSemaphorePermit), PoolError> {
        let permit = self.semaphore.clone().try_acquire_owned().map_err(|_| {
            self.error_count.fetch_add(1, Ordering::Relaxed);
            PoolError::PoolExhausted {
                max: self.config.max_connections_per_provider,
            }
        })?;

        self.borrow_count.fetch_add(1, Ordering::Relaxed);

        let client = {
            let borrowed = {
                let mut clients = self.clients.lock();
                match borrow_client(&mut clients, &self.config) {
                    BorrowFromIdle::Found(mut pooled) => {
                        validate_pooled_client(pooled.client.as_ref())?;
                        pooled.meta.mark_used();
                        Ok(pooled.client)
                    }
                    BorrowFromIdle::Exhausted => Err(()),
                }
            };
            match borrowed {
                Ok(client) => client,
                Err(()) => {
                    self.total_created.fetch_add(1, Ordering::Relaxed);
                    let created = self.provider.create_client().await.map_err(|e| {
                        self.error_count.fetch_add(1, Ordering::Relaxed);
                        PoolError::InvalidClient {
                            reason: e.to_string(),
                        }
                    })?;
                    validate_pooled_client(created.as_ref())?;
                    created
                }
            }
        };

        self.in_use_count.fetch_add(1, Ordering::Relaxed);
        Ok((client, permit))
    }

    /// Return a client to the pool
    fn release(&self, client: Box<dyn LLMClient>) {
        self.in_use_count.fetch_sub(1, Ordering::Relaxed);
        let mut clients = self.clients.lock();
        let _ = return_client(&mut clients, client, &self.config);
    }

    /// Remove stale connections from the pool
    fn cleanup_stale(&self) -> usize {
        let mut clients = self.clients.lock();
        let before = clients.len();
        clients.retain(|c| !c.meta.is_stale(&self.config));
        before - clients.len()
    }

    /// Get pool statistics
    fn stats(&self) -> ProviderPoolStats {
        let clients = self.clients.lock();
        pool_stats(
            clients.len(),
            self.in_use_count.load(Ordering::Relaxed),
            self.total_created.load(Ordering::Relaxed),
            self.config.max_connections_per_provider,
            self.borrow_count.load(Ordering::Relaxed),
            self.error_count.load(Ordering::Relaxed),
        )
    }

    /// Drain all connections (for shutdown)
    fn drain(&self) {
        let mut clients = self.clients.lock();
        clients.clear();
    }
}

/// Statistics for a provider pool
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ProviderPoolStats {
    pub available: usize,
    pub in_use: usize,
    pub total: usize,
    pub total_created: u64,
    pub max_size: usize,
    pub borrow_count: u64,
    pub error_count: u64,
}

impl std::fmt::Display for ProviderPoolStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "idle={} active={} total={} created={} max={} borrows={} errors={}",
            self.available,
            self.in_use,
            self.total,
            self.total_created,
            self.max_size,
            self.borrow_count,
            self.error_count
        )
    }
}

/// Overall pool statistics
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct PoolStats {
    pub providers: HashMap<String, ProviderPoolStats>,
    pub total_available: usize,
    pub total_in_use: usize,
    pub total_connections: usize,
    pub borrow_count: u64,
    pub error_count: u64,
}

impl std::fmt::Display for PoolStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "providers={} idle={} active={} total={} borrows={} errors={}",
            self.providers.len(),
            self.total_available,
            self.total_in_use,
            self.total_connections,
            self.borrow_count,
            self.error_count
        )
    }
}

/// Guard that returns a client to the pool when dropped
pub struct PooledClientGuard {
    client: Option<Box<dyn LLMClient>>,
    pool: Arc<ProviderPool>,
    _permit: OwnedSemaphorePermit,
}

impl std::fmt::Debug for PooledClientGuard {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PooledClientGuard")
            .field("has_client", &self.client.is_some())
            .field("pool", &self.pool)
            .finish()
    }
}

impl PooledClientGuard {
    /// Get a reference to the underlying client
    pub fn client(&self) -> &dyn LLMClient {
        self.client.as_ref().expect("Client already taken").as_ref()
    }

    /// Get a mutable reference to the underlying client
    pub fn client_mut(&mut self) -> &mut dyn LLMClient {
        self.client.as_mut().expect("Client already taken").as_mut()
    }

    /// Take ownership of the client, preventing it from being returned to the pool
    ///
    /// This is useful if you need to move the client elsewhere, but be aware that
    /// it won't be returned to the pool.
    pub fn take(mut self) -> Box<dyn LLMClient> {
        self.client.take().expect("Client already taken")
    }
}

impl Drop for PooledClientGuard {
    fn drop(&mut self) {
        if let Some(client) = self.client.take() {
            self.pool.release(client);
        }
    }
}

impl std::ops::Deref for PooledClientGuard {
    type Target = Box<dyn LLMClient>;

    fn deref(&self) -> &Self::Target {
        self.client.as_ref().expect("Client already taken")
    }
}


#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct LLMPoolSnapshot {
    pub config: PoolConfig,
    pub stats: PoolStats,
    pub providers: Vec<String>,
    pub shutdown: bool,
}

pub type LLMPool = ClientPool;

/// LLM Client Pool for managing reusable client connections
///
/// The pool maintains separate sub-pools for each registered provider,
/// allowing efficient reuse of HTTP connections and client state.
pub struct ClientPool {
    config: PoolConfig,
    providers: RwLock<HashMap<String, Arc<ProviderPool>>>,
    shutdown: std::sync::atomic::AtomicBool,
}

impl ClientPool {
    /// Create a new client pool with the given configuration
    pub fn new(config: PoolConfig) -> Self {
        Self {
            config,
            providers: RwLock::new(HashMap::new()),
            shutdown: std::sync::atomic::AtomicBool::new(false),
        }
    }

    /// Create a new client pool with default configuration
    pub fn with_defaults() -> Self {
        Self::new(PoolConfig::default())
    }

    /// Register a provider with the pool
    ///
    /// This creates a sub-pool for the given provider that will manage
    /// client instances for that provider.
    #[allow(unreachable_code, unused_variables)]
    pub fn register_provider(&self, name: &str, provider: Provider) {
        let pool = Arc::new(ProviderPool::new(provider, self.config.clone()));
        let mut providers = self.providers.write();
        providers.insert(name.to_string(), pool);
    }

    /// Check if a provider is registered
    pub fn has_provider(&self, name: &str) -> bool {
        self.providers.read().contains_key(name)
    }

    /// List all registered provider names
    pub fn provider_names(&self) -> Vec<String> {
        self.providers.read().keys().cloned().collect()
    }

    /// Get a client from the pool for the specified provider
    ///
    /// The returned guard will automatically return the client to the pool
    /// when dropped.
    pub async fn get(&self, provider_name: &str) -> Result<PooledClientGuard> {
        self.get_with_error(provider_name).await.map_err(Into::into)
    }

    pub async fn get_with_error(
        &self,
        provider_name: &str,
    ) -> std::result::Result<PooledClientGuard, PoolError> {
        if self.shutdown.load(Ordering::Relaxed) {
            return Err(PoolError::InvalidClient {
                reason: "pool is shutting down".to_string(),
            });
        }

        let pool = {
            let providers = self.providers.read();
            providers.get(provider_name).cloned().ok_or_else(|| PoolError::InvalidClient {
                reason: format!("provider '{provider_name}' not registered in pool"),
            })?
        };

        let (client, permit) = pool.acquire().await?;

        Ok(PooledClientGuard {
            client: Some(client),
            pool,
            _permit: permit,
        })
    }

    pub async fn try_get(
        &self,
        provider_name: &str,
    ) -> std::result::Result<PooledClientGuard, PoolError> {
        if self.shutdown.load(Ordering::Relaxed) {
            return Err(PoolError::InvalidClient {
                reason: "pool is shutting down".to_string(),
            });
        }

        let pool = {
            let providers = self.providers.read();
            providers.get(provider_name).cloned().ok_or_else(|| PoolError::InvalidClient {
                reason: format!("provider '{provider_name}' not registered in pool"),
            })?
        };

        let (client, permit) = pool.try_acquire().await?;

        Ok(PooledClientGuard {
            client: Some(client),
            pool,
            _permit: permit,
        })
    }

    /// Get pool statistics
    pub fn stats(&self) -> PoolStats {
        let providers = self.providers.read();
        let mut stats = PoolStats {
            providers: HashMap::new(),
            total_available: 0,
            total_in_use: 0,
            total_connections: 0,
            borrow_count: 0,
            error_count: 0,
        };

        for (name, pool) in providers.iter() {
            let provider_stats = pool.stats();
            stats.total_available += provider_stats.available;
            stats.total_in_use += provider_stats.in_use;
            stats.total_connections += provider_stats.total;
            stats.borrow_count += provider_stats.borrow_count;
            stats.error_count += provider_stats.error_count;
            stats.providers.insert(name.clone(), provider_stats);
        }

        stats
    }

    /// Clean up stale connections across all providers
    ///
    /// Returns the total number of connections removed.
    pub fn cleanup_stale(&self) -> usize {
        let providers = self.providers.read();
        providers.values().map(|p| p.cleanup_stale()).sum()
    }

    /// Start a background task that periodically cleans up stale connections
    ///
    /// The task runs until the pool is shut down.
    pub fn start_cleanup_task(self: &Arc<Self>) -> tokio::task::JoinHandle<()> {
        let pool = Arc::clone(self);
        let interval = pool.config.health_check_interval;

        tokio::spawn(async move {
            let mut interval_timer = tokio::time::interval(interval);
            loop {
                interval_timer.tick().await;

                if pool.shutdown.load(Ordering::Relaxed) {
                    break;
                }

                let removed = pool.cleanup_stale();
                if removed > 0 {
                    tracing::debug!("Pool cleanup: removed {} stale connections", removed);
                }
            }
        })
    }

    /// Gracefully shut down the pool
    ///
    /// This prevents new clients from being acquired and drains all existing
    /// connections.
    pub fn shutdown(&self) {
        self.shutdown.store(true, Ordering::Relaxed);

        let providers = self.providers.read();
        for pool in providers.values() {
            pool.drain();
        }
    }

    /// Check if the pool is shut down
    pub fn is_shutdown(&self) -> bool {
        self.shutdown.load(Ordering::Relaxed)
    }

    pub fn snapshot(&self) -> LLMPoolSnapshot {
        LLMPoolSnapshot {
            config: self.config.clone(),
            stats: self.stats(),
            providers: self.provider_names(),
            shutdown: self.is_shutdown(),
        }
    }
}


impl std::fmt::Display for ClientPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let snap = self.snapshot();
        write!(f, "LLMPool(shutdown={}, {})", snap.shutdown, snap.stats)
    }
}

impl Default for ClientPool {
    fn default() -> Self {
        Self::with_defaults()
    }
}

/// Builder for creating a `ClientPool` with registered providers
pub struct ClientPoolBuilder {
    config: PoolConfig,
    providers: Vec<(String, Provider)>,
}

impl ClientPoolBuilder {
    /// Create a new builder with default configuration
    pub fn new() -> Self {
        Self {
            config: PoolConfig::default(),
            providers: Vec::new(),
        }
    }

    /// Set the pool configuration
    pub fn config(mut self, config: PoolConfig) -> Self {
        self.config = config;
        self
    }

    /// Add a provider to the pool
    pub fn provider(mut self, name: impl Into<String>, provider: Provider) -> Self {
        self.providers.push((name.into(), provider));
        self
    }

    /// Build the client pool
    pub fn build(self) -> ClientPool {
        let pool = ClientPool::new(self.config);
        for (name, provider) in self.providers {
            pool.register_provider(&name, provider);
        }
        pool
    }

    /// Build the client pool wrapped in an Arc
    pub fn build_arc(self) -> Arc<ClientPool> {
        Arc::new(self.build())
    }
}

impl Default for ClientPoolBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::client::test_support::MockLLMClient;
    use std::sync::atomic::Ordering as AtomicOrdering;
    fn mock_client(model: &str) -> Box<dyn LLMClient> {
        Box::new(MockLLMClient::new(model))
    }

    fn test_stub_provider() -> Provider {
        Provider::TestStub {
            model: "mock".to_string(),
        }
    }

    fn provider_pool(config: PoolConfig) -> Arc<ProviderPool> {
        Arc::new(ProviderPool::new(test_stub_provider(), config))
    }

    fn seed_pool(pool: &ProviderPool, clients: Vec<Box<dyn LLMClient>>) {
        let mut guard = pool.clients.lock();
        for client in clients {
            guard.push(PooledClient {
                client,
                meta: PooledClientMeta::new(),
            });
        }
    }

    fn register_seeded_pool(
        client_pool: &ClientPool,
        name: &str,
        config: PoolConfig,
        clients: Vec<Box<dyn LLMClient>>,
    ) {
        let sub = provider_pool(config);
        seed_pool(&sub, clients);
        client_pool.providers.write().insert(name.to_string(), sub);
    }

    #[test]
    fn test_pool_config_defaults() {
        let config = PoolConfig::default();
        assert_eq!(config.max_connections_per_provider, 10);
        assert_eq!(config.min_idle_connections, 2);
        assert_eq!(config.idle_timeout, Duration::from_secs(300));
        assert_eq!(config.max_lifetime, Duration::from_secs(1800));
        assert_eq!(config.health_check_interval, Duration::from_secs(60));
        assert_eq!(config.acquire_timeout, Duration::from_secs(30));
        assert!(config.enable_health_check);
    }

    #[test]
    fn test_pool_config_builder() {
        let config = PoolConfig::default()
            .with_max_connections(20)
            .with_idle_timeout(Duration::from_secs(60))
            .without_health_check();

        assert_eq!(config.max_connections_per_provider, 20);
        assert_eq!(config.idle_timeout, Duration::from_secs(60));
        assert!(!config.enable_health_check);
    }

    #[test]
    fn test_pool_config_with_max_lifetime() {
        let config = PoolConfig::default().with_max_lifetime(Duration::from_secs(42));
        assert_eq!(config.max_lifetime, Duration::from_secs(42));
    }

    #[test]
    fn test_pool_config_clone_preserves_fields() {
        let config = PoolConfig::default()
            .with_max_connections(7)
            .with_idle_timeout(Duration::from_secs(11))
            .with_max_lifetime(Duration::from_secs(22))
            .without_health_check();
        let cloned = config.clone();
        assert_eq!(cloned.max_connections_per_provider, 7);
        assert_eq!(cloned.idle_timeout, Duration::from_secs(11));
        assert_eq!(cloned.max_lifetime, Duration::from_secs(22));
        assert!(!cloned.enable_health_check);
    }

    #[test]
    fn test_pool_config_debug_format() {
        let debug = format!("{:?}", PoolConfig::default().with_max_connections(3));
        assert!(debug.contains("max_connections_per_provider"));
        assert!(debug.contains('3'));
    }

    #[test]
    fn test_pool_config_max_lifetime_stale() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(1))
            .with_max_lifetime(Duration::from_millis(5));
        let meta = PooledClientMeta::new();
        std::thread::sleep(Duration::from_millis(6));
        assert!(meta.is_stale(&config));
    }

    #[test]
    fn test_pooled_client_meta_stale_detection() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(10))
            .with_max_lifetime(Duration::from_millis(50));

        let meta = PooledClientMeta::new();
        assert!(!meta.is_stale(&config));

        std::thread::sleep(Duration::from_millis(15));
        assert!(meta.is_stale(&config));
    }

    #[test]
    fn test_pooled_client_meta_mark_used_increments() {
        let mut meta = PooledClientMeta::new();
        meta.mark_used();
        meta.mark_used();
        assert_eq!(meta.use_count.load(AtomicOrdering::Relaxed), 2);
    }

    #[test]
    fn test_pooled_client_meta_mark_used_refreshes_idle_timer() {
        let config = PoolConfig::default().with_idle_timeout(Duration::from_millis(30));
        let mut meta = PooledClientMeta::new();
        std::thread::sleep(Duration::from_millis(20));
        meta.mark_used();
        assert!(!meta.is_stale(&config));
    }

    #[test]
    fn test_pooled_client_debug_output() {
        let pooled = PooledClient {
            client: mock_client("debug-model"),
            meta: PooledClientMeta::new(),
        };
        let debug = format!("{pooled:?}");
        assert!(debug.contains("PooledClient"));
        assert!(debug.contains("meta"));
    }

    #[test]
    fn test_provider_pool_stats_clone_and_debug() {
        let stats = pool_stats(2, 1, 5, 10, 3, 1);
        let cloned = stats.clone();
        assert_eq!(cloned.available, 2);
        assert_eq!(cloned.in_use, 1);
        let debug = format!("{stats:?}");
        assert!(debug.contains("available"));
    }

    #[test]
    fn test_pool_stats_clone_and_debug() {
        let mut providers = HashMap::new();
        providers.insert(
            "p1".to_string(),
            pool_stats(1, 0, 1, 3, 0, 0),
        );
        let stats = PoolStats {
            providers,
            total_available: 1,
            total_in_use: 0,
            total_connections: 1,
            borrow_count: 0,
            error_count: 0,
        };
        let cloned = stats.clone();
        assert_eq!(cloned.total_available, 1);
        assert_eq!(cloned.providers.len(), 1);
        assert!(format!("{stats:?}").contains("total_in_use"));
    }

    #[test]
    fn test_client_pool_default_impl() {
        let pool = ClientPool::default();
        assert!(!pool.is_shutdown());
        assert_eq!(pool.stats().total_available, 0);
    }

    #[test]
    fn test_client_pool_builder_default_impl() {
        let builder = ClientPoolBuilder::default();
        let pool = builder.build();
        assert!(!pool.has_provider("anything"));
    }

    #[test]
    fn test_builder_empty_build() {
        let pool = ClientPoolBuilder::new().build();
        assert!(pool.provider_names().is_empty());
        assert!(!pool.has_provider("missing"));
    }

    #[test]
    fn test_builder_build_arc() {
        let pool = ClientPoolBuilder::new()
            .config(PoolConfig::default().with_max_connections(4))
            .build_arc();
        assert_eq!(pool.stats().total_available, 0);
        assert!(!pool.is_shutdown());
    }

    #[test]
    fn test_pool_stats() {
        let pool = ClientPool::with_defaults();
        let stats = pool.stats();
        assert_eq!(stats.total_available, 0);
        assert_eq!(stats.total_in_use, 0);
        assert!(stats.providers.is_empty());
    }

    #[test]
    fn test_cleanup_stale_on_empty_pool() {
        let pool = ClientPool::with_defaults();
        assert_eq!(pool.cleanup_stale(), 0);
    }

    #[test]
    fn test_pool_shutdown() {
        let pool = ClientPool::with_defaults();
        assert!(!pool.is_shutdown());
        pool.shutdown();
        assert!(pool.is_shutdown());
    }

    #[test]
    fn test_pool_double_shutdown_is_safe() {
        let pool = ClientPool::with_defaults();
        pool.shutdown();
        pool.shutdown();
        assert!(pool.is_shutdown());
    }

    #[test]
    fn test_provider_registration() {
        let pool = ClientPool::with_defaults();
        pool.register_provider("ollama", test_stub_provider());
        assert!(pool.has_provider("ollama"));
        assert!(!pool.has_provider("openai"));
        assert_eq!(pool.provider_names(), vec!["ollama"]);
    }

    #[test]
    fn test_builder_pattern() {
        let pool = ClientPoolBuilder::new()
            .config(PoolConfig::default().with_max_connections(5))
            .provider("ollama", test_stub_provider())
            .build();
        assert!(pool.has_provider("ollama"));
    }

    #[test]
    fn test_builder_multiple_providers() {
        let pool = ClientPoolBuilder::new()
            .provider("a", test_stub_provider())
            .provider("b", test_stub_provider())
            .build();
        let mut names = pool.provider_names();
        names.sort();
        assert_eq!(names, vec!["a", "b"]);
    }

    #[tokio::test]
    async fn test_get_unregistered_provider_error() {
        let pool = ClientPool::with_defaults();
        let result = pool.get("nonexistent").await;
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), AppError::LLM(_)));
    }

    #[tokio::test]
    async fn test_acquire_reuses_seeded_mock_without_network() {
        let config = PoolConfig::default()
            .with_max_connections(2)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("seeded")]);

        let guard = pool.get("mock").await.expect("seeded client should be available");
        assert_eq!(guard.client().model_name(), "seeded");
        assert_eq!(pool.stats().providers["mock"].total_created, 0);
    }

    #[tokio::test]
    async fn test_acquire_skips_stale_prefers_first_fresh() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(5))
            .with_max_lifetime(Duration::from_secs(60));
        let sub = provider_pool(config);

        {
            let mut guard = sub.clients.lock();
            guard.push(PooledClient {
                client: mock_client("stale"),
                meta: PooledClientMeta::new(),
            });
            std::thread::sleep(Duration::from_millis(8));
            guard.push(PooledClient {
                client: mock_client("fresh"),
                meta: PooledClientMeta::new(),
            });
        }

        let (client, permit) = sub.acquire().await.expect("acquire fresh client");
        assert_eq!(client.model_name(), "fresh");
        drop(client);
        drop(permit);
    }

    #[tokio::test]
    async fn test_release_and_reacquire_reuses_pooled_client() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("reusable")]);

        {
            let guard = pool.get("mock").await.expect("first acquire");
            assert_eq!(guard.client().model_name(), "reusable");
        }

        let guard = pool.get("mock").await.expect("second acquire");
        assert_eq!(guard.client().model_name(), "reusable");
        let stats = pool.stats().providers["mock"].clone();
        assert_eq!(stats.total_created, 0);
        assert_eq!(stats.in_use, 1);
    }

    #[tokio::test]
    async fn test_release_drops_client_when_idle_pool_full() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let sub = provider_pool(config);
        seed_pool(&sub, vec![mock_client("idle-0")]);

        let (client, permit) = sub.acquire().await.expect("acquire seeded client");
        sub.release(client);
        drop(permit);
        assert_eq!(sub.stats().available, 1);
        assert_eq!(sub.stats().in_use, 0);

        let (overflow, permit2) = sub.acquire().await.expect("acquire again");
        sub.release(overflow);
        drop(permit2);

        let stats = sub.stats();
        assert_eq!(stats.available, 1);
        assert_eq!(stats.in_use, 0);
        assert_eq!(stats.total_created, 0);
    }

    #[tokio::test]
    async fn test_provider_pool_in_use_accounting() {
        let config = PoolConfig::default()
            .with_max_connections(2)
            .without_health_check();
        let sub = provider_pool(config);
        seed_pool(&sub, vec![mock_client("a")]);

        let (_client, permit) = sub.acquire().await.expect("acquire");
        let stats = sub.stats();
        assert_eq!(stats.in_use, 1);
        assert_eq!(stats.available, 0);
        drop(permit);
    }

    #[tokio::test]
    async fn test_cleanup_stale_removes_idle_clients() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(5))
            .without_health_check();
        let sub = provider_pool(config);
        {
            let mut guard = sub.clients.lock();
            guard.push(PooledClient {
                client: mock_client("old"),
                meta: PooledClientMeta::new(),
            });
            std::thread::sleep(Duration::from_millis(8));
        }
        let removed = sub.cleanup_stale();
        assert_eq!(removed, 1);
        assert_eq!(sub.stats().available, 0);
    }

    #[tokio::test]
    async fn test_pooled_client_guard_debug_and_deref() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("guard")]);

        let guard = pool.get("mock").await.expect("guard acquire");
        let debug = format!("{guard:?}");
        assert!(debug.contains("PooledClientGuard"));
        assert!(debug.contains("has_client"));
        assert_eq!(guard.model_name(), "guard");
    }

    #[test]
    fn test_register_provider_overwrites_existing_name() {
        let pool = ClientPool::with_defaults();
        pool.register_provider("ollama", test_stub_provider());
        pool.register_provider("ollama", test_stub_provider());
        assert_eq!(pool.provider_names(), vec!["ollama"]);
    }

    #[tokio::test]
    async fn test_stats_aggregate_multiple_providers() {
        let config = PoolConfig::default()
            .with_max_connections(2)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "a", config.clone(), vec![mock_client("a1")]);
        register_seeded_pool(&pool, "b", config, vec![mock_client("b1")]);

        let _ga = pool.get("a").await.expect("provider a");
        let stats = pool.stats();
        assert_eq!(stats.providers.len(), 2);
        assert_eq!(stats.total_in_use, 1);
        assert_eq!(stats.total_available, 1);
    }

    #[tokio::test]
    async fn test_shutdown_drains_seeded_clients() {
        let config = PoolConfig::default()
            .with_max_connections(2)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("seeded")]);
        assert_eq!(pool.stats().total_available, 1);

        pool.shutdown();
        assert!(pool.is_shutdown());
        assert_eq!(pool.stats().total_available, 0);
    }


    #[tokio::test]
    async fn test_acquire_creates_client_when_pool_empty() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let sub = provider_pool(config);

        let (client, permit) = sub.acquire().await.expect("create via TestStub");
        assert_eq!(client.model_name(), "mock");
        drop(client);
        drop(permit);
        assert_eq!(sub.stats().total_created, 1);
    }

    #[tokio::test]
    async fn test_pooled_client_guard_client_mut_and_take() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("guard-mut")]);

        let mut guard = pool.get("mock").await.expect("guard acquire");
        assert_eq!(guard.client_mut().model_name(), "guard-mut");
        let taken = guard.take();
        assert_eq!(taken.model_name(), "guard-mut");
    }

    fn serde_roundtrip<T>(value: &T) -> T
    where
        T: serde::Serialize + for<'de> serde::Deserialize<'de> + PartialEq + std::fmt::Debug,
    {
        let json = serde_json::to_string(value).unwrap();
        let parsed: T = serde_json::from_str(&json).unwrap();
        assert_eq!(*value, parsed);
        parsed
    }

    #[test]
    fn test_pool_config_serde_roundtrip() {
        let config = PoolConfig::default()
            .with_max_connections(4)
            .with_idle_timeout(Duration::from_secs(90))
            .without_health_check();
        serde_roundtrip(&config);
    }

    #[test]
    fn test_pool_stats_serde_roundtrip() {
        let mut providers = HashMap::new();
        providers.insert("mock".into(), pool_stats(1, 2, 3, 4, 5, 6));
        let stats = PoolStats {
            providers,
            total_available: 1,
            total_in_use: 2,
            total_connections: 3,
            borrow_count: 5,
            error_count: 6,
        };
        serde_roundtrip(&stats);
    }

    #[test]
    fn test_provider_pool_stats_serde_roundtrip() {
        serde_roundtrip(&pool_stats(2, 1, 9, 10, 7, 2));
    }

    #[test]
    fn test_llm_pool_snapshot_serde_roundtrip() {
        let pool = ClientPoolBuilder::new()
            .provider("mock", test_stub_provider())
            .build();
        serde_roundtrip(&pool.snapshot());
    }

    #[test]
    fn test_llm_pool_type_alias() {
        let pool: LLMPool = ClientPool::with_defaults();
        assert_eq!(pool.stats().total_available, 0);
    }

    #[test]
    fn test_pool_error_serde_roundtrip_exhausted() {
        serde_roundtrip(&PoolError::PoolExhausted { max: 3 });
    }

    #[test]
    fn test_pool_error_serde_roundtrip_timeout() {
        serde_roundtrip(&PoolError::Timeout { timeout_ms: 250 });
    }

    #[test]
    fn test_pool_error_serde_roundtrip_invalid_client() {
        serde_roundtrip(&PoolError::InvalidClient {
            reason: "bad".into(),
        });
    }

    #[test]
    fn test_pool_error_display_variants() {
        assert!(PoolError::PoolExhausted { max: 2 }
            .to_string()
            .contains("pool exhausted"));
        assert!(PoolError::Timeout { timeout_ms: 10 }
            .to_string()
            .contains("timeout"));
        assert!(PoolError::InvalidClient {
            reason: "x".into()
        }
        .to_string()
        .contains("invalid"));
    }

    #[test]
    fn test_pool_error_clone_debug() {
        let err = PoolError::PoolExhausted { max: 1 };
        assert_eq!(err, err.clone());
        assert!(format!("{err:?}").contains("PoolExhausted"));
    }

    #[test]
    fn test_pool_config_display() {
        let s = PoolConfig::default().with_max_connections(6).to_string();
        assert!(s.contains("max=6"));
    }

    #[test]
    fn test_pool_stats_display() {
        let stats = PoolStats {
            providers: HashMap::new(),
            total_available: 1,
            total_in_use: 2,
            total_connections: 3,
            borrow_count: 4,
            error_count: 5,
        };
        let s = stats.to_string();
        assert!(s.contains("idle=1"));
        assert!(s.contains("errors=5"));
    }

    #[test]
    fn test_provider_pool_stats_display() {
        let s = pool_stats(1, 2, 0, 4, 9, 1).to_string();
        assert!(s.contains("borrows=9"));
        assert!(s.contains("total=3"));
    }

    #[test]
    fn test_client_pool_display() {
        let pool = ClientPool::with_defaults();
        let s = pool.to_string();
        assert!(s.contains("LLMPool"));
        assert!(s.contains("shutdown=false"));
    }

    #[test]
    fn test_borrow_client_returns_first_healthy() {
        let config = PoolConfig::default().without_health_check();
        let mut idle = vec![PooledClient {
            client: mock_client("a"),
            meta: PooledClientMeta::new(),
        }];
        match borrow_client(&mut idle, &config) {
            BorrowFromIdle::Found(p) => assert_eq!(p.client.model_name(), "a"),
            _ => panic!("expected found"),
        }
        assert!(idle.is_empty());
    }

    #[test]
    fn test_borrow_client_purges_stale_entries() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(1));
        let mut idle = vec![PooledClient {
            client: mock_client("stale"),
            meta: PooledClientMeta::new(),
        }];
        std::thread::sleep(Duration::from_millis(3));
        assert!(matches!(borrow_client(&mut idle, &config), BorrowFromIdle::Exhausted));
        assert!(idle.is_empty());
    }

    #[test]
    fn test_return_client_respects_capacity() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let mut idle = vec![];
        assert_eq!(
            return_client(&mut idle, mock_client("a"), &config),
            ReturnDisposition::Returned
        );
        assert_eq!(
            return_client(&mut idle, mock_client("b"), &config),
            ReturnDisposition::Dropped
        );
        assert_eq!(idle.len(), 1);
    }

    #[test]
    fn test_health_check_disabled_ignores_stale_meta() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(1))
            .without_health_check();
        let meta = PooledClientMeta::new();
        std::thread::sleep(Duration::from_millis(3));
        assert!(health_check(&meta, &config));
    }

    #[test]
    fn test_health_check_enabled_rejects_stale_meta() {
        let config = PoolConfig::default().with_idle_timeout(Duration::from_millis(1));
        let meta = PooledClientMeta::new();
        std::thread::sleep(Duration::from_millis(3));
        assert!(!health_check(&meta, &config));
    }

    #[test]
    fn test_pool_stats_helper_totals() {
        let stats = pool_stats(2, 3, 10, 8, 4, 1);
        assert_eq!(stats.total, 5);
        assert_eq!(stats.borrow_count, 4);
        assert_eq!(stats.error_count, 1);
    }

    #[test]
    fn test_validate_pooled_client_rejects_empty_model() {
        let client = mock_client("");
        let err = validate_pooled_client(client.as_ref()).unwrap_err();
        assert!(matches!(err, PoolError::InvalidClient { .. }));
    }

    #[test]
    fn test_validate_pooled_client_accepts_named_model() {
        validate_pooled_client(mock_client("ok").as_ref()).unwrap();
    }

    #[tokio::test]
    async fn test_try_get_pool_exhausted_when_at_capacity() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("only")]);

        let _guard = pool.try_get("mock").await.expect("first borrow");
        let err = pool.try_get("mock").await.unwrap_err();
        assert!(matches!(err, PoolError::PoolExhausted { max: 1 }));
        assert_eq!(pool.stats().providers["mock"].error_count, 1);
    }

    #[tokio::test]
    async fn test_acquire_timeout_increments_error_count() {
        let mut config = PoolConfig::default()
            .with_max_connections(1)
            .with_idle_timeout(Duration::from_secs(60))
            .without_health_check();
        config.acquire_timeout = Duration::from_millis(50);
        let sub = provider_pool(config);
        seed_pool(&sub, vec![mock_client("held")]);

        let (_c, permit) = sub.acquire().await.expect("hold permit");
        let err = match sub.acquire().await {
            Err(e) => e,
            Ok(_) => panic!("expected pool acquire timeout"),
        };
        assert!(matches!(err, PoolError::Timeout { .. }));
        assert_eq!(sub.stats().error_count, 1);
        drop(permit);
    }

    #[tokio::test]
    async fn test_borrow_count_increments_on_success() {
        let config = PoolConfig::default()
            .with_max_connections(2)
            .without_health_check();
        let sub = provider_pool(config);
        seed_pool(&sub, vec![mock_client("x")]);
        let (_c, permit) = sub.acquire().await.unwrap();
        drop(permit);
        assert_eq!(sub.stats().borrow_count, 1);
    }

    #[tokio::test]
    async fn test_concurrent_borrows_respect_max_connections() {
        let config = PoolConfig::default()
            .with_max_connections(2)
            .without_health_check();
        let pool = Arc::new(ClientPool::new(config.clone()));
        register_seeded_pool(&pool, "mock", config, vec![mock_client("c1"), mock_client("c2")]);

        let (tx, mut rx) = tokio::sync::oneshot::channel();
        let pool_bg = Arc::clone(&pool);
        tokio::spawn(async move {
            let result = pool_bg.get("mock").await;
            let _ = tx.send(result);
        });

        let g1 = pool.get("mock").await.expect("first");
        let g2 = pool.get("mock").await.expect("second");
        assert!(
            tokio::time::timeout(Duration::from_millis(100), &mut rx)
                .await
                .is_err(),
            "third borrow should still be waiting"
        );
        drop(g1);
        drop(g2);
        let third_result = rx.await.expect("channel").expect("third completes after release");
        drop(third_result);
    }

    #[tokio::test]
    async fn test_get_with_error_unregistered_invalid_client() {
        let pool = ClientPool::with_defaults();
        let err = pool.get_with_error("missing").await.unwrap_err();
        assert!(matches!(err, PoolError::InvalidClient { .. }));
    }

    #[tokio::test]
    async fn test_get_with_error_shutdown_invalid_client() {
        let pool = ClientPool::with_defaults();
        pool.shutdown();
        let err = pool.get_with_error("anything").await.unwrap_err();
        assert!(matches!(err, PoolError::InvalidClient { .. }));
    }

    #[tokio::test]
    async fn test_cleanup_stale_client_pool_aggregates() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(5))
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("old")]);
        std::thread::sleep(Duration::from_millis(8));
        assert_eq!(pool.cleanup_stale(), 1);
        assert_eq!(pool.stats().total_available, 0);
    }

    #[tokio::test]
    async fn test_stats_track_aggregate_borrow_and_error_counts() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("one")]);
        let _g = pool.try_get("mock").await.unwrap();
        let _ = pool.try_get("mock").await;
        let stats = pool.stats();
        assert_eq!(stats.borrow_count, 1);
        assert_eq!(stats.error_count, 1);
    }

    #[tokio::test]
    async fn test_acquire_removes_stale_before_creating_client() {
        let config = PoolConfig::default()
            .with_idle_timeout(Duration::from_millis(5));
        let sub = provider_pool(config);
        {
            let mut guard = sub.clients.lock();
            guard.push(PooledClient {
                client: mock_client("stale-only"),
                meta: PooledClientMeta::new(),
            });
            std::thread::sleep(Duration::from_millis(8));
        }
        let (client, permit) = sub.acquire().await.expect("creates fresh client");
        assert_eq!(client.model_name(), "mock");
        assert_eq!(sub.stats().total_created, 1);
        drop(client);
        drop(permit);
    }

    #[tokio::test]
    async fn test_race_release_then_reacquire() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = Arc::new(ClientPool::new(config.clone()));
        register_seeded_pool(&pool, "mock", config, vec![mock_client("race")]);

        let g1 = pool.get("mock").await.expect("first");
        let pool2 = Arc::clone(&pool);
        let j = tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(5)).await;
            pool2.get("mock").await
        });
        drop(g1);
        let g2 = j.await.expect("join").expect("second acquire after release");
        assert_eq!(g2.model_name(), "race");
    }

    #[test]
    fn test_return_disposition_debug_clone() {
        let d = ReturnDisposition::Returned;
        assert_eq!(d, d);
        assert!(format!("{d:?}").contains("Returned"));
    }

    #[test]
    fn test_pool_stats_clone_preserves_aggregate_fields() {
        let stats = PoolStats {
            providers: HashMap::new(),
            total_available: 0,
            total_in_use: 0,
            total_connections: 0,
            borrow_count: 2,
            error_count: 3,
        };
        let cloned = stats.clone();
        assert_eq!(cloned.borrow_count, 2);
        assert_eq!(cloned.error_count, 3);
    }

    #[tokio::test]
    async fn test_provider_pool_stats_after_release_shows_idle() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let sub = provider_pool(config);
        let (client, permit) = sub.acquire().await.unwrap();
        sub.release(client);
        drop(permit);
        let stats = sub.stats();
        assert_eq!(stats.available, 1);
        assert_eq!(stats.in_use, 0);
        assert_eq!(stats.total, 1);
    }




    #[test]
    fn test_pool_stats_default_aggregate_fields() {
        let stats = ClientPool::with_defaults().stats();
        assert_eq!(stats.total_connections, 0);
        assert_eq!(stats.borrow_count, 0);
        assert_eq!(stats.error_count, 0);
    }

    #[test]
    fn test_return_disposition_variants() {
        assert_ne!(ReturnDisposition::Returned, ReturnDisposition::Dropped);
    }

    #[tokio::test]
    async fn test_try_get_success_returns_guard() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("ok")]);
        let guard = pool.try_get("mock").await.expect("try_get ok");
        assert_eq!(guard.model_name(), "ok");
    }

    #[tokio::test]
    async fn test_get_with_error_success_path() {
        let config = PoolConfig::default()
            .with_max_connections(1)
            .without_health_check();
        let pool = ClientPool::new(config.clone());
        register_seeded_pool(&pool, "mock", config, vec![mock_client("ok2")]);
        let guard = pool.get_with_error("mock").await.expect("get ok");
        assert_eq!(guard.model_name(), "ok2");
    }

    #[test]
    fn test_llm_pool_snapshot_lists_providers() {
        let pool = ClientPoolBuilder::new()
            .provider("a", test_stub_provider())
            .provider("b", test_stub_provider())
            .build();
        let snap = pool.snapshot();
        let mut names = snap.providers;
        names.sort();
        assert_eq!(names, vec!["a", "b"]);
    }

    #[test]
    fn test_health_check_fresh_client_meta() {
        let config = PoolConfig::default().with_idle_timeout(Duration::from_secs(60));
        let meta = PooledClientMeta::new();
        assert!(health_check(&meta, &config));
    }

    #[tokio::test]
    async fn test_cleanup_stale_on_provider_pool() {
        let config = PoolConfig::default().with_idle_timeout(Duration::from_millis(5));
        let sub = provider_pool(config);
        {
            let mut guard = sub.clients.lock();
            guard.push(PooledClient {
                client: mock_client("gone"),
                meta: PooledClientMeta::new(),
            });
            std::thread::sleep(Duration::from_millis(8));
        }
        assert_eq!(sub.cleanup_stale(), 1);
    }

    #[tokio::test]
    async fn test_get_after_shutdown() {
        let pool = ClientPool::with_defaults();
        pool.shutdown();

        let result = pool.get("anything").await;
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), AppError::LLM(_)));
    }
}