orch/

control.rs

//! Control session — context assembly, one-shot agent invocation, response parsing.
//!
//! The control session provides an interactive ops assistant that assembles context
//! from SQLite (memories, summaries, live state), invokes an agent CLI one-shot
//! with `--session-id` for conversation continuity, and parses the response for
//! summary extraction.
//!
//! ## Session continuity
//!
//! Each named session (e.g., "default", "ops") gets a UUID stored in SQLite KV
//! (`control:session_id:{name}`). Session agents (claude, kimi, minimax, grok)
//! receive `--session-id <uuid>` so they resume prior conversation state.
//! The UUID is regenerated when the agent or model changes.
//!
//! ## Model selection
//!
//! `/model <spec>` where spec is:
//! - `agent:model` — explicit agent + model (e.g., `minimax:sonnet`, `opencode:minimax-m2.5-free`)
//! - `model` — infer agent from model name (e.g., `sonnet` → claude, `gpt-4o` → codex)
//!
//! All agents: claude, codex, opencode, kimi, minimax, grok
//! - kimi and minimax are claude-compatible wrappers (separate binaries in PATH)
//! - opencode supports `opencode models` for listing available models

use std::collections::HashMap;
use std::process;
use std::sync::PoisonError;
use std::sync::{Arc, LazyLock, Mutex};
use std::time;

use crate::cmd_cache;
use crate::engine;
use crate::engine::router::RouterConfig;
use crate::engine::runner::agents;
use crate::engine::runner::agents::opencode;
use crate::engine::runner::direct;
use crate::home;
use crate::store::TaskStore;
use anyhow::{Context, Result};
use regex::Regex;
use tokio::fs;
use tokio::process::Command;
use tokio::sync;
use tokio::time as tokio_time;
use tokio::time::{Duration, timeout};
use uuid::Uuid;

const SUBPROCESS_TIMEOUT: Duration = Duration::from_secs(10);

/// Per-session invocation locks — ensures only one agent invocation runs at a time per session.
///
/// Keyed by `session_id`. A new `Arc<tokio::sync::Mutex<()>>` is inserted on first use and
/// reused for all subsequent calls with the same session_id. Entries are evicted after each
/// `send_message` call when no concurrent waiters remain. Note: because callers keep a
/// local `Arc` clone while executing, the map entry + local clone equal 2 strong refs when
/// there are no other waiters. We therefore test for `Arc::strong_count() == 2` when
/// deciding to evict the map entry (map entry + local clone). If other waiters exist the
/// count will be > 2 and we leave the entry in place.
static SESSION_LOCKS: LazyLock<Mutex<HashMap<String, Arc<sync::Mutex<()>>>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

/// System prompt template, loaded at compile time.
const SYSTEM_TEMPLATE: &str = include_str!("../prompts/control_system.md");

/// Timeout for agent invocations (2 minutes).
const AGENT_TIMEOUT: time::Duration = time::Duration::from_secs(120);

/// Maximum bytes for a single control-session turn payload.
/// When the combined state + user message exceeds this limit, the payload is split
/// across multiple sequential turns so the user's question is never truncated.
const MAX_CONTROL_MESSAGE_BYTES: usize = 32 * 1024;

/// Delivery mode for control session messages.
enum DeliveryMode {
    /// Single message delivery (under-limit messages)
    Single,
    /// Chunked delivery with session continuity (claude, kimi, minimax) — uses --resume
    Chunked,
    /// Chunked delivery without session continuity (codex, opencode) — independent invocations
    ChunkedNonSession,
}

/// Parsed model specification from `/model` command.
#[derive(Debug, Clone, PartialEq)]
pub struct ModelSpec {
    pub agent: String,
    pub model: String,
}

/// Parsed control command from the chat session.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ControlCommand {
    Model(Option<String>),
    Agent(Option<String>),
    Restart,
}

/// The agent's first model under `models:` in the config. Empty means orch passes no `--model`
/// and the agent's CLI uses its own default.
fn default_model_for_agent(agent: &str) -> String {
    first_model(&RouterConfig::configured_models(agent))
}

// A pool keyword such as `opencode:free` names no model, so it is skipped
fn first_model(models: &[String]) -> String {
    models
        .iter()
        .find(|m| !m.contains(':'))
        .cloned()
        .unwrap_or_default()
}

/// Parse a control-session command.
pub fn parse_control_command(message: &str) -> Option<ControlCommand> {
    let trimmed = message.trim();
    let rest = trimmed.strip_prefix('/')?;
    let (command, args) = rest
        .split_once(' ')
        .map_or((rest, ""), |(cmd, tail)| (cmd, tail.trim()));

    match command {
        "model" => Some(ControlCommand::Model(
            (!args.is_empty()).then(|| args.to_string()),
        )),
        "agent" => Some(ControlCommand::Agent(
            (!args.is_empty()).then(|| args.to_string()),
        )),
        "restart" => Some(ControlCommand::Restart),
        _ => None,
    }
}

/// Parse a model spec string into agent + model.
///
/// Formats:
/// - `agent:model` → explicit (e.g., `minimax:sonnet`, `opencode:minimax-m2.5-free`)
/// - `model` → infer agent from model name
pub fn parse_model_spec(spec: &str, current_agent: &str) -> ModelSpec {
    if let Some((agent, model)) = spec.split_once(':') {
        ModelSpec {
            agent: agent.to_string(),
            model: RouterConfig::normalize_model_identifier(model),
        }
    } else {
        let model = RouterConfig::normalize_model_identifier(spec);
        let configured: Vec<(String, Vec<String>)> = engine::configured_agents()
            .into_iter()
            .map(|agent| {
                let models = RouterConfig::configured_models(&agent);
                (agent, models)
            })
            .collect();
        ModelSpec {
            agent: infer_agent(&model, &configured, current_agent),
            model,
        }
    }
}

/// The agent whose `models:` entry lists the model, or the current agent when none does.
fn infer_agent(model: &str, configured: &[(String, Vec<String>)], current_agent: &str) -> String {
    configured
        .iter()
        .find(|(_, models)| models.iter().any(|m| m == model))
        .map_or(current_agent, |(agent, _)| agent.as_str())
        .to_string()
}

/// Validate that an agent name is known and installed.
///
/// Performs lightweight checks (known agent + binary in PATH) without
/// a test invocation. Used by `/agent` before persisting.
pub fn validate_agent(agent: &str) -> Result<()> {
    let agents = engine::configured_agents();
    if !agents.iter().any(|a| a == agent) {
        anyhow::bail!(
            "unknown agent '{}'. Available: {}",
            agent,
            agents.join(", ")
        );
    }
    if !cmd_cache::command_exists(agent) {
        anyhow::bail!("agent '{agent}' not found in PATH. Is it installed?");
    }
    Ok(())
}

