claude_pool/

pool.rs

//! Core pool engine for managing Claude CLI slots.
//!
//! The [`Pool`] struct is the main entry point. It manages slot lifecycle,
//! task assignment, budget tracking, and shared context.
//!
//! # Example
//!
//! ```no_run
//! # async fn example() -> claude_pool::Result<()> {
//! use claude_pool::{Pool, PoolConfig};
//!
//! let claude = claude_wrapper::Claude::builder().build()?;
//! let pool = Pool::builder(claude)
//!     .slots(4)
//!     .build()
//!     .await?;
//!
//! let result = pool.run("write a haiku about rust").await?;
//! println!("{}", result.output);
//!
//! pool.drain().await?;
//! # Ok(())
//! # }
//! ```

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};

use tokio::sync::Mutex;

use claude_wrapper::Claude;
use claude_wrapper::types::OutputFormat;

use crate::config::ResolvedConfig;
use crate::error::{Error, Result};
use crate::skill::SkillRegistry;
use crate::store::PoolStore;
use crate::types::*;

/// Shared pool state behind an `Arc`.
struct PoolInner<S: PoolStore> {
    claude: Claude,
    config: PoolConfig,
    store: S,
    total_spend: AtomicU64,
    shutdown: AtomicBool,
    /// Context key-value pairs injected into slot system prompts.
    context: dashmap::DashMap<String, String>,
    /// Mutex for slot assignment to avoid races.
    assignment_lock: Mutex<()>,
    /// Worktree manager, if worktree isolation is enabled.
    worktree_manager: Option<crate::worktree::WorktreeManager>,
    /// In-flight chain progress, keyed by task ID.
    chain_progress: dashmap::DashMap<String, crate::chain::ChainProgress>,
}

/// A pool of Claude CLI slots.
///
/// Created via [`Pool::builder`]. Manages slot lifecycle, task routing,
/// and budget enforcement.
pub struct Pool<S: PoolStore> {
    inner: Arc<PoolInner<S>>,
}

// Manual Clone so we don't require S: Clone
impl<S: PoolStore> Clone for Pool<S> {
    fn clone(&self) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
        }
    }
}

/// Builder for constructing a [`Pool`].
pub struct PoolBuilder<S: PoolStore> {
    claude: Claude,
    slot_count: usize,
    config: PoolConfig,
    store: S,
    slot_configs: Vec<SlotConfig>,
}

impl<S: PoolStore + 'static> PoolBuilder<S> {
    /// Set the number of slots to spawn.
    pub fn slots(mut self, count: usize) -> Self {
        self.slot_count = count;
        self
    }

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

    /// Add a per-slot configuration override.
    ///
    /// Call multiple times for multiple slots. Slot configs are applied
    /// in order: the first call sets slot-0's config, the second slot-1's, etc.
    /// Slots without an explicit config get [`SlotConfig::default()`].
    pub fn slot_config(mut self, config: SlotConfig) -> Self {
        self.slot_configs.push(config);
        self
    }

    /// Build and initialize the pool, registering slots in the store.
    pub async fn build(self) -> Result<Pool<S>> {
        // Set up worktree manager if isolation is enabled.
        let worktree_manager = if self.config.worktree_isolation {
            let repo_dir = self
                .claude
                .working_dir()
                .map(|p| p.to_path_buf())
                .unwrap_or_else(|| std::env::current_dir().unwrap_or_default());
            Some(crate::worktree::WorktreeManager::new(repo_dir, None))
        } else {
            None
        };

        let inner = Arc::new(PoolInner {
            claude: self.claude,
            config: self.config,
            store: self.store,
            total_spend: AtomicU64::new(0),
            shutdown: AtomicBool::new(false),
            context: dashmap::DashMap::new(),
            assignment_lock: Mutex::new(()),
            worktree_manager,
            chain_progress: dashmap::DashMap::new(),
        });

        // Register slots in the store.
        for i in 0..self.slot_count {
            let slot_config = self.slot_configs.get(i).cloned().unwrap_or_default();

            let slot_id = SlotId(format!("slot-{i}"));

            // Create worktree if isolation is enabled.
            let worktree_path = if let Some(ref mgr) = inner.worktree_manager {
                let path = mgr.create(&slot_id).await?;
                Some(path.to_string_lossy().into_owned())
            } else {
                None
            };

            let record = SlotRecord {
                id: slot_id,
                state: SlotState::Idle,
                config: slot_config,
                current_task: None,
                session_id: None,
                tasks_completed: 0,
                cost_microdollars: 0,
                restart_count: 0,
                worktree_path,
            };
            inner.store.put_slot(record).await?;
        }

        Ok(Pool { inner })
    }
}

impl Pool<crate::store::InMemoryStore> {
    /// Create a builder with the default in-memory store.
    pub fn builder(claude: Claude) -> PoolBuilder<crate::store::InMemoryStore> {
        PoolBuilder {
            claude,
            slot_count: 1,
            config: PoolConfig::default(),
            store: crate::store::InMemoryStore::new(),
            slot_configs: Vec::new(),
        }
    }
}