/// Validate a model spec by running a quick test invocation.
///
/// Sends "hello" to the agent with the specified model to verify
/// the agent binary exists and the model is available.
#[expect(clippy::print_stderr, reason = "progress line for `orch chat /model`")]
pub async fn validate_model(spec: &ModelSpec) -> Result<()> {
    // 1. Check agent is known and installed
    validate_agent(&spec.agent)?;

    // 3. For opencode, pre-check against `opencode models` list (fast rejection)
    if spec.agent == "opencode"
        && let Ok(models) = list_opencode_models().await
        && !models.iter().any(|m| m == &spec.model)
    {
        let suggestions: Vec<_> = models
            .iter()
            .filter(|m| m.contains(&spec.model) || spec.model.contains(m.as_str()))
            .take(10)
            .cloned()
            .collect();
        let shown = if suggestions.is_empty() {
            models.iter().take(20).cloned().collect()
        } else {
            suggestions
        };
        anyhow::bail!(
            "model '{}' not found in opencode. Available models:\n{}",
            spec.model,
            shown.join("\n")
        );
    }

    // 4. Test invocation — always run to verify the model actually works
    eprintln!("Testing {}:{} ...", spec.agent, spec.model);
    let result = test_invoke(&spec.agent, &spec.model).await;
    match result {
        Ok(()) => Ok(()),
        Err(e) => anyhow::bail!(
            "model '{}' on agent '{}' is not available: {e}",
            spec.model,
            spec.agent
        ),
    }
}

/// List available opencode models via `opencode models`.
async fn list_opencode_models() -> Result<Vec<String>> {
    match opencode::run_opencode_models_cli(time::Duration::from_secs(10)).await {
        (models, "ok" | "empty_output") => Ok(models),
        (_, reason) => anyhow::bail!("opencode models failed: {reason}"),
    }
}

/// Run a minimal test invocation to verify agent+model availability.
async fn test_invoke(agent: &str, model: &str) -> Result<()> {
    let msg = "Reply with just the word 'ok'.";
    // Reuse invoke_agent — it already handles temp files, timeout, and error checking
    invoke_agent(agent, model, "", msg).await?;
    Ok(())
}

/// Assembled context for a chat invocation.
#[derive(Debug)]
pub struct ChatContext {
    /// System prompt (instructions, role, commands)
    pub system: String,
    /// User-visible context prepended to the message (live state, memories)
    pub state: String,
}

/// Assemble the chat context from SQLite state and live system info.
pub async fn assemble_context(store: &TaskStore, session_id: &str) -> Result<ChatContext> {
    // 1. Gather memories from KV (keys matching control:memory:{session}:*)
    // Use kv_list_prefix which correctly escapes LIKE metacharacters so that
    // session names containing '_' or '%' don't leak memories from other sessions.
    let memory_prefix = format!("control:memory:{session_id}:");
    let raw_memories = store
        .kv_list_prefix(&memory_prefix)
        .await
        .unwrap_or_else(|e| {
            tracing::warn!(session_id, error = %e, "failed to load control memories from store");
            Vec::new()
        });
    let memories = if raw_memories.is_empty() {
        "(none)".to_string()
    } else {
        raw_memories
            .iter()
            .map(|(key, value)| {
                let short_key = key.strip_prefix("control:memory:").unwrap_or(key);
                format!("- **{short_key}**: {value}")
            })
            .collect::<Vec<_>>()
            .join("\n")
    };

    // 2. Recent conversation summaries
    let summaries = store
        .control_recent_summaries(session_id, 20)
        .await
        .unwrap_or_else(|e| {
            tracing::warn!(session_id, error = %e, "failed to load control summaries from store");
            Vec::new()
        });
    let recent_summaries = if summaries.is_empty() {
        "(no recent conversation)".to_string()
    } else {
        summaries
            .iter()
            .map(|s| format!("- {s}"))
            .collect::<Vec<_>>()
            .join("\n")
    };

    // 3-5. Run independent subprocess calls in parallel (best-effort, 10s timeout each)
    let (service_status, task_list, job_list) = tokio::join!(
        async {
            let mut cmd = Command::new("brew");
            cmd.args(["services", "info", "orch"]).kill_on_drop(true);
            match timeout(SUBPROCESS_TIMEOUT, cmd.output()).await {
                Ok(Ok(output)) if output.status.success() => {
                    String::from_utf8_lossy(&output.stdout).trim().to_string()
                }
                _ => "(could not check service status)".to_string(),
            }
        },
        async {
            let mut cmd = Command::new("orch");
            cmd.args(["task", "list"]).kill_on_drop(true);
            match timeout(SUBPROCESS_TIMEOUT, cmd.output()).await {
                Ok(Ok(output)) if output.status.success() => {
                    String::from_utf8_lossy(&output.stdout).trim().to_string()
                }
                _ => "(could not fetch live state)".to_string(),
            }
        },
        async {
            let mut cmd = Command::new("orch");
            cmd.args(["job", "list"]).kill_on_drop(true);
            match timeout(SUBPROCESS_TIMEOUT, cmd.output()).await {
                Ok(Ok(output)) if output.status.success() => {
                    String::from_utf8_lossy(&output.stdout).trim().to_string()
                }
                _ => "(no jobs configured)".to_string(),
            }
        },
    );

    let version = env!("ORCH_VERSION");
    let state = format!(
        "## Current Live State\n### Version\norch {version}\n\n### Service\n{service_status}\n\n### Tasks\n{task_list}\n\n### Scheduled Jobs\n{job_list}\n\n### Memories\n{memories}\n\n### Recent Conversation\n{recent_summaries}"
    );

    // 6. Build system prompt (instructions only, no live state in system prompt)
    let system = SYSTEM_TEMPLATE.to_string();

    Ok(ChatContext { system, state })
}

/// Get the current model from KV, defaulting to the current agent's configured model when the
/// key is absent. Returns an error if the store read fails.
pub async fn get_model(store: &TaskStore) -> Result<String> {
    match store
        .kv_get("control:model")
        .await
        .context("reading control:model from store")?
    {
        Some(model) => Ok(model),
        None => Ok(default_model_for_agent(&get_agent(store).await?)),
    }
}

/// Get the current agent from KV, defaulting to `"claude"` when the key is absent.
/// Returns an error if the store read fails.
pub async fn get_agent(store: &TaskStore) -> Result<String> {
    Ok(store
        .kv_get("control:agent")
        .await
        .context("reading control:agent from store")?
        .unwrap_or_else(|| "claude".to_string()))
}

/// Get the current control session spec from KV.
pub async fn get_current_spec(store: &TaskStore) -> Result<ModelSpec> {
    Ok(ModelSpec {
        agent: get_agent(store).await?,
        model: get_model(store).await?,
    })
}

/// Set the current model in KV.
pub async fn set_model(store: &TaskStore, model: &str) -> Result<()> {
    store
        .kv_set("control:model", model)
        .await
        .context("setting control model")?;
    Ok(())
}

/// Set the current agent in KV.
pub async fn set_agent(store: &TaskStore, agent: &str) -> Result<()> {
    store
        .kv_set("control:agent", agent)
        .await
        .context("setting control agent")?;
    Ok(())
}

/// Set the current model and agent in KV.
pub async fn set_model_spec(store: &TaskStore, spec: &ModelSpec) -> Result<()> {
    set_model(store, &spec.model).await?;
    set_agent(store, &spec.agent).await?;
    Ok(())
}