impl<S: PoolStore + 'static> Pool<S> {
    /// Create a builder with a custom store.
    pub fn builder_with_store(claude: Claude, store: S) -> PoolBuilder<S> {
        PoolBuilder {
            claude,
            slot_count: 1,
            config: PoolConfig::default(),
            store,
            slot_configs: Vec::new(),
        }
    }

    /// Run a task synchronously, blocking until completion.
    ///
    /// Assigns the task to the next idle slot, executes the prompt,
    /// and returns the result.
    pub async fn run(&self, prompt: &str) -> Result<TaskResult> {
        self.run_with_config(prompt, None).await
    }

    /// Run a task with per-task config overrides.
    pub async fn run_with_config(
        &self,
        prompt: &str,
        task_config: Option<SlotConfig>,
    ) -> Result<TaskResult> {
        self.check_shutdown()?;
        self.check_budget()?;

        let task_id = TaskId(format!("task-{}", new_id()));

        let record = TaskRecord {
            id: task_id.clone(),
            prompt: prompt.to_string(),
            state: TaskState::Pending,
            slot_id: None,
            result: None,
            tags: vec![],
            config: task_config,
        };
        self.inner.store.put_task(record).await?;

        let (slot_id, slot_config) = self.assign_slot(&task_id).await?;
        let result = self
            .execute_task(&task_id, prompt, &slot_id, &slot_config)
            .await;

        self.release_slot(&slot_id, &task_id, &result).await?;

        let task_result = result?;
        // Update task record with result.
        let mut task = self
            .inner
            .store
            .get_task(&task_id)
            .await?
            .ok_or_else(|| Error::TaskNotFound(task_id.0.clone()))?;
        task.state = TaskState::Completed;
        task.result = Some(task_result.clone());
        self.inner.store.put_task(task).await?;

        Ok(task_result)
    }

    /// Submit a task for async execution, returning the task ID immediately.
    ///
    /// Use [`Pool::result`] to poll for completion.
    pub async fn submit(&self, prompt: &str) -> Result<TaskId> {
        self.submit_with_config(prompt, None, vec![]).await
    }

    /// Submit a task with config overrides and tags.
    pub async fn submit_with_config(
        &self,
        prompt: &str,
        task_config: Option<SlotConfig>,
        tags: Vec<String>,
    ) -> Result<TaskId> {
        self.check_shutdown()?;
        self.check_budget()?;

        let task_id = TaskId(format!("task-{}", new_id()));
        let prompt = prompt.to_string();

        let record = TaskRecord {
            id: task_id.clone(),
            prompt: prompt.clone(),
            state: TaskState::Pending,
            slot_id: None,
            result: None,
            tags,
            config: task_config,
        };
        self.inner.store.put_task(record).await?;

        // Spawn the task execution in the background.
        let pool = self.clone();
        let tid = task_id.clone();
        tokio::spawn(async move {
            let task = match pool.inner.store.get_task(&tid).await {
                Ok(Some(t)) => t,
                _ => return,
            };

            match pool.assign_slot(&tid).await {
                Ok((slot_id, slot_config)) => {
                    let result = pool
                        .execute_task(&tid, &prompt, &slot_id, &slot_config)
                        .await;

                    let _ = pool.release_slot(&slot_id, &tid, &result).await;

                    let mut updated = task;
                    match result {
                        Ok(task_result) => {
                            updated.state = TaskState::Completed;
                            updated.result = Some(task_result);
                        }
                        Err(e) => {
                            updated.state = TaskState::Failed;
                            updated.result = Some(TaskResult {
                                output: e.to_string(),
                                success: false,
                                cost_microdollars: 0,
                                turns_used: 0,
                                session_id: None,
                            });
                        }
                    }
                    let _ = pool.inner.store.put_task(updated).await;
                }
                Err(e) => {
                    let mut updated = task;
                    updated.state = TaskState::Failed;
                    updated.result = Some(TaskResult {
                        output: e.to_string(),
                        success: false,
                        cost_microdollars: 0,
                        turns_used: 0,
                        session_id: None,
                    });
                    let _ = pool.inner.store.put_task(updated).await;
                }
            }
        });

        Ok(task_id)
    }

    /// Get the result of a submitted task.
    ///
    /// Returns `None` if the task is still pending/running.
    pub async fn result(&self, task_id: &TaskId) -> Result<Option<TaskResult>> {
        let task = self
            .inner
            .store
            .get_task(task_id)
            .await?
            .ok_or_else(|| Error::TaskNotFound(task_id.0.clone()))?;

        match task.state {
            TaskState::Completed | TaskState::Failed => Ok(task.result),
            _ => Ok(None),
        }
    }

    /// Cancel a pending or running task.
    pub async fn cancel(&self, task_id: &TaskId) -> Result<()> {
        let mut task = self
            .inner
            .store
            .get_task(task_id)
            .await?
            .ok_or_else(|| Error::TaskNotFound(task_id.0.clone()))?;

        match task.state {
            TaskState::Pending => {
                task.state = TaskState::Cancelled;
                self.inner.store.put_task(task).await?;
                Ok(())
            }
            TaskState::Running => {
                // Mark as cancelled; the executing task will check on completion.
                task.state = TaskState::Cancelled;
                self.inner.store.put_task(task).await?;
                Ok(())
            }
            _ => Ok(()), // already terminal
        }
    }

    /// Execute tasks in parallel across available slots, collecting all results.
    ///
    /// Queues excess prompts until a slot becomes idle. Returns once all
    /// prompts complete or timeout waiting for slot availability.
    pub async fn fan_out(&self, prompts: &[&str]) -> Result<Vec<TaskResult>> {
        self.check_shutdown()?;
        self.check_budget()?;

        let mut handles = Vec::with_capacity(prompts.len());

        for prompt in prompts {
            let pool = self.clone();
            let prompt = prompt.to_string();
            handles.push(tokio::spawn(async move { pool.run(&prompt).await }));
        }

        let mut results = Vec::with_capacity(handles.len());
        for handle in handles {
            results.push(
                handle
                    .await
                    .map_err(|e| Error::Store(format!("task join error: {e}")))?,
            );
        }

        results.into_iter().collect()
    }

    /// Submit a chain for async execution, returning a task ID immediately.
    ///
    /// Use [`Pool::chain_progress`] to check per-step progress, or
    /// [`Pool::result`] to get the final [`crate::ChainResult`] (serialized as JSON)
    /// once complete.
    pub async fn submit_chain(
        &self,
        steps: Vec<crate::chain::ChainStep>,
        skills: &SkillRegistry,
        options: crate::chain::ChainOptions,
    ) -> Result<TaskId> {
        self.check_shutdown()?;
        self.check_budget()?;

        let task_id = TaskId(format!("chain-{}", new_id()));

        let record = TaskRecord {
            id: task_id.clone(),
            prompt: format!("chain: {} steps", steps.len()),
            state: TaskState::Pending,
            slot_id: None,
            result: None,
            tags: options.tags,
            config: None,
        };
        self.inner.store.put_task(record).await?;

        // Initialize progress.
        let progress = crate::chain::ChainProgress {
            total_steps: steps.len(),
            current_step: None,
            current_step_name: None,
            completed_steps: vec![],
            status: crate::chain::ChainStatus::Running,
        };
        self.inner
            .chain_progress
            .insert(task_id.0.clone(), progress);

        // Mark as running.
        if let Some(mut task) = self.inner.store.get_task(&task_id).await? {
            task.state = TaskState::Running;
            self.inner.store.put_task(task).await?;
        }

        let pool = self.clone();
        let tid = task_id.clone();
        let skills = skills.clone();
        tokio::spawn(async move {
            let result =
                crate::chain::execute_chain_with_progress(&pool, &skills, &steps, Some(&tid)).await;

            // Store the chain result as the task result.
            if let Some(mut task) = pool.inner.store.get_task(&tid).await.ok().flatten() {
                match result {
                    Ok(chain_result) => {
                        let success = chain_result.success;
                        task.state = if success {
                            TaskState::Completed
                        } else {
                            TaskState::Failed
                        };
                        task.result = Some(TaskResult {
                            output: serde_json::to_string(&chain_result).unwrap_or_default(),
                            success,
                            cost_microdollars: chain_result.total_cost_microdollars,
                            turns_used: 0,
                            session_id: None,
                        });
                    }
                    Err(e) => {
                        task.state = TaskState::Failed;
                        task.result = Some(TaskResult {
                            output: e.to_string(),
                            success: false,
                            cost_microdollars: 0,
                            turns_used: 0,
                            session_id: None,
                        });
                    }
                }
                let _ = pool.inner.store.put_task(task).await;
            }
        });

        Ok(task_id)
    }

    /// Submit multiple chains for parallel execution, returning all task IDs immediately.
    ///
    /// Each chain runs on its own slot concurrently. Use [`Pool::chain_progress`] to check
    /// per-step progress, or [`Pool::result`] to get the final result once complete.
    pub async fn fan_out_chains(
        &self,
        chains: Vec<Vec<crate::chain::ChainStep>>,
        skills: &SkillRegistry,
        options: crate::chain::ChainOptions,
    ) -> Result<Vec<TaskId>> {
        self.check_shutdown()?;
        self.check_budget()?;

        let mut handles = Vec::with_capacity(chains.len());

        for chain_steps in chains {
            let pool = self.clone();
            let skills = skills.clone();
            let options = options.clone();
            handles.push(tokio::spawn(async move {
                pool.submit_chain(chain_steps, &skills, options).await
            }));
        }

        let mut task_ids = Vec::with_capacity(handles.len());
        for handle in handles {
            match handle.await {
                Ok(Ok(task_id)) => task_ids.push(task_id),
                Ok(Err(e)) => {
                    // Log the error but continue collecting other task IDs
                    tracing::warn!("failed to submit chain: {}", e);
                }
                Err(e) => {
                    tracing::warn!("chain submission task panicked: {}", e);
                }
            }
        }

        Ok(task_ids)
    }

    /// Submit a workflow template for async execution.
    ///
    /// Instantiates the workflow by substituting placeholders with arguments,
    /// then submits the resulting chain. Returns the task ID immediately.
    pub async fn submit_workflow(
        &self,
        workflow_name: &str,
        arguments: std::collections::HashMap<String, String>,
        skills: &SkillRegistry,
        workflows: &crate::workflow::WorkflowRegistry,
        tags: Vec<String>,
    ) -> Result<TaskId> {
        // Get the workflow and instantiate it
        let workflow = workflows
            .get(workflow_name)
            .ok_or_else(|| Error::Store(format!("workflow '{}' not found", workflow_name)))?;

        let steps = workflow.instantiate(&arguments)?;

        // Submit the instantiated chain with tags
        let options = crate::chain::ChainOptions { tags };
        self.submit_chain(steps, skills, options).await
    }

    /// Get the progress of an in-flight chain.
    ///
    /// Returns `None` if no chain is tracked for this task ID.
    pub fn chain_progress(&self, task_id: &TaskId) -> Option<crate::chain::ChainProgress> {
        self.inner
            .chain_progress
            .get(&task_id.0)
            .map(|v| v.value().clone())
    }

    /// Store chain progress (called internally by `execute_chain_with_progress`).
    pub(crate) async fn set_chain_progress(
        &self,
        task_id: &TaskId,
        progress: crate::chain::ChainProgress,
    ) {
        self.inner
            .chain_progress
            .insert(task_id.0.clone(), progress);
    }

    /// Set a shared context value.
    ///
    /// Context is injected into slot system prompts at task start.
    pub fn set_context(&self, key: impl Into<String>, value: impl Into<String>) {
        self.inner.context.insert(key.into(), value.into());
    }

    /// Get a shared context value.
    pub fn get_context(&self, key: &str) -> Option<String> {
        self.inner.context.get(key).map(|v| v.value().clone())
    }

    /// Remove a shared context value.
    pub fn delete_context(&self, key: &str) -> Option<String> {
        self.inner.context.remove(key).map(|(_, v)| v)
    }

    /// List all context keys and values.
    pub fn list_context(&self) -> Vec<(String, String)> {
        self.inner
            .context
            .iter()
            .map(|r| (r.key().clone(), r.value().clone()))
            .collect()
    }

    /// Gracefully shut down the pool.
    ///
    /// Marks the pool as shut down so no new tasks are accepted,
    /// then waits for in-flight tasks to complete.
    pub async fn drain(&self) -> Result<DrainSummary> {
        self.inner.shutdown.store(true, Ordering::SeqCst);

        // Wait for all running tasks to finish.
        loop {
            let running = self
                .inner
                .store
                .list_tasks(&TaskFilter {
                    state: Some(TaskState::Running),
                    ..Default::default()
                })
                .await?;
            if running.is_empty() {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        }

        // Mark all slots as stopped.
        let slots = self.inner.store.list_slots().await?;
        let mut total_cost = 0u64;
        let mut total_tasks = 0u64;
        let slot_ids: Vec<_> = slots.iter().map(|w| w.id.clone()).collect();

        for mut slot in slots {
            total_cost += slot.cost_microdollars;
            total_tasks += slot.tasks_completed;
            slot.state = SlotState::Stopped;
            self.inner.store.put_slot(slot).await?;
        }

        // Clean up worktrees if isolation was enabled.
        if let Some(ref mgr) = self.inner.worktree_manager {
            mgr.cleanup_all(&slot_ids).await?;
        }

        Ok(DrainSummary {
            total_cost_microdollars: total_cost,
            total_tasks_completed: total_tasks,
        })
    }

    /// Get a snapshot of pool status.
    pub async fn status(&self) -> Result<PoolStatus> {
        let slots = self.inner.store.list_slots().await?;
        let idle = slots.iter().filter(|w| w.state == SlotState::Idle).count();
        let busy = slots.iter().filter(|w| w.state == SlotState::Busy).count();

        let running_tasks = self
            .inner
            .store
            .list_tasks(&TaskFilter {
                state: Some(TaskState::Running),
                ..Default::default()
            })
            .await?
            .len();

        let pending_tasks = self
            .inner
            .store
            .list_tasks(&TaskFilter {
                state: Some(TaskState::Pending),
                ..Default::default()
            })
            .await?
            .len();

        Ok(PoolStatus {
            total_slots: slots.len(),
            idle_slots: idle,
            busy_slots: busy,
            running_tasks,
            pending_tasks,
            total_spend_microdollars: self.inner.total_spend.load(Ordering::Relaxed),
            budget_microdollars: self.inner.config.budget_microdollars,
            shutdown: self.inner.shutdown.load(Ordering::Relaxed),
        })
    }

    /// Get a reference to the store.
    pub fn store(&self) -> &S {
        &self.inner.store
    }

    /// Scale up the pool by adding N new slots.
    ///
    /// Returns the new total slot count.
    /// Fails if the new count exceeds max_slots.
    pub async fn scale_up(&self, count: usize) -> Result<usize> {
        if count == 0 {
            return Ok(self.inner.store.list_slots().await?.len());
        }

        let current_slots = self.inner.store.list_slots().await?;
        let current_count = current_slots.len();
        let new_count = current_count + count;

        if new_count > self.inner.config.scaling.max_slots {
            return Err(Error::Store(format!(
                "cannot scale up to {} slots: exceeds max_slots ({})",
                new_count, self.inner.config.scaling.max_slots
            )));
        }

        // Find the next available slot ID.
        let existing_ids: Vec<usize> = current_slots
            .iter()
            .filter_map(|w| w.id.0.strip_prefix("slot-").and_then(|s| s.parse().ok()))
            .collect();
        let mut next_id = existing_ids.iter().max().unwrap_or(&0) + 1;

        // Create and register new slots.
        for _ in 0..count {
            let slot_id = SlotId(format!("slot-{next_id}"));
            next_id += 1;

            // Create worktree if isolation is enabled.
            let worktree_path = if let Some(ref mgr) = self.inner.worktree_manager {
                let path = mgr.create(&slot_id).await?;
                Some(path.to_string_lossy().into_owned())
            } else {
                None
            };

            let record = SlotRecord {
                id: slot_id,
                state: SlotState::Idle,
                config: SlotConfig::default(),
                current_task: None,
                session_id: None,
                tasks_completed: 0,
                cost_microdollars: 0,
                restart_count: 0,
                worktree_path,
            };
            self.inner.store.put_slot(record).await?;
        }

        Ok(new_count)
    }

    /// Scale down the pool by removing N slots.
    ///
    /// Removes idle slots first. If not enough idle slots are available,
    /// waits for busy slots to complete (with timeout) before removing them.
    /// Returns the new total slot count.
    /// Fails if the new count drops below min_slots.
    pub async fn scale_down(&self, count: usize) -> Result<usize> {
        if count == 0 {
            return Ok(self.inner.store.list_slots().await?.len());
        }

        let mut slots = self.inner.store.list_slots().await?;
        let current_count = slots.len();
        let new_count = current_count.saturating_sub(count);

        if new_count < self.inner.config.scaling.min_slots {
            return Err(Error::Store(format!(
                "cannot scale down to {} slots: below min_slots ({})",
                new_count, self.inner.config.scaling.min_slots
            )));
        }

        // Sort to prioritize removing least-active slots.
        slots.sort_by_key(|w| std::cmp::Reverse(w.tasks_completed));

        let slots_to_remove = &slots[..count];
        let timeout = std::time::Duration::from_secs(30);

        for slot in slots_to_remove {
            // Wait for slot to finish any running task (with timeout).
            let deadline = std::time::Instant::now() + timeout;
            loop {
                if let Some(w) = self.inner.store.get_slot(&slot.id).await? {
                    if w.state != SlotState::Busy {
                        break;
                    }
                    if std::time::Instant::now() >= deadline {
                        // Timeout: still busy, but proceed with removal anyway.
                        break;
                    }
                } else {
                    break;
                }
                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
            }

            // Cleanup worktree if applicable.
            if let Some(ref mgr) = self.inner.worktree_manager
                && slot.worktree_path.is_some()
            {
                let _ = mgr.cleanup_all(std::slice::from_ref(&slot.id)).await;
            }

            // Delete slot record.
            self.inner.store.delete_slot(&slot.id).await?;
        }

        Ok(new_count)
    }

    /// Set the target number of slots, scaling up or down as needed.
    pub async fn set_target_slots(&self, target: usize) -> Result<usize> {
        let current = self.inner.store.list_slots().await?.len();
        if target > current {
            self.scale_up(target - current).await
        } else if target < current {
            self.scale_down(current - target).await
        } else {
            Ok(current)
        }
    }

    // ── Internal helpers ─────────────────────────────────────────────

    fn check_shutdown(&self) -> Result<()> {
        if self.inner.shutdown.load(Ordering::SeqCst) {
            Err(Error::PoolShutdown)
        } else {
            Ok(())
        }
    }

    fn check_budget(&self) -> Result<()> {
        if let Some(limit) = self.inner.config.budget_microdollars {
            let spent = self.inner.total_spend.load(Ordering::Relaxed);
            if spent >= limit {
                return Err(Error::BudgetExhausted {
                    spent_microdollars: spent,
                    limit_microdollars: limit,
                });
            }
        }
        Ok(())
    }

    /// Wait for an idle slot to become available, with exponential backoff.
    async fn wait_for_idle_slot_with_timeout(&self, timeout_secs: u64) -> Result<SlotRecord> {
        use std::time::{Duration, Instant};

        let deadline = Instant::now() + Duration::from_secs(timeout_secs);
        let mut backoff_ms = 10u64;
        const MAX_BACKOFF_MS: u64 = 500;

        loop {
            self.check_shutdown()?;

            let slots = self.inner.store.list_slots().await?;
            for slot in slots {
                if slot.state == SlotState::Idle {
                    return Ok(slot);
                }
            }

            if Instant::now() >= deadline {
                return Err(Error::NoSlotAvailable { timeout_secs });
            }

            tokio::time::sleep(Duration::from_millis(backoff_ms)).await;
            backoff_ms = std::cmp::min((backoff_ms as f64 * 1.5) as u64, MAX_BACKOFF_MS);
        }
    }

    /// Find an idle slot and assign the task to it, waiting if necessary.
    async fn assign_slot(&self, task_id: &TaskId) -> Result<(SlotId, SlotConfig)> {
        let _lock = self.inner.assignment_lock.lock().await;

        let timeout = self.inner.config.slot_assignment_timeout_secs;
        let mut slot = self.wait_for_idle_slot_with_timeout(timeout).await?;
        let config = slot.config.clone();

        slot.state = SlotState::Busy;
        slot.current_task = Some(task_id.clone());
        self.inner.store.put_slot(slot.clone()).await?;

        // Update task with assigned slot.
        if let Some(mut task) = self.inner.store.get_task(task_id).await? {
            task.state = TaskState::Running;
            task.slot_id = Some(slot.id.clone());
            self.inner.store.put_task(task).await?;
        }

        Ok((slot.id, config))
    }

    /// Release a slot back to idle after task completion.
    async fn release_slot(
        &self,
        slot_id: &SlotId,
        _task_id: &TaskId,
        result: &std::result::Result<TaskResult, Error>,
    ) -> Result<()> {
        if let Some(mut slot) = self.inner.store.get_slot(slot_id).await? {
            slot.state = SlotState::Idle;
            slot.current_task = None;

            if let Ok(task_result) = result {
                slot.tasks_completed += 1;
                slot.cost_microdollars += task_result.cost_microdollars;
                slot.session_id = task_result.session_id.clone();

                // Update global spend tracker.
                self.inner
                    .total_spend
                    .fetch_add(task_result.cost_microdollars, Ordering::Relaxed);
            }

            self.inner.store.put_slot(slot).await?;
        }
        Ok(())
    }

    /// Execute a task on a specific slot by invoking the Claude CLI.
    async fn execute_task(
        &self,
        _task_id: &TaskId,
        prompt: &str,
        slot_id: &SlotId,
        slot_config: &SlotConfig,
    ) -> Result<TaskResult> {
        let task_record = self.inner.store.get_task(_task_id).await?;
        let task_cfg = task_record.as_ref().and_then(|t| t.config.as_ref());

        let resolved = ResolvedConfig::resolve(&self.inner.config, slot_config, task_cfg);

        // Build the system prompt with identity and context injection.
        let system_prompt = self.build_system_prompt(&resolved, slot_config);

        // Build and execute the query.
        let mut cmd = claude_wrapper::QueryCommand::new(prompt)
            .output_format(OutputFormat::Json)
            .permission_mode(resolved.permission_mode);

        if resolved.permission_mode == PermissionMode::BypassPermissions {
            cmd = cmd.dangerously_skip_permissions();
        }

        if let Some(ref model) = resolved.model {
            cmd = cmd.model(model);
        }
        if let Some(max_turns) = resolved.max_turns {
            cmd = cmd.max_turns(max_turns);
        }
        if let Some(ref sp) = system_prompt {
            cmd = cmd.system_prompt(sp);
        }
        if let Some(effort) = resolved.effort {
            cmd = cmd.effort(effort);
        }
        if !resolved.allowed_tools.is_empty() {
            cmd = cmd.allowed_tools(&resolved.allowed_tools);
        }

        // Use worktree working dir if the slot has one, otherwise use default.
        let claude_instance = if let Some(slot) = self.inner.store.get_slot(slot_id).await? {
            // Resume session if the slot has one.
            if let Some(ref session_id) = slot.session_id {
                cmd = cmd.resume(session_id);
            }

            if let Some(ref wt_path) = slot.worktree_path {
                self.inner.claude.with_working_dir(wt_path)
            } else {
                self.inner.claude.clone()
            }
        } else {
            self.inner.claude.clone()
        };

        tracing::debug!(
            slot_id = %slot_id.0,
            model = ?resolved.model,
            effort = ?resolved.effort,
            "executing task"
        );

        let query_result = cmd.execute_json(&claude_instance).await?;

        let cost_microdollars = query_result
            .cost_usd
            .map(|c| (c * 1_000_000.0) as u64)
            .unwrap_or(0);

        Ok(TaskResult {
            output: query_result.result,
            success: !query_result.is_error,
            cost_microdollars,
            turns_used: 0, // TODO: extract from query result when available
            session_id: Some(query_result.session_id),
        })
    }

    /// Build the system prompt by combining slot identity, resolved config and context.
    fn build_system_prompt(
        &self,
        resolved: &ResolvedConfig,
        slot_config: &SlotConfig,
    ) -> Option<String> {
        let context_entries: Vec<_> = self.list_context();

        // Check if there's any content to include
        let has_identity = slot_config.name.is_some()
            || slot_config.role.is_some()
            || slot_config.description.is_some();

        if resolved.system_prompt.is_none() && context_entries.is_empty() && !has_identity {
            return None;
        }

        let mut parts = Vec::new();

        // Inject slot identity
        if has_identity {
            let mut identity = String::new();
            identity.push_str("You are ");

            if let Some(ref name) = slot_config.name {
                identity.push_str(name);
            } else {
                identity.push_str("a slot");
            }

            if let Some(ref role) = slot_config.role {
                identity.push_str(", a ");
                identity.push_str(role);
            }

            if let Some(ref description) = slot_config.description {
                identity.push_str(". ");
                identity.push_str(description);
            } else if slot_config.role.is_some() {
                identity.push('.');
            }

            parts.push(identity);
        }

        if let Some(ref sp) = resolved.system_prompt {
            parts.push(sp.clone());
        }

        if !context_entries.is_empty() {
            parts.push("\n\n## Shared Context\n".to_string());
            for (key, value) in &context_entries {
                parts.push(format!("- **{key}**: {value}"));
            }
        }

        Some(parts.join("\n"))
    }
}