/// Set the current agent and its default model in KV.
pub async fn set_agent_spec(store: &TaskStore, agent: &str) -> Result<ModelSpec> {
    let spec = ModelSpec {
        agent: agent.to_string(),
        model: default_model_for_agent(agent),
    };
    set_model_spec(store, &spec).await?;
    Ok(spec)
}

fn display_model(model: &str) -> &str {
    if model.is_empty() { "default" } else { model }
}

fn format_current_spec(spec: &ModelSpec) -> String {
    format!("Current: {}:{}", spec.agent, display_model(&spec.model))
}

fn format_switched_spec(spec: &ModelSpec) -> String {
    format!(
        "Switched to **{}:{}** (validated)",
        spec.agent,
        display_model(&spec.model)
    )
}

async fn handle_control_command(
    store: &TaskStore,
    session_id: &str,
    channel_thread: Option<&str>,
    command: ControlCommand,
) -> Result<String> {
    match command {
        ControlCommand::Model(None) | ControlCommand::Agent(None) => {
            Ok(format_current_spec(&get_current_spec(store).await?))
        }
        ControlCommand::Model(Some(new_spec)) => {
            let spec = parse_model_spec(&new_spec, &get_agent(store).await?);
            validate_model(&spec).await?;
            set_model_spec(store, &spec).await?;
            // Reset session UUID so next message starts a fresh conversation
            reset_session_uuid(store, session_id).await?;
            Ok(format_switched_spec(&spec))
        }
        ControlCommand::Agent(Some(agent)) => {
            validate_agent(&agent)?;
            let spec = set_agent_spec(store, &agent).await?;
            // Reset session UUID so next message starts a fresh conversation
            reset_session_uuid(store, session_id).await?;
            Ok(format_switched_spec(&spec))
        }
        ControlCommand::Restart => {
            // Store the conversation key so the post-restart startup check can send a confirmation.
            if let Some(ct) = channel_thread {
                store.kv_set("control:restart_pending", ct).await.ok();
            }
            // Delay the SIGTERM slightly so the acknowledgment reply is delivered first.
            tokio::spawn(async {
                tokio_time::sleep(Duration::from_secs(2)).await;
                let pid = process::id().to_string();
                let _ = Command::new("kill").args(["-TERM", &pid]).status().await;
            });
            Ok("Restarting orch service\u{2026} I'll confirm here when back online.".to_string())
        }
    }
}

/// Handle a control-session command if present.
pub async fn maybe_handle_control_command(
    store: &TaskStore,
    session_id: &str,
    channel_thread: Option<&str>,
    message: &str,
) -> Result<Option<String>> {
    Ok(match parse_control_command(message) {
        Some(command) => {
            Some(handle_control_command(store, session_id, channel_thread, command).await?)
        }
        None => None,
    })
}

/// Memory entry parsed from an agent response.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemoryEntry {
    pub key: String,
    pub value: String,
}

/// Parsed agent response metadata.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedResponse {
    pub clean_text: String,
    pub summary: Option<String>,
    pub memories: Vec<MemoryEntry>,
}

static MEMORY_RE: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(r#"(?s)<memory\s+key="([^"]+)">(.+?)</memory>"#).expect("valid memory regex")
});

static SUMMARY_RE: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?s)<summary>(.+?)</summary>").expect("valid summary regex"));

/// Parse an agent response, extracting optional `<summary>` and `<memory>` tags.
pub fn parse_response(raw: &str) -> ParsedResponse {
    let mut clean = raw.to_string();
    let memories = MEMORY_RE
        .captures_iter(raw)
        .map(|caps| MemoryEntry {
            key: caps[1].trim().to_string(),
            value: caps[2].trim().to_string(),
        })
        .collect::<Vec<_>>();
    clean = MEMORY_RE.replace_all(&clean, "").to_string();

    let summary = SUMMARY_RE
        .captures(raw)
        .map(|caps| caps[1].trim().to_string());
    clean = SUMMARY_RE.replace_all(&clean, "").to_string();

    ParsedResponse {
        clean_text: clean.trim().to_string(),
        summary,
        memories,
    }
}

async fn persist_memories(
    store: &TaskStore,
    session_id: &str,
    memories: &[MemoryEntry],
) -> Result<()> {
    for memory in memories {
        store
            .kv_set(
                &format!("control:memory:{session_id}:{}", memory.key),
                &memory.value,
            )
            .await
            .with_context(|| format!("storing control memory {}", memory.key))?;
    }
    Ok(())
}

/// Result of an agent invocation, including response text and token usage.
#[derive(Debug)]
pub struct InvokeResult {
    pub text: String,
    pub input_tokens: Option<u64>,
    pub output_tokens: Option<u64>,
    /// Cost in USD as the agent reported it
    pub cost_usd: Option<f64>,
}

// An empty model omits `--model`, so the agent's CLI uses its own default
fn model_arg(model: &str) -> Option<&str> {
    (!model.is_empty()).then_some(model)
}

fn token_u64_to_i64_saturating(tokens: u64) -> i64 {
    i64::try_from(tokens).unwrap_or(i64::MAX)
}

fn opt_token_u64_to_i64_saturating(tokens: Option<u64>) -> Option<i64> {
    tokens.map(token_u64_to_i64_saturating)
}

fn total_tokens_u64_to_i64_saturating(
    input_tokens: Option<u64>,
    output_tokens: Option<u64>,
) -> Option<i64> {
    match (input_tokens, output_tokens) {
        (Some(input), Some(output)) => {
            Some(token_u64_to_i64_saturating(input.saturating_add(output)))
        }
        (Some(tokens), None) | (None, Some(tokens)) => Some(token_u64_to_i64_saturating(tokens)),
        _ => None,
    }
}

/// Prepare temp directory for agent invocation files.
///
/// Uses `~/.orch/state/control/{pid}/` for isolation between concurrent processes.
async fn prepare_temp_dir() -> Result<String> {
    let home = home::orch_home()?;
    let dir = format!("{}/state/control/{}", home.display(), process::id());
    fs::create_dir_all(&dir).await?;
    Ok(dir)
}

/// Invoke an agent CLI one-shot and return structured result with token usage.
///
/// Delegates to `runner::direct::run_direct()` for command building, execution,
/// parsing, and error classification.
///
/// Writes temp files to `~/.orch/state/control/{pid}/`.
pub async fn invoke_agent(
    agent: &str,
    model: &str,
    context: &str,
    message: &str,
) -> Result<InvokeResult> {
    let dir = prepare_temp_dir().await?;

    let result = direct::run_direct(
        agent,
        model_arg(model),
        context,
        message,
        AGENT_TIMEOUT,
        &dir,
    )
    .await?;

    Ok(InvokeResult {
        text: result.text,
        input_tokens: result.input_tokens,
        output_tokens: result.output_tokens,
        cost_usd: result.cost_usd,
    })
}

/// KV key for storing the session UUID for a named session.
fn session_uuid_key(session_name: &str) -> String {
    format!("control:session_id:{session_name}")
}

/// Get or create a session UUID for the given session name.
///
/// On first call for a session, generates a new UUID and persists it to KV.
/// Subsequent calls return the same UUID, giving agents conversation continuity
/// via `--session-id`.
/// Get or create a session UUID for the given session name.
///
/// On first call for a session, generates a new UUID and persists it to KV.
/// Subsequent calls return the same UUID, giving agents conversation continuity
/// via `--session-id` (new) or `--resume` (existing).
///
/// Returns `(uuid, is_new)` so callers know whether to create or resume the session.
pub async fn get_or_create_session_uuid(
    store: &TaskStore,
    session_name: &str,
) -> Result<(String, bool)> {
    let key = session_uuid_key(session_name);
    let candidate = Uuid::new_v4().to_string();
    let stored = store
        .kv_insert_if_absent(&key, &candidate)
        .await
        .context("storing session UUID")?;
    let is_new = stored == candidate;
    if is_new {
        tracing::info!(session_name, uuid = %stored, "created new chat session UUID");
    }
    Ok((stored, is_new))
}

/// Reset the session UUID so the next message starts a fresh conversation.
///
/// Called when agent or model changes, since the new agent won't share
/// the previous agent's session state. Deletes the stored UUID so the next
/// `get_or_create_session_uuid` call generates a fresh one.
pub async fn reset_session_uuid(store: &TaskStore, session_name: &str) -> Result<()> {
    let key = session_uuid_key(session_name);
    store
        .kv_delete(&key)
        .await
        .context("resetting session UUID")?;
    tracing::info!(session_name, "reset chat session UUID");
    Ok(())
}

/// Invoke an agent one-shot with `--session-id` or `--resume` for conversation continuity.
///
/// For claude-compatible agents (claude, kimi, minimax), passes `--session-id <uuid>`
/// when creating a new session, or `--resume <uuid>` when continuing an existing one.
/// Other agents get a plain one-shot call.
pub async fn invoke_agent_with_session(
    agent: &str,
    model: &str,
    context: &str,
    message: &str,
    session_uuid: Option<&str>,
    resume: bool,
) -> Result<InvokeResult> {
    let dir = prepare_temp_dir().await?;

    let session = session_uuid.map(|sid| crate::engine::runner::direct::SessionContinuation {
        session_id: sid,
        resume,
    });

    let result = direct::run_direct_with_session(
        agent,
        model_arg(model),
        context,
        message,
        AGENT_TIMEOUT,
        &dir,
        session,
    )
    .await?;

    Ok(InvokeResult {
        text: result.text,
        input_tokens: result.input_tokens,
        output_tokens: result.output_tokens,
        cost_usd: result.cost_usd,
    })
}