/// Summary returned by [`Pool::drain`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DrainSummary {
    /// Total cost across all slots in microdollars.
    pub total_cost_microdollars: u64,
    /// Total number of tasks completed.
    pub total_tasks_completed: u64,
}

/// Snapshot of pool status.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PoolStatus {
    /// Total number of slots.
    pub total_slots: usize,
    /// Number of idle slots.
    pub idle_slots: usize,
    /// Number of busy slots.
    pub busy_slots: usize,
    /// Number of currently running tasks.
    pub running_tasks: usize,
    /// Number of pending (queued) tasks.
    pub pending_tasks: usize,
    /// Total spend in microdollars.
    pub total_spend_microdollars: u64,
    /// Budget cap in microdollars, if set.
    pub budget_microdollars: Option<u64>,
    /// Whether the pool is shutting down.
    pub shutdown: bool,
}

use serde::{Deserialize, Serialize};

/// Generate a short unique ID.
fn new_id() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    format!("{nanos:x}")
}

#[cfg(test)]
mod tests {
    use super::*;

    fn mock_claude() -> Claude {
        // Build a Claude instance pointing at a non-existent binary.
        // Tests that don't actually execute tasks can use this.
        Claude::builder().binary("/usr/bin/false").build().unwrap()
    }

    #[tokio::test]
    async fn build_pool_registers_slots() {
        let pool = Pool::builder(mock_claude()).slots(3).build().await.unwrap();

        let slots = pool.store().list_slots().await.unwrap();
        assert_eq!(slots.len(), 3);

        for slot in &slots {
            assert_eq!(slot.state, SlotState::Idle);
        }
    }