/// High-level entry point: process a user message and return the assistant response.
///
/// Handles `/model` and `/agent` commands, assembles context, invokes the agent,
/// parses the response for summaries, and stores both user and assistant messages.
pub async fn send_message(
    store: &TaskStore,
    session_id: &str,
    channel: &str,
    channel_thread: Option<&str>,
    message: &str,
) -> Result<String> {
    // Handle control commands — don't store as messages.
    if let Some(response) =
        maybe_handle_control_command(store, session_id, channel_thread, message).await?
    {
        return Ok(response);
    }

    // Acquire per-session lock — serializes concurrent invocations for the same session.
    // Messages that arrive while an invocation is running will queue here and execute in order.
    let session_lock = {
        let mut map = SESSION_LOCKS.lock().unwrap_or_else(PoisonError::into_inner);
        Arc::clone(
            map.entry(session_id.to_string())
                .or_insert_with(|| Arc::new(sync::Mutex::new(()))),
        )
    };
    let guard = session_lock.lock().await;

    // Store user message
    store
        .insert_control_message(
            session_id,
            "user",
            channel,
            channel_thread,
            message,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
        )
        .await?;

    // Assemble context
    let ctx = assemble_context(store, session_id).await?;

    // Resolve model and agent from KV
    let model = get_model(store).await?;
    let agent = get_agent(store).await?;

    // Prepend live state to user message so the LLM sees it directly
    let full_message = format!("{}\n\n---\n\n{}", ctx.state, message);

    // Get or create a session UUID for conversation continuity
    let (session_uuid, is_new) = get_or_create_session_uuid(store, session_id).await?;

    // When the combined payload exceeds the per-turn limit, split it across multiple
    // sequential turns so the user's question (trailing content) is never truncated.
    // For session-supporting agents (claude, kimi, minimax), use --resume for continuity
    // and persist ALL intermediate assistant responses so local transcript matches provider.
    // For other agents (codex, opencode), split into chunks and send each as independent
    // invocation (no --resume), each getting its own persisted assistant response.
    let agent_has_sessions = agents::supports_session(&agent);
    let (chunks, delivery_mode) = if full_message.len() > MAX_CONTROL_MESSAGE_BYTES {
        tracing::info!(
            bytes = full_message.len(),
            limit = MAX_CONTROL_MESSAGE_BYTES,
            agent = %agent,
            agent_has_sessions,
            "control session payload exceeds limit"
        );
        if agent_has_sessions {
            // Session agents: split into chunks for sequential delivery with --resume
            (
                build_control_chunks(&ctx.state, message, true),
                DeliveryMode::Chunked,
            )
        } else {
            // Non-session agents: split into chunks, send each independently
            (
                build_control_chunks(&ctx.state, message, false),
                DeliveryMode::ChunkedNonSession,
            )
        }
    } else {
        (vec![full_message], DeliveryMode::Single)
    };

    let _num_chunks = chunks.len();

    // Deliver chunks based on mode
    let result = match delivery_mode {
        DeliveryMode::Chunked => {
            // Session-based delivery (claude, kimi, minimax) — use --resume for continuity
            // Persist ALL intermediate responses so local DB matches provider transcript.
            let mut active_uuid = session_uuid.clone();
            let mut is_first_chunk = !is_new;
            let mut last_result: Option<InvokeResult> = None;

            for chunk in &chunks {
                let r = invoke_agent_with_session(
                    &agent,
                    &model,
                    &ctx.system,
                    chunk,
                    Some(&active_uuid),
                    is_first_chunk,
                )
                .await;
                let chunk_result = match r {
                    Ok(result) => result,
                    Err(e) => {
                        let err_str = e.to_string();
                        let is_stale = err_str.contains("stale session");
                        let is_timeout = err_str.contains("timeout after")
                            || err_str.contains("timed out after");
                        if is_stale || is_timeout {
                            if is_stale {
                                tracing::info!(
                                    session_id,
                                    "chat session UUID expired; resetting and retrying with fresh session"
                                );
                            } else {
                                tracing::info!(
                                    session_id,
                                    "chat session timed out; resetting and retrying with fresh session"
                                );
                            }
                            reset_session_uuid(store, session_id).await?;
                            let (fresh_uuid, _) =
                                get_or_create_session_uuid(store, session_id).await?;
                            active_uuid = fresh_uuid;
                            is_first_chunk = false;
                            let fresh_result = invoke_agent_with_session(
                                &agent,
                                &model,
                                &ctx.system,
                                chunk,
                                Some(&active_uuid),
                                false,
                            )
                            .await?;

                            if is_timeout {
                                InvokeResult {
                                    text: format!(
                                        "_(New session started.)_\n\n{}",
                                        fresh_result.text
                                    ),
                                    ..fresh_result
                                }
                            } else {
                                fresh_result
                            }
                        } else {
                            return Err(e);
                        }
                    }
                };

                // Persist this chunk's assistant response (even intermediate ones)
                let parsed = parse_response(&chunk_result.text);
                let total_tokens = total_tokens_u64_to_i64_saturating(
                    chunk_result.input_tokens,
                    chunk_result.output_tokens,
                );
                let input_tokens = opt_token_u64_to_i64_saturating(chunk_result.input_tokens);
                let output_tokens = opt_token_u64_to_i64_saturating(chunk_result.output_tokens);
                let cost_usd = chunk_result.cost_usd;
                store
                    .insert_control_message(
                        session_id,
                        "assistant",
                        channel,
                        channel_thread,
                        &parsed.clean_text,
                        parsed.summary.as_deref(),
                        Some(&model),
                        Some(&agent),
                        input_tokens,
                        output_tokens,
                        total_tokens,
                        cost_usd,
                    )
                    .await?;
                persist_memories(store, session_id, &parsed.memories).await?;

                last_result = Some(chunk_result);
            }

            // Return final clean_text directly — all chunks already persisted inline
            let final_parsed = parse_response(&last_result.expect("at least one chunk").text);
            final_parsed.clean_text
        }
        DeliveryMode::ChunkedNonSession => {
            // Non-session agents (codex, opencode) — each chunk is independent invocation
            // No --resume, but each chunk's response is persisted to keep transcript complete.
            let mut last_text = String::new();

            for chunk in &chunks {
                let chunk_result = invoke_agent(&agent, &model, &ctx.system, chunk).await?;

                // Persist this chunk's assistant response
                let parsed = parse_response(&chunk_result.text);
                let total_tokens = total_tokens_u64_to_i64_saturating(
                    chunk_result.input_tokens,
                    chunk_result.output_tokens,
                );
                let input_tokens = opt_token_u64_to_i64_saturating(chunk_result.input_tokens);
                let output_tokens = opt_token_u64_to_i64_saturating(chunk_result.output_tokens);
                let cost_usd = chunk_result.cost_usd;
                store
                    .insert_control_message(
                        session_id,
                        "assistant",
                        channel,
                        channel_thread,
                        &parsed.clean_text,
                        parsed.summary.as_deref(),
                        Some(&model),
                        Some(&agent),
                        input_tokens,
                        output_tokens,
                        total_tokens,
                        cost_usd,
                    )
                    .await?;
                persist_memories(store, session_id, &parsed.memories).await?;

                last_text = parsed.clean_text;
            }

            last_text
        }
        DeliveryMode::Single => {
            // Single message delivery (under-limit messages) — use invoke_agent_with_session
            // for conversation continuity with --session-id / --resume
            let mut active_uuid = session_uuid.clone();
            let mut is_first_chunk = !is_new;
            let chunk_result = loop {
                let r = invoke_agent_with_session(
                    &agent,
                    &model,
                    &ctx.system,
                    &chunks[0],
                    Some(&active_uuid),
                    is_first_chunk,
                )
                .await;
                match r {
                    Ok(result) => break result,
                    Err(e) => {
                        let err_str = e.to_string();
                        let is_stale = err_str.contains("stale session");
                        let is_timeout = err_str.contains("timeout after")
                            || err_str.contains("timed out after");
                        if is_stale || is_timeout {
                            if is_stale {
                                tracing::info!(
                                    session_id,
                                    "chat session UUID expired; resetting and retrying with fresh session"
                                );
                            } else {
                                tracing::info!(
                                    session_id,
                                    "chat session timed out; resetting and retrying with fresh session"
                                );
                            }
                            reset_session_uuid(store, session_id).await?;
                            let (fresh_uuid, _) =
                                get_or_create_session_uuid(store, session_id).await?;
                            active_uuid = fresh_uuid;
                            is_first_chunk = false;
                            continue;
                        }
                        return Err(e);
                    }
                }
            };

            // Parse response for summary and memory tags
            let parsed = parse_response(&chunk_result.text);

            // Store assistant message with token usage.
            // Must happen before persist_memories so that if the DB write fails the
            // KV store is not left with orphaned memory keys that have no corresponding
            // conversation history entry.
            let total_tokens = total_tokens_u64_to_i64_saturating(
                chunk_result.input_tokens,
                chunk_result.output_tokens,
            );
            let input_tokens = opt_token_u64_to_i64_saturating(chunk_result.input_tokens);
            let output_tokens = opt_token_u64_to_i64_saturating(chunk_result.output_tokens);
            let cost_usd = chunk_result.cost_usd;
            store
                .insert_control_message(
                    session_id,
                    "assistant",
                    channel,
                    channel_thread,
                    &parsed.clean_text,
                    parsed.summary.as_deref(),
                    Some(&model),
                    Some(&agent),
                    input_tokens,
                    output_tokens,
                    total_tokens,
                    cost_usd,
                )
                .await?;
            persist_memories(store, session_id, &parsed.memories).await?;

            parsed.clean_text
        }
    };

    // Release the per-session lock before eviction check.
    drop(guard);

    // Evict the session lock entry if no concurrent waiters are holding a reference.
    // After drop(guard) above the remaining strong references are:
    // - the map entry
    // - the local `session_lock` clone held by this function.
    // That sums to 2 when there are no other waiters. If another waiter entered
    // concurrently the count will be > 2 and we should not evict. Only remove the
    // map entry when the strong_count == 2.
    if Arc::strong_count(&session_lock) == 2 {
        let mut map = SESSION_LOCKS.lock().unwrap_or_else(PoisonError::into_inner);
        // Re-check under the map lock: a concurrent caller may have cloned the Arc
        // between our strong_count check above and acquiring the map lock.
        if map
            .get(session_id)
            .is_some_and(|l| Arc::strong_count(l) == 2)
        {
            map.remove(session_id);
        }
    }

    Ok(result)
}

/// Split `text` into chunks of at most `max_bytes` bytes each.
/// Prefers splitting at newline boundaries to preserve readability.
fn chunk_text(text: &str, max_bytes: usize) -> Vec<String> {
    if text.len() <= max_bytes {
        return vec![text.to_string()];
    }
    let mut chunks = Vec::new();
    let mut start = 0;
    let total = text.len();

    while start < total {
        let mut end = (start + max_bytes).min(total);
        // Respect UTF-8 char boundaries
        while end > start && !text.is_char_boundary(end) {
            end -= 1;
        }
        if end == start {
            // Single char exceeds limit — advance past it
            if let Some(ch) = text[start..].chars().next() {
                end = start + ch.len_utf8();
            } else {
                break;
            }
        }
        // Prefer splitting at a newline within the chunk
        if end < total
            && let Some(nl_offset) = text[start..end].rfind('\n')
        {
            end = start + nl_offset + 1;
        }
        chunks.push(text[start..end].to_string());
        start = end;
    }
    chunks
}

/// Build control session message chunks for agents with or without session support.
///
/// For session-supporting agents (claude, kimi, minimax), chunks are raw text pieces
/// that will be sent sequentially with `--resume` for conversation continuity.
///
/// For non-session agents (codex, opencode), each chunk is prefixed with a header
/// indicating its position in the sequence, since each invocation is independent.
fn build_control_chunks(state: &str, message: &str, agent_has_sessions: bool) -> Vec<String> {
    let header_reserve = if agent_has_sessions { 0 } else { 200 };
    let effective_limit = MAX_CONTROL_MESSAGE_BYTES.saturating_sub(header_reserve);
    let mut parts = Vec::new();

    // Split state into chunks
    let state_chunks = chunk_text(state, effective_limit);
    for (i, chunk) in state_chunks.iter().enumerate() {
        let prefixed = if agent_has_sessions {
            chunk.clone()
        } else {
            format!(
                "[Context chunk {}/{}]\n{}",
                i + 1,
                state_chunks.len(),
                chunk
            )
        };
        parts.push(prefixed);
    }

    // Split user message if needed, then add as final chunk(s)
    let msg_chunks = chunk_text(message, effective_limit);
    for (i, chunk) in msg_chunks.iter().enumerate() {
        let prefixed = if agent_has_sessions {
            chunk.clone()
        } else if msg_chunks.len() == 1 {
            format!("[Context complete. User question:]\n{chunk}")
        } else {
            format!(
                "[User question chunk {}/{}]\n{}",
                i + 1,
                msg_chunks.len(),
                chunk
            )
        };
        parts.push(prefixed);
    }
    parts
}

#[cfg(test)]
mod tests {
    use tokio::runtime::Runtime;

    use std::collections::BTreeSet;

    use super::*;
    use crate::store::TaskStore;

    #[test]
    fn parse_spec_agent_model() {
        let spec = parse_model_spec("minimax:sonnet", "claude");
        assert_eq!(spec.agent, "minimax");
        assert_eq!(spec.model, "sonnet");
    }

    #[test]
    fn parse_spec_opencode_model() {
        let spec = parse_model_spec("opencode:minimax-m2.5-free", "claude");
        assert_eq!(spec.agent, "opencode");
        assert_eq!(spec.model, "minimax-m2.5-free");
    }

    fn configured(entries: &[(&str, &[&str])]) -> Vec<(String, Vec<String>)> {
        entries
            .iter()
            .map(|(agent, models)| {
                (
                    agent.to_string(),
                    models.iter().map(ToString::to_string).collect(),
                )
            })
            .collect()
    }

    #[test]
    fn infer_agent_finds_the_agent_whose_config_lists_the_model() {
        let configured = configured(&[("claude", &["sonnet"]), ("codex", &["gpt-4o"])]);
        assert_eq!(infer_agent("gpt-4o", &configured, "claude"), "codex");
        assert_eq!(infer_agent("sonnet", &configured, "codex"), "claude");
    }

    #[test]
    fn infer_agent_keeps_the_current_agent_for_an_unlisted_model() {
        let configured = configured(&[("claude", &["sonnet"])]);
        assert_eq!(
            infer_agent("deepseek-r1", &configured, "opencode"),
            "opencode"
        );
    }

    #[test]
    fn first_model_skips_pool_keywords() {
        let models = vec!["opencode:free".to_string(), "provider/model".to_string()];
        assert_eq!(first_model(&models), "provider/model");
    }

    #[test]
    fn first_model_is_empty_without_models() {
        assert_eq!(first_model(&[]), "");
        assert_eq!(model_arg(""), None);
        assert_eq!(model_arg("sonnet"), Some("sonnet"));
    }

    #[tokio::test]
    async fn assemble_context_includes_memories_and_summaries() {
        let store = TaskStore::open_memory().await.unwrap();
        store
            .kv_set("control:memory:default:tz", "User timezone is UTC-3")
            .await
            .unwrap();
        store
            .insert_control_message(
                "default",
                "assistant",
                "cli",
                None,
                "response",
                Some("checked status"),
                Some("sonnet"),
                Some("claude"),
                None,
                None,
                None,
                None,
            )
            .await
            .unwrap();

        let ctx = assemble_context(&store, "default").await.unwrap();
        assert!(ctx.state.contains("UTC-3"), "should include memories");
        assert!(
            ctx.state.contains("checked status"),
            "should include summaries"
        );
    }

    #[test]
    fn extract_summary_from_response() {
        let response = "Here are your tasks.\n\n<summary>listed 3 active tasks</summary>";
        let parsed = parse_response(response);
        assert_eq!(parsed.summary.as_deref(), Some("listed 3 active tasks"));
        assert!(!parsed.clean_text.contains("<summary>"));
    }

    #[test]
    fn extract_summary_missing() {
        let response = "No tasks found.";
        let parsed = parse_response(response);
        assert_eq!(parsed.summary, None);
        assert_eq!(parsed.clean_text, "No tasks found.");
    }

    #[test]
    fn extract_summary_closing_tag_before_opening_does_not_panic() {
        // </summary> appears before <summary> — must not panic
        let response = "Example: </summary> is a closing tag.\n<summary>actual summary</summary>";
        let parsed = parse_response(response);
        assert_eq!(parsed.summary.as_deref(), Some("actual summary"));
        assert!(!parsed.clean_text.contains("<summary>"));
    }

    #[test]
    fn extract_memory_tags() {
        let response = r#"Remember these.

<memory key="tz">User timezone is UTC-3</memory>
<memory key="theme">Prefer concise answers</memory>

<summary>stored two memories</summary>"#;
        let parsed = parse_response(response);
        assert_eq!(parsed.memories.len(), 2);
        assert_eq!(parsed.memories[0].key, "tz");
        assert_eq!(parsed.memories[0].value, "User timezone is UTC-3");
        assert_eq!(parsed.memories[1].key, "theme");
        assert_eq!(parsed.memories[1].value, "Prefer concise answers");
        assert_eq!(parsed.summary.as_deref(), Some("stored two memories"));
        assert!(!parsed.clean_text.contains("<memory"));
    }

    #[tokio::test]
    async fn persist_memories_writes_to_kv() {
        let store = TaskStore::open_memory().await.unwrap();
        let memories = vec![MemoryEntry {
            key: "tz".to_string(),
            value: "User timezone is UTC-3".to_string(),
        }];

        persist_memories(&store, "default", &memories)
            .await
            .unwrap();

        assert_eq!(
            store
                .kv_get("control:memory:default:tz")
                .await
                .unwrap()
                .as_deref(),
            Some("User timezone is UTC-3")
        );
    }

    #[tokio::test]
    async fn send_message_model_switch() {
        let store = TaskStore::open_memory().await.unwrap();
        set_agent(&store, "unlisted-agent").await.unwrap();
        let model = get_model(&store).await.unwrap();
        assert_eq!(model, "", "an agent with no models: entry gets no --model");

        // /model with explicit agent:model — will fail validation (no binary in test)
        // so test parse_model_spec + set_model_spec directly
        let spec = parse_model_spec("minimax:sonnet", "claude");
        set_model_spec(&store, &spec).await.unwrap();
        assert_eq!(get_model(&store).await.unwrap(), "sonnet");
        assert_eq!(get_agent(&store).await.unwrap(), "minimax");
    }