    #[tokio::test]
    async fn pool_with_slot_configs() {
        let pool = Pool::builder(mock_claude())
            .slots(2)
            .slot_config(SlotConfig {
                model: Some("opus".into()),
                role: Some("reviewer".into()),
                ..Default::default()
            })
            .build()
            .await
            .unwrap();

        let slots = pool.store().list_slots().await.unwrap();
        let w0 = slots.iter().find(|w| w.id.0 == "slot-0").unwrap();
        let w1 = slots.iter().find(|w| w.id.0 == "slot-1").unwrap();
        assert_eq!(w0.config.model.as_deref(), Some("opus"));
        assert_eq!(w0.config.role.as_deref(), Some("reviewer"));
        // Slot 1 gets default config.
        assert!(w1.config.model.is_none());
    }

    #[tokio::test]
    async fn context_operations() {
        let pool = Pool::builder(mock_claude()).slots(1).build().await.unwrap();

        pool.set_context("repo", "claude-wrapper");
        pool.set_context("branch", "main");

        assert_eq!(pool.get_context("repo").as_deref(), Some("claude-wrapper"));
        assert_eq!(pool.list_context().len(), 2);

        pool.delete_context("branch");
        assert!(pool.get_context("branch").is_none());
    }

    #[tokio::test]
    async fn drain_marks_slots_stopped() {
        let pool = Pool::builder(mock_claude()).slots(2).build().await.unwrap();

        let summary = pool.drain().await.unwrap();
        assert_eq!(summary.total_tasks_completed, 0);

        let slots = pool.store().list_slots().await.unwrap();
        for w in &slots {
            assert_eq!(w.state, SlotState::Stopped);
        }

        // Pool rejects new work after drain.
        assert!(pool.run("hello").await.is_err());
    }