    #[tokio::test]
    async fn send_message_model_show_current() {
        let store = TaskStore::open_memory().await.unwrap();
        set_agent(&store, "unlisted-agent").await.unwrap();
        let response = send_message(&store, "default", "cli", None, "/model")
            .await
            .unwrap();
        assert_eq!(response, "Current: unlisted-agent:default");
    }

    #[tokio::test]
    async fn send_message_agent_show_current() {
        let store = TaskStore::open_memory().await.unwrap();
        set_agent(&store, "unlisted-agent").await.unwrap();
        let response = send_message(&store, "default", "cli", None, "/agent")
            .await
            .unwrap();
        assert_eq!(response, "Current: unlisted-agent:default");
    }

    #[tokio::test]
    async fn send_message_agent_switches_with_default_model() {
        let store = TaskStore::open_memory().await.unwrap();
        // Directly set agent spec to bypass validation (no binary in test env)
        let spec = set_agent_spec(&store, "unlisted-agent").await.unwrap();
        assert_eq!(spec.agent, "unlisted-agent");
        assert_eq!(
            spec.model, "",
            "an agent with no models: entry gets no --model"
        );
        assert_eq!(get_agent(&store).await.unwrap(), "unlisted-agent");
        assert_eq!(get_model(&store).await.unwrap(), "");
    }

    #[test]
    fn parse_control_command_model_and_agent() {
        assert_eq!(
            parse_control_command("/model"),
            Some(ControlCommand::Model(None))
        );
        assert_eq!(
            parse_control_command("/agent"),
            Some(ControlCommand::Agent(None))
        );
        assert_eq!(
            parse_control_command("/agent codex"),
            Some(ControlCommand::Agent(Some("codex".to_string())))
        );
        assert_eq!(
            parse_control_command("/model opencode:deepseek-r1"),
            Some(ControlCommand::Model(Some(
                "opencode:deepseek-r1".to_string()
            )))
        );
        assert_eq!(
            parse_control_command("/restart"),
            Some(ControlCommand::Restart)
        );
        assert_eq!(parse_control_command("/kill"), None);
        assert_eq!(parse_control_command("/session"), None);
    }

    #[test]
    fn control_prompt_slash_commands_match_supported_surface() {
        let documented = Regex::new(r"`/([a-z]+)(?:[^`]*)`")
            .unwrap()
            .captures_iter(SYSTEM_TEMPLATE)
            .map(|caps| caps[1].to_string())
            .collect::<BTreeSet<_>>();

        let expected = ["agent", "model", "restart"]
            .into_iter()
            .map(str::to_string)
            .collect::<BTreeSet<_>>();

        assert_eq!(documented, expected);
    }

    #[tokio::test]
    async fn set_agent_spec_persists_default_model() {
        let store = TaskStore::open_memory().await.unwrap();
        let spec = set_agent_spec(&store, "unlisted-agent").await.unwrap();
        assert_eq!(spec.agent, "unlisted-agent");
        assert_eq!(spec.model, "");
        assert_eq!(get_agent(&store).await.unwrap(), "unlisted-agent");
        assert_eq!(get_model(&store).await.unwrap(), "");
    }

    #[tokio::test]
    async fn send_message_model_switch_does_not_store() {
        let store = TaskStore::open_memory().await.unwrap();
        // set directly to avoid validation in test env
        let spec = parse_model_spec("claude:haiku", "claude");
        set_model_spec(&store, &spec).await.unwrap();
        let messages = store
            .list_control_messages("default", None, 10)
            .await
            .unwrap();
        assert_eq!(messages.len(), 0);
    }