    #[tokio::test]
    async fn budget_enforcement() {
        let pool = Pool::builder(mock_claude())
            .slots(1)
            .config(PoolConfig {
                budget_microdollars: Some(100),
                ..Default::default()
            })
            .build()
            .await
            .unwrap();

        // Simulate spending past the budget.
        pool.inner.total_spend.store(100, Ordering::Relaxed);

        let err = pool.run("hello").await.unwrap_err();
        assert!(matches!(err, Error::BudgetExhausted { .. }));
    }

    #[tokio::test]
    async fn status_snapshot() {
        let pool = Pool::builder(mock_claude())
            .slots(3)
            .config(PoolConfig {
                budget_microdollars: Some(1_000_000),
                ..Default::default()
            })
            .build()
            .await
            .unwrap();

        let status = pool.status().await.unwrap();
        assert_eq!(status.total_slots, 3);
        assert_eq!(status.idle_slots, 3);
        assert_eq!(status.busy_slots, 0);
        assert_eq!(status.budget_microdollars, Some(1_000_000));
        assert!(!status.shutdown);
    }

    #[tokio::test]
    async fn no_idle_slots_timeout() {
        let pool = Pool::builder(mock_claude())
            .slots(1)
            .config(PoolConfig {
                slot_assignment_timeout_secs: 1,
                ..Default::default()
            })
            .build()
            .await
            .unwrap();

        // Manually mark the slot as busy.
        let mut slots = pool.store().list_slots().await.unwrap();
        slots[0].state = SlotState::Busy;
        pool.store().put_slot(slots[0].clone()).await.unwrap();

        let err = pool.run("hello").await.unwrap_err();
        assert!(matches!(err, Error::NoSlotAvailable { timeout_secs: 1 }));
    }