    #[test]
    fn validate_unknown_agent_fails() {
        let spec = ModelSpec {
            agent: "nonexistent".to_string(),
            model: "test".to_string(),
        };
        let rt = Runtime::new().unwrap();
        let result = rt.block_on(validate_model(&spec));
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("unknown agent"));
    }

    #[test]
    fn validate_agent_rejects_unknown() {
        let result = validate_agent("nonexistent");
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("unknown agent"), "got: {err}");
    }

    #[test]
    fn validate_agent_rejects_empty() {
        let result = validate_agent("");
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("unknown agent"));
    }

    #[tokio::test]
    async fn agent_command_rejects_unknown_before_persisting() {
        let store = TaskStore::open_memory().await.unwrap();
        // Set a known-good agent first
        set_agent(&store, "claude").await.unwrap();

        // Try switching to an unknown agent — should fail
        let result =
            maybe_handle_control_command(&store, "default", None, "/agent bogusagent").await;
        assert!(result.is_err(), "should reject unknown agent");
        assert!(result.unwrap_err().to_string().contains("unknown agent"));

        // Verify the original agent is still set (not overwritten)
        assert_eq!(get_agent(&store).await.unwrap(), "claude");
    }

    #[test]
    fn token_u64_to_i64_saturating_caps_large_values() {
        assert_eq!(token_u64_to_i64_saturating(u64::MAX), i64::MAX);
    }

    #[test]
    fn total_tokens_u64_to_i64_saturating_caps_single_large_value() {
        assert_eq!(
            total_tokens_u64_to_i64_saturating(Some(u64::MAX), None),
            Some(i64::MAX)
        );
    }

    #[test]
    fn total_tokens_u64_to_i64_saturating_caps_large_sum() {
        assert_eq!(
            total_tokens_u64_to_i64_saturating(Some(u64::MAX), Some(1)),
            Some(i64::MAX)
        );
    }

    #[tokio::test]
    async fn insert_control_message_stores_cost_usd() {
        let store = TaskStore::open_memory().await.unwrap();

        let input_tokens: u64 = 500_000;
        let output_tokens: u64 = 100_000;
        let cost_usd = Some(0.25);

        store
            .insert_control_message(
                "default",
                "assistant",
                "cli",
                None,
                "hello",
                Some("greeted"),
                Some("sonnet"),
                Some("claude"),
                opt_token_u64_to_i64_saturating(Some(input_tokens)),
                opt_token_u64_to_i64_saturating(Some(output_tokens)),
                total_tokens_u64_to_i64_saturating(Some(input_tokens), Some(output_tokens)),
                cost_usd,
            )
            .await
            .unwrap();

        let messages = store
            .list_control_messages("default", None, 10)
            .await
            .unwrap();
        assert_eq!(messages.len(), 1);
        let msg = &messages[0];
        assert!(msg.cost_usd.is_some(), "cost_usd must be stored, not NULL");
        assert_eq!(
            msg.cost_usd,
            Some(0.25),
            "the reported cost is stored as is"
        );
    }

    #[tokio::test]
    async fn session_uuid_created_and_stable() {
        let store = TaskStore::open_memory().await.unwrap();
        let (uuid1, is_new1) = get_or_create_session_uuid(&store, "default").await.unwrap();
        let (uuid2, is_new2) = get_or_create_session_uuid(&store, "default").await.unwrap();
        assert_eq!(uuid1, uuid2, "same session should return same UUID");
        assert!(is_new1, "first call should be new");
        assert!(!is_new2, "second call should be existing");
        assert_eq!(uuid1.len(), 36, "should be a valid UUID string");
    }

    #[tokio::test]
    async fn session_uuid_differs_per_session() {
        let store = TaskStore::open_memory().await.unwrap();
        let (uuid_default, _) = get_or_create_session_uuid(&store, "default").await.unwrap();
        let (uuid_ops, _) = get_or_create_session_uuid(&store, "ops").await.unwrap();
        assert_ne!(uuid_default, uuid_ops);
    }

    #[tokio::test]
    async fn session_uuid_reset_creates_new() {
        let store = TaskStore::open_memory().await.unwrap();
        let (uuid1, _) = get_or_create_session_uuid(&store, "default").await.unwrap();
        reset_session_uuid(&store, "default").await.unwrap();
        let (uuid2, is_new2) = get_or_create_session_uuid(&store, "default").await.unwrap();
        assert_ne!(uuid1, uuid2, "reset should produce a new UUID");
        assert!(is_new2, "after reset should be new");
    }

    /// Verify that the timeout error patterns used in send_message recovery match
    /// the actual error strings produced by the agent runner (AgentError::Timeout
    /// formats as "timeout after Xs"; run_shell_command formats as "agent timed out
    /// after Xs"). Both must trigger session reset rather than bubbling up as errors.
    #[test]
    fn timeout_error_patterns_match_runner_output() {
        let patterns = ["timeout after 1800s", "agent timed out after 120s"];
        for &msg in &patterns {
            let is_timeout = msg.contains("timeout after") || msg.contains("timed out after");
            assert!(is_timeout, "pattern should be detected as timeout: {msg:?}");
        }
        // Unrelated errors must not be misidentified as timeouts.
        let non_timeout = [
            "rate limit: 429",
            "stale session (session expired): abc",
            "auth error: invalid key",
        ];
        for &msg in &non_timeout {
            let is_timeout = msg.contains("timeout after") || msg.contains("timed out after");
            assert!(
                !is_timeout,
                "pattern must not be detected as timeout: {msg:?}"
            );
        }
    }

    // ── chunk_text ──────────────────────────────────────────────────────────────

    #[test]
    fn chunk_text_under_limit_returns_single_chunk() {
        let text = "hello world";
        let chunks = chunk_text(text, 100);
        assert_eq!(chunks, vec!["hello world".to_string()]);
    }

    #[test]
    fn chunk_text_exact_limit_returns_single_chunk() {
        let text = "abcde";
        let chunks = chunk_text(text, 5);
        assert_eq!(chunks.len(), 1);
        assert_eq!(chunks[0], "abcde");
    }

    #[test]
    fn chunk_text_splits_and_reassembles() {
        let text = "line1\nline2\nline3\nline4\nline5";
        let chunks = chunk_text(text, 12);
        let reassembled: String = chunks.join("");
        assert_eq!(reassembled, text);
    }

    #[test]
    fn chunk_text_respects_utf8_boundaries() {
        // 4-byte emoji placed at the split boundary
        let content = format!("{}😀", "x".repeat(97)); // 101 bytes; limit=100 cuts inside emoji

        let chunks = chunk_text(&content, 100);
        let reassembled: String = chunks.join("");
        assert_eq!(reassembled, content);
        for c in &chunks {
            assert!(
                str::from_utf8(c.as_bytes()).is_ok(),
                "chunk is not valid UTF-8"
            );
        }
    }

    #[test]
    fn chunk_text_prefers_newline_boundary() {
        // "aaaa\nbbbb" with limit=6 — should split after the '\n', not mid-word
        let text = "aaaa\nbbbb";
        let chunks = chunk_text(text, 6);
        // First chunk should end with '\n' (split at newline, not byte 6 into "aaaa\nb")
        assert!(
            chunks[0].ends_with('\n'),
            "expected split at newline, got: {:?}",
            chunks[0]
        );
        let reassembled: String = chunks.join("");
        assert_eq!(reassembled, text);
    }

    #[test]
    fn chunk_text_empty_string() {
        let chunks = chunk_text("", 32);
        assert!(chunks.is_empty() || chunks == vec![String::new()]);
    }

    // ── build_control_chunks ────────────────────────────────────────────────

    #[test]
    fn build_control_chunks_session_agent_splits_state_and_message() {
        // build_control_chunks always splits state and message into separate chunks
        // (it's only called when full_message exceeds the limit)
        let state = "state content";
        let message = "user question";
        let chunks = build_control_chunks(state, message, true);
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0], "state content");
        assert_eq!(chunks[1], "user question");
    }

    #[test]
    fn build_control_chunks_session_agent_long_state_splits() {
        // Long state gets split into multiple chunks
        let state = "line\n".repeat(10000); // ~60KB
        let message = "question";
        let chunks = build_control_chunks(&state, message, true);
        // Should have multiple state chunks + 1 message chunk
        assert!(chunks.len() >= 3);
        // All but last are state chunks
        for chunk in &chunks[..chunks.len() - 1] {
            assert!(chunk.contains("line"));
        }
        // Last is message
        assert_eq!(chunks.last().unwrap(), "question");
    }

    #[test]
    fn build_control_chunks_non_session_agent_adds_headers() {
        let state = "state content";
        let message = "user question";
        let chunks = build_control_chunks(state, message, false);
        assert_eq!(chunks.len(), 2);
        assert!(chunks[0].starts_with("[Context chunk 1/1]"));
        assert!(chunks[0].contains("state content"));
        assert!(chunks[1].starts_with("[Context complete. User question:]"));
        assert!(chunks[1].contains("user question"));
    }

    #[test]
    fn build_control_chunks_non_session_agent_splits_long_message() {
        let state = "state";
        let message = "x".repeat(50000); // Exceeds 32KB limit
        let chunks = build_control_chunks(state, &message, false);
        // Should have 1 state chunk + multiple message chunks
        assert!(chunks.len() >= 2);
        assert!(chunks[0].starts_with("[Context chunk 1/1]"));
        // Message chunks should have question headers
        for chunk in &chunks[1..] {
            assert!(chunk.starts_with("[User question chunk"));
        }
        // Verify reassembly (without headers)
        let reassembled_state = chunks[0].split_once('\n').map_or("", |(_, r)| r);
        assert_eq!(reassembled_state, "state");
    }

    #[test]
    fn build_control_chunks_non_session_agent_long_state_splits() {
        let state = "line\n".repeat(10000); // ~60KB
        let message = "question";
        let chunks = build_control_chunks(&state, message, false);
        // Multiple state chunks + 1 message chunk
        assert!(chunks.len() >= 3);
        for (i, chunk) in chunks.iter().enumerate() {
            if i < chunks.len() - 1 {
                assert!(chunk.starts_with(&format!("[Context chunk {}/", i + 1)));
                assert!(chunk.contains("line"));
            } else {
                assert!(chunk.starts_with("[Context complete. User question:]"));
                assert_eq!(chunk.split_once('\n').map_or("", |(_, r)| r), "question");
            }
        }
    }
}