    #[tokio::test]
    async fn fan_out_with_excess_prompts() {
        // This test verifies that fan_out can queue excess prompts.
        // With 2 slots and 4 prompts, all 4 should eventually complete.
        // Since we use mock_claude (non-existent binary), actual execution will fail,
        // but we're testing that the queueing mechanism works (assignment tries to get a slot).
        let pool = Pool::builder(mock_claude()).slots(2).build().await.unwrap();

        let prompts = vec!["prompt1", "prompt2", "prompt3", "prompt4"];

        // This will fail due to mock binary, but the key point is that
        // it tries to execute all prompts even though we only have 2 slots.
        // Before the fix, excess prompts would fail with "no idle slots" immediately.
        // After the fix, they queue and wait.
        let results = pool.fan_out(&prompts).await;

        // We expect all 4 tasks to be attempted (the mock binary failure is expected).
        // The test is that we get 4 results (not an immediate failure due to slot count).
        match results {
            Ok(_) | Err(_) => {
                // Both outcomes are ok; we're testing that fan_out doesn't fail
                // with immediate "no idle slots" error when prompts > slots.
            }
        }
    }

    #[tokio::test]
    async fn slot_identity_fields_persisted() {
        let pool = Pool::builder(mock_claude())
            .slots(1)
            .slot_config(SlotConfig {
                name: Some("reviewer".into()),
                role: Some("code_review".into()),
                description: Some("Reviews PRs for correctness and style".into()),
                ..Default::default()
            })
            .build()
            .await
            .unwrap();

        let slots = pool.store().list_slots().await.unwrap();
        let slot = slots.iter().find(|w| w.id.0 == "slot-0").unwrap();

        assert_eq!(slot.config.name.as_deref(), Some("reviewer"));
        assert_eq!(slot.config.role.as_deref(), Some("code_review"));
        assert_eq!(
            slot.config.description.as_deref(),
            Some("Reviews PRs for correctness and style")
        );
    }

    #[tokio::test]
    async fn scale_up_increases_slot_count() {
        let pool = Pool::builder(mock_claude()).slots(2).build().await.unwrap();

        let initial_count = pool.store().list_slots().await.unwrap().len();
        assert_eq!(initial_count, 2);

        let new_count = pool.scale_up(3).await.unwrap();
        assert_eq!(new_count, 5);

        let slots = pool.store().list_slots().await.unwrap();
        assert_eq!(slots.len(), 5);

        // Verify new slots are idle.
        for slot in slots.iter().skip(2) {
            assert_eq!(slot.state, SlotState::Idle);
        }
    }

    #[tokio::test]
    async fn scale_up_respects_max_slots() {
        let mut config = PoolConfig::default();
        config.scaling.max_slots = 4;

        let pool = Pool::builder(mock_claude())
            .slots(2)
            .config(config)
            .build()
            .await
            .unwrap();

        // Try to scale beyond max.
        let result = pool.scale_up(5).await;
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("exceeds max_slots")
        );

        // Verify count unchanged.
        assert_eq!(pool.store().list_slots().await.unwrap().len(), 2);
    }

    #[tokio::test]
    async fn scale_down_reduces_slot_count() {
        let pool = Pool::builder(mock_claude()).slots(4).build().await.unwrap();

        let initial = pool.store().list_slots().await.unwrap().len();
        assert_eq!(initial, 4);

        let new_count = pool.scale_down(2).await.unwrap();
        assert_eq!(new_count, 2);

        assert_eq!(pool.store().list_slots().await.unwrap().len(), 2);
    }

    #[tokio::test]
    async fn scale_down_respects_min_slots() {
        let mut config = PoolConfig::default();
        config.scaling.min_slots = 2;

        let pool = Pool::builder(mock_claude())
            .slots(3)
            .config(config)
            .build()
            .await
            .unwrap();

        // Try to scale below min.
        let result = pool.scale_down(2).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("below min_slots"));

        // Verify count unchanged.
        assert_eq!(pool.store().list_slots().await.unwrap().len(), 3);
    }

    #[tokio::test]
    async fn set_target_slots_scales_up() {
        let pool = Pool::builder(mock_claude()).slots(2).build().await.unwrap();

        let new_count = pool.set_target_slots(5).await.unwrap();
        assert_eq!(new_count, 5);
        assert_eq!(pool.store().list_slots().await.unwrap().len(), 5);
    }

    #[tokio::test]
    async fn set_target_slots_scales_down() {
        let pool = Pool::builder(mock_claude()).slots(5).build().await.unwrap();

        let new_count = pool.set_target_slots(2).await.unwrap();
        assert_eq!(new_count, 2);
        assert_eq!(pool.store().list_slots().await.unwrap().len(), 2);
    }

    #[tokio::test]
    async fn set_target_slots_no_op_when_equal() {
        let pool = Pool::builder(mock_claude()).slots(3).build().await.unwrap();

        let new_count = pool.set_target_slots(3).await.unwrap();
        assert_eq!(new_count, 3);
    }

    #[tokio::test]
    async fn fan_out_chains_submits_all_chains() {
        let pool = Pool::builder(mock_claude()).slots(2).build().await.unwrap();

        let skills = crate::skill::SkillRegistry::new();
        let options = crate::chain::ChainOptions { tags: vec![] };

        // Create two chains, each with one prompt step.
        let chain1 = vec![crate::chain::ChainStep {
            name: "step1".into(),
            action: crate::chain::StepAction::Prompt {
                prompt: "prompt 1".into(),
            },
            config: None,
            failure_policy: crate::chain::StepFailurePolicy {
                retries: 0,
                recovery_prompt: None,
            },
        }];

        let chain2 = vec![crate::chain::ChainStep {
            name: "step1".into(),
            action: crate::chain::StepAction::Prompt {
                prompt: "prompt 2".into(),
            },
            config: None,
            failure_policy: crate::chain::StepFailurePolicy {
                retries: 0,
                recovery_prompt: None,
            },
        }];

        let chains = vec![chain1, chain2];

        // Submit both chains in parallel.
        let task_ids = pool.fan_out_chains(chains, &skills, options).await.unwrap();

        // Should have 2 task IDs.
        assert_eq!(task_ids.len(), 2);

        // Verify task IDs are different.
        assert_ne!(task_ids[0].0, task_ids[1].0);

        // Verify tasks exist in the store.
        for task_id in &task_ids {
            let task = pool.store().get_task(task_id).await.unwrap();
            assert!(task.is_some());
        }
    }
